From 543bfb8e0f473dc3141e80cca917f89f18a945a8 Mon Sep 17 00:00:00 2001 From: Xu Zhao Date: Fri, 11 Oct 2024 09:54:25 -0400 Subject: [PATCH] Install nvidia driver in docker build file --- docker/tritonbench-nightly.dockerfile | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/docker/tritonbench-nightly.dockerfile b/docker/tritonbench-nightly.dockerfile index 6d669391..f27da7d2 100644 --- 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