zenml-nightly 0.61.0.dev20240711__py3-none-any.whl → 0.61.0.dev20240713__py3-none-any.whl

zenml/stack_deployments/gcp_stack_deployment.py

@@ -0,0 +1,260 @@
+#  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."""
+
+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",
+        }
+
+    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",
+            # TODO: remove this before the branch is merged
+            cloudshell_git_branch="feature/prd-482-gcp-stack-deployment",
+        )
+        # 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

@@ -15,13 +15,16 @@
 
 import datetime
 from abc import abstractmethod
-from typing import ClassVar, Dict, List, Optional, Tuple
+from typing import ClassVar, Dict, List, Optional
 
 from pydantic import BaseModel
 
-from zenml.enums import StackDeploymentProvider
+from zenml.client import Client
+from zenml.enums import StackComponentType, StackDeploymentProvider
 from zenml.models import (
     DeployedStack,
+    StackDeploymentConfig,
+    StackDeploymentInfo,
 )
 
 
@@ -29,7 +32,10 @@ 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
@@ -69,6 +75,16 @@ class ZenMLCloudStackDeployment(BaseModel):
             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]]:
@@ -89,42 +105,104 @@ class ZenMLCloudStackDeployment(BaseModel):
             names to region descriptions.
         """
 
-    @abstractmethod
-    def deploy_url(
-        self,
-        zenml_server_url: str,
-        zenml_server_api_token: str,
-    ) -> Tuple[str, str]:
-        """Return the URL to deploy the ZenML stack to the specified cloud provider.
+    @classmethod
+    def get_deployment_info(cls) -> StackDeploymentInfo:
+        """Return information about the ZenML Cloud Stack Deployment.
 
-        The URL should point to a cloud provider console where the user can
-        deploy the ZenML stack and should include as many pre-filled parameters
-        as possible.
+        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(),
+        )
 
-        Args:
-            zenml_server_url: The URL of the ZenML server.
-            zenml_server_api_token: The API token to authenticate with the ZenML
-                server.
+    @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 URL to deploy the ZenML stack to the specified cloud provider
-            and a text description of the URL.
+            The configuration to deploy the ZenML stack to the specified cloud
+            provider.
         """
 
-    @abstractmethod
     def get_stack(
-        self,
-        date_start: Optional[datetime.datetime] = None,
+        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 and stack name.
+        provider.
 
         Args:
-            date_start: The start date of the deployment.
+            date_start: The date when the deployment started.
 
         Returns:
-            The ZenML stack that was deployed and registered or None if the
-            stack was not found.
+            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

@@ -19,10 +19,14 @@ from zenml.enums import StackDeploymentProvider
 from zenml.stack_deployments.aws_stack_deployment import (
     AWSZenMLCloudStackDeployment,
 )
+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,
 }
 
 

zenml/zen_server/routers/devices_endpoints.py

@@ -27,7 +27,7 @@ from zenml.constants import (
     DEVICES,
     VERSION_1,
 )
-from zenml.enums import OAuthDeviceStatus
+from zenml.enums import OAuthDeviceStatus, OnboardingStep
 from zenml.models import (
     OAuthDeviceFilter,
     OAuthDeviceInternalUpdate,
@@ -270,6 +270,9 @@ def verify_authorized_device(
             update=update,
         )
 
+    store.update_onboarding_state(
+        completed_steps={OnboardingStep.DEVICE_VERIFIED}
+    )
     return device_model
 
 

zenml/zen_server/routers/server_endpoints.py

@@ -13,12 +13,19 @@
 #  permissions and limitations under the License.
 """Endpoint definitions for authentication (login)."""
 
-from typing import Optional
+from typing import List, Optional
 
 from fastapi import APIRouter, Security
 
 import zenml
-from zenml.constants import ACTIVATE, API, INFO, SERVER_SETTINGS, VERSION_1
+from zenml.constants import (
+    ACTIVATE,
+    API,
+    INFO,
+    ONBOARDING_STATE,
+    SERVER_SETTINGS,
+    VERSION_1,
+)
 from zenml.enums import AuthScheme
 from zenml.exceptions import IllegalOperationError
 from zenml.models import (
@@ -64,6 +71,26 @@ def server_info() -> ServerModel:
     return zen_store().get_store_info()
 
 
+@router.get(
+    ONBOARDING_STATE,
+    responses={
+        401: error_response,
+        404: error_response,
+        422: error_response,
+    },
+)
+@handle_exceptions
+def get_onboarding_state(
+    _: AuthContext = Security(authorize),
+) -> List[str]:
+    """Get the onboarding state of the server.
+
+    Returns:
+        The onboarding state of the server.
+    """
+    return zen_store().get_onboarding_state()
+
+
 # We don't have any concrete value that tells us whether a server is a cloud
 # tenant, so we use `external_server_id` as the best proxy option.
 # For cloud tenants, we don't add these endpoints as the server settings don't

zenml/zen_server/routers/stack_deployment_endpoints.py

@@ -14,13 +14,25 @@
 """Endpoint definitions for stack deployments."""
 
 import datetime
-from typing import Optional, Tuple
+from typing import Optional
 
 from fastapi import APIRouter, Request, Security
 
-from zenml.constants import API, INFO, STACK, STACK_DEPLOYMENT, URL, VERSION_1
+from zenml.constants import (
+    API,
+    CONFIG,
+    INFO,
+    STACK,
+    STACK_DEPLOYMENT,
+    STACK_DEPLOYMENT_API_TOKEN_EXPIRATION,
+    VERSION_1,
+)
 from zenml.enums import StackDeploymentProvider
-from zenml.models import DeployedStack, StackDeploymentInfo
+from zenml.models import (
+    DeployedStack,
+    StackDeploymentConfig,
+    StackDeploymentInfo,
+)
 from zenml.stack_deployments.utils import get_stack_deployment_class
 from zenml.zen_server.auth import AuthContext, authorize
 from zenml.zen_server.exceptions import error_response
@@ -54,27 +66,20 @@ def get_stack_deployment_info(
         Information about the stack deployment provider.
     """
     stack_deployment_class = get_stack_deployment_class(provider)
-    return StackDeploymentInfo(
-        provider=provider,
-        description=stack_deployment_class.description(),
-        instructions=stack_deployment_class.instructions(),
-        post_deploy_instructions=stack_deployment_class.post_deploy_instructions(),
-        permissions=stack_deployment_class.permissions(),
-        locations=stack_deployment_class.locations(),
-    )
+    return stack_deployment_class.get_deployment_info()
 
 
 @router.get(
-    URL,
+    CONFIG,
 )
 @handle_exceptions
-def get_stack_deployment_url(
+def get_stack_deployment_config(
     request: Request,
     provider: StackDeploymentProvider,
     stack_name: str,
     location: Optional[str] = None,
     auth_context: AuthContext = Security(authorize),
-) -> Tuple[str, str]:
+) -> StackDeploymentConfig:
     """Return the URL to deploy the ZenML stack to the specified cloud provider.
 
     Args:
@@ -85,8 +90,8 @@ def get_stack_deployment_url(
         auth_context: The authentication context.
 
     Returns:
-        The URL to deploy the ZenML stack to the specified cloud provider
-        and a text description of the URL.
+        The cloud provider console URL where the stack will be deployed and
+        the configuration for the stack deployment.
     """
     verify_permission(
         resource_type=ResourceType.SERVICE_CONNECTOR, action=Action.CREATE
@@ -107,12 +112,17 @@ def get_stack_deployment_url(
     assert token is not None
 
     # A new API token is generated for the stack deployment
-    expires = datetime.datetime.utcnow() + datetime.timedelta(minutes=60)
+    expires = datetime.datetime.utcnow() + datetime.timedelta(
+        minutes=STACK_DEPLOYMENT_API_TOKEN_EXPIRATION
+    )
     api_token = token.encode(expires=expires)
 
     return stack_deployment_class(
-        stack_name=stack_name, location=location
-    ).deploy_url(zenml_server_url=str(url), zenml_server_api_token=api_token)
+        stack_name=stack_name,
+        location=location,
+        zenml_server_url=str(url),
+        zenml_server_api_token=api_token,
+    ).get_deployment_config()
 
 
 @router.get(
@@ -140,5 +150,9 @@ def get_deployed_stack(
     """
     stack_deployment_class = get_stack_deployment_class(provider)
     return stack_deployment_class(
-        stack_name=stack_name, location=location
+        stack_name=stack_name,
+        location=location,
+        # These fields are not needed for this operation
+        zenml_server_url="",
+        zenml_server_api_token="",
     ).get_stack(date_start=date_start)