zenml-nightly 0.70.0.dev20241202__py3-none-any.whl → 0.70.0.dev20241204__py3-none-any.whl

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/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()

zenml/integrations/neptune/experiment_trackers/run_state.py

@@ -20,7 +20,6 @@ import neptune
 
 import zenml
 from zenml.client import Client
-from zenml.integrations.constants import NEPTUNE
 from zenml.utils.singleton import SingletonMetaClass
 
 if TYPE_CHECKING:
@@ -29,20 +28,38 @@ if TYPE_CHECKING:
 _INTEGRATION_VERSION_KEY = "source_code/integrations/zenml"
 
 
-class InvalidExperimentTrackerSelected(Exception):
-    """Raised if a Neptune run is fetched while using a different experiment tracker."""
-
-
 class RunProvider(metaclass=SingletonMetaClass):
     """Singleton object used to store and persist a Neptune run state across the pipeline."""
 
     def __init__(self) -> None:
         """Initialize RunProvider. Called with no arguments."""
         self._active_run: Optional["Run"] = None
-        self._project: Optional[str]
-        self._run_name: Optional[str]
-        self._token: Optional[str]
-        self._tags: Optional[List[str]]
+        self._project: Optional[str] = None
+        self._run_name: Optional[str] = None
+        self._token: Optional[str] = None
+        self._tags: Optional[List[str]] = None
+        self._initialized = False
+
+    def initialize(
+        self,
+        project: Optional[str] = None,
+        token: Optional[str] = None,
+        run_name: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+    ) -> None:
+        """Initialize the run state.
+
+        Args:
+            project: The neptune project.
+            token: The neptune token.
+            run_name: The neptune run name.
+            tags: Tags for the neptune run.
+        """
+        self._project = project
+        self._token = token
+        self._run_name = run_name
+        self._tags = tags
+        self._initialized = True
 
     @property
     def project(self) -> Optional[Any]:
@@ -53,15 +70,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._project
 
-    @project.setter
-    def project(self, project: str) -> None:
-        """Setter for project name.
-
-        Args:
-            project: Neptune project name
-        """
-        self._project = project
-
     @property
     def token(self) -> Optional[Any]:
         """Getter for API token.
@@ -71,15 +79,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._token
 
-    @token.setter
-    def token(self, token: str) -> None:
-        """Setter for API token.
-
-        Args:
-            token: Neptune API token
-        """
-        self._token = token
-
     @property
     def run_name(self) -> Optional[Any]:
         """Getter for run name.
@@ -89,15 +88,6 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._run_name
 
-    @run_name.setter
-    def run_name(self, run_name: str) -> None:
-        """Setter for run name.
-
-        Args:
-            run_name: name of the pipeline run
-        """
-        self._run_name = run_name
-
     @property
     def tags(self) -> Optional[Any]:
         """Getter for run tags.
@@ -107,14 +97,14 @@ class RunProvider(metaclass=SingletonMetaClass):
         """
         return self._tags
 
-    @tags.setter
-    def tags(self, tags: List[str]) -> None:
-        """Setter for run tags.
+    @property
+    def initialized(self) -> bool:
+        """If the run state is initialized.
 
-        Args:
-            tags: list of tags associated with a Neptune run
+        Returns:
+            If the run state is initialized.
         """
-        self._tags = tags
+        return self._initialized
 
     @property
     def active_run(self) -> "Run":
@@ -137,9 +127,14 @@ class RunProvider(metaclass=SingletonMetaClass):
             self._active_run = run
         return self._active_run
 
-    def reset_active_run(self) -> None:
-        """Resets the active run state to None."""
+    def reset(self) -> None:
+        """Reset the run state."""
         self._active_run = None
+        self._project = None
+        self._run_name = None
+        self._token = None
+        self._tags = None
+        self._initialized = False
 
 
 def get_neptune_run() -> "Run":
@@ -149,14 +144,35 @@ def get_neptune_run() -> "Run":
         Neptune run object
 
     Raises:
-        InvalidExperimentTrackerSelected: when called while using an experiment tracker other than Neptune
+        RuntimeError: When unable to fetch the active neptune run.
     """
-    client = Client()
-    experiment_tracker = client.active_stack.experiment_tracker
-    if experiment_tracker.flavor == NEPTUNE:  # type: ignore
-        return experiment_tracker.run_state.active_run  # type: ignore
-    raise InvalidExperimentTrackerSelected(
-        "Fetching a Neptune run works only with the 'neptune' flavor of "
-        "the experiment tracker. The flavor currently selected is %s"
-        % experiment_tracker.flavor  # type: ignore
+    from zenml.integrations.neptune.experiment_trackers import (
+        NeptuneExperimentTracker,
     )
+
+    experiment_tracker = Client().active_stack.experiment_tracker
+
+    if not experiment_tracker:
+        raise RuntimeError(
+            "Unable to get neptune run: Missing experiment tracker in the "
+            "active stack."
+        )
+
+    if not isinstance(experiment_tracker, NeptuneExperimentTracker):
+        raise RuntimeError(
+            "Unable to get neptune run: Experiment tracker in the active "
+            f"stack ({experiment_tracker.flavor}) is not a neptune experiment "
+            "tracker."
+        )
+
+    run_state = experiment_tracker.run_state
+    if not run_state.initialized:
+        raise RuntimeError(
+            "Unable to get neptune run: The experiment tracker has not been "
+            "initialized. To solve this, make sure you use the experiment "
+            "tracker in your step. See "
+            "https://docs.zenml.io/stack-components/experiment-trackers/neptune#how-do-you-use-it "
+            "for more information."
+        )
+
+    return experiment_tracker.run_state.active_run

zenml/integrations/registry.py

@@ -111,7 +111,7 @@ class IntegrationRegistry(object):
                 )
             else:
                 raise KeyError(
-                    f"Version {integration_name} does not exist. "
+                    f"Integration {integration_name} does not exist. "
                     f"Currently the following integrations are implemented. "
                     f"{self.list_integration_names}"
                 )
@@ -148,7 +148,7 @@ class IntegrationRegistry(object):
                 ].get_uninstall_requirements(target_os=target_os)
             else:
                 raise KeyError(
-                    f"Version {integration_name} does not exist. "
+                    f"Integration {integration_name} does not exist. "
                     f"Currently the following integrations are implemented. "
                     f"{self.list_integration_names}"
                 )

zenml/integrations/skypilot/flavors/skypilot_orchestrator_base_vm_config.py

@@ -144,3 +144,15 @@ class SkypilotBaseOrchestratorConfig(
             True if this config is for a local component, False otherwise.
         """
         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 Skypilot orchestrator runs the entire pipeline in a single VM, or
+        # starts additional VMs from the root VM. Both of those cases are
+        # currently not supported when using client-side caching.
+        return False

zenml/materializers/built_in_materializer.py

@@ -429,7 +429,7 @@ class BuiltInContainerMaterializer(BaseMaterializer):
         # doesn't work for non-serializable types as they
         # are saved as list of lists in different files
         if _is_serializable(data):
-            return {self.data_path: VisualizationType.JSON}
+            return {self.data_path.replace("\\", "/"): VisualizationType.JSON}
         return {}
 
     def extract_metadata(self, data: Any) -> Dict[str, "MetadataType"]:

zenml/orchestrators/base_orchestrator.py

@@ -61,7 +61,7 @@ class BaseOrchestratorConfig(StackComponentConfig):
                     "The 'custom_docker_base_image_name' field has been "
                     "deprecated. To use a custom base container image with your "
                     "orchestrators, please use the DockerSettings in your "
-                    "pipeline (see https://docs.zenml.io/how-to/customize-docker-builds)."
+                    "pipeline (see https://docs.zenml.io/how-to/infrastructure-deployment/customize-docker-builds)."
                 )
 
         return data
@@ -84,6 +84,15 @@ class BaseOrchestratorConfig(StackComponentConfig):
         """
         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.
+        """
+        return True
+
 
 class BaseOrchestrator(StackComponent, ABC):
     """Base class for all orchestrators.
@@ -205,6 +214,7 @@ class BaseOrchestrator(StackComponent, ABC):
 
         if (
             placeholder_run
+            and self.config.supports_client_side_caching
             and not deployment.schedule
             and not prevent_client_side_caching
         ):
@@ -232,6 +242,8 @@ class BaseOrchestrator(StackComponent, ABC):
                 self._cleanup_run()
                 logger.info("All steps of the pipeline run were cached.")
                 return
+        else:
+            logger.debug("Skipping client-side caching.")
 
         try:
             if metadata_iterator := self.prepare_or_run_pipeline(