Add new NVIDIA CUDA feature

src/nvidia/devcontainer-feature.json

@@ -0,0 +1,17 @@
+{
+  "id": "nvidia",
+  "name": "Nvidia",
+  "description": "Installs Nvidia's CUDA libraries.",
+  "options": {
+    "install_cuda": {
+      "type": "boolean",
+      "default": true,
+      "description": "Install CUDA libraries"
+    },
+    "install_cudnn": {
+      "type": "boolean",
+      "default": true,
+      "description": "Install cuDNN libraries"
+    }
+  }
+}

src/nvidia/install.sh

@@ -0,0 +1,36 @@
+#!/usr/bin/env bash
+
+set -e
+
+install_cuda=${INSTALL_CUDA:-"true"}
+install_cudnn=${INSTALL_CUDNN:-"true"}
+
+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
+
+# Download Nvidia's keyring package
+ubuntu_version="$(lsb_release -sr | sed 's/\.//g')"
+keyring_repo="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu$ubuntu_version/x86_64"
+keyring_package="cuda-keyring_1.0-1_all.deb"
+keyring_package_url="$keyring_repo/$keyring_package"
+keyring_package_path="$(mktemp -d)"
+keyring_package_file="$keyring_package_path/$keyring_package"
+wget -O "$keyring_package_file" "$keyring_package_url"
+
+# Install Nvidia's keyring package
+apt-get install -yq "$keyring_package_file"
+apt-get update -yq
+
+if [ "$install_cuda" = "true" ]; then
+    echo "Installing CUDA libraries..."
+    apt-get install -yq cuda-libraries-11-7
+    echo "Successfully installed CUDA libraries!"
+fi
+
+if [ "$install_cudnn" = "true" ]; then
+    echo "Installing cuDNN libraries..."
+    apt-get install -yq libcudnn8
+    echo "Successfully installed cuDNN libraries!"
+fi

test/nvidia/test.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+set -e
+
+# Optional: Import test library
+source dev-container-features-test-lib
+
+# Definition specific tests
+check "libcudart.so.11.0" [ 1 -eq "$(find /usr -name 'libcudart.so.11.0' | wc -l)" ]
+check "libcublas.so.11" [ 1 -eq "$(find /usr -name 'libcublas.so.11' | wc -l)" ]
+check "libcublasLt.so.11" [ 1 -eq "$(find /usr -name 'libcublasLt.so.11' | wc -l)" ]
+check "libcufft.so.10" [ 1 -eq "$(find /usr -name 'libcufft.so.10' | wc -l)" ]
+check "libcurand.so.10" [ 1 -eq "$(find /usr -name 'libcurand.so.10' | wc -l)" ]
+check "libcusolver.so.11" [ 1 -eq "$(find /usr -name 'libcusolver.so.11' | wc -l)" ]
+check "libcusparse.so.11" [ 1 -eq "$(find /usr -name 'libcusparse.so.11' | wc -l)" ]
+check "libcudnn.so.8" [ 1 -eq "$(find /usr -name 'libcudnn.so.8' | wc -l)" ]
+
+# Report result
+reportResults