Add an option to resume only until a certain step

The goal is to be complementary to spin and allow incremental development

Author: romain-intel

metaflow/cli_components/run_cmds.py

@@ -191,15 +191,24 @@ def wrapper(*args, **kwargs):
     hidden=True,
     help="If specified, it identifies the task that started this resume call. It is in the form of {step_name}-{task_id}",
 )
-@click.argument("step-to-rerun", required=False)
+@click.option(
+    "--step-only/--no-step-only",
+    default=False,
+    show_default=True,
+    help="If specified, runs up to the specified step(s) (inclusive) and stops. "
+    "If the steps are not reachable in the cloned run (ie: parent steps were "
+    "not executed), an error will be raised.",
+)
+@click.argument("steps-to-rerun", required=False, nargs=-1)
 @click.command(help="Resume execution of a previous run of this flow.")
 @tracing.cli("cli/resume")
 @common_run_options
 @click.pass_obj
 def resume(
     obj,
     tags=None,
-    step_to_rerun=None,
+    steps_to_rerun=None,
+    step_only=False,
     origin_run_id=None,
     run_id=None,
     clone_only=False,
@@ -221,17 +230,22 @@ def resume(
                 "A previous run id was not found. Specify --origin-run-id."
             )
 
-    if step_to_rerun is None:
+    if steps_to_rerun is None:
         steps_to_rerun = set()
+        if step_only:
+            raise CommandException(
+                "Cannot step-only resume without specifying at least one step to execute"
+            )
     else:
         # validate step name
-        if step_to_rerun not in obj.graph.nodes:
-            raise CommandException(
-                "invalid step name {0} specified, must be step present in "
-                "current form of execution graph. Valid step names include: {1}".format(
-                    step_to_rerun, ",".join(list(obj.graph.nodes.keys()))
+        for step_to_rerun in steps_to_rerun:
+            if step_to_rerun not in obj.graph.nodes:
+                raise CommandException(
+                    "invalid step name {0} specified, must be step present in "
+                    "current form of execution graph. Valid step names include: {1}".format(
+                        step_to_rerun, ",".join(list(obj.graph.nodes.keys()))
+                    )
                 )
-            )
 
         ## TODO: instead of checking execution path here, can add a warning later
         ## instead of throwing an error. This is for resuming a step which was not
@@ -245,8 +259,10 @@ def resume(
         #         f"part of the original execution path for run '{origin_run_id}'."
         #     )
 
-        steps_to_rerun = {step_to_rerun}
+        steps_to_rerun = set(steps_to_rerun)
 
+    if step_only:
+        clone_only = False
     if run_id:
         # Run-ids that are provided by the metadata service are always integers.
         # External providers or run-ids (like external schedulers) always need to
@@ -274,12 +290,16 @@ def resume(
         clone_only=clone_only,
         reentrant=reentrant,
         steps_to_rerun=steps_to_rerun,
+        step_only=step_only,
         max_workers=max_workers,
         max_num_splits=max_num_splits,
         max_log_size=max_log_size * 1024 * 1024,
         resume_identifier=resume_identifier,
     )
     write_file(run_id_file, runtime.run_id)
+    if step_only:
+        write_latest_run_id(obj, runtime.run_id)
+
     runtime.print_workflow_info()
 
     runtime.persist_constants()

metaflow/runtime.py

@@ -103,6 +103,7 @@ def __init__(
         clone_only=False,
         reentrant=False,
         steps_to_rerun=None,
+        step_only=False,
         max_workers=MAX_WORKERS,
         max_num_splits=MAX_NUM_SPLITS,
         max_log_size=MAX_LOG_SIZE,
@@ -145,14 +146,51 @@ def __init__(
         self._skip_decorator_hooks = skip_decorator_hooks
 
         # If steps_to_rerun is specified, we will not clone them in resume mode.
-        self._steps_to_rerun = steps_to_rerun or {}
+        self._steps_to_rerun = steps_to_rerun or set()
+        self._steps_can_clone = set()
+        self._steps_ran = set()
+        self._step_only = step_only
+        all_steps = set()
+        cannot_clone_steps = set(self._steps_to_rerun)
         # sorted_nodes are in topological order already, so we only need to
         # iterate through the nodes once to get a stable set of rerun steps.
+        # A few modes:
+        #  - no steps_to_rerun:
+        #    - not clone_only and not step_only: clone all previously executed steps and
+        #      continue execution.
+        #    - clone_only and not step_only: clone all steps that have previously executed
+        #      and stop
+        #    - not clone_only and step_only: NOT possible (requires a steps_to_rerun)
+        #    - clone_only and step_only: NOT possible (requires a steps_to_rerun)
+        #    => in all these cases, _steps_to_rerun is empty and so _steps_can_clone is
+        #      all_steps
+        #  - steps_to_rerun:
+        #    - not clone_only and not step_only: clone all previously executed steps *except*
+        #      any of the steps in steps_to_rerun and the subsequent steps. Continue execution.
+        #    => _steps_to_rerun contains the steps to rerun and all descendants. _steps_can_clone
+        #      contains all other steps
+        #    - clone_only and not step_only: clone all steps that have previously executed
+        #      up to (but not including) any of the steps in steps_to_rerun and
+        #      subsequent steps.
+        #    => same as above but steps_to_rerun is not used to run anything
+        #    - not clone_only and step_only: clone all steps that have previously executed
+        #      up to (but not including) any of the steps in steps_to_rerun and
+        #      subsequent steps. Execute *only* the steps in steps_to_rerun if possible
+        #      and stop.
+        #    - clone_only and step_only: NOT possible (if step_only is specified, we turn
+        #      off clone_only -- clone_only implies no further execution since task
+        #      objects will not be generated).
+        #    => _steps_to_rerun contains *only* the initially passed steps to run and
+        #      _steps_can_clone contains the same as in the other cases.
         for step_name in self._graph.sorted_nodes:
-            if step_name in self._steps_to_rerun:
+            all_steps.add(step_name)
+            if step_name in cannot_clone_steps:
                 out_funcs = self._graph[step_name].out_funcs or []
                 for next_step in out_funcs:
-                    self._steps_to_rerun.add(next_step)
+                    cannot_clone_steps.add(next_step)
+        self._steps_can_clone = all_steps - cannot_clone_steps
+        if not self._step_only:
+            self._steps_to_rerun = cannot_clone_steps
 
         self._origin_ds_set = None
         if clone_run_id:
@@ -399,7 +437,7 @@ def clone_original_run(self, generate_task_obj=False, verbose=True):
             if (
                 task_ds["_task_ok"]
                 and step_name != "_parameters"
-                and (step_name not in self._steps_to_rerun)
+                and (step_name in self._steps_can_clone)
             ):
                 # "_unbounded_foreach" is a special flag to indicate that the transition
                 # is an unbounded foreach.
@@ -677,6 +715,19 @@ def execute(self):
                 system_msg=True,
             )
             self._params_task.mark_resume_done()
+        elif self._step_only:
+            # Check that we ran all the steps in self._steps_to_rerun
+            steps_missing = self._steps_to_rerun - self._steps_ran
+            if steps_missing:
+                raise MetaflowInternalError(
+                    "The following steps were not executed: {0}".format(
+                        ", ".join(steps_missing)
+                    )
+                )
+            self._logger(
+                "Step-only resume complete -- all specified steps were executed!",
+                system_msg=True,
+            )
         else:
             raise MetaflowInternalError(
                 "The *end* step was not successful by the end of flow."
@@ -1073,6 +1124,8 @@ def _queue_task_foreach(self, task, next_steps):
     def _queue_tasks(self, finished_tasks):
         # finished tasks include only successful tasks
         for task in finished_tasks:
+            step_name, _, _ = task.finished_id
+            self._steps_ran.add(step_name)
             self._finished[task.finished_id] = task.path
             self._is_cloned[task.path] = task.is_cloned
 
@@ -1137,6 +1190,15 @@ def _queue_tasks(self, finished_tasks):
                     )
                 )
 
+            if self._step_only:
+                # We need to filter next_steps to only include steps that are in
+                # self._steps_to_rerun
+                next_steps = [
+                    step for step in next_steps if step in self._steps_to_rerun
+                ]
+                if not next_steps:
+                    # No steps to execute, so we can stop
+                    return
             # Different transition types require different treatment
             if any(self._graph[f].type == "join" for f in next_steps):
                 # Next step is a join