Migrate AWS EKS AMI from Amazon Linux 2 to Amazon Linux 2023 (#3166)

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Vinicius D. Cerutti <51954708+viniciusdc@users.noreply.github.com>

Author: 3 people

src/_nebari/stages/infrastructure/__init__.py

@@ -169,8 +169,8 @@ class AzureInputVars(schema.Base):
 
 
 class AWSAmiTypes(str, enum.Enum):
-    AL2_x86_64 = "AL2_x86_64"
-    AL2_x86_64_GPU = "AL2_x86_64_GPU"
+    AL2023_x86_64_STANDARD = "AL2023_x86_64_STANDARD"
+    AL2023_x86_64_NVIDIA = "AL2023_x86_64_NVIDIA"
     CUSTOM = "CUSTOM"
 
 
@@ -219,19 +219,19 @@ def construct_aws_ami_type(
 
     Returns the AMI type (str) determined by the following rules:
         - Returns "CUSTOM" if a `launch_template` is provided and it includes a valid `ami_id`.
-        - Returns "AL2_x86_64_GPU" if `gpu_enabled` is True and no valid
+        - Returns "AL2023_x86_64_NVIDIA" if `gpu_enabled` is True and no valid
           `launch_template` is provided (None).
-        - Returns "AL2_x86_64" as the default AMI type if `gpu_enabled` is False and no
+        - Returns "AL2023_x86_64_STANDARD" as the default AMI type if `gpu_enabled` is False and no
           valid `launch_template` is provided (None).
     """
 
     if launch_template and getattr(launch_template, "ami_id", None):
         return "CUSTOM"
 
     if gpu_enabled:
-        return "AL2_x86_64_GPU"
+        return "AL2023_x86_64_NVIDIA"
 
-    return "AL2_x86_64"
+    return "AL2023_x86_64_STANDARD"
 
 
 class AWSInputVars(schema.Base):

src/_nebari/stages/infrastructure/template/aws/main.tf

@@ -87,6 +87,7 @@ module "kubernetes" {
   tags               = local.additional_tags
   region             = var.region
   kubernetes_version = var.kubernetes_version
+  environment        = var.environment
 
   cluster_subnets         = local.subnet_ids
   cluster_security_groups = [local.security_group_id]

src/_nebari/stages/infrastructure/template/aws/modules/kubernetes/autoscaling.tf

@@ -14,7 +14,11 @@ data "aws_iam_policy_document" "worker_autoscaling" {
       "autoscaling:DescribeAutoScalingInstances",
       "autoscaling:DescribeLaunchConfigurations",
       "autoscaling:DescribeTags",
+      "ec2:DescribeImages",
+      "ec2:DescribeInstanceTypes",
       "ec2:DescribeLaunchTemplateVersions",
+      "ec2:GetInstanceTypesFromInstanceRequirements",
+      "eks:DescribeNodegroup",
     ]
 
     resources = ["*"]

src/_nebari/stages/infrastructure/template/aws/modules/kubernetes/locals.tf

@@ -8,8 +8,7 @@ locals {
   node_group_policies = concat([
     "arn:${local.partition}:iam::aws:policy/AmazonEKSWorkerNodePolicy",
     "arn:${local.partition}:iam::aws:policy/AmazonEKS_CNI_Policy",
-    "arn:${local.partition}:iam::aws:policy/service-role/AmazonEBSCSIDriverPolicy",
-    aws_iam_policy.worker_autoscaling.arn
+    "arn:${local.partition}:iam::aws:policy/service-role/AmazonEBSCSIDriverPolicy"
   ], var.node_group_additional_policies)
 
   gpu_node_group_names = [for node_group in var.node_groups : node_group.name if node_group.gpu == true]

src/_nebari/stages/infrastructure/template/aws/modules/kubernetes/main.tf

@@ -206,3 +206,71 @@ resource "aws_iam_openid_connect_provider" "oidc_provider" {
     var.tags
   )
 }
+
+# IAM role for EBS CSI driver using IRSA
+resource "aws_iam_role" "ebs_csi_driver" {
+  name = "${var.name}-ebs-csi-driver"
+
+  # Trust policy - allows the Kubernetes service account to assume this role via OIDC
+  assume_role_policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [{
+      Effect = "Allow"
+      Principal = {
+        Federated = aws_iam_openid_connect_provider.oidc_provider.arn
+      }
+      Action = "sts:AssumeRoleWithWebIdentity"
+      Condition = {
+        StringEquals = {
+          "${replace(aws_eks_cluster.main.identity[0].oidc[0].issuer, "https://", "")}:sub" = "system:serviceaccount:kube-system:ebs-csi-controller-sa"
+          "${replace(aws_eks_cluster.main.identity[0].oidc[0].issuer, "https://", "")}:aud" = "sts.amazonaws.com"
+        }
+      }
+    }]
+  })
+
+  tags = merge(
+    { Name = "${var.name}-ebs-csi-driver" },
+    var.tags
+  )
+}
+
+# Attach the AWS managed policy for EBS CSI driver
+resource "aws_iam_role_policy_attachment" "ebs_csi_driver" {
+  role       = aws_iam_role.ebs_csi_driver.name
+  policy_arn = "arn:aws:iam::aws:policy/service-role/AmazonEBSCSIDriverPolicy"
+}
+
+# IAM role for Cluster Autoscaler using IRSA
+resource "aws_iam_role" "cluster_autoscaler" {
+  name = "${var.name}-cluster-autoscaler"
+
+  # Trust policy - allows the Kubernetes service account to assume this role via OIDC
+  assume_role_policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [{
+      Effect = "Allow"
+      Principal = {
+        Federated = aws_iam_openid_connect_provider.oidc_provider.arn
+      }
+      Action = "sts:AssumeRoleWithWebIdentity"
+      Condition = {
+        StringEquals = {
+          "${replace(aws_eks_cluster.main.identity[0].oidc[0].issuer, "https://", "")}:sub" = "system:serviceaccount:${var.environment}:cluster-autoscaler"
+          "${replace(aws_eks_cluster.main.identity[0].oidc[0].issuer, "https://", "")}:aud" = "sts.amazonaws.com"
+        }
+      }
+    }]
+  })
+
+  tags = merge(
+    { Name = "${var.name}-cluster-autoscaler" },
+    var.tags
+  )
+}
+
+# Attach the autoscaling policy to Cluster Autoscaler role
+resource "aws_iam_role_policy_attachment" "cluster_autoscaler" {
+  role       = aws_iam_role.cluster_autoscaler.name
+  policy_arn = aws_iam_policy.worker_autoscaling.arn
+}

src/_nebari/stages/infrastructure/template/aws/modules/kubernetes/outputs.tf

@@ -23,6 +23,11 @@ output "oidc_provider_arn" {
   value       = aws_iam_openid_connect_provider.oidc_provider.arn
 }
 
+output "cluster_autoscaler_role_arn" {
+  description = "IAM role ARN for Cluster Autoscaler (IRSA)"
+  value       = aws_iam_role.cluster_autoscaler.arn
+}
+
 # https://github.com/terraform-aws-modules/terraform-aws-eks/blob/16f46db94b7158fd762d9133119206aaa7cf6d63/examples/self_managed_node_group/main.tf
 output "kubeconfig" {
   description = "Kubernetes connection configuration kubeconfig"

src/_nebari/stages/infrastructure/template/aws/modules/kubernetes/variables.tf

@@ -94,3 +94,8 @@ variable "permissions_boundary" {
   type        = string
   default     = null
 }
+
+variable "environment" {
+  description = "Namespace/environment for Kubernetes resources (used in IRSA trust policies)"
+  type        = string
+}

src/_nebari/stages/infrastructure/template/aws/outputs.tf

@@ -34,3 +34,8 @@ output "oidc_provider_arn" {
   description = "The ARN of the OIDC Provider"
   value       = module.kubernetes.oidc_provider_arn
 }
+
+output "cluster_autoscaler_role_arn" {
+  description = "IAM role ARN for Cluster Autoscaler (IRSA)"
+  value       = module.kubernetes.cluster_autoscaler_role_arn
+}

src/_nebari/stages/infrastructure/template/aws/versions.tf

@@ -2,7 +2,7 @@ terraform {
   required_providers {
     aws = {
       source  = "hashicorp/aws"
-      version = "5.33.0"
+      version = "6.18.0"
     }
   }
   required_version = ">= 1.0"

src/_nebari/stages/kubernetes_initialize/__init__.py

@@ -45,6 +45,7 @@ class InputVars(schema.Base):
     external_container_reg: Optional[ExtContainerReg] = None
     gpu_enabled: bool = False
     gpu_node_group_names: List[str] = []
+    cluster_autoscaler_role_arn: Optional[str] = None
 
 
 class InputSchema(schema.Base):
@@ -94,6 +95,13 @@ def input_vars(self, stage_outputs: Dict[str, Dict[str, Any]]):
                 if self.config.amazon_web_services.node_groups[group].gpu
             ]
             input_vars.aws_region = self.config.amazon_web_services.region
+            # Get the Cluster Autoscaler IAM role ARN from infrastructure stage output
+            if "stages/02-infrastructure" in stage_outputs:
+                input_vars.cluster_autoscaler_role_arn = (
+                    stage_outputs["stages/02-infrastructure"]
+                    .get("cluster_autoscaler_role_arn", {})
+                    .get("value", "")
+                )
 
         return input_vars.model_dump()