Feature: Support spinning a single metaflow step (Rebased) (#2506)

To test `spin` on a new flow you can do the following:

### Simple case:
```
python <flow_name.py> --environment=conda spin <step_name>
```

### Pass in specific pathspec:
```
python runtime_dag_flow.py --environment=conda spin RuntimeDAGFlow/13/step_c/275232971
```

### Pass in custom artifacts via module:
```
python runtime_dag_flow.py spin RuntimeDAGFlow/13/step_d/275233082 --artifacts-module ./my_artifacts.py
```

### Skip decorators (including the whitelisted ones):
```
python complex_dag_flow.py --environment=conda spin step_d --skip-decorators
```

### Use with Runner API:
```
    with Runner('complex_dag_flow.py', environment="conda").spin(
        "<Some Val>",
        artifacts_module='./artifacts/complex_dag_step_d.py',
    ) as spin:
        print("-" * 50)
        print(f"Running test for step: step_a")
        spin_task = spin.task
        print(f"my_output: {spin_task['my_output']}")
        assert spin_task['my_output'].data == [-1]
```

See the tests for more examples on hot to use this command.

---------

Co-authored-by: Romain Cledat <rcledat@netflix.com>

Author: talsperre

.gitignore

@@ -2,6 +2,8 @@ __pycache__/
 *.py[cod]
 *$py.class
 *.metaflow
+*.metaflow_spin
+metaflow_card_cache/
 
 build/
 dist/

README.md

@@ -60,5 +60,3 @@ We'd love to hear from you. Join our community [Slack workspace](http://slack.ou
 
 ## Contributing
 We welcome contributions to Metaflow. Please see our [contribution guide](https://docs.metaflow.org/introduction/contributing-to-metaflow) for more details.
-
-

metaflow/__init__.py

@@ -146,6 +146,7 @@ class and related decorators.
     metadata,
     get_metadata,
     default_metadata,
+    inspect_spin,
     Metaflow,
     Flow,
     Run,

metaflow/cli.py

@@ -1,3 +1,4 @@
+import os
 import functools
 import inspect
 import os
@@ -7,7 +8,6 @@
 
 import metaflow.tracing as tracing
 from metaflow._vendor import click
-from metaflow.system import _system_logger, _system_monitor
 
 from . import decorators, lint, metaflow_version, parameters, plugins
 from .cli_args import cli_args
@@ -27,6 +27,8 @@
     DEFAULT_PACKAGE_SUFFIXES,
 )
 from .metaflow_current import current
+from .metaflow_profile import from_start
+from metaflow.system import _system_monitor, _system_logger
 from .metaflow_environment import MetaflowEnvironment
 from .packaging_sys import MetaflowCodeContent
 from .plugins import (
@@ -38,9 +40,9 @@
 )
 from .pylint_wrapper import PyLint
 from .R import metaflow_r_version, use_r
+from .util import get_latest_run_id, resolve_identity, decompress_list
 from .user_configs.config_options import LocalFileInput, config_options
 from .user_configs.config_parameters import ConfigValue
-from .util import get_latest_run_id, resolve_identity
 
 ERASE_TO_EOL = "\033[K"
 HIGHLIGHT = "red"
@@ -125,6 +127,8 @@ def logger(body="", system_msg=False, head="", bad=False, timestamp=True, nl=Tru
         "step": "metaflow.cli_components.step_cmd.step",
         "run": "metaflow.cli_components.run_cmds.run",
         "resume": "metaflow.cli_components.run_cmds.resume",
+        "spin": "metaflow.cli_components.run_cmds.spin",
+        "spin-step": "metaflow.cli_components.step_cmd.spin_step",
     },
 )
 def cli(ctx):
@@ -318,6 +322,13 @@ def version(obj):
     hidden=True,
     is_eager=True,
 )
+@click.option(
+    "--mode",
+    type=click.Choice(["spin"]),
+    default=None,
+    help="Execution mode for metaflow CLI commands. Use 'spin' to enable "
+    "spin metadata and spin datastore for executions",
+)
 @click.pass_context
 def start(
     ctx,
@@ -335,6 +346,7 @@ def start(
     local_config_file=None,
     config=None,
     config_value=None,
+    mode=None,
     **deco_options
 ):
     if quiet:
@@ -347,6 +359,7 @@ def start(
     if use_r():
         version = metaflow_r_version()
 
+    from_start("MetaflowCLI: Starting")
     echo("Metaflow %s" % version, fg="magenta", bold=True, nl=False)
     echo(" executing *%s*" % ctx.obj.flow.name, fg="magenta", nl=False)
     echo(" for *%s*" % resolve_identity(), fg="magenta")
@@ -366,6 +379,7 @@ def start(
     ctx.obj.check = functools.partial(_check, echo)
     ctx.obj.top_cli = cli
     ctx.obj.package_suffixes = package_suffixes.split(",")
+    ctx.obj.spin_mode = mode == "spin"
 
     ctx.obj.datastore_impl = [d for d in DATASTORES if d.TYPE == datastore][0]
 
@@ -472,19 +486,12 @@ def start(
     # set force rebuild flag for environments that support it.
     ctx.obj.environment._force_rebuild = force_rebuild_environments
     ctx.obj.environment.validate_environment(ctx.obj.logger, datastore)
-
     ctx.obj.event_logger = LOGGING_SIDECARS[event_logger](
         flow=ctx.obj.flow, env=ctx.obj.environment
     )
-    ctx.obj.event_logger.start()
-    _system_logger.init_system_logger(ctx.obj.flow.name, ctx.obj.event_logger)
-
     ctx.obj.monitor = MONITOR_SIDECARS[monitor](
         flow=ctx.obj.flow, env=ctx.obj.environment
     )
-    ctx.obj.monitor.start()
-    _system_monitor.init_system_monitor(ctx.obj.flow.name, ctx.obj.monitor)
-
     ctx.obj.metadata = [m for m in METADATA_PROVIDERS if m.TYPE == metadata][0](
         ctx.obj.environment, ctx.obj.flow, ctx.obj.event_logger, ctx.obj.monitor
     )
@@ -498,6 +505,57 @@ def start(
     )
 
     ctx.obj.config_options = config_options
+    ctx.obj.is_spin = False
+    ctx.obj.skip_decorators = False
+
+    # Override values for spin steps, or if we are in spin mode
+    if (
+        hasattr(ctx, "saved_args")
+        and ctx.saved_args
+        and "spin" in ctx.saved_args[0]
+        or ctx.obj.spin_mode
+    ):
+        # To minimize side effects for spin, we will only use the following:
+        # - local metadata provider,
+        # - local datastore,
+        # - local environment,
+        # - null event logger,
+        # - null monitor
+        ctx.obj.is_spin = True
+        if "--skip-decorators" in ctx.saved_args:
+            ctx.obj.skip_decorators = True
+
+        ctx.obj.event_logger = LOGGING_SIDECARS["nullSidecarLogger"](
+            flow=ctx.obj.flow, env=ctx.obj.environment
+        )
+        ctx.obj.monitor = MONITOR_SIDECARS["nullSidecarMonitor"](
+            flow=ctx.obj.flow, env=ctx.obj.environment
+        )
+        # Use spin metadata, spin datastore, and spin datastore root
+        ctx.obj.metadata = [m for m in METADATA_PROVIDERS if m.TYPE == "spin"][0](
+            ctx.obj.environment, ctx.obj.flow, ctx.obj.event_logger, ctx.obj.monitor
+        )
+        ctx.obj.datastore_impl = [d for d in DATASTORES if d.TYPE == "spin"][0]
+        datastore_root = ctx.obj.datastore_impl.get_datastore_root_from_config(
+            ctx.obj.echo, create_on_absent=True
+        )
+        ctx.obj.datastore_impl.datastore_root = datastore_root
+
+        ctx.obj.flow_datastore = FlowDataStore(
+            ctx.obj.flow.name,
+            ctx.obj.environment,  # Same environment as run/resume
+            ctx.obj.metadata,  # local metadata
+            ctx.obj.event_logger,  # null event logger
+            ctx.obj.monitor,  # null monitor
+            storage_impl=ctx.obj.datastore_impl,
+        )
+
+    # Start event logger and monitor
+    ctx.obj.event_logger.start()
+    _system_logger.init_system_logger(ctx.obj.flow.name, ctx.obj.event_logger)
+
+    ctx.obj.monitor.start()
+    _system_monitor.init_system_monitor(ctx.obj.flow.name, ctx.obj.monitor)
 
     decorators._init(ctx.obj.flow)
 
@@ -512,9 +570,11 @@ def start(
         ctx.obj.logger,
         echo,
         deco_options,
+        ctx.obj.is_spin,
+        ctx.obj.skip_decorators,
     )
 
-    # In the case of run/resume, we will want to apply the TL decospecs
+    # In the case of run/resume/spin, we will want to apply the TL decospecs
     # *after* the run decospecs so that they don't take precedence. In other
     # words, for the same decorator, we want `myflow.py run --with foo` to
     # take precedence over any other `foo` decospec
@@ -542,11 +602,10 @@ def start(
     if (
         hasattr(ctx, "saved_args")
         and ctx.saved_args
-        and ctx.saved_args[0] not in ("run", "resume")
+        and ctx.saved_args[0] not in ("run", "resume", "spin")
     ):
-        # run/resume are special cases because they can add more decorators with --with,
+        # run/resume/spin are special cases because they can add more decorators with --with,
         # so they have to take care of themselves.
-
         all_decospecs = ctx.obj.tl_decospecs + list(
             ctx.obj.environment.decospecs() or []
         )
@@ -556,6 +615,9 @@ def start(
         # or a scheduler setting them up in their own way.
         if ctx.saved_args[0] not in ("step", "init"):
             all_decospecs += DEFAULT_DECOSPECS.split()
+        elif ctx.saved_args[0] == "spin-step":
+            # If we are in spin-args, we will not attach any decorators
+            all_decospecs = []
         if all_decospecs:
             decorators._attach_decorators(ctx.obj.flow, all_decospecs)
             decorators._init(ctx.obj.flow)
@@ -569,6 +631,9 @@ def start(
             ctx.obj.environment,
             ctx.obj.flow_datastore,
             ctx.obj.logger,
+            # The last two arguments are only used for spin steps
+            ctx.obj.is_spin,
+            ctx.obj.skip_decorators,
         )
 
         # Check the graph again (mutators may have changed it)