AWS Image Builder implementation

Author: stefannica

src/zenml/image_builders/base_image_builder.py

@@ -25,6 +25,7 @@
 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 @@ def build(
     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 @@ def _upload_build_context(
                 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,15 +122,15 @@ def _upload_build_context(
 
         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)
                 if not data:
                     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)

src/zenml/image_builders/build_context.py

@@ -20,7 +20,7 @@
 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 @@ def dockerignore_file(self) -> Optional[str]:
         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

src/zenml/integrations/aws/__init__.py

@@ -33,6 +33,7 @@
 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 @@ def flavors(cls) -> List[Type[Flavor]]:
         """
         from zenml.integrations.aws.flavors import (
             AWSContainerRegistryFlavor,
+            AWSImageBuilderFlavor,
             SagemakerOrchestratorFlavor,
             SagemakerStepOperatorFlavor,
         )
 
         return [
             AWSContainerRegistryFlavor,
+            AWSImageBuilderFlavor,
             SagemakerStepOperatorFlavor,
             SagemakerOrchestratorFlavor,
         ]

src/zenml/integrations/aws/flavors/__init__.py

@@ -17,6 +17,10 @@
     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 @@
 __all__ = [
     "AWSContainerRegistryFlavor",
     "AWSContainerRegistryConfig",
+    "AWSImageBuilderConfig",
+    "AWSImageBuilderFlavor",
     "SagemakerStepOperatorFlavor",
     "SagemakerStepOperatorConfig",
     "SagemakerOrchestratorFlavor",

src/zenml/integrations/aws/flavors/aws_image_builder_flavor.py

@@ -0,0 +1,141 @@
+#  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, 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
+from zenml.utils.secret_utils import SecretField
+
+if TYPE_CHECKING:
+    from zenml.integrations.aws.image_builders import AWSImageBuilder
+
+
+class AWSImageBuilderConfig(BaseImageBuilderConfig):
+    """AWS Code Build image builder configuration.
+
+    Attributes:
+        code_build_project: The name of the AWS CodeBuild project to use to
+            build the image.
+        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.
+        aws_secret_access_key: The AWS secret access key to use to authenticate
+            to AWS. If not provided, the value from the default AWS config will
+            be used.
+        aws_auth_role_arn: The ARN of an intermediate IAM role to assume when
+            authenticating to AWS.
+        region: The AWS region where the processing job will be run. If not
+            provided, the value from the default AWS config will be used.
+        implicit_auth: Whether to use implicit authentication to authenticate
+            the AWS Code Build build to the container registry. 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_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
+    aws_access_key_id: Optional[str] = SecretField(default=None)
+    aws_secret_access_key: Optional[str] = SecretField(default=None)
+    aws_auth_role_arn: Optional[str] = None
+    region: Optional[str] = None
+    implicit_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

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