fix(gpu): don't wipe a customer's driver when cleaning the prebake on --gpu-driver None by xuexu6666 · Pull Request #9257 · Azure/AgentBaker

Xu Xue (xuexu6666) · 2026-08-19T21:09:23Z

Problem

Since #8786 + #8803, the CUDA driver kernel module is pre-baked into the shared x86 gen2 Ubuntu VHD, so every x86 gen2 Ubuntu node — GPU-managed, --gpu-driver None, and non-GPU — boots with an installed, DKMS-registered driver + the marker /opt/azure/aks-gpu/dkms-marker. cleanUpPrebakedGPUDriver tears it down on nodes that don't install the AKS-managed driver.

That teardown is gated only on the marker's presence and then blanket-removes any nvidia driver. So on a --gpu-driver None GPU node where the customer baked their own driver (via PreparedImageSpecification / a custom image on top of the shared VHD), it wipes the customer's nvidia-smi, DKMS module and libs at CSE — leaving a driverless node. Workaround today: rm -f /opt/azure/aks-gpu/dkms-marker in the bake.

Fix — scope the teardown to AKS's own baked version (don't weaken it)

The marker records the version AKS baked: aks-gpu writes driver_version= together with the DKMS module, so the marker version always equals AKS's DKMS-registered version. Use it to tell AKS's driver apart from a customer's:

Installed DKMS version == marker (or nothing else present) → AKS's pre-bake is what's here → full teardown runs, unchanged. This covers non-GPU, GPU-Operator / no-own-driver, and grid nodes — so the safety teardowns stay intact: fix(gpu): unload resident prebaked nvidia module in cleanUpPrebakedGPUDriver #8933 (unload the resident auto-loaded module or the GPU Operator crashloops), fix(gpu): tear down cuda-lts VHD prebake on GRID nodes before managed install #8919 (grid cuda-prebake NVML mismatch), and dead-weight / attack-surface removal. AKS's own baked driver is always cleaned up.
A different version is DKMS-registered → a customer replaced the pre-bake with their own. Remove only AKS's baked version's DKMS registration if it lingers, and preserve the customer's module + userspace.

This keeps the teardown's protective default for every SKU (it is never skipped for AKS's own driver) while no longer wiping a customer-owned driver. No new API/tag, no heuristic that could suppress a needed cleanup.

Testing

ShellSpec: 2 new cases — preserve a customer-replaced driver while still cleaning AKS's baked version, and tear down on match; the 5 existing cleanUpPrebakedGPUDriver cases pass unchanged (their markers carry no driver_version=, so the new path is inert). 11 examples, 0 failures.
shellcheck -s bash clean; bash -n parses.
The install script isn't embedded in generated testdata (verified plain + gzip/base64-decoded across pkg/, apiserver/, aks-node-controller/), so no make generate diff.

Edge cases

If a customer re-bakes the exact same version AKS baked, it looks AKS-owned and is torn down (rare — no reason to re-bake an identical version).
Relies on the customer driver being DKMS-registered (the .run --dkms / GPU-Operator path). A non-DKMS .run install isn't detected; noted as a follow-up if needed.

cc Ganeshkumar Ashokavardhanan (@ganeshkumarashok) — this refines the --gpu-driver None teardown from #8786/#8803 (and keeps #8933/#8919 intact). Draft pending GPU-team review.

github-actions · 2026-08-19T21:13:32Z

Windows Unit Test Results

3 files 13 suites 50s ⏱️
404 tests 404 ✅ 0 💤 0 ❌
407 runs 407 ✅ 0 💤 0 ❌

Results for commit bb96abd.

♻️ This comment has been updated with latest results.

… --gpu-driver None cleanUpPrebakedGPUDriver tears down the CUDA driver pre-baked into the shared Ubuntu VHD on nodes that don't install the AKS-managed driver -- non-GPU VMs and GPU nodes opted out via --gpu-driver None. It is gated only on AKS's marker /opt/azure/aks-gpu/dkms-marker (present on the shared VHD regardless of the pool's driver setting) and then blanket-removes ANY nvidia driver. So on a --gpu-driver None GPU node where the customer baked their OWN driver (via PreparedImageSpecification or a custom image), it wipes the customer's nvidia-smi, DKMS module and libs at CSE, leaving a driverless node. Scope the teardown to AKS's own baked version. The marker records the version AKS baked -- aks-gpu writes driver_version= together with the DKMS module, so it equals AKS's DKMS-registered version. If a DIFFERENT version is DKMS-registered, a customer replaced the pre-bake with their own: remove only AKS's baked version's DKMS registration (if it lingers) and preserve the customer's module + userspace. When AKS's baked version is still the installed one (non-GPU, GPU-Operator / no-own-driver, grid), the full teardown runs unchanged -- so the safety teardowns (#8933 resident nvidia module unload, #8919 grid cuda-prebake teardown, dead-weight / attack-surface removal) are never skipped for AKS's own driver. ShellSpec: 2 new cases (preserve customer driver while still cleaning AKS's version; tear down on match); the 5 existing cleanUpPrebakedGPUDriver cases are unchanged. shellcheck clean; the install script is not embedded in generated testdata, so there is no `make generate` diff. Signed-off-by: xuexu6666 <xuex@microsoft.com>

…r scan Identify a customer-replaced driver by comparing the marker's baked version against the driver actually effective on the node (loaded module -> on-disk .ko -> DKMS registration) via new getInstalledNvidiaDriverVersion, so .run and custom-image installs are covered, not just DKMS-registered ones. On a version mismatch, clean only AKS-attributable residue (version-scoped DKMS dir + /usr/bin/lib64) and leave the customer's driver intact; same-version or undetectable falls through to full teardown. Grid caller passes force_full_teardown so the cuda-vs-grid KIND teardown stays unconditional.

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Suppressed comments (1)

parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh:320

🔴 High Risk — Script Logic / Backward Compatibility: The loaded module is not a reliable ownership signal here because the AKS-prebaked module can auto-load before CSE (as noted below in this function). If that loaded module is the marker version while a customer has registered a different DKMS version, this helper returns the marker version and the caller falls through to the blanket teardown, deleting the customer's /var/lib/dkms/nvidia registration and userspace—the exact case this PR intends to preserve. Determine ownership by examining all DKMS registrations for a version different from the marker before treating a loaded/on-disk marker-version module as AKS-only, and add coverage for loaded=marker plus DKMS=customer.

    if [ -r /sys/module/nvidia/version ]; then
        v="$(cat /sys/module/nvidia/version 2>/dev/null)"
    fi
    if [ -z "${v}" ]; then
        v="$(modinfo -F version nvidia 2>/dev/null | head -n1)"

Addresses Copilot review: when AKS's marker-version residue and a customer's different-version driver coexist (a lingering AKS DKMS dir next to the customer's, or AKS's module auto-loaded at boot while the customer's build is DKMS-registered), returning the first/most-authoritative version could yield the marker version and trigger a full teardown that wipes the customer driver. Replace getInstalledNvidiaDriverVersion (single effective version) with findCustomerDriverVersion, which scans every signal -- loaded module, on-disk .ko, and ALL DKMS registrations -- and returns any version differing from the marker. The decision is now order-independent. Add ShellSpec coverage for both glob orderings and the authoritative-source-reports-marker case.

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

Addresses Copilot high-risk review: when a customer's different DKMS version triggers preserve but the loaded/on-disk module is still AKS's own (== marker), the early return left AKS's module resident or on disk next to the customer's userspace -- recreating the NVML version mismatch the teardown exists to prevent. The preserve branch now, gated strictly to the marker version, also rmmod's AKS's idle loaded module and removes its on-disk .ko, while leaving the customer's different-version driver untouched. Split out loadedNvidiaModuleVersion so the loaded-module signal is mockable; add ShellSpec coverage.

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

Addresses two Copilot medium-risk findings: - Removing /usr/bin/lib64 left stale ld.so.cache entries a preserved customer nvidia-smi could still resolve AKS's NVML libs through; run ldconfig after, as the full teardown does. - The preserve path always reported preserved_customer_driver even when AKS's marker-version module was busy / the rmmod failed and stayed resident. Recheck the loaded marker version after the scoped strip and emit status=incomplete (the #8933 security-coverage alert + retry signal) when AKS residue remains; the marker is left in place so the next provision retries. Add ShellSpec for the busy/failing-unload case.

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

…rve path Addresses Copilot medium-risk: deleting AKS's marker-version module from updates/dkms leaves modules.dep pointing at the removed path, so a later modprobe nvidia could fail to resolve the PRESERVED customer module. Run depmod for the running kernel after the scoped .ko removal; unlike the full teardown (driver-free, no depmod needed), this path expects an nvidia module to remain usable. Assert depmod in the ShellSpec case.

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

…serve Addresses Copilot medium-risk: the .ko removal was scoped to $(uname -r), so AKS's marker-version module survived under other installed kernels (e.g. a patched VHD), and booting/falling back to one would load the stale module against the customer's userspace. Sweep every /lib/modules/*/updates/dkms tree, remove only files whose embedded version equals the marker (customer's different-version module preserved), and depmod each affected kernel. Add a multi-kernel ShellSpec case (GPU_MODULES_DIR seam) proving per-file version gating keeps the customer's module.

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

Addresses two Copilot medium-risk findings on the preserve path: - depmod ran after the first module removed, but more nvidia*.ko for the same kernel were deleted afterward, so modules.dep could reference deleted files. Collect touched kernels and depmod each ONCE after the whole removal sweep. - The completeness gate checked only the loaded module; a failed best-effort removal still reported preserved_customer_driver while AKS residue remained. Track removal-command failures (DKMS dir + per-kernel .ko) and, combined with the resident-module check, emit status=incomplete when any AKS residue survives (renamed the field aks_module_after -> aks_residue). Added ShellSpec for the multi-module-single-depmod and removal-failure cases.

Copilot

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

Xu Xue (xuexu6666) force-pushed the xuxue/gpu-driver-none-preserve-custom branch from fc2ce1c to c61d65e Compare August 19, 2026 21:24

Xu Xue (xuexu6666) changed the title ~~fix(gpu): preserve a customer-replaced driver on --gpu-driver None nodes~~ fix(gpu): don't wipe a customer's driver when cleaning the prebake on --gpu-driver None Aug 20, 2026

Xu Xue (xuexu6666) force-pushed the xuxue/gpu-driver-none-preserve-custom branch from c61d65e to 1a54b2c Compare August 20, 2026 00:01

Xu Xue (xuexu6666) marked this pull request as ready for review August 26, 2026 15:50

Copilot AI balanced review requested due to automatic review settings August 26, 2026 15:50

Copilot started reviewing on behalf of Xu Xue (xuexu6666) August 26, 2026 16:02 View session