Services or capabilities described in Amazon Web Services documentation might vary by Region. To see the differences applicable to the China Regions, see Getting Started with Amazon Web Services in China (PDF).

Help improve this page

To contribute to this user guide, choose the Edit this page on GitHub link that is located in the right pane of every page.

Manage Neuron devices on Amazon EKS

Amazon Trainium and Amazon Inferentia are purpose-built machine learning chips designed by Amazon. Amazon EKS supports two mechanisms for managing Neuron devices in EKS clusters: the Neuron DRA driver and the Neuron Kubernetes device plugin.

It’s recommended to use the Neuron DRA driver for new deployments on EKS clusters running Kubernetes version 1.34 or later. The Neuron DRA driver provides topology-aware allocation, connected device subset scheduling, Logical NeuronCore (LNC) configuration, and UltraServer multi-node allocation without requiring custom scheduler extensions. The Neuron device plugin remains supported.

Neuron DRA driver vs. Neuron device plugin

Install the Neuron DRA driver

The Neuron DRA driver advertises Neuron devices as ResourceSlice objects with the DeviceClass name neuron.aws.com. The driver runs as a DaemonSet and automatically discovers Neuron devices and their topology attributes.

Detailed information about the Neuron DRA driver is available in the Neuron DRA documentation.

Using the Neuron DRA driver with Bottlerocket is not currently supported.

Prerequisites

Procedure

See the Neuron DRA documentation for information on the available ResourceSlice object attributes.

Request Neuron devices in a Pod

To request Neuron devices using the DRA driver, create a ResourceClaimTemplate that references the neuron.aws.com DeviceClass and reference it in your Pod specification.

The following example requests all Neuron devices on a trn2.48xlarge instance:

apiVersion: resource.k8s.io/v1
kind: ResourceClaimTemplate
metadata:
  name: all-neurons
spec:
  spec:
    devices:
      requests:
      - name: neurons
        exactly:
          deviceClassName: neuron.aws.com
          selectors:
          - cel:
              expression: "device.attributes['neuron.aws.com'].instanceType == 'trn2.48xlarge'"
          allocationMode: All
---
apiVersion: v1
kind: Pod
metadata:
  name: neuron-workload
spec:
  containers:
  - name: app
    ...
    resources:
      claims:
      - name: neurons
  resourceClaims:
  - name: neurons
    resourceClaimTemplateName: all-neurons

Allocate connected device subsets

The Neuron DRA driver can allocate subsets of connected Neuron devices without requiring the Neuron scheduler extension. Supported subset sizes are 1, 4, 8, or 16 devices. Use the matchAttribute constraint with a topology group ID to ensure devices are connected.

The following example requests 4 connected Neuron devices:

apiVersion: resource.k8s.io/v1
kind: ResourceClaimTemplate
metadata:
  name: 1x4-connected-neurons
spec:
  spec:
    devices:
      requests:
      - name: neurons
        exactly:
          deviceClassName: neuron.aws.com
          allocationMode: ExactCount
          count: 4
          selectors:
          - cel:
              expression: "device.attributes['neuron.aws.com'].instanceType == 'trn2.48xlarge'"
      constraints:
      - requests: ["neurons"]
        matchAttribute: "resource.aws.com/devicegroup4_id"

The supported matchAttribute values for connected subsets are resource.aws.com/devicegroup1_id, resource.aws.com/devicegroup4_id, resource.aws.com/devicegroup8_id, and resource.aws.com/devicegroup16_id.

Configure Logical NeuronCores (LNC)

The Neuron DRA driver allows per-workload Logical NeuronCore configuration through ResourceClaimTemplate parameters. This eliminates the need to pre-configure LNC in EC2 Launch Templates.

The following example requests all Neuron devices with LNC set to 1:

apiVersion: resource.k8s.io/v1
kind: ResourceClaimTemplate
metadata:
  name: all-neurons-lnc-1
spec:
  spec:
    devices:
      requests:
      - name: neurons
        exactly:
          deviceClassName: neuron.aws.com
          selectors:
          - cel:
              expression: "device.attributes['neuron.aws.com'].instanceType == 'trn2.48xlarge'"
          allocationMode: All
      config:
      - requests: ["neurons"]
        opaque:
          driver: neuron.aws.com
          parameters:
            apiVersion: neuron.aws.com/v1
            kind: NeuronConfig
            logicalNeuronCore: 1

Install the Neuron Kubernetes device plugin

The Neuron Kubernetes device plugin advertises Neuron devices as aws.amazon.com/neuron and NeuronCores as aws.amazon.com/neuroncore extended resources. You request Neuron devices in container resource requests and limits.

Prerequisites

Procedure

Verify Neuron devices with a test Pod

You can verify that Neuron devices are accessible by running the neuron-ls tool in a test Pod.

For more information about using Neuron devices with Amazon EKS, see the Neuron documentation for running on EKS.