ob-metaflow 2.17.3.1__py2.py3-none-any.whl → 2.18.0.1__py2.py3-none-any.whl

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/runtime.py

@@ -15,11 +15,13 @@ import tempfile
 import time
 import subprocess
 from datetime import datetime
+from enum import Enum
 from io import BytesIO
 from itertools import chain
 from functools import partial
 from concurrent import futures
 
+from typing import Dict, Tuple
 from metaflow.datastore.exceptions import DataException
 from contextlib import contextmanager
 
@@ -60,6 +62,7 @@ PROGRESS_INTERVAL = 300  # s
 # leveraging the TaskDataStoreSet.
 PREFETCH_DATA_ARTIFACTS = [
     "_foreach_stack",
+    "_iteration_stack",
     "_task_ok",
     "_transition",
     "_control_mapper_tasks",
@@ -67,6 +70,14 @@ PREFETCH_DATA_ARTIFACTS = [
 ]
 RESUME_POLL_SECONDS = 60
 
+
+class LoopBehavior(Enum):
+    NONE = "none"
+    ENTERING = "entering"
+    EXITING = "exiting"
+    LOOPING = "looping"
+
+
 # Runtime must use logsource=RUNTIME_LOG_SOURCE for all loglines that it
 # formats according to mflog. See a comment in mflog.__init__
 mflog_msg = partial(mflog.decorate, RUNTIME_LOG_SOURCE)
@@ -290,6 +301,7 @@ class NativeRuntime(object):
         pathspec_index,
         cloned_task_pathspec_index,
         finished_tuple,
+        iteration_tuple,
         ubf_context,
         generate_task_obj,
         verbose=False,
@@ -334,7 +346,7 @@ class NativeRuntime(object):
                 self._metadata,
                 origin_ds_set=self._origin_ds_set,
             )
-            self._finished[(step_name, finished_tuple)] = task_pathspec
+            self._finished[(step_name, finished_tuple, iteration_tuple)] = task_pathspec
             self._is_cloned[task_pathspec] = True
         except Exception as e:
             self._logger(
@@ -415,6 +427,7 @@ class NativeRuntime(object):
                 finished_tuple = tuple(
                     [s._replace(value=0) for s in task_ds.get("_foreach_stack", ())]
                 )
+                iteration_tuple = tuple(task_ds.get("_iteration_stack", ()))
                 cloned_task_pathspec_index = pathspec_index.split("/")[1]
                 if task_ds.get("_control_task_is_mapper_zero", False):
                     # Replace None with index 0 for control task as it is part of the
@@ -440,6 +453,7 @@ class NativeRuntime(object):
                         pathspec_index,
                         cloned_task_pathspec_index,
                         finished_tuple,
+                        iteration_tuple,
                         is_ubf_mapper_task,
                         ubf_context,
                     )
@@ -454,6 +468,7 @@ class NativeRuntime(object):
                     pathspec_index,
                     cloned_task_pathspec_index,
                     finished_tuple,
+                    iteration_tuple,
                     ubf_context=ubf_context,
                     generate_task_obj=generate_task_obj and (not is_ubf_mapper_task),
                     verbose=verbose,
@@ -464,6 +479,7 @@ class NativeRuntime(object):
                     pathspec_index,
                     cloned_task_pathspec_index,
                     finished_tuple,
+                    iteration_tuple,
                     is_ubf_mapper_task,
                     ubf_context,
                 ) in inputs
@@ -484,6 +500,7 @@ class NativeRuntime(object):
                 self._queue_push("start", {"input_paths": [self._params_task.path]})
             else:
                 self._queue_push("start", {})
+
         progress_tstamp = time.time()
         with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8") as config_file:
             # Configurations are passed through a file to avoid overloading the
@@ -504,7 +521,74 @@ class NativeRuntime(object):
                 ):
                     # 1. are any of the current workers finished?
                     if self._cloned_tasks:
-                        finished_tasks = self._cloned_tasks
+                        finished_tasks = []
+
+                        # For loops (right now just recursive steps), we need to find
+                        # the exact frontier because if we queue all "successors" to all
+                        # the finished iterations, we would incorrectly launch multiple
+                        # successors. We therefore have to strip out all non-last
+                        # iterations *per* foreach branch.
+                        idx_per_finished_id = (
+                            {}
+                        )  # type: Dict[Tuple[str, Tuple[int, ...], Tuple[int, Tuple[int, ...]]]]
+                        for task in self._cloned_tasks:
+                            step_name, foreach_stack, iteration_stack = task.finished_id
+                            existing_task_idx = idx_per_finished_id.get(
+                                (step_name, foreach_stack), None
+                            )
+                            if existing_task_idx is not None:
+                                len_diff = len(iteration_stack) - len(
+                                    existing_task_idx[1]
+                                )
+                                # In this case, we need to keep only the latest iteration
+                                if (
+                                    len_diff == 0
+                                    and iteration_stack > existing_task_idx[1]
+                                ) or len_diff == -1:
+                                    # We remove the one we currently have and replace
+                                    # by this one. The second option means that we are
+                                    # adding the finished iteration marker.
+                                    existing_task = finished_tasks[existing_task_idx[0]]
+                                    # These are the first two lines of _queue_tasks
+                                    # We still consider the tasks finished so we need
+                                    # to update state to be clean.
+                                    self._finished[existing_task.finished_id] = (
+                                        existing_task.path
+                                    )
+                                    self._is_cloned[existing_task.path] = (
+                                        existing_task.is_cloned
+                                    )
+
+                                    finished_tasks[existing_task_idx[0]] = task
+                                    idx_per_finished_id[(step_name, foreach_stack)] = (
+                                        existing_task_idx[0],
+                                        iteration_stack,
+                                    )
+                                elif (
+                                    len_diff == 0
+                                    and iteration_stack < existing_task_idx[1]
+                                ) or len_diff == 1:
+                                    # The second option is when we have already marked
+                                    # the end of the iteration in self._finished and
+                                    # are now seeing a previous iteration.
+                                    # We just mark the task as finished but we don't
+                                    # put it in the finished_tasks list to pass to
+                                    # the _queue_tasks function
+                                    self._finished[task.finished_id] = task.path
+                                    self._is_cloned[task.path] = task.is_cloned
+                                else:
+                                    raise MetaflowInternalError(
+                                        "Unexpected recursive cloned tasks -- "
+                                        "this is a bug, please report it."
+                                    )
+                            else:
+                                # New entry
+                                finished_tasks.append(task)
+                                idx_per_finished_id[(step_name, foreach_stack)] = (
+                                    len(finished_tasks) - 1,
+                                    iteration_stack,
+                                )
+
                         # reset the list of cloned tasks and let poll_workers handle
                         # the remaining transition
                         self._cloned_tasks = []
@@ -578,7 +662,7 @@ class NativeRuntime(object):
                 self._run_exit_hooks()
 
         # assert that end was executed and it was successful
-        if ("end", ()) in self._finished:
+        if ("end", (), ()) in self._finished:
             if self._run_url:
                 self._logger(
                     "Done! See the run in the UI at %s" % self._run_url,
@@ -604,7 +688,7 @@ class NativeRuntime(object):
             if not exit_hook_decos:
                 return
 
-            successful = ("end", ()) in self._finished or self._clone_only
+            successful = ("end", (), ()) in self._finished or self._clone_only
             pathspec = f"{self._graph.name}/{self._run_id}"
             flow_file = self._environment.get_environment_info()["script"]
 
@@ -672,29 +756,60 @@ class NativeRuntime(object):
 
     # Given the current task information (task_index), the type of transition,
     # and the split index, return the new task index.
-    def _translate_index(self, task, next_step, type, split_index=None):
-        match = re.match(r"^(.+)\[(.*)\]$", task.task_index)
+    def _translate_index(
+        self, task, next_step, type, split_index=None, loop_mode=LoopBehavior.NONE
+    ):
+        match = re.match(r"^(.+)\[(.*)\]\[(.*)\]$", task.task_index)
         if match:
-            _, foreach_index = match.groups()
+            _, foreach_index, iteration_index = match.groups()
             # Convert foreach_index to a list of integers
             if len(foreach_index) > 0:
                 foreach_index = foreach_index.split(",")
             else:
                 foreach_index = []
+            # Ditto for iteration_index
+            if len(iteration_index) > 0:
+                iteration_index = iteration_index.split(",")
+            else:
+                iteration_index = []
         else:
             raise ValueError(
-                "Index not in the format of {run_id}/{step_name}[{foreach_index}]"
+                "Index not in the format of {run_id}/{step_name}[{foreach_index}][{iteration_index}]"
             )
+        if loop_mode == LoopBehavior.NONE:
+            # Check if we are entering a looping construct. Right now, only recursive
+            # steps are looping constructs
+            next_step_node = self._graph[next_step]
+            if (
+                next_step_node.type == "split-switch"
+                and next_step in next_step_node.out_funcs
+            ):
+                loop_mode = LoopBehavior.ENTERING
+
+        # Update iteration_index
+        if loop_mode == LoopBehavior.ENTERING:
+            # We are entering a loop, so we add a new iteration level
+            iteration_index.append("0")
+        elif loop_mode == LoopBehavior.EXITING:
+            iteration_index = iteration_index[:-1]
+        elif loop_mode == LoopBehavior.LOOPING:
+            if len(iteration_index) == 0:
+                raise MetaflowInternalError(
+                    "In looping mode but there is no iteration index"
+                )
+            iteration_index[-1] = str(int(iteration_index[-1]) + 1)
+        iteration_index = ",".join(iteration_index)
+
         if type == "linear":
-            return "%s[%s]" % (next_step, ",".join(foreach_index))
+            return "%s[%s][%s]" % (next_step, ",".join(foreach_index), iteration_index)
         elif type == "join":
             indices = []
             if len(foreach_index) > 0:
                 indices = foreach_index[:-1]
-            return "%s[%s]" % (next_step, ",".join(indices))
+            return "%s[%s][%s]" % (next_step, ",".join(indices), iteration_index)
         elif type == "split":
             foreach_index.append(str(split_index))
-            return "%s[%s]" % (next_step, ",".join(foreach_index))
+            return "%s[%s][%s]" % (next_step, ",".join(foreach_index), iteration_index)
 
     # Store the parameters needed for task creation, so that pushing on items
     # onto the run_queue is an inexpensive operation.
@@ -778,17 +893,19 @@ class NativeRuntime(object):
                 # tasks is incorrect and contains the pathspec of the *cloned* run
                 # but we don't use it for anything. We could look to clean it up though
                 if not task.is_cloned:
-                    _, foreach_stack = task.finished_id
+                    _, foreach_stack, iteration_stack = task.finished_id
                     top = foreach_stack[-1]
                     bottom = list(foreach_stack[:-1])
                     for i in range(num_splits):
                         s = tuple(bottom + [top._replace(index=i)])
-                        self._finished[(task.step, s)] = mapper_tasks[i]
+                        self._finished[(task.step, s, iteration_stack)] = mapper_tasks[
+                            i
+                        ]
                         self._is_cloned[mapper_tasks[i]] = False
 
             # Find and check status of control task and retrieve its pathspec
             # for retrieving unbounded foreach cardinality.
-            _, foreach_stack = task.finished_id
+            _, foreach_stack, iteration_stack = task.finished_id
             top = foreach_stack[-1]
             bottom = list(foreach_stack[:-1])
             s = tuple(bottom + [top._replace(index=None)])
@@ -797,7 +914,7 @@ class NativeRuntime(object):
             # it will have index=0 instead of index=None.
             if task.results.get("_control_task_is_mapper_zero", False):
                 s = tuple(bottom + [top._replace(index=0)])
-            control_path = self._finished.get((task.step, s))
+            control_path = self._finished.get((task.step, s, iteration_stack))
             if control_path:
                 # Control task was successful.
                 # Additionally check the state of (sibling) mapper tasks as well
@@ -806,7 +923,9 @@ class NativeRuntime(object):
                 required_tasks = []
                 for i in range(num_splits):
                     s = tuple(bottom + [top._replace(index=i)])
-                    required_tasks.append(self._finished.get((task.step, s)))
+                    required_tasks.append(
+                        self._finished.get((task.step, s, iteration_stack))
+                    )
 
                 if all(required_tasks):
                     index = self._translate_index(task, next_step, "join")
@@ -819,7 +938,7 @@ class NativeRuntime(object):
         else:
             # matching_split is the split-parent of the finished task
             matching_split = self._graph[self._graph[next_step].split_parents[-1]]
-            _, foreach_stack = task.finished_id
+            _, foreach_stack, iteration_stack = task.finished_id
 
             direct_parents = set(self._graph[next_step].in_funcs)
 
@@ -837,7 +956,7 @@ class NativeRuntime(object):
                     filter(
                         lambda x: x is not None,
                         [
-                            self._finished.get((p, s))
+                            self._finished.get((p, s, iteration_stack))
                             for p in direct_parents
                             for s in siblings(foreach_stack)
                         ],
@@ -852,11 +971,12 @@ class NativeRuntime(object):
                     filter(
                         lambda x: x is not None,
                         [
-                            self._finished.get((p, foreach_stack))
+                            self._finished.get((p, foreach_stack, iteration_stack))
                             for p in direct_parents
                         ],
                     )
                 )
+
                 required_count = len(matching_split.out_funcs)
                 join_type = "linear"
                 index = self._translate_index(task, next_step, "linear")
@@ -868,9 +988,18 @@ class NativeRuntime(object):
                     index,
                 )
 
-    def _queue_task_switch(self, task, next_steps):
+    def _queue_task_switch(self, task, next_steps, is_recursive):
         chosen_step = next_steps[0]
-        index = self._translate_index(task, chosen_step, "linear")
+
+        loop_mode = LoopBehavior.NONE
+        if is_recursive:
+            if chosen_step != task.step:
+                # We are exiting a loop
+                loop_mode = LoopBehavior.EXITING
+            else:
+                # We are staying in the loop
+                loop_mode = LoopBehavior.LOOPING
+        index = self._translate_index(task, chosen_step, "linear", None, loop_mode)
         self._queue_push(chosen_step, {"input_paths": [task.path]}, index)
 
     def _queue_task_foreach(self, task, next_steps):
@@ -949,7 +1078,9 @@ class NativeRuntime(object):
                 next_steps = []
                 foreach = None
             expected = self._graph[task.step].out_funcs
+
             if self._graph[task.step].type == "split-switch":
+                is_recursive = task.step in self._graph[task.step].out_funcs
                 if len(next_steps) != 1:
                     msg = (
                         "Switch step *{step}* should transition to exactly "
@@ -970,6 +1101,15 @@ class NativeRuntime(object):
                             expected=", ".join(expected),
                         )
                     )
+                # When exiting a recursive loop, we mark that the loop itself has
+                # finished by adding a special entry in self._finished which has
+                # an iteration stack that is shorter (ie: we are out of the loop) so
+                # that we can then find it when looking at successor tasks to launch.
+                if is_recursive and next_steps[0] != task.step:
+                    step_name, finished_tuple, iteration_tuple = task.finished_id
+                    self._finished[
+                        (step_name, finished_tuple, iteration_tuple[:-1])
+                    ] = task.path
             elif next_steps != expected:
                 msg = (
                     "Based on static analysis of the code, step *{step}* "
@@ -995,8 +1135,8 @@ class NativeRuntime(object):
                 # Next step is a foreach child
                 self._queue_task_foreach(task, next_steps)
             elif self._graph[task.step].type == "split-switch":
-                # Next step is switch - queue the chosen step
-                self._queue_task_switch(task, next_steps)
+                # Current step is switch - queue the chosen step
+                self._queue_task_switch(task, next_steps, is_recursive)
             else:
                 # Next steps are normal linear steps
                 for step in next_steps:
@@ -1537,13 +1677,13 @@ class Task(object):
     @property
     def finished_id(self):
         # note: id is not available before the task has finished.
-        # Index already identifies the task within the foreach,
-        # we will remove foreach value so that it is easier to
+        # Index already identifies the task within the foreach and loop.
+        # We will remove foreach value so that it is easier to
         # identify siblings within a foreach.
         foreach_stack_tuple = tuple(
             [s._replace(value=0) for s in self.results["_foreach_stack"]]
         )
-        return (self.step, foreach_stack_tuple)
+        return (self.step, foreach_stack_tuple, tuple(self.results["_iteration_stack"]))
 
     @property
     def is_cloned(self):

metaflow/task.py

@@ -117,11 +117,21 @@ class MetaflowTask(object):
 
         # We back out of the stack of generators
         for w in reversed(wrappers_stack):
-            r = w.post_step(orig_step_func.name, self.flow, raised_exception)
-            if r is None or isinstance(r, Exception):
+            try:
+                r = w.post_step(orig_step_func.name, self.flow, raised_exception)
+            except Exception as ex:
+                r = ex
+            if r is None:
                 raised_exception = None
+            elif isinstance(r, Exception):
+                raised_exception = r
             elif isinstance(r, tuple):
-                raised_exception, fake_next_call_args = r
+                if len(r) == 2:
+                    raised_exception, fake_next_call_args = r
+                else:
+                    # The last argument is an exception to be re-raised. Used in
+                    # user_step_decorator's post_step
+                    raise r[2]
             else:
                 raise RuntimeError(
                     "Invalid return value from a UserStepDecorator. Expected an"
@@ -239,6 +249,7 @@ class MetaflowTask(object):
                 # Prefetch 'foreach' related artifacts to improve time taken by
                 # _init_foreach.
                 prefetch_data_artifacts = [
+                    "_iteration_stack",
                     "_foreach_stack",
                     "_foreach_num_splits",
                     "_foreach_var",
@@ -375,6 +386,56 @@ class MetaflowTask(object):
         elif "_foreach_stack" in inputs[0]:
             self.flow._foreach_stack = inputs[0]["_foreach_stack"]
 
+    def _init_iteration(self, step_name, inputs, is_recursive_step):
+        # We track the iteration "stack" for loops. At this time, we
+        # only support one type of "looping" which is a recursive step but
+        # this can generalize to arbitrary well-scoped loops in the future.
+
+        # _iteration_stack will contain the iteration count for each loop
+        # level. Currently, there will be only no elements (no loops) or
+        # a single element (a single recursive step).
+
+        # We just need to determine the rules to add a new looping level,
+        # increment the looping level or pop the looping level. In our
+        # current support for only recursive steps, this is pretty straightforward:
+        # 1) if is_recursive_step:
+        #    - we are entering a loop -- we are either entering for the first time
+        #      or we are continuing the loop. Note that a recursive step CANNOT
+        #      be a join step so there is always a single input
+        #    1a) If inputs[0]["_iteration_stack"] contains an element, we are looping
+        #      so we increment the count
+        #    1b) If inputs[0]["_iteration_stack"] is empty, this is the first time we
+        #      are entering the loop so we set the iteration count to 0
+        # 2) if it is not a recursive step, we need to determine if this is the step
+        #    *after* the recursive step. The easiest way to determine that is to
+        #    look at all inputs (there can be multiple in case of a join) and pop
+        #    _iteration_stack if it is set. However, since we know that non recursive
+        #    steps are *never* part of an iteration, we can simplify and just set it
+        #    to [] without even checking anything. We will have to revisit this if/when
+        #    more complex loop structures are supported.
+
+        # Note that just like _foreach_stack, we need to set _iteration_stack to *something*
+        # so that it doesn't get clobbered weirdly by merge_artifacts.
+
+        if is_recursive_step:
+            # Case 1)
+            if len(inputs) != 1:
+                raise MetaflowInternalError(
+                    "Step *%s* is a recursive step but got multiple inputs." % step_name
+                )
+            inp = inputs[0]
+            if "_iteration_stack" not in inp or not inp["_iteration_stack"]:
+                # Case 1b)
+                self.flow._iteration_stack = [0]
+            else:
+                # Case 1a)
+                stack = inp["_iteration_stack"]
+                stack[-1] += 1
+                self.flow._iteration_stack = stack
+        else:
+            # Case 2)
+            self.flow._iteration_stack = []
+
     def _clone_flow(self, datastore):
         x = self.flow.__class__(use_cli=False)
         x._set_datastore(datastore)
@@ -563,6 +624,12 @@ class MetaflowTask(object):
             # 3. initialize foreach state
             self._init_foreach(step_name, join_type, inputs, split_index)
 
+            # 4. initialize the iteration state
+            is_recursive_step = (
+                node.type == "split-switch" and step_name in node.out_funcs
+            )
+            self._init_iteration(step_name, inputs, is_recursive_step)
+
             # Add foreach stack to metadata of the task
 
             foreach_stack = (

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/user_decorators/user_step_decorator.py

@@ -658,8 +658,14 @@ def user_step_decorator(*args, **kwargs):
                         self._generator.send(None)
                 except StopIteration as e:
                     to_return = e.value
+                except Exception as e:
+                    return e
                 else:
-                    raise MetaflowException(" %s should only yield once" % self)
+                    return (
+                        None,
+                        None,
+                        MetaflowException(" %s should only yield once" % self),
+                    )
                 return to_return
 
         return WrapClass

metaflow/version.py

@@ -1 +1 @@
-metaflow_version = "2.17.3.1"
+metaflow_version = "2.18.0.1"

{ob_metaflow-2.17.3.1.dist-info → ob_metaflow-2.18.0.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ob-metaflow
-Version: 2.17.3.1
+Version: 2.18.0.1
 Summary: Metaflow: More AI and ML, Less Engineering
 Author: Netflix, Outerbounds & the Metaflow Community
 Author-email: help@outerbounds.co
@@ -12,7 +12,7 @@ Requires-Dist: boto3
 Requires-Dist: pylint
 Requires-Dist: kubernetes
 Provides-Extra: stubs
-Requires-Dist: metaflow-stubs==2.17.3.1; extra == "stubs"
+Requires-Dist: metaflow-stubs==2.18.0.1; extra == "stubs"
 Dynamic: author
 Dynamic: author-email
 Dynamic: description