zenml-nightly 0.70.0.dev20241201__py3-none-any.whl → 0.71.0.dev20241223__py3-none-any.whl

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/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"]

zenml/integrations/aws/image_builders/aws_image_builder.py

@@ -0,0 +1,307 @@
+#  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 implementation."""
+
+import time
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, cast
+from urllib.parse import urlparse
+from uuid import uuid4
+
+import boto3
+
+from zenml.enums import StackComponentType
+from zenml.image_builders import BaseImageBuilder
+from zenml.integrations.aws import (
+    AWS_CONTAINER_REGISTRY_FLAVOR,
+)
+from zenml.integrations.aws.flavors import AWSImageBuilderConfig
+from zenml.logger import get_logger
+from zenml.stack import StackValidator
+from zenml.utils.archivable import ArchiveType
+
+if TYPE_CHECKING:
+    from zenml.container_registries import BaseContainerRegistry
+    from zenml.image_builders import BuildContext
+    from zenml.stack import Stack
+
+logger = get_logger(__name__)
+
+
+class AWSImageBuilder(BaseImageBuilder):
+    """AWS Code Build image builder implementation."""
+
+    _code_build_client: Optional[Any] = None
+
+    @property
+    def config(self) -> AWSImageBuilderConfig:
+        """The stack component configuration.
+
+        Returns:
+            The configuration.
+        """
+        return cast(AWSImageBuilderConfig, self._config)
+
+    @property
+    def is_building_locally(self) -> bool:
+        """Whether the image builder builds the images on the client machine.
+
+        Returns:
+            True if the image builder builds locally, False otherwise.
+        """
+        return False
+
+    @property
+    def validator(self) -> Optional["StackValidator"]:
+        """Validates the stack for the AWS Code Build Image Builder.
+
+        The AWS Code Build Image Builder requires a container registry to
+        push the image to and an S3 Artifact Store to upload the build context,
+        so AWS Code Build can access it.
+
+        Returns:
+            Stack validator.
+        """
+
+        def _validate_remote_components(stack: "Stack") -> Tuple[bool, str]:
+            if stack.artifact_store.flavor != "s3":
+                return False, (
+                    "The AWS Image Builder requires an S3 Artifact Store to "
+                    "upload the build context, so AWS Code Build can access it."
+                    "Please update your stack to include an S3 Artifact Store "
+                    "and try again."
+                )
+
+            return True, ""
+
+        return StackValidator(
+            required_components={StackComponentType.CONTAINER_REGISTRY},
+            custom_validation_function=_validate_remote_components,
+        )
+
+    @property
+    def code_build_client(self) -> Any:
+        """The authenticated AWS Code Build client to use for interacting with AWS services.
+
+        Returns:
+            The authenticated AWS Code Build client.
+
+        Raises:
+            RuntimeError: If the AWS Code Build client cannot be created.
+        """
+        if (
+            self._code_build_client is not None
+            and self.connector_has_expired()
+        ):
+            self._code_build_client = None
+        if self._code_build_client is not None:
+            return self._code_build_client
+
+        # Option 1: Service connector
+        if connector := self.get_connector():
+            boto_session = connector.connect()
+            if not isinstance(boto_session, boto3.Session):
+                raise RuntimeError(
+                    f"Expected to receive a `boto3.Session` object from the "
+                    f"linked connector, but got type `{type(boto_session)}`."
+                )
+        # Option 2: Implicit configuration
+        else:
+            boto_session = boto3.Session()
+
+        self._code_build_client = boto_session.client("codebuild")
+        return self._code_build_client
+
+    def build(
+        self,
+        image_name: str,
+        build_context: "BuildContext",
+        docker_build_options: Dict[str, Any],
+        container_registry: Optional["BaseContainerRegistry"] = None,
+    ) -> str:
+        """Builds and pushes a Docker image.
+
+        Args:
+            image_name: Name of the image to build and push.
+            build_context: The build context to use for the image.
+            docker_build_options: Docker build options.
+            container_registry: Optional container registry to push to.
+
+        Returns:
+            The Docker image name with digest.
+
+        Raises:
+            RuntimeError: If no container registry is passed.
+            RuntimeError: If the Cloud Build build fails.
+        """
+        if not container_registry:
+            raise RuntimeError(
+                "The AWS Image Builder requires a container registry to push "
+                "the image to. Please provide one and try again."
+            )
+
+        logger.info("Using AWS Code Build to build image `%s`", image_name)
+        cloud_build_context = self._upload_build_context(
+            build_context=build_context,
+            parent_path_directory_name=f"code-build-contexts/{str(self.id)}",
+            archive_type=ArchiveType.ZIP,
+        )
+
+        url_parts = urlparse(cloud_build_context)
+        bucket = url_parts.netloc
+        object_path = url_parts.path.lstrip("/")
+        logger.info(
+            "Build context located in bucket `%s` and object path `%s`",
+            bucket,
+            object_path,
+        )
+
+        # Pass authentication credentials as environment variables, if
+        # the container registry has credentials and if implicit authentication
+        # is disabled
+        environment_variables_override: Dict[str, str] = {}
+        pre_build_commands = []
+        if not self.config.implicit_container_registry_auth:
+            credentials = container_registry.credentials
+            if credentials:
+                environment_variables_override = {
+                    "CONTAINER_REGISTRY_USERNAME": credentials[0],
+                    "CONTAINER_REGISTRY_PASSWORD": credentials[1],
+                }
+                pre_build_commands = [
+                    "echo Logging in to container registry",
+                    'echo "$CONTAINER_REGISTRY_PASSWORD" | docker login --username "$CONTAINER_REGISTRY_USERNAME" --password-stdin '
+                    f"{container_registry.config.uri}",
+                ]
+        elif container_registry.flavor == AWS_CONTAINER_REGISTRY_FLAVOR:
+            pre_build_commands = [
+                "echo Logging in to EKS",
+                f"aws ecr get-login-password --region {self.code_build_client._client_config.region_name} | docker login --username AWS --password-stdin {container_registry.config.uri}",
+            ]
+
+        # Convert the docker_build_options dictionary to a list of strings
+        docker_build_args = ""
+        for key, value in docker_build_options.items():
+            option = f"--{key}"
+            if isinstance(value, list):
+                for val in value:
+                    docker_build_args += f"{option} {val} "
+            elif value is not None and not isinstance(value, bool):
+                docker_build_args += f"{option} {value} "
+            elif value is not False:
+                docker_build_args += f"{option} "
+
+        pre_build_commands_str = "\n".join(
+            [f"            - {command}" for command in pre_build_commands]
+        )
+
+        # Generate and use a unique tag for the Docker image. This is easier
+        # than trying to parse the image digest from the Code Build logs.
+        build_id = str(uuid4())
+        # Replace the tag in the image name with the unique build ID
+        repo_name = image_name.split(":")[0]
+        alt_image_name = f"{repo_name}:{build_id}"
+
+        buildspec = f"""
+version: 0.2
+phases:
+    pre_build:
+        commands:
+{pre_build_commands_str}
+    build:
+        commands:
+            - echo Build started on `date`
+            - echo Building the Docker image...
+            - docker build -t {image_name} . {docker_build_args}
+            - echo Build completed on `date`
+    post_build:
+        commands:
+            - echo Pushing the Docker image...
+            - docker push {image_name}
+            - docker tag {image_name} {alt_image_name}
+            - docker push {alt_image_name}
+            - echo Pushed the Docker image
+artifacts:
+    files:
+        - '**/*'
+"""
+
+        if self.config.custom_env_vars:
+            environment_variables_override.update(self.config.custom_env_vars)
+
+        environment_variables_override_list = [
+            {
+                "name": key,
+                "value": value,
+                "type": "PLAINTEXT",
+            }
+            for key, value in environment_variables_override.items()
+        ]
+
+        # Override the build project with the parameters needed to run a
+        # docker-in-docker build, as covered here: https://docs.aws.amazon.com/codebuild/latest/userguide/sample-docker-section.html
+        response = self.code_build_client.start_build(
+            projectName=self.config.code_build_project,
+            environmentTypeOverride="LINUX_CONTAINER",
+            imageOverride=self.config.build_image,
+            computeTypeOverride=self.config.compute_type,
+            privilegedModeOverride=False,
+            sourceTypeOverride="S3",
+            sourceLocationOverride=f"{bucket}/{object_path}",
+            buildspecOverride=buildspec,
+            environmentVariablesOverride=environment_variables_override_list,
+            # no artifacts
+            artifactsOverride={"type": "NO_ARTIFACTS"},
+        )
+
+        build_arn = response["build"]["arn"]
+
+        # Parse the AWS region, account, codebuild project and build name from the ARN
+        aws_region, aws_account, build = build_arn.split(":", maxsplit=5)[3:6]
+        codebuild_project = build.split("/")[1].split(":")[0]
+
+        logs_url = f"https://{aws_region}.console.aws.amazon.com/codesuite/codebuild/{aws_account}/projects/{codebuild_project}/{build}/log"
+        logger.info(
+            f"Running Code Build to build the Docker image. Cloud Build logs: `{logs_url}`",
+        )
+
+        # Wait for the build to complete
+        code_build_id = response["build"]["id"]
+        while True:
+            build_status = self.code_build_client.batch_get_builds(
+                ids=[code_build_id]
+            )
+            build = build_status["builds"][0]
+            status = build["buildStatus"]
+            if status in [
+                "SUCCEEDED",
+                "FAILED",
+                "FAULT",
+                "TIMED_OUT",
+                "STOPPED",
+            ]:
+                break
+            time.sleep(10)
+
+        if status != "SUCCEEDED":
+            raise RuntimeError(
+                f"The Code Build run to build the Docker image has failed. More "
+                f"information can be found in the Cloud Build logs: {logs_url}."
+            )
+
+        logger.info(
+            f"The Docker image has been built successfully. More information can "
+            f"be found in the Cloud Build logs: `{logs_url}`."
+        )
+
+        return alt_image_name

zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -793,7 +793,7 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                     "set or set to 0. The accelerator type will be ignored. "
                     "To fix this warning, either remove the specified "
                     "accelerator type or set the `gpu_count` using the "
-                    "ResourceSettings (https://docs.zenml.io/how-to/training-with-gpus#specify-resource-requirements-for-steps)."
+                    "ResourceSettings (https://docs.zenml.io/how-to/advanced-topics/training-with-gpus)."
                 )
 
         return dynamic_component

zenml/integrations/kaniko/image_builders/kaniko_image_builder.py

@@ -25,6 +25,7 @@ from zenml.image_builders import BaseImageBuilder
 from zenml.integrations.kaniko.flavors import KanikoImageBuilderConfig
 from zenml.logger import get_logger
 from zenml.stack import StackValidator
+from zenml.utils.archivable import ArchiveType
 
 if TYPE_CHECKING:
     from zenml.container_registries import BaseContainerRegistry
@@ -295,7 +296,7 @@ class KanikoImageBuilder(BaseImageBuilder):
         logger.debug("Writing build context to process stdin.")
         assert process.stdin
         with process.stdin as _, tempfile.TemporaryFile(mode="w+b") as f:
-            build_context.write_archive(f, use_gzip=True)
+            build_context.write_archive(f, archive_type=ArchiveType.TAR_GZ)
             while True:
                 data = f.read(1024)
                 if not data:

zenml/integrations/kubernetes/flavors/kubernetes_orchestrator_flavor.py

@@ -134,6 +134,17 @@ class KubernetesOrchestratorConfig(
         """
         return True
 
+    @property
+    def supports_client_side_caching(self) -> bool:
+        """Whether the orchestrator supports client side caching.
+
+        Returns:
+            Whether the orchestrator supports client side caching.
+        """
+        # The Kubernetes orchestrator starts step pods from a pipeline pod.
+        # This is currently not supported when using client-side caching.
+        return False
+
 
 class KubernetesOrchestratorFlavor(BaseOrchestratorFlavor):
     """Kubernetes orchestrator flavor."""

zenml/integrations/kubernetes/step_operators/kubernetes_step_operator.py

@@ -33,7 +33,6 @@ from zenml.stack import Stack, StackValidator
 from zenml.step_operators import BaseStepOperator
 
 if TYPE_CHECKING:
-    from zenml.config.base_settings import BaseSettings
     from zenml.config.step_run_info import StepRunInfo
     from zenml.models import PipelineDeploymentBase
 

zenml/integrations/lightning/flavors/lightning_orchestrator_flavor.py

@@ -94,6 +94,17 @@ class LightningOrchestratorConfig(
         """
         return False
 
+    @property
+    def supports_client_side_caching(self) -> bool:
+        """Whether the orchestrator supports client side caching.
+
+        Returns:
+            Whether the orchestrator supports client side caching.
+        """
+        # The Lightning orchestrator starts step studios from a pipeline studio.
+        # This is currently not supported when using client-side caching.
+        return False
+
 
 class LightningOrchestratorFlavor(BaseOrchestratorFlavor):
     """Lightning orchestrator flavor."""

zenml/integrations/neptune/experiment_trackers/neptune_experiment_tracker.py

@@ -77,10 +77,12 @@ class NeptuneExperimentTracker(BaseExperimentTracker):
             NeptuneExperimentTrackerSettings, self.get_settings(info)
         )
 
-        self.run_state.token = self.config.api_token
-        self.run_state.project = self.config.project
-        self.run_state.run_name = info.run_name
-        self.run_state.tags = list(settings.tags)
+        self.run_state.initialize(
+            project=self.config.project,
+            token=self.config.api_token,
+            run_name=info.run_name,
+            tags=list(settings.tags),
+        )
 
     def get_step_run_metadata(
         self, info: "StepRunInfo"
@@ -107,4 +109,4 @@ class NeptuneExperimentTracker(BaseExperimentTracker):
         """
         self.run_state.active_run.sync()
         self.run_state.active_run.stop()
-        self.run_state.reset_active_run()
+        self.run_state.reset()