From eee8dfbaad97d74ddc6ad438e51a5cfccfe97a98 Mon Sep 17 00:00:00 2001 From: Shashank Srikanth Date: Tue, 25 Aug 2026 19:55:33 +0000 Subject: [PATCH 1/3] fix: normalize Argo conditional execution --- metaflow/plugins/argo/argo_workflows.py | 508 +++++++----------- .../unit/test_argo_conditional_input_paths.py | 15 + test/unit/test_argo_conditional_wrappers.py | 200 +++++++ .../unit/test_argo_nested_conditional_join.py | 128 ++++- 4 files changed, 538 insertions(+), 313 deletions(-) create mode 100644 test/unit/test_argo_conditional_wrappers.py diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index a061eb83e5e..d66e8b7f21a 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/metaflow/plugins/argo/argo_workflows.py @@ -1031,7 +1031,6 @@ def _parse_conditional_branches(self): self.conditional_join_nodes = set() self.matching_conditional_join_dict = {} self.recursive_nodes = set() - self.wrapped_conditional_nodes = set() node_conditional_parents = {} node_conditional_branches = {} @@ -1245,6 +1244,25 @@ def _input_path_ref(self, node_name): def _is_recursive_node(self, node): return node.name in self.recursive_nodes + def _conditional_control_parameters(self, node): + parameters = [] + for in_func in node.in_funcs: + predecessor = self.graph[in_func] + sanitized = self._sanitize(in_func) + if predecessor.type == "split-switch": + parameters.append( + Parameter("switch-step-value-%s" % sanitized).value( + "{{tasks.%s.outputs.parameters.switch-step}}" % sanitized + ) + ) + if self._is_conditional_node(predecessor) and predecessor.type != "foreach": + parameters.append( + Parameter("should-run-%s" % sanitized).value( + "{{tasks.%s.outputs.parameters.should-run}}" % sanitized + ) + ) + return parameters + def _matching_conditional_join(self, node): # If no earlier conditional join step is found during parsing, # fall back to the graph's terminal step. @@ -1253,10 +1271,9 @@ def _matching_conditional_join(self, node): def _build_conditional_wrapper(self, node, dag_task_parameters): """Build a Steps wrapper template for a conditional node. - The wrapper always runs and always produces outputs (task-id, should-run). - The inner step uses a `when` clause to conditionally skip execution. - This ensures that task-id references from downstream steps always resolve, - even when the conditional branch is not taken. + The wrapper produces stable outputs when the logical node succeeds or is + inactive. If selected execution fails, the wrapper fails and downstream + dependencies prevent consumers from resolving its missing outputs. """ sanitized = self._sanitize(node.name) inner_template = self._sanitize("cond-%s" % node.name) @@ -1270,7 +1287,7 @@ def _build_conditional_wrapper(self, node, dag_task_parameters): in_func for in_func in node.in_funcs if self._is_conditional_node(self.graph[in_func]) - and self.graph[in_func].type not in ("foreach",) + and self.graph[in_func].type != "foreach" ] # Build wrapper input declarations and inner step arguments. @@ -1319,13 +1336,6 @@ def _build_conditional_wrapper(self, node, dag_task_parameters): inner_step.when(inner_when) wrapper_outputs = [ - Parameter("task-id").valueFrom( - { - "expression": "steps['inner']?.status == 'Succeeded'" - " ? steps['inner'].outputs.parameters['task-id']" - " : 'SKIPPED'" - } - ), Parameter("should-run").valueFrom( { "expression": "steps['inner']?.status == 'Succeeded'" @@ -1333,6 +1343,17 @@ def _build_conditional_wrapper(self, node, dag_task_parameters): } ), ] + if not node.parallel_step and node.name != self.graph.end_step: + wrapper_outputs.insert( + 0, + Parameter("task-id").valueFrom( + { + "expression": "steps['inner']?.status == 'Succeeded'" + " ? steps['inner'].outputs.parameters['task-id']" + " : 'SKIPPED'" + } + ), + ) # Forward additional outputs based on node type so that # downstream foreach/switch DAG tasks can reference them. @@ -1346,7 +1367,7 @@ def _build_conditional_wrapper(self, node, dag_task_parameters): } ) ) - if node.type == "foreach": + if node.type == "foreach" and not node.parallel_step: wrapper_outputs.extend( [ Parameter("num-splits").valueFrom( @@ -1365,7 +1386,7 @@ def _build_conditional_wrapper(self, node, dag_task_parameters): ), ] ) - if getattr(node, "parallel_foreach", False): + if node.parallel_step or getattr(node, "parallel_foreach", False): wrapper_outputs.extend( [ Parameter("num-parallel").valueFrom( @@ -1392,6 +1413,51 @@ def _build_conditional_wrapper(self, node, dag_task_parameters): .outputs(Outputs().parameters(wrapper_outputs)) ) + def _build_foreach_join_wrapper( + self, node, dag_task_parameters, should_run_parameter + ): + wrapper_input_params = [ + Parameter(parameter.payload["name"]) for parameter in dag_task_parameters + ] + inner_params = [ + Parameter(parameter.payload["name"]).value( + "{{inputs.parameters.%s}}" % parameter.payload["name"] + ) + for parameter in dag_task_parameters + if parameter.payload["name"] != should_run_parameter + ] + inner_step = ( + WorkflowStep() + .name("inner") + .template(self._sanitize("cond-%s" % node.name)) + .arguments(Arguments().parameters(inner_params)) + .when("{{inputs.parameters.%s}} == true" % should_run_parameter) + ) + return ( + Template(self._sanitize(node.name)) + .steps([inner_step]) + .inputs(Inputs().parameters(wrapper_input_params)) + .outputs( + Outputs().parameters( + [ + Parameter("task-id").valueFrom( + { + "expression": "steps['inner']?.status == 'Succeeded'" + " ? steps['inner'].outputs.parameters['task-id']" + " : 'SKIPPED'" + } + ), + Parameter("should-run").valueFrom( + { + "expression": "steps['inner']?.status == 'Succeeded'" + " ? 'true' : 'false'" + } + ), + ] + ) + ) + ) + # Visit every node and yield the uber DAGTemplate(s). def _dag_templates(self): def _visit( @@ -1448,7 +1514,7 @@ def _visit( Parameter("split-index").value("{{inputs.parameters.split-index}}"), ] if self._is_conditional_node(node): - self.wrapped_conditional_nodes.add(node.name) + parameters.extend(self._conditional_control_parameters(node)) templates.append(self._build_conditional_wrapper(node, parameters)) dag_task = ( DAGTask(self._sanitize(node.name)) @@ -1530,6 +1596,10 @@ def _visit( ] ) + if self._is_conditional_node(node): + parameters.extend(self._conditional_control_parameters(node)) + templates.append(self._build_conditional_wrapper(node, parameters)) + dag_task = ( DAGTask(self._sanitize(node.name)) .template(self._sanitize(node.name)) @@ -1538,10 +1608,7 @@ def _visit( else: # Every other node needs only input-paths has_wrapped_conditional_pred = any( - self._is_conditional_node(self.graph[n]) - and self.graph[n].type not in ("foreach",) - and not self._is_recursive_node(self.graph[n]) - for n in node.in_funcs + self._is_conditional_node(self.graph[n]) for n in node.in_funcs ) if has_wrapped_conditional_pred: # Build input-paths as an Argo expression that only @@ -1558,20 +1625,11 @@ def _visit( "+ tasks['%s'].outputs.parameters['task-id']" % (n, sanitized) ) - if ( - self._is_conditional_node(pred) - and pred.type not in ("foreach",) - and not self._is_recursive_node(pred) - ): - if n in self.wrapped_conditional_nodes: - ran_check = ( - "tasks['%s'].outputs.parameters['should-run'] == 'true'" - % sanitized - ) - else: - ran_check = ( - "tasks['%s']?.status == 'Succeeded'" % sanitized - ) + if self._is_conditional_node(pred): + ran_check = ( + "tasks['%s'].outputs.parameters['should-run'] == 'true'" + % sanitized + ) expr_parts.append( "(%s ? %s + ',' : '')" % (ran_check, path_expr) ) @@ -1619,158 +1677,17 @@ def _visit( ] ) - # Wrap conditional nodes in Steps templates so their - # task-id output is always resolvable (even when skipped). - is_wrapped_conditional = self._is_conditional_node( - node - ) and not self._is_recursive_node(node) + # Wrap conditional nodes so inactive branches expose stable outputs. + is_wrapped_conditional = self._is_conditional_node(node) if is_wrapped_conditional: - for sf in node.in_funcs: - if self.graph[sf].type == "split-switch": - parameters.append( - Parameter( - "switch-step-value-%s" % self._sanitize(sf) - ).value( - "{{tasks.%s.outputs.parameters.switch-step}}" - % self._sanitize(sf) - ) - ) - for cp in node.in_funcs: - if self._is_conditional_node(self.graph[cp]) and self.graph[ - cp - ].type not in ("foreach",): - sanitized_cp = self._sanitize(cp) - if cp in self.wrapped_conditional_nodes: - param_value = ( - "{{tasks.%s.outputs.parameters.should-run}}" - % sanitized_cp - ) - else: - param_value = ( - "{{=tasks['%s']?.status == 'Succeeded'" - " ? 'true' : 'false'}}" % sanitized_cp - ) - parameters.append( - Parameter("should-run-%s" % sanitized_cp).value( - param_value - ) - ) + parameters.extend(self._conditional_control_parameters(node)) - conditional_deps = [ + # Conditional graph nodes are public wrappers: an inactive node + # succeeds, an active successful node succeeds, and an active + # failed node does not. Every graph edge is therefore a barrier. + depends_str = " && ".join( "%s.Succeeded" % self._sanitize(in_func) for in_func in node.in_funcs - if self._is_conditional_node(self.graph[in_func]) - or self.graph[in_func].type == "split-switch" - ] - required_deps = [ - "%s.Succeeded" % self._sanitize(in_func) - for in_func in node.in_funcs - if not self._is_conditional_node(self.graph[in_func]) - and self.graph[in_func].type != "split-switch" - ] - if self._is_conditional_skip_node( - node - ) or self._many_in_funcs_all_conditional(node): - # skip nodes need unique condition handling - conditional_deps = [ - "%s.Succeeded" % self._sanitize(in_func) - for in_func in node.in_funcs - ] - required_deps = [] - - # join steps in_funcs need special handling, as there can be disjoint sets of always-executing and conditional branches. - if node.type == "join" and any( - self._is_conditional_node(self.graph[fn]) for fn in node.in_funcs - ): - - def _split_switch_ancestors(step_name, first_ancestor): - acc = [] - for in_fn in self.graph[step_name].in_funcs: - if self.graph[in_fn].type == "split-switch": - acc.append(in_fn) - if not in_fn == first_ancestor: - acc.extend( - _split_switch_ancestors(in_fn, first_ancestor) - ) - - return acc - - node_groups = {} - node_switch_ancestors = {} - for fn in node.in_funcs: - if self.graph[fn].split_branches: - # This is the latest split in the DAG. - last_split = self.graph[fn].split_branches[-1] - switch_ancestors = _split_switch_ancestors( - fn, node.split_parents[-1] - ) - if switch_ancestors: - node_switch_ancestors[fn] = switch_ancestors - new_funcs = node_groups.get(last_split, []) - new_funcs.append(fn) - node_groups[last_split] = new_funcs - - def build_ancestor_tree(node_groups, switch_ancestors): - result = {} - for parent, children in node_groups.items(): - nodes = [ - n - for g in children - for n in (g if isinstance(g, list) else [g]) - ] - - # Group nodes by their ancestor set - by_anc = defaultdict(list) - for n in nodes: - by_anc[frozenset(switch_ancestors.get(n, []))].append(n) - - # Sort from most specific (most ancestors) to least - groups = sorted( - by_anc.items(), key=lambda x: len(x[0]), reverse=True - ) - - # Greedily build chains: add to a chain if this key is a subset of its first (largest) key - chains = [] - for key, grp in groups: - for chain in chains: - if key <= chain[0][0]: - chain.append((key, grp)) - break - else: - chains.append([(key, grp)]) - - result[parent] = [[g for _, g in chain] for chain in chains] - return result - - if node_groups: - conditional_deps = [] - required_deps = [] - - for parent, chains in build_ancestor_tree( - node_groups, node_switch_ancestors - ).items(): - parts = [] - for chain in chains: - groups = [ - "({})".format( - " || ".join( - "%s.Succeeded" % self._sanitize(g) - for g in grp - ) - ) - for grp in chain - ] - parts.append("({})".format(" || ".join(groups))) - required_deps.append("&&".join(parts)) - - both_conditions = required_deps and conditional_deps - - depends_str = "{required}{_and}{conditional}".format( - required=("(%s)" if both_conditions else "%s") - % " && ".join(required_deps), - _and=" && " if both_conditions else "", - conditional=("(%s)" if both_conditions else "%s") - % " || ".join(conditional_deps), ) dag_task = ( DAGTask(self._sanitize(node.name)) @@ -1781,11 +1698,9 @@ def build_ancestor_tree(node_groups, switch_ancestors): if is_wrapped_conditional: # Create a Steps wrapper template for this conditional - # node. The wrapper always runs (no `when` on the DAG - # task) and always produces a task-id output, preventing - # Argo v3.7.11+ from requeuing downstream tasks that - # reference potentially-skipped predecessors. - self.wrapped_conditional_nodes.add(node.name) + # node. The wrapper always runs after successful graph + # dependencies and emits stable outputs when its inner + # task succeeds or is inactive. templates.append(self._build_conditional_wrapper(node, parameters)) else: # Non-wrapped conditional/join nodes keep the original @@ -1800,15 +1715,24 @@ def build_ancestor_tree(node_groups, switch_ancestors): or self._is_conditional_skip_node(node) or self._is_conditional_join_node(node) ) and switch_in_funcs: - # It is possible that the some of the leading steps did not execute at all. In this case the switch-step output would be missing and needs to be accounted for. - # Use safe navigation (?.) so the expression resolves to nil instead of causing requeuing on Argo v3.7.11+. + # Some non-recursive leading steps may not have executed. + # Use safe navigation for those predecessors so a missing + # switch-step resolves to nil instead of causing requeuing + # on Argo v3.7.11+. conditional_when = "||".join( [ - "({{=(tasks['%s']?.status == 'Succeeded' ? tasks['%s']?.outputs?.parameters['switch-step'] : nil) == '%s'}})" - % ( - self._sanitize(switch_in_func), - self._sanitize(switch_in_func), - node.name, + ( + "{{tasks.%s.outputs.parameters.switch-step}}==%s" + % (self._sanitize(switch_in_func), node.name) + if self._is_recursive_node( + self.graph[switch_in_func] + ) + else "({{=(tasks['%s']?.status == 'Succeeded' ? tasks['%s']?.outputs?.parameters['switch-step'] : nil) == '%s'}})" + % ( + self._sanitize(switch_in_func), + self._sanitize(switch_in_func), + node.name, + ) ) for switch_in_func in switch_in_funcs ] @@ -1872,8 +1796,20 @@ def build_ancestor_tree(node_groups, switch_ancestors): # - 'example-step-internal' which uses the metaflow step executing template 'recursive-example-step' # - 'example-step-recursion' which calls the parent template 'example-step' if switch-step output from 'example-step-internal' matches the condition. sanitized_name = self._sanitize(node.name) + recursive_template_name = ( + self._sanitize("cond-%s" % node.name) + if self._is_conditional_node(node) + else sanitized_name + ) + recursive_parameters = [ + Parameter(parameter.payload["name"]) + for parameter in parameters + if not parameter.payload["name"].startswith( + ("switch-step-value-", "should-run-") + ) + ] templates.append( - Template(sanitized_name) + Template(recursive_template_name) .steps( [ WorkflowStep() @@ -1886,8 +1822,6 @@ def build_ancestor_tree(node_groups, switch_ancestors): "{{inputs.parameters.input-paths}}" ) ] - # Add the additional inputs required by specific node types. - # We do not need to cover joins or @parallel, as a split-switch step can not be either one of these. + ( [ Parameter("split-index").value( @@ -1905,7 +1839,7 @@ def build_ancestor_tree(node_groups, switch_ancestors): [ WorkflowStep() .name("%s-recursion" % sanitized_name) - .template(sanitized_name) + .template(recursive_template_name) .when( "{{steps.%s-internal.outputs.parameters.switch-step}}==%s" % (sanitized_name, node.name) @@ -1931,7 +1865,7 @@ def build_ancestor_tree(node_groups, switch_ancestors): ), ] ) - .inputs(Inputs().parameters(parameters)) + .inputs(Inputs().parameters(recursive_parameters)) .outputs( # NOTE: We try to read the output parameters from the recursive template call first (-recursion), and the internal step second (-internal). # This guarantees that we always get the output parameters of the last recursive step that executed. @@ -1939,14 +1873,20 @@ def build_ancestor_tree(node_groups, switch_ancestors): [ Parameter("task-id").valueFrom( { - "expression": "(steps['%s-recursion']?.outputs ?? steps['%s-internal']?.outputs).parameters['task-id']" - % (sanitized_name, sanitized_name) + "expression": "steps['{0}-recursion']?.status == 'Succeeded'" + " ? steps['{0}-recursion'].outputs.parameters['task-id']" + " : steps['{0}-internal'].outputs.parameters['task-id']".format( + sanitized_name + ) } ), Parameter("switch-step").valueFrom( { - "expression": "(steps['%s-recursion']?.outputs ?? steps['%s-internal']?.outputs).parameters['switch-step']" - % (sanitized_name, sanitized_name) + "expression": "steps['{0}-recursion']?.status == 'Succeeded'" + " ? steps['{0}-recursion'].outputs.parameters['switch-step']" + " : steps['{0}-internal'].outputs.parameters['switch-step']".format( + sanitized_name + ) } ), ] @@ -2085,106 +2025,70 @@ def build_ancestor_tree(node_groups, switch_ancestors): ) ) ) - .outputs( - Outputs().parameters( - [ - # non @parallel tasks set task-ids as outputs - Parameter("task-id").valueFrom( - { - "parameter": "{{tasks.%s.outputs.parameters.task-id}}" - % self._sanitize( - self.graph[node.matching_join].in_funcs[0] - ) - } - if not self._is_conditional_join_node( - self.graph[node.matching_join] - ) - else - # Note: If the nodes leading to the join are conditional, then we need to use an expression to pick the outputs from the task that executed. - # ref for operators: https://github.com/expr-lang/expr/blob/master/docs/language-definition.md - { - "expression": "get((%s)?.parameters, 'task-id')" - % " ?? ".join( - f"tasks['{self._sanitize(func)}']?.outputs" - for func in self.graph[ - node.matching_join - ].in_funcs - ) - } - ), - ] - if not node.parallel_foreach - else [ - # @parallel tasks set `task-id-entropy` and `num-parallel` - # as outputs so task-ids can be derived in the join step. - # Both of these values should be propagated from the - # jobset labels. - Parameter("num-parallel").valueFrom( - { - "parameter": "{{tasks.%s.outputs.parameters.num-parallel}}" - % self._sanitize( - self.graph[node.matching_join].in_funcs[0] - ) - } - ), - Parameter("task-id-entropy").valueFrom( - { - "parameter": "{{tasks.%s.outputs.parameters.task-id-entropy}}" - % self._sanitize( - self.graph[node.matching_join].in_funcs[0] - ) - } - ), - ] - ) - ) .dag(DAGTemplate().fail_fast().tasks(dag_tasks_1)) ) - join_foreach_task = ( - DAGTask(self._sanitize(self.graph[node.matching_join].name)) - .template(self._sanitize(self.graph[node.matching_join].name)) - .depends(f"{foreach_template_name}.Succeeded") - .arguments( - Arguments().parameters( - ( - [ - Parameter("input-paths").value( - "argo-{{workflow.name}}/%s/{{tasks.%s.outputs.parameters.task-id}}" - % (node.name, self._sanitize(node.name)) - ), - Parameter("split-cardinality").value( - "{{tasks.%s.outputs.parameters.split-cardinality}}" - % self._sanitize(node.name) - ), - ] - if not node.parallel_foreach - else [ - Parameter("num-parallel").value( - "{{tasks.%s.outputs.parameters.num-parallel}}" - % self._sanitize(node.name) - ), - Parameter("task-id-entropy").value( - "{{tasks.%s.outputs.parameters.task-id-entropy}}" - % self._sanitize(node.name) - ), - ] - ) - + ( - [ - Parameter("split-index").value( - # TODO : Pass down these parameters to the jobset stuff. - "{{inputs.parameters.split-index}}" - ), - Parameter("root-input-path").value( - "{{inputs.parameters.input-paths}}" - ), - ] - if parent_foreach - else [] - ) + join_node = self.graph[node.matching_join] + join_parameters = ( + [ + Parameter("input-paths").value( + "argo-{{workflow.name}}/%s/{{tasks.%s.outputs.parameters.task-id}}" + % (node.name, self._sanitize(node.name)) + ), + Parameter("split-cardinality").value( + "{{tasks.%s.outputs.parameters.split-cardinality}}" + % self._sanitize(node.name) + ), + ] + if not node.parallel_foreach + else [ + Parameter("num-parallel").value( + "{{tasks.%s.outputs.parameters.num-parallel}}" + % self._sanitize(node.name) + ), + Parameter("task-id-entropy").value( + "{{tasks.%s.outputs.parameters.task-id-entropy}}" + % self._sanitize(node.name) + ), + ] + ) + ( + [ + Parameter("split-index").value( + # TODO : Pass down these parameters to the jobset stuff. + "{{inputs.parameters.split-index}}" + ), + Parameter("root-input-path").value( + "{{inputs.parameters.input-paths}}" + ), + ] + if parent_foreach + else [] + ) + join_depends = f"{foreach_template_name}.Succeeded" + if self._is_conditional_node(join_node): + should_run_parameter = "should-run-%s" % self._sanitize(node.name) + join_parameters.append( + Parameter(should_run_parameter).value( + "{{tasks.%s.outputs.parameters.should-run}}" + % self._sanitize(node.name) ) ) + join_depends = "%s.Succeeded && (%s.Succeeded || %s.Skipped)" % ( + self._sanitize(node.name), + foreach_template_name, + foreach_template_name, + ) + templates.append( + self._build_foreach_join_wrapper( + join_node, join_parameters, should_run_parameter + ) + ) + + join_foreach_task = ( + DAGTask(self._sanitize(join_node.name)) + .template(self._sanitize(join_node.name)) + .depends(join_depends) + .arguments(Arguments().parameters(join_parameters)) ) dag_tasks.append(join_foreach_task) seen.append(self.graph[node.matching_join].name) @@ -3033,7 +2937,11 @@ def _container_templates(self): jobset.worker.environment_variable("TASK_ID_PREFIX", "worker") yield ( - Template(ArgoWorkflows._sanitize(node.name)) + Template( + self._sanitize("cond-%s" % node.name) + if self._is_conditional_node(node) + else self._sanitize(node.name) + ) .resource( "create", jobset.dump(), @@ -3064,7 +2972,7 @@ def _container_templates(self): # The recursive template has the original step name, # this becomes a template within the recursive ones 'steps' template_name = self._sanitize("recursive-%s" % node.name) - elif node.name in self.wrapped_conditional_nodes: + elif self._is_conditional_node(node): # Wrapped conditional nodes have a Steps template that # takes the original name; the container template gets a # "cond-" prefix so the wrapper can reference it. diff --git a/test/unit/test_argo_conditional_input_paths.py b/test/unit/test_argo_conditional_input_paths.py index d793e767aa6..5e7ef15573d 100644 --- a/test/unit/test_argo_conditional_input_paths.py +++ b/test/unit/test_argo_conditional_input_paths.py @@ -50,6 +50,7 @@ def chain_skip_argo(mocker): event_logger=None, monitor=None, username="test-user", + enable_heartbeat_daemon=False, ) @@ -79,6 +80,20 @@ def _skipped_task_path(step_name): return "%s/%s/SKIPPED" % (RUN_ID, step_name) +def _depends(argo, node_name): + dag = argo._dag_templates()[-1].payload["dag"] + sanitized = ArgoWorkflows._sanitize(node_name) + return next( + task.get("depends", "") for task in dag["tasks"] if task["name"] == sanitized + ) + + +def test_chain_skip_dependencies_are_failure_barriers(chain_skip_argo): + assert _depends(chain_skip_argo, "end") == ( + "start.Succeeded && step2.Succeeded && step3.Succeeded" + ) + + def test_chain_skip_fallback_uses_latest_executed_split_switch(chain_skip_argo): node = chain_skip_argo.graph["end"] assert node.in_funcs == ["start", "step2", "step3"] diff --git a/test/unit/test_argo_conditional_wrappers.py b/test/unit/test_argo_conditional_wrappers.py new file mode 100644 index 00000000000..3d7e3834ed0 --- /dev/null +++ b/test/unit/test_argo_conditional_wrappers.py @@ -0,0 +1,200 @@ +import pytest + +from metaflow import FlowSpec, parallel, retry, step +from metaflow.plugins.argo.argo_workflows import ArgoWorkflows + + +class ConditionalForeachFlow(FlowSpec): + @step + def start(self): + self.route = "fanout" + self.next( + {"fanout": self.fanout, "shortcut": self.shortcut}, + condition="route", + ) + + @step + def fanout(self): + self.items = [1, 2] + self.next(self.worker, foreach="items") + + @step + def worker(self): + self.next(self.fanout_join) + + @step + def fanout_join(self, inputs): + self.next(self.outer_join) + + @step + def shortcut(self): + self.next(self.outer_join) + + @step + def outer_join(self): + self.next(self.end) + + @step + def end(self): + pass + + +class ConditionalParallelFlow(FlowSpec): + @step + def start(self): + self.route = "fanout" + self.next( + {"fanout": self.fanout, "shortcut": self.shortcut}, + condition="route", + ) + + @step + def fanout(self): + self.next(self.parallel_worker, num_parallel=2) + + @retry(times=2) + @parallel + @step + def parallel_worker(self): + self.next(self.parallel_join) + + @step + def parallel_join(self, inputs): + self.next(self.outer_join) + + @step + def shortcut(self): + self.next(self.outer_join) + + @step + def outer_join(self): + self.next(self.end) + + @step + def end(self): + pass + + +@pytest.fixture +def conditional_foreach_argo(mocker): + mocker.patch.object(ArgoWorkflows, "_compile_workflow_template", return_value=None) + mocker.patch.object(ArgoWorkflows, "_compile_sensor", return_value=None) + return ArgoWorkflows( + name="conditional-foreach", + graph=ConditionalForeachFlow._graph, + flow=ConditionalForeachFlow(use_cli=False), + code_package_metadata={}, + code_package_sha="sha", + code_package_url="s3://metaflow/test", + production_token="token", + metadata=None, + flow_datastore=None, + environment=None, + event_logger=None, + monitor=None, + username="test-user", + enable_heartbeat_daemon=False, + ) + + +@pytest.fixture +def conditional_parallel_argo(mocker): + mocker.patch.object(ArgoWorkflows, "_compile_workflow_template", return_value=None) + mocker.patch.object(ArgoWorkflows, "_compile_sensor", return_value=None) + return ArgoWorkflows( + name="conditional-parallel", + graph=ConditionalParallelFlow._graph, + flow=ConditionalParallelFlow(use_cli=False), + code_package_metadata={}, + code_package_sha="sha", + code_package_url="s3://metaflow/test", + production_token="token", + metadata=None, + flow_datastore=None, + environment=None, + event_logger=None, + monitor=None, + username="test-user", + enable_heartbeat_daemon=False, + ) + + +def test_conditional_foreach_completion_is_normalized(conditional_foreach_argo): + templates = conditional_foreach_argo._dag_templates() + top_level = templates[-1].payload["dag"]["tasks"] + tasks = {task["name"]: task for task in top_level} + assert tasks["fanout-join"]["depends"] == ( + "fanout.Succeeded && " + "(fanout-foreach-items.Succeeded || fanout-foreach-items.Skipped)" + ) + assert tasks["outer-join"]["depends"] == ( + "fanout-join.Succeeded && shortcut.Succeeded" + ) + + by_name = {template.payload["name"]: template.payload for template in templates} + wrapper = by_name["fanout-join"] + assert wrapper["steps"][0][0]["template"] == "cond-fanout-join" + assert wrapper["steps"][0][0]["when"] == ( + "{{inputs.parameters.should-run-fanout}} == true" + ) + + body_task = by_name["fanout-foreach-items"]["dag"]["tasks"][0] + assert body_task["name"] == "worker" + assert "should-run-fanout" not in { + parameter["name"] for parameter in body_task["arguments"]["parameters"] + } + body_wrapper = by_name["worker"] + assert "should-run-fanout" not in { + parameter["name"] for parameter in body_wrapper["inputs"]["parameters"] + } + assert "when" not in body_wrapper["steps"][0][0] + assert "outputs" not in by_name["fanout-foreach-items"] + + +def test_conditional_parallel_resource_is_wrapped(conditional_parallel_argo): + templates = conditional_parallel_argo._dag_templates() + by_name = {template.payload["name"]: template.payload for template in templates} + wrapper = by_name["parallel-worker"] + assert wrapper["steps"][0][0]["template"] == "cond-parallel-worker" + assert "when" not in wrapper["steps"][0][0] + assert "should-run-fanout" not in { + parameter["name"] for parameter in wrapper["inputs"]["parameters"] + } + assert {parameter["name"] for parameter in wrapper["outputs"]["parameters"]} == { + "should-run", + "num-parallel", + "task-id-entropy", + } + + body_task = by_name["fanout-foreach-parallel"]["dag"]["tasks"][0] + body_parameters = { + parameter["name"]: parameter["value"] + for parameter in body_task["arguments"]["parameters"] + } + assert "should-run-fanout" not in body_parameters + assert body_parameters["retryCount"] == "{{retries}}" + assert body_parameters["jobset-name"] == ( + "js-{{inputs.parameters.task-id-entropy}}{{retries}}" + ) + inner_parameters = { + parameter["name"]: parameter["value"] + for parameter in wrapper["steps"][0][0]["arguments"]["parameters"] + } + assert inner_parameters["retryCount"] == "{{inputs.parameters.retryCount}}" + assert inner_parameters["jobset-name"] == "{{inputs.parameters.jobset-name}}" + assert "outputs" not in by_name["fanout-foreach-parallel"] + + join_wrapper = by_name["parallel-join"] + assert join_wrapper["steps"][0][0]["when"] == ( + "{{inputs.parameters.should-run-fanout}} == true" + ) + + top_level = templates[-1].payload["dag"]["tasks"] + tasks = {task["name"]: task for task in top_level} + assert tasks["parallel-join"]["depends"] == ( + "fanout.Succeeded && " + "(fanout-foreach-parallel.Succeeded || fanout-foreach-parallel.Skipped)" + ) + assert tasks["outer-join"]["depends"] == ( + "parallel-join.Succeeded && shortcut.Succeeded" + ) diff --git a/test/unit/test_argo_nested_conditional_join.py b/test/unit/test_argo_nested_conditional_join.py index 5a73ec4152a..ffa872d32ab 100644 --- a/test/unit/test_argo_nested_conditional_join.py +++ b/test/unit/test_argo_nested_conditional_join.py @@ -5,9 +5,8 @@ on the actual generated Argo `depends` string, not just on the intermediate `conditional_nodes` / `conditional_join_nodes` / `matching_conditional_join_dict` bookkeeping — the bookkeeping can look -right while the emitted `depends` field still uses `&&` between -mutually exclusive branches, which is what actually breaks deployed -workflows (the join gets stuck in `Omitted`). +right while the emitted `depends` field does not match the wrapper +contract used by the deployed workflow. """ import pytest @@ -15,7 +14,6 @@ from metaflow import FlowSpec, step from metaflow.plugins.argo.argo_workflows import ArgoWorkflows - # ── Flows ──────────────────────────────────────────────────────────────────── @@ -230,6 +228,30 @@ def end(self): pass +class RecursiveSwitchInForeachFlow(FlowSpec): + @step + def start(self): + self.items = ["item"] + self.next(self.loop, foreach="items") + + @step + def loop(self): + self.route = "done" + self.next({"done": self.after_loop, "again": self.loop}, condition="route") + + @step + def after_loop(self): + self.next(self.join) + + @step + def join(self, inputs): + self.next(self.end) + + @step + def end(self): + pass + + # ── Fixtures ───────────────────────────────────────────────────────────────── @@ -292,6 +314,13 @@ def recursive_switch_argo(mocker): return _make_argo(mocker, RecursiveSwitchJoinFlow, "recursive-switch") +@pytest.fixture +def recursive_switch_in_foreach_argo(mocker): + return _make_argo( + mocker, RecursiveSwitchInForeachFlow, "recursive-switch-in-foreach" + ) + + # ── Tests ──────────────────────────────────────────────────────────────────── @@ -307,8 +336,8 @@ def test_nested_switch_alpha_order(nested_alpha_argo): assert aw.matching_conditional_join_dict["start"] == "outer_join" assert aw.matching_conditional_join_dict["a_branch_a"] == "inner_join" - assert _depends(aw, "inner_join") == "b-sub-a.Succeeded || b-sub-b.Succeeded" - assert _depends(aw, "outer_join") == "a-branch-b.Succeeded || inner-join.Succeeded" + assert _depends(aw, "inner_join") == "b-sub-a.Succeeded && b-sub-b.Succeeded" + assert _depends(aw, "outer_join") == "a-branch-b.Succeeded && inner-join.Succeeded" def test_nested_switch_reverse_order(nested_reverse_argo): @@ -323,8 +352,8 @@ def test_nested_switch_reverse_order(nested_reverse_argo): assert aw.matching_conditional_join_dict["start"] == "outer_join" assert aw.matching_conditional_join_dict["z_branch_a"] == "inner_join" - assert _depends(aw, "inner_join") == "x-sub-a.Succeeded || x-sub-b.Succeeded" - assert _depends(aw, "outer_join") == "inner-join.Succeeded || z-branch-b.Succeeded" + assert _depends(aw, "inner_join") == "x-sub-a.Succeeded && x-sub-b.Succeeded" + assert _depends(aw, "outer_join") == "inner-join.Succeeded && z-branch-b.Succeeded" def test_simple_switch_regression(simple_switch_argo): @@ -336,7 +365,7 @@ def test_simple_switch_regression(simple_switch_argo): assert "join" in aw.conditional_join_nodes assert aw.matching_conditional_join_dict["start"] == "join" - assert _depends(aw, "join") == "left.Succeeded || right.Succeeded" + assert _depends(aw, "join") == "left.Succeeded && right.Succeeded" def test_sequential_switch_regression(sequential_switch_argo): @@ -344,14 +373,87 @@ def test_sequential_switch_regression(sequential_switch_argo): assert aw.matching_conditional_join_dict["join1"] == "join2" - assert _depends(aw, "join1") == "left.Succeeded || right.Succeeded" - assert _depends(aw, "join2") == "down.Succeeded || up.Succeeded" + assert _depends(aw, "join1") == "left.Succeeded && right.Succeeded" + assert _depends(aw, "join2") == "down.Succeeded && up.Succeeded" -def test_recursive_switch_join_depends_or(recursive_switch_argo): +def test_recursive_switch_join_is_wrapped(recursive_switch_argo): aw = recursive_switch_argo assert "step_b_loop" in aw.recursive_nodes assert aw.matching_conditional_join_dict["start"] == "merge" - assert _depends(aw, "merge") == "shortcut.Succeeded || step-c.Succeeded" + assert _depends(aw, "merge") == "shortcut.Succeeded && step-c.Succeeded" + + templates = aw._dag_templates() + by_name = {template.payload["name"]: template.payload for template in templates} + wrapper = by_name["step-b-loop"] + driver = by_name["cond-step-b-loop"] + assert wrapper["steps"][0][0]["template"] == "cond-step-b-loop" + assert driver["steps"][1][0]["template"] == "cond-step-b-loop" + assert driver["inputs"]["parameters"] == [{"name": "input-paths"}] + assert by_name["step-c"]["steps"][0][0]["when"] == ( + "({{inputs.parameters.switch-step-value-step-b-loop}} == step_c && " + "{{inputs.parameters.should-run-step-b-loop}} == true)" + ) + + outputs = { + parameter["name"]: parameter["valueFrom"]["expression"] + for parameter in driver["outputs"]["parameters"] + } + assert outputs["task-id"] == ( + "steps['step-b-loop-recursion']?.status == 'Succeeded'" + " ? steps['step-b-loop-recursion'].outputs.parameters['task-id']" + " : steps['step-b-loop-internal'].outputs.parameters['task-id']" + ) + assert outputs["switch-step"] == ( + "steps['step-b-loop-recursion']?.status == 'Succeeded'" + " ? steps['step-b-loop-recursion'].outputs.parameters['switch-step']" + " : steps['step-b-loop-internal'].outputs.parameters['switch-step']" + ) + + +def test_recursive_switch_in_foreach_driver_and_exit_condition( + recursive_switch_in_foreach_argo, +): + templates = recursive_switch_in_foreach_argo._dag_templates() + by_name = {template.payload["name"]: template.payload for template in templates} + driver = by_name["loop"] + + assert driver["inputs"]["parameters"] == [ + {"name": "input-paths"}, + {"name": "split-index"}, + ] + assert driver["steps"][0][0]["arguments"]["parameters"] == [ + { + "name": "input-paths", + "value": "{{inputs.parameters.input-paths}}", + }, + { + "name": "split-index", + "value": "{{inputs.parameters.split-index}}", + }, + ] + assert driver["steps"][1][0]["arguments"]["parameters"] == [ + { + "name": "input-paths", + "value": ( + "argo-{{workflow.name}}/loop/" + "{{steps.loop-internal.outputs.parameters.task-id}}" + ), + }, + { + "name": "split-index", + "value": "{{inputs.parameters.split-index}}", + }, + ] + + foreach = by_name["start-foreach-items"] + after_loop = next( + task for task in foreach["dag"]["tasks"] if task["name"] == "after-loop" + ) + assert after_loop["depends"] == "loop.Succeeded" + assert ( + after_loop["when"] + == "{{tasks.loop.outputs.parameters.switch-step}}==after_loop" + ) From d7aeb1c82f90c8cea23677c4e98fd42febd455bd Mon Sep 17 00:00:00 2001 From: Shashank Srikanth Date: Tue, 25 Aug 2026 22:15:58 +0000 Subject: [PATCH 2/3] refactor: simplify Argo conditional wrappers --- metaflow/plugins/argo/argo_workflows.py | 140 ++++++++---------------- 1 file changed, 45 insertions(+), 95 deletions(-) diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index d66e8b7f21a..dc8d9532d25 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/metaflow/plugins/argo/argo_workflows.py @@ -1268,12 +1268,19 @@ def _matching_conditional_join(self, node): # fall back to the graph's terminal step. return self.matching_conditional_join_dict.get(node.name, self.graph.end_step) - def _build_conditional_wrapper(self, node, dag_task_parameters): + def _build_conditional_wrapper( + self, node, dag_task_parameters, when_parameter=None + ): """Build a Steps wrapper template for a conditional node. The wrapper produces stable outputs when the logical node succeeds or is inactive. If selected execution fails, the wrapper fails and downstream dependencies prevent consumers from resolving its missing outputs. + + `when_parameter` overrides the activity condition derived from + `node.in_funcs`. Foreach joins need this because their only graph + predecessor is inside the sub-DAG, so activity is carried by the foreach + split's `should-run` instead. """ sanitized = self._sanitize(node.name) inner_template = self._sanitize("cond-%s" % node.name) @@ -1324,7 +1331,10 @@ def _build_conditional_wrapper(self, node, dag_task_parameters): when_parts.append( "{{inputs.parameters.should-run-%s}} == true" % self._sanitize(cp) ) - inner_when = " || ".join(when_parts) if when_parts else None + if when_parameter is not None: + inner_when = "{{inputs.parameters.%s}} == true" % when_parameter + else: + inner_when = " || ".join(when_parts) if when_parts else None inner_step = ( WorkflowStep() @@ -1413,51 +1423,6 @@ def _build_conditional_wrapper(self, node, dag_task_parameters): .outputs(Outputs().parameters(wrapper_outputs)) ) - def _build_foreach_join_wrapper( - self, node, dag_task_parameters, should_run_parameter - ): - wrapper_input_params = [ - Parameter(parameter.payload["name"]) for parameter in dag_task_parameters - ] - inner_params = [ - Parameter(parameter.payload["name"]).value( - "{{inputs.parameters.%s}}" % parameter.payload["name"] - ) - for parameter in dag_task_parameters - if parameter.payload["name"] != should_run_parameter - ] - inner_step = ( - WorkflowStep() - .name("inner") - .template(self._sanitize("cond-%s" % node.name)) - .arguments(Arguments().parameters(inner_params)) - .when("{{inputs.parameters.%s}} == true" % should_run_parameter) - ) - return ( - Template(self._sanitize(node.name)) - .steps([inner_step]) - .inputs(Inputs().parameters(wrapper_input_params)) - .outputs( - Outputs().parameters( - [ - Parameter("task-id").valueFrom( - { - "expression": "steps['inner']?.status == 'Succeeded'" - " ? steps['inner'].outputs.parameters['task-id']" - " : 'SKIPPED'" - } - ), - Parameter("should-run").valueFrom( - { - "expression": "steps['inner']?.status == 'Succeeded'" - " ? 'true' : 'false'" - } - ), - ] - ) - ) - ) - # Visit every node and yield the uber DAGTemplate(s). def _dag_templates(self): def _visit( @@ -1703,62 +1668,42 @@ def _visit( # task succeeds or is inactive. templates.append(self._build_conditional_wrapper(node, parameters)) else: - # Non-wrapped conditional/join nodes keep the original - # `when` clause on the DAG task. + # Mixed-predecessor joins need no `when`: `depends` already + # requires every predecessor to succeed, making the old status + # arm tautological. All-switch joins retain explicit selection. switch_in_funcs = [ in_func for in_func in node.in_funcs if self.graph[in_func].type == "split-switch" ] if ( - self._is_conditional_node(node) - or self._is_conditional_skip_node(node) - or self._is_conditional_join_node(node) - ) and switch_in_funcs: - # Some non-recursive leading steps may not have executed. - # Use safe navigation for those predecessors so a missing - # switch-step resolves to nil instead of causing requeuing - # on Argo v3.7.11+. - conditional_when = "||".join( - [ - ( - "{{tasks.%s.outputs.parameters.switch-step}}==%s" - % (self._sanitize(switch_in_func), node.name) - if self._is_recursive_node( - self.graph[switch_in_func] - ) - else "({{=(tasks['%s']?.status == 'Succeeded' ? tasks['%s']?.outputs?.parameters['switch-step'] : nil) == '%s'}})" - % ( - self._sanitize(switch_in_func), - self._sanitize(switch_in_func), - node.name, - ) - ) - for switch_in_func in switch_in_funcs - ] - ) - - non_switch_in_funcs = [ - in_func - for in_func in node.in_funcs - if in_func not in switch_in_funcs - ] - status_when = "" - if non_switch_in_funcs: - status_when = "||".join( + self._is_conditional_join_node(node) + and switch_in_funcs + and len(switch_in_funcs) == len(node.in_funcs) + ): + # Preserve #3344's safe-navigation form for non-recursive + # switches: a missing switch output resolves to nil instead + # of causing Argo v3.7.11+ to requeue the task. + dag_task.when( + "||".join( [ - "{{tasks.%s.status}}==Succeeded" - % self._sanitize(in_func) - for in_func in non_switch_in_funcs + ( + "{{tasks.%s.outputs.parameters.switch-step}}==%s" + % (self._sanitize(switch_in_func), node.name) + if self._is_recursive_node( + self.graph[switch_in_func] + ) + else "({{=(tasks['%s']?.status == 'Succeeded' ? tasks['%s']?.outputs?.parameters['switch-step'] : nil) == '%s'}})" + % ( + self._sanitize(switch_in_func), + self._sanitize(switch_in_func), + node.name, + ) + ) + for switch_in_func in switch_in_funcs ] ) - - total_when = ( - f"({status_when}) || ({conditional_when})" - if status_when - else conditional_when ) - dag_task.when(total_when) dag_tasks.append(dag_task) # End the workflow if we have reached the end of the flow @@ -2008,6 +1953,9 @@ def _visit( # (start [sets num-splits]) --> (task-a-foreach-(0,0) [dummy task]) --> (task-a) --> (join) --> (end) # The (task-a-foreach-(0,0) [dummy task]) propagates the values of the `split-index` and the input paths. # to the actual foreach task. + # This sub-DAG deliberately declares no outputs: an inactive + # conditional foreach expands to zero tasks, so output references + # would be unresolvable. templates.append( Template(foreach_template_name) .inputs( @@ -2079,8 +2027,10 @@ def _visit( foreach_template_name, ) templates.append( - self._build_foreach_join_wrapper( - join_node, join_parameters, should_run_parameter + self._build_conditional_wrapper( + join_node, + join_parameters, + when_parameter=should_run_parameter, ) ) From 1c51ab189530dfd66a81100a47b5a207039fd471 Mon Sep 17 00:00:00 2001 From: Shashank Srikanth Date: Tue, 25 Aug 2026 22:59:30 +0000 Subject: [PATCH 3/3] docs: explain Argo conditional wrapper logic --- metaflow/plugins/argo/argo_workflows.py | 86 ++++++++++++++++--------- 1 file changed, 56 insertions(+), 30 deletions(-) diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index dc8d9532d25..e7c4868b558 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/metaflow/plugins/argo/argo_workflows.py @@ -1245,6 +1245,13 @@ def _is_recursive_node(self, node): return node.name in self.recursive_nodes def _conditional_control_parameters(self, node): + """Return parameters used to decide whether this step should run. + + For each switch before this step, pass which branch it selected. For + each conditional step before this one, pass whether it ran or was + skipped. Foreach joins handle this separately because the foreach body + runs in a child DAG. + """ parameters = [] for in_func in node.in_funcs: predecessor = self.graph[in_func] @@ -1271,16 +1278,17 @@ def _matching_conditional_join(self, node): def _build_conditional_wrapper( self, node, dag_task_parameters, when_parameter=None ): - """Build a Steps wrapper template for a conditional node. + """Create a wrapper around a step that may be skipped. - The wrapper produces stable outputs when the logical node succeeds or is - inactive. If selected execution fails, the wrapper fails and downstream - dependencies prevent consumers from resolving its missing outputs. + The wrapper runs after its dependencies succeed. It runs the real step + only when its branch was selected. If the branch was not selected, the + wrapper still succeeds and returns placeholder outputs that downstream + steps can safely reference. If the selected step fails, the wrapper + fails too. - `when_parameter` overrides the activity condition derived from - `node.in_funcs`. Foreach joins need this because their only graph - predecessor is inside the sub-DAG, so activity is carried by the foreach - split's `should-run` instead. + Foreach joins can pass `when_parameter` because the decision to run the + join comes from the foreach split rather than the steps inside its child + DAG. """ sanitized = self._sanitize(node.name) inner_template = self._sanitize("cond-%s" % node.name) @@ -1297,9 +1305,9 @@ def _build_conditional_wrapper( and self.graph[in_func].type != "foreach" ] - # Build wrapper input declarations and inner step arguments. - # The wrapper forwards all original parameters to the inner step, - # and keeps the conditional-control parameters for the when clause. + # The wrapper receives every parameter, but passes only normal step + # inputs to the real template. Branch-selection parameters are used only + # by the wrapper's `when` expression. wrapper_input_params = [] inner_params = [] for p in dag_task_parameters: @@ -1576,11 +1584,9 @@ def _visit( self._is_conditional_node(self.graph[n]) for n in node.in_funcs ) if has_wrapped_conditional_pred: - # Build input-paths as an Argo expression that only - # includes paths from predecessors whose wrapper - # reported should-run == 'true'. This avoids SKIPPED - # task-ids in input-paths entirely, so the downstream - # filter works regardless of metaflow version. + # Include paths only from conditional inputs that actually + # ran. A skipped wrapper returns a placeholder task ID, not + # a real datastore path. expr_parts = [] for n in node.in_funcs: pred = self.graph[n] @@ -1647,9 +1653,10 @@ def _visit( if is_wrapped_conditional: parameters.extend(self._conditional_control_parameters(node)) - # Conditional graph nodes are public wrappers: an inactive node - # succeeds, an active successful node succeeds, and an active - # failed node does not. Every graph edge is therefore a barrier. + # Every conditional step is represented by a wrapper. The + # wrapper succeeds when its branch was skipped or its step + # completed, and fails when a selected step fails. Downstream + # tasks can therefore require every input wrapper to succeed. depends_str = " && ".join( "%s.Succeeded" % self._sanitize(in_func) for in_func in node.in_funcs @@ -1668,9 +1675,12 @@ def _visit( # task succeeds or is inactive. templates.append(self._build_conditional_wrapper(node, parameters)) else: - # Mixed-predecessor joins need no `when`: `depends` already - # requires every predecessor to succeed, making the old status - # arm tautological. All-switch joins retain explicit selection. + # `depends` decides when this join is eligible to run; `when` + # decides whether it runs or is skipped. Because `depends` + # already requires every input task to succeed, checking a + # non-switch input's status again would always be true. Keep + # `when` only when all inputs are switches and the selected + # branch still matters. switch_in_funcs = [ in_func for in_func in node.in_funcs @@ -1681,9 +1691,9 @@ def _visit( and switch_in_funcs and len(switch_in_funcs) == len(node.in_funcs) ): - # Preserve #3344's safe-navigation form for non-recursive - # switches: a missing switch output resolves to nil instead - # of causing Argo v3.7.11+ to requeue the task. + # For a non-recursive switch, treat a missing output as no + # branch match. Safe navigation resolves it to nil instead + # of making Argo v3.7.11+ retry the unresolved expression. dag_task.when( "||".join( [ @@ -1741,11 +1751,18 @@ def _visit( # - 'example-step-internal' which uses the metaflow step executing template 'recursive-example-step' # - 'example-step-recursion' which calls the parent template 'example-step' if switch-step output from 'example-step-internal' matches the condition. sanitized_name = self._sanitize(node.name) + # A conditional recursive step uses its original template + # name for the wrapper. Give the recursive driver a "cond-" + # name so the wrapper can call it. recursive_template_name = ( self._sanitize("cond-%s" % node.name) if self._is_conditional_node(node) else sanitized_name ) + # Declare recursive template inputs using names only. + # Copying caller values would leave `tasks.*` references that + # do not exist inside the recursive template. Branch-selection + # parameters are needed only by the outer wrapper. recursive_parameters = [ Parameter(parameter.payload["name"]) for parameter in parameters @@ -1812,8 +1829,11 @@ def _visit( ) .inputs(Inputs().parameters(recursive_parameters)) .outputs( - # NOTE: We try to read the output parameters from the recursive template call first (-recursion), and the internal step second (-internal). - # This guarantees that we always get the output parameters of the last recursive step that executed. + # If recursion continued, use the recursive call's + # outputs. Otherwise use the current iteration's + # outputs. Check the status explicitly because a + # skipped recursive call may still have an outputs + # object. Outputs().parameters( [ Parameter("task-id").valueFrom( @@ -1953,9 +1973,9 @@ def _visit( # (start [sets num-splits]) --> (task-a-foreach-(0,0) [dummy task]) --> (task-a) --> (join) --> (end) # The (task-a-foreach-(0,0) [dummy task]) propagates the values of the `split-index` and the input paths. # to the actual foreach task. - # This sub-DAG deliberately declares no outputs: an inactive - # conditional foreach expands to zero tasks, so output references - # would be unresolvable. + # Do not declare outputs for this child DAG. When the branch is + # not selected, the fanout creates no child tasks, so there is no + # task output to reference. templates.append( Template(foreach_template_name) .inputs( @@ -2021,6 +2041,9 @@ def _visit( % self._sanitize(node.name) ) ) + # When the fanout was not selected, Argo marks it Skipped. + # That is valid only when the split wrapper succeeded. A + # failure in the selected split must still stop the join. join_depends = "%s.Succeeded && (%s.Succeeded || %s.Skipped)" % ( self._sanitize(node.name), foreach_template_name, @@ -2886,6 +2909,9 @@ def _container_templates(self): jobset.control.environment_variable("TASK_ID_PREFIX", "control") jobset.worker.environment_variable("TASK_ID_PREFIX", "worker") + # A conditional @parallel step uses its original template name + # for the wrapper. Give the underlying JobSet template a + # "cond-" name so the wrapper can call it. yield ( Template( self._sanitize("cond-%s" % node.name)