From 3318db970f5ccfd7a29a3f44e7da07adf78f8552 Mon Sep 17 00:00:00 2001 From: Ganeshkumar Ashokavardhanan Date: Tue, 14 Jul 2026 17:10:10 -0700 Subject: [PATCH] fix(gpu): tear down prebaked GPU driver on version mismatch, not just kind mismatch Generalizes the managed-GPU prebake teardown from grid-only kind mismatch (cleanUpGridNodeCudaPrebake, #8919) to any prebake/node driver mismatch (cleanUpMismatchedPrebakedGPUDriver): 1. driver-KIND mismatch (cuda prebake on a GRID node) -- unchanged behavior. 2. driver-VERSION mismatch, same kind -- NEW. If the baked cuda(-lts) prebake's driver version differs from the version this node installs (GPU_DV), the stale prebaked module + /usr/bin/lib64 libs would collide with the fresh install -> "Failed to initialize NVML: Driver/library version mismatch". Now torn down before install. The dkms-marker previously recorded only driver_kind=, so CSE could not detect same-kind version drift (the gap #8915 flagged and deferred). This adds driver_version= to the marker at VHD-build time (install-dependencies.sh, from NVIDIA_DRIVER_IMAGE_TAG with the image-SHA suffix stripped to match CSE's GPU_DV) and compares it at node boot. Backward compatible (VHDs live ~6 months): markers written before this change have no driver_version= line; when absent, CSE cannot compare versions and falls back to kind-only behavior -- identical to the prior cleanUpGridNodeCudaPrebake. No teardown is ever triggered by a missing version field. Adds normalizeGPUDriverKind helper (cuda-lts->cuda, grid-v20->grid) and shellspec coverage for kind mismatch, version mismatch, version match no-op, missing-version fallback, unknown-GPU_DV no-op, and the existing kind/OS/marker-absence no-ops. make shellspec: 142 examples, 0 failures. make generate: no testdata drift (cse_config.sh is VHD-build-uploaded, not embedded in CSE-gen snapshots). --- .../linux/cloud-init/artifacts/cse_config.sh | 76 +++++++++---- .../cloud-init/artifacts/cse_config_spec.sh | 101 ++++++++++++++++-- vhdbuilder/packer/install-dependencies.sh | 8 ++ 3 files changed, 151 insertions(+), 34 deletions(-) diff --git a/parts/linux/cloud-init/artifacts/cse_config.sh b/parts/linux/cloud-init/artifacts/cse_config.sh index 357bebbf056..9b643a193be 100755 --- a/parts/linux/cloud-init/artifacts/cse_config.sh +++ b/parts/linux/cloud-init/artifacts/cse_config.sh @@ -1321,33 +1321,62 @@ logGPUDriverPrebakeReadiness() { echo "AKS_GPU_PREBAKE event=managed_gpu driver_type=${NVIDIA_GPU_DRIVER_TYPE:-} marker_present=${marker_present} driver_kind_match=${driver_kind_match}" } -# cleanUpGridNodeCudaPrebake tears down a cuda-lts driver pre-baked into the shared Ubuntu VHD when -# THIS node installs a GRID driver. The shared VHD bakes only the cuda(-lts) driver + a DKMS marker; -# a GRID/converged (A10, NVv5) node then installs the grid driver on top, and the stale prebaked -# cuda module + its /usr/bin/lib64 userspace libs collide with the grid driver -> nvidia-smi fails -# with "Failed to initialize NVML: Driver/library version mismatch". This is a pure driver-KIND -# mismatch (grid vs cuda), so no version comparison is needed. A legacy marker with no driver_kind= -# line is treated as a cuda prebake (the only kind the VHD bakes today). No-op unless the node is -# GRID and a prebake marker exists. Reuses cleanUpPrebakedGPUDriver (from cse_install_ubuntu.sh) for -# the actual removal. NAP/agentpool cuda nodes are intentionally untouched here. -cleanUpGridNodeCudaPrebake() { +# normalizeGPUDriverKind maps an AgentBaker driver-type (grid, grid-v20, cuda, cuda-lts) to the +# aks-gpu marker's driver_kind family (grid or cuda) -- the container's DRIVER_KIND build arg. Image +# variants "cuda-lts"/"grid-v20" bake markers as "cuda"/"grid". Unknown values pass through unchanged +# so a mismatch is still detectable rather than silently normalized to a match. +normalizeGPUDriverKind() { + case "${1:-}" in + cuda*) echo "cuda" ;; + grid*) echo "grid" ;; + *) echo "${1:-}" ;; + esac +} + +# cleanUpMismatchedPrebakedGPUDriver tears down a driver pre-baked into the shared Ubuntu VHD when it +# does NOT match what THIS managed GPU node installs, so the stale prebaked module + its +# /usr/bin/lib64 userspace libs cannot collide with the freshly-installed driver (which otherwise +# surfaces as "Failed to initialize NVML: Driver/library version mismatch"). Two mismatch classes: +# +# 1. KIND mismatch (e.g. cuda prebake on a GRID node): the shared VHD bakes a cuda(-lts) driver; +# a GRID/converged (A10, NVv5) node installs the grid driver on top. Different driver families +# always collide -> tear down. (This is the case #8919 originally fixed.) +# 2. VERSION mismatch, same kind (e.g. cuda-lts R580 prebake vs a cuda R470 install on NCv1, or any +# future cuda-vs-cuda version skew where the baked VHD driver lags the CSE-requested version): +# same family but different driver version -> stale module still collides -> tear down. +# +# Backward compatibility (VHDs live ~6 months): a marker written before this change has no +# driver_version= line. When the marker version is absent we CANNOT compare versions, so we fall +# back to kind-only behavior (tear down on kind mismatch, keep on kind match) -- identical to the +# prior cleanUpGridNodeCudaPrebake. A legacy marker with no driver_kind= line is treated as a cuda +# prebake (the only kind the VHD bakes today). No-op unless a prebake marker exists. Reuses +# cleanUpPrebakedGPUDriver (from cse_install_ubuntu.sh) for the actual removal. +cleanUpMismatchedPrebakedGPUDriver() { [ "$OS" = "$UBUNTU_OS_NAME" ] || return 0 local marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}" [ -f "${marker}" ] || return 0 - local node_kind m_kind - case "${NVIDIA_GPU_DRIVER_TYPE:-}" in - grid*) node_kind=grid ;; - *) return 0 ;; - esac + local node_kind m_kind m_version reason + node_kind="$(normalizeGPUDriverKind "${NVIDIA_GPU_DRIVER_TYPE:-}")" + # An empty driver_kind= (legacy marker) is treated as cuda -- the only family the VHD bakes today. m_kind="$(sed -n 's/^driver_kind=//p' "${marker}" | head -n1)" - - # Keep only when the prebake is explicitly grid (matches this grid node). An empty marker kind is - # a legacy cuda prebake; a "cuda" marker is a cuda prebake -- both mismatch a grid node, tear down. - if [ "${m_kind}" = "grid" ]; then + [ -n "${m_kind}" ] || m_kind=cuda + # driver_version= is absent on markers written before this field was added; empty => skip the + # version comparison and fall back to kind-only behavior (no false teardown on old VHDs). + m_version="$(sed -n 's/^driver_version=//p' "${marker}" | head -n1)" + + if [ -n "${node_kind}" ] && [ "${m_kind}" != "${node_kind}" ]; then + reason=driver_kind_mismatch + elif [ -n "${m_version}" ] && [ -n "${GPU_DV:-}" ] && [ "${m_version}" != "${GPU_DV}" ]; then + # Same kind but the baked driver version differs from what this node installs (GPU_DV is the + # AgentBaker-requested driver version, cse_helpers.sh). Only compared when both are known. + reason=driver_version_mismatch + else + # Match (or not enough info to declare a mismatch): keep the prebake for the skip-build path. return 0 fi - echo "AKS_GPU_PREBAKE event=grid_cuda_prebake_teardown driver_type=${NVIDIA_GPU_DRIVER_TYPE:-} marker_kind=${m_kind:-none} node_kind=${node_kind} action=teardown" + + echo "AKS_GPU_PREBAKE event=prebake_mismatch_teardown driver_type=${NVIDIA_GPU_DRIVER_TYPE:-} node_kind=${node_kind:-none} marker_kind=${m_kind:-none} marker_version=${m_version:-none} requested_version=${GPU_DV:-none} reason=${reason} action=teardown" cleanUpPrebakedGPUDriver } @@ -1356,11 +1385,12 @@ ensureGPUDrivers() { return fi - # Tear down a mismatched cuda-lts VHD prebake before a GRID node installs its own driver, or the - # stale module/libs collide with the grid driver (NVML version mismatch). Runs before the dispatch + # Tear down a mismatched VHD prebake before this node installs its own managed driver, or the + # stale prebaked module/libs collide with it (NVML version mismatch). Covers driver-kind mismatch + # (e.g. cuda prebake on a GRID node) and same-kind driver-version skew. Runs before the dispatch # below so it covers both the configGPUDrivers and validateGPUDrivers paths. if [ "$OS" = "$UBUNTU_OS_NAME" ]; then - logs_to_events "AKS.CSE.ensureGPUDrivers.cleanUpGridNodeCudaPrebake" cleanUpGridNodeCudaPrebake || exit $ERR_GPU_DRIVERS_START_FAIL + logs_to_events "AKS.CSE.ensureGPUDrivers.cleanUpMismatchedPrebakedGPUDriver" cleanUpMismatchedPrebakedGPUDriver || exit $ERR_GPU_DRIVERS_START_FAIL fi if [ "${CONFIG_GPU_DRIVER_IF_NEEDED}" = true ]; then diff --git a/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh index 39c19e39355..17a3c529413 100755 --- a/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh +++ b/spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh @@ -74,20 +74,23 @@ Describe 'cse_config.sh' End End - Describe 'cleanUpGridNodeCudaPrebake' + Describe 'cleanUpMismatchedPrebakedGPUDriver' # Stub the actual removal so tests assert the keep-vs-teardown DECISION without touching the # real filesystem. OS defaults to Ubuntu (the only path this function acts on). OS="$UBUNTU_OS_NAME" - # shellcheck disable=SC2329 # invoked dynamically by cleanUpGridNodeCudaPrebake. + # shellcheck disable=SC2329 # invoked dynamically by cleanUpMismatchedPrebakedGPUDriver. cleanUpPrebakedGPUDriver() { echo "STUB_TEARDOWN_CALLED"; } + # --- driver-KIND mismatch (the original #8919 grid-vs-cuda behavior, preserved) --- + It 'tears down a cuda prebake before installing a GRID driver (A10/GRID outage path)' marker="$(mktemp)" printf 'driver_kind=cuda\n' > "$marker" GPU_DKMS_MARKER_FILE="$marker" NVIDIA_GPU_DRIVER_TYPE="grid" - When call cleanUpGridNodeCudaPrebake + When call cleanUpMismatchedPrebakedGPUDriver The output should include "action=teardown" + The output should include "reason=driver_kind_mismatch" The output should include "marker_kind=cuda" The output should include "node_kind=grid" The output should include "STUB_TEARDOWN_CALLED" @@ -99,8 +102,9 @@ Describe 'cse_config.sh' printf 'driver_kind=cuda\n' > "$marker" GPU_DKMS_MARKER_FILE="$marker" NVIDIA_GPU_DRIVER_TYPE="grid-v20" - When call cleanUpGridNodeCudaPrebake + When call cleanUpMismatchedPrebakedGPUDriver The output should include "action=teardown" + The output should include "reason=driver_kind_mismatch" The output should include "STUB_TEARDOWN_CALLED" rm -f "$marker" End @@ -110,19 +114,79 @@ Describe 'cse_config.sh' printf 'kernel=5.15.0-1114-azure\n' > "$marker" # no driver_kind= line GPU_DKMS_MARKER_FILE="$marker" NVIDIA_GPU_DRIVER_TYPE="grid" - When call cleanUpGridNodeCudaPrebake + When call cleanUpMismatchedPrebakedGPUDriver The output should include "action=teardown" - The output should include "marker_kind=none" + The output should include "reason=driver_kind_mismatch" The output should include "STUB_TEARDOWN_CALLED" rm -f "$marker" End - It 'is a no-op on a CUDA node (leaves the cuda prebake for the version-match/library-bump path)' + # --- driver-VERSION mismatch, same kind (new behavior) --- + + It 'tears down a same-kind cuda prebake whose baked version differs from the requested version' + marker="$(mktemp)" + printf 'driver_kind=cuda\ndriver_version=580.159.04\n' > "$marker" + GPU_DKMS_MARKER_FILE="$marker" + NVIDIA_GPU_DRIVER_TYPE="cuda-lts" + GPU_DV="470.256.02" # node installs a different cuda driver version than baked + When call cleanUpMismatchedPrebakedGPUDriver + The output should include "action=teardown" + The output should include "reason=driver_version_mismatch" + The output should include "marker_version=580.159.04" + The output should include "requested_version=470.256.02" + The output should include "STUB_TEARDOWN_CALLED" + GPU_DV="" + rm -f "$marker" + End + + It 'is a no-op when the baked version matches the requested version' + marker="$(mktemp)" + printf 'driver_kind=cuda\ndriver_version=580.159.04\n' > "$marker" + GPU_DKMS_MARKER_FILE="$marker" + NVIDIA_GPU_DRIVER_TYPE="cuda-lts" + GPU_DV="580.159.04" + When call cleanUpMismatchedPrebakedGPUDriver + The output should not include "STUB_TEARDOWN_CALLED" + The status should be success + GPU_DV="" + rm -f "$marker" + End + + # --- backward compatibility: markers with no driver_version= (VHDs built before this change) --- + + It 'falls back to kind-only (keeps a same-kind cuda prebake) when the marker has no version' + marker="$(mktemp)" + printf 'driver_kind=cuda\n' > "$marker" # no driver_version= line (old VHD) + GPU_DKMS_MARKER_FILE="$marker" + NVIDIA_GPU_DRIVER_TYPE="cuda-lts" + GPU_DV="470.256.02" # even a different requested version cannot be compared -> keep + When call cleanUpMismatchedPrebakedGPUDriver + The output should not include "STUB_TEARDOWN_CALLED" + The status should be success + GPU_DV="" + rm -f "$marker" + End + + It 'does not compare versions when the requested version (GPU_DV) is unknown' + marker="$(mktemp)" + printf 'driver_kind=cuda\ndriver_version=580.159.04\n' > "$marker" + GPU_DKMS_MARKER_FILE="$marker" + NVIDIA_GPU_DRIVER_TYPE="cuda-lts" + GPU_DV="" # requested version not set -> cannot compare -> keep + When call cleanUpMismatchedPrebakedGPUDriver + The output should not include "STUB_TEARDOWN_CALLED" + The status should be success + rm -f "$marker" + End + + # --- matches and no-ops --- + + It 'is a no-op on a CUDA node whose kind matches (leaves the cuda prebake for skip-build)' marker="$(mktemp)" printf 'driver_kind=cuda\n' > "$marker" GPU_DKMS_MARKER_FILE="$marker" NVIDIA_GPU_DRIVER_TYPE="cuda-lts" - When call cleanUpGridNodeCudaPrebake + When call cleanUpMismatchedPrebakedGPUDriver The output should not include "STUB_TEARDOWN_CALLED" The status should be success rm -f "$marker" @@ -133,7 +197,7 @@ Describe 'cse_config.sh' printf 'driver_kind=grid\n' > "$marker" GPU_DKMS_MARKER_FILE="$marker" NVIDIA_GPU_DRIVER_TYPE="grid" - When call cleanUpGridNodeCudaPrebake + When call cleanUpMismatchedPrebakedGPUDriver The output should not include "STUB_TEARDOWN_CALLED" The status should be success rm -f "$marker" @@ -142,7 +206,7 @@ Describe 'cse_config.sh' It 'is a no-op when no prebake marker exists' GPU_DKMS_MARKER_FILE="$(mktemp)"; rm -f "${GPU_DKMS_MARKER_FILE}" NVIDIA_GPU_DRIVER_TYPE="grid" - When call cleanUpGridNodeCudaPrebake + When call cleanUpMismatchedPrebakedGPUDriver The output should not include "STUB_TEARDOWN_CALLED" The status should be success End @@ -153,7 +217,7 @@ Describe 'cse_config.sh' GPU_DKMS_MARKER_FILE="$marker" OS="MARINER" # override the Ubuntu default set at the Describe level NVIDIA_GPU_DRIVER_TYPE="grid" - When call cleanUpGridNodeCudaPrebake + When call cleanUpMismatchedPrebakedGPUDriver The output should not include "STUB_TEARDOWN_CALLED" The status should be success OS="$UBUNTU_OS_NAME" # restore for any subsequent examples @@ -161,6 +225,21 @@ Describe 'cse_config.sh' End End + Describe 'normalizeGPUDriverKind' + It 'maps cuda and cuda-lts to cuda' + When call normalizeGPUDriverKind "cuda-lts" + The output should equal "cuda" + End + It 'maps grid and grid-v20 to grid' + When call normalizeGPUDriverKind "grid-v20" + The output should equal "grid" + End + It 'passes through an unknown driver type unchanged' + When call normalizeGPUDriverKind "something-else" + The output should equal "something-else" + End + End + Describe 'configureAzureJson' AZURE_JSON_PATH="azure.json" AKS_CUSTOM_CLOUD_JSON_PATH="customcloud.json" diff --git a/vhdbuilder/packer/install-dependencies.sh b/vhdbuilder/packer/install-dependencies.sh index 3762d1141c5..2f8399b37bb 100644 --- a/vhdbuilder/packer/install-dependencies.sh +++ b/vhdbuilder/packer/install-dependencies.sh @@ -760,6 +760,14 @@ EOF echo "Error: NVIDIA CUDA prebake did not produce /opt/azure/aks-gpu/dkms-marker" exit 1 fi + # Record the baked driver VERSION in the marker so node-boot CSE can detect a version mismatch + # (baked prebake vs the version this node installs) and tear down a stale prebake before it + # collides -- see cleanUpMismatchedPrebakedGPUDriver in cse_config.sh. The aks-gpu container tag + # is "-" (e.g. 580.159.04-20260629214430); strip the trailing image + # SHA so the recorded value matches CSE's GPU_DV (GPUDriverVersion). Appended, not overwriting the + # container-written driver_kind= line. Older VHDs omit this line; CSE falls back to kind-only there. + PREBAKED_DRIVER_VERSION="${NVIDIA_DRIVER_IMAGE_TAG%-*}" + echo "driver_version=${PREBAKED_DRIVER_VERSION}" >> /opt/azure/aks-gpu/dkms-marker cat << EOF >> ${VHD_LOGS_FILEPATH} - nvidia-cuda-driver-prebaked=${NVIDIA_DRIVER_IMAGE_TAG} (kernel $(uname -r)) EOF