Add support for GPU enabled bottlerocket nodes (#298)

cookpad · Apr 12, 2022 · 7339634 · 7339634
1 parent a58487d
commit 7339634
Show file tree

Hide file tree

Showing 7 changed files with 74 additions and 6 deletions.
diff --git a/examples/cluster/bottlerocket_gpu_node_group/cluster.tf b/examples/cluster/bottlerocket_gpu_node_group/cluster.tf
@@ -0,0 +1,7 @@
+data "terraform_remote_state" "cluster" {
+  backend = "local"
+
+  config = {
+    path = "${path.module}/../terraform.tfstate"
+  }
+}
diff --git a/examples/cluster/bottlerocket_gpu_node_group/main.tf b/examples/cluster/bottlerocket_gpu_node_group/main.tf
@@ -0,0 +1,30 @@
+terraform {
+  required_providers {
+    aws = {
+      source  = "hashicorp/aws"
+      version = "3.53.0"
+    }
+  }
+}
+
+provider "aws" {
+  region              = "us-east-1"
+  allowed_account_ids = ["214219211678"]
+}
+
+module "node_group" {
+  source = "../../../modules/asg_node_group"
+
+  cluster_config = data.terraform_remote_state.cluster.outputs.cluster_config
+
+  name           = "bottlerocket-gpu-nodes"
+  key_name       = "development"
+  bottlerocket   = true
+  gpu            = true
+  instance_types = ["g4dn.xlarge"]
+  min_size       = 1
+
+  labels = {
+    "cookpad.com/terraform-aws-eks-test-environment" = data.terraform_remote_state.cluster.outputs.cluster_name
+  }
+}
diff --git a/modules/asg_node_group/README.md b/modules/asg_node_group/README.md
@@ -218,7 +218,7 @@ module "bottlerocket_nodes" {
 enable the [AWS EBS CSI driver](https://github.com/kubernetes-sigs/aws-ebs-csi-driver) by setting `aws_ebs_csi_driver = true` on the cluster module.
 see: https://github.com/bottlerocket-os/bottlerocket/blob/develop/QUICKSTART-EKS.md#csi-plugin
 
-⚠️ Bottlerocket does not yet [support GPU nodes](https://github.com/bottlerocket-os/bottlerocket/issues/769), do not set `gpu = true` when `bottlerocket = true`, as this may result in an invalid configuration!
+⚠️ Bottlerocket now [supports GPU nodes](https://github.com/bottlerocket-os/bottlerocket/blob/develop/QUICKSTART-EKS.md#aws-k8s--nvidia-variants), set `gpu = true` to enable them. Ensure that you set `instance_types` to a GPU instance type.
 
 📝 If you want to get a shell session on your instances via Bottlerocket's SSM agent
 you will need to attach the `arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore` policy

diff --git a/modules/asg_node_group/main.tf b/modules/asg_node_group/main.tf
@@ -43,6 +43,7 @@ locals {
   labels = merge(
     { "node-group.k8s.cookpad.com/name" = local.node_group_label },
     var.gpu ? { "nvidia.com/gpu" = "true" } : {},
+    var.bottlerocket ? { "bottlerocket" = "true" } : {},
     var.labels,
   )
 }
@@ -65,7 +66,7 @@ data "aws_ami" "image" {
 }
 
 data "aws_ssm_parameter" "bottlerocket_image_id" {
-  name = "/aws/service/bottlerocket/aws-k8s-${local.k8s_version}/x86_64/latest/image_id"
+  name = "/aws/service/bottlerocket/aws-k8s-${local.k8s_version}${var.gpu ? "-nvidia" : ""}/x86_64/latest/image_id"
 }
 
 data "aws_ami" "bottlerocket_image" {

diff --git a/modules/cluster/addons/helm/nvidia-device-plugin.yaml b/modules/cluster/addons/helm/nvidia-device-plugin.yaml
@@ -1,2 +1,11 @@
-nodeSelector: 
-  nvidia.com/gpu: "true"
+affinity:
+  nodeAffinity:
+    requiredDuringSchedulingIgnoredDuringExecution:
+      nodeSelectorTerms:
+      - matchExpressions:
+        - key: nvidia.com/gpu
+          operator: In
+          values:
+          - "true"
+        - key: bottlerocket
+          operator: DoesNotExist
diff --git a/modules/cluster/addons/nvidia-device-plugin.yaml b/modules/cluster/addons/nvidia-device-plugin.yaml
@@ -69,8 +69,17 @@ spec:
         - name: device-plugin
           hostPath:
             path: /var/lib/kubelet/device-plugins
-      nodeSelector:
-        nvidia.com/gpu: "true"
+      affinity:
+        nodeAffinity:
+          requiredDuringSchedulingIgnoredDuringExecution:
+            nodeSelectorTerms:
+            - matchExpressions:
+              - key: nvidia.com/gpu
+                operator: In
+                values:
+                - "true"
+              - key: bottlerocket
+                operator: DoesNotExist
       tolerations:
         - key: CriticalAddonsOnly
           operator: Exists

diff --git a/test/cluster_test.go b/test/cluster_test.go
@@ -119,6 +119,18 @@ func TestTerraformAwsEksCluster(t *testing.T) {
 		validateKubeBench(t, kubeconfig)
 		validateNodeTerminationHandler(t, kubeconfig)
 	})
+
+	test_structure.RunTestStage(t, "validate_bottlerocket_gpu_node_group", func() {
+		terraformOptions := test_structure.LoadTerraformOptions(t, workingDir)
+		kubeconfig := writeKubeconfig(t, terraform.Output(t, terraformOptions, "cluster_name"))
+		defer os.Remove(kubeconfig)
+		gpuNodeGroupDir := "../examples/cluster/bottlerocket_gpu_node_group"
+		deployTerraform(t, gpuNodeGroupDir, map[string]interface{}{})
+		defer cleanupTerraform(t, gpuNodeGroupDir)
+		validateGPUNodes(t, kubeconfig)
+		validateKubeBench(t, kubeconfig)
+		validateNodeTerminationHandler(t, kubeconfig)
+	})
 }
 
 func validateNodeLabels(t *testing.T, kubeconfig string, clusterName string) {