zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/integrations/aws/container_registries/aws_container_registry.py

@@ -17,15 +17,23 @@ import re
 from typing import List, Optional, cast
 
 import boto3
+from botocore.client import BaseClient
 from botocore.exceptions import BotoCoreError, ClientError
 
+from zenml.client import Client
 from zenml.container_registries.base_container_registry import (
     BaseContainerRegistry,
 )
 from zenml.integrations.aws.flavors.aws_container_registry_flavor import (
     AWSContainerRegistryConfig,
 )
+from zenml.integrations.aws.service_connectors import (
+    AWSServiceConnector,
+)
 from zenml.logger import get_logger
+from zenml.service_connectors.service_connector_registry import (
+    service_connector_registry,
+)
 
 logger = get_logger(__name__)
 
@@ -61,6 +69,35 @@ class AWSContainerRegistry(BaseContainerRegistry):
 
         return match.group(1)
 
+    def _get_ecr_client(self) -> BaseClient:
+        """Get an ECR client.
+
+        If this container registry is configured with an AWS service connector,
+        we use that connector to create an authenticated client. Otherwise
+        local AWS credentials will be used.
+
+        Returns:
+            An ECR client.
+        """
+        if self.connector:
+            try:
+                model = Client().get_service_connector(self.connector)
+                connector = service_connector_registry.instantiate_connector(
+                    model=model
+                )
+                assert isinstance(connector, AWSServiceConnector)
+                return connector.get_ecr_client()
+            except Exception as e:
+                logger.error(
+                    "Unable to get ECR client from service connector: %s",
+                    str(e),
+                )
+
+        return boto3.Session().client(
+            "ecr",
+            region_name=self._get_region(),
+        )
+
     def prepare_image_push(self, image_name: str) -> None:
         """Logs warning message if trying to push an image for which no repository exists.
 
@@ -76,10 +113,9 @@ class AWSContainerRegistry(BaseContainerRegistry):
             raise ValueError(f"Invalid docker image name '{image_name}'.")
         repo_name = match.group(1)
 
+        client = self._get_ecr_client()
         try:
-            response = boto3.client(
-                "ecr", region_name=self._get_region()
-            ).describe_repositories()
+            response = client.describe_repositories()
         except (BotoCoreError, ClientError):
             logger.warning(
                 "Amazon ECR requires you to create a repository before you can "
@@ -123,9 +159,9 @@ class AWSContainerRegistry(BaseContainerRegistry):
         """
         return (
             "Amazon ECR requires you to create a repository before you can "
-            "push an image to it. If you want to for example run a pipeline "
-            "using our Kubeflow orchestrator, ZenML will automatically build a "
-            f"docker image called `{self.config.uri}/zenml-kubeflow:<PIPELINE_NAME>` "
-            f"and try to push it. This will fail unless you create the "
-            f"repository `zenml-kubeflow` inside your amazon registry."
+            "push an image to it. If you want to for run a pipeline "
+            "using a remote orchestrator, ZenML will automatically build a "
+            f"docker image called `{self.config.uri}/zenml:<PIPELINE_NAME>` "
+            f"and try to push it. This will fail unless you create a "
+            f"repository called `zenml` inside your Amazon ECR."
         )

zenml/integrations/aws/flavors/__init__.py

@@ -17,6 +17,10 @@ from zenml.integrations.aws.flavors.aws_container_registry_flavor import (
     AWSContainerRegistryConfig,
     AWSContainerRegistryFlavor,
 )
+from zenml.integrations.aws.flavors.aws_image_builder_flavor import (
+    AWSImageBuilderConfig,
+    AWSImageBuilderFlavor,
+)
 from zenml.integrations.aws.flavors.sagemaker_orchestrator_flavor import (
     SagemakerOrchestratorConfig,
     SagemakerOrchestratorFlavor,
@@ -29,6 +33,8 @@ from zenml.integrations.aws.flavors.sagemaker_step_operator_flavor import (
 __all__ = [
     "AWSContainerRegistryFlavor",
     "AWSContainerRegistryConfig",
+    "AWSImageBuilderConfig",
+    "AWSImageBuilderFlavor",
     "SagemakerStepOperatorFlavor",
     "SagemakerStepOperatorConfig",
     "SagemakerOrchestratorFlavor",

zenml/integrations/aws/flavors/aws_container_registry_flavor.py

@@ -15,7 +15,7 @@
 
 from typing import TYPE_CHECKING, Optional, Type
 
-from pydantic import validator
+from pydantic import field_validator
 
 from zenml.constants import DOCKER_REGISTRY_RESOURCE_TYPE
 from zenml.container_registries.base_container_registry import (
@@ -37,7 +37,8 @@ if TYPE_CHECKING:
 class AWSContainerRegistryConfig(BaseContainerRegistryConfig):
     """Configuration for AWS Container Registry."""
 
-    @validator("uri")
+    @field_validator("uri")
+    @classmethod
     def validate_aws_uri(cls, uri: str) -> str:
         """Validates that the URI is in the correct format.
 

zenml/integrations/aws/flavors/aws_image_builder_flavor.py

@@ -0,0 +1,146 @@
+#  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.
+"""AWS Code Build image builder flavor."""
+
+from typing import TYPE_CHECKING, Dict, Optional, Type
+
+from zenml.image_builders import BaseImageBuilderConfig, BaseImageBuilderFlavor
+from zenml.integrations.aws import (
+    AWS_CONNECTOR_TYPE,
+    AWS_IMAGE_BUILDER_FLAVOR,
+    AWS_RESOURCE_TYPE,
+)
+from zenml.models import ServiceConnectorRequirements
+
+if TYPE_CHECKING:
+    from zenml.integrations.aws.image_builders import AWSImageBuilder
+
+
+DEFAULT_CLOUDBUILD_IMAGE = "bentolor/docker-dind-awscli"
+DEFAULT_CLOUDBUILD_COMPUTE_TYPE = "BUILD_GENERAL1_SMALL"
+
+
+class AWSImageBuilderConfig(BaseImageBuilderConfig):
+    """AWS Code Build image builder configuration.
+
+    Attributes:
+        code_build_project: The name of an existing AWS CodeBuild project to use
+            to build the image. The CodeBuild project must exist in the AWS
+            account and region inferred from the AWS service connector
+            credentials or implicitly from the local AWS config.
+        build_image: The Docker image to use for the AWS CodeBuild environment.
+            The image must have Docker installed and be able to run Docker
+            commands. The default image is bentolor/docker-dind-awscli.
+            This can be customized to use a mirror, if needed, in case the
+            Dockerhub image is not accessible or rate-limited.
+        custom_env_vars: Custom environment variables to pass to the AWS
+            CodeBuild build.
+        compute_type: The compute type to use for the AWS CodeBuild build.
+            The default is BUILD_GENERAL1_SMALL.
+        implicit_container_registry_auth: Whether to use implicit authentication
+            to authenticate the AWS Code Build build to the container registry
+            when pushing container images. If set to False, the container
+            registry credentials must be explicitly configured for the container
+            registry stack component or the container registry stack component
+            must be linked to a service connector.
+            NOTE: When implicit_container_registry_auth is set to False, the
+            container registry credentials will be passed to the AWS Code Build
+            build as environment variables. This is not recommended for
+            production use unless your service connector is configured to
+            generate short-lived credentials.
+    """
+
+    code_build_project: str
+    build_image: str = DEFAULT_CLOUDBUILD_IMAGE
+    custom_env_vars: Optional[Dict[str, str]] = None
+    compute_type: str = DEFAULT_CLOUDBUILD_COMPUTE_TYPE
+    implicit_container_registry_auth: bool = True
+
+
+class AWSImageBuilderFlavor(BaseImageBuilderFlavor):
+    """AWS Code Build image builder flavor."""
+
+    @property
+    def name(self) -> str:
+        """The flavor name.
+
+        Returns:
+            The name of the flavor.
+        """
+        return AWS_IMAGE_BUILDER_FLAVOR
+
+    @property
+    def service_connector_requirements(
+        self,
+    ) -> Optional[ServiceConnectorRequirements]:
+        """Service connector resource requirements for service connectors.
+
+        Specifies resource requirements that are used to filter the available
+        service connector types that are compatible with this flavor.
+
+        Returns:
+            Requirements for compatible service connectors, if a service
+            connector is required for this flavor.
+        """
+        return ServiceConnectorRequirements(
+            connector_type=AWS_CONNECTOR_TYPE,
+            resource_type=AWS_RESOURCE_TYPE,
+        )
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A url to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A url to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A url to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/image_builder/aws.png"
+
+    @property
+    def config_class(self) -> Type[BaseImageBuilderConfig]:
+        """The config class.
+
+        Returns:
+            The config class.
+        """
+        return AWSImageBuilderConfig
+
+    @property
+    def implementation_class(self) -> Type["AWSImageBuilder"]:
+        """Implementation class.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.aws.image_builders import AWSImageBuilder
+
+        return AWSImageBuilder

zenml/integrations/aws/flavors/sagemaker_orchestrator_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
 
+from pydantic import Field, model_validator
+
 from zenml.config.base_settings import BaseSettings
 from zenml.integrations.aws import (
     AWS_RESOURCE_TYPE,
@@ -23,23 +25,44 @@ from zenml.integrations.aws import (
 from zenml.models import ServiceConnectorRequirements
 from zenml.orchestrators import BaseOrchestratorConfig
 from zenml.orchestrators.base_orchestrator import BaseOrchestratorFlavor
+from zenml.utils import deprecation_utils
 from zenml.utils.secret_utils import SecretField
 
 if TYPE_CHECKING:
     from zenml.integrations.aws.orchestrators import SagemakerOrchestrator
 
+DEFAULT_TRAINING_INSTANCE_TYPE = "ml.m5.xlarge"
+DEFAULT_PROCESSING_INSTANCE_TYPE = "ml.t3.medium"
+DEFAULT_OUTPUT_DATA_S3_MODE = "EndOfJob"
+
 
 class SagemakerOrchestratorSettings(BaseSettings):
     """Settings for the Sagemaker orchestrator.
 
     Attributes:
+        synchronous: If `True`, the client running a pipeline using this
+            orchestrator waits until all steps finish running. If `False`,
+            the client returns immediately and the pipeline is executed
+            asynchronously. Defaults to `True`.
         instance_type: The instance type to use for the processing job.
-        processor_role: The IAM role to use for the step execution on a Processor.
+        execution_role: The IAM role to use for the step execution.
+        processor_role: DEPRECATED: use `execution_role` instead.
         volume_size_in_gb: The size of the EBS volume to use for the processing
             job.
         max_runtime_in_seconds: The maximum runtime in seconds for the
             processing job.
-        processor_tags: Tags to apply to the Processor assigned to the step.
+        tags: Tags to apply to the Processor/Estimator assigned to the step.
+        pipeline_tags: Tags to apply to the pipeline via the
+            sagemaker.workflow.pipeline.Pipeline.create method.
+        processor_tags: DEPRECATED: use `tags` instead.
+        keep_alive_period_in_seconds: The time in seconds after which the
+            provisioned instance will be terminated if not used. This is only
+            applicable for TrainingStep type and it is not possible to use
+            TrainingStep type if the `output_data_s3_uri` is set to Dict[str, str].
+        use_training_step: Whether to use the TrainingStep type.
+            It is not possible to use TrainingStep type
+            if the `output_data_s3_uri` is set to Dict[str, str] or if the
+            `output_data_s3_mode` != "EndOfJob".
         processor_args: Arguments that are directly passed to the SageMaker
             Processor for a specific step, allowing for overriding the default
             settings provided when configuring the component. See
@@ -48,6 +71,13 @@ class SagemakerOrchestratorSettings(BaseSettings):
             For processor_args.instance_type, check
             https://docs.aws.amazon.com/sagemaker/latest/dg/notebooks-available-instance-types.html
             for a list of available instance types.
+        estimator_args: Arguments that are directly passed to the SageMaker
+            Estimator for a specific step, allowing for overriding the default
+            settings provided when configuring the component. See
+            https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html#sagemaker.estimator.Estimator
+            for a full list of arguments.
+            For a list of available instance types, check
+            https://docs.aws.amazon.com/sagemaker/latest/dg/cmn-info-instance-types.html.
         input_data_s3_mode: How data is made available to the container.
             Two possible input modes: File, Pipe.
         input_data_s3_uri: S3 URI where data is located if not locally,
@@ -72,21 +102,75 @@ class SagemakerOrchestratorSettings(BaseSettings):
                     Data must be available locally in /opt/ml/processing/output/data/<ChannelName>.
     """
 
-    instance_type: str = "ml.t3.medium"
-    processor_role: Optional[str] = None
+    synchronous: bool = True
+
+    instance_type: Optional[str] = None
+    execution_role: Optional[str] = None
     volume_size_in_gb: int = 30
     max_runtime_in_seconds: int = 86400
-    processor_tags: Dict[str, str] = {}
+    tags: Dict[str, str] = {}
+    pipeline_tags: Dict[str, str] = {}
+    keep_alive_period_in_seconds: Optional[int] = 300  # 5 minutes
+    use_training_step: Optional[bool] = None
 
     processor_args: Dict[str, Any] = {}
+    estimator_args: Dict[str, Any] = {}
+
     input_data_s3_mode: str = "File"
-    input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = None
+    input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(
+        default=None, union_mode="left_to_right"
+    )
+
+    output_data_s3_mode: str = DEFAULT_OUTPUT_DATA_S3_MODE
+    output_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(
+        default=None, union_mode="left_to_right"
+    )
+
+    processor_role: Optional[str] = None
+    processor_tags: Dict[str, str] = {}
+    _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
+        ("processor_role", "execution_role"), ("processor_tags", "tags")
+    )
 
-    output_data_s3_mode: str = "EndOfJob"
-    output_data_s3_uri: Optional[Union[str, Dict[str, str]]] = None
+    @model_validator(mode="before")
+    def validate_model(cls, data: Dict[str, Any]) -> Dict[str, Any]:
+        """Check if model is configured correctly.
 
+        Args:
+            data: The model data.
+
+        Returns:
+            The validated model data.
 
-class SagemakerOrchestratorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+        Raises:
+            ValueError: If the model is configured incorrectly.
+        """
+        use_training_step = data.get("use_training_step", True)
+        output_data_s3_uri = data.get("output_data_s3_uri", None)
+        output_data_s3_mode = data.get(
+            "output_data_s3_mode", DEFAULT_OUTPUT_DATA_S3_MODE
+        )
+        if use_training_step and (
+            isinstance(output_data_s3_uri, dict)
+            or (
+                isinstance(output_data_s3_uri, str)
+                and (output_data_s3_mode != DEFAULT_OUTPUT_DATA_S3_MODE)
+            )
+        ):
+            raise ValueError(
+                "`use_training_step=True` is not supported when `output_data_s3_uri` is a dict or "
+                f"when `output_data_s3_mode` is not '{DEFAULT_OUTPUT_DATA_S3_MODE}'."
+            )
+        instance_type = data.get("instance_type", None)
+        if instance_type is None:
+            if use_training_step:
+                data["instance_type"] = DEFAULT_TRAINING_INSTANCE_TYPE
+            else:
+                data["instance_type"] = DEFAULT_PROCESSING_INSTANCE_TYPE
+        return data
+
+
+class SagemakerOrchestratorConfig(
     BaseOrchestratorConfig, SagemakerOrchestratorSettings
 ):
     """Config for the Sagemaker orchestrator.
@@ -99,10 +183,6 @@ class SagemakerOrchestratorConfig(  # type: ignore[misc] # https://github.com/py
         loaded from the default AWS config.
 
     Attributes:
-        synchronous: If `True`, the client running a pipeline using this
-            orchestrator waits until all steps finish running. If `False`,
-            the client returns immediately and the pipeline is executed
-            asynchronously. Defaults to `True`.
         execution_role: The IAM role ARN to use for the pipeline.
         aws_access_key_id: The AWS access key ID to use to authenticate to AWS.
             If not provided, the value from the default AWS config will be used.
@@ -122,10 +202,9 @@ class SagemakerOrchestratorConfig(  # type: ignore[misc] # https://github.com/py
             "sagemaker-{region}-{aws-account-id}".
     """
 
-    synchronous: bool = True
     execution_role: str
-    aws_access_key_id: Optional[str] = SecretField()
-    aws_secret_access_key: Optional[str] = SecretField()
+    aws_access_key_id: Optional[str] = SecretField(default=None)
+    aws_secret_access_key: Optional[str] = SecretField(default=None)
     aws_profile: Optional[str] = None
     aws_auth_role_arn: Optional[str] = None
     region: Optional[str] = None
@@ -144,6 +223,15 @@ class SagemakerOrchestratorConfig(  # type: ignore[misc] # https://github.com/py
         """
         return True
 
+    @property
+    def is_synchronous(self) -> bool:
+        """Whether the orchestrator runs synchronous or not.
+
+        Returns:
+            Whether the orchestrator runs synchronous or not.
+        """
+        return self.synchronous
+
 
 class SagemakerOrchestratorFlavor(BaseOrchestratorFlavor):
     """Flavor for the Sagemaker orchestrator."""

zenml/integrations/aws/flavors/sagemaker_step_operator_flavor.py

@@ -15,6 +15,8 @@
 
 from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
 
+from pydantic import Field
+
 from zenml.config.base_settings import BaseSettings
 from zenml.integrations.aws import (
     AWS_RESOURCE_TYPE,
@@ -58,7 +60,9 @@ class SagemakerStepOperatorSettings(BaseSettings):
 
     instance_type: Optional[str] = None
     experiment_name: Optional[str] = None
-    input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = None
+    input_data_s3_uri: Optional[Union[str, Dict[str, str]]] = Field(
+        default=None, union_mode="left_to_right"
+    )
     estimator_args: Dict[str, Any] = {}
 
     _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
@@ -66,7 +70,7 @@ class SagemakerStepOperatorSettings(BaseSettings):
     )
 
 
-class SagemakerStepOperatorConfig(  # type: ignore[misc] # https://github.com/pydantic/pydantic/issues/4173
+class SagemakerStepOperatorConfig(
     BaseStepOperatorConfig, SagemakerStepOperatorSettings
 ):
     """Config for the Sagemaker step operator.

zenml/integrations/aws/image_builders/__init__.py

@@ -0,0 +1,20 @@
+#  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.
+"""Initialization for the AWS image builder."""
+
+from zenml.integrations.aws.image_builders.aws_image_builder import (
+    AWSImageBuilder,
+)
+
+__all__ = ["AWSImageBuilder"]