Override nvidia env var for 11.8 (mosaicml#2722)

mvpatel2000 · Nov 16, 2023 · bfbb89a · bfbb89a
1 parent 3cf73cc
commit bfbb89a
Show file tree

Hide file tree

Showing 5 changed files with 39 additions and 0 deletions.
diff --git a/.github/workflows/pr-docker.yaml b/.github/workflows/pr-docker.yaml
@@ -56,6 +56,7 @@ jobs:
         TORCHVISION_VERSION=${{ matrix.TORCHVISION_VERSION }}
         PYTORCH_NIGHTLY_URL=${{ matrix.PYTORCH_NIGHTLY_URL }}
         PYTORCH_NIGHTLY_VERSION=${{ matrix.PYTORCH_NIGHTLY_VERSION }}
+        NVIDIA_REQUIRE_CUDA_OVERRIDE=${{ matrix.NVIDIA_REQUIRE_CUDA_OVERRIDE }}
       context: ./docker
       image-name: ${{ matrix.IMAGE_NAME }}
       image-uuid: ${{ matrix.UUID }}

diff --git a/.github/workflows/release-docker.yaml b/.github/workflows/release-docker.yaml
@@ -53,6 +53,7 @@ jobs:
         TORCHVISION_VERSION=${{ matrix.TORCHVISION_VERSION }}
         PYTORCH_NIGHTLY_URL=${{ matrix.PYTORCH_NIGHTLY_URL }}
         PYTORCH_NIGHTLY_VERSION=${{ matrix.PYTORCH_NIGHTLY_VERSION }}
+        NVIDIA_REQUIRE_CUDA_OVERRIDE=${{ matrix.NVIDIA_REQUIRE_CUDA_OVERRIDE }}
       context: ./docker
       image-name: ${{ matrix.IMAGE_NAME }}
       image-uuid: ${{ matrix.UUID }}

diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -354,6 +354,11 @@ RUN pip install --no-cache-dir --upgrade \
         ipython${IPYTHON_VERSION} \
         urllib3${URLLIB3_VERSION}
 
+##################################################
+# Override NVIDIA mistaken env var for 11.8 images
+##################################################
+ARG NVIDIA_REQUIRE_CUDA_OVERRIDE
+ENV NVIDIA_REQUIRE_CUDA=${NVIDIA_REQUIRE_CUDA_OVERRIDE:-$NVIDIA_REQUIRE_CUDA}
 
 ################
 # Composer Image

diff --git a/docker/build_matrix.yaml b/docker/build_matrix.yaml
@@ -4,6 +4,7 @@
   CUDA_VERSION: 12.1.0
   IMAGE_NAME: torch-2-1-0-cu121
   MOFED_VERSION: 5.5-1.0.3.2
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
@@ -18,6 +19,7 @@
   CUDA_VERSION: 12.1.0
   IMAGE_NAME: torch-2-1-0-cu121-aws
   MOFED_VERSION: ''
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
@@ -32,6 +34,7 @@
   CUDA_VERSION: ''
   IMAGE_NAME: torch-2-1-0-cpu
   MOFED_VERSION: ''
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
@@ -46,6 +49,12 @@
   CUDA_VERSION: 11.8.0
   IMAGE_NAME: torch-2-0-1-cu118
   MOFED_VERSION: 5.5-1.0.3.2
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: cuda>=11.8 brand=tesla,driver>=470,driver<471 brand=tesla,driver>=515,driver<516
+    brand=unknown,driver>=470,driver<471 brand=unknown,driver>=515,driver<516 brand=nvidia,driver>=470,driver<471
+    brand=nvidia,driver>=515,driver<516 brand=nvidiartx,driver>=470,driver<471 brand=nvidiartx,driver>=515,driver<516
+    brand=geforce,driver>=470,driver<471 brand=geforce,driver>=515,driver<516 brand=quadro,driver>=470,driver<471
+    brand=quadro,driver>=515,driver<516 brand=titan,driver>=470,driver<471 brand=titan,driver>=515,driver<516
+    brand=titanrtx,driver>=470,driver<471 brand=titanrtx,driver>=515,driver<516
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
@@ -59,6 +68,12 @@
   CUDA_VERSION: 11.8.0
   IMAGE_NAME: torch-2-0-1-cu118-aws
   MOFED_VERSION: ''
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: cuda>=11.8 brand=tesla,driver>=470,driver<471 brand=tesla,driver>=515,driver<516
+    brand=unknown,driver>=470,driver<471 brand=unknown,driver>=515,driver<516 brand=nvidia,driver>=470,driver<471
+    brand=nvidia,driver>=515,driver<516 brand=nvidiartx,driver>=470,driver<471 brand=nvidiartx,driver>=515,driver<516
+    brand=geforce,driver>=470,driver<471 brand=geforce,driver>=515,driver<516 brand=quadro,driver>=470,driver<471
+    brand=quadro,driver>=515,driver<516 brand=titan,driver>=470,driver<471 brand=titan,driver>=515,driver<516
+    brand=titanrtx,driver>=470,driver<471 brand=titanrtx,driver>=515,driver<516
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
@@ -72,6 +87,7 @@
   CUDA_VERSION: ''
   IMAGE_NAME: torch-2-0-1-cpu
   MOFED_VERSION: ''
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
@@ -85,6 +101,7 @@
   CUDA_VERSION: 11.7.1
   IMAGE_NAME: torch-1-13-1-cu117
   MOFED_VERSION: 5.5-1.0.3.2
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
@@ -98,6 +115,7 @@
   CUDA_VERSION: 11.7.1
   IMAGE_NAME: torch-1-13-1-cu117-aws
   MOFED_VERSION: ''
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''
@@ -111,6 +129,7 @@
   CUDA_VERSION: ''
   IMAGE_NAME: torch-1-13-1-cpu
   MOFED_VERSION: ''
+  NVIDIA_REQUIRE_CUDA_OVERRIDE: ''
   PYTHON_VERSION: '3.10'
   PYTORCH_NIGHTLY_URL: ''
   PYTORCH_NIGHTLY_VERSION: ''

diff --git a/docker/generate_build_matrix.py b/docker/generate_build_matrix.py
@@ -136,6 +136,17 @@ def _main():
 
         cuda_version = _get_cuda_version(pytorch_version=pytorch_version, use_cuda=use_cuda)
 
+        override_string = ('cuda>=11.8 brand=tesla,driver>=470,driver<471 '
+                           'brand=tesla,driver>=515,driver<516 brand=unknown,driver>=470,driver<471 '
+                           'brand=unknown,driver>=515,driver<516 brand=nvidia,driver>=470,driver<471 '
+                           'brand=nvidia,driver>=515,driver<516 brand=nvidiartx,driver>=470,driver<471 '
+                           'brand=nvidiartx,driver>=515,driver<516 brand=geforce,driver>=470,driver<471 '
+                           'brand=geforce,driver>=515,driver<516 brand=quadro,driver>=470,driver<471 '
+                           'brand=quadro,driver>=515,driver<516 brand=titan,driver>=470,driver<471 '
+                           'brand=titan,driver>=515,driver<516 brand=titanrtx,driver>=470,driver<471 '
+                           'brand=titanrtx,driver>=515,driver<516')
+        nvidia_require_cuda_override = '' if cuda_version != '11.8.0' else override_string
+
         entry = {
             'IMAGE_NAME':
                 _get_image_name(pytorch_version, cuda_version, stage, interconnect),
@@ -163,6 +174,8 @@ def _main():
                 '',
             'PYTORCH_NIGHTLY_VERSION':
                 '',
+            'NVIDIA_REQUIRE_CUDA_OVERRIDE':
+                nvidia_require_cuda_override,
         }
 
         # Only build EFA image on latest python with cuda on pytorch_stage