feat(downsample): use memory limit shape when doing limited downsamples

igneous/downsample_scales.py

@@ -214,12 +214,15 @@ def prec(x):
 def create_downsample_scales(
     layer_path, mip, ds_shape, axis='z', 
     preserve_chunk_size=False, chunk_size=None,
-    encoding=None, factor=None
+    encoding=None, factor=None, max_mips=None,
   ):
   vol = CloudVolume(layer_path, mip)
 
   resolutions = compute_scales(vol, mip, ds_shape, axis, factor, chunk_size)
 
+  if max_mips is not None:
+    resolutions = resolutions[:max_mips]
+
   if len(resolutions) == 0:
     print("WARNING: No scales generated.")
 

igneous/task_creation/image.py

@@ -245,14 +245,18 @@ def ds_shape(mip, chunk_size=None, factor=None):
     if factor is None:
       factor = downsample_scales.axis_to_factor(axis)
 
-    if num_mips is None:
-      num_mips = num_mips_from_memory_target(
-        memory_target, vol.dtype, shape, factor
-      )
+    viable_mips = num_mips_from_memory_target(
+      memory_target, vol.dtype, shape, factor
+    )
 
-    shape.x *= factor[0] ** num_mips
-    shape.y *= factor[1] ** num_mips
-    shape.z *= factor[2] ** num_mips
+    if viable_mips < num_mips:
+      raise ValueError(
+        f"Memory limit ({memory_target} bytes) too low to "
+        "compute {num_mips} mips at a time. {viable_mips} mips possible.")
+
+    shape.x *= factor[0] ** viable_mips
+    shape.y *= factor[1] ** viable_mips
+    shape.z *= factor[2] ** viable_mips
 
     return shape, num_mips
 
@@ -262,7 +266,7 @@ def ds_shape(mip, chunk_size=None, factor=None):
   vol = downsample_scales.create_downsample_scales(
     layer_path, mip, shape, 
     preserve_chunk_size=preserve_chunk_size, chunk_size=chunk_size,
-    encoding=encoding, factor=factor
+    encoding=encoding, factor=factor, max_mips=num_mips,
   )
 
   for mip_i in range(mip+1, min(mip + num_mips, len(vol.available_mips))):
@@ -293,6 +297,7 @@ def task(self, shape, offset):
         dest_path=dest_path,
         compress=compress,
         factor=factor,
+        max_mips=num_mips,
       )
 
     def on_finish(self):

igneous/tasks/image/image.py

@@ -35,7 +35,7 @@
 def downsample_and_upload(
     image, bounds, vol, ds_shape,
     mip=0, axis='z', skip_first=False,
-    sparse=False, factor=None
+    sparse=False, factor=None, max_mips=None
   ):
     ds_shape = min2(vol.volume_size, ds_shape[:3])
     underlying_mip = (mip + 1) if (mip + 1) in vol.available_mips else mip
@@ -45,6 +45,9 @@ def downsample_and_upload(
       factor = downsample_scales.axis_to_factor(axis)
     factors = downsample_scales.compute_factors(ds_shape, factor, chunk_size, vol.volume_size)
 
+    if max_mips is not None:
+      factors = factors[:max_mips]
+
     if len(factors) == 0:
       print("No factors generated. Image Shape: {}, Downsample Shape: {}, Volume Shape: {}, Bounds: {}".format(
           image.shape, ds_shape, vol.volume_size, bounds)
@@ -365,20 +368,21 @@ def select_bounding_boxes(self, dataset_bounds):
 
 @queueable
 def TransferTask(
-  src_path, dest_path,
-  mip, shape, offset,
+  src_path:str, dest_path:str,
+  mip:int, shape, offset,
   translate=(0,0,0),
-  fill_missing=False,
-  skip_first=False,
-  skip_downsamples=False,
-  delete_black_uploads=False,
-  background_color=0,
-  sparse=False,
-  axis='z',
-  agglomerate=False,
-  timestamp=None,
+  fill_missing:bool = False,
+  skip_first:bool = False,
+  skip_downsamples:bool = False,
+  delete_black_uploads:bool = False,
+  background_color:int = 0,
+  sparse:bool = False,
+  axis:chr = 'z',
+  agglomerate:bool = False,
+  timestamp:Optional[int] = None,
   compress='gzip',
   factor=None,
+  max_mips:Optional[int] = None,
 ):
   """
   Transfer an image to a new location while enabling
@@ -436,15 +440,16 @@ def TransferTask(
       shape, mip=mip,
       skip_first=skip_first,
       sparse=sparse, axis=axis,
-      factor=factor
+      factor=factor, max_mips=max_mips,
     )
 
 @queueable
 def DownsampleTask(
   layer_path, mip, shape, offset,
   fill_missing=False, axis='z', sparse=False,
   delete_black_uploads=False, background_color=0,
-  dest_path=None, compress="gzip", factor=None
+  dest_path=None, compress="gzip", factor=None,
+  max_mips=None,
 ):
   """
   Downsamples a cutout of the volume. By default it performs
@@ -467,6 +472,7 @@ def DownsampleTask(
     axis=axis,
     compress=compress,
     factor=factor,
+    max_mips=max_mips,
   )
 
 @queueable