From 0409c1fa592ac3a56c883499e7039b01617ec204 Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Thu, 30 Nov 2023 21:16:56 -0800
Subject: [PATCH] docker: set PATH, LD_LIBRARY_PATH, and capabilities (#1336)

* docker: set PATH, LD_LIBRARY_PATH, and capabilities

* example: update k8s gpu manifest
---
 Dockerfile                   | 6 ++++++
 examples/kubernetes/gpu.yaml | 6 ++++--
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 713f8a42..c50665b6 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -19,5 +19,11 @@ RUN apt-get update && apt-get install -y ca-certificates
 COPY --from=0 /go/src/github.com/jmorganca/ollama/ollama /bin/ollama
 EXPOSE 11434
 ENV OLLAMA_HOST 0.0.0.0
+
+# set some environment variable for better NVIDIA compatibility
+ENV PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
+ENV LD_LIBRARY_PATH=/usr/local/nvidia/lib:/usr/local/nvidia/lib64
+ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility
+
 ENTRYPOINT ["/bin/ollama"]
 CMD ["serve"]
diff --git a/examples/kubernetes/gpu.yaml b/examples/kubernetes/gpu.yaml
index 4ee5f07d..ba90abb6 100644
--- a/examples/kubernetes/gpu.yaml
+++ b/examples/kubernetes/gpu.yaml
@@ -25,9 +25,11 @@ spec:
         image: ollama/ollama:latest
         env:
         - 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
         ports:
         - name: http
           containerPort: 11434