Using taints and tolerations in EKS to run machine learning pods on GPU nodes

One of our clients uses EKS and they had a problem: running machine learning pods on a special Kubernetes node group that has GPU nodes.

We solved this problem using Kubernetes taints, tolerations, labels, and node selectors.

This is how to configure it:

Choose an Amazon EKS-optimized Arm Amazon Linux AMI and install the NVIDIA device plugin for Kubernetes.

The link below describes how to configure this part properly:
https://docs.aws.amazon.com/eks/latest/userguide/eks-optimized-ami.html

Before proceeding to the next section, make sure you have:

• nodes using EKS optimised Arm Amazon Linux AMI
• a node group that contains these nodes; and
• NVIDIA device plugin for Kubernetes installed.

Prevent pods from running on the GPU nodes

To achieve this, add a taint to the node group that has the GPU nodes.

If you are using the AWS Console, follow these instructions:

• Go to EKS
• Select the cluster
• Click on the Compute tab
• Click on the target Node Group
• Click on the Kubernetes Taints tab
• Add a Kubernetes taint like this:
  • Key: GPU
  • Value: true
  • Effect: NoSchedule

When this taint is added, basically it means that no pod can be scheduled to execute on any of the nodes of this node group unless the pod tolerates this taint.

Allow the target pods to run on GPU nodes

Add the tolerations to the target Kubernetes pod/deployment. This is a sample of tolerations in a Kubernetes deployment yaml file:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: my-app
spec:
  template:
    spec:
      tolerations:
        - key: "GPU"
          operator: "Equal"
          value: "true"
          effect: "NoSchedule"

When the tolerations tag is added, it means this pod ‘tolerates’ being scheduled to nodes of the node group above.

But how can we prevent these pods from being scheduled in other node groups?

Add the label to the node group

If you are using the AWS Console, follow these instructions:

• Go to EKS
• Select the cluster
• Click on the Compute tab
• Click on the target Node Group
• Click on the Kubernetes Labels
• Add a new label with the following values:
  
  • Key = GPU
  • Value = true

This step does not change anything yet, but is a prerequisite for the next step.

Ensure these pods can only be scheduled on this node group

To do this, it’s necessary to add a node selector.

This is a sample of nodeSelector and tolerations together in a Kubernetes deployment yaml file:

apiVersion: apps/v1
kind: Deployment
metadata:
  name: my-app
spec:
  template:
    spec:
      nodeSelector:
        GPU: "true"
      tolerations:
        - key: "GPU"
          operator: "Equal"
          value: "true"
          effect: "NoSchedule"

Since the node group also has the label “GPU: true”, then it guarantees that the pods created by this deployment can only be scheduled in the GPU node group.

Wrapping up

In this blog post, we discussed how to schedule pods in different node groups.

First, we learned that we can use taints and tolerations to repel other pods from running on node groups, then we looked at using node selectors and labels to ensure pods will only run in a specific node group.

The combination of these two techniques allows us to control where pods are scheduled in a very flexible way and, more importantly, solve our client’s problem.