Add security context to Ray clusters

modules/kuberay-cluster/kuberay-autopilot-values.yaml

@@ -42,7 +42,8 @@ head:
   autoscalerOptions:
     upscalingMode: Default
     idleTimeoutSeconds: 60
-    securityContext: {}
+    securityContext:
+      ${indent(4, security_context)}
     env: []
     envFrom: []
     # resources specifies optional resource request and limit overrides for the autoscaler container.
@@ -111,7 +112,8 @@ head:
   tolerations: []
   affinity: {}
   # Ray container security context.
-  securityContext: {}
+  securityContext:
+    ${indent(4, security_context)}
   volumes:
     - name: ray-logs
       emptyDir: {}
@@ -213,7 +215,8 @@ additionalWorkerGroups:
     tolerations: []
     affinity: {}
     # Ray container security context.
-    securityContext: {}
+    securityContext:
+      ${indent(4, security_context)}
     volumes:
       - name: ray-logs
         emptyDir: {}
@@ -309,7 +312,8 @@ additionalWorkerGroups:
     tolerations: []
     affinity: {}
     # Ray container security context.
-    securityContext: {}
+    securityContext:
+      ${indent(4, security_context)}
     volumes:
       - name: ray-logs
         emptyDir: {}

modules/kuberay-cluster/kuberay-gpu-values.yaml

@@ -105,7 +105,8 @@ head:
   tolerations: []
   affinity: {}
   # Ray container security context.
-  securityContext: {}
+  securityContext:
+    ${indent(4, security_context)}
   volumes:
     - name: ray-logs
       emptyDir: {}
@@ -202,7 +203,8 @@ worker:
   tolerations: []
   affinity: {}
   # Ray container security context.
-  securityContext: {}
+  securityContext:
+    ${indent(4, security_context)}
   volumes:
     - name: ray-logs
       emptyDir: {}
@@ -291,7 +293,8 @@ additionalWorkerGroups:
     tolerations: []
     affinity: {}
     # Ray container security context.
-    securityContext: {}
+    securityContext:
+      ${indent(4, security_context)}
     volumes:
       - name: ray-logs
         emptyDir: {}

modules/kuberay-cluster/kuberay-values.yaml

@@ -102,7 +102,8 @@ head:
   tolerations: []
   affinity: {}
   # Ray container security context.
-  securityContext: {}
+  securityContext:
+    ${indent(4, security_context)}
   volumes:
     - name: ray-logs
       emptyDir: {}
@@ -196,7 +197,8 @@ worker:
   tolerations: []
   affinity: {}
   # Ray container security context.
-  securityContext: {}
+  securityContext:
+    ${indent(4, security_context)}
   volumes:
     - name: ray-logs
       emptyDir: {}

modules/kuberay-cluster/kuberay.tf

@@ -30,18 +30,22 @@ resource "helm_release" "ray-cluster" {
       gcs_bucket          = var.gcs_bucket
       k8s_service_account = var.google_service_account
       grafana_host        = var.grafana_host
+      security_context    = chomp(yamlencode({for k, v in var.security_context : k => v if v != null }))
       }) : var.enable_tpu ? templatefile("${path.module}/kuberay-tpu-values.yaml", {
       gcs_bucket          = var.gcs_bucket
       k8s_service_account = var.google_service_account
       grafana_host        = var.grafana_host
+      security_context    = chomp(yamlencode({for k, v in var.security_context : k => v if v != null }))
       }) : var.enable_gpu ? templatefile("${path.module}/kuberay-gpu-values.yaml", {
       gcs_bucket          = var.gcs_bucket
       k8s_service_account = var.google_service_account
       grafana_host        = var.grafana_host
+      security_context    = chomp(yamlencode({for k, v in var.security_context : k => v if v != null }))
       }) : templatefile("${path.module}/kuberay-values.yaml", {
       gcs_bucket          = var.gcs_bucket
       k8s_service_account = var.google_service_account
       grafana_host        = var.grafana_host
+      security_context    = chomp(yamlencode({for k, v in var.security_context : k => v if v != null }))
     })
   ]
 }

modules/kuberay-cluster/variables.tf

@@ -57,3 +57,55 @@ variable "gcs_bucket" {
 variable "grafana_host" {
   type = string
 }
+
+# Commenting out underutilized, nested variables because they're harder to
+# strip from the YAML encoding. Add these back if we can easily strip out
+# the null values.
+variable "security_context" {
+  description = "Kubernetes security context to set on all ray cluster pods"
+  type = object({
+    allowPrivilegeEscalation = optional(bool)
+    capabilities = optional(object({
+    # Not typically used
+    # add  = optional(list(string))
+      drop = optional(list(string))
+    }))
+    privileged = optional(bool)
+    procMount  = optional(string)
+    readOnlyRootFilesystem = optional(bool)
+    runAsGroup             = optional(number)
+    runAsNonRoot           = optional(bool)
+    runAsUser              = optional(number)
+    seLinuxOptions         = optional(object({
+      level = optional(string)
+      role  = optional(string)
+      type  = optional(string)
+      user  = optional(string)
+    }))
+    seccompProfile         = optional(object({
+    # Not typically used
+    # localhostProfile = optional(string)
+      type             = optional(string)
+    }))
+    windowsOptions = optional(object({
+      gmsaCredentialSpec    = optional(string)
+      gmsaCredentiaSpecName = optional(string)
+      hostProcess           = bool
+      runAsUserName         = string
+    }))
+  })
+
+  default = {
+    allowPrivilegeEscalation = false
+    capabilities = {
+      drop = ["ALL"]
+    }
+    # GKE will automatically mount GPUs and TPUs into unprivileged pods
+    privileged = false
+    readOnlyFileSystem = true
+    runAsNonRoot = true
+    seccompProfile = {
+      type = "RuntimeDefault"
+    }
+  }
+}