ob-metaflow 2.15.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/cli_components/run_cmds.py

@@ -8,19 +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
+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".
@@ -34,26 +41,38 @@ 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_attrs()
-        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
 
     obj.metadata.add_sticky_tags(tags=tags)
 
@@ -61,10 +80,37 @@ def before_run(obj, tags, decospecs):
     # We explicitly avoid doing this in `start` since it is invoked for every
     # step in the run.
     obj.package = MetaflowPackage(
-        obj.flow, obj.environment, obj.echo, obj.package_suffixes
+        obj.flow,
+        obj.environment,
+        obj.echo,
+        suffixes=obj.package_suffixes,
+        flow_datastore=obj.flow_datastore if FEAT_ALWAYS_UPLOAD_CODE_PACKAGE else None,
     )
 
 
+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:
@@ -129,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)
@@ -187,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,
@@ -224,6 +257,19 @@ def resume(
                     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
+        ## taken inside a branch i.e. not present in the execution path.
+
+        # origin_run = Run(f"{obj.flow.name}/{origin_run_id}", _namespace_check=False)
+        # executed_steps = {step.path_components[-1] for step in origin_run}
+        # if step_to_rerun not in executed_steps:
+        #     raise CommandException(
+        #         f"Cannot resume from step '{step_to_rerun}'. This step was not "
+        #         f"part of the original execution path for run '{origin_run_id}'."
+        #     )
+
         steps_to_rerun = {step_to_rerun}
 
     if run_id:
@@ -305,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",
@@ -327,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)
@@ -380,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,