Auto-Scaling GPU Nodes in EKS for Deep Learning Workloads

Question

Pulumi · Accepted Answer

When deploying a Kubernetes cluster intended for deep learning workloads on AWS, it's common to leverage Amazon EKS (Elastic Kubernetes Service) for the management of the cluster and to use GPU-accelerated instances for the execution of workloads.

For Auto-Scaling GPU Nodes in EKS, we typically need the following components:

1. An EKS cluster
2. A Node Group with GPU instances
3. AWS Identity and Access Management (IAM) roles
4. Amazon Machine Images (AMIs) that are optimized for EKS and support GPUs
5. An Auto Scaling configuration to automatically adjust the number of GPU instances based on demand
6. Appropriate VPC and subnet configuration to ensure network connectivity

Below is a Pulumi program written in Python that provisions an EKS cluster with an auto-scaling node group configured with GPU instances to support deep learning workloads.

```python
import pulumi
import pulumi_aws as aws
import pulumi_eks as eks

# Create an IAM role that can be used by the EKS service.
eks_role = aws.iam.Role("eks_role", assume_role_policy=aws.iam.get_policy_document(statements=[
    aws.iam.get_policy_document_statement(principal={
        "type": "Service",
        "identifiers": ["eks.amazonaws.com"],
    })
]).json)

# Attach the required policies to the IAM role.
# These policies are necessary for the EKS cluster to manage resources.
policies = ["arn:aws:iam::aws:policy/AmazonEKSClusterPolicy", "arn:aws:iam::aws:policy/AmazonEKSServicePolicy"]
for policy in policies:
    aws.iam.RolePolicyAttachment(f"eks_role_{policy.split('/')[-1]}",
        policy_arn=policy,
        role=eks_role.name)

# Create the VPC for the EKS cluster.
vpc = aws.ec2.Vpc("vpc",
    cidr_block="10.100.0.0/16",
    enable_dns_hostnames=True,
    enable_dns_support=True)

# Create the subnet for the EKS cluster.
subnet = aws.ec2.Subnet("subnet",
    vpc_id=vpc.id,
    cidr_block="10.100.0.0/24",
    map_public_ip_on_launch=True)

# Create the EKS cluster.
cluster = eks.Cluster("cluster",
    role_arn=eks_role.arn,
    vpc_id=vpc.id,
    private_subnet_ids=[subnet.id])

# Create an IAM role that can be used by EKS worker nodes.
node_role = aws.iam.Role("node_role", assume_role_policy=aws.iam.get_policy_document(statements=[
    aws.iam.get_policy_document_statement(principal={
        "type": "Service",
        "identifiers": ["ec2.amazonaws.com"],
    })
]).json)

# Attach the required policy to the IAM role for worker nodes.
# This policy allows worker nodes to register with the cluster.
worker_role_policy = aws.iam.RolePolicyAttachment("node_role_AmazonEKSWorkerNodePolicy",
    policy_arn="arn:aws:iam::aws:policy/AmazonEKSWorkerNodePolicy",
    role=node_role.name)

# Create the subnet for the EKS worker nodes.
node_subnet = aws.ec2.Subnet("node_subnet",
    vpc_id=vpc.id,
    cidr_block="10.100.1.0/24",
    map_public_ip_on_launch=True)

# Choose an appropriate AMI for EKS nodes with GPU support.
gpu_ami_id = aws.ssm.get_parameter(name="/aws/service/eks/optimized-ami/1.18/amazon-linux-2-gpu/recommended/image_id").value

# Define the auto-scaling group's scaling settings for the GPU node group.
scaling_config = eks.NodeGroupScalingConfigArgs(
    desired_size=1,
    min_size=1,
    max_size=2,
)

# Create a managed node group for GPU instances within the EKS cluster.
gpu_node_group = eks.ManagedNodeGroup("gpu_node_group",
    cluster=cluster.eks_cluster,
    node_role_arn=node_role.arn,
    subnet_ids=[node_subnet.id],
    ami_type="AL2_x86_64_GPU",
    instance_types=["p2.xlarge"],
    scaling_config=scaling_config)

# Export the cluster's kubeconfig.
pulumi.export("kubeconfig", cluster.kubeconfig)
```

This program sets up an EKS cluster and a GPU-enabled node group. The node group uses the Amazon p2.xlarge instance type, which is suitable for general-purpose GPU compute tasks. The scaling configuration for the node group specifies a minimum of 1 node and a maximum of 2 nodes. Thus, the cluster will scale automatically based on the provided min and max sizes. Of course, these numbers and instance types should be adjusted based on the actual workload requirements and cost considerations.

Keep in mind that optimal setup may require additional tweaks for specific use cases, such as configuring the right VPC networking or setting taints and tolerations for the pod scheduling. For more information on advanced configurations and how to work with Pulumi and EKS, you can refer to the [Pulumi EKS documentation](https://www.pulumi.com/registry/packages/eks/).