PLAT-7660 Adds default gpu labels/taints (#183)

dominodatalab · Nov 29, 2023 · b6ff53e · b6ff53e
1 parent 9f76fe5
commit b6ff53e
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 1 deletion.
diff --git a/examples/tfvars/custom-node-pool.tfvars b/examples/tfvars/custom-node-pool.tfvars
@@ -34,4 +34,68 @@ additional_node_groups = {
       type = "gp3"
     }
   }
+  custom-gpu-1 = {
+    gpu = true
+    ## Just for testing the `gpu` bool behavior
+    instance_types = [
+      "m5.2xlarge"
+    ],
+    min_per_az     = 0,
+    max_per_az     = 10,
+    desired_per_az = 0,
+    availability_zone_ids = [
+      "usw2-az1",
+      "usw2-az2"
+    ],
+    labels = {
+      "dominodatalab.com/node-pool" = "custom-gpu-1 "
+    },
+    volume = {
+      size = 100,
+      type = "gp3"
+    }
+  }
+  custom-gpu-2 = {
+    instance_types = [
+      "g4dn.xlarge"
+    ],
+    min_per_az     = 0,
+    max_per_az     = 10,
+    desired_per_az = 0,
+    availability_zone_ids = [
+      "usw2-az1",
+      "usw2-az2"
+    ],
+    labels = {
+      "dominodatalab.com/node-pool" = "custom-gpu-2 "
+    },
+    volume = {
+      size = 100,
+      type = "gp3"
+    }
+  }
+  custom-gpu-3 = {
+    instance_types = [
+      "g4dn.xlarge"
+    ],
+    min_per_az     = 0,
+    max_per_az     = 10,
+    desired_per_az = 0,
+    availability_zone_ids = [
+      "usw2-az1",
+      "usw2-az2"
+    ],
+    labels = {
+      "dominodatalab.com/node-pool" = "custom-gpu-3 "
+    },
+    volume = {
+      size = 100,
+      type = "gp3"
+    }
+    taints = [{
+      key    = "nvidia.com/gpu"
+      value  = "true"
+      effect = "NO_SCHEDULE"
+    }]
+  }
 }
diff --git a/modules/nodes/main.tf b/modules/nodes/main.tf
@@ -6,12 +6,20 @@ data "aws_ec2_instance_type" "all" {
 }
 
 locals {
+  gpu_labels = { "nvidia.com/gpu" = true }
+  gpu_taints = [{
+    key    = "nvidia.com/gpu"
+    value  = "true"
+    effect = "NO_SCHEDULE"
+  }]
   node_groups = {
     for name, ng in
     merge(var.additional_node_groups, var.default_node_groups) :
     name => merge(ng, {
-      gpu           = ng.gpu != null ? ng.gpu : anytrue([for itype in ng.instance_types : length(data.aws_ec2_instance_type.all[itype].gpus) > 0]),
+      gpu           = coalesce(ng.gpu, false) || anytrue([for itype in ng.instance_types : length(data.aws_ec2_instance_type.all[itype].gpus) > 0])
       instance_tags = merge(data.aws_default_tags.this.tags, ng.tags)
+      labels        = coalesce(ng.gpu, false) || anytrue([for itype in ng.instance_types : length(data.aws_ec2_instance_type.all[itype].gpus) > 0]) ? merge(local.gpu_labels, ng.labels) : ng.labels
+      taints        = coalesce(ng.gpu, false) || anytrue([for itype in ng.instance_types : length(data.aws_ec2_instance_type.all[itype].gpus) > 0]) ? distinct(concat(local.gpu_taints, ng.taints)) : ng.taints
     })
   }
 }