Add new NVIDIA CUDA feature

.github/workflows/test-all.yaml

@@ -29,6 +29,7 @@ jobs:
             "java",
             "kubectl-helm-minikube",
             "node",
+            "nvidia-cuda",
             "oryx",
             "php",
             "powershell",

.github/workflows/test-pr.yaml

@@ -29,6 +29,7 @@ jobs:
             java: ./**/java/**
             kubectl-helm-minikube: ./**/kubectl-helm-minikube/**
             node: ./**/node/**
+            nvidia-cuda: ./**/nvidia-cuda/**
             oryx: ./**/oryx/**
             php: ./**/php/**
             powershell: ./**/powershell/**

src/nvidia-cuda/devcontainer-feature.json

@@ -0,0 +1,28 @@
+{
+  "id": "nvidia-cuda",
+  "name": "NVIDIA CUDA",
+  "description": "Installs NVIDIA CUDA libraries.",
+  "version": "1.0.0",
+  "options": {
+    "installCudnn": {
+      "type": "boolean",
+      "default": false,
+      "description": "Install CUDA Deep Neural Network (cuDNN) library"
+    },
+    "installNvtx": {
+      "type": "boolean",
+      "default": false,
+      "description": "Install NVIDIA Tools Extension (NVTX)"
+    },
+    "version": {
+      "type": "string",
+      "default": "latest",
+      "description": "Version of CUDA to install, formatted as MAJOR.MINOR (for example, 11.7)"
+    },
+    "cudnnVersion": {
+      "type": "string",
+      "default": "latest",
+      "description": "Full version of cuDNN to install (for example, 8.5.0.96-1)"
+    }
+  }
+}

src/nvidia-cuda/install.sh

@@ -0,0 +1,75 @@
+#!/usr/bin/env bash
+
+set -e
+
+INSTALL_CUDNN=${INSTALLCUDNN:-"false"}
+INSTALL_NVTX=${INSTALLNVTX:-"false"}
+CUDA_VERSION=${VERSION:-"latest"}
+CUDNN_VERSION=${CUDNNVERSION:-"latest"}
+
+# NVIDIA's package names include this information
+LATEST_CUDA_VERSION="11.7"
+LATEST_CUDNN_VERSION="8.5.0.96-1"
+if [ "$CUDA_VERSION" = "latest" ]; then CUDA_VERSION="$LATEST_CUDA_VERSION"; fi
+if [ "$CUDNN_VERSION" = "latest" ]; then CUDNN_VERSION="$LATEST_CUDNN_VERSION"; fi
+
+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
+
+# Function to run apt-get if needed
+apt_get_update_if_needed() {
+    if [ ! -d "/var/lib/apt/lists" ] || [ "$(ls /var/lib/apt/lists/ | wc -l)" = "0" ]; then
+        echo "Running apt-get update..."
+        apt-get update
+    else
+        echo "Skipping apt-get update."
+    fi
+}
+
+# Prints a more helpful error message when installation fails
+apt_try_install() {
+    apt-get install -yq "$1" || {
+        local exit_code=$?
+        echo "Failed to install $1"
+        echo "See $NVIDIA_REPO_URL for all available packages and versions"
+        return $exit_code
+    }
+}
+
+# Checks if packages are installed and installs them if not
+check_packages() {
+    if ! dpkg -s "$@" > /dev/null 2>&1; then
+        apt_get_update_if_needed
+        apt-get -y install --no-install-recommends "$@"
+    fi
+}
+
+check_packages 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
+
+echo "Installing CUDA libraries..."
+apt_try_install "cuda-libraries-${CUDA_VERSION/./-}"
+
+if [ "$INSTALL_CUDNN" = "true" ]; then
+    echo "Installing cuDNN libraries..."
+    apt_try_install "libcudnn8=${CUDNN_VERSION}+cuda${CUDA_VERSION}"
+fi
+
+if [ "$INSTALL_NVTX" = "true" ]; then
+    echo "Installing NVTX..."
+    apt_try_install "cuda-nvtx-${CUDA_VERSION/./-}"
+fi
+
+echo "Done!"

test-scenarios/install_cudnn_nvxt.sh

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

test-scenarios/scenarios.json

@@ -91,5 +91,16 @@
                 "version": "3"
             }
         }
+    },
+    "install_cudnn_nvxt": {
+        "image": "debian",
+        "features": {
+            "nvidia-cuda": {
+                "installCudnn": true,
+                "installNvtx": true,
+                "version": "11.5",
+                "cudnnVersion": "8.3.2.44-1"
+            }
+        }
     }
 }

test/nvidia-cuda/test.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+
+set -e
+
+# Optional: Import test library
+source dev-container-features-test-lib
+
+# Definition specific tests
+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