metaflow 2.17.3__py2.py3-none-any.whl → 2.17.5__py2.py3-none-any.whl

metaflow/datastore/task_datastore.py

@@ -222,6 +222,9 @@ class TaskDataStore(object):
     @property
     def pathspec_index(self):
         idxstr = ",".join(map(str, (f.index for f in self["_foreach_stack"])))
+        if "_iteration_stack" in self:
+            itrstr = ",".join(map(str, (f for f in self["_iteration_stack"])))
+            return "%s/%s[%s][%s]" % (self._run_id, self._step_name, idxstr, itrstr)
         return "%s/%s[%s]" % (self._run_id, self._step_name, idxstr)
 
     @property

metaflow/graph.py

@@ -478,7 +478,9 @@ class FlowGraph(object):
                     cur_name = cur_node.matching_join
                 elif node_type == "split-switch":
                     all_paths = [
-                        populate_block(s, end_name) for s in cur_node.out_funcs
+                        populate_block(s, end_name)
+                        for s in cur_node.out_funcs
+                        if s != cur_name
                     ]
                     resulting_list.append(all_paths)
                     cur_name = end_name

metaflow/lint.py

@@ -175,6 +175,8 @@ def check_for_acyclicity(graph):
 
     def check_path(node, seen):
         for n in node.out_funcs:
+            if node.type == "split-switch" and n == node.name:
+                continue
             if n in seen:
                 path = "->".join(seen + [n])
                 raise LintWarn(
@@ -241,6 +243,8 @@ def check_split_join_balance(graph):
         elif node.type == "split-switch":
             # For a switch, continue traversal down each path with the same stack
             for n in node.out_funcs:
+                if node.type == "split-switch" and n == node.name:
+                    continue
                 traverse(graph[n], split_stack)
             return
         elif node.type == "end":
@@ -293,6 +297,8 @@ def check_split_join_balance(graph):
             new_stack = split_stack
 
         for n in node.out_funcs:
+            if node.type == "split-switch" and n == node.name:
+                continue
             traverse(graph[n], new_stack)
 
     traverse(graph["start"], [])

metaflow/plugins/argo/argo_workflows.py

@@ -926,6 +926,7 @@ class ArgoWorkflows(object):
         self.conditional_nodes = set()
         self.conditional_join_nodes = set()
         self.matching_conditional_join_dict = {}
+        self.recursive_nodes = set()
 
         node_conditional_parents = {}
         node_conditional_branches = {}
@@ -948,6 +949,12 @@ class ArgoWorkflows(object):
                 )
                 node_conditional_parents[node.name] = conditional_parents
 
+                # check for recursion. this split is recursive if any of its out functions are itself.
+                if any(
+                    out_func for out_func in node.out_funcs if out_func == node.name
+                ):
+                    self.recursive_nodes.add(node.name)
+
             if conditional_parents and not node.type == "split-switch":
                 node_conditional_parents[node.name] = conditional_parents
                 conditional_branch = conditional_branch + [node.name]
@@ -1033,6 +1040,9 @@ class ArgoWorkflows(object):
     def _is_conditional_join_node(self, node):
         return node.name in self.conditional_join_nodes
 
+    def _is_recursive_node(self, node):
+        return node.name in self.recursive_nodes
+
     def _matching_conditional_join(self, node):
         return self.matching_conditional_join_dict.get(node.name, None)
 
@@ -1044,6 +1054,7 @@ class ArgoWorkflows(object):
             templates=None,
             dag_tasks=None,
             parent_foreach=None,
+            seen=None,
         ):  # Returns Tuple[List[Template], List[DAGTask]]
             """ """
             # Every for-each node results in a separate subDAG and an equivalent
@@ -1053,6 +1064,8 @@ class ArgoWorkflows(object):
             # of the for-each node.
 
             # Emit if we have reached the end of the sub workflow
+            if seen is None:
+                seen = []
             if dag_tasks is None:
                 dag_tasks = []
             if templates is None:
@@ -1060,6 +1073,13 @@ class ArgoWorkflows(object):
 
             if exit_node is not None and exit_node is node.name:
                 return templates, dag_tasks
+            if node.name in seen:
+                return templates, dag_tasks
+
+            seen.append(node.name)
+
+            # helper variable for recursive conditional inputs
+            has_foreach_inputs = False
             if node.name == "start":
                 # Start node has no dependencies.
                 dag_task = DAGTask(self._sanitize(node.name)).template(
@@ -1073,9 +1093,10 @@ class ArgoWorkflows(object):
                 # vs what is a "num_parallel" based foreach (i.e. something that follows gang semantics.)
                 # A `regular` foreach is basically any arbitrary kind of foreach.
             ):
+                # helper variable for recursive conditional inputs
+                has_foreach_inputs = True
                 # Child of a foreach node needs input-paths as well as split-index
                 # This child is the first node of the sub workflow and has no dependency
-
                 parameters = [
                     Parameter("input-paths").value("{{inputs.parameters.input-paths}}"),
                     Parameter("split-index").value("{{inputs.parameters.split-index}}"),
@@ -1253,6 +1274,7 @@ class ArgoWorkflows(object):
                         templates,
                         dag_tasks,
                         parent_foreach,
+                        seen,
                     )
                 return _visit(
                     self.graph[node.matching_join],
@@ -1260,8 +1282,102 @@ class ArgoWorkflows(object):
                     templates,
                     dag_tasks,
                     parent_foreach,
+                    seen,
                 )
             elif node.type == "split-switch":
+                if self._is_recursive_node(node):
+                    # we need an additional recursive template if the step is recursive
+                    # NOTE: in the recursive case, the original step is renamed in the container templates to 'recursive-<step_name>'
+                    # so that we do not have to touch the step references in the DAG.
+                    #
+                    # NOTE: The way that recursion in Argo Workflows is achieved is with the following structure:
+                    # - the usual 'example-step' template which would match example_step in flow code is renamed to 'recursive-example-step'
+                    # - templates has another template with the original task name: 'example-step'
+                    # - the template 'example-step' in turn has steps
+                    #   - '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)
+                    templates.append(
+                        Template(sanitized_name)
+                        .steps(
+                            [
+                                WorkflowStep()
+                                .name("%s-internal" % sanitized_name)
+                                .template("recursive-%s" % sanitized_name)
+                                .arguments(
+                                    Arguments().parameters(
+                                        [
+                                            Parameter("input-paths").value(
+                                                "{{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(
+                                                    "{{inputs.parameters.split-index}}"
+                                                )
+                                            ]
+                                            if has_foreach_inputs
+                                            else []
+                                        )
+                                    )
+                                )
+                            ]
+                        )
+                        .steps(
+                            [
+                                WorkflowStep()
+                                .name("%s-recursion" % sanitized_name)
+                                .template(sanitized_name)
+                                .when(
+                                    "{{steps.%s-internal.outputs.parameters.switch-step}}==%s"
+                                    % (sanitized_name, node.name)
+                                )
+                                .arguments(
+                                    Arguments().parameters(
+                                        [
+                                            Parameter("input-paths").value(
+                                                "argo-{{workflow.name}}/%s/{{steps.%s-internal.outputs.parameters.task-id}}"
+                                                % (node.name, sanitized_name)
+                                            )
+                                        ]
+                                        + (
+                                            [
+                                                Parameter("split-index").value(
+                                                    "{{inputs.parameters.split-index}}"
+                                                )
+                                            ]
+                                            if has_foreach_inputs
+                                            else []
+                                        )
+                                    )
+                                ),
+                            ]
+                        )
+                        .inputs(Inputs().parameters(parameters))
+                        .outputs(
+                            # NOTE: We try to read the output parameters from the recursive template call first (<step>-recursion), and the internal step second (<step>-internal).
+                            # This guarantees that we always get the output parameters of the last recursive step that executed.
+                            Outputs().parameters(
+                                [
+                                    Parameter("task-id").valueFrom(
+                                        {
+                                            "expression": "(steps['%s-recursion']?.outputs ?? steps['%s-internal']?.outputs).parameters['task-id']"
+                                            % (sanitized_name, sanitized_name)
+                                        }
+                                    ),
+                                    Parameter("switch-step").valueFrom(
+                                        {
+                                            "expression": "(steps['%s-recursion']?.outputs ?? steps['%s-internal']?.outputs).parameters['switch-step']"
+                                            % (sanitized_name, sanitized_name)
+                                        }
+                                    ),
+                                ]
+                            )
+                        )
+                    )
                 for n in node.out_funcs:
                     _visit(
                         self.graph[n],
@@ -1269,6 +1385,7 @@ class ArgoWorkflows(object):
                         templates,
                         dag_tasks,
                         parent_foreach,
+                        seen,
                     )
 
                 return _visit(
@@ -1277,6 +1394,7 @@ class ArgoWorkflows(object):
                     templates,
                     dag_tasks,
                     parent_foreach,
+                    seen,
                 )
             # For foreach nodes generate a new sub DAGTemplate
             # We do this for "regular" foreaches (ie. `self.next(self.a, foreach=)`)
@@ -1367,6 +1485,7 @@ class ArgoWorkflows(object):
                     templates,
                     [],
                     node.name,
+                    seen,
                 )
 
                 # How do foreach's work on Argo:
@@ -1500,6 +1619,7 @@ class ArgoWorkflows(object):
                     templates,
                     dag_tasks,
                     parent_foreach,
+                    seen,
                 )
             # For linear nodes continue traversing to the next node
             if node.type in ("linear", "join", "start"):
@@ -1509,6 +1629,7 @@ class ArgoWorkflows(object):
                     templates,
                     dag_tasks,
                     parent_foreach,
+                    seen,
                 )
             else:
                 raise ArgoWorkflowsException(
@@ -2290,8 +2411,13 @@ class ArgoWorkflows(object):
                     )
                 )
             else:
+                template_name = self._sanitize(node.name)
+                if self._is_recursive_node(node):
+                    # 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)
                 yield (
-                    Template(self._sanitize(node.name))
+                    Template(template_name)
                     # Set @timeout values
                     .active_deadline_seconds(run_time_limit)
                     # Set service account
@@ -3750,6 +3876,10 @@ class WorkflowStep(object):
         self.payload["template"] = str(template)
         return self
 
+    def arguments(self, arguments):
+        self.payload["arguments"] = arguments.to_json()
+        return self
+
     def when(self, condition):
         self.payload["when"] = str(condition)
         return self