From d06dc7139a1aea7fcf3030518ef82c0e91b16d14 Mon Sep 17 00:00:00 2001
From: Jean-Baptiste DETROYES <jean-baptiste@detroyes.fr>
Date: Fri, 15 Dec 2023 15:06:17 +0100
Subject: [PATCH] feat: update GPU vars and resources limits

---
 templates/deployment.yaml | 43 +++++++++++++++++++++------------------
 1 file changed, 23 insertions(+), 20 deletions(-)

diff --git a/templates/deployment.yaml b/templates/deployment.yaml
index 734532d..cc2fbe9 100644
--- a/templates/deployment.yaml
+++ b/templates/deployment.yaml
@@ -19,7 +19,7 @@ spec:
       {{- end }}
       labels:
         {{- include "ollama.labels" . | nindent 8 }}
- {{- with .Values.podLabels }}
+        {{- with .Values.podLabels }}
         {{- toYaml . | nindent 8 }}
         {{- end }}
     spec:
@@ -41,24 +41,30 @@ spec:
               containerPort: {{ .Values.service.port }}
               protocol: TCP
           env:
-          {{- if .Values.ollama.gpu.enable }}
+          {{- if or .Values.ollama.gpu.enabled .Values.ollama.gpu.enable }}
             - name: PATH
-              value: /usr/local/nvidia/bin:/usr/local/nvidia/lib64:/usr/bin:/usr/sbin:/bin:/sbin
+              value: /usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
             - name: LD_LIBRARY_PATH
-              value: /usr/local/nvidia/lib64
+              value: /usr/local/nvidia/lib:/usr/local/nvidia/lib64
+            - name: NVIDIA_DRIVER_CAPABILITIES
+              value: compute,utility
           {{- end}}
+          {{- if .Values.resources }}
           resources:
-            {{- toYaml .Values.resources | nindent 12 }}
-            {{- if .Values.ollama.gpu.enable }}
-            limits:
-              nvidia.com/gpu: "{{ .Values.ollama.gpu.nbrGpu }}:{{ .Values.ollama.gpu.nbrGpu | default 1 }}"
-            {{- end}}
+            {{- $limits := default dict .Values.resources.limits }}
+            {{- if or .Values.ollama.gpu.enabled .Values.ollama.gpu.enable }}
+              {{- $gpuLimit := dict "nvidia.com/gpu" (.Values.ollama.gpu.number | .Values.ollama.gpu.nbrGpu | default 1) }}
+              {{- $limits = merge $limits $gpuLimit }}
+            {{- end }}
+            {{- $ressources := merge .Values.resources (dict "limits" $limits) }}
+            {{- toYaml $ressources | nindent 12 }}
+          {{- end}}
           volumeMounts:
             - name: ollama-data
               mountPath: /root/.ollama
-            {{- with .Values.volumeMounts }}
-              {{- toYaml . | nindent 12 }}
-            {{- end }}
+          {{- with .Values.volumeMounts }}
+            {{- toYaml . | nindent 12 }}
+          {{- end }}
           {{- if .Values.livenessProbe.enabled }}
           livenessProbe:
             httpGet:
@@ -83,15 +89,12 @@ spec:
           {{- end }}
       volumes:
         - name: ollama-data
-        {{- if .Values.persistentVolume.enabled }}
-          {{- if .Values.persistentVolume.existingClaim }}
+          {{- if .Values.persistentVolume.enabled }}
           persistentVolumeClaim:
-            claimName: {{ .Values.persistentVolume.existingClaim }}
+            claimName: {{ .Values.persistentVolume.existingClaim |  default (printf "%s" (include "ollama.fullname" .)) }}
           {{- else }}
-          persistentVolumeClaim:
-            claimName: {{ template "ollama.fullname" . }}
+          emptyDir: { }
           {{- end }}
-        {{- end }}
         {{- with .Values.volumes }}
           {{- toYaml . | nindent 8 }}
         {{- end }}
@@ -103,9 +106,9 @@ spec:
       affinity:
         {{- toYaml . | nindent 8 }}
       {{- end }}
-      {{- if or (.Values.ollama.gpu.enable) (.Values.tolerations) }}
+      {{- if or .Values.ollama.gpu.enabled .Values.tolerations }}
       tolerations:
-        {{- if .Values.ollama.gpu.enable }}
+        {{- if or .Values.ollama.gpu.enabled .Values.ollama.gpu.enable }}
         - key: nvidia.com/gpu
           operator: Exists
           effect: NoSchedule