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

metaflow/cli_components/step_cmd.py

@@ -0,0 +1,334 @@
+from metaflow._vendor import click
+
+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, read_artifacts_module
+import metaflow.tracing as tracing
+
+
+@click.command(help="Internal command to execute a single task.", hidden=True)
+@tracing.cli("cli/step")
+@click.argument("step-name")
+@click.option(
+    "--run-id",
+    default=None,
+    required=True,
+    help="ID for one execution of all steps in the flow.",
+)
+@click.option(
+    "--task-id",
+    default=None,
+    required=True,
+    show_default=True,
+    help="ID for this instance of the step.",
+)
+@click.option(
+    "--input-paths",
+    help="A comma-separated list of pathspecs specifying inputs for this step.",
+)
+@click.option(
+    "--input-paths-filename",
+    type=click.Path(exists=True, readable=True, dir_okay=False, resolve_path=True),
+    help="A filename containing the argument typically passed to `input-paths`",
+    hidden=True,
+)
+@click.option(
+    "--split-index",
+    type=int,
+    default=None,
+    show_default=True,
+    help="Index of this foreach split.",
+)
+@click.option(
+    "--tag",
+    "opt_tag",
+    multiple=True,
+    default=None,
+    help="Annotate this run with the given tag. You can specify "
+    "this option multiple times to attach multiple tags in "
+    "the task.",
+)
+@click.option(
+    "--namespace",
+    "opt_namespace",
+    default=None,
+    help="Change namespace from the default (your username) to the specified tag.",
+)
+@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(
+    "--clone-only",
+    default=None,
+    help="Pathspec of the origin task for this task to clone. Do "
+    "not execute anything.",
+)
+@click.option(
+    "--clone-run-id",
+    default=None,
+    help="Run id of the origin flow, if this task is part of a flow being resumed.",
+)
+@click.option(
+    "--ubf-context",
+    default="none",
+    type=click.Choice(["none", UBF_CONTROL, UBF_TASK]),
+    help="Provides additional context if this task is of type unbounded foreach.",
+)
+@click.option(
+    "--num-parallel",
+    default=0,
+    type=int,
+    help="Number of parallel instances of a step. Ignored in local mode (see parallel decorator code).",
+)
+@click.pass_context
+def step(
+    ctx,
+    step_name,
+    opt_tag=None,
+    run_id=None,
+    task_id=None,
+    input_paths=None,
+    input_paths_filename=None,
+    split_index=None,
+    opt_namespace=None,
+    retry_count=None,
+    max_user_code_retries=None,
+    clone_only=None,
+    clone_run_id=None,
+    ubf_context="none",
+    num_parallel=None,
+):
+    if ctx.obj.is_quiet:
+        echo = echo_dev_null
+    else:
+        echo = echo_always
+
+    if ubf_context == "none":
+        ubf_context = None
+    if opt_namespace is not None:
+        namespace(opt_namespace)
+
+    func = None
+    try:
+        func = getattr(ctx.obj.flow, step_name)
+    except:
+        raise CommandException("Step *%s* doesn't exist." % step_name)
+    if not func.is_step:
+        raise CommandException("Function *%s* is not a step." % step_name)
+    echo("Executing a step, *%s*" % step_name, fg="magenta", bold=False)
+
+    step_kwargs = ctx.params
+    # Remove argument `step_name` from `step_kwargs`.
+    step_kwargs.pop("step_name", None)
+    # Remove `opt_*` prefix from (some) option keys.
+    step_kwargs = dict(
+        [(k[4:], v) if k.startswith("opt_") else (k, v) for k, v in step_kwargs.items()]
+    )
+    cli_args._set_step_kwargs(step_kwargs)
+
+    ctx.obj.metadata.add_sticky_tags(tags=opt_tag)
+    if not input_paths and input_paths_filename:
+        with open(input_paths_filename, mode="r", encoding="utf-8") as f:
+            input_paths = f.read().strip(" \n\"'")
+
+    paths = decompress_list(input_paths) if input_paths else []
+
+    task = MetaflowTask(
+        ctx.obj.flow,
+        ctx.obj.flow_datastore,
+        ctx.obj.metadata,
+        ctx.obj.environment,
+        ctx.obj.echo,
+        ctx.obj.event_logger,
+        ctx.obj.monitor,
+        ubf_context,
+    )
+    if clone_only:
+        task.clone_only(
+            step_name,
+            run_id,
+            task_id,
+            clone_only,
+            retry_count,
+        )
+    else:
+        task.run_step(
+            step_name,
+            run_id,
+            task_id,
+            clone_run_id,
+            paths,
+            split_index,
+            retry_count,
+            max_user_code_retries,
+        )
+
+    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/cli_components/utils.py

@@ -0,0 +1,140 @@
+import importlib
+from metaflow._vendor import click
+from metaflow.extension_support.plugins import get_plugin
+
+
+class LazyPluginCommandCollection(click.CommandCollection):
+    # lazy_source should only point to things that are resolved as CLI plugins.
+    def __init__(self, *args, lazy_sources=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        # lazy_sources is a list of strings in the form
+        # "{plugin_name}" -> "{module-name}.{command-object-name}"
+        self.lazy_sources = lazy_sources or {}
+        self._lazy_loaded = {}
+
+    def invoke(self, ctx):
+        # NOTE: This is copied from MultiCommand.invoke. The change is that we
+        # behave like chain in the sense that we evaluate the subcommand *after*
+        # invoking the base command but we don't chain the commands like self.chain
+        # would otherwise indicate.
+        # The goal of this is to make sure that the first command is properly executed
+        # *first* prior to loading the other subcommands. It's more a lazy_subcommand_load
+        # than a chain.
+        # Look for CHANGE HERE in this code to see where the changes are made.
+        # If click is updated, this may also need to be updated. This version is for
+        # click 7.1.2.
+        def _process_result(value):
+            if self.result_callback is not None:
+                value = ctx.invoke(self.result_callback, value, **ctx.params)
+            return value
+
+        if not ctx.protected_args:
+            # If we are invoked without command the chain flag controls
+            # how this happens.  If we are not in chain mode, the return
+            # value here is the return value of the command.
+            # If however we are in chain mode, the return value is the
+            # return value of the result processor invoked with an empty
+            # list (which means that no subcommand actually was executed).
+            if self.invoke_without_command:
+                # CHANGE HERE: We behave like self.chain = False here
+
+                # if not self.chain:
+                return click.Command.invoke(self, ctx)
+                # with ctx:
+                #    click.Command.invoke(self, ctx)
+                #    return _process_result([])
+
+            ctx.fail("Missing command.")
+
+        # Fetch args back out
+        args = ctx.protected_args + ctx.args
+        ctx.args = []
+        ctx.protected_args = []
+        # CHANGE HERE: Add saved_args so we have access to it in the command to be
+        # able to infer what we are calling next
+        ctx.saved_args = args
+
+        # If we're not in chain mode, we only allow the invocation of a
+        # single command but we also inform the current context about the
+        # name of the command to invoke.
+        # CHANGE HERE: We change this block to do the invoke *before* the resolve_command
+        # Make sure the context is entered so we do not clean up
+        # resources until the result processor has worked.
+        with ctx:
+            ctx.invoked_subcommand = "*" if args else None
+            click.Command.invoke(self, ctx)
+            cmd_name, cmd, args = self.resolve_command(ctx, args)
+            sub_ctx = cmd.make_context(cmd_name, args, parent=ctx)
+            with sub_ctx:
+                return _process_result(sub_ctx.command.invoke(sub_ctx))
+
+        # CHANGE HERE: Removed all the part of chain mode.
+
+    def list_commands(self, ctx):
+        base = super().list_commands(ctx)
+        for source_name, source in self.lazy_sources.items():
+            subgroup = self._lazy_load(source_name, source)
+            base.extend(subgroup.list_commands(ctx))
+        return base
+
+    def get_command(self, ctx, cmd_name):
+        base_cmd = super().get_command(ctx, cmd_name)
+        if base_cmd is not None:
+            return base_cmd
+        for source_name, source in self.lazy_sources.items():
+            subgroup = self._lazy_load(source_name, source)
+            cmd = subgroup.get_command(ctx, cmd_name)
+            if cmd is not None:
+                return cmd
+        return None
+
+    def _lazy_load(self, source_name, source_path):
+        if source_name in self._lazy_loaded:
+            return self._lazy_loaded[source_name]
+        cmd_object = get_plugin("cli", source_path, source_name)
+        if not isinstance(cmd_object, click.Group):
+            raise ValueError(
+                f"Lazy loading of {source_name} failed by returning "
+                "a non-group object"
+            )
+        self._lazy_loaded[source_name] = cmd_object
+        return cmd_object
+
+
+class LazyGroup(click.Group):
+    def __init__(self, *args, lazy_subcommands=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        # lazy_subcommands is a list of strings in the form
+        # "{command} -> "{module-name}.{command-object-name}"
+        self.lazy_subcommands = lazy_subcommands or {}
+        self._lazy_loaded = {}
+
+    def list_commands(self, ctx):
+        base = super().list_commands(ctx)
+        lazy = sorted(self.lazy_subcommands.keys())
+        return base + lazy
+
+    def get_command(self, ctx, cmd_name):
+        if cmd_name in self.lazy_subcommands:
+            return self._lazy_load(cmd_name)
+        return super().get_command(ctx, cmd_name)
+
+    def _lazy_load(self, cmd_name):
+        if cmd_name in self._lazy_loaded:
+            return self._lazy_loaded[cmd_name]
+
+        import_path = self.lazy_subcommands[cmd_name]
+        modname, cmd = import_path.rsplit(".", 1)
+        # do the import
+        mod = importlib.import_module(modname)
+        # get the Command object from that module
+        cmd_object = getattr(mod, cmd)
+        # check the result to make debugging easier. note that wrapped BaseCommand
+        # can be functions
+        if not isinstance(cmd_object, click.BaseCommand):
+            raise ValueError(
+                f"Lazy loading of {import_path} failed by returning "
+                f"a non-command object {type(cmd_object)}"
+            )
+        self._lazy_loaded[cmd_name] = cmd_object
+        return cmd_object

metaflow/client/__init__.py

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