Enabling GPU Nodes for PyTorch Workloads on EKS with Autoscaling
Published: 2023-09-12 18:24:40

## Enabling GPU Nodes for PyTorch Workloads on EKS with Autoscaling

Amazon Elastic Kubernetes Service (EKS) provides a managed Kubernetes service that makes it easier for users to run Kubernetes on AWS without needing to install, operate, and maintain their own Kubernetes control plane or nodes. When running machine learning workloads, especially those that require GPU acceleration like PyTorch, it's essential to set up GPU nodes. This article will guide you through the process of setting up GPU nodes for PyTorch workloads on EKS with autoscaling.

### Prerequisites

* Basic knowledge of Kubernetes, EKS, and Terraform.
* AWS CLI and kubectl installed and configured.
* Docker installed.
* Helm installed.

### 1. Setting up the EKS Cluster with GPU Nodes using Terraform

Before applying the Terraform configuration:

* Ensure you have the AWS provider configured in your Terraform setup.
* Initialize the Terraform directory using `terraform init`.

Here's a snippet from a Terraform configuration that sets up an EKS cluster and a self-managed GPU node group:

```hcl
## (https://github.com/aws-ia/terraform-aws-eks-blueprints)
## ... [other Terraform code]

## Cluster Configuration
module "eks" {
  # ... [other configuration]

  self_managed_node_groups = {
    gpu_node_group = {
      node_group_name = "gpu-node-group"
      ami_type        = "AL2_x86_64_GPU"
      capacity_type   = "ON_DEMAND"
      instance_types = [
        "g4dn.xlarge",
        "g4dn.2xlarge",
      ]
      # ... [other configuration]

      taints = {
        dedicated = {
          key    = "nvidia.com/gpu"
          value  = "true"
          effect = "NO_SCHEDULE"
        }
      }
      # ... [other configuration]
    }
  }
}
```

### 2. Building the PyTorch Container with GPU Support

Your `requirements.txt` should contain:

```bash
torch
torchvision
```

To run PyTorch workloads on the GPU nodes, you need a container image with the necessary dependencies. Here's a Dockerfile that sets up a PyTorch environment with GPU support:

Your Dockerfile should contain:

```Dockerfile
FROM nvidia/cuda:12.2.0-devel-ubuntu20.04
WORKDIR /app

# Install Python and pip
RUN apt-get update && \
    apt-get install -y python3 python3-pip && \
    rm -rf /var/lib/apt/lists/*

# Install requirements
COPY requirements.txt ./
RUN pip install --no-cache-dir -r requirements.txt
COPY app.py ./
# Run app
CMD ["python3", "./app.py"]
```

Your `app.py` can be as follows:

```python
import torch
import time

def main():
    # Check if CUDA (GPU support) is available
    if not torch.cuda.is_available():
        print("GPU not available. Exiting.")
        return

    # Set the device to GPU
    device = torch.device("cuda:0")

    iteration = 0
    while True:
        iteration += 1

        # Create two random tensors
        a = torch.randn(1000, 1000, device=device)
        b = torch.randn(1000, 1000, device=device)

        # Perform a matrix multiplication on GPU
        c = torch.matmul(a, b)

        print(f"Iteration {iteration}: Matrix multiplication completed on GPU!")
        print(f"Result tensor shape: {c.shape}")

        # Pause for a short duration before the next iteration
        time.sleep(2)

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

### 3. Pushing PyTorch GPU App to AWS ECR

#### Prerequisites:

* Docker installed and running.
* AWS CLI installed and configured with the necessary permissions.

#### Steps:

1. **Set Environment Variables**:
    
    Set your AWS account ID and region as environment variables:
    
    ```bash
    export AWS_ACCOUNT=<your-account-id>
    export AWS_REGION=<your-region>
    ```
    
    Replace `<your-account-id>` with your AWS account ID and `<your-region>` with your AWS region (e.g., `us-west-1`).
    
2. **Docker Build**:
    
    Navigate to the directory containing your Dockerfile and app.js. Build the Docker image:
    
    ```bash
    docker buildx build --platform linux/amd64 -t pytorch-gpu-app .
    ```
    
3. **Create an ECR Repository** (if you haven't already):
    
    ```bash
    aws ecr create-repository --repository-name pytorch-gpu-app --region $AWS_REGION
    ```
    
4. **Authenticate Docker to the ECR Registry**:
    
    ```bash
    aws ecr get-login-password --region $AWS_REGION | docker login --username AWS --password-stdin $AWS_ACCOUNT.dkr.ecr.$AWS_REGION.amazonaws.com
    ```
    
5. **Tag the Docker Image**:
    
    ```bash
    docker tag pytorch-gpu-app:latest $AWS_ACCOUNT.dkr.ecr.$AWS_REGION.amazonaws.com/pytorch-gpu-app:latest
    ```
    
6. **Push the Docker Image to ECR**:
    
    ```bash
    docker push $AWS_ACCOUNT.dkr.ecr.$AWS_REGION.amazonaws.com/pytorch-gpu-app:latest
    ```
    

### 3.5 Creating an imagePullSecret for AWS ECR

Before deploying your PyTorch workload on EKS, if your Docker image is stored in a private AWS ECR repository, you'll need to create a Kubernetes secret (`imagePullSecret`) to allow your EKS nodes to pull the image.

1. **Retrieve an authentication token** to use to authenticate your Docker client to your registry:
    
    ```bash
    TOKEN=$(aws ecr get-login-password --region $AWS_REGION)
    ```
    
2. **Create the `imagePullSecret`**:
    
    ```bash
kubectl create secret docker-registry ecr-secret \
--docker-server=AWS_ACCOUNT.dkr.ecr.$AWS_REDION.amazonaws.com \
--docker-username=AWS \
--docker-password="${TOKEN}"
secret/ecr-secret created
    ```
    
    Replace `<your-account-id>` with your AWS account ID and `<your-region>` with your AWS region.
    

### 4. Deploying the PyTorch Workload on EKS

Before deploying our workloads we need to enable the nvidia k8s-device-plugin: 

```bash
kubectl create -f https://raw.githubusercontent.com/NVIDIA/k8s-device-plugin/v0.14.1/nvidia-device-plugin.yml
```

To deploy the PyTorch workload on the GPU nodes in EKS, use the following **deployment.yaml**:

```yaml
apiVersion: apps/v1
kind: Deployment
metadata:
  name: pytorch
spec:
  replicas: 1
  selector:
    matchLabels:
      app: pytorch
  template:
    metadata:
      labels:
        app: pytorch
    spec:
      tolerations:
        - key: "nvidia.com/gpu"
          operator: "Equal"
          value: "true"
          effect: "NoSchedule"
      nodeSelector:
        NodeGroup: gpu-node-group
      imagePullSecrets:
        - name: ecr-secret
      containers:
        - name: pytorch
          image: <your-account-id>.dkr.ecr.<your-region>.amazonaws.com/pytorch-gpu-app:latest
          imagePullPolicy: Always
          resources:
            requests:
              cpu: 2
              memory: 2Gi
              nvidia.com/gpu: 1
            limits:
              cpu: 4
              memory: 4Gi
              nvidia.com/gpu: 1
```

* * *


### 4. Installing the Cluster Autoscaler

The Cluster Autoscaler automatically adjusts the size of the cluster, adding or removing nodes based on resource requirements and constraints. To install the Cluster Autoscaler on your EKS cluster, you can use the Helm package manager.

Ensure that the necessary RBAC roles and permissions are set up for the Cluster Autoscaler.

First, add the Cluster Autoscaler Helm repository:

```bash
$ helm repo add autoscaler https://kubernetes.github.io/autoscaler
```

#### Using Autodiscovery

This method allows the Cluster Autoscaler to discover node groups automatically:

```bash
$ helm install my-release autoscaler/cluster-autoscaler \
    --set 'autoDiscovery.clusterName'=<CLUSTER NAME>
```

Replace `<CLUSTER NAME>` with the name of your EKS cluster.

###5. Scaling Up GPU Nodes in EKS
Scaling up the number of GPU nodes in your EKS cluster ensures that you have sufficient resources to handle increased workloads. Here's how you can scale up the GPU nodes:

####Update the Terraform Configuration:

Modify your existing Terraform configuration to increase the desired number of nodes in the gpu_node_group. For instance, if you initially set up 2 nodes and want to scale up to 4, you'd update the desired_capacity attribute.

```hcl
self_managed_node_groups = {
  gpu_node_group = {
    # ... [other configuration]

    desired_capacity = 4
    min_size         = 2
    max_size         = 6

    # ... [other configuration]
  }
}
```

The min_size and max_size attributes define the minimum and maximum number of nodes in the node group, respectively. Adjust these values based on your requirements.

Apply the Updated Configuration:

Run the following command to apply the updated configuration:

```bash
terraform apply
```

Terraform will show you a plan of the changes it will make. Review the plan to ensure it's making the desired changes, then confirm to apply.


####Monitor the Scaling Process:

You can monitor the scaling process using the AWS Management Console or the kubectl command. To check the nodes being added to your cluster using kubectl, run:

```bash
kubectl get nodes
```

This will display a list of all nodes in your cluster. You should see the new nodes being added.

####Optimizing Costs:

Remember, adding more GPU nodes will increase your AWS bill. To optimize costs, consider using a mix of On-Demand and Spot Instances, or setting up Auto Scaling policies that scale down the number of nodes during off-peak hours.

---
In conclusion, by following these steps, you can efficiently set up GPU nodes for PyTorch workloads on EKS with autoscaling. This setup ensures that your machine learning workloads can leverage the power of GPUs for faster processing and better performance, while also benefiting from the elasticity provided by the Cluster Autoscaler. Remember to tear down resources after you're done to avoid incurring unnecessary costs using commands like terraform destroy and deleting the EKS cluster.