fix helm chart for ray integration

Author: Yevet

applications/hcc/modules/nemo/main.tf

@@ -99,6 +99,18 @@ resource "helm_release" "kuberay_operator" {
   version    = "1.1.0"
   namespace = "default"
   reset_values = true
+  set {
+    name  = "tolerations[0].key"
+    value = "nvidia.com/gpu"
+  }
+  set {
+    name  = "tolerations[0].operator"
+    value = "Exists"
+  }
+  set {
+    name  = "tolerations[0].effect"
+    value = "NoSchedule"
+  }
 }
 
 resource "helm_release" "raycluster" {
@@ -108,7 +120,44 @@ resource "helm_release" "raycluster" {
   chart      = "ray-cluster"           
   version    = "1.1.0"
   namespace  = "default"
+  reset_values = true
   depends_on = [
     helm_release.kuberay_operator
   ]
+  set {
+    name  = "tolerations[0].key"
+    value = "nvidia.com/gpu"
+  }
+  set {
+    name  = "tolerations[0].operator"
+    value = "Exists"
+  }
+  set {
+    name  = "tolerations[0].effect"
+    value = "NoSchedule"
+  }
+  set {
+    name  = "head.tolerations[0].key"
+    value = "nvidia.com/gpu"
+  }
+  set {
+    name  = "head.tolerations[0].operator"
+    value = "Exists"
+  }
+  set {
+    name  = "head.tolerations[0].effect"
+    value = "NoSchedule"
+  }
+  set {
+    name  = "worker.tolerations[0].key"
+    value = "nvidia.com/gpu"
+  }
+  set {
+    name  = "worker.tolerations[0].operator"
+    value = "Exists"
+  }
+  set {
+    name  = "worker.tolerations[0].effect"
+    value = "NoSchedule"
+  }
 }

applications/hcc/variables.tf

@@ -118,14 +118,14 @@ variable "a3_mega_consumption_model" {
 }
 
 locals {
-  placement_policy_valid = var.gpu_type != "A3 Mega" || local.recipe == "gke" || length(var.placement_policy_name) > 0
+  placement_policy_valid = var.gpu_type != "A3 Mega" || local.recipe == "gke" || local.recipe == "gke-ray" || length(var.placement_policy_name) > 0
   a3_consumption_model_check = length(var.a3_ultra_consumption_model) > 0 || length(var.a3_mega_consumption_model) > 0
   recipe = {
     "A3 Mega" = var.a3mega_recipe
     "A3 Ultra" = var.a3ultra_recipe
   }[var.gpu_type]
   recipes_not_empty = length(local.recipe) > 0
-  reservation_valid = !(local.recipe != "gke") || length(var.reservation) > 0
+  reservation_valid = local.recipe == "gke" || local.recipe == "gke-ray" || length(var.reservation) > 0
 }
 
 resource "null_resource" "input_validation" {
@@ -146,7 +146,7 @@ resource "null_resource" "input_validation" {
     }
     precondition {
       condition     = local.reservation_valid
-      error_message = "The 'reservation' variable must not be empty when recipe is not 'gke'."
+      error_message = "The 'reservation' variable must not be empty when recipe is not 'gke' or 'gke-ray'."
     }
   }
 }