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

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, FlowStateItems
 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_config import SPIN_ALLOWED_DECORATORS
 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"
@@ -42,7 +40,7 @@ class BadStepDecoratorException(MetaflowException):
             "not declared as a @step. Make sure you apply this decorator "
             "on a function which has @step on the line just before the "
             "function name and @{deco} is above @step.".format(
-                deco=deco, func=func.__name__
+                deco=deco, func=getattr(func, "__name__", str(func))
             )
         )
         super(BadStepDecoratorException, self).__init__(msg)
@@ -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:
@@ -152,14 +154,16 @@ class Decorator(object):
         # Note that by design, later values override previous ones.
         self.attributes, new_user_attributes = unpack_delayed_evaluator(self.attributes)
         self._user_defined_attributes.update(new_user_attributes)
-        self.attributes = resolve_delayed_evaluator(self.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
@@ -266,7 +294,11 @@ def add_decorator_options(cmd):
 
 
 def flow_decorators(flow_cls):
-    return [d for deco_list in flow_cls._flow_decorators.values() for d in deco_list]
+    return [
+        d
+        for deco_list in flow_cls._flow_state[FlowStateItems.FLOW_DECORATORS].values()
+        for d in deco_list
+    ]
 
 
 class StepDecorator(Decorator):
@@ -315,15 +347,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 []
 
@@ -443,12 +496,20 @@ def _base_flow_decorator(decofunc, *args, **kwargs):
         cls = args[0]
         if isinstance(cls, type) and issubclass(cls, FlowSpec):
             # flow decorators add attributes in the class dictionary,
-            # _flow_decorators. _flow_decorators is of type `{key:[decos]}`
-            if decofunc.name in cls._flow_decorators and not decofunc.allow_multiple:
+            # cls._flow_state[FlowStateItems.FLOW_DECORATORS]. This is of type `{key:[decos]}`
+            self_flow_decos = cls._flow_state.self_data[FlowStateItems.FLOW_DECORATORS]
+            inherited_flow_decos = cls._flow_state.inherited_data.get(
+                FlowStateItems.FLOW_DECORATORS, {}
+            )
+
+            if (
+                decofunc.name in self_flow_decos
+                or decofunc.name in inherited_flow_decos
+            ) and not decofunc.allow_multiple:
                 raise DuplicateFlowDecoratorException(decofunc.name)
             else:
                 deco_instance = decofunc(attributes=kwargs, statically_defined=True)
-                cls._flow_decorators.setdefault(decofunc.name, []).append(deco_instance)
+                self_flow_decos.setdefault(decofunc.name, []).append(deco_instance)
         else:
             raise BadFlowDecoratorException(decofunc.name)
         return cls
@@ -472,7 +533,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, UserStepDecoratorBase)):
             func = func._my_step
         if not hasattr(func, "is_step"):
             raise BadStepDecoratorException(decotype.name, func)
@@ -498,9 +559,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 +571,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,49 +655,95 @@ def _attach_decorators_to_step(step, decospecs):
     effect as if you defined the decorators statically in the source for
     the step.
     """
+    for decospec in decospecs:
+        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
 
-    decos = _get_all_step_decos()
+        else:
+            step_deco.add_or_raise(step, False, 1, None)
 
-    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)
+
+def _should_skip_decorator_for_spin(
+    deco, is_spin, skip_decorators, logger, decorator_type="decorator"
+):
+    """
+    Determine if a decorator should be skipped for spin steps.
+
+    Parameters:
+    -----------
+    deco : Decorator
+        The decorator instance to check
+    is_spin : bool
+        Whether this is a spin step
+    skip_decorators : bool
+        Whether to skip all decorators
+    logger : callable
+        Logger function for warnings
+    decorator_type : str
+        Type of decorator ("Flow decorator" or "Step decorator") for logging
+
+    Returns:
+    --------
+    bool
+        True if the decorator should be skipped, False otherwise
+    """
+    if not is_spin:
+        return False
+
+    # Skip all decorator hooks if skip_decorators is True
+    if skip_decorators:
+        return True
+
+    # Run decorator hooks for spin steps only if they are in the whitelist
+    if deco.name not in SPIN_ALLOWED_DECORATORS:
+        logger(
+            f"[Warning] Ignoring {decorator_type} '{deco.name}' as it is not supported in spin steps.",
+            system_msg=True,
+            timestamp=False,
+            bad=True,
+        )
+        return True
+
+    return False
 
 
 def _init(flow, only_non_static=False):
-    for decorators in flow._flow_decorators.values():
+    flow_decos = flow._flow_state[FlowStateItems.FLOW_DECORATORS]
+    for decorators in flow_decos.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(
-    flow, graph, environment, flow_datastore, metadata, logger, echo, deco_options
+    flow,
+    graph,
+    environment,
+    flow_datastore,
+    metadata,
+    logger,
+    echo,
+    deco_options,
+    is_spin=False,
+    skip_decorators=False,
 ):
     # Since all flow decorators are stored as `{key:[deco]}` we iterate through each of them.
-    for decorators in flow._flow_decorators.values():
+    flow_decos = flow._flow_state[FlowStateItems.FLOW_DECORATORS]
+    for decorators in flow_decos.values():
         # First resolve the `options` for the flow decorator.
         # Options are passed from cli.
         # For example `@project` can take a `--name` / `--branch` from the cli as options.
@@ -600,6 +769,10 @@ def _init_flow_decorators(
                 for option, option_info in deco.options.items()
             }
         for deco in decorators:
+            if _should_skip_decorator_for_spin(
+                deco, is_spin, skip_decorators, logger, "Flow decorator"
+            ):
+                continue
             deco.flow_init(
                 flow,
                 graph,
@@ -612,9 +785,89 @@ def _init_flow_decorators(
             )
 
 
-def _init_step_decorators(flow, graph, environment, flow_datastore, logger):
+def _init_step_decorators(
+    flow,
+    graph,
+    environment,
+    flow_datastore,
+    logger,
+    is_spin=False,
+    skip_decorators=False,
+):
+    # NOTE: We don't need the graph but keeping it for backwards compatibility with
+    # extensions that use it directly. We will remove it at some point.
+
+    # 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[FlowStateItems.FLOW_MUTATORS]:
+        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 (mutate)" % 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 cls._flow_state[FlowStateItems.CACHED_PARAMETERS] is not None:
+                cls._flow_state[FlowStateItems.CACHED_PARAMETERS] = None
+        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 for %s (mutate)"
+                    % (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)
+
+    cls._init_graph()
+    graph = flow._graph
+
     for step in flow:
         for deco in step.decorators:
+            if _should_skip_decorator_for_spin(
+                deco, is_spin, skip_decorators, logger, "Step decorator"
+            ):
+                continue
             deco.step_init(
                 flow,
                 graph,
@@ -638,7 +891,7 @@ StepFlag = NewType("StepFlag", bool)
 
 @overload
 def step(
-    f: Callable[[FlowSpecDerived], None]
+    f: Callable[[FlowSpecDerived], None],
 ) -> Callable[[FlowSpecDerived, StepFlag], None]: ...
 
 
@@ -649,7 +902,7 @@ def step(
 
 
 def step(
-    f: Union[Callable[[FlowSpecDerived], None], Callable[[FlowSpecDerived, Any], None]]
+    f: Union[Callable[[FlowSpecDerived], None], Callable[[FlowSpecDerived, Any], None]],
 ):
     """
     Marks a method in a FlowSpec as a Metaflow Step. Note that this
@@ -685,12 +938,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