zenml-nightly 0.70.0.dev20241203__py3-none-any.whl → 0.70.0.dev20241205__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.70.0.dev20241203
+0.70.0.dev20241205

zenml/config/server_config.py

@@ -27,6 +27,7 @@ from zenml.constants import (
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_POLLING,
     DEFAULT_ZENML_SERVER_DEVICE_AUTH_TIMEOUT,
     DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME,
+    DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_MAX_LIFETIME,
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY,
     DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE,
     DEFAULT_ZENML_SERVER_MAX_DEVICE_AUTH_ATTEMPTS,
@@ -269,6 +270,9 @@ class ServerConfiguration(BaseModel):
     generic_api_token_lifetime: PositiveInt = (
         DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME
     )
+    generic_api_token_max_lifetime: PositiveInt = (
+        DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_MAX_LIFETIME
+    )
 
     external_login_url: Optional[str] = None
     external_user_info_url: Optional[str] = None

zenml/constants.py

@@ -170,6 +170,9 @@ ENV_ZENML_DISABLE_STEP_LOGS_STORAGE = "ZENML_DISABLE_STEP_LOGS_STORAGE"
 ENV_ZENML_IGNORE_FAILURE_HOOK = "ZENML_IGNORE_FAILURE_HOOK"
 ENV_ZENML_CUSTOM_SOURCE_ROOT = "ZENML_CUSTOM_SOURCE_ROOT"
 ENV_ZENML_WHEEL_PACKAGE_NAME = "ZENML_WHEEL_PACKAGE_NAME"
+ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = (
+    "ZENML_PIPELINE_API_TOKEN_EXPIRATION"
+)
 
 # ZenML Server environment variables
 ENV_ZENML_SERVER_PREFIX = "ZENML_SERVER_"
@@ -268,6 +271,9 @@ DEFAULT_ZENML_SERVER_PIPELINE_RUN_AUTH_WINDOW = 60 * 48  # 48 hours
 DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_MINUTE = 5
 DEFAULT_ZENML_SERVER_LOGIN_RATE_LIMIT_DAY = 1000
 DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_LIFETIME = 60 * 60  # 1 hour
+DEFAULT_ZENML_SERVER_GENERIC_API_TOKEN_MAX_LIFETIME = (
+    60 * 60 * 24 * 7
+)  # 7 days
 
 DEFAULT_ZENML_SERVER_SECURE_HEADERS_HSTS = (
     "max-age=63072000; includeSubdomains"
@@ -466,3 +472,7 @@ BANNED_NAME_CHARACTERS = "\t\n\r\v\f"
 
 
 STACK_DEPLOYMENT_API_TOKEN_EXPIRATION = 60 * 6  # 6 hours
+
+ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = handle_int_env_var(
+    ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION, default=0
+)

zenml/image_builders/base_image_builder.py

@@ -25,6 +25,7 @@ from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.stack import Flavor, StackComponent
 from zenml.stack.stack_component import StackComponentConfig
+from zenml.utils.archivable import ArchiveType
 
 if TYPE_CHECKING:
     from zenml.container_registries import BaseContainerRegistry
@@ -100,6 +101,7 @@ class BaseImageBuilder(StackComponent, ABC):
     def _upload_build_context(
         build_context: "BuildContext",
         parent_path_directory_name: str,
+        archive_type: ArchiveType = ArchiveType.TAR_GZ,
     ) -> str:
         """Uploads a Docker image build context to a remote location.
 
@@ -109,6 +111,7 @@ class BaseImageBuilder(StackComponent, ABC):
                 the build context to. It will be appended to the artifact
                 store path to create the parent path where the build context
                 will be uploaded to.
+            archive_type: The type of archive to create.
 
         Returns:
             The path to the uploaded build context.
@@ -119,7 +122,7 @@ class BaseImageBuilder(StackComponent, ABC):
 
         hash_ = hashlib.sha1()  # nosec
         with tempfile.NamedTemporaryFile(mode="w+b", delete=False) as f:
-            build_context.write_archive(f, use_gzip=True)
+            build_context.write_archive(f, archive_type)
 
             while True:
                 data = f.read(64 * 1024)
@@ -127,7 +130,7 @@ class BaseImageBuilder(StackComponent, ABC):
                     break
                 hash_.update(data)
 
-            filename = f"{hash_.hexdigest()}.tar.gz"
+            filename = f"{hash_.hexdigest()}.{archive_type.value}"
             filepath = f"{parent_path}/{filename}"
             if not fileio.exists(filepath):
                 logger.info("Uploading build context to `%s`.", filepath)

zenml/image_builders/build_context.py

@@ -20,7 +20,7 @@ from zenml.constants import REPOSITORY_DIRECTORY_NAME
 from zenml.io import fileio
 from zenml.logger import get_logger
 from zenml.utils import io_utils, string_utils
-from zenml.utils.archivable import Archivable
+from zenml.utils.archivable import Archivable, ArchiveType
 
 logger = get_logger(__name__)
 
@@ -69,28 +69,19 @@ class BuildContext(Archivable):
         return None
 
     def write_archive(
-        self, output_file: IO[bytes], use_gzip: bool = True
+        self,
+        output_file: IO[bytes],
+        archive_type: ArchiveType = ArchiveType.TAR_GZ,
     ) -> None:
         """Writes an archive of the build context to the given file.
 
         Args:
             output_file: The file to write the archive to.
-            use_gzip: Whether to use `gzip` to compress the file.
+            archive_type: The type of archive to create.
         """
-        from docker.utils import build as docker_build_utils
-
-        files = self.get_files()
-        extra_files = self.get_extra_files()
-
-        context_archive = docker_build_utils.create_archive(
-            fileobj=output_file,
-            root=self._root,
-            files=sorted(files.keys()),
-            gzip=use_gzip,
-            extra_files=list(extra_files.items()),
-        )
+        super().write_archive(output_file, archive_type)
 
-        build_context_size = os.path.getsize(context_archive.name)
+        build_context_size = os.path.getsize(output_file.name)
         if (
             self._root
             and build_context_size > 50 * 1024 * 1024

zenml/integrations/aws/__init__.py

@@ -33,6 +33,7 @@ AWS_SAGEMAKER_ORCHESTRATOR_FLAVOR = "sagemaker"
 AWS_CONNECTOR_TYPE = "aws"
 AWS_RESOURCE_TYPE = "aws-generic"
 S3_RESOURCE_TYPE = "s3-bucket"
+AWS_IMAGE_BUILDER_FLAVOR = "aws"
 
 class AWSIntegration(Integration):
     """Definition of AWS integration for ZenML."""
@@ -59,12 +60,14 @@ class AWSIntegration(Integration):
         """
         from zenml.integrations.aws.flavors import (
             AWSContainerRegistryFlavor,
+            AWSImageBuilderFlavor,
             SagemakerOrchestratorFlavor,
             SagemakerStepOperatorFlavor,
         )
 
         return [
             AWSContainerRegistryFlavor,
+            AWSImageBuilderFlavor,
             SagemakerStepOperatorFlavor,
             SagemakerOrchestratorFlavor,
         ]

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_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/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: