metaflow 2.18.12__py2.py3-none-any.whl → 2.19.0__py2.py3-none-any.whl

metaflow/__init__.py

@@ -146,6 +146,7 @@ from .client import (
     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 @@ from datetime import datetime
 
 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 @@ from .metaflow_config import (
     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 .plugins import (
 )
 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)

metaflow/cli_components/run_cmds.py

@@ -8,21 +8,26 @@ from .. import decorators, namespace, parameters, tracing
 from ..exception import CommandException
 from ..graph import FlowGraph
 from ..metaflow_current import current
-from ..metaflow_config import DEFAULT_DECOSPECS, FEAT_ALWAYS_UPLOAD_CODE_PACKAGE
+from ..metaflow_config import (
+    DEFAULT_DECOSPECS,
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
+    SPIN_PERSIST,
+)
+from ..metaflow_profile import from_start
 from ..package import MetaflowPackage
-from ..runtime import NativeRuntime
+from ..runtime import NativeRuntime, SpinRuntime
 from ..system import _system_logger
 
 # from ..client.core import Run
 
 from ..tagging_util import validate_tags
-from ..util import get_latest_run_id, write_latest_run_id
+from ..util import get_latest_run_id, write_latest_run_id, parse_spin_pathspec
 
 
-def before_run(obj, tags, decospecs):
+def before_run(obj, tags, decospecs, skip_decorators=False):
     validate_tags(tags)
 
-    # There's a --with option both at the top-level and for the run
+    # There's a --with option both at the top-level and for the run/resume/spin
     # subcommand. Why?
     #
     # "run --with shoes" looks so much better than "--with shoes run".
@@ -36,26 +41,36 @@ def before_run(obj, tags, decospecs):
     # - run level decospecs
     # - top level decospecs
     # - environment decospecs
-    all_decospecs = (
-        list(decospecs or [])
-        + obj.tl_decospecs
-        + list(obj.environment.decospecs() or [])
-    )
-    if all_decospecs:
-        # These decospecs are the ones from run/resume PLUS the ones from the
-        # environment (for example the @conda)
-        decorators._attach_decorators(obj.flow, all_decospecs)
-        decorators._init(obj.flow)
-        # Regenerate graph if we attached more decorators
-        obj.flow.__class__._init_graph()
-        obj.graph = obj.flow._graph
-
-    obj.check(obj.graph, obj.flow, obj.environment, pylint=obj.pylint)
-    # obj.environment.init_environment(obj.logger)
-
-    decorators._init_step_decorators(
-        obj.flow, obj.graph, obj.environment, obj.flow_datastore, obj.logger
+    from_start(
+        f"Inside before_run, skip_decorators={skip_decorators}, is_spin={obj.is_spin}"
     )
+    if not skip_decorators:
+        all_decospecs = (
+            list(decospecs or [])
+            + obj.tl_decospecs
+            + list(obj.environment.decospecs() or [])
+        )
+        if all_decospecs:
+            # These decospecs are the ones from run/resume/spin PLUS the ones from the
+            # environment (for example the @conda)
+            decorators._attach_decorators(obj.flow, all_decospecs)
+            decorators._init(obj.flow)
+            # Regenerate graph if we attached more decorators
+            obj.flow.__class__._init_graph()
+            obj.graph = obj.flow._graph
+
+        obj.check(obj.graph, obj.flow, obj.environment, pylint=obj.pylint)
+        # obj.environment.init_environment(obj.logger)
+
+        decorators._init_step_decorators(
+            obj.flow,
+            obj.graph,
+            obj.environment,
+            obj.flow_datastore,
+            obj.logger,
+            obj.is_spin,
+            skip_decorators,
+        )
     # Re-read graph since it may have been modified by mutators
     obj.graph = obj.flow._graph
 
@@ -73,6 +88,29 @@ def before_run(obj, tags, decospecs):
     )
 
 
+def common_runner_options(func):
+    @click.option(
+        "--run-id-file",
+        default=None,
+        show_default=True,
+        type=str,
+        help="Write the ID of this run to the file specified.",
+    )
+    @click.option(
+        "--runner-attribute-file",
+        default=None,
+        show_default=True,
+        type=str,
+        help="Write the metadata and pathspec of this run to the file specified. Used internally "
+        "for Metaflow's Runner API.",
+    )
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
 def write_file(file_path, content):
     if file_path is not None:
         with open(file_path, "w", encoding="utf-8") as f:
@@ -137,20 +175,6 @@ def common_run_options(func):
         "in steps.",
         callback=config_callback,
     )
-    @click.option(
-        "--run-id-file",
-        default=None,
-        show_default=True,
-        type=str,
-        help="Write the ID of this run to the file specified.",
-    )
-    @click.option(
-        "--runner-attribute-file",
-        default=None,
-        show_default=True,
-        type=str,
-        help="Write the metadata and pathspec of this run to the file specified. Used internally for Metaflow's Runner API.",
-    )
     @wraps(func)
     def wrapper(*args, **kwargs):
         return func(*args, **kwargs)
@@ -195,6 +219,7 @@ def common_run_options(func):
 @click.command(help="Resume execution of a previous run of this flow.")
 @tracing.cli("cli/resume")
 @common_run_options
+@common_runner_options
 @click.pass_obj
 def resume(
     obj,
@@ -326,6 +351,7 @@ def resume(
 @click.command(help="Run the workflow locally.")
 @tracing.cli("cli/run")
 @common_run_options
+@common_runner_options
 @click.option(
     "--namespace",
     "user_namespace",
@@ -348,7 +374,7 @@ def run(
     run_id_file=None,
     runner_attribute_file=None,
     user_namespace=None,
-    **kwargs
+    **kwargs,
 ):
     if user_namespace is not None:
         namespace(user_namespace or None)
@@ -401,3 +427,120 @@ def run(
             )
     with runtime.run_heartbeat():
         runtime.execute()
+
+
+# @parameters.add_custom_parameters(deploy_mode=True)
+@click.command(help="Spins up a task for a given step from a previous run locally.")
+@tracing.cli("cli/spin")
+@click.argument("pathspec")
+@click.option(
+    "--skip-decorators/--no-skip-decorators",
+    is_flag=True,
+    # Default False matches the saved_args check in cli.py for spin steps - skip_decorators
+    # only becomes True when explicitly passed, otherwise decorators are applied by default
+    default=False,
+    show_default=True,
+    help="Skip decorators attached to the step or flow.",
+)
+@click.option(
+    "--artifacts-module",
+    default=None,
+    show_default=True,
+    help="Path to a module that contains artifacts to be used in the spun step. "
+    "The artifacts should be defined as a dictionary called ARTIFACTS with keys as "
+    "the artifact names and values as the artifact values. The artifact values will "
+    "overwrite the default values of the artifacts used in the spun step.",
+)
+@click.option(
+    "--persist/--no-persist",
+    "persist",
+    default=SPIN_PERSIST,
+    show_default=True,
+    help="Whether to persist the artifacts in the spun step. If set to False, "
+    "the artifacts will not be persisted and will not be available in the spun step's "
+    "datastore.",
+)
+@click.option(
+    "--max-log-size",
+    default=10,
+    show_default=True,
+    help="Maximum size of stdout and stderr captured in "
+    "megabytes. If a step outputs more than this to "
+    "stdout/stderr, its output will be truncated.",
+)
+@common_runner_options
+@click.pass_obj
+def spin(
+    obj,
+    pathspec,
+    persist=True,
+    artifacts_module=None,
+    skip_decorators=False,
+    max_log_size=None,
+    run_id_file=None,
+    runner_attribute_file=None,
+    **kwargs,
+):
+    # Parse the pathspec argument to extract step name and full pathspec
+    step_name, parsed_pathspec = parse_spin_pathspec(pathspec, obj.flow.name)
+
+    before_run(obj, [], [], skip_decorators)
+    obj.echo(f"Spinning up step *{step_name}* locally for flow *{obj.flow.name}*")
+    # For spin, flow parameters come from the original run, but _set_constants
+    # requires them in kwargs. Use parameter defaults as placeholders - they'll be
+    # overwritten when the spin step loads artifacts from the original run.
+    flow_param_defaults = {}
+    for var, param in obj.flow._get_parameters():
+        if not param.IS_CONFIG_PARAMETER:
+            default_value = param.kwargs.get("default")
+            # Use None for required parameters without defaults
+            flow_param_defaults[param.name.replace("-", "_").lower()] = default_value
+    obj.flow._set_constants(obj.graph, flow_param_defaults, obj.config_options)
+    step_func = getattr(obj.flow, step_name, None)
+    if step_func is None:
+        raise CommandException(
+            f"Step '{step_name}' not found in flow '{obj.flow.name}'. "
+            "Please provide a valid step name."
+        )
+    from_start("Spin: before spin runtime init")
+    spin_runtime = SpinRuntime(
+        obj.flow,
+        obj.graph,
+        obj.flow_datastore,
+        obj.metadata,
+        obj.environment,
+        obj.package,
+        obj.logger,
+        obj.entrypoint,
+        obj.event_logger,
+        obj.monitor,
+        step_func,
+        step_name,
+        parsed_pathspec,
+        skip_decorators,
+        artifacts_module,
+        persist,
+        max_log_size * 1024 * 1024,
+    )
+    write_latest_run_id(obj, spin_runtime.run_id)
+    write_file(run_id_file, spin_runtime.run_id)
+    # We only need the root for the metadata, i.e. the portion before DATASTORE_LOCAL_DIR
+    datastore_root = spin_runtime._flow_datastore._storage_impl.datastore_root
+    orig_task_metadata_root = datastore_root.rsplit("/", 1)[0]
+    from_start("Spin: going to execute")
+    spin_runtime.execute()
+    from_start("Spin: after spin runtime execute")
+
+    if runner_attribute_file:
+        with open(runner_attribute_file, "w") as f:
+            json.dump(
+                {
+                    "task_id": spin_runtime.task.task_id,
+                    "step_name": step_name,
+                    "run_id": spin_runtime.run_id,
+                    "flow_name": obj.flow.name,
+                    # Store metadata in a format that can be used by the Runner API
+                    "metadata": f"{obj.metadata.__class__.TYPE}@{orig_task_metadata_root}",
+                },
+                f,
+            )

metaflow/cli_components/step_cmd.py

@@ -1,12 +1,17 @@
 from metaflow._vendor import click
 
-from .. import decorators, namespace
+from .. import namespace
 from ..cli import echo_always, echo_dev_null
 from ..cli_args import cli_args
+from ..datastore.flow_datastore import FlowDataStore
 from ..exception import CommandException
+from ..client.filecache import FileCache, FileBlobCache, TaskMetadataCache
+from ..metaflow_config import SPIN_ALLOWED_DECORATORS
+from ..metaflow_profile import from_start
+from ..plugins import DATASTORES
 from ..task import MetaflowTask
 from ..unbounded_foreach import UBF_CONTROL, UBF_TASK
-from ..util import decompress_list
+from ..util import decompress_list, read_artifacts_module
 import metaflow.tracing as tracing
 
 
@@ -109,7 +114,6 @@ def step(
     ubf_context="none",
     num_parallel=None,
 ):
-
     if ctx.obj.is_quiet:
         echo = echo_dev_null
     else:
@@ -118,7 +122,7 @@ def step(
     if ubf_context == "none":
         ubf_context = None
     if opt_namespace is not None:
-        namespace(opt_namespace or None)
+        namespace(opt_namespace)
 
     func = None
     try:
@@ -176,3 +180,155 @@ def step(
         )
 
     echo("Success", fg="green", bold=True, indent=True)
+
+
+@click.command(help="Internal command to spin a single task.", hidden=True)
+@click.argument("step-name")
+@click.option(
+    "--run-id",
+    default=None,
+    required=True,
+    help="Original run ID for the step that will be spun",
+)
+@click.option(
+    "--task-id",
+    default=None,
+    required=True,
+    help="Original Task ID for the step that will be spun",
+)
+@click.option(
+    "--orig-flow-datastore",
+    show_default=True,
+    help="Original datastore for the flow from which a task is being spun",
+)
+@click.option(
+    "--input-paths",
+    help="A comma-separated list of pathspecs specifying inputs for this step.",
+)
+@click.option(
+    "--split-index",
+    type=int,
+    default=None,
+    show_default=True,
+    help="Index of this foreach split.",
+)
+@click.option(
+    "--retry-count",
+    default=0,
+    help="How many times we have attempted to run this task.",
+)
+@click.option(
+    "--max-user-code-retries",
+    default=0,
+    help="How many times we should attempt running the user code.",
+)
+@click.option(
+    "--namespace",
+    "opt_namespace",
+    default=None,
+    help="Change namespace from the default (your username) to the specified tag.",
+)
+@click.option(
+    "--skip-decorators/--no-skip-decorators",
+    is_flag=True,
+    default=False,
+    show_default=True,
+    help="Skip decorators attached to the step or flow.",
+)
+@click.option(
+    "--persist/--no-persist",
+    "persist",
+    default=True,
+    show_default=True,
+    help="Whether to persist the artifacts in the spun step. If set to false, the artifacts will not"
+    " be persisted and will not be available in the spun step's datastore.",
+)
+@click.option(
+    "--artifacts-module",
+    default=None,
+    show_default=True,
+    help="Path to a module that contains artifacts to be used in the spun step. The artifacts should "
+    "be defined as a dictionary called ARTIFACTS with keys as the artifact names and values as the "
+    "artifact values. The artifact values will overwrite the default values of the artifacts used in "
+    "the spun step.",
+)
+@click.pass_context
+def spin_step(
+    ctx,
+    step_name,
+    orig_flow_datastore,
+    run_id=None,
+    task_id=None,
+    input_paths=None,
+    split_index=None,
+    retry_count=None,
+    max_user_code_retries=None,
+    opt_namespace=None,
+    skip_decorators=False,
+    artifacts_module=None,
+    persist=True,
+):
+    import time
+
+    if ctx.obj.is_quiet:
+        echo = echo_dev_null
+    else:
+        echo = echo_always
+
+    if opt_namespace is not None:
+        namespace(opt_namespace)
+
+    input_paths = decompress_list(input_paths) if input_paths else []
+
+    skip_decorators = skip_decorators
+    whitelist_decorators = [] if skip_decorators else SPIN_ALLOWED_DECORATORS
+    from_start("SpinStep: initialized decorators")
+    spin_artifacts = read_artifacts_module(artifacts_module) if artifacts_module else {}
+    from_start("SpinStep: read artifacts module")
+
+    ds_type, ds_root = orig_flow_datastore.split("@")
+    orig_datastore_impl = [d for d in DATASTORES if d.TYPE == ds_type][0]
+    orig_datastore_impl.datastore_root = ds_root
+    orig_flow_datastore = FlowDataStore(
+        ctx.obj.flow.name,
+        environment=None,
+        storage_impl=orig_datastore_impl,
+        ds_root=ds_root,
+    )
+
+    filecache = FileCache()
+    orig_flow_datastore.set_metadata_cache(
+        TaskMetadataCache(filecache, ds_type, ds_root, ctx.obj.flow.name)
+    )
+    orig_flow_datastore.ca_store.set_blob_cache(
+        FileBlobCache(
+            filecache, FileCache.flow_ds_id(ds_type, ds_root, ctx.obj.flow.name)
+        )
+    )
+
+    task = MetaflowTask(
+        ctx.obj.flow,
+        ctx.obj.flow_datastore,
+        ctx.obj.metadata,
+        ctx.obj.environment,
+        echo,
+        ctx.obj.event_logger,
+        ctx.obj.monitor,
+        None,  # no unbounded foreach context
+        orig_flow_datastore=orig_flow_datastore,
+        spin_artifacts=spin_artifacts,
+    )
+    from_start("SpinStep: initialized task")
+    task.run_step(
+        step_name,
+        run_id,
+        task_id,
+        None,
+        input_paths,
+        split_index,
+        retry_count,
+        max_user_code_retries,
+        whitelist_decorators,
+        persist,
+    )
+    from_start("SpinStep: ran step")

metaflow/client/__init__.py

@@ -6,6 +6,7 @@ from .core import (
     metadata,
     get_metadata,
     default_metadata,
+    inspect_spin,
     Metaflow,
     Flow,
     Run,