From 9f7d289f06925e5b3ab893d5f389e42d796baf94 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Wed, 28 Oct 2026 13:44:38 +1300
Subject: [PATCH 1/3] fix(gpu): clean AKS drivers before PIS capture

Move strict Ubuntu cleanup into a VHD-baked helper so traditional CustomData stays under its size limit. Fail closed before customer customization; do not clean customer drivers on PIS-derived nodes.

Refs Azure/AgentBaker#9702

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 .../cloud-init/artifacts/cse_config_gpu.sh    |   6 +
 parts/linux/cloud-init/artifacts/cse_main.sh  |  22 +-
 parts/linux/cloud-init/artifacts/cse_start.sh |  10 +-
 .../artifacts/mariner/cse_install_mariner.sh  |  16 +
 .../ubuntu/cse_install_ubuntu_gpu.sh          | 299 ++++++++++++++++++
 vhdbuilder/packer/packer_source.sh            |   7 +
 .../packer/vhd-image-builder-arm64-gb.json    |   5 +
 .../packer/vhd-image-builder-arm64-gen2.json  |   5 +
 vhdbuilder/packer/vhd-image-builder-base.json |   5 +
 vhdbuilder/packer/vhd-image-builder-cvm.json  |   5 +
 10 files changed, 374 insertions(+), 6 deletions(-)
 create mode 100755 parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh

diff --git a/parts/linux/cloud-init/artifacts/cse_config_gpu.sh b/parts/linux/cloud-init/artifacts/cse_config_gpu.sh
index 77e3b52b2e7..07d2b493c84 100755
--- a/parts/linux/cloud-init/artifacts/cse_config_gpu.sh
+++ b/parts/linux/cloud-init/artifacts/cse_config_gpu.sh
@@ -254,6 +254,7 @@ ensureGPUDrivers() {
 }
 
 # Install AMD AMA core SW package for MA35D (Supernova GPU SKU)
+#{{if or IsMariner IsCustomImage}}
 dnf_install_amd_ama_core_package() {
     ver=$1; retries=$2; wait_sleep=$3; timeout=$4; shift && shift && shift && shift
     # Currently version 1.5.0 is supported.  Add more versions as they become available.
@@ -274,6 +275,7 @@ dnf_install_amd_ama_core_package() {
     done
     echo Executed dnf install AMD AMA core package $i times;
 }
+#{{end}}
 
 # Install AMD AMA drivers/SW for MA35D (Supernova GPU SKU)
 # Note that this depends on access to download.microsoft.com, so network-isolated clusters are not supported
@@ -283,6 +285,9 @@ setupAmdAma() {
     fi
 
     if isAzureLinux "$OS"; then
+        :
+        # Template guards leave raw VHD scripts intact; unknown custom images retain the implementation.
+        #{{if or IsMariner IsCustomImage}}
         # Install MA35D packages - currently version 1.5 and above are supported
         # This install script will extract the driver version/build number to find
         # the corresponding core/FW packages
@@ -348,6 +353,7 @@ setupAmdAma() {
         if [ "$(systemctl is-active kubelet)" = "active" ]; then
             systemctl restart kubelet
         fi
+        #{{end}}
     fi
 }
 
diff --git a/parts/linux/cloud-init/artifacts/cse_main.sh b/parts/linux/cloud-init/artifacts/cse_main.sh
index 25e750f832f..561f1a11938 100755
--- a/parts/linux/cloud-init/artifacts/cse_main.sh
+++ b/parts/linux/cloud-init/artifacts/cse_main.sh
@@ -402,6 +402,21 @@ EOF
         logs_to_events "AKS.CSE.ensureContainerd.ensureArtifactStreaming" ensureArtifactStreaming || exit $ERR_ARTIFACT_STREAMING_INSTALL
     fi
 
+    # Bake cleanup uses the explicit GPU_NODE=false opt-out in CustomData, even
+    # on GPU VMs. SkipGpuDriverInstall is a nodePrep VM tag; do not infer bake
+    # eligibility from cached IMDS or tags on derived nodes.
+    if [ "${PRE_PROVISION_ONLY}" = "true" ] && [ "${GPU_NODE}" != "true" ] &&
+       { [ "${OS}" = "${UBUNTU_OS_NAME}" ] || { isMarinerOrAzureLinux "$OS" "$OS_VARIANT" && ! isAzureLinuxOSGuard "$OS" "$OS_VARIANT"; }; }; then
+        if [ "${OS}" = "${UBUNTU_OS_NAME}" ]; then
+            if [ ! -r /opt/azure/containers/provision_installs_ubuntu_gpu.sh ]; then
+                echo "image containing PIS GPU cleanup helper required"
+                exit $ERR_NVIDIA_DRIVER_INSTALL
+            fi
+            source /opt/azure/containers/provision_installs_ubuntu_gpu.sh || exit $ERR_NVIDIA_DRIVER_INSTALL
+        fi
+        logs_to_events "AKS.CSE.cleanUpGPUDriversForBasePrep" cleanUpGPUDriversForBasePrep || exit $ERR_NVIDIA_DRIVER_INSTALL
+    fi
+
     # Enable localdns to handle node and pod DNS traffic via a local CoreDNS instance.
     # If hosts plugin is enabled, enableAKSLocalDNSHostsSetup() was already called at the
     # very start of basePrep (before disableSystemdResolved) to give the timer a head start
@@ -639,8 +654,9 @@ function nodePrep {
     logs_to_events "AKS.CSE.configureNodeProblemDetector" configureNodeProblemDetector
 
 
-    # Clean up GPU drivers if not a GPU node or if skipping driver install
-    if [ "${GPU_NODE}" != "true" ] || [ "${skip_nvidia_driver_install}" = "true" ]; then
+    # Only clean up on executions that ran basePrep. Derived PIS nodes skip it
+    # and must retain any customer driver installed after the image bake.
+    if [ "${BASE_PREP_RAN}" = "true" ] && { [ "${GPU_NODE}" != "true" ] || [ "${skip_nvidia_driver_install}" = "true" ]; }; then
         logs_to_events "AKS.CSE.cleanUpGPUDrivers" cleanUpGPUDrivers
     fi
 
@@ -743,8 +759,10 @@ function nodePrep {
 # when nodes are created from that VHD image.
 logs_to_events "AKS.CSE.configureNodeTimeSync" configure_node_time_sync_or_report_error || exit $?
 
+BASE_PREP_RAN=false
 if [ ! -f /opt/azure/containers/base_prep.complete ]; then
     basePrep
+    BASE_PREP_RAN=true
 else
     echo "Skipping basePrep - base_prep.complete file exists"
 fi
diff --git a/parts/linux/cloud-init/artifacts/cse_start.sh b/parts/linux/cloud-init/artifacts/cse_start.sh
index 1b28fd5abfc..ae931668d9b 100644
--- a/parts/linux/cloud-init/artifacts/cse_start.sh
+++ b/parts/linux/cloud-init/artifacts/cse_start.sh
@@ -115,10 +115,12 @@ upload_logs() {
 }
 # Create the marker for the completed provisioning stage.
 if [ "${PRE_PROVISION_ONLY}" = "true" ]; then
-    # Stage 1: Create marker indicating Stage 2 is needed
-    mkdir -p /opt/azure/containers && touch /opt/azure/containers/base_prep.complete
-    echo "Stage 1 complete - kubelet configuration skipped, Stage 2 required" >> /var/log/azure/cluster-provision.log
-    echo "Created base_prep.complete marker file" >> /var/log/azure/cluster-provision.log
+    if [ "$EXIT_CODE" -eq 0 ]; then
+        # Stage 1: Create marker indicating Stage 2 is needed
+        mkdir -p /opt/azure/containers && touch /opt/azure/containers/base_prep.complete
+        echo "Stage 1 complete - kubelet configuration skipped, Stage 2 required" >> /var/log/azure/cluster-provision.log
+        echo "Created base_prep.complete marker file" >> /var/log/azure/cluster-provision.log
+    fi
 else
     # provision.complete signals that a normal provisioning attempt finished.
     mkdir -p /opt/azure/containers && touch /opt/azure/containers/provision.complete
diff --git a/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh b/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh
index 15e6deb56b2..5c33b96d58e 100755
--- a/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh
+++ b/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh
@@ -830,6 +830,22 @@ cleanUpGPUDrivers() {
   done
 }
 
+cleanUpGPUDriversForBasePrep() {
+  local packageName cache_dir
+  cleanUpGPUDrivers
+  if [ -e "${GPU_DEST}" ] || [ -e /opt/gpu ]; then
+    echo "GPU basePrep cleanup incomplete: AKS driver artifacts remain" >&2
+    return 1
+  fi
+  for packageName in $(managedGPUPackageList); do
+    cache_dir=$(getPackageCacheDir "${packageName}")
+    if [ -e "${cache_dir}" ]; then
+      echo "GPU basePrep cleanup incomplete: ${packageName} cache remains" >&2
+      return 1
+    fi
+  done
+}
+
 installMinimalBuildDeps() {
     echo "installMinimalBuildDeps not implemented for mariner"
 }
diff --git a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh
new file mode 100755
index 00000000000..092af8ea50f
--- /dev/null
+++ b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh
@@ -0,0 +1,299 @@
+#!/bin/bash
+
+prebakedGPUDriverArtifactsRemain() {
+    local binary modules
+    [ -d /var/lib/dkms/nvidia ] || [ -e /usr/bin/lib64 ] || [ -e /etc/ld.so.conf.d/nvidia.conf ] && return 0
+    for binary 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
+        [ -e "/usr/bin/${binary}" ] && return 0
+    done
+    compgen -G '/lib/modules/*/updates/dkms/nvidia*.ko*' >/dev/null && return 0
+    modules=$(lsmod) || return 0
+    grep -q '^nvidia' <<< "${modules}"
+}
+
+graceBlackwellDriverBOMEntries() {
+    # Packer stages this BOM only on NVIDIA_GB images; package versions must still match it before removal.
+    local bom="${1:-${GB_MAI_BOM_FILE:-/opt/azure/containers/gb-mai-bom.json}}"
+    local wave="${2:-versions-wave2}"
+    [ -s "${bom}" ] || return 1
+
+    jq -er --arg wave "${wave}" '
+        .[$wave] as $packages |
+        select(($packages | type) == "object") |
+        select(any($packages | keys[]; test("^nvidia-dkms-[0-9]+-open$"))) |
+        select(any($packages | keys[]; test("^nvidia-driver-[0-9]+-open$"))) |
+        select(all($packages | keys[]; test("^(libxnvctrl0|libnvidia-[A-Za-z0-9-]+|nvidia-[A-Za-z0-9-]+|xserver-xorg-video-nvidia-[0-9]+)$"))) |
+        select(all($packages | values[]; type == "string" and length > 0)) |
+        $packages | to_entries[] | "\(.key)=\(.value)"
+    ' "${bom}"
+}
+
+# The optional allowlist is newline-separated package=version entries. Return 0 and
+# print the first unowned package (or an inspection error) on stdout; return 1 only
+# when inspection succeeds and no unowned package remains.
+graceBlackwellPackageManagedDriverRemains() {
+    local allowed="${1:-}" package_states
+    package_states=$(dpkg-query -W -f='${binary:Package}\t${db:Status-Status}\n') || {
+        echo "unable to inspect installed NVIDIA driver packages"
+        return 0
+    }
+    awk -F '\t' -v allowed="${allowed}" '
+        $2 == "installed" {
+            name = $1
+            sub(/:.*/, "", name)
+            if (name !~ /^(nvidia-(dkms|driver|utils|compute-utils|firmware|kernel-)|nvidia-(modprobe|persistenced)$|xserver-xorg-video-nvidia-|libxnvctrl0$|libnvidia-)/ || name ~ /^libnvidia-container/) next
+            count = split(allowed, owned, "\n")
+            for (i = 1; i <= count; i++) {
+                sub(/=.*/, "", owned[i])
+                if (owned[i] == name) next
+            }
+            print $1
+            found = 1
+            exit
+        }
+        END { if (!found) exit 1 }
+    ' <<< "${package_states}"
+}
+
+graceBlackwellInstalledDriverPackages() {
+    local entries="${1}" package expected state version series="" auto_packages bom_package
+    local -a installed=()
+
+    while IFS='=' read -r package expected; do
+        [ -n "${package}" ] || continue
+        case "${package}" in
+            nvidia-driver-[0-9]*-open)
+                series="${package#nvidia-driver-}"
+                series="${series%-open}"
+                ;;
+        esac
+        state=$(dpkg-query -W -f='${db:Status-Status}\t${Version}' "${package}" 2>/dev/null) || state=""
+        [ "${state%%$'\t'*}" = installed ] || continue
+        version="${state#*$'\t'}"
+        if [ "${version}" != "${expected}" ]; then
+            echo "Grace Blackwell package ${package} version ${version} does not match the GB BOM ${expected}" >&2
+            return 1
+        fi
+        installed+=("${package}=${version}")
+    done <<< "${entries}"
+
+    [ -n "${series}" ] || return 1
+    auto_packages=$(apt-mark showauto) || return 1
+    # BOMs do not always list apt dependencies; use each dependency's exact dpkg version, never the driver's version.
+    for package in "nvidia-utils-${series}" "nvidia-compute-utils-${series}" nvidia-modprobe nvidia-persistenced; do
+        while IFS='=' read -r bom_package expected; do
+            [ "${bom_package}" = "${package}" ] && continue 2
+        done <<< "${entries}"
+        state=$(dpkg-query -W -f='${db:Status-Status}\t${Version}' "${package}" 2>/dev/null) || state=""
+        [ "${state%%$'\t'*}" = installed ] || continue
+        version="${state#*$'\t'}"
+        if [ -z "${version}" ] || ! grep -Eq "^${package}(:arm64)?$" <<< "${auto_packages}"; then
+            echo "Grace Blackwell dependency ${package} is not an auto-installed package with a verifiable version" >&2
+            return 1
+        fi
+        installed+=("${package}=${version}")
+    done
+
+    printf '%s\n' "${installed[@]}"
+}
+
+graceBlackwellConfigIsUnmodified() {
+    local config="${1}" expected
+    shift
+    [ -e "${config}" ] || return 0
+
+    expected=$(printf '%s\n' "$@")
+    if ! cmp -s "${config}" <(printf '%s\n' "$@") &&
+        ! cmp -s "${config}" <(printf '%s' "${expected}"); then
+        echo "Grace Blackwell module configuration ${config} was modified; refusing teardown" >&2
+        return 1
+    fi
+}
+
+# Fresh cleanup requires AKS-owned config and exact BOM versions; an interrupted
+# cleanup uses the pending marker to retry after config removal. The marker is
+# written after purge and removed only after initramfs succeeds. The caller checks
+# for remaining packages and artifacts before basePrep can be marked complete.
+cleanUpGraceBlackwellGPUDriver() {
+    local entries="${1}" wave="${2}" pending="${3:-false}"
+    local packages package version modules service module remaining_package pending_file
+    pending_file="${GB_DRIVER_CLEANUP_PENDING_FILE:-/opt/azure/containers/gb-driver-cleanup.pending}"
+    local nouveau_config="${GB_NOUVEAU_MODPROBE_CONFIG_FILE:-/etc/modprobe.d/blacklist-nouveau.conf}"
+    local peermem_config="${GB_NVIDIA_PEERMEM_CONFIG_FILE:-/etc/modules-load.d/nvidia-peermem.conf}"
+    local nvidia_config="${GB_NVIDIA_MODPROBE_CONFIG_FILE:-/etc/modprobe.d/nvidia.conf}"
+    packages=$(graceBlackwellInstalledDriverPackages "${entries}") || return 1
+    if remaining_package=$(graceBlackwellPackageManagedDriverRemains "${packages}"); then
+        echo "NVIDIA package ${remaining_package} is not owned by the Grace Blackwell BOM" >&2
+        return 1
+    fi
+    if [ "${pending}" != true ] && [ ! -f "${nvidia_config}" ]; then
+        modules=$(lsmod) || return 1
+        # BasePrep may retry after teardown succeeded but a later step failed.
+        if [ -z "${packages}" ] && ! grep -q '^nvidia' <<< "${modules}" &&
+            ! prebakedGPUDriverArtifactsRemain &&
+            [ ! -e "${nvidia_config}" ] && [ ! -L "${nvidia_config}" ] &&
+            [ ! -e "${peermem_config}" ] && [ ! -L "${peermem_config}" ] &&
+            [ ! -e "${nouveau_config}" ] && [ ! -L "${nouveau_config}" ]; then
+            return 0
+        fi
+        echo "Grace Blackwell NVIDIA module configuration ${nvidia_config} is unavailable; refusing teardown" >&2
+        return 1
+    fi
+    [ "${pending}" = true ] || [ -f "${nvidia_config}" ] || {
+        echo "Grace Blackwell NVIDIA module configuration ${nvidia_config} is unavailable; refusing teardown" >&2
+        return 1
+    }
+    graceBlackwellConfigIsUnmodified "${peermem_config}" \
+        '# Explicitly load nvidia_peermem at boot for GPUDirect RDMA.' \
+        "# We cannot rely on a \`softdep nvidia post: nvidia_peermem\` in modprobe.d because" \
+        "# \`nvidia\` is autoloaded by the kernel via PCI modalias before the OFED/RDMA" \
+        '# stack is up, so the post-dep fails silently. systemd-modules-load.service' \
+        '# runs later in boot, by which point modprobe can resolve the full dep chain' \
+        '# (ib_core, mlx5_ib, nvidia, ...).' \
+        'nvidia_peermem' || return 1
+    graceBlackwellConfigIsUnmodified "${nvidia_config}" \
+        'options nvidia NVreg_RestrictProfilingToAdminUsers=0' \
+        'options nvidia NVreg_CreateImexChannel0=1' \
+        'options nvidia NVreg_CoherentGPUMemoryMode=driver' \
+        'options nvidia NVreg_RegistryDwords="RMBug5172204War=4"' || return 1
+    graceBlackwellConfigIsUnmodified "${nouveau_config}" \
+        'blacklist nouveau' 'options nouveau modeset=0' || return 1
+
+    local -a purge_packages=()
+    while IFS='=' read -r package version; do
+        [ -n "${package}" ] && purge_packages+=("${package}")
+    done <<< "${packages}"
+
+    for service in nvidia-device-plugin nvidia-dcgm-exporter nvidia-dcgm nvidia-imex nvidia-persistenced; do
+        systemctlDisableAndStop "${service}" || return 1
+        if systemctl is-active --quiet "${service}" || systemctl is-enabled --quiet "${service}"; then
+            echo "Grace Blackwell service ${service} remains active or enabled" >&2
+            return 1
+        fi
+    done
+
+    modules=$(lsmod) || return 1
+    if grep -q '^nvidia' <<< "${modules}" && systemctl cat openibd &>/dev/null; then
+        systemctl_stop 20 5 25 openibd || {
+            echo "Grace Blackwell service openibd could not be stopped" >&2
+            return 1
+        }
+        if systemctl is-active --quiet openibd; then
+            echo "Grace Blackwell service openibd remains active" >&2
+            return 1
+        fi
+        modules=$(lsmod) || return 1
+    fi
+    for module in nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia; do
+        grep -q "^${module}[[:space:]]" <<< "${modules}" || continue
+        rmmod "${module}" || { echo "Failed to unload Grace Blackwell module ${module}" >&2; return 1; }
+        modules=$(lsmod) || return 1
+        grep -q "^${module}[[:space:]]" <<< "${modules}" && {
+            echo "Grace Blackwell module ${module} remains loaded after removal" >&2
+            return 1
+        }
+    done
+    grep -q '^nvidia' <<< "${modules}" && { echo "NVIDIA module remains loaded after Grace Blackwell teardown" >&2; return 1; }
+
+    if [ "${#purge_packages[@]}" -gt 0 ]; then
+        apt_get_purge 10 5 300 "${purge_packages[@]}" || {
+            echo "Grace Blackwell driver package purge failed" >&2
+            return 1
+        }
+    fi
+
+    if [ "${pending}" != true ]; then
+        printf '%s\n' "${wave}" > "${pending_file}" || {
+            rm -f "${pending_file}"
+            echo "Unable to preserve Grace Blackwell teardown provenance" >&2
+            return 1
+        }
+    fi
+    rm -f "${peermem_config}" "${nvidia_config}" "${nouveau_config}" || return 1
+    for package in "${peermem_config}" "${nvidia_config}" "${nouveau_config}"; do
+        if [ -e "${package}" ] || [ -L "${package}" ]; then
+            echo "Grace Blackwell module configuration ${package} remains after removal" >&2
+            return 1
+        fi
+    done
+    update-initramfs -u || return 1
+    rm -f "${pending_file}" || return 1
+    [ ! -e "${pending_file}" ] || return 1
+}
+
+# Sourced during Ubuntu PIS basePrep after the CSE helpers and distro installer:
+# GPU_DEST, managedGPUPackageList, cleanUpGPUDrivers, apt_get_purge,
+# systemctlDisableAndStop, and systemctl_stop must be available. Fresh or pending
+# GB teardown verifies ownership; any remaining artifact blocks the basePrep marker.
+cleanUpGPUDriversForBasePrep() {
+    local packageName marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}" remaining_package
+    local peermem_config="${GB_NVIDIA_PEERMEM_CONFIG_FILE:-/etc/modules-load.d/nvidia-peermem.conf}"
+    local bom="${GB_MAI_BOM_FILE:-/opt/azure/containers/gb-mai-bom.json}"
+    local legacy_bom="${GB200_MAI_BOM_FILE:-/opt/azure/containers/gb200-mai-bom.json}"
+    local pending_file="${GB_DRIVER_CLEANUP_PENDING_FILE:-/opt/azure/containers/gb-driver-cleanup.pending}"
+    local wave=versions-wave2 entries gb_layout=false pending_teardown=false pending_wave
+
+    if [ -e "${pending_file}" ]; then
+        IFS= read -r pending_wave < "${pending_file}" || {
+            echo "GPU basePrep cleanup incomplete: invalid Grace Blackwell pending marker" >&2
+            return 1
+        }
+        case "${pending_wave}" in
+            versions-wave2) ;;
+            versions-wave1)
+                bom="${legacy_bom}"
+                wave=versions-wave1
+                ;;
+            *)
+            echo "GPU basePrep cleanup incomplete: invalid Grace Blackwell pending marker" >&2
+            return 1
+            ;;
+        esac
+        gb_layout=true
+        pending_teardown=true
+    elif [ -e "${peermem_config}" ]; then
+        gb_layout=true
+    elif [ -e "${legacy_bom}" ] && [ ! -e "${bom}" ]; then
+        bom="${legacy_bom}"
+        wave=versions-wave1
+        gb_layout=true
+    fi
+
+    entries=$(graceBlackwellDriverBOMEntries "${bom}" "${wave}" 2>/dev/null) || entries=""
+    if [ "${gb_layout}" = true ]; then
+        if [ -z "${entries}" ]; then
+            echo "GPU basePrep cleanup incomplete: Grace Blackwell BOM is unavailable" >&2
+            return 1
+        fi
+        cleanUpGraceBlackwellGPUDriver "${entries}" "${wave}" "${pending_teardown}" || {
+            echo "GPU basePrep cleanup incomplete: Grace Blackwell driver teardown failed" >&2
+            return 1
+        }
+    elif remaining_package=$(graceBlackwellPackageManagedDriverRemains); then
+        echo "GPU basePrep cleanup incomplete: NVIDIA package ${remaining_package} has no verifiable GB BOM ownership" >&2
+        return 1
+    fi
+
+    if remaining_package=$(graceBlackwellPackageManagedDriverRemains); then
+        echo "GPU basePrep cleanup incomplete: NVIDIA package ${remaining_package} remains after Grace Blackwell teardown" >&2
+        return 1
+    fi
+
+    cleanUpGPUDrivers
+    if prebakedGPUDriverArtifactsRemain; then
+        echo "GPU basePrep cleanup incomplete: pre-baked GPU driver artifacts remain" >&2
+        return 1
+    fi
+    if [ -e "${GPU_DEST}" ] || [ -e /opt/gpu ] || [ -e "${marker}" ]; then
+        echo "GPU basePrep cleanup incomplete: AKS driver artifacts remain" >&2
+        return 1
+    fi
+    for packageName in $(managedGPUPackageList); do
+        if [ -e "/opt/${packageName}" ]; then
+            echo "GPU basePrep cleanup incomplete: ${packageName} cache remains" >&2
+            return 1
+        fi
+    done
+}
diff --git a/vhdbuilder/packer/packer_source.sh b/vhdbuilder/packer/packer_source.sh
index 3b088e23268..8330fafc1b1 100644
--- a/vhdbuilder/packer/packer_source.sh
+++ b/vhdbuilder/packer/packer_source.sh
@@ -22,6 +22,12 @@ copyNetworkdRecoveryFiles() {
   fi
 }
 
+copyUbuntuGPUCleanupHelper() {
+  if isUbuntu "$OS"; then
+    cpAndMode /home/packer/provision_installs_ubuntu_gpu.sh /opt/azure/containers/provision_installs_ubuntu_gpu.sh 0744
+  fi
+}
+
 copyPackerFiles() {
   SYSCTL_CONFIG_SRC=/home/packer/sysctl-d-60-CIS.conf
   SYSCTL_CONFIG_DEST=/etc/sysctl.d/60-CIS.conf
@@ -320,6 +326,7 @@ copyPackerFiles() {
   CSE_INSTALL_DISTRO_SRC=/home/packer/provision_installs_distro.sh
   CSE_INSTALL_DISTRO_DEST=/opt/azure/containers/provision_installs_distro.sh
   cpAndMode $CSE_INSTALL_DISTRO_SRC $CSE_INSTALL_DISTRO_DEST 0744
+  copyUbuntuGPUCleanupHelper
 
   CSE_HELPERS_SRC=/home/packer/provision_source.sh
   CSE_HELPERS_DEST=/opt/azure/containers/provision_source.sh
diff --git a/vhdbuilder/packer/vhd-image-builder-arm64-gb.json b/vhdbuilder/packer/vhd-image-builder-arm64-gb.json
index a74db09e3c7..231d5382d5c 100644
--- a/vhdbuilder/packer/vhd-image-builder-arm64-gb.json
+++ b/vhdbuilder/packer/vhd-image-builder-arm64-gb.json
@@ -169,6 +169,11 @@
       "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh",
       "destination": "/home/packer/provision_installs_distro.sh"
     },
+    {
+      "type": "file",
+      "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh",
+      "destination": "/home/packer/provision_installs_ubuntu_gpu.sh"
+    },
     {
       "type": "file",
       "source": "parts/linux/cloud-init/artifacts/cse_helpers.sh",
diff --git a/vhdbuilder/packer/vhd-image-builder-arm64-gen2.json b/vhdbuilder/packer/vhd-image-builder-arm64-gen2.json
index d60eba4569f..e3456826cf0 100644
--- a/vhdbuilder/packer/vhd-image-builder-arm64-gen2.json
+++ b/vhdbuilder/packer/vhd-image-builder-arm64-gen2.json
@@ -154,6 +154,11 @@
       "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh",
       "destination": "/home/packer/provision_installs_distro.sh"
     },
+    {
+      "type": "file",
+      "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh",
+      "destination": "/home/packer/provision_installs_ubuntu_gpu.sh"
+    },
     {
       "type": "file",
       "source": "parts/linux/cloud-init/artifacts/cse_helpers.sh",
diff --git a/vhdbuilder/packer/vhd-image-builder-base.json b/vhdbuilder/packer/vhd-image-builder-base.json
index c1a7cfe1c23..d2be462828f 100644
--- a/vhdbuilder/packer/vhd-image-builder-base.json
+++ b/vhdbuilder/packer/vhd-image-builder-base.json
@@ -157,6 +157,11 @@
       "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh",
       "destination": "/home/packer/provision_installs_distro.sh"
     },
+    {
+      "type": "file",
+      "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh",
+      "destination": "/home/packer/provision_installs_ubuntu_gpu.sh"
+    },
     {
       "type": "file",
       "source": "parts/linux/cloud-init/artifacts/cse_helpers.sh",
diff --git a/vhdbuilder/packer/vhd-image-builder-cvm.json b/vhdbuilder/packer/vhd-image-builder-cvm.json
index 228b2187700..d0f9387317c 100644
--- a/vhdbuilder/packer/vhd-image-builder-cvm.json
+++ b/vhdbuilder/packer/vhd-image-builder-cvm.json
@@ -160,6 +160,11 @@
       "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh",
       "destination": "/home/packer/provision_installs_distro.sh"
     },
+    {
+      "type": "file",
+      "source": "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh",
+      "destination": "/home/packer/provision_installs_ubuntu_gpu.sh"
+    },
     {
       "type": "file",
       "source": "parts/linux/cloud-init/artifacts/cse_helpers.sh",

From 6985d158492466f6cde904bfe19a7996f406c060 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Wed, 28 Oct 2026 13:44:43 +1300
Subject: [PATCH 2/3] test(gpu): cover PIS cleanup and payload limits

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 pkg/agent/gpu_pis_delivery_test.go            |  85 ++++++++
 .../artifacts/cse_install_mariner_spec.sh     |  28 ++-
 .../cse_install_ubuntu_gpu_pis_spec.sh        | 189 ++++++++++++++++++
 .../artifacts/cse_main_gpu_pis_spec.sh        | 102 ++++++++++
 spec/vhdbuilder/packer/pis_gpu_helper_spec.sh |  19 ++
 5 files changed, 422 insertions(+), 1 deletion(-)
 create mode 100644 pkg/agent/gpu_pis_delivery_test.go
 create mode 100644 spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_gpu_pis_spec.sh
 create mode 100644 spec/parts/linux/cloud-init/artifacts/cse_main_gpu_pis_spec.sh
 create mode 100644 spec/vhdbuilder/packer/pis_gpu_helper_spec.sh

diff --git a/pkg/agent/gpu_pis_delivery_test.go b/pkg/agent/gpu_pis_delivery_test.go
new file mode 100644
index 00000000000..3a9212448cf
--- /dev/null
+++ b/pkg/agent/gpu_pis_delivery_test.go
@@ -0,0 +1,85 @@
+package agent
+
+import (
+	"encoding/base64"
+	"os/exec"
+	"path/filepath"
+	"testing"
+
+	"github.com/Azure/agentbaker/pkg/agent/datamodel"
+	"github.com/stretchr/testify/require"
+	"gopkg.in/yaml.v3"
+)
+
+func TestPISGPUCustomDataFitsAndUsesBakedHelper(t *testing.T) {
+	for _, mode := range []struct {
+		preProvisionOnly bool
+		phase1           bool
+	}{
+		{preProvisionOnly: true, phase1: true},
+		{preProvisionOnly: true},
+		{},
+	} {
+		config := newNodeCustomDataRenderConfig(datamodel.AKSUbuntuContainerd2404Gen2)
+		pool := config.AgentPoolProfile
+		pool.Name, pool.VMSize = "agent2", "Standard_NC4as_T4_v3"
+		pool.StorageProfile = "ManagedDisks"
+		pool.AvailabilityProfile = datamodel.VirtualMachineScaleSets
+		pool.VnetSubnetID = "/subscriptions/test/resourceGroups/test/providers/Microsoft.Network/virtualNetworks/test/subnets/default"
+		properties := config.ContainerService.Properties
+		properties.LinuxProfile = &datamodel.LinuxProfile{AdminUsername: "azureuser"}
+		properties.LinuxProfile.SSH.PublicKeys = []datamodel.PublicKey{{KeyData: "fixture-key"}}
+		properties.ServicePrincipalProfile = &datamodel.ServicePrincipalProfile{ClientID: "fixture", Secret: "fixture"}
+		config.TenantID, config.SubscriptionID = "tenantID", "subID"
+		config.ResourceGroupName, config.UserAssignedIdentityClientID = "resourceGroupName", "userAssignedID"
+		config.PrimaryScaleSetName = "aks-agent2-36873793-vmss"
+		config.KubeletConfig = map[string]string{"--cluster-domain": "cluster.local", "--cluster-dns": "10.0.0.10", "--cloud-provider": "azure", "--max-pods": "110"}
+		config.PreProvisionOnly = mode.preProvisionOnly
+		config.EnableScriptlessCSECmd = mode.phase1
+
+		payload := InitializeTemplateGenerator().getLinuxNodeBootstrappingPayload(config)
+		compressed, err := base64.StdEncoding.DecodeString(payload)
+		require.NoError(t, err)
+		t.Logf("preProvisionOnly=%t phase1=%t: gzip=%d base64=%d", mode.preProvisionOnly, mode.phase1, len(compressed), len(payload))
+		decoded, err := getGzipDecodedValue(compressed)
+		require.NoError(t, err)
+		require.LessOrEqual(t, len(compressed), 65535)
+		require.LessOrEqual(t, len(payload), 87380)
+
+		var document cloudInit
+		require.NoError(t, yaml.Unmarshal(decoded, &document))
+		var installer bool
+		for _, file := range document.WriteFiles {
+			require.NotEqual(t, "/opt/azure/containers/provision_installs_ubuntu_gpu.sh", file.Path)
+			if file.Path == "/opt/azure/containers/provision_installs_distro.sh" {
+				installer = true
+				script, err := getGzipDecodedValue([]byte(file.Content))
+				require.NoError(t, err)
+				require.Contains(t, string(script), "cleanUpGPUDrivers()")
+			}
+		}
+		if mode.phase1 {
+			require.Equal(t, "/opt/azure/containers/scriptless-cse-overrides.txt", document.WriteFiles[0].Path)
+		} else {
+			require.True(t, installer)
+		}
+	}
+}
+
+func TestGPUAMARenderingKeepsAzureLinuxImplementation(t *testing.T) {
+	config := newNodeCustomDataRenderConfig(datamodel.AKSAzureLinuxV3Gen2)
+	payload := getBase64EncodedGzippedCustomScript(kubernetesCSEConfigGPU, config)
+	compressed, err := base64.StdEncoding.DecodeString(payload)
+	require.NoError(t, err)
+	rendered, err := getGzipDecodedValue(compressed)
+	require.NoError(t, err)
+
+	raw := exec.Command("bash", "-c", `source "$1"; declare -f setupAmdAma dnf_install_amd_ama_core_package`,
+		"gpu-raw-test", filepath.Join("..", "..", "parts", kubernetesCSEConfigGPU))
+	rawFunctions, err := raw.CombinedOutput()
+	require.NoError(t, err, string(rawFunctions))
+	renderedCmd := exec.Command("bash", "-c", string(rendered)+"\ndeclare -f setupAmdAma dnf_install_amd_ama_core_package")
+	renderedFunctions, err := renderedCmd.CombinedOutput()
+	require.NoError(t, err, string(renderedFunctions))
+	require.Equal(t, string(rawFunctions), string(renderedFunctions))
+}
diff --git a/spec/parts/linux/cloud-init/artifacts/cse_install_mariner_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_install_mariner_spec.sh
index 3e3962c4dc8..4a17d7141dc 100644
--- a/spec/parts/linux/cloud-init/artifacts/cse_install_mariner_spec.sh
+++ b/spec/parts/linux/cloud-init/artifacts/cse_install_mariner_spec.sh
@@ -28,6 +28,32 @@ Describe 'cse_install_mariner.sh'
     Include "./parts/linux/cloud-init/artifacts/cse_install.sh"
     Include "./parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh"
 
+    Describe 'cleanUpGPUDriversForBasePrep'
+        GPU_DEST="${PWD}/.shellspec-nonexistent-gpu-dest"
+        managedGPUPackageList() { echo "driver-cache"; }
+        getPackageCacheDir() { echo "$gpu_test_root/driver-cache"; }
+        setup_gpu_cache() { gpu_test_root="${PWD}/.shellspec-gpu-cache-$$"; mkdir -p "$gpu_test_root"; }
+        teardown_gpu_cache() { rm -rf "$gpu_test_root"; }
+        BeforeEach 'setup_gpu_cache'
+        AfterEach 'teardown_gpu_cache'
+
+        It 'fails when an AKS cache survives cleanup'
+            mkdir -p "$gpu_test_root/driver-cache"
+            cleanUpGPUDrivers() { :; }
+            When call cleanUpGPUDriversForBasePrep
+            The status should be failure
+            The stderr should include "driver-cache cache remains"
+        End
+
+        It 'succeeds when AKS caches are removed'
+            mkdir -p "$gpu_test_root/driver-cache"
+            cleanUpGPUDrivers() { rm -rf "$gpu_test_root/driver-cache"; }
+            When call cleanUpGPUDriversForBasePrep
+            The status should be success
+            The path "$gpu_test_root/driver-cache" should not be exist
+        End
+    End
+
     Describe 'installDeps'
         It 'installs the required packages with installDeps for Mariner 2.0'
             OS_VERSION="2.0"
@@ -772,7 +798,7 @@ EOF
     End
 
     Describe 'installPackageFromCache version matching'
-        rpm_version_cache="/tmp/shellspec-rpm-version-cache-$$"
+        rpm_version_cache="${PWD}/.shellspec-rpm-version-cache-$$"
 
         setup_version_cache() {
             RPM_PACKAGE_CACHE_BASE_DIR="$rpm_version_cache"
diff --git a/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_gpu_pis_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_gpu_pis_spec.sh
new file mode 100644
index 00000000000..b387244ac4f
--- /dev/null
+++ b/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_gpu_pis_spec.sh
@@ -0,0 +1,189 @@
+#!/bin/bash
+
+Describe 'Ubuntu PIS GPU basePrep cleanup'
+    Include "./parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh"
+    fixture="${PWD}/.shellspec-gpu-cleanup-$$"
+    GPU_DEST="${fixture}/gpu-dest"
+    managedGPUPackageList() { :; }
+
+    setup_cleanup() {
+        mkdir -p "${fixture}"
+        GPU_DKMS_MARKER_FILE="${fixture}/marker"
+        GB_NVIDIA_PEERMEM_CONFIG_FILE="${fixture}/peermem.conf"
+        GB_NVIDIA_MODPROBE_CONFIG_FILE="${fixture}/nvidia.conf"
+        GB_NOUVEAU_MODPROBE_CONFIG_FILE="${fixture}/nouveau.conf"
+        GB_DRIVER_CLEANUP_PENDING_FILE="${fixture}/cleanup.pending"
+    }
+
+    gb_entries() {
+        printf '%s\n' 'nvidia-dkms-580-open=580.159.04-1ubuntu1' \
+            'nvidia-driver-580-open=580.159.04-1ubuntu1' \
+            'libnvidia-compute-580=580.159.04-1ubuntu1'
+    }
+
+    setup_gb_package_mocks() {
+        dpkg-query() {
+            if [ "$#" -eq 2 ]; then
+                [ "${GB_TEST_INSPECTION_ERROR:-false}" = true ] && return 1
+                printf '%s\tinstalled\n' nvidia-dkms-580-open nvidia-driver-580-open \
+                    libnvidia-compute-580 nvidia-utils-580 nvidia-modprobe
+                return
+            fi
+            case "${@: -1}" in
+                nvidia-dkms-580-open|libnvidia-compute-580)
+                    printf 'installed\t580.159.04-1ubuntu1\n' ;;
+                nvidia-driver-580-open)
+                    printf 'installed\t%s\n' "${GB_TEST_DRIVER_VERSION:-580.159.04-1ubuntu1}" ;;
+                nvidia-utils-580)
+                    printf 'installed\t580.159.04-1ubuntu2\n' ;;
+                nvidia-modprobe)
+                    printf 'installed\t580.159.04-1ubuntu3\n' ;;
+                *)
+                    return 1 ;;
+            esac
+        }
+        apt-mark() {
+            [ "$1" = showauto ] || return 1
+            if [ "${GB_TEST_NONAUTO:-false}" != true ]; then
+                printf 'nvidia-utils-580\n'
+            fi
+            printf 'nvidia-modprobe\n'
+        }
+        systemctlDisableAndStop() { touch "${fixture}/teardown"; }
+        apt_get_purge() { touch "${fixture}/teardown"; }
+    }
+
+    stage_gb_configs() {
+        cp parts/linux/cloud-init/artifacts/ubuntu/gb/nvidia-peermem.conf "${GB_NVIDIA_PEERMEM_CONFIG_FILE}"
+        cp parts/linux/cloud-init/artifacts/ubuntu/gb/modprobe-nvidia-parameters.conf "${GB_NVIDIA_MODPROBE_CONFIG_FILE}"
+        printf 'blacklist nouveau\noptions nouveau modeset=0\n' > "${GB_NOUVEAU_MODPROBE_CONFIG_FILE}"
+    }
+
+    BeforeEach 'setup_cleanup'
+    AfterEach 'rm -rf "$fixture"'
+
+    It 'selects only BOM-matched packages and verified auto-installed dependencies'
+        setup_gb_package_mocks
+        When call graceBlackwellInstalledDriverPackages "$(gb_entries)"
+        The status should be success
+        The output should equal "$(printf '%s\n' \
+            'nvidia-dkms-580-open=580.159.04-1ubuntu1' \
+            'nvidia-driver-580-open=580.159.04-1ubuntu1' \
+            'libnvidia-compute-580=580.159.04-1ubuntu1' \
+            'nvidia-utils-580=580.159.04-1ubuntu2' \
+            'nvidia-modprobe=580.159.04-1ubuntu3')"
+    End
+
+    It 'rejects a BOM version mismatch before teardown'
+        setup_gb_package_mocks
+        GB_TEST_DRIVER_VERSION=580.159.04-1ubuntu9
+        When call cleanUpGraceBlackwellGPUDriver "$(gb_entries)" versions-wave2
+        The status should be failure
+        The stderr should include 'does not match the GB BOM'
+        The path "${fixture}/teardown" should not be exist
+        The path "${GB_DRIVER_CLEANUP_PENDING_FILE}" should not be exist
+    End
+
+    It 'rejects an unverified dependency before teardown'
+        setup_gb_package_mocks
+        GB_TEST_NONAUTO=true
+        When call cleanUpGraceBlackwellGPUDriver "$(gb_entries)" versions-wave2
+        The status should be failure
+        The stderr should include 'nvidia-utils-580 is not an auto-installed package'
+        The path "${fixture}/teardown" should not be exist
+        The path "${GB_DRIVER_CLEANUP_PENDING_FILE}" should not be exist
+    End
+
+    It 'fails closed when installed package inspection fails'
+        setup_gb_package_mocks
+        GB_TEST_INSPECTION_ERROR=true
+        When call cleanUpGraceBlackwellGPUDriver "$(gb_entries)" versions-wave2
+        The status should be failure
+        The stderr should include 'unable to inspect installed NVIDIA driver packages'
+        The path "${fixture}/teardown" should not be exist
+    End
+
+    It 'accepts the checked-in GB configuration templates without weakening exact matching'
+        setup_gb_package_mocks
+        stage_gb_configs
+        systemctlDisableAndStop() { :; }
+        systemctl() { return 1; }
+        lsmod() { printf 'Module Size Used by\n'; }
+        apt_get_purge() { printf '%s\n' "$@" > "${fixture}/purged"; }
+        update-initramfs() { :; }
+        When call cleanUpGraceBlackwellGPUDriver "$(gb_entries)" versions-wave2
+        The status should be success
+        The contents of file "${fixture}/purged" should equal "$(printf '%s\n' 10 5 300 \
+            nvidia-dkms-580-open nvidia-driver-580-open libnvidia-compute-580 \
+            nvidia-utils-580 nvidia-modprobe)"
+        The path "${GB_NVIDIA_PEERMEM_CONFIG_FILE}" should not be exist
+        The path "${GB_NVIDIA_MODPROBE_CONFIG_FILE}" should not be exist
+        The path "${GB_DRIVER_CLEANUP_PENDING_FILE}" should not be exist
+    End
+
+    It 'rejects changed GB configuration before teardown'
+        setup_gb_package_mocks
+        stage_gb_configs
+        printf '# customer edit\n' >> "${GB_NVIDIA_PEERMEM_CONFIG_FILE}"
+        When call cleanUpGraceBlackwellGPUDriver "$(gb_entries)" versions-wave2
+        The status should be failure
+        The stderr should include 'was modified; refusing teardown'
+        The path "${fixture}/teardown" should not be exist
+        The path "${GB_NVIDIA_PEERMEM_CONFIG_FILE}" should be exist
+    End
+
+    It 'fails CUDA cleanup when the prebake marker remains'
+        touch "${GPU_DKMS_MARKER_FILE}"
+        dpkg-query() { :; }
+        lsmod() { :; }
+        cleanUpGPUDrivers() { :; }
+        prebakedGPUDriverArtifactsRemain() { return 1; }
+        When call cleanUpGPUDriversForBasePrep
+        The status should be failure
+        The stderr should include "AKS driver artifacts remain"
+        The path "${GPU_DKMS_MARKER_FILE}" should be exist
+    End
+
+    It 'refuses GB teardown when installed NVIDIA packages are not BOM-owned'
+        touch "${GPU_DKMS_MARKER_FILE}"
+        graceBlackwellInstalledDriverPackages() { printf '%s\n' 'nvidia-driver-580-open=580.159.04-1ubuntu1'; }
+        dpkg-query() { printf 'nvidia-driver-580-open\tinstalled\nnvidia-driver-535\tinstalled\n'; }
+        When call cleanUpGraceBlackwellGPUDriver 'nvidia-driver-580-open=580.159.04-1ubuntu1' versions-wave2
+        The status should be failure
+        The stderr should include 'NVIDIA package nvidia-driver-535 is not owned by the Grace Blackwell BOM'
+        The path "${GPU_DKMS_MARKER_FILE}" should be exist
+    End
+
+    It 'retains GB cleanup provenance after initramfs failure and retries'
+        cp parts/linux/cloud-init/artifacts/ubuntu/gb/nvidia-peermem.conf "${GB_NVIDIA_PEERMEM_CONFIG_FILE}"
+        cp parts/linux/cloud-init/artifacts/ubuntu/gb/modprobe-nvidia-parameters.conf "${GB_NVIDIA_MODPROBE_CONFIG_FILE}"
+        printf 'blacklist nouveau\noptions nouveau modeset=0\n' > "${GB_NOUVEAU_MODPROBE_CONFIG_FILE}"
+        touch "${GPU_DKMS_MARKER_FILE}"
+        graceBlackwellDriverBOMEntries() { printf '%s\n' 'nvidia-driver-580-open=580.159.04-1ubuntu1'; }
+        graceBlackwellInstalledDriverPackages() { :; }
+        graceBlackwellPackageManagedDriverRemains() { return 1; }
+        systemctlDisableAndStop() { :; }
+        systemctl() { return 1; }
+        lsmod() { printf 'Module Size Used by\n'; }
+        update_attempts=0
+        update-initramfs() {
+            update_attempts=$((update_attempts + 1))
+            [ "${update_attempts}" -gt 1 ]
+        }
+        cleanUpGPUDrivers() { rm -f "${GPU_DKMS_MARKER_FILE}"; }
+        prebakedGPUDriverArtifactsRemain() { return 1; }
+        retry_cleanup() {
+            local status=0 pending_wave
+            cleanUpGPUDriversForBasePrep > "${fixture}/first-attempt.log" 2>&1 && status=1
+            grep -q 'GPU basePrep cleanup incomplete: Grace Blackwell driver teardown failed' "${fixture}/first-attempt.log" || status=1
+            IFS= read -r pending_wave < "${GB_DRIVER_CLEANUP_PENDING_FILE}" || status=1
+            [ "${pending_wave}" = versions-wave2 ] || status=1
+            cleanUpGPUDriversForBasePrep || status=1
+            [ ! -e "${GB_DRIVER_CLEANUP_PENDING_FILE}" ] || status=1
+            [ ! -e "${GPU_DKMS_MARKER_FILE}" ] || status=1
+            return "${status}"
+        }
+        When call retry_cleanup
+        The status should be success
+    End
+End
diff --git a/spec/parts/linux/cloud-init/artifacts/cse_main_gpu_pis_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_main_gpu_pis_spec.sh
new file mode 100644
index 00000000000..866aea302b2
--- /dev/null
+++ b/spec/parts/linux/cloud-init/artifacts/cse_main_gpu_pis_spec.sh
@@ -0,0 +1,102 @@
+#!/bin/bash
+
+Describe 'GPU cleanup across PIS phases'
+    CSE_MAIN="./parts/linux/cloud-init/artifacts/cse_main.sh"
+    fixture="${PWD}/.shellspec-cse-main-gpu-pis-$$"
+
+    setup_driver() {
+        driver_root="${fixture}/driver"
+        marker="${fixture}/base_prep.complete"
+        mkdir -p "${driver_root}/usr/local/nvidia"
+        GPU_NODE=false
+        PRE_PROVISION_ONLY=false
+        skip_nvidia_driver_install=false
+    }
+
+    run_phases() {
+        local dispatch cleanup
+        dispatch=$(awk '
+            /^BASE_PREP_RAN=false$/ { reading = 1 }
+            reading && /^echo "Custom script finished."/ { exit }
+            reading { print }
+        ' "$CSE_MAIN")
+        cleanup=$(sed -n '/    if \[ "\${BASE_PREP_RAN}" = "true" \]/,/    fi/p' "$CSE_MAIN")
+        dispatch=${dispatch//\/opt\/azure\/containers\/base_prep.complete/$marker}
+        basePrep() { echo "basePrep"; }
+        nodePrep() { eval "$cleanup"; }
+        logs_to_events() { shift; "$@"; }
+        cleanUpGPUDrivers() {
+            rm -rf "${driver_root}/usr/local/nvidia"
+            touch "${fixture}/cleanup-called"
+        }
+        eval "$dispatch"
+    }
+
+    run_cse_marker() {
+        local marker_script
+        PRE_PROVISION_ONLY=true
+        marker_script=$(awk '
+            /^if \[ "\$\{PRE_PROVISION_ONLY\}" = "true" \]; then/ { reading = 1 }
+            reading && /^if \[ "\$EXIT_CODE" -ne 0 \]; then/ { exit }
+            reading { print }
+        ' ./parts/linux/cloud-init/artifacts/cse_start.sh)
+        marker_script=${marker_script//\/opt\/azure\/containers/$fixture}
+        marker_script=${marker_script//\/var\/log\/azure\/cluster-provision.log/$fixture/provision.log}
+        eval "$marker_script"
+    }
+
+    BeforeEach 'setup_driver'
+    AfterEach 'rm -rf "$fixture"'
+
+    It 'preserves customer drivers on a PIS-derived opt-out node'
+        touch "$marker"
+        When call run_phases
+        The status should be success
+        The output should include "Skipping basePrep"
+        The path "$driver_root/usr/local/nvidia" should be exist
+        The path "$fixture/cleanup-called" should not be exist
+    End
+
+    It 'preserves a customer driver on a PIS-derived GPU with a skip tag'
+        GPU_NODE=true
+        skip_nvidia_driver_install=true
+        touch "$marker"
+        When call run_phases
+        The status should be success
+        The output should include "Skipping basePrep"
+        The path "$driver_root/usr/local/nvidia" should be exist
+        The path "$fixture/cleanup-called" should not be exist
+    End
+
+    It 'cleans CUDA drivers only after basePrep ran in this execution'
+        When call run_phases
+        The status should be success
+        The output should include "basePrep"
+        The path "$driver_root/usr/local/nvidia" should not be exist
+        The path "$fixture/cleanup-called" should be exist
+    End
+
+    It 'still cleans a non-PIS GPU opted out with a skip tag'
+        GPU_NODE=true
+        skip_nvidia_driver_install=true
+        When call run_phases
+        The status should be success
+        The output should include "basePrep"
+        The path "$driver_root/usr/local/nvidia" should not be exist
+        The path "$fixture/cleanup-called" should be exist
+    End
+
+    It 'does not write the bake marker after failed basePrep'
+        EXIT_CODE=224
+        When call run_cse_marker
+        The status should be success
+        The path "$fixture/base_prep.complete" should not be exist
+    End
+
+    It 'writes the bake marker after successful basePrep'
+        EXIT_CODE=0
+        When call run_cse_marker
+        The status should be success
+        The path "$fixture/base_prep.complete" should be exist
+    End
+End
diff --git a/spec/vhdbuilder/packer/pis_gpu_helper_spec.sh b/spec/vhdbuilder/packer/pis_gpu_helper_spec.sh
new file mode 100644
index 00000000000..cb556b84757
--- /dev/null
+++ b/spec/vhdbuilder/packer/pis_gpu_helper_spec.sh
@@ -0,0 +1,19 @@
+#!/bin/bash
+
+Describe 'Ubuntu PIS GPU cleanup helper VHD installation'
+  Include './vhdbuilder/packer/packer_source.sh'
+
+  isUbuntu() { [ "$1" = UBUNTU ]; }
+  cpAndMode() { printf '%s\n' "$*"; }
+
+  It 'copies the helper into Ubuntu images'
+    OS=UBUNTU
+    When call copyUbuntuGPUCleanupHelper
+    The output should equal '/home/packer/provision_installs_ubuntu_gpu.sh /opt/azure/containers/provision_installs_ubuntu_gpu.sh 0744'
+  End
+
+  It 'calls the helper copy during Packer file installation'
+    When call sed -n '/^copyPackerFiles()/,/^}/p' vhdbuilder/packer/packer_source.sh
+    The output should include '  copyUbuntuGPUCleanupHelper'
+  End
+End

From cea1cead071770d3d6bbcfbfabe7997d5b15ef88 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Wed, 28 Oct 2026 13:44:46 +1300
Subject: [PATCH 3/3] test(e2e): verify driver survives PIS capture

Exercise a customer T4 installation after basePrep using a branch-built image; validate the installed driver on the derived VM.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 .pipelines/e2e-gpu.yaml                    |  6 +-
 e2e/README.md                              | 11 ++++
 e2e/scenario/provision.go                  | 67 +++++++++++++++++++++-
 e2e/scenario/scenario.go                   | 32 +++++++++++
 e2e/scenario/types.go                      |  4 ++
 e2e/scenario/validators.go                 | 37 ++++++++++++
 e2e/scenario/validators_gpu_driver_test.go | 28 +++++++++
 e2e/scenario/validators_test.go            | 39 +++++++++++++
 e2e/scenario/vmss.go                       |  2 +-
 e2e/scenario/vmss_test.go                  | 61 ++++++++++++++++++++
 10 files changed, 284 insertions(+), 3 deletions(-)

diff --git a/.pipelines/e2e-gpu.yaml b/.pipelines/e2e-gpu.yaml
index 6c7e3e39ac5..375d80c5faf 100644
--- a/.pipelines/e2e-gpu.yaml
+++ b/.pipelines/e2e-gpu.yaml
@@ -1,6 +1,10 @@
 name: $(Date:yyyyMMdd)$(Rev:.r)
+parameters:
+  - name: scenarioFilter
+    type: string
+    default: "gpu=true"
 variables:
-  TAGS_TO_RUN: "gpu=true"
+  TAGS_TO_RUN: ${{ parameters.scenarioFilter }}
   TAGS_TO_SKIP: "os=windows,os=azurelinux,os=azurecontainerlinux"
   SKIP_E2E_TESTS: false
   # PRs skip quota/capacity-constrained scenarios; main and scheduled runs fail them.
diff --git a/e2e/README.md b/e2e/README.md
index 2d7c56c02ee..cafd31c2814 100644
--- a/e2e/README.md
+++ b/e2e/README.md
@@ -94,6 +94,17 @@ coverage matching production. A Scriptless tag or an AKSNodeConfig mutator alone
 does not opt into Phase 3. Existing scriptless opt-outs, image restrictions, and
 VHD-caching behavior remain in effect.
 
+### Ubuntu PIS GPU cleanup
+
+`Ubuntu2404_GPUDriverOptOut_VHDCaching` needs a new Ubuntu 24.04 Gen2 image
+containing the baked cleanup helper; use same-run E2E VHD metadata or
+`--vhd-metadata-file` with `2404gen2containerd`. The helper is baked, not
+transferred in CustomData. An opt-out PIS bake without it fails explicitly;
+ordinary nodes and derived PIS nodes do not require it.
+Opt out of the managed driver in bake CustomData (`EnableNvidia=false`,
+rendering `GPU_NODE=false`) before capture; a tag only on the derived VM does
+not request bake cleanup.
+
 ### LocalDNS fallback checks
 
 The existing LocalDNS scenarios use short-lived DNS pods on the scenario node
diff --git a/e2e/scenario/provision.go b/e2e/scenario/provision.go
index c0044fb1a2c..699a943fca9 100644
--- a/e2e/scenario/provision.go
+++ b/e2e/scenario/provision.go
@@ -70,7 +70,13 @@ func runVHDCachingScenario(ctx context.Context, name string, original *Scenario)
 				ValidateSystemdUnitIsNotRunning(ctx, scenario, "kubelet"),
 			)
 		}
-		return validationErr
+		if validationErr != nil {
+			return validationErr
+		}
+		if original.Config.VHDCachingPostBasePrep != nil {
+			return original.Config.VHDCachingPostBasePrep(ctx, scenario)
+		}
+		return nil
 	}
 
 	bakeScenario.Config.VMConfigMutator = func(vmss *armcompute.VirtualMachineScaleSet) {
@@ -832,6 +838,65 @@ func RunCommand(ctx context.Context, s *Scenario, command string) (armcompute.Vi
 	return view, runCommandScriptError(view)
 }
 
+const (
+	customerGPUDriverImage = "mcr.microsoft.com/aks/aks-gpu-cuda-lts"
+)
+
+func customerGPUDriverImageReference() (string, string, error) {
+	if datamodel.NvidiaCudaLTSDriverVersion == "" || datamodel.AKSGPUCudaLTSVersionSuffix == "" {
+		return "", "", fmt.Errorf("approved NVIDIA CUDA LTS driver image version is unavailable")
+	}
+	version := datamodel.NvidiaCudaLTSDriverVersion
+	tag := version + "-" + datamodel.AKSGPUCudaLTSVersionSuffix
+	return customerGPUDriverImage + ":" + tag, version, nil
+}
+
+// installCustomerNvidiaDriver manually runs the VHD-cached, components-approved installer after
+// basePrep. This models a customer RunCommand install; the scenario leaves managed GPU installation
+// opted out, so nodePrep does not invoke the installer.
+func installCustomerNvidiaDriver(ctx context.Context, s *Scenario) error {
+	imageRef, expectedVersion, err := customerGPUDriverImageReference()
+	if err != nil {
+		return err
+	}
+	command := fmt.Sprintf(`set -euo pipefail
+image_ref=%s
+expected_version=%s
+sudo mkdir -p /opt/gpu /opt/actions
+sudo ctr -n k8s.io images ls -q | grep -Fxq "${image_ref}"
+sudo 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 \
+  "${image_ref}" gpuinstall /entrypoint.sh copy
+sudo ctr -n k8s.io images rm "${image_ref}"
+
+installer_dir="/opt/gpu/NVIDIA-Linux-x86_64-${expected_version}"
+test -x "${installer_dir}/nvidia-installer"
+cd "${installer_dir}"
+sudo ./nvidia-installer -s -k="$(uname -r)" --no-drm --dkms
+
+sudo nvidia-modprobe -u -c0
+gpu_info="$(sudo nvidia-smi --query-gpu=driver_version,name --format=csv,noheader | head -n1)"
+driver_version="${gpu_info%%,*}"
+driver_version="$(printf '%%s' "${driver_version}" | tr -d '[:space:]')"
+gpu_name="${gpu_info#*,}"
+case "${gpu_name}" in *T4*) ;; *) echo "expected a T4 GPU, got: ${gpu_name}" >&2; exit 1 ;; esac
+test "${driver_version}" = "${expected_version}"
+sudo nvidia-smi -L | grep -qi 'T4'
+lsmod | grep -q '^nvidia[[:space:]]'
+test "$(sudo modinfo -F version nvidia)" = "${expected_version}"
+`,
+		shellSingleQuote(imageRef),
+		shellSingleQuote(expectedVersion),
+	)
+	_, err = RunCommand(ctx, s, command)
+	if err != nil {
+		return fmt.Errorf("install customer NVIDIA driver from cached approved artifact: %w", err)
+	}
+	return nil
+}
+
 func RestartVMSSVM(ctx context.Context, s *Scenario) error {
 	if s.Runtime == nil || s.Runtime.Cluster == nil || s.Runtime.Cluster.Model == nil ||
 		s.Runtime.Cluster.Model.Properties == nil || s.Runtime.Cluster.Model.Properties.NodeResourceGroup == nil ||
diff --git a/e2e/scenario/scenario.go b/e2e/scenario/scenario.go
index 2d5121a16fd..f8675e75148 100755
--- a/e2e/scenario/scenario.go
+++ b/e2e/scenario/scenario.go
@@ -2981,6 +2981,7 @@ var _ = Register(&Scenario{
 			runcVersions := components.GetExpectedPackageVersions("runc", "ubuntu", "r2404")
 			return errors.Join(
 				ValidateNvidiaSMINotInstalled(ctx, s),
+				ValidateNvidiaModulesNotLoaded(ctx, s),
 				ValidateContainerd2Properties(ctx, s, containerdVersions),
 				ValidateRuncVersion(ctx, s, runcVersions),
 			)
@@ -3106,6 +3107,37 @@ var _ = Register(&Scenario{
 	},
 })
 
+var _ = Register(&Scenario{
+	Name:             "Ubuntu2404_GPUDriverOptOut_VHDCaching",
+	Description:      "Installs the VHD-cached, components-approved NVIDIA driver artifact with a post-basePrep customer RunCommand while managed driver installation remains opted out, then verifies the real driver survives PIS VHD-cached provisioning",
+	K8sSystemPoolSKU: "Standard_D2s_v3",
+	Tags: Tags{
+		GPU: true,
+	},
+	Config: Config{
+		Cluster:                ClusterKubenet,
+		VHD:                    config.VHDUbuntu2404Gen2Containerd,
+		VHDCaching:             true,
+		VHDCachingPostBasePrep: installCustomerNvidiaDriver,
+		PreProvisionBootstrapConfigMutator: func(_ *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) {
+			nbc.EnableScriptlessCSECmd = !config.Config.DisableScriptless
+		},
+		BootstrapConfigMutator: func(_ *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) {
+			nbc.EnableScriptlessCSECmd = false
+			nbc.AgentPoolProfile.VMSize = "Standard_NC4as_T4_v3"
+			nbc.ConfigGPUDriverIfNeeded = false
+			nbc.EnableGPUDevicePluginIfNeeded = false
+			nbc.EnableNvidia = false
+		},
+		VMConfigMutator: func(vmss *armcompute.VirtualMachineScaleSet) {
+			vmss.SKU.Name = to.Ptr("Standard_NC4as_T4_v3")
+		},
+		Validator: func(ctx context.Context, s *Scenario) error {
+			return ValidateCustomerNvidiaDriver(ctx, s)
+		},
+	},
+})
+
 var _ = Register(&Scenario{
 	Name:        "Ubuntu2404_AKSVMExtension_FilesystemCorruption",
 	Description: "Tests Ubuntu 24.04 NPD filesystem-corruption reporting with the AKS VM Extension and hotfix target resolution from the NBC command when ANC JSON is absent",
diff --git a/e2e/scenario/types.go b/e2e/scenario/types.go
index 27a2c4456dd..cefcf180b58 100644
--- a/e2e/scenario/types.go
+++ b/e2e/scenario/types.go
@@ -198,6 +198,10 @@ type Config struct {
 	// The main purpose is to validate VHD Caching logic and ensure a reboot step between basePrep and nodePrep doesn't break anything.
 	VHDCaching bool
 
+	// VHDCachingPostBasePrep runs on the bake VM after basePrep completes and before its disk is captured.
+	// Use it only for VHD-caching scenarios that need to seed state between the two provisioning phases.
+	VHDCachingPostBasePrep func(context.Context, *Scenario) error
+
 	// ExpectedError, when set, indicates that VMSS creation is expected to fail with an error containing this substring.
 	// The assertion is performed during the scenario run.
 	ExpectedError string
diff --git a/e2e/scenario/validators.go b/e2e/scenario/validators.go
index bd318d124a6..160706a009c 100644
--- a/e2e/scenario/validators.go
+++ b/e2e/scenario/validators.go
@@ -621,12 +621,49 @@ func ValidateNvidiaSMINotInstalled(ctx context.Context, s *Scenario) error {
 	return assert.Contains(execResult.stderr, "nvidia-smi: command not found", "expected stderr to contain 'nvidia-smi: command not found', but got %q", execResult.stderr)
 }
 
+func ValidateNvidiaModulesNotLoaded(ctx context.Context, s *Scenario) error {
+	_, err := execScriptOnVMForScenarioValidateExitCode(ctx, s, nvidiaModulesNotLoadedScript(), 0, "NVIDIA kernel modules remain loaded")
+	return err
+}
+
+func nvidiaModulesNotLoadedScript() string {
+	return `set -euo pipefail
+lsmod | awk '$1 == "nvidia" || $1 ~ /^nvidia_/ { print "unexpected loaded NVIDIA module: " $1 > "/dev/stderr"; exit 1 }'`
+}
+
 func ValidateNvidiaSMIInstalled(ctx context.Context, s *Scenario) error {
 	command := []string{"set -ex", "sudo nvidia-smi"}
 	_, err := execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "could not execute nvidia-smi command")
 	return err
 }
 
+func ValidateCustomerNvidiaDriver(ctx context.Context, s *Scenario) error {
+	_, expectedVersion, err := customerGPUDriverImageReference()
+	if err != nil {
+		return err
+	}
+	installerPath := fmt.Sprintf("/opt/gpu/NVIDIA-Linux-x86_64-%s/nvidia-installer", expectedVersion)
+	command := customerNvidiaDriverValidationScript(expectedVersion, installerPath)
+	_, err = execScriptOnVMForScenarioValidateExitCode(ctx, s, command, 0, "customer-installed NVIDIA driver did not survive PIS image capture")
+	return err
+}
+
+func customerNvidiaDriverValidationScript(expectedVersion, installerPath string) string {
+	return fmt.Sprintf(`set -euo pipefail
+expected_version=%s
+test -x %s
+lsmod | grep -q '^nvidia[[:space:]]'
+module_version="$(sudo modinfo -F version nvidia)"
+test "${module_version}" = "${expected_version}"
+gpu_version="$(sudo nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n1 | tr -d '[:space:]')"
+test "${gpu_version}" = "${expected_version}"
+sudo nvidia-smi
+`,
+		shellSingleQuote(expectedVersion),
+		shellSingleQuote(installerPath),
+	)
+}
+
 func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) error {
 	command := []string{
 		"set -ex",
diff --git a/e2e/scenario/validators_gpu_driver_test.go b/e2e/scenario/validators_gpu_driver_test.go
index d2ba0b3add2..cc6a888a473 100644
--- a/e2e/scenario/validators_gpu_driver_test.go
+++ b/e2e/scenario/validators_gpu_driver_test.go
@@ -65,3 +65,31 @@ func TestNvidiaDriverVersionValidationScriptRejectsInvalidVersion(t *testing.T)
 		})
 	}
 }
+
+func TestNvidiaModulesNotLoadedScript(t *testing.T) {
+	tests := []struct {
+		name      string
+		modules   string
+		lsmodExit string
+		wantExit  int
+	}{
+		{"unrelated modules", "Module Size Used by\nx_tables 123 0", "0", 0},
+		{"core module", "Module Size Used by\nnvidia 123 0", "0", 1},
+		{"auxiliary module", "Module Size Used by\nnvidia_uvm 123 0", "0", 1},
+		{"inspection failure", "Module Size Used by", "17", 17},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			command := exec.Command("bash", "-c", `lsmod() { printf '%s\n' "$TEST_LSMOD_OUTPUT"; return "$TEST_LSMOD_EXIT"; }
+`+nvidiaModulesNotLoadedScript())
+			command.Env = append(os.Environ(), "TEST_LSMOD_OUTPUT="+tt.modules, "TEST_LSMOD_EXIT="+tt.lsmodExit)
+			output, err := command.CombinedOutput()
+			if tt.wantExit == 0 {
+				require.NoError(t, err, string(output))
+			} else {
+				require.Error(t, err)
+			}
+			require.Equal(t, tt.wantExit, command.ProcessState.ExitCode(), string(output))
+		})
+	}
+}
diff --git a/e2e/scenario/validators_test.go b/e2e/scenario/validators_test.go
index 1b23beebaad..e4b25f18a77 100644
--- a/e2e/scenario/validators_test.go
+++ b/e2e/scenario/validators_test.go
@@ -3,6 +3,10 @@ package scenario
 import (
 	"crypto/sha256"
 	"encoding/hex"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
@@ -10,6 +14,41 @@ import (
 	resourcev1 "k8s.io/api/resource/v1"
 )
 
+func TestCustomerNvidiaDriverValidationRequiresStagedInstaller(t *testing.T) {
+	const expectedVersion = "580.159.04"
+	binDir := t.TempDir()
+	installerPath := filepath.Join(binDir, "nvidia-installer")
+	// PATH-local lsmod and sudo stand in for GPU hardware while the real
+	// validation script checks the staged installer and driver state.
+	require.NoError(t, os.WriteFile(filepath.Join(binDir, "lsmod"), []byte("#!/bin/sh\nprintf 'nvidia 123 0\\n'\n"), 0o755))
+	sudo := fmt.Sprintf(`#!/bin/sh
+case "$1" in
+  modinfo) printf '%%s\n' '%s'; exit 0 ;;
+  nvidia-smi)
+    if [ "$2" = "--query-gpu=driver_version" ]; then
+      printf '%%s\n' '%s'
+    fi
+    exit 0
+    ;;
+esac
+exit 1
+`, expectedVersion, expectedVersion)
+	require.NoError(t, os.WriteFile(filepath.Join(binDir, "sudo"), []byte(sudo), 0o755))
+	t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH"))
+
+	runValidation := func() error {
+		output, err := exec.Command("bash", "-x", "-c", customerNvidiaDriverValidationScript(expectedVersion, installerPath)).CombinedOutput()
+		if err != nil {
+			return fmt.Errorf("%w: %s", err, output)
+		}
+		return nil
+	}
+	require.Error(t, runValidation(), "validation must reject correct driver state when the staged installer was removed")
+
+	require.NoError(t, os.WriteFile(installerPath, []byte("#!/bin/sh\n"), 0o755))
+	require.NoError(t, runValidation(), "validation must pass when the staged installer and driver state are present")
+}
+
 func TestValidateDRAMIGAllocation(t *testing.T) {
 	tests := []struct {
 		name    string
diff --git a/e2e/scenario/vmss.go b/e2e/scenario/vmss.go
index bc0b1785188..be232f77f6b 100644
--- a/e2e/scenario/vmss.go
+++ b/e2e/scenario/vmss.go
@@ -526,7 +526,7 @@ func createVMSSModel(ctx context.Context, s *Scenario) (armcompute.VirtualMachin
 			return armcompute.VirtualMachineScaleSet{}, fmt.Errorf("inject customData write_files entries: %w", err)
 		}
 	}
-	if !config.Config.DisableScriptless && !scriptlessProvisioningEnabled && s.VHD.SupportsScriptless() {
+	if s.Runtime.NBC.EnableScriptlessCSECmd && !scriptlessProvisioningEnabled && s.VHD.SupportsScriptless() {
 		// Validate that the custom data doesn't contain any script content,
 		// which indicates that the scriptless CSE is working as intended
 		decodedCustomData, err := base64.StdEncoding.DecodeString(customData)
diff --git a/e2e/scenario/vmss_test.go b/e2e/scenario/vmss_test.go
index 6b77e77c270..2c2cc1be07d 100644
--- a/e2e/scenario/vmss_test.go
+++ b/e2e/scenario/vmss_test.go
@@ -44,6 +44,67 @@ func (f vmssCreationTestPolicy) Do(req *policy.Request) (*http.Response, error)
 	return f(req.Raw()), nil
 }
 
+func TestCreateVMSSModelHonorsScriptlessNBCOverride(t *testing.T) {
+	oldScriptless := config.Config.DisableScriptless
+	oldPrepareVHD := CachedPrepareVHD
+	config.Config.DisableScriptless = false
+	CachedPrepareVHD = func(context.Context, GetVHDRequest) (config.VHDResourceID, error) {
+		return "fixture-image", nil
+	}
+	t.Cleanup(func() {
+		config.Config.DisableScriptless = oldScriptless
+		CachedPrepareVHD = oldPrepareVHD
+	})
+
+	for _, enableScriptlessCSECmd := range []bool{false, true} {
+		name := "traditional override"
+		if enableScriptlessCSECmd {
+			name = "scriptless CSE"
+		}
+		t.Run(name, func(t *testing.T) {
+			vhd := *config.VHDUbuntu2404Gen2Containerd
+			cluster := &Cluster{
+				Model: &armcontainerservice.ManagedCluster{
+					Name:     to.Ptr("cluster"),
+					Location: to.Ptr("westus3"),
+					Properties: &armcontainerservice.ManagedClusterProperties{
+						CurrentKubernetesVersion: to.Ptr("1.34.0"),
+						NodeResourceGroup:        to.Ptr("node-rg"),
+						NetworkProfile: &armcontainerservice.NetworkProfile{
+							NetworkPlugin: to.Ptr(armcontainerservice.NetworkPluginKubenet),
+						},
+					},
+				},
+				KubeletIdentity: &armcontainerservice.UserAssignedIdentity{
+					ClientID:   to.Ptr("client-id"),
+					ResourceID: to.Ptr("/subscriptions/sub/resourceGroups/rg/providers/Microsoft.ManagedIdentity/userAssignedIdentities/kubelet"),
+				},
+				ClusterParams: &ClusterParams{
+					BootstrapToken: "token",
+					CACert:         []byte("ca-cert"),
+					FQDN:           "cluster.example.test",
+				},
+			}
+			nbc, err := getBaseNBC(t.Context(), cluster, &vhd)
+			require.NoError(t, err)
+			nbc.EnableScriptlessCSECmd = enableScriptlessCSECmd
+			s := &Scenario{
+				Config:   Config{VHD: &vhd},
+				Location: "westus3",
+				Runtime: &ScenarioRuntime{
+					NBC:      nbc,
+					Cluster:  cluster,
+					VMSSName: "vmss",
+				},
+			}
+
+			model, err := createVMSSModel(t.Context(), s)
+			require.NoError(t, err)
+			require.NotEmpty(t, *model.Properties.VirtualMachineProfile.OSProfile.CustomData)
+		})
+	}
+}
+
 type vmssCreationTestCase struct {
 	name          string
 	provisionCode string