metaflow 2.11.16__py2.py3-none-any.whl → 2.12.1__py2.py3-none-any.whl

metaflow/cli.py

@@ -1,42 +1,21 @@
 import inspect
+import os
 import sys
 import traceback
 from datetime import datetime
 from functools import wraps
-import metaflow.tracing as tracing
 
+import metaflow.tracing as tracing
 from metaflow._vendor import click
 
-from . import lint
-from . import plugins
-from . import parameters
-from . import decorators
-from . import metaflow_version
-from . import namespace
-from .metaflow_current import current
+from . import decorators, lint, metaflow_version, namespace, parameters, plugins
 from .cli_args import cli_args
-from .tagging_util import validate_tags
-from .util import (
-    resolve_identity,
-    decompress_list,
-    write_latest_run_id,
-    get_latest_run_id,
-)
-from .task import MetaflowTask
+from .client.core import get_metadata
+from .datastore import FlowDataStore, TaskDataStore, TaskDataStoreSet
 from .exception import CommandException, MetaflowException
 from .graph import FlowGraph
-from .datastore import FlowDataStore, TaskDataStoreSet, TaskDataStore
-
-from .runtime import NativeRuntime
-from .package import MetaflowPackage
-from .plugins import (
-    DATASTORES,
-    ENVIRONMENTS,
-    LOGGING_SIDECARS,
-    METADATA_PROVIDERS,
-    MONITOR_SIDECARS,
-)
 from .metaflow_config import (
+    DECOSPECS,
     DEFAULT_DATASTORE,
     DEFAULT_ENVIRONMENT,
     DEFAULT_EVENT_LOGGER,
@@ -44,12 +23,29 @@ from .metaflow_config import (
     DEFAULT_MONITOR,
     DEFAULT_PACKAGE_SUFFIXES,
 )
+from .metaflow_current import current
 from .metaflow_environment import MetaflowEnvironment
+from .mflog import LOG_SOURCES, mflog
+from .package import MetaflowPackage
+from .plugins import (
+    DATASTORES,
+    ENVIRONMENTS,
+    LOGGING_SIDECARS,
+    METADATA_PROVIDERS,
+    MONITOR_SIDECARS,
+)
 from .pylint_wrapper import PyLint
-from .R import use_r, metaflow_r_version
-from .mflog import mflog, LOG_SOURCES
+from .R import metaflow_r_version, use_r
+from .runtime import NativeRuntime
+from .tagging_util import validate_tags
+from .task import MetaflowTask
 from .unbounded_foreach import UBF_CONTROL, UBF_TASK
-
+from .util import (
+    decompress_list,
+    get_latest_run_id,
+    resolve_identity,
+    write_latest_run_id,
+)
 
 ERASE_TO_EOL = "\033[K"
 HIGHLIGHT = "red"
@@ -124,6 +120,26 @@ def logger(body="", system_msg=False, head="", bad=False, timestamp=True, nl=Tru
     click.secho(body, bold=system_msg, fg=LOGGER_BAD_COLOR if bad else None, nl=nl)
 
 
+def config_merge_cb(ctx, param, value):
+    # Callback to:
+    #  - read  the Click auto_envvar variable from both the
+    #    environment AND the configuration
+    #  - merge that value with the value passed in the command line (value)
+    #  - return the value as a tuple
+    # Note that this function gets called even if there is no option passed on the
+    # command line.
+    # NOTE: Assumes that ctx.auto_envvar_prefix is set to METAFLOW (same as in
+    # from_conf)
+
+    # Special case where DECOSPECS and value are the same. This happens
+    # when there is no --with option at the TL and DECOSPECS is read from
+    # the env var. In this case, click also passes it as value
+    splits = DECOSPECS.split()
+    if len(splits) == len(value) and all([a == b for (a, b) in zip(splits, value)]):
+        return value
+    return tuple(list(value) + DECOSPECS.split())
+
+
 @click.group()
 def cli(ctx):
     pass
@@ -557,6 +573,13 @@ def common_run_options(func):
         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)
@@ -615,8 +638,9 @@ def resume(
     decospecs=None,
     run_id_file=None,
     resume_identifier=None,
+    runner_attribute_file=None,
 ):
-    before_run(obj, tags, decospecs + obj.environment.decospecs())
+    before_run(obj, tags, decospecs)
 
     if origin_run_id is None:
         origin_run_id = get_latest_run_id(obj.echo, obj.flow.name)
@@ -670,9 +694,14 @@ def resume(
         max_log_size=max_log_size * 1024 * 1024,
         resume_identifier=resume_identifier,
     )
-    write_run_id(run_id_file, runtime.run_id)
+    write_file(run_id_file, runtime.run_id)
     runtime.print_workflow_info()
+
     runtime.persist_constants()
+    write_file(
+        runner_attribute_file,
+        "%s:%s" % (get_metadata(), "/".join((obj.flow.name, runtime.run_id))),
+    )
     if clone_only:
         runtime.clone_original_run()
     else:
@@ -703,12 +732,13 @@ def run(
     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 + obj.environment.decospecs())
+    before_run(obj, tags, decospecs)
 
     runtime = NativeRuntime(
         obj.flow,
@@ -726,18 +756,22 @@ def run(
         max_log_size=max_log_size * 1024 * 1024,
     )
     write_latest_run_id(obj, runtime.run_id)
-    write_run_id(run_id_file, runtime.run_id)
+    write_file(run_id_file, runtime.run_id)
 
     obj.flow._set_constants(obj.graph, kwargs)
     runtime.print_workflow_info()
     runtime.persist_constants()
+    write_file(
+        runner_attribute_file,
+        "%s:%s" % (get_metadata(), "/".join((obj.flow.name, runtime.run_id))),
+    )
     runtime.execute()
 
 
-def write_run_id(run_id_file, run_id):
-    if run_id_file is not None:
-        with open(run_id_file, "w") as f:
-            f.write(str(run_id))
+def write_file(file_path, content):
+    if file_path is not None:
+        with open(file_path, "w") as f:
+            f.write(str(content))
 
 
 def before_run(obj, tags, decospecs):
@@ -752,9 +786,20 @@ def before_run(obj, tags, decospecs):
     # 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.
-    if decospecs:
-        decorators._attach_decorators(obj.flow, decospecs)
+
+    # 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)
         obj.graph = FlowGraph(obj.flow.__class__)
+
     obj.check(obj.graph, obj.flow, obj.environment, pylint=obj.pylint)
     # obj.environment.init_environment(obj.logger)
 
@@ -825,6 +870,7 @@ def version(obj):
     multiple=True,
     help="Add a decorator to all steps. You can specify this option "
     "multiple times to attach multiple decorators in steps.",
+    callback=config_merge_cb,
 )
 @click.option(
     "--pylint/--no-pylint",
@@ -943,8 +989,11 @@ def start(
         deco_options,
     )
 
-    if decospecs:
-        decorators._attach_decorators(ctx.obj.flow, decospecs)
+    # In the case of run/resume, 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
+    ctx.obj.tl_decospecs = list(decospecs or [])
 
     # initialize current and parameter context for deploy-time parameters
     current._set_env(flow=ctx.obj.flow, is_running=False)
@@ -955,7 +1004,14 @@ def start(
     if ctx.invoked_subcommand not in ("run", "resume"):
         # run/resume are special cases because they can add more decorators with --with,
         # so they have to take care of themselves.
-        decorators._attach_decorators(ctx.obj.flow, ctx.obj.environment.decospecs())
+        all_decospecs = ctx.obj.tl_decospecs + list(
+            ctx.obj.environment.decospecs() or []
+        )
+        if all_decospecs:
+            decorators._attach_decorators(ctx.obj.flow, all_decospecs)
+            # Regenerate graph if we attached more decorators
+            ctx.obj.graph = FlowGraph(ctx.obj.flow.__class__)
+
         decorators._init_step_decorators(
             ctx.obj.flow,
             ctx.obj.graph,
@@ -963,6 +1019,7 @@ def start(
             ctx.obj.flow_datastore,
             ctx.obj.logger,
         )
+
         # TODO (savin): Enable lazy instantiation of package
         ctx.obj.package = None
     if ctx.invoked_subcommand is None:

metaflow/cmd/develop/stubs.py

@@ -23,6 +23,8 @@ def _check_stubs_supported():
         if _py_ver >= (3, 4):
             if _py_ver >= (3, 8):
                 from importlib import metadata
+            elif _py_ver >= (3, 7):
+                from metaflow._vendor import importlib_metadata as metadata
             elif _py_ver >= (3, 6):
                 from metaflow._vendor.v3_6 import importlib_metadata as metadata
             else:

metaflow/decorators.py

@@ -158,13 +158,14 @@ class Decorator(object):
                     attr_list.append("%s=%s" % (k, str(v)))
                 else:
                     attr_list.append("%s=%s" % (k, json.dumps(v).replace('"', '\\"')))
+
             attrstr = ",".join(attr_list)
             return "%s:%s" % (self.name, attrstr)
         else:
             return self.name
 
     def __str__(self):
-        mode = "decorated" if self.statically_defined else "cli"
+        mode = "static" if self.statically_defined else "dynamic"
         attrs = " ".join("%s=%s" % x for x in self.attributes.items())
         if attrs:
             attrs = " " + attrs
@@ -450,6 +451,18 @@ def _base_step_decorator(decotype, *args, **kwargs):
         return wrap
 
 
+_all_step_decos = None
+
+
+def _get_all_step_decos():
+    global _all_step_decos
+    if _all_step_decos is None:
+        from .plugins import STEP_DECORATORS
+
+        _all_step_decos = {decotype.name: decotype for decotype in STEP_DECORATORS}
+    return _all_step_decos
+
+
 def _attach_decorators(flow, decospecs):
     """
     Attach decorators to all steps during runtime. This has the same
@@ -462,6 +475,7 @@ def _attach_decorators(flow, decospecs):
     #
     # Note that each step gets its own instance of the decorator class,
     # so decorator can maintain step-specific state.
+
     for step in flow:
         _attach_decorators_to_step(step, decospecs)
 
@@ -472,9 +486,8 @@ def _attach_decorators_to_step(step, decospecs):
     effect as if you defined the decorators statically in the source for
     the step.
     """
-    from .plugins import STEP_DECORATORS
 
-    decos = {decotype.name: decotype for decotype in STEP_DECORATORS}
+    decos = _get_all_step_decos()
 
     for decospec in decospecs:
         splits = decospec.split(":", 1)

metaflow/extension_support/__init__.py

@@ -262,6 +262,8 @@ if _py_ver >= (3, 4):
 
     if _py_ver >= (3, 8):
         from importlib import metadata
+    elif _py_ver >= (3, 7):
+        from metaflow._vendor import importlib_metadata as metadata
     elif _py_ver >= (3, 6):
         from metaflow._vendor.v3_6 import importlib_metadata as metadata
     else:

metaflow/metaflow_config.py

@@ -249,6 +249,14 @@ CONTACT_INFO = from_conf(
     },
 )
 
+
+###
+# Decorators
+###
+# Format is a space separated string of decospecs (what is passed
+# using --with)
+DECOSPECS = from_conf("DECOSPECS", "")
+
 ###
 # AWS Batch configuration
 ###
@@ -496,6 +504,8 @@ def get_pinned_conda_libs(python_version, datastore_type):
 try:
     from metaflow.extension_support import get_modules
 
+    _TOGGLE_DECOSPECS = []
+
     ext_modules = get_modules("config")
     for m in ext_modules:
         # We load into globals whatever we have in extension_module
@@ -519,8 +529,18 @@ try:
                     return d1
 
                 globals()[n] = _new_get_pinned_conda_libs
+            elif n == "TOGGLE_DECOSPECS":
+                if any([x.startswith("-") for x in o]):
+                    raise ValueError("Removing decospecs is not currently supported")
+                if any(" " in x for x in o):
+                    raise ValueError("Decospecs cannot contain spaces")
+                _TOGGLE_DECOSPECS.extend(o)
             elif not n.startswith("__") and not isinstance(o, types.ModuleType):
                 globals()[n] = o
+    # If DECOSPECS is set, use that, else extrapolate from extensions
+    if not DECOSPECS:
+        DECOSPECS = " ".join(_TOGGLE_DECOSPECS)
+
 finally:
     # Erase all temporary names to avoid leaking things
     for _n in [
@@ -537,6 +557,7 @@ finally:
         "v",
         "f1",
         "f2",
+        "_TOGGLE_DECOSPECS",
     ]:
         try:
             del globals()[_n]

metaflow/parameters.py

@@ -388,6 +388,7 @@ def add_custom_parameters(deploy_mode=False):
     # deploy_mode determines whether deploy-time functions should or should
     # not be evaluated for this command
     def wrapper(cmd):
+        cmd.has_flow_params = True
         # Iterate over parameters in reverse order so cmd.params lists options
         # in the order they are defined in the FlowSpec subclass
         for arg in parameters[::-1]:

metaflow/plugins/argo/argo_workflows.py

@@ -184,6 +184,12 @@ class ArgoWorkflows(object):
         # allowed by Metaflow - guaranteeing uniqueness.
         return name.replace("_", "-")
 
+    @staticmethod
+    def _sensor_name(name):
+        # Unfortunately, Argo Events Sensor names don't allow for
+        # dots (sensors run into an error) which rules out self.name :(
+        return name.replace(".", "-")
+
     @staticmethod
     def list_templates(flow_name, all=False):
         client = ArgoClient(namespace=KUBERNETES_NAMESPACE)
@@ -216,7 +222,7 @@ class ArgoWorkflows(object):
 
         # The workflow might have sensors attached to it, which consume actual resources.
         # Try to delete these as well.
-        sensor_deleted = client.delete_sensor(name)
+        sensor_deleted = client.delete_sensor(ArgoWorkflows._sensor_name(name))
 
         # After cleaning up related resources, delete the workflow in question.
         # Failure in deleting is treated as critical and will be made visible to the user
@@ -333,10 +339,9 @@ class ArgoWorkflows(object):
             argo_client.schedule_workflow_template(
                 self.name, self._schedule, self._timezone
             )
-            # Register sensor. Unfortunately, Argo Events Sensor names don't allow for
-            # dots (sensors run into an error) which rules out self.name :(
+            # Register sensor.
             # Metaflow will overwrite any existing sensor.
-            sensor_name = self.name.replace(".", "-")
+            sensor_name = ArgoWorkflows._sensor_name(self.name)
             if self._sensor:
                 argo_client.register_sensor(sensor_name, self._sensor.to_json())
             else:
@@ -2039,7 +2044,7 @@ class ArgoWorkflows(object):
             .metadata(
                 # Sensor metadata.
                 ObjectMeta()
-                .name(self.name.replace(".", "-"))
+                .name(ArgoWorkflows._sensor_name(self.name))
                 .namespace(KUBERNETES_NAMESPACE)
                 .label("app.kubernetes.io/name", "metaflow-sensor")
                 .label("app.kubernetes.io/part-of", "metaflow")

metaflow/plugins/aws/batch/batch_decorator.py

@@ -88,15 +88,15 @@ class BatchDecorator(StepDecorator):
         Alias for inferentia. Use only one of the two.
     efa : int, default 0
         Number of elastic fabric adapter network devices to attach to container
-    ephemeral_storage: int, default None
-        The total amount, in GiB, of ephemeral storage to set for the task (21-200)
+    ephemeral_storage : int, default None
+        The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
         This is only relevant for Fargate compute environments
     log_driver: str, optional, default None
         The log driver to use for the Amazon ECS container.
     log_options: List[str], optional, default None
         List of strings containing options for the chosen log driver. The configurable values
         depend on the `log driver` chosen. Validation of these options is not supported yet.
-        Example usage: ["awslogs-group:aws/batch/job"]
+        Example: [`awslogs-group:aws/batch/job`]
     """
 
     name = "batch"

metaflow/plugins/kubernetes/kubernetes_job.py

@@ -275,11 +275,6 @@ class KubernetesJob(object):
         # (unique UID) per Metaflow task attempt.
         client = self._client.get()
 
-        # tmpfs variables
-        use_tmpfs = self._kwargs["use_tmpfs"]
-        tmpfs_size = self._kwargs["tmpfs_size"]
-        tmpfs_enabled = use_tmpfs or (tmpfs_size and not use_tmpfs)
-
         self._job = client.V1Job(
             api_version="batch/v1",
             kind="Job",