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

metaflow/graph.py

@@ -2,6 +2,8 @@ import inspect
 import ast
 import re
 
+from itertools import chain
+
 
 from .util import to_pod
 
@@ -45,13 +47,17 @@ def deindent_docstring(doc):
 
 
 class DAGNode(object):
-    def __init__(self, func_ast, decos, doc, source_file, lineno):
+    def __init__(
+        self, func_ast, decos, wrappers, config_decorators, doc, source_file, lineno
+    ):
         self.name = func_ast.name
         self.source_file = source_file
         # lineno is the start line of decorators in source_file
         # func_ast.lineno is lines from decorators start to def of function
         self.func_lineno = lineno + func_ast.lineno - 1
         self.decorators = decos
+        self.wrappers = wrappers
+        self.config_decorators = config_decorators
         self.doc = deindent_docstring(doc)
         self.parallel_step = any(getattr(deco, "IS_PARALLEL", False) for deco in decos)
 
@@ -62,6 +68,8 @@ class DAGNode(object):
         self.has_tail_next = False
         self.invalid_tail_next = False
         self.num_args = 0
+        self.switch_cases = {}
+        self.condition = None
         self.foreach_param = None
         self.num_parallel = 0
         self.parallel_foreach = False
@@ -70,6 +78,7 @@ class DAGNode(object):
         # these attributes are populated by _traverse_graph
         self.in_funcs = set()
         self.split_parents = []
+        self.split_branches = []
         self.matching_join = None
         # these attributes are populated by _postprocess
         self.is_inside_foreach = False
@@ -77,6 +86,56 @@ class DAGNode(object):
     def _expr_str(self, expr):
         return "%s.%s" % (expr.value.id, expr.attr)
 
+    def _parse_switch_dict(self, dict_node):
+        switch_cases = {}
+
+        if isinstance(dict_node, ast.Dict):
+            for key, value in zip(dict_node.keys, dict_node.values):
+                case_key = None
+
+                # handle string literals
+                if hasattr(ast, "Str") and isinstance(key, ast.Str):
+                    case_key = key.s
+                elif isinstance(key, ast.Constant):
+                    case_key = key.value
+                elif isinstance(key, ast.Attribute):
+                    if isinstance(key.value, ast.Attribute) and isinstance(
+                        key.value.value, ast.Name
+                    ):
+                        # This handles self.config.some_key
+                        if key.value.value.id == "self":
+                            config_var = key.value.attr
+                            config_key = key.attr
+                            case_key = f"config:{config_var}.{config_key}"
+                        else:
+                            return None
+                    else:
+                        return None
+
+                # handle variables or other dynamic expressions - not allowed
+                elif isinstance(key, ast.Name):
+                    return None
+                else:
+                    # can't statically analyze this key
+                    return None
+
+                if case_key is None:
+                    return None
+
+                # extract the step name from the value
+                if isinstance(value, ast.Attribute) and isinstance(
+                    value.value, ast.Name
+                ):
+                    if value.value.id == "self":
+                        step_name = value.attr
+                        switch_cases[case_key] = step_name
+                    else:
+                        return None
+                else:
+                    return None
+
+        return switch_cases if switch_cases else None
+
     def _parse(self, func_ast, lineno):
         self.num_args = len(func_ast.args.args)
         tail = func_ast.body[-1]
@@ -98,7 +157,38 @@ class DAGNode(object):
             self.has_tail_next = True
             self.invalid_tail_next = True
             self.tail_next_lineno = lineno + tail.lineno - 1
-            self.out_funcs = [e.attr for e in tail.value.args]
+
+            # Check if first argument is a dictionary (switch case)
+            if (
+                len(tail.value.args) == 1
+                and isinstance(tail.value.args[0], ast.Dict)
+                and any(k.arg == "condition" for k in tail.value.keywords)
+            ):
+                # This is a switch statement
+                switch_cases = self._parse_switch_dict(tail.value.args[0])
+                condition_name = None
+
+                # Get condition parameter
+                for keyword in tail.value.keywords:
+                    if keyword.arg == "condition":
+                        if hasattr(ast, "Str") and isinstance(keyword.value, ast.Str):
+                            condition_name = keyword.value.s
+                        elif isinstance(keyword.value, ast.Constant) and isinstance(
+                            keyword.value.value, str
+                        ):
+                            condition_name = keyword.value.value
+                        break
+
+                if switch_cases and condition_name:
+                    self.type = "split-switch"
+                    self.condition = condition_name
+                    self.switch_cases = switch_cases
+                    self.out_funcs = list(switch_cases.values())
+                    self.invalid_tail_next = False
+                    return
+
+            else:
+                self.out_funcs = [e.attr for e in tail.value.args]
 
             keywords = dict(
                 (k.arg, getattr(k.value, "s", None)) for k in tail.value.keywords
@@ -138,6 +228,7 @@ class DAGNode(object):
     in_funcs={in_funcs}
     out_funcs={out_funcs}
     split_parents={parents}
+    split_branches={branches}
     matching_join={matching_join}
     is_inside_foreach={is_inside_foreach}
     decorators={decos}
@@ -145,6 +236,7 @@ class DAGNode(object):
     has_tail_next={0.has_tail_next} (line {0.tail_next_lineno})
     invalid_tail_next={0.invalid_tail_next}
     foreach_param={0.foreach_param}
+    condition={0.condition}
     parallel_step={0.parallel_step}
     parallel_foreach={0.parallel_foreach}
     -> {out}""".format(
@@ -154,6 +246,7 @@ class DAGNode(object):
             out_funcs=", ".join("[%s]" % x for x in self.out_funcs),
             in_funcs=", ".join("[%s]" % x for x in self.in_funcs),
             parents=", ".join("[%s]" % x for x in self.split_parents),
+            branches=", ".join("[%s]" % x for x in self.split_branches),
             decos=" | ".join(map(str, self.decorators)),
             out=", ".join("[%s]" % x for x in self.out_funcs),
         )
@@ -181,7 +274,13 @@ class FlowGraph(object):
                 source_code = deindent_docstring("".join(source_lines))
                 function_ast = ast.parse(source_code).body[0]
                 node = DAGNode(
-                    function_ast, func.decorators, func.__doc__, source_file, lineno
+                    function_ast,
+                    func.decorators,
+                    func.wrappers,
+                    func.config_decorators,
+                    func.__doc__,
+                    source_file,
+                    lineno,
                 )
                 nodes[element] = node
         return nodes
@@ -198,7 +297,8 @@ class FlowGraph(object):
                 node.is_inside_foreach = True
 
     def _traverse_graph(self):
-        def traverse(node, seen, split_parents):
+        def traverse(node, seen, split_parents, split_branches):
+            add_split_branch = False
             try:
                 self.sorted_nodes.remove(node.name)
             except ValueError:
@@ -206,15 +306,23 @@ class FlowGraph(object):
             self.sorted_nodes.append(node.name)
             if node.type in ("split", "foreach"):
                 node.split_parents = split_parents
+                node.split_branches = split_branches
+                add_split_branch = True
                 split_parents = split_parents + [node.name]
+            elif node.type == "split-switch":
+                node.split_parents = split_parents
+                node.split_branches = split_branches
             elif node.type == "join":
                 # ignore joins without splits
                 if split_parents:
                     self[split_parents[-1]].matching_join = node.name
                     node.split_parents = split_parents
+                    node.split_branches = split_branches[:-1]
                     split_parents = split_parents[:-1]
+                    split_branches = split_branches[:-1]
             else:
                 node.split_parents = split_parents
+                node.split_branches = split_branches
 
             for n in node.out_funcs:
                 # graph may contain loops - ignore them
@@ -223,10 +331,15 @@ class FlowGraph(object):
                     if n in self:
                         child = self[n]
                         child.in_funcs.add(node.name)
-                        traverse(child, seen + [n], split_parents)
+                        traverse(
+                            child,
+                            seen + [n],
+                            split_parents,
+                            split_branches + ([n] if add_split_branch else []),
+                        )
 
         if "start" in self:
-            traverse(self["start"], [], [])
+            traverse(self["start"], [], [], [])
 
         # fix the order of in_funcs
         for node in self.nodes.values():
@@ -247,15 +360,37 @@ class FlowGraph(object):
     def output_dot(self):
         def edge_specs():
             for node in self.nodes.values():
-                for edge in node.out_funcs:
-                    yield "%s -> %s;" % (node.name, edge)
+                if node.type == "split-switch":
+                    # Label edges for switch cases
+                    for case_value, step_name in node.switch_cases.items():
+                        yield (
+                            '{0} -> {1} [label="{2}" color="blue" fontcolor="blue"];'.format(
+                                node.name, step_name, case_value
+                            )
+                        )
+                else:
+                    for edge in node.out_funcs:
+                        yield "%s -> %s;" % (node.name, edge)
 
         def node_specs():
             for node in self.nodes.values():
-                nodetype = "join" if node.num_args > 1 else node.type
-                yield '"{0.name}"' '[ label = <<b>{0.name}</b> | <font point-size="10">{type}</font>> ' '  fontname = "Helvetica" ' '  shape = "record" ];'.format(
-                    node, type=nodetype
-                )
+                if node.type == "split-switch":
+                    # Hexagon shape for switch nodes
+                    condition_label = (
+                        f"switch: {node.condition}" if node.condition else "switch"
+                    )
+                    yield (
+                        '"{0.name}" '
+                        '[ label = <<b>{0.name}</b><br/><font point-size="9">{condition}</font>> '
+                        '  fontname = "Helvetica" '
+                        '  shape = "hexagon" '
+                        '  style = "filled" fillcolor = "lightgreen" ];'
+                    ).format(node, condition=condition_label)
+                else:
+                    nodetype = "join" if node.num_args > 1 else node.type
+                    yield '"{0.name}"' '[ label = <<b>{0.name}</b> | <font point-size="10">{type}</font>> ' '  fontname = "Helvetica" ' '  shape = "record" ];'.format(
+                        node, type=nodetype
+                    )
 
         return (
             "digraph {0.name} {{\n"
@@ -279,6 +414,8 @@ class FlowGraph(object):
                 if node.parallel_foreach:
                     return "split-parallel"
                 return "split-foreach"
+            elif node.type == "split-switch":
+                return "split-switch"
             return "unknown"  # Should never happen
 
         def node_to_dict(name, node):
@@ -293,9 +430,19 @@ class FlowGraph(object):
                         "name": deco.name,
                         "attributes": to_pod(deco.attributes),
                         "statically_defined": deco.statically_defined,
+                        "inserted_by": deco.inserted_by,
                     }
                     for deco in node.decorators
                     if not deco.name.startswith("_")
+                ]
+                + [
+                    {
+                        "name": deco.decorator_name,
+                        "attributes": {"_args": deco._args, **deco._kwargs},
+                        "statically_defined": deco.statically_defined,
+                        "inserted_by": deco.inserted_by,
+                    }
+                    for deco in chain(node.wrappers, node.config_decorators)
                 ],
                 "next": node.out_funcs,
             }
@@ -303,6 +450,9 @@ class FlowGraph(object):
                 d["foreach_artifact"] = node.foreach_param
             elif d["type"] == "split-parallel":
                 d["num_parallel"] = node.num_parallel
+            elif d["type"] == "split-switch":
+                d["condition"] = node.condition
+                d["switch_cases"] = node.switch_cases
             if node.matching_join:
                 d["matching_join"] = node.matching_join
             return d
@@ -317,8 +467,8 @@ class FlowGraph(object):
                 steps_info[cur_name] = node_dict
                 resulting_list.append(cur_name)
 
-                if cur_node.type not in ("start", "linear", "join"):
-                    # We need to look at the different branches for this
+                node_type = node_to_type(cur_node)
+                if node_type in ("split-static", "split-foreach"):
                     resulting_list.append(
                         [
                             populate_block(s, cur_node.matching_join)
@@ -326,8 +476,21 @@ class FlowGraph(object):
                         ]
                     )
                     cur_name = cur_node.matching_join
+                elif node_type == "split-switch":
+                    all_paths = [
+                        populate_block(s, end_name)
+                        for s in cur_node.out_funcs
+                        if s != cur_name
+                    ]
+                    resulting_list.append(all_paths)
+                    cur_name = end_name
                 else:
-                    cur_name = cur_node.out_funcs[0]
+                    # handles only linear, start, and join steps.
+                    if cur_node.out_funcs:
+                        cur_name = cur_node.out_funcs[0]
+                    else:
+                        # handles terminal nodes or when we jump to 'end_name'.
+                        break
             return resulting_list
 
         graph_structure = populate_block("start", "end")

metaflow/includefile.py

@@ -7,9 +7,10 @@ import json
 import os
 
 from hashlib import sha1
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Dict, Optional, Union
 
 from metaflow._vendor import click
+from metaflow._vendor import yaml
 
 from .exception import MetaflowException
 from .parameters import (
@@ -20,7 +21,7 @@ from .parameters import (
 )
 
 from .plugins import DATACLIENTS
-from .user_configs.config_parameters import ConfigValue
+from .user_configs.config_options import ConfigInput
 from .util import get_username
 
 import functools
@@ -261,6 +262,12 @@ class IncludeFile(Parameter):
     show_default : bool, default True
         If True, show the default value in the help text. A value of None is equivalent
         to True.
+    parser : Union[str, Callable[[str], Any]], optional, default None
+        If a callable, it is a function that can parse the file contents
+        into any desired format. If a string, the string should refer to
+        a function (like "my_parser_package.my_parser.my_parser_function") which should
+        be able to parse the file contents. If the name starts with a ".", it is assumed
+        to be relative to "metaflow".
     """
 
     def __init__(
@@ -270,6 +277,7 @@ class IncludeFile(Parameter):
         is_text: Optional[bool] = None,
         encoding: Optional[str] = None,
         help: Optional[str] = None,
+        parser: Optional[Union[str, Callable[[str], Any]]] = None,
         **kwargs: Dict[str, str]
     ):
         self._includefile_overrides = {}
@@ -277,6 +285,7 @@ class IncludeFile(Parameter):
             self._includefile_overrides["is_text"] = is_text
         if encoding is not None:
             self._includefile_overrides["encoding"] = encoding
+        self._parser = parser
         # NOTA: Right now, there is an issue where these can't be overridden by config
         # in all circumstances. Ignoring for now.
         super(IncludeFile, self).__init__(
@@ -336,7 +345,20 @@ class IncludeFile(Parameter):
     def load_parameter(self, v):
         if v is None:
             return v
-        return v.decode(self.name, var_type="Parameter")
+
+        # Get the raw content from the file
+        content = v.decode(self.name, var_type="Parameter")
+        # If a parser is specified, use it to parse the content
+        if self._parser is not None:
+            try:
+                return ConfigInput._call_parser(self._parser, content)
+            except Exception as e:
+                raise MetaflowException(
+                    "Failed to parse content in parameter '%s' using parser: %s"
+                    % (self.name, str(e))
+                ) from e
+
+        return content
 
     @staticmethod
     def _eval_default(is_text, encoding, default_path):

metaflow/lint.py

@@ -134,7 +134,13 @@ def check_valid_transitions(graph):
     msg = (
         "Step *{0.name}* specifies an invalid self.next() transition. "
         "Make sure the self.next() expression matches with one of the "
-        "supported transition types."
+        "supported transition types:\n"
+        "  • Linear: self.next(self.step_name)\n"
+        "  • Fan-out: self.next(self.step1, self.step2, ...)\n"
+        "  • Foreach: self.next(self.step, foreach='variable')\n"
+        "  • Switch: self.next({{\"key\": self.step, ...}}, condition='variable')\n\n"
+        "For switch statements, keys must be string literals, numbers or config expressions "
+        "(self.config.key_name), not variables."
     )
     for node in graph:
         if node.type != "end" and node.has_tail_next and node.invalid_tail_next:
@@ -169,6 +175,8 @@ def check_for_acyclicity(graph):
 
     def check_path(node, seen):
         for n in node.out_funcs:
+            if node.type == "split-switch" and n == node.name:
+                continue
             if n in seen:
                 path = "->".join(seen + [n])
                 raise LintWarn(
@@ -232,7 +240,15 @@ def check_split_join_balance(graph):
             new_stack = split_stack
         elif node.type in ("split", "foreach"):
             new_stack = split_stack + [("split", node.out_funcs)]
+        elif node.type == "split-switch":
+            # For a switch, continue traversal down each path with the same stack
+            for n in node.out_funcs:
+                if node.type == "split-switch" and n == node.name:
+                    continue
+                traverse(graph[n], split_stack)
+            return
         elif node.type == "end":
+            new_stack = split_stack
             if split_stack:
                 _, split_roots = split_stack.pop()
                 roots = ", ".join(split_roots)
@@ -240,11 +256,22 @@ def check_split_join_balance(graph):
                     msg0.format(roots=roots), node.func_lineno, node.source_file
                 )
         elif node.type == "join":
+            new_stack = split_stack
             if split_stack:
                 _, split_roots = split_stack[-1]
                 new_stack = split_stack[:-1]
-                if len(node.in_funcs) != len(split_roots):
-                    paths = ", ".join(node.in_funcs)
+
+                # Resolve each incoming function to its root branch from the split.
+                resolved_branches = set(
+                    graph[n].split_branches[-1] for n in node.in_funcs
+                )
+
+                # compares the set of resolved branches against the expected branches
+                # from the split.
+                if len(resolved_branches) != len(
+                    split_roots
+                ) or resolved_branches ^ set(split_roots):
+                    paths = ", ".join(resolved_branches)
                     roots = ", ".join(split_roots)
                     raise LintWarn(
                         msg1.format(
@@ -266,13 +293,55 @@ def check_split_join_balance(graph):
 
             if not all_equal(map(parents, node.in_funcs)):
                 raise LintWarn(msg3.format(node), node.func_lineno, node.source_file)
+        else:
+            new_stack = split_stack
 
         for n in node.out_funcs:
+            if node.type == "split-switch" and n == node.name:
+                continue
             traverse(graph[n], new_stack)
 
     traverse(graph["start"], [])
 
 
+@linter.ensure_static_graph
+@linter.check
+def check_switch_splits(graph):
+    """Check conditional split constraints"""
+    msg0 = (
+        "Step *{0.name}* is a switch split but defines {num} transitions. "
+        "Switch splits must define at least 2 transitions."
+    )
+    msg1 = "Step *{0.name}* is a switch split but has no condition variable."
+    msg2 = "Step *{0.name}* is a switch split but has no switch cases defined."
+
+    for node in graph:
+        if node.type == "split-switch":
+            # Check at least 2 outputs
+            if len(node.out_funcs) < 2:
+                raise LintWarn(
+                    msg0.format(node, num=len(node.out_funcs)),
+                    node.func_lineno,
+                    node.source_file,
+                )
+
+            # Check condition exists
+            if not node.condition:
+                raise LintWarn(
+                    msg1.format(node),
+                    node.func_lineno,
+                    node.source_file,
+                )
+
+            # Check switch cases exist
+            if not node.switch_cases:
+                raise LintWarn(
+                    msg2.format(node),
+                    node.func_lineno,
+                    node.source_file,
+                )
+
+
 @linter.ensure_static_graph
 @linter.check
 def check_empty_foreaches(graph):
@@ -347,3 +416,25 @@ def check_nested_foreach(graph):
         if node.type == "foreach":
             if any(graph[p].type == "foreach" for p in node.split_parents):
                 raise LintWarn(msg.format(node), node.func_lineno, node.source_file)
+
+
+@linter.ensure_static_graph
+@linter.check
+def check_ambiguous_joins(graph):
+    for node in graph:
+        if node.type == "join":
+            problematic_parents = [
+                p_name
+                for p_name in node.in_funcs
+                if graph[p_name].type == "split-switch"
+            ]
+            if problematic_parents:
+                msg = (
+                    "A conditional path cannot lead directly to a join step.\n"
+                    "In your conditional step(s) {parents}, one or more of the possible paths transition directly to the join step {join_name}.\n"
+                    "As a workaround, please introduce an intermediate, unconditional step on that specific path before joining."
+                ).format(
+                    parents=", ".join("*%s*" % p for p in problematic_parents),
+                    join_name="*%s*" % node.name,
+                )
+                raise LintWarn(msg, node.func_lineno, node.source_file)

metaflow/meta_files.py

@@ -0,0 +1,13 @@
+_UNINITIALIZED = object()
+_info_file_content = _UNINITIALIZED
+
+
+def read_info_file():
+    # Prevent circular import
+    from .packaging_sys import MetaflowCodeContent
+
+    global _info_file_content
+
+    if id(_info_file_content) == id(_UNINITIALIZED):
+        _info_file_content = MetaflowCodeContent.get_info()
+    return _info_file_content

metaflow/metadata_provider/metadata.py

@@ -632,7 +632,12 @@ class MetadataProvider(object):
 
     def _get_git_info_as_dict(self):
         git_info = {}
-        env = self._environment.get_environment_info()
+        # NOTE: For flows executing remotely, we want to read from the INFO file of the code package that contains
+        # information on the original environment that deployed the flow.
+        # Otherwise git related info will be missing, as the repository is not part of the codepackage.
+        from metaflow.packaging_sys import MetaflowCodeContent
+
+        env = MetaflowCodeContent.get_info() or self._environment.get_environment_info()
         for key in [
             "repo_url",
             "branch_name",
@@ -674,11 +679,17 @@ class MetadataProvider(object):
         if code_sha:
             code_url = os.environ.get("METAFLOW_CODE_URL")
             code_ds = os.environ.get("METAFLOW_CODE_DS")
+            code_metadata = os.environ.get("METAFLOW_CODE_METADATA")
             metadata.append(
                 MetaDatum(
                     field="code-package",
                     value=json.dumps(
-                        {"ds_type": code_ds, "sha": code_sha, "location": code_url}
+                        {
+                            "ds_type": code_ds,
+                            "sha": code_sha,
+                            "location": code_url,
+                            "metadata": code_metadata,
+                        }
                     ),
                     type="code-package",
                     tags=["attempt_id:{0}".format(attempt)],