fix(gpu): tear down mismatched VHD-prebaked driver before managed install

parts/linux/cloud-init/artifacts/cse_config.sh

@@ -1219,6 +1219,13 @@ installGPUDriverImage() {
 configGPUDrivers() {
     if [ "$OS" = "$UBUNTU_OS_NAME" ]; then
         waitForContainerdReady || exit $ERR_GPU_DRIVERS_START_FAIL
+        # If the shared VHD prebaked a driver that is not an EXACT match (kind AND version) for what
+        # this node installs -- e.g. the cuda-lts prebake on a NAP cuda node (same kind, different
+        # version) or on a grid/converged A10 node (different kind) -- tear it down first, otherwise
+        # the stale prebaked module collides with the driver installed below and nvidia-smi fails with
+        # an NVML driver/library version mismatch. No-op when there is no marker or on an exact match
+        # (the consume fast path).
+        cleanUpMismatchedPrebakedGPUDriver
         mkdir -p /opt/{actions,gpu}
         # The driver image is normally pre-pulled into the VHD; only hit the registry when it is
         # actually missing so provisioning doesn't pay a redundant manifest/layer round trip.
@@ -1295,20 +1302,68 @@ validateGPUDrivers() {
     fi
 }
 
+# gpuDriverKindFromType maps the AgentBaker driver-type ($NVIDIA_GPU_DRIVER_TYPE) to the aks-gpu
+# marker's driver_kind (the container's DRIVER_KIND build arg): image variants "cuda-lts" and
+# "grid-v20" bake markers as "cuda"/"grid" respectively. Echoes the mapped kind on stdout.
+gpuDriverKindFromType() {
+    case "${1}" in
+        cuda*) echo cuda ;;
+        grid*) echo grid ;;
+        *) echo "${1}" ;;
+    esac
+}
+
+# prebakedGPUDriverVersion echoes the version of the driver DKMS-registered by the VHD prebake, read
+# from the on-disk DKMS source tree (/var/lib/dkms/nvidia/<version>). Empty if none is registered.
+# The prebake marker itself records only driver_kind, not the version, so the DKMS tree is the
+# authoritative source for what version is actually baked in.
+prebakedGPUDriverVersion() {
+    local d
+    for d in /var/lib/dkms/nvidia/*/; do
+        [ -d "${d}" ] || continue
+        basename "${d}"
+        return 0
+    done
+}
+
+# cleanUpMismatchedPrebakedGPUDriver tears down a VHD-prebaked driver on a managed GPU node unless it
+# EXACTLY matches the driver this node is about to install -- same kind (marker driver_kind vs mapped
+# $NVIDIA_GPU_DRIVER_TYPE) AND same version (on-disk DKMS version vs $GPU_DV). A kind-only match is
+# not enough: a NAP "cuda" node inherits the shared VHD's "cuda-lts" prebake (both kind=cuda) but at
+# a different version (e.g. 580.126.09 vs 580.159.04), so the stale libs in /usr/bin/lib64 collide
+# with the freshly installed driver and nvidia-smi fails with "Failed to initialize NVML:
+# Driver/library version mismatch". GRID/converged A10 nodes mismatch on kind. Only an exact
+# kind+version match (the consume fast path this node is meant to take) is kept. No-op when the
+# marker is absent (VHD without prebake). Reuses cleanUpPrebakedGPUDriver (sourced from
+# cse_install_ubuntu.sh) for the actual removal.
+cleanUpMismatchedPrebakedGPUDriver() {
+    local marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}"
+    [ -f "${marker}" ] || return 0
+    local m_kind node_kind prebaked_ver node_ver
+    m_kind="$(sed -n 's/^driver_kind=//p' "${marker}" | head -n1)"
+    node_kind="$(gpuDriverKindFromType "${NVIDIA_GPU_DRIVER_TYPE}")"
+    prebaked_ver="$(prebakedGPUDriverVersion)"
+    node_ver="${GPU_DV:-}"
+    # Keep ONLY on an exact kind+version match, and only when every field is known -- an unprovable
+    # match (empty kind or empty version on either side) falls through to teardown, since an unusable
+    # prebaked module is worse than a redundant rebuild.
+    if [ -n "${m_kind}" ] && [ "${m_kind}" = "${node_kind}" ] && \
+       [ -n "${prebaked_ver}" ] && [ -n "${node_ver}" ] && [ "${prebaked_ver}" = "${node_ver}" ]; then
+        echo "AKS_GPU_PREBAKE event=mismatch_teardown driver_type=${NVIDIA_GPU_DRIVER_TYPE:-} marker_kind=${m_kind} node_kind=${node_kind} prebaked_version=${prebaked_ver} node_version=${node_ver} action=keep"
+        return 0
+    fi
+    echo "AKS_GPU_PREBAKE event=mismatch_teardown driver_type=${NVIDIA_GPU_DRIVER_TYPE:-} marker_kind=${m_kind:-} node_kind=${node_kind} prebaked_version=${prebaked_ver:-} node_version=${node_ver:-} action=teardown"
+    cleanUpPrebakedGPUDriver
+}
+
 # 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() {
     local marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}"
     local marker_present=false driver_kind_match=false m_kind node_kind
-    # Map the AgentBaker driver-type to the aks-gpu marker's driver_kind (the container's DRIVER_KIND
-    # build arg): image variants "cuda-lts" and "grid-v20" bake markers as "cuda"/"grid" respectively.
-    case "${NVIDIA_GPU_DRIVER_TYPE}" in
-        cuda*) node_kind=cuda ;;
-        grid*) node_kind=grid ;;
-        *) node_kind="${NVIDIA_GPU_DRIVER_TYPE}" ;;
-    esac
+    node_kind="$(gpuDriverKindFromType "${NVIDIA_GPU_DRIVER_TYPE}")"
     if [ -f "${marker}" ]; then
         marker_present=true
         m_kind="$(sed -n 's/^driver_kind=//p' "${marker}" | head -n1)"

parts/linux/cloud-init/nodecustomdata.yml

@@ -34,6 +34,16 @@ write_files:
     {{GetVariableProperty "cloudInitData" "initAKSCustomCloud"}}
 {{end}}
 
+
+# ---- hotfix: auto-generated by hotfix-generate GH Action ----
+- path: {{GetCSEConfigScriptFilepath}}
+  permissions: "0744"
+  encoding: gzip
+  owner: root
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionConfigs"}}
+
+# ---- end hotfix ----
 {{- else }}
 - path: {{GetCSEHelpersScriptFilepath}}
   permissions: "0744"

spec/parts/linux/cloud-init/artifacts/cse_config_spec.sh

@@ -74,6 +74,77 @@ Describe 'cse_config.sh'
         End
     End
 
+    Describe 'cleanUpMismatchedPrebakedGPUDriver'
+        # Stub the actual removal so tests assert the keep-vs-teardown DECISION without touching the
+        # real filesystem, and stub the on-disk prebaked version so tests control it without a real
+        # DKMS tree. Both stubs print sentinels / values we can match on.
+        cleanUpPrebakedGPUDriver() { echo "STUB_TEARDOWN_CALLED"; }
+        prebakedGPUDriverVersion() { echo "${STUB_PREBAKED_VER:-}"; }
+
+        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="cuda"; GPU_DV="580.159.04"; STUB_PREBAKED_VER="580.159.04"
+            When call cleanUpMismatchedPrebakedGPUDriver
+            The output should not include "STUB_TEARDOWN_CALLED"
+            The status should be success
+        End
+
+        It 'keeps the prebaked driver on an exact kind+version match (agentpool cuda-lts consume path)'
+            marker="$(mktemp)"
+            printf 'driver_kind=cuda\n' > "$marker"
+            GPU_DKMS_MARKER_FILE="$marker"
+            NVIDIA_GPU_DRIVER_TYPE="cuda-lts"; GPU_DV="580.159.04"; STUB_PREBAKED_VER="580.159.04"
+            When call cleanUpMismatchedPrebakedGPUDriver
+            The output should include "action=keep"
+            The output should not include "STUB_TEARDOWN_CALLED"
+            rm -f "$marker"
+        End
+
+        It 'tears down on same kind but different version (NAP cuda node inheriting cuda-lts prebake)'
+            marker="$(mktemp)"
+            printf 'driver_kind=cuda\n' > "$marker"
+            GPU_DKMS_MARKER_FILE="$marker"
+            NVIDIA_GPU_DRIVER_TYPE="cuda"; GPU_DV="580.126.09"; STUB_PREBAKED_VER="580.159.04"
+            When call cleanUpMismatchedPrebakedGPUDriver
+            The output should include "action=teardown"
+            The output should include "STUB_TEARDOWN_CALLED"
+            rm -f "$marker"
+        End
+
+        It 'tears down a cuda prebaked driver on a GRID node (kind mismatch; A10/GRID outage path)'
+            marker="$(mktemp)"
+            printf 'driver_kind=cuda\n' > "$marker"
+            GPU_DKMS_MARKER_FILE="$marker"
+            NVIDIA_GPU_DRIVER_TYPE="grid"; GPU_DV="570.00.00"; STUB_PREBAKED_VER="580.159.04"
+            When call cleanUpMismatchedPrebakedGPUDriver
+            The output should include "action=teardown"
+            The output should include "STUB_TEARDOWN_CALLED"
+            rm -f "$marker"
+        End
+
+        It 'tears down when the marker lacks driver_kind (cannot prove a match)'
+            marker="$(mktemp)"
+            printf 'kernel=6.8.0-1059-azure\n' > "$marker"   # no driver_kind= line
+            GPU_DKMS_MARKER_FILE="$marker"
+            NVIDIA_GPU_DRIVER_TYPE="cuda"; GPU_DV="580.159.04"; STUB_PREBAKED_VER="580.159.04"
+            When call cleanUpMismatchedPrebakedGPUDriver
+            The output should include "action=teardown"
+            The output should include "STUB_TEARDOWN_CALLED"
+            rm -f "$marker"
+        End
+
+        It 'tears down when the prebaked on-disk version cannot be determined (empty)'
+            marker="$(mktemp)"
+            printf 'driver_kind=cuda\n' > "$marker"
+            GPU_DKMS_MARKER_FILE="$marker"
+            NVIDIA_GPU_DRIVER_TYPE="cuda"; GPU_DV="580.159.04"; STUB_PREBAKED_VER=""
+            When call cleanUpMismatchedPrebakedGPUDriver
+            The output should include "action=teardown"
+            The output should include "STUB_TEARDOWN_CALLED"
+            rm -f "$marker"
+        End
+    End
+
     Describe 'configureAzureJson'
         AZURE_JSON_PATH="azure.json"
         AKS_CUSTOM_CLOUD_JSON_PATH="customcloud.json"