ob-metaflow 2.15.18.1__py2.py3-none-any.whl → 2.16.0.1__py2.py3-none-any.whl

metaflow/user_decorators/common.py

@@ -0,0 +1,144 @@
+from typing import Dict, Optional, List, Tuple
+
+
+class _TrieNode:
+    def __init__(
+        self, parent: Optional["_TrieNode"] = None, component: Optional[str] = None
+    ):
+        self.parent = parent
+        self.component = component
+        self.children = {}  # type: Dict[str, "_TrieNode"]
+        self.total_children = 0
+        self.value = None
+        self.end_value = None
+
+    def traverse(self, value: type) -> Optional["_TrieNode"]:
+        if self.total_children == 0:
+            self.end_value = value
+        else:
+            self.end_value = None
+        self.total_children += 1
+
+    def remove_child(self, child_name: str) -> bool:
+        if child_name in self.children:
+            del self.children[child_name]
+            self.total_children -= 1
+            return True
+        return False
+
+
+class ClassPath_Trie:
+    def __init__(self):
+        self.root = _TrieNode(None, None)
+        self.inited = False
+        self._value_to_node = {}  # type: Dict[type, _TrieNode]
+
+    def init(self, initial_nodes: Optional[List[Tuple[str, type]]] = None):
+        # We need to do this so we can delay import of STEP_DECORATORS
+        self.inited = True
+        for classpath_name, value in initial_nodes or []:
+            self.insert(classpath_name, value)
+
+    def insert(self, classpath_name: str, value: type):
+        node = self.root
+        components = reversed(classpath_name.split("."))
+        for c in components:
+            node = node.children.setdefault(c, _TrieNode(node, c))
+            node.traverse(value)
+        node.total_children -= (
+            1  # We do not count the last node as having itself as a child
+        )
+        node.value = value
+        self._value_to_node[value] = node
+
+    def search(self, classpath_name: str) -> Optional[type]:
+        node = self.root
+        components = reversed(classpath_name.split("."))
+        for c in components:
+            if c not in node.children:
+                return None
+            node = node.children[c]
+        return node.value
+
+    def remove(self, classpath_name: str):
+        components = list(reversed(classpath_name.split(".")))
+
+        def _remove(node: _TrieNode, components, depth):
+            if depth == len(components):
+                if node.value is not None:
+                    del self._value_to_node[node.value]
+                    node.value = None
+                    return len(node.children) == 0
+                return False
+            c = components[depth]
+            if c not in node.children:
+                return False
+            did_delete_child = _remove(node.children[c], components, depth + 1)
+            if did_delete_child:
+                node.remove_child(c)
+                if node.total_children == 1:
+                    # If we have one total child left, we have at least one
+                    # child and that one has an end_value
+                    for child in node.children.values():
+                        assert (
+                            child.end_value
+                        ), "Node with one child must have an end_value"
+                        node.end_value = child.end_value
+                return node.total_children == 0
+            return False
+
+        _remove(self.root, components, 0)
+
+    def unique_prefix_value(self, classpath_name: str) -> Optional[type]:
+        node = self.root
+        components = reversed(classpath_name.split("."))
+        for c in components:
+            if c not in node.children:
+                return None
+            node = node.children[c]
+        # If we reach here, it means the classpath_name is a prefix.
+        # We check if it has only one path forward (end_value will be non None)
+        # If value is not None, we also consider this to be a unique "prefix"
+        # This happens since this trie is also filled with metaflow default decorators
+        return node.end_value or node.value
+
+    def unique_prefix_for_type(self, value: type) -> Optional[str]:
+        node = self._value_to_node.get(value, None)
+        if node is None:
+            return None
+        components = []
+        while node:
+            if node.end_value == value:
+                components = []
+            if node.component is not None:
+                components.append(node.component)
+            node = node.parent
+        return ".".join(components)
+
+    def get_unique_prefixes(self) -> Dict[str, type]:
+        """
+        Get all unique prefixes in the trie.
+
+        Returns
+        -------
+        List[str]
+            A list of unique prefixes.
+        """
+        to_return = {}
+
+        def _collect(node, current_prefix):
+            if node.end_value is not None:
+                to_return[current_prefix] = node.end_value
+                # We stop there and don't look further since we found the unique prefix
+                return
+            if node.value is not None:
+                to_return[current_prefix] = node.value
+                # We continue to look for more unique prefixes
+            for child_name, child_node in node.children.items():
+                _collect(
+                    child_node,
+                    f"{current_prefix}.{child_name}" if current_prefix else child_name,
+                )
+
+        _collect(self.root, "")
+        return {".".join(reversed(k.split("."))): v for k, v in to_return.items()}

metaflow/user_decorators/mutable_flow.py

@@ -0,0 +1,499 @@
+from functools import partial
+from typing import Any, Dict, Generator, List, Optional, Tuple, TYPE_CHECKING, Union
+
+from metaflow.debug import debug
+from metaflow.exception import MetaflowException
+from metaflow.user_configs.config_parameters import ConfigValue
+
+if TYPE_CHECKING:
+    import metaflow.flowspec
+    import metaflow.parameters
+    import metaflow.user_decorators.mutable_step
+
+
+class MutableFlow:
+    IGNORE = 1
+    ERROR = 2
+    OVERRIDE = 3
+
+    def __init__(
+        self,
+        flow_spec: "metaflow.flowspec.FlowSpec",
+        pre_mutate: bool = False,
+        statically_defined: bool = False,
+        inserted_by: Optional[str] = None,
+    ):
+        self._flow_cls = flow_spec
+        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(
+                "MutableFlow should only be created by StepMutators or FlowMutators. "
+                "This is an internal error."
+            )
+
+    @property
+    def decorator_specs(
+        self,
+    ) -> Generator[Tuple[str, str, List[Any], Dict[str, Any]], None, None]:
+        """
+        Iterate over all the decorator specifications of this flow. 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 decorator specification to remove a decorator from the flow
+        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 decos in self._flow_cls._flow_decorators.values():
+            for deco in decos:
+                # 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)
+
+    @property
+    def configs(self) -> Generator[Tuple[str, ConfigValue], None, None]:
+        """
+        Iterate over all user configurations in this flow
+
+        Use this to parameterize your flow based on configuration. As an example, the
+        `pre_mutate`/`mutate` methods can add decorators to steps in the flow that
+        depend on values in the configuration.
+
+        ```
+        class MyDecorator(FlowMutator):
+            def mutate(flow: MutableFlow):
+                val = next(flow.configs)[1].steps.start.cpu
+                flow.start.add_decorator(environment, vars={'mycpu': val})
+                return flow
+
+        @MyDecorator()
+        class TestFlow(FlowSpec):
+            config = Config('myconfig.json')
+
+            @step
+            def start(self):
+                pass
+        ```
+        can be used to add an environment decorator to the `start` step.
+
+        Yields
+        ------
+        Tuple[str, ConfigValue]
+            Iterates over the configurations of the flow
+        """
+        from metaflow.flowspec import _FlowState
+
+        # When configs are parsed, they are loaded in _flow_state[_FlowState.CONFIGS]
+        for name, value in self._flow_cls._flow_state.get(
+            _FlowState.CONFIGS, {}
+        ).items():
+            r = name, ConfigValue(value)
+            debug.userconf_exec("Mutable flow yielding config: %s" % str(r))
+            yield r
+
+    @property
+    def parameters(
+        self,
+    ) -> Generator[Tuple[str, "metaflow.parameters.Parameter"], None, None]:
+        """
+        Iterate over all the parameters in this flow.
+
+        Yields
+        ------
+        Tuple[str, Parameter]
+            Name of the parameter and parameter in the flow
+        """
+        for var, param in self._flow_cls._get_parameters():
+            if param.IS_CONFIG_PARAMETER:
+                continue
+            debug.userconf_exec(
+                "Mutable flow yielding parameter: %s" % str((var, param))
+            )
+            yield var, param
+
+    @property
+    def steps(
+        self,
+    ) -> Generator[
+        Tuple[str, "metaflow.user_decorators.mutable_step.MutableStep"], None, None
+    ]:
+        """
+        Iterate over all the steps in this flow. The order of the steps
+        returned is not guaranteed.
+
+        Yields
+        ------
+        Tuple[str, MutableStep]
+            A tuple with the step name and the step proxy
+        """
+        from .mutable_step import MutableStep
+
+        for var in dir(self._flow_cls):
+            potential_step = getattr(self._flow_cls, var)
+            if callable(potential_step) and hasattr(potential_step, "is_step"):
+                debug.userconf_exec("Mutable flow yielding step: %s" % var)
+                yield var, MutableStep(
+                    self._flow_cls,
+                    potential_step,
+                    pre_mutate=self._pre_mutate,
+                    statically_defined=self._statically_defined,
+                    inserted_by=self._inserted_by,
+                )
+
+    def add_parameter(
+        self, name: str, value: "metaflow.parameters.Parameter", overwrite: bool = False
+    ) -> None:
+        """
+        Add a parameter to the flow. You can only add parameters in the `pre_mutate`
+        method.
+
+        Parameters
+        ----------
+        name : str
+            Name of the parameter
+        value : Parameter
+            Parameter to add to the flow
+        overwrite : bool, default False
+            If True, overwrite the parameter if it already exists
+        """
+        if not self._pre_mutate:
+            raise MetaflowException(
+                "Adding parameter '%s' from %s is only allowed in the `pre_mutate` "
+                "method and not the `mutate` method" % (name, self._inserted_by)
+            )
+        from metaflow.parameters import Parameter
+
+        if hasattr(self._flow_cls, name) and not overwrite:
+            raise MetaflowException(
+                "Flow '%s' already has a class member '%s' -- "
+                "set overwrite=True in add_parameter to overwrite it."
+                % (self._flow_cls.__name__, name)
+            )
+        if not isinstance(value, Parameter) or value.IS_CONFIG_PARAMETER:
+            raise MetaflowException(
+                "Only a Parameter or an IncludeFile can be added using `add_parameter`"
+                "; got %s" % type(value)
+            )
+        debug.userconf_exec("Mutable flow adding parameter %s to flow" % name)
+        setattr(self._flow_cls, name, value)
+
+    def remove_parameter(self, parameter_name: str) -> bool:
+        """
+        Remove a parameter from the flow.
+
+        The name given should match the name of the parameter (can be different
+        from the name of the parameter in the flow. You can not remove config parameters.
+        You can only remove parameters in the `pre_mutate` method.
+
+        Parameters
+        ----------
+        parameter_name : str
+            Name of the parameter
+
+        Returns
+        -------
+        bool
+            Returns True if the parameter was removed
+        """
+        if not self._pre_mutate:
+            raise MetaflowException(
+                "Removing parameter '%s' from %s is only allowed in the `pre_mutate` "
+                "method and not the `mutate` method"
+                % (parameter_name, " from ".join(self._inserted_by))
+            )
+        from metaflow.flowspec import _FlowState
+
+        for var, param in self._flow_cls._get_parameters():
+            if param.IS_CONFIG_PARAMETER:
+                continue
+            if param.name == parameter_name:
+                delattr(self._flow_cls, var)
+                debug.userconf_exec(
+                    "Mutable flow removing parameter %s from flow" % var
+                )
+                # Reset so that we don't list it again
+                del self._flow_cls._flow_state[_FlowState.CACHED_PARAMETERS]
+                return True
+        debug.userconf_exec(
+            "Mutable flow failed to remove parameter %s from flow" % parameter_name
+        )
+        return False
+
+    def add_decorator(
+        self,
+        deco_type: Union[partial, str],
+        deco_args: Optional[List[Any]] = None,
+        deco_kwargs: Optional[Dict[str, Any]] = None,
+        duplicates: int = IGNORE,
+    ) -> None:
+        """
+        Add a Metaflow flow-decorator to a flow. You can only add decorators in the
+        `pre_mutate` method.
+
+        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 project
+
+        ...
+        my_flow.add_decorator(project, deco_kwargs={"name":"my_project"})
+        ```
+
+        is equivalent to:
+        ```
+        my_flow.add_decorator("project:name=my_project")
+        ```
+
+        Note in the later case, there is no need to import the flow decorator.
+
+        The latter syntax is useful to, for example, allow decorators to be stored as
+        strings in a configuration file.
+
+        In terms of precedence for decorators:
+          - if a decorator can be applied multiple times all decorators
+            added are kept (this is rare for flow-decorators).
+          - if `duplicates` is set to `MutableFlow.IGNORE`, then the decorator
+            being added is ignored (in other words, the existing decorator has precedence).
+          - if `duplicates` is set to `MutableFlow.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 `MutableFlow.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, str]
+            The decorator class to add to this flow. If using a string, you cannot specify
+            additional arguments as all argument will be parsed from the decorator
+            specification.
+        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 MutableFlow.IGNORE
+            Instruction on how to handle duplicates. It can be one of:
+            - `MutableFlow.IGNORE`: Ignore the decorator if it already exists.
+            - `MutableFlow.ERROR`: Raise an error if the decorator already exists.
+            - `MutableFlow.OVERRIDE`: Remove the existing decorator and add this one.
+
+        """
+        if not self._pre_mutate:
+            raise MetaflowException(
+                "Adding flow-decorator '%s' from %s is only allowed in the `pre_mutate` "
+                "method and not the `mutate` method"
+                % (
+                    deco_type if isinstance(deco_type, str) else deco_type.name,
+                    self._inserted_by,
+                )
+            )
+        # Prevent circular import
+        from metaflow.decorators import (
+            DuplicateFlowDecoratorException,
+            FlowDecorator,
+            extract_flow_decorator_from_decospec,
+        )
+
+        deco_args = deco_args or []
+        deco_kwargs = deco_kwargs or {}
+
+        def _add_flow_decorator(flow_deco):
+            if deco_args:
+                raise MetaflowException(
+                    "Flow decorators do not take additional positional arguments"
+                )
+            # Update kwargs:
+            flow_deco.attributes.update(deco_kwargs)
+
+            # Check duplicates
+            def _do_add():
+                flow_deco.statically_defined = self._statically_defined
+                flow_deco.inserted_by = self._inserted_by
+                self._flow_cls._flow_decorators.setdefault(flow_deco.name, []).append(
+                    flow_deco
+                )
+                debug.userconf_exec(
+                    "Mutable flow adding flow decorator '%s'" % deco_type
+                )
+
+            existing_deco = [
+                d for d in self._flow_cls._flow_decorators if d.name == flow_deco.name
+            ]
+
+            if flow_deco.allow_multiple or not existing_deco:
+                _do_add()
+            elif duplicates == MutableFlow.IGNORE:
+                # If we ignore, we do not add the decorator
+                debug.userconf_exec(
+                    "Mutable flow ignoring flow decorator '%s'"
+                    "(already exists and duplicates are ignored)" % flow_deco.name
+                )
+            elif duplicates == MutableFlow.OVERRIDE:
+                # If we override, we remove the existing decorator and add this one
+                debug.userconf_exec(
+                    "Mutable flow overriding flow decorator '%s' "
+                    "(removing existing decorator and adding new one)" % flow_deco.name
+                )
+                self._flow_cls._flow_decorators = [
+                    d
+                    for d in self._flow_cls._flow_decorators
+                    if d.name != flow_deco.name
+                ]
+                _do_add()
+            elif duplicates == MutableFlow.ERROR:
+                # If we error, we raise an exception
+                if self._statically_defined:
+                    raise DuplicateFlowDecoratorException(flow_deco.name)
+                else:
+                    debug.userconf_exec(
+                        "Mutable flow ignoring flow decorator '%s' "
+                        "(already exists and non statically defined)" % flow_deco.name
+                    )
+            else:
+                raise ValueError("Invalid duplicates value: %s" % duplicates)
+
+        # If deco_type is a string, we want to parse it to a decospec
+        if isinstance(deco_type, str):
+            flow_deco, has_args_kwargs = extract_flow_decorator_from_decospec(deco_type)
+            if (deco_args or deco_kwargs) and has_args_kwargs:
+                raise MetaflowException(
+                    "Cannot specify additional arguments when adding a flow decorator "
+                    "using a decospec that already contains arguments"
+                )
+            _add_flow_decorator(flow_deco)
+            return
+
+        # Validate deco_type
+        if (
+            not isinstance(deco_type, partial)
+            or len(deco_type.args) != 1
+            or not issubclass(deco_type.args[0], FlowDecorator)
+        ):
+            raise TypeError("add_decorator takes a FlowDecorator")
+
+        deco_type = deco_type.args[0]
+        _add_flow_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 flow-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 decorators 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.
+        """
+
+        if not self._pre_mutate:
+            raise MetaflowException(
+                "Removing flow-decorator '%s' from %s is only allowed in the `pre_mutate` "
+                "method and not the `mutate` method" % (deco_name, self._inserted_by)
+            )
+
+        do_all = deco_args is None and deco_kwargs is None
+        did_remove = False
+        if do_all and deco_name in self._flow_cls._flow_decorators:
+            del self._flow_cls._flow_decorators[deco_name]
+            return True
+        old_deco_list = self._flow_cls._flow_decorators.get(deco_name)
+        if not old_deco_list:
+            debug.userconf_exec(
+                "Mutable flow failed to remove decorator '%s' from flow (non present)"
+                % deco_name
+            )
+            return False
+        new_deco_list = []
+        for deco in old_deco_list:
+            if deco.get_args_kwargs() == (deco_args or [], deco_kwargs or {}):
+                did_remove = True
+            else:
+                new_deco_list.append(deco)
+        debug.userconf_exec(
+            "Mutable flow removed %d decorators from flow"
+            % (len(old_deco_list) - len(new_deco_list))
+        )
+        if new_deco_list:
+            self._flow_cls._flow_decorators[deconame] = new_deco_list
+        else:
+            del self._flow_cls._flow_decorators[deco_name]
+        return did_remove
+
+    def __getattr__(self, name):
+        # We allow direct access to the steps, configs and parameters but nothing else
+        from metaflow.parameters import Parameter
+
+        from .mutable_step import MutableStep
+
+        attr = getattr(self._flow_cls, name)
+        if attr:
+            # Steps
+            if callable(attr) and hasattr(attr, "is_step"):
+                return MutableStep(
+                    self._flow_cls,
+                    attr,
+                    pre_mutate=self._pre_mutate,
+                    statically_defined=self._statically_defined,
+                    inserted_by=self._inserted_by,
+                )
+            if name[0] == "_" or name in self._flow_cls._NON_PARAMETERS:
+                raise AttributeError(self, name)
+            if isinstance(attr, (Parameter, ConfigValue)):
+                return attr
+        raise AttributeError(self, name)