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

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

@@ -240,9 +240,9 @@ removeNvidiaRepos() {
 # The prebaked module MAY already be loaded when we run: cuda(-lts) prebakes auto-load nvidia.ko at
 # boot (~5s, well before CSE), so on a cuda/cuda-lts SKU opted out via --gpu-driver None the module
 # is resident even though ensureGPUDrivers never ran. (grid prebakes do not auto-load, so grid nodes
-# arrive here with no module.) Deleting the on-disk .ko then leaves a stale loaded module -- unused
-# (refcnt 0, no /dev/nvidia*) but resident until reboot, and a landmine for a subsequent GPU Operator
-# install. So we rmmod it first, when idle, before removing the files. No-op unless the marker exists.
+# arrive here with no module.) Deleting the on-disk .ko then leaves a stale loaded module -- a landmine
+# for a subsequent GPU Operator install (its driver-container's unload-then-load fails against the
+# resident module). So we rmmod it first, before removing the files. No-op unless the marker exists.
 cleanUpPrebakedGPUDriver() {
     local marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}"
     if [ ! -f "${marker}" ]; then
@@ -252,18 +252,22 @@ cleanUpPrebakedGPUDriver() {
     local dkms_before=false module_before=false module_after=false
     [ -d /var/lib/dkms/nvidia ] && dkms_before=true
 
-    # Unload the prebaked nvidia module if it auto-loaded at boot (cuda/cuda-lts SKUs). Only when idle
-    # (refcnt 0 and no device nodes) -- this node doesn't install a driver, so nothing should be using
-    # it; if something is, leave it and let module_after=true flag an incomplete teardown. Unload the
-    # dependent modules first (modeset/uvm/drm) so nvidia's refcnt drops to 0. Best-effort; failures
-    # do not abort provisioning.
+    # Unload the prebaked nvidia module stack if it auto-loaded at boot (cuda/cuda-lts SKUs). The
+    # dependent modules (nvidia_modeset/uvm/drm/peermem) hold a reference on nvidia -- i.e. they are
+    # exactly what keeps /sys/module/nvidia/refcnt nonzero -- so we must unload THEM first, then
+    # nvidia. (A refcnt==0 pre-check would be self-defeating: on a normal auto-loaded stack nvidia's
+    # refcnt is already >=1 from those dependents, so it would skip the very unloads that free it.)
+    # rmmod without -f refuses to unload a module in genuine use (a running GPU process pinning it),
+    # so this is safe: a truly busy stack simply stays loaded and is reported as incomplete below.
+    # nvidia-persistenced (if the prebake started it) holds an open device handle -- not a module
+    # dependency, so refcnt would not reflect it -- so stop it first, mirroring NVIDIA's own
+    # gpu-driver-container _unload_driver. Best-effort throughout; failures do not abort provisioning.
     if lsmod | grep -q '^nvidia'; then
         module_before=true
-        if [ "$(cat /sys/module/nvidia/refcnt 2>/dev/null || echo 0)" = "0" ] && ! ls /dev/nvidia* >/dev/null 2>&1; then
-            for mod in nvidia_uvm nvidia_drm nvidia_modeset nvidia_peermem nvidia; do
-                rmmod "${mod}" 2>/dev/null || true
-            done
-        fi
+        systemctl stop nvidia-persistenced 2>/dev/null || true
+        for mod in nvidia_drm nvidia_uvm nvidia_modeset nvidia_peermem nvidia; do
+            rmmod "${mod}" 2>/dev/null || true
+        done
     fi
     lsmod | grep -q '^nvidia' && module_after=true
 

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

@@ -52,38 +52,66 @@ Describe 'cse_install_ubuntu.sh'
             The output should include "module_after=false"
         End
 
-        It 'unloads an idle prebaked nvidia module that auto-loaded at boot (cuda/cuda-lts SKUs)'
+        It 'unloads a full auto-loaded module stack: dependents (modeset/uvm/drm) before nvidia'
+            # Regression guard for the refcnt trap: a real cuda/cuda-lts auto-load leaves nvidia with
+            # refcnt>=1 because nvidia_modeset/uvm/drm depend on it. The unload must remove those
+            # dependents FIRST, then nvidia. (A refcnt==0 pre-check would skip the whole stack and
+            # leave nvidia resident -- the exact bug this replaces.)
             marker="$(mktemp)"
             GPU_DKMS_MARKER_FILE="${marker}"
             ldconfig() { echo "mock ldconfig"; }
-            # simulate a loaded-but-idle module: lsmod shows nvidia until rmmod is "run"
-            _nvidia_loaded=true
-            lsmod() { if [ "${_nvidia_loaded}" = true ]; then echo "nvidia 104165376 0"; else echo ""; fi; }
-            cat() { echo "0"; }        # /sys/module/nvidia/refcnt = 0 (idle)
-            ls() { return 1; }          # no /dev/nvidia* device nodes
-            rmmod() { _nvidia_loaded=false; echo "mock rmmod $*"; }
+            systemctl() { echo "mock systemctl $*"; }
+            # Realistic kernel model: dependents each hold a ref on nvidia. rmmod of a dependent
+            # succeeds; rmmod of nvidia FAILS (returns 1, module stays) while any dependent is still
+            # loaded -- so if the code attempted nvidia first, it would fail and nvidia_after=true, and
+            # the ordering assertions below would not see the required sequence. Space-delimited set
+            # with guard spaces so removal is exact-token (avoids the "nvidia" substring stripping
+            # "nvidia_modeset"). Each rmmod line is emitted so line-based order can be asserted.
+            _loaded=" nvidia nvidia_modeset nvidia_uvm nvidia_drm nvidia_peermem "
+            _deps_loaded() { case "${_loaded}" in *" nvidia_modeset "*|*" nvidia_uvm "*|*" nvidia_drm "*|*" nvidia_peermem "*) return 0;; *) return 1;; esac; }
+            lsmod() { case "${_loaded}" in *" nvidia "*) echo "nvidia 104165376 4";; *) echo "";; esac; }
+            rmmod() {
+                if [ "$1" = "nvidia" ] && _deps_loaded; then
+                    echo "mock rmmod nvidia FAILED (deps still loaded)"; return 1
+                fi
+                _loaded="${_loaded/ $1 / }"; echo "mock rmmod $1"
+            }
             When call cleanUpPrebakedGPUDriver
             The status should be success
-            The output should include "mock rmmod nvidia"
             The output should include "module_before=true"
+            # persistenced stopped first (holds a device handle, not a module dependency)
+            The output should include "mock systemctl stop nvidia-persistenced"
+            # ORDER-SENSITIVE: the loop is nvidia_drm, nvidia_uvm, nvidia_modeset, nvidia_peermem, nvidia.
+            # A wrong order (nvidia before its deps) would emit "FAILED" and leave module_after=true.
+            The line 2 of output should equal "mock systemctl stop nvidia-persistenced"
+            The line 3 of output should equal "mock rmmod nvidia_drm"
+            The line 4 of output should equal "mock rmmod nvidia_uvm"
+            The line 5 of output should equal "mock rmmod nvidia_modeset"
+            The line 6 of output should equal "mock rmmod nvidia_peermem"
+            The line 7 of output should equal "mock rmmod nvidia"
+            The output should not include "FAILED"
+            # after the full-stack unload the module is gone -> complete
             The output should include "module_after=false"
             The output should include "status=cleaned"
         End
 
-        It 'keeps the marker (incomplete) when a busy nvidia module cannot be unloaded'
+        It 'keeps the marker (incomplete) when a genuinely busy module refuses to unload'
+            # rmmod without -f refuses an in-use module and returns nonzero; nvidia stays loaded. The
+            # teardown must report module_after=true and KEEP the marker so the next provision retries.
             marker="$(mktemp)"
             GPU_DKMS_MARKER_FILE="${marker}"
             ldconfig() { echo "mock ldconfig"; }
-            lsmod() { echo "nvidia 104165376 2"; }  # stays loaded (refcnt shows in-use)
-            cat() { echo "2"; }                       # refcnt != 0 -> do not rmmod
-            ls() { return 1; }
-            rmmod() { echo "mock rmmod $*"; }         # should NOT be called
+            systemctl() { echo "mock systemctl $*"; }
+            lsmod() { echo "nvidia 104165376 1"; }        # stays loaded regardless of rmmod attempts
+            rmmod() { echo "mock rmmod $1 (in use)"; return 1; }
             When call cleanUpPrebakedGPUDriver
             The status should be success
-            The output should not include "mock rmmod"
             The output should include "module_before=true"
+            # rmmod IS attempted now (no self-defeating refcnt pre-check), it just fails on a busy module
+            The output should include "mock rmmod nvidia"
             The output should include "module_after=true"
             The output should include "status=incomplete"
+            The output should include "marker_after=true"
         End
     End
 End