From ca381bb97dbf1820fdea8a4ed5da8cfdf6ceff8e Mon Sep 17 00:00:00 2001 From: Sakari Ikonen Date: Wed, 12 Aug 2026 14:30:42 +0300 Subject: [PATCH 1/3] fix argo conditionals issue with new argo --- metaflow/plugins/argo/argo_workflows.py | 20 ++++++--- .../plugins/argo/conditional_input_paths.py | 4 +- .../unit/test_argo_conditional_input_paths.py | 42 ++++++++++++++++++- 3 files changed, 59 insertions(+), 7 deletions(-) diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index b033c6d638f..67411b18dac 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/metaflow/plugins/argo/argo_workflows.py @@ -1230,6 +1230,20 @@ def _skippable_input_steps_in_dag_order(self, node): reverse=True, ) + def _input_path_ref(self, node_name): + sanitized = self._sanitize(node_name) + pred_node = self.graph[node_name] + if self._is_conditional_node(pred_node) or pred_node.type == "split-switch": + return ( + "argo-{{workflow.name}}/%s/" + "{{=tasks['%s']?.outputs?.parameters['task-id'] ?? 'SKIPPED'}}" + % (node_name, sanitized) + ) + return "argo-{{workflow.name}}/%s/{{tasks.%s.outputs.parameters.task-id}}" % ( + node_name, + sanitized, + ) + def _is_recursive_node(self, node): return node.name in self.recursive_nodes @@ -1383,11 +1397,7 @@ def _visit( parameters = [ Parameter("input-paths").value( compress_list( - [ - "argo-{{workflow.name}}/%s/{{tasks.%s.outputs.parameters.task-id}}" - % (n, self._sanitize(n)) - for n in node.in_funcs - ], + [self._input_path_ref(n) for n in node.in_funcs], # NOTE: We set zlibmin to infinite because zlib compression for the Argo input-paths breaks template value substitution. zlibmin=inf, ) diff --git a/metaflow/plugins/argo/conditional_input_paths.py b/metaflow/plugins/argo/conditional_input_paths.py index fda796d1aac..4330fb5ac0d 100644 --- a/metaflow/plugins/argo/conditional_input_paths.py +++ b/metaflow/plugins/argo/conditional_input_paths.py @@ -25,7 +25,9 @@ def generate_input_paths(input_paths, skippable_steps): # strip these out of the list. # all pathspecs of leading steps that executed. - trimmed = [path for path in paths if not "{{" in path] + trimmed = [ + path for path in paths if "{{" not in path and not path.endswith("/SKIPPED") + ] skippable_steps = [step for step in skippable_steps if step] skippable_step_set = set(skippable_steps) diff --git a/test/unit/test_argo_conditional_input_paths.py b/test/unit/test_argo_conditional_input_paths.py index a52217aebcc..d793e767aa6 100644 --- a/test/unit/test_argo_conditional_input_paths.py +++ b/test/unit/test_argo_conditional_input_paths.py @@ -75,6 +75,10 @@ def _unresolved_task_path(step_name): ) +def _skipped_task_path(step_name): + return "%s/%s/SKIPPED" % (RUN_ID, step_name) + + 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"] @@ -91,6 +95,20 @@ def test_chain_skip_fallback_uses_latest_executed_split_switch(chain_skip_argo): assert _decode_input_paths(result) == [_task_path("step2")] +def test_chain_skip_with_skipped_sentinel(chain_skip_argo): + """Same as above but with SKIPPED sentinel (Argo v3.7.11+ behavior).""" + node = chain_skip_argo.graph["end"] + skippable_steps = chain_skip_argo._skippable_input_steps_in_dag_order(node) + + input_paths = _encode_input_paths( + [_task_path("start"), _task_path("step2"), _skipped_task_path("step3")] + ) + + result = generate_input_paths(input_paths, skippable_steps) + + assert _decode_input_paths(result) == [_task_path("step2")] + + @pytest.mark.parametrize( "paths, skippable_steps, expected", [ @@ -105,8 +123,30 @@ def test_chain_skip_fallback_uses_latest_executed_split_switch(chain_skip_argo): [_task_path("branch")], ), ([_task_path("step"), _task_path("step2")], ["step"], [_task_path("step2")]), + ( + [_task_path("start"), _skipped_task_path("branch")], + ["start"], + [_task_path("start")], + ), + ( + [_task_path("start"), _task_path("step2"), _skipped_task_path("step3")], + ["step2", "start"], + [_task_path("step2")], + ), + ( + [_skipped_task_path("left"), _task_path("right")], + [], + [_task_path("right")], + ), + ], + ids=[ + "normal_join", + "non_skippable_executed", + "exact_step_name", + "skipped_non_skippable", + "skipped_with_skippable_fallback", + "skipped_no_skippable_steps", ], - ids=["normal_join", "non_skippable_executed", "exact_step_name"], ) def test_generate_input_paths_filters_by_exact_step_name( paths, skippable_steps, expected From 920f9faeee8a91c66c50e562de7e701be31c9be7 Mon Sep 17 00:00:00 2001 From: Sakari Ikonen Date: Tue, 18 Aug 2026 11:40:34 +0300 Subject: [PATCH 2/3] bump devtool argo to 3.7.11 and implement fix to argo conditionals input-paths issue --- devtools/Tiltfile | 4 +- metaflow/plugins/argo/argo_workflows.py | 354 ++++++++++++++++++++---- 2 files changed, 302 insertions(+), 56 deletions(-) diff --git a/devtools/Tiltfile b/devtools/Tiltfile index a58f2cb8a28..c585f314127 100644 --- a/devtools/Tiltfile +++ b/devtools/Tiltfile @@ -12,8 +12,8 @@ allow_k8s_contexts('minikube') # --------------------------------------------------------------------------- # Version configuration # --------------------------------------------------------------------------- -ARGO_WORKFLOWS_HELM_CHART_VERSION = os.getenv("ARGO_WORKFLOWS_HELM_CHART_VERSION", "0.45.2") -ARGO_WORKFLOWS_IMAGE_TAG = os.getenv("ARGO_WORKFLOWS_IMAGE_TAG", "v3.6.0") +ARGO_WORKFLOWS_HELM_CHART_VERSION = os.getenv("ARGO_WORKFLOWS_HELM_CHART_VERSION", "0.45.8") +ARGO_WORKFLOWS_IMAGE_TAG = os.getenv("ARGO_WORKFLOWS_IMAGE_TAG", "v3.7.11") AIRFLOW_HELM_CHART_VERSION = os.getenv("AIRFLOW_HELM_CHART_VERSION", "1.15.0") AIRFLOW_IMAGE_TAG = os.getenv("AIRFLOW_IMAGE_TAG", "2.10.4") diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index 67411b18dac..a061eb83e5e 100644 --- a/metaflow/plugins/argo/argo_workflows.py +++ b/metaflow/plugins/argo/argo_workflows.py @@ -1031,6 +1031,7 @@ 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 = {} @@ -1042,6 +1043,8 @@ def _visit(node, conditional_branch, conditional_parents=None): # skip regular non-conditional nodes entirely return + had_conditional_parents = bool(conditional_parents) + if node.type == "split-switch": conditional_branch = conditional_branch + [node.name] c_br = node_conditional_branches.get(node.name, []) @@ -1062,7 +1065,9 @@ def _visit(node, conditional_branch, conditional_parents=None): ): self.recursive_nodes.add(node.name) - if conditional_parents and not node.type == "split-switch": + if conditional_parents and ( + node.type != "split-switch" or had_conditional_parents + ): node_conditional_parents[node.name] = conditional_parents conditional_branch = conditional_branch + [node.name] c_br = node_conditional_branches.get(node.name, []) @@ -1232,13 +1237,6 @@ def _skippable_input_steps_in_dag_order(self, node): def _input_path_ref(self, node_name): sanitized = self._sanitize(node_name) - pred_node = self.graph[node_name] - if self._is_conditional_node(pred_node) or pred_node.type == "split-switch": - return ( - "argo-{{workflow.name}}/%s/" - "{{=tasks['%s']?.outputs?.parameters['task-id'] ?? 'SKIPPED'}}" - % (node_name, sanitized) - ) return "argo-{{workflow.name}}/%s/{{tasks.%s.outputs.parameters.task-id}}" % ( node_name, sanitized, @@ -1252,6 +1250,148 @@ 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): + """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. + """ + sanitized = self._sanitize(node.name) + inner_template = self._sanitize("cond-%s" % node.name) + + switch_in_funcs = [ + in_func + for in_func in node.in_funcs + if self.graph[in_func].type == "split-switch" + ] + conditional_preds = [ + 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",) + ] + + # 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. + wrapper_input_params = [] + inner_params = [] + for p in dag_task_parameters: + name = p.payload["name"] + wrapper_input_params.append(Parameter(name)) + if name.startswith("switch-step-value-") or name.startswith("should-run-"): + continue + inner_params.append( + Parameter(name).value("{{inputs.parameters.%s}}" % name) + ) + + # Build when clause for the inner step + when_parts = [] + for sf in switch_in_funcs: + switch_check = "{{inputs.parameters.switch-step-value-%s}} == %s" % ( + self._sanitize(sf), + node.name, + ) + if self._is_conditional_node(self.graph[sf]): + should_run_check = ( + "{{inputs.parameters.should-run-%s}} == true" % self._sanitize(sf) + ) + when_parts.append("(%s && %s)" % (switch_check, should_run_check)) + else: + when_parts.append(switch_check) + for cp in conditional_preds: + if self.graph[cp].type == "split-switch": + continue + when_parts.append( + "{{inputs.parameters.should-run-%s}} == true" % self._sanitize(cp) + ) + inner_when = " || ".join(when_parts) if when_parts else None + + inner_step = ( + WorkflowStep() + .name("inner") + .template(inner_template) + .arguments(Arguments().parameters(inner_params)) + ) + if inner_when: + 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'" + " ? 'true' : 'false'" + } + ), + ] + + # Forward additional outputs based on node type so that + # downstream foreach/switch DAG tasks can reference them. + if node.type == "split-switch": + wrapper_outputs.append( + Parameter("switch-step").valueFrom( + { + "expression": "steps['inner']?.status == 'Succeeded'" + " ? steps['inner'].outputs.parameters['switch-step']" + " : 'SKIPPED'" + } + ) + ) + if node.type == "foreach": + wrapper_outputs.extend( + [ + Parameter("num-splits").valueFrom( + { + "expression": "steps['inner']?.status == 'Succeeded'" + " ? steps['inner'].outputs.parameters['num-splits']" + " : '[]'" + } + ), + Parameter("split-cardinality").valueFrom( + { + "expression": "steps['inner']?.status == 'Succeeded'" + " ? steps['inner'].outputs.parameters['split-cardinality']" + " : '0'" + } + ), + ] + ) + if getattr(node, "parallel_foreach", False): + wrapper_outputs.extend( + [ + Parameter("num-parallel").valueFrom( + { + "expression": "steps['inner']?.status == 'Succeeded'" + " ? steps['inner'].outputs.parameters['num-parallel']" + " : '0'" + } + ), + Parameter("task-id-entropy").valueFrom( + { + "expression": "steps['inner']?.status == 'Succeeded'" + " ? steps['inner'].outputs.parameters['task-id-entropy']" + " : ''" + } + ), + ] + ) + + return ( + Template(sanitized) + .steps([inner_step]) + .inputs(Inputs().parameters(wrapper_input_params)) + .outputs(Outputs().parameters(wrapper_outputs)) + ) + # Visit every node and yield the uber DAGTemplate(s). def _dag_templates(self): def _visit( @@ -1307,6 +1447,9 @@ def _visit( Parameter("input-paths").value("{{inputs.parameters.input-paths}}"), Parameter("split-index").value("{{inputs.parameters.split-index}}"), ] + if self._is_conditional_node(node): + self.wrapped_conditional_nodes.add(node.name) + templates.append(self._build_conditional_wrapper(node, parameters)) dag_task = ( DAGTask(self._sanitize(node.name)) .template(self._sanitize(node.name)) @@ -1394,15 +1537,57 @@ def _visit( ) else: # Every other node needs only input-paths - parameters = [ - Parameter("input-paths").value( - compress_list( - [self._input_path_ref(n) for n in node.in_funcs], - # NOTE: We set zlibmin to infinite because zlib compression for the Argo input-paths breaks template value substitution. - zlibmin=inf, + 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 + ) + 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. + expr_parts = [] + for n in node.in_funcs: + pred = self.graph[n] + sanitized = self._sanitize(n) + path_expr = ( + "'argo-' + workflow.name + '/%s/' " + "+ 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 + ) + expr_parts.append( + "(%s ? %s + ',' : '')" % (ran_check, path_expr) + ) + else: + expr_parts.append("%s + ','" % path_expr) + input_paths_value = "{{=sprig.trimSuffix(',', %s)}}" % " + ".join( + expr_parts ) - ] + else: + input_path_refs = [self._input_path_ref(n) for n in node.in_funcs] + input_paths_value = compress_list( + input_path_refs, + # NOTE: We set zlibmin to infinite because zlib compression for the Argo input-paths breaks template value substitution. + zlibmin=inf, + ) + parameters = [Parameter("input-paths").value(input_paths_value)] # NOTE: Due to limitations with Argo Workflows Parameter size we # can not pass arbitrarily large lists of task id's to join tasks. # Instead we ensure that task id's for foreach tasks can be @@ -1434,6 +1619,43 @@ 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) + 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 + ) + ) + conditional_deps = [ "%s.Succeeded" % self._sanitize(in_func) for in_func in node.in_funcs @@ -1523,6 +1745,7 @@ def build_ancestor_tree(node_groups, switch_ancestors): if node_groups: conditional_deps = [] required_deps = [] + for parent, chains in build_ancestor_tree( node_groups, node_switch_ancestors ).items(): @@ -1556,52 +1779,62 @@ def build_ancestor_tree(node_groups, switch_ancestors): .arguments(Arguments().parameters(parameters)) ) - # Add conditional if this is the first step in a conditional branch - 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: - # 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. - # 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'}})" - % ( - self._sanitize(switch_in_func), - self._sanitize(switch_in_func), - node.name, - ) - for switch_in_func in switch_in_funcs - ] - ) - - non_switch_in_funcs = [ + 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) + templates.append(self._build_conditional_wrapper(node, parameters)) + else: + # Non-wrapped conditional/join nodes keep the original + # `when` clause on the DAG task. + switch_in_funcs = [ in_func for in_func in node.in_funcs - if in_func not in switch_in_funcs + if self.graph[in_func].type == "split-switch" ] - status_when = "" - if non_switch_in_funcs: - status_when = "||".join( + if ( + self._is_conditional_node(node) + 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+. + conditional_when = "||".join( [ - "{{tasks.%s.status}}==Succeeded" - % self._sanitize(in_func) - for in_func in non_switch_in_funcs + "({{=(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) + 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( + [ + "{{tasks.%s.status}}==Succeeded" + % self._sanitize(in_func) + for in_func in non_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 @@ -2052,6 +2285,10 @@ def _container_templates(self): self._is_conditional_join_node(node) or self._many_in_funcs_all_conditional(node) or self._is_conditional_skip_node(node) + or any( + self._is_conditional_node(self.graph[in_func]) + for in_func in node.in_funcs + ) ) and not ( node.type == "join" @@ -2270,6 +2507,10 @@ def _container_templates(self): self._is_conditional_join_node(node) or self._many_in_funcs_all_conditional(node) or self._is_conditional_skip_node(node) + or any( + self._is_conditional_node(self.graph[in_func]) + for in_func in node.in_funcs + ) ) and not ( node.type == "join" and self.graph[node.split_parents[-1]].type == "foreach" @@ -2823,6 +3064,11 @@ 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: + # 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. + template_name = self._sanitize("cond-%s" % node.name) yield ( Template(template_name) # Set @timeout values From 326d90a8a7c455ad90f725ef8b9f1e08cf4ff93d Mon Sep 17 00:00:00 2001 From: Sakari Ikonen Date: Tue, 25 Aug 2026 22:41:30 +0300 Subject: [PATCH 3/3] fix: exclude foreach joins from wrapped_conditional_nodes precompute Foreach-join DAGTasks are built via a separate code path in _dag_templates() that never wraps the join node, so including them in the precomputed wrapped_conditional_nodes set caused their container template to be renamed to a cond-* name the DAGTask never references, producing "template name ... undefined" deploy errors. Co-Authored-By: Claude Sonnet 5 --- metaflow/plugins/argo/argo_workflows.py | 290 ++++++++++++------ .../unit/test_argo_nested_conditional_join.py | 98 +++++- 2 files changed, 283 insertions(+), 105 deletions(-) diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index a061eb83e5e..e9af3303f1d 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 = {} @@ -1190,9 +1189,36 @@ def _cleanup_conditional_status(node_name, seen): ]: _cleanup_conditional_status(node.name, []) + # Precompute which conditional nodes get wrapped in a Steps template + # (see _build_conditional_wrapper) up front, rather than relying on + # incremental registration during the _dag_templates() DAG walk. + # A node can be visited (and thus need to know whether one of its + # *own* in_funcs is wrapped) before that in_func's own wrapping + # decision has been recorded - e.g. a join step that also starts a + # new split-switch is fully processed via whichever incoming branch + # reaches it first in the traversal, which can happen before its + # other sibling branch(es) have been visited at all. + # Foreach joins are excluded: their DAGTask/template are built via + # the dedicated foreach-handling branch of _dag_templates(), which + # never calls _build_conditional_wrapper() for the join node - so + # there is no Steps wrapper template for the container-template + # rename below (see _container_templates()) to point to. + self.wrapped_conditional_nodes = { + node.name + for node in self.graph + if self._is_conditional_node(node) + and not self._is_recursive_node(node) + and not self._is_foreach_join_node(node) + } + def _is_conditional_node(self, node): return node.name in self.conditional_nodes + def _is_foreach_join_node(self, node): + return node.type == "join" and self.graph[node.split_parents[-1]].type == ( + "foreach" + ) + def _is_conditional_skip_node(self, node): return ( self._is_conditional_node(node) @@ -1245,6 +1271,22 @@ def _input_path_ref(self, node_name): def _is_recursive_node(self, node): return node.name in self.recursive_nodes + def _predecessor_ran_expr(self, in_func): + # Wrapped conditional nodes always run their wrapper template to + # completion (and thus always report DAGTask status 'Succeeded'), + # regardless of whether their `inner` step actually executed. So + # a wrapped predecessor's task status can no longer be used to + # determine whether its branch was actually taken - use its + # `should-run` output instead, which reflects the `inner` step's + # actual outcome. Non-wrapped predecessors (e.g. recursive nodes) + # retain their original Skipped/Succeeded task-status semantics. + sanitized = self._sanitize(in_func) + if in_func in self.wrapped_conditional_nodes: + return "tasks['%s']?.outputs?.parameters['should-run'] == 'true'" % ( + sanitized + ) + return "tasks['%s']?.status == 'Succeeded'" % sanitized + def _matching_conditional_join(self, node): # If no earlier conditional join step is found during parsing, # fall back to the graph's terminal step. @@ -1448,7 +1490,6 @@ def _visit( Parameter("split-index").value("{{inputs.parameters.split-index}}"), ] if self._is_conditional_node(node): - self.wrapped_conditional_nodes.add(node.name) templates.append(self._build_conditional_wrapper(node, parameters)) dag_task = ( DAGTask(self._sanitize(node.name)) @@ -1656,122 +1697,157 @@ def _visit( ) ) - conditional_deps = [ - "%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 + def _build_dep_lists(leaf_fn): conditional_deps = [ - "%s.Succeeded" % self._sanitize(in_func) + leaf_fn(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 = [ + leaf_fn(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" ] - required_deps = [] + if self._is_conditional_skip_node( + node + ) or self._many_in_funcs_all_conditional(node): + # skip nodes need unique condition handling + conditional_deps = [ + leaf_fn(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 - ): + # 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) - ) + 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 + 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]) - ] + 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) + # 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, + ) - # 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 - # Greedily build chains: add to a chain if this key is a subset of its first (largest) key - chains = [] - for key, grp in groups: + 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: - if key <= chain[0][0]: - chain.append((key, grp)) - break - else: - chains.append([(key, grp)]) + groups = [ + "({})".format( + " || ".join(leaf_fn(g) for g in grp) + ) + for grp in chain + ] + parts.append("({})".format(" || ".join(groups))) + required_deps.append("&&".join(parts)) - result[parent] = [[g for _, g in chain] for chain in chains] - return result + return required_deps, conditional_deps - if node_groups: - conditional_deps = [] - required_deps = [] + def _format_dep_expr( + required_deps, conditional_deps, req_sep, cond_sep + ): + both_conditions = required_deps and conditional_deps + return "{required}{_and}{conditional}".format( + required=("(%s)" if both_conditions else "%s") + % req_sep.join(required_deps), + _and=" && " if both_conditions else "", + conditional=("(%s)" if both_conditions else "%s") + % cond_sep.join(conditional_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)) + required_deps, conditional_deps = _build_dep_lists( + lambda in_func: "%s.Succeeded" % self._sanitize(in_func) + ) - both_conditions = required_deps and conditional_deps + depends_str = _format_dep_expr( + required_deps, 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), + # Mirror the depends() boolean structure above, but replace + # each leaf's task-status check with a check of the + # predecessor's actual outcome (see _predecessor_ran_expr). + # depends() only tells us that predecessor DAGTasks have + # completed (wrapped conditional nodes always complete + # with status Succeeded, whether or not their `inner` step + # actually ran) - it can no longer tell us whether any + # conditional branch was really taken. This expression is + # used to build a `when` gate for conditional/join nodes + # that don't already get one via switch_in_funcs below. + ran_required_deps, ran_conditional_deps = _build_dep_lists( + self._predecessor_ran_expr ) + conditional_ran_when = None + if ran_required_deps or ran_conditional_deps: + conditional_ran_when = "{{=%s}}" % _format_dep_expr( + ran_required_deps, ran_conditional_deps, " && ", " || " + ) dag_task = ( DAGTask(self._sanitize(node.name)) .depends(depends_str) @@ -1785,7 +1861,6 @@ def build_ancestor_tree(node_groups, switch_ancestors): # 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) templates.append(self._build_conditional_wrapper(node, parameters)) else: # Non-wrapped conditional/join nodes keep the original @@ -1835,6 +1910,17 @@ def build_ancestor_tree(node_groups, switch_ancestors): else conditional_when ) dag_task.when(total_when) + elif conditional_ran_when: + # This node's conditional predecessors are not + # split-switch nodes themselves (e.g. this is a + # join closing out branch/join steps further down + # a conditional chain). Since such predecessors may + # be wrapped conditional nodes whose DAGTask always + # reports 'Succeeded' regardless of whether their + # branch was actually taken, depends() alone is not + # enough to gate execution here - use should-run + # based checks instead. + dag_task.when(conditional_ran_when) dag_tasks.append(dag_task) # End the workflow if we have reached the end of the flow diff --git a/test/unit/test_argo_nested_conditional_join.py b/test/unit/test_argo_nested_conditional_join.py index 5a73ec4152a..b05d2c673e1 100644 --- a/test/unit/test_argo_nested_conditional_join.py +++ b/test/unit/test_argo_nested_conditional_join.py @@ -256,17 +256,36 @@ def _make_argo(mocker, flow_cls, name): ) -def _depends(aw, node_name): - """Return the Argo `depends` string generated for a given step name.""" +def _task(aw, node_name): + """Return the raw Argo DAG task dict generated for a given step name.""" templates = aw._dag_templates() dag = templates[-1].payload["dag"] sanitized = ArgoWorkflows._sanitize(node_name) for task in dag["tasks"]: if task["name"] == sanitized: - return task.get("depends", "") + return task raise AssertionError(f"no DAG task found for step {node_name!r}") +def _depends(aw, node_name): + """Return the Argo `depends` string generated for a given step name.""" + return _task(aw, node_name).get("depends", "") + + +def _when(aw, node_name): + """Return the Argo `when` string generated for a given step name.""" + return _task(aw, node_name).get("when") + + +def _param(aw, node_name, param_name): + """Return the value of an input parameter on a given step's DAG task.""" + task = _task(aw, node_name) + for p in task["arguments"]["parameters"]: + if p["name"] == param_name: + return p["value"] + raise AssertionError(f"no parameter {param_name!r} found for step {node_name!r}") + + @pytest.fixture def nested_alpha_argo(mocker): return _make_argo(mocker, NestedSwitchAlphaFlow, "nested-alpha") @@ -355,3 +374,76 @@ def test_recursive_switch_join_depends_or(recursive_switch_argo): assert aw.matching_conditional_join_dict["start"] == "merge" assert _depends(aw, "merge") == "shortcut.Succeeded || step-c.Succeeded" + + +# ── Regression tests ──────────────────────────────────────────────────────── +# +# Conditional branch nodes are wrapped in a Steps template (see +# `_build_conditional_wrapper`) so that their `task-id` output always +# resolves, even when their branch is skipped. This means a wrapped node's +# own DAGTask always completes with status 'Succeeded', regardless of +# whether its `inner` step actually ran - so depends()'s `X.Succeeded` +# checks can no longer tell a join whether a conditional branch was really +# taken. A closing join whose conditional predecessors aren't themselves +# split-switch nodes (so it never got a `when` clause from the +# switch-based path) must instead gate on each predecessor's `should-run` +# output. + + +def test_nested_switch_alpha_closing_join_when_gate(nested_alpha_argo): + aw = nested_alpha_argo + + assert _when(aw, "inner_join") is None # gated inside its own wrapper instead + when = _when(aw, "outer_join") + assert when is not None + assert "tasks['a-branch-b']?.outputs?.parameters['should-run'] == 'true'" in when + assert "tasks['inner-join']?.outputs?.parameters['should-run'] == 'true'" in when + + +def test_simple_switch_closing_join_when_gate(simple_switch_argo): + aw = simple_switch_argo + + when = _when(aw, "join") + assert when is not None + assert "tasks['left']?.outputs?.parameters['should-run'] == 'true'" in when + assert "tasks['right']?.outputs?.parameters['should-run'] == 'true'" in when + + +def test_sequential_switch_closing_join_when_gate(sequential_switch_argo): + aw = sequential_switch_argo + + # join1 is itself a wrapped conditional node (it also starts a second + # split-switch), so it is gated inside its own wrapper's `inner` step. + assert _when(aw, "join1") is None + when = _when(aw, "join2") + assert when is not None + assert "tasks['down']?.outputs?.parameters['should-run'] == 'true'" in when + assert "tasks['up']?.outputs?.parameters['should-run'] == 'true'" in when + + +def test_recursive_switch_closing_join_when_gate(recursive_switch_argo): + aw = recursive_switch_argo + + when = _when(aw, "merge") + assert when is not None + assert "tasks['shortcut']?.outputs?.parameters['should-run'] == 'true'" in when + assert "tasks['step-c']?.outputs?.parameters['should-run'] == 'true'" in when + + +def test_dual_join_switch_should_run_params_use_output_not_status( + sequential_switch_argo, +): + """join1 both closes the first switch's branches and opens a second one, + so it is a wrapped conditional node whose `inner` step gating depends on + should-run-left/should-run-right input parameters. Regardless of which + of {left, right} the DAG walk finalizes join1's task from first, both + parameters must reference the predecessor's should-run *output*, not a + status fallback (which would be trivially true for wrapped nodes).""" + aw = sequential_switch_argo + + assert _param(aw, "join1", "should-run-left") == ( + "{{tasks.left.outputs.parameters.should-run}}" + ) + assert _param(aw, "join1", "should-run-right") == ( + "{{tasks.right.outputs.parameters.should-run}}" + )