metaflow 2.12.28__py2.py3-none-any.whl → 2.12.30__py2.py3-none-any.whl

metaflow/plugins/datatools/local.py

@@ -81,6 +81,8 @@ class Local(object):
     In the future, we may want to allow it to be used in a way similar to the S3() client.
     """
 
+    TYPE = "local"
+
     @staticmethod
     def _makedirs(path):
         try:

metaflow/plugins/datatools/s3/s3.py

@@ -504,6 +504,8 @@ class S3(object):
         If `run` is not specified, use this as the S3 prefix.
     """
 
+    TYPE = "s3"
+
     @classmethod
     def get_root_from_config(cls, echo, create_on_absent=True):
         return DATATOOLS_S3ROOT

metaflow/plugins/env_escape/__init__.py

@@ -124,9 +124,9 @@ def load():
     cur_path = os.path.dirname(__file__)
     sys.path = [p for p in old_paths if p != cur_path]
     # Handle special case where we launch a shell (including with a command)
-    # and we are in the CWD (searched if '' is the first element of sys.path)
-    if cur_path == os.getcwd() and sys.path[0] == '':
-        sys.path = sys.path[1:]
+    # and we are in the CWD (searched if '' is present in sys.path)
+    if cur_path == os.getcwd() and '' in sys.path:
+        sys.path.remove("")
 
     # Remove the module (this file) to reload it properly. Do *NOT* update sys.modules but
     # modify directly since it may be referenced elsewhere

metaflow/plugins/gcp/includefile_support.py

@@ -8,6 +8,9 @@ from metaflow.exception import MetaflowException, MetaflowInternalError
 
 
 class GS(object):
+
+    TYPE = "gs"
+
     @classmethod
     def get_root_from_config(cls, echo, create_on_absent=True):
         from metaflow.metaflow_config import DATATOOLS_GSROOT

metaflow/plugins/kubernetes/kubernetes_cli.py

@@ -9,7 +9,7 @@ import metaflow.tracing as tracing
 from metaflow import JSONTypeClass, util
 from metaflow._vendor import click
 from metaflow.exception import METAFLOW_EXIT_DISALLOW_RETRY, MetaflowException
-from metaflow.metadata.util import sync_local_metadata_from_datastore
+from metaflow.metadata_provider.util import sync_local_metadata_from_datastore
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR, KUBERNETES_LABELS
 from metaflow.mflog import TASK_LOG_SOURCE
 from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -7,8 +7,8 @@ import time
 from metaflow import current
 from metaflow.decorators import StepDecorator
 from metaflow.exception import MetaflowException
-from metaflow.metadata import MetaDatum
-from metaflow.metadata.util import sync_local_metadata_to_datastore
+from metaflow.metadata_provider import MetaDatum
+from metaflow.metadata_provider.util import sync_local_metadata_to_datastore
 from metaflow.metaflow_config import (
     DATASTORE_LOCAL_DIR,
     KUBERNETES_CONTAINER_IMAGE,
@@ -73,8 +73,9 @@ class KubernetesDecorator(StepDecorator):
         in Metaflow configuration.
     node_selector: Union[Dict[str,str], str], optional, default None
         Kubernetes node selector(s) to apply to the pod running the task.
-        Can be passed in as a comma separated string of values e.g. "kubernetes.io/os=linux,kubernetes.io/arch=amd64"
-        or as a dictionary {"kubernetes.io/os": "linux", "kubernetes.io/arch": "amd64"}
+        Can be passed in as a comma separated string of values e.g.
+        'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
+        {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
     namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
         Kubernetes namespace to use when launching pod in Kubernetes.
     gpu : int, optional, default None

metaflow/plugins/{metadata → metadata_providers}/local.py

@@ -8,9 +8,9 @@ import time
 from collections import namedtuple
 
 from metaflow.exception import MetaflowInternalError, MetaflowTaggingError
-from metaflow.metadata.metadata import ObjectOrder
+from metaflow.metadata_provider.metadata import ObjectOrder
 from metaflow.metaflow_config import DATASTORE_LOCAL_DIR
-from metaflow.metadata import MetadataProvider
+from metaflow.metadata_provider import MetadataProvider
 from metaflow.tagging_util import MAX_USER_TAG_SET_SIZE, validate_tags
 
 

metaflow/plugins/{metadata → metadata_providers}/service.py

@@ -14,8 +14,8 @@ from metaflow.metaflow_config import (
     SERVICE_HEADERS,
     SERVICE_URL,
 )
-from metaflow.metadata import MetadataProvider
-from metaflow.metadata.heartbeat import HB_URL_KEY
+from metaflow.metadata_provider import MetadataProvider
+from metaflow.metadata_provider.heartbeat import HB_URL_KEY
 from metaflow.sidecar import Message, MessageTypes, Sidecar
 
 from metaflow.util import version_parse

metaflow/plugins/parallel_decorator.py

@@ -2,7 +2,7 @@ from collections import namedtuple
 from metaflow.decorators import StepDecorator
 from metaflow.unbounded_foreach import UBF_CONTROL, CONTROL_TASK_TAG
 from metaflow.exception import MetaflowException
-from metaflow.metadata import MetaDatum
+from metaflow.metadata_provider import MetaDatum
 from metaflow.metaflow_current import current, Parallel
 import os
 import sys

metaflow/plugins/pypi/conda_decorator.py

@@ -8,7 +8,7 @@ import tempfile
 
 from metaflow.decorators import FlowDecorator, StepDecorator
 from metaflow.extension_support import EXT_PKG
-from metaflow.metadata import MetaDatum
+from metaflow.metadata_provider import MetaDatum
 from metaflow.metaflow_environment import InvalidEnvironmentException
 from metaflow.util import get_metaflow_root
 

metaflow/plugins/test_unbounded_foreach_decorator.py

@@ -15,7 +15,7 @@ from metaflow.unbounded_foreach import (
     CONTROL_TASK_TAG,
 )
 from metaflow.util import to_unicode
-from metaflow.metadata import MetaDatum
+from metaflow.metadata_provider import MetaDatum
 
 
 class InternalTestUnboundedForeachInput(UnboundedForeachInput):

metaflow/runner/click_api.py

@@ -193,6 +193,10 @@ class MetaflowAPI(object):
     def chain(self):
         return self._chain
 
+    @property
+    def name(self):
+        return self._API_NAME
+
     @classmethod
     def from_cli(cls, flow_file: str, cli_collection: Callable) -> Callable:
         flow_cls = extract_flow_class_from_file(flow_file)

metaflow/runner/deployer.py

@@ -1,53 +1,50 @@
-import os
-import sys
 import json
 import time
-import importlib
-import functools
-import tempfile
 
-from typing import Optional, Dict, ClassVar
+from typing import ClassVar, Dict, Optional, TYPE_CHECKING
 
 from metaflow.exception import MetaflowNotFound
-from metaflow.runner.subprocess_manager import SubprocessManager
-from metaflow.runner.utils import handle_timeout
+from metaflow.metaflow_config import DEFAULT_FROM_DEPLOYMENT_IMPL
 
+if TYPE_CHECKING:
+    import metaflow
+    import metaflow.runner.deployer_impl
 
-def get_lower_level_group(
-    api, top_level_kwargs: Dict, _type: Optional[str], deployer_kwargs: Dict
-):
-    """
-    Retrieve a lower-level group from the API based on the type and provided arguments.
 
-    Parameters
-    ----------
-    api : MetaflowAPI
-        Metaflow API instance.
-    top_level_kwargs : Dict
-        Top-level keyword arguments to pass to the API.
-    _type : str
-        Type of the deployer implementation to target.
-    deployer_kwargs : Dict
-        Keyword arguments specific to the deployer.
-
-    Returns
-    -------
-    Any
-        The lower-level group object retrieved from the API.
-
-    Raises
-    ------
-    ValueError
-        If the `_type` is None.
-    """
-    if _type is None:
-        raise ValueError(
-            "DeployerImpl doesn't have a 'TYPE' to target. Please use a sub-class of DeployerImpl."
-        )
-    return getattr(api(**top_level_kwargs), _type)(**deployer_kwargs)
+class DeployerMeta(type):
+    def __new__(mcs, name, bases, dct):
+        cls = super().__new__(mcs, name, bases, dct)
+
+        from metaflow.plugins import DEPLOYER_IMPL_PROVIDERS
+
+        def _injected_method(method_name, deployer_class):
+            def f(self, **deployer_kwargs):
+                return deployer_class(
+                    deployer_kwargs=deployer_kwargs,
+                    flow_file=self.flow_file,
+                    show_output=self.show_output,
+                    profile=self.profile,
+                    env=self.env,
+                    cwd=self.cwd,
+                    file_read_timeout=self.file_read_timeout,
+                    **self.top_level_kwargs,
+                )
+
+            f.__doc__ = provider_class.__doc__ or ""
+            f.__name__ = method_name
+            return f
+
+        for provider_class in DEPLOYER_IMPL_PROVIDERS:
+            # TYPE is the name of the CLI groups i.e.
+            # `argo-workflows` instead of `argo_workflows`
+            # The injected method names replace '-' by '_' though.
+            method_name = provider_class.TYPE.replace("-", "_")
+            setattr(cls, method_name, _injected_method(method_name, provider_class))
+
+        return cls
 
 
-class Deployer(object):
+class Deployer(metaclass=DeployerMeta):
     """
     Use the `Deployer` class to configure and access one of the production
     orchestrators supported by Metaflow.
@@ -81,7 +78,7 @@ class Deployer(object):
         env: Optional[Dict] = None,
         cwd: Optional[str] = None,
         file_read_timeout: int = 3600,
-        **kwargs
+        **kwargs,
     ):
         self.flow_file = flow_file
         self.show_output = show_output
@@ -91,56 +88,16 @@ class Deployer(object):
         self.file_read_timeout = file_read_timeout
         self.top_level_kwargs = kwargs
 
-        from metaflow.plugins import DEPLOYER_IMPL_PROVIDERS
-
-        for provider_class in DEPLOYER_IMPL_PROVIDERS:
-            # TYPE is the name of the CLI groups i.e.
-            # `argo-workflows` instead of `argo_workflows`
-            # The injected method names replace '-' by '_' though.
-            method_name = provider_class.TYPE.replace("-", "_")
-            setattr(Deployer, method_name, self.__make_function(provider_class))
-
-    def __make_function(self, deployer_class):
-        """
-        Create a function for the given deployer class.
-
-        Parameters
-        ----------
-        deployer_class : Type[DeployerImpl]
-            Deployer implementation class.
-
-        Returns
-        -------
-        Callable
-            Function that initializes and returns an instance of the deployer class.
-        """
-
-        def f(self, **deployer_kwargs):
-            return deployer_class(
-                deployer_kwargs=deployer_kwargs,
-                flow_file=self.flow_file,
-                show_output=self.show_output,
-                profile=self.profile,
-                env=self.env,
-                cwd=self.cwd,
-                file_read_timeout=self.file_read_timeout,
-                **self.top_level_kwargs
-            )
-
-        return f
-
 
 class TriggeredRun(object):
     """
-    TriggeredRun class represents a run that has been triggered on a production orchestrator.
-
-    Only when the `start` task starts running, the `run` object corresponding to the run
-    becomes available.
+    TriggeredRun class represents a run that has been triggered on a
+    production orchestrator.
     """
 
     def __init__(
         self,
-        deployer: "DeployerImpl",
+        deployer: "metaflow.runner.deployer_impl.DeployerImpl",
         content: str,
     ):
         self.deployer = deployer
@@ -149,31 +106,18 @@ class TriggeredRun(object):
         self.pathspec = content_json.get("pathspec")
         self.name = content_json.get("name")
 
-    def _enrich_object(self, env):
-        """
-        Enrich the TriggeredRun object with additional properties and methods.
-
-        Parameters
-        ----------
-        env : dict
-            Environment dictionary containing properties and methods to add.
-        """
-        for k, v in env.items():
-            if isinstance(v, property):
-                setattr(self.__class__, k, v)
-            elif callable(v):
-                setattr(self, k, functools.partial(v, self))
-            else:
-                setattr(self, k, v)
-
-    def wait_for_run(self, timeout=None):
+    def wait_for_run(self, timeout: Optional[int] = None):
         """
         Wait for the `run` property to become available.
 
+        The `run` property becomes available only after the `start` task of the triggered
+        flow starts running.
+
         Parameters
         ----------
-        timeout : int, optional
-            Maximum time to wait for the `run` to become available, in seconds. If None, wait indefinitely.
+        timeout : int, optional, default None
+            Maximum time to wait for the `run` to become available, in seconds. If
+            None, wait indefinitely.
 
         Raises
         ------
@@ -194,7 +138,7 @@ class TriggeredRun(object):
             time.sleep(check_interval)
 
     @property
-    def run(self):
+    def run(self) -> Optional["metaflow.Run"]:
         """
         Retrieve the `Run` object for the triggered run.
 
@@ -214,178 +158,104 @@ class TriggeredRun(object):
             return None
 
 
-class DeployedFlow(object):
-    """
-    DeployedFlow class represents a flow that has been deployed.
+class DeployedFlowMeta(type):
+    def __new__(mcs, name, bases, dct):
+        cls = super().__new__(mcs, name, bases, dct)
+        if not bases:
+            # Inject methods only in DeployedFlow and not any of its
+            # subclasses
+            from metaflow.plugins import DEPLOYER_IMPL_PROVIDERS
 
-    Parameters
-    ----------
-    deployer : DeployerImpl
-        Instance of the deployer implementation.
-    """
-
-    def __init__(self, deployer: "DeployerImpl"):
-        self.deployer = deployer
+            allowed_providers = dict(
+                {
+                    provider.TYPE.replace("-", "_"): provider
+                    for provider in DEPLOYER_IMPL_PROVIDERS
+                }
+            )
 
-    def _enrich_object(self, env):
-        """
-        Enrich the DeployedFlow object with additional properties and methods.
+            def _default_injected_method():
+                def f(
+                    cls,
+                    identifier: str,
+                    metadata: Optional[str] = None,
+                    impl: str = DEFAULT_FROM_DEPLOYMENT_IMPL.replace("-", "_"),
+                ) -> "DeployedFlow":
+                    """
+                    Retrieves a `DeployedFlow` object from an identifier and optional
+                    metadata. The `impl` parameter specifies the deployer implementation
+                    to use (like `argo-workflows`).
+
+                    Parameters
+                    ----------
+                    identifier : str
+                        Deployer specific identifier for the workflow to retrieve
+                    metadata : str, optional, default None
+                        Optional deployer specific metadata.
+                    impl : str, optional, default given by METAFLOW_DEFAULT_FROM_DEPLOYMENT_IMPL
+                        The default implementation to use if not specified
+
+                    Returns
+                    -------
+                    DeployedFlow
+                        A `DeployedFlow` object representing the deployed flow corresponding
+                        to the identifier
+                    """
+                    if impl in allowed_providers:
+                        return (
+                            allowed_providers[impl]
+                            .deployed_flow_type()
+                            .from_deployment(identifier, metadata)
+                        )
+                    else:
+                        raise ValueError(
+                            f"No deployer '{impl}' exists; valid deployers are: "
+                            f"{list(allowed_providers.keys())}"
+                        )
+
+                f.__name__ = "from_deployment"
+                return f
+
+            def _per_type_injected_method(method_name, impl):
+                def f(
+                    cls,
+                    identifier: str,
+                    metadata: Optional[str] = None,
+                ):
+                    return (
+                        allowed_providers[impl]
+                        .deployed_flow_type()
+                        .from_deployment(identifier, metadata)
+                    )
+
+                f.__name__ = method_name
+                return f
+
+            setattr(cls, "from_deployment", classmethod(_default_injected_method()))
+
+            for impl in allowed_providers:
+                method_name = f"from_{impl}"
+                setattr(
+                    cls,
+                    method_name,
+                    classmethod(_per_type_injected_method(method_name, impl)),
+                )
 
-        Parameters
-        ----------
-        env : dict
-            Environment dictionary containing properties and methods to add.
-        """
-        for k, v in env.items():
-            if isinstance(v, property):
-                setattr(self.__class__, k, v)
-            elif callable(v):
-                setattr(self, k, functools.partial(v, self))
-            else:
-                setattr(self, k, v)
+        return cls
 
 
-class DeployerImpl(object):
+class DeployedFlow(metaclass=DeployedFlowMeta):
     """
-    Base class for deployer implementations. Each implementation should define a TYPE
-    class variable that matches the name of the CLI group.
+    DeployedFlow class represents a flow that has been deployed.
 
-    Parameters
-    ----------
-    flow_file : str
-        Path to the flow file to deploy.
-    show_output : bool, default True
-        Show the 'stdout' and 'stderr' to the console by default.
-    profile : Optional[str], default None
-        Metaflow profile to use for the deployment. If not specified, the default
-        profile is used.
-    env : Optional[Dict], default None
-        Additional environment variables to set for the deployment.
-    cwd : Optional[str], default None
-        The directory to run the subprocess in; if not specified, the current
-        directory is used.
-    file_read_timeout : int, default 3600
-        The timeout until which we try to read the deployer attribute file.
-    **kwargs : Any
-        Additional arguments that you would pass to `python myflow.py` before
-        the deployment command.
+    This class is not meant to be instantiated directly. Instead, it is returned from
+    methods of `Deployer`.
     """
 
+    # This should match the TYPE value in DeployerImpl for proper stub generation
     TYPE: ClassVar[Optional[str]] = None
 
-    def __init__(
-        self,
-        flow_file: str,
-        show_output: bool = True,
-        profile: Optional[str] = None,
-        env: Optional[Dict] = None,
-        cwd: Optional[str] = None,
-        file_read_timeout: int = 3600,
-        **kwargs
-    ):
-        if self.TYPE is None:
-            raise ValueError(
-                "DeployerImpl doesn't have a 'TYPE' to target. Please use a sub-class of DeployerImpl."
-            )
-
-        if "metaflow.cli" in sys.modules:
-            importlib.reload(sys.modules["metaflow.cli"])
-        from metaflow.cli import start
-        from metaflow.runner.click_api import MetaflowAPI
-
-        self.flow_file = flow_file
-        self.show_output = show_output
-        self.profile = profile
-        self.env = env
-        self.cwd = cwd
-        self.file_read_timeout = file_read_timeout
-
-        self.env_vars = os.environ.copy()
-        self.env_vars.update(self.env or {})
-        if self.profile:
-            self.env_vars["METAFLOW_PROFILE"] = profile
-
-        self.spm = SubprocessManager()
-        self.top_level_kwargs = kwargs
-        self.api = MetaflowAPI.from_cli(self.flow_file, start)
-
-    def __enter__(self) -> "DeployerImpl":
-        return self
-
-    def create(self, **kwargs) -> DeployedFlow:
-        """
-        Create a deployed flow using the deployer implementation.
-
-        Parameters
-        ----------
-        **kwargs : Any
-            Additional arguments to pass to `create` corresponding to the
-            command line arguments of `create`
-
-        Returns
-        -------
-        DeployedFlow
-            DeployedFlow object representing the deployed flow.
-
-        Raises
-        ------
-        Exception
-            If there is an error during deployment.
-        """
-        with tempfile.TemporaryDirectory() as temp_dir:
-            tfp_runner_attribute = tempfile.NamedTemporaryFile(
-                dir=temp_dir, delete=False
-            )
-            # every subclass needs to have `self.deployer_kwargs`
-            command = get_lower_level_group(
-                self.api, self.top_level_kwargs, self.TYPE, self.deployer_kwargs
-            ).create(deployer_attribute_file=tfp_runner_attribute.name, **kwargs)
-
-            pid = self.spm.run_command(
-                [sys.executable, *command],
-                env=self.env_vars,
-                cwd=self.cwd,
-                show_output=self.show_output,
-            )
-
-            command_obj = self.spm.get(pid)
-            content = handle_timeout(
-                tfp_runner_attribute, command_obj, self.file_read_timeout
-            )
-            content = json.loads(content)
-            self.name = content.get("name")
-            self.flow_name = content.get("flow_name")
-            self.metadata = content.get("metadata")
-            # Additional info is used to pass additional deployer specific information.
-            # It is used in non-OSS deployers (extensions).
-            self.additional_info = content.get("additional_info", {})
-
-            if command_obj.process.returncode == 0:
-                deployed_flow = DeployedFlow(deployer=self)
-                self._enrich_deployed_flow(deployed_flow)
-                return deployed_flow
-
-        raise Exception("Error deploying %s to %s" % (self.flow_file, self.TYPE))
-
-    def _enrich_deployed_flow(self, deployed_flow: DeployedFlow):
-        """
-        Enrich the DeployedFlow object with additional properties and methods.
-
-        Parameters
-        ----------
-        deployed_flow : DeployedFlow
-            The DeployedFlow object to enrich.
-        """
-        raise NotImplementedError
-
-    def __exit__(self, exc_type, exc_value, traceback):
-        """
-        Cleanup resources on exit.
-        """
-        self.cleanup()
-
-    def cleanup(self):
-        """
-        Cleanup resources.
-        """
-        self.spm.cleanup()
+    def __init__(self, deployer: "metaflow.runner.deployer_impl.DeployerImpl"):
+        self.deployer = deployer
+        self.name = self.deployer.name
+        self.flow_name = self.deployer.flow_name
+        self.metadata = self.deployer.metadata