Skip to content

Karpenter Blueprint: Deploy an AWS Trainium or AWS Inferentia workload

This blueprint describes the components to get started with using AWS Trainium or AWS Inferentia with Amazon EKS, from AMIs to device plugins. Karpenter streamlines managing accelerated instance node lifecycle management. When you specify accelerated instance types in your Karpenter NodePool, Karpenter automatically selects the appropriate Amazon EKS AMI. Karpenter's provisioning also enables efficient use of Spot Instances across a diverse range of instance types - you can specify multiple accelerator options in your NodePool configuration. This flexibility allows you to balance performance and cost-effectiveness while running accelerator workloads in your Amazon EKS cluster.

When using AL2023 and Bottlerocket you need to deploy a Neuron Kubernetes device plugin to advertise neuron devices from the host, you also need the neuron driver (aws-neuronx-dkms). For AL2023 there is a Neuron EKS Accelerated AMI which is packaged with the Neuron driver. Bottlerocket also has the neuron driver packaged though it is part of the standard EKS Bottlerocket AMI.

Requirements

  • A Kubernetes cluster with Karpenter installed. You can use the blueprint we've used to test this pattern at the cluster folder in the root of this repository.

Deploy Neuron Helm Chart for Kubernetes

The neuron helm chart simplifies installation of the Kubernetes device plugin for Neuron, Neuron Scheduler and Neuron Node Problem Detector. The Neuron Scheduler Extension is disabled by default. The Neuron scheduler extension is required for scheduling pods that require more than one Neuron core or device resource. For a graphical depiction of how the Neuron scheduler extension works, see Neuron Scheduler Extension Flow Diagram. The Neuron scheduler extension finds sets of directly connected devices with minimal communication latency when scheduling containers.

If you require the neuron scheduler, this can be deployed on a general purpose or system NodePool with general purpose compute. You can override the scheduler name as part of the helm chart install scheduler.customScheduler.fullnameOverride property. To use the neuron scheduler you must specify the schedulerName in your Pod specification.

For configuring the Neuron Node Problem Detector see the following AWS Neuron documentation.

To install the neuron helm chart run the following:

helm install neuron-helm-chart \
  oci://public.ecr.aws/neuron/neuron-helm-chart \
  --set "scheduler.enabled=true" \
  --set "scheduler.customScheduler.fullnameOverride=neuron-scheduler" \
  --set "npd.enabled=false" \
  --namespace kube-system

By default, the neuron device plugin and scheduler are deployed in kube-system.

The neuron device plugin advertises both neuron cores aws.amazon.com/neuroncore and neuron devices aws.amazon.com/neuron to the kubelet. When scheduling a workload on aws.amazon.com/neuron all cores associated with that neuron device will be allocated to the Pod.

Confirm installation: 

helm ls

Now that you have the device set-up, let’s enable Karpenter to launch AWS Trainium / Inferentia Amazon EC2 Instances.

Create a NodeClass and NodePool with AWS Trainium / Inferentia Amazon EC2 Instances (AL2023)

The following NodeClass, specify the Security Group and Subnet selector, along with AMI. We are using AL2023 here, and when launching an accelerated instance Karpenter will pick the respective EKS optimized AMI.

Before applying the neuron-nodeclass.yaml replace KARPENTER_NODE_IAM_ROLE_NAME and CLUSTER_NAME in the file with your specific cluster details. If you're using the Terraform template provided in this repo, run the following commands to get the EKS cluster name and the IAM Role name for the Karpenter nodes:

export CLUSTER_NAME=$(terraform -chdir="../../cluster/terraform" output -raw cluster_name)
export KARPENTER_NODE_IAM_ROLE_NAME=$(terraform -chdir="../../cluster/terraform" output -raw node_instance_role_name)

NOTE: If you're not using Terraform, you need to get those values manually. CLUSTER_NAME is the name of your EKS cluster (not the ARN). Karpenter auto-generates the instance profile in your EC2NodeClass given the role that you specify in spec.role with the placeholder KARPENTER_NODE_IAM_ROLE_NAME, which is a way to pass a single IAM role to the EC2 instance launched by the Karpenter NodePool. Typically, the instance profile name is the same as the IAM role(not the ARN).

The EC2NodeClass we’ll deploy looks like this, execute the following command to create the EC2NodeClass file:

cat << EOF > neuron-nodeclass.yaml
apiVersion: karpenter.k8s.aws/v1
kind: EC2NodeClass
metadata:
  name: neuron
spec:
  amiSelectorTerms:
  - alias: al2023@latest
  role: "$KARPENTER_NODE_IAM_ROLE_NAME"
  blockDeviceMappings:
  - deviceName: /dev/xvda
    ebs:
      deleteOnTermination: true
      iops: 10000
      throughput: 125
      volumeSize: 100Gi
      volumeType: gp3
  securityGroupSelectorTerms:
  - tags:
      karpenter.sh/discovery: $CLUSTER_NAME
  subnetSelectorTerms:
  - tags:
      karpenter.sh/discovery: $CLUSTER_NAME
EOF

kubectl apply -f neuron-nodeclass.yaml

If you want to tune node properties it is suggested to create a new EC2NodeClass. For example, you may want to tune node properties such as ephemeral storage size, block device mappings, capacity reservations selector.

Create a dedicated NodePool, states provision instances from inf and trn category, and only allow workloads that tolerate the aws.amazon.com/neuron taint to be scheduled. Apply the following NodePool.

cat << EOF > neuron-nodepool.yaml
apiVersion: karpenter.sh/v1
kind: NodePool
metadata:
  name: neuron
spec:
  limits:
    cpu: 100
    memory: 100Gi
    aws.amazon.com/neuron: 5
  template:
    metadata:
      labels:
        intent: neuron
    spec:
      nodeClassRef:
        group: karpenter.k8s.aws
        name: neuron 
        kind: EC2NodeClass
      requirements:
        - key: karpenter.sh/capacity-type
          operator: In
          values: ["on-demand", "spot"]
        - key: karpenter.k8s.aws/instance-category
          operator: In
          values: ["inf", "trn"]
      expireAfter: 720h
      taints:
         - key: aws.amazon.com/neuron
           effect: NoSchedule
  disruption:
    consolidationPolicy: WhenEmpty
    consolidateAfter: 5m
EOF

kubectl apply -f neuron-nodepool.yaml

We’ve added the aws.amazon.com/neuron taint in the NodePool to prevent workloads that do not tolerate this taint being scheduled on nodes managed by this NodePool (they might not take advantage of it).

Now let’s deploy a test workload.

Deploy a test workload to test neuron drivers are loaded

The following Pod manifest launches a pod and calls the Neuron CLI to check if the accelerator is detected and the driver versions printed to standard output, use kubectl logs pod/neuron-ls to observe. Create workload.yaml from the following Pod specification: 

cat << EOF > workload.yaml
apiVersion: v1
kind: Pod
metadata:
  name: neuron-ls
spec:
  nodeSelector:
    intent: neuron
    karpenter.k8s.aws/instance-accelerator-name: inferentia
  restartPolicy: OnFailure
  schedulerName: neuron-scheduler
  containers:
  - name: neuron-ls
    image: public.ecr.aws/neuron/pytorch-inference-vllm-neuronx:0.13.0-neuronx-py312-sdk2.27.1-ubuntu24.04
    args:
    - "neuron-ls"
    resources:
      requests:
        memory: "30Gi"
        cpu: "3500m"
      limits:
        memory: "30Gi"
        aws.amazon.com/neuron: 2
  tolerations:
  - key: aws.amazon.com/neuron
    effect: NoSchedule
    operator: Exists
EOF

Besides the other constraints to match the workload with the NodePool constraints, we set a node selector karpenter.k8s.aws/instance-accelerator-name to specify the accelerator. We also set the schedulerName to neuron-scheduler.

To deploy the workload execute the following:

$> kubectl apply -f workload.yaml
pod/neuron-ls created

After sometime to list nodes, neuron devices and cores run the following:

$> kubectl get nodes "-o=custom-columns=NAME:.metadata.name,NeuronDevices:.status.allocatable.aws\.amazon\.com/neuron,NeuronCores:.status.allocatable.aws\.amazon\.com/neuroncore"

NAME                        NeuronDevices   NeuronCores
ip-xx-x-x-xxx.ec2.internal  4               16

You can check the pods status by executing:

$> kubectl get pods
NAME        READY   STATUS    RESTARTS   AGE
neuron-ls   1/1     Running   0          4m36s

You can view the pods neuron-smi logs by executing:

$> kubectl logs pod/neuron-ls

instance-type: inf1.6xlarge
instance-id: <redacted>
+--------+--------+----------+--------+-----------+--------------+----------+------+
| NEURON | NEURON |  NEURON  | NEURON | CONNECTED |     PCI      |   CPU    | NUMA |
| DEVICE | CORES  | CORE IDS | MEMORY |  DEVICES  |     BDF      | AFFINITY | NODE |
+--------+--------+----------+--------+-----------+--------------+----------+------+
| 0      | 4      | 0-3      | 8 GB   | 1         | 0000:00:1c.0 | 0-23     | -1   |
| 1      | 4      | 4-7      | 8 GB   | 2, 0      | 0000:00:1d.0 | 0-23     | -1   |
+--------+--------+----------+--------+-----------+--------------+----------+------+
EKS Auto Mode **Prerequisite:** an EKS cluster with Auto Mode enabled, and an EKS Access Entry granting `AmazonEKSAutoNodePolicy` to the node IAM role used by Auto Mode. > If you're using the Terraform template under [`cluster/automode/`](https://github.com/aws-samples/karpenter-blueprints/tree/main/cluster/automode) in this repo, the cluster, node IAM role, and Access Entry are all created for you — you can skip the manual access entry steps below. Trainium / Inferentia support is simpler on Auto Mode: the managed Bottlerocket AMIs include the AWS Neuron drivers, **and Auto Mode also includes the Neuron Kubernetes device plugin automatically** ([docs](https://docs.aws.amazon.com/eks/latest/userguide/auto-accelerated.html)). You do **not** need to install the `neuron-helm-chart` — skip that entire step. To deploy on Auto Mode, replace `<>` and apply the Auto Mode manifests: 
sed -i "s/<<CLUSTER_NAME>>/$CLUSTER_NAME/g" neuron-nodeclass-automode.yaml
kubectl apply -f neuron-nodeclass-automode.yaml
kubectl apply -f neuron-nodepool-automode.yaml
Differences from the OSS version: - **Skip the `neuron-helm-chart` install** — Auto Mode includes the Neuron device plugin and drivers - `NodeClass` (`eks.amazonaws.com/v1`) replaces `EC2NodeClass` (`karpenter.k8s.aws/v1`) - `amiSelectorTerms`, `role`, and `blockDeviceMappings` are removed — Auto Mode manages these - Instance label keys use the `eks.amazonaws.com/` prefix (`eks.amazonaws.com/instance-category`) The example workload from the OSS section above uses `nodeSelector: karpenter.k8s.aws/instance-accelerator-name: inferentia` and `schedulerName: neuron-scheduler`. On Auto Mode: - The label key becomes `eks.amazonaws.com/instance-accelerator-name` - **Drop `schedulerName: neuron-scheduler`** — the Neuron scheduler extension is not installed by Auto Mode. Workloads requiring more than one Neuron core/device can still run, but won't get the topology-aware scheduling that the optional scheduler extension provides. Updated workload nodeSelector: 
nodeSelector:
  intent: neuron
  eks.amazonaws.com/instance-accelerator-name: inferentia
If you are **not** using the `cluster/automode/` Terraform template, configure the Access Entry manually: 
aws eks create-access-entry \
  --cluster-name $CLUSTER_NAME \
  --principal-arn <node-role-arn> \
  --type EC2

aws eks associate-access-policy \
  --cluster-name $CLUSTER_NAME \
  --principal-arn <node-role-arn> \
  --policy-arn arn:aws:eks::aws:cluster-access-policy/AmazonEKSAutoNodePolicy \
  --access-scope type=cluster

Clean-up

To clean-up execute the following commands:

kubectl delete -f workload.yaml
kubectl delete -f neuron-nodepool.yaml
kubectl delete -f neuron-nodeclass.yaml
helm delete -n kube-system neuron-helm-chart