ob-metaflow 2.12.27.1__py2.py3-none-any.whl → 2.12.30.2__py2.py3-none-any.whl

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,217 +158,104 @@ class TriggeredRun(object):
             return None
 
 
-class LazyDeploymentMethod:
-    def __init__(self, module_path, func_name):
-        self.module_path = module_path
-        self.func_name = func_name
-        self.func = None
-
-    def __call__(self, *args, **kwargs):
-        if self.func is None:
-            module = importlib.import_module(self.module_path)
-            self.func = getattr(module, self.func_name)
-        return self.func(*args, **kwargs)
-
-
-class DeploymentMethodsMeta(type):
-    from metaflow.plugins import FROM_DEPLOYMENT_PROVIDERS
-    from metaflow.metaflow_config import FROM_DEPLOYMENT_IMPL
-
+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
+
+            allowed_providers = dict(
+                {
+                    provider.TYPE.replace("-", "_"): provider
+                    for provider in DEPLOYER_IMPL_PROVIDERS
+                }
+            )
 
-        def from_deployment(identifier, metadata=None, impl=None):
-            if impl is None:
-                impl = mcs.FROM_DEPLOYMENT_IMPL
-
-            if impl not in mcs.FROM_DEPLOYMENT_PROVIDERS:
-                raise ValueError("This method is not available for: %s" % impl)
-
-            module_path = mcs.FROM_DEPLOYMENT_PROVIDERS[impl]
-            lazy_method = LazyDeploymentMethod(module_path, "from_deployment")
-            return lazy_method(identifier=identifier, metadata=metadata)
-
-        setattr(cls, "from_deployment", staticmethod(from_deployment))
+            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)),
+                )
 
         return cls
 
 
-class DeployedFlow(metaclass=DeploymentMethodsMeta):
+class DeployedFlow(metaclass=DeployedFlowMeta):
     """
     DeployedFlow class represents a flow that has been deployed.
 
-    Parameters
-    ----------
-    deployer : DeployerImpl
-        Instance of the deployer implementation.
+    This class is not meant to be instantiated directly. Instead, it is returned from
+    methods of `Deployer`.
     """
 
-    def __init__(self, deployer: "DeployerImpl"):
+    # This should match the TYPE value in DeployerImpl for proper stub generation
+    TYPE: ClassVar[Optional[str]] = None
+
+    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
-
-    def _enrich_object(self, env):
-        """
-        Enrich the DeployedFlow 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)
-
-
-class DeployerImpl(object):
-    """
-    Base class for deployer implementations. Each implementation should define a TYPE
-    class variable that matches the name of the CLI group.
-
-    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.
-    """
-
-    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()

metaflow/runner/deployer_impl.py

@@ -0,0 +1,167 @@
+import importlib
+import json
+import os
+import sys
+import tempfile
+
+from typing import Any, ClassVar, Dict, Optional, TYPE_CHECKING, Type
+
+from .subprocess_manager import SubprocessManager
+from .utils import get_lower_level_group, handle_timeout
+
+if TYPE_CHECKING:
+    import metaflow.runner.deployer
+
+# NOTE: This file is separate from the deployer.py file to prevent circular imports.
+# This file is needed in any of the DeployerImpl implementations
+# (like argo_workflows_deployer.py) which is in turn needed to create the Deployer
+# class (ie: it uses ArgoWorkflowsDeployer to create the Deployer class).
+
+
+class DeployerImpl(object):
+    """
+    Base class for deployer implementations. Each implementation should define a TYPE
+    class variable that matches the name of the CLI group.
+
+    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.
+    """
+
+    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)
+
+    @property
+    def deployer_kwargs(self) -> Dict[str, Any]:
+        raise NotImplementedError
+
+    @staticmethod
+    def deployed_flow_type() -> Type["metaflow.runner.deployer.DeployedFlow"]:
+        raise NotImplementedError
+
+    def __enter__(self) -> "DeployerImpl":
+        return self
+
+    def create(self, **kwargs) -> "metaflow.runner.deployer.DeployedFlow":
+        """
+        Create a sub-class of a `DeployedFlow` depending on 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.
+        """
+        # Sub-classes should implement this by simply calling _create and pass the
+        # proper class as the DeployedFlow to return.
+        raise NotImplementedError
+
+    def _create(
+        self, create_class: Type["metaflow.runner.deployer.DeployedFlow"], **kwargs
+    ) -> "metaflow.runner.deployer.DeployedFlow":
+        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:
+                return create_class(deployer=self)
+
+        raise RuntimeError("Error deploying %s to %s" % (self.flow_file, self.TYPE))
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """
+        Cleanup resources on exit.
+        """
+        self.cleanup()
+
+    def cleanup(self):
+        """
+        Cleanup resources.
+        """
+        self.spm.cleanup()