diff --git a/packages/aws-cdk-lib/aws-eks/README.md b/packages/aws-cdk-lib/aws-eks/README.md index 707b298bf560c..dc36eaed4c9b3 100644 --- a/packages/aws-cdk-lib/aws-eks/README.md +++ b/packages/aws-cdk-lib/aws-eks/README.md @@ -282,7 +282,7 @@ cluster.addNodegroupCapacity('custom-node-group', { }); ``` -> **NOTE:** If you add instances with the inferentia class (`inf1` or `inf2`) or trainium class (`trn1` or `trn1n`) +> **NOTE:** If you add instances with the inferentia class (`inf1` or `inf2`) or trainium class (`trn1`, `trn1n`, or `trn2`) > the [neuron plugin](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/containers/dlc-then-eks-devflow.html) > will be automatically installed in the kubernetes cluster. diff --git a/packages/aws-cdk-lib/aws-eks/lib/private/nodegroup.ts b/packages/aws-cdk-lib/aws-eks/lib/private/nodegroup.ts index 633537bb33c00..e463761e6284b 100644 --- a/packages/aws-cdk-lib/aws-eks/lib/private/nodegroup.ts +++ b/packages/aws-cdk-lib/aws-eks/lib/private/nodegroup.ts @@ -5,8 +5,15 @@ import { InstanceClass, InstanceSize, InstanceType } from '../../../aws-ec2'; */ export function isGpuInstanceType(instanceType: InstanceType): boolean { // compare instanceType to known GPU InstanceTypes + // P-series: NVIDIA GPU instances (P2, P3, P4, P5 with H100/H200 GPUs) + // G-series: NVIDIA GPU instances for graphics and ML inference + // INF-series: AWS Inferentia chips for ML inference + // TRN-series: AWS Trainium chips for ML training and inference + // See: https://aws.amazon.com/ec2/instance-types/p5/ + // See: https://aws.amazon.com/ec2/instance-types/trn2/ const knownGpuInstanceTypes = [InstanceClass.P2, InstanceClass.P3, InstanceClass.P3DN, InstanceClass.P4DE, InstanceClass.P4D, + InstanceClass.P5, InstanceClass.P5E, InstanceClass.P5EN, InstanceClass.G3S, InstanceClass.G3, InstanceClass.G4DN, InstanceClass.G4AD, InstanceClass.G5, InstanceClass.G5G, InstanceClass.G6, - InstanceClass.G6E, InstanceClass.INF1, InstanceClass.INF2]; + InstanceClass.G6E, InstanceClass.INF1, InstanceClass.INF2, InstanceClass.TRN1, InstanceClass.TRN1N, InstanceClass.TRN2]; return knownGpuInstanceTypes.some((c) => instanceType.sameInstanceClassAs(InstanceType.of(c, InstanceSize.LARGE))); } diff --git a/packages/aws-cdk-lib/aws-eks/test/nodegroup.test.ts b/packages/aws-cdk-lib/aws-eks/test/nodegroup.test.ts index 8572a9b008f29..c190ee2354a8a 100644 --- a/packages/aws-cdk-lib/aws-eks/test/nodegroup.test.ts +++ b/packages/aws-cdk-lib/aws-eks/test/nodegroup.test.ts @@ -1841,6 +1841,12 @@ describe('isGpuInstanceType', () => { ec2.InstanceType.of(ec2.InstanceClass.G3S, ec2.InstanceSize.XLARGE), ec2.InstanceType.of(ec2.InstanceClass.G5, ec2.InstanceSize.XLARGE), ec2.InstanceType.of(ec2.InstanceClass.G5G, ec2.InstanceSize.XLARGE), + ec2.InstanceType.of(ec2.InstanceClass.TRN1, ec2.InstanceSize.XLARGE2), + ec2.InstanceType.of(ec2.InstanceClass.TRN1N, ec2.InstanceSize.XLARGE32), + ec2.InstanceType.of(ec2.InstanceClass.TRN2, ec2.InstanceSize.XLARGE48), + ec2.InstanceType.of(ec2.InstanceClass.P5, ec2.InstanceSize.XLARGE48), + ec2.InstanceType.of(ec2.InstanceClass.P5E, ec2.InstanceSize.XLARGE48), + ec2.InstanceType.of(ec2.InstanceClass.P5EN, ec2.InstanceSize.XLARGE48), ]; gpuInstanceTypes.forEach(instanceType => { expect(isGpuInstanceType(instanceType)).toBe(true);