auto find cudnn version

src/nvidia-cuda/devcontainer-feature.json

@@ -1,6 +1,6 @@
 {
   "id": "nvidia-cuda",
-  "version": "1.1.3",
+  "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",
@@ -48,6 +48,7 @@
     "cudnnVersion": {
       "type": "string",
       "proposals": [
+        "automatic",
         "8.9.5.29",
         "8.9.4.25",
         "8.9.3.28",
@@ -79,7 +80,7 @@
         "9.3.0.75",
         "9.4.0.58"
       ],
-      "default": "8.6.0.163",
+      "default": "automatic",
       "description": "Version of cuDNN to install"
     }
   },

src/nvidia-cuda/install.sh

@@ -59,20 +59,30 @@ 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
 fi
 
 echo "Installing CUDA libraries..."
 apt-get install -yq "$cuda_pkg"
+apt-get update -yq --fix-missing
+
+# auto find recent cudnn version
+major_cuda_version=$(echo "${CUDA_VERSION}" | cut -d '.' -f 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)
 
 if [ "$INSTALL_CUDNN" = "true" ]; then
     # Ensure that the requested version of cuDNN is available AND compatible
-    #if major cudnn version is 9, then we need to install libcudnn9-cuda-<major_version> package
-    #else we need to install libcudnn8-cuda-<major_version> package
+    #if major cudnn version is 9, then we need to install libcudnn9-cuda-<major_cuda_version>_<CUDNN_VERSION>-1 package
+    #else we need to install libcudnn8_<CUDNN_VERSION>-1+cuda<CUDA_VERSION>" package
     if [[ $major_cudnn_version -ge "9" ]]
     then
         cudnn_pkg_version="libcudnn9-cuda-${major_cuda_version}=${CUDNN_VERSION}-1"
@@ -91,13 +101,14 @@ fi
 
 if [ "$INSTALL_CUDNNDEV" = "true" ]; then
     # Ensure that the requested version of cuDNN development package is available AND compatible
+    #if major cudnn version is 9, then we need to install libcudnn9-dev-cuda-<major_cuda_version>_<CUDNN_VERSION>-1 package
+    #else we need to install libcudnn8-dev_<CUDNN_VERSION>-1+cuda<CUDA_VERSION>" package
     if [[ $major_cudnn_version -ge "9" ]]
     then
         cudnn_dev_pkg_version="libcudnn9-dev-cuda-${major_cuda_version}=${CUDNN_VERSION}-1"
     else
         cudnn_dev_pkg_version="libcudnn8-dev=${CUDNN_VERSION}-1+cuda${CUDA_VERSION}"
     fi
-
     if ! apt-cache show "$cudnn_dev_pkg_version"; then
         echo "The requested version of cuDNN development package is not available: cuDNN $CUDNN_VERSION for CUDA $CUDA_VERSION"
         exit 1

test/nvidia-cuda/install_cuda_12_3_version.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+set -e
+
+# Optional: Import test library
+source dev-container-features-test-lib
+
+# # Check installation of libcudnn9-cuda-12 (9.4.0)
+check "libcudnn.so.9.5.0" test 1 -eq "$(find /usr -name 'libcudnn.so.9.5.0' | wc -l)"
+
+# Check installation of cuda-nvtx-12-3 (12.3)
+check "cuda-12-3+nvtx" test -e '/usr/local/cuda-12.3/targets/x86_64-linux/include/nvtx3/'
+
+# Report result
+reportResults

test/nvidia-cuda/install_cuda_12_4_version.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+set -e
+
+# Optional: Import test library
+source dev-container-features-test-lib
+
+# # Check installation of libcudnn9-cuda-12 (9.4.0)
+check "libcudnn.so.9.5.0" test 1 -eq "$(find /usr -name 'libcudnn.so.9.5.0' | wc -l)"
+
+# Check installation of cuda-nvtx-12-4 (12.4)
+check "cuda-12-4+nvtx" test -e '/usr/local/cuda-12.4/targets/x86_64-linux/include/nvtx3/'
+
+# Report result
+reportResults

test/nvidia-cuda/install_cuda_12_5_version.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+set -e
+
+# Optional: Import test library
+source dev-container-features-test-lib
+
+# # Check installation of libcudnn9-cuda-12 (9.4.0)
+check "libcudnn.so.9.5.0" test 1 -eq "$(find /usr -name 'libcudnn.so.9.5.0' | wc -l)"
+
+# Check installation of cuda-nvtx-12-5 (12.5)
+check "cuda-12-5+nvtx" test -e '/usr/local/cuda-12.5/targets/x86_64-linux/include/nvtx3/'
+
+# Report result
+reportResults

test/nvidia-cuda/scenarios.json

@@ -16,47 +16,41 @@
             "nvidia-cuda": {
                 "installCudnn": true,
                 "installNvtx": true,
-                "cudaVersion": "11.7",
-                "cudnnVersion": "8.5.0.96"
+                "cudaVersion": "11.7"
             }
         }
     },
-    "install_cudnn_123_version": {
+    "install_cuda_12_3_version": {
         "image": "debian",
         "features": {
             "nvidia-cuda": {
                 "installCudnn": true,
-                "installNvtx": true,
                 "installCudnnDev": true,
-                "installToolkit": true,
-                "cudaVersion": "12.3",
-                "cudnnVersion": "9.4.0.58"
+                "installNvtx": true,
+                "cudaVersion": "12.3"
             }
         }
     },
-    "install_cudnn_124_version": {
+    "install_cuda_12_4_version": {
         "image": "debian",
         "features": {
             "nvidia-cuda": {
                 "installCudnn": true,
-                "installNvtx": true,
                 "installCudnnDev": true,
-                "installToolkit": true,
-                "cudaVersion": "12.4",
-                "cudnnVersion": "9.4.0.58"
+                "installNvtx": true,
+                "cudaVersion": "12.4"
             }
         }
     },
-    "install_cudnn_latest_version": {
+    "install_cuda_12_5_version": {
         "image": "debian",
         "features": {
             "nvidia-cuda": {
                 "installCudnn": true,
-                "installNvtx": true,
                 "installCudnnDev": true,
-                "installToolkit": true,
+                "installNvtx": true,
                 "cudaVersion": "12.5",
-                "cudnnVersion": "9.4.0.58"
+                "cudnnVersion": "9.5.0.50"
             }
         }
     }