From 0f113540b7f2611a2b145c77f41e3311943b9c54 Mon Sep 17 00:00:00 2001 From: sulixu <24964493+sulixu@users.noreply.github.com> Date: Thu, 17 Sep 2026 15:20:27 -0700 Subject: [PATCH] fix(gpu): keep CUDA prebakes unregistered until node installation --- .../cloud-init/artifacts/cse_config_gpu.sh | 43 ++- .../artifacts/gpu_prebake_dkms_spec.sh | 129 +++++++ .../cloud-init/artifacts/gpu_prebake_spec.sh | 361 ++++++++++++++++++ vhdbuilder/packer/install-dependencies.sh | 18 +- vhdbuilder/packer/nvidia-prebake.md | 71 ++++ .../packer/test/linux-vhd-content-test.sh | 23 ++ 6 files changed, 633 insertions(+), 12 deletions(-) create mode 100644 spec/parts/linux/cloud-init/artifacts/gpu_prebake_dkms_spec.sh create mode 100644 spec/parts/linux/cloud-init/artifacts/gpu_prebake_spec.sh create mode 100644 vhdbuilder/packer/nvidia-prebake.md diff --git a/parts/linux/cloud-init/artifacts/cse_config_gpu.sh b/parts/linux/cloud-init/artifacts/cse_config_gpu.sh index 4fd20718e5b..197062d5703 100644 --- a/parts/linux/cloud-init/artifacts/cse_config_gpu.sh +++ b/parts/linux/cloud-init/artifacts/cse_config_gpu.sh @@ -132,10 +132,8 @@ validateGPUDrivers() { fi } -# logGPUDriverPrebakeReadiness emits a stage-1 observability signal on a managed GPU node: whether -# the aks-gpu prebake marker is present and matches this node's driver kind -- i.e. whether stage-2 -# (skip-build) would take the fast path. Lets the rollout confirm managed CUDA GPU nodes are ready -# before enabling consume. Observability only; no behavior change. +# logGPUDriverPrebakeReadiness reports the marker and driver kind on a managed GPU node. +# This is an observability signal, not permission to skip the normal installer. logGPUDriverPrebakeReadiness() { local marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}" local marker_present=false driver_kind_match=false m_kind node_kind @@ -188,6 +186,23 @@ cleanUpGridNodeCudaPrebake() { cleanUpPrebakedGPUDriver } +# Return success only if DKMS records the on-disk NVIDIA driver as installed for +# the running kernel. This does not register, build, install, or load a module. +isNvidiaDKMSInstalledForCurrentKernel() { + local version kernel arch status + kernel="$(uname -r)" || return 1 + arch="$(uname -m)" || return 1 + version="$(modinfo -k "${kernel}" -F version nvidia)" || return 1 + [ -n "${version}" ] || return 1 + status="$(dkms status -m nvidia -v "${version}" -k "${kernel}" -a "${arch}")" || return 1 + # DKMS may retain a backup of the non-DKMS prebake. That is not an error. + # Do not accept added/built records, other kernels, or module mismatch warnings. + case "${status% (original_module exists)}" in + "nvidia/${version}, ${kernel}, ${arch}: installed"|"nvidia, ${version}, ${kernel}, ${arch}: installed") return 0 ;; + esac + return 1 +} + ensureGPUDrivers() { if [ "$(isARM64)" -eq 1 ]; then return @@ -200,10 +215,24 @@ ensureGPUDrivers() { logs_to_events "AKS.CSE.ensureGPUDrivers.cleanUpGridNodeCudaPrebake" cleanUpGridNodeCudaPrebake || exit $ERR_GPU_DRIVERS_START_FAIL fi - if [ "${CONFIG_GPU_DRIVER_IF_NEEDED}" = true ]; then - logs_to_events "AKS.CSE.ensureGPUDrivers.configGPUDrivers" configGPUDrivers + local cuda_prebake=false + case "${OS}/${NVIDIA_GPU_DRIVER_TYPE:-}" in + "${UBUNTU_OS_NAME}"/cuda*) + if [ -f "${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}" ]; then + cuda_prebake=true + fi + ;; + esac + # A loadable prebake is not enough: managed CUDA nodes need a full DKMS + # installation for kernel updates. Keep this in nodePrep to cover PIS too. + if [ "${CONFIG_GPU_DRIVER_IF_NEEDED}" = true ] || { [ "${cuda_prebake}" = true ] && ! isNvidiaDKMSInstalledForCurrentKernel; }; then + logs_to_events "AKS.CSE.ensureGPUDrivers.configGPUDrivers" configGPUDrivers || exit $ERR_GPU_DRIVERS_START_FAIL else - logs_to_events "AKS.CSE.ensureGPUDrivers.validateGPUDrivers" validateGPUDrivers + logs_to_events "AKS.CSE.ensureGPUDrivers.validateGPUDrivers" validateGPUDrivers || exit $ERR_GPU_DRIVERS_START_FAIL + fi + if [ "${cuda_prebake}" = true ] && ! isNvidiaDKMSInstalledForCurrentKernel; then + echo "NVIDIA DKMS installation is incomplete for the running kernel" + exit $ERR_GPU_DRIVERS_START_FAIL fi if [ "$OS" = "$UBUNTU_OS_NAME" ]; then logs_to_events "AKS.CSE.ensureGPUDrivers.nvidia-modprobe" "systemctlEnableAndStart nvidia-modprobe 30" || exit $ERR_GPU_DRIVERS_START_FAIL diff --git a/spec/parts/linux/cloud-init/artifacts/gpu_prebake_dkms_spec.sh b/spec/parts/linux/cloud-init/artifacts/gpu_prebake_dkms_spec.sh new file mode 100644 index 00000000000..2526f10dc80 --- /dev/null +++ b/spec/parts/linux/cloud-init/artifacts/gpu_prebake_dkms_spec.sh @@ -0,0 +1,129 @@ +#!/bin/bash + +# Run in a disposable Linux container with DKMS installed, or set DKMS_TEST_BIN +# to a DKMS script. All driver state is a fixture; no NVIDIA module is compiled. +Describe 'CUDA prebake with real DKMS' + Include ./parts/linux/cloud-init/artifacts/cse_config_gpu.sh + DKMS_TEST_BIN="${DKMS_TEST_BIN:-$(command -v dkms || true)}" + Skip if 'requires root in a disposable Linux test environment' [ "$(id -u)" -ne 0 ] + Skip if 'requires DKMS or DKMS_TEST_BIN' [ ! -r "${DKMS_TEST_BIN}" ] + + setup() { + TEST_DIR="$(mktemp -d)" + DRIVER_VERSION=580.126.09 + RUNNING_KERNEL=6.8.0-prebaked + DKMS_TREE="${TEST_DIR}/dkms" + MODULE_TREE="${TEST_DIR}/modules" + SOURCE_TREE="${TEST_DIR}/sources" + mkdir -p "${DKMS_TREE}" "${SOURCE_TREE}/nvidia-${DRIVER_VERSION}" \ + "${TEST_DIR}/headers/include" "${MODULE_TREE}/${RUNNING_KERNEL}/updates/dkms" + cat > "${SOURCE_TREE}/nvidia-${DRIVER_VERSION}/dkms.conf" < "${SOURCE_TREE}/nvidia-${DRIVER_VERSION}/nvidia.c" + echo 'precompiled-module-fixture' > "${MODULE_TREE}/${RUNNING_KERNEL}/updates/dkms/nvidia.ko" + } + cleanup() { rm -rf "${TEST_DIR}"; } + BeforeEach setup + AfterEach cleanup + + uname() { case "$1" in -r) echo "${RUNNING_KERNEL}" ;; -m) echo x86_64 ;; esac; } + modinfo() { echo "${DRIVER_VERSION}"; } + dkms() { + bash "${DKMS_TEST_BIN}" --dkmstree "${DKMS_TREE}" \ + --sourcetree "${SOURCE_TREE}" --installtree "${MODULE_TREE}" \ + --kernelsourcedir "${TEST_DIR}/headers" "$@" + } + + # Create the files and links that DKMS status uses for an installed record. + # This verifies status parsing/discovery, not a real NVIDIA installation. + installed_fixture() { + dkms add -m nvidia -v "${DRIVER_VERSION}" > "${TEST_DIR}/add.log" 2>&1 || return 1 + local built="${DKMS_TREE}/nvidia/${DRIVER_VERSION}/${RUNNING_KERNEL}/x86_64" + mkdir -p "${built}/module" + cp "${MODULE_TREE}/${RUNNING_KERNEL}/updates/dkms/nvidia.ko" "${built}/module/nvidia.ko" + ln -s "${built}" "${DKMS_TREE}/nvidia/kernel-${RUNNING_KERNEL}-x86_64" + } + + failed_cleanup_then_update() { + cleanup_driver() { return 1; } + cleanup_driver || true + dkms autoinstall -k 6.8.0-security-update + } + + It 'does not select NVIDIA after failed CPU cleanup even when module and source files remain' + When call failed_cleanup_then_update + The status should be success + The output should not include 'nvidia' + The path "${DKMS_TREE}/nvidia" should not be exist + The path "${TEST_DIR}/build-attempted" should not be exist + The path "${MODULE_TREE}/${RUNNING_KERNEL}/updates/dkms/nvidia.ko" should be file + End + + It 'rejects a source-only added record' + dkms add -m nvidia -v "${DRIVER_VERSION}" > "${TEST_DIR}/add.log" 2>&1 + When call isNvidiaDKMSInstalledForCurrentKernel + The status should be failure + The output should equal '' + The path "${TEST_DIR}/build-attempted" should not be exist + End + + It 'accepts an installed record for the current kernel' + installed_fixture + When call isNvidiaDKMSInstalledForCurrentKernel + The status should be success + The output should equal '' + End + + It 'accepts an installed record with a backup of the original module' + installed_fixture + mkdir -p "${DKMS_TREE}/nvidia/original_module/${RUNNING_KERNEL}/x86_64" + When call isNvidiaDKMSInstalledForCurrentKernel + The status should be success + The output should equal '' + End + + It 'rejects installation for another kernel' + installed_fixture + RUNNING_KERNEL=6.8.0-security-update + When call isNvidiaDKMSInstalledForCurrentKernel + The status should be failure + The output should equal '' + End + + It 'rejects a mismatch between the DKMS cache and installed module' + installed_fixture + echo 'different-module' > "${MODULE_TREE}/${RUNNING_KERNEL}/updates/dkms/nvidia.ko" + When call isNvidiaDKMSInstalledForCurrentKernel + The status should be failure + The output should equal '' + End + + It 'does not build again for the already-installed kernel' + installed_fixture + When call dkms autoinstall -k "${RUNNING_KERNEL}" + The status should be success + The output should not include 'Building module' + The path "${TEST_DIR}/build-attempted" should not be exist + End + + attempt_kernel_update() { + # The fixture deliberately fails compilation. DKMS versions differ in + # their exit status here; the build sentinel is the assertion. + dkms autoinstall -k 6.8.0-security-update > "${TEST_DIR}/update.log" 2>&1 || true + } + + It 'selects NVIDIA for a new kernel after full installed state exists' + installed_fixture + When call attempt_kernel_update + The status should be success + The output should equal '' + The path "${TEST_DIR}/build-attempted" should be file + End +End diff --git a/spec/parts/linux/cloud-init/artifacts/gpu_prebake_spec.sh b/spec/parts/linux/cloud-init/artifacts/gpu_prebake_spec.sh new file mode 100644 index 00000000000..9a407b40490 --- /dev/null +++ b/spec/parts/linux/cloud-init/artifacts/gpu_prebake_spec.sh @@ -0,0 +1,361 @@ +#!/bin/bash + +Describe 'unregistered CUDA prebakes' + Include ./parts/linux/cloud-init/artifacts/cse_config_gpu.sh + + setup() { + TEST_DIR="$(mktemp -d)" + GPU_DKMS_MARKER_FILE="${TEST_DIR}/dkms-marker" + printf 'driver_kind=cuda\n' > "${GPU_DKMS_MARKER_FILE}" + } + cleanup() { command rm -rf "${TEST_DIR}"; } + BeforeEach setup + AfterEach cleanup + + OS=UBUNTU + UBUNTU_OS_NAME=UBUNTU + GPU_NODE=true + skip_nvidia_driver_install=false + NVIDIA_GPU_DRIVER_TYPE=cuda-lts + CONFIG_GPU_DRIVER_IF_NEEDED=true + NVIDIA_DRIVER_IMAGE=example/aks-gpu + NVIDIA_DRIVER_IMAGE_TAG=test + CTR_GPU_INSTALL_CMD='ctr run' + ERR_GPU_DRIVERS_START_FAIL=84 + + isARM64() { echo 0; } + logs_to_events() { shift; eval "$*"; } + systemctlEnableAndStart() { :; } + retrycmd_if_failure() { + shift 3 + echo "$*" + DKMS_STATUS="nvidia/580.126.09, 6.8.0-test, x86_64: installed" + } + uname() { case "$1" in -r) echo 6.8.0-test ;; -m) echo x86_64 ;; esac; } + modinfo() { echo 580.126.09; } + dkms() { echo "${DKMS_STATUS:-}"; } + configGPUDrivers() { installGPUDriverImage; } + validateGPUDrivers() { echo 'validation only'; } + cleanUpPrebakedGPUDriver() { echo 'CUDA cleanup'; } + + Describe 'GPU dispatch' + Parameters + true + false + End + It 'uses normal install for unregistered CUDA in both installation and validation-only mode' + CONFIG_GPU_DRIVER_IF_NEEDED="$1" + When call ensureGPUDrivers + The status should be success + The output should include '/entrypoint.sh install' + The output should not include 'install-skip-build' + The output should not include 'validation only' + End + End + + Describe 'unchanged paths' + It 'rejects an installer that returns success without completing DKMS installation' + retrycmd_if_failure() { :; } + When run ensureGPUDrivers + The status should equal 84 + The output should include 'NVIDIA DKMS installation is incomplete' + End + + It 'propagates container installation failure' + retrycmd_if_failure() { return 1; } + When call installGPUDriverImage + The status should be failure + The output should equal '' + End + + It 'propagates installation failure even with an older installed DKMS record' + DKMS_STATUS='nvidia/580.126.09, 6.8.0-test, x86_64: installed' + retrycmd_if_failure() { return 1; } + When run ensureGPUDrivers + The status should equal 84 + The output should equal '' + End + + It 'does not reinstall a complete DKMS installation in validation-only mode' + DKMS_STATUS='nvidia/580.126.09, 6.8.0-test, x86_64: installed' + CONFIG_GPU_DRIVER_IF_NEEDED=false + When call ensureGPUDrivers + The status should be success + The output should include 'validation only' + The output should not include '/entrypoint.sh' + End + + It 'keeps ordinary install on images without a prebake' + command rm "${GPU_DKMS_MARKER_FILE}" + When call ensureGPUDrivers + The status should be success + The output should include '/entrypoint.sh install' + The output should not include 'install-skip-build' + End + + It 'keeps validation-only behavior without a prebake' + command rm "${GPU_DKMS_MARKER_FILE}" + CONFIG_GPU_DRIVER_IF_NEEDED=false + When call ensureGPUDrivers + The status should be success + The output should include 'validation only' + The output should not include '/entrypoint.sh' + End + + It 'cleans CUDA before ordinary GRID installation' + NVIDIA_GPU_DRIVER_TYPE=grid-v20 + When call ensureGPUDrivers + The status should be success + The output should include 'CUDA cleanup' + The output should include '/entrypoint.sh install' + The output should not include 'install-skip-build' + End + + It 'does not change non-Ubuntu validation' + OS=AZURELINUX + CONFIG_GPU_DRIVER_IF_NEEDED=false + When call ensureGPUDrivers + The status should be success + The output should equal 'validation only' + End + End + + Describe 'current-kernel DKMS installation check' + Describe 'incomplete or mismatched installation' + Parameters + '' + 'nvidia/580.126.09: added' + 'nvidia/580.126.09, 6.8.0-test, x86_64: built' + 'nvidia/580.126.09, 6.8.0-old, x86_64: installed' + 'nvidia/580.126.09, 6.8.0-test, aarch64: installed' + 'nvidia/570.0.0, 6.8.0-test, x86_64: installed' + 'nvidia/580.126.09: broken' + 'nvidia/580.126.09, 6.8.0-test, x86_64: installed (WARNING! Diff between built and installed module!)' + End + It 'does not accept the record as installed' + DKMS_STATUS="$1" + When call isNvidiaDKMSInstalledForCurrentKernel + The status should be failure + The output should equal '' + End + + It 'repairs the record through normal installation on a validation-only node' + DKMS_STATUS="$1" + CONFIG_GPU_DRIVER_IF_NEEDED=false + When call ensureGPUDrivers + The status should be success + The output should include '/entrypoint.sh install' + The output should not include 'install-skip-build' + End + End + + Describe 'complete installation' + Parameters + 'nvidia/580.126.09, 6.8.0-test, x86_64: installed' + 'nvidia/580.126.09, 6.8.0-test, x86_64: installed (original_module exists)' + 'nvidia, 580.126.09, 6.8.0-test, x86_64: installed' + End + It 'accepts a current-kernel installed record including a retained prebake backup' + DKMS_STATUS="$1" + When call isNvidiaDKMSInstalledForCurrentKernel + The status should be success + The output should equal '' + End + End + + It 'queries the exact driver version, running kernel, and architecture' + dkms() { + echo "$*" > "${TEST_DIR}/dkms-call" + echo 'nvidia/580.126.09, 6.8.0-test, x86_64: installed' + } + When call isNvidiaDKMSInstalledForCurrentKernel + The status should be success + The contents of file "${TEST_DIR}/dkms-call" should equal 'status -m nvidia -v 580.126.09 -k 6.8.0-test -a x86_64' + End + + It 'rejects an unsuccessful DKMS query' + dkms() { return 1; } + When call isNvidiaDKMSInstalledForCurrentKernel + The status should be failure + The output should equal '' + End + + It 'rejects an unreadable module' + modinfo() { return 1; } + When call isNvidiaDKMSInstalledForCurrentKernel + The status should be failure + The output should equal '' + End + End + + It 'keeps the existing ARM64 early return' + isARM64() { echo 1; } + When call ensureGPUDrivers + The status should be success + The output should equal '' + End + + Describe 'PIS stage gates' + eval "$(sed -n '/^function nodePrep {/,/^}/p' ./parts/linux/cloud-init/artifacts/cse_main.sh)" + should_skip_nvidia_drivers() { echo "${LIVE_SKIP}"; } + basePrep() { echo 'unexpected basePrep'; return 1; } + reconcileVulnerableKernelModuleMitigation() { :; } + isAmdAmaEnabledNode() { return 1; } + checkServiceHealth() { :; } + systemctl() { return 1; } + retrycmd_if_failure() { + shift 3 + case "$*" in + *'/entrypoint.sh'*) + echo "$*" + DKMS_STATUS='nvidia/580.126.09, 6.8.0-test, x86_64: installed' + ;; + esac + } + logs_to_events() { + case "$1" in + AKS.CSE.ensureGPUDrivers) ensureGPUDrivers; exit "$?" ;; + AKS.CSE.ensureGPUDrivers.*) shift; eval "$*" ;; + AKS.CSE.cleanUpGPUDrivers) echo 'cleanup without registration'; exit 0 ;; + esac + } + provision_cached_image() { + touch "${TEST_DIR}/base_prep.complete" + eval "$(sed -n '/^if \[ ! -f \/opt\/azure\/containers\/base_prep.complete \]/,/^echo "Custom script finished."/p' \ + ./parts/linux/cloud-init/artifacts/cse_main.sh | + sed "s|/opt/azure/containers/base_prep.complete|${TEST_DIR}/base_prep.complete|g")" + } + API_SERVER_NAME=127.0.0.1 + PRE_PROVISION_ONLY=false + LIVE_SKIP=false + + It 'uses the live opt-out decision and installs with DKMS when basePrep is skipped' + skip_nvidia_driver_install=true + CONFIG_GPU_DRIVER_IF_NEEDED=false + When run provision_cached_image + The status should be success + The output should include 'Skipping basePrep' + The output should include '/entrypoint.sh install' + The output should not include 'install-skip-build' + The output should not include 'unexpected basePrep' + End + + It 'keeps validation when a PIS node already has a complete DKMS installation' + DKMS_STATUS='nvidia/580.126.09, 6.8.0-test, x86_64: installed' + CONFIG_GPU_DRIVER_IF_NEEDED=false + When run provision_cached_image + The status should be success + The output should include 'Skipping basePrep' + The output should include 'validation only' + The output should not include '/entrypoint.sh' + End + + It 'does not register during pre-provision-only execution' + PRE_PROVISION_ONLY=true + When run provision_cached_image + The status should be success + The output should include 'Skipping nodePrep' + The output should not include '/entrypoint.sh' + End + + Describe 'PIS CPU and opt-out' + Parameters + false false + true true + End + It 'takes cleanup without calling ensureGPUDrivers' + GPU_NODE="$1" + LIVE_SKIP="$2" + ensureGPUDrivers() { echo 'unexpected ensureGPUDrivers'; return 1; } + When run provision_cached_image + The status should be success + The output should include 'cleanup without registration' + The output should not include 'unexpected ensureGPUDrivers' + The output should not include '/entrypoint.sh' + End + End + End + + Describe 'VHD build contract' + FEATURE_FLAGS=NVIDIA_CUDA_PREBAKE + NVIDIA_DRIVER_IMAGE=example/aks-gpu + NVIDIA_DRIVER_IMAGE_TAG=test + apt_get_install() { :; } + retrycmd_if_failure() { touch "${TEST_DIR}/dkms-marker"; } + dkms() { echo "${DKMS_STATUS:-}"; return "${DKMS_EXIT:-0}"; } + run_prebake() { + VHD_LOGS_FILEPATH="${TEST_DIR}/vhd.log" + eval "$(sed -n '/^buildNVIDIAKernelModule() {/,/^}/p' ./vhdbuilder/packer/install-dependencies.sh | + sed "s|/opt/azure/aks-gpu/dkms-marker|${TEST_DIR}/dkms-marker|g; s|/var/lib/dkms/nvidia|${TEST_DIR}/nvidia|g")" + buildNVIDIAKernelModule + } + + It 'accepts an unregistered prebake' + When run run_prebake + The status should be success + The output should include 'Pre-building NVIDIA CUDA' + End + + It 'rejects an old container that leaves a registration' + DKMS_STATUS='nvidia/580.126.09, added' + When run run_prebake + The status should be failure + The output should include 'must not register NVIDIA' + End + + It 'rejects a dangling registration even when dkms status is empty' + ln -s "${TEST_DIR}/missing" "${TEST_DIR}/nvidia" + When run run_prebake + The status should be failure + The output should include 'must not register NVIDIA' + End + + It 'rejects an unsuccessful dkms status check' + DKMS_EXIT=1 + When run run_prebake + The status should be failure + The output should include 'Pre-building NVIDIA CUDA' + End + End + + Describe 'final VHD content validation' + FEATURE_FLAGS=NVIDIA_CUDA_PREBAKE + OS_SKU=Ubuntu + dkms() { echo "${DKMS_STATUS:-}"; return "${DKMS_EXIT:-0}"; } + modinfo() { return "${MODINFO_EXIT:-0}"; } + err() { echo "$*" >&2; } + validate_image() { + eval "$(sed -n '/^testNvidiaPrebakeUnregistered() {/,/^}/p' ./vhdbuilder/packer/test/linux-vhd-content-test.sh | + sed "s|/opt/azure/aks-gpu/dkms-marker|${TEST_DIR}/dkms-marker|g; s|/var/lib/dkms/nvidia|${TEST_DIR}/nvidia|g")" + testNvidiaPrebakeUnregistered + } + + It 'accepts compiled modules with no registration' + When call validate_image + The status should be success + The output should equal '' + End + + It 'rejects registration introduced after the build-only step' + mkdir -p "${TEST_DIR}/nvidia" + When call validate_image + The status should be failure + The stderr should include 'active NVIDIA DKMS registration' + End + + It 'rejects a missing compiled module' + MODINFO_EXIT=1 + When call validate_image + The status should be failure + The stderr should include 'compiled module is missing' + End + + It 'does not require a prebake on other image variants' + FEATURE_FLAGS=None + command rm "${GPU_DKMS_MARKER_FILE}" + When call validate_image + The status should be success + The output should equal '' + End + End +End diff --git a/vhdbuilder/packer/install-dependencies.sh b/vhdbuilder/packer/install-dependencies.sh index dd4d4cd18a8..218c0fae827 100644 --- a/vhdbuilder/packer/install-dependencies.sh +++ b/vhdbuilder/packer/install-dependencies.sh @@ -879,12 +879,12 @@ buildNVIDIAKernelModule() { if [ $OS = $UBUNTU_OS_NAME ] && [ "$(isARM64)" -ne 1 ]; then # No ARM64 SKU with GPU now gpu_action="copy" - # Opt-in: pre-build the NVIDIA kernel module into the VHD so node provisioning skips the - # ~100s in-CSE DKMS compile. The aks-gpu container is run in "build-only" mode: it compiles - # and DKMS-registers the kernel module + stages userspace libs against THIS VHD's kernel, + # Opt-in: pre-build the NVIDIA kernel module into the VHD. This safety-only stage + # does not enable skip-build at node boot. In "build-only" mode, aks-gpu compiles + # the kernel module + stages userspace libs without registering NVIDIA with DKMS, # performs NO device access (safe on the GPU-less Packer builder), and writes the marker - # /opt/azure/aks-gpu/dkms-marker. At node boot, configGPUDrivers passes "install-skip-build" - # when that marker matches, running only the device-dependent steps. + # /opt/azure/aks-gpu/dkms-marker. Managed CUDA nodes still use the normal + # "install" action with --dkms to establish a complete installation. # The driver image is intentionally LEFT in the VHD: boot-time device init still sources the # container toolkit debs, fabric manager, containerd runtime config and udev rules from it. # Dropping the image is a separate, deferred size optimization. @@ -901,6 +901,14 @@ buildNVIDIAKernelModule() { echo "Error: NVIDIA CUDA prebake did not produce /opt/azure/aks-gpu/dkms-marker" exit 1 fi + # Also enforce this in the builder: an older aks-gpu image can still register + # NVIDIA. Shared-image safety must not depend on successful node-time cleanup. + local nvidia_dkms_status + nvidia_dkms_status="$(dkms status -m nvidia)" || exit 1 + if [ -n "${nvidia_dkms_status}" ] || [ -e /var/lib/dkms/nvidia ] || [ -L /var/lib/dkms/nvidia ]; then + echo "Error: CUDA prebake must not register NVIDIA with DKMS; update the aks-gpu build image" + exit 1 + fi cat << EOF >> ${VHD_LOGS_FILEPATH} - nvidia-cuda-driver-prebaked=${NVIDIA_DRIVER_IMAGE_TAG} (kernel $(uname -r)) EOF diff --git a/vhdbuilder/packer/nvidia-prebake.md b/vhdbuilder/packer/nvidia-prebake.md new file mode 100644 index 00000000000..baccb8c3006 --- /dev/null +++ b/vhdbuilder/packer/nvidia-prebake.md @@ -0,0 +1,71 @@ +# NVIDIA prebakes without image-time DKMS registration + +The shared Ubuntu x86-64 VHD contains the compiled CUDA module and libraries, +but no NVIDIA registration in `/var/lib/dkms`. `aks-gpu` +`build-only` uses `--no-dkms` and installs modules under +`/lib/modules//updates/dkms`, the existing CPU/GRID cleanup path. +Both the build step and the final VHD content test reject active registration. +An older build container that still registers NVIDIA fails the build. +The companion container change is [Azure/aks-gpu#191](https://github.com/Azure/aks-gpu/pull/191). + +CPU and opted-out nodes do not invoke the GPU installer. Even if artifact cleanup +fails, normal DKMS kernel-update discovery has no NVIDIA registration to rebuild. +This does not remove all driver files, prevent module loading, or block an +explicit administrator request to add/install the remaining NVIDIA sources. +Existing registered images and PIS caches still need their existing cleanup or +replacement; changing the image build does not repair deployed nodes. + +## Managed CUDA node flow + +`nodePrep` resolves the live GPU opt-out setting, including on PIS nodes that skip +`basePrep`. For a marked CUDA prebake, CSE checks whether DKMS records the on-disk +driver as installed for the running kernel and architecture. A loadable module, +an `added` record, or a `built` record is not enough. + +If installation is requested, or the prebake has no complete DKMS installation, +CSE runs the existing container `install` action. The normal NVIDIA installer +uses `--dkms`; it compiles, installs, and registers the driver. Device and +container-runtime setup remain unchanged. If validation-only mode already has +a complete installation, CSE keeps the existing validation path. + +CSE checks the installed state after either path. An incomplete installation or +a DKMS module mismatch fails provisioning with the existing GPU error code. +The check permits DKMS's harmless `original_module exists` backup notice. +GRID, non-Ubuntu, ARM64, and unmarked validation paths keep their existing behavior. + +There is an existing GRID cleanup limit. Cleanup can return success while leaving +CUDA files in place. In validation-only mode, if the remaining driver passes +`nvidia-modprobe` and `nvidia-smi`, CSE can finish without running the GRID installer. +This change does not fix that path. A mocked control-flow test shows the path is +possible; it has not been verified on a real GRID node. + +This is a safety-only change. It does not use `dkms add`, move DKMS caches, or +enable `install-skip-build`. A fresh unregistered CUDA image requires a full +installation, even in validation-only mode. This can remove the startup-time +saving from validating an already-loadable prebake. No startup-time measurement +has been made. Nodes with a complete DKMS installation can still validate it. The +container's unused skip-build mode is unchanged and needs separate work before +AgentBaker can enable it for these unregistered images. + +## Release order and validation + +1. Ship the updated CSE to all consumers first, including validation-only and + PIS paths. The new unregistered VHD needs this check on those paths. Runtime + containers can keep their existing normal `install --dkms` behavior. +2. Publish the updated `aks-gpu` build image and update its builder reference, then + enable these builder checks and publish the new VHDs. This source change does + not invent a new image tag or publish an image. The existing + `NVIDIA_CUDA_PREBAKE` flag is retained; no parking/layout flag is introduced. +3. Rebuild PIS caches from compatible images. For rollback, restore compatible + VHD/PIS image selection before rolling CSE back. + +Before release, test actual VHD builds on Ubuntu 22.04 and 24.04; CPU security +patching with intentionally failed cleanup; managed CUDA install and +validation-only/PIS paths; kernel/driver mismatches; GRID and opt-out paths; and +GPU kernel update plus reboot. Unit tests cannot establish hardware correctness. + +Local tests cover the installer arguments, no-registration build checks, normal +CUDA dispatch, PIS gates, and rejection of incomplete DKMS states. The separate +`gpu_prebake_dkms_spec.sh` uses actual DKMS with fixture source/module trees to +check status parsing and kernel-update discovery. It requires root and DKMS in +a disposable Linux test environment; it does not compile a real NVIDIA driver. diff --git a/vhdbuilder/packer/test/linux-vhd-content-test.sh b/vhdbuilder/packer/test/linux-vhd-content-test.sh index 6a9b6c185cc..91faf5e78b1 100644 --- a/vhdbuilder/packer/test/linux-vhd-content-test.sh +++ b/vhdbuilder/packer/test/linux-vhd-content-test.sh @@ -2726,6 +2726,28 @@ testContainerNetworkingPluginsInstalled() { return 0 } +# Reject an active registration at final image validation too, after all build steps. +# CPU/opt-out cleanup must not be needed to make kernel security updates safe. +testNvidiaPrebakeUnregistered() { + local test="testNvidiaPrebakeUnregistered" + if [ "$OS_SKU" != "Ubuntu" ] || ! grep -q NVIDIA_CUDA_PREBAKE <<< "$FEATURE_FLAGS"; then + return 0 + fi + local status + if ! status="$(dkms status -m nvidia)"; then + err "$test" "Could not verify NVIDIA DKMS state" + return 1 + fi + if [ -n "$status" ] || [ -e /var/lib/dkms/nvidia ] || [ -L /var/lib/dkms/nvidia ]; then + err "$test" "Shared CUDA prebake contains an active NVIDIA DKMS registration" + return 1 + fi + if [ ! -f /opt/azure/aks-gpu/dkms-marker ] || ! modinfo nvidia >/dev/null 2>&1; then + err "$test" "CUDA prebake marker or compiled module is missing" + return 1 + fi +} + # As we call these tests, we need to bear in mind how the test results are processed by the # the caller in run-tests.sh. That code uses az vm run-command invoke to run this script # on a VM. It then looks at stderr to see if any errors were reported. Notably it doesn't @@ -2740,6 +2762,7 @@ testContainerNetworkingPluginsInstalled() { checkPerformanceData testBccTools $OS_SKU $OS_VERSION testVHDBuildLogsExist +testNvidiaPrebakeUnregistered testAzureLinuxNvidiaGPUDriverReleaseNotes testCriticalTools testPackagesInstalled