metaflow 2.15.21__py2.py3-none-any.whl → 2.16.0__py2.py3-none-any.whl

metaflow/__init__.py

@@ -104,7 +104,13 @@ from .flowspec import FlowSpec
 from .parameters import Parameter, JSONTypeClass, JSONType
 
 from .user_configs.config_parameters import Config, ConfigValue, config_expr
-from .user_configs.config_decorators import CustomFlowDecorator, CustomStepDecorator
+from .user_decorators.user_step_decorator import (
+    UserStepDecorator,
+    StepMutator,
+    user_step_decorator,
+    USER_SKIP_STEP,
+)
+from .user_decorators.user_flow_decorator import FlowMutator
 
 # data layer
 # For historical reasons, we make metaflow.plugins.datatools accessible as

metaflow/cli.py

@@ -28,6 +28,7 @@ from .metaflow_config import (
 from .metaflow_current import current
 from metaflow.system import _system_monitor, _system_logger
 from .metaflow_environment import MetaflowEnvironment
+from .packaging_sys import MetaflowCodeContent
 from .plugins import (
     DATASTORES,
     ENVIRONMENTS,
@@ -152,8 +153,13 @@ def check(obj, warnings=False):
 def show(obj):
     echo_always("\n%s" % obj.graph.doc)
     for node_name in obj.graph.sorted_nodes:
+        echo_always("")
         node = obj.graph[node_name]
-        echo_always("\nStep *%s*" % node.name, err=False)
+        for deco in node.decorators:
+            echo_always("@%s" % deco.name, err=False)
+        for deco in node.wrappers:
+            echo_always("@%s" % deco.decorator_name, err=False)
+        echo_always("Step *%s*" % node.name, err=False)
         echo_always(node.doc if node.doc else "?", indent=True, err=False)
         if node.type != "end":
             echo_always(
@@ -336,6 +342,11 @@ def start(
     echo(" executing *%s*" % ctx.obj.flow.name, fg="magenta", nl=False)
     echo(" for *%s*" % resolve_identity(), fg="magenta")
 
+    # Check if we need to setup the distribution finder (if running )
+    dist_info = MetaflowCodeContent.get_distribution_finder()
+    if dist_info:
+        sys.meta_path.append(dist_info)
+
     # Setup the context
     cli_args._set_top_kwargs(ctx.params)
     ctx.obj.echo = echo
@@ -436,6 +447,10 @@ def start(
         # be raised. For resume, since we ignore those options, we ignore the error.
         raise ctx.obj.delayed_config_exception
 
+    # Init all values in the config decorators and then process them
+    for decorator in ctx.obj.flow._flow_state.get(_FlowState.CONFIG_DECORATORS, []):
+        decorator.external_init()
+
     new_cls = ctx.obj.flow._process_config_decorators(config_options)
     if new_cls:
         ctx.obj.flow = new_cls(use_cli=False)

metaflow/cli_components/init_cmd.py

@@ -46,6 +46,7 @@ def init(obj, run_id=None, task_id=None, tags=None, **kwargs):
         obj.event_logger,
         obj.monitor,
         run_id=run_id,
+        skip_decorator_hooks=True,
     )
     obj.flow._set_constants(obj.graph, kwargs, obj.config_options)
     runtime.persist_constants(task_id=task_id)

metaflow/cli_components/run_cmds.py

@@ -8,7 +8,7 @@ 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
 from ..package import MetaflowPackage
 from ..runtime import NativeRuntime
 from ..system import _system_logger
@@ -61,7 +61,11 @@ 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,
     )
 
 

metaflow/client/core.py

@@ -32,11 +32,12 @@ from metaflow.exception import (
 from metaflow.includefile import IncludedFile
 from metaflow.metaflow_config import DEFAULT_METADATA, MAX_ATTEMPTS
 from metaflow.metaflow_environment import MetaflowEnvironment
+from metaflow.package import MetaflowPackage
+from metaflow.packaging_sys import ContentType
 from metaflow.plugins import ENVIRONMENTS, METADATA_PROVIDERS
 from metaflow.unbounded_foreach import CONTROL_TASK_TAG
 from metaflow.util import cached_property, is_stringish, resolve_identity, to_unicode
 
-from ..info_file import INFO_FILE
 from .filecache import FileCache
 
 if TYPE_CHECKING:
@@ -816,20 +817,26 @@ class MetaflowCode(object):
         self._path = info["location"]
         self._ds_type = info["ds_type"]
         self._sha = info["sha"]
+        self._code_metadata = info.get(
+            "metadata",
+            '{"version": 0, "archive_format": "tgz", "mfcontent_version": 0}',
+        )
+
+        self._backend = MetaflowPackage.get_backend(self._code_metadata)
 
         if filecache is None:
             filecache = FileCache()
         _, blobdata = filecache.get_data(
             self._ds_type, self._flow_name, self._path, self._sha
         )
-        code_obj = BytesIO(blobdata)
-        self._tar = tarfile.open(fileobj=code_obj, mode="r:gz")
-        # The JSON module in Python3 deals with Unicode. Tar gives bytes.
-        info_str = (
-            self._tar.extractfile(os.path.basename(INFO_FILE)).read().decode("utf-8")
-        )
-        self._info = json.loads(info_str)
-        self._flowspec = self._tar.extractfile(self._info["script"]).read()
+        self._code_obj = BytesIO(blobdata)
+        self._info = MetaflowPackage.cls_get_info(self._code_metadata, self._code_obj)
+        if self._info:
+            self._flowspec = MetaflowPackage.cls_get_content(
+                self._code_metadata, self._code_obj, self._info["script"]
+            )
+        else:
+            raise MetaflowInternalError("Code package metadata is invalid.")
 
     @property
     def path(self) -> str:
@@ -877,7 +884,9 @@ class MetaflowCode(object):
         TarFile
             TarFile for everything in this code package
         """
-        return self._tar
+        if self._backend.type == "tgz":
+            return self._backend.cls_open(self._code_obj)
+        raise RuntimeError("Archive is not a tarball")
 
     def extract(self) -> TemporaryDirectory:
         """
@@ -908,27 +917,10 @@ class MetaflowCode(object):
             The directory and its contents are automatically deleted when
             this object is garbage collected.
         """
-        exclusions = [
-            "metaflow/",
-            "metaflow_extensions/",
-            "INFO",
-            "CONFIG_PARAMETERS",
-            "conda.manifest",
-            # This file is created when using the conda/pypi features available in
-            # nflx-metaflow-extensions: https://github.com/Netflix/metaflow-nflx-extensions
-            "condav2-1.cnd",
-        ]
-        members = [
-            m
-            for m in self.tarball.getmembers()
-            if not any(
-                (x.endswith("/") and m.name.startswith(x)) or (m.name == x)
-                for x in exclusions
-            )
-        ]
-
         tmp = TemporaryDirectory()
-        self.tarball.extractall(tmp.name, members)
+        MetaflowPackage.cls_extract_into(
+            self._code_metadata, self._code_obj, tmp.name, ContentType.USER_CONTENT
+        )
         return tmp
 
     @property

metaflow/datastore/task_datastore.py

@@ -99,7 +99,6 @@ class TaskDataStore(object):
         mode="r",
         allow_not_done=False,
     ):
-
         self._storage_impl = flow_datastore._storage_impl
         self.TYPE = self._storage_impl.TYPE
         self._ca_store = flow_datastore.ca_store

metaflow/debug.py

@@ -42,6 +42,11 @@ class Debug(object):
         filename = inspect.stack()[1][1]
         print("debug[%s %s:%s]: %s" % (typ, filename, lineno, s), file=sys.stderr)
 
+    def __getattr__(self, name):
+        # Small piece of code to get pyright and other linters to recognize that there
+        # are dynamic attributes.
+        return getattr(self, name)
+
     def noop(self, args):
         pass
 

metaflow/decorators.py

@@ -1,37 +1,35 @@
-from functools import partial
+import importlib
 import json
 import re
 
-from typing import Any, Callable, NewType, TypeVar, Union, overload
+from functools import partial
+from typing import Any, Callable, Dict, List, NewType, Tuple, TypeVar, Union, overload
 
-from .flowspec import FlowSpec
+from .flowspec import FlowSpec, _FlowState
 from .exception import (
     MetaflowInternalError,
     MetaflowException,
     InvalidDecoratorAttribute,
 )
 
+from .debug import debug
 from .parameters import current_flow
-from .user_configs.config_decorators import CustomStepDecorator
 from .user_configs.config_parameters import (
     UNPACK_KEY,
     resolve_delayed_evaluator,
     unpack_delayed_evaluator,
 )
+from .user_decorators.mutable_flow import MutableFlow
+from .user_decorators.mutable_step import MutableStep
+from .user_decorators.user_flow_decorator import FlowMutator, FlowMutatorMeta
+from .user_decorators.user_step_decorator import (
+    StepMutator,
+    UserStepDecoratorBase,
+    UserStepDecoratorMeta,
+)
 
 from metaflow._vendor import click
 
-try:
-    unicode
-except NameError:
-    unicode = str
-    basestring = str
-
-# Contains the decorators on which _init was called. We want to ensure it is called
-# only once on each decorator and, as the _init() function below can be called in
-# several places, we need to track which decorator had their init function called
-_inited_decorators = set()
-
 
 class BadStepDecoratorException(MetaflowException):
     headline = "Syntax error"
@@ -63,11 +61,14 @@ class UnknownStepDecoratorException(MetaflowException):
     headline = "Unknown step decorator"
 
     def __init__(self, deconame):
-        from .plugins import STEP_DECORATORS
-
         decos = ", ".join(
-            t.name for t in STEP_DECORATORS if not t.name.endswith("_internal")
+            [
+                x
+                for x in UserStepDecoratorMeta.all_decorators().keys()
+                if not x.endswith("_internal")
+            ]
         )
+
         msg = (
             "Unknown step decorator *{deconame}*. The following decorators are "
             "supported: *{decos}*".format(deconame=deconame, decos=decos)
@@ -92,9 +93,7 @@ class UnknownFlowDecoratorException(MetaflowException):
     headline = "Unknown flow decorator"
 
     def __init__(self, deconame):
-        from .plugins import FLOW_DECORATORS
-
-        decos = ", ".join(t.name for t in FLOW_DECORATORS)
+        decos = ", ".join(FlowMutatorMeta.all_decorators().keys())
         msg = (
             "Unknown flow decorator *{deconame}*. The following decorators are "
             "supported: *{decos}*".format(deconame=deconame, decos=decos)
@@ -123,9 +122,10 @@ class Decorator(object):
     # `allow_multiple` allows setting many decorators of the same type to a step/flow.
     allow_multiple = False
 
-    def __init__(self, attributes=None, statically_defined=False):
+    def __init__(self, attributes=None, statically_defined=False, inserted_by=None):
         self.attributes = self.defaults.copy()
         self.statically_defined = statically_defined
+        self.inserted_by = inserted_by
         self._user_defined_attributes = set()
         self._ran_init = False
 
@@ -143,7 +143,9 @@ class Decorator(object):
         Initializes the decorator. In general, any operation you would do in __init__
         should be done here.
         """
+        pass
 
+    def external_init(self):
         # In some cases (specifically when using remove_decorator), we may need to call
         # init multiple times. Short-circuit re-evaluating.
         if self._ran_init:
@@ -154,12 +156,14 @@ class Decorator(object):
         self._user_defined_attributes.update(new_user_attributes)
         self.attributes = resolve_delayed_evaluator(self.attributes, to_dict=True)
 
+        if "init" in self.__class__.__dict__:
+            self.init()
         self._ran_init = True
 
     @classmethod
-    def _parse_decorator_spec(cls, deco_spec):
+    def extract_args_kwargs_from_decorator_spec(cls, deco_spec):
         if len(deco_spec) == 0:
-            return cls()
+            return [], {}
 
         attrs = {}
         # TODO: Do we really want to allow spaces in the names of attributes?!?
@@ -179,9 +183,20 @@ class Decorator(object):
                         val_parsed = val.strip()
 
             attrs[name.strip()] = val_parsed
-        return cls(attributes=attrs)
+
+            return [], attrs
+
+    @classmethod
+    def parse_decorator_spec(cls, deco_spec):
+        if len(deco_spec) == 0:
+            return cls()
+
+        _, kwargs = cls.extract_args_kwargs_from_decorator_spec(deco_spec)
+        return cls(attributes=kwargs)
 
     def make_decorator_spec(self):
+        # Make sure all attributes are evaluated
+        self.external_init()
         attrs = {k: v for k, v in self.attributes.items() if v is not None}
         if attrs:
             attr_list = []
@@ -189,7 +204,7 @@ class Decorator(object):
             # escaping but for more complex types (typically dictionaries or lists),
             # we dump using JSON.
             for k, v in attrs.items():
-                if isinstance(v, (int, float, unicode, basestring)):
+                if isinstance(v, (int, float, str)):
                     attr_list.append("%s=%s" % (k, str(v)))
                 else:
                     attr_list.append("%s=%s" % (k, json.dumps(v).replace('"', '\\"')))
@@ -199,8 +214,21 @@ class Decorator(object):
         else:
             return self.name
 
+    def get_args_kwargs(self) -> Tuple[List[Any], Dict[str, Any]]:
+        """
+        Get the arguments and keyword arguments of the decorator.
+
+        Returns
+        -------
+        Tuple[List[Any], Dict[str, Any]]
+            A tuple containing a list of arguments and a dictionary of keyword arguments.
+        """
+        return [], dict(self.attributes)
+
     def __str__(self):
         mode = "static" if self.statically_defined else "dynamic"
+        if self.inserted_by:
+            mode += " (inserted by %s)" % " from ".join(self.inserted_by)
         attrs = " ".join("%s=%s" % x for x in self.attributes.items())
         if attrs:
             attrs = " " + attrs
@@ -315,15 +343,36 @@ class StepDecorator(Decorator):
 
     def add_to_package(self):
         """
-        Called to add custom packages needed for a decorator. This hook will be
+        Called to add custom files needed for this environment. This hook will be
         called in the `MetaflowPackage` class where metaflow compiles the code package
-        tarball. This hook is invoked in the `MetaflowPackage`'s `path_tuples`
-        function. The `path_tuples` function is a generator that yields a tuple of
-        `(file_path, arcname)`.`file_path` is the path of the file in the local file system;
-        the `arcname` is the path of the file in the constructed tarball or the path of the file
-        after decompressing the tarball.
-
-        Returns a list of tuples where each tuple represents (file_path, arcname)
+        tarball. This hook can return one of two things (the first is for backwards
+        compatibility -- move to the second):
+          - a generator yielding a tuple of `(file_path, arcname)` to add files to
+            the code package. `file_path` is the path to the file on the local filesystem
+            and `arcname` is the path relative to the packaged code.
+          - a generator yielding a tuple of `(content, arcname, type)` where:
+            - type is one of
+            ContentType.{USER_CONTENT, CODE_CONTENT, MODULE_CONTENT, OTHER_CONTENT}
+            - for USER_CONTENT:
+              - the file will be included relative to the directory containing the
+                user's flow file.
+              - content: path to the file to include
+              - arcname: path relative to the directory containing the user's flow file
+            - for CODE_CONTENT:
+              - the file will be included relative to the code directory in the package.
+                This will be the directory containing `metaflow`.
+              - content: path to the file to include
+              - arcname: path relative to the code directory in the package
+            - for MODULE_CONTENT:
+              - the module will be added to the code package as a python module. It will
+                be accessible as usual (import <module_name>)
+              - content: name of the module
+              - arcname: None (ignored)
+            - for OTHER_CONTENT:
+              - the file will be included relative to any other configuration/metadata
+                files for the flow
+              - content: path to the file to include
+              - arcname: path relative to the config directory in the package
         """
         return []
 
@@ -472,7 +521,7 @@ def _base_step_decorator(decotype, *args, **kwargs):
         # No keyword arguments specified for the decorator, e.g. @foobar.
         # The first argument is the function to be decorated.
         func = args[0]
-        if isinstance(func, CustomStepDecorator):
+        if isinstance(func, StepMutator):
             func = func._my_step
         if not hasattr(func, "is_step"):
             raise BadStepDecoratorException(decotype.name, func)
@@ -498,9 +547,10 @@ def _base_step_decorator(decotype, *args, **kwargs):
 
 
 _all_step_decos = None
+_all_flow_decos = None
 
 
-def _get_all_step_decos():
+def get_all_step_decos():
     global _all_step_decos
     if _all_step_decos is None:
         from .plugins import STEP_DECORATORS
@@ -509,6 +559,67 @@ def _get_all_step_decos():
     return _all_step_decos
 
 
+def get_all_flow_decos():
+    global _all_flow_decos
+    if _all_flow_decos is None:
+        from .plugins import FLOW_DECORATORS
+
+        _all_flow_decos = {decotype.name: decotype for decotype in FLOW_DECORATORS}
+    return _all_flow_decos
+
+
+def extract_step_decorator_from_decospec(decospec: str):
+    splits = decospec.split(":", 1)
+    deconame = splits[0]
+
+    # Check if it is a user-defined decorator or metaflow decorator
+    deco_cls = UserStepDecoratorMeta.get_decorator_by_name(deconame)
+    if deco_cls is not None:
+        return (
+            deco_cls.parse_decorator_spec(splits[1] if len(splits) > 1 else ""),
+            len(splits) > 1,
+        )
+
+    # Check if this is a decorator we can import
+    if "." in deconame:
+        # We consider this to be a import path to a user decorator so
+        # something like "my_package.my_decorator"
+        module_name, class_name = deconame.rsplit(".", 1)
+        try:
+            module = importlib.import_module(module_name)
+        except ImportError as e:
+            raise MetaflowException(
+                "Could not import user decorator %s" % deconame
+            ) from e
+        deco_cls = getattr(module, class_name, None)
+        if (
+            deco_cls is None
+            or not isinstance(deco_cls, type)
+            or not issubclass(deco_cls, UserStepDecoratorBase)
+        ):
+            raise UnknownStepDecoratorException(deconame)
+        return (
+            deco_cls.parse_decorator_spec(splits[1] if len(splits) > 1 else ""),
+            len(splits) > 1,
+        )
+
+    raise UnknownStepDecoratorException(deconame)
+
+
+def extract_flow_decorator_from_decospec(decospec: str):
+    splits = decospec.split(":", 1)
+    deconame = splits[0]
+    # Check if it is a user-defined decorator or metaflow decorator
+    deco_cls = FlowMutatorMeta.get_decorator_by_name(deconame)
+    if deco_cls is not None:
+        return (
+            deco_cls.parse_decorator_spec(splits[1] if len(splits) > 1 else ""),
+            len(splits) > 1,
+        )
+    else:
+        raise UnknownFlowDecoratorException(deconame)
+
+
 def _attach_decorators(flow, decospecs):
     """
     Attach decorators to all steps during runtime. This has the same
@@ -532,42 +643,33 @@ def _attach_decorators_to_step(step, decospecs):
     effect as if you defined the decorators statically in the source for
     the step.
     """
-
-    decos = _get_all_step_decos()
-
     for decospec in decospecs:
-        splits = decospec.split(":", 1)
-        deconame = splits[0]
-        if deconame not in decos:
-            raise UnknownStepDecoratorException(deconame)
-        # Attach the decorator to step if it doesn't have the decorator
-        # already. This means that statically defined decorators are always
-        # preferred over runtime decorators.
-        if (
-            deconame not in [deco.name for deco in step.decorators]
-            or decos[deconame].allow_multiple
-        ):
-            # if the decorator is present in a step and is of type allow_multiple
-            # then add the decorator to the step
-            deco = decos[deconame]._parse_decorator_spec(
-                splits[1] if len(splits) > 1 else ""
-            )
-            step.decorators.append(deco)
+        step_deco, _ = extract_step_decorator_from_decospec(decospec)
+        if isinstance(step_deco, StepDecorator):
+            # Check multiple
+            if (
+                step_deco.name not in [deco.name for deco in step.decorators]
+                or step_deco.allow_multiple
+            ):
+                step.decorators.append(step_deco)
+            # Else it is ignored -- this is a non-static decorator
+
+        else:
+            step_deco.add_or_raise(step, False, 1, None)
 
 
 def _init(flow, only_non_static=False):
     for decorators in flow._flow_decorators.values():
         for deco in decorators:
-            if deco in _inited_decorators:
-                continue
-            deco.init()
-            _inited_decorators.add(deco)
+            deco.external_init()
+
     for flowstep in flow:
         for deco in flowstep.decorators:
-            if deco in _inited_decorators:
-                continue
-            deco.init()
-            _inited_decorators.add(deco)
+            deco.external_init()
+        for deco in flowstep.config_decorators or []:
+            deco.external_init()
+        for deco in flowstep.wrappers or []:
+            deco.external_init()
 
 
 def _init_flow_decorators(
@@ -613,6 +715,68 @@ def _init_flow_decorators(
 
 
 def _init_step_decorators(flow, graph, environment, flow_datastore, logger):
+    # We call the mutate method for both the flow and step mutators.
+    cls = flow.__class__
+    # Run all the decorators. We first run the flow-level decorators
+    # and then the step level ones to maintain a consistent order with how
+    # other decorators are run.
+
+    for deco in cls._flow_state.get(_FlowState.CONFIG_DECORATORS, []):
+        if isinstance(deco, FlowMutator):
+            inserted_by_value = [deco.decorator_name] + (deco.inserted_by or [])
+            mutable_flow = MutableFlow(
+                cls,
+                pre_mutate=False,
+                statically_defined=deco.statically_defined,
+                inserted_by=inserted_by_value,
+            )
+            # Sanity check to make sure we are applying the decorator to the right
+            # class
+            if not deco._flow_cls == cls and not issubclass(cls, deco._flow_cls):
+                raise MetaflowInternalError(
+                    "FlowMutator registered on the wrong flow -- "
+                    "expected %s but got %s" % (deco._flow_cls.__name__, cls.__name__)
+                )
+            debug.userconf_exec(
+                "Evaluating flow level decorator %s (post)" % deco.__class__.__name__
+            )
+            deco.mutate(mutable_flow)
+            # We reset cached_parameters on the very off chance that the user added
+            # more configurations based on the configuration
+            if _FlowState.CACHED_PARAMETERS in cls._flow_state:
+                del cls._flow_state[_FlowState.CACHED_PARAMETERS]
+        else:
+            raise MetaflowInternalError(
+                "A non FlowMutator found in flow custom decorators"
+            )
+
+    for step in cls._steps:
+        for deco in step.config_decorators:
+            inserted_by_value = [deco.decorator_name] + (deco.inserted_by or [])
+
+            if isinstance(deco, StepMutator):
+                debug.userconf_exec(
+                    "Evaluating step level decorator %s (post) for %s"
+                    % (deco.__class__.__name__, step.name)
+                )
+                deco.mutate(
+                    MutableStep(
+                        cls,
+                        step,
+                        pre_mutate=False,
+                        statically_defined=deco.statically_defined,
+                        inserted_by=inserted_by_value,
+                    )
+                )
+            else:
+                raise MetaflowInternalError(
+                    "A non StepMutator found in step custom decorators"
+                )
+
+        if step.config_decorators:
+            # We remove all mention of the custom step decorator
+            setattr(cls, step.name, step)
+
     for step in flow:
         for deco in step.decorators:
             deco.step_init(
@@ -685,12 +849,8 @@ def step(
     f.is_step = True
     f.decorators = []
     f.config_decorators = []
-    try:
-        # python 3
-        f.name = f.__name__
-    except:
-        # python 2
-        f.name = f.__func__.func_name
+    f.wrappers = []
+    f.name = f.__name__
     return f