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 b014eb55c09..c217e99fa8d 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/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index eb5fb1adc3d..b81fb76c82f 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/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) diff --git a/metaflow/plugins/argo/argo_workflows_decorator.py b/metaflow/plugins/argo/argo_workflows_decorator.py index 172d5c77ab4..6521cf2d2bd 100644 --- a/metaflow/plugins/argo/argo_workflows_decorator.py +++ b/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 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/core/graphs/switch_fanout.json b/test/core/graphs/switch_fanout.json new file mode 100644 index 00000000000..c13ddb97e17 --- /dev/null +++ b/test/core/graphs/switch_fanout.json @@ -0,0 +1,19 @@ +{ + "name": "switch_fanout", + "graph": { + "start": { "linear": "switch_step", "quals": ["start"] }, + "switch_step": { + "switch": { + "hit": "finalize", + "miss": ["clip", "face"] + }, + "condition": "condition", + "quals": ["switch-simple"] + }, + "finalize": { "linear": "end", "quals": ["path-hit"] }, + "clip": { "linear": "join_miss", "quals": ["path-clip"] }, + "face": { "linear": "join_miss", "quals": ["path-face"] }, + "join_miss": { "join": true, "linear": "end", "quals": ["fanout-join"] }, + "end": { "quals": ["end"] } + } +} diff --git a/test/core/metaflow_test/formatter.py b/test/core/metaflow_test/formatter.py index f6b505377d1..dce3b5cbaf7 100644 --- a/test/core/metaflow_test/formatter.py +++ b/test/core/metaflow_test/formatter.py @@ -164,13 +164,20 @@ def _flow_lines(self): branches = ",".join("self.%s" % x for x in node["branch"]) yield 2, "self.next(%s)" % branches elif "switch" in node: - # Handle switch nodes - generate the switch dictionary and condition + # Handle switch nodes - generate the switch dictionary and condition. + # A case value can be a string (single target) or a list (fanout). switch_dict = node["switch"] condition = node["condition"] + + def _fmt_case_value(branch): + if isinstance(branch, list): + return "[%s]" % ", ".join("self.%s" % s for s in branch) + return "self.%s" % branch + switch_branches = ( "{" + ", ".join( - '"%s": self.%s' % (key, branch) + '"%s": %s' % (key, _fmt_case_value(branch)) for key, branch in switch_dict.items() ) + "}" diff --git a/test/core/tests/switch_fanout_case.py b/test/core/tests/switch_fanout_case.py new file mode 100644 index 00000000000..ae940e74504 --- /dev/null +++ b/test/core/tests/switch_fanout_case.py @@ -0,0 +1,81 @@ +from metaflow_test import MetaflowTest, steps, assert_equals + + +class SwitchFanoutHitTest(MetaflowTest): + """ + Tests a switch where the selected case is a single-target ('hit' -> finalize). + """ + + PRIORITY = 2 + ONLY_GRAPHS = ["switch_fanout"] + + @steps(0, ["start"], required=True) + def step_start(self): + self.condition = "hit" + + @steps(0, ["switch-simple"], required=True) + def step_switch(self): + pass + + @steps(0, ["path-hit"], required=True) + def step_finalize(self): + self.result = "hit" + + @steps(0, ["path-clip"], required=True) + def step_clip(self): + self.result = "clip" + + @steps(0, ["path-face"], required=True) + def step_face(self): + self.result = "face" + + @steps(0, ["fanout-join"], required=True) + def step_join_miss(self, inputs): + self.result = ",".join(sorted(inp.result for inp in inputs)) + + @steps(1, ["end"], required=True) + def step_end(self): + assert_equals("hit", self.result) + + def check_results(self, flow, checker): + checker.assert_artifact("finalize", "result", "hit") + + +class SwitchFanoutMissTest(MetaflowTest): + """ + Tests a switch where the selected case is a list-valued fanout ('miss' -> [clip, face]). + """ + + PRIORITY = 2 + ONLY_GRAPHS = ["switch_fanout"] + + @steps(0, ["start"], required=True) + def step_start(self): + self.condition = "miss" + + @steps(0, ["switch-simple"], required=True) + def step_switch(self): + pass + + @steps(0, ["path-hit"], required=True) + def step_finalize(self): + self.result = "hit" + + @steps(0, ["path-clip"], required=True) + def step_clip(self): + self.result = "clip" + + @steps(0, ["path-face"], required=True) + def step_face(self): + self.result = "face" + + @steps(0, ["fanout-join"], required=True) + def step_join_miss(self, inputs): + self.result = ",".join(sorted(inp.result for inp in inputs)) + + @steps(1, ["end"], required=True) + def step_end(self): + assert_equals("clip,face", self.result) + + def check_results(self, flow, checker): + checker.assert_artifact("join_miss", "result", "clip,face") 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)