fix(gpu): e2e-validated hardening for CUDA prebake

Incorporates fixes surfaced by multi-SKU AgentBaker e2e validation of the
prebake feature, and integrates with main's logs_to_events-wrapped GPU install:

- cse_config.sh: only request install-skip-build when the prebake marker's
  driver_kind matches THIS node's NVIDIA_GPU_DRIVER_TYPE. A CUDA-prebaked marker
  on a shared VHD must NOT short-circuit a GRID node's install: the GRID image
  may not support install-skip-build and would fail to stage its userspace files
  (observed as CSE exit 84 in e2e). Pass the action through the timed
  installGPUDriverImage wrapper added on main.
- cse_install_ubuntu.sh (cleanUpPrebakedGPUDriver): drop the slow per-version
  'dkms remove --all' (~35s on the non-GPU provisioning critical path) in favor
  of removing the DKMS source tree + built module; also remove the driver
  userspace BINARIES (nvidia-smi etc.) so a non-GPU node is genuinely
  driver-free instead of leaving nvidia-smi on PATH erroring on missing libs.
- install-dependencies.sh: install gcc/make/libc6-dev before the build-only
  bake -- the standard non-GPU VHD builder ships gcc/make but not libc6-dev, so
  nvidia-installer cannot compile the module without it. The boot-time fallback
  recompile already gets these via installDeps, so it stays intact.
- specs: cover the driver_kind guard (match -> skip-build, mismatch -> full
  install) and the faster binaries+marker teardown.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: ganeshkumarashok

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

@@ -1221,11 +1221,16 @@ configGPUDrivers() {
         mkdir -p /opt/{actions,gpu}
         # When the kernel module was pre-built into the VHD (build-only at image-bake time),
         # a marker is present. Ask aks-gpu to skip the ~100s DKMS recompile and run only the
-        # device-dependent steps. aks-gpu independently re-validates the marker (kernel +
-        # driver_version + driver_kind) and falls back to a full build on any mismatch, so this
-        # is safe even after a kernel upgrade or on a shared VHD with a different driver kind.
+        # device-dependent steps -- but ONLY when the marker's driver_kind matches THIS node's
+        # driver (NVIDIA_GPU_DRIVER_TYPE). A CUDA-prebaked marker on a GRID node (or vice-versa)
+        # must request a full "install": the other driver image may not even support
+        # install-skip-build and would fail to stage its userspace files (e.g. /opt/gpu/config.sh).
+        # aks-gpu still independently re-validates the marker (kernel + driver_version +
+        # driver_kind) and falls back to a full build on any remaining mismatch (e.g. kernel drift).
         GPU_INSTALL_ACTION="install"
-        if [ -f "${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}" ]; then
+        GPU_DKMS_MARKER="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}"
+        if [ -f "$GPU_DKMS_MARKER" ] && \
+           [ "$(sed -n 's/^driver_kind=//p' "$GPU_DKMS_MARKER" | head -n1)" = "$NVIDIA_GPU_DRIVER_TYPE" ]; then
             GPU_INSTALL_ACTION="install-skip-build"
         fi
         # The driver image is normally pre-pulled into the VHD; only hit the registry when it is

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

@@ -240,17 +240,23 @@ cleanUpPrebakedGPUDriver() {
     fi
     echo "Removing pre-baked NVIDIA driver inherited from shared VHD on non-GPU node"
 
-    # Deregister the nvidia DKMS module so future kernel upgrades stop rebuilding it. Parsing stops
-    # at the first ',', ':' or space so "nvidia/<ver>: added" and "nvidia/<ver>, <kernel>: installed"
-    # both yield just <ver>.
-    if command -v dkms >/dev/null 2>&1; then
-        dkms status 2>/dev/null | sed -n 's#^nvidia/\([^,: ]*\).*#\1#p' | sort -u | while read -r nvidiaVersion; do
-            [ -n "${nvidiaVersion}" ] && dkms remove "nvidia/${nvidiaVersion}" --all || true
-        done
-    fi
+    # Deregister the nvidia DKMS module so future kernel upgrades stop rebuilding it, WITHOUT the
+    # slow `dkms remove --all` (it dominated CSE duration on the non-GPU provisioning path, ~35s).
+    # Removing the DKMS source tree deregisters it (dkms autoinstall iterates /var/lib/dkms/*), and
+    # removing the built module reclaims disk. The module is never loaded on a non-GPU node, so no
+    # depmod/initramfs refresh is required.
     rm -rf /var/lib/dkms/nvidia || true
+    rm -f /lib/modules/*/updates/dkms/nvidia*.ko* 2>/dev/null || true
     # aks-gpu relocates the userspace libs under GPU_DEST/lib64; on Ubuntu GPU_DEST=/usr/bin.
     rm -rf /usr/bin/lib64 || true
+    # nvidia-installer also drops driver userspace BINARIES under GPU_DEST (=/usr/bin on Ubuntu).
+    # Remove them too so a non-GPU node looks genuinely driver-free: otherwise e.g. `nvidia-smi`
+    # remains on PATH and, with its libs (lib64) gone, errors instead of being "command not found".
+    for nvidiaBin in nvidia-smi nvidia-debugdump nvidia-persistenced nvidia-cuda-mps-control \
+                     nvidia-cuda-mps-server nvidia-modprobe nvidia-bug-report.sh nvidia-powerd \
+                     nvidia-ngx-updater nvidia-sleep.sh; do
+        rm -f "/usr/bin/${nvidiaBin}" || true
+    done
     rm -f /etc/ld.so.conf.d/nvidia.conf || true
     ldconfig || true
     rm -f "${marker}" || true

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

@@ -1890,35 +1890,53 @@ SETUP_EOF
 
     Describe 'configGPUDrivers'
         # Mock everything the Ubuntu path touches so the test exercises only the
-        # marker -> aks-gpu action selection (install vs install-skip-build).
+        # marker -> aks-gpu action selection (install vs install-skip-build), including the
+        # driver_kind guard (a CUDA-baked marker on a GRID node must NOT skip the build).
+        # logs_to_events is mocked to faithfully dispatch the wrapped command (dropping the
+        # event-name arg) so the real installGPUDriverImage runs and surfaces the action.
+        logs_to_events() { shift; $@; }
         waitForContainerdReady() { return 0; }
         mkdir() { :; }
         ctr() { echo "ctr $*"; }
         nvidia-modprobe() { return 0; }
         nvidia-smi() { return 0; }
         ldconfig() { return 0; }
         isMarinerOrAzureLinux() { return 1; }
+        isAzureLinuxOSGuard() { return 1; }
         isACL() { return 1; }
         systemctlEnableAndStart() { return 0; }
         systemctl() { return 0; }
         # Capture the action passed to the install container.
         retrycmd_if_failure() { shift 3; echo "INSTALL_CMD: $*"; return 0; }
 
-        BeforeEach 'OS="$UBUNTU_OS_NAME"; NVIDIA_DRIVER_IMAGE="mcr.microsoft.com/aks/aks-gpu-cuda"; NVIDIA_DRIVER_IMAGE_TAG="580.0.0"; CTR_GPU_INSTALL_CMD="ctr-run"; GPU_DKMS_MARKER_FILE="$(mktemp -u)"'
+        BeforeEach 'OS="$UBUNTU_OS_NAME"; NVIDIA_GPU_DRIVER_TYPE="cuda"; NVIDIA_DRIVER_IMAGE="mcr.microsoft.com/aks/aks-gpu-cuda"; NVIDIA_DRIVER_IMAGE_TAG="580.0.0"; CTR_GPU_INSTALL_CMD="ctr-run"; GPU_DKMS_MARKER_FILE="$(mktemp -u)"'
 
         It 'uses the full install action when no prebake marker is present'
             When call configGPUDrivers
             The output should include "/entrypoint.sh install"
             The output should not include "install-skip-build"
         End
 
-        It 'uses install-skip-build when the prebake marker is present'
+        It 'uses install-skip-build when the prebake marker matches the node driver kind'
             marker="$(mktemp)"
+            printf 'driver_kind=cuda\n' > "$marker"
             GPU_DKMS_MARKER_FILE="$marker"
             When call configGPUDrivers
             The output should include "/entrypoint.sh install-skip-build"
             rm -f "$marker"
         End
+
+        It 'falls back to full install when the marker driver_kind does not match the node (CUDA marker on GRID node)'
+            marker="$(mktemp)"
+            printf 'driver_kind=cuda\n' > "$marker"
+            GPU_DKMS_MARKER_FILE="$marker"
+            NVIDIA_GPU_DRIVER_TYPE="grid"
+            NVIDIA_DRIVER_IMAGE="mcr.microsoft.com/aks/aks-gpu-grid"
+            When call configGPUDrivers
+            The output should include "/entrypoint.sh install"
+            The output should not include "install-skip-build"
+            rm -f "$marker"
+        End
     End
 
     Describe 'configureManagedGPUExperience'

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

@@ -11,43 +11,24 @@ Describe 'cse_install_ubuntu.sh'
             The output should equal ""
         End
 
-        It 'deregisters the nvidia DKMS module and removes baked artifacts when the marker is present'
+        It 'deregisters the nvidia DKMS module and removes baked artifacts (libs, binaries, marker) when present'
             marker="$(mktemp)"
             GPU_DKMS_MARKER_FILE="${marker}"
-            # mock dkms so `command -v dkms` succeeds and `dkms status` returns the installed form
-            dkms() {
-                if [ "$1" = "status" ]; then
-                    echo "nvidia/580.126.09, 6.8.0-1029-azure, x86_64: installed"
-                else
-                    echo "mock dkms $*"
-                fi
-            }
             rm() { echo "mock rm $*"; }
             ldconfig() { echo "mock ldconfig"; }
             When call cleanUpPrebakedGPUDriver
             The status should be success
             The output should include "Removing pre-baked NVIDIA driver"
-            The output should include "mock dkms remove nvidia/580.126.09 --all"
+            # deregisters via the DKMS source tree + built module removal (no slow dkms remove)
             The output should include "mock rm -rf /var/lib/dkms/nvidia"
+            The output should include "mock rm -f /lib/modules"
+            # relocated userspace libs
             The output should include "mock rm -rf /usr/bin/lib64"
+            # driver userspace binaries so nvidia-smi becomes "command not found" on non-GPU nodes
+            The output should include "mock rm -f /usr/bin/nvidia-smi"
             The output should include "mock ldconfig"
-        End
-
-        It 'parses the bare "nvidia/<ver>: added" dkms status form to a clean version'
-            marker="$(mktemp)"
-            GPU_DKMS_MARKER_FILE="${marker}"
-            dkms() {
-                if [ "$1" = "status" ]; then
-                    echo "nvidia/570.86.15: added"
-                else
-                    echo "mock dkms $*"
-                fi
-            }
-            rm() { echo "mock rm $*"; }
-            ldconfig() { echo "mock ldconfig"; }
-            When call cleanUpPrebakedGPUDriver
-            The status should be success
-            The output should include "mock dkms remove nvidia/570.86.15 --all"
+            # the slow per-version dkms remove --all must NOT be on the critical path anymore
+            The output should not include "dkms remove"
         End
     End
 End

vhdbuilder/packer/install-dependencies.sh

@@ -749,6 +749,11 @@ EOF
   # Dropping the image is a separate, deferred size optimization.
   if grep -q "NVIDIA_CUDA_PREBAKE" <<< "$FEATURE_FLAGS"; then
     echo "Pre-building NVIDIA CUDA kernel module into the VHD (build-only) for kernel $(uname -r)"
+    # nvidia-installer compiles the kernel module and needs the libc development headers (libc6-dev),
+    # which the standard (non-GPU) VHD builder image does not ship by default (gcc/make are present
+    # but libc6-dev is not). Ensure the kernel-module build toolchain before the bake; the boot-time
+    # fallback path already gets these via installDeps, so the runtime recompile stays intact.
+    apt_get_install 10 2 300 gcc make libc6-dev || exit 1
     CTR_GPU_PREBUILD_CMD="ctr -n k8s.io run --privileged --rm --net-host --with-ns pid:/proc/1/ns/pid --mount type=bind,src=/opt/gpu,dst=/mnt/gpu,options=rbind --mount type=bind,src=/opt/actions,dst=/mnt/actions,options=rbind"
     retrycmd_if_failure 3 10 600 bash -c "$CTR_GPU_PREBUILD_CMD $NVIDIA_DRIVER_IMAGE:$NVIDIA_DRIVER_IMAGE_TAG gpuprebuild /entrypoint.sh build-only" || exit 1
     if [ ! -f /opt/azure/aks-gpu/dkms-marker ]; then