Don't count virtual arrays towards max_total_source_arrays in optimization (#838)

Author: tomwhite

cubed/core/optimization.py

@@ -8,6 +8,7 @@
     fuse,
     fuse_multiple,
 )
+from cubed.storage.virtual import VirtualArray
 
 logger = logging.getLogger(__name__)
 
@@ -152,15 +153,15 @@ def is_fusable_with_predecessors(node_dict):
 
 
 def num_source_arrays(dag, name):
-    """Return the number of (non-hidden) arrays that are inputs to an op.
+    """Return the number of (non-virtual) arrays that are inputs to an op.
 
-    Hidden arrays are used for internal bookkeeping, are very small virtual arrays
-    (empty, or offsets for example), and are not shown on the plan visualization.
-    For these reasons they shouldn't count towards ``max_total_source_arrays``.
+    Virtual arrays are very small arrays that are held in memory and not read from storage,
+    so they shouldn't count towards ``max_total_source_arrays``.
     """
     nodes = dict(dag.nodes(data=True))
     return sum(
-        not nodes[array]["hidden"] for array in predecessors_unordered(dag, name)
+        not isinstance(nodes[array]["target"], VirtualArray)
+        for array in predecessors_unordered(dag, name)
     )
 
 

cubed/storage/virtual.py

@@ -11,7 +11,11 @@
 from cubed.utils import array_memory, broadcast_trick, memory_repr
 
 
-class VirtualEmptyArray(ArrayMetadata):
+class VirtualArray(ArrayMetadata):
+    pass
+
+
+class VirtualEmptyArray(VirtualArray):
     """An array that is never materialized (in memory or on disk) and contains empty values."""
 
     def __init__(
@@ -34,7 +38,7 @@ def chunkmem(self):
         return array_memory(self.dtype, (1,))
 
 
-class VirtualFullArray(ArrayMetadata):
+class VirtualFullArray(VirtualArray):
     """An array that is never materialized (in memory or on disk) and contains a single fill value."""
 
     def __init__(
@@ -61,7 +65,7 @@ def chunkmem(self):
         return array_memory(self.dtype, (1,))
 
 
-class VirtualOffsetsArray(ArrayMetadata):
+class VirtualOffsetsArray(VirtualArray):
     """An array that is never materialized (in memory or on disk) and contains sequentially incrementing integers."""
 
     def __init__(self, shape: T_Shape):
@@ -77,7 +81,7 @@ def __getitem__(self, key):
         )
 
 
-class VirtualInMemoryArray(ArrayMetadata):
+class VirtualInMemoryArray(VirtualArray):
     """A small array that is held in memory but never materialized on disk."""
 
     def __init__(

cubed/tests/test_optimization.py

@@ -22,7 +22,7 @@
 )
 from cubed.core.plan import arrays_to_plan
 from cubed.tests.test_core import sqrts
-from cubed.tests.utils import TaskCounter
+from cubed.tests.utils import TaskCounter, create_zarr
 
 
 @pytest.fixture
@@ -679,15 +679,19 @@ def stack_add(*a):
 #          \     /
 #             p
 #
-def test_fuse_large_fan_in_default(spec):
-    a = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    b = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    c = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    d = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    e = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    f = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    g = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    h = xp.ones((2, 2), chunks=(2, 2), spec=spec)
+def test_fuse_large_fan_in_default(tmp_path, spec):
+    # use zarr input so that they count towards max_total_source_arrays
+    store = tmp_path / "source.zarr"
+    create_zarr(np.ones((2, 2)), chunks=(2, 2), store=store)
+
+    a = cubed.from_zarr(store, spec=spec)
+    b = cubed.from_zarr(store, spec=spec)
+    c = cubed.from_zarr(store, spec=spec)
+    d = cubed.from_zarr(store, spec=spec)
+    e = cubed.from_zarr(store, spec=spec)
+    f = cubed.from_zarr(store, spec=spec)
+    g = cubed.from_zarr(store, spec=spec)
+    h = cubed.from_zarr(store, spec=spec)
 
     i = xp.add(a, b)
     j = xp.add(c, d)
@@ -738,15 +742,19 @@ def test_fuse_large_fan_in_default(spec):
 #          \     /
 #             p
 #
-def test_fuse_large_fan_in_override(spec):
-    a = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    b = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    c = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    d = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    e = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    f = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    g = xp.ones((2, 2), chunks=(2, 2), spec=spec)
-    h = xp.ones((2, 2), chunks=(2, 2), spec=spec)
+def test_fuse_large_fan_in_override(tmp_path, spec):
+    # use zarr input so that they count towards max_total_source_arrays
+    store = tmp_path / "source.zarr"
+    create_zarr(np.ones((2, 2)), chunks=(2, 2), store=store)
+
+    a = cubed.from_zarr(store, spec=spec)
+    b = cubed.from_zarr(store, spec=spec)
+    c = cubed.from_zarr(store, spec=spec)
+    d = cubed.from_zarr(store, spec=spec)
+    e = cubed.from_zarr(store, spec=spec)
+    f = cubed.from_zarr(store, spec=spec)
+    g = cubed.from_zarr(store, spec=spec)
+    h = cubed.from_zarr(store, spec=spec)
 
     i = xp.add(a, b)
     j = xp.add(c, d)