From 698b8945202015e5bca592c7ed89e3cc27297cff Mon Sep 17 00:00:00 2001
From: Tyler Potts <49161327+tylerpotts@users.noreply.github.com>
Date: Mon, 15 Jun 2026 13:46:21 -0500
Subject: [PATCH 1/2] Add nvidia.com/gpu toleration to GPU Ray Serve workloads

nebari-infrastructure-core auto-taints AWS GPU node groups with
nvidia.com/gpu=true:NoSchedule, and EKS has no admission controller to
inject a matching toleration. Without one, GPU-requesting Ray Serve
head/worker pods stop scheduling onto GPU nodes once the taint lands.

Inject the nvidia.com/gpu toleration (operator: Exists, so it matches
any taint value) on a Ray group's pod template when its resources
request nvidia.com/gpu, via a new nebari-rayserve.tolerations helper.
Pods that don't request a GPU render no tolerations and are unchanged.
Explicit head.tolerations / worker.tolerations are appended.

Closes https://github.com/nebari-dev/nebari-rayserve-pack/issues/14
---
 chart/templates/_helpers.tpl    | 27 +++++++++++++++++++++++++++
 chart/templates/rayservice.yaml |  8 ++++++++
 chart/values.yaml               | 12 ++++++++++++
 3 files changed, 47 insertions(+)

diff --git a/chart/templates/_helpers.tpl b/chart/templates/_helpers.tpl
index d8f8d69..81995b8 100644
--- a/chart/templates/_helpers.tpl
+++ b/chart/templates/_helpers.tpl
@@ -61,3 +61,30 @@ 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. 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") -}}
+{{- $tolerations = append $tolerations (dict "key" "nvidia.com/gpu" "operator" "Exists" "effect" "NoSchedule") -}}
+{{- end -}}
+{{- if $tolerations -}}
+{{- toYaml $tolerations -}}
+{{- end -}}
+{{- end }}
diff --git a/chart/templates/rayservice.yaml b/chart/templates/rayservice.yaml
index c05ce9f..ebfec82 100644
--- a/chart/templates/rayservice.yaml
+++ b/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 }}"
diff --git a/chart/values.yaml b/chart/values.yaml
index 2a4c92f..a51ca38 100644
--- a/chart/values.yaml
+++ b/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

From e4514c70a463f09f9d47a5887e4e01d444b1d1ab Mon Sep 17 00:00:00 2001
From: Tyler Potts <49161327+tylerpotts@users.noreply.github.com>
Date: Thu, 18 Jun 2026 13:42:51 -0500
Subject: [PATCH 2/2] fix: avoid duplicate nvidia.com/gpu toleration when user
 defines one

Only inject the nvidia.com/gpu toleration when the group config does not
already define a toleration for that key, so a user-provided toleration
acts as an intentional override rather than producing a duplicate entry.
---
 chart/templates/_helpers.tpl | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/chart/templates/_helpers.tpl b/chart/templates/_helpers.tpl
index 81995b8..6d5c4b2 100644
--- a/chart/templates/_helpers.tpl
+++ b/chart/templates/_helpers.tpl
@@ -70,7 +70,10 @@ 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. Any explicit tolerations from the group config are appended.
+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.
@@ -82,8 +85,16 @@ set, so non-GPU pods are unchanged.
 {{- $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 -}}