Fix ManagedNodeGroups with custom launch templates always using x86_6…

…4 AMI (#1324)

A custom launch template is created for a Node Group when
`kubeletExtraArgs`, `bootstrapExtraArgs` or `enableIMDSv2` is set. In
that case, the recommended AMI is retrieved from SSM Parameter Store.

The input property `instanceType` is used to select the right
architecture. But `ManagedNodeGroup` doesn't have an `instanceType`
input property, it has a list of `instanceTypes`.

This change fixes this by checking the `instanceTypes` property of
Managed Node Groups as well.

Fixes #1323

examples/managed-nodegroups/index.ts

@@ -76,3 +76,12 @@ const managedNodeGroup3 = eks.createManagedNodeGroup("example-managed-ng3", {
   nodeRole: role2,
   enableIMDSv2: true,
 });
+
+// Node Group with graviton instances
+const managedNodeGroup4 = eks.createManagedNodeGroup("example-managed-ng4", {
+  cluster: cluster,
+  nodeRole: role0,
+  kubeletExtraArgs: "--max-pods=500",
+  enableIMDSv2: true,
+  instanceTypes: ["t4g.medium"],
+});

nodejs/eks/nodegroup.test.ts

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 import { isGravitonInstance } from "./nodegroup";
+import { getArchitecture } from "./nodegroup";
 
 const gravitonInstances = [
     "c6g.12xlarge",
@@ -311,4 +312,32 @@ describe("isGravitonInstance", () => {
             expect(isGravitonInstance(instanceType)).toBe(false);
         });
     });
+    describe("getArchitecture", () => {
+        test("should return 'x86_64' when only x86_64 instances are provided", () => {
+            const instanceTypes = ["c5.large", "m5.large", "t3.large"];
+            const architecture = getArchitecture(instanceTypes);
+            expect(architecture).toBe("x86_64");
+        });
+
+        test("should return 'arm64' when only arm64 instances are provided", () => {
+            const instanceTypes = ["c6g.large", "m6g.large", "t4g.large"];
+            const architecture = getArchitecture(instanceTypes);
+            expect(architecture).toBe("arm64");
+        });
+
+        test("should throw an error when both x86_64 and arm64 instances are provided", () => {
+            const instanceTypes = ["c5.large", "c6g.large"];
+            expect(() => {
+                getArchitecture(instanceTypes);
+            }).toThrowError(
+                "Cannot determine architecture of instance types. The provided instance types do not share a common architecture",
+            );
+        });
+
+        test("should return 'x86_64' when no instance types are provided", () => {
+            const instanceTypes: string[] = [];
+            const architecture = getArchitecture(instanceTypes);
+            expect(architecture).toBe("x86_64");
+        });
+    });
 });

nodejs/eks/nodegroup.ts

@@ -1844,15 +1844,17 @@ Content-Type: text/x-shellscript; charset="us-ascii"
         ? { httpTokens: "required", httpPutResponseHopLimit: 2, httpEndpoint: "enabled" }
         : undefined;
 
-    const blockDeviceMappings = args.diskSize ? [
-        {
-            // /dev/xvda is the default device name for the root volume on an Amazon Linux 2 & AL2023 instance.
-            deviceName: "/dev/xvda",
-            ebs: {
-                volumeSize: args.diskSize
-            },
-        },
-    ] : undefined;
+    const blockDeviceMappings = args.diskSize
+        ? [
+              {
+                  // /dev/xvda is the default device name for the root volume on an Amazon Linux 2 & AL2023 instance.
+                  deviceName: "/dev/xvda",
+                  ebs: {
+                      volumeSize: args.diskSize,
+                  },
+              },
+          ]
+        : undefined;
 
     return new aws.ec2.LaunchTemplate(
         `${name}-launchTemplate`,
@@ -1883,9 +1885,15 @@ function getRecommendedAMI(
     parent: pulumi.Resource | undefined,
 ): pulumi.Input<string> {
     const gpu = "gpu" in args ? args.gpu : undefined;
-    const instanceType = "instanceType" in args ? args.instanceType : undefined;
 
-    const amiType = getAMIType(args.amiType, gpu, instanceType);
+    let instanceTypes: pulumi.Input<pulumi.Input<string>[]> | undefined;
+    if ("instanceType" in args && args.instanceType) {
+        instanceTypes = [args.instanceType];
+    } else if ("instanceTypes" in args) {
+        instanceTypes = args.instanceTypes;
+    }
+
+    const amiType = getAMIType(args.amiType, gpu, instanceTypes);
     // if specified use the version from the args, otherwise use the version from the cluster.
     const version = args.version ? args.version : k8sVersion;
 
@@ -1930,14 +1938,20 @@ export function isGravitonInstance(instanceType: string): boolean {
 
 /**
  * getAMIType returns the AMI type to use for the given configuration based on the
- * architecutre of the instance type.
+ * architecture of the instance type.
  */
 function getAMIType(
     amiType: pulumi.Input<string> | undefined,
     gpu: pulumi.Input<boolean> | undefined,
-    instanceType: pulumi.Input<string> | undefined,
+    instanceTypes: pulumi.Input<pulumi.Input<string>[]> | undefined,
 ): pulumi.Output<string> {
-    return pulumi.all([amiType, gpu, instanceType]).apply(([amiType, gpu, instanceType]) => {
+    const architecture = pulumi.output(instanceTypes).apply((instanceTypes) => {
+        return pulumi
+            .all(instanceTypes ?? [])
+            .apply((instanceTypes) => getArchitecture(instanceTypes));
+    });
+
+    return pulumi.all([amiType, gpu, architecture]).apply(([amiType, gpu, architecture]) => {
         if (amiType) {
             // Return the user-specified AMI type.
             return amiType;
@@ -1948,11 +1962,42 @@ function getAMIType(
             return "amazon-linux-2-gpu";
         }
 
-        if (instanceType && isGravitonInstance(instanceType)) {
+        if (architecture === "arm64") {
             // Return the Amazon Linux 2 ARM64 AMI type.
             return "amazon-linux-2-arm64";
         }
 
         return "amazon-linux-2";
     });
 }
+
+type NodeArchitecture = "arm64" | "x86_64";
+
+/**
+ * Determines the architecture based on the provided instance types. Defaults to "x86_64" if no instance types are provided.
+ *
+ * @param instanceTypes - An array of instance types.
+ * @returns The architecture of the instance types, either "arm64" or "x86_64".
+ * @throws {pulumi.ResourceError} If the provided instance types do not share a common architecture.
+ */
+export function getArchitecture(instanceTypes: string[]): NodeArchitecture {
+    let hasGravitonInstances = false;
+    let hasX64Instances = false;
+
+    instanceTypes.forEach((instanceType) => {
+        if (isGravitonInstance(instanceType)) {
+            hasGravitonInstances = true;
+        } else {
+            hasX64Instances = true;
+        }
+
+        if (hasGravitonInstances && hasX64Instances) {
+            throw new pulumi.ResourceError(
+                "Cannot determine architecture of instance types. The provided instance types do not share a common architecture",
+                undefined,
+            );
+        }
+    });
+
+    return hasGravitonInstances ? "arm64" : "x86_64";
+}

nodejs/eks/tsconfig.json

@@ -37,6 +37,6 @@
 
         "authenticationMode.test.ts",
         "dependencies.test.ts",
-        "nodegroup.test.ts",
+        "nodegroup.test.ts"
     ]
 }