Scalable AI Model Serving with Amazon ECS

Question

Pulumi · Accepted Answer

Deploying a scalable AI model serving solution on Amazon ECS involves several steps. The key components of such a solution typically include an Amazon Elastic Container Registry (ECR) to store the container images of your model, an ECS cluster to run your containers, and a load balancer to distribute requests among the containers. Below I outline a Pulumi program that will create these resources.

The program does the following:

1. It creates an ECR repository to hold the Docker image of our AI model.
2. It sets up an ECS cluster where our containers will be deployed.
3. It creates an ECS task definition that describes how the containers will run, including the Docker image to use and the required resources.
4. It establishes an ECS service that maintains the desired number of instances of the task definition running.
5. It sets up a load balancer to distribute traffic between the instances of our task for high availability and scalability.

```python
import pulumi
import pulumi_aws as aws

# Create an ECR repository to store our AI model's container image.
ai_model_repo = aws.ecr.Repository("aiModelRepo")

# Create an ECS cluster that will host our AI model containers.
ecs_cluster = aws.ecs.Cluster("ecsCluster")

# Assuming you've already built and pushed your Docker image to the ECR repository
# generated from a previous build process, and the Docker image URL is stored in a Pulumi config.
# For this example, we assume the Docker image URL is `123456789012.dkr.ecr.us-west-2.amazonaws.com/ai-model:v1`.
image_url = "123456789012.dkr.ecr.us-west-2.amazonaws.com/ai-model:v1"

# Define an ECS task definition for our AI model container.
task_definition = aws.ecs.TaskDefinition("taskDefinition",
    family="aiModelFamily",
    cpu="256",
    memory="512",
    network_mode="awsvpc",
    requires_compatibilities=["FARGATE"],
    execution_role_arn=aws.iam.Role("ecsTaskExecutionRole",
        assume_role_policy=json.dumps({
            "Version": "2012-10-17",
            "Statement": [{
                "Action": "sts:AssumeRole",
                "Principal": {
                    "Service": "ecs-tasks.amazonaws.com",
                },
                "Effect": "Allow",
                "Sid": "",
            }],
        }),
    ).arn,
    container_definitions=pulumi.Output.all(ai_model_repo.repository_url, image_url).apply(
        lambda args: json.dumps([{
            "name": "aiModelContainer",
            "image": args[1],
            "cpu": 256,
            "memory": 512,
            "essential": True,
            "portMappings": [{
                "containerPort": 80,
                "hostPort": 80
            }],
        }])
    )
)

# Set up an Application Load Balancer (ALB) to distribute traffic to our AI model containers.
# Create the load balancer.
lb = aws.lb.LoadBalancer("appLb",
    internal=False,
    load_balancer_type="application",
    security_groups=[aws_security_group_id],
    subnets=aws_subnet_ids
)

# Create the target group.
target_group = aws.lb.TargetGroup("appTg",
    port=80,
    protocol="HTTP",
    target_type="ip",
    vpc_id=aws_vpc_id
)

# Create the listener.
listener = aws.lb.Listener("appListener",
    load_balancer_arn=lb.arn,
    port=80,
    default_actions=[{
        "type": "forward",
        "target_group_arn": target_group.arn
    }]
)

# Create an ECS service to run and maintain our containers.
ecs_service = aws.ecs.Service("ecsService",
    cluster=ecs_cluster.arn,
    desired_count=2, # Start with 2 instances for high availability.
    launch_type="FARGATE",
    task_definition=task_definition.arn,
    network_configuration={
        "awsvpcConfiguration": {
            "subnets": aws_subnet_ids,
            "security_groups": [aws_security_group_id],
            "assign_public_ip": "ENABLED",
        }
    },
    load_balancers=[{
        "target_group_arn": target_group.arn,
        "container_name": "aiModelContainer",
        "container_port": 80,
    }],
    opts=pulumi.ResourceOptions(depends_on=[listener])
)

# Export the URLs of the load balancer to access the services.
pulumi.export("load_balancer_url", lb.dns_name)
```

In this Pulumi program, replace `aws_security_group_id`, `aws_subnet_ids`, and `aws_vpc_id` with your specific AWS security group ID, subnet IDs, and VPC ID respectively. These are required to configure the network where your ECS service will be deployed.

In the `task_definition`, we've specified the CPU and memory that each task instance should use. This would, of course, need to be tailored to the needs of the specific AI model you are serving. We've also chosen a networking mode compatible with AWS Fargate, which is a serverless compute engine for containers and simplifies the process of running containers.

Finally, the ECS service is configured with a network configuration that specifies the subnets and security groups to use, as well as enabling a public IP assignment, which can then be accessed via the ALB. The `desired_count` parameter in the ECS service specifies the number of task instances we want to start with.

With this configuration, we can further adjust the `desired_count` or use AWS's auto-scaling features to scale the number of tasks up or down based on the load, making this setup highly scalable.