feat(eks): add neuron plugin in addNodegroupCapacity()

Previously, the neuron plugin was only added in addAutoScalingGroupCapacity() when the instanceType was some of the INFERENTIA types. Let's also add the plugin if at least one (some/any) of the instanceTypes in addNodegroupCapacity() is of the INFERENTIA types.
aws · Nov 9, 2023 · 0583de9 · 0583de9
1 parent 820bb99
commit 0583de9
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/packages/aws-cdk-lib/aws-eks/lib/cluster.ts b/packages/aws-cdk-lib/aws-eks/lib/cluster.ts
@@ -1800,6 +1800,9 @@ export class Cluster extends ClusterBase {
    * @param options options for creating a new nodegroup
    */
   public addNodegroupCapacity(id: string, options?: NodegroupOptions): Nodegroup {
+    if (options?.instanceTypes?.some(i => nodeTypeForInstanceType(i) === NodeType.INFERENTIA)) {
+      this.addNeuronDevicePlugin();
+    }
     return new Nodegroup(this, `Nodegroup${id}`, {
       cluster: this,
       ...options,

diff --git a/packages/aws-cdk-lib/aws-eks/test/cluster.test.ts b/packages/aws-cdk-lib/aws-eks/test/cluster.test.ts
@@ -2210,6 +2210,42 @@ describe('cluster', () => {
       });
     });
 
+    test('inf1 instances are supported in addNodegroupCapacity', () => {
+      // GIVEN
+      const { stack } = testFixtureNoVpc();
+      const cluster = new eks.Cluster(stack, 'Cluster', { defaultCapacity: 0, version: CLUSTER_VERSION, prune: false });
+
+      // WHEN
+      cluster.addNodegroupCapacity('InferenceInstances', {
+        instanceTypes: [new ec2.InstanceType('inf1.2xlarge')],
+      });
+      const fileContents = fs.readFileSync(path.join(__dirname, '../lib', 'addons/neuron-device-plugin.yaml'), 'utf8');
+      const sanitized = YAML.parse(fileContents);
+
+      // THEN
+      Template.fromStack(stack).hasResourceProperties(eks.KubernetesManifest.RESOURCE_TYPE, {
+        Manifest: JSON.stringify([sanitized]),
+      });
+    });
+    test('inf2 instances are supported in addNodegroupCapacity', () => {
+      // GIVEN
+      const { stack } = testFixtureNoVpc();
+      const cluster = new eks.Cluster(stack, 'Cluster', { defaultCapacity: 0, version: CLUSTER_VERSION, prune: false });
+
+      // WHEN
+      cluster.addNodegroupCapacity('InferenceInstances', {
+        instanceTypes: [new ec2.InstanceType('inf2.xlarge')],
+      });
+      const fileContents = fs.readFileSync(path.join(__dirname, '../lib', 'addons/neuron-device-plugin.yaml'), 'utf8');
+      const sanitized = YAML.parse(fileContents);
+
+      // THEN
+      Template.fromStack(stack).hasResourceProperties(eks.KubernetesManifest.RESOURCE_TYPE, {
+        Manifest: JSON.stringify([sanitized]),
+      });
+    });
+
+
     test('kubectl resources are always created after all fargate profiles', () => {
       // GIVEN
       const { stack, app } = testFixture();