chore(eks): add isGpuInstanceType tests

aws · Oct 11, 2024 · 7f82fd5 · 7f82fd5
1 parent cf48414
commit 7f82fd5
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 5 deletions.
diff --git a/packages/aws-cdk-lib/aws-ec2/lib/instance-types.ts b/packages/aws-cdk-lib/aws-ec2/lib/instance-types.ts
@@ -1599,7 +1599,6 @@ export class InstanceType {
       [InstanceClass.GRAPHICS5_GRAVITON2]: 'g5g',
       [InstanceClass.G5G]: 'g5g',
       [InstanceClass.GRAPHICS6]: 'g6',
-      [InstanceClass.G6E]: 'g6e',
       [InstanceClass.G6]: 'g6',
       [InstanceClass.GRAPHICS6_EFFICIENT]: 'g6e',
       [InstanceClass.G6E]: 'g6e',

diff --git a/packages/aws-cdk-lib/aws-eks/lib/managed-nodegroup.ts b/packages/aws-cdk-lib/aws-eks/lib/managed-nodegroup.ts
@@ -602,7 +602,7 @@ const gpuAmiTypes: NodegroupAmiType[] = [
  * This function check if the instanceType is GPU instance.
  * @param instanceType The EC2 instance type
  */
-function isGpuInstanceType(instanceType: InstanceType): boolean {
+export function isGpuInstanceType(instanceType: InstanceType): boolean {
   //compare instanceType to known GPU InstanceTypes
   const knownGpuInstanceTypes = [InstanceClass.P2, InstanceClass.P3, InstanceClass.P3DN, InstanceClass.P4DE, InstanceClass.P4D,
     InstanceClass.G3S, InstanceClass.G3, InstanceClass.G4DN, InstanceClass.G4AD, InstanceClass.G5, InstanceClass.G5G, InstanceClass.G6,

diff --git a/packages/aws-cdk-lib/aws-eks/test/nodegroup.test.ts b/packages/aws-cdk-lib/aws-eks/test/nodegroup.test.ts
@@ -6,7 +6,7 @@ import * as iam from '../../aws-iam';
 import * as cdk from '../../core';
 import * as cxapi from '../../cx-api';
 import * as eks from '../lib';
-import { NodegroupAmiType, TaintEffect } from '../lib';
+import { isGpuInstanceType, NodegroupAmiType, TaintEffect } from '../lib';
 
 /* eslint-disable max-len */
 
@@ -617,8 +617,8 @@ describe('node group', () => {
     new eks.Nodegroup(stack, 'Nodegroup', {
       cluster,
       instanceTypes: [
-        new ec2.InstanceType('p3.large'),
-        new ec2.InstanceType('g3.large'),
+        new ec2.InstanceType('g6e.large'),
+        new ec2.InstanceType('g5.large'),
       ],
     });
 
@@ -1735,3 +1735,38 @@ describe('node group', () => {
     expect(() => cluster.addNodegroupCapacity('ng', { maxUnavailablePercentage: 101 })).toThrow(/maxUnavailablePercentage must be between 1 and 100/);
   });
 });
+
+describe('isGpuInstanceType', () => {
+  it('should return true for known GPU instance types', () => {
+    const gpuInstanceTypes = [
+      ec2.InstanceType.of(ec2.InstanceClass.P2, ec2.InstanceSize.XLARGE),
+      ec2.InstanceType.of(ec2.InstanceClass.G3, ec2.InstanceSize.XLARGE),
+      ec2.InstanceType.of(ec2.InstanceClass.P4D, ec2.InstanceSize.LARGE),
+      ec2.InstanceType.of(ec2.InstanceClass.G6, ec2.InstanceSize.MEDIUM),
+      ec2.InstanceType.of(ec2.InstanceClass.G6E, ec2.InstanceSize.XLARGE2),
+    ];
+    gpuInstanceTypes.forEach(instanceType => {
+      expect(isGpuInstanceType(instanceType)).toBe(true);
+    });
+  });
+  it('should return false for non-GPU instance types', () => {
+    const nonGpuInstanceTypes = [
+      ec2.InstanceType.of(ec2.InstanceClass.T3, ec2.InstanceSize.MICRO),
+      ec2.InstanceType.of(ec2.InstanceClass.M5, ec2.InstanceSize.LARGE),
+      ec2.InstanceType.of(ec2.InstanceClass.C5, ec2.InstanceSize.XLARGE),
+    ];
+    nonGpuInstanceTypes.forEach(instanceType => {
+      expect(isGpuInstanceType(instanceType)).toBe(false);
+    });
+  });
+  it('should return true for different sizes of GPU instance types', () => {
+    const gpuInstanceTypes = [
+      ec2.InstanceType.of(ec2.InstanceClass.G6, ec2.InstanceSize.XLARGE),
+      ec2.InstanceType.of(ec2.InstanceClass.G6, ec2.InstanceSize.XLARGE16),
+      ec2.InstanceType.of(ec2.InstanceClass.G6, ec2.InstanceSize.XLARGE48),
+    ];
+    gpuInstanceTypes.forEach(instanceType => {
+      expect(isGpuInstanceType(instanceType)).toBe(true);
+    });
+  });
+});