Fix reserved_mem in create_lazy_zarr_arrays (found in slow tests …

…failure) (#273)
cubed-dev · Jul 24, 2023 · 8b939dd · 8b939dd
1 parent 6f9d958
commit 8b939dd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/cubed/core/plan.py b/cubed/core/plan.py
@@ -122,8 +122,8 @@ def create_lazy_zarr_arrays(self, dag):
         lazy_zarr_arrays = []
         reserved_mem_values = []
         for n, d in dag.nodes(data=True):
-            if "reserved_mem" in d and d["reserved_mem"] is not None:
-                reserved_mem_values.append(d["reserved_mem"])
+            if "pipeline" in d and d["pipeline"].reserved_mem is not None:
+                reserved_mem_values.append(d["pipeline"].reserved_mem)
             if isinstance(d["target"], LazyZarrArray):
                 all_array_nodes.append(n)
                 lazy_zarr_arrays.append(d["target"])