diff --git a/chart/templates/_helpers.tpl b/chart/templates/_helpers.tpl index d8f8d69..6d5c4b2 100644 --- a/chart/templates/_helpers.tpl +++ b/chart/templates/_helpers.tpl @@ -61,3 +61,41 @@ Ray serve service name - RayService creates a service named -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 }} 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