EKS: Determine ami_type for node pool based on instance_type

Author: pst

aws/_modules/eks/node_pool.tf

@@ -20,6 +20,8 @@ module "node_pool" {
 
   taints = var.taints
 
+  ami_type = null
+
   # force node_pool to depend on aws-auth configmap
   depends-on-aws-auth = {
     name      = kubernetes_config_map.current.metadata[0].name

aws/_modules/eks/node_pool/main.tf

@@ -1,3 +1,13 @@
+data "aws_ec2_instance_type" "current" {
+  # ami_type is always determined by the first instance_type in the list
+  instance_type = element(tolist(var.instance_types), 0)
+}
+
+locals {
+  cpu_ami_type = data.aws_ec2_instance_type.current.supported_architectures[0] == "arm64" ? "AL2_ARM_64" : "AL2_x86_64"
+  ami_type     = length(data.aws_ec2_instance_type.current.gpus) > 0 ? "AL2_x86_64_GPU" : local.cpu_ami_type
+}
+
 resource "aws_eks_node_group" "nodes" {
   cluster_name    = var.cluster_name
   node_group_name = var.node_group_name
@@ -10,6 +20,7 @@ resource "aws_eks_node_group" "nodes" {
     min_size     = var.min_size
   }
 
+  ami_type       = var.ami_type == null ? local.ami_type : var.ami_type
   instance_types = var.instance_types
   disk_size      = var.disk_size
 

aws/_modules/eks/node_pool/variables.tf

@@ -79,3 +79,8 @@ variable "labels" {
   description = "Kubernetes labels to set on the nodes created by the node pool. Merged with Kubestack default labels."
   default     = {}
 }
+
+variable "ami_type" {
+  type        = string
+  description = "AMI type to use for nodes of the node pool."
+}

aws/cluster/node-pool/configuration.tf

@@ -16,6 +16,8 @@ locals {
   max_size              = lookup(local.cfg, "max_size")
   disk_size             = lookup(local.cfg, "disk_size", null)
 
+  ami_type = lookup(local.cfg, "ami_type")
+
   availability_zones_lookup = local.cfg["availability_zones"] == null ? "" : local.cfg["availability_zones"]
   availability_zones        = compact(split(",", local.availability_zones_lookup))
 

aws/cluster/node-pool/main.tf

@@ -15,6 +15,8 @@ module "node_pool" {
   max_size       = local.max_size
   min_size       = local.min_size
 
+  ami_type = local.ami_type
+
   disk_size = local.disk_size
 
   taints = local.taints

aws/cluster/node-pool/variables.tf

@@ -9,6 +9,8 @@ variable "configuration" {
     max_size         = optional(string)
     disk_size        = optional(string)
 
+    ami_type = optional(string)
+
     availability_zones = optional(string)
 
     vpc_subnet_ids = optional(string)

tests/eks_zero_node_pools.tf

@@ -54,6 +54,8 @@ module "eks_zero_node_pool_existing_subnets_one_az_only" {
       desired_capacity = 1
       min_size         = 1
       max_size         = 3
+
+      ami_type = "AL2_x86_64_GPU"
     }
 
     # Settings for Ops-cluster