diff --git a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh index e9e5fd95107..59331660532 100755 --- a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh +++ b/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 diff --git a/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_spec.sh index f846c8f7ddc..95d86d2295b 100644 --- a/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_spec.sh +++ b/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