Repository: tomwhite/cubed
Pull request #344 (2024-01-08 14:25:00)

Show nbytes for each array and total on visualization
This helps in understanding what the amount of data being stored in intermediate arrays is.


Patch:
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
index 66265697..e8a7ef97 100644
--- a/cubed/core/plan.py
+++ b/cubed/core/plan.py
@@ -232,6 +232,17 @@ class Plan:
         ]
         return max(projected_mem_values) if len(projected_mem_values) > 0 else 0
 
+    def total_nbytes(self, optimize_graph: bool = True, optimize_function=None) -> int:
+        """Return the total number of bytes for all materialized arrays in this plan."""
+        dag = self._finalize_dag(optimize_graph, optimize_function)
+        nbytes = 0
+        for _, d in dag.nodes(data=True):
+            if d.get("type") == "array":
+                target = d["target"]
+                if isinstance(target, (LazyZarrArray, zarr.Array)):
+                    nbytes += target.nbytes
+        return nbytes
+
     def visualize(
         self,
         filename="cubed",
@@ -255,7 +266,8 @@ class Plan:
             "rankdir": rankdir,
             "label": (
                 f"num tasks: {self.num_tasks(optimize_graph, optimize_function)}\n"
-                f"max projected memory: {memory_repr(self.max_projected_mem(optimize_graph, optimize_function))}"
+                f"max projected memory: {memory_repr(self.max_projected_mem(optimize_graph, optimize_function))}\n"
+                f"total nbytes: {memory_repr(self.total_nbytes(optimize_graph, optimize_function))}"
             ),
             "labelloc": "bottom",
             "labeljust": "left",
@@ -341,11 +353,13 @@ class Plan:
             elif node_type == "array":
                 target = d["target"]
                 chunkmem = memory_repr(chunk_memory(target.dtype, target.chunks))
+                nbytes = None
 
                 # materialized arrays are light orange, virtual arrays are white
                 if isinstance(target, (LazyZarrArray, zarr.Array)):
                     d["style"] = "filled"
                     d["fillcolor"] = "#ffd8b1"
+                    nbytes = memory_repr(target.nbytes)
                 if n in array_display_names:
                     var_name = array_display_names[n]
                     d["label"] = f"{n} ({var_name})"
@@ -356,6 +370,8 @@ class Plan:
                 tooltip += f"chunks: {target.chunks}\n"
                 tooltip += f"dtype: {target.dtype}\n"
                 tooltip += f"chunk memory: {chunkmem}\n"
+                if nbytes is not None:
+                    tooltip += f"nbytes: {nbytes}\n"
 
                 del d["target"]
 
diff --git a/cubed/storage/zarr.py b/cubed/storage/zarr.py
index d1f17e05..8c467e69 100644
--- a/cubed/storage/zarr.py
+++ b/cubed/storage/zarr.py
@@ -30,6 +30,7 @@ class LazyZarrArray:
         self.shape = template.shape
         self.dtype = template.dtype
         self.chunks = template.chunks
+        self.nbytes = template.nbytes
 
         self.store = store
         self.fill_value = fill_value


Test patch:
diff --git a/cubed/tests/test_core.py b/cubed/tests/test_core.py
index dff0fd6e..18299ae0 100644
--- a/cubed/tests/test_core.py
+++ b/cubed/tests/test_core.py
@@ -508,3 +508,16 @@ def test_plan_scaling(tmp_path, factor):
 
     assert c.plan.num_tasks() > 0
     c.visualize(filename=tmp_path / "c")
+
+
+@pytest.mark.parametrize("t_length", [50, 500, 5000, 50000])
+def test_plan_quad_means(tmp_path, t_length):
+    # based on sizes from https://gist.github.com/TomNicholas/c6a28f7c22c6981f75bce280d3e28283
+    spec = cubed.Spec(tmp_path, allowed_mem="2GB", reserved_mem="100MB")
+    u = cubed.random.random((t_length, 1, 987, 1920), chunks=(10, 1, -1, -1), spec=spec)
+    v = cubed.random.random((t_length, 1, 987, 1920), chunks=(10, 1, -1, -1), spec=spec)
+    uv = u * v
+    m = xp.mean(uv, axis=0)
+
+    assert m.plan.num_tasks() > 0
+    m.visualize(filename=tmp_path / "quad_means")
diff --git a/cubed/tests/test_optimization.py b/cubed/tests/test_optimization.py
index 740748b5..e406064a 100644
--- a/cubed/tests/test_optimization.py
+++ b/cubed/tests/test_optimization.py
@@ -22,10 +22,12 @@ def test_fusion(spec):
     num_created_arrays = 3  # b, c, d (a is not created on disk)
     assert d.plan.num_arrays(optimize_graph=False) == num_arrays
     assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 12
+    assert d.plan.total_nbytes(optimize_graph=False) == b.nbytes + c.nbytes + d.nbytes
     num_arrays = 2  # a, d
     num_created_arrays = 1  # d (a is not created on disk)
     assert d.plan.num_arrays(optimize_graph=True) == num_arrays
     assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 4
+    assert d.plan.total_nbytes(optimize_graph=True) == d.nbytes
 
     task_counter = TaskCounter()
     result = d.compute(callbacks=[task_counter])