Merge pull request #80 from jeromekelleher/catch-memory-error

A few minor usability improvements

Author: jeromekelleher

bio2zarr/core.py

@@ -53,7 +53,12 @@ def wait_on_futures(futures):
     for future in cf.as_completed(futures):
         exception = future.exception()
         if exception is not None:
-            raise exception
+            cancel_futures(futures)
+            if isinstance(exception, cf.process.BrokenProcessPool):
+                raise RuntimeError(
+                    "Worker process died: you may have run out of memory") from exception
+            else:
+                raise exception
 
 
 def cancel_futures(futures):
@@ -74,7 +79,10 @@ def __init__(self, array, offset):
         assert offset % array.chunks[0] == 0
         dims = list(array.shape)
         dims[0] = min(array.chunks[0], array.shape[0])
-        self.buff = np.zeros(dims, dtype=array.dtype)
+        self.buff = np.empty(dims, dtype=array.dtype)
+        # Explicitly Fill with zeros here to make any out-of-memory errors happen
+        # quickly.
+        self.buff[:] = 0
         self.buffer_row = 0
 
     @property

bio2zarr/vcf.py

@@ -1529,13 +1529,27 @@ def encode(
             )
             with core.ParallelWorkManager(worker_processes, progress_config) as pwm:
                 if "call_genotype" in conversion_spec.columns:
-                    logger.info(f"Submit encode call_genotypes in {len(slices)} slices")
+                    arrays = [
+                        sgvcf.root["call_genotype"],
+                        sgvcf.root["call_genotype_phased"],
+                        sgvcf.root["call_genotype_mask"],
+                    ]
+                    min_mem = sum(array.blocks[0].nbytes for array in arrays)
+                    logger.info(
+                        f"Submit encode call_genotypes in {len(slices)} slices. "
+                        f"Min per-worker mem={display_size(min_mem)}"
+                    )
                     for start, stop in slices:
                         pwm.submit(sgvcf.encode_genotypes_slice, pcvcf, start, stop)
 
                 for col in chunked_2d:
                     if col.vcf_field is not None:
-                        logger.info(f"Submit encode {col.name} in {len(slices)} slices")
+                        array = sgvcf.root[col.name]
+                        min_mem = array.blocks[0].nbytes
+                        logger.info(
+                            f"Submit encode {col.name} in {len(slices)} slices. "
+                            f"Min per-worker mem={display_size(min_mem)}"
+                        )
                         for start, stop in slices:
                             pwm.submit(
                                 sgvcf.encode_column_slice, pcvcf, col, start, stop