Fix directory progress

Author: jeromekelleher

bio2zarr/core.py

@@ -50,7 +50,8 @@ def wait_on_futures(futures):
             cancel_futures(futures)
             if isinstance(exception, cf.process.BrokenProcessPool):
                 raise RuntimeError(
-                    "Worker process died: you may have run out of memory") from exception
+                    "Worker process died: you may have run out of memory"
+                ) from exception
             else:
                 raise exception
 

bio2zarr/vcf.py

@@ -890,6 +890,15 @@ def num_columns(self):
         return len(self.columns)
 
 
+
+def mkdir_with_progress(path):
+    logger.debug(f"mkdir f{path}")
+    # NOTE we may have race-conditions here, I'm not sure. Hopefully allowing
+    # parents=True will take care of it.
+    path.mkdir(parents=True)
+    core.update_progress(1)
+
+
 class IntermediateColumnarFormatWriter:
     def __init__(self, path):
         self.path = pathlib.Path(path)
@@ -932,7 +941,7 @@ def init(
         # dependencies as well.
         self.metadata.provenance = {"source": f"bio2zarr-{provenance.__version__}"}
 
-        self.mkdirs(worker_processes)
+        self.mkdirs(worker_processes, show_progress=show_progress)
 
         # Note: this is needed for the current version of the vcfzarr spec, but it's
         # probably going to be dropped.
@@ -947,30 +956,30 @@ def init(
             json.dump(self.metadata.asdict(), f, indent=4)
         return self.num_partitions
 
-    def mkdirs(self, worker_processes=1):
-        logger.info(
-            f"Creating {len(self.metadata.fields) * self.num_partitions} directories"
-        )
+    def mkdirs(self, worker_processes=1, show_progress=False):
+        num_dirs = len(self.metadata.fields) * self.num_partitions
+        logger.info(f"Creating {num_dirs} directories")
         self.path.mkdir()
         self.wip_path.mkdir()
         # Due to high latency batch system filesystems, we create all the directories in
         # parallel
         progress_config = core.ProgressConfig(
-            total=len(self.metadata.fields) * self.num_partitions,
-            units="dir",
-            title="Creating directories",
-            show=True
+            total=num_dirs,
+            units="dirs",
+            title="Mkdirs",
+            show=show_progress,
         )
         with core.ParallelWorkManager(
-                worker_processes=worker_processes,
-                progress_config=progress_config
+            worker_processes=worker_processes, progress_config=progress_config
         ) as manager:
             for field in self.metadata.fields:
                 col_path = get_vcf_field_path(self.path, field)
+                # Don't bother trying to count the intermediate directories towards
+                # progress
                 manager.submit(col_path.mkdir, parents=True)
                 for j in range(self.num_partitions):
                     part_path = col_path / f"p{j}"
-                    manager.submit(part_path.mkdir, parents=True)
+                    manager.submit(mkdir_with_progress, part_path)
 
     def load_partition_summaries(self):
         summaries = []