Simplify resume implementation by taking out of runtime (it's an array-level feature) (#725)

Author: tomwhite

cubed/core/plan.py

@@ -6,7 +6,7 @@
 import uuid
 from datetime import datetime
 from functools import lru_cache
-from typing import Callable, Optional, Tuple
+from typing import Any, Callable, Dict, Optional, Tuple
 
 import networkx as nx
 import zarr
@@ -16,7 +16,7 @@
 from cubed.primitive.types import PrimitiveOperation
 from cubed.runtime.pipeline import visit_nodes
 from cubed.runtime.types import ComputeEndEvent, ComputeStartEvent, CubedPipeline
-from cubed.storage.zarr import LazyZarrArray
+from cubed.storage.zarr import LazyZarrArray, open_if_lazy_zarr_array
 from cubed.utils import (
     chunk_memory,
     extract_stack_summaries,
@@ -25,6 +25,11 @@
     memory_repr,
 )
 
+try:
+    from zarr.errors import ArrayNotFoundError  # type: ignore
+except ImportError:
+    ArrayNotFoundError = FileNotFoundError  # zarr-python 3
+
 # 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()}"
 
@@ -297,16 +302,22 @@ def execute(
         )
         dag = finalized_plan.dag
 
+        if resume:
+            # mark nodes as computed so they are not visited by visit_nodes
+            dag = dag.copy()
+            nodes = {n: d for (n, d) in dag.nodes(data=True)}
+            for name in list(nx.topological_sort(dag)):
+                dag.nodes[name]["computed"] = already_computed(name, dag, nodes)
+
         compute_id = f"compute-{datetime.now().strftime('%Y%m%dT%H%M%S.%f')}"
 
         if callbacks is not None:
-            event = ComputeStartEvent(compute_id, dag, resume)
+            event = ComputeStartEvent(compute_id, dag)
             [callback.on_compute_start(event) for callback in callbacks]
         executor.execute_dag(
             dag,
             compute_id=compute_id,
             callbacks=callbacks,
-            resume=resume,
             spec=spec,
             **kwargs,
         )
@@ -498,11 +509,10 @@ class FinalizedPlan:
     def __init__(self, dag):
         self.dag = dag
 
-    def max_projected_mem(self, resume=None):
+    def max_projected_mem(self):
         """Return the maximum projected memory across all tasks to execute this plan."""
         projected_mem_values = [
-            node["primitive_op"].projected_mem
-            for _, node in visit_nodes(self.dag, resume=resume)
+            node["primitive_op"].projected_mem for _, node in visit_nodes(self.dag)
         ]
         return max(projected_mem_values) if len(projected_mem_values) > 0 else 0
 
@@ -514,10 +524,10 @@ def num_primitive_ops(self) -> int:
         """Return the number of primitive operations in this plan."""
         return len(list(visit_nodes(self.dag)))
 
-    def num_tasks(self, resume=None):
+    def num_tasks(self):
         """Return the number of tasks needed to execute this plan."""
         tasks = 0
-        for _, node in visit_nodes(self.dag, resume=resume):
+        for _, node in visit_nodes(self.dag):
             tasks += node["primitive_op"].num_tasks
         return tasks
 
@@ -589,3 +599,35 @@ def create_zarr_arrays(lazy_zarr_arrays, allowed_mem, reserved_mem):
         num_tasks=num_tasks,
         fusable_with_predecessors=False,
     )
+
+
+def already_computed(name, dag, nodes: Dict[str, Any]) -> bool:
+    """
+    Return True if the array for a node doesn't have a pipeline to compute it,
+    or if all its outputs have already been computed (all chunks are present).
+    """
+    pipeline = nodes[name].get("pipeline", None)
+    if pipeline is None:
+        return True
+
+    # if no outputs have targets then need to compute (this is the create-arrays case)
+    if all(
+        [nodes[output].get("target", None) is None for output in dag.successors(name)]
+    ):
+        return False
+
+    for output in dag.successors(name):
+        target = nodes[output].get("target", None)
+        if target is not None:
+            try:
+                target = open_if_lazy_zarr_array(target)
+                if not hasattr(target, "nchunks_initialized"):
+                    raise NotImplementedError(
+                        f"Zarr array type {type(target)} does not support resume since it doesn't have a 'nchunks_initialized' property"
+                    )
+                # this check can be expensive since it has to list the directory to find nchunks_initialized
+                if target.ndim == 0 or target.nchunks_initialized != target.nchunks:
+                    return False
+            except ArrayNotFoundError:
+                return False
+    return True

cubed/diagnostics/history.py

@@ -10,7 +10,7 @@
 class HistoryCallback(Callback):
     def on_compute_start(self, event):
         plan = []
-        for name, node in visit_nodes(event.dag, event.resume):
+        for name, node in visit_nodes(event.dag):
             primitive_op = node["primitive_op"]
             plan.append(
                 dict(

cubed/diagnostics/mem_warn.py

@@ -9,7 +9,7 @@ class MemoryWarningCallback(Callback):
     def on_compute_start(self, event):
         # store ops keyed by name
         self.ops = {}
-        for name, node in visit_nodes(event.dag, event.resume):
+        for name, node in visit_nodes(event.dag):
             primitive_op = node["primitive_op"]
             self.ops[name] = primitive_op
 

cubed/diagnostics/rich.py

@@ -42,7 +42,7 @@ def on_compute_start(self, event):
         progress = logger_aware_progress.__enter__()
 
         progress_tasks = {}
-        for name, node in visit_nodes(event.dag, event.resume):
+        for name, node in visit_nodes(event.dag):
             num_tasks = node["primitive_op"].num_tasks
             op_display_name = node["op_display_name"].replace("\n", " ")
             progress_task = progress.add_task(

cubed/diagnostics/tqdm.py

@@ -21,13 +21,13 @@ def on_compute_start(self, event):
         max_op_display_name = (
             max(
                 len(node["op_display_name"].replace("\n", " "))
-                for _, node in visit_nodes(event.dag, event.resume)
+                for _, node in visit_nodes(event.dag)
             )
             + 1  # for the colon
         )
 
         self.pbars = {}
-        for i, (name, node) in enumerate(visit_nodes(event.dag, event.resume)):
+        for i, (name, node) in enumerate(visit_nodes(event.dag)):
             num_tasks = node["primitive_op"].num_tasks
             op_display_name = node["op_display_name"].replace("\n", " ") + ":"
             # note double curlies to get literal { and } for tqdm bar format

cubed/runtime/asyncio.py

@@ -126,7 +126,6 @@ async def async_map_dag(
     create_futures_func: Callable,
     dag: MultiDiGraph,
     callbacks: Optional[Sequence[Callback]] = None,
-    resume: Optional[bool] = None,
     compute_arrays_in_parallel: Optional[bool] = None,
     **kwargs,
 ) -> None:
@@ -135,7 +134,7 @@ async def async_map_dag(
     """
     if not compute_arrays_in_parallel:
         # run one pipeline at a time
-        for name, node in visit_nodes(dag, resume=resume):
+        for name, node in visit_nodes(dag):
             handle_operation_start_callbacks(callbacks, name)
             st = pipeline_to_stream(
                 create_futures_func, name, node["pipeline"], **kwargs
@@ -144,7 +143,7 @@ async def async_map_dag(
                 async for result, stats in streamer:
                     handle_callbacks(callbacks, result, stats)
     else:
-        for gen in visit_node_generations(dag, resume=resume):
+        for gen in visit_node_generations(dag):
             # run pipelines in the same topological generation in parallel by merging their streams
             streams = [
                 pipeline_to_stream(

cubed/runtime/executors/beam.py

@@ -86,13 +86,11 @@ class BeamExecutor(DagExecutor):
     def name(self) -> str:
         return "beam"
 
-    def execute_dag(
-        self, dag, callbacks=None, resume=None, spec=None, compute_id=None, **kwargs
-    ):
+    def execute_dag(self, dag, callbacks=None, spec=None, compute_id=None, **kwargs):
         dag = dag.copy()
         pipeline = beam.Pipeline(**kwargs)
 
-        for name, node in visit_nodes(dag, resume=resume):
+        for name, node in visit_nodes(dag):
             cubed_pipeline = node["pipeline"]
 
             inputs = list(dag.predecessors(name))

cubed/runtime/executors/coiled.py

@@ -31,15 +31,14 @@ def execute_dag(
         self,
         dag: MultiDiGraph,
         callbacks: Optional[Sequence[Callback]] = None,
-        resume: Optional[bool] = None,
         spec: Optional[Spec] = None,
         compute_id: Optional[str] = None,
         **coiled_kwargs: Mapping[str, Any],
     ) -> None:
         merged_kwargs = {**self.kwargs, **coiled_kwargs}
         minimum_workers = merged_kwargs.pop("minimum_workers", None)
         # Note this currently only builds the task graph for each stage once it gets to that stage in computation
-        for name, node in visit_nodes(dag, resume=resume):
+        for name, node in visit_nodes(dag):
             handle_operation_start_callbacks(callbacks, name)
             pipeline = node["pipeline"]
             coiled_function = make_coiled_function(pipeline.function, merged_kwargs)

cubed/runtime/executors/dask.py

@@ -69,7 +69,6 @@ def execute_dag(
         self,
         dag: MultiDiGraph,
         callbacks: Optional[Sequence[Callback]] = None,
-        resume: Optional[bool] = None,
         spec: Optional[Spec] = None,
         compute_id: Optional[str] = None,
         compute_kwargs: Optional[Dict[str, Any]] = None,
@@ -80,7 +79,6 @@ def execute_dag(
             self._async_execute_dag(
                 dag,
                 callbacks=callbacks,
-                resume=resume,
                 spec=spec,
                 compute_kwargs=compute_kwargs,
                 compute_id=compute_id,
@@ -92,7 +90,6 @@ async def _async_execute_dag(
         self,
         dag: MultiDiGraph,
         callbacks: Optional[Sequence[Callback]] = None,
-        resume: Optional[bool] = None,
         spec: Optional[Spec] = None,
         compute_kwargs: Optional[Dict[str, Any]] = None,
         compute_arrays_in_parallel: Optional[bool] = None,
@@ -118,7 +115,6 @@ async def _async_execute_dag(
                 create_futures_func,
                 dag=dag,
                 callbacks=callbacks,
-                resume=resume,
                 compute_arrays_in_parallel=compute_arrays_in_parallel,
                 create_backup_futures_func=create_backup_futures_func,
                 **kwargs,

cubed/runtime/executors/lithops.py

@@ -167,7 +167,6 @@ def map_unordered(
 def execute_dag(
     dag: MultiDiGraph,
     callbacks: Optional[Sequence[Callback]] = None,
-    resume: Optional[bool] = None,
     spec: Optional[Spec] = None,
     compute_arrays_in_parallel: Optional[bool] = None,
     **kwargs,
@@ -188,7 +187,7 @@ def execute_dag(
             )
     with RetryingFunctionExecutor(function_executor) as executor:
         if not compute_arrays_in_parallel:
-            for name, node in visit_nodes(dag, resume=resume):
+            for name, node in visit_nodes(dag):
                 handle_operation_start_callbacks(callbacks, name)
                 pipeline = node["pipeline"]
                 for result, stats in map_unordered(
@@ -207,7 +206,7 @@ def execute_dag(
                 ):
                     handle_callbacks(callbacks, result, stats)
         else:
-            for gen in visit_node_generations(dag, resume=resume):
+            for gen in visit_node_generations(dag):
                 group_map_functions = []
                 group_map_iterdata = []
                 group_names = []
@@ -262,7 +261,6 @@ def execute_dag(
         self,
         dag: MultiDiGraph,
         callbacks: Optional[Sequence[Callback]] = None,
-        resume: Optional[bool] = None,
         spec: Optional[Spec] = None,
         compute_id: Optional[str] = None,
         **kwargs,
@@ -271,7 +269,6 @@ def execute_dag(
         execute_dag(
             dag,
             callbacks=callbacks,
-            resume=resume,
             spec=spec,
             compute_id=compute_id,
             **merged_kwargs,