install nvidia driver in train image

mozilla · Jan 20, 2025 · bae1fc9 · bae1fc9
1 parent fd98c03
commit bae1fc9
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/taskcluster/docker/train/Dockerfile b/taskcluster/docker/train/Dockerfile
@@ -18,5 +18,12 @@ RUN apt-get update -qq \
                           libicu-dev \
     && apt-get clean
 
+
+RUN curl -sSO https://developer.download.nvidia.com/compute/nvidia-driver/565.57.01/local_installers/nvidia-driver-local-repo-ubuntu2204-565.57.01_1.0-1_amd64.deb \
+    && dpkg -i nvidia-driver-local-repo-ubuntu2204-565.57.01_1.0-1_amd64.deb \
+    && cp /var/nvidia-driver-local-repo-ubuntu2204-565.57.01/nvidia-driver-*-keyring.gpg /usr/share/keyrings/ \
+    && apt-get update \
+    && apt-get install -y cuda-drivers-565
+
 VOLUME /builds/worker/checkouts
 VOLUME /builds/worker/.cache