Make HDF5 dataset creation script compute-parallel

Author: amorehead

project/datasets/builder/create_hdf5_dataset.py

@@ -3,14 +3,15 @@
 import os
 
 from pathlib import Path
-from tqdm import tqdm
+from parallel import submit_jobs
 
 from project.utils.utils import convert_pair_pickle_to_hdf5
 
 
 @click.command()
 @click.argument('raw_data_dir', default='../DIPS/final/raw', type=click.Path(exists=True))
-def main(raw_data_dir: str):
+@click.option('--num_cpus', '-c', default=1)
+def main(raw_data_dir: str, num_cpus: int):
     raw_data_dir = Path(raw_data_dir)
     raw_data_pickle_filepaths = []
     for root, dirs, files in os.walk(raw_data_dir):
@@ -19,11 +20,9 @@ def main(raw_data_dir: str):
                 for file in subfiles:
                     if file.endswith('.dill'):
                         raw_data_pickle_filepaths.append(raw_data_dir / dir / file)
-    for pickle_filepath in tqdm(raw_data_pickle_filepaths):
-        convert_pair_pickle_to_hdf5(
-            pickle_filepath=pickle_filepath,
-            hdf5_filepath=Path(pickle_filepath).with_suffix(".hdf5")
-        )
+    inputs = [(pickle_filepath, Path(pickle_filepath).with_suffix(".hdf5")) for pickle_filepath in raw_data_pickle_filepaths]
+    submit_jobs(convert_pair_pickle_to_hdf5, inputs, num_cpus)
+
     # filepath = Path("project/datasets/DIPS/final/raw/0g/10gs.pdb1_0.dill")
     # pickle_example = convert_pair_hdf5_to_pickle(
     #     hdf5_filepath=Path(filepath).with_suffix(".hdf5")