Add nvidia.com/gpu toleration to Ray Serve GPU workloads

chart/templates/_helpers.tpl

@@ -61,3 +61,41 @@ Ray serve service name - RayService creates a service named <rayservice-name>-se
 {{- define "nebari-rayserve.serve-service-name" -}}
 {{- printf "%s-serve-svc" (include "nebari-rayserve.fullname" .) }}
 {{- end }}
+
+{{/*
+Tolerations for a Ray group (head or worker). Pass the group config
+(.Values.head or .Values.worker).
+
+When the group's resources request an NVIDIA GPU, the nvidia.com/gpu
+toleration is injected automatically so GPU pods can schedule onto nodes
+tainted nvidia.com/gpu=...:NoSchedule (e.g. nebari-infrastructure-core's
+auto-tainted AWS GPU node groups). operator: Exists matches any taint
+value. The toleration is only injected when the group config does not
+already define one for the nvidia.com/gpu key, so a user-provided
+toleration acts as an intentional override. Any explicit tolerations from
+the group config are appended.
+
+Renders nothing when no GPU is requested and no explicit tolerations are
+set, so non-GPU pods are unchanged.
+*/}}
+{{- define "nebari-rayserve.tolerations" -}}
+{{- $config := . -}}
+{{- $tolerations := $config.tolerations | default list -}}
+{{- $resources := $config.resources | default dict -}}
+{{- $limits := $resources.limits | default dict -}}
+{{- $requests := $resources.requests | default dict -}}
+{{- if or (hasKey $limits "nvidia.com/gpu") (hasKey $requests "nvidia.com/gpu") -}}
+{{- $hasGpuToleration := false -}}
+{{- range $tolerations -}}
+{{- if eq (.key | default "") "nvidia.com/gpu" -}}
+{{- $hasGpuToleration = true -}}
+{{- end -}}
+{{- end -}}
+{{- if not $hasGpuToleration -}}
+{{- $tolerations = append $tolerations (dict "key" "nvidia.com/gpu" "operator" "Exists" "effect" "NoSchedule") -}}
+{{- end -}}
+{{- end -}}
+{{- if $tolerations -}}
+{{- toYaml $tolerations -}}
+{{- end -}}
+{{- end }}

chart/templates/rayservice.yaml

@@ -27,6 +27,10 @@ spec:
           {{- with .Values.head.runtimeClassName }}
           runtimeClassName: {{ . }}
           {{- end }}
+          {{- with (include "nebari-rayserve.tolerations" .Values.head) }}
+          tolerations:
+            {{- . | nindent 12 }}
+          {{- end }}
           containers:
             - name: ray-head
               image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}"
@@ -64,6 +68,10 @@ spec:
             {{- with .Values.worker.runtimeClassName }}
             runtimeClassName: {{ . }}
             {{- end }}
+            {{- with (include "nebari-rayserve.tolerations" .Values.worker) }}
+            tolerations:
+              {{- . | nindent 14 }}
+            {{- end }}
             containers:
               - name: ray-worker
                 image: "{{ .Values.image.repository }}:{{ .Values.image.tag }}"

chart/values.yaml

@@ -95,6 +95,12 @@ head:
       cpu: "1"
       memory: "2Gi"
 
+  # Extra tolerations for the head pod. The nvidia.com/gpu toleration is
+  # injected automatically when resources above request nvidia.com/gpu, so
+  # GPU pods schedule onto nodes tainted nvidia.com/gpu=...:NoSchedule. Add
+  # any additional tolerations here.
+  tolerations: []
+
   # Head pod readiness/liveness probes. Defaults to {} which lets KubeRay
   # apply its built-in probes. Override here to set explicit probes; see
   # the worker.readinessProbe / worker.livenessProbe block below for the
@@ -119,6 +125,12 @@ worker:
       cpu: "1"
       memory: "2Gi"
 
+  # Extra tolerations for worker pods. The nvidia.com/gpu toleration is
+  # injected automatically when resources above request nvidia.com/gpu, so
+  # GPU workers schedule onto nodes tainted nvidia.com/gpu=...:NoSchedule.
+  # Add any additional tolerations here.
+  tolerations: []
+
   # Worker pod readiness/liveness probes.
   #
   # KubeRay's default worker probes chain a raylet healthz check with