ob-metaflow 2.16.8.2rc0__py2.py3-none-any.whl → 2.16.8.2rc2__py2.py3-none-any.whl

metaflow/_vendor/click/core.py

@@ -719,7 +719,7 @@ class BaseCommand(object):
         prog_name=None,
         complete_var=None,
         standalone_mode=True,
-        **extra
+        **extra,
     ):
         """This is the way to invoke a script with all the bells and
         whistles as a command line application.  This will always terminate
@@ -1101,7 +1101,7 @@ class MultiCommand(Command):
         subcommand_metavar=None,
         chain=False,
         result_callback=None,
-        **attrs
+        **attrs,
     ):
         Command.__init__(self, name, **attrs)
         if no_args_is_help is None:
@@ -1463,6 +1463,7 @@ class Parameter(object):
         parameter. The old callback format will still work, but it will
         raise a warning to give you a chance to migrate the code easier.
     """
+
     param_type_name = "parameter"
 
     def __init__(
@@ -1708,7 +1709,7 @@ class Option(Parameter):
         hidden=False,
         show_choices=True,
         show_envvar=False,
-        **attrs
+        **attrs,
     ):
         default_is_missing = attrs.get("default", _missing) is _missing
         Parameter.__init__(self, param_decls, type=type, **attrs)

metaflow/_vendor/imghdr/__init__.py

@@ -6,13 +6,7 @@ import warnings
 __all__ = ["what"]
 
 
-# python-deadlib: Replace deprecation warning not to raise exception
-warnings.warn(
-    f"{__name__} was removed in Python 3.13. "
-    f"Please be aware that you are currently NOT using standard '{__name__}', "
-    f"but instead a separately installed 'standard-{__name__}'.",
-    DeprecationWarning, stacklevel=2
-)
+warnings._deprecated(__name__, remove=(3, 13))
 
 
 #-------------------------#

metaflow/cli.py

@@ -7,7 +7,6 @@ from datetime import datetime
 
 import metaflow.tracing as tracing
 from metaflow._vendor import click
-from metaflow.system import _system_logger, _system_monitor
 
 from . import decorators, lint, metaflow_version, parameters, plugins
 from .cli_args import cli_args
@@ -27,6 +26,7 @@ from .metaflow_config import (
     DEFAULT_PACKAGE_SUFFIXES,
 )
 from .metaflow_current import current
+from metaflow.system import _system_monitor, _system_logger
 from .metaflow_environment import MetaflowEnvironment
 from .packaging_sys import MetaflowCodeContent
 from .plugins import (
@@ -38,9 +38,9 @@ from .plugins import (
 )
 from .pylint_wrapper import PyLint
 from .R import metaflow_r_version, use_r
+from .util import get_latest_run_id, resolve_identity
 from .user_configs.config_options import LocalFileInput, config_options
 from .user_configs.config_parameters import ConfigValue
-from .util import get_latest_run_id, resolve_identity
 
 ERASE_TO_EOL = "\033[K"
 HIGHLIGHT = "red"
@@ -56,15 +56,6 @@ def echo_dev_null(*args, **kwargs):
 
 
 def echo_always(line, **kwargs):
-    if kwargs.pop("wrap", False):
-        import textwrap
-
-        indent_str = INDENT if kwargs.get("indent", None) else ""
-        effective_width = 80 - len(indent_str)
-        wrapped = textwrap.wrap(line, width=effective_width, break_long_words=False)
-        line = "\n".join(indent_str + l for l in wrapped)
-        kwargs["indent"] = False
-
     kwargs["err"] = kwargs.get("err", True)
     if kwargs.pop("indent", None):
         line = "\n".join(INDENT + x for x in line.splitlines())

metaflow/cli_components/run_cmds.py

@@ -13,6 +13,8 @@ from ..package import MetaflowPackage
 from ..runtime import NativeRuntime
 from ..system import _system_logger
 
+# from ..client.core import Run
+
 from ..tagging_util import validate_tags
 from ..util import get_latest_run_id, write_latest_run_id
 
@@ -230,6 +232,19 @@ def resume(
                     step_to_rerun, ",".join(list(obj.graph.nodes.keys()))
                 )
             )
+
+        ## TODO: instead of checking execution path here, can add a warning later
+        ## instead of throwing an error. This is for resuming a step which was not
+        ## taken inside a branch i.e. not present in the execution path.
+
+        # origin_run = Run(f"{obj.flow.name}/{origin_run_id}", _namespace_check=False)
+        # executed_steps = {step.path_components[-1] for step in origin_run}
+        # if step_to_rerun not in executed_steps:
+        #     raise CommandException(
+        #         f"Cannot resume from step '{step_to_rerun}'. This step was not "
+        #         f"part of the original execution path for run '{origin_run_id}'."
+        #     )
+
         steps_to_rerun = {step_to_rerun}
 
     if run_id:

metaflow/client/core.py

@@ -831,12 +831,10 @@ class MetaflowCode(object):
         )
         self._code_obj = BytesIO(blobdata)
         self._info = MetaflowPackage.cls_get_info(self._code_metadata, self._code_obj)
-        self._code_obj.seek(0)
         if self._info:
             self._flowspec = MetaflowPackage.cls_get_content(
                 self._code_metadata, self._code_obj, self._info["script"]
             )
-            self._code_obj.seek(0)
         else:
             raise MetaflowInternalError("Code package metadata is invalid.")
 
@@ -887,9 +885,7 @@ class MetaflowCode(object):
             TarFile for everything in this code package
         """
         if self._backend.type == "tgz":
-            to_return = self._backend.cls_open(self._code_obj)
-            self._code_obj.seek(0)
-            return to_return
+            return self._backend.cls_open(self._code_obj)
         raise RuntimeError("Archive is not a tarball")
 
     def extract(self) -> TemporaryDirectory:
@@ -925,7 +921,6 @@ class MetaflowCode(object):
         MetaflowPackage.cls_extract_into(
             self._code_metadata, self._code_obj, tmp.name, ContentType.USER_CONTENT
         )
-        self._code_obj.seek(0)
         return tmp
 
     @property

metaflow/extension_support/__init__.py

@@ -205,10 +205,9 @@ def package_mfext_all():
     # the packaged metaflow_extensions directory "self-contained" so that
     # python doesn't go and search other parts of the system for more
     # metaflow_extensions.
-    if _all_packages:
-        yield os.path.join(
-            os.path.dirname(os.path.abspath(__file__)), "_empty_file.py"
-        ), os.path.join(EXT_PKG, "__init__.py")
+    yield os.path.join(
+        os.path.dirname(os.path.abspath(__file__)), "_empty_file.py"
+    ), os.path.join(EXT_PKG, "__init__.py")
 
     for p in _all_packages:
         for path_tuple in package_mfext_package(p):

metaflow/flowspec.py

@@ -788,6 +788,35 @@ class FlowSpec(metaclass=FlowSpecMeta):
         value = item if _is_primitive_type(item) else reprlib.Repr().repr(item)
         return basestring(value)[:MAXIMUM_FOREACH_VALUE_CHARS]
 
+    def _validate_switch_cases(self, switch_cases, step):
+        resolved_cases = {}
+        for case_key, step_method in switch_cases.items():
+            if isinstance(case_key, str) and case_key.startswith("config:"):
+                full_path = case_key[len("config:") :]
+                parts = full_path.split(".", 1)
+                if len(parts) == 2:
+                    config_var_name, config_key_name = parts
+                    try:
+                        config_obj = getattr(self, config_var_name)
+                        resolved_key = str(getattr(config_obj, config_key_name))
+                    except AttributeError:
+                        msg = (
+                            "Step *{step}* references unknown config '{path}' "
+                            "in switch case.".format(step=step, path=full_path)
+                        )
+                        raise InvalidNextException(msg)
+                else:
+                    raise MetaflowInternalError(
+                        "Invalid config path format in switch case."
+                    )
+            else:
+                resolved_key = case_key
+
+            func_name = step_method.__func__.__name__
+            resolved_cases[resolved_key] = func_name
+
+        return resolved_cases
+
     def next(self, *dsts: Callable[..., None], **kwargs) -> None:
         """
         Indicates the next step to execute after this step has completed.
@@ -812,6 +841,15 @@ class FlowSpec(metaclass=FlowSpecMeta):
           evaluates to an iterator. A task will be launched for each value in the iterator and
           each task will execute the code specified by the step `foreach_step`.
 
+        - Switch statement:
+          ```
+          self.next({"case1": self.step_a, "case2": self.step_b}, condition='condition_variable')
+          ```
+          In this situation, `step_a` and `step_b` are methods in the current class decorated
+          with the `@step` decorator and `condition_variable` is a variable name in the current
+          class. The value of the condition variable determines which step to execute. If the
+          value doesn't match any of the dictionary keys, a RuntimeError is raised.
+
         Parameters
         ----------
         dsts : Callable[..., None]
@@ -827,6 +865,7 @@ class FlowSpec(metaclass=FlowSpecMeta):
 
         foreach = kwargs.pop("foreach", None)
         num_parallel = kwargs.pop("num_parallel", None)
+        condition = kwargs.pop("condition", None)
         if kwargs:
             kw = next(iter(kwargs))
             msg = (
@@ -843,6 +882,79 @@ class FlowSpec(metaclass=FlowSpecMeta):
             )
             raise InvalidNextException(msg)
 
+        # check: switch case using condition
+        if condition is not None:
+            if len(dsts) != 1 or not isinstance(dsts[0], dict) or not dsts[0]:
+                msg = (
+                    "Step *{step}* has an invalid self.next() transition. "
+                    "When using 'condition', the transition must be to a single, "
+                    "non-empty dictionary mapping condition values to step methods.".format(
+                        step=step
+                    )
+                )
+                raise InvalidNextException(msg)
+
+            if not isinstance(condition, basestring):
+                msg = (
+                    "Step *{step}* has an invalid self.next() transition. "
+                    "The argument to 'condition' must be a string.".format(step=step)
+                )
+                raise InvalidNextException(msg)
+
+            if foreach is not None or num_parallel is not None:
+                msg = (
+                    "Step *{step}* has an invalid self.next() transition. "
+                    "Switch statements cannot be combined with foreach or num_parallel.".format(
+                        step=step
+                    )
+                )
+                raise InvalidNextException(msg)
+
+            switch_cases = dsts[0]
+
+            # Validate that condition variable exists
+            try:
+                condition_value = getattr(self, condition)
+            except AttributeError:
+                msg = (
+                    "Condition variable *self.{var}* in step *{step}* "
+                    "does not exist. Make sure you set self.{var} in this step.".format(
+                        step=step, var=condition
+                    )
+                )
+                raise InvalidNextException(msg)
+
+            resolved_switch_cases = self._validate_switch_cases(switch_cases, step)
+
+            if str(condition_value) not in resolved_switch_cases:
+                available_cases = list(resolved_switch_cases.keys())
+                raise RuntimeError(
+                    f"Switch condition variable '{condition}' has value '{condition_value}' "
+                    f"which is not in the available cases: {available_cases}"
+                )
+
+            # Get the chosen step and set transition directly
+            chosen_step = resolved_switch_cases[str(condition_value)]
+
+            # Validate that the chosen step exists
+            if not hasattr(self, chosen_step):
+                msg = (
+                    "Step *{step}* specifies a switch transition to an "
+                    "unknown step, *{name}*.".format(step=step, name=chosen_step)
+                )
+                raise InvalidNextException(msg)
+
+            self._transition = ([chosen_step], None)
+            return
+
+        # Check for an invalid transition: a dictionary used without a 'condition' parameter.
+        if len(dsts) == 1 and isinstance(dsts[0], dict):
+            msg = (
+                "Step *{step}* has an invalid self.next() transition. "
+                "Dictionary argument requires 'condition' parameter.".format(step=step)
+            )
+            raise InvalidNextException(msg)
+
         # check: all destinations are methods of this object
         funcs = []
         for i, dst in enumerate(dsts):
@@ -933,7 +1045,7 @@ class FlowSpec(metaclass=FlowSpecMeta):
             self._foreach_var = foreach
 
         # check: non-keyword transitions are valid
-        if foreach is None:
+        if foreach is None and condition is None:
             if len(dsts) < 1:
                 msg = (
                     "Step *{step}* has an invalid self.next() transition. "

metaflow/graph.py

@@ -68,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
@@ -83,6 +85,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 isinstance(key, ast.Str):
+                    case_key = key.s
+                elif isinstance(key, ast.Constant) and isinstance(key.value, str):
+                    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]
@@ -104,7 +156,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 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
@@ -151,6 +234,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(
@@ -219,6 +303,8 @@ class FlowGraph(object):
             if node.type in ("split", "foreach"):
                 node.split_parents = split_parents
                 split_parents = split_parents + [node.name]
+            elif node.type == "split-switch":
+                node.split_parents = split_parents
             elif node.type == "join":
                 # ignore joins without splits
                 if split_parents:
@@ -259,15 +345,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"
@@ -291,6 +399,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):
@@ -325,6 +435,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
@@ -339,8 +452,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)
@@ -348,8 +461,19 @@ 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
+                    ]
+                    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/lint.py

@@ -1,6 +1,6 @@
 import re
 from .exception import MetaflowException
-from .util import all_equal
+from .util import all_equal, get_split_branch_for_node
 
 
 class LintWarn(MetaflowException):
@@ -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 or config expressions "
+        "(self.config.key_name), not variables or numbers."
     )
     for node in graph:
         if node.type != "end" and node.has_tail_next and node.invalid_tail_next:
@@ -232,7 +238,13 @@ 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:
+                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,10 +252,25 @@ 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):
+
+                # Identify the split this join corresponds to from its parentage.
+                split_node_name = node.split_parents[-1]
+
+                # Resolve each incoming function to its root branch from the split.
+                resolved_branches = set()
+                for in_node_name in node.in_funcs:
+                    branch = get_split_branch_for_node(
+                        graph, in_node_name, split_node_name
+                    )
+                    if branch:
+                        resolved_branches.add(branch)
+
+                # compares the set of resolved branches against the expected branches from the split.
+                if len(resolved_branches) != len(split_roots):
                     paths = ", ".join(node.in_funcs)
                     roots = ", ".join(split_roots)
                     raise LintWarn(
@@ -266,6 +293,8 @@ 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:
             traverse(graph[n], new_stack)
@@ -273,6 +302,44 @@ def check_split_join_balance(graph):
     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):