Add new NVIDIA CUDA feature (#80)

* Add new Nvidia feature

* Remove random empty files

* Update comments

* Rename feature to nvidia-cuda

* Add feature to tests

* Add version

* Move test to match new name

* Add final output message

* Fix capitalization of NVIDIA

* Remove option for base CUDA

* Use camelCase

* Check for required packages

* Use os-release instead of lsb_release

* Clean up keyring variables

* Collapse keyring lines

* Always install CUDA libraries

* Add option to install NVTX

* Always use ubuntu2004 repo

* Use test instead of brackets

* Add default values to feature

* Add version options for CUDA and cuDNN

* Rename CUDA version option

* Add scenario to test specific CUDA/cuDNN version

* Rename cuDNN scenario

* Fix typo in test scenario

* Update variable casing

* Add more helpful error messages

* Remove default values from script

* Use enum for version option

* Polish new scenarios

* Remove apt_get_update_if_needed and check_packages

* Add more versions

* Improve error messages

* Comments and feature description

Co-authored-by: Josh Spicer <joshspicer@github.com>

8 files changed, 180 insertions, 0 deletions

diff --git a/.github/workflows/test-all.yaml b/.github/workflows/test-all.yaml
index d04c534..644a4ec 100644
--- a/.github/workflows/test-all.yaml
+++ b/.github/workflows/test-all.yaml
@@ -28,6 +28,7 @@ jobs:
             "java",
             "kubectl-helm-minikube",
             "node",
+            "nvidia-cuda",
             "oryx",
             "php",
             "powershell",
diff --git a/.github/workflows/test-pr.yaml b/.github/workflows/test-pr.yaml
index 444422a..b960039 100644
--- a/.github/workflows/test-pr.yaml
+++ b/.github/workflows/test-pr.yaml
@@ -28,6 +28,7 @@ jobs:
             java: ./**/java/**
             kubectl-helm-minikube: ./**/kubectl-helm-minikube/**
             node: ./**/node/**
+            nvidia-cuda: ./**/nvidia-cuda/**
             oryx: ./**/oryx/**
             php: ./**/php/**
             powershell: ./**/powershell/**
diff --git a/src/nvidia-cuda/devcontainer-feature.json b/src/nvidia-cuda/devcontainer-feature.json
new file mode 100644
index 0000000..711374b
--- /dev/null
+++ b/src/nvidia-cuda/devcontainer-feature.json
@@ -0,0 +1,51 @@
+{
+  "id": "nvidia-cuda",
+  "name": "NVIDIA CUDA",
+  "description": "Installs shared libraries for NVIDIA CUDA.",
+  "version": "1.0.0",
+  "options": {
+    "installCudnn": {
+      "type": "boolean",
+      "default": false,
+      "description": "Additionally install CUDA Deep Neural Network (cuDNN) shared library"
+    },
+    "installNvtx": {
+      "type": "boolean",
+      "default": false,
+      "description": "Additionally install NVIDIA Tools Extension (NVTX)"
+    },
+    "cudaVersion": {
+      "type": "string",
+      "enum": [
+        "11.7",
+        "11.6",
+        "11.5",
+        "11.4",
+        "11.3",
+        "11.2"
+      ],
+      "default": "11.7",
+      "description": "Version of CUDA to install"
+    },
+    "cudnnVersion": {
+      "type": "string",
+      "enum": [
+        "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"
+      ],
+      "default": "8.5.0.96",
+      "description": "Version of cuDNN to install"
+    }
+  }
+}
diff --git a/src/nvidia-cuda/install.sh b/src/nvidia-cuda/install.sh
new file mode 100644
index 0000000..f017ecb
--- /dev/null
+++ b/src/nvidia-cuda/install.sh
@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+
+set -e
+
+INSTALL_CUDNN=${INSTALLCUDNN}
+INSTALL_NVTX=${INSTALLNVTX}
+CUDA_VERSION=${CUDAVERSION}
+CUDNN_VERSION=${CUDNNVERSION}
+
+if [ "$(id -u)" -ne 0 ]; then
+    echo -e 'Script must be run as root. Use sudo, su, or add "USER root" to your Dockerfile before running this script.'
+    exit 1
+fi
+
+# Install dependencies
+apt-get update -yq
+apt-get install -yq wget ca-certificates
+
+# Add NVIDIA's package repository to apt so that we can download packages
+# Always use the ubuntu2004 repo because the other repos (e.g., debian11) are missing packages
+NVIDIA_REPO_URL="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64"
+KEYRING_PACKAGE="cuda-keyring_1.0-1_all.deb"
+KEYRING_PACKAGE_URL="$NVIDIA_REPO_URL/$KEYRING_PACKAGE"
+KEYRING_PACKAGE_PATH="$(mktemp -d)"
+KEYRING_PACKAGE_FILE="$KEYRING_PACKAGE_PATH/$KEYRING_PACKAGE"
+wget -O "$KEYRING_PACKAGE_FILE" "$KEYRING_PACKAGE_URL"
+apt-get install -yq "$KEYRING_PACKAGE_FILE"
+apt-get update -yq
+
+# Ensure that the requested version of CUDA is available
+cuda_pkg="cuda-libraries-${CUDA_VERSION/./-}"
+nvtx_pkg="cuda-nvtx-${CUDA_VERSION/./-}"
+if ! apt-cache show "$cuda_pkg"; then
+    echo "The requested version of CUDA is not available: CUDA $CUDA_VERSION"
+    exit 1
+fi
+
+# Ensure that the requested version of cuDNN is available AND compatible
+cudnn_pkg_version="libcudnn8=${CUDNN_VERSION}-1+cuda${CUDA_VERSION}"
+if ! apt-cache show "$cudnn_pkg_version"; then
+    echo "The requested version of cuDNN is not available: cuDNN $CUDNN_VERSION for CUDA $CUDA_VERSION"
+    exit 1
+fi
+
+echo "Installing CUDA libraries..."
+apt-get install -yq "$cuda_pkg"
+
+if [ "$INSTALL_CUDNN" = "true" ]; then
+    echo "Installing cuDNN libraries..."
+    apt-get install -yq "$cudnn_pkg_version"
+fi
+
+if [ "$INSTALL_NVTX" = "true" ]; then
+    echo "Installing NVTX..."
+    apt-get install -yq "$nvtx_pkg"
+fi
+
+echo "Done!"
diff --git a/test-scenarios/install_cudnn_nvxt.sh b/test-scenarios/install_cudnn_nvxt.sh
new file mode 100644
index 0000000..7b20f43
--- /dev/null
+++ b/test-scenarios/install_cudnn_nvxt.sh
@@ -0,0 +1,15 @@
+#!/bin/bash
+
+set -e
+
+# Optional: Import test library
+source dev-container-features-test-lib
+
+# Check installation of libcudnn8
+check "libcudnn.so.8" test 1 -eq "$(find /usr -name 'libcudnn.so.8' | wc -l)"
+
+# Check installation of cuda-nvtx-11-<version>
+check "cuda-11+nvtx" test -e '/usr/local/cuda-11/targets/x86_64-linux/include/nvtx3'
+
+# Report result
+reportResults
diff --git a/test-scenarios/install_cudnn_nvxt_version.sh b/test-scenarios/install_cudnn_nvxt_version.sh
new file mode 100644
index 0000000..a7f46bd
--- /dev/null
+++ b/test-scenarios/install_cudnn_nvxt_version.sh
@@ -0,0 +1,16 @@
+#!/bin/bash
+
+set -e
+
+# Optional: Import test library
+source dev-container-features-test-lib
+
+# Check installation of libcudnn8 (8.3.2)
+check "libcudnn.so.8.3.2" test 1 -eq "$(find /usr -name 'libcudnn.so.8.3.2' | wc -l)"
+
+# Check installation of cuda-nvtx-11-5 (11.5)
+check "cuda-11-5+nvtx" test -e '/usr/local/cuda-11.5/targets/x86_64-linux/include/nvtx3'
+
+# Report result
+reportResults
+    
\ No newline at end of file
diff --git a/test-scenarios/scenarios.json b/test-scenarios/scenarios.json
index 853da2f..4f45aba 100644
--- a/test-scenarios/scenarios.json
+++ b/test-scenarios/scenarios.json
@@ -91,5 +91,25 @@
                 "version": "3"
             }
         }
+    },
+    "install_cudnn_nvxt": {
+        "image": "debian",
+        "features": {
+            "nvidia-cuda": {
+                "installCudnn": true,
+                "installNvtx": true
+            }
+        }
+    },
+    "install_cudnn_nvxt_version": {
+        "image": "debian",
+        "features": {
+            "nvidia-cuda": {
+                "installCudnn": true,
+                "installNvtx": true,
+                "cudaVersion": "11.5",
+                "cudnnVersion": "8.3.2.44"
+            }
+        }
     }
 }
diff --git a/test/nvidia-cuda/test.sh b/test/nvidia-cuda/test.sh
new file mode 100644
index 0000000..ef50d7b
--- /dev/null
+++ b/test/nvidia-cuda/test.sh
@@ -0,0 +1,18 @@
+#!/bin/bash
+
+set -e
+
+# Optional: Import test library
+source dev-container-features-test-lib
+
+# Check installation of cuda-libraries-11-<version>
+check "libcudart.so.11.0" test 1 -eq "$(find /usr -name 'libcudart.so.11.0' | wc -l)"
+check "libcublas.so.11" test 1 -eq "$(find /usr -name 'libcublas.so.11' | wc -l)"
+check "libcublasLt.so.11" test 1 -eq "$(find /usr -name 'libcublasLt.so.11' | wc -l)"
+check "libcufft.so.10" test 1 -eq "$(find /usr -name 'libcufft.so.10' | wc -l)"
+check "libcurand.so.10" test 1 -eq "$(find /usr -name 'libcurand.so.10' | wc -l)"
+check "libcusolver.so.11" test 1 -eq "$(find /usr -name 'libcusolver.so.11' | wc -l)"
+check "libcusparse.so.11" test 1 -eq "$(find /usr -name 'libcusparse.so.11' | wc -l)"
+
+# Report result
+reportResults