zenml-nightly 0.58.2.dev20240626__py3-none-any.whl → 0.61.0.dev20240712__py3-none-any.whl

zenml/models/v2/misc/full_stack.py

@@ -0,0 +1,97 @@
+#  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.
+"""Models representing full stack requests."""
+
+from typing import Any, Dict, List, Optional, Union
+from uuid import UUID
+
+from pydantic import BaseModel, Field, model_validator
+
+from zenml.constants import STR_FIELD_MAX_LENGTH
+from zenml.enums import StackComponentType
+from zenml.models.v2.base.base import BaseRequest
+
+
+class ServiceConnectorInfo(BaseModel):
+    """Information about the service connector when creating a full stack."""
+
+    type: str
+    auth_method: str
+    configuration: Dict[str, Any] = {}
+
+
+class ComponentInfo(BaseModel):
+    """Information about each stack components when creating a full stack."""
+
+    flavor: str
+    service_connector_index: Optional[int] = Field(
+        default=None,
+        title="The id of the service connector from the list "
+        "`service_connectors`.",
+        description="The id of the service connector from the list "
+        "`service_connectors` from `FullStackRequest`.",
+    )
+    service_connector_resource_id: Optional[str] = None
+    configuration: Dict[str, Any] = {}
+
+
+class FullStackRequest(BaseRequest):
+    """Request model for a full-stack."""
+
+    user: Optional[UUID] = None
+    workspace: Optional[UUID] = None
+
+    name: str = Field(
+        title="The name of the stack.", max_length=STR_FIELD_MAX_LENGTH
+    )
+    description: Optional[str] = Field(
+        default="",
+        title="The description of the stack",
+        max_length=STR_FIELD_MAX_LENGTH,
+    )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
+    service_connectors: List[Union[UUID, ServiceConnectorInfo]] = Field(
+        default=[],
+        title="The service connectors dictionary for the full stack "
+        "registration.",
+        description="The UUID of an already existing service connector or "
+        "request information to create a service connector from "
+        "scratch.",
+    )
+    components: Dict[StackComponentType, Union[UUID, ComponentInfo]] = Field(
+        title="The mapping for the components of the full stack registration.",
+        description="The mapping from component types to either UUIDs of "
+        "existing components or request information for brand new "
+        "components.",
+    )
+
+    @model_validator(mode="after")
+    def _validate_indexes_in_components(self) -> "FullStackRequest":
+        for component in self.components.values():
+            if isinstance(component, ComponentInfo):
+                if component.service_connector_index is not None:
+                    if (
+                        component.service_connector_index < 0
+                        or component.service_connector_index
+                        >= len(self.service_connectors)
+                    ):
+                        raise ValueError(
+                            f"Service connector index {component.service_connector_index} "
+                            "is out of range. Please provide a valid index referring to "
+                            "the position in the list of service connectors."
+                        )
+        return self

zenml/models/v2/misc/stack_deployment.py

@@ -0,0 +1,86 @@
+#  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.
+"""Models related to cloud stack deployments."""
+
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+from zenml.enums import StackDeploymentProvider
+from zenml.models.v2.core.service_connector import ServiceConnectorResponse
+from zenml.models.v2.core.stack import StackResponse
+
+
+class StackDeploymentInfo(BaseModel):
+    """Information about a stack deployment."""
+
+    provider: StackDeploymentProvider = Field(
+        title="The provider of the stack deployment."
+    )
+    description: str = Field(
+        title="The description of the stack deployment.",
+        description="The description of the stack deployment.",
+    )
+    instructions: str = Field(
+        title="The instructions for deploying the stack.",
+        description="The instructions for deploying the stack.",
+    )
+    post_deploy_instructions: str = Field(
+        title="The instructions for post-deployment.",
+        description="The instructions for post-deployment.",
+    )
+    integrations: List[str] = Field(
+        title="ZenML integrations required for the stack.",
+        description="The list of ZenML integrations that need to be installed "
+        "for the stack to be usable.",
+    )
+    permissions: Dict[str, List[str]] = Field(
+        title="The permissions granted to ZenML to access the cloud resources.",
+        description="The permissions granted to ZenML to access the cloud "
+        "resources, as a dictionary grouping permissions by resource.",
+    )
+    locations: Dict[str, str] = Field(
+        title="The locations where the stack can be deployed.",
+        description="The locations where the stack can be deployed, as a "
+        "dictionary mapping location names to descriptions.",
+    )
+
+
+class StackDeploymentConfig(BaseModel):
+    """Configuration about a stack deployment."""
+
+    deployment_url: str = Field(
+        title="The cloud provider console URL where the stack will be deployed.",
+    )
+    deployment_url_text: str = Field(
+        title="A textual description for the cloud provider console URL.",
+    )
+    configuration: Optional[str] = Field(
+        title="Configuration for the stack deployment that the user must "
+        "manually configure into the cloud provider console.",
+    )
+
+
+class DeployedStack(BaseModel):
+    """Information about a deployed stack."""
+
+    stack: StackResponse = Field(
+        title="The stack that was deployed.",
+        description="The stack that was deployed.",
+    )
+    service_connector: Optional[ServiceConnectorResponse] = Field(
+        default=None,
+        title="The service connector for the deployed stack.",
+        description="The service connector for the deployed stack.",
+    )

zenml/new/pipelines/pipeline.py

@@ -281,7 +281,7 @@ class Pipeline:
         raise RuntimeError(
             f"Cannot get the model of pipeline '{self.name}' because it has "
             f"not been registered yet. Please ensure that the pipeline has "
-            f"been run and that at least one step has been executed."
+            f"been run or built and try again."
         )
 
     @contextmanager

zenml/orchestrators/input_utils.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Utilities for inputs."""
 
-import functools
 from typing import TYPE_CHECKING, Dict, List, Tuple
 from uuid import UUID
 
@@ -48,13 +47,11 @@ def resolve_step_inputs(
     """
     from zenml.models import ArtifactVersionResponse, RunMetadataResponse
 
-    list_run_steps = functools.partial(
-        Client().list_run_steps, pipeline_run_id=run_id
-    )
-
     current_run_steps = {
         run_step.name: run_step
-        for run_step in pagination_utils.depaginate(list_run_steps)
+        for run_step in pagination_utils.depaginate(
+            Client().list_run_steps, pipeline_run_id=run_id
+        )
     }
 
     input_artifacts: Dict[str, "ArtifactVersionResponse"] = {}

zenml/stack/stack.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Implementation of the ZenML Stack class."""
 
-import functools
 import itertools
 import json
 import os
@@ -155,11 +154,9 @@ class Stack:
 
         # Run a hydrated list call once to avoid one request per component
         component_models = pagination_utils.depaginate(
-            list_method=functools.partial(
-                Client().list_stack_components,
-                stack_id=stack_model.id,
-                hydrate=True,
-            )
+            Client().list_stack_components,
+            stack_id=stack_model.id,
+            hydrate=True,
         )
 
         stack_components = {

zenml/stack_deployments/__init__.py

@@ -0,0 +1,14 @@
+#  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.
+"""ZenML Stack Deployments."""

zenml/stack_deployments/aws_stack_deployment.py

@@ -0,0 +1,254 @@
+#  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 AWS."""
+
+from typing import ClassVar, Dict, List
+
+from zenml.enums import StackDeploymentProvider
+from zenml.models import StackDeploymentConfig
+from zenml.stack_deployments.stack_deployment import ZenMLCloudStackDeployment
+from zenml.utils.string_utils import random_str
+
+AWS_DEPLOYMENT_TYPE = "cloud-formation"
+
+
+class AWSZenMLCloudStackDeployment(ZenMLCloudStackDeployment):
+    """AWS ZenML Cloud Stack Deployment."""
+
+    provider: ClassVar[StackDeploymentProvider] = StackDeploymentProvider.AWS
+    deployment: ClassVar[str] = AWS_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 AWS ZenML stack authenticated and connected to
+all the necessary cloud infrastructure resources required to run pipelines in
+AWS.
+"""
+
+    @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 the AWS console in your browser where you'll be asked
+to log into your AWS account and create a CloudFormation ZenML stack. The stack
+parameters will be pre-filled with the necessary information to connect ZenML to
+your AWS account, so you should only need to review and confirm the stack.
+
+**NOTE**: The CloudFormation stack will create the following new resources in
+your AWS account. Please ensure you have the necessary permissions and are aware
+of any potential costs:
+
+- An S3 bucket registered as a [ZenML artifact store](https://docs.zenml.io/stack-components/artifact-stores/s3).
+- An ECR repository registered as a [ZenML container registry](https://docs.zenml.io/stack-components/container-registries/aws).
+- Sagemaker registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/sagemaker).
+- An IAM user and IAM role with the minimum necessary permissions to access the
+above resources.
+- An AWS 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/aws-service-connector).
+
+The CloudFormation stack will automatically create an AWS 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
+CloudFormation stack.
+
+**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 AWS pricing. 
+Some services may be eligible for the AWS Free Tier. Use [the AWS Pricing Calculator](https://calculator.aws)
+for a detailed estimate based on your usage.
+
+💡 **After the CloudFormation stack is deployed, you can 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 CloudFormation at any time to revoke ZenML's access to your AWS account and
+to clean up the resources created by the stack by using the AWS CloudFormation
+console.
+"""
+
+    @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 [
+            "aws",
+            "s3",
+        ]
+
+    @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 {
+            "S3 Bucket": [
+                "s3:ListBucket",
+                "s3:GetObject",
+                "s3:PutObject",
+                "s3:DeleteObject",
+            ],
+            "ECR Repository": [
+                "ecr:DescribeRepositories",
+                "ecr:ListRepositories",
+                "ecr:DescribeRegistry",
+                "ecr:BatchGetImage",
+                "ecr:DescribeImages",
+                "ecr:BatchCheckLayerAvailability",
+                "ecr:GetDownloadUrlForLayer",
+                "ecr:InitiateLayerUpload",
+                "ecr:UploadLayerPart",
+                "ecr:CompleteLayerUpload",
+                "ecr:PutImage",
+                "ecr:GetAuthorizationToken",
+            ],
+            "SageMaker (Client)": [
+                "sagemaker:CreatePipeline",
+                "sagemaker:StartPipelineExecution",
+                "sagemaker:DescribePipeline",
+                "sagemaker:DescribePipelineExecution",
+            ],
+            "SageMaker (Jobs)": [
+                "AmazonSageMakerFullAccess",
+            ],
+        }
+
+    @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 AWS regions
+
+        # Based on the AWS regions listed at
+        # https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/using-regions-availability-zones.html
+        return {
+            "US East (Ohio)": "us-east-2",
+            "US East (N. Virginia)": "us-east-1",
+            "US West (N. California)": "us-west-1",
+            "US West (Oregon)": "us-west-2",
+            "Africa (Cape Town)": "af-south-1",
+            "Asia Pacific (Hong Kong)": "ap-east-1",
+            "Asia Pacific (Hyderabad)": "ap-south-2",
+            "Asia Pacific (Jakarta)": "ap-southeast-3",
+            "Asia Pacific (Melbourne)": "ap-southeast-4",
+            "Asia Pacific (Mumbai)": "ap-south-1",
+            "Asia Pacific (Osaka)": "ap-northeast-3",
+            "Asia Pacific (Seoul)": "ap-northeast-2",
+            "Asia Pacific (Singapore)": "ap-southeast-1",
+            "Asia Pacific (Sydney)": "ap-southeast-2",
+            "Asia Pacific (Tokyo)": "ap-northeast-1",
+            "Canada (Central)": "ca-central-1",
+            "Canada West (Calgary)": "ca-west-1",
+            "Europe (Frankfurt)": "eu-central-1",
+            "Europe (Ireland)": "eu-west-1",
+            "Europe (London)": "eu-west-2",
+            "Europe (Milan)": "eu-south-1",
+            "Europe (Paris)": "eu-west-3",
+            "Europe (Spain)": "eu-south-2",
+            "Europe (Stockholm)": "eu-north-1",
+            "Europe (Zurich)": "eu-central-2",
+            "Israel (Tel Aviv)": "il-central-1",
+            "Middle East (Bahrain)": "me-south-1",
+            "Middle East (UAE)": "me-central-1",
+            "South America (São Paulo)": "sa-east-1",
+        }
+
+    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(
+            stackName=self.stack_name,
+            templateURL="https://zenml-cf-templates.s3.eu-central-1.amazonaws.com/aws-ecr-s3-sagemaker.yaml",
+            param_ResourceName=f"zenml-{random_str(6).lower()}",
+            param_ZenMLServerURL=self.zenml_server_url,
+            param_ZenMLServerAPIToken=self.zenml_server_api_token,
+        )
+        # Encode the parameters as URL query parameters
+        query_params = "&".join([f"{k}={v}" for k, v in params.items()])
+
+        region = ""
+        if self.location:
+            region = f"region={self.location}"
+
+        url = (
+            f"https://console.aws.amazon.com/cloudformation/home?"
+            f"{region}#/stacks/create/review?{query_params}"
+        )
+
+        return StackDeploymentConfig(
+            deployment_url=url,
+            deployment_url_text="AWS CloudFormation Console",
+            configuration=None,
+        )