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

metaflow/client/core.py

@@ -32,11 +32,12 @@ from metaflow.exception import (
 from metaflow.includefile import IncludedFile
 from metaflow.metaflow_config import DEFAULT_METADATA, MAX_ATTEMPTS
 from metaflow.metaflow_environment import MetaflowEnvironment
+from metaflow.package import MetaflowPackage
+from metaflow.packaging_sys import ContentType
 from metaflow.plugins import ENVIRONMENTS, METADATA_PROVIDERS
 from metaflow.unbounded_foreach import CONTROL_TASK_TAG
 from metaflow.util import cached_property, is_stringish, resolve_identity, to_unicode
 
-from ..info_file import INFO_FILE
 from .filecache import FileCache
 
 if TYPE_CHECKING:
@@ -206,6 +207,20 @@ def default_namespace() -> str:
     return get_namespace()
 
 
+def inspect_spin(datastore_root: str = "."):
+    """
+    Set metadata provider to spin metadata so that users can inspect spin
+    steps, tasks, and artifacts.
+
+    Parameters
+    ----------
+    datastore_root : str, default "."
+        The root path to the spin datastore.
+    """
+    metadata_str = f"spin@{datastore_root}"
+    metadata(metadata_str)
+
+
 MetaflowArtifacts = NamedTuple
 
 
@@ -276,6 +291,7 @@ class MetaflowObject(object):
         self._attempt = attempt
         self._current_namespace = _current_namespace or get_namespace()
         self._namespace_check = _namespace_check
+
         # If the current namespace is False, we disable checking for namespace for this
         # and all children objects. Not setting namespace_check to False has the consequence
         # of preventing access to children objects after the namespace changes
@@ -302,7 +318,7 @@ class MetaflowObject(object):
             # distinguish between "attempt will happen" and "no such
             # attempt exists".
 
-        if pathspec:
+        if pathspec and _object is None:
             ids = pathspec.split("/")
 
             if self._NAME == "flow" and len(ids) != 1:
@@ -816,20 +832,29 @@ class MetaflowCode(object):
         self._path = info["location"]
         self._ds_type = info["ds_type"]
         self._sha = info["sha"]
+        self._code_metadata = info.get(
+            "metadata",
+            '{"version": 0, "archive_format": "tgz", "mfcontent_version": 0}',
+        )
+
+        self._backend = MetaflowPackage.get_backend(self._code_metadata)
 
         if filecache is None:
             filecache = FileCache()
         _, blobdata = filecache.get_data(
             self._ds_type, self._flow_name, self._path, self._sha
         )
-        code_obj = BytesIO(blobdata)
-        self._tar = tarfile.open(fileobj=code_obj, mode="r:gz")
-        # The JSON module in Python3 deals with Unicode. Tar gives bytes.
-        info_str = (
-            self._tar.extractfile(os.path.basename(INFO_FILE)).read().decode("utf-8")
-        )
-        self._info = json.loads(info_str)
-        self._flowspec = self._tar.extractfile(self._info["script"]).read()
+        self._code_obj = BytesIO(blobdata)
+        self._info = MetaflowPackage.cls_get_info(self._code_metadata, self._code_obj)
+        self._code_obj.seek(0)
+        if self._info:
+            self._flowspec = MetaflowPackage.cls_get_content(
+                self._code_metadata, self._code_obj, self._info["script"]
+            )
+            self._code_obj.seek(0)
+        else:
+            raise MetaflowInternalError("Code package metadata is invalid.")
+        self._tarball = None
 
     @property
     def path(self) -> str:
@@ -877,7 +902,14 @@ class MetaflowCode(object):
         TarFile
             TarFile for everything in this code package
         """
-        return self._tar
+        # We only return one tarball because the different TarFile objects share
+        # a common bytes buffer (self._code_obj).
+        if self._tarball is not None:
+            return self._tarball
+        if self._backend.type == "tgz":
+            self._tarball = self._backend.cls_open(self._code_obj)
+            return self._tarball
+        raise RuntimeError("Archive is not a tarball")
 
     def extract(self) -> TemporaryDirectory:
         """
@@ -908,27 +940,15 @@ class MetaflowCode(object):
             The directory and its contents are automatically deleted when
             this object is garbage collected.
         """
-        exclusions = [
-            "metaflow/",
-            "metaflow_extensions/",
-            "INFO",
-            "CONFIG_PARAMETERS",
-            "conda.manifest",
-            # This file is created when using the conda/pypi features available in
-            # nflx-metaflow-extensions: https://github.com/Netflix/metaflow-nflx-extensions
-            "condav2-1.cnd",
-        ]
-        members = [
-            m
-            for m in self.tarball.getmembers()
-            if not any(
-                (x.endswith("/") and m.name.startswith(x)) or (m.name == x)
-                for x in exclusions
-            )
-        ]
-
         tmp = TemporaryDirectory()
-        self.tarball.extractall(tmp.name, members)
+        # We save the position we are in _code_obj -- in case tarball is using it at
+        # the same time -- so we can reset it to not perturb tarball.
+        pos = self._code_obj.tell()
+        self._code_obj.seek(0)
+        MetaflowPackage.cls_extract_into(
+            self._code_metadata, self._code_obj, tmp.name, ContentType.USER_CONTENT
+        )
+        self._code_obj.seek(pos)
         return tmp
 
     @property
@@ -1184,149 +1204,197 @@ class Task(MetaflowObject):
     _PARENT_CLASS = "step"
     _CHILD_CLASS = "artifact"
 
-    def __init__(self, *args, **kwargs):
-        super(Task, self).__init__(*args, **kwargs)
-
     def _iter_filter(self, x):
         # exclude private data artifacts
         return x.id[0] != "_"
 
-    def _iter_matching_tasks(self, steps, metadata_key, metadata_pattern):
+    def _get_matching_pathspecs(self, steps, metadata_key, metadata_pattern):
         """
-        Yield tasks from specified steps matching a foreach path pattern.
+        Yield pathspecs of tasks from specified steps that match a given metadata pattern.
 
         Parameters
         ----------
         steps : List[str]
-            List of step names to search for tasks
-        pattern : str
-            Regex pattern to match foreach-indices metadata
+            List of Step objects to search for tasks.
+        metadata_key : str
+            Metadata key to filter tasks on (e.g., 'foreach-execution-path').
+        metadata_pattern : str
+            Regular expression pattern to match against the metadata value.
 
-        Returns
-        -------
-        Iterator[Task]
-            Tasks matching the foreach path pattern
+        Yields
+        ------
+        str
+            Pathspec of each task whose metadata value for the specified key matches the pattern.
         """
         flow_id, run_id, _, _ = self.path_components
-
         for step in steps:
             task_pathspecs = self._metaflow.metadata.filter_tasks_by_metadata(
-                flow_id, run_id, step.id, metadata_key, metadata_pattern
+                flow_id, run_id, step, metadata_key, metadata_pattern
             )
             for task_pathspec in task_pathspecs:
-                yield Task(pathspec=task_pathspec, _namespace_check=False)
+                yield task_pathspec
+
+    @staticmethod
+    def _get_previous_steps(graph_info, step_name):
+        # Get the parent steps
+        steps = []
+        for node_name, attributes in graph_info["steps"].items():
+            if step_name in attributes["next"]:
+                steps.append(node_name)
+        return steps
 
     @property
-    def parent_tasks(self) -> Iterator["Task"]:
+    def parent_task_pathspecs(self) -> Iterator[str]:
         """
-        Yields all parent tasks of the current task if one exists.
+        Yields pathspecs of all parent tasks of the current task.
 
         Yields
         ------
-        Task
-            Parent task of the current task
-
+        str
+            Pathspec of the parent task of the current task
         """
-        flow_id, run_id, _, _ = self.path_components
+        _, _, step_name, _ = self.path_components
+        metadata_dict = self.metadata_dict
+        graph_info = self["_graph_info"].data
 
-        steps = list(self.parent.parent_steps)
-        if not steps:
-            return []
-
-        current_path = self.metadata_dict.get("foreach-execution-path", "")
+        # Get the parent steps
+        steps = self._get_previous_steps(graph_info, step_name)
+        node_type = graph_info["steps"][step_name]["type"]
+        metadata_key = "foreach-execution-path"
+        current_path = metadata_dict.get(metadata_key)
 
         if len(steps) > 1:
             # Static join - use exact path matching
             pattern = current_path or ".*"
-            yield from self._iter_matching_tasks(
-                steps, "foreach-execution-path", pattern
-            )
-            return
-
-        # Handle single step case
-        target_task = Step(
-            f"{flow_id}/{run_id}/{steps[0].id}", _namespace_check=False
-        ).task
-        target_path = target_task.metadata_dict.get("foreach-execution-path")
-
-        if not target_path or not current_path:
-            # (Current task, "A:10") and (Parent task, "")
-            # Pattern: ".*"
-            pattern = ".*"
         else:
-            current_depth = len(current_path.split(","))
-            target_depth = len(target_path.split(","))
-
-            if current_depth < target_depth:
-                # Foreach join
-                # (Current task, "A:10,B:13") and (Parent task, "A:10,B:13,C:21")
-                # Pattern: "A:10,B:13,.*"
-                pattern = f"{current_path},.*"
-            else:
-                # Foreach split or linear step
-                # Option 1:
-                # (Current task, "A:10,B:13,C:21") and (Parent task, "A:10,B:13")
-                # Option 2:
-                # (Current task, "A:10,B:13") and (Parent task, "A:10,B:13")
-                # Pattern: "A:10,B:13"
-                pattern = ",".join(current_path.split(",")[:target_depth])
+            if not steps:
+                return  # No parent steps, yield nothing
 
-        yield from self._iter_matching_tasks(steps, "foreach-execution-path", pattern)
+            if not current_path:
+                # Current task is not part of a foreach
+                # Pattern: ".*"
+                pattern = ".*"
+            else:
+                current_depth = len(current_path.split(","))
+                if node_type == "join":
+                    # Foreach join
+                    # (Current task, "A:10,B:13") and (Parent task, "A:10,B:13,C:21")
+                    # Pattern: "A:10,B:13,.*"
+                    pattern = f"{current_path},.*"
+                else:
+                    # Foreach split or linear step
+                    # Pattern: "A:10,B:13"
+                    parent_step_type = graph_info["steps"][steps[0]]["type"]
+                    target_depth = current_depth
+                    if (
+                        parent_step_type == "split-foreach"
+                        or parent_step_type == "split-parallel"
+                    ) and current_depth == 1:
+                        # (Current task, "A:10") and (Parent task, "")
+                        pattern = ".*"
+                    else:
+                        # (Current task, "A:10,B:13,C:21") and (Parent task, "A:10,B:13")
+                        # (Current task, "A:10,B:13") and (Parent task, "A:10,B:13")
+                        if (
+                            parent_step_type == "split-foreach"
+                            or parent_step_type == "split-parallel"
+                        ):
+                            target_depth = current_depth - 1
+                        pattern = ",".join(current_path.split(",")[:target_depth])
+
+        for pathspec in self._get_matching_pathspecs(steps, metadata_key, pattern):
+            yield pathspec
 
     @property
-    def child_tasks(self) -> Iterator["Task"]:
+    def child_task_pathspecs(self) -> Iterator[str]:
         """
-        Yield all child tasks of the current task if one exists.
+        Yields pathspecs of all child tasks of the current task.
 
         Yields
         ------
-        Task
-            Child task of the current task
+        str
+            Pathspec of the child task of the current task
         """
-        flow_id, run_id, _, _ = self.path_components
-        steps = list(self.parent.child_steps)
-        if not steps:
-            return []
+        flow_id, run_id, step_name, _ = self.path_components
+        metadata_dict = self.metadata_dict
+        graph_info = self["_graph_info"].data
 
-        current_path = self.metadata_dict.get("foreach-execution-path", "")
+        # Get the child steps
+        steps = graph_info["steps"][step_name]["next"]
+
+        node_type = graph_info["steps"][step_name]["type"]
+        metadata_key = "foreach-execution-path"
+        current_path = metadata_dict.get(metadata_key)
 
         if len(steps) > 1:
             # Static split - use exact path matching
             pattern = current_path or ".*"
-            yield from self._iter_matching_tasks(
-                steps, "foreach-execution-path", pattern
-            )
-            return
-
-        # Handle single step case
-        target_task = Step(
-            f"{flow_id}/{run_id}/{steps[0].id}", _namespace_check=False
-        ).task
-        target_path = target_task.metadata_dict.get("foreach-execution-path")
-
-        if not target_path or not current_path:
-            # (Current task, "A:10") and (Child task, "")
-            # Pattern: ".*"
-            pattern = ".*"
         else:
-            current_depth = len(current_path.split(","))
-            target_depth = len(target_path.split(","))
-
-            if current_depth < target_depth:
-                # Foreach split
-                # (Current task, "A:10,B:13") and (Child task, "A:10,B:13,C:21")
-                # Pattern: "A:10,B:13,.*"
-                pattern = f"{current_path},.*"
+            if not steps:
+                return  # No child steps, yield nothing
+
+            if not current_path:
+                # Current task is not part of a foreach
+                # Pattern: ".*"
+                pattern = ".*"
             else:
-                # Foreach join or linear step
-                # Option 1:
-                # (Current task, "A:10,B:13,C:21") and (Child task, "A:10,B:13")
-                # Option 2:
-                # (Current task, "A:10,B:13") and (Child task, "A:10,B:13")
-                # Pattern: "A:10,B:13"
-                pattern = ",".join(current_path.split(",")[:target_depth])
-
-        yield from self._iter_matching_tasks(steps, "foreach-execution-path", pattern)
+                current_depth = len(current_path.split(","))
+                if node_type == "split-foreach" or node_type == "split-parallel":
+                    # Foreach split
+                    # (Current task, "A:10,B:13") and (Child task, "A:10,B:13,C:21")
+                    # Pattern: "A:10,B:13,.*"
+                    pattern = f"{current_path},.*"
+                else:
+                    # Foreach join or linear step
+                    # Pattern: "A:10,B:13"
+                    child_step_type = graph_info["steps"][steps[0]]["type"]
+
+                    # We need to know if the child step is a foreach join or a static join
+                    child_step_prev_steps = self._get_previous_steps(
+                        graph_info, steps[0]
+                    )
+                    if len(child_step_prev_steps) > 1:
+                        child_step_type = "static-join"
+                    target_depth = current_depth
+                    if child_step_type == "join" and current_depth == 1:
+                        # (Current task, "A:10") and (Child task, "")
+                        pattern = ".*"
+                    else:
+                        # (Current task, "A:10,B:13,C:21") and (Child task, "A:10,B:13")
+                        # (Current task, "A:10,B:13") and (Child task, "A:10,B:13")
+                        if child_step_type == "join":
+                            target_depth = current_depth - 1
+                        pattern = ",".join(current_path.split(",")[:target_depth])
+
+        for pathspec in self._get_matching_pathspecs(steps, metadata_key, pattern):
+            yield pathspec
+
+    @property
+    def parent_tasks(self) -> Iterator["Task"]:
+        """
+        Yields all parent tasks of the current task if one exists.
+
+        Yields
+        ------
+        Task
+            Parent task of the current task
+        """
+        parent_task_pathspecs = self.parent_task_pathspecs
+        for pathspec in parent_task_pathspecs:
+            yield Task(pathspec=pathspec, _namespace_check=False)
+
+    @property
+    def child_tasks(self) -> Iterator["Task"]:
+        """
+        Yields all child tasks of the current task if one exists.
+
+        Yields
+        ------
+        Task
+            Child task of the current task
+        """
+        for pathspec in self.child_task_pathspecs:
+            yield Task(pathspec=pathspec, _namespace_check=False)
 
     @property
     def metadata(self) -> List[Metadata]:

metaflow/client/filecache.py

@@ -1,5 +1,6 @@
 from __future__ import print_function
 from collections import OrderedDict
+import json
 import os
 import sys
 import time
@@ -10,13 +11,14 @@ from urllib.parse import urlparse
 
 from metaflow.datastore import FlowDataStore
 from metaflow.datastore.content_addressed_store import BlobCache
+from metaflow.datastore.flow_datastore import MetadataCache
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import (
     CLIENT_CACHE_PATH,
     CLIENT_CACHE_MAX_SIZE,
     CLIENT_CACHE_MAX_FLOWDATASTORE_COUNT,
-    CLIENT_CACHE_MAX_TASKDATASTORE_COUNT,
 )
+from metaflow.metaflow_profile import from_start
 
 from metaflow.plugins import DATASTORES
 
@@ -63,8 +65,8 @@ class FileCache(object):
         # when querying for sizes of artifacts. Once we have queried for the size
         # of one artifact in a TaskDatastore, caching this means that any
         # queries on that same TaskDatastore will be quick (since we already
-        # have all the metadata)
-        self._task_metadata_caches = OrderedDict()
+        # have all the metadata). We keep track of this in a file so it persists
+        # across processes.
 
     @property
     def cache_dir(self):
@@ -87,7 +89,7 @@ class FileCache(object):
     ):
         ds_cls = self._get_datastore_storage_impl(ds_type)
         ds_root = ds_cls.path_join(*ds_cls.path_split(location)[:-5])
-        cache_id = self._flow_ds_id(ds_type, ds_root, flow_name)
+        cache_id = self.flow_ds_id(ds_type, ds_root, flow_name)
 
         token = (
             "%s.cached"
@@ -311,13 +313,13 @@ class FileCache(object):
         self._objects = sorted(objects, reverse=False)
 
     @staticmethod
-    def _flow_ds_id(ds_type, ds_root, flow_name):
+    def flow_ds_id(ds_type, ds_root, flow_name):
         p = urlparse(ds_root)
         sanitized_root = (p.netloc + p.path).replace("/", "_")
         return ".".join([ds_type, sanitized_root, flow_name])
 
     @staticmethod
-    def _task_ds_id(ds_type, ds_root, flow_name, run_id, step_name, task_id, attempt):
+    def task_ds_id(ds_type, ds_root, flow_name, run_id, step_name, task_id, attempt):
         p = urlparse(ds_root)
         sanitized_root = (p.netloc + p.path).replace("/", "_")
         return ".".join(
@@ -365,7 +367,7 @@ class FileCache(object):
         return storage_impl[0]
 
     def _get_flow_datastore(self, ds_type, ds_root, flow_name):
-        cache_id = self._flow_ds_id(ds_type, ds_root, flow_name)
+        cache_id = self.flow_ds_id(ds_type, ds_root, flow_name)
         cached_flow_datastore = self._store_caches.get(cache_id)
 
         if cached_flow_datastore:
@@ -380,9 +382,14 @@ class FileCache(object):
                 ds_root=ds_root,
             )
             blob_cache = self._blob_caches.setdefault(
-                cache_id, FileBlobCache(self, cache_id)
+                cache_id,
+                (
+                    FileBlobCache(self, cache_id),
+                    TaskMetadataCache(self, ds_type, ds_root, flow_name),
+                ),
             )
-            cached_flow_datastore.ca_store.set_blob_cache(blob_cache)
+            cached_flow_datastore.ca_store.set_blob_cache(blob_cache[0])
+            cached_flow_datastore.set_metadata_cache(blob_cache[1])
             self._store_caches[cache_id] = cached_flow_datastore
             if len(self._store_caches) > CLIENT_CACHE_MAX_FLOWDATASTORE_COUNT:
                 cache_id_to_remove, _ = self._store_caches.popitem(last=False)
@@ -393,32 +400,52 @@ class FileCache(object):
         self, ds_type, ds_root, flow_name, run_id, step_name, task_id, attempt
     ):
         flow_ds = self._get_flow_datastore(ds_type, ds_root, flow_name)
-        cached_metadata = None
-        if attempt is not None:
-            cache_id = self._task_ds_id(
-                ds_type, ds_root, flow_name, run_id, step_name, task_id, attempt
+
+        return flow_ds.get_task_datastore(run_id, step_name, task_id, attempt=attempt)
+
+
+class TaskMetadataCache(MetadataCache):
+    def __init__(self, filecache, ds_type, ds_root, flow_name):
+        self._filecache = filecache
+        self._ds_type = ds_type
+        self._ds_root = ds_root
+        self._flow_name = flow_name
+
+    def _path(self, run_id, step_name, task_id, attempt):
+        if attempt is None:
+            raise MetaflowException(
+                "Attempt number must be specified to use task metadata cache. Raise an issue "
+                "on Metaflow GitHub if you see this message.",
             )
-            cached_metadata = self._task_metadata_caches.get(cache_id)
-            if cached_metadata:
-                od_move_to_end(self._task_metadata_caches, cache_id)
-                return flow_ds.get_task_datastore(
-                    run_id,
-                    step_name,
-                    task_id,
-                    attempt=attempt,
-                    data_metadata=cached_metadata,
-                )
-        # If we are here, we either have attempt=None or nothing in the cache
-        task_ds = flow_ds.get_task_datastore(
-            run_id, step_name, task_id, attempt=attempt
+        cache_id = self._filecache.task_ds_id(
+            self._ds_type,
+            self._ds_root,
+            self._flow_name,
+            run_id,
+            step_name,
+            task_id,
+            attempt,
         )
-        cache_id = self._task_ds_id(
-            ds_type, ds_root, flow_name, run_id, step_name, task_id, task_ds.attempt
+        token = (
+            "%s.cached"
+            % sha1(
+                os.path.join(
+                    run_id, step_name, task_id, str(attempt), "metadata"
+                ).encode("utf-8")
+            ).hexdigest()
+        )
+        return os.path.join(self._filecache.cache_dir, cache_id, token[:2], token)
+
+    def load_metadata(self, run_id, step_name, task_id, attempt):
+        d = self._filecache.read_file(self._path(run_id, step_name, task_id, attempt))
+        if d:
+            return json.loads(d)
+
+    def store_metadata(self, run_id, step_name, task_id, attempt, metadata_dict):
+        self._filecache.create_file(
+            self._path(run_id, step_name, task_id, attempt),
+            json.dumps(metadata_dict).encode("utf-8"),
         )
-        self._task_metadata_caches[cache_id] = task_ds.ds_metadata
-        if len(self._task_metadata_caches) > CLIENT_CACHE_MAX_TASKDATASTORE_COUNT:
-            self._task_metadata_caches.popitem(last=False)
-        return task_ds
 
 
 class FileBlobCache(BlobCache):

metaflow/cmd/code/__init__.py

@@ -6,6 +6,7 @@ from tempfile import TemporaryDirectory
 from typing import Any, Callable, List, Mapping, Optional, cast
 
 from metaflow import Run
+from metaflow.util import walk_without_cycles
 from metaflow._vendor import click
 from metaflow.cli import echo_always
 
@@ -51,7 +52,7 @@ def perform_diff(
         target_dir = os.getcwd()
 
     diffs = []
-    for dirpath, dirnames, filenames in os.walk(source_dir, followlinks=True):
+    for dirpath, _, filenames in walk_without_cycles(source_dir):
         for fname in filenames:
             # NOTE: the paths below need to be set up carefully
             # for the `patch` command to work. Better not to touch