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

metaflow/cmd/develop/stubs.py

@@ -12,23 +12,13 @@ from . import develop
 from .stub_generator import StubGenerator
 
 _py_ver = sys.version_info[:2]
-_metadata_package = None
 
-
-def _check_stubs_supported():
-    global _metadata_package
-    if _metadata_package is not None:
-        return _metadata_package
-    else:
-        if _py_ver >= (3, 4):
-            if _py_ver >= (3, 8):
-                from importlib import metadata
-            elif _py_ver >= (3, 6):
-                from metaflow._vendor.v3_6 import importlib_metadata as metadata
-            else:
-                from metaflow._vendor.v3_5 import importlib_metadata as metadata
-            _metadata_package = metadata
-        return _metadata_package
+if _py_ver >= (3, 8):
+    from importlib import metadata
+elif _py_ver >= (3, 7):
+    from metaflow._vendor.v3_7 import importlib_metadata as metadata
+else:
+    from metaflow._vendor.v3_6 import importlib_metadata as metadata
 
 
 @develop.group(short_help="Stubs management")
@@ -43,12 +33,6 @@ def stubs(ctx: Any):
     This CLI provides utilities to check and generate stubs for your current Metaflow
     installation.
     """
-    if _check_stubs_supported() is None:
-        raise click.UsageError(
-            "Building and installing stubs are not supported on Python %d.%d "
-            "(3.4 minimum required)" % _py_ver,
-            ctx=ctx,
-        )
 
 
 @stubs.command(short_help="Check validity of stubs")
@@ -168,7 +152,7 @@ def install(ctx: Any, force: bool):
                     "Metaflow stubs are already installed and valid -- use --force to reinstall"
                 )
             return
-    mf_version, _ = get_mf_version()
+    mf_version, _ = get_mf_version(True)
     with tempfile.TemporaryDirectory() as tmp_dir:
         with open(os.path.join(tmp_dir, "setup.py"), "w") as f:
             f.write(
@@ -185,7 +169,7 @@ setup(
     packages=find_namespace_packages(),
     package_data={{"metaflow-stubs": ["generated_for.txt", "py.typed", "**/*.pyi"]}},
     install_requires=["metaflow=={mf_version}"],
-    python_requires=">=3.5.2",
+    python_requires=">=3.6.1",
 )
                 """
             )
@@ -259,10 +243,10 @@ def split_version(vers: str) -> Tuple[str, Optional[str]]:
     return vers_split[0], vers_split[1]
 
 
-def get_mf_version() -> Tuple[str, Optional[str]]:
+def get_mf_version(public: bool = False) -> Tuple[str, Optional[str]]:
     from metaflow.metaflow_version import get_version
 
-    return split_version(get_version())
+    return split_version(get_version(public))
 
 
 def get_stubs_version(stubs_root_path: Optional[str]) -> Tuple[str, Optional[str]]:
@@ -328,14 +312,14 @@ def get_packages_for_stubs() -> Tuple[List[Tuple[str, str]], List[str]]:
     # some reason it shows up multiple times.
     interesting_dists = [
         d
-        for d in _metadata_package.distributions()
+        for d in metadata.distributions()
         if any(
             [
                 p == "metaflow-stubs"
                 for p in (d.read_text("top_level.txt") or "").split()
             ]
         )
-        and isinstance(d, _metadata_package.PathDistribution)
+        and isinstance(d, metadata.PathDistribution)
     ]
 
     for dist in interesting_dists:

metaflow/cmd/main_cli.py

@@ -12,7 +12,7 @@ import metaflow.tracing as tracing
 
 
 @click.group()
-@tracing.cli_entrypoint("cli/main")
+@tracing.cli("cli/main")
 def main():
     pass
 
@@ -67,6 +67,7 @@ CMDS_DESC = [
     ("configure", ".configure_cmd.cli"),
     ("tutorials", ".tutorials_cmd.cli"),
     ("develop", ".develop.cli"),
+    ("code", ".code.cli"),
 ]
 
 process_cmds(globals())
@@ -84,15 +85,16 @@ def start(ctx):
 
     import metaflow
 
+    version = get_version()
     echo("Metaflow ", fg="magenta", bold=True, nl=False)
 
     if ctx.invoked_subcommand is None:
-        echo("(%s): " % get_version(), fg="magenta", bold=False, nl=False)
+        echo("(%s): " % version, fg="magenta", bold=False, nl=False)
     else:
-        echo("(%s)\n" % get_version(), fg="magenta", bold=False)
+        echo("(%s)\n" % version, fg="magenta", bold=False)
 
     if ctx.invoked_subcommand is None:
-        echo("More data science, less engineering\n", fg="magenta")
+        echo("More AI, less engineering\n", fg="magenta")
 
         lnk_sz = max(len(lnk) for lnk in CONTACT_INFO.values()) + 1
         for what, lnk in CONTACT_INFO.items():

metaflow/cmd/make_wrapper.py

@@ -0,0 +1,78 @@
+import sys
+import subprocess
+from pathlib import Path
+import sysconfig
+import site
+
+
+def find_makefile():
+    possible_dirs = []
+
+    # 1) The standard sysconfig-based location
+    data_dir = sysconfig.get_paths()["data"]
+    possible_dirs.append(Path(data_dir) / "share" / "metaflow" / "devtools")
+
+    # 2) The user base (e.g. ~/.local on many systems)
+    user_base = site.getuserbase()  # e.g. /home/runner/.local
+    possible_dirs.append(Path(user_base) / "share" / "metaflow" / "devtools")
+
+    # 3) site-packages can vary, we can guess share/.. near each site-packages
+    # (Works if pip actually placed devtools near site-packages.)
+    for p in site.getsitepackages():
+        possible_dirs.append(Path(p).parent / "share" / "metaflow" / "devtools")
+    user_site = site.getusersitepackages()
+    possible_dirs.append(Path(user_site).parent / "share" / "metaflow" / "devtools")
+
+    for candidate_dir in possible_dirs:
+        makefile_candidate = candidate_dir / "Makefile"
+        if makefile_candidate.is_file():
+            return makefile_candidate
+
+    # 4) When developing, Metaflow might be installed with --editable, which means the devtools will not be located within site-packages.
+    # We read the actual location from package metadata in this case, but only do this heavier operation if the above lookups fail.
+    try:
+        import json
+        from importlib.metadata import Distribution
+
+        direct_url = Distribution.from_name("metaflow").read_text("direct_url.json")
+        if direct_url:
+            content = json.loads(direct_url)
+            url = content.get("url", "")
+            if not url.startswith("file://"):
+                return None
+
+            makefile_candidate = (
+                Path(url.replace("file://", "")) / "devtools" / "Makefile"
+            )
+            if makefile_candidate.is_file():
+                return makefile_candidate
+        else:
+            # No dist metadata found. This is tied to the version of pip being used
+            # Do not bother with .egg-link installs due to the handling of the file contents being a headache due to lack of a unified spec.
+            print(
+                "Could not locate an installation of Metaflow. No package metadata found."
+            )
+            print(
+                "If Metaflow is installed as editable, try upgrading the version of pip and reinstalling in order to generate proper package metadata.\n"
+            )
+    except Exception:
+        return None
+
+    return None
+
+
+def main():
+    makefile_path = find_makefile()
+    if not makefile_path:
+        print("ERROR: Could not find executable in any known location.")
+        sys.exit(1)
+    cmd = ["make", "-f", str(makefile_path)] + sys.argv[1:]
+
+    try:
+        completed = subprocess.run(cmd, check=True)
+        sys.exit(completed.returncode)
+    except subprocess.CalledProcessError as ex:
+        sys.exit(ex.returncode)
+    except KeyboardInterrupt:
+        print("Process interrupted by user. Exiting cleanly.")
+        sys.exit(1)

metaflow/datastore/__init__.py

@@ -2,3 +2,4 @@ from .inputs import Inputs
 from .flow_datastore import FlowDataStore
 from .datastore_set import TaskDataStoreSet
 from .task_datastore import TaskDataStore
+from .spin_datastore import SpinTaskDatastore

metaflow/datastore/content_addressed_store.py

@@ -38,7 +38,7 @@ class ContentAddressedStore(object):
     def set_blob_cache(self, blob_cache):
         self._blob_cache = blob_cache
 
-    def save_blobs(self, blob_iter, raw=False, len_hint=0):
+    def save_blobs(self, blob_iter, raw=False, len_hint=0, is_transfer=False):
         """
         Saves blobs of data to the datastore
 
@@ -60,11 +60,16 @@ class ContentAddressedStore(object):
 
         Parameters
         ----------
-        blob_iter : Iterator over bytes objects to save
-        raw : bool, optional
+        blob_iter : Iterator
+            Iterator over bytes objects to save
+        raw : bool, default False
             Whether to save the bytes directly or process them, by default False
-        len_hint : Hint of the number of blobs that will be produced by the
+        len_hint : int, default 0
+            Hint of the number of blobs that will be produced by the
             iterator, by default 0
+        is_transfer : bool, default False
+            If True, this indicates we are saving blobs directly from the output of another
+            content addressed store's
 
         Returns
         -------
@@ -76,6 +81,20 @@ class ContentAddressedStore(object):
 
         def packing_iter():
             for blob in blob_iter:
+                if is_transfer:
+                    key, blob_data, meta = blob
+                    path = self._storage_impl.path_join(self._prefix, key[:2], key)
+                    # Transfer data is always raw/decompressed, so mark it as such
+                    meta_corrected = {"cas_raw": True, "cas_version": 1}
+
+                    results.append(
+                        self.save_blobs_result(
+                            uri=self._storage_impl.full_uri(path),
+                            key=key,
+                        )
+                    )
+                    yield path, (BytesIO(blob_data), meta_corrected)
+                    continue
                 sha = sha1(blob).hexdigest()
                 path = self._storage_impl.path_join(self._prefix, sha[:2], sha)
                 results.append(
@@ -100,7 +119,7 @@ class ContentAddressedStore(object):
         self._storage_impl.save_bytes(packing_iter(), overwrite=True, len_hint=len_hint)
         return results
 
-    def load_blobs(self, keys, force_raw=False):
+    def load_blobs(self, keys, force_raw=False, is_transfer=False):
         """
         Mirror function of save_blobs
 
@@ -111,15 +130,20 @@ class ContentAddressedStore(object):
         ----------
         keys : List of string
             Key describing the object to load
-        force_raw : bool, optional
+        force_raw : bool, default False
             Support for backward compatibility with previous datastores. If
             True, this will force the key to be loaded as is (raw). By default,
             False
+        is_transfer : bool, default False
+            If True, this indicates we are loading blobs to transfer them directly
+            to another datastore. We will, in this case, also transfer the metadata
+            and do minimal processing. This is for internal use only.
 
         Returns
         -------
         Returns an iterator of (string, bytes) tuples; the iterator may return keys
-        in a different order than were passed in.
+        in a different order than were passed in. If is_transfer is True, the tuple
+        has three elements with the third one being the metadata.
         """
         load_paths = []
         for key in keys:
@@ -127,13 +151,17 @@ class ContentAddressedStore(object):
             if self._blob_cache:
                 blob = self._blob_cache.load_key(key)
             if blob is not None:
-                yield key, blob
+                if is_transfer:
+                    # Cached blobs are decompressed/processed bytes regardless of original format
+                    yield key, blob, {"cas_raw": False, "cas_version": 1}
+                else:
+                    yield key, blob
             else:
                 path = self._storage_impl.path_join(self._prefix, key[:2], key)
                 load_paths.append((key, path))
 
         with self._storage_impl.load_bytes([p for _, p in load_paths]) as loaded:
-            for (path_key, file_path, meta) in loaded:
+            for path_key, file_path, meta in loaded:
                 key = self._storage_impl.path_split(path_key)[-1]
                 # At this point, we either return the object as is (if raw) or
                 # decode it according to the encoding version
@@ -169,7 +197,10 @@ class ContentAddressedStore(object):
                 if self._blob_cache:
                     self._blob_cache.store_key(key, blob)
 
-                yield key, blob
+                if is_transfer:
+                    yield key, blob, meta  # Preserve exact original metadata from storage
+                else:
+                    yield key, blob
 
     def _unpack_backward_compatible(self, blob):
         # This is the backward compatible unpack

metaflow/datastore/datastore_set.py

@@ -21,9 +21,18 @@ class TaskDataStoreSet(object):
         pathspecs=None,
         prefetch_data_artifacts=None,
         allow_not_done=False,
+        join_type=None,
+        orig_flow_datastore=None,
+        spin_artifacts=None,
     ):
-        self.task_datastores = flow_datastore.get_latest_task_datastores(
-            run_id, steps=steps, pathspecs=pathspecs, allow_not_done=allow_not_done
+        self.task_datastores = flow_datastore.get_task_datastores(
+            run_id,
+            steps=steps,
+            pathspecs=pathspecs,
+            allow_not_done=allow_not_done,
+            join_type=join_type,
+            orig_flow_datastore=orig_flow_datastore,
+            spin_artifacts=spin_artifacts,
         )
 
         if prefetch_data_artifacts:

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):
@@ -13,7 +16,7 @@ class FlowDataStore(object):
     def __init__(
         self,
         flow_name,
-        environment,
+        environment=None,
         metadata=None,
         event_logger=None,
         monitor=None,
@@ -31,7 +34,7 @@ class FlowDataStore(object):
         ----------
         flow_name : str
             The name of the flow
-        environment : MetaflowEnvironment
+        environment : MetaflowEnvironment, optional
             Environment this datastore is operating in
         metadata : MetadataProvider, optional
             The metadata provider to use and update if needed, by default None
@@ -63,12 +66,28 @@ 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 get_latest_task_datastores(
-        self, run_id=None, steps=None, pathspecs=None, allow_not_done=False
+    def set_metadata_cache(self, cache):
+        self._metadata_cache = cache
+
+    def get_task_datastores(
+        self,
+        run_id=None,
+        steps=None,
+        pathspecs=None,
+        allow_not_done=False,
+        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.
@@ -88,11 +107,27 @@ 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
             wasn't marked as done, by default False
+        attempt : int, optional
+            Attempt number of the tasks to return.  If not provided, returns latest attempt.
+        include_prior : boolean, default False
+            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
         -------
@@ -126,8 +161,20 @@ class FlowDataStore(object):
                 if task.is_file is False
             ]
         urls = []
+        # parse content urls for specific attempt only, or for all attempts in max range
+        attempt_range = range(metaflow_config.MAX_ATTEMPTS)
+        # we have no reason to check for attempts greater than MAX_ATTEMPTS, as they do not exist.
+        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 range(metaflow_config.MAX_ATTEMPTS):
+            # task_url can have a trailing slash, so strip this to avoid empty strings in the split
+            task_splits = task_url.rstrip("/").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,
@@ -168,11 +215,30 @@ class FlowDataStore(object):
             for (run, step, task), attempt in latest_started_attempts.items()
         )
         if allow_not_done:
-            latest_to_fetch = latest_started_attempts
+            latest_to_fetch = (
+                done_attempts.union(latest_started_attempts)
+                if include_prior
+                else latest_started_attempts
+            )
         else:
-            latest_to_fetch = latest_started_attempts & done_attempts
+            latest_to_fetch = (
+                done_attempts
+                if include_prior
+                else (latest_started_attempts & done_attempts)
+            )
         latest_to_fetch = [
-            (v[0], v[1], v[2], v[3], data_objs.get(v), "r", 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))
@@ -186,8 +252,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,
@@ -196,8 +317,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
 
@@ -239,3 +375,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 {}