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

metaflow/user_decorators/mutable_step.py

@@ -0,0 +1,424 @@
+from functools import partial
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    List,
+    Optional,
+    TYPE_CHECKING,
+    Tuple,
+    Union,
+)
+
+from metaflow.debug import debug
+from metaflow.exception import MetaflowException
+
+from .user_step_decorator import StepMutator, UserStepDecoratorBase
+
+if TYPE_CHECKING:
+    import metaflow.decorators
+    import metaflow.flowspec
+
+
+class MutableStep:
+    IGNORE = 1
+    ERROR = 2
+    OVERRIDE = 3
+
+    def __init__(
+        self,
+        flow_spec: "metaflow.flowspec.FlowSpec",
+        step: Union[
+            Callable[["metaflow.decorators.FlowSpecDerived"], None],
+            Callable[["metaflow.decorators.FlowSpecDerived", Any], None],
+        ],
+        pre_mutate: bool = False,
+        statically_defined: bool = False,
+        inserted_by: Optional[str] = None,
+    ):
+        from .mutable_flow import MutableFlow
+
+        self._mutable_container = MutableFlow(
+            flow_spec,
+            pre_mutate=pre_mutate,
+            statically_defined=statically_defined,
+            inserted_by=inserted_by,
+        )
+        self._flow_cls = flow_spec.__class__
+        self._my_step = step
+        self._pre_mutate = pre_mutate
+        self._statically_defined = statically_defined
+        self._inserted_by = inserted_by
+        if self._inserted_by is None:
+            # This is an error because MutableSteps should only be created by
+            # StepMutators or FlowMutators. We need to catch it now because otherwise
+            # we may put stuff on the command line (with --with) that would get added
+            # twice and weird behavior may ensue.
+            raise MetaflowException(
+                "MutableStep should only be created by StepMutators or FlowMutators. "
+                "This is an internal error."
+            )
+
+    @property
+    def flow(self) -> "metaflow.user_decorator.mutable_flow.MutableFlow":
+        """
+        The flow that contains this step
+
+        Returns
+        -------
+        MutableFlow
+            The flow that contains this step
+        """
+        return self._mutable_container
+
+    @property
+    def decorator_specs(
+        self,
+    ) -> Generator[Tuple[str, str, List[Any], Dict[str, Any]], None, None]:
+        """
+        Iterate over all the decorator specifications of this step. Note that the same
+        type of decorator may be present multiple times and no order is guaranteed.
+
+        The returned tuple contains:
+        - The decorator's name (shortest possible)
+        - The decorator's fully qualified name (in the case of Metaflow decorators, this
+          will indicate which extension the decorator comes from)
+        - A list of positional arguments
+        - A dictionary of keyword arguments
+
+        You can use the resulting tuple to remove a decorator for example
+
+        Yields
+        ------
+        str, str, List[Any], Dict[str, Any]
+            A tuple containing the decorator name, it's fully qualified name,
+            a list of positional arguments, and a dictionary of keyword arguments.
+        """
+        for deco in self._my_step.decorators:
+            # 3.7 does not support yield foo, *bar syntax so we
+            # work around
+            r = [
+                deco.name,
+                "%s.%s"
+                % (
+                    deco.__class__.__module__,
+                    deco.__class__.__name__,
+                ),
+            ]
+            r.extend(deco.get_args_kwargs())
+            yield tuple(r)
+
+        for deco in self._my_step.wrappers:
+            r = [
+                UserStepDecoratorBase.get_decorator_name(deco.__class__),
+                deco.decorator_name,
+            ]
+            r.extend(deco.get_args_kwargs())
+            yield tuple(r)
+
+        for deco in self._my_step.config_decorators:
+            r = [
+                UserStepDecoratorBase.get_decorator_name(deco.__class__),
+                deco.decorator_name,
+            ]
+            r.extend(deco.get_args_kwargs())
+            yield tuple(r)
+
+    def add_decorator(
+        self,
+        deco_type: Union[partial, UserStepDecoratorBase, str],
+        deco_args: Optional[List[Any]] = None,
+        deco_kwargs: Optional[Dict[str, Any]] = None,
+        duplicates: int = IGNORE,
+    ) -> None:
+        """
+        Add a Metaflow step-decorator or a user step-decorator to a step.
+
+        You can either add the decorator itself or its decorator specification for it
+        (the same you would get back from decorator_specs). You can also mix and match
+        but you cannot provide arguments both through the string and the
+        deco_args/deco_kwargs.
+
+        As an example:
+        ```
+        from metaflow import environment
+        ...
+        my_step.add_decorator(environment, deco_kwargs={"vars": {"foo": 42})}
+        ```
+
+        is equivalent to:
+        ```
+        my_step.add_decorator('environment:vars={"foo": 42}')
+        ```
+
+        is equivalent to:
+        ```
+        my_step.add_decorator('environment', deco_kwargs={"vars":{"foo": 42}})
+        ```
+
+        but this is not allowed:
+        ```
+        my_step.add_decorator('environment:vars={"bar" 43}', deco_kwargs={"vars":{"foo": 42}})
+        ```
+
+        Note in the case where you specify a
+        string for the decorator, there is no need to import the decorator.
+
+        The string syntax is useful to, for example, allow decorators to be stored as
+        strings in a configuration file.
+
+        You can only add StepMutators in the pre_mutate stage.
+
+        In terms of precedence for decorators:
+          - if a decorator can be applied multiple times (like `@card`) all decorators
+            added are kept.
+          - if `duplicates` is set to `MutableStep.IGNORE`, then the decorator
+            being added is ignored (in other words, the existing decorator has precedence).
+          - if `duplicates` is set to `MutableStep.OVERRIDE`, then the *existing*
+            decorator is removed and this newly added one replaces it (in other
+            words, the newly added decorator has precedence).
+          - if `duplicates` is set to `MutableStep.ERROR`, then an error is raised but only
+            if the newly added decorator is *static* (ie: defined directly in the code).
+            If not, it is ignored.
+
+        Parameters
+        ----------
+        deco_type : Union[partial, UserStepDecoratorBase, str]
+            The decorator class to add to this step.
+        deco_args : List[Any], optional, default None
+            Positional arguments to pass to the decorator.
+        deco_kwargs : Dict[str, Any], optional, default None
+            Keyword arguments to pass to the decorator.
+        duplicates : int, default MutableStep.IGNORE
+            Instruction on how to handle duplicates. It can be one of:
+            - `MutableStep.IGNORE`: Ignore the decorator if it already exists.
+            - `MutableStep.ERROR`: Raise an error if the decorator already exists.
+            - `MutableStep.OVERRIDE`: Remove the existing decorator and add this one.
+        """
+        # Prevent circular import
+        from metaflow.decorators import (
+            DuplicateStepDecoratorException,
+            StepDecorator,
+            extract_step_decorator_from_decospec,
+        )
+
+        deco_args = deco_args or []
+        deco_kwargs = deco_kwargs or {}
+
+        def _add_step_decorator(step_deco):
+            if deco_args:
+                raise MetaflowException(
+                    "Step decorators do not take additional positional arguments"
+                )
+            # Update kwargs:
+            step_deco.attributes.update(deco_kwargs)
+
+            # Check duplicates
+            def _do_add():
+                step_deco.statically_defined = self._statically_defined
+                step_deco.inserted_by = self._inserted_by
+                self._my_step.decorators.append(step_deco)
+                debug.userconf_exec(
+                    "Mutable step adding step decorator '%s' to step '%s'"
+                    % (deco_type, self._my_step.name)
+                )
+
+            existing_deco = [
+                d for d in self._my_step.decorators if d.name == step_deco.name
+            ]
+
+            if step_deco.allow_multiple or not existing_deco:
+                _do_add()
+            elif duplicates == MutableStep.IGNORE:
+                # If we ignore, we do not add the decorator
+                debug.userconf_exec(
+                    "Mutable step ignoring step decorator '%s' on step '%s' "
+                    "(already exists and duplicates are ignored)"
+                    % (step_deco.name, self._my_step.name)
+                )
+            elif duplicates == MutableStep.OVERRIDE:
+                # If we override, we remove the existing decorator and add this one
+                debug.userconf_exec(
+                    "Mutable step overriding step decorator '%s' on step '%s' "
+                    "(removing existing decorator and adding new one)"
+                    % (step_deco.name, self._my_step.name)
+                )
+                self._my_step.decorators = [
+                    d for d in self._my_step.decorators if d.name != step_deco.name
+                ]
+                _do_add()
+            elif duplicates == MutableStep.ERROR:
+                # If we error, we raise an exception
+                if self._statically_defined:
+                    raise DuplicateStepDecoratorException(step_deco.name, self._my_step)
+                else:
+                    debug.userconf_exec(
+                        "Mutable step ignoring step decorator '%s' on step '%s' "
+                        "(already exists and non statically defined)"
+                        % (step_deco.name, self._my_step.name)
+                    )
+            else:
+                raise ValueError("Invalid duplicates value: %s" % duplicates)
+
+        if isinstance(deco_type, str):
+            step_deco, has_args_kwargs = extract_step_decorator_from_decospec(deco_type)
+            if (deco_args or deco_kwargs) and has_args_kwargs:
+                raise MetaflowException(
+                    "Cannot specify additional arguments when adding a user step "
+                    "decorator using a decospec that already has arguments"
+                )
+
+            if isinstance(step_deco, StepDecorator):
+                _add_step_decorator(step_deco)
+            else:
+                # User defined decorator.
+                if not self._pre_mutate and isinstance(step_deco, StepMutator):
+                    raise MetaflowException(
+                        "Adding step mutator '%s' from %s is only allowed in the "
+                        "`pre_mutate` method and not the `mutate` method"
+                        % (step_deco.decorator_name, self._inserted_by)
+                    )
+
+                if deco_args or deco_kwargs:
+                    # We need to recreate the object if there were args or kwargs
+                    # since they were not in the string
+                    step_deco = step_deco.__class__(*deco_args, **deco_kwargs)
+
+                step_deco.add_or_raise(
+                    self._my_step,
+                    self._statically_defined,
+                    duplicates,
+                    self._inserted_by,
+                )
+            return
+
+        if isinstance(deco_type, type) and issubclass(deco_type, UserStepDecoratorBase):
+            # We can only add step mutators in the pre mutate stage.
+            if not self._pre_mutate and issubclass(deco_type, StepMutator):
+                raise MetaflowException(
+                    "Adding step mutator '%s' from %s is only allowed in the "
+                    "`pre_mutate` method and not the `mutate` method"
+                    % (step_deco.decorator_name, self._inserted_by)
+                )
+            debug.userconf_exec(
+                "Mutable step adding decorator %s to step %s"
+                % (deco_type, self._my_step.name)
+            )
+
+            d = deco_type(*deco_args, **deco_kwargs)
+            # add_or_raise properly registers the decorator
+            d.add_or_raise(
+                self._my_step, self._statically_defined, duplicates, self._inserted_by
+            )
+            return
+
+        # At this point, it should be a regular Metaflow step decorator
+        if (
+            not isinstance(deco_type, partial)
+            or len(deco_type.args) != 1
+            or not issubclass(deco_type.args[0], StepDecorator)
+        ):
+            raise TypeError(
+                "add_decorator takes a metaflow decorator or user StepDecorator"
+            )
+
+        deco_type = deco_type.args[0]
+        _add_step_decorator(
+            deco_type(
+                attributes=deco_kwargs,
+                statically_defined=self._statically_defined,
+                inserted_by=self._inserted_by,
+            )
+        )
+
+    def remove_decorator(
+        self,
+        deco_name: str,
+        deco_args: Optional[List[Any]] = None,
+        deco_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> bool:
+        """
+        Remove a step-level decorator. To remove a decorator, you can pass the decorator
+        specification (obtained from `decorator_specs` for example).
+        Note that if multiple decorators share the same decorator specification
+        (very rare), they will all be removed.
+
+        You can only remove StepMutators in the `pre_mutate` method.
+
+        Parameters
+        ----------
+        deco_name : str
+            Decorator specification of the decorator to remove. If nothing else is
+            specified, all decorators matching that name will be removed.
+        deco_args : List[Any], optional, default None
+            Positional arguments to match the decorator specification.
+        deco_kwargs : Dict[str, Any], optional, default None
+            Keyword arguments to match the decorator specification.
+
+        Returns
+        -------
+        bool
+            Returns True if a decorator was removed.
+        """
+
+        do_all = deco_args is None and deco_kwargs is None
+        did_remove = False
+        canonical_deco_type = UserStepDecoratorBase.get_decorator_by_name(deco_name)
+        if issubclass(canonical_deco_type, UserStepDecoratorBase):
+            for attr in ["config_decorators", "wrappers"]:
+                new_deco_list = []
+
+                for deco in getattr(self._my_step, attr):
+                    if deco.decorator_name == canonical_deco_type.decorator_name:
+                        if do_all:
+                            continue  # We remove all decorators with this name
+                        if deco.get_args_kwargs() == (
+                            deco_args or [],
+                            deco_kwargs or {},
+                        ):
+                            if not self._pre_mutate and isinstance(deco, StepMutator):
+                                raise MetaflowException(
+                                    "Removing step mutator '%s' from %s is only allowed in the "
+                                    "`pre_mutate` method and not the `mutate` method"
+                                    % (deco.decorator_name, self._inserted_by)
+                                )
+                            did_remove = True
+                            debug.userconf_exec(
+                                "Mutable step removing user step decorator '%s' from step '%s'"
+                                % (deco.decorator_name, self._my_step.name)
+                            )
+                        else:
+                            new_deco_list.append(deco)
+                    else:
+                        new_deco_list.append(deco)
+                setattr(self._my_step, attr, new_deco_list)
+
+        if did_remove:
+            return True
+        for deco in self._my_step.decorators:
+            if deco.name == deco_name:
+                if do_all:
+                    continue  # We remove all decorators with this name
+                # Check if the decorator specification matches
+                if deco.get_args_kwargs() == (deco_args, deco_kwargs):
+                    did_remove = True
+                    debug.userconf_exec(
+                        "Mutable step removing step decorator '%s' from step '%s'"
+                        % (deco.name, self._my_step.name)
+                    )
+                else:
+                    new_deco_list.append(deco)
+            else:
+                new_deco_list.append(deco)
+
+        self._my_step.decorators = new_deco_list
+
+        if did_remove:
+            return True
+
+        debug.userconf_exec(
+            "Mutable step did not find decorator '%s' to remove from step '%s'"
+            % (deco_name, self._my_step.name)
+        )
+        return False

metaflow/user_decorators/user_flow_decorator.py

@@ -0,0 +1,263 @@
+from typing import Dict, Optional, Union, TYPE_CHECKING
+
+from metaflow.exception import MetaflowException
+from metaflow.user_configs.config_parameters import (
+    resolve_delayed_evaluator,
+    unpack_delayed_evaluator,
+)
+
+from .common import ClassPath_Trie
+
+if TYPE_CHECKING:
+    import metaflow.flowspec
+    import metaflow.user_decorators.mutable_flow
+
+
+class FlowMutatorMeta(type):
+    _all_registered_decorators = ClassPath_Trie()
+    _do_not_register = set()
+    _import_modules = set()
+
+    def __new__(mcs, name, bases, namespace):
+        cls = super().__new__(mcs, name, bases, namespace)
+        cls.decorator_name = getattr(
+            cls, "_decorator_name", f"{cls.__module__}.{cls.__name__}"
+        )
+        if not cls.__module__.startswith("metaflow.") and not cls.__module__.startswith(
+            "metaflow_extensions."
+        ):
+            mcs._import_modules.add(cls.__module__)
+
+        if name == "FlowMutator" or cls.decorator_name in mcs._do_not_register:
+            return cls
+
+        # We inject a __init_subclass__ method so we can figure out if there
+        # are subclasses. We want to register as decorators only the ones that do
+        # not have a subclass. The logic is that everything is registered and if
+        # a subclass shows up, we will unregister the parent class leaving only those
+        # classes that do not have any subclasses registered.
+        @classmethod
+        def do_unregister(cls_, **_kwargs):
+            for base in cls_.__bases__:
+                if isinstance(base, FlowMutatorMeta):
+                    # If the base is a FlowMutatorMeta, we unregister it
+                    # so that we don't have any decorators that are not the
+                    # most derived one.
+                    mcs._all_registered_decorators.remove(base.decorator_name)
+                    # Also make sure we don't register again
+                    mcs._do_not_register.add(base.decorator_name)
+
+        cls.__init_subclass__ = do_unregister
+        mcs._all_registered_decorators.insert(cls.decorator_name, cls)
+        return cls
+
+    @classmethod
+    def all_decorators(mcs) -> Dict[str, "FlowMutatorMeta"]:
+        mcs._check_init()
+        return mcs._all_registered_decorators.get_unique_prefixes()
+
+    def __str__(cls):
+        return "FlowMutator(%s)" % cls.decorator_name
+
+    @classmethod
+    def get_decorator_by_name(
+        mcs, decorator_name: str
+    ) -> Optional[Union["FlowDecoratorMeta", "metaflow.decorators.Decorator"]]:
+        """
+        Get a decorator by its name.
+
+        Parameters
+        ----------
+        decorator_name: str
+            The name of the decorator to retrieve.
+
+        Returns
+        -------
+        Optional[FlowDecoratorMeta]
+            The decorator class if found, None otherwise.
+        """
+        mcs._check_init()
+        return mcs._all_registered_decorators.unique_prefix_value(decorator_name)
+
+    @classmethod
+    def get_decorator_name(mcs, decorator_type: type) -> Optional[str]:
+        """
+        Get the minimally unique classpath name for a decorator type.
+
+        Parameters
+        ----------
+        decorator_type: type
+            The type of the decorator to retrieve the name for.
+
+        Returns
+        -------
+        Optional[str]
+            The minimally unique classpath name if found, None otherwise.
+        """
+        mcs._check_init()
+        return mcs._all_registered_decorators.unique_prefix_for_type(decorator_type)
+
+    @classmethod
+    def _check_init(mcs):
+        # Delay importing STEP_DECORATORS until we actually need it
+        if not mcs._all_registered_decorators.inited:
+            from metaflow.plugins import FLOW_DECORATORS
+
+            mcs._all_registered_decorators.init([(t.name, t) for t in FLOW_DECORATORS])
+
+
+class FlowMutator(metaclass=FlowMutatorMeta):
+    """
+    Derive from this class to implement a flow mutator.
+
+    A flow mutator allows you to introspect a flow and its included steps. You can
+    then add parameters, configurations and decorators to the flow as well as modify
+    any of its steps.
+    use values available through configurations to determine how to mutate the flow.
+
+    There are two main methods provided:
+      - pre_mutate: called as early as possible right after configuration values are read.
+      - mutate: called right after all the command line is parsed but before any
+        Metaflow decorators are applied.
+
+    The `mutate` method does not allow you to modify the flow itself but you can still
+    modify the steps.
+    """
+
+    def __init__(self, *args, **kwargs):
+        from ..flowspec import FlowSpecMeta
+
+        self._flow_cls = None
+        # Flow mutators are always statically defined (no way of passing them
+        # on the command line or adding them via configs)
+        self.statically_defined = True
+        self.inserted_by = None
+
+        # The arguments are actually passed to the init function for this decorator
+        # and used in _graph_info
+        self._args = args
+        self._kwargs = kwargs
+        if args and isinstance(args[0], (FlowMutator, FlowSpecMeta)):
+            # This means the decorator is bare like @MyDecorator
+            # and the first argument is the FlowSpec or another decorator (they
+            # can be stacked)
+
+            # Now set the flow class we apply to
+            if isinstance(args[0], FlowSpecMeta):
+                self._set_flow_cls(args[0])
+            else:
+                self._set_flow_cls(args[0]._flow_cls)
+            self._args = self._args[1:]  # Remove the first argument
+
+    def __mro_entries__(self, bases):
+        # This is called in the following case:
+        # @MyMutator
+        # class MyBaseFlowSpec(FlowSpec):
+        #   pass
+        #
+        # class MyFlow(MyBaseFlowSpec):
+        #   pass
+        #
+        # MyBaseFlowSpec will be an object of type MyMutator which is not
+        # great when inheriting from it. With this method, we ensure that the
+        # base class will actually be MyBaseFlowSpec
+        return (self._flow_cls,)
+
+    def __call__(
+        self, flow_spec: Optional["metaflow.flowspec.FlowSpecMeta"] = None
+    ) -> "metaflow.flowspec.FlowSpecMeta":
+        if flow_spec:
+            if isinstance(flow_spec, FlowMutator):
+                flow_spec = flow_spec._flow_cls
+            return self._set_flow_cls(flow_spec)
+        elif not self._flow_cls:
+            # This means that somehow the initialization did not happen properly
+            # so this may have been applied to a non flow
+            raise MetaflowException("A FlowMutator can only be applied to a FlowSpec")
+        # NOTA: This returns self._flow_cls() because the object in the case of
+        # @FlowDecorator
+        # class MyFlow(FlowSpec):
+        #     pass
+        # the object is a FlowDecorator and when the main function calls it, we end up
+        # here and need to actually call the FlowSpec. This is not the case when using
+        # a decorator with arguments because in the line above, we will have returned a
+        # FlowSpec object. Previous solution was to use __get__ but this does not seem
+        # to work properly.
+        return self._flow_cls()
+
+    def _set_flow_cls(
+        self, flow_spec: "metaflow.flowspec.FlowSpecMeta"
+    ) -> "metaflow.flowspec.FlowSpecMeta":
+        from ..flowspec import _FlowState
+
+        flow_spec._flow_state.setdefault(_FlowState.CONFIG_DECORATORS, []).append(self)
+        self._flow_cls = flow_spec
+        return flow_spec
+
+    def __str__(self):
+        return str(self.__class__)
+
+    def init(self, *args, **kwargs):
+        """
+        Implement this method if you wish for your FlowMutator to take in arguments.
+
+        Your flow-mutator can then look like:
+
+        @MyMutator(arg1, arg2)
+        class MyFlow(FlowSpec):
+            pass
+
+        It is an error to use your mutator with arguments but not implement this method.
+
+        """
+        pass
+
+    def external_init(self):
+        # You can use config values in the arguments to a FlowMutator
+        # so we resolve those as well
+        self._args = [resolve_delayed_evaluator(arg) for arg in self._args]
+        self._kwargs, _ = unpack_delayed_evaluator(self._kwargs)
+        self._kwargs = {
+            k: resolve_delayed_evaluator(v) for k, v in self._kwargs.items()
+        }
+        if self._args or self._kwargs:
+            if "init" not in self.__class__.__dict__:
+                raise MetaflowException(
+                    "%s is used with arguments but does not implement init" % self
+                )
+            self.init(*self._args, **self._kwargs)
+
+    def pre_mutate(
+        self, mutable_flow: "metaflow.user_decorators.mutable_flow.MutableFlow"
+    ) -> None:
+        """
+        Method called right after all configuration values are read.
+
+        Parameters
+        ----------
+        mutable_flow : metaflow.user_decorators.mutable_flow.MutableFlow
+            A representation of this flow
+        """
+        return None
+
+    def mutate(
+        self, mutable_flow: "metaflow.user_decorators.mutable_flow.MutableFlow"
+    ) -> None:
+        """
+        Method called right before the first Metaflow step decorator is applied. This
+        means that the command line, including all `--with` options has been parsed.
+
+        Given how late this function is called, there are a few restrictions on what
+        you can do; the following methods on MutableFlow are not allowed and calling
+        them will result in an error:
+          - add_parameter/remove_parameter
+          - add_decorator/remove_decorator
+
+        To call these methods, use the `pre_mutate` method instead.
+
+        Parameters
+        ----------
+        mutable_flow : metaflow.user_decorators.mutable_flow.MutableFlow
+            A representation of this flow
+        """
+        return None