fix(gpu): unload full nvidia module stack in prebake cleanup, not just idle bare module

[ganeshkumarashok]

Summary

Follow-up to #8933. That PR added an nvidia-module unload to cleanUpPrebakedGPUDriver (the --gpu-driver None / skip-driver path) but guarded it behind refcnt==0 && no /dev/nvidia* — which only works for the exact bare-nvidia-only case that was reproduced, and silently no-ops on the common full-stack case.

The bug

if [ "$(cat /sys/module/nvidia/refcnt)" = "0" ] && ! ls /dev/nvidia* ; then
    for mod in nvidia_uvm nvidia_drm nvidia_modeset nvidia_peermem nvidia; do rmmod ...; done
fi

/sys/module/nvidia/refcnt counts the modules that depend on nvidia (nvidia_modeset, nvidia_uvm, nvidia_drm). On a normal cuda/cuda-lts auto-load the whole stack is present, so nvidia's refcnt is 1–3, not 0 → the guard fails → the loop never runs. The dependents in that loop are exactly what keep nvidia's refcnt nonzero, so the guard blocks the very unloads that would free it. Net: on the realistic case the stale stack stays fully resident — the thing #8933 set out to remove.

Second defect: ! ls /dev/nvidia* as a "busy" proxy is wrong — udev / nvidia-modprobe create /dev/nvidia* on an idle module, so this also wrongly blocks legitimately-idle nodes.

(Credit: caught in review — the existing test only exercised a bare nvidia module, so it passed unrealistically.)

The fix

• Drop both pre-checks. Unload unconditionally in dependency order — dependents (nvidia_drm, nvidia_uvm, nvidia_modeset, nvidia_peermem) first, then nvidia. rmmod without -f already refuses to unload a module in genuine use (returns nonzero, module stays), so no guard is needed to stay safe: a truly busy stack simply remains loaded and is reported below.
• Stop nvidia-persistenced first. If the prebake started it, it holds an open device handle — not a module dependency, so refcnt wouldn't even reflect it. This mirrors NVIDIA's own gpu-driver-container _unload_driver.
• Best-effort throughout; never aborts provisioning. If something is genuinely busy, module_after=true → status=incomplete → the marker is retained so the next provision retries (existing mechanism, unchanged).

Why not fail CSE on incomplete cleanup?

Considered and rejected as the default: a --gpu-driver None node doesn't need the driver to join and run non-GPU pods; the stale module only bites a later GPU-Operator install (whose own _unload_driver || exit 1 already surfaces a loud CrashLoop signal). Hard-failing CSE would turn a not-yet-a-problem into provisioning failure / autoscaler churn. With this fix, "incomplete" should be essentially unreachable for idle stacks. If stricter behavior is wanted for specific customers, it should be an opt-in toggle, not the default.

Tests

• Replaces the bare-module test (passed unrealistically) with:
  • full auto-loaded stack (nvidia + modeset + uvm + drm, refcnt shown as 3): asserts persistenced stopped first, dependents unloaded before nvidia, and the module is gone → status=cleaned.
  • genuinely busy module (rmmod returns nonzero): asserts rmmod is attempted (no self-defeating guard), module stays → module_after=true → status=incomplete → marker_after=true (retained).
• make shellspec (docker/bash): 5 examples, 0 failures.
• make generate: no testdata drift (cse_install_ubuntu.sh is VHD-build-uploaded, not embedded in CSE-gen snapshots). shellcheck: clean in the changed region.

Risk

🟢 Low. Scoped to the opted-out / non-GPU cleanup path (never runs when AKS installs a driver). Strictly more correct than the merged guard; rmmod's in-use protection bounds the blast radius, and the change cannot abort provisioning.

Related: #8933 (this fixes its guard), #8919 (grid teardown), #8945 (version-drift teardown).

[Copilot]

Pull request overview

This PR fixes GPU prebake cleanup on Ubuntu nodes that opt out of the AKS-managed NVIDIA driver (--gpu-driver None / skip-driver path) by reliably unloading an auto-loaded full NVIDIA kernel module stack, rather than silently no-op’ing due to a self-defeating refcnt==0 + /dev/nvidia* guard.

Changes:

• Make cleanUpPrebakedGPUDriver stop nvidia-persistenced (best-effort) and attempt to rmmod the NVIDIA module chain unconditionally (dependents first), relying on rmmod’s in-use protection instead of pre-check guards.
• Update ShellSpec coverage to exercise a realistic “full stack auto-loaded” scenario and a “busy module refuses to unload” scenario.

Reviewed changes

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

File	Description
parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh	Removes the refcnt/device-node guard and unloads the full NVIDIA module stack in dependency order, stopping nvidia-persistenced first.
spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_spec.sh	Updates tests to cover full-stack unload and busy-module retry/marker retention behavior.

[ganeshkumarashok]

Good catch — fixed in 4dc18b6. The test now enforces order two ways: (1) the rmmod mock models real kernel semantics (rmmod nvidia FAILS while any dependent is still loaded, exact-token removal instead of the substring //nvidia/ that was corrupting nvidia_modeset→_modeset), and (2) explicit The line N of output should equal assertions for the full sequence (persistenced stop → nvidia_drm/uvm/modeset/peermem → nvidia). Verified with a negative control: reversing the code's loop to put nvidia first makes this test fail.

[Copilot]

Pull request overview

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