ob-metaflow 2.15.13.1__py2.py3-none-any.whl → 2.19.7.1rc0__py2.py3-none-any.whl

metaflow/task.py

@@ -11,8 +11,10 @@ from types import MethodType, FunctionType
 from metaflow.sidecar import Message, MessageTypes
 from metaflow.datastore.exceptions import DataException
 
+from metaflow.plugins import METADATA_PROVIDERS
 from .metaflow_config import MAX_ATTEMPTS
 from .metadata_provider import MetaDatum
+from .metaflow_profile import from_start
 from .mflog import TASK_LOG_SOURCE
 from .datastore import Inputs, TaskDataStoreSet
 from .exception import (
@@ -24,6 +26,7 @@ from .unbounded_foreach import UBF_CONTROL
 from .util import all_equal, get_username, resolve_identity, unicode_type
 from .clone_util import clone_task_helper
 from .metaflow_current import current
+from metaflow.user_configs.config_parameters import ConfigValue
 from metaflow.system import _system_logger, _system_monitor
 from metaflow.tracing import get_trace_id
 from metaflow.tuple_util import ForeachFrame
@@ -47,6 +50,8 @@ class MetaflowTask(object):
         event_logger,
         monitor,
         ubf_context,
+        orig_flow_datastore=None,
+        spin_artifacts=None,
     ):
         self.flow = flow
         self.flow_datastore = flow_datastore
@@ -56,12 +61,126 @@ class MetaflowTask(object):
         self.event_logger = event_logger
         self.monitor = monitor
         self.ubf_context = ubf_context
+        self.orig_flow_datastore = orig_flow_datastore
+        self.spin_artifacts = spin_artifacts
+
+    def _exec_step_function(self, step_function, orig_step_func, input_obj=None):
+        wrappers_stack = []
+        wrapped_func = None
+
+        # Will set to non-Falsy if we need to fake calling `self.next`
+        # This is used when skipping the step.
+        # If a dictionary, it will
+        # contain the arguments to pass to `self.next`. If
+        # True, it means we are using whatever the usual
+        # arguments to `self.next` are for this step.
+        fake_next_call_args = False
+        raised_exception = None
+        had_raised_exception = False
+
+        # If we have wrappers w1, w2 and w3, we need to execute
+        #  - w3_pre
+        #  - w2_pre
+        #  - w1_pre
+        #  - step_function
+        #  - w1_post
+        #  - w2_post
+        #  - w3_post
+        # in that order. We do this by maintaining a stack of generators.
+        # Note that if any of the pre functions returns a function, we execute that
+        # instead of the rest of the inside part. This is useful if you want to create
+        # no-op function for example.
+        for w in reversed(orig_step_func.wrappers):
+            wrapped_func = w.pre_step(orig_step_func.name, self.flow, input_obj)
+            wrappers_stack.append(w)
+            if w.skip_step:
+                # We are not going to run anything so we will have to fake calling
+                # next.
+                fake_next_call_args = w.skip_step
+                break
+            if wrapped_func:
+                break  # We have nothing left to do since we now execute the
+                # wrapped function
+            # Else, we continue down the list of wrappers
+        try:
+            # fake_next_call is used here to also indicate that the step was skipped
+            # so we do not execute anything.
+            if not fake_next_call_args:
+                if input_obj is None:
+                    if wrapped_func:
+                        fake_next_call_args = wrapped_func(self.flow)
+                    else:
+                        step_function()
+                else:
+                    if wrapped_func:
+                        fake_next_call_args = wrapped_func(self.flow, input_obj)
+                    else:
+                        step_function(input_obj)
+        except Exception as ex:
+            raised_exception = ex
+            had_raised_exception = True
 
-    def _exec_step_function(self, step_function, input_obj=None):
-        if input_obj is None:
-            step_function()
-        else:
-            step_function(input_obj)
+        # We back out of the stack of generators
+        for w in reversed(wrappers_stack):
+            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):
+                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"
+                    "exception or an exception and arguments for self.next, got: %s" % r
+                )
+        if raised_exception:
+            # We have an exception that we need to propagate
+            raise raised_exception
+
+        if fake_next_call_args or had_raised_exception:
+            # We want to override the next call or we caught an exception (in which
+            # case the regular step code didn't call self.next). In this case,
+            # we need to set the transition variables
+            # properly. We call the next function as needed
+            # We also do this in case we want to gobble the exception.
+            graph_node = self.flow._graph[orig_step_func.name]
+            out_funcs = [getattr(self.flow, f) for f in graph_node.out_funcs]
+            if out_funcs:
+                self.flow._transition = None
+                if isinstance(fake_next_call_args, dict) and fake_next_call_args:
+                    # Not an empty dictionary -- we use this as arguments for the next
+                    # call
+                    self.flow.next(*out_funcs, **fake_next_call_args)
+                elif (
+                    fake_next_call_args == True
+                    or fake_next_call_args == {}
+                    or had_raised_exception
+                ):
+                    # We need to extract things from the self.next. This is not possible
+                    # in the case where there was a num_parallel.
+                    if graph_node.parallel_foreach:
+                        raise RuntimeError(
+                            "Skipping a parallel foreach step without providing "
+                            "the arguments to the self.next call is not supported. "
+                        )
+                    if graph_node.foreach_param:
+                        self.flow.next(*out_funcs, foreach=graph_node.foreach_param)
+                    else:
+                        self.flow.next(*out_funcs)
+                else:
+                    raise RuntimeError(
+                        "Invalid value passed to self.next; expected "
+                        " bool of a dictionary; got: %s" % fake_next_call_args
+                    )
 
     def _init_parameters(self, parameter_ds, passdown=True):
         cls = self.flow.__class__
@@ -120,7 +239,6 @@ class MetaflowTask(object):
             lambda _, parameter_ds=parameter_ds: parameter_ds["_graph_info"],
         )
         all_vars.append("_graph_info")
-
         if passdown:
             self.flow._datastore.passdown_partial(parameter_ds, all_vars)
         return param_only_vars
@@ -136,6 +254,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",
@@ -147,6 +266,9 @@ class MetaflowTask(object):
                 run_id,
                 pathspecs=input_paths,
                 prefetch_data_artifacts=prefetch_data_artifacts,
+                join_type=join_type,
+                orig_flow_datastore=self.orig_flow_datastore,
+                spin_artifacts=self.spin_artifacts,
             )
             ds_list = [ds for ds in datastore_set]
             if len(ds_list) != len(input_paths):
@@ -158,10 +280,27 @@ class MetaflowTask(object):
             # initialize directly in the single input case.
             ds_list = []
             for input_path in input_paths:
-                run_id, step_name, task_id = input_path.split("/")
+                parts = input_path.split("/")
+                if len(parts) == 3:
+                    run_id, step_name, task_id = parts
+                    attempt = None
+                else:
+                    run_id, step_name, task_id, attempt = parts
+                    attempt = int(attempt)
+
                 ds_list.append(
-                    self.flow_datastore.get_task_datastore(run_id, step_name, task_id)
+                    self.flow_datastore.get_task_datastore(
+                        run_id,
+                        step_name,
+                        task_id,
+                        attempt=attempt,
+                        join_type=join_type,
+                        orig_flow_datastore=self.orig_flow_datastore,
+                        spin_artifacts=self.spin_artifacts,
+                    )
                 )
+                from_start("MetaflowTask: got datastore for input path %s" % input_path)
+
         if not ds_list:
             # this guards against errors in input paths
             raise MetaflowDataMissing(
@@ -272,6 +411,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)
@@ -382,6 +571,8 @@ class MetaflowTask(object):
         split_index,
         retry_count,
         max_user_code_retries,
+        whitelist_decorators=None,
+        persist=True,
     ):
         if run_id and task_id:
             self.metadata.register_run_id(run_id)
@@ -440,7 +631,14 @@ class MetaflowTask(object):
 
         step_func = getattr(self.flow, step_name)
         decorators = step_func.decorators
-
+        if self.orig_flow_datastore:
+            # We filter only the whitelisted decorators in case of spin step.
+            decorators = (
+                []
+                if not whitelist_decorators
+                else [deco for deco in decorators if deco.name in whitelist_decorators]
+            )
+        from_start("MetaflowTask: decorators initialized")
         node = self.flow._graph[step_name]
         join_type = None
         if node.type == "join":
@@ -448,17 +646,26 @@ class MetaflowTask(object):
 
         # 1. initialize output datastore
         output = self.flow_datastore.get_task_datastore(
-            run_id, step_name, task_id, attempt=retry_count, mode="w"
+            run_id, step_name, task_id, attempt=retry_count, mode="w", persist=persist
         )
 
         output.init_task()
+        from_start("MetaflowTask: output datastore initialized")
 
         if input_paths:
             # 2. initialize input datastores
             inputs = self._init_data(run_id, join_type, input_paths)
+            from_start("MetaflowTask: input datastores initialized")
 
             # 3. initialize foreach state
             self._init_foreach(step_name, join_type, inputs, split_index)
+            from_start("MetaflowTask: foreach state initialized")
+
+            # 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
 
@@ -511,7 +718,7 @@ class MetaflowTask(object):
                         ),
                     ]
                 )
-
+        from_start("MetaflowTask: finished input processing")
         self.metadata.register_metadata(
             run_id,
             step_name,
@@ -538,6 +745,9 @@ class MetaflowTask(object):
         output.save_metadata(
             {
                 "task_begin": {
+                    "code_package_metadata": os.environ.get(
+                        "METAFLOW_CODE_METADATA", ""
+                    ),
                     "code_package_sha": os.environ.get("METAFLOW_CODE_SHA"),
                     "code_package_ds": os.environ.get("METAFLOW_CODE_DS"),
                     "code_package_url": os.environ.get("METAFLOW_CODE_URL"),
@@ -562,8 +772,11 @@ class MetaflowTask(object):
             "project_flow_name": current.get("project_flow_name"),
             "trace_id": trace_id or None,
         }
+
+        from_start("MetaflowTask: task metadata initialized")
         start = time.time()
         self.metadata.start_task_heartbeat(self.flow.name, run_id, step_name, task_id)
+        from_start("MetaflowTask: heartbeat started")
         with self.monitor.measure("metaflow.task.duration"):
             try:
                 with self.monitor.count("metaflow.task.start"):
@@ -583,18 +796,23 @@ class MetaflowTask(object):
                 # should either be set prior to running the user code or listed in
                 # FlowSpec._EPHEMERAL to allow for proper merging/importing of
                 # user artifacts in the user's step code.
-
                 if join_type:
                     # Join step:
 
-                    # Ensure that we have the right number of inputs. The
-                    # foreach case is checked above.
-                    if join_type != "foreach" and len(inputs) != len(node.in_funcs):
-                        raise MetaflowDataMissing(
-                            "Join *%s* expected %d "
-                            "inputs but only %d inputs "
-                            "were found" % (step_name, len(node.in_funcs), len(inputs))
-                        )
+                    # Ensure that we have the right number of inputs.
+                    if join_type != "foreach":
+                        # Find the corresponding split node from the graph.
+                        split_node = self.flow._graph[node.split_parents[-1]]
+                        # The number of expected inputs is the number of branches
+                        # from that split -- we can't use in_funcs because there may
+                        # be more due to split-switch branches that all converge here.
+                        expected_inputs = len(split_node.out_funcs)
+
+                        if len(inputs) != expected_inputs:
+                            raise MetaflowDataMissing(
+                                "Join *%s* expected %d inputs but only %d inputs "
+                                "were found" % (step_name, expected_inputs, len(inputs))
+                            )
 
                     # Multiple input contexts are passed in as an argument
                     # to the step function.
@@ -636,11 +854,19 @@ class MetaflowTask(object):
                                 "graph_info": self.flow._graph_info,
                             }
                         )
+                from_start("MetaflowTask: before pre-step decorators")
                 for deco in decorators:
+                    if deco.name == "card" and self.orig_flow_datastore:
+                        # if spin step and card decorator, pass spin metadata
+                        metadata = [m for m in METADATA_PROVIDERS if m.TYPE == "spin"][
+                            0
+                        ](self.environment, self.flow, self.event_logger, self.monitor)
+                    else:
+                        metadata = self.metadata
                     deco.task_pre_step(
                         step_name,
                         output,
-                        self.metadata,
+                        metadata,
                         run_id,
                         task_id,
                         self.flow,
@@ -651,6 +877,7 @@ class MetaflowTask(object):
                         inputs,
                     )
 
+                orig_step_func = step_func
                 for deco in decorators:
                     # decorators can actually decorate the step function,
                     # or they can replace it altogether. This functionality
@@ -665,12 +892,12 @@ class MetaflowTask(object):
                         max_user_code_retries,
                         self.ubf_context,
                     )
-
+                from_start("MetaflowTask: finished decorator processing")
                 if join_type:
-                    self._exec_step_function(step_func, input_obj)
+                    self._exec_step_function(step_func, orig_step_func, input_obj)
                 else:
-                    self._exec_step_function(step_func)
-
+                    self._exec_step_function(step_func, orig_step_func)
+                from_start("MetaflowTask: step function executed")
                 for deco in decorators:
                     deco.task_post_step(
                         step_name,
@@ -713,6 +940,7 @@ class MetaflowTask(object):
                     raise
 
             finally:
+                from_start("MetaflowTask: decorators finalized")
                 if self.ubf_context == UBF_CONTROL:
                     self._finalize_control_task()
 
@@ -752,7 +980,7 @@ class MetaflowTask(object):
                     )
 
                 output.save_metadata({"task_end": {}})
-
+                from_start("MetaflowTask: output persisted")
                 # this writes a success marker indicating that the
                 # "transaction" is done
                 output.done()
@@ -781,3 +1009,4 @@ class MetaflowTask(object):
                     name="duration",
                     payload={**task_payload, "msg": str(duration)},
                 )
+                from_start("MetaflowTask: task run completed")

metaflow/user_configs/config_options.py

@@ -7,8 +7,9 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 from metaflow._vendor import click
 from metaflow.debug import debug
 
-from .config_parameters import CONFIG_FILE, ConfigValue
+from .config_parameters import ConfigValue
 from ..exception import MetaflowException, MetaflowInternalError
+from ..packaging_sys import MetaflowCodeContent
 from ..parameters import DeployTimeField, ParameterContext, current_flow
 from ..util import get_username
 
@@ -24,12 +25,16 @@ _CONVERTED_DEFAULT_NO_FILE = _CONVERTED_DEFAULT + _NO_FILE
 
 def _load_config_values(info_file: Optional[str] = None) -> Optional[Dict[Any, Any]]:
     if info_file is None:
-        info_file = os.path.basename(CONFIG_FILE)
-    try:
-        with open(info_file, encoding="utf-8") as contents:
-            return json.load(contents).get("user_configs", {})
-    except IOError:
-        return None
+        config_content = MetaflowCodeContent.get_config()
+    else:
+        try:
+            with open(info_file, encoding="utf-8") as f:
+                config_content = json.load(f)
+        except IOError:
+            return None
+    if config_content:
+        return config_content.get("user_configs", {})
+    return None
 
 
 class ConvertPath(click.Path):
@@ -181,7 +186,7 @@ class ConfigInput:
         click_obj: Optional[Any] = None,
     ):
         from ..cli import echo_always, echo_dev_null  # Prevent circular import
-        from ..flowspec import _FlowState  # Prevent circular import
+        from ..flowspec import FlowStateItems  # Prevent circular import
 
         flow_cls = getattr(current_flow, "flow_cls", None)
         if flow_cls is None:
@@ -221,13 +226,13 @@ class ConfigInput:
         if param_name == "config_value":
             self._value_values = {
                 k.lower(): v
-                for k, v in param_value
+                for k, v in param_value.items()
                 if v is not None and not v.startswith(_CONVERTED_DEFAULT)
             }
         else:
             self._path_values = {
                 k.lower(): v
-                for k, v in param_value
+                for k, v in param_value.items()
                 if v is not None and not v.startswith(_CONVERTED_DEFAULT)
             }
         if do_return:
@@ -255,7 +260,6 @@ class ConfigInput:
             for k in all_keys
         )
 
-        flow_cls._flow_state[_FlowState.CONFIGS] = {}
         to_return = {}
 
         if not has_all_kv:
@@ -327,14 +331,14 @@ class ConfigInput:
             if val is None:
                 missing_configs.add(name)
                 to_return[name] = None
-                flow_cls._flow_state[_FlowState.CONFIGS][name] = None
-                continue
-            if val.startswith(_CONVERTED_DEFAULT_NO_FILE):
-                no_default_file.append(name)
+                flow_cls._flow_state.self_data[FlowStateItems.CONFIGS][name] = None
                 continue
             if val.startswith(_CONVERTED_NO_FILE):
                 no_file.append(name)
                 continue
+            if val.startswith(_CONVERTED_DEFAULT_NO_FILE):
+                no_default_file.append(name)
+                continue
 
             val = val[len(_CONVERT_PREFIX) :]  # Remove the _CONVERT_PREFIX
             if val.startswith(_DEFAULT_PREFIX):  # Remove the _DEFAULT_PREFIX if needed
@@ -351,8 +355,12 @@ class ConfigInput:
                         click_obj.delayed_config_exception = exc
                         return None
                     raise exc from e
-                flow_cls._flow_state[_FlowState.CONFIGS][name] = read_value
-                to_return[name] = ConfigValue(read_value)
+                flow_cls._flow_state.self_data[FlowStateItems.CONFIGS][
+                    name
+                ] = read_value
+                to_return[name] = (
+                    ConfigValue(read_value) if read_value is not None else None
+                )
             else:
                 if self._parsers[name]:
                     read_value = self._call_parser(self._parsers[name], val)
@@ -366,8 +374,12 @@ class ConfigInput:
                         )
                         continue
                     # TODO: Support YAML
-                flow_cls._flow_state[_FlowState.CONFIGS][name] = read_value
-                to_return[name] = ConfigValue(read_value)
+                flow_cls._flow_state.self_data[FlowStateItems.CONFIGS][
+                    name
+                ] = read_value
+                to_return[name] = (
+                    ConfigValue(read_value) if read_value is not None else None
+                )
 
         reqs = missing_configs.intersection(self._req_configs)
         for missing in reqs:
@@ -398,7 +410,7 @@ class ConfigInput:
         return self.process_configs(
             ctx.obj.flow.name,
             param.name,
-            value,
+            dict(value),
             ctx.params["quiet"],
             ctx.params["datastore"],
             click_obj=ctx.obj,
@@ -433,7 +445,7 @@ class LocalFileInput(click.Path):
     # Small wrapper around click.Path to set the value from which to read configuration
     # values. This is set immediately upon processing the --local-config-file
     # option and will therefore then be available when processing any of the other
-    # --config options (which will call ConfigInput.process_configs
+    # --config options (which will call ConfigInput.process_configs)
     name = "LocalFileInput"
 
     def convert(self, value, param, ctx):