ob-metaflow 2.18.13.1__py2.py3-none-any.whl → 2.19.0.1__py2.py3-none-any.whl

metaflow/flowspec.py

@@ -307,7 +307,8 @@ class FlowSpec(metaclass=FlowSpecMeta):
                         % (deco._flow_cls.__name__, cls.__name__)
                     )
                 debug.userconf_exec(
-                    "Evaluating flow level decorator %s" % deco.__class__.__name__
+                    "Evaluating flow level decorator %s (pre-mutate)"
+                    % deco.__class__.__name__
                 )
                 deco.pre_mutate(mutable_flow)
                 # We reset cached_parameters on the very off chance that the user added
@@ -324,7 +325,7 @@ class FlowSpec(metaclass=FlowSpecMeta):
                 if isinstance(deco, StepMutator):
                     inserted_by_value = [deco.decorator_name] + (deco.inserted_by or [])
                     debug.userconf_exec(
-                        "Evaluating step level decorator %s for %s"
+                        "Evaluating step level decorator %s for %s (pre-mutate)"
                         % (deco.__class__.__name__, step.name)
                     )
                     deco.pre_mutate(

metaflow/metaflow_config.py

@@ -21,6 +21,7 @@ if sys.platform == "darwin":
 
 # Path to the local directory to store artifacts for 'local' datastore.
 DATASTORE_LOCAL_DIR = ".metaflow"
+DATASTORE_SPIN_LOCAL_DIR = ".metaflow_spin"
 
 # Local configuration file (in .metaflow) containing overrides per-project
 LOCAL_CONFIG_FILE = "config.json"
@@ -47,6 +48,38 @@ DEFAULT_FROM_DEPLOYMENT_IMPL = from_conf(
     "DEFAULT_FROM_DEPLOYMENT_IMPL", "argo-workflows"
 )
 
+###
+# Spin configuration
+###
+# Essentially a whitelist of decorators that are allowed in Spin steps
+SPIN_ALLOWED_DECORATORS = from_conf(
+    "SPIN_ALLOWED_DECORATORS",
+    [
+        "conda",
+        "pypi",
+        "conda_base",
+        "pypi_base",
+        "environment",
+        "project",
+        "timeout",
+        "conda_env_internal",
+        "card",
+    ],
+)
+
+# Essentially a blacklist of decorators that are not allowed in Spin steps
+# Note: decorators not in either SPIN_ALLOWED_DECORATORS or SPIN_DISALLOWED_DECORATORS
+# are simply ignored in Spin steps
+SPIN_DISALLOWED_DECORATORS = from_conf(
+    "SPIN_DISALLOWED_DECORATORS",
+    [
+        "parallel",
+    ],
+)
+
+# Default value for persist option in spin command
+SPIN_PERSIST = from_conf("SPIN_PERSIST", False)
+
 ###
 # User configuration
 ###
@@ -57,6 +90,7 @@ USER = from_conf("USER")
 # Datastore configuration
 ###
 DATASTORE_SYSROOT_LOCAL = from_conf("DATASTORE_SYSROOT_LOCAL")
+DATASTORE_SYSROOT_SPIN = from_conf("DATASTORE_SYSROOT_SPIN")
 # S3 bucket and prefix to store artifacts for 's3' datastore.
 DATASTORE_SYSROOT_S3 = from_conf("DATASTORE_SYSROOT_S3")
 # Azure Blob Storage container and blob prefix
@@ -109,6 +143,9 @@ S3_WORKER_COUNT = from_conf("S3_WORKER_COUNT", 64)
 # top-level retries)
 S3_TRANSIENT_RETRY_COUNT = from_conf("S3_TRANSIENT_RETRY_COUNT", 20)
 
+# Whether to log transient retry messages to stdout
+S3_LOG_TRANSIENT_RETRIES = from_conf("S3_LOG_TRANSIENT_RETRIES", False)
+
 # S3 retry configuration used in the aws client
 # Use the adaptive retry strategy by default
 S3_CLIENT_RETRY_CONFIG = from_conf(
@@ -462,6 +499,10 @@ ESCAPE_HATCH_WARNING = from_conf("ESCAPE_HATCH_WARNING", True)
 ###
 FEAT_ALWAYS_UPLOAD_CODE_PACKAGE = from_conf("FEAT_ALWAYS_UPLOAD_CODE_PACKAGE", False)
 ###
+# Profile
+###
+PROFILE_FROM_START = from_conf("PROFILE_FROM_START", False)
+###
 # Debug configuration
 ###
 DEBUG_OPTIONS = [

metaflow/metaflow_profile.py

@@ -2,6 +2,24 @@ import time
 
 from contextlib import contextmanager
 
+from .metaflow_config import PROFILE_FROM_START
+
+init_time = None
+
+
+if PROFILE_FROM_START:
+
+    def from_start(msg: str):
+        global init_time
+        if init_time is None:
+            init_time = time.time()
+        print("From start: %s took %dms" % (msg, int((time.time() - init_time) * 1000)))
+
+else:
+
+    def from_start(_msg: str):
+        pass
+
 
 @contextmanager
 def profile(label, stats_dict=None):

metaflow/packaging_sys/utils.py

@@ -2,45 +2,7 @@ import os
 from contextlib import contextmanager
 from typing import Callable, Generator, List, Optional, Tuple
 
-from ..util import to_unicode
-
-
-# this is os.walk(follow_symlinks=True) with cycle detection
-def walk_without_cycles(
-    top_root: str,
-    exclude_dirs: Optional[List[str]] = None,
-) -> Generator[Tuple[str, List[str]], None, None]:
-    seen = set()
-
-    default_skip_dirs = ["__pycache__"]
-
-    def _recurse(root, skip_dirs):
-        for parent, dirs, files in os.walk(root):
-            dirs[:] = [d for d in dirs if d not in skip_dirs]
-            for d in dirs:
-                path = os.path.join(parent, d)
-                if os.path.islink(path):
-                    # Breaking loops: never follow the same symlink twice
-                    #
-                    # NOTE: this also means that links to sibling links are
-                    # not followed. In this case:
-                    #
-                    #   x -> y
-                    #   y -> oo
-                    #   oo/real_file
-                    #
-                    # real_file is only included twice, not three times
-                    reallink = os.path.realpath(path)
-                    if reallink not in seen:
-                        seen.add(reallink)
-                        for x in _recurse(path, default_skip_dirs):
-                            yield x
-            yield parent, files
-
-    skip_dirs = set(default_skip_dirs + (exclude_dirs or []))
-    for x in _recurse(top_root, skip_dirs):
-        skip_dirs = default_skip_dirs
-        yield x
+from ..util import to_unicode, walk_without_cycles
 
 
 def walk(
@@ -53,6 +15,7 @@ def walk(
     prefixlen = len("%s/" % os.path.dirname(root))
     for (
         path,
+        _,
         files,
     ) in walk_without_cycles(root, exclude_tl_dirs):
         if exclude_hidden and "/." in path:

metaflow/packaging_sys/v1.py

@@ -16,7 +16,7 @@ from ..exception import MetaflowException
 from ..metaflow_version import get_version
 from ..user_decorators.user_flow_decorator import FlowMutatorMeta
 from ..user_decorators.user_step_decorator import UserStepDecoratorMeta
-from ..util import get_metaflow_root
+from ..util import get_metaflow_root, walk_without_cycles
 from . import ContentType, MFCONTENT_MARKER, MetaflowCodeContentV1Base
 from .distribution_support import _ModuleInfo, modules_to_distributions
 from .utils import suffix_filter, walk
@@ -269,12 +269,50 @@ class MetaflowCodeContentV1(MetaflowCodeContentV1Base):
         # If the module is a single file, we handle this here by looking at __file__
         # which will point to the single file. If it is an actual module, __path__
         # will contain the path(s) to the module
+        if hasattr(module, "__file__") and module.__file__:
+            root_paths = [Path(module.__file__).resolve().as_posix()]
+        else:
+            root_paths = []
+            seen_path_values = set()
+            new_paths = module.__spec__.submodule_search_locations
+            while new_paths:
+                paths = new_paths
+                new_paths = []
+                for p in paths:
+                    if p in seen_path_values:
+                        continue
+                    if os.path.isdir(p):
+                        root_paths.append(Path(p).resolve().as_posix())
+                    elif p in sys.path_importer_cache:
+                        # We have a path hook that we likely need to call to get the actual path
+                        addl_spec = sys.path_importer_cache[p].find_spec(name)
+                        if (
+                            addl_spec is not None
+                            and addl_spec.submodule_search_locations
+                        ):
+                            new_paths.extend(addl_spec.submodule_search_locations)
+                    else:
+                        # This may not be as required since it is likely the importer cache has
+                        # everything already but just in case, we will also go through the
+                        # path hooks and see if we find another one
+                        for path_hook in sys.path_hooks:
+                            try:
+                                finder = path_hook(p)
+                                addl_spec = finder.find_spec(name)
+                                if (
+                                    addl_spec is not None
+                                    and addl_spec.submodule_search_locations
+                                ):
+                                    new_paths.extend(
+                                        addl_spec.submodule_search_locations
+                                    )
+                                    break
+                            except ImportError:
+                                continue
+                    seen_path_values.add(p)
         self._modules[name] = _ModuleInfo(
             name,
-            set(
-                Path(p).resolve().as_posix()
-                for p in getattr(module, "__path__", [module.__file__])
-            ),
+            set(root_paths),
             module,
             False,  # This is not a Metaflow module (added by the user manually)
         )
@@ -417,15 +455,7 @@ class MetaflowCodeContentV1(MetaflowCodeContentV1Base):
                     % (dist_name, name)
                 )
                 dist_root = str(dist.locate_file(name))
-                if dist_root not in paths:
-                    # This is an error because it means that this distribution is
-                    # not contributing to the module.
-                    raise RuntimeError(
-                        "Distribution '%s' is not contributing to module '%s' as "
-                        "expected (got '%s' when expected one of %s)"
-                        % (dist.metadata["Name"], name, dist_root, paths)
-                    )
-                paths.discard(dist_root)
+                has_file_in_root = False
                 if dist_name not in self._distmetainfo:
                     # Possible that a distribution contributes to multiple modules
                     self._distmetainfo[dist_name] = {
@@ -438,13 +468,30 @@ class MetaflowCodeContentV1(MetaflowCodeContentV1Base):
                 for file in dist.files or []:
                     # Skip files that do not belong to this module (distribution may
                     # provide multiple modules)
-                    if file.parts[: len(prefix_parts)] != prefix_parts:
+                    if (
+                        file.parts[: len(prefix_parts)] != prefix_parts
+                        or file.suffix == ".pth"
+                        or str(file).startswith("__editable__")
+                    ):
                         continue
                     if file.parts[len(prefix_parts)] == "__init__.py":
                         has_init = True
+                    has_file_in_root = True
+                    # At this point, we know that we are seeing actual files in the
+                    # dist_root so we make sure it is as expected
+                    if dist_root not in paths:
+                        # This is an error because it means that this distribution is
+                        # not contributing to the module.
+                        raise RuntimeError(
+                            "Distribution '%s' is not contributing to module '%s' as "
+                            "expected (got '%s' when expected one of %s)"
+                            % (dist.metadata["Name"], name, dist_root, paths)
+                        )
                     yield str(
                         dist.locate_file(file).resolve().as_posix()
                     ), os.path.join(self._code_dir, *prefix_parts, *file.parts[1:])
+                if has_file_in_root:
+                    paths.discard(dist_root)
 
         # Now if there are more paths left in paths, it means there is a non-distribution
         # component to this package which we also include.
@@ -460,7 +507,7 @@ class MetaflowCodeContentV1(MetaflowCodeContentV1Base):
                 )
                 has_init = True
             else:
-                for root, _, files in os.walk(path):
+                for root, _, files in walk_without_cycles(path):
                     for file in files:
                         if any(file.endswith(x) for x in EXT_EXCLUDE_SUFFIXES):
                             continue

metaflow/plugins/__init__.py

@@ -83,11 +83,13 @@ ENVIRONMENTS_DESC = [
 METADATA_PROVIDERS_DESC = [
     ("service", ".metadata_providers.service.ServiceMetadataProvider"),
     ("local", ".metadata_providers.local.LocalMetadataProvider"),
+    ("spin", ".metadata_providers.spin.SpinMetadataProvider"),
 ]
 
 # Add datastore here
 DATASTORES_DESC = [
     ("local", ".datastores.local_storage.LocalStorage"),
+    ("spin", ".datastores.spin_storage.SpinStorage"),
     ("s3", ".datastores.s3_storage.S3Storage"),
     ("azure", ".datastores.azure_storage.AzureStorage"),
     ("gs", ".datastores.gs_storage.GSStorage"),

metaflow/plugins/cards/card_datastore.py

@@ -9,6 +9,7 @@ from metaflow.metaflow_config import (
     CARD_S3ROOT,
     CARD_LOCALROOT,
     DATASTORE_LOCAL_DIR,
+    DATASTORE_SPIN_LOCAL_DIR,
     CARD_SUFFIX,
     CARD_AZUREROOT,
     CARD_GSROOT,
@@ -58,12 +59,17 @@ class CardDatastore(object):
             return CARD_AZUREROOT
         elif storage_type == "gs":
             return CARD_GSROOT
-        elif storage_type == "local":
+        elif storage_type == "local" or storage_type == "spin":
             # Borrowing some of the logic from LocalStorage.get_storage_root
             result = CARD_LOCALROOT
+            local_dir = (
+                DATASTORE_SPIN_LOCAL_DIR
+                if storage_type == "spin"
+                else DATASTORE_LOCAL_DIR
+            )
             if result is None:
                 current_path = os.getcwd()
-                check_dir = os.path.join(current_path, DATASTORE_LOCAL_DIR)
+                check_dir = os.path.join(current_path, local_dir)
                 check_dir = os.path.realpath(check_dir)
                 orig_path = check_dir
                 while not os.path.isdir(check_dir):
@@ -73,7 +79,7 @@ class CardDatastore(object):
                         # return the top level path
                         return os.path.join(orig_path, CARD_SUFFIX)
                     current_path = new_path
-                    check_dir = os.path.join(current_path, DATASTORE_LOCAL_DIR)
+                    check_dir = os.path.join(current_path, local_dir)
                 return os.path.join(check_dir, CARD_SUFFIX)
         else:
             # Let's make it obvious we need to update this block for each new datastore backend...

metaflow/plugins/cards/card_decorator.py

@@ -171,6 +171,7 @@ class CardDecorator(StepDecorator):
         self._flow_datastore = flow_datastore
         self._environment = environment
         self._logger = logger
+
         self.card_options = None
 
         # We check for configuration options. We do this here before they are

metaflow/plugins/cards/card_modules/basic.py

@@ -496,9 +496,15 @@ class TaskInfoComponent(MetaflowCardComponent):
             )
 
         # ignore the name as a parameter
-        param_ids = [
-            p.id for p in self._task.parent.parent["_parameters"].task if p.id != "name"
-        ]
+        if "_parameters" not in self._task.parent.parent:
+            # In case of spin steps, there is no _parameters task
+            param_ids = []
+        else:
+            param_ids = [
+                p.id
+                for p in self._task.parent.parent["_parameters"].task
+                if p.id != "name"
+            ]
         if len(param_ids) > 0:
             # Extract parameter from the Parameter Task. That is less brittle.
             parameter_data = TaskToDict(

metaflow/plugins/datastores/local_storage.py

@@ -1,24 +1,29 @@
 import json
 import os
 
-from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, DATASTORE_SYSROOT_LOCAL
+from metaflow.metaflow_config import (
+    DATASTORE_LOCAL_DIR,
+    DATASTORE_SYSROOT_LOCAL,
+)
 from metaflow.datastore.datastore_storage import CloseAfterUse, DataStoreStorage
 
 
 class LocalStorage(DataStoreStorage):
     TYPE = "local"
     METADATA_DIR = "_meta"
+    DATASTORE_DIR = DATASTORE_LOCAL_DIR  # ".metaflow"
+    SYSROOT_VAR = DATASTORE_SYSROOT_LOCAL
 
     @classmethod
     def get_datastore_root_from_config(cls, echo, create_on_absent=True):
-        result = DATASTORE_SYSROOT_LOCAL
+        result = cls.SYSROOT_VAR
         if result is None:
             try:
                 # Python2
                 current_path = os.getcwdu()
             except:  # noqa E722
                 current_path = os.getcwd()
-            check_dir = os.path.join(current_path, DATASTORE_LOCAL_DIR)
+            check_dir = os.path.join(current_path, cls.DATASTORE_DIR)
             check_dir = os.path.realpath(check_dir)
             orig_path = check_dir
             top_level_reached = False
@@ -28,12 +33,13 @@ class LocalStorage(DataStoreStorage):
                     top_level_reached = True
                     break  # We are no longer making upward progress
                 current_path = new_path
-                check_dir = os.path.join(current_path, DATASTORE_LOCAL_DIR)
+                check_dir = os.path.join(current_path, cls.DATASTORE_DIR)
             if top_level_reached:
                 if create_on_absent:
                     # Could not find any directory to use so create a new one
                     echo(
-                        "Creating local datastore in current directory (%s)" % orig_path
+                        "Creating %s datastore in current directory (%s)"
+                        % (cls.TYPE, orig_path)
                     )
                     os.mkdir(orig_path)
                     result = orig_path
@@ -42,7 +48,7 @@ class LocalStorage(DataStoreStorage):
             else:
                 result = check_dir
         else:
-            result = os.path.join(result, DATASTORE_LOCAL_DIR)
+            result = os.path.join(result, cls.DATASTORE_DIR)
         return result
 
     @staticmethod

metaflow/plugins/datastores/spin_storage.py

@@ -0,0 +1,12 @@
+from metaflow.metaflow_config import (
+    DATASTORE_SPIN_LOCAL_DIR,
+    DATASTORE_SYSROOT_SPIN,
+)
+from metaflow.plugins.datastores.local_storage import LocalStorage
+
+
+class SpinStorage(LocalStorage):
+    TYPE = "spin"
+    METADATA_DIR = "_meta"
+    DATASTORE_DIR = DATASTORE_SPIN_LOCAL_DIR  # ".metaflow_spin"
+    SYSROOT_VAR = DATASTORE_SYSROOT_SPIN

metaflow/plugins/datatools/s3/s3.py

@@ -18,6 +18,7 @@ from metaflow.metaflow_config import (
     DATATOOLS_S3ROOT,
     S3_RETRY_COUNT,
     S3_TRANSIENT_RETRY_COUNT,
+    S3_LOG_TRANSIENT_RETRIES,
     S3_SERVER_SIDE_ENCRYPTION,
     S3_WORKER_COUNT,
     TEMPDIR,
@@ -1760,17 +1761,35 @@ class S3(object):
                 # due to a transient failure so we try again.
                 transient_retry_count += 1
                 total_ok_count += last_ok_count
-                print(
-                    "Transient S3 failure (attempt #%d) -- total success: %d, "
-                    "last attempt %d/%d -- remaining: %d"
-                    % (
-                        transient_retry_count,
-                        total_ok_count,
-                        last_ok_count,
-                        last_ok_count + last_retry_count,
-                        len(pending_retries),
+
+                if S3_LOG_TRANSIENT_RETRIES:
+                    # Extract transient error type from pending retry lines
+                    error_info = ""
+                    if pending_retries:
+                        try:
+                            # Parse the first line to get transient error type
+                            first_retry = json.loads(
+                                pending_retries[0].decode("utf-8").strip()
+                            )
+                            if "transient_error_type" in first_retry:
+                                error_info = (
+                                    " (%s)" % first_retry["transient_error_type"]
+                                )
+                        except (json.JSONDecodeError, IndexError, KeyError):
+                            pass
+
+                    print(
+                        "Transient S3 failure (attempt #%d) -- total success: %d, "
+                        "last attempt %d/%d -- remaining: %d%s"
+                        % (
+                            transient_retry_count,
+                            total_ok_count,
+                            last_ok_count,
+                            last_ok_count + last_retry_count,
+                            len(pending_retries),
+                            error_info,
+                        )
                     )
-                )
                 if inject_failures == 0:
                     # Don't sleep when we are "faking" the failures
                     self._jitter_sleep(transient_retry_count)