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

metaflow/datastore/flow_datastore.py

@@ -1,10 +1,13 @@
 import itertools
 import json
+from abc import ABC, abstractmethod
 
 from .. import metaflow_config
 
 from .content_addressed_store import ContentAddressedStore
 from .task_datastore import TaskDataStore
+from .spin_datastore import SpinTaskDatastore
+from ..metaflow_profile import from_start
 
 
 class FlowDataStore(object):
@@ -63,10 +66,16 @@ class FlowDataStore(object):
             self._storage_impl.path_join(self.flow_name, "data"), self._storage_impl
         )
 
+        # Private
+        self._metadata_cache = None
+
     @property
     def datastore_root(self):
         return self._storage_impl.datastore_root
 
+    def set_metadata_cache(self, cache):
+        self._metadata_cache = cache
+
     def get_task_datastores(
         self,
         run_id=None,
@@ -76,6 +85,9 @@ class FlowDataStore(object):
         attempt=None,
         include_prior=False,
         mode="r",
+        join_type=None,
+        orig_flow_datastore=None,
+        spin_artifacts=None,
     ):
         """
         Return a list of TaskDataStore for a subset of the tasks.
@@ -95,7 +107,7 @@ class FlowDataStore(object):
             Steps to get the tasks from. If run_id is specified, this
             must also be specified, by default None
         pathspecs : List[str], optional
-            Full task specs (run_id/step_name/task_id). Can be used instead of
+            Full task specs (run_id/step_name/task_id[/attempt]). Can be used instead of
             specifying run_id and steps, by default None
         allow_not_done : bool, optional
             If True, returns the latest attempt of a task even if that attempt
@@ -106,6 +118,16 @@ class FlowDataStore(object):
             If True, returns all attempts up to and including attempt.
         mode : str, default "r"
             Mode to initialize the returned TaskDataStores in.
+        join_type : str, optional, default None
+            If specified, the join type for the task. This is used to determine
+            the user specified artifacts for the task in case of a spin task.
+        orig_flow_datastore : MetadataProvider, optional, default None
+            The metadata provider in case of a spin task. If provided, the
+            returned TaskDataStore will be a SpinTaskDatastore instead of a
+            TaskDataStore.
+        spin_artifacts : Dict[str, Any], optional, default None
+            Artifacts provided by user that can override the artifacts fetched via the
+            spin pathspec.
 
         Returns
         -------
@@ -145,7 +167,13 @@ class FlowDataStore(object):
         if attempt is not None and attempt <= metaflow_config.MAX_ATTEMPTS - 1:
             attempt_range = range(attempt + 1) if include_prior else [attempt]
         for task_url in task_urls:
-            for attempt in attempt_range:
+            task_splits = task_url.split("/")
+            # Usually it is flow, run, step, task (so 4 components) -- if we have a
+            # fifth one, there is a specific attempt number listed as well.
+            task_attempt_range = attempt_range
+            if len(task_splits) == 5:
+                task_attempt_range = [int(task_splits[4])]
+            for attempt in task_attempt_range:
                 for suffix in [
                     TaskDataStore.METADATA_DATA_SUFFIX,
                     TaskDataStore.METADATA_ATTEMPT_SUFFIX,
@@ -198,7 +226,18 @@ class FlowDataStore(object):
                 else (latest_started_attempts & done_attempts)
             )
         latest_to_fetch = [
-            (v[0], v[1], v[2], v[3], data_objs.get(v), mode, allow_not_done)
+            (
+                v[0],
+                v[1],
+                v[2],
+                v[3],
+                data_objs.get(v),
+                mode,
+                allow_not_done,
+                join_type,
+                orig_flow_datastore,
+                spin_artifacts,
+            )
             for v in latest_to_fetch
         ]
         return list(itertools.starmap(self.get_task_datastore, latest_to_fetch))
@@ -212,8 +251,63 @@ class FlowDataStore(object):
         data_metadata=None,
         mode="r",
         allow_not_done=False,
+        join_type=None,
+        orig_flow_datastore=None,
+        spin_artifacts=None,
+        persist=True,
     ):
-        return TaskDataStore(
+        if orig_flow_datastore is not None:
+            # In spin step subprocess, use SpinTaskDatastore for accessing artifacts
+            if join_type is not None:
+                # If join_type is specified, we need to use the artifacts corresponding
+                # to that particular join index, specified by the parent task pathspec.
+                spin_artifacts = spin_artifacts.get(
+                    f"{run_id}/{step_name}/{task_id}", {}
+                )
+            from_start(
+                "FlowDataStore: get_task_datastore for spin task for type %s %s metadata"
+                % (self.TYPE, "without" if data_metadata is None else "with")
+            )
+            # Get the task datastore for the spun task.
+            orig_datastore = orig_flow_datastore.get_task_datastore(
+                run_id,
+                step_name,
+                task_id,
+                attempt=attempt,
+                data_metadata=data_metadata,
+                mode=mode,
+                allow_not_done=allow_not_done,
+                persist=persist,
+            )
+
+            return SpinTaskDatastore(
+                self.flow_name,
+                run_id,
+                step_name,
+                task_id,
+                orig_datastore,
+                spin_artifacts,
+            )
+
+        cache_hit = False
+        if (
+            self._metadata_cache is not None
+            and data_metadata is None
+            and attempt is not None
+            and allow_not_done is False
+        ):
+            # If we have a metadata cache, we can try to load the metadata
+            # from the cache if it is not provided.
+            data_metadata = self._metadata_cache.load_metadata(
+                run_id, step_name, task_id, attempt
+            )
+            cache_hit = data_metadata is not None
+
+        from_start(
+            "FlowDataStore: get_task_datastore for regular task for type %s %s metadata"
+            % (self.TYPE, "without" if data_metadata is None else "with")
+        )
+        task_datastore = TaskDataStore(
             self,
             run_id,
             step_name,
@@ -222,8 +316,23 @@ class FlowDataStore(object):
             data_metadata=data_metadata,
             mode=mode,
             allow_not_done=allow_not_done,
+            persist=persist,
         )
 
+        # Only persist in cache if it is non-changing (so done only) and we have
+        # a non-None attempt
+        if (
+            not cache_hit
+            and self._metadata_cache is not None
+            and allow_not_done is False
+            and attempt is not None
+        ):
+            self._metadata_cache.store_metadata(
+                run_id, step_name, task_id, attempt, task_datastore.ds_metadata
+            )
+
+        return task_datastore
+
     def save_data(self, data_iter, len_hint=0):
         """Saves data to the underlying content-addressed store
 
@@ -265,3 +374,13 @@ class FlowDataStore(object):
         """
         for key, blob in self.ca_store.load_blobs(keys, force_raw=force_raw):
             yield key, blob
+
+
+class MetadataCache(ABC):
+    @abstractmethod
+    def load_metadata(self, run_id, step_name, task_id, attempt):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def store_metadata(self, run_id, step_name, task_id, attempt, metadata_dict):
+        raise NotImplementedError()

metaflow/datastore/spin_datastore.py

@@ -0,0 +1,91 @@
+from typing import Dict, Any
+from .task_datastore import TaskDataStore, require_mode
+from ..metaflow_profile import from_start
+
+
+class SpinTaskDatastore(object):
+    def __init__(
+        self,
+        flow_name: str,
+        run_id: str,
+        step_name: str,
+        task_id: str,
+        orig_datastore: TaskDataStore,
+        spin_artifacts: Dict[str, Any],
+    ):
+        """
+        SpinTaskDatastore is a datastore for a task that is used to retrieve
+        artifacts and attributes for a spin step. It uses the task pathspec
+        from a previous execution of the step to access the artifacts and attributes.
+
+        Parameters:
+        -----------
+        flow_name : str
+            Name of the flow
+        run_id : str
+            Run ID of the flow
+        step_name : str
+            Name of the step
+        task_id : str
+            Task ID of the step
+        orig_datastore : TaskDataStore
+            The datastore for the underlying task that is being spun.
+        spin_artifacts : Dict[str, Any]
+            User provided artifacts that are to be used in the spin task. This is a dictionary
+            where keys are artifact names and values are the actual data or metadata.
+        """
+        self.flow_name = flow_name
+        self.run_id = run_id
+        self.step_name = step_name
+        self.task_id = task_id
+        self.orig_datastore = orig_datastore
+        self.spin_artifacts = spin_artifacts
+        self._task = None
+
+        # Update _objects and _info in order to persist artifacts
+        # See `persist` method in `TaskDatastore` for more details
+        self._objects = self.orig_datastore._objects.copy()
+        self._info = self.orig_datastore._info.copy()
+
+        # We strip out some of the control ones
+        for key in ("_transition",):
+            if key in self._objects:
+                del self._objects[key]
+                del self._info[key]
+
+        from_start("SpinTaskDatastore: Initialized artifacts")
+
+    @require_mode(None)
+    def __getitem__(self, name):
+        try:
+            # Check if it's an artifact in the spin_artifacts
+            return self.spin_artifacts[name]
+        except KeyError:
+            try:
+                # Check if it's an attribute of the task
+                # _foreach_stack, _foreach_index, ...
+                return self.orig_datastore[name]
+            except (KeyError, AttributeError) as e:
+                raise KeyError(
+                    f"Attribute '{name}' not found in the previous execution "
+                    f"of the tasks for `{self.step_name}`."
+                ) from e
+
+    @require_mode(None)
+    def is_none(self, name):
+        val = self.__getitem__(name)
+        return val is None
+
+    @require_mode(None)
+    def __contains__(self, name):
+        try:
+            _ = self.__getitem__(name)
+            return True
+        except KeyError:
+            return False
+
+    @require_mode(None)
+    def items(self):
+        if self._objects:
+            return self._objects.items()
+        return {}

metaflow/datastore/task_datastore.py

@@ -6,6 +6,7 @@ import time
 
 from functools import wraps
 from io import BufferedIOBase, FileIO, RawIOBase
+from typing import List, Optional
 from types import MethodType, FunctionType
 
 from .. import metaflow_config
@@ -98,6 +99,7 @@ class TaskDataStore(object):
         data_metadata=None,
         mode="r",
         allow_not_done=False,
+        persist=True,
     ):
         self._storage_impl = flow_datastore._storage_impl
         self.TYPE = self._storage_impl.TYPE
@@ -113,6 +115,7 @@ class TaskDataStore(object):
         self._attempt = attempt
         self._metadata = flow_datastore.metadata
         self._parent = flow_datastore
+        self._persist = persist
 
         # The GZIP encodings are for backward compatibility
         self._encodings = {"pickle-v2", "gzip+pickle-v2"}
@@ -148,6 +151,8 @@ class TaskDataStore(object):
                     )
                     if self.has_metadata(check_meta, add_attempt=False):
                         max_attempt = i
+                    elif max_attempt is not None:
+                        break
                 if self._attempt is None:
                     self._attempt = max_attempt
                 elif max_attempt is None or self._attempt > max_attempt:
@@ -253,6 +258,72 @@ class TaskDataStore(object):
         """
         self.save_metadata({self.METADATA_ATTEMPT_SUFFIX: {"time": time.time()}})
 
+    @only_if_not_done
+    @require_mode("w")
+    def transfer_artifacts(
+        self, other_datastore: "TaskDataStore", names: Optional[List[str]] = None
+    ):
+        """
+        Copies the blobs from other_datastore to this datastore if the datastore roots
+        are different.
+
+        This is used specifically for spin so we can bring in artifacts from the original
+        datastore.
+
+        Parameters
+        ----------
+        other_datastore : TaskDataStore
+            Other datastore from which to copy artifacts from
+        names : List[str], optional, default None
+            If provided, only transfer the artifacts with these names. If None,
+            transfer all artifacts from the other datastore.
+        """
+        if (
+            other_datastore.TYPE == self.TYPE
+            and other_datastore._storage_impl.datastore_root
+            == self._storage_impl.datastore_root
+        ):
+            # Nothing to transfer -- artifacts are already saved properly
+            return
+
+        # Determine which artifacts need to be transferred
+        if names is None:
+            # Transfer all artifacts from other datastore
+            artifacts_to_transfer = list(other_datastore._objects.keys())
+        else:
+            # Transfer only specified artifacts
+            artifacts_to_transfer = [
+                name for name in names if name in other_datastore._objects
+            ]
+
+        if not artifacts_to_transfer:
+            return
+
+        # Get SHA keys for artifacts to transfer
+        shas_to_transfer = [
+            other_datastore._objects[name] for name in artifacts_to_transfer
+        ]
+
+        # Check which blobs are missing locally
+        missing_shas = []
+        for sha in shas_to_transfer:
+            local_path = self._ca_store._storage_impl.path_join(
+                self._ca_store._prefix, sha[:2], sha
+            )
+            if not self._ca_store._storage_impl.is_file([local_path])[0]:
+                missing_shas.append(sha)
+
+        if not missing_shas:
+            return  # All blobs already exist locally
+
+        # Load blobs from other datastore in transfer mode
+        transfer_blobs = other_datastore._ca_store.load_blobs(
+            missing_shas, is_transfer=True
+        )
+
+        # Save blobs to local datastore in transfer mode
+        self._ca_store.save_blobs(transfer_blobs, is_transfer=True)
+
     @only_if_not_done
     @require_mode("w")
     def save_artifacts(self, artifacts_iter, len_hint=0):
@@ -683,14 +754,16 @@ class TaskDataStore(object):
         flow : FlowSpec
             Flow to persist
         """
+        if not self._persist:
+            return
 
         if flow._datastore:
             self._objects.update(flow._datastore._objects)
             self._info.update(flow._datastore._info)
 
-        # we create a list of valid_artifacts in advance, outside of
-        # artifacts_iter, so we can provide a len_hint below
+        # Scan flow object FIRST
         valid_artifacts = []
+        current_artifact_names = set()
         for var in dir(flow):
             if var.startswith("__") or var in flow._EPHEMERAL:
                 continue
@@ -707,6 +780,16 @@ class TaskDataStore(object):
                 or isinstance(val, Parameter)
             ):
                 valid_artifacts.append((var, val))
+                current_artifact_names.add(var)
+
+        # Transfer ONLY artifacts that aren't being overridden
+        if hasattr(flow._datastore, "orig_datastore"):
+            parent_artifacts = set(flow._datastore._objects.keys())
+            unchanged_artifacts = parent_artifacts - current_artifact_names
+            if unchanged_artifacts:
+                self.transfer_artifacts(
+                    flow._datastore.orig_datastore, names=list(unchanged_artifacts)
+                )
 
         def artifacts_iter():
             # we consume the valid_artifacts list destructively to
@@ -722,6 +805,7 @@ class TaskDataStore(object):
                     delattr(flow, var)
                 yield var, val
 
+        # Save current artifacts
         self.save_artifacts(artifacts_iter(), len_hint=len(valid_artifacts))
 
     @only_if_not_done

metaflow/decorators.py

@@ -27,7 +27,7 @@ from .user_decorators.user_step_decorator import (
     UserStepDecoratorBase,
     UserStepDecoratorMeta,
 )
-
+from .metaflow_config import SPIN_ALLOWED_DECORATORS
 from metaflow._vendor import click
 
 
@@ -658,6 +658,50 @@ def _attach_decorators_to_step(step, decospecs):
             step_deco.add_or_raise(step, False, 1, None)
 
 
+def _should_skip_decorator_for_spin(
+    deco, is_spin, skip_decorators, logger, decorator_type="decorator"
+):
+    """
+    Determine if a decorator should be skipped for spin steps.
+
+    Parameters:
+    -----------
+    deco : Decorator
+        The decorator instance to check
+    is_spin : bool
+        Whether this is a spin step
+    skip_decorators : bool
+        Whether to skip all decorators
+    logger : callable
+        Logger function for warnings
+    decorator_type : str
+        Type of decorator ("Flow decorator" or "Step decorator") for logging
+
+    Returns:
+    --------
+    bool
+        True if the decorator should be skipped, False otherwise
+    """
+    if not is_spin:
+        return False
+
+    # Skip all decorator hooks if skip_decorators is True
+    if skip_decorators:
+        return True
+
+    # Run decorator hooks for spin steps only if they are in the whitelist
+    if deco.name not in SPIN_ALLOWED_DECORATORS:
+        logger(
+            f"[Warning] Ignoring {decorator_type} '{deco.name}' as it is not supported in spin steps.",
+            system_msg=True,
+            timestamp=False,
+            bad=True,
+        )
+        return True
+
+    return False
+
+
 def _init(flow, only_non_static=False):
     for decorators in flow._flow_decorators.values():
         for deco in decorators:
@@ -673,7 +717,16 @@ def _init(flow, only_non_static=False):
 
 
 def _init_flow_decorators(
-    flow, graph, environment, flow_datastore, metadata, logger, echo, deco_options
+    flow,
+    graph,
+    environment,
+    flow_datastore,
+    metadata,
+    logger,
+    echo,
+    deco_options,
+    is_spin=False,
+    skip_decorators=False,
 ):
     # Since all flow decorators are stored as `{key:[deco]}` we iterate through each of them.
     for decorators in flow._flow_decorators.values():
@@ -702,6 +755,10 @@ def _init_flow_decorators(
                 for option, option_info in deco.options.items()
             }
         for deco in decorators:
+            if _should_skip_decorator_for_spin(
+                deco, is_spin, skip_decorators, logger, "Flow decorator"
+            ):
+                continue
             deco.flow_init(
                 flow,
                 graph,
@@ -714,8 +771,16 @@ def _init_flow_decorators(
             )
 
 
-def _init_step_decorators(flow, graph, environment, flow_datastore, logger):
-    # NOTE: We don't need graph but keeping it for backwards compatibility with
+def _init_step_decorators(
+    flow,
+    graph,
+    environment,
+    flow_datastore,
+    logger,
+    is_spin=False,
+    skip_decorators=False,
+):
+    # NOTE: We don't need the graph but keeping it for backwards compatibility with
     # extensions that use it directly. We will remove it at some point.
 
     # We call the mutate method for both the flow and step mutators.
@@ -741,7 +806,7 @@ def _init_step_decorators(flow, graph, environment, flow_datastore, logger):
                     "expected %s but got %s" % (deco._flow_cls.__name__, cls.__name__)
                 )
             debug.userconf_exec(
-                "Evaluating flow level decorator %s (post)" % deco.__class__.__name__
+                "Evaluating flow level decorator %s (mutate)" % deco.__class__.__name__
             )
             deco.mutate(mutable_flow)
             # We reset cached_parameters on the very off chance that the user added
@@ -759,7 +824,7 @@ def _init_step_decorators(flow, graph, environment, flow_datastore, logger):
 
             if isinstance(deco, StepMutator):
                 debug.userconf_exec(
-                    "Evaluating step level decorator %s (post) for %s"
+                    "Evaluating step level decorator %s for %s (mutate)"
                     % (deco.__class__.__name__, step.name)
                 )
                 deco.mutate(
@@ -785,6 +850,10 @@ def _init_step_decorators(flow, graph, environment, flow_datastore, logger):
 
     for step in flow:
         for deco in step.decorators:
+            if _should_skip_decorator_for_spin(
+                deco, is_spin, skip_decorators, logger, "Step decorator"
+            ):
+                continue
             deco.step_init(
                 flow,
                 graph,