metaflow 2.17.4__tar.gz → 2.18.0__tar.gz

{metaflow-2.17.4/metaflow.egg-info → metaflow-2.18.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metaflow
-Version: 2.17.4
+Version: 2.18.0
 Summary: Metaflow: More AI and ML, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -26,7 +26,7 @@ License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: boto3
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.17.4; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.18.0; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

{metaflow-2.17.4 → metaflow-2.18.0}/metaflow/lint.py

@@ -416,3 +416,25 @@ def check_nested_foreach(graph):
         if node.type == "foreach":
             if any(graph[p].type == "foreach" for p in node.split_parents):
                 raise LintWarn(msg.format(node), node.func_lineno, node.source_file)
+
+
+@linter.ensure_static_graph
+@linter.check
+def check_ambiguous_joins(graph):
+    for node in graph:
+        if node.type == "join":
+            problematic_parents = [
+                p_name
+                for p_name in node.in_funcs
+                if graph[p_name].type == "split-switch"
+            ]
+            if problematic_parents:
+                msg = (
+                    "A conditional path cannot lead directly to a join step.\n"
+                    "In your conditional step(s) {parents}, one or more of the possible paths transition directly to the join step {join_name}.\n"
+                    "As a workaround, please introduce an intermediate, unconditional step on that specific path before joining."
+                ).format(
+                    parents=", ".join("*%s*" % p for p in problematic_parents),
+                    join_name="*%s*" % node.name,
+                )
+                raise LintWarn(msg, node.func_lineno, node.source_file)

{metaflow-2.17.4 → metaflow-2.18.0}/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(
@@ -1773,8 +1894,10 @@ class ArgoWorkflows(object):
                 # foreach-joins straight out of conditional branches are not yet supported
                 if self._is_conditional_join_node(node):
                     raise ArgoWorkflowsException(
-                        "Conditionals steps that transition directly into a join step are not currently supported. "
-                        "As a workaround, you can add a normal step after the conditional steps that transitions to a join step."
+                        "Conditional steps inside a foreach that transition directly into a join step are not currently supported.\n"
+                        "As a workaround, add a common step after the conditional steps %s "
+                        "that will transition to a join."
+                        % ", ".join("*%s*" % f for f in node.in_funcs)
                     )
                 # Set aggregated input-paths for a for-each join
                 foreach_step = next(
@@ -2290,8 +2413,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 +3878,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

{metaflow-2.17.4 → metaflow-2.18.0}/metaflow/plugins/package_cli.py

@@ -60,7 +60,7 @@ def list(obj, archive=False):
 @click.pass_obj
 def save(obj, path):
     with open(path, "wb") as f:
-        f.write(obj.package.blob())
+        f.write(obj.package.blob)
     obj.echo(
         "Code package saved in *%s* with metadata: %s"
         % (path, obj.package.package_metadata),

{metaflow-2.17.4 → metaflow-2.18.0}/metaflow/task.py

@@ -121,8 +121,10 @@ class MetaflowTask(object):
                 r = w.post_step(orig_step_func.name, self.flow, raised_exception)
             except Exception as ex:
                 r = ex
-            if r is None or isinstance(r, Exception):
+            if r is None:
                 raised_exception = None
+            elif isinstance(r, Exception):
+                raised_exception = r
             elif isinstance(r, tuple):
                 if len(r) == 2:
                     raised_exception, fake_next_call_args = r

{metaflow-2.17.4 → metaflow-2.18.0}/metaflow/user_configs/config_parameters.py

@@ -271,6 +271,9 @@ class ConfigValue(collections.abc.Mapping, dict):
             v = obj
         return v
 
+    def __reduce__(self):
+        return (self.__class__, (self.to_dict(),))
+
 
 class DelayEvaluator(collections.abc.Mapping):
     """
@@ -484,7 +487,6 @@ class Config(Parameter, collections.abc.Mapping):
         parser: Optional[Union[str, Callable[[str], Dict[Any, Any]]]] = None,
         **kwargs: Dict[str, str]
     ):
-
         if default is not None and default_value is not None:
             raise MetaflowException(
                 "For config '%s', you can only specify default or default_value, not both"

metaflow-2.18.0/metaflow/version.py

@@ -0,0 +1 @@
+metaflow_version = "2.18.0"

{metaflow-2.17.4 → metaflow-2.18.0/metaflow.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: metaflow
-Version: 2.17.4
+Version: 2.18.0
 Summary: Metaflow: More AI and ML, Less Engineering
 Author: Metaflow Developers
 Author-email: help@metaflow.org
@@ -26,7 +26,7 @@ License-File: LICENSE
 Requires-Dist: requests
 Requires-Dist: boto3
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.17.4; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.18.0; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: classifier

{metaflow-2.17.4 → metaflow-2.18.0}/metaflow.egg-info/requires.txt

@@ -2,4 +2,4 @@ requests
 boto3
 
 [stubs]
-metaflow-stubs==2.17.4
+metaflow-stubs==2.18.0

metaflow-2.17.4/metaflow/version.py

@@ -1 +0,0 @@
-metaflow_version = "2.17.4"