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

metaflow/datastore/task_datastore.py

@@ -6,11 +6,12 @@ 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
 from ..exception import MetaflowInternalError
-from ..metadata import DataArtifact, MetaDatum
+from ..metadata_provider import DataArtifact, MetaDatum
 from ..parameters import Parameter
 from ..util import Path, is_stringish, to_fileobj
 
@@ -98,8 +99,8 @@ 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
         self._ca_store = flow_datastore.ca_store
@@ -114,11 +115,12 @@ 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"}
         ver = sys.version_info[0] * 10 + sys.version_info[1]
-        if ver >= 34:
+        if ver >= 36:
             self._encodings.add("pickle-v4")
             self._encodings.add("gzip+pickle-v4")
 
@@ -149,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:
@@ -173,6 +177,26 @@ class TaskDataStore(object):
                 if data_obj is not None:
                     self._objects = data_obj.get("objects", {})
                     self._info = data_obj.get("info", {})
+        elif self._mode == "d":
+            self._objects = {}
+            self._info = {}
+
+            if self._attempt is None:
+                for i in range(metaflow_config.MAX_ATTEMPTS):
+                    check_meta = self._metadata_name_for_attempt(
+                        self.METADATA_ATTEMPT_SUFFIX, i
+                    )
+                    if self.has_metadata(check_meta, add_attempt=False):
+                        self._attempt = i
+
+            # Do not allow destructive operations on the datastore if attempt is still in flight
+            # and we explicitly did not allow operating on running tasks.
+            if not allow_not_done and not self.has_metadata(self.METADATA_DONE_SUFFIX):
+                raise DataException(
+                    "No completed attempts of the task was found for task '%s'"
+                    % self._path
+                )
+
         else:
             raise DataException("Unknown datastore mode: '%s'" % self._mode)
 
@@ -203,6 +227,9 @@ class TaskDataStore(object):
     @property
     def pathspec_index(self):
         idxstr = ",".join(map(str, (f.index for f in self["_foreach_stack"])))
+        if "_iteration_stack" in self:
+            itrstr = ",".join(map(str, (f for f in self["_iteration_stack"])))
+            return "%s/%s[%s][%s]" % (self._run_id, self._step_name, idxstr, itrstr)
         return "%s/%s[%s]" % (self._run_id, self._step_name, idxstr)
 
     @property
@@ -233,7 +260,73 @@ class TaskDataStore(object):
 
     @only_if_not_done
     @require_mode("w")
-    def save_artifacts(self, artifacts_iter, force_v4=False, len_hint=0):
+    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):
         """
         Saves Metaflow Artifacts (Python objects) to the datastore and stores
         any relevant metadata needed to retrieve them.
@@ -249,11 +342,6 @@ class TaskDataStore(object):
         artifacts : Iterator[(string, object)]
             Iterator over the human-readable name of the object to save
             and the object itself
-        force_v4 : boolean or Dict[string -> boolean]
-            Indicates whether the artifact should be pickled using the v4
-            version of pickle. If a single boolean, applies to all artifacts.
-            If a dictionary, applies to the object named only. Defaults to False
-            if not present or not specified
         len_hint: integer
             Estimated number of items in artifacts_iter
         """
@@ -261,40 +349,24 @@ class TaskDataStore(object):
 
         def pickle_iter():
             for name, obj in artifacts_iter:
-                do_v4 = (
-                    force_v4 and force_v4
-                    if isinstance(force_v4, bool)
-                    else force_v4.get(name, False)
-                )
-                if do_v4:
-                    encode_type = "gzip+pickle-v4"
-                    if encode_type not in self._encodings:
-                        raise DataException(
-                            "Artifact *%s* requires a serialization encoding that "
-                            "requires Python 3.4 or newer." % name
-                        )
+                encode_type = "gzip+pickle-v4"
+                if encode_type in self._encodings:
                     try:
                         blob = pickle.dumps(obj, protocol=4)
                     except TypeError as e:
-                        raise UnpicklableArtifactException(name)
+                        raise UnpicklableArtifactException(name) from e
                 else:
                     try:
                         blob = pickle.dumps(obj, protocol=2)
                         encode_type = "gzip+pickle-v2"
-                    except (SystemError, OverflowError):
-                        encode_type = "gzip+pickle-v4"
-                        if encode_type not in self._encodings:
-                            raise DataException(
-                                "Artifact *%s* is very large (over 2GB). "
-                                "You need to use Python 3.4 or newer if you want to "
-                                "serialize large objects." % name
-                            )
-                        try:
-                            blob = pickle.dumps(obj, protocol=4)
-                        except TypeError as e:
-                            raise UnpicklableArtifactException(name)
+                    except (SystemError, OverflowError) as e:
+                        raise DataException(
+                            "Artifact *%s* is very large (over 2GB). "
+                            "You need to use Python 3.6 or newer if you want to "
+                            "serialize large objects." % name
+                        ) from e
                     except TypeError as e:
-                        raise UnpicklableArtifactException(name)
+                        raise UnpicklableArtifactException(name) from e
 
                 self._info[name] = {
                     "size": len(blob),
@@ -353,7 +425,7 @@ class TaskDataStore(object):
                 encode_type = "gzip+pickle-v2"
             if encode_type not in self._encodings:
                 raise DataException(
-                    "Python 3.4 or later is required to load artifact '%s'" % name
+                    "Python 3.6 or later is required to load artifact '%s'" % name
                 )
             else:
                 to_load[self._objects[name]].append(name)
@@ -361,7 +433,7 @@ class TaskDataStore(object):
         # We assume that if we have one "old" style artifact, all of them are
         # like that which is an easy assumption to make since artifacts are all
         # stored by the same implementation of the datastore for a given task.
-        for (key, blob) in self._ca_store.load_blobs(to_load.keys()):
+        for key, blob in self._ca_store.load_blobs(to_load.keys()):
             names = to_load[key]
             for name in names:
                 # We unpickle everytime to have fully distinct objects (the user
@@ -682,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
@@ -706,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
@@ -721,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
@@ -750,6 +835,36 @@ class TaskDataStore(object):
                 to_store_dict[n] = data
         self._save_file(to_store_dict)
 
+    @require_mode("d")
+    def scrub_logs(self, logsources, stream, attempt_override=None):
+        path_logsources = {
+            self._metadata_name_for_attempt(
+                self._get_log_location(s, stream),
+                attempt_override=attempt_override,
+            ): s
+            for s in logsources
+        }
+
+        # Legacy log paths
+        legacy_log = self._metadata_name_for_attempt(
+            "%s.log" % stream, attempt_override
+        )
+        path_logsources[legacy_log] = stream
+
+        existing_paths = [
+            path
+            for path in path_logsources.keys()
+            if self.has_metadata(path, add_attempt=False)
+        ]
+
+        # Replace log contents with [REDACTED source stream]
+        to_store_dict = {
+            path: bytes("[REDACTED %s %s]" % (path_logsources[path], stream), "utf-8")
+            for path in existing_paths
+        }
+
+        self._save_file(to_store_dict, add_attempt=False, allow_overwrite=True)
+
     @require_mode("r")
     def load_log_legacy(self, stream, attempt_override=None):
         """

metaflow/debug.py

@@ -42,6 +42,11 @@ class Debug(object):
         filename = inspect.stack()[1][1]
         print("debug[%s %s:%s]: %s" % (typ, filename, lineno, s), file=sys.stderr)
 
+    def __getattr__(self, name):
+        # Small piece of code to get pyright and other linters to recognize that there
+        # are dynamic attributes.
+        return getattr(self, name)
+
     def noop(self, args):
         pass