Scale your machine studying workloads on Amazon ECS powered by AWS Trainium situations

Operating machine studying (ML) workloads with containers is turning into a standard observe. Containers can totally encapsulate not simply your coaching code, however your complete dependency stack right down to the {hardware} libraries and drivers. What you get is an ML improvement surroundings that’s constant and transportable. With containers, scaling on a cluster turns into a lot simpler.

In late 2022, AWS introduced the overall availability of Amazon EC2 Trn1 instances powered by AWS Trainium accelerators, that are goal constructed for high-performance deep studying coaching. Trn1 situations ship as much as 50% financial savings on coaching prices over different comparable Amazon Elastic Compute Cloud (Amazon EC2) situations. Additionally, the AWS Neuron SDK was launched to enhance this acceleration, giving builders instruments to work together with this expertise comparable to to compile, runtime, and profile to realize high-performance and cost-effective mannequin trainings.

Amazon Elastic Container Service (Amazon ECS) is a completely managed container orchestration service that simplifies your deployment, administration, and scaling of containerized purposes. Merely describe your utility and the sources required, and Amazon ECS will launch, monitor, and scale your utility throughout versatile compute choices with automated integrations to different supporting AWS companies that your utility wants.

On this publish, we present you the way to run your ML coaching jobs in a container utilizing Amazon ECS to deploy, handle, and scale your ML workload.

Answer overview

We stroll you thru the next high-level steps:

1. Provision an ECS cluster of Trn1 situations with AWS CloudFormation.
2. Construct a customized container picture with the Neuron SDK and push it to Amazon Elastic Container Registry (Amazon ECR).
3. Create a job definition to outline an ML coaching job to be run by Amazon ECS.
4. Run the ML job on Amazon ECS.

Conditions

To comply with alongside, familiarity with core AWS companies comparable to Amazon EC2 and Amazon ECS is implied.

Provision an ECS cluster of Trn1 situations

To get began, launch the supplied CloudFormation template, which is able to provision required sources comparable to a VPC, ECS cluster, and EC2 Trainium occasion.

We use the Neuron SDK to run deep studying workloads on AWS Inferentia and Trainium-based situations. It helps you in your end-to-end ML improvement lifecycle to create new fashions, optimize them, then deploy them for manufacturing. To coach your mannequin with Trainium, you could set up the Neuron SDK on the EC2 situations the place the ECS duties will run to map the NeuronDevice related to the {hardware}, in addition to the Docker picture that can be pushed to Amazon ECR to entry the instructions to coach your mannequin.

Customary variations of Amazon Linux 2 or Ubuntu 20 don’t include AWS Neuron drivers put in. Subsequently, we now have two totally different choices.

The primary possibility is to make use of a Deep Studying Amazon Machine Picture (DLAMI) that has the Neuron SDK already put in. A pattern is obtainable on the GitHub repo. You may select a DLAMI primarily based on the opereating system. Then run the next command to get the AMI ID:

aws ec2 describe-images --region us-east-1 --owners amazon --filters 'Identify=title,Values=Deep Studying AMI Neuron PyTorch 1.13.? (Amazon Linux 2) ????????' 'Identify=state,Values=accessible' --query 'reverse(sort_by(Photographs, &CreationDate))[:1].ImageId' --output textual content

The output can be as follows:

ami-06c40dd4f80434809

This AMI ID can change over time, so be sure that to make use of the command to get the best AMI ID.

Now you possibly can change this AMI ID within the CloudFormation script and use the ready-to-use Neuron SDK. To do that, search for EcsAmiId in Parameters:

"EcsAmiId": { 
    "Kind": "String", 
    "Description": "AMI ID", 
    "Default": "ami-09def9404c46ac27c" 
}

The second possibility is to create an occasion filling the userdata discipline throughout stack creation. You don’t want to put in it as a result of CloudFormation will set this up. For extra data, discuss with the Neuron Setup Guide.

For this publish, we use possibility 2, in case you could use a customized picture. Full the next steps:

1. Launch the supplied CloudFormation template.
2. For KeyName, enter a reputation of your required key pair, and it’ll preload the parameters. For this publish, we use trainium-key.
3. Enter a reputation to your stack.
4. If you happen to’re operating within the us-east-1 Area, you possibly can hold the values for ALBName and AZIds at their default.

To verify what Availability Zone within the Area has Trn1 accessible, run the next command:

aws ec2 describe-instance-type-offerings --region us-east1 --location-type availability-zone --filter Identify=instance-type,Values=trn1.2xlarge

5. Select Subsequent and end creating the stack.

When the stack is full, you possibly can transfer to the subsequent step.

Put together and push an ECR picture with the Neuron SDK

Amazon ECR is a completely managed container registry providing high-performance internet hosting, so you possibly can reliably deploy utility photos and artifacts anyplace. We use Amazon ECR to retailer a customized Docker picture containing our scripts and Neuron packages wanted to coach a mannequin with ECS jobs operating on Trn1 situations. You may create an ECR repository utilizing the AWS Command Line Interface (AWS CLI) or AWS Management Console. For this publish, we use the console. Full the next steps:

1. On the Amazon ECR console, create a brand new repository.
2. For Visibility settings¸ choose Non-public.
3. For Repository title, enter a reputation.
4. Select Create repository.

Now that you’ve a repository, let’s construct and push a picture, which might be constructed regionally (into your laptop computer) or in a AWS Cloud9 surroundings. We’re coaching a multi-layer perceptron (MLP) mannequin. For the unique code, discuss with Multi-Layer Perceptron Training Tutorial.

5. Copy the train.py and model.py recordsdata right into a mission.

It’s already appropriate with Neuron, so that you don’t want to vary any code.

7. 5. Create a Dockerfile that has the instructions to put in the Neuron SDK and coaching scripts:

FROM amazonlinux:2

RUN echo $'[neuron] n
title=Neuron YUM Repository n
baseurl=https://yum.repos.neuron.amazonaws.com n
enabled=1' > /and so forth/yum.repos.d/neuron.repo

RUN rpm --import https://yum.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB

RUN yum set up aws-neuronx-collectives-2.* -y
RUN yum set up aws-neuronx-runtime-lib-2.* -y
RUN yum set up aws-neuronx-tools-2.* -y
RUN yum set up -y tar gzip pip
RUN yum set up -y python3 python3-pip
RUN yum set up -y python3.7-venv gcc-c++
RUN python3.7 -m venv aws_neuron_venv_pytorch

# Activate Python venv
ENV PATH="/aws_neuron_venv_pytorch/bin:$PATH"
RUN python -m pip set up -U pip
RUN python -m pip set up wget
RUN python -m pip set up awscli

RUN python -m pip config set world.extra-index-url https://pip.repos.neuron.amazonaws.com
RUN python -m pip set up torchvision tqdm torch-neuronx neuronx-cc==2.* pillow
RUN mkdir -p /choose/ml/mnist_mlp
COPY mannequin.py /choose/ml/mnist_mlp/mannequin.py
COPY prepare.py /choose/ml/mnist_mlp/prepare.py
RUN chmod +x /choose/ml/mnist_mlp/prepare.py
CMD ["python3", "/opt/ml/mnist_mlp/train.py"]

aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin {your-account-id}.dkr.ecr.{your-region}.amazonaws.com

docker construct -t mlp_trainium .

docker tag mlp_trainium:newest {your-account-id}.dkr.ecr.us-east-1.amazonaws.com/mlp_trainium:newest

docker push {your-account-id}.dkr.ecr.{your-region}.amazonaws.com/{your-ecr-repo-name}:newest

{
    "household": "mlp_trainium",
    "containerDefinitions": [
        {
            "name": "mlp_trainium",
            "image": "{your-account-id}.dkr.ecr.us-east-1.amazonaws.com/{your-ecr-repo-name}",
            "cpu": 0,
            "memoryReservation": 1000,
            "portMappings": [],
            "important": true,
            "surroundings": [],
            "mountPoints": [],
            "volumesFrom": [],
            "linuxParameters": {
                "capabilities": {
                    "add": [
                        "IPC_LOCK"
                    ]
                },
                "units": [
                    {
                        "hostPath": "/dev/neuron0",
                        "containerPath": "/dev/neuron0",
                        "permissions": [
                            "read",
                            "write"
                        ]
                    }
                ]
            },
            ,
            "logConfiguration": {
                "logDriver": "awslogs",
                "choices": {
                    "awslogs-create-group": "true",
                    "awslogs-group": "/ecs/task-logs",
                    "awslogs-region": "us-east-1",
                    "awslogs-stream-prefix": "ecs"
                }
            }
        }
    ],
    "networkMode": "awsvpc",
    "placementConstraints": [
        {
            "type": "memberOf",
            "expression": "attribute:ecs.os-type == linux"
        },
        {
            "type": "memberOf",
            "expression": "attribute:ecs.instance-type == trn1.2xlarge"
        }
    ],
    "requiresCompatibilities": [
        "EC2"
    ],
    "cpu": "1024",
    "reminiscence": "3072"
}

aws ecs register-task-definition 
--family mlp-trainium 
--container-definitions '[{    
    "name": "my-container-1",    
    "image": "{your-account-id}.dkr.ecr.us-east-1.amazonaws.com/{your-ecr-repo-name}",
    "cpu": 0,
    "memoryReservation": 1000,
    "portMappings": [],
    "important": true,
    "surroundings": [],
    "mountPoints": [],
    "volumesFrom": [],
    "logConfiguration": {
        "logDriver": "awslogs",
        "choices": {
            "awslogs-create-group": "true",
            "awslogs-group": "/ecs/task-logs",
            "awslogs-region": "us-east-1",
            "awslogs-stream-prefix": "ecs"
        }
    },
    "linuxParameters": {
        "capabilities": {
            "add": [
                "IPC_LOCK"
            ]
        },
        "units": [{
            "hostPath": "/dev/neuron0",
            "containerPath": "/dev/neuron0",
            "permissions": ["read", "write"]
        }]
    }
}]' 
--requires-compatibilities EC2
--cpu "8192" 
--memory "16384" 
--placement-constraints '[{
    "type": "memberOf",
    "expression": "attribute:ecs.instance-type == trn1.2xlarge"
}, {
    "type": "memberOf",
    "expression": "attribute:ecs.os-type == linux"
}]'

aws ecs run-task --cluster <your-cluster-name> --task-definition <your-task-name> --count 1 --network-configuration '{"awsvpcConfiguration": {"subnets": ["<your-subnet-name> "], "securityGroups": ["<your-sg-name> "] }}'