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

metaflow/task.py

@@ -6,6 +6,7 @@ import os
 import time
 import traceback
 
+
 from types import MethodType, FunctionType
 
 from metaflow.sidecar import Message, MessageTypes
@@ -24,6 +25,7 @@ from .unbounded_foreach import UBF_CONTROL
 from .util import all_equal, get_username, resolve_identity, unicode_type
 from .clone_util import clone_task_helper
 from .metaflow_current import current
+from metaflow.user_configs.config_parameters import ConfigValue
 from metaflow.system import _system_logger, _system_monitor
 from metaflow.tracing import get_trace_id
 from metaflow.tuple_util import ForeachFrame
@@ -57,11 +59,89 @@ class MetaflowTask(object):
         self.monitor = monitor
         self.ubf_context = ubf_context
 
-    def _exec_step_function(self, step_function, input_obj=None):
-        if input_obj is None:
-            step_function()
-        else:
-            step_function(input_obj)
+    def _exec_step_function(self, step_function, orig_step_func, input_obj=None):
+        wrappers_stack = []
+        wrapped_func = None
+        do_next = False
+        raised_exception = None
+        # If we have wrappers w1, w2 and w3, we need to execute
+        #  - w3_pre
+        #  - w2_pre
+        #  - w1_pre
+        #  - step_function
+        #  - w1_post
+        #  - w2_post
+        #  - w3_post
+        # in that order. We do this by maintaining a stack of generators.
+        # Note that if any of the pre functions returns a function, we execute that
+        # instead of the rest of the inside part. This is useful if you want to create
+        # no-op function for example.
+        for w in reversed(orig_step_func.wrappers):
+            wrapped_func = w.pre_step(orig_step_func.name, self.flow, input_obj)
+            wrappers_stack.append(w)
+            if w.skip_step:
+                # We have nothing to run
+                do_next = w.skip_step
+                break
+            if wrapped_func:
+                break  # We have nothing left to do since we now execute the
+                # wrapped function
+            # Else, we continue down the list of wrappers
+        try:
+            if not do_next:
+                if input_obj is None:
+                    if wrapped_func:
+                        do_next = wrapped_func(self.flow)
+                        if not do_next:
+                            do_next = True
+                    else:
+                        step_function()
+                else:
+                    if wrapped_func:
+                        do_next = wrapped_func(self.flow, input_obj)
+                        if not do_next:
+                            do_next = True
+                    else:
+                        step_function(input_obj)
+        except Exception as ex:
+            raised_exception = ex
+
+        if do_next:
+            # If we are skipping the step, or executed a wrapped function,
+            # we need to set the transition variables
+            # properly. We call the next function as needed
+            graph_node = self.flow._graph[step_function.name]
+            out_funcs = [getattr(self.flow, f) for f in graph_node.out_funcs]
+            if out_funcs:
+                if isinstance(do_next, bool):
+                    # We need to extract things from the self.next. This is not possible
+                    # in the case where there was a num_parallel.
+                    if graph_node.parallel_foreach:
+                        raise RuntimeError(
+                            "Skipping a parallel foreach step without providing "
+                            "the arguments to the self.next call is not supported. "
+                        )
+                    if graph_node.foreach_param:
+                        self.flow.next(*out_funcs, foreach=graph_node.foreach_param)
+                    else:
+                        self.flow.next(*out_funcs)
+                elif isinstance(do_next, dict):
+                    # Here it is a dictionary so we just call the next method with
+                    # those arguments
+                    self.flow.next(*out_funcs, **do_next)
+                else:
+                    raise RuntimeError(
+                        "Invalid value passed to self.next; expected "
+                        " bool of a dictionary; got: %s" % do_next
+                    )
+        # We back out of the stack of generators
+        for w in reversed(wrappers_stack):
+            raised_exception = w.post_step(
+                step_function.name, self.flow, raised_exception
+            )
+        if raised_exception:
+            # We have an exception that we need to propagate
+            raise raised_exception
 
     def _init_parameters(self, parameter_ds, passdown=True):
         cls = self.flow.__class__
@@ -538,6 +618,9 @@ class MetaflowTask(object):
         output.save_metadata(
             {
                 "task_begin": {
+                    "code_package_metadata": os.environ.get(
+                        "METAFLOW_CODE_METADATA", ""
+                    ),
                     "code_package_sha": os.environ.get("METAFLOW_CODE_SHA"),
                     "code_package_ds": os.environ.get("METAFLOW_CODE_DS"),
                     "code_package_url": os.environ.get("METAFLOW_CODE_URL"),
@@ -651,6 +734,7 @@ class MetaflowTask(object):
                         inputs,
                     )
 
+                orig_step_func = step_func
                 for deco in decorators:
                     # decorators can actually decorate the step function,
                     # or they can replace it altogether. This functionality
@@ -667,9 +751,9 @@ class MetaflowTask(object):
                     )
 
                 if join_type:
-                    self._exec_step_function(step_func, input_obj)
+                    self._exec_step_function(step_func, orig_step_func, input_obj)
                 else:
-                    self._exec_step_function(step_func)
+                    self._exec_step_function(step_func, orig_step_func)
 
                 for deco in decorators:
                     deco.task_post_step(

metaflow/user_configs/config_options.py

@@ -7,8 +7,9 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from metaflow._vendor import click
 from metaflow.debug import debug
 
-from .config_parameters import CONFIG_FILE, ConfigValue
+from .config_parameters import ConfigValue
 from ..exception import MetaflowException, MetaflowInternalError
+from ..packaging_sys import MetaflowCodeContent
 from ..parameters import DeployTimeField, ParameterContext, current_flow
 from ..util import get_username
 
@@ -24,12 +25,16 @@ _CONVERTED_DEFAULT_NO_FILE = _CONVERTED_DEFAULT + _NO_FILE
 
 def _load_config_values(info_file: Optional[str] = None) -> Optional[Dict[Any, Any]]:
     if info_file is None:
-        info_file = os.path.basename(CONFIG_FILE)
-    try:
-        with open(info_file, encoding="utf-8") as contents:
-            return json.load(contents).get("user_configs", {})
-    except IOError:
-        return None
+        config_content = MetaflowCodeContent.get_config()
+    else:
+        try:
+            with open(info_file, encoding="utf-8") as f:
+                config_content = json.load(f)
+        except IOError:
+            return None
+    if config_content:
+        return config_content.get("user_configs", {})
+    return None
 
 
 class ConvertPath(click.Path):
@@ -437,7 +442,7 @@ class LocalFileInput(click.Path):
     # Small wrapper around click.Path to set the value from which to read configuration
     # values. This is set immediately upon processing the --local-config-file
     # option and will therefore then be available when processing any of the other
-    # --config options (which will call ConfigInput.process_configs
+    # --config options (which will call ConfigInput.process_configs)
     name = "LocalFileInput"
 
     def convert(self, value, param, ctx):

metaflow/user_configs/config_parameters.py

@@ -49,10 +49,6 @@ if TYPE_CHECKING:
 
 #     return tracefunc_closure
 
-CONFIG_FILE = os.path.join(
-    os.path.dirname(os.path.abspath(__file__)), "CONFIG_PARAMETERS"
-)
-
 ID_PATTERN = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
 
 UNPACK_KEY = "_unpacked_delayed_"

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()}