metaflow 2.18.12__py2.py3-none-any.whl → 2.19.0__py2.py3-none-any.whl

metaflow/runtime.py

@@ -26,20 +26,29 @@ from metaflow.datastore.exceptions import DataException
 from contextlib import contextmanager
 
 from . import get_namespace
+from .client.filecache import FileCache, FileBlobCache, TaskMetadataCache
 from .metadata_provider import MetaDatum
-from .metaflow_config import FEAT_ALWAYS_UPLOAD_CODE_PACKAGE, MAX_ATTEMPTS, UI_URL
+from .metaflow_config import (
+    FEAT_ALWAYS_UPLOAD_CODE_PACKAGE,
+    MAX_ATTEMPTS,
+    UI_URL,
+    SPIN_ALLOWED_DECORATORS,
+    SPIN_DISALLOWED_DECORATORS,
+)
+from .metaflow_profile import from_start
+from .plugins import DATASTORES
 from .exception import (
     MetaflowException,
     MetaflowInternalError,
     METAFLOW_EXIT_DISALLOW_RETRY,
 )
 from . import procpoll
-from .datastore import TaskDataStoreSet
+from .datastore import FlowDataStore, TaskDataStoreSet
 from .debug import debug
 from .decorators import flow_decorators
 from .flowspec import _FlowState
 from .mflog import mflog, RUNTIME_LOG_SOURCE
-from .util import to_unicode, compress_list, unicode_type
+from .util import to_unicode, compress_list, unicode_type, get_latest_task_pathspec
 from .clone_util import clone_task_helper
 from .unbounded_foreach import (
     CONTROL_TASK_TAG,
@@ -85,6 +94,253 @@ mflog_msg = partial(mflog.decorate, RUNTIME_LOG_SOURCE)
 # TODO option: output dot graph periodically about execution
 
 
+class SpinRuntime(object):
+    def __init__(
+        self,
+        flow,
+        graph,
+        flow_datastore,
+        metadata,
+        environment,
+        package,
+        logger,
+        entrypoint,
+        event_logger,
+        monitor,
+        step_func,
+        step_name,
+        spin_pathspec,
+        skip_decorators=False,
+        artifacts_module=None,
+        persist=True,
+        max_log_size=MAX_LOG_SIZE,
+    ):
+        from metaflow import Task
+
+        self._flow = flow
+        self._graph = graph
+        self._flow_datastore = flow_datastore
+        self._metadata = metadata
+        self._environment = environment
+        self._package = package
+        self._logger = logger
+        self._entrypoint = entrypoint
+        self._event_logger = event_logger
+        self._monitor = monitor
+
+        self._step_func = step_func
+
+        # Determine if we have a complete pathspec or need to get the task
+        if spin_pathspec:
+            parts = spin_pathspec.split("/")
+            if len(parts) == 4:
+                # Complete pathspec: flow/run/step/task_id
+                try:
+                    # If user provides whole pathspec, we do not need to check namespace
+                    task = Task(spin_pathspec, _namespace_check=False)
+                except Exception:
+                    raise MetaflowException(
+                        f"Invalid pathspec: {spin_pathspec} for step: {step_name}"
+                    )
+            elif len(parts) == 3:
+                # Partial pathspec: flow/run/step - need to get the task
+                _, run_id, _ = parts
+                task = get_latest_task_pathspec(flow.name, step_name, run_id=run_id)
+                logger(
+                    f"To make spin even faster, provide complete pathspec with task_id: {task.pathspec}",
+                    system_msg=True,
+                )
+            else:
+                raise MetaflowException(
+                    f"Invalid pathspec format: {spin_pathspec}. Expected flow/run/step or flow/run/step/task_id"
+                )
+        else:
+            # No pathspec provided, get latest task for this step
+            task = get_latest_task_pathspec(flow.name, step_name)
+            logger(
+                f"To make spin even faster, provide complete pathspec {task.pathspec}",
+                system_msg=True,
+            )
+        from_start("SpinRuntime: after getting task")
+
+        # Get the original FlowDatastore so we can use it to access artifacts from the
+        # spun task
+        meta_dict = task.metadata_dict
+        ds_type = meta_dict["ds-type"]
+        ds_root = meta_dict["ds-root"]
+        orig_datastore_impl = [d for d in DATASTORES if d.TYPE == ds_type][0]
+        orig_datastore_impl.datastore_root = ds_root
+        spin_pathspec = task.pathspec
+        orig_flow_datastore = FlowDataStore(
+            flow.name,
+            environment=None,
+            storage_impl=orig_datastore_impl,
+            ds_root=ds_root,
+        )
+
+        self._filecache = FileCache()
+        orig_flow_datastore.set_metadata_cache(
+            TaskMetadataCache(self._filecache, ds_type, ds_root, flow.name)
+        )
+        orig_flow_datastore.ca_store.set_blob_cache(
+            FileBlobCache(
+                self._filecache, FileCache.flow_ds_id(ds_type, ds_root, flow.name)
+            )
+        )
+
+        self._orig_flow_datastore = orig_flow_datastore
+        self._spin_pathspec = spin_pathspec
+        self._persist = persist
+        self._spin_task = task
+        self._input_paths = None
+        self._split_index = None
+        self._whitelist_decorators = None
+        self._config_file_name = None
+        self._skip_decorators = skip_decorators
+        self._artifacts_module = artifacts_module
+        self._max_log_size = max_log_size
+        self._encoding = sys.stdout.encoding or "UTF-8"
+
+        # Create a new run_id for the spin task
+        self.run_id = self._metadata.new_run_id()
+        # Raise exception if we have a black listed decorator
+        for deco in self._step_func.decorators:
+            if deco.name in SPIN_DISALLOWED_DECORATORS:
+                raise MetaflowException(
+                    f"Spinning steps with @{deco.name} decorator is not supported."
+                )
+
+        for deco in self.whitelist_decorators:
+            deco.runtime_init(flow, graph, package, self.run_id)
+        from_start("SpinRuntime: after init decorators")
+
+    @property
+    def split_index(self):
+        """
+        Returns the split index, caching the result after the first access.
+        """
+        if self._split_index is None:
+            self._split_index = getattr(self._spin_task, "index", None)
+
+        return self._split_index
+
+    @property
+    def input_paths(self):
+        def _format_input_paths(task_pathspec, attempt):
+            _, run_id, step_name, task_id = task_pathspec.split("/")
+            return f"{run_id}/{step_name}/{task_id}/{attempt}"
+
+        if self._input_paths:
+            return self._input_paths
+
+        if self._step_func.name == "start":
+            from metaflow import Step
+
+            flow_name, run_id, _, _ = self._spin_pathspec.split("/")
+            task = Step(
+                f"{flow_name}/{run_id}/_parameters", _namespace_check=False
+            ).task
+            self._input_paths = [
+                _format_input_paths(task.pathspec, task.current_attempt)
+            ]
+        else:
+            parent_tasks = self._spin_task.parent_tasks
+            self._input_paths = [
+                _format_input_paths(t.pathspec, t.current_attempt) for t in parent_tasks
+            ]
+        return self._input_paths
+
+    @property
+    def whitelist_decorators(self):
+        if self._skip_decorators:
+            self._whitelist_decorators = []
+            return self._whitelist_decorators
+        if self._whitelist_decorators:
+            return self._whitelist_decorators
+        self._whitelist_decorators = [
+            deco
+            for deco in self._step_func.decorators
+            if any(deco.name.startswith(prefix) for prefix in SPIN_ALLOWED_DECORATORS)
+        ]
+        return self._whitelist_decorators
+
+    def _new_task(self, step, input_paths=None, **kwargs):
+        return Task(
+            flow_datastore=self._flow_datastore,
+            flow=self._flow,
+            step=step,
+            run_id=self.run_id,
+            metadata=self._metadata,
+            environment=self._environment,
+            entrypoint=self._entrypoint,
+            event_logger=self._event_logger,
+            monitor=self._monitor,
+            input_paths=input_paths,
+            decos=self.whitelist_decorators,
+            logger=self._logger,
+            split_index=self.split_index,
+            **kwargs,
+        )
+
+    def execute(self):
+        exception = None
+        with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8") as config_file:
+            config_value = dump_config_values(self._flow)
+            if config_value:
+                json.dump(config_value, config_file)
+                config_file.flush()
+                self._config_file_name = config_file.name
+            else:
+                self._config_file_name = None
+            from_start("SpinRuntime: config values processed")
+            self.task = self._new_task(self._step_func.name, self.input_paths)
+            try:
+                self._launch_and_monitor_task()
+            except Exception as ex:
+                self._logger("Task failed.", system_msg=True, bad=True)
+                exception = ex
+                raise
+            finally:
+                for deco in self.whitelist_decorators:
+                    deco.runtime_finished(exception)
+
+    def _launch_and_monitor_task(self):
+        worker = Worker(
+            self.task,
+            self._max_log_size,
+            self._config_file_name,
+            orig_flow_datastore=self._orig_flow_datastore,
+            spin_pathspec=self._spin_pathspec,
+            artifacts_module=self._artifacts_module,
+            persist=self._persist,
+            skip_decorators=self._skip_decorators,
+        )
+        from_start("SpinRuntime: created worker")
+
+        poll = procpoll.make_poll()
+        fds = worker.fds()
+        for fd in fds:
+            poll.add(fd)
+
+        active_fds = set(fds)
+
+        while active_fds:
+            events = poll.poll(POLL_TIMEOUT)
+            for event in events:
+                if event.can_read:
+                    worker.read_logline(event.fd)
+                if event.is_terminated:
+                    poll.remove(event.fd)
+                    active_fds.remove(event.fd)
+        from_start("SpinRuntime: read loglines")
+        returncode = worker.terminate()
+        from_start("SpinRuntime: worker terminated")
+        if returncode != 0:
+            raise TaskFailed(self.task, f"Task failed with return code {returncode}")
+        else:
+            self._logger("Task finished successfully.", system_msg=True)
+
+
 class NativeRuntime(object):
     def __init__(
         self,
@@ -1769,8 +2025,27 @@ class CLIArgs(object):
     for step execution in StepDecorator.runtime_step_cli().
     """
 
-    def __init__(self, task):
+    def __init__(
+        self,
+        task,
+        orig_flow_datastore=None,
+        spin_pathspec=None,
+        artifacts_module=None,
+        persist=True,
+        skip_decorators=False,
+    ):
         self.task = task
+        if orig_flow_datastore is not None:
+            self.orig_flow_datastore = "%s@%s" % (
+                orig_flow_datastore.TYPE,
+                orig_flow_datastore.datastore_root,
+            )
+        else:
+            self.orig_flow_datastore = None
+        self.spin_pathspec = spin_pathspec
+        self.artifacts_module = artifacts_module
+        self.persist = persist
+        self.skip_decorators = skip_decorators
         self.entrypoint = list(task.entrypoint)
         step_obj = getattr(self.task.flow, self.task.step)
         self.top_level_options = {
@@ -1808,21 +2083,49 @@ class CLIArgs(object):
                 (k, ConfigInput.make_key_name(k)) for k in configs
             ]
 
+        if spin_pathspec:
+            self.spin_args()
+        else:
+            self.default_args()
+
+    def default_args(self):
         self.commands = ["step"]
         self.command_args = [self.task.step]
         self.command_options = {
-            "run-id": task.run_id,
-            "task-id": task.task_id,
-            "input-paths": compress_list(task.input_paths),
-            "split-index": task.split_index,
-            "retry-count": task.retries,
-            "max-user-code-retries": task.user_code_retries,
-            "tag": task.tags,
+            "run-id": self.task.run_id,
+            "task-id": self.task.task_id,
+            "input-paths": compress_list(self.task.input_paths),
+            "split-index": self.task.split_index,
+            "retry-count": self.task.retries,
+            "max-user-code-retries": self.task.user_code_retries,
+            "tag": self.task.tags,
             "namespace": get_namespace() or "",
-            "ubf-context": task.ubf_context,
+            "ubf-context": self.task.ubf_context,
         }
         self.env = {}
 
+    def spin_args(self):
+        self.commands = ["spin-step"]
+        self.command_args = [self.task.step]
+
+        self.command_options = {
+            "run-id": self.task.run_id,
+            "task-id": self.task.task_id,
+            "input-paths": compress_list(self.task.input_paths),
+            "split-index": self.task.split_index,
+            "retry-count": self.task.retries,
+            "max-user-code-retries": self.task.user_code_retries,
+            "namespace": get_namespace() or "",
+            "orig-flow-datastore": self.orig_flow_datastore,
+            "artifacts-module": self.artifacts_module,
+            "skip-decorators": self.skip_decorators,
+        }
+        if self.persist:
+            self.command_options["persist"] = True
+        else:
+            self.command_options["no-persist"] = True
+        self.env = {}
+
     def get_args(self):
         # TODO: Make one with dict_to_cli_options; see cli_args.py for more detail
         def _options(mapping):
@@ -1861,9 +2164,24 @@ class CLIArgs(object):
 
 
 class Worker(object):
-    def __init__(self, task, max_logs_size, config_file_name):
+    def __init__(
+        self,
+        task,
+        max_logs_size,
+        config_file_name,
+        orig_flow_datastore=None,
+        spin_pathspec=None,
+        artifacts_module=None,
+        persist=True,
+        skip_decorators=False,
+    ):
         self.task = task
         self._config_file_name = config_file_name
+        self._orig_flow_datastore = orig_flow_datastore
+        self._spin_pathspec = spin_pathspec
+        self._artifacts_module = artifacts_module
+        self._skip_decorators = skip_decorators
+        self._persist = persist
         self._proc = self._launch()
 
         if task.retries > task.user_code_retries:
@@ -1895,7 +2213,14 @@ class Worker(object):
         # not it is properly shut down)
 
     def _launch(self):
-        args = CLIArgs(self.task)
+        args = CLIArgs(
+            self.task,
+            orig_flow_datastore=self._orig_flow_datastore,
+            spin_pathspec=self._spin_pathspec,
+            artifacts_module=self._artifacts_module,
+            persist=self._persist,
+            skip_decorators=self._skip_decorators,
+        )
         env = dict(os.environ)
 
         if self.task.clone_run_id:
@@ -1928,6 +2253,7 @@ class Worker(object):
         # by read_logline() below that relies on readline() not blocking
         # print('running', args)
         cmdline = args.get_args()
+        from_start(f"Command line: {' '.join(cmdline)}")
         debug.subcommand_exec(cmdline)
         return subprocess.Popen(
             cmdline,
@@ -2050,13 +2376,14 @@ class Worker(object):
                     else:
                         self.emit_log(b"Task failed.", self._stderr, system_msg=True)
             else:
-                num = self.task.results["_foreach_num_splits"]
-                if num:
-                    self.task.log(
-                        "Foreach yields %d child steps." % num,
-                        system_msg=True,
-                        pid=self._proc.pid,
-                    )
+                if not self._spin_pathspec:
+                    num = self.task.results["_foreach_num_splits"]
+                    if num:
+                        self.task.log(
+                            "Foreach yields %d child steps." % num,
+                            system_msg=True,
+                            pid=self._proc.pid,
+                        )
                 self.task.log(
                     "Task finished successfully.", system_msg=True, pid=self._proc.pid
                 )

metaflow/task.py

@@ -6,14 +6,15 @@ import os
 import time
 import traceback
 
-
 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 (
@@ -49,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
@@ -58,6 +61,8 @@ 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 = []
@@ -150,6 +155,7 @@ class MetaflowTask(object):
             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
@@ -233,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
@@ -261,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):
@@ -272,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(
@@ -546,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)
@@ -604,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":
@@ -612,17 +646,20 @@ 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 = (
@@ -681,7 +718,7 @@ class MetaflowTask(object):
                         ),
                     ]
                 )
-
+        from_start("MetaflowTask: finished input processing")
         self.metadata.register_metadata(
             run_id,
             step_name,
@@ -735,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"):
@@ -756,7 +796,6 @@ 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:
 
@@ -815,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,
@@ -845,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, orig_step_func, input_obj)
                 else:
                     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,
@@ -893,6 +940,7 @@ class MetaflowTask(object):
                     raise
 
             finally:
+                from_start("MetaflowTask: decorators finalized")
                 if self.ubf_context == UBF_CONTROL:
                     self._finalize_control_task()
 
@@ -932,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()
@@ -961,3 +1009,4 @@ class MetaflowTask(object):
                     name="duration",
                     payload={**task_payload, "msg": str(duration)},
                 )
+                from_start("MetaflowTask: task run completed")

metaflow/user_configs/config_parameters.py

@@ -379,7 +379,7 @@ class DelayEvaluator(collections.abc.Mapping):
                 to_eval_expr,
                 self._globals or globals(),
                 {
-                    k: ConfigValue(v)
+                    k: ConfigValue(v) if v is not None else None
                     for k, v in flow_cls._flow_state.get(_FlowState.CONFIGS, {}).items()
                 },
             )
@@ -507,9 +507,7 @@ class Config(Parameter, collections.abc.Mapping):
         self._delayed_evaluator = None
 
     def load_parameter(self, v):
-        if v is None:
-            return None
-        return ConfigValue(v)
+        return ConfigValue(v) if v is not None else None
 
     def _store_value(self, v: Any) -> None:
         self._computed_value = v

metaflow/user_decorators/mutable_flow.py

@@ -114,7 +114,7 @@ class MutableFlow:
         for name, value in self._flow_cls._flow_state.get(
             _FlowState.CONFIGS, {}
         ).items():
-            r = name, ConfigValue(value)
+            r = name, ConfigValue(value) if value is not None else None
             debug.userconf_exec("Mutable flow yielding config: %s" % str(r))
             yield r