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

zenml/new/pipelines/pipeline.py

@@ -19,7 +19,6 @@ import inspect
 from contextlib import contextmanager
 from datetime import datetime
 from pathlib import Path
-from types import FunctionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -98,11 +97,11 @@ if TYPE_CHECKING:
     from zenml.config.source import Source
     from zenml.model.lazy_load import ModelVersionDataLazyLoader
     from zenml.model.model import Model
+    from zenml.types import HookSpecification
 
     StepConfigurationUpdateOrDict = Union[
         Dict[str, Any], StepConfigurationUpdate
     ]
-    HookSpecification = Union[str, "Source", FunctionType]
 
 logger = get_logger(__name__)
 
@@ -282,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
@@ -1307,6 +1306,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
         step_configurations: Optional[
             Mapping[str, "StepConfigurationUpdateOrDict"]
         ] = None,
+        steps: Optional[Mapping[str, "StepConfigurationUpdateOrDict"]] = None,
         config_path: Optional[str] = None,
         unlisted: bool = False,
         prevent_build_reuse: bool = False,
@@ -1319,6 +1319,9 @@ To avoid this consider setting pipeline parameters only in one place (config or
             schedule: Optional schedule to use for the run.
             build: Optional build to use for the run.
             step_configurations: Configurations for steps of the pipeline.
+            steps: Configurations for steps of the pipeline. This is equivalent
+                to `step_configurations`, and will be ignored if
+                `step_configurations` is set as well.
             config_path: Path to a yaml configuration file. This file will
                 be parsed as a
                 `zenml.config.pipeline_configurations.PipelineRunConfiguration`
@@ -1334,6 +1337,13 @@ To avoid this consider setting pipeline parameters only in one place (config or
         Returns:
             The copied pipeline instance.
         """
+        if steps and step_configurations:
+            logger.warning(
+                "Step configurations were passed using both the "
+                "`step_configurations` and `steps` keywords, ignoring the "
+                "values passed using the `steps` keyword."
+            )
+
         pipeline_copy = self.copy()
 
         pipeline_copy._parse_config_file(
@@ -1353,7 +1363,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 "run_name": run_name,
                 "schedule": schedule,
                 "build": build,
-                "step_configurations": step_configurations,
+                "step_configurations": step_configurations or steps,
                 "config_path": config_path,
                 "unlisted": unlisted,
                 "prevent_build_reuse": prevent_build_reuse,

zenml/new/pipelines/pipeline_decorator.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """ZenML pipeline decorator definition."""
 
-from types import FunctionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -31,8 +30,8 @@ if TYPE_CHECKING:
     from zenml.config.base_settings import SettingsOrDict
     from zenml.model.model import Model
     from zenml.new.pipelines.pipeline import Pipeline
+    from zenml.types import HookSpecification
 
-    HookSpecification = Union[str, FunctionType]
     F = TypeVar("F", bound=Callable[..., None])
 
 logger = get_logger(__name__)

zenml/new/pipelines/run_utils.py

@@ -3,16 +3,7 @@
 import time
 from collections import defaultdict
 from datetime import datetime
-from types import FunctionType
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
-    Optional,
-    Set,
-    Tuple,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, Dict, Optional, Set, Tuple, Union
 from uuid import UUID
 
 from zenml import constants
@@ -36,13 +27,11 @@ from zenml.utils import cloud_utils
 from zenml.zen_stores.base_zen_store import BaseZenStore
 
 if TYPE_CHECKING:
-    from zenml.config.source import Source
     from zenml.model.model import Model
 
     StepConfigurationUpdateOrDict = Union[
         Dict[str, Any], StepConfigurationUpdate
     ]
-    HookSpecification = Union[str, "Source", FunctionType]
 
 logger = get_logger(__name__)
 

zenml/new/steps/step_decorator.py

@@ -30,17 +30,16 @@ from typing import (
 from zenml.logger import get_logger
 
 if TYPE_CHECKING:
-    from types import FunctionType
-
     from zenml.config.base_settings import SettingsOrDict
     from zenml.config.retry_config import StepRetryConfig
     from zenml.config.source import Source
     from zenml.materializers.base_materializer import BaseMaterializer
     from zenml.model.model import Model
     from zenml.steps import BaseStep
+    from zenml.types import HookSpecification
 
     MaterializerClassOrSource = Union[str, Source, Type[BaseMaterializer]]
-    HookSpecification = Union[str, Source, FunctionType]
+
     OutputMaterializersSpecification = Union[
         MaterializerClassOrSource,
         Sequence[MaterializerClassOrSource],

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/pipelines/base_pipeline.py

@@ -15,7 +15,6 @@
 
 import inspect
 from abc import ABC, abstractmethod
-from types import FunctionType
 from typing import TYPE_CHECKING, Any, ClassVar, Dict, Mapping, Optional, Union
 from uuid import UUID
 
@@ -35,7 +34,6 @@ if TYPE_CHECKING:
     StepConfigurationUpdateOrDict = Union[
         Dict[str, Any], StepConfigurationUpdate
     ]
-    HookSpecification = Union[str, FunctionType]
 
 logger = get_logger(__name__)
 

zenml/pipelines/pipeline_decorator.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Legacy ZenML pipeline decorator definition."""
 
-from types import FunctionType
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -46,8 +45,8 @@ from zenml.pipelines.base_pipeline import (
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import SettingsOrDict
+    from zenml.types import HookSpecification
 
-    HookSpecification = Union[str, FunctionType]
 
 logger = get_logger(__name__)
 

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/stack_component.py

@@ -827,6 +827,10 @@ class StackComponent:
             f"Suspending provisioned resources not implemented for {self}."
         )
 
+    def cleanup(self) -> None:
+        """Cleans up the component after it has been used."""
+        pass
+
     def __repr__(self) -> str:
         """String representation of the stack component.
 

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,
+        )