From bdb41f52a5899f70e2f324236dec383907b7b006 Mon Sep 17 00:00:00 2001 From: Nissan Pow Date: Tue, 4 Aug 2026 13:38:44 -0700 Subject: [PATCH 1/6] Support fanout switch cases --- metaflow/flowspec.py | 53 +++++++++---- metaflow/graph.py | 87 ++++++++++++++++---- metaflow/lint.py | 18 +++-- metaflow/runtime.py | 77 ++++++++++++------ metaflow/task.py | 52 ++++++++++-- test/unit/test_switch_fanout_cases.py | 109 ++++++++++++++++++++++++++ 6 files changed, 325 insertions(+), 71 deletions(-) create mode 100644 test/unit/test_switch_fanout_cases.py diff --git a/metaflow/flowspec.py b/metaflow/flowspec.py index 307c19df2c5..ce1ffedbcad 100644 --- a/metaflow/flowspec.py +++ b/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. diff --git a/metaflow/graph.py b/metaflow/graph.py index 3e4eb5c85d5..fe295e796e0 100644 --- a/metaflow/graph.py +++ b/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) diff --git a/metaflow/lint.py b/metaflow/lint.py index b95e7ba0e95..a88be88a2ed 100644 --- a/metaflow/lint.py +++ b/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 diff --git a/metaflow/runtime.py b/metaflow/runtime.py index 33ee71b95f4..a9e13f103f2 100644 --- a/metaflow/runtime.py +++ b/metaflow/runtime.py @@ -47,6 +47,7 @@ from .debug import debug from .decorators import flow_decorators from .flowspec import FlowStateItems +from .graph import switch_case_target_lists from .mflog import mflog, RUNTIME_LOG_SOURCE from .util import to_unicode, compress_list, unicode_type, get_latest_task_pathspec from .clone_util import clone_task_helper @@ -1304,7 +1305,19 @@ def siblings(foreach_stack): ) ) - required_count = len(matching_split.out_funcs) + if matching_split.type == "split-switch": + case_targets = self._switch_case_targets_for_task( + matching_split, task.step + ) + if case_targets is None: + raise MetaflowInternalError( + "Step *%s* is joining a switch fanout from *%s*, " + "but the runtime could not determine the selected " + "switch case." % (next_step, matching_split.name) + ) + required_count = len(case_targets) + else: + required_count = len(matching_split.out_funcs) join_type = "linear" index = self._translate_index(task, next_step, "linear") if len(required_tasks) == required_count: @@ -1315,19 +1328,33 @@ def siblings(foreach_stack): index, ) - def _queue_task_switch(self, task, next_steps, is_recursive): - chosen_step = next_steps[0] + def _switch_case_targets_for_task(self, switch_node, step_name): + node = self._graph[step_name] + branch_root = None + if ( + node.split_parents + and node.split_parents[-1] == switch_node.name + and node.split_branches + ): + branch_root = node.split_branches[-1] + + for targets in switch_case_target_lists(switch_node.switch_cases): + if branch_root in targets or step_name in targets: + return targets + return None - loop_mode = LoopBehavior.NONE - if is_recursive: - if chosen_step != task.step: - # We are exiting a loop - loop_mode = LoopBehavior.EXITING - else: - # We are staying in the loop - loop_mode = LoopBehavior.LOOPING - index = self._translate_index(task, chosen_step, "linear", None, loop_mode) - self._queue_push(chosen_step, {"input_paths": [task.path]}, index) + def _queue_task_switch(self, task, next_steps, is_recursive): + for chosen_step in next_steps: + loop_mode = LoopBehavior.NONE + if is_recursive: + if chosen_step != task.step: + # We are exiting a loop + loop_mode = LoopBehavior.EXITING + else: + # We are staying in the loop + loop_mode = LoopBehavior.LOOPING + index = self._translate_index(task, chosen_step, "linear", None, loop_mode) + self._queue_push(chosen_step, {"input_paths": [task.path]}, index) def _queue_task_foreach(self, task, next_steps): # CHECK: this condition should be enforced by the linter but @@ -1408,31 +1435,29 @@ def _queue_tasks(self, finished_tasks): if self._graph[task.step].type == "split-switch": is_recursive = task.step in self._graph[task.step].out_funcs - if len(next_steps) != 1: - msg = ( - "Switch step *{step}* should transition to exactly " - "one step at runtime, but got: {actual}" - ) - raise MetaflowInternalError( - msg.format(step=task.step, actual=", ".join(next_steps)) - ) - if next_steps[0] not in expected: + expected_cases = switch_case_target_lists( + self._graph[task.step].switch_cases + ) + if next_steps not in expected_cases: msg = ( "Switch step *{step}* transitioned to unexpected " - "step *{actual}*. Expected one of: {expected}" + "step(s) *{actual}*. Expected one of: {expected}" ) raise MetaflowInternalError( msg.format( step=task.step, - actual=next_steps[0], - expected=", ".join(expected), + actual=", ".join(next_steps), + expected=", ".join( + "[%s]" % ", ".join(targets) + for targets in expected_cases + ), ) ) # When exiting a recursive loop, we mark that the loop itself has # finished by adding a special entry in self._finished which has # an iteration stack that is shorter (ie: we are out of the loop) so # that we can then find it when looking at successor tasks to launch. - if is_recursive and next_steps[0] != task.step: + if is_recursive and task.step not in next_steps: step_name, finished_tuple, iteration_tuple = task.finished_id self._finished[ (step_name, finished_tuple, iteration_tuple[:-1]) diff --git a/metaflow/task.py b/metaflow/task.py index 74ca92a6bdb..8071c1652d6 100644 --- a/metaflow/task.py +++ b/metaflow/task.py @@ -17,6 +17,7 @@ from .metaflow_profile import from_start from .mflog import TASK_LOG_SOURCE from .datastore import Inputs, TaskDataStoreSet +from .graph import switch_case_target_lists from .exception import ( MetaflowInternalError, MetaflowDataMissing, @@ -822,12 +823,51 @@ def run_step( if join_type != "foreach": # Find the corresponding split node from the graph. split_node = self.flow._graph[node.split_parents[-1]] - # The number of expected inputs is the number of branches - # from that split -- we can't use in_funcs because there may - # be more due to split-switch branches that all converge here. - expected_inputs = len(split_node.out_funcs) - - if len(inputs) != expected_inputs: + expected_inputs = None + expected_counts = None + + if split_node.type == "split-switch": + branch_roots = set() + for in_func in node.in_funcs: + in_node = self.flow._graph[in_func] + if ( + in_node.split_parents + and in_node.split_parents[-1] == split_node.name + and in_node.split_branches + ): + branch_roots.add(in_node.split_branches[-1]) + + expected_counts = set( + len(targets) + for targets in switch_case_target_lists( + split_node.switch_cases + ) + if len(targets) > 1 + and set(targets).issubset(branch_roots) + ) + if not expected_counts: + expected_inputs = len(split_node.out_funcs) + else: + # The number of expected inputs is the number of branches + # from that split -- we can't use in_funcs because there may + # be more due to split-switch branches that all converge here. + expected_inputs = len(split_node.out_funcs) + + if expected_counts: + if len(inputs) not in expected_counts: + raise MetaflowDataMissing( + "Join *%s* expected one of %s inputs but only %d " + "inputs were found" + % ( + step_name, + ", ".join( + str(count) + for count in sorted(expected_counts) + ), + len(inputs), + ) + ) + elif len(inputs) != expected_inputs: raise MetaflowDataMissing( "Join *%s* expected %d inputs but only %d inputs " "were found" % (step_name, expected_inputs, len(inputs)) diff --git a/test/unit/test_switch_fanout_cases.py b/test/unit/test_switch_fanout_cases.py new file mode 100644 index 00000000000..f411905f4e6 --- /dev/null +++ b/test/unit/test_switch_fanout_cases.py @@ -0,0 +1,109 @@ +import pytest + +from metaflow import FlowSpec, step +from metaflow.flowspec import InvalidNextException +from metaflow.lint import LintWarn, linter + + +class SwitchFanoutCaseFlow(FlowSpec): + @step + def start(self): + self.route = "miss" + self.next( + {"hit": self.finalize, "miss": [self.clip, self.face]}, + condition="route", + ) + + @step + def finalize(self): + self.next(self.end) + + @step + def clip(self): + self.next(self.join_miss) + + @step + def face(self): + self.next(self.join_miss) + + @step + def join_miss(self, inputs): + self.next(self.end) + + @step + def end(self): + pass + + +class SwitchFanoutToEndJoinFlow(FlowSpec): + @step + def start(self): + self.route = "miss" + self.next( + {"hit": self.finalize, "miss": [self.clip, self.face]}, + condition="route", + ) + + @step + def finalize(self): + self.next(self.end) + + @step + def clip(self): + self.next(self.end) + + @step + def face(self): + self.next(self.end) + + @step + def end(self, inputs): + pass + + +def test_graph_parses_switch_fanout_case(): + graph = SwitchFanoutCaseFlow._graph + + assert graph["start"].type == "split-switch" + assert graph["start"].switch_cases == { + "hit": "finalize", + "miss": ["clip", "face"], + } + assert graph["start"].out_funcs == ["finalize", "clip", "face"] + assert graph["clip"].split_parents == ["start"] + assert graph["clip"].split_branches == ["clip"] + assert graph["face"].split_parents == ["start"] + assert graph["face"].split_branches == ["face"] + assert graph["join_miss"].type == "join" + assert graph["join_miss"].split_parents == ["start"] + + +def test_switch_fanout_case_passes_lint(): + linter.run_checks(SwitchFanoutCaseFlow._graph) + + +def test_runtime_switch_fanout_transition_uses_selected_case_targets(): + flow = SwitchFanoutCaseFlow(use_cli=False) + flow._current_step = "start" + flow.route = "miss" + + flow.next( + {"hit": flow.finalize, "miss": [flow.clip, flow.face]}, + condition="route", + ) + + assert flow._transition == (["clip", "face"], None) + + +def test_runtime_switch_fanout_rejects_empty_case(): + flow = SwitchFanoutCaseFlow(use_cli=False) + flow._current_step = "start" + flow.route = "miss" + + with pytest.raises(InvalidNextException, match="empty switch transition"): + flow.next({"hit": flow.finalize, "miss": []}, condition="route") + + +def test_switch_fanout_case_cannot_join_at_terminal_step(): + with pytest.raises(LintWarn, match="terminal step .* should not be a join step"): + linter.run_checks(SwitchFanoutToEndJoinFlow._graph) From 19d84511d1ce51bccba36906f34a0c55be24f67b Mon Sep 17 00:00:00 2001 From: npow Date: Wed, 5 Aug 2026 22:24:02 +0000 Subject: [PATCH 2/6] argo: raise error for fanout switch cases (not yet supported) Co-Authored-By: Claude Sonnet 4.6 (1M context) --- metaflow/plugins/argo/argo_workflows.py | 13 ++++++++++++- metaflow/plugins/argo/argo_workflows_decorator.py | 7 +++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index b033c6d638f..bd6fa227db8 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/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 list-valued 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 = {} diff --git a/metaflow/plugins/argo/argo_workflows_decorator.py b/metaflow/plugins/argo/argo_workflows_decorator.py index 172d5c77ab4..059ad4953cb 100644 --- a/metaflow/plugins/argo/argo_workflows_decorator.py +++ b/metaflow/plugins/argo/argo_workflows_decorator.py @@ -4,6 +4,7 @@ from metaflow import current from metaflow.decorators import StepDecorator +from metaflow.exception import MetaflowException from metaflow.events import Trigger from metaflow.metadata_provider import MetaDatum from metaflow.graph import FlowGraph @@ -128,6 +129,12 @@ def task_finished( if graph[step_name].type == "split-switch": # TODO: A nicer way to access the chosen step? _out_funcs, _ = flow._transition + if len(_out_funcs) > 1: + raise MetaflowException( + "Step *%s* selected a switch case that fans out to multiple targets %s, " + "which is not yet supported on Argo Workflows." + % (step_name, _out_funcs) + ) chosen_step = _out_funcs[0] with open("/mnt/out/switch_step", "w") as file: file.write(chosen_step) From 38ba3f88f1c89bcdd6742c7e7f07f1b607356eaa Mon Sep 17 00:00:00 2001 From: npow Date: Mon, 10 Aug 2026 17:05:05 +0000 Subject: [PATCH 3/6] Fix nested and shared switch fanout joins --- metaflow/graph.py | 9 ++ metaflow/lint.py | 20 ++-- metaflow/plugins/argo/argo_workflows.py | 2 +- metaflow/plugins/cards/ui/src/types.ts | 2 +- metaflow/runtime.py | 10 +- metaflow/task.py | 13 +- .../unit/test_argo_conditional_input_paths.py | 52 +++++++- test/unit/test_switch_fanout_cases.py | 111 ++++++++++++++++++ 8 files changed, 194 insertions(+), 25 deletions(-) diff --git a/metaflow/graph.py b/metaflow/graph.py index fe295e796e0..4082f33ebbf 100644 --- a/metaflow/graph.py +++ b/metaflow/graph.py @@ -84,6 +84,15 @@ def flatten_switch_cases(switch_cases): 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) # --------------------------------------------------------------------------- diff --git a/metaflow/lint.py b/metaflow/lint.py index a88be88a2ed..debb855570c 100644 --- a/metaflow/lint.py +++ b/metaflow/lint.py @@ -1,6 +1,6 @@ import re from .exception import MetaflowException -from .graph import switch_case_target_lists +from .graph import split_branch_for_node, switch_case_target_lists from .util import all_equal @@ -320,13 +320,13 @@ 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": # 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)] + split_stack + [(node.name, targets)] if len(targets) > 1 else split_stack ) @@ -348,13 +348,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. diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index bd6fa227db8..19180091d82 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/metaflow/plugins/argo/argo_workflows.py @@ -1032,7 +1032,7 @@ def _parse_conditional_branches(self): for targets in switch_case_target_lists(node.switch_cases): if len(targets) > 1: raise ArgoWorkflowsException( - "Step *%s* uses a list-valued switch case (fanout), " + "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 diff --git a/metaflow/plugins/cards/ui/src/types.ts b/metaflow/plugins/cards/ui/src/types.ts index 241dd987ea5..011048d5ac7 100644 --- a/metaflow/plugins/cards/ui/src/types.ts +++ b/metaflow/plugins/cards/ui/src/types.ts @@ -83,7 +83,7 @@ export interface DagStep { failed?: boolean; num_failed?: number; condition?: string; - switch_cases?: Record; + switch_cases?: Record; pathToStep?: string; connections?: string[]; } diff --git a/metaflow/runtime.py b/metaflow/runtime.py index a9e13f103f2..e9fdf5743b2 100644 --- a/metaflow/runtime.py +++ b/metaflow/runtime.py @@ -47,7 +47,7 @@ from .debug import debug from .decorators import flow_decorators from .flowspec import FlowStateItems -from .graph import switch_case_target_lists +from .graph import split_branch_for_node, switch_case_target_lists from .mflog import mflog, RUNTIME_LOG_SOURCE from .util import to_unicode, compress_list, unicode_type, get_latest_task_pathspec from .clone_util import clone_task_helper @@ -1330,13 +1330,7 @@ def siblings(foreach_stack): def _switch_case_targets_for_task(self, switch_node, step_name): node = self._graph[step_name] - branch_root = None - if ( - node.split_parents - and node.split_parents[-1] == switch_node.name - and node.split_branches - ): - branch_root = node.split_branches[-1] + branch_root = split_branch_for_node(node, switch_node.name) for targets in switch_case_target_lists(switch_node.switch_cases): if branch_root in targets or step_name in targets: diff --git a/metaflow/task.py b/metaflow/task.py index 8071c1652d6..40df325c16b 100644 --- a/metaflow/task.py +++ b/metaflow/task.py @@ -17,7 +17,7 @@ from .metaflow_profile import from_start from .mflog import TASK_LOG_SOURCE from .datastore import Inputs, TaskDataStoreSet -from .graph import switch_case_target_lists +from .graph import split_branch_for_node, switch_case_target_lists from .exception import ( MetaflowInternalError, MetaflowDataMissing, @@ -830,12 +830,11 @@ def run_step( branch_roots = set() for in_func in node.in_funcs: in_node = self.flow._graph[in_func] - if ( - in_node.split_parents - and in_node.split_parents[-1] == split_node.name - and in_node.split_branches - ): - branch_roots.add(in_node.split_branches[-1]) + branch_root = split_branch_for_node( + in_node, split_node.name + ) + if branch_root is not None: + branch_roots.add(branch_root) expected_counts = set( len(targets) diff --git a/test/unit/test_argo_conditional_input_paths.py b/test/unit/test_argo_conditional_input_paths.py index a52217aebcc..ad0a7eb8540 100644 --- a/test/unit/test_argo_conditional_input_paths.py +++ b/test/unit/test_argo_conditional_input_paths.py @@ -4,7 +4,10 @@ import pytest from metaflow import FlowSpec, step -from metaflow.plugins.argo.argo_workflows import ArgoWorkflows +from metaflow.plugins.argo.argo_workflows import ( + ArgoWorkflows, + ArgoWorkflowsException, +) from metaflow.plugins.argo.conditional_input_paths import generate_input_paths from metaflow.util import compress_list, decompress_list @@ -31,6 +34,32 @@ def end(self): pass +class SwitchFanoutFlow(FlowSpec): + @step + def start(self): + self.route = "fanout" + self.next( + {"fanout": (self.left, self.right), "linear": self.end}, + condition="route", + ) + + @step + def left(self): + self.next(self.join) + + @step + def right(self): + self.next(self.join) + + @step + def join(self, inputs): + self.next(self.end) + + @step + def end(self): + pass + + @pytest.fixture def chain_skip_argo(mocker): mocker.patch.object(ArgoWorkflows, "_compile_workflow_template", return_value=None) @@ -114,3 +143,24 @@ def test_generate_input_paths_filters_by_exact_step_name( result = generate_input_paths(_encode_input_paths(paths), skippable_steps) assert _decode_input_paths(result) == expected + + +def test_argo_rejects_multi_target_switch_case(mocker): + mocker.patch.object(ArgoWorkflows, "_compile_workflow_template", return_value=None) + mocker.patch.object(ArgoWorkflows, "_compile_sensor", return_value=None) + with pytest.raises(ArgoWorkflowsException, match="multi-target switch case"): + ArgoWorkflows( + name="switch-fanout", + graph=SwitchFanoutFlow._graph, + flow=SwitchFanoutFlow(use_cli=False), + code_package_metadata={}, + code_package_sha="sha", + code_package_url="s3://metaflow/switch-fanout", + production_token="token", + metadata=None, + flow_datastore=None, + environment=None, + event_logger=None, + monitor=None, + username="test-user", + ) diff --git a/test/unit/test_switch_fanout_cases.py b/test/unit/test_switch_fanout_cases.py index f411905f4e6..eae33035795 100644 --- a/test/unit/test_switch_fanout_cases.py +++ b/test/unit/test_switch_fanout_cases.py @@ -1,3 +1,8 @@ +import os +import subprocess +import sys +import textwrap + import pytest from metaflow import FlowSpec, step @@ -5,6 +10,77 @@ from metaflow.lint import LintWarn, linter +RUNTIME_FLOW = r""" +from metaflow import FlowSpec, Parameter, step + + +class SwitchFanoutRuntimeFlow(FlowSpec): + route = Parameter("route") + + @step + def start(self): + self.next( + {"a": (self.a_split, self.a_foreach), + "b": [self.b_one, self.b_two, self.b_three]}, + condition="route", + ) + + @step + def a_split(self): + self.next(self.a_left, self.a_right) + + @step + def a_left(self): + self.next(self.a_split_join) + + @step + def a_right(self): + self.next(self.a_split_join) + + @step + def a_split_join(self, inputs): + self.next(self.shared_join) + + @step + def a_foreach(self): + self.values = [1, 2] + self.next(self.a_worker, foreach="values") + + @step + def a_worker(self): + self.next(self.a_foreach_join) + + @step + def a_foreach_join(self, inputs): + self.next(self.shared_join) + + @step + def b_one(self): + self.next(self.shared_join) + + @step + def b_two(self): + self.next(self.shared_join) + + @step + def b_three(self): + self.next(self.shared_join) + + @step + def shared_join(self, inputs): + self.input_count = sum(1 for _ in inputs) + self.next(self.end) + + @step + def end(self): + print("INPUT_COUNT=%d" % self.input_count) + + +if __name__ == "__main__": + SwitchFanoutRuntimeFlow() +""" + + class SwitchFanoutCaseFlow(FlowSpec): @step def start(self): @@ -107,3 +183,38 @@ def test_runtime_switch_fanout_rejects_empty_case(): def test_switch_fanout_case_cannot_join_at_terminal_step(): with pytest.raises(LintWarn, match="terminal step .* should not be a join step"): linter.run_checks(SwitchFanoutToEndJoinFlow._graph) + + +@pytest.mark.parametrize("route, expected_count", [("a", 2), ("b", 3)]) +def test_switch_fanout_executes_nested_splits_and_shared_join( + tmp_path, route, expected_count +): + flow_file = tmp_path / "switch_fanout_runtime_flow.py" + flow_file.write_text(textwrap.dedent(RUNTIME_FLOW)) + env = os.environ.copy() + env["METAFLOW_USER"] = "switch-fanout-test" + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + env["PYTHONPATH"] = os.pathsep.join( + path for path in (repo_root, env.get("PYTHONPATH")) if path + ) + + result = subprocess.run( + [ + sys.executable, + str(flow_file), + "--datastore=local", + "--metadata=local", + "run", + "--route", + route, + ], + cwd=str(tmp_path), + env=env, + capture_output=True, + text=True, + timeout=60, + ) + + output = result.stdout + result.stderr + assert result.returncode == 0, output + assert "INPUT_COUNT=%d" % expected_count in output From 69cb4e5fc1ac54bf6312753b67706e63996e80a0 Mon Sep 17 00:00:00 2001 From: npow Date: Mon, 10 Aug 2026 18:52:31 +0000 Subject: [PATCH 4/6] Reject overlapping switch fanout targets --- metaflow/flowspec.py | 3 +- metaflow/lint.py | 26 ++++++++++++++++ test/unit/test_switch_fanout_cases.py | 45 +++++++++++++++++++++++++++ 3 files changed, 73 insertions(+), 1 deletion(-) diff --git a/metaflow/flowspec.py b/metaflow/flowspec.py index ce1ffedbcad..58bab63240a 100644 --- a/metaflow/flowspec.py +++ b/metaflow/flowspec.py @@ -939,7 +939,8 @@ def next(self, *dsts: Callable[..., None], **kwargs) -> None: 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. + after the condition has selected that case. Multi-target cases must not share + target steps with other cases. Parameters ---------- diff --git a/metaflow/lint.py b/metaflow/lint.py index debb855570c..d1cbb4e2f77 100644 --- a/metaflow/lint.py +++ b/metaflow/lint.py @@ -410,6 +410,11 @@ def check_switch_splits(graph): ) 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": @@ -437,6 +442,27 @@ def check_switch_splits(graph): node.source_file, ) + cases = [ + (case_value, set(targets)) + for case_value, targets in zip( + node.switch_cases, switch_case_target_lists(node.switch_cases) + ) + ] + 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 diff --git a/test/unit/test_switch_fanout_cases.py b/test/unit/test_switch_fanout_cases.py index eae33035795..256349c17b0 100644 --- a/test/unit/test_switch_fanout_cases.py +++ b/test/unit/test_switch_fanout_cases.py @@ -137,6 +137,43 @@ def end(self, inputs): pass +class SwitchFanoutOverlappingTargetsFlow(FlowSpec): + @step + def start(self): + self.route = "b" + self.next( + { + "a": [self.shared, self.a_only], + "b": [self.shared, self.b_one, self.b_two], + }, + condition="route", + ) + + @step + def shared(self): + self.next(self.join_case) + + @step + def a_only(self): + self.next(self.join_case) + + @step + def b_one(self): + self.next(self.join_case) + + @step + def b_two(self): + self.next(self.join_case) + + @step + def join_case(self, inputs): + self.next(self.end) + + @step + def end(self): + pass + + def test_graph_parses_switch_fanout_case(): graph = SwitchFanoutCaseFlow._graph @@ -185,6 +222,14 @@ def test_switch_fanout_case_cannot_join_at_terminal_step(): linter.run_checks(SwitchFanoutToEndJoinFlow._graph) +def test_switch_fanout_cases_cannot_share_targets(): + with pytest.raises( + LintWarn, + match="multi-target switch cases .* share target step.*shared", + ): + linter.run_checks(SwitchFanoutOverlappingTargetsFlow._graph) + + @pytest.mark.parametrize("route, expected_count", [("a", 2), ("b", 3)]) def test_switch_fanout_executes_nested_splits_and_shared_join( tmp_path, route, expected_count From 9738e9136cdc518b22b6c02e7486496bc65392f8 Mon Sep 17 00:00:00 2001 From: Shashank Srikanth Date: Wed, 12 Aug 2026 20:13:49 +0000 Subject: [PATCH 5/6] Fix switch fanout case validation --- metaflow/flowspec.py | 5 +- metaflow/lint.py | 25 ++- metaflow/runtime.py | 2 +- metaflow/task.py | 69 +++--- test/unit/flows/switch_fanout_runtime_flow.py | 69 ++++++ .../unit/test_argo_conditional_input_paths.py | 22 ++ test/unit/test_switch_fanout_cases.py | 212 ++++++++++-------- 7 files changed, 256 insertions(+), 148 deletions(-) create mode 100644 test/unit/flows/switch_fanout_runtime_flow.py diff --git a/metaflow/flowspec.py b/metaflow/flowspec.py index 58bab63240a..f305b86f087 100644 --- a/metaflow/flowspec.py +++ b/metaflow/flowspec.py @@ -980,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) diff --git a/metaflow/lint.py b/metaflow/lint.py index d1cbb4e2f77..7b608afdd14 100644 --- a/metaflow/lint.py +++ b/metaflow/lint.py @@ -1,6 +1,10 @@ import re from .exception import MetaflowException -from .graph import split_branch_for_node, switch_case_target_lists +from .graph import ( + split_branch_for_node, + switch_case_targets, + switch_case_target_lists, +) from .util import all_equal @@ -213,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, ' + "\"fanout\": [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." ) @@ -405,8 +410,8 @@ 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." @@ -418,10 +423,10 @@ def check_switch_splits(graph): 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, ) @@ -443,10 +448,8 @@ def check_switch_splits(graph): ) cases = [ - (case_value, set(targets)) - for case_value, targets in zip( - node.switch_cases, switch_case_target_lists(node.switch_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]: diff --git a/metaflow/runtime.py b/metaflow/runtime.py index e9fdf5743b2..7479c62eb65 100644 --- a/metaflow/runtime.py +++ b/metaflow/runtime.py @@ -1333,7 +1333,7 @@ def _switch_case_targets_for_task(self, switch_node, step_name): branch_root = split_branch_for_node(node, switch_node.name) for targets in switch_case_target_lists(switch_node.switch_cases): - if branch_root in targets or step_name in targets: + if branch_root in targets: return targets return None diff --git a/metaflow/task.py b/metaflow/task.py index 40df325c16b..fd5e1c4006b 100644 --- a/metaflow/task.py +++ b/metaflow/task.py @@ -36,6 +36,28 @@ MAX_FOREACH_PATH_LENGTH = 256 +def _switch_case_targets_for_input_steps(graph, switch_node, input_step_names): + """Resolve one switch case from the branch roots of actual join inputs.""" + branch_roots = [ + split_branch_for_node(graph[input_step_name], switch_node.name) + for input_step_name in input_step_names + ] + unique_roots = set(branch_roots) + if ( + not branch_roots + or None in unique_roots + or len(branch_roots) != len(unique_roots) + ): + return None + + matching_cases = [ + targets + for targets in switch_case_target_lists(switch_node.switch_cases) + if unique_roots.issubset(targets) + ] + return matching_cases[0] if len(matching_cases) == 1 else None + + class MetaflowTask(object): """ MetaflowTask prepares a Flow instance for execution of a single step. @@ -823,50 +845,27 @@ def run_step( if join_type != "foreach": # Find the corresponding split node from the graph. split_node = self.flow._graph[node.split_parents[-1]] - expected_inputs = None - expected_counts = None if split_node.type == "split-switch": - branch_roots = set() - for in_func in node.in_funcs: - in_node = self.flow._graph[in_func] - branch_root = split_branch_for_node( - in_node, split_node.name - ) - if branch_root is not None: - branch_roots.add(branch_root) - - expected_counts = set( - len(targets) - for targets in switch_case_target_lists( - split_node.switch_cases - ) - if len(targets) > 1 - and set(targets).issubset(branch_roots) + case_targets = _switch_case_targets_for_input_steps( + self.flow._graph, + split_node, + (input_ds.step_name for input_ds in inputs), ) - if not expected_counts: - expected_inputs = len(split_node.out_funcs) + if case_targets is None: + raise MetaflowInternalError( + "Join *%s* could not determine exactly one " + "selected switch case of *%s* from its input " + "branches." % (step_name, split_node.name) + ) + expected_inputs = len(case_targets) else: # The number of expected inputs is the number of branches # from that split -- we can't use in_funcs because there may # be more due to split-switch branches that all converge here. expected_inputs = len(split_node.out_funcs) - if expected_counts: - if len(inputs) not in expected_counts: - raise MetaflowDataMissing( - "Join *%s* expected one of %s inputs but only %d " - "inputs were found" - % ( - step_name, - ", ".join( - str(count) - for count in sorted(expected_counts) - ), - len(inputs), - ) - ) - elif len(inputs) != expected_inputs: + if len(inputs) != expected_inputs: raise MetaflowDataMissing( "Join *%s* expected %d inputs but only %d inputs " "were found" % (step_name, expected_inputs, len(inputs)) diff --git a/test/unit/flows/switch_fanout_runtime_flow.py b/test/unit/flows/switch_fanout_runtime_flow.py new file mode 100644 index 00000000000..a70f3d3cd0d --- /dev/null +++ b/test/unit/flows/switch_fanout_runtime_flow.py @@ -0,0 +1,69 @@ +from metaflow import FlowSpec, Parameter, step + + +class SwitchFanoutRuntimeFlow(FlowSpec): + route = Parameter("route") + + @step + def start(self): + self.next( + { + "a": (self.a_split, self.a_foreach), + "b": [self.b_one, self.b_two, self.b_three], + }, + condition="route", + ) + + @step + def a_split(self): + self.next(self.a_left, self.a_right) + + @step + def a_left(self): + self.next(self.a_split_join) + + @step + def a_right(self): + self.next(self.a_split_join) + + @step + def a_split_join(self, inputs): + self.next(self.shared_join) + + @step + def a_foreach(self): + self.values = [1, 2] + self.next(self.a_worker, foreach="values") + + @step + def a_worker(self): + self.next(self.a_foreach_join) + + @step + def a_foreach_join(self, inputs): + self.next(self.shared_join) + + @step + def b_one(self): + self.next(self.shared_join) + + @step + def b_two(self): + self.next(self.shared_join) + + @step + def b_three(self): + self.next(self.shared_join) + + @step + def shared_join(self, inputs): + self.input_count = sum(1 for _ in inputs) + self.next(self.end) + + @step + def end(self): + print("INPUT_COUNT=%d" % self.input_count) + + +if __name__ == "__main__": + SwitchFanoutRuntimeFlow() diff --git a/test/unit/test_argo_conditional_input_paths.py b/test/unit/test_argo_conditional_input_paths.py index ad0a7eb8540..6bfcb920de2 100644 --- a/test/unit/test_argo_conditional_input_paths.py +++ b/test/unit/test_argo_conditional_input_paths.py @@ -4,10 +4,14 @@ import pytest from metaflow import FlowSpec, step +from metaflow.exception import MetaflowException from metaflow.plugins.argo.argo_workflows import ( ArgoWorkflows, ArgoWorkflowsException, ) +from metaflow.plugins.argo.argo_workflows_decorator import ( + ArgoWorkflowsInternalDecorator, +) from metaflow.plugins.argo.conditional_input_paths import generate_input_paths from metaflow.util import compress_list, decompress_list @@ -164,3 +168,21 @@ def test_argo_rejects_multi_target_switch_case(mocker): monitor=None, username="test-user", ) + + +def test_argo_runtime_rejects_multi_target_switch_case(): + flow = SwitchFanoutFlow(use_cli=False) + flow._transition = (["left", "right"], None) + + with pytest.raises( + MetaflowException, + match="selected a switch case that fans out to multiple targets", + ): + ArgoWorkflowsInternalDecorator().task_finished( + "start", + flow, + SwitchFanoutFlow._graph, + is_task_ok=True, + retry_count=0, + max_user_code_retries=0, + ) diff --git a/test/unit/test_switch_fanout_cases.py b/test/unit/test_switch_fanout_cases.py index 256349c17b0..80b88b1435d 100644 --- a/test/unit/test_switch_fanout_cases.py +++ b/test/unit/test_switch_fanout_cases.py @@ -1,87 +1,60 @@ import os -import subprocess -import sys -import textwrap import pytest -from metaflow import FlowSpec, step +from metaflow import FlowSpec, Runner, step from metaflow.flowspec import InvalidNextException from metaflow.lint import LintWarn, linter +from metaflow.task import _switch_case_targets_for_input_steps -RUNTIME_FLOW = r""" -from metaflow import FlowSpec, Parameter, step +SWITCH_FANOUT_RUNTIME_FLOW_FILE = os.path.join( + os.path.dirname(__file__), "flows", "switch_fanout_runtime_flow.py" +) -class SwitchFanoutRuntimeFlow(FlowSpec): - route = Parameter("route") - +class SwitchFanoutCaseFlow(FlowSpec): @step def start(self): + self.route = "miss" self.next( - {"a": (self.a_split, self.a_foreach), - "b": [self.b_one, self.b_two, self.b_three]}, + { + "hit": [self.finalize, self.hit_second, self.hit_third], + "miss": [self.clip, self.face], + }, condition="route", ) @step - def a_split(self): - self.next(self.a_left, self.a_right) - - @step - def a_left(self): - self.next(self.a_split_join) - - @step - def a_right(self): - self.next(self.a_split_join) - - @step - def a_split_join(self, inputs): - self.next(self.shared_join) - - @step - def a_foreach(self): - self.values = [1, 2] - self.next(self.a_worker, foreach="values") - - @step - def a_worker(self): - self.next(self.a_foreach_join) + def finalize(self): + self.next(self.join_case) @step - def a_foreach_join(self, inputs): - self.next(self.shared_join) + def hit_second(self): + self.next(self.join_case) @step - def b_one(self): - self.next(self.shared_join) + def hit_third(self): + self.next(self.join_case) @step - def b_two(self): - self.next(self.shared_join) + def clip(self): + self.next(self.join_case) @step - def b_three(self): - self.next(self.shared_join) + def face(self): + self.next(self.join_case) @step - def shared_join(self, inputs): - self.input_count = sum(1 for _ in inputs) + def join_case(self, inputs): self.next(self.end) @step def end(self): - print("INPUT_COUNT=%d" % self.input_count) - - -if __name__ == "__main__": - SwitchFanoutRuntimeFlow() -""" + pass -class SwitchFanoutCaseFlow(FlowSpec): +class SwitchFanoutToEndJoinFlow(FlowSpec): @step def start(self): self.route = "miss" @@ -96,14 +69,25 @@ def finalize(self): @step def clip(self): - self.next(self.join_miss) + self.next(self.end) @step def face(self): - self.next(self.join_miss) + self.next(self.end) @step - def join_miss(self, inputs): + def end(self, inputs): + pass + + +class SwitchSharedScalarTargetFlow(FlowSpec): + @step + def start(self): + self.route = "a" + self.next({"a": self.shared, "b": self.shared}, condition="route") + + @step + def shared(self): self.next(self.end) @step @@ -111,29 +95,26 @@ def end(self): pass -class SwitchFanoutToEndJoinFlow(FlowSpec): +class SingleCaseSwitchFanoutFlow(FlowSpec): @step def start(self): - self.route = "miss" - self.next( - {"hit": self.finalize, "miss": [self.clip, self.face]}, - condition="route", - ) + self.route = "only" + self.next({"only": [self.left, self.right]}, condition="route") @step - def finalize(self): - self.next(self.end) + def left(self): + self.next(self.join) @step - def clip(self): - self.next(self.end) + def right(self): + self.next(self.join) @step - def face(self): + def join(self, inputs): self.next(self.end) @step - def end(self, inputs): + def end(self): pass @@ -179,23 +160,29 @@ def test_graph_parses_switch_fanout_case(): assert graph["start"].type == "split-switch" assert graph["start"].switch_cases == { - "hit": "finalize", + "hit": ["finalize", "hit_second", "hit_third"], "miss": ["clip", "face"], } - assert graph["start"].out_funcs == ["finalize", "clip", "face"] + assert graph["start"].out_funcs == [ + "finalize", + "hit_second", + "hit_third", + "clip", + "face", + ] assert graph["clip"].split_parents == ["start"] assert graph["clip"].split_branches == ["clip"] assert graph["face"].split_parents == ["start"] assert graph["face"].split_branches == ["face"] - assert graph["join_miss"].type == "join" - assert graph["join_miss"].split_parents == ["start"] + assert graph["join_case"].type == "join" + assert graph["join_case"].split_parents == ["start"] def test_switch_fanout_case_passes_lint(): linter.run_checks(SwitchFanoutCaseFlow._graph) -def test_runtime_switch_fanout_transition_uses_selected_case_targets(): +def test_switch_fanout_transition_uses_selected_case_targets(): flow = SwitchFanoutCaseFlow(use_cli=False) flow._current_step = "start" flow.route = "miss" @@ -208,7 +195,7 @@ def test_runtime_switch_fanout_transition_uses_selected_case_targets(): assert flow._transition == (["clip", "face"], None) -def test_runtime_switch_fanout_rejects_empty_case(): +def test_switch_fanout_transition_rejects_empty_case(): flow = SwitchFanoutCaseFlow(use_cli=False) flow._current_step = "start" flow.route = "miss" @@ -222,6 +209,45 @@ def test_switch_fanout_case_cannot_join_at_terminal_step(): linter.run_checks(SwitchFanoutToEndJoinFlow._graph) +def test_scalar_switch_cases_can_share_target(): + linter.run_checks(SwitchSharedScalarTargetFlow._graph) + + +def test_single_switch_case_is_rejected_even_when_it_has_multiple_targets(): + with pytest.raises(LintWarn, match="1 found, at least 2 required"): + linter.run_checks(SingleCaseSwitchFanoutFlow._graph) + + +@pytest.mark.parametrize( + "input_step_names, expected_targets", + [ + (("clip",), ["clip", "face"]), + (("clip", "face"), ["clip", "face"]), + (("finalize", "hit_second"), ["finalize", "hit_second", "hit_third"]), + (("clip", "finalize"), None), + (("clip", "clip"), None), + ], + ids=[ + "missing-branch", + "complete-case", + "partial-wider-case", + "mixed-cases", + "duplicate-branch", + ], +) +def test_task_resolves_switch_case_from_actual_input_branches( + input_step_names, expected_targets +): + assert ( + _switch_case_targets_for_input_steps( + SwitchFanoutCaseFlow._graph, + SwitchFanoutCaseFlow._graph["start"], + input_step_names, + ) + == expected_targets + ) + + def test_switch_fanout_cases_cannot_share_targets(): with pytest.raises( LintWarn, @@ -234,32 +260,22 @@ def test_switch_fanout_cases_cannot_share_targets(): def test_switch_fanout_executes_nested_splits_and_shared_join( tmp_path, route, expected_count ): - flow_file = tmp_path / "switch_fanout_runtime_flow.py" - flow_file.write_text(textwrap.dedent(RUNTIME_FLOW)) - env = os.environ.copy() - env["METAFLOW_USER"] = "switch-fanout-test" repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) - env["PYTHONPATH"] = os.pathsep.join( - path for path in (repo_root, env.get("PYTHONPATH")) if path - ) - - result = subprocess.run( - [ - sys.executable, - str(flow_file), - "--datastore=local", - "--metadata=local", - "run", - "--route", - route, - ], + env = { + "METAFLOW_USER": "switch-fanout-test", + "PYTHONPATH": os.pathsep.join( + path for path in (repo_root, os.environ.get("PYTHONPATH")) if path + ), + } + with Runner( + SWITCH_FANOUT_RUNTIME_FLOW_FILE, + show_output=False, cwd=str(tmp_path), env=env, - capture_output=True, - text=True, - timeout=60, - ) - - output = result.stdout + result.stderr - assert result.returncode == 0, output - assert "INPUT_COUNT=%d" % expected_count in output + datastore="local", + metadata="local", + file_read_timeout=60, + ).run(route=route) as running: + output = running.stdout + running.stderr + assert running.returncode == 0, output + assert "INPUT_COUNT=%d" % expected_count in output From 91f39e69168b1d944eac8fd9523168a358d35739 Mon Sep 17 00:00:00 2001 From: Shashank Srikanth Date: Wed, 12 Aug 2026 21:01:14 +0000 Subject: [PATCH 6/6] Clean up switch lint handling --- metaflow/lint.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/metaflow/lint.py b/metaflow/lint.py index 7b608afdd14..8fe70054690 100644 --- a/metaflow/lint.py +++ b/metaflow/lint.py @@ -218,7 +218,7 @@ def check_valid_transitions(graph): " • Fan-out: self.next(self.step1, self.step2, ...)\n" " • Foreach: self.next(self.step, foreach='variable')\n" ' • Switch: self.next({{"key": self.step, ' - "\"fanout\": [self.step1, self.step2]}}, condition='variable')\n\n" + "\"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." ) @@ -398,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], [])