From e9b01b34333a9e2e4e0c54e440616a91a2d971f4 Mon Sep 17 00:00:00 2001
From: OlivierDehaene <olivier@huggingface.co>
Date: Tue, 9 May 2023 12:36:02 +0200
Subject: [PATCH] fix(dockerfile): fix nvidia env vars (#297)

Fixes #291
---
 Dockerfile | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index d59acc10..8810f27e 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -120,6 +120,16 @@ ENV HUGGINGFACE_HUB_CACHE=/data \
     NUM_SHARD=1 \
     PORT=80
 
+# NVIDIA env vars
+ENV NVIDIA_REQUIRE_CUDA cuda>=11.8
+ENV NVIDIA_VISIBLE_DEVICES all
+ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
+# Required for nvidia-docker v1
+RUN /bin/bash -c echo "/usr/local/nvidia/lib" >> /etc/ld.so.conf.d/nvidia.conf && \
+    echo "/usr/local/nvidia/lib64" >> /etc/ld.so.conf.d/nvidia.conf
+ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
+ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
+
 LABEL com.nvidia.volumes.needed="nvidia_driver"
 
 WORKDIR /usr/src
@@ -165,11 +175,6 @@ FROM base as sagemaker
 COPY sagemaker-entrypoint.sh entrypoint.sh
 RUN chmod +x entrypoint.sh
 
-# NVIDIA env vars
-ENV NVIDIA_VISIBLE_DEVICES all
-ENV NVIDIA_DRIVER_CAPABILITIES compute,utility
-ENV LD_LIBRARY_PATH /usr/local/nvidia/lib:/usr/local/nvidia/lib64
-
 ENTRYPOINT ["./entrypoint.sh"]
 
 # Final image