zenml-nightly 0.58.2.dev20240626__py3-none-any.whl → 0.62.0.dev20240726__py3-none-any.whl

zenml/stack_deployments/gcp_stack_deployment.py

@@ -0,0 +1,269 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+
+#       https://www.apache.org/licenses/LICENSE-2.0
+
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Functionality to deploy a ZenML stack to GCP."""
+
+import re
+from typing import ClassVar, Dict, List
+
+from zenml.enums import StackDeploymentProvider
+from zenml.models import StackDeploymentConfig
+from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
+
+GCP_DEPLOYMENT_TYPE = "deployment-manager"
+
+
+class GCPZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
+    """GCP ZenML Cloud Stack Deployment."""
+
+    provider: ClassVar[StackDeploymentProvider] = StackDeploymentProvider.GCP
+    deployment: ClassVar[str] = GCP_DEPLOYMENT_TYPE
+
+    @classmethod
+    def description(cls) -> str:
+        """Return a description of the ZenML Cloud Stack Deployment.
+
+        This will be displayed when the user is prompted to deploy
+        the ZenML stack.
+
+        Returns:
+            A MarkDown description of the ZenML Cloud Stack Deployment.
+        """
+        return """
+Provision and register a basic GCP ZenML stack authenticated and connected to
+all the necessary cloud infrastructure resources required to run pipelines in
+GCP.
+"""
+
+    @classmethod
+    def instructions(cls) -> str:
+        """Return instructions on how to deploy the ZenML stack to the specified cloud provider.
+
+        This will be displayed before the user is prompted to deploy the ZenML
+        stack.
+
+        Returns:
+            MarkDown instructions on how to deploy the ZenML stack to the
+            specified cloud provider.
+        """
+        return """
+You will be redirected to a GCP Cloud Shell console in your browser where you'll
+be asked to log into your GCP project and then create a Deployment Manager
+deployment to provision the necessary cloud resources for ZenML.
+
+**NOTE**: The Deployment Manager deployment will create the following new
+resources in your GCP project. Please ensure you have the necessary permissions
+and are aware of any potential costs:
+
+- A GCS bucket registered as a [ZenML artifact store](https://docs.zenml.io/stack-components/artifact-stores/gcp).
+- A Google Artifact Registry registered as a [ZenML container registry](https://docs.zenml.io/stack-components/container-registries/gcp).
+- Vertex AI registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/vertex).
+- GCP Cloud Build registered as a [ZenML image builder](https://docs.zenml.io/stack-components/image-builders/gcp).
+- A GCP Service Account with the minimum necessary permissions to access the
+above resources.
+- An GCP Service Account access key used to give access to ZenML to connect to
+the above resources through a [ZenML service connector](https://docs.zenml.io/how-to/auth-management/gcp-service-connector).
+
+The Deployment Manager deployment will automatically create a GCP Service
+Account secret key and will share it with ZenML to give it permission to access
+the resources created by the stack. You can revoke these permissions at any time
+by deleting the Deployment Manager deployment in the GCP Cloud Console.
+
+**Estimated costs**
+
+A small training job would cost around: $0.60
+
+These are rough estimates and actual costs may vary based on your usage and specific GCP pricing. 
+Some services may be eligible for the GCP Free Tier. Use [the GCP Pricing Calculator](https://cloud.google.com/products/calculator)
+for a detailed estimate based on your usage.
+
+⚠️  **The Cloud Shell session will warn you that the ZenML GitHub repository is
+untrusted. We recommend that you review [the contents of the repository](https://github.com/zenml-io/zenml/tree/main/infra/gcp)
+and then check the `Trust repo` checkbox to proceed with the deployment,
+otherwise the Cloud Shell session will not be authenticated to access your
+GCP projects. You will also get a chance to review the scripts that will be
+executed in the Cloud Shell session before proceeding.**
+
+💡 **After the Deployment Manager deployment is complete, you can close the Cloud
+Shell session and return to the CLI to view details about the associated ZenML
+stack automatically registered with ZenML.**
+"""
+
+    @classmethod
+    def post_deploy_instructions(cls) -> str:
+        """Return instructions on what to do after the deployment is complete.
+
+        This will be displayed after the deployment is complete.
+
+        Returns:
+            MarkDown instructions on what to do after the deployment is
+            complete.
+        """
+        return """
+The ZenML stack has been successfully deployed and registered. You can delete
+the Deployment Manager deployment at any time to revoke ZenML's access to your
+GCP project and to clean up the resources created by the stack by using
+[the GCP Cloud Console](https://console.cloud.google.com/dm/deployments).
+"""
+
+    @classmethod
+    def integrations(cls) -> List[str]:
+        """Return the ZenML integrations required for the stack.
+
+        Returns:
+            The list of ZenML integrations that need to be installed for the
+            stack to be usable.
+        """
+        return [
+            "gcp",
+        ]
+
+    @classmethod
+    def permissions(cls) -> Dict[str, List[str]]:
+        """Return the permissions granted to ZenML to access the cloud resources.
+
+        Returns:
+            The permissions granted to ZenML to access the cloud resources, as
+            a dictionary grouping permissions by resource.
+        """
+        return {
+            "GCS Bucket": [
+                "roles/storage.objectUser",
+            ],
+            "GCP Artifact Registry": [
+                "roles/artifactregistry.createOnPushWriter",
+            ],
+            "Vertex AI (Client)": [
+                "roles/aiplatform.user",
+            ],
+            "Vertex AI (Jobs)": [
+                "roles/aiplatform.serviceAgent",
+            ],
+            "Cloud Build (Client)": [
+                "roles/cloudbuild.builds.editor",
+            ],
+        }
+
+    @classmethod
+    def locations(cls) -> Dict[str, str]:
+        """Return the locations where the ZenML stack can be deployed.
+
+        Returns:
+            The regions where the ZenML stack can be deployed as a map of region
+            names to region descriptions.
+        """
+        # Return a list of all possible GCP regions
+
+        # Based on the AWS regions listed at
+        # https://cloud.google.com/about/locations
+        return {
+            "Africa (Johannesburg)": "africa-south1",
+            "Asia Pacific (Taiwan)": "asia-east1",
+            "Asia Pacific (Hong Kong)": "asia-east2",
+            "Asia Pacific (Tokyo)": "asia-northeast1",
+            "Asia Pacific (Osaka)": "asia-northeast2",
+            "Asia Pacific (Seoul)": "asia-northeast3",
+            "Asia Pacific (Mumbai)": "asia-south1",
+            "Asia Pacific (Delhi)": "asia-south2",
+            "Asia Pacific (Singapore)": "asia-southeast1",
+            "Asia Pacific (Jakarta)": "asia-southeast2",
+            "Australia (Sydney)": "australia-southeast1",
+            "Australia (Melbourne)": "australia-southeast2",
+            "Europe (Belgium)": "europe-west1",
+            "Europe (London)": "europe-west2",
+            "Europe (Frankfurt)": "europe-west3",
+            "Europe (Netherlands)": "europe-west4",
+            "Europe (Zurich)": "europe-west6",
+            "Europe (Milan)": "europe-west8",
+            "Europe (Paris)": "europe-west9",
+            "Europe (Berlin)": "europe-west10",
+            "Europe (Turin)": "europe-west12",
+            "Europe (Warsaw)": "europe-central2",
+            "Europe (Finland)": "europe-north1",
+            "Europe (Madrid)": "europe-southwest1",
+            "Middle East (Doha)": "me-central1",
+            "Middle East (Dubai)": "me-central2",
+            "Middle East (Tel Aviv)": "me-west1",
+            "North America (Montreal)": "northamerica-northeast1",
+            "North America (Toronto)": "northamerica-northeast2",
+            "South America (Sao Paulo)": "southamerica-east1",
+            "South America (Santiago)": "southamerica-west1",
+            "US Central (Iowa)": "us-central1",
+            "US East (South Carolina)": "us-east1",
+            "US East (Northern Virginia)": "us-east4",
+            "US East (Columbus)": "us-east5",
+            "US South (Dallas)": "us-south1",
+            "US West (Oregon)": "us-west1",
+            "US West (Los Angeles)": "us-west2",
+            "US West (Salt Lake City)": "us-west3",
+            "US West (Las Vegas)": "us-west4",
+        }
+
+    @classmethod
+    def skypilot_default_regions(cls) -> Dict[str, str]:
+        """Returns the regions supported by default for the Skypilot.
+
+        Returns:
+            The regions supported by default for the Skypilot.
+        """
+        matcher = re.compile(r"us-.*")
+        return {k: v for k, v in cls.locations().items() if matcher.match(v)}
+
+    def get_deployment_config(
+        self,
+    ) -> StackDeploymentConfig:
+        """Return the configuration to deploy the ZenML stack to the specified cloud provider.
+
+        The configuration should include:
+
+        * a cloud provider console URL where the user will be redirected to
+        deploy the ZenML stack. The URL should include as many pre-filled
+        URL query parameters as possible.
+        * a textual description of the URL
+        * some deployment providers may require additional configuration
+        parameters to be passed to the cloud provider in addition to the
+        deployment URL query parameters. Where that is the case, this method
+        should also return a string that the user can copy and paste into the
+        cloud provider console to deploy the ZenML stack (e.g. a set of
+        environment variables, or YAML configuration snippet etc.).
+
+        Returns:
+            The configuration to deploy the ZenML stack to the specified cloud
+            provider.
+        """
+        params = dict(
+            cloudshell_git_repo="https://github.com/zenml-io/zenml",
+            cloudshell_workspace="infra/gcp",
+            cloudshell_open_in_editor="gcp-gar-gcs-vertex.jinja,gcp-gar-gcs-vertex-deploy.sh",
+            cloudshell_tutorial="gcp-gar-gcs-vertex.md",
+            ephemeral="true",
+        )
+        # Encode the parameters as URL query parameters
+        query_params = "&".join([f"{k}={v}" for k, v in params.items()])
+        url = (
+            f"https://shell.cloud.google.com/cloudshell/editor?{query_params}"
+        )
+
+        config = f"""
+### BEGIN CONFIGURATION ###
+ZENML_STACK_NAME={self.stack_name}
+ZENML_STACK_REGION={self.location or "europe-west3"}
+ZENML_SERVER_URL={self.zenml_server_url}
+ZENML_SERVER_API_TOKEN={self.zenml_server_api_token}
+### END CONFIGURATION ###"""
+
+        return StackDeploymentConfig(
+            deployment_url=url,
+            deployment_url_text="GCP Cloud Shell Console",
+            configuration=config,
+        )

zenml/stack_deployments/stack_deployment.py

@@ -0,0 +1,218 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+
+#       https://www.apache.org/licenses/LICENSE-2.0
+
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Functionality to deploy a ZenML stack to a cloud provider."""
+
+import datetime
+from abc import abstractmethod
+from typing import ClassVar, Dict, List, Optional
+
+from pydantic import BaseModel
+
+from zenml.client import Client
+from zenml.enums import StackComponentType, StackDeploymentProvider
+from zenml.models import (
+    DeployedStack,
+    StackDeploymentConfig,
+    StackDeploymentInfo,
+)
+
+
+class ZenMLCloudStackDeployment(BaseModel):
+    """ZenML Cloud Stack CLI Deployment base class."""
+
+    provider: ClassVar[StackDeploymentProvider]
+    deployment: ClassVar[str]
+    stack_name: str
+    zenml_server_url: str
+    zenml_server_api_token: str
+    location: Optional[str] = None
+
+    @classmethod
+    @abstractmethod
+    def description(cls) -> str:
+        """Return a description of the ZenML Cloud Stack Deployment.
+
+        This will be displayed when the user is prompted to deploy
+        the ZenML stack.
+
+        Returns:
+            A MarkDown description of the ZenML Cloud Stack Deployment.
+        """
+
+    @classmethod
+    @abstractmethod
+    def instructions(cls) -> str:
+        """Return instructions on how to deploy the ZenML stack to the specified cloud provider.
+
+        This will be displayed before the user is prompted to deploy the ZenML
+        stack.
+
+        Returns:
+            MarkDown instructions on how to deploy the ZenML stack to the
+            specified cloud provider.
+        """
+
+    @classmethod
+    @abstractmethod
+    def post_deploy_instructions(cls) -> str:
+        """Return instructions on what to do after the deployment is complete.
+
+        This will be displayed after the deployment is complete.
+
+        Returns:
+            MarkDown instructions on what to do after the deployment is
+            complete.
+        """
+
+    @classmethod
+    @abstractmethod
+    def integrations(cls) -> List[str]:
+        """Return the ZenML integrations required for the stack.
+
+        Returns:
+            The list of ZenML integrations that need to be installed for the
+            stack to be usable.
+        """
+
+    @classmethod
+    @abstractmethod
+    def permissions(cls) -> Dict[str, List[str]]:
+        """Return the permissions granted to ZenML to access the cloud resources.
+
+        Returns:
+            The permissions granted to ZenML to access the cloud resources, as
+            a dictionary grouping permissions by resource.
+        """
+
+    @classmethod
+    @abstractmethod
+    def locations(cls) -> Dict[str, str]:
+        """Return the locations where the ZenML stack can be deployed.
+
+        Returns:
+            The regions where the ZenML stack can be deployed as a map of region
+            names to region descriptions.
+        """
+
+    @classmethod
+    def skypilot_default_regions(cls) -> Dict[str, str]:
+        """Returns the regions supported by default for the Skypilot.
+
+        Returns:
+            The regions supported by default for the Skypilot.
+        """
+        return cls.locations()
+
+    @classmethod
+    def get_deployment_info(cls) -> StackDeploymentInfo:
+        """Return information about the ZenML Cloud Stack Deployment.
+
+        Returns:
+            Information about the ZenML Cloud Stack Deployment.
+        """
+        return StackDeploymentInfo(
+            provider=cls.provider,
+            description=cls.description(),
+            instructions=cls.instructions(),
+            post_deploy_instructions=cls.post_deploy_instructions(),
+            integrations=cls.integrations(),
+            permissions=cls.permissions(),
+            locations=cls.locations(),
+            skypilot_default_regions=cls.skypilot_default_regions(),
+        )
+
+    @abstractmethod
+    def get_deployment_config(
+        self,
+    ) -> StackDeploymentConfig:
+        """Return the configuration to deploy the ZenML stack to the specified cloud provider.
+
+        The configuration should include:
+
+        * a cloud provider console URL where the user will be redirected to
+        deploy the ZenML stack. The URL should include as many pre-filled
+        URL query parameters as possible.
+        * a textual description of the URL
+        * some deployment providers may require additional configuration
+        parameters to be passed to the cloud provider in addition to the
+        deployment URL query parameters. Where that is the case, this method
+        should also return a string that the user can copy and paste into the
+        cloud provider console to deploy the ZenML stack (e.g. a set of
+        environment variables, or YAML configuration snippet etc.).
+
+        Returns:
+            The configuration to deploy the ZenML stack to the specified cloud
+            provider.
+        """
+
+    def get_stack(
+        self, date_start: Optional[datetime.datetime] = None
+    ) -> Optional[DeployedStack]:
+        """Return the ZenML stack that was deployed and registered.
+
+        This method is called to retrieve a ZenML stack matching the deployment
+        provider.
+
+        Args:
+            date_start: The date when the deployment started.
+
+        Returns:
+            The ZenML stack that was deployed and registered or None if a
+            matching stack was not found.
+        """
+        client = Client()
+
+        # It's difficult to find a stack that matches the CloudFormation
+        # deployment 100% because the user can change the stack name before they
+        # deploy the stack in GCP.
+        #
+        # We try to find a full GCP stack that matches the deployment provider
+        # that was registered after this deployment was created.
+
+        # Get all stacks created after the start date
+        stacks = client.list_stacks(
+            created=f"gt:{str(date_start.replace(microsecond=0))}"
+            if date_start
+            else None,
+            sort_by="desc:created",
+            size=50,
+        )
+
+        if not stacks.items:
+            return None
+
+        # Find a stack that best matches the deployment provider
+        for stack in stacks.items:
+            if not stack.labels:
+                continue
+
+            if stack.labels.get("zenml:provider") != self.provider.value:
+                continue
+
+            if stack.labels.get("zenml:deployment") != self.deployment:
+                continue
+
+            artifact_store = stack.components[
+                StackComponentType.ARTIFACT_STORE
+            ][0]
+
+            if not artifact_store.connector:
+                continue
+
+            return DeployedStack(
+                stack=stack,
+                service_connector=artifact_store.connector,
+            )
+
+        return None

zenml/stack_deployments/utils.py

@@ -0,0 +1,48 @@
+#  Copyright (c) ZenML GmbH 2024. All Rights Reserved.
+
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+
+#       https://www.apache.org/licenses/LICENSE-2.0
+
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Functionality to deploy a ZenML stack to a cloud provider."""
+
+from typing import Type
+
+from zenml.enums import StackDeploymentProvider
+from zenml.stack_deployments.aws_stack_deployment import (
+    AWSZenMLCloudStackDeployment,
+)
+from zenml.stack_deployments.azure_stack_deployment import (
+    AZUREZenMLCloudStackDeployment,
+)
+from zenml.stack_deployments.gcp_stack_deployment import (
+    GCPZenMLCloudStackDeployment,
+)
+from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
+
+STACK_DEPLOYMENT_PROVIDERS = {
+    StackDeploymentProvider.AWS: AWSZenMLCloudStackDeployment,
+    StackDeploymentProvider.GCP: GCPZenMLCloudStackDeployment,
+    StackDeploymentProvider.AZURE: AZUREZenMLCloudStackDeployment,
+}
+
+
+def get_stack_deployment_class(
+    provider: StackDeploymentProvider,
+) -> Type[ZenMLCloudStackDeployment]:
+    """Get the ZenML Cloud Stack Deployment class for the specified provider.
+
+    Args:
+        provider: The stack deployment provider.
+
+    Returns:
+        The ZenML Cloud Stack Deployment class for the specified provider.
+    """
+    return STACK_DEPLOYMENT_PROVIDERS[provider]

zenml/steps/base_step.py

@@ -479,7 +479,7 @@ class BaseStep(metaclass=BaseStepMeta):
             bound_args = signature.bind_partial(*args, **kwargs)
         except TypeError as e:
             raise StepInterfaceError(
-                f"Wrong arguments when calling step `{self.name}`: {e}"
+                f"Wrong arguments when calling step '{self.name}': {e}"
             ) from e
 
         artifacts = {}
@@ -980,7 +980,7 @@ class BaseStep(metaclass=BaseStepMeta):
                 )
         if conflicting_parameters:
             is_plural = "s" if len(conflicting_parameters) > 1 else ""
-            msg = f"Configured parameter{is_plural} for the step `{self.name}` conflict{'' if not is_plural else 's'} with parameter{is_plural} passed in runtime:\n"
+            msg = f"Configured parameter{is_plural} for the step '{self.name}' conflict{'' if not is_plural else 's'} with parameter{is_plural} passed in runtime:\n"
             for key, values in conflicting_parameters.items():
                 msg += (
                     f"`{key}`: config=`{values[0]}` | runtime=`{values[1]}`\n"
@@ -1023,7 +1023,7 @@ To avoid this consider setting step parameters only in one place (config or code
             if output_name not in allowed_output_names:
                 raise StepInterfaceError(
                     f"Got unexpected materializers for non-existent "
-                    f"output '{output_name}' in step `{self.name}`. "
+                    f"output '{output_name}' in step '{self.name}'. "
                     f"Only materializers for the outputs "
                     f"{allowed_output_names} of this step can"
                     f" be registered."
@@ -1036,7 +1036,7 @@ To avoid this consider setting step parameters only in one place (config or code
                     ):
                         raise StepInterfaceError(
                             f"Materializer source `{source}` "
-                            f"for output '{output_name}' of step `{self.name}` "
+                            f"for output '{output_name}' of step '{self.name}' "
                             "does not resolve to a `BaseMaterializer` subclass."
                         )
 
@@ -1070,7 +1070,9 @@ To avoid this consider setting step parameters only in one place (config or code
                 or key in client_lazy_loaders
             ):
                 continue
-            raise StepInterfaceError(f"Missing entrypoint input {key}.")
+            raise StepInterfaceError(
+                f"Missing entrypoint input '{key}' in step '{self.name}'."
+            )
 
     def _finalize_configuration(
         self,

zenml/utils/function_utils.py

@@ -107,7 +107,7 @@ def _is_valid_optional_arg(arg_type: Any) -> bool:
             if (
                 args[0] not in _ALLOWED_TYPES
                 and not _is_valid_collection_arg(args[0])
-            ) or args[1] != type(None):
+            ) or args[1] is not type(None):
                 return False
         return True
     return False
@@ -146,7 +146,7 @@ def _cli_wrapped_function(func: F) -> F:
         if _is_valid_optional_arg(arg_type):
             arg_type = arg_type.__args__[0]
         arg_name = _cli_arg_name(arg_name)
-        if arg_type == bool:
+        if arg_type is bool:
             options.append(
                 click.option(
                     f"--{arg_name}",

zenml/utils/pagination_utils.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Pagination utilities."""
 
-from typing import Callable, List, TypeVar
+from typing import Any, Callable, List, TypeVar
 
 from zenml.models import BaseIdentifiedResponse, Page
 
@@ -21,20 +21,22 @@ AnyResponse = TypeVar("AnyResponse", bound=BaseIdentifiedResponse)  # type: igno
 
 
 def depaginate(
-    list_method: Callable[..., Page[AnyResponse]],
+    list_method: Callable[..., Page[AnyResponse]], **kwargs: Any
 ) -> List[AnyResponse]:
     """Depaginate the results from a client or store method that returns pages.
 
     Args:
-        list_method: The list method to wrap around.
+        list_method: The list method to depaginate.
+        **kwargs: Arguments for the list method.
 
     Returns:
         A list of the corresponding Response Models.
     """
-    page = list_method()
+    page = list_method(**kwargs)
     items = list(page.items)
     while page.index < page.total_pages:
-        page = list_method(page=page.index + 1)
+        kwargs["page"] = page.index + 1
+        page = list_method(**kwargs)
         items += list(page.items)
 
     return items