Export generalizations

src/spyglass/common/common_dandi.py

@@ -111,14 +111,18 @@ def compile_dandiset(
         dandi_api_key: Optional[str] = None,
         dandi_instance: Optional[str] = "dandi",
         skip_raw_files: Optional[bool] = False,
+        n_compile_processes: Optional[int] = 1,
+        n_upload_processes: Optional[int] = None,
+        n_organize_processes: Optional[int] = None,
+        n_validate_processes: Optional[int] = 1,
     ):
         """Compile a Dandiset from the export.
         Parameters
         ----------
         key : dict
             ExportSelection key
         dandiset_id : str
-            Dandiset ID generated by the user on the dadndi server
+            Dandiset ID generated by the user on the dandi server
         dandi_api_key : str, optional
             API key for the dandi server. Optional if the environment variable
             DANDI_API_KEY is set.
@@ -162,19 +166,44 @@ def compile_dandiset(
                 )
 
         os.makedirs(destination_dir, exist_ok=False)
-        for file in source_files:
-            if os.path.exists(f"{destination_dir}/{os.path.basename(file)}"):
-                continue
-            if skip_raw_files and raw_dir in file:
-                continue
-            # copy the file if it has external links so can be safely edited
-            if nwb_has_external_links(file):
-                shutil.copy(file, f"{destination_dir}/{os.path.basename(file)}")
-            else:
-                os.symlink(file, f"{destination_dir}/{os.path.basename(file)}")
+        # for file in source_files:
+        #     if os.path.exists(f"{destination_dir}/{os.path.basename(file)}"):
+        #         continue
+        #     if skip_raw_files and raw_dir in file:
+        #         continue
+        #     # copy the file if it has external links so can be safely edited
+        #     if nwb_has_external_links(file):
+        #         shutil.copy(file, f"{destination_dir}/{os.path.basename(file)}")
+        #     else:
+        #         os.symlink(file, f"{destination_dir}/{os.path.basename(file)}")
+        logger.info(
+            f"Compiling dandiset in {destination_dir} from {len(source_files)} files"
+        )
+        if n_compile_processes == 1:
+            for file in source_files:
+                _make_file_in_dandi_dir(file, destination_dir, skip_raw_files)
+        else:
+            from multiprocessing import Pool
+
+            print(
+                f"Using multiprocessing to compile dandi export. {n_compile_processes} processes"
+            )
+            with Pool(processes=n_compile_processes) as pool:
+                pool.starmap(
+                    _make_file_in_dandi_dir,
+                    [
+                        (file, destination_dir, skip_raw_files)
+                        for file in source_files
+                    ],
+                )
 
         # validate the dandiset
-        validate_dandiset(destination_dir, ignore_external_files=True)
+        logger.info("Validating dandiset before organization")
+        validate_dandiset(
+            destination_dir,
+            ignore_external_files=True,
+            n_processes=n_validate_processes,
+        )
 
         # given dandiset_id, download the dandiset to the export_dir
         url = (
@@ -184,24 +213,28 @@ def compile_dandiset(
         dandi.download.download(url, output_dir=paper_dir)
 
         # organize the files in the dandiset directory
+        logger.info("Organizing dandiset")
         dandi.organize.organize(
             destination_dir,
             dandiset_dir,
             update_external_file_paths=True,
             invalid=OrganizeInvalid.FAIL,
             media_files_mode=CopyMode.SYMLINK,
             files_mode=FileOperationMode.COPY,
+            jobs=n_organize_processes,
         )
 
         # get the dandi name translations
         translations = lookup_dandi_translation(destination_dir, dandiset_dir)
 
         # upload the dandiset to the dandi server
+        logger.info("Uploading dandiset")
         if dandi_api_key:
             os.environ["DANDI_API_KEY"] = dandi_api_key
         dandi.upload.upload(
             [dandiset_dir],
             dandi_instance=dandi_instance,
+            jobs=n_upload_processes,
         )
         logger.info(f"Dandiset {dandiset_id} uploaded")
         # insert the translations into the dandi table
@@ -239,6 +272,18 @@ def write_mysqldump(self, export_key: dict):
         sql_dump.write_mysqldump([self & key], file_suffix="_dandi")
 
 
+def _make_file_in_dandi_dir(file, destination_dir, skip_raw_files):
+    if os.path.exists(f"{destination_dir}/{os.path.basename(file)}"):
+        return
+    if skip_raw_files and raw_dir in file:
+        return
+    # copy the file if it has external links so can be safely edited
+    if nwb_has_external_links(file):
+        shutil.copy(file, f"{destination_dir}/{os.path.basename(file)}")
+    else:
+        os.symlink(file, f"{destination_dir}/{os.path.basename(file)}")
+
+
 def _get_metadata(path):
     # taken from definition within dandi.organize.organize
     try:
@@ -314,7 +359,7 @@ def lookup_dandi_translation(source_dir: str, dandiset_dir: str):
 
 
 def validate_dandiset(
-    folder, min_severity="ERROR", ignore_external_files=False
+    folder, min_severity="ERROR", ignore_external_files=False, n_processes=1
 ):
     """Validate the dandiset directory
 
@@ -329,7 +374,21 @@ def validate_dandiset(
         whether to ignore external file errors. Used if validating
         before the organize step
     """
-    validator_result = dandi.validate.validate(folder)
+    if n_processes == 1:
+        validator_result = dandi.validate.validate(folder)
+    else:
+        from multiprocessing import Pool
+
+        from dandi.files import find_dandi_files
+
+        files_to_validate = [x.filepath for x in find_dandi_files(folder)]
+
+        print(
+            f"Using multiprocessing to validate dandi export. {n_processes} processes"
+        )
+        with Pool(processes=n_processes) as pool:
+            per_file_results = pool.map(validate_1, files_to_validate)
+        validator_result = [item for sub in per_file_results for item in sub]
     min_severity_value = Severity[min_severity].value
 
     filtered_results = [
@@ -356,3 +415,7 @@ def validate_dandiset(
                 ]
             )
         )
+
+
+def validate_1(path):
+    return list(dandi.validate.validate(path))

src/spyglass/common/common_usage.py

@@ -560,7 +560,7 @@ def make(self, key):
         self.Table().insert(table_inserts)
         self.File().insert(file_inserts)
 
-    def prepare_files_for_export(self, key, **kwargs):
+    def prepare_files_for_export(self, key, n_processes=1, **kwargs):
         """Resolve common known errors to make a set of analysis
         files dandi compliant
 
@@ -570,12 +570,19 @@ def prepare_files_for_export(self, key, **kwargs):
             restriction for a single entry of the Export table
         """
         key = (self & key).fetch1("KEY")
-        self._make_fileset_ids_unique(key)
         file_list = (self.File() & key).fetch("file_path")
-        for file in file_list:
-            update_analysis_for_dandi_standard(file, **kwargs)
 
-    def _make_fileset_ids_unique(self, key):
+        if n_processes == 1:
+            self._make_fileset_ids_unique(key)
+            for file in file_list:
+                update_analysis_for_dandi_standard(file, **kwargs)
+            return
+        with Pool(processes=n_processes) as pool:
+            # pool.map(make_file_obj_id_unique, file_list)
+            print("SKIPPING make_file_obj_id_unique in parallel")
+            pool.map(update_analysis_for_dandi_standard, file_list)
+
+    def _make_fileset_ids_unique(self, key, n_processes=1):
         """Make the object_id of each nwb in a dataset unique"""
         key = (self & key).fetch1("KEY")
         file_list = (self.File() & key).fetch("file_path")