# Source: https://docs.zenml.io/stacks/popular-stacks/aws-guide.md

# AWS

This page aims to quickly set up a minimal production stack on AWS. With just a few simple steps, you will set up an IAM role with specifically-scoped permissions that ZenML can use to authenticate with the relevant AWS resources.

{% hint style="info" %}
Would you like to skip ahead and deploy a full AWS ZenML cloud stack already?

Check out the [in-browser stack deployment wizard](https://docs.zenml.io/how-to/infrastructure-deployment/stack-deployment/deploy-a-cloud-stack),\
the [stack registration wizard](https://docs.zenml.io/how-to/infrastructure-deployment/stack-deployment/register-a-cloud-stack),\
or [the ZenML AWS Terraform module](https://docs.zenml.io/how-to/infrastructure-deployment/stack-deployment/deploy-a-cloud-stack-with-terraform)\
for a shortcut on how to deploy & register this stack.
{% endhint %}

## 1) Set up credentials and local environment

To follow this guide, you need:

* An active AWS account with necessary permissions for AWS S3, SageMaker, ECR, and ECS.
* ZenML [installed](https://docs.zenml.io/getting-started/installation)
* AWS CLI installed and configured with your AWS credentials. You can follow the instructions [here](https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html).

Once ready, navigate to the AWS console:

1. Choose an AWS region: In the AWS console, choose the region where you want to deploy your ZenML stack resources. Make note of the region name (e.g., `us-east-1`, `eu-west-2`, etc.) as you will need it in subsequent steps.
2. Create an IAM role:

For this, you'll need to find out your AWS account ID. You can find this by running:

```shell
aws sts get-caller-identity --query Account --output text
```

This will output your AWS account ID. Make a note of this as you will need it in the next steps. (If you're doing anything more esoteric with your AWS account and IAM roles, this might not work for you. The account ID here that we're trying to get is the root account ID that you use to log in to the AWS console.)

Then create a file named `assume-role-policy.json` with the following content:

```json
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Principal": {
        "AWS": "arn:aws:iam::<YOUR_ACCOUNT_ID>:root",
        "Service": "sagemaker.amazonaws.com"
      },
      "Action": "sts:AssumeRole"
    }
  ]
}
```

Make sure to replace the placeholder `<YOUR_ACCOUNT_ID>` with your actual AWS account ID that we found earlier.

Now create a new IAM role that ZenML will use to access AWS resources. We'll use `zenml-role` as a role name in this example, but you can feel free to choose something else if you prefer. Run the following command to create the role:

```shell
aws iam create-role --role-name zenml-role --assume-role-policy-document file://assume-role-policy.json
```

Be sure to take note of the information that is output to the terminal, as you will need it in the next steps, especially the Role ARN.

3. Create and attach least-privilege policies to the role:

Instead of using broad managed policies, create custom policies that follow the principle of least privilege. First, create the necessary policy documents:

**Create S3 policy document (`s3-policy.json`):**

```json
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "s3:ListBucket",
        "s3:GetObject",
        "s3:PutObject",
        "s3:DeleteObject",
        "s3:GetBucketVersioning",
        "s3:ListBucketVersions",
        "s3:DeleteObjectVersion"
      ],
      "Resource": [
        "arn:aws:s3:::your-bucket-name",
        "arn:aws:s3:::your-bucket-name/*"
      ]
    },
    {
      "Effect": "Allow",
      "Action": "s3:ListAllMyBuckets",
      "Resource": "*"
    }
  ]
}
```

**Create ECR policy document (`ecr-policy.json`):**

```json
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "ecr:BatchGetImage",
        "ecr:BatchCheckLayerAvailability",
        "ecr:GetDownloadUrlForLayer",
        "ecr:GetAuthorizationToken",
        "ecr:InitiateLayerUpload",
        "ecr:UploadLayerPart",
        "ecr:CompleteLayerUpload",
        "ecr:PutImage",
        "ecr:DescribeRepositories",
        "ecr:ListRepositories",
        "ecr:DescribeImages"
      ],
      "Resource": "*"
    }
  ]
}
```

**Create SageMaker policy document (`sagemaker-policy.json`):**

```json
{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "sagemaker:CreatePipeline",
        "sagemaker:StartPipelineExecution",
        "sagemaker:StopPipelineExecution",
        "sagemaker:DescribePipeline",
        "sagemaker:DescribePipelineExecution",
        "sagemaker:ListPipelineExecutions",
        "sagemaker:ListPipelineExecutionSteps",
        "sagemaker:UpdatePipeline",
        "sagemaker:DeletePipeline",
        "sagemaker:CreateProcessingJob",
        "sagemaker:DescribeProcessingJob",
        "sagemaker:StopProcessingJob",
        "sagemaker:CreateTrainingJob",
        "sagemaker:DescribeTrainingJob",
        "sagemaker:StopTrainingJob"
      ],
      "Resource": "*"
    },
    {
      "Effect": "Allow",
      "Action": "iam:PassRole",
      "Resource": "arn:aws:iam::<YOUR_ACCOUNT_ID>:role/zenml-role",
      "Condition": {
        "StringEquals": {
          "iam:PassedToService": "sagemaker.amazonaws.com"
        }
      }
    }
  ]
}
```

Replace `<YOUR_ACCOUNT_ID>` and `your-bucket-name` with your actual values, then create and attach the policies:

```shell
# Create the custom policies
aws iam create-policy --policy-name ZenML-S3-Policy --policy-document file://s3-policy.json
aws iam create-policy --policy-name ZenML-ECR-Policy --policy-document file://ecr-policy.json
aws iam create-policy --policy-name ZenML-SageMaker-Policy --policy-document file://sagemaker-policy.json

# Attach the custom policies to the role
aws iam attach-role-policy --role-name zenml-role --policy-arn arn:aws:iam::<YOUR_ACCOUNT_ID>:policy/ZenML-S3-Policy
aws iam attach-role-policy --role-name zenml-role --policy-arn arn:aws:iam::<YOUR_ACCOUNT_ID>:policy/ZenML-ECR-Policy
aws iam attach-role-policy --role-name zenml-role --policy-arn arn:aws:iam::<YOUR_ACCOUNT_ID>:policy/ZenML-SageMaker-Policy
```

4. If you have not already, install the AWS and S3 ZenML integrations:

```shell
zenml integration install aws s3 -y
```

## 2) Create a Service Connector within ZenML

Create an AWS Service Connector within ZenML. The service connector will allow ZenML and other ZenML components to authenticate themselves with AWS using the IAM role.

{% tabs %}
{% tab title="CLI" %}

```shell
zenml service-connector register aws_connector \
  --type aws \
  --auth-method iam-role \
  --role_arn=<ROLE_ARN> \
  --region=<YOUR_REGION> \
  --aws_access_key_id=<YOUR_ACCESS_KEY_ID> \
  --aws_secret_access_key=<YOUR_SECRET_ACCESS_KEY>
```

Replace `<ROLE_ARN>` with the ARN of the IAM role you created in the previous step, `<YOUR_REGION>` with the respective value and use your AWS access key ID and secret access key that we noted down earlier.
{% endtab %}
{% endtabs %}

## 3) Create Stack Components

### Artifact Store (S3)

An [artifact store](https://docs.zenml.io/user-guides/production-guide/remote-storage) is used for storing and versioning data flowing through your pipelines.

1. Before you run anything within the ZenML CLI, create an AWS S3 bucket. If you already have one, you can skip this step. (Note: the bucket name should be unique, so you might need to try a few times to find a unique name.)

```shell
aws s3api create-bucket --bucket your-bucket-name
```

Once this is done, you can create the ZenML stack component as follows:

2. Register an S3 Artifact Store with the connector:

```shell
zenml artifact-store register cloud_artifact_store -f s3 --path=s3://bucket-name --connector aws_connector
```

More details [here](https://docs.zenml.io/stacks/artifact-stores/s3).

### Orchestrator (SageMaker Pipelines)

An [orchestrator](https://docs.zenml.io/user-guides/production-guide/cloud-orchestration) is the compute backend to run your pipelines.

1. Before you run anything within the ZenML CLI, head on over to AWS and create a SageMaker domain (Skip this if you already have one). The instructions for creating a domain can be found [in the AWS core documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/onboard-quick-start.html).

A SageMaker domain is a central management unit for all SageMaker users and resources within a region. It provides a single sign-on (SSO) experience and enables users to create and manage SageMaker resources, such as notebooks, training jobs, and endpoints, within a collaborative environment.

When you create a SageMaker domain, you specify the configuration settings, such as the domain name, user profiles, and security settings. Each user within a domain gets their own isolated workspace, which includes a JupyterLab interface, a set of compute resources, and persistent storage.

The SageMaker orchestrator in ZenML requires a SageMaker domain to run pipelines because it leverages the SageMaker Pipelines service, which is part of the SageMaker ecosystem. SageMaker Pipelines allows you to define, execute, and manage end-to-end machine learning workflows using a declarative approach.

By creating a SageMaker domain, you establish the necessary environment and permissions for the SageMaker orchestrator to interact with SageMaker Pipelines and other SageMaker resources seamlessly. The domain acts as a prerequisite for using the SageMaker orchestrator in ZenML.

Once this is done, you can create the ZenML stack component as follows:

2. Register a SageMaker Pipelines orchestrator stack component:

You'll need the IAM role ARN that we noted down earlier to register the orchestrator. This is the 'execution role' ARN you need to pass to the orchestrator.

```shell
zenml orchestrator register sagemaker-orchestrator --flavor=sagemaker --region=<YOUR_REGION> --execution_role=<ROLE_ARN>
```

**Note**: The SageMaker orchestrator utilizes the AWS configuration for operation and does not require direct connection via a service connector for authentication, as it relies on your AWS CLI configurations or environment variables.

More details [here](https://docs.zenml.io/stacks/orchestrators/sagemaker).

### Container Registry (ECR)

A [container registry](https://docs.zenml.io/stacks/container-registries) is used to store Docker images for your pipelines.

1. You'll need to create a repository in ECR. If you already have one, you can skip this step.

```shell
aws ecr create-repository --repository-name zenml --region <YOUR_REGION>
```

Once this is done, you can create the ZenML stack component as follows:

2. Register an ECR container registry stack component:

```shell
zenml container-registry register ecr-registry --flavor=aws --uri=<ACCOUNT_ID>.dkr.ecr.<YOUR_REGION>.amazonaws.com --connector aws-connector
```

More details [here](https://docs.zenml.io/stacks/container-registries/aws).

## 4) Create stack

{% tabs %}
{% tab title="CLI" %}

```shell
export STACK_NAME=aws_stack

zenml stack register ${STACK_NAME} -o ${ORCHESTRATOR_NAME} \
    -a ${ARTIFACT_STORE_NAME} -c ${CONTAINER_REGISTRY_NAME} --set
```

{% hint style="info" %}
In case you want to also add any other stack components to this stack, feel free to do so.
{% endhint %}
{% endtab %}

{% tab title="Dashboard" %}

{% endtab %}
{% endtabs %}

## 5) And you're already done!

Just like that, you now have a fully working AWS stack ready to go. Feel free to take it for a spin by running a pipeline on it.

Define a ZenML pipeline:

```python
from zenml import pipeline, step

@step
def hello_world() -> str:
    return "Hello from SageMaker!"

@pipeline
def aws_sagemaker_pipeline():
    hello_world()

if __name__ == "__main__":
    aws_sagemaker_pipeline()
```

Save this code to run.py and execute it. The pipeline will use AWS S3 for artifact storage, Amazon SageMaker Pipelines for orchestration, and Amazon ECR for container registry.

```shell
python run.py
```

<figure><img src="https://1559531010-files.gitbook.io/~/files/v0/b/gitbook-x-prod.appspot.com/o/spaces%2Fu4tnWimk4Ev9z09qY15M%2Fuploads%2Fgit-blob-bce622fd61f1356f1622280b02cc8ea2b74b305c%2Frun_with_repository.png?alt=media" alt=""><figcaption><p>Sequence of events that happen when running a pipeline on a remote stack with a code repository</p></figcaption></figure>

Read more in the [production guide](https://docs.zenml.io/user-guides/production-guide).

## Cleanup

{% hint style="warning" %}
Make sure you no longer need the resources before deleting them. The instructions and commands that follow are DESTRUCTIVE.
{% endhint %}

Delete any AWS resources you no longer use to avoid additional charges. You'll want to do the following:

```shell
# delete the S3 bucket
aws s3 rm s3://your-bucket-name --recursive
aws s3api delete-bucket --bucket your-bucket-name

# delete the SageMaker domain
aws sagemaker delete-domain --domain-id <DOMAIN_ID>

# delete the ECR repository
aws ecr delete-repository --repository-name zenml-repository --force

# detach custom policies from the IAM role
aws iam detach-role-policy --role-name zenml-role --policy-arn arn:aws:iam::<YOUR_ACCOUNT_ID>:policy/ZenML-S3-Policy
aws iam detach-role-policy --role-name zenml-role --policy-arn arn:aws:iam::<YOUR_ACCOUNT_ID>:policy/ZenML-ECR-Policy
aws iam detach-role-policy --role-name zenml-role --policy-arn arn:aws:iam::<YOUR_ACCOUNT_ID>:policy/ZenML-SageMaker-Policy

# delete the custom policies
aws iam delete-policy --policy-arn arn:aws:iam::<YOUR_ACCOUNT_ID>:policy/ZenML-S3-Policy
aws iam delete-policy --policy-arn arn:aws:iam::<YOUR_ACCOUNT_ID>:policy/ZenML-ECR-Policy
aws iam delete-policy --policy-arn arn:aws:iam::<YOUR_ACCOUNT_ID>:policy/ZenML-SageMaker-Policy

# delete the IAM role
aws iam delete-role --role-name zenml-role
```

Make sure to run these commands in the same AWS region where you created the resources.

By running these cleanup commands, you will delete the S3 bucket, SageMaker domain, ECR repository, and IAM role, along with their associated policies. This will help you avoid any unnecessary charges for resources you no longer need.

Remember to be cautious when deleting resources and ensure that you no longer require them before running the deletion commands.

## Conclusion

In this guide, we walked through the process of setting up an AWS stack with ZenML to run your machine learning pipelines in a scalable and production-ready environment. The key steps included:

1. Setting up credentials and the local environment by creating an IAM role with the necessary permissions.
2. Creating a ZenML service connector to authenticate with AWS services using the IAM role.
3. Configuring stack components, including an S3 artifact store, a SageMaker Pipelines orchestrator, and an ECR container registry.
4. Registering the stack components and creating a ZenML stack.

By following these steps, you can leverage the power of AWS services, such as S3 for artifact storage, SageMaker Pipelines for orchestration, and ECR for container management, all within the ZenML framework. This setup allows you to build, deploy, and manage machine learning pipelines efficiently and scale your workloads based on your requirements.

The benefits of using an AWS stack with ZenML include:

* Scalability: Leverage the scalability of AWS services to handle large-scale machine learning workloads.
* Reproducibility: Ensure reproducibility of your pipelines with versioned artifacts and containerized environments.
* Collaboration: Enable collaboration among team members by using a centralized stack and shared resources.
* Flexibility: Customize and extend your stack components based on your specific needs and preferences.

Now that you have a functional AWS stack set up with ZenML, you can explore more advanced features and capabilities offered by ZenML. Some next steps to consider:

* Dive deeper into ZenML's [production guide](https://docs.zenml.io/user-guides/production-guide) to learn best practices for deploying and managing production-ready pipelines.
* Explore ZenML's [integrations](https://docs.zenml.io/stacks) with other popular tools and frameworks in the machine learning ecosystem.
* Join the [ZenML community](https://zenml.io/slack) to connect with other users, ask questions, and get support.

By leveraging the power of AWS and ZenML, you can streamline your machine learning workflows, improve collaboration, and deploy production-ready pipelines with ease. What follows is a set of best practices for using your AWS stack with ZenML.

## Best Practices for Using an AWS Stack with ZenML

When working with an AWS stack in ZenML, consider the following best practices to optimize your workflow, enhance security, and improve cost-efficiency. These are all things you might want to do or amend in your own setup once you have tried running some pipelines on your AWS stack.

### Use IAM Roles and Least Privilege Principle

Always adhere to the principle of least privilege when setting up IAM roles. The guide above provides specific custom IAM policies with minimal required permissions instead of broad managed policies. This approach significantly reduces security risks by:

* Limiting S3 access to only your specific bucket
* Restricting SageMaker permissions to pipeline operations only
* Scoping ECR access to container operations only
* Including proper IAM PassRole conditions

Regularly review and audit your [IAM roles](https://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles.html) to ensure they remain appropriate and secure. Consider using AWS CloudTrail to monitor which permissions are actually being used and remove any unnecessary ones.

### Leverage AWS Resource Tagging

Implement a [consistent tagging strategy](https://aws.amazon.com/solutions/guidance/tagging-on-aws/) for all of your AWS resources that you use for your pipelines. For example, if you have S3 as an artifact store in your stack, you should tag it like shown below:

```shell
aws s3api put-bucket-tagging --bucket your-bucket-name --tagging 'TagSet=[{Key=Project,Value=ZenML},{Key=Environment,Value=Production}]'
```

These tags will help you with billing and cost allocation tracking and also with any cleanup efforts.

### Implement Cost Management Strategies

Use [AWS Cost Explorer](https://aws.amazon.com/aws-cost-management/aws-cost-explorer/) and [AWS Budgets](https://aws.amazon.com/aws-cost-management/aws-budgets/) to monitor and manage your spending. To create a cost budget:

1. Create a JSON file (e.g., `budget-config.json`) defining the budget:

```json
{
  "BudgetLimit": {
    "Amount": "100",
    "Unit": "USD"
  },
  "BudgetName": "ZenML Monthly Budget",
  "BudgetType": "COST",
  "CostFilters": {
    "TagKeyValue": [
      "user:Project$ZenML"
    ]
  },
  "CostTypes": {
    "IncludeTax": true,
    "IncludeSubscription": true,
    "UseBlended": false
  },
  "TimeUnit": "MONTHLY"
}
```

2. Create the cost budget:

```shell
aws budgets create-budget --account-id your-account-id --budget file://budget-config.json
```

Set up cost allocation tags to track expenses related to your ZenML projects:

```shell
aws ce create-cost-category-definition --name ZenML-Projects --rules-version 1 --rules file://rules.json
```

### Use Warm Pools for your SageMaker Pipelines

[Warm Pools in SageMaker](https://docs.zenml.io/stacks/orchestrators/sagemaker#using-warm-pools-for-your-pipelines) can significantly reduce the startup time of your pipeline steps, leading to faster iterations and improved development efficiency. This feature keeps compute instances in a "warm" state, ready to quickly start new jobs.

To enable Warm Pools, use the `SagemakerOrchestratorSettings` class:

```python
from zenml.integrations.aws.orchestrators.sagemaker import SagemakerOrchestratorSettings

sagemaker_orchestrator_settings = SagemakerOrchestratorSettings(
    keep_alive_period_in_seconds = 300, # 5 minutes, default value
)
```

This configuration keeps instances warm for 5 minutes after each job completes, allowing subsequent jobs to start faster if initiated within this timeframe. The reduced startup time can be particularly beneficial for iterative development processes or frequently run pipelines.

### Implement a Robust Backup Strategy

Regularly backup your critical data and configurations. For S3, enable versioning and consider using [cross-region replication](https://docs.aws.amazon.com/AmazonS3/latest/userguide/replication.html) for disaster recovery.

By following these best practices and implementing the provided examples, you can create a more secure, efficient, and cost-effective AWS stack for your ZenML projects. Remember to regularly review and update your practices as your projects evolve and as AWS introduces new features and services.

<figure><img src="https://static.scarf.sh/a.png?x-pxid=f0b4f458-0a54-4fcd-aa95-d5ee424815bc" alt="ZenML Scarf"><figcaption></figcaption></figure>