review comments addressed

devcontainers · Oct 1, 2024 · 6d28519 · 6d28519
1 parent 022948e
commit 6d28519
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 7 deletions.
diff --git a/src/nvidia-cuda/devcontainer-feature.json b/src/nvidia-cuda/devcontainer-feature.json
@@ -1,6 +1,6 @@
 {
   "id": "nvidia-cuda",
-  "version": "1.1.4",
+  "version": "1.2.0",
   "name": "NVIDIA CUDA",
   "description": "Installs shared libraries for NVIDIA CUDA.",
   "documentationURL": "https://github.com/devcontainers/features/tree/main/src/nvidia-cuda",
@@ -44,6 +44,44 @@
       ],
       "default": "11.8",
       "description": "Version of CUDA to install"
+    },
+    "cudnnVersion": {
+      "type": "string",
+      "proposals": [
+        "8.9.5.29",
+        "8.9.4.25",
+        "8.9.3.28",
+        "8.9.2.26",
+        "8.9.1.23",
+        "8.9.0.131",
+        "8.8.1.3",
+        "8.8.0.121",
+        "8.7.0.84",
+        "8.6.0.163",
+        "8.5.0.96",
+        "8.4.1.50",
+        "8.4.0.27",
+        "8.3.3.40",
+        "8.3.2.44",
+        "8.3.1.22",
+        "8.3.0.98",
+        "8.2.4.15",
+        "8.2.2.26",
+        "8.2.1.32",
+        "8.2.0.53",
+        "8.1.1.33",
+        "8.1.0.77",
+        "9.0.0.312",
+        "9.1.0.70",
+        "9.1.1.17",
+        "9.2.0.82",
+        "9.2.1.18",
+        "9.3.0.75",
+        "9.4.0.58",
+        "automatic"
+      ],
+      "default": "automatic",
+      "description": "Version of cuDNN to install"
     }
   },
   "installsAfter": [

diff --git a/src/nvidia-cuda/install.sh b/src/nvidia-cuda/install.sh
@@ -10,6 +10,7 @@ INSTALL_CUDNNDEV=${INSTALLCUDNNDEV}
 INSTALL_NVTX=${INSTALLNVTX}
 INSTALL_TOOLKIT=${INSTALLTOOLKIT}
 CUDA_VERSION=${CUDAVERSION}
+CUDNN_VERSION=${CUDNNVERSION}
 
 . /etc/os-release 
 
@@ -58,8 +59,6 @@ apt-get update -yq
 cuda_pkg="cuda-libraries-${CUDA_VERSION/./-}"
 nvtx_pkg="cuda-nvtx-${CUDA_VERSION/./-}"
 toolkit_pkg="cuda-toolkit-${CUDA_VERSION/./-}"
-major_cudnn_version=$(echo "${CUDNN_VERSION}" | cut -d '.' -f 1)
-major_cuda_version=$(echo "${CUDA_VERSION}" | cut -d '.' -f 1)
 if ! apt-cache show "$cuda_pkg"; then
     echo "The requested version of CUDA is not available: CUDA $CUDA_VERSION"
     exit 1
@@ -71,10 +70,12 @@ apt-get update -yq
 
 # auto find recent cudnn version
 major_cuda_version=$(echo "${CUDA_VERSION}" | cut -d '.' -f 1)
-if [[ "$CUDA_VERSION" < "12.3" ]]; then
-    CUDNN_VERSION=$(apt-cache policy libcudnn8 | grep "$CUDA_VERSION" | grep -Eo '^[^-1+]*' | sort -V | tail -n1 | xargs)
-else
-    CUDNN_VERSION=$(apt-cache policy libcudnn9-cuda-$major_cuda_version | grep "Candidate" | awk '{print $2}' | grep -Eo '^[^-1+]*')
+if [ "$CUDNN_VERSION" = "automatic" ]; then
+    if [[ "$CUDA_VERSION" < "12.3" ]]; then
+        CUDNN_VERSION=$(apt-cache policy libcudnn8 | grep "$CUDA_VERSION" | grep -Eo '^[^-1+]*' | sort -V | tail -n1 | xargs)
+    else
+        CUDNN_VERSION=$(apt-cache policy libcudnn9-cuda-$major_cuda_version | grep "Candidate" | awk '{print $2}' | grep -Eo '^[^-1+]*')
+    fi
 fi
 major_cudnn_version=$(echo "${CUDNN_VERSION}" | cut -d '.' -f 1)