Repository: tomwhite/cubed (python)

Issue:
Lazy Zarr array creation
Currently Cubed will create empty Zarr arrays for the intermediate data as operations are called, but before `compute()` is called. This can be slow since for some operations (like reductions) a lot of arrays need to be created on a cloud store by the client.

It would be nice if nothing was written to storage until the computation starts. This would have the following advantages:

1. Running operations before compute would be much faster. (This would make it faster and easier to try plans with different chunk sizes, etc, before running them.)
2. Only the arrays in the optimised plan need to be created.
3. Creation of the arrays could be done in parallel.
4. Creation of the arrays could be done by workers rather than the client. This would typically be faster since workers are closer to storage. (It would also open up the possibility of using a private filesystem only available to workers, and not to the client, in the future. Modal has something like this, for example.)

Note that the arrays need to be created strictly before running write operations - it wouldn't be OK to just lazily create them as a part of the write operation as the write operations are highly parallel and there would be conflicts.

Pull request description:
Make executors create Zarr arrays
Fixes #205 

The way this works is to add a new node to the DAG that creates the Zarr arrays, and a new edge from it to every Zarr array. 

Running on the examples I now see that the progress bars get shown straight away, and the first one is called "create-arrays" and has one task per array:

```
create-arrays: 100%|███████████████████████████████████████████████████████████| 3/3 [00:20<00:00,  6.82s/it]
array-003: 100%|███████████████████████████████████████████████████████████████| 4/4 [00:20<00:00,  5.12s/it]
[[ 2  4  6  8]
 [10 12 14 16]
 [18 20 22 24]                                                                                               
 [26 28 30 32]]                                                                                              
```

For the visualization there is an extra node called "create-arrays". Its dependencies are not shown since it points to every array (which would clutter the diagram), but it's useful to see it for accounting for the total number of tasks:

![cubed](https://github.com/tomwhite/cubed/assets/85085/5bc349b1-e711-4c59-8958-e076c7ac5eb4)


Patch:
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
index 1e83229a..37ee6bef 100644
--- a/cubed/core/plan.py
+++ b/cubed/core/plan.py
@@ -6,10 +6,11 @@ from datetime import datetime
 import networkx as nx
 
 from cubed.primitive.blockwise import can_fuse_pipelines, fuse
+from cubed.primitive.types import CubedPipeline
 from cubed.runtime.pipeline import already_computed
 from cubed.storage.zarr import LazyZarrArray
 from cubed.utils import chunk_memory, extract_stack_summaries, join_path, memory_repr
-from cubed.vendor.rechunker.types import PipelineExecutor
+from cubed.vendor.rechunker.types import PipelineExecutor, Stage
 
 # A unique ID with sensible ordering, used for making directory names
 CONTEXT_ID = f"cubed-{datetime.now().strftime('%Y%m%dT%H%M%S')}-{uuid.uuid4()}"
@@ -132,13 +133,35 @@ class Plan:
 
         return Plan(dag)
 
-    def create_lazy_zarr_arrays(self, dag, executor=None, **kwargs):
-        # TODO: make this op a node in the dag and run using the executor
-
-        # find all arrays in dag and create them
-        for _, d in dag.nodes(data=True):
+    def create_lazy_zarr_arrays(self, dag):
+        # find all lazy zarr arrays in dag
+        all_array_nodes = []
+        lazy_zarr_arrays = []
+        for n, d in dag.nodes(data=True):
             if isinstance(d["target"], LazyZarrArray):
-                d["target"].create(mode="a")
+                all_array_nodes.append(n)
+                lazy_zarr_arrays.append(d["target"])
+
+        if len(lazy_zarr_arrays) > 0:
+            # add new node and edges
+            name = "create-arrays"
+            op_name = name
+            pipeline = create_zarr_arrays(lazy_zarr_arrays)
+            dag.add_node(
+                name,
+                name=name,
+                op_name=op_name,
+                target=None,
+                pipeline=pipeline,
+                projected_mem=pipeline.projected_mem,
+                reserved_mem=0,  # TODO: get from Spec
+                num_tasks=pipeline.num_tasks,
+            )
+            # make create arrays node a dependency of all lazy array nodes
+            for n in all_array_nodes:
+                dag.add_edge(name, n)
+
+        return dag
 
     def execute(
         self,
@@ -150,8 +173,7 @@ class Plan:
         **kwargs,
     ):
         dag = self.optimize().dag if optimize_graph else self.dag.copy()
-
-        self.create_lazy_zarr_arrays(dag, executor=executor, **kwargs)
+        dag = self.create_lazy_zarr_arrays(dag)
 
         if isinstance(executor, PipelineExecutor):
             while len(dag) > 0:
@@ -190,6 +212,7 @@ class Plan:
     def num_tasks(self, optimize_graph=True, resume=None):
         """Return the number of tasks needed to execute this plan."""
         dag = self.optimize().dag if optimize_graph else self.dag.copy()
+        dag = self.create_lazy_zarr_arrays(dag)
         tasks = 0
         for _, node in visit_nodes(dag, resume=resume):
             pipeline = node["pipeline"]
@@ -203,6 +226,7 @@ class Plan:
     def max_projected_mem(self, optimize_graph=True, resume=None):
         """Return the maximum projected memory across all tasks to execute this plan."""
         dag = self.optimize().dag if optimize_graph else self.dag.copy()
+        dag = self.create_lazy_zarr_arrays(dag)
         projected_mem_values = [
             node["pipeline"].projected_mem
             for _, node in visit_nodes(dag, resume=resume)
@@ -213,6 +237,10 @@ class Plan:
         self, filename="cubed", format=None, rankdir="TB", optimize_graph=True
     ):
         dag = self.optimize().dag if optimize_graph else self.dag.copy()
+        dag = self.create_lazy_zarr_arrays(dag)
+
+        # remove edges from create-arrays node to avoid cluttering the diagram
+        dag.remove_edges_from(list(dag.out_edges("create-arrays")))
 
         # remove hidden nodes
         dag.remove_nodes_from(
@@ -261,19 +289,22 @@ class Plan:
             else:  # creation function
                 op_name_summary = ""
             target = d["target"]
-            chunkmem = memory_repr(chunk_memory(target.dtype, target.chunks))
-            tooltip = (
-                f"name: {n}\n"
-                f"shape: {target.shape}\n"
-                f"chunks: {target.chunks}\n"
-                f"dtype: {target.dtype}\n"
-                f"chunk memory: {chunkmem}\n"
-            )
+            if target is not None:
+                chunkmem = memory_repr(chunk_memory(target.dtype, target.chunks))
+                tooltip = (
+                    f"name: {n}\n"
+                    f"shape: {target.shape}\n"
+                    f"chunks: {target.chunks}\n"
+                    f"dtype: {target.dtype}\n"
+                    f"chunk memory: {chunkmem}\n"
+                )
+            else:
+                tooltip = ""
             if "pipeline" in d:
                 pipeline = d["pipeline"]
                 tooltip += f"\nprojected memory: {memory_repr(pipeline.projected_mem)}"
                 tooltip += f"\ntasks: {pipeline.num_tasks}"
-            if "stack_summaries" in d:
+            if "stack_summaries" in d and d["stack_summaries"] is not None:
                 # add call stack information
                 stack_summaries = d["stack_summaries"]
 
@@ -301,7 +332,7 @@ class Plan:
                 tooltip += f"\nline: {line}"
                 del d["stack_summaries"]
 
-            d["tooltip"] = tooltip
+            d["tooltip"] = tooltip.strip()
 
             # remove pipeline attribute since it is a long string that causes graphviz to fail
             if "pipeline" in d:
@@ -370,3 +401,30 @@ def visit_node_generations(dag, resume=None):
         ]
         if len(gen) > 0:
             yield gen
+
+
+def create_zarr_array(lazy_zarr_array, *, config=None):
+    """Stage function for create."""
+    lazy_zarr_array.create(mode="a")
+
+
+def create_zarr_arrays(lazy_zarr_arrays):
+    stages = [
+        Stage(
+            create_zarr_array,
+            "create_zarr_array",
+            mappable=lazy_zarr_arrays,
+        )
+    ]
+
+    # projected memory is size of largest initial values, or 0 if there aren't any
+    projected_mem = max(
+        [
+            lza.initial_values.nbytes if lza.initial_values is not None else 0
+            for lza in lazy_zarr_arrays
+        ],
+        default=0,
+    )
+    num_tasks = len(lazy_zarr_arrays)
+
+    return CubedPipeline(stages, None, None, None, projected_mem, num_tasks)


Test patch:
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index 243b6103..f54758c0 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -9,6 +9,8 @@ from numpy.testing import assert_array_equal
 
 import cubed
 import cubed.array_api as xp
+from cubed.extensions.history import HistoryCallback
+from cubed.extensions.timeline import TimelineVisualizationCallback
 from cubed.extensions.tqdm import TqdmProgressBar
 from cubed.primitive.blockwise import apply_blockwise
 from cubed.runtime.types import DagExecutor
@@ -209,8 +211,9 @@ def test_rechunk_same_chunks(spec):
     b = a.rechunk((2, 1))
     task_counter = TaskCounter()
     res = b.compute(callbacks=[task_counter])
-    # no tasks should have run since chunks are same
-    assert task_counter.value == 0
+    # no tasks except array creation task should have run since chunks are same
+    num_created_arrays = 1
+    assert task_counter.value == num_created_arrays
 
     assert_array_equal(res, np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]]))
 
@@ -427,22 +430,31 @@ def test_retries(mocker, spec):
     assert_array_equal(c.compute(), np.array([[2, 3, 4], [5, 6, 7], [8, 9, 10]]))
 
 
+@pytest.mark.skipif(
+    platform.system() == "Windows", reason="measuring memory does not run on windows"
+)
 def test_callbacks(spec, executor):
     if not isinstance(executor, DagExecutor):
         pytest.skip(f"{type(executor)} does not support callbacks")
 
     task_counter = TaskCounter()
-    progress = TqdmProgressBar()  # test indirectly (doesn't fail)
+    # test following indirectly by checking they don't cause a failure
+    progress = TqdmProgressBar()
+    hist = HistoryCallback()
+    timeline_viz = TimelineVisualizationCallback()
 
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)
     b = xp.asarray([[1, 1, 1], [1, 1, 1], [1, 1, 1]], chunks=(2, 2), spec=spec)
     c = xp.add(a, b)
     assert_array_equal(
-        c.compute(executor=executor, callbacks=[task_counter, progress]),
+        c.compute(
+            executor=executor, callbacks=[task_counter, progress, hist, timeline_viz]
+        ),
         np.array([[2, 3, 4], [5, 6, 7], [8, 9, 10]]),
     )
 
-    assert task_counter.value == 4
+    num_created_arrays = 3
+    assert task_counter.value == num_created_arrays + 4
 
 
 @pytest.mark.cloud
@@ -459,7 +471,8 @@ def test_callbacks_modal(spec, modal_executor):
             np.array([[2, 3, 4], [5, 6, 7], [8, 9, 10]]),
         )
 
-        assert task_counter.value == 4
+        num_created_arrays = 3
+        assert task_counter.value == num_created_arrays + 4
     finally:
         fs = fsspec.open(tmp_path).fs
         fs.rm(tmp_path, recursive=True)
@@ -471,16 +484,20 @@ def test_already_computed(spec):
     c = xp.add(a, b)
     d = xp.negative(c)
 
-    assert d.plan.num_tasks(optimize_graph=False) == 8
+    num_created_arrays = 4  # a, b, c, d
+    assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 8
 
     task_counter = TaskCounter()
     c.compute(callbacks=[task_counter], optimize_graph=False)
-    assert task_counter.value == 4
+    num_created_arrays = 3  # a, b, c
+    assert task_counter.value == num_created_arrays + 4
 
     # since c has already been computed, when computing d only 4 tasks are run, instead of 8
     task_counter = TaskCounter()
     d.compute(callbacks=[task_counter], optimize_graph=False, resume=True)
-    assert task_counter.value == 4
+    # the create arrays tasks are run again, even though they exist
+    num_created_arrays = 4  # a, b, c, d
+    assert task_counter.value == num_created_arrays + 4
 
 
 @pytest.mark.skipif(platform.system() == "Windows", reason="does not run on windows")
diff --git a/cubed/tests/test_optimization.py b/cubed/tests/test_optimization.py
index 01be76ee..b772ce1d 100644
--- a/cubed/tests/test_optimization.py
+++ b/cubed/tests/test_optimization.py
@@ -18,12 +18,14 @@ def test_fusion(spec):
     c = xp.astype(b, np.float32)
     d = xp.negative(c)
 
-    assert d.plan.num_tasks(optimize_graph=False) == 12
-    assert d.plan.num_tasks(optimize_graph=True) == 4
+    num_created_arrays = 4  # a, b, c, d
+    assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 12
+    num_created_arrays = 2  # a, d
+    assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 4
 
     task_counter = TaskCounter()
     result = d.compute(callbacks=[task_counter])
-    assert task_counter.value == 4
+    assert task_counter.value == num_created_arrays + 4
 
     assert_array_equal(
         result,
@@ -37,12 +39,14 @@ def test_fusion_transpose(spec):
     c = xp.astype(b, np.float32)
     d = c.T
 
-    assert d.plan.num_tasks(optimize_graph=False) == 12
-    assert d.plan.num_tasks(optimize_graph=True) == 4
+    num_created_arrays = 4  # a, b, c, d
+    assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 12
+    num_created_arrays = 2  # a, d
+    assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 4
 
     task_counter = TaskCounter()
     result = d.compute(callbacks=[task_counter])
-    assert task_counter.value == 4
+    assert task_counter.value == num_created_arrays + 4
 
     assert_array_equal(
         result,
@@ -57,12 +61,13 @@ def test_no_fusion(spec):
     c = xp.positive(b)
     d = xp.equal(b, c)
 
-    assert d.plan.num_tasks(optimize_graph=False) == 3
-    assert d.plan.num_tasks(optimize_graph=True) == 3
+    num_created_arrays = 4  # a, b, c, d
+    assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 3
+    assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 3
 
     task_counter = TaskCounter()
     result = d.compute(callbacks=[task_counter])
-    assert task_counter.value == 3
+    assert task_counter.value == num_created_arrays + 3
 
     assert_array_equal(result, np.ones((2, 2)))
 
@@ -76,11 +81,12 @@ def test_no_fusion_multiple_edges(spec):
     # this should not be fused under the current logic
     d = xp.equal(b, c)
 
-    assert d.plan.num_tasks(optimize_graph=False) == 2
-    assert d.plan.num_tasks(optimize_graph=True) == 2
+    num_created_arrays = 3  # a, c, d
+    assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 2
+    assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 2
 
     task_counter = TaskCounter()
     result = d.compute(callbacks=[task_counter])
-    assert task_counter.value == 2
+    assert task_counter.value == num_created_arrays + 2
 
     assert_array_equal(result, np.full((2, 2), True))