num_tasks is wrong when region is specified (#854)

Author: tomwhite

cubed/core/ops.py

@@ -243,6 +243,7 @@ def key_function(out_key):
             chunkss=[chunks],
             target_stores=[target],
             output_blocks=output_blocks,
+            num_tasks=source.npartitions,
             **blockwise_kwargs,
         )
         from cubed import Array

cubed/primitive/blockwise.py

@@ -296,6 +296,7 @@ def general_blockwise(
     target_chunks_: Optional[T_RegularChunks] = None,
     return_writes_stores: bool = False,
     output_blocks: Optional[Iterator[List[int]]] = None,
+    num_tasks=None,
     **kwargs,
 ) -> PrimitiveOperation:
     """A more general form of ``blockwise`` that uses a function to specify the block
@@ -418,7 +419,8 @@ def general_blockwise(
         output_blocks = map(
             list, itertools.product(*[range(len(c)) for c in chunks_normal])
         )
-    num_tasks = math.prod(len(c) for c in chunks_normal)
+    if num_tasks is None:
+        num_tasks = math.prod(len(c) for c in chunks_normal)
 
     pipeline = CubedPipeline(
         apply_blockwise,

cubed/tests/test_core.py

@@ -13,7 +13,13 @@
 from cubed._testing import assert_array_equal
 from cubed.array_api.dtypes import _floating_dtypes
 from cubed.backend_array_api import namespace as nxp
-from cubed.core.ops import general_blockwise, merge_chunks, partial_reduce, tree_reduce
+from cubed.core.ops import (
+    _store_array,
+    general_blockwise,
+    merge_chunks,
+    partial_reduce,
+    tree_reduce,
+)
 from cubed.core.optimization import fuse_all_optimize_dag, multiple_inputs_optimize_dag
 from cubed.core.plan import ArrayRole
 from cubed.storage.store import open_storage_array
@@ -197,6 +203,11 @@ def test_to_zarr_region(tmp_path, spec, executor):
     # note that the same zarr store is overwritten in the following tests
 
     region = (slice(0, 2), slice(0, 2))
+
+    # need to use internal _store_array function to access plan
+    # since to_zarr is eager
+    assert _store_array(a[:2, :2], z, region=region).plan().num_tasks == 1
+
     cubed.to_zarr(a[:2, :2], z, region=region, executor=executor)
     res = open_storage_array(store, mode="r")
     assert_array_equal(res[region], np.array([[1, 2], [5, 6]]))