Pass array names to optimize functions

Author: tomwhite

cubed/core/array.py

@@ -285,6 +285,7 @@ def compute(
         optimize_graph=optimize_graph,
         optimize_function=optimize_function,
         resume=resume,
+        array_names=tuple(a.name for a in arrays),
         spec=spec,
         **kwargs,
     )
@@ -335,6 +336,7 @@ def visualize(
         optimize_graph=optimize_graph,
         optimize_function=optimize_function,
         show_hidden=show_hidden,
+        array_names=tuple(a.name for a in arrays),
     )
 
 

cubed/core/optimization.py

@@ -12,7 +12,7 @@
 logger = logging.getLogger(__name__)
 
 
-def simple_optimize_dag(dag):
+def simple_optimize_dag(dag, array_names=None):
     """Apply map blocks fusion."""
 
     # note there is no need to prune the dag, since the way it is built
@@ -205,6 +205,7 @@ def fuse_predecessors(
     dag,
     name,
     *,
+    array_names=None,
     max_total_source_arrays=4,
     max_total_num_input_blocks=None,
     always_fuse=None,
@@ -258,6 +259,7 @@ def fuse_predecessors(
 def multiple_inputs_optimize_dag(
     dag,
     *,
+    array_names=None,
     max_total_source_arrays=4,
     max_total_num_input_blocks=None,
     always_fuse=None,
@@ -270,6 +272,7 @@ def multiple_inputs_optimize_dag(
         dag = fuse_predecessors(
             dag,
             name,
+            array_names=array_names,
             max_total_source_arrays=max_total_source_arrays,
             max_total_num_input_blocks=max_total_num_input_blocks,
             always_fuse=always_fuse,
@@ -278,18 +281,20 @@ def multiple_inputs_optimize_dag(
     return dag
 
 
-def fuse_all_optimize_dag(dag):
+def fuse_all_optimize_dag(dag, array_names=None):
     """Force all operations to be fused."""
     dag = dag.copy()
     always_fuse = [op for op in dag.nodes() if op.startswith("op-")]
-    return multiple_inputs_optimize_dag(dag, always_fuse=always_fuse)
+    return multiple_inputs_optimize_dag(
+        dag, array_names=array_names, always_fuse=always_fuse
+    )
 
 
-def fuse_only_optimize_dag(dag, *, only_fuse=None):
+def fuse_only_optimize_dag(dag, *, array_names=None, only_fuse=None):
     """Force only specified operations to be fused, all others will be left even if they are suitable for fusion."""
     dag = dag.copy()
     always_fuse = only_fuse
     never_fuse = set(op for op in dag.nodes() if op.startswith("op-")) - set(only_fuse)
     return multiple_inputs_optimize_dag(
-        dag, always_fuse=always_fuse, never_fuse=never_fuse
+        dag, array_names=array_names, always_fuse=always_fuse, never_fuse=never_fuse
     )

cubed/core/plan.py

@@ -6,7 +6,7 @@
 import uuid
 from datetime import datetime
 from functools import lru_cache
-from typing import Callable, Optional
+from typing import Callable, Optional, Tuple
 
 import networkx as nx
 import zarr
@@ -154,10 +154,11 @@ def arrays_to_plan(cls, *arrays):
     def optimize(
         self,
         optimize_function: Optional[Callable[..., nx.MultiDiGraph]] = None,
+        array_names: Optional[Tuple[str]] = None,
     ):
         if optimize_function is None:
             optimize_function = multiple_inputs_optimize_dag
-        dag = optimize_function(self.dag)
+        dag = optimize_function(self.dag, array_names=array_names)
         return Plan(dag)
 
     def _create_lazy_zarr_arrays(self, dag):
@@ -243,8 +244,9 @@ def _finalize(
         optimize_graph: bool = True,
         optimize_function=None,
         compile_function: Optional[Decorator] = None,
+        array_names=None,
     ) -> "FinalizedPlan":
-        dag = self.optimize(optimize_function).dag if optimize_graph else self.dag
+        dag = self.optimize(optimize_function, array_names).dag if optimize_graph else self.dag
         # create a copy since _create_lazy_zarr_arrays mutates the dag
         dag = dag.copy()
         if callable(compile_function):
@@ -260,11 +262,12 @@ def execute(
         optimize_function=None,
         compile_function=None,
         resume=None,
+        array_names=None,
         spec=None,
         **kwargs,
     ):
         finalized_plan = self._finalize(
-            optimize_graph, optimize_function, compile_function
+            optimize_graph, optimize_function, compile_function, array_names=array_names
         )
         dag = finalized_plan.dag
 
@@ -293,8 +296,11 @@ def visualize(
         optimize_graph=True,
         optimize_function=None,
         show_hidden=False,
+        array_names=None,
     ):
-        finalized_plan = self._finalize(optimize_graph, optimize_function)
+        finalized_plan = self._finalize(
+            optimize_graph, optimize_function, array_names=array_names
+        )
         dag = finalized_plan.dag
         dag = dag.copy()  # make a copy since we mutate the DAG below
 

cubed/tests/test_optimization.py

@@ -198,7 +198,7 @@ def test_custom_optimize_function(spec):
         < num_tasks_with_no_optimization
     )
 
-    def custom_optimize_function(dag):
+    def custom_optimize_function(dag, array_names=None):
         # leave DAG unchanged
         return dag