diff --git a/metaflow/plugins/argo/argo_workflows.py b/metaflow/plugins/argo/argo_workflows.py index a061eb83e5e..e7c4868b558 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,18 +1244,51 @@ 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): + """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] + 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. 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. + def _build_conditional_wrapper( + self, node, dag_task_parameters, when_parameter=None + ): + """Create a wrapper around a step that may be skipped. - 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 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. + + 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) @@ -1270,12 +1302,12 @@ 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. - # 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: @@ -1307,7 +1339,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() @@ -1319,13 +1354,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 +1361,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 +1385,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 +1404,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( @@ -1448,7 +1487,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 +1569,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,17 +1581,12 @@ 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 - # 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] @@ -1558,20 +1596,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 +1648,18 @@ 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 = [ + # 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 - 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,60 +1670,50 @@ 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 - # `when` clause on the DAG task. + # `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 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. - # 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' ? 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) + ): + # 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( [ - "{{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 @@ -1872,8 +1751,27 @@ 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) + # 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 + if not parameter.payload["name"].startswith( + ("switch-step-value-", "should-run-") + ) + ] templates.append( - Template(sanitized_name) + Template(recursive_template_name) .steps( [ WorkflowStep() @@ -1886,8 +1784,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 +1801,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,22 +1827,31 @@ 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. + # 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( { - "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 + ) } ), ] @@ -2068,6 +1973,9 @@ def build_ancestor_tree(node_groups, switch_ancestors): # (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. + # 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( @@ -2085,106 +1993,75 @@ 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) ) ) + # 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, + foreach_template_name, + ) + templates.append( + self._build_conditional_wrapper( + join_node, + join_parameters, + when_parameter=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) @@ -3032,8 +2909,15 @@ 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(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 +2948,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" + )