metaflow 2.12.36__py2.py3-none-any.whl → 2.12.37__py2.py3-none-any.whl

metaflow/cli_args.py

@@ -12,7 +12,14 @@
 # well as the converting of options in runtime.py. We should make it so that we
 # can properly shlex things and un-shlex when using. Ideally this should all be
 # done in one place.
+#
+# NOTE: There is an important between these two as well:
+#  - this one will include local_config_file whereas the other one WILL NOT.
+#    This is because this is used when constructing the parallel UBF command which
+#    executes locally and therefore needs the local_config_file but the other (remote)
+#    commands do not.
 
+from .user_configs.config_options import ConfigInput
 from .util import to_unicode
 
 
@@ -65,6 +72,16 @@ class CLIArgs(object):
             # keyword in Python, so we call it 'decospecs' in click args
             if k == "decospecs":
                 k = "with"
+            if k in ("config_file_options", "config_value_options"):
+                # Special handling here since we gather them all in one option but actually
+                # need to send them one at a time using --config-value <name> kv.<name>.
+                # Note it can be either config_file_options or config_value_options depending
+                # on click processing order.
+                for config_name in v.keys():
+                    yield "--config-value"
+                    yield to_unicode(config_name)
+                    yield to_unicode(ConfigInput.make_key_name(config_name))
+                continue
             k = k.replace("_", "-")
             v = v if isinstance(v, (list, tuple, set)) else [v]
             for value in v:

metaflow/cli_components/dump_cmd.py

@@ -0,0 +1,96 @@
+import pickle
+
+from metaflow._vendor import click
+
+from ..cli import echo_always, echo_dev_null
+from ..datastore import TaskDataStoreSet
+from ..exception import CommandException
+
+
+@click.command(
+    help="Get data artifacts of a task or all tasks in a step. "
+    "The format for input-path is either <run_id>/<step_name> or "
+    "<run_id>/<step_name>/<task_id>."
+)
+@click.argument("input-path")
+@click.option(
+    "--private/--no-private",
+    default=False,
+    show_default=True,
+    help="Show also private attributes.",
+)
+@click.option(
+    "--max-value-size",
+    default=1000,
+    show_default=True,
+    type=int,
+    help="Show only values that are smaller than this number. "
+    "Set to 0 to see only keys.",
+)
+@click.option(
+    "--include",
+    type=str,
+    default="",
+    help="Include only artifacts in the given comma-separated list.",
+)
+@click.option(
+    "--file", type=str, default=None, help="Serialize artifacts in the given file."
+)
+@click.pass_obj
+def dump(obj, input_path, private=None, max_value_size=None, include=None, file=None):
+
+    if obj.is_quiet:
+        echo = echo_dev_null
+    else:
+        echo = echo_always
+
+    output = {}
+    kwargs = {
+        "show_private": private,
+        "max_value_size": max_value_size,
+        "include": {t for t in include.split(",") if t},
+    }
+
+    # Pathspec can either be run_id/step_name or run_id/step_name/task_id.
+    parts = input_path.split("/")
+    if len(parts) == 2:
+        run_id, step_name = parts
+        task_id = None
+    elif len(parts) == 3:
+        run_id, step_name, task_id = parts
+    else:
+        raise CommandException(
+            "input_path should either be run_id/step_name or run_id/step_name/task_id"
+        )
+
+    datastore_set = TaskDataStoreSet(
+        obj.flow_datastore,
+        run_id,
+        steps=[step_name],
+        prefetch_data_artifacts=kwargs.get("include"),
+    )
+    if task_id:
+        ds_list = [datastore_set.get_with_pathspec(input_path)]
+    else:
+        ds_list = list(datastore_set)  # get all tasks
+
+    for ds in ds_list:
+        echo(
+            "Dumping output of run_id=*{run_id}* "
+            "step=*{step}* task_id=*{task_id}*".format(
+                run_id=ds.run_id, step=ds.step_name, task_id=ds.task_id
+            ),
+            fg="magenta",
+        )
+
+        if file is None:
+            echo_always(
+                ds.format(**kwargs), highlight="green", highlight_bold=False, err=False
+            )
+        else:
+            output[ds.pathspec] = ds.to_dict(**kwargs)
+
+    if file is not None:
+        with open(file, "wb") as f:
+            pickle.dump(output, f, protocol=pickle.HIGHEST_PROTOCOL)
+        echo("Artifacts written to *%s*" % file)

metaflow/cli_components/init_cmd.py

@@ -0,0 +1,51 @@
+from metaflow._vendor import click
+
+from .. import parameters
+from ..runtime import NativeRuntime
+
+
+@parameters.add_custom_parameters(deploy_mode=False)
+@click.command(help="Internal command to initialize a run.", hidden=True)
+@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, help="ID for this instance of the step."
+)
+@click.option(
+    "--tag",
+    "tags",
+    multiple=True,
+    default=None,
+    help="Tags for this instance of the step.",
+)
+@click.pass_obj
+def init(obj, run_id=None, task_id=None, tags=None, **kwargs):
+    # init is a separate command instead of an option in 'step'
+    # since we need to capture user-specified parameters with
+    # @add_custom_parameters. Adding custom parameters to 'step'
+    # is not desirable due to the possibility of name clashes between
+    # user-specified parameters and our internal options. Note that
+    # user-specified parameters are often defined as environment
+    # variables.
+
+    obj.metadata.add_sticky_tags(tags=tags)
+
+    runtime = NativeRuntime(
+        obj.flow,
+        obj.graph,
+        obj.flow_datastore,
+        obj.metadata,
+        obj.environment,
+        obj.package,
+        obj.logger,
+        obj.entrypoint,
+        obj.event_logger,
+        obj.monitor,
+        run_id=run_id,
+    )
+    obj.flow._set_constants(obj.graph, kwargs, obj.config_options)
+    runtime.persist_constants(task_id=task_id)

metaflow/cli_components/run_cmds.py

@@ -0,0 +1,358 @@
+import json
+
+from functools import wraps
+
+from metaflow._vendor import click
+
+from .. import decorators, namespace, parameters, tracing
+from ..exception import CommandException
+from ..graph import FlowGraph
+from ..metaflow_current import current
+from ..package import MetaflowPackage
+from ..runtime import NativeRuntime
+from ..system import _system_logger
+
+from ..tagging_util import validate_tags
+from ..util import get_latest_run_id, write_latest_run_id
+
+
+def before_run(obj, tags, decospecs):
+    validate_tags(tags)
+
+    # There's a --with option both at the top-level and for the run
+    # subcommand. Why?
+    #
+    # "run --with shoes" looks so much better than "--with shoes run".
+    # This is a very common use case of --with.
+    #
+    # A downside is that we need to have the following decorators handling
+    # in two places in this module and make sure _init_step_decorators
+    # doesn't get called twice.
+
+    # We want the order to be the following:
+    # - run level decospecs
+    # - top level decospecs
+    # - environment decospecs
+    all_decospecs = (
+        list(decospecs or [])
+        + obj.tl_decospecs
+        + list(obj.environment.decospecs() or [])
+    )
+    if all_decospecs:
+        decorators._attach_decorators(obj.flow, all_decospecs)
+        decorators._init(obj.flow, only_non_static=True)
+        obj.graph = FlowGraph(obj.flow.__class__)
+
+    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.metadata.add_sticky_tags(tags=tags)
+
+    # Package working directory only once per run.
+    # 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
+    )
+
+
+def write_file(file_path, content):
+    if file_path is not None:
+        with open(file_path, "w", encoding="utf-8") as f:
+            f.write(str(content))
+
+
+def common_run_options(func):
+    @click.option(
+        "--tag",
+        "tags",
+        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 run.",
+    )
+    @click.option(
+        "--max-workers",
+        default=16,
+        show_default=True,
+        help="Maximum number of parallel processes.",
+    )
+    @click.option(
+        "--max-num-splits",
+        default=100,
+        show_default=True,
+        help="Maximum number of splits allowed in a foreach. This "
+        "is a safety check preventing bugs from triggering "
+        "thousands of steps inadvertently.",
+    )
+    @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.",
+    )
+    @click.option(
+        "--with",
+        "decospecs",
+        multiple=True,
+        help="Add a decorator to all steps. You can specify this "
+        "option multiple times to attach multiple decorators "
+        "in steps.",
+    )
+    @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
+
+
+@click.option(
+    "--origin-run-id",
+    default=None,
+    help="ID of the run that should be resumed. By default, the "
+    "last run executed locally.",
+)
+@click.option(
+    "--run-id",
+    default=None,
+    help="Run ID for the new run. By default, a new run-id will be generated",
+    hidden=True,
+)
+@click.option(
+    "--clone-only/--no-clone-only",
+    default=False,
+    show_default=True,
+    help="Only clone tasks without continuing execution",
+    hidden=True,
+)
+@click.option(
+    "--reentrant/--no-reentrant",
+    default=False,
+    show_default=True,
+    hidden=True,
+    help="If specified, allows this call to be called in parallel",
+)
+@click.option(
+    "--resume-identifier",
+    default=None,
+    show_default=True,
+    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.command(help="Resume execution of a previous run of this flow.")
+@common_run_options
+@click.pass_obj
+def resume(
+    obj,
+    tags=None,
+    step_to_rerun=None,
+    origin_run_id=None,
+    run_id=None,
+    clone_only=False,
+    reentrant=False,
+    max_workers=None,
+    max_num_splits=None,
+    max_log_size=None,
+    decospecs=None,
+    run_id_file=None,
+    resume_identifier=None,
+    runner_attribute_file=None,
+):
+    before_run(obj, tags, decospecs)
+
+    if origin_run_id is None:
+        origin_run_id = get_latest_run_id(obj.echo, obj.flow.name)
+        if origin_run_id is None:
+            raise CommandException(
+                "A previous run id was not found. Specify --origin-run-id."
+            )
+
+    if step_to_rerun is None:
+        steps_to_rerun = set()
+    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()))
+                )
+            )
+        steps_to_rerun = {step_to_rerun}
+
+    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
+        # be non-integers to avoid any clashes. This condition ensures this.
+        try:
+            int(run_id)
+        except:
+            pass
+        else:
+            raise CommandException("run-id %s cannot be an integer" % run_id)
+
+    runtime = NativeRuntime(
+        obj.flow,
+        obj.graph,
+        obj.flow_datastore,
+        obj.metadata,
+        obj.environment,
+        obj.package,
+        obj.logger,
+        obj.entrypoint,
+        obj.event_logger,
+        obj.monitor,
+        run_id=run_id,
+        clone_run_id=origin_run_id,
+        clone_only=clone_only,
+        reentrant=reentrant,
+        steps_to_rerun=steps_to_rerun,
+        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)
+    runtime.print_workflow_info()
+
+    runtime.persist_constants()
+
+    if runner_attribute_file:
+        with open(runner_attribute_file, "w", encoding="utf-8") as f:
+            json.dump(
+                {
+                    "run_id": runtime.run_id,
+                    "flow_name": obj.flow.name,
+                    "metadata": obj.metadata.metadata_str(),
+                },
+                f,
+            )
+
+    # We may skip clone-only resume if this is not a resume leader,
+    # and clone is already complete.
+    if runtime.should_skip_clone_only_execution():
+        return
+
+    current._update_env(
+        {
+            "run_id": runtime.run_id,
+        }
+    )
+    _system_logger.log_event(
+        level="info",
+        module="metaflow.resume",
+        name="start",
+        payload={
+            "msg": "Resuming run",
+        },
+    )
+
+    with runtime.run_heartbeat():
+        if clone_only:
+            runtime.clone_original_run()
+        else:
+            runtime.clone_original_run(generate_task_obj=True, verbose=False)
+            runtime.execute()
+
+
+@parameters.add_custom_parameters(deploy_mode=True)
+@click.command(help="Run the workflow locally.")
+@tracing.cli_entrypoint("cli/run")
+@common_run_options
+@click.option(
+    "--namespace",
+    "user_namespace",
+    default=None,
+    help="Change namespace from the default (your username) to "
+    "the specified tag. Note that this option does not alter "
+    "tags assigned to the objects produced by this run, just "
+    "what existing objects are visible in the client API. You "
+    "can enable the global namespace with an empty string."
+    "--namespace=",
+)
+@click.pass_obj
+def run(
+    obj,
+    tags=None,
+    max_workers=None,
+    max_num_splits=None,
+    max_log_size=None,
+    decospecs=None,
+    run_id_file=None,
+    runner_attribute_file=None,
+    user_namespace=None,
+    **kwargs
+):
+    if user_namespace is not None:
+        namespace(user_namespace or None)
+    before_run(obj, tags, decospecs)
+
+    runtime = NativeRuntime(
+        obj.flow,
+        obj.graph,
+        obj.flow_datastore,
+        obj.metadata,
+        obj.environment,
+        obj.package,
+        obj.logger,
+        obj.entrypoint,
+        obj.event_logger,
+        obj.monitor,
+        max_workers=max_workers,
+        max_num_splits=max_num_splits,
+        max_log_size=max_log_size * 1024 * 1024,
+    )
+    write_latest_run_id(obj, runtime.run_id)
+    write_file(run_id_file, runtime.run_id)
+
+    obj.flow._set_constants(obj.graph, kwargs, obj.config_options)
+    current._update_env(
+        {
+            "run_id": runtime.run_id,
+        }
+    )
+    _system_logger.log_event(
+        level="info",
+        module="metaflow.run",
+        name="start",
+        payload={
+            "msg": "Starting run",
+        },
+    )
+    with runtime.run_heartbeat():
+        runtime.print_workflow_info()
+        runtime.persist_constants()
+
+        if runner_attribute_file:
+            with open(runner_attribute_file, "w", encoding="utf-8") as f:
+                json.dump(
+                    {
+                        "run_id": runtime.run_id,
+                        "flow_name": obj.flow.name,
+                        "metadata": obj.metadata.metadata_str(),
+                    },
+                    f,
+                )
+        runtime.execute()