fix(gpu): tear down prebaked GPU driver on version mismatch, not just kind mismatch

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

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,14 +217,29 @@ 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
             rm -f "$marker"
         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"

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 "<driver_version>-<image_sha>" (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