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

metaflow/cmd/develop/stub_generator.py

@@ -7,7 +7,6 @@ import pathlib
 import re
 import time
 import typing
-
 from datetime import datetime
 from io import StringIO
 from types import ModuleType
@@ -335,6 +334,8 @@ class StubGenerator:
 
         # Imports that are needed at the top of the file
         self._imports = set()  # type: Set[str]
+
+        self._sub_module_imports = set()  # type: Set[Tuple[str, str]]``
         # Typing imports (behind if TYPE_CHECKING) that are needed at the top of the file
         self._typing_imports = set()  # type: Set[str]
         # Typevars that are defined
@@ -488,9 +489,6 @@ class StubGenerator:
                 self._imports.add(name)
 
         def _add_to_typing_check(name, is_module=False):
-            # if name != self._current_module_name:
-            #    self._typing_imports.add(name)
-            #
             if name == "None":
                 return
             if is_module:
@@ -504,6 +502,24 @@ class StubGenerator:
                     # the current file
                     self._typing_imports.add(splits[0])
 
+        def _format_qualified_class_name(cls: type) -> str:
+            """Helper to format a class with its qualified module name"""
+            # Special case for NoneType - return None
+            if cls.__name__ == "NoneType":
+                return "None"
+
+            module = inspect.getmodule(cls)
+            if (
+                module
+                and module.__name__ != "builtins"
+                and module.__name__ != "__main__"
+            ):
+                module_name = self._get_module_name_alias(module.__name__)
+                _add_to_typing_check(module_name, is_module=True)
+                return f"{module_name}.{cls.__name__}"
+            else:
+                return cls.__name__
+
         if isinstance(element, str):
             # Special case for self referential things (particularly in a class)
             if element == self._current_name:
@@ -557,19 +573,15 @@ class StubGenerator:
                 return element.__name__
         elif isinstance(element, type(Ellipsis)):
             return "..."
-        # elif (
-        #    isinstance(element, typing._GenericAlias)
-        #    and hasattr(element, "_name")
-        #    and element._name in ("List", "Tuple", "Dict", "Set")
-        # ):
-        #    # 3.7 has these as _GenericAlias but they don't behave like the ones in 3.10
-        #    _add_to_import("typing")
-        #    return str(element)
         elif isinstance(element, typing._GenericAlias):
             # We need to check things recursively in __args__ if it exists
             args_str = []
             for arg in getattr(element, "__args__", []):
-                args_str.append(self._get_element_name_with_module(arg))
+                # Special handling for class objects in type arguments
+                if isinstance(arg, type):
+                    args_str.append(_format_qualified_class_name(arg))
+                else:
+                    args_str.append(self._get_element_name_with_module(arg))
 
             _add_to_import("typing")
             if element._name:
@@ -584,12 +596,15 @@ class StubGenerator:
                         args_str = [call_args, args_str[-1]]
                 return "typing.%s[%s]" % (element._name, ", ".join(args_str))
             else:
-                return "%s[%s]" % (element.__origin__, ", ".join(args_str))
+                # Handle the case where we have a generic type without a _name
+                origin = element.__origin__
+                if isinstance(origin, type):
+                    origin_str = _format_qualified_class_name(origin)
+                else:
+                    origin_str = str(origin)
+                return "%s[%s]" % (origin_str, ", ".join(args_str))
         elif isinstance(element, ForwardRef):
             f_arg = self._get_module_name_alias(element.__forward_arg__)
-            # if f_arg in ("Run", "Task"):  # HACK -- forward references in current.py
-            #    _add_to_import("metaflow")
-            #    f_arg = "metaflow.%s" % f_arg
             _add_to_typing_check(f_arg)
             return '"%s"' % f_arg
         elif inspect.getmodule(element) == inspect.getmodule(typing):
@@ -629,6 +644,21 @@ class StubGenerator:
                     "deployer"
                 ] = (self._current_module_name + "." + name)
 
+        # Handle TypedDict gracefully for Python 3.7 compatibility
+        # _TypedDictMeta is not available in Python 3.7
+        typed_dict_meta = getattr(typing, "_TypedDictMeta", None)
+        if typed_dict_meta is not None and isinstance(clazz, typed_dict_meta):
+            self._sub_module_imports.add(("typing", "TypedDict"))
+            total_flag = getattr(clazz, "__total__", False)
+            buff = StringIO()
+            # Emit the TypedDict base and total flag
+            buff.write(f"class {name}(TypedDict, total={total_flag}):\n")
+            # Write out each field from __annotations__
+            for field_name, field_type in clazz.__annotations__.items():
+                ann = self._get_element_name_with_module(field_type)
+                buff.write(f"{TAB}{field_name}: {ann}\n")
+            return buff.getvalue()
+
         buff = StringIO()
         # Class prototype
         buff.write("class " + name.split(".")[-1] + "(")
@@ -973,7 +1003,6 @@ class StubGenerator:
         ]
 
         docs = split_docs(raw_doc, section_boundaries)
-
         parameters, no_arg_version = parse_params_from_doc(docs["param_doc"])
 
         if docs["add_to_current_doc"]:
@@ -1501,6 +1530,8 @@ class StubGenerator:
                     f.write("import " + module + "\n")
                     if module == "typing":
                         imported_typing = True
+                for module, sub_module in self._sub_module_imports:
+                    f.write(f"from {module} import {sub_module}\n")
                 if self._typing_imports:
                     if not imported_typing:
                         f.write("import typing\n")

metaflow/cmd/develop/stubs.py

@@ -12,25 +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, 7):
-                from metaflow._vendor.v3_7 import importlib_metadata as 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")
@@ -45,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")
@@ -187,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",
 )
                 """
             )
@@ -330,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/make_wrapper.py

@@ -28,6 +28,36 @@ def find_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
 
 

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,7 +151,11 @@ 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))
@@ -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_task_datastores(
-            run_id, steps=steps, pathspecs=pathspecs, allow_not_done=allow_not_done
+            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):
@@ -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,14 @@ 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_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,
@@ -198,7 +227,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 +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,
@@ -222,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
 
@@ -265,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 {}