argo: implement fanout switch case support

metaflow/flowspec.py

@@ -938,6 +938,8 @@ 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.
 
         Parameters
         ----------
@@ -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,25 @@ 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
+
+
 # ---------------------------------------------------------------------------
 # Note on "sourceless" DAGNodes (used by FunctionSpec)
 # ---------------------------------------------------------------------------
@@ -168,6 +187,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 +242,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 +292,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 +533,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 +543,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 +601,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,6 @@
 import re
 from .exception import MetaflowException
+from .graph import switch_case_target_lists
 from .util import all_equal
 
 
@@ -321,11 +322,18 @@ def traverse(node, 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)
+            # 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 + [("split", 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

metaflow/plugins/argo/argo_workflows.py

@@ -1556,7 +1556,10 @@ def build_ancestor_tree(node_groups, switch_ancestors):
                     # NOTE: Due to an issue in Argo Workflows 'when' clauses, we can not use ternaries or 'safe' getters directly on a tasks['step-name'] due to this leading to errors when the step has not executed.
                     conditional_when = "||".join(
                         [
-                            "({{=(tasks['%s'].status == 'Succeeded' ? tasks['%s'].outputs.parameters['switch-step'] : nil) == '%s'}})"
+                            # switch-step holds a comma-separated list of chosen targets
+                            # (single-target cases have no comma; fanout cases have one).
+                            # CEL split+exists handles both uniformly.
+                            "({{=(tasks['%s'].status == 'Succeeded' ? tasks['%s'].outputs.parameters['switch-step'].split(',').exists(x, x == '%s') : false)}})"
                             % (
                                 self._sanitize(switch_in_func),
                                 self._sanitize(switch_in_func),
@@ -1659,7 +1662,7 @@ def build_ancestor_tree(node_groups, switch_ancestors):
                                 .name("%s-recursion" % sanitized_name)
                                 .template(sanitized_name)
                                 .when(
-                                    "{{steps.%s-internal.outputs.parameters.switch-step}}==%s"
+                                    "{{=steps['%s-internal'].outputs.parameters['switch-step'].split(',').exists(x, x == '%s')}}"
                                     % (sanitized_name, node.name)
                                 )
                                 .arguments(
@@ -1801,7 +1804,7 @@ def build_ancestor_tree(node_groups, switch_ancestors):
                 ):
                     in_func = node.in_funcs[0]
                     foreach_task.when(
-                        "{{tasks.%s.outputs.parameters.switch-step}}==%s"
+                        "{{=tasks['%s'].outputs.parameters['switch-step'].split(',').exists(x, x == '%s')}}"
                         % (self._sanitize(in_func), node.name)
                     )
                 dag_tasks.append(foreach_task)

metaflow/plugins/argo/argo_workflows_decorator.py

@@ -128,9 +128,11 @@ def task_finished(
         if graph[step_name].type == "split-switch":
             # TODO: A nicer way to access the chosen step?
             _out_funcs, _ = flow._transition
-            chosen_step = _out_funcs[0]
             with open("/mnt/out/switch_step", "w") as file:
-                file.write(chosen_step)
+                # Comma-separated so fanout cases (list-valued) and single-target
+                # cases both fit in the same scalar Argo output parameter.
+                # The DAG/step conditions use CEL `.split(',').exists(...)`.
+                file.write(",".join(_out_funcs))
 
         # For steps that have a `@parallel` decorator set to them, we will be relying on Jobsets
         # to run the task. In this case, we cannot set anything in the