Support multi-target conditional switch cases

metaflow/flowspec.py

@@ -938,6 +938,9 @@ def next(self, *dsts: Callable[..., None], **kwargs) -> None:
           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.
+          A case value may also be a non-empty list or tuple of step methods to fan out
+          after the condition has selected that case. Multi-target cases must not share
+          target steps with other cases.
 
         Parameters
         ----------
@@ -977,9 +980,8 @@ def next(self, *dsts: Callable[..., None], **kwargs) -> None:
                 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
-                    )
+                    "non-empty dictionary mapping condition values to a step method "
+                    "or a non-empty list or tuple of step methods.".format(step=step)
                 )
                 raise InvalidNextException(msg)
 
@@ -1021,26 +1023,43 @@ def next(self, *dsts: Callable[..., None], **kwargs) -> None:
                 )
 
             # Get the chosen step and set transition directly
-            chosen_step_func = switch_cases[condition_value]
+            chosen_step_funcs = switch_cases[condition_value]
+            if isinstance(chosen_step_funcs, (list, tuple)):
+                if not chosen_step_funcs:
+                    msg = (
+                        "Step *{step}* specifies an empty switch transition. "
+                        "Make sure the value in the dictionary is a step method "
+                        "or a non-empty list or tuple of step methods.".format(
+                            step=step
+                        )
+                    )
+                    raise InvalidNextException(msg)
+            else:
+                chosen_step_funcs = [chosen_step_funcs]
 
             # Validate that the chosen step exists
-            try:
-                name = chosen_step_func.__func__.__name__
-            except:
-                msg = (
-                    "Step *{step}* specifies a switch transition that is not a function. "
-                    "Make sure the value in the dictionary is a method "
-                    "of the Flow class.".format(step=step)
-                )
-                raise InvalidNextException(msg)
-            if not hasattr(self, name):
-                msg = (
-                    "Step *{step}* specifies a switch transition to an "
-                    "unknown step, *{name}*.".format(step=step, name=name)
-                )
-                raise InvalidNextException(msg)
+            names = []
+            for chosen_step_func in chosen_step_funcs:
+                try:
+                    name = chosen_step_func.__func__.__name__
+                except AttributeError:
+                    msg = (
+                        "Step *{step}* specifies a switch transition that is not a function. "
+                        "Make sure the value in the dictionary is a step method "
+                        "or a non-empty list or tuple of step methods.".format(
+                            step=step
+                        )
+                    )
+                    raise InvalidNextException(msg)
+                if not hasattr(self, name):
+                    msg = (
+                        "Step *{step}* specifies a switch transition to an "
+                        "unknown step, *{name}*.".format(step=step, name=name)
+                    )
+                    raise InvalidNextException(msg)
+                names.append(name)
 
-            self._transition = ([name], None)
+            self._transition = (names, None)
             return
 
         # Check for an invalid transition: a dictionary used without a 'condition' parameter.

metaflow/graph.py

@@ -65,6 +65,34 @@ def _ast_literal_value(node):
     return None
 
 
+def switch_case_targets(case_value):
+    if isinstance(case_value, (list, tuple)):
+        return list(case_value)
+    return [case_value]
+
+
+def switch_case_target_lists(switch_cases):
+    return [switch_case_targets(case_value) for case_value in switch_cases.values()]
+
+
+def flatten_switch_cases(switch_cases):
+    out_funcs = []
+    for targets in switch_case_target_lists(switch_cases):
+        for target in targets:
+            if target not in out_funcs:
+                out_funcs.append(target)
+    return out_funcs
+
+
+def split_branch_for_node(node, split_name):
+    """Return the branch root for ``node`` at the named enclosing split."""
+    try:
+        split_index = node.split_parents.index(split_name)
+        return node.split_branches[split_index]
+    except (ValueError, IndexError):
+        return None
+
+
 # ---------------------------------------------------------------------------
 # Note on "sourceless" DAGNodes (used by FunctionSpec)
 # ---------------------------------------------------------------------------
@@ -168,6 +196,25 @@ def __init__(
     def _expr_str(self, expr):
         return "%s.%s" % (expr.value.id, expr.attr)
 
+    def _parse_switch_target(self, value):
+        if isinstance(value, ast.Attribute) and isinstance(value.value, ast.Name):
+            if value.value.id == "self":
+                return value.attr
+            return None
+        return None
+
+    def _parse_switch_value(self, value):
+        target = self._parse_switch_target(value)
+        if target is not None:
+            return target
+
+        if isinstance(value, (ast.List, ast.Tuple)) and value.elts:
+            targets = [self._parse_switch_target(elt) for elt in value.elts]
+            if all(target is not None for target in targets):
+                return targets
+
+        return None
+
     def _parse_switch_dict(self, dict_node):
         switch_cases = {}
 
@@ -204,17 +251,10 @@ def _parse_switch_dict(self, dict_node):
                 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:
+                case_value = self._parse_switch_value(value)
+                if case_value is None:
                     return None
+                switch_cases[case_key] = case_value
 
         return switch_cases if switch_cases else None
 
@@ -261,7 +301,7 @@ def _parse(self, func_ast, lineno):
                     self.type = "split-switch"
                     self.condition = condition_name
                     self.switch_cases = switch_cases
-                    self.out_funcs = list(switch_cases.values())
+                    self.out_funcs = flatten_switch_cases(switch_cases)
                     self.invalid_tail_next = False
                     return
 
@@ -502,7 +542,8 @@ def traverse(node, seen, split_parents, split_branches):
             elif node.type == "join":
                 # ignore joins without splits
                 if split_parents:
-                    self[split_parents[-1]].matching_join = node.name
+                    if self[split_parents[-1]].type != "split-switch":
+                        self[split_parents[-1]].matching_join = node.name
                     node.split_parents = split_parents
                     node.split_branches = split_branches[:-1]
                     split_parents = split_parents[:-1]
@@ -511,6 +552,26 @@ def traverse(node, seen, split_parents, split_branches):
                 node.split_parents = split_parents
                 node.split_branches = split_branches
 
+            if node.type == "split-switch":
+                for targets in switch_case_target_lists(node.switch_cases):
+                    case_is_fanout = len(targets) > 1
+                    child_split_parents = (
+                        split_parents + [node.name] if case_is_fanout else split_parents
+                    )
+                    for n in targets:
+                        if n == node.name:
+                            continue
+                        if n not in seen and n in self:
+                            child = self[n]
+                            child.in_funcs.add(node.name)
+                            traverse(
+                                child,
+                                seen + [n],
+                                child_split_parents,
+                                split_branches + ([n] if case_is_fanout else []),
+                            )
+                return
+
             for n in node.out_funcs:
                 # graph may contain loops - ignore them
                 if n not in seen:
@@ -549,12 +610,13 @@ def edge_specs():
             for node in self.nodes.values():
                 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
+                    for case_value, case_target in node.switch_cases.items():
+                        for step_name in switch_case_targets(case_target):
+                            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)

metaflow/lint.py

@@ -1,5 +1,10 @@
 import re
 from .exception import MetaflowException
+from .graph import (
+    split_branch_for_node,
+    switch_case_targets,
+    switch_case_target_lists,
+)
 from .util import all_equal
 
 
@@ -212,7 +217,8 @@ def check_valid_transitions(graph):
         "  • 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"
+        '  • Switch: self.next({{"key": self.step, '
+        "\"key2\": [self.step1, self.step2]}}, condition='variable')\n\n"
         "For switch statements, keys must be string literals, numbers or config expressions "
         "(self.config.key_name), not variables."
     )
@@ -319,13 +325,20 @@ def traverse(node, split_stack):
         if node.type in ("start", "linear"):
             new_stack = split_stack
         elif node.type in ("split", "foreach"):
-            new_stack = split_stack + [("split", node.out_funcs)]
+            new_stack = split_stack + [(node.name, 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)
+            # A switch selects exactly one case. A list-valued case then behaves
+            # like a split, but only within that selected case.
+            for targets in switch_case_target_lists(node.switch_cases):
+                case_stack = (
+                    split_stack + [(node.name, targets)]
+                    if len(targets) > 1
+                    else split_stack
+                )
+                for n in targets:
+                    if n == node.name:
+                        continue
+                    traverse(graph[n], case_stack)
             return
         elif node.type == "end":
             new_stack = split_stack
@@ -340,13 +353,19 @@ def traverse(node, split_stack):
         elif node.type == "join":
             new_stack = split_stack
             if split_stack:
-                _, split_roots = split_stack[-1]
+                split_name, split_roots = split_stack[-1]
                 new_stack = split_stack[:-1]
 
                 # 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
-                )
+                resolved_branches = {
+                    split_branch_for_node(graph[n], split_name) for n in node.in_funcs
+                }
+                resolved_branches.discard(None)
+                # A shared switch join has static predecessors from every case,
+                # though only the selected case runs. Validate this traversal's
+                # case against its own predecessors.
+                if graph[split_name].type == "split-switch":
+                    resolved_branches.intersection_update(split_roots)
 
                 # compares the set of resolved branches against the expected branches
                 # from the split.
@@ -379,8 +398,6 @@ def parents(n):
             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[graph.start_step], [])
@@ -391,18 +408,23 @@ def parents(n):
 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."
+        "Step *{0.name}* is a switch split with too few cases: "
+        "{num} found, at least 2 required."
     )
     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."
+    msg3 = (
+        "Step *{0.name}* has multi-target switch cases *{case_a}* and "
+        "*{case_b}* that share target step(s) *{targets}*. Multi-target switch "
+        "cases must have disjoint targets."
+    )
 
     for node in graph:
         if node.type == "split-switch":
-            # Check at least 2 outputs
-            if len(node.out_funcs) < 2:
+            # out_funcs contains unique graph edges, not switch choices.
+            if len(node.switch_cases) < 2:
                 raise LintWarn(
-                    msg0.format(node, num=len(node.out_funcs)),
+                    msg0.format(node, num=len(node.switch_cases)),
                     node.func_lineno,
                     node.source_file,
                 )
@@ -423,6 +445,25 @@ def check_switch_splits(graph):
                     node.source_file,
                 )
 
+            cases = [
+                (case_value, set(switch_case_targets(case_targets)))
+                for case_value, case_targets in node.switch_cases.items()
+            ]
+            for case_index, (case_value, targets) in enumerate(cases):
+                for other_value, other_targets in cases[:case_index]:
+                    overlap = targets & other_targets
+                    if overlap and (len(targets) > 1 or len(other_targets) > 1):
+                        raise LintWarn(
+                            msg3.format(
+                                node,
+                                case_a=repr(other_value),
+                                case_b=repr(case_value),
+                                targets=", ".join(sorted(overlap)),
+                            ),
+                            node.func_lineno,
+                            node.source_file,
+                        )
+
 
 @linter.ensure_static_graph
 @linter.check

metaflow/plugins/argo/argo_workflows.py

@@ -14,7 +14,7 @@
 from metaflow import JSONType, current
 from metaflow.decorators import flow_decorators
 from metaflow.exception import MetaflowException
-from metaflow.graph import FlowGraph
+from metaflow.graph import FlowGraph, switch_case_target_lists
 from metaflow.includefile import FilePathClass
 from metaflow.metaflow_config import (
     ARGO_EVENTS_EVENT,
@@ -1027,6 +1027,17 @@ def _compile_workflow_template(self):
 
     # Visit every node and record information on conditional step structure
     def _parse_conditional_branches(self):
+        for node in self.graph.nodes.values():
+            if node.type == "split-switch":
+                for targets in switch_case_target_lists(node.switch_cases):
+                    if len(targets) > 1:
+                        raise ArgoWorkflowsException(
+                            "Step *%s* uses a multi-target switch case (fanout), "
+                            "which is not yet supported on Argo Workflows. "
+                            "Use a dedicated step to fan out after the condition instead."
+                            % node.name
+                        )
+
         self.conditional_nodes = set()
         self.conditional_join_nodes = set()
         self.matching_conditional_join_dict = {}