Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

AWS Image Builder implementation #2904

Merged
merged 17 commits into from
Dec 3, 2024
Merged
Show file tree
Hide file tree
Changes from 10 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
235 changes: 235 additions & 0 deletions docs/book/component-guide/image-builders/aws.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,235 @@
---
bcdurak marked this conversation as resolved.
Show resolved Hide resolved
description: Building container images with AWS CodeBuild
---

# AWS Image Builder

The AWS image builder is an [image builder](./image-builders.md) flavor provided by the ZenML `aws` integration that uses [AWS CodeBuild](https://aws.amazon.com/codebuild) to build container images.

### When to use it

You should use the AWS image builder if:

* you're **unable** to install or use [Docker](https://www.docker.com) on your client machine.
* you're already using AWS.
* your stack is mainly composed of other AWS components such as the [S3 Artifact Store](../artifact-stores/s3.md) or the [Sagemaker Orchestrator](../orchestrators/sagemaker.md).
stefannica marked this conversation as resolved.
Show resolved Hide resolved
stefannica marked this conversation as resolved.
Show resolved Hide resolved

### How to deploy it

{% hint style="info" %}
Would you like to skip ahead and deploy a full ZenML cloud stack already,
including the AWS image builder? Check out the
[in-browser stack deployment wizard](../../how-to/infrastructure-deployment/stack-deployment/deploy-a-cloud-stack.md),
the [stack registration wizard](../../how-to/infrastructure-deployment/stack-deployment/register-a-cloud-stack.md),
or [the ZenML AWS Terraform module](../../how-to/infrastructure-deployment/stack-deployment/deploy-a-cloud-stack-with-terraform.md)
for a shortcut on how to deploy & register this stack component.
{% endhint %}
bcdurak marked this conversation as resolved.
Show resolved Hide resolved

### How to use it

To use the AWS image builder, you need:

* The ZenML `aws` integration installed. If you haven't done so, run:

```shell
zenml integration install aws
```
* An [S3 Artifact Store](../artifact-stores/s3.md) where the build context will be uploaded, so AWS CodeBuild can access it.
* Recommended: an [AWS container registry](../container-registries/aws.md) where the built image will be pushed. The AWS CodeBuild service can also work with other container registries, but [explicit authentication](#authentication-methods) must be enabled in this case.
* An [AWS CodeBuild project](https://aws.amazon.com/codebuild) created in the AWS account and region where you want to build the Docker images, preferably in the same region as the ECR container registry where images will be pushed (if applicable). The CodeBuild project configuration is largely irrelevant, as ZenML will override most of the default settings for each build according to the [AWS Docker build guide](https://docs.aws.amazon.com/codebuild/latest/userguide/sample-docker-section.html). Some example default configuration values are:
* **Source Type**: `Amazon S3`
* **Bucket**: The same S3 bucket used by the ZenML S3 Artifact Store.
* **S3 folder**: any value (e.g. `codebuild`);
* **Environment Type**: `Linux Container`
* **Environment Image**: `bentolor/docker-dind-awscli`
* **Privileged Mode**: `false`

The user must take care that the **Service Role** attached to the CodeBuild project also has the necessary permissions to access the S3 bucket to read objects and the ECR container registry to push images (if applicable):

```json
{
"Version": "2012-10-17",
"Statement": [
{
"Effect": "Allow",
"Action": [
"s3:GetObject",
"s3:GetObjectVersion"
],
"Resource": "arn:aws:s3:::<BUCKET_NAME>/*"
},
{
"Effect": "Allow",
"Action": [
"ecr:BatchGetImage",
"ecr:DescribeImages",
"ecr:BatchCheckLayerAvailability",
"ecr:GetDownloadUrlForLayer",
"ecr:InitiateLayerUpload",
"ecr:UploadLayerPart",
"ecr:CompleteLayerUpload",
"ecr:PutImage"
],
"Resource": "arn:aws:ecr:<REGION>:<ACCOUNT_ID>:repository/<REPOSITORY_NAME>"
},
{
"Effect": "Allow",
"Action": [
"ecr:GetAuthorizationToken"
],
"Resource": "*"
},
]
}
```

* Recommended: grant ZenML access to trigger AWS CodeBuild builds by registering an [AWS Service Connector](../../how-to/infrastructure-deployment/auth-management/aws-service-connector.md) with the proper credentials and permissions, as covered in the [Authentication Methods](aws.md#authentication-methods) section. If not provided, then the AWS credentials will be inferred from the environment where the pipeline is triggered.
stefannica marked this conversation as resolved.
Show resolved Hide resolved

We can register the image builder and use it in our active stack:

```shell
zenml image-builder register <IMAGE_BUILDER_NAME> \
--flavor=aws \
--code_build_project=<CODEBUILD_PROJECT_NAME>

# Register and activate a stack with the new image builder
zenml stack register <STACK_NAME> -i <IMAGE_BUILDER_NAME> ... --set
```

You also need to set up [authentication](aws.md#authentication-methods) required to access the CodeBuild AWS service.

#### Authentication Methods

Integrating and using an AWS Image Builder in your pipelines is not possible without employing some form of authentication. If you're looking for a quick way to get started locally, you can use the _Local Authentication_ method. However, the recommended way to authenticate to the AWS cloud platform is through [an AWS Service Connector](../../how-to/infrastructure-deployment/auth-management/aws-service-connector.md). This is particularly useful if you are configuring ZenML stacks that combine the AWS Image Builder with other remote stack components also running in AWS.

{% tabs %}
{% tab title="Implicit Authentication" %}
This method uses the implicit AWS authentication available _in the environment where the ZenML code is running_. On your local machine, this is the quickest way to configure an AWS Image Builder. You don't need to supply credentials explicitly when you register the AWS Image Builder, as it leverages the local credentials and configuration that the AWS CLI stores on your local machine. However, you will need to install and set up the AWS CLI on your machine as a prerequisite, as covered in [the AWS CLI documentation](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html), before you register the AWS Image Builder.

{% hint style="warning" %}
Stacks using the AWS Image Builder set up with local authentication are not portable across environments. To make ZenML pipelines fully portable, it is recommended to use [an AWS Service Connector](../../how-to/infrastructure-deployment/auth-management/aws-service-connector.md) to authenticate your AWS Image Builder to the AWS cloud platform.
{% endhint %}
{% endtab %}

{% tab title="AWS Service Connector (recommended)" %}
To set up the AWS Image Builder to authenticate to AWS and access the AWS CodeBuild services, it is recommended to leverage the many features provided by [the AWS Service Connector](../../how-to/infrastructure-deployment/auth-management/aws-service-connector.md) such as auto-configuration, best security practices regarding long-lived credentials and reusing the same credentials across multiple stack components.

If you don't already have an AWS Service Connector configured in your ZenML deployment, you can register one using the interactive CLI command. You also have the option to configure an AWS Service Connector that can be used to access more than just the AWS CodeBuild service:

```sh
zenml service-connector register --type aws -i
```

A non-interactive CLI example that leverages [the AWS CLI configuration](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html) on your local machine to auto-configure an AWS Service Connector for the AWS CodeBuild service:

```sh
zenml service-connector register <CONNECTOR_NAME> --type aws --resource-type aws-generic --auto-configure
```

{% code title="Example Command Output" %}
```
$ zenml service-connector register aws-generic --type aws --resource-type aws-generic --auto-configure
Successfully registered service connector `aws-generic` with access to the following resources:
┏━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━┓
┃ RESOURCE TYPE │ RESOURCE NAMES ┃
┠────────────────┼────────────────┨
┃ 🔶 aws-generic │ eu-central-1 ┃
┗━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━┛
```
{% endcode %}

> **Note**: Please remember to grant the entity associated with your AWS credentials permissions to access the CodeBuild API and to run CodeBuilder builds:
>
> ```json
> {
> "Version": "2012-10-17",
> "Statement": [
> {
> "Effect": "Allow",
> "Action": [
> "codebuild:StartBuild",
> "codebuild:BatchGetBuilds",
> ],
> "Resource": "arn:aws:codebuild:<REGION>:<ACCOUNT_ID>:project/<CODEBUILD_PROJECT_NAME>"
> },
> ]
> }
> ```
>

The AWS Service Connector supports [many different authentication methods](../../how-to/infrastructure-deployment/auth-management/aws-service-connector.md#authentication-methods) with different levels of security and convenience. You should pick the one that best fits your use case.

If you already have one or more AWS Service Connectors configured in your ZenML deployment, you can check which of them can be used to access generic AWS resources like the one required for your AWS Image Builder by running e.g.:

```sh
zenml service-connector list-resources --resource-type aws-generic
```

{% code title="Example Command Output" %}
```
The following 'aws-generic' resources can be accessed by service connectors configured in your workspace:
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━┓
┃ CONNECTOR ID │ CONNECTOR NAME │ CONNECTOR TYPE │ RESOURCE TYPE │ RESOURCE NAMES ┃
┠──────────────────────────────────────┼────────────────┼────────────────┼────────────────┼────────────────┨
┃ 7113ba9b-efdd-4a0a-94dc-fb67926e58a1 │ aws-generic │ 🔶 aws │ 🔶 aws-generic │ eu-central-1 ┃
┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━┛
```
{% endcode %}

After having set up or decided on an AWS Service Connector to use to authenticate to AWS, you can register the AWS Image Builder as follows:

```sh
zenml image-builder register <IMAGE_BUILDER_NAME> \
--flavor=aws \
--code_build_project=<CODEBUILD_PROJECT_NAME> \
--connector <CONNECTOR_ID>
```

To connect an AWS Image Builder to an AWS Service Connector at a later point, you can use the following command:

```sh
zenml image-builder connect <IMAGE_BUILDER_NAME> --connector <CONNECTOR_ID>
```

{% code title="Example Command Output" %}
```
$ zenml image-builder connect aws-image-builder --connector aws-generic
Successfully connected image builder `aws-image-builder` to the following resources:
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━┯━━━━━━━━━━━━━━━━┓
┃ CONNECTOR ID │ CONNECTOR NAME │ CONNECTOR TYPE │ RESOURCE TYPE │ RESOURCE NAMES ┃
┠──────────────────────────────────────┼────────────────┼────────────────┼────────────────┼────────────────┨
┃ 7113ba9b-efdd-4a0a-94dc-fb67926e58a1 │ aws-generic │ 🔶 aws │ 🔶 aws-generic │ eu-central-1 ┃
┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━┷━━━━━━━━━━━━━━━━┛
```
{% endcode %}

As a final step, you can use the AWS Image Builder in a ZenML Stack:

```sh
# Register and set a stack with the new image builder
zenml stack register <STACK_NAME> -i <IMAGE_BUILDER_NAME> ... --set
```
{% endtab %}
{% endtabs %}

#### Customizing AWS CodeBuild builds

The AWS Image Builder can be customized to a certain extent by providing additional configuration options when registering the image builder. The following additional attributes can be set:

* `build_image`: The Docker image used to build the Docker image. The default is `bentolor/docker-dind-awscli`,
which is a Docker image that includes both Docker-in-Docker and the AWS CLI.


{% hint style="info" %}
If you are running into Dockerhub rate-limits, it might be a good idea to copy this image to your own container registry and customize the `build_image` attribute to point to your own image.
stefannica marked this conversation as resolved.
Show resolved Hide resolved
{% endhint %}

* `compute_type`: The compute type used for the CodeBuild project. The default is `BUILD_GENERAL1_SMALL`.
* `custom_env_vars`: A dictionary of custom environment variables to be set in the CodeBuild project.
* `implicit_container_registry_auth`: A boolean flag that indicates whether to use implicit or explicit authentication when authenticating the AWS CodeBuild build to the target container registry:

* when this is set to `true` (default), the builds will be configured to use whatever implicit authentication credentials are already available within the build container. As a special case for ECR registries, the service IAM role attached to the CodeBuild project is used to authenticate to the target ECR container registry and therefore the service role must include the necessary permissions to push images to the target ECR registry.
* when set to `false`, the credentials attached to the ZenML Container Registry stack component in the active stack will be set as build environment variables and used to authenticate to the target container registry. This is useful when the target container registry is not an ECR registry or when the service role attached to the CodeBuild project does not have the necessary permissions to push images to the target ECR registry. This works best when the ZenML Container Registry stack component is also linked to the external container registry via a Service Connector.


<figure><img src="https://static.scarf.sh/a.png?x-pxid=f0b4f458-0a54-4fcd-aa95-d5ee424815bc" alt="ZenML Scarf"><figcaption></figcaption></figure>
1 change: 1 addition & 0 deletions docs/book/component-guide/image-builders/image-builders.md
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ image builders are provided by integrations:
| [LocalImageBuilder](local.md) | `local` | _built-in_ | Builds your Docker images locally. |
| [KanikoImageBuilder](kaniko.md) | `kaniko` | `kaniko` | Builds your Docker images in Kubernetes using Kaniko. |
| [GCPImageBuilder](gcp.md) | `gcp` | `gcp` | Builds your Docker images using Google Cloud Build. |
| [AWSImageBuilder](aws.md) | `aws` | `aws` | Builds your Docker images using AWS Code Build. |
| [Custom Implementation](custom.md) | _custom_ | | Extend the image builder abstraction and provide your own implementation |

If you would like to see the available flavors of image builders, you can use the command:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -249,7 +249,8 @@ following components:
* a local Orchestrator, if `orchestrator` is set to `local`. This can be used in combination with the SageMaker Step Operator to selectively run some steps locally and some on SageMaker.
* if `orchestrator` is set to `sagemaker` (default): a SageMaker Orchestrator linked to the AWS account via an AWS Service Connector configured with IAM role credentials
* if `orchestrator` is set to `skypilot`: a SkyPilot Orchestrator linked to the AWS account via an AWS Service Connector configured with IAM role credentials
4. a SageMaker Step Operator linked to the AWS account via an AWS Service Connector configured with IAM role credentials
4. an AWS CodeBuild Image Builder linked to the AWS account via an AWS Service Connector configured with IAM role credentials
5. a SageMaker Step Operator linked to the AWS account via an AWS Service Connector configured with IAM role credentials

To use the ZenML stack, you will need to install the required integrations:

Expand Down
2 changes: 1 addition & 1 deletion examples/e2e/.copier-answers.yml
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
# Changes here will be overwritten by Copier
_commit: 2024.11.20-2-g760142f
_commit: 2024.11.28
_src_path: gh:zenml-io/template-e2e-batch
data_quality_checks: true
email: [email protected]
Expand Down
2 changes: 1 addition & 1 deletion examples/e2e_nlp/.copier-answers.yml
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
# Changes here will be overwritten by Copier
_commit: 2024.10.30-2-g1ae14e3
_commit: 2024.11.28
_src_path: gh:zenml-io/template-nlp
accelerator: cpu
cloud_of_choice: aws
Expand Down
2 changes: 1 addition & 1 deletion examples/llm_finetuning/.copier-answers.yml
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
# Changes here will be overwritten by Copier
_commit: 2024.11.08-2-gece1d46
_commit: 2024.11.28
_src_path: gh:zenml-io/template-llm-finetuning
bf16: true
cuda_version: cuda11.8
Expand Down
2 changes: 1 addition & 1 deletion examples/mlops_starter/.copier-answers.yml
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
# Changes here will be overwritten by Copier
_commit: 2024.10.30-7-gb60e441
_commit: 2024.11.28
_src_path: gh:zenml-io/template-starter
email: [email protected]
full_name: ZenML GmbH
Expand Down
7 changes: 5 additions & 2 deletions src/zenml/image_builders/base_image_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand All @@ -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.
Expand All @@ -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)
Expand Down
Loading
Loading