Show nbytes for each array and total on visualization (#344)

cubed-dev · Jan 8, 2024 · 41a1225 · 41a1225
1 parent f6a4e9f
commit 41a1225
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 1 deletion.
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
@@ -232,6 +232,17 @@ def max_projected_mem(
         ]
         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 @@ def visualize(
             "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 @@ def visualize(
             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 @@ def visualize(
                 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
@@ -30,6 +30,7 @@ def __init__(
         self.shape = template.shape
         self.dtype = template.dtype
         self.chunks = template.chunks
+        self.nbytes = template.nbytes
 
         self.store = store
         self.fill_value = fill_value

diff --git 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])