Install nvidia driver in docker build file

pytorch-labs · Oct 11, 2024 · 543bfb8 · 543bfb8
1 parent 280d0cc
commit 543bfb8
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/docker/tritonbench-nightly.dockerfile b/docker/tritonbench-nightly.dockerfile
@@ -41,12 +41,17 @@ RUN cd /workspace/tritonbench && \
         python utils/cuda_utils.py --check-torch-nightly-version --force-date "${FORCE_DATE}"; \
     fi
 
+# Tritonbench library build and test require libcuda.so.1
+# which is from NVIDIA driver
+RUN sudo apt update && sudo apt-get install -y libnvidia-compute-550
+
 # Install Tritonbench
 RUN cd /workspace/tritonbench && \
     bash .ci/tritonbench/install.sh
 
-# Test Tritonbench (libcuda.so.1 is required for fbgemm test, so install libnvidia-compute-550 as a hack)
-RUN sudo apt update && sudo apt-get install -y libnvidia-compute-550 && \
-    cd /workspace/tritonbench && \
-    bash .ci/tritonbench/test-install.sh && \
-    sudo apt-get purge -y libnvidia-compute-550
+# Test Tritonbench
+RUN cd /workspace/tritonbench && \
+    bash .ci/tritonbench/test-install.sh
+
+# Remove NVIDIA driver library - they are supposed to be mapped at runtime
+RUN sudo apt-get purge -y libnvidia-compute-550