From d9ddd90d82498082cdb69dbe6742390973c829c0 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Mon, 5 Oct 2026 17:52:42 +1300
Subject: [PATCH 01/24] fix(gpu): preserve customer-installed PIS drivers

Clean AKS-owned artifacts before PIS customization and verify teardown independently of the prebake marker. Preserve customer drivers on derived nodes while retaining non-PIS cleanup.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 e2e/scenario/provision.go                     |  22 ++-
 e2e/scenario/scenario.go                      |  29 +++
 e2e/scenario/types.go                         |   4 +
 e2e/scenario/vmss.go                          |   2 +-
 parts/linux/cloud-init/artifacts/cse_main.sh  |   8 +-
 parts/linux/cloud-init/artifacts/cse_start.sh |  10 +-
 .../artifacts/mariner/cse_install_mariner.sh  |  16 ++
 .../artifacts/ubuntu/cse_install_ubuntu.sh    |  34 +++-
 .../artifacts/cse_install_mariner_spec.sh     |  26 +++
 .../artifacts/cse_install_ubuntu_spec.sh      |  69 +++++++
 .../artifacts/cse_main_gpu_pis_spec.sh        | 181 ++++++++++++++++++
 11 files changed, 393 insertions(+), 8 deletions(-)
 create mode 100644 spec/parts/linux/cloud-init/artifacts/cse_main_gpu_pis_spec.sh

diff --git a/e2e/scenario/provision.go b/e2e/scenario/provision.go
index b9494bbd346..8ac452ef523 100644
--- a/e2e/scenario/provision.go
+++ b/e2e/scenario/provision.go
@@ -68,7 +68,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) {
@@ -755,6 +761,20 @@ func RunCommand(ctx context.Context, s *Scenario, command string) (armcompute.Vi
 	return view, runCommandScriptError(view)
 }
 
+const customerDriverSentinelPath = "/usr/local/nvidia/customer-driver-sentinel"
+
+func installCustomerDriverSentinel(ctx context.Context, s *Scenario) error {
+	command := fmt.Sprintf(
+		"sudo mkdir -p /usr/local/nvidia && printf '%%s\\n' '#!/bin/sh' '# customer-installed driver sentinel' 'exit 0' | sudo tee %s >/dev/null && sudo chmod 0755 %s",
+		customerDriverSentinelPath,
+		customerDriverSentinelPath,
+	)
+	if _, err := RunCommand(ctx, s, command); err != nil {
+		return fmt.Errorf("install customer driver sentinel: %w", err)
+	}
+	return nil
+}
+
 // runCommandScriptError converts a RunCommand instance view into an error if the
 // script itself failed. The ARM CreateOrUpdate operation reports success as long as
 // the extension was able to run the script — a non-zero exit, throw, or timeout
diff --git a/e2e/scenario/scenario.go b/e2e/scenario/scenario.go
index dd4b1177435..b6dd736685b 100644
--- a/e2e/scenario/scenario.go
+++ b/e2e/scenario/scenario.go
@@ -3049,6 +3049,35 @@ var _ = Register(&Scenario{
 	},
 })
 
+var _ = Register(&Scenario{
+	Name:             "Ubuntu2404_GPUDriverOptOut_VHDCaching",
+	Description:      "Tests that a customer-installed NVIDIA driver sentinel survives PIS VHD-cached provisioning on an opted-out GPU node",
+	Location:         "westus2",
+	K8sSystemPoolSKU: "Standard_D2s_v3",
+	Tags: Tags{
+		GPU: true,
+	},
+	Config: Config{
+		Cluster:                ClusterKubenet,
+		VHD:                    config.VHDUbuntu2404Gen2Containerd,
+		VHDCaching:             true,
+		VHDCachingPostBasePrep: installCustomerDriverSentinel,
+		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 ValidateFileHasContent(ctx, s, customerDriverSentinelPath, "customer-installed driver sentinel")
+		},
+	},
+})
+
 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 af5145c58e5..50284766455 100644
--- a/e2e/scenario/types.go
+++ b/e2e/scenario/types.go
@@ -192,6 +192,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/vmss.go b/e2e/scenario/vmss.go
index bd86761a4e7..764b1c92d9a 100644
--- a/e2e/scenario/vmss.go
+++ b/e2e/scenario/vmss.go
@@ -420,7 +420,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 && !scriptlessNBCCSECmdEnabled && s.VHD.SupportsScriptless() {
+	if !config.Config.DisableScriptless && !scriptlessNBCCSECmdEnabled && s.VHD.SupportsScriptless() && s.Runtime.NBC.EnableScriptlessCSECmd {
 		// 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/parts/linux/cloud-init/artifacts/cse_main.sh b/parts/linux/cloud-init/artifacts/cse_main.sh
index 9a801488766..96421b63075 100755
--- a/parts/linux/cloud-init/artifacts/cse_main.sh
+++ b/parts/linux/cloud-init/artifacts/cse_main.sh
@@ -407,6 +407,11 @@ EOF
         logs_to_events "AKS.CSE.ensureContainerd.ensureArtifactStreaming" ensureArtifactStreaming || exit $ERR_ARTIFACT_STREAMING_INSTALL
     fi
 
+    if [ "${PRE_PROVISION_ONLY}" = "true" ] && [ "${GPU_NODE}" != "true" ] &&
+       { [ "${OS}" = "${UBUNTU_OS_NAME}" ] || { isMarinerOrAzureLinux "$OS" "$OS_VARIANT" && ! isAzureLinuxOSGuard "$OS" "$OS_VARIANT"; }; }; then
+        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
@@ -644,7 +649,7 @@ function nodePrep {
 
 
     # Clean up GPU drivers if not a GPU node or if skipping driver install
-    if [ "${GPU_NODE}" != "true" ] || [ "${skip_nvidia_driver_install}" = "true" ]; then
+    if [ "${BASE_PREP_RAN}" = "true" ] && { [ "${GPU_NODE}" != "true" ] || [ "${skip_nvidia_driver_install}" = "true" ]; }; then
         logs_to_events "AKS.CSE.cleanUpGPUDrivers" cleanUpGPUDrivers
     fi
 
@@ -750,6 +755,7 @@ function nodePrep {
 # when nodes are created from that VHD image.
 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 bece0b6081f..e5a4eda3d82 100644
--- a/parts/linux/cloud-init/artifacts/cse_start.sh
+++ b/parts/linux/cloud-init/artifacts/cse_start.sh
@@ -114,10 +114,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 8d19ce74b2c..1b861f33107 100755
--- a/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh
+++ b/parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh
@@ -833,6 +833,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.sh b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
index 0d86742de72..9886e25b002 100755
--- a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
+++ b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
@@ -369,7 +369,7 @@ cleanUpPrebakedGPUDriver() {
     local dkms_after=false modprobe_after=false marker_after=true status=cleaned
     [ -d /var/lib/dkms/nvidia ] && dkms_after=true
     [ -e /usr/bin/nvidia-modprobe ] && modprobe_after=true
-    if [ "${dkms_after}" = false ] && [ "${modprobe_after}" = false ] && [ "${module_after}" = false ]; then
+    if ! prebakedGPUDriverArtifactsRemain; then
         rm -f "${marker}" || true
         [ -f "${marker}" ] || marker_after=false
     fi
@@ -379,6 +379,19 @@ cleanUpPrebakedGPUDriver() {
     echo "AKS_GPU_PREBAKE event=teardown gpu_node=${GPU_NODE:-} status=${status} dkms_before=${dkms_before} module_before=${module_before} module_after=${module_after} marker_after=${marker_after} dkms_after=${dkms_after} modprobe_after=${modprobe_after}"
 }
 
+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}"
+}
+
 cleanUpGPUDrivers() {
     rm -Rf $GPU_DEST /opt/gpu
 
@@ -393,6 +406,25 @@ cleanUpGPUDrivers() {
     cleanUpPrebakedGPUDriver
 }
 
+cleanUpGPUDriversForBasePrep() {
+    local packageName marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}"
+    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
+}
+
 installCriCtlPackage() {
     local version="${1:-}"
     local fullPackageVersion
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..77c7c9f0171 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="/tmp/nonexistent-shellspec-gpu-dest"
+        managedGPUPackageList() { echo "driver-cache"; }
+        getPackageCacheDir() { echo "$gpu_test_root/driver-cache"; }
+        setup_gpu_cache() { gpu_test_root=$(mktemp -d); }
+        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"
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 8c18fb3badd..1665a3cddf8 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
@@ -85,6 +85,75 @@ Describe 'cse_install_ubuntu.sh'
             The output should include "module_after=true"
             The output should include "status=incomplete"
         End
+
+        It 'keeps the marker if any prebaked userspace artifact remains'
+            marker="$(mktemp)"
+            GPU_DKMS_MARKER_FILE="$marker"
+            rm() { echo "mock rm $*"; }
+            ldconfig() { :; }
+            lsmod() { :; }
+            prebakedGPUDriverArtifactsRemain() { return 0; }
+            When call cleanUpPrebakedGPUDriver
+            The status should be success
+            The output should include "status=incomplete"
+            The output should not include "mock rm -f $marker"
+        End
+
+        It 'keeps the marker when loaded-module inspection fails'
+            marker="$(mktemp)"
+            GPU_DKMS_MARKER_FILE="$marker"
+            rm() { echo "mock rm $*"; }
+            ldconfig() { :; }
+            lsmod() { return 1; }
+            When call cleanUpPrebakedGPUDriver
+            The status should be success
+            The output should include "status=incomplete"
+            The output should not include "mock rm -f $marker"
+        End
+    End
+
+    Describe 'cleanUpGPUDriversForBasePrep'
+        GPU_DEST="/tmp/nonexistent-shellspec-gpu-dest"
+        managedGPUPackageList() { :; }
+
+        It 'fails when the prebake marker survives cleanup'
+            GPU_DKMS_MARKER_FILE="$(mktemp)"
+            cleanUpGPUDrivers() { :; }
+            lsmod() { echo ""; }
+            When call cleanUpGPUDriversForBasePrep
+            The status should be failure
+            The stderr should include "GPU basePrep cleanup incomplete"
+            The path "$GPU_DKMS_MARKER_FILE" should be exist
+            AfterRun 'rm -f "$GPU_DKMS_MARKER_FILE"'
+        End
+
+        It 'fails when nvidia remains loaded without a prebake marker'
+            GPU_DKMS_MARKER_FILE="${PWD}/.shellspec-absent-gpu-marker-$$"
+            cleanUpGPUDrivers() { :; }
+            lsmod() { echo "nvidia 104165376 0"; }
+            When call cleanUpGPUDriversForBasePrep
+            The status should be failure
+            The stderr should include "GPU basePrep cleanup incomplete"
+        End
+
+        It 'fails when an unmarked prebake artifact remains'
+            GPU_DKMS_MARKER_FILE="${PWD}/.shellspec-absent-gpu-marker-$$"
+            cleanUpGPUDrivers() { :; }
+            compgen() { [ "$2" = '/lib/modules/*/updates/dkms/nvidia*.ko*' ]; }
+            lsmod() { echo ""; }
+            When call cleanUpGPUDriversForBasePrep
+            The status should be failure
+            The stderr should include "GPU basePrep cleanup incomplete"
+        End
+
+        It 'succeeds when the prebake marker and AKS paths are gone'
+            GPU_DKMS_MARKER_FILE="$(mktemp)"
+            cleanUpGPUDrivers() { rm -f "$GPU_DKMS_MARKER_FILE"; }
+            lsmod() { echo ""; }
+            When call cleanUpGPUDriversForBasePrep
+            The status should be success
+            The path "$GPU_DKMS_MARKER_FILE" should not be exist
+        End
     End
 
     Describe 'installPackageFromCache version matching'
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..a6e2ead837c
--- /dev/null
+++ b/spec/parts/linux/cloud-init/artifacts/cse_main_gpu_pis_spec.sh
@@ -0,0 +1,181 @@
+#!/bin/bash
+
+Describe 'GPU cleanup across PIS phases'
+    CSE_MAIN="./parts/linux/cloud-init/artifacts/cse_main.sh"
+
+    run_phases() {
+        local dispatch cleanup
+        dispatch=$(awk '
+            /^if \[ ! -f \/opt\/azure\/containers\/base_prep.complete \]; then/ { reading = 1 }
+            reading && /^echo "Custom script finished."/ { exit }
+            reading { print }
+        ' "$CSE_MAIN")
+        cleanup=$(sed -n '/    # Clean up GPU drivers if not a GPU node/,/    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" "$driver_root/var/lib/dkms/nvidia" \
+                "$driver_root/lib/modules/test/updates/dkms/nvidia.ko" "$driver_root/usr/bin/nvidia-smi"
+            rm -f "$driver_root/dkms-marker"
+        }
+        eval "$dispatch"
+    }
+
+    driver_survives() {
+        run_phases
+        test -d "$driver_root/usr/local/nvidia" &&
+            test -d "$driver_root/var/lib/dkms/nvidia" &&
+            test -f "$driver_root/lib/modules/test/updates/dkms/nvidia.ko" &&
+            test -f "$driver_root/usr/bin/nvidia-smi" &&
+            test -f "$driver_root/dkms-marker"
+    }
+
+    setup_driver() {
+        driver_root=$(mktemp -d)
+        marker="$driver_root/base_prep.complete"
+        mkdir -p "$driver_root/usr/local/nvidia" "$driver_root/var/lib/dkms/nvidia" \
+            "$driver_root/lib/modules/test/updates/dkms" "$driver_root/usr/bin"
+        touch "$driver_root/dkms-marker" \
+            "$driver_root/lib/modules/test/updates/dkms/nvidia.ko" "$driver_root/usr/bin/nvidia-smi"
+        GPU_NODE=false
+        PRE_PROVISION_ONLY=false
+        skip_nvidia_driver_install=false
+    }
+    BeforeEach 'setup_driver'
+    AfterEach 'rm -rf "$driver_root"'
+
+    It 'preserves customer-installed driver artifacts on a derived opt-out node'
+        touch "$marker"
+        When call driver_survives
+        The status should be success
+        The output should include "Skipping basePrep"
+    End
+
+    It 'preserves customer drivers on a derived GPU node with the legacy skip tag'
+        touch "$marker"
+        GPU_NODE=true
+        skip_nvidia_driver_install=true
+        When call driver_survives
+        The status should be success
+        The output should include "Skipping basePrep"
+    End
+
+    Describe 'PIS basePrep completion marker'
+        run_cse_marker() {
+            local marker_script
+            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/$marker_root}
+            marker_script=${marker_script//\/var\/log\/azure\/cluster-provision.log/$marker_root/provision.log}
+            eval "$marker_script"
+        }
+
+        setup_marker() { marker_root=$(mktemp -d); PRE_PROVISION_ONLY=true; }
+        BeforeEach 'setup_marker'
+        AfterEach 'rm -rf "$marker_root"'
+
+        It 'does not capture a failed GPU teardown as a completed image'
+            EXIT_CODE=224
+            When call run_cse_marker
+            The status should be success
+            The path "$marker_root/base_prep.complete" should not be exist
+        End
+
+        It 'marks a successful bake for derived nodes'
+            EXIT_CODE=0
+            When call run_cse_marker
+            The status should be success
+            The path "$marker_root/base_prep.complete" should be exist
+        End
+    End
+
+    It 'still cleans an opt-out node when 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
+    End
+
+    It 'still cleans a non-PIS node with the legacy skip-driver 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
+    End
+
+    It 'does not clean a GPU bake VM'
+        PRE_PROVISION_ONLY=true
+        When call driver_survives
+        The status should be success
+        The output should include "basePrep"
+    End
+
+    base_cleanup() {
+        local cleanup
+        cleanup=$(sed -n '/    if \[ "${PRE_PROVISION_ONLY}" = "true" \] && \[ "${GPU_NODE}" != "true" \] &&/,/    fi/p' "$CSE_MAIN")
+        logs_to_events() { shift; "$@"; }
+        cleanUpGPUDriversForBasePrep() { echo "basePrep driver cleanup"; return "${CLEANUP_STATUS:-0}"; }
+        isMarinerOrAzureLinux() { [ "$OS" = azurelinux ]; }
+        isAzureLinuxOSGuard() { [ "$OS_VARIANT" = osguard ]; }
+        UBUNTU_OS_NAME=ubuntu
+        ERR_NVIDIA_DRIVER_INSTALL=224
+        ( eval "$cleanup" )
+    }
+
+    It 'tears down AKS drivers on a non-GPU bake VM before image capture'
+        PRE_PROVISION_ONLY=true
+        OS=ubuntu
+        When call base_cleanup
+        The status should be success
+        The output should include "basePrep driver cleanup"
+    End
+
+    It 'fails basePrep when bake-time teardown cannot be verified'
+        PRE_PROVISION_ONLY=true
+        OS=ubuntu
+        CLEANUP_STATUS=1
+        When call base_cleanup
+        The status should equal 224
+        The output should include "basePrep driver cleanup"
+    End
+
+    It 'does not tear down AKS drivers on a GPU bake VM'
+        PRE_PROVISION_ONLY=true
+        OS=ubuntu
+        GPU_NODE=true
+        When call base_cleanup
+        The status should be success
+        The output should not include "basePrep driver cleanup"
+    End
+
+    It 'does not tear down AKS drivers in non-PIS basePrep'
+        OS=ubuntu
+        When call base_cleanup
+        The status should be success
+        The output should not include "basePrep driver cleanup"
+    End
+
+    It 'tears down AKS caches on an Azure Linux bake VM'
+        PRE_PROVISION_ONLY=true
+        OS=azurelinux
+        When call base_cleanup
+        The status should be success
+        The output should include "basePrep driver cleanup"
+    End
+
+    It 'does not call unavailable cleanup on Azure Container Linux'
+        PRE_PROVISION_ONLY=true
+        OS=azurelinux
+        OS_VARIANT=osguard
+        When call base_cleanup
+        The status should be success
+        The output should not include "basePrep driver cleanup"
+    End
+End

From 24ac59ed5981c7450ac87a105ae5c0639fe955ff Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Mon, 5 Oct 2026 22:58:20 +1300
Subject: [PATCH 02/24] test(e2e): diagnose image capture stalls

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 e2e/scenario/cleanup.go        | 16 +++++++++++++++-
 e2e/scenario/execution_test.go |  6 +++++-
 e2e/scenario/provision.go      |  6 +++++-
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/e2e/scenario/cleanup.go b/e2e/scenario/cleanup.go
index a3fa1e3d18e..3b2881ee105 100644
--- a/e2e/scenario/cleanup.go
+++ b/e2e/scenario/cleanup.go
@@ -4,9 +4,13 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"reflect"
+	"runtime"
 	"runtime/debug"
 	"sync"
 	"time"
+
+	"github.com/Azure/agentbaker/e2e/logging"
 )
 
 const (
@@ -50,8 +54,18 @@ func (c *scenarioCleanup) runCleanups(ctx context.Context) error {
 		batchErrs := make([]error, len(cleanups))
 		var wg sync.WaitGroup
 		for i, fn := range cleanups {
+			i, fn := i, fn
+			name := fmt.Sprintf("cleanup[%d]", i+1)
+			if callback := runtime.FuncForPC(reflect.ValueOf(fn).Pointer()); callback != nil {
+				name = callback.Name()
+			}
+			deadline, _ := ctx.Deadline()
+			started := time.Now()
+			logging.Logf(ctx, "scenario cleanup %s started (deadline %s)", name, deadline.Format(time.RFC3339))
 			wg.Go(func() {
-				batchErrs[i] = runWithPanicRecovery(ctx, fn)
+				err := runWithPanicRecovery(ctx, fn)
+				batchErrs[i] = err
+				logging.Logf(ctx, "scenario cleanup %s finished in %s (error: %v, context error: %v)", name, time.Since(started), err, ctx.Err())
 			})
 		}
 		wg.Wait()
diff --git a/e2e/scenario/execution_test.go b/e2e/scenario/execution_test.go
index bae464aa6eb..bee67c162d9 100644
--- a/e2e/scenario/execution_test.go
+++ b/e2e/scenario/execution_test.go
@@ -57,7 +57,11 @@ func TestExecutionKeepsLoggerThroughSkipRunAndCleanup(t *testing.T) {
 			} else {
 				require.NoError(t, outcome.Error)
 			}
-			assert.Equal(t, []string{"skip check", "running scenario", "cleanup"}, logger.logs)
+			require.Len(t, logger.logs, 5)
+			assert.Equal(t, []string{"skip check", "running scenario"}, logger.logs[:2])
+			assert.Contains(t, logger.logs[2], "scenario cleanup ")
+			assert.Equal(t, "cleanup", logger.logs[3])
+			assert.Contains(t, logger.logs[4], "scenario cleanup ")
 		})
 	}
 }
diff --git a/e2e/scenario/provision.go b/e2e/scenario/provision.go
index 8ac452ef523..eca48da42d6 100644
--- a/e2e/scenario/provision.go
+++ b/e2e/scenario/provision.go
@@ -934,9 +934,13 @@ func CreateImage(ctx context.Context, s *Scenario) (*config.Image, error) {
 }
 
 // CreateSIGImageVersionFromDisk creates a new SIG image version directly from a VM disk
-func CreateSIGImageVersionFromDisk(ctx context.Context, s *Scenario, version string, diskResourceID string) (*config.Image, error) {
+func CreateSIGImageVersionFromDisk(ctx context.Context, s *Scenario, version string, diskResourceID string) (result *config.Image, returnErr error) {
 	startTime := time.Now()
 	defer func() {
+		if returnErr != nil {
+			logging.Logf(ctx, "SIG image version %s from disk %s failed after %s: %v", version, diskResourceID, time.Since(startTime), returnErr)
+			return
+		}
 		logging.Logf(ctx, "Created SIG image version %s from disk %s in %s", version, diskResourceID, time.Since(startTime))
 	}()
 	if s.Runtime == nil || s.Runtime.VM == nil || s.Runtime.VM.VM == nil ||

From b9d6f120df6974d0e5d1c93538c8834dd7669dcb Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Mon, 5 Oct 2026 23:46:56 +1300
Subject: [PATCH 03/24] test(gpu): verify real driver survives PIS

Install a cached NVIDIA driver through customer RunCommand, then check
version, binary/module hashes and GPU functionality on the derived node.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 e2e/scenario/provision.go  | 102 ++++++++++++++++++++++++++++++++++---
 e2e/scenario/scenario.go   |   9 ++--
 e2e/scenario/validators.go |  46 +++++++++++++++++
 3 files changed, 147 insertions(+), 10 deletions(-)

diff --git a/e2e/scenario/provision.go b/e2e/scenario/provision.go
index eca48da42d6..fb532bcc3b7 100644
--- a/e2e/scenario/provision.go
+++ b/e2e/scenario/provision.go
@@ -761,16 +761,104 @@ func RunCommand(ctx context.Context, s *Scenario, command string) (armcompute.Vi
 	return view, runCommandScriptError(view)
 }
 
-const customerDriverSentinelPath = "/usr/local/nvidia/customer-driver-sentinel"
+const (
+	customerDriverSentinelPath = "/usr/local/nvidia/customer-driver-sentinel"
+	customerDriverReceiptPath  = "/opt/azure/containers/customer-gpu-driver.receipt"
+	customerNvidiaCDIDropIn    = "/etc/systemd/system/nvidia-cdi-refresh.service.d/10-aks-tolerate-generate-failure.conf"
+	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
+}
 
-func installCustomerDriverSentinel(ctx context.Context, s *Scenario) error {
-	command := fmt.Sprintf(
-		"sudo mkdir -p /usr/local/nvidia && printf '%%s\\n' '#!/bin/sh' '# customer-installed driver sentinel' 'exit 0' | sudo tee %s >/dev/null && sudo chmod 0755 %s",
-		customerDriverSentinelPath,
-		customerDriverSentinelPath,
+// 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, _, err := customerGPUDriverImageReference()
+	if err != nil {
+		return err
+	}
+	command := fmt.Sprintf(`set -euo pipefail
+preferred_image_ref=%s
+image_name=%s
+receipt=%s
+
+sudo mkdir -p /opt/gpu /opt/actions
+cached_image_refs="$(sudo ctr -n k8s.io images ls -q | awk -v prefix="${image_name}:" 'index($0, prefix) == 1 {
+  tag = substr($0, length(prefix) + 1)
+  if (tag ~ /^[0-9]+(\.[0-9]+)+-[0-9]+$/) print $0
+}' | sort -V)"
+if [ -z "${cached_image_refs}" ]; then
+  echo "no versioned cached approved NVIDIA driver image found for ${image_name}" >&2
+  exit 1
+fi
+if printf '%%s\n' "${cached_image_refs}" | grep -Fxq "${preferred_image_ref}"; then
+  image_ref="${preferred_image_ref}"
+else
+  image_ref="$(printf '%%s\n' "${cached_image_refs}" | tail -n1)"
+  echo "Using cached approved NVIDIA driver image ${image_ref}; current preferred image ${preferred_image_ref} is not present"
+fi
+image_tag="${image_ref##*:}"
+expected_version="${image_tag%%-*}"
+# Match configureNvidiaCDIRefresh in CSE before the NVIDIA toolkit package starts this unit.
+sudo install -d -m 0755 /etc/systemd/system/nvidia-cdi-refresh.service.d
+printf '[Service]\nSuccessExitStatus=1 2 127\n' | sudo tee %s >/dev/null
+sudo systemctl daemon-reload
+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 install
+sudo ctr -n k8s.io images rm "${image_ref}"
+
+sudo nvidia-modprobe -u -c0
+gpu_info="$(sudo nvidia-smi --query-gpu=driver_version,name --format=csv,noheader | head -n1)"
+driver_version="$(printf '%%s\n' "${gpu_info}" | awk -F, '{ gsub(/^[[:space:]]+|[[:space:]]+$/, "", $1); print $1 }')"
+gpu_name="$(printf '%%s\n' "${gpu_info}" | awk -F, '{ gsub(/^[[:space:]]+|[[:space:]]+$/, "", $2); print $2 }')"
+test "${driver_version}" = "${expected_version}"
+case "${gpu_name}" in *T4*) ;; *) echo "expected a T4 GPU, got: ${gpu_name}" >&2; exit 1 ;; esac
+sudo nvidia-smi -L | grep -qi 'T4'
+lsmod | grep -q '^nvidia[[:space:]]'
+test "$(sudo modinfo -F version nvidia)" = "${expected_version}"
+
+nvidia_smi_path="$(readlink -f "$(command -v nvidia-smi)")"
+nvidia_module_path="$(sudo modinfo -n nvidia)"
+test -x "${nvidia_smi_path}"
+test -f "${nvidia_module_path}"
+nvidia_smi_sha256="$(sha256sum "${nvidia_smi_path}" | awk '{ print $1 }')"
+nvidia_module_sha256="$(sha256sum "${nvidia_module_path}" | awk '{ print $1 }')"
+sudo install -d -m 0755 "$(dirname "${receipt}")"
+{
+  printf 'image_ref=%%s\n' "${image_ref}"
+  printf 'driver_version=%%s\n' "${driver_version}"
+  printf 'gpu_name=%%s\n' "${gpu_name}"
+  printf 'nvidia_smi_path=%%s\n' "${nvidia_smi_path}"
+  printf 'nvidia_smi_sha256=%%s\n' "${nvidia_smi_sha256}"
+  printf 'nvidia_module_path=%%s\n' "${nvidia_module_path}"
+  printf 'nvidia_module_sha256=%%s\n' "${nvidia_module_sha256}"
+} | sudo tee "${receipt}" >/dev/null
+sudo chmod 0444 "${receipt}"
+
+sudo mkdir -p /usr/local/nvidia
+printf '%%s\n' '#!/bin/sh' '# customer-installed driver sentinel' 'exit 0' | sudo tee %s >/dev/null
+sudo chmod 0755 %s
+`,
+		shellSingleQuote(imageRef),
+		shellSingleQuote(customerGPUDriverImage),
+		shellSingleQuote(customerDriverReceiptPath),
+		shellSingleQuote(customerNvidiaCDIDropIn),
+		shellSingleQuote(customerDriverSentinelPath),
+		shellSingleQuote(customerDriverSentinelPath),
 	)
 	if _, err := RunCommand(ctx, s, command); err != nil {
-		return fmt.Errorf("install customer driver sentinel: %w", err)
+		return fmt.Errorf("install customer NVIDIA driver from cached approved artifact: %w", err)
 	}
 	return nil
 }
diff --git a/e2e/scenario/scenario.go b/e2e/scenario/scenario.go
index b6dd736685b..1c58b418e8a 100644
--- a/e2e/scenario/scenario.go
+++ b/e2e/scenario/scenario.go
@@ -3051,7 +3051,7 @@ var _ = Register(&Scenario{
 
 var _ = Register(&Scenario{
 	Name:             "Ubuntu2404_GPUDriverOptOut_VHDCaching",
-	Description:      "Tests that a customer-installed NVIDIA driver sentinel survives PIS VHD-cached provisioning on an opted-out GPU node",
+	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",
 	Location:         "westus2",
 	K8sSystemPoolSKU: "Standard_D2s_v3",
 	Tags: Tags{
@@ -3061,7 +3061,7 @@ var _ = Register(&Scenario{
 		Cluster:                ClusterKubenet,
 		VHD:                    config.VHDUbuntu2404Gen2Containerd,
 		VHDCaching:             true,
-		VHDCachingPostBasePrep: installCustomerDriverSentinel,
+		VHDCachingPostBasePrep: installCustomerNvidiaDriver,
 		BootstrapConfigMutator: func(_ *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) {
 			nbc.EnableScriptlessCSECmd = false
 			nbc.AgentPoolProfile.VMSize = "Standard_NC4as_T4_v3"
@@ -3073,7 +3073,10 @@ var _ = Register(&Scenario{
 			vmss.SKU.Name = to.Ptr("Standard_NC4as_T4_v3")
 		},
 		Validator: func(ctx context.Context, s *Scenario) error {
-			return ValidateFileHasContent(ctx, s, customerDriverSentinelPath, "customer-installed driver sentinel")
+			return errors.Join(
+				ValidateCustomerNvidiaDriver(ctx, s),
+				ValidateFileHasContent(ctx, s, customerDriverSentinelPath, "customer-installed driver sentinel"),
+			)
 		},
 	},
 })
diff --git a/e2e/scenario/validators.go b/e2e/scenario/validators.go
index c2290641c05..ae22b4911e8 100644
--- a/e2e/scenario/validators.go
+++ b/e2e/scenario/validators.go
@@ -580,6 +580,52 @@ func ValidateNvidiaSMIInstalled(ctx context.Context, s *Scenario) error {
 	return err
 }
 
+func ValidateCustomerNvidiaDriver(ctx context.Context, s *Scenario) error {
+	if _, _, err := customerGPUDriverImageReference(); err != nil {
+		return err
+	}
+	command := fmt.Sprintf(`set -euo pipefail
+receipt=%s
+expected_image_name=%s
+test -s "${receipt}"
+receipt_value() { sed -n "s/^$1=//p" "${receipt}"; }
+
+receipt_image_ref="$(receipt_value image_ref)"
+case "${receipt_image_ref}" in "${expected_image_name}:"*) ;; *) echo "unexpected NVIDIA driver artifact: ${receipt_image_ref}" >&2; exit 1 ;; esac
+image_tag="${receipt_image_ref##*:}"
+if ! [[ "${image_tag}" =~ ^[0-9]+(\.[0-9]+)+-[0-9]+$ ]]; then
+  echo "invalid cached NVIDIA driver image tag: ${image_tag}" >&2
+  exit 1
+fi
+expected_version="${image_tag%%-*}"
+receipt_driver_version="$(receipt_value driver_version)"
+test "${receipt_driver_version}" = "${expected_version}"
+receipt_gpu_name="$(receipt_value gpu_name)"
+case "${receipt_gpu_name}" in *T4*) ;; *) echo "receipt does not identify a T4 GPU: ${receipt_gpu_name}" >&2; exit 1 ;; esac
+
+gpu_info="$(sudo nvidia-smi --query-gpu=driver_version,name --format=csv,noheader | head -n1)"
+driver_version="$(printf '%%s\n' "${gpu_info}" | awk -F, '{ gsub(/^[[:space:]]+|[[:space:]]+$/, "", $1); print $1 }')"
+gpu_name="$(printf '%%s\n' "${gpu_info}" | awk -F, '{ gsub(/^[[:space:]]+|[[:space:]]+$/, "", $2); print $2 }')"
+test "${driver_version}" = "${receipt_driver_version}"
+test "${gpu_name}" = "${receipt_gpu_name}"
+sudo nvidia-smi -L | grep -qi 'T4'
+lsmod | grep -q '^nvidia[[:space:]]'
+
+nvidia_smi_path="$(readlink -f "$(command -v nvidia-smi)")"
+nvidia_module_path="$(sudo modinfo -n nvidia)"
+test "${nvidia_smi_path}" = "$(receipt_value nvidia_smi_path)"
+test "${nvidia_module_path}" = "$(receipt_value nvidia_module_path)"
+test "$(sudo modinfo -F version nvidia)" = "${receipt_driver_version}"
+test "$(sha256sum "${nvidia_smi_path}" | awk '{ print $1 }')" = "$(receipt_value nvidia_smi_sha256)"
+test "$(sha256sum "${nvidia_module_path}" | awk '{ print $1 }')" = "$(receipt_value nvidia_module_sha256)"
+`,
+		shellSingleQuote(customerDriverReceiptPath),
+		shellSingleQuote(customerGPUDriverImage),
+	)
+	_, err := execScriptOnVMForScenarioValidateExitCode(ctx, s, command, 0, "customer-installed NVIDIA driver did not survive PIS image capture")
+	return err
+}
+
 func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) error {
 	command := []string{
 		"set -ex",

From 9f9a275e46ff94f213434998a2179b2c0ce28798 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Mon, 5 Oct 2026 23:48:46 +1300
Subject: [PATCH 04/24] fix(gpu): remove GB bake-owned drivers safely

Use staged GB BOM versions to identify package ownership before PIS
customization. Support both GB200 and current layouts; retain failure
state when module teardown or initramfs rebuild cannot complete.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 .../artifacts/ubuntu/cse_install_ubuntu.sh    | 233 +++++++++++-
 .../artifacts/cse_install_ubuntu_spec.sh      | 334 ++++++++++++++++++
 2 files changed, 566 insertions(+), 1 deletion(-)

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 9886e25b002..21766ac5a0c 100755
--- a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
+++ b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
@@ -392,6 +392,185 @@ prebakedGPUDriverArtifactsRemain() {
     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}"
+}
+
+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
+}
+
+cleanUpGraceBlackwellGPUDriver() {
+    local entries="${1}" bom="${2}" wave="${3}" pending="${4:-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
+    [ "${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 openibd; 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
+    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
+}
+
 cleanUpGPUDrivers() {
     rm -Rf $GPU_DEST /opt/gpu
 
@@ -407,7 +586,59 @@ cleanUpGPUDrivers() {
 }
 
 cleanUpGPUDriversForBasePrep() {
-    local packageName marker="${GPU_DKMS_MARKER_FILE:-/opt/azure/aks-gpu/dkms-marker}"
+    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}" "${bom}" "${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
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 1665a3cddf8..f4c147a54c4 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
@@ -116,6 +116,80 @@ Describe 'cse_install_ubuntu.sh'
         GPU_DEST="/tmp/nonexistent-shellspec-gpu-dest"
         managedGPUPackageList() { :; }
 
+        createGraceBlackwellFixture() {
+            GB_MAI_BOM_FILE="${PWD}/.shellspec-gb-bom-$$.json"
+            GB200_MAI_BOM_FILE=""
+            GPU_DKMS_MARKER_FILE="${PWD}/.shellspec-gpu-marker-$$"
+            GB_NVIDIA_PEERMEM_CONFIG_FILE="${PWD}/.shellspec-nvidia-peermem-$$.conf"
+            GB_NVIDIA_MODPROBE_CONFIG_FILE="${PWD}/.shellspec-nvidia-modprobe-$$.conf"
+            GB_NOUVEAU_MODPROBE_CONFIG_FILE="${PWD}/.shellspec-nouveau-$$.conf"
+            GB_DRIVER_CLEANUP_PENDING_FILE="${PWD}/.shellspec-gb-cleanup-pending-$$"
+            printf '%s\n' '{"versions-wave2":{"libnvidia-common-580":"580.159.04-1ubuntu1","nvidia-dkms-580-open":"580.159.04-1ubuntu1","nvidia-driver-580-open":"580.159.04-1ubuntu1"}}' > "$GB_MAI_BOM_FILE"
+            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"
+        }
+
+        createLegacyGraceBlackwellFixture() {
+            createGraceBlackwellFixture
+            rm -f "$GB_MAI_BOM_FILE" "$GB_NVIDIA_PEERMEM_CONFIG_FILE" "$GB_NOUVEAU_MODPROBE_CONFIG_FILE"
+            GB200_MAI_BOM_FILE="${PWD}/.shellspec-gb200-bom-$$.json"
+            printf '%s\n' '{"versions-wave1":{"libnvidia-common-580":"580.105.08-0ubuntu1","nvidia-dkms-580-open":"580.105.08-0ubuntu1","nvidia-driver-580-open":"580.105.08-0ubuntu1"}}' > "$GB200_MAI_BOM_FILE"
+            GB_NVIDIA_MODPROBE_CONFIG_FILE="${PWD}/.shellspec-nvidia-modprobe-$$.conf"
+            printf '%s\n' 'options nvidia NVreg_RestrictProfilingToAdminUsers=0' \
+                'options nvidia NVreg_CreateImexChannel0=1' \
+                'options nvidia NVreg_CoherentGPUMemoryMode=driver' \
+                'options nvidia NVreg_RegistryDwords="RMBug5172204War=4"' > "$GB_NVIDIA_MODPROBE_CONFIG_FILE"
+        }
+
+        cleanupGraceBlackwellFixture() {
+            rm -f "${GB_MAI_BOM_FILE:-}" "${GB200_MAI_BOM_FILE:-}" "${GB_DRIVER_CLEANUP_PENDING_FILE:-}" "${GPU_DKMS_MARKER_FILE:-}" \
+                "${GB_NVIDIA_PEERMEM_CONFIG_FILE:-}" "${GB_NVIDIA_MODPROBE_CONFIG_FILE:-}" \
+                "${GB_NOUVEAU_MODPROBE_CONFIG_FILE:-}"
+        }
+
+        runGraceBlackwellFailureAndRetainState() {
+            cleanUpGPUDriversForBasePrep
+            status=$?
+            [ -f "$GPU_DKMS_MARKER_FILE" ] && echo "aks marker retained" || status=0
+            [ -f "$GB_NVIDIA_PEERMEM_CONFIG_FILE" ] && echo "GB boot config retained"
+            cleanupGraceBlackwellFixture
+            return "$status"
+        }
+
+        mockGraceBlackwellPackageQueries() {
+            dpkg-query() {
+                if [ "$#" -eq 2 ]; then
+                    [ "${gb_packages_purged:-false}" = true ] && return 0
+                    printf 'libnvidia-common-580\tinstalled\nnvidia-dkms-580-open\tinstalled\nnvidia-driver-580-open\tinstalled\nnvidia-utils-580\tinstalled\nnvidia-modprobe\tinstalled\nnvidia-persistenced\tinstalled\n'
+                    [ -n "${gb_extra_driver_package:-}" ] && printf '%s\tinstalled\n' "$gb_extra_driver_package"
+                    return 0
+                fi
+                case "$3" in
+                    libnvidia-common-580|nvidia-dkms-580-open|nvidia-driver-580-open)
+                        if [ "${gb_packages_purged:-false}" = true ]; then
+                            printf 'not-installed\t\n'
+                        elif [ "${gb_version_mismatch:-false}" = true ] && [ "$3" = nvidia-driver-580-open ]; then
+                            printf 'installed\t580.200.01-1ubuntu1\n'
+                        else
+                            printf 'installed\t%s\n' "${gb_driver_version:-580.159.04-1ubuntu1}"
+                        fi
+                        ;;
+                    nvidia-utils-580)
+                        [ "${gb_packages_purged:-false}" = true ] && printf 'not-installed\t\n' || printf 'installed\t%s\n' "${gb_utils_version:-580.159.04-1}"
+                        ;;
+                    nvidia-modprobe)
+                        [ "${gb_packages_purged:-false}" = true ] && printf 'not-installed\t\n' || printf 'installed\t%s\n' "${gb_modprobe_version:-1.10.0-1ubuntu1}"
+                        ;;
+                    nvidia-persistenced)
+                        [ "${gb_packages_purged:-false}" = true ] && printf 'not-installed\t\n' || printf 'installed\t%s\n' "${gb_persistenced_version:-1.10.0-1ubuntu1}"
+                        ;;
+                    *) printf 'not-installed\t\n' ;;
+                esac
+            }
+            apt-mark() { printf 'nvidia-utils-580\nnvidia-modprobe\nnvidia-persistenced\n'; }
+        }
+
         It 'fails when the prebake marker survives cleanup'
             GPU_DKMS_MARKER_FILE="$(mktemp)"
             cleanUpGPUDrivers() { :; }
@@ -146,6 +220,266 @@ Describe 'cse_install_ubuntu.sh'
             The stderr should include "GPU basePrep cleanup incomplete"
         End
 
+        It 'tears down current and legacy BOM-owned Grace Blackwell drivers before marker cleanup'
+            runGraceBlackwellTeardown() {
+                local layout="${1:-current}"
+                if [ "$layout" = legacy ]; then
+                    createLegacyGraceBlackwellFixture
+                    gb_driver_version=580.105.08-0ubuntu1
+                else
+                    createGraceBlackwellFixture
+                    gb_driver_version=580.159.04-1ubuntu1
+                    gb_utils_version=580.159.04-1
+                    gb_modprobe_version=1.10.0-1ubuntu1
+                    gb_persistenced_version=1.10.0-1ubuntu1
+                fi
+                if [ "$layout" = legacy ]; then
+                    gb_utils_version=580.105.08-1
+                    gb_modprobe_version=1.10.0-1ubuntu1
+                    gb_persistenced_version=1.10.0-1ubuntu1
+                fi
+                : > "$GPU_DKMS_MARKER_FILE"
+                gb_loaded_modules='nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia '
+                gb_expected_modules='nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia '
+                gb_services_stopped=false
+                gb_packages_purged=false
+                mockGraceBlackwellPackageQueries
+                systemctlDisableAndStop() {
+                    echo "stop:$1"
+                    if [ "$1" = openibd ]; then
+                        gb_services_stopped=true
+                    fi
+                    return 0
+                }
+                systemctl() { return 1; }
+                lsmod() {
+                    echo "Module Size Used by"
+                    for gb_module in $gb_loaded_modules; do
+                        echo "$gb_module 123 0"
+                    done
+                }
+                rmmod() {
+                    [ "$gb_services_stopped" = true ] || return 1
+                    [ "${gb_expected_modules%% *}" = "$1" ] || return 1
+                    echo "rmmod:$1"
+                    gb_expected_modules="${gb_expected_modules#* }"
+                    gb_loaded_modules="${gb_loaded_modules#* }"
+                }
+                update-initramfs() {
+                    [ "$gb_packages_purged" = true ] || return 1
+                    [ ! -e "$GB_NVIDIA_PEERMEM_CONFIG_FILE" ] || return 1
+                    [ ! -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || return 1
+                    [ ! -e "$GB_NOUVEAU_MODPROBE_CONFIG_FILE" ] || return 1
+                    [ -f "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
+                    echo "update-initramfs"
+                }
+                apt_get_purge() {
+                    [ "$1 $2 $3" = '10 5 300' ] || return 1
+                    shift 3
+                    local expected='libnvidia-common-580 nvidia-dkms-580-open nvidia-driver-580-open nvidia-utils-580 nvidia-modprobe nvidia-persistenced'
+                    [ "$*" = "$expected" ] || return 1
+                    [ -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || return 1
+                    [ "$gb_services_stopped" = true ] || return 1
+                    [ -z "$gb_loaded_modules" ] || return 1
+                    [ -z "$gb_expected_modules" ] || return 1
+                    echo "purge:$*"
+                    gb_packages_purged=true
+                }
+                cleanUpGPUDrivers() {
+                    [ "$gb_packages_purged" = true ] || return 1
+                    [ ! -e "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
+                    echo "aks-marker-cleanup"
+                    rm -f "$GPU_DKMS_MARKER_FILE"
+                }
+                prebakedGPUDriverArtifactsRemain() { [ "$gb_packages_purged" != true ]; }
+                cleanUpGPUDriversForBasePrep
+                status=$?
+                cleanupGraceBlackwellFixture
+                return "$status"
+            }
+
+            runBothGraceBlackwellTeardowns() {
+                runGraceBlackwellTeardown current || return 1
+                runGraceBlackwellTeardown legacy
+            }
+
+            When call runBothGraceBlackwellTeardowns
+            The status should be success
+            The output should include "rmmod:nvidia_peermem"
+            The output should include "stop:nvidia-imex"
+            The output should include "stop:openibd"
+            The output should include "purge:"
+            The output should include "aks-marker-cleanup"
+        End
+
+        It 'uses a BOM-listed dependency version instead of comparing it to the driver version'
+            runGraceBlackwellBOMDependencyVersion() {
+                createGraceBlackwellFixture
+                gb_driver_version=580.159.04-1ubuntu1
+                gb_modprobe_version=1.10.0-1ubuntu1
+                printf '%s\n' '{"versions-wave2":{"libnvidia-common-580":"580.159.04-1ubuntu1","nvidia-dkms-580-open":"580.159.04-1ubuntu1","nvidia-driver-580-open":"580.159.04-1ubuntu1","nvidia-modprobe":"1.10.0-1ubuntu1"}}' > "$GB_MAI_BOM_FILE"
+                mockGraceBlackwellPackageQueries
+                entries=$(graceBlackwellDriverBOMEntries "$GB_MAI_BOM_FILE" versions-wave2) || return 1
+                graceBlackwellInstalledDriverPackages "$entries"
+                status=$?
+                cleanupGraceBlackwellFixture
+                return "$status"
+            }
+            When call runGraceBlackwellBOMDependencyVersion
+            The status should be success
+            The output should include "nvidia-modprobe=1.10.0-1ubuntu1"
+        End
+
+        It 'retains GB provenance and retries after config removal or initramfs failure'
+            runGraceBlackwellTeardownRetry() {
+                local status=0
+                createGraceBlackwellFixture
+                : > "$GPU_DKMS_MARKER_FILE"
+                gb_driver_version=580.159.04-1ubuntu1
+                gb_utils_version=580.159.04-1
+                gb_modprobe_version=1.10.0-1ubuntu1
+                gb_persistenced_version=1.10.0-1ubuntu1
+                gb_packages_purged=false
+                gb_loaded_modules=""
+                mockGraceBlackwellPackageQueries
+                systemctlDisableAndStop() { :; }
+                systemctl() { return 1; }
+                lsmod() { echo "Module Size Used by"; }
+                apt_get_purge() { gb_packages_purged=true; }
+                gb_fail_config_remove=true
+                rm() {
+                    if [ "${gb_fail_config_remove:-false}" = true ] &&
+                        [ "$*" = "-f $GB_NVIDIA_PEERMEM_CONFIG_FILE $GB_NVIDIA_MODPROBE_CONFIG_FILE $GB_NOUVEAU_MODPROBE_CONFIG_FILE" ]; then
+                        gb_fail_config_remove=false
+                        return 1
+                    fi
+                    command rm "$@"
+                }
+                gb_initramfs_failures=1
+                update-initramfs() {
+                    [ ! -e "$GB_NVIDIA_PEERMEM_CONFIG_FILE" ] || return 1
+                    [ ! -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || return 1
+                    [ ! -e "$GB_NOUVEAU_MODPROBE_CONFIG_FILE" ] || return 1
+                    [ -f "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
+                    if [ "$gb_initramfs_failures" -gt 0 ]; then
+                        gb_initramfs_failures=$((gb_initramfs_failures - 1))
+                        return 1
+                    fi
+                }
+                cleanUpGPUDrivers() {
+                    [ "$gb_packages_purged" = true ] || return 1
+                    [ ! -e "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
+                    rm -f "$GPU_DKMS_MARKER_FILE"
+                }
+                prebakedGPUDriverArtifactsRemain() { [ "$gb_packages_purged" != true ]; }
+
+                cleanUpGPUDriversForBasePrep && status=1
+                [ -f "$GPU_DKMS_MARKER_FILE" ] || status=1
+                [ -f "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || status=1
+                [ -f "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || status=1
+                cleanUpGPUDriversForBasePrep && status=1
+                [ -f "$GPU_DKMS_MARKER_FILE" ] || status=1
+                [ -f "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || status=1
+                [ ! -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || status=1
+                cleanUpGPUDriversForBasePrep || status=1
+                [ ! -e "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || status=1
+                [ ! -e "$GPU_DKMS_MARKER_FILE" ] || status=1
+                cleanupGraceBlackwellFixture
+                return "$status"
+            }
+            When call runGraceBlackwellTeardownRetry
+            The status should be success
+            The stderr should include "GPU basePrep cleanup incomplete: Grace Blackwell driver teardown failed"
+        End
+
+        It 'retains the marker on a legacy GB VHD without a BOM'
+            runGraceBlackwellWithoutBOM() {
+                status=0
+                createGraceBlackwellFixture
+                rm -f "$GB_MAI_BOM_FILE"
+                : > "$GPU_DKMS_MARKER_FILE"
+                cleanUpGPUDrivers() { echo "unexpected marker cleanup"; }
+                runGraceBlackwellFailureAndRetainState || status=1
+                createGraceBlackwellFixture
+                rm -f "$GB_MAI_BOM_FILE" "$GB_NVIDIA_PEERMEM_CONFIG_FILE"
+                : > "$GPU_DKMS_MARKER_FILE"
+                dpkg-query() {
+                    [ "$#" -eq 2 ] && printf 'nvidia-dkms-580-open\tinstalled\n'
+                }
+                runGraceBlackwellFailureAndRetainState || status=1
+                return "$status"
+            }
+            When call runGraceBlackwellWithoutBOM
+            The status should be failure
+            The stderr should include "Grace Blackwell BOM is unavailable"
+            The stderr should include "has no verifiable GB BOM ownership"
+            The output should include "aks marker retained"
+            The output should include "GB boot config retained"
+            The output should not include "unexpected marker cleanup"
+        End
+
+        It 'retains the marker for mixed ownership and package-purge failures'
+            runUnsafeGraceBlackwellTeardown() {
+                status=0
+                createGraceBlackwellFixture
+                : > "$GPU_DKMS_MARKER_FILE"
+                gb_packages_purged=false
+                gb_extra_driver_package=nvidia-driver-535-server
+                mockGraceBlackwellPackageQueries
+                systemctlDisableAndStop() { echo "unexpected service stop"; }
+                apt_get_purge() { echo "unexpected package purge"; }
+                cleanUpGPUDrivers() { echo "unexpected marker cleanup"; }
+                runGraceBlackwellFailureAndRetainState || status=1
+                createGraceBlackwellFixture
+                : > "$GPU_DKMS_MARKER_FILE"
+                gb_packages_purged=false
+                gb_extra_driver_package=""
+                mockGraceBlackwellPackageQueries
+                printf '\n# customer change\n' >> "$GB_NVIDIA_MODPROBE_CONFIG_FILE"
+                systemctlDisableAndStop() { echo "unexpected service stop"; }
+                systemctl() { return 1; }
+                lsmod() { echo "nvidia 123 0"; }
+                rmmod() { echo "unexpected module unload"; }
+                apt_get_purge() { echo "unexpected package purge"; }
+                cleanUpGPUDrivers() { echo "unexpected marker cleanup"; }
+                runGraceBlackwellFailureAndRetainState || status=1
+                createGraceBlackwellFixture
+                : > "$GPU_DKMS_MARKER_FILE"
+                gb_extra_driver_package=""
+                gb_packages_purged=false
+                mockGraceBlackwellPackageQueries
+                lsmod() { echo "Module Size Used by"; }
+                systemctlDisableAndStop() { :; }
+                systemctl() { return 1; }
+                systemctlDisableAndStop() { :; }
+                systemctl() { return 1; }
+                update-initramfs() { :; }
+                apt_get_purge() { return 1; }
+                runGraceBlackwellFailureAndRetainState || status=1
+                createGraceBlackwellFixture
+                gb_version_mismatch=true
+                mockGraceBlackwellPackageQueries
+                : > "$GPU_DKMS_MARKER_FILE"
+                apt_get_purge() { echo "unexpected package purge"; }
+                cleanUpGPUDrivers() { echo "unexpected marker cleanup"; }
+                cleanUpGPUDriversForBasePrep || status=1
+                cleanupGraceBlackwellFixture
+                return "$status"
+            }
+            When call runUnsafeGraceBlackwellTeardown
+            The status should be failure
+            The stderr should include "is not owned by the Grace Blackwell BOM"
+            The stderr should include "module configuration"
+            The stderr should include "Grace Blackwell driver package purge failed"
+            The stderr should include "does not match the GB BOM"
+            The output should include "aks marker retained"
+            The output should include "GB boot config retained"
+            The output should not include "unexpected service stop"
+            The output should not include "unexpected module unload"
+            The output should not include "unexpected package purge"
+            The output should not include "unexpected marker cleanup"
+        End
+
         It 'succeeds when the prebake marker and AKS paths are gone'
             GPU_DKMS_MARKER_FILE="$(mktemp)"
             cleanUpGPUDrivers() { rm -f "$GPU_DKMS_MARKER_FILE"; }

From a58444b7bc25045b47f2eeac02a1b261731dd3ee Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Tue, 6 Oct 2026 00:14:13 +1300
Subject: [PATCH 05/24] test(gpu): install only customer T4 driver

Avoid Fabric Manager, which requires NVSwitch hardware. Log driver
receipts on bake and derived nodes to retain preservation evidence.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 e2e/scenario/provision.go  | 50 +++++++++++++++++++++++++++++---------
 e2e/scenario/validators.go | 18 ++++++++++++--
 2 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/e2e/scenario/provision.go b/e2e/scenario/provision.go
index fb532bcc3b7..544cce8e07c 100644
--- a/e2e/scenario/provision.go
+++ b/e2e/scenario/provision.go
@@ -762,10 +762,11 @@ func RunCommand(ctx context.Context, s *Scenario, command string) (armcompute.Vi
 }
 
 const (
-	customerDriverSentinelPath = "/usr/local/nvidia/customer-driver-sentinel"
-	customerDriverReceiptPath  = "/opt/azure/containers/customer-gpu-driver.receipt"
-	customerNvidiaCDIDropIn    = "/etc/systemd/system/nvidia-cdi-refresh.service.d/10-aks-tolerate-generate-failure.conf"
-	customerGPUDriverImage     = "mcr.microsoft.com/aks/aks-gpu-cuda-lts"
+	customerDriverSentinelPath     = "/usr/local/nvidia/customer-driver-sentinel"
+	customerDriverReceiptPath      = "/opt/azure/containers/customer-gpu-driver.receipt"
+	customerGPUDriverImage         = "mcr.microsoft.com/aks/aks-gpu-cuda-lts"
+	customerGPUDriverEvidenceStart = "CUSTOM_NVIDIA_DRIVER_EVIDENCE_BEGIN"
+	customerGPUDriverEvidenceEnd   = "CUSTOM_NVIDIA_DRIVER_EVIDENCE_END"
 )
 
 func customerGPUDriverImageReference() (string, string, error) {
@@ -807,17 +808,18 @@ else
 fi
 image_tag="${image_ref##*:}"
 expected_version="${image_tag%%-*}"
-# Match configureNvidiaCDIRefresh in CSE before the NVIDIA toolkit package starts this unit.
-sudo install -d -m 0755 /etc/systemd/system/nvidia-cdi-refresh.service.d
-printf '[Service]\nSuccessExitStatus=1 2 127\n' | sudo tee %s >/dev/null
-sudo systemctl daemon-reload
 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 install
+  "${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)" --log-file-name=/var/log/nvidia-installer-customer-e2e.log -a --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="$(printf '%%s\n' "${gpu_info}" | awk -F, '{ gsub(/^[[:space:]]+|[[:space:]]+$/, "", $1); print $1 }')"
@@ -846,6 +848,12 @@ sudo install -d -m 0755 "$(dirname "${receipt}")"
 } | sudo tee "${receipt}" >/dev/null
 sudo chmod 0444 "${receipt}"
 
+printf '%%s\n' %s
+sudo cat "${receipt}"
+printf 'nvidia_smi_gpu_info=%%s\n' "${gpu_info}"
+sudo nvidia-smi -L
+printf '%%s\n' %s
+
 sudo mkdir -p /usr/local/nvidia
 printf '%%s\n' '#!/bin/sh' '# customer-installed driver sentinel' 'exit 0' | sudo tee %s >/dev/null
 sudo chmod 0755 %s
@@ -853,16 +861,36 @@ sudo chmod 0755 %s
 		shellSingleQuote(imageRef),
 		shellSingleQuote(customerGPUDriverImage),
 		shellSingleQuote(customerDriverReceiptPath),
-		shellSingleQuote(customerNvidiaCDIDropIn),
+		shellSingleQuote(customerGPUDriverEvidenceStart),
+		shellSingleQuote(customerGPUDriverEvidenceEnd),
 		shellSingleQuote(customerDriverSentinelPath),
 		shellSingleQuote(customerDriverSentinelPath),
 	)
-	if _, err := RunCommand(ctx, s, command); err != nil {
+	view, err := RunCommand(ctx, s, command)
+	if err != nil {
 		return fmt.Errorf("install customer NVIDIA driver from cached approved artifact: %w", err)
 	}
+	if view.Output != nil {
+		logCustomerNvidiaDriverEvidence(ctx, *view.Output)
+	}
 	return nil
 }
 
+func logCustomerNvidiaDriverEvidence(ctx context.Context, output string) {
+	start := strings.LastIndex(output, customerGPUDriverEvidenceStart)
+	if start < 0 {
+		return
+	}
+	start += len(customerGPUDriverEvidenceStart)
+	end := strings.Index(output[start:], customerGPUDriverEvidenceEnd)
+	if end < 0 {
+		return
+	}
+	if evidence := strings.TrimSpace(output[start : start+end]); evidence != "" {
+		logging.Logf(ctx, "customer-installed NVIDIA driver evidence:\n%s", evidence)
+	}
+}
+
 // runCommandScriptError converts a RunCommand instance view into an error if the
 // script itself failed. The ARM CreateOrUpdate operation reports success as long as
 // the extension was able to run the script — a non-zero exit, throw, or timeout
diff --git a/e2e/scenario/validators.go b/e2e/scenario/validators.go
index ae22b4911e8..0e0ab2af2d5 100644
--- a/e2e/scenario/validators.go
+++ b/e2e/scenario/validators.go
@@ -618,12 +618,26 @@ test "${nvidia_module_path}" = "$(receipt_value nvidia_module_path)"
 test "$(sudo modinfo -F version nvidia)" = "${receipt_driver_version}"
 test "$(sha256sum "${nvidia_smi_path}" | awk '{ print $1 }')" = "$(receipt_value nvidia_smi_sha256)"
 test "$(sha256sum "${nvidia_module_path}" | awk '{ print $1 }')" = "$(receipt_value nvidia_module_sha256)"
+printf '%%s\n' %s
+cat "${receipt}"
+printf 'nvidia_smi_gpu_info=%%s\n' "${gpu_info}"
+printf 'nvidia_module_version=%%s\n' "$(sudo modinfo -F version nvidia)"
+printf 'nvidia_smi_current_sha256=%%s\n' "$(sha256sum "${nvidia_smi_path}" | awk '{ print $1 }')"
+printf 'nvidia_module_current_sha256=%%s\n' "$(sha256sum "${nvidia_module_path}" | awk '{ print $1 }')"
+sudo nvidia-smi -L
+printf '%%s\n' %s
 `,
 		shellSingleQuote(customerDriverReceiptPath),
 		shellSingleQuote(customerGPUDriverImage),
+		shellSingleQuote(customerGPUDriverEvidenceStart),
+		shellSingleQuote(customerGPUDriverEvidenceEnd),
 	)
-	_, err := execScriptOnVMForScenarioValidateExitCode(ctx, s, command, 0, "customer-installed NVIDIA driver did not survive PIS image capture")
-	return err
+	result, err := execScriptOnVMForScenarioValidateExitCode(ctx, s, command, 0, "customer-installed NVIDIA driver did not survive PIS image capture")
+	if err != nil {
+		return err
+	}
+	logCustomerNvidiaDriverEvidence(ctx, result.stdout)
+	return nil
 }
 
 func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) error {

From 5497dfdf0176b2efec49558a80d39ba972542402 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Tue, 6 Oct 2026 02:02:03 +1300
Subject: [PATCH 06/24] test(gpu): use default PIS capture gallery

Use the CI-supported default region instead of the westus2 gallery,
where the pipeline identity cannot capture the PIS image.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 e2e/scenario/scenario.go | 1 -
 1 file changed, 1 deletion(-)

diff --git a/e2e/scenario/scenario.go b/e2e/scenario/scenario.go
index 1c58b418e8a..38312c33190 100644
--- a/e2e/scenario/scenario.go
+++ b/e2e/scenario/scenario.go
@@ -3052,7 +3052,6 @@ 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",
-	Location:         "westus2",
 	K8sSystemPoolSKU: "Standard_D2s_v3",
 	Tags: Tags{
 		GPU: true,

From 8fcd0d9f2a3ba2b1baa606b6e4e84d25d7c1e18d Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Thu, 8 Oct 2026 12:14:06 +1300
Subject: [PATCH 07/24] test(gpu): slim PIS driver regression

Reduce the PIS E2E proof to module/version and nvidia-smi checks, and keep only lifecycle-critical ShellSpec cases. Share Grace Blackwell fixture setup without weakening teardown coverage.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 e2e/scenario/cleanup.go                       |  16 +-
 e2e/scenario/execution_test.go                |   6 +-
 e2e/scenario/provision.go                     | 101 ++--------
 e2e/scenario/scenario.go                      |   5 +-
 e2e/scenario/validators.go                    |  62 ++----
 .../artifacts/cse_install_mariner_spec.sh     |   4 +-
 .../artifacts/cse_install_ubuntu_spec.sh      | 112 +++++------
 .../artifacts/cse_main_gpu_pis_spec.sh        | 186 ++++++------------
 8 files changed, 135 insertions(+), 357 deletions(-)

diff --git a/e2e/scenario/cleanup.go b/e2e/scenario/cleanup.go
index 3b2881ee105..a3fa1e3d18e 100644
--- a/e2e/scenario/cleanup.go
+++ b/e2e/scenario/cleanup.go
@@ -4,13 +4,9 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"reflect"
-	"runtime"
 	"runtime/debug"
 	"sync"
 	"time"
-
-	"github.com/Azure/agentbaker/e2e/logging"
 )
 
 const (
@@ -54,18 +50,8 @@ func (c *scenarioCleanup) runCleanups(ctx context.Context) error {
 		batchErrs := make([]error, len(cleanups))
 		var wg sync.WaitGroup
 		for i, fn := range cleanups {
-			i, fn := i, fn
-			name := fmt.Sprintf("cleanup[%d]", i+1)
-			if callback := runtime.FuncForPC(reflect.ValueOf(fn).Pointer()); callback != nil {
-				name = callback.Name()
-			}
-			deadline, _ := ctx.Deadline()
-			started := time.Now()
-			logging.Logf(ctx, "scenario cleanup %s started (deadline %s)", name, deadline.Format(time.RFC3339))
 			wg.Go(func() {
-				err := runWithPanicRecovery(ctx, fn)
-				batchErrs[i] = err
-				logging.Logf(ctx, "scenario cleanup %s finished in %s (error: %v, context error: %v)", name, time.Since(started), err, ctx.Err())
+				batchErrs[i] = runWithPanicRecovery(ctx, fn)
 			})
 		}
 		wg.Wait()
diff --git a/e2e/scenario/execution_test.go b/e2e/scenario/execution_test.go
index bee67c162d9..bae464aa6eb 100644
--- a/e2e/scenario/execution_test.go
+++ b/e2e/scenario/execution_test.go
@@ -57,11 +57,7 @@ func TestExecutionKeepsLoggerThroughSkipRunAndCleanup(t *testing.T) {
 			} else {
 				require.NoError(t, outcome.Error)
 			}
-			require.Len(t, logger.logs, 5)
-			assert.Equal(t, []string{"skip check", "running scenario"}, logger.logs[:2])
-			assert.Contains(t, logger.logs[2], "scenario cleanup ")
-			assert.Equal(t, "cleanup", logger.logs[3])
-			assert.Contains(t, logger.logs[4], "scenario cleanup ")
+			assert.Equal(t, []string{"skip check", "running scenario", "cleanup"}, logger.logs)
 		})
 	}
 }
diff --git a/e2e/scenario/provision.go b/e2e/scenario/provision.go
index 544cce8e07c..51b5dc7af0f 100644
--- a/e2e/scenario/provision.go
+++ b/e2e/scenario/provision.go
@@ -762,11 +762,7 @@ func RunCommand(ctx context.Context, s *Scenario, command string) (armcompute.Vi
 }
 
 const (
-	customerDriverSentinelPath     = "/usr/local/nvidia/customer-driver-sentinel"
-	customerDriverReceiptPath      = "/opt/azure/containers/customer-gpu-driver.receipt"
-	customerGPUDriverImage         = "mcr.microsoft.com/aks/aks-gpu-cuda-lts"
-	customerGPUDriverEvidenceStart = "CUSTOM_NVIDIA_DRIVER_EVIDENCE_BEGIN"
-	customerGPUDriverEvidenceEnd   = "CUSTOM_NVIDIA_DRIVER_EVIDENCE_END"
+	customerGPUDriverImage = "mcr.microsoft.com/aks/aks-gpu-cuda-lts"
 )
 
 func customerGPUDriverImageReference() (string, string, error) {
@@ -782,32 +778,15 @@ func customerGPUDriverImageReference() (string, string, error) {
 // 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, _, err := customerGPUDriverImageReference()
+	imageRef, expectedVersion, err := customerGPUDriverImageReference()
 	if err != nil {
 		return err
 	}
 	command := fmt.Sprintf(`set -euo pipefail
-preferred_image_ref=%s
-image_name=%s
-receipt=%s
-
+image_ref=%s
+expected_version=%s
 sudo mkdir -p /opt/gpu /opt/actions
-cached_image_refs="$(sudo ctr -n k8s.io images ls -q | awk -v prefix="${image_name}:" 'index($0, prefix) == 1 {
-  tag = substr($0, length(prefix) + 1)
-  if (tag ~ /^[0-9]+(\.[0-9]+)+-[0-9]+$/) print $0
-}' | sort -V)"
-if [ -z "${cached_image_refs}" ]; then
-  echo "no versioned cached approved NVIDIA driver image found for ${image_name}" >&2
-  exit 1
-fi
-if printf '%%s\n' "${cached_image_refs}" | grep -Fxq "${preferred_image_ref}"; then
-  image_ref="${preferred_image_ref}"
-else
-  image_ref="$(printf '%%s\n' "${cached_image_refs}" | tail -n1)"
-  echo "Using cached approved NVIDIA driver image ${image_ref}; current preferred image ${preferred_image_ref} is not present"
-fi
-image_tag="${image_ref##*:}"
-expected_version="${image_tag%%-*}"
+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 \
@@ -818,79 +797,29 @@ 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)" --log-file-name=/var/log/nvidia-installer-customer-e2e.log -a --no-drm --dkms
+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="$(printf '%%s\n' "${gpu_info}" | awk -F, '{ gsub(/^[[:space:]]+|[[:space:]]+$/, "", $1); print $1 }')"
-gpu_name="$(printf '%%s\n' "${gpu_info}" | awk -F, '{ gsub(/^[[:space:]]+|[[:space:]]+$/, "", $2); print $2 }')"
-test "${driver_version}" = "${expected_version}"
+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}"
-
-nvidia_smi_path="$(readlink -f "$(command -v nvidia-smi)")"
-nvidia_module_path="$(sudo modinfo -n nvidia)"
-test -x "${nvidia_smi_path}"
-test -f "${nvidia_module_path}"
-nvidia_smi_sha256="$(sha256sum "${nvidia_smi_path}" | awk '{ print $1 }')"
-nvidia_module_sha256="$(sha256sum "${nvidia_module_path}" | awk '{ print $1 }')"
-sudo install -d -m 0755 "$(dirname "${receipt}")"
-{
-  printf 'image_ref=%%s\n' "${image_ref}"
-  printf 'driver_version=%%s\n' "${driver_version}"
-  printf 'gpu_name=%%s\n' "${gpu_name}"
-  printf 'nvidia_smi_path=%%s\n' "${nvidia_smi_path}"
-  printf 'nvidia_smi_sha256=%%s\n' "${nvidia_smi_sha256}"
-  printf 'nvidia_module_path=%%s\n' "${nvidia_module_path}"
-  printf 'nvidia_module_sha256=%%s\n' "${nvidia_module_sha256}"
-} | sudo tee "${receipt}" >/dev/null
-sudo chmod 0444 "${receipt}"
-
-printf '%%s\n' %s
-sudo cat "${receipt}"
-printf 'nvidia_smi_gpu_info=%%s\n' "${gpu_info}"
-sudo nvidia-smi -L
-printf '%%s\n' %s
-
-sudo mkdir -p /usr/local/nvidia
-printf '%%s\n' '#!/bin/sh' '# customer-installed driver sentinel' 'exit 0' | sudo tee %s >/dev/null
-sudo chmod 0755 %s
 `,
 		shellSingleQuote(imageRef),
-		shellSingleQuote(customerGPUDriverImage),
-		shellSingleQuote(customerDriverReceiptPath),
-		shellSingleQuote(customerGPUDriverEvidenceStart),
-		shellSingleQuote(customerGPUDriverEvidenceEnd),
-		shellSingleQuote(customerDriverSentinelPath),
-		shellSingleQuote(customerDriverSentinelPath),
+		shellSingleQuote(expectedVersion),
 	)
-	view, err := RunCommand(ctx, s, command)
+	_, err = RunCommand(ctx, s, command)
 	if err != nil {
 		return fmt.Errorf("install customer NVIDIA driver from cached approved artifact: %w", err)
 	}
-	if view.Output != nil {
-		logCustomerNvidiaDriverEvidence(ctx, *view.Output)
-	}
 	return nil
 }
 
-func logCustomerNvidiaDriverEvidence(ctx context.Context, output string) {
-	start := strings.LastIndex(output, customerGPUDriverEvidenceStart)
-	if start < 0 {
-		return
-	}
-	start += len(customerGPUDriverEvidenceStart)
-	end := strings.Index(output[start:], customerGPUDriverEvidenceEnd)
-	if end < 0 {
-		return
-	}
-	if evidence := strings.TrimSpace(output[start : start+end]); evidence != "" {
-		logging.Logf(ctx, "customer-installed NVIDIA driver evidence:\n%s", evidence)
-	}
-}
-
 // runCommandScriptError converts a RunCommand instance view into an error if the
 // script itself failed. The ARM CreateOrUpdate operation reports success as long as
 // the extension was able to run the script — a non-zero exit, throw, or timeout
@@ -1050,13 +979,9 @@ func CreateImage(ctx context.Context, s *Scenario) (*config.Image, error) {
 }
 
 // CreateSIGImageVersionFromDisk creates a new SIG image version directly from a VM disk
-func CreateSIGImageVersionFromDisk(ctx context.Context, s *Scenario, version string, diskResourceID string) (result *config.Image, returnErr error) {
+func CreateSIGImageVersionFromDisk(ctx context.Context, s *Scenario, version string, diskResourceID string) (*config.Image, error) {
 	startTime := time.Now()
 	defer func() {
-		if returnErr != nil {
-			logging.Logf(ctx, "SIG image version %s from disk %s failed after %s: %v", version, diskResourceID, time.Since(startTime), returnErr)
-			return
-		}
 		logging.Logf(ctx, "Created SIG image version %s from disk %s in %s", version, diskResourceID, time.Since(startTime))
 	}()
 	if s.Runtime == nil || s.Runtime.VM == nil || s.Runtime.VM.VM == nil ||
diff --git a/e2e/scenario/scenario.go b/e2e/scenario/scenario.go
index 38312c33190..503d40bfeb6 100644
--- a/e2e/scenario/scenario.go
+++ b/e2e/scenario/scenario.go
@@ -3072,10 +3072,7 @@ var _ = Register(&Scenario{
 			vmss.SKU.Name = to.Ptr("Standard_NC4as_T4_v3")
 		},
 		Validator: func(ctx context.Context, s *Scenario) error {
-			return errors.Join(
-				ValidateCustomerNvidiaDriver(ctx, s),
-				ValidateFileHasContent(ctx, s, customerDriverSentinelPath, "customer-installed driver sentinel"),
-			)
+			return ValidateCustomerNvidiaDriver(ctx, s)
 		},
 	},
 })
diff --git a/e2e/scenario/validators.go b/e2e/scenario/validators.go
index 0e0ab2af2d5..8cdff2a2277 100644
--- a/e2e/scenario/validators.go
+++ b/e2e/scenario/validators.go
@@ -581,63 +581,23 @@ func ValidateNvidiaSMIInstalled(ctx context.Context, s *Scenario) error {
 }
 
 func ValidateCustomerNvidiaDriver(ctx context.Context, s *Scenario) error {
-	if _, _, err := customerGPUDriverImageReference(); err != nil {
+	_, expectedVersion, err := customerGPUDriverImageReference()
+	if err != nil {
 		return err
 	}
 	command := fmt.Sprintf(`set -euo pipefail
-receipt=%s
-expected_image_name=%s
-test -s "${receipt}"
-receipt_value() { sed -n "s/^$1=//p" "${receipt}"; }
-
-receipt_image_ref="$(receipt_value image_ref)"
-case "${receipt_image_ref}" in "${expected_image_name}:"*) ;; *) echo "unexpected NVIDIA driver artifact: ${receipt_image_ref}" >&2; exit 1 ;; esac
-image_tag="${receipt_image_ref##*:}"
-if ! [[ "${image_tag}" =~ ^[0-9]+(\.[0-9]+)+-[0-9]+$ ]]; then
-  echo "invalid cached NVIDIA driver image tag: ${image_tag}" >&2
-  exit 1
-fi
-expected_version="${image_tag%%-*}"
-receipt_driver_version="$(receipt_value driver_version)"
-test "${receipt_driver_version}" = "${expected_version}"
-receipt_gpu_name="$(receipt_value gpu_name)"
-case "${receipt_gpu_name}" in *T4*) ;; *) echo "receipt does not identify a T4 GPU: ${receipt_gpu_name}" >&2; exit 1 ;; esac
-
-gpu_info="$(sudo nvidia-smi --query-gpu=driver_version,name --format=csv,noheader | head -n1)"
-driver_version="$(printf '%%s\n' "${gpu_info}" | awk -F, '{ gsub(/^[[:space:]]+|[[:space:]]+$/, "", $1); print $1 }')"
-gpu_name="$(printf '%%s\n' "${gpu_info}" | awk -F, '{ gsub(/^[[:space:]]+|[[:space:]]+$/, "", $2); print $2 }')"
-test "${driver_version}" = "${receipt_driver_version}"
-test "${gpu_name}" = "${receipt_gpu_name}"
-sudo nvidia-smi -L | grep -qi 'T4'
+expected_version=%s
 lsmod | grep -q '^nvidia[[:space:]]'
-
-nvidia_smi_path="$(readlink -f "$(command -v nvidia-smi)")"
-nvidia_module_path="$(sudo modinfo -n nvidia)"
-test "${nvidia_smi_path}" = "$(receipt_value nvidia_smi_path)"
-test "${nvidia_module_path}" = "$(receipt_value nvidia_module_path)"
-test "$(sudo modinfo -F version nvidia)" = "${receipt_driver_version}"
-test "$(sha256sum "${nvidia_smi_path}" | awk '{ print $1 }')" = "$(receipt_value nvidia_smi_sha256)"
-test "$(sha256sum "${nvidia_module_path}" | awk '{ print $1 }')" = "$(receipt_value nvidia_module_sha256)"
-printf '%%s\n' %s
-cat "${receipt}"
-printf 'nvidia_smi_gpu_info=%%s\n' "${gpu_info}"
-printf 'nvidia_module_version=%%s\n' "$(sudo modinfo -F version nvidia)"
-printf 'nvidia_smi_current_sha256=%%s\n' "$(sha256sum "${nvidia_smi_path}" | awk '{ print $1 }')"
-printf 'nvidia_module_current_sha256=%%s\n' "$(sha256sum "${nvidia_module_path}" | awk '{ print $1 }')"
-sudo nvidia-smi -L
-printf '%%s\n' %s
+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(customerDriverReceiptPath),
-		shellSingleQuote(customerGPUDriverImage),
-		shellSingleQuote(customerGPUDriverEvidenceStart),
-		shellSingleQuote(customerGPUDriverEvidenceEnd),
+		shellSingleQuote(expectedVersion),
 	)
-	result, err := execScriptOnVMForScenarioValidateExitCode(ctx, s, command, 0, "customer-installed NVIDIA driver did not survive PIS image capture")
-	if err != nil {
-		return err
-	}
-	logCustomerNvidiaDriverEvidence(ctx, result.stdout)
-	return nil
+	_, err = execScriptOnVMForScenarioValidateExitCode(ctx, s, command, 0, "customer-installed NVIDIA driver did not survive PIS image capture")
+	return err
 }
 
 func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) error {
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 77c7c9f0171..c243d5e7d88 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
@@ -29,7 +29,7 @@ Describe 'cse_install_mariner.sh'
     Include "./parts/linux/cloud-init/artifacts/mariner/cse_install_mariner.sh"
 
     Describe 'cleanUpGPUDriversForBasePrep'
-        GPU_DEST="/tmp/nonexistent-shellspec-gpu-dest"
+        GPU_DEST="${PWD}/.shellspec-nonexistent-gpu-dest"
         managedGPUPackageList() { echo "driver-cache"; }
         getPackageCacheDir() { echo "$gpu_test_root/driver-cache"; }
         setup_gpu_cache() { gpu_test_root=$(mktemp -d); }
@@ -798,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_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_spec.sh
index f4c147a54c4..91a87be7d79 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
@@ -113,33 +113,30 @@ Describe 'cse_install_ubuntu.sh'
     End
 
     Describe 'cleanUpGPUDriversForBasePrep'
-        GPU_DEST="/tmp/nonexistent-shellspec-gpu-dest"
+        GPU_DEST="${PWD}/.shellspec-nonexistent-gpu-dest"
         managedGPUPackageList() { :; }
 
         createGraceBlackwellFixture() {
+            local wave="${1:-wave2}"
             GB_MAI_BOM_FILE="${PWD}/.shellspec-gb-bom-$$.json"
-            GB200_MAI_BOM_FILE=""
+            GB200_MAI_BOM_FILE="${PWD}/.shellspec-gb200-bom-$$.json"
             GPU_DKMS_MARKER_FILE="${PWD}/.shellspec-gpu-marker-$$"
             GB_NVIDIA_PEERMEM_CONFIG_FILE="${PWD}/.shellspec-nvidia-peermem-$$.conf"
             GB_NVIDIA_MODPROBE_CONFIG_FILE="${PWD}/.shellspec-nvidia-modprobe-$$.conf"
             GB_NOUVEAU_MODPROBE_CONFIG_FILE="${PWD}/.shellspec-nouveau-$$.conf"
             GB_DRIVER_CLEANUP_PENDING_FILE="${PWD}/.shellspec-gb-cleanup-pending-$$"
-            printf '%s\n' '{"versions-wave2":{"libnvidia-common-580":"580.159.04-1ubuntu1","nvidia-dkms-580-open":"580.159.04-1ubuntu1","nvidia-driver-580-open":"580.159.04-1ubuntu1"}}' > "$GB_MAI_BOM_FILE"
-            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"
-        }
-
-        createLegacyGraceBlackwellFixture() {
-            createGraceBlackwellFixture
-            rm -f "$GB_MAI_BOM_FILE" "$GB_NVIDIA_PEERMEM_CONFIG_FILE" "$GB_NOUVEAU_MODPROBE_CONFIG_FILE"
-            GB200_MAI_BOM_FILE="${PWD}/.shellspec-gb200-bom-$$.json"
-            printf '%s\n' '{"versions-wave1":{"libnvidia-common-580":"580.105.08-0ubuntu1","nvidia-dkms-580-open":"580.105.08-0ubuntu1","nvidia-driver-580-open":"580.105.08-0ubuntu1"}}' > "$GB200_MAI_BOM_FILE"
-            GB_NVIDIA_MODPROBE_CONFIG_FILE="${PWD}/.shellspec-nvidia-modprobe-$$.conf"
-            printf '%s\n' 'options nvidia NVreg_RestrictProfilingToAdminUsers=0' \
-                'options nvidia NVreg_CreateImexChannel0=1' \
-                'options nvidia NVreg_CoherentGPUMemoryMode=driver' \
-                'options nvidia NVreg_RegistryDwords="RMBug5172204War=4"' > "$GB_NVIDIA_MODPROBE_CONFIG_FILE"
+            if [ "$wave" = wave1 ]; then
+                printf '%s\n' '{"versions-wave1":{"libnvidia-common-580":"580.105.08-0ubuntu1","nvidia-dkms-580-open":"580.105.08-0ubuntu1","nvidia-driver-580-open":"580.105.08-0ubuntu1"}}' > "$GB200_MAI_BOM_FILE"
+                printf '%s\n' 'options nvidia NVreg_RestrictProfilingToAdminUsers=0' \
+                    'options nvidia NVreg_CreateImexChannel0=1' \
+                    'options nvidia NVreg_CoherentGPUMemoryMode=driver' \
+                    'options nvidia NVreg_RegistryDwords="RMBug5172204War=4"' > "$GB_NVIDIA_MODPROBE_CONFIG_FILE"
+            else
+                printf '%s\n' '{"versions-wave2":{"libnvidia-common-580":"580.159.04-1ubuntu1","nvidia-dkms-580-open":"580.159.04-1ubuntu1","nvidia-driver-580-open":"580.159.04-1ubuntu1"}}' > "$GB_MAI_BOM_FILE"
+                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"
+            fi
         }
 
         cleanupGraceBlackwellFixture() {
@@ -190,6 +187,25 @@ Describe 'cse_install_ubuntu.sh'
             apt-mark() { printf 'nvidia-utils-580\nnvidia-modprobe\nnvidia-persistenced\n'; }
         }
 
+        setupGraceBlackwellTeardown() {
+            local wave="${1:-wave2}"
+            createGraceBlackwellFixture "$wave"
+            : > "$GPU_DKMS_MARKER_FILE"
+            gb_driver_version=580.159.04-1ubuntu1
+            gb_utils_version=580.159.04-1
+            gb_modprobe_version=1.10.0-1ubuntu1
+            gb_persistenced_version=1.10.0-1ubuntu1
+            if [ "$wave" = wave1 ]; then
+                gb_driver_version=580.105.08-0ubuntu1
+                gb_utils_version=580.105.08-1
+            fi
+            gb_loaded_modules='nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia '
+            gb_expected_modules='nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia '
+            gb_services_stopped=false
+            gb_packages_purged=false
+            mockGraceBlackwellPackageQueries
+        }
+
         It 'fails when the prebake marker survives cleanup'
             GPU_DKMS_MARKER_FILE="$(mktemp)"
             cleanUpGPUDrivers() { :; }
@@ -222,28 +238,8 @@ Describe 'cse_install_ubuntu.sh'
 
         It 'tears down current and legacy BOM-owned Grace Blackwell drivers before marker cleanup'
             runGraceBlackwellTeardown() {
-                local layout="${1:-current}"
-                if [ "$layout" = legacy ]; then
-                    createLegacyGraceBlackwellFixture
-                    gb_driver_version=580.105.08-0ubuntu1
-                else
-                    createGraceBlackwellFixture
-                    gb_driver_version=580.159.04-1ubuntu1
-                    gb_utils_version=580.159.04-1
-                    gb_modprobe_version=1.10.0-1ubuntu1
-                    gb_persistenced_version=1.10.0-1ubuntu1
-                fi
-                if [ "$layout" = legacy ]; then
-                    gb_utils_version=580.105.08-1
-                    gb_modprobe_version=1.10.0-1ubuntu1
-                    gb_persistenced_version=1.10.0-1ubuntu1
-                fi
-                : > "$GPU_DKMS_MARKER_FILE"
-                gb_loaded_modules='nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia '
-                gb_expected_modules='nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia '
-                gb_services_stopped=false
-                gb_packages_purged=false
-                mockGraceBlackwellPackageQueries
+                local wave="${1:-wave2}"
+                setupGraceBlackwellTeardown "$wave"
                 systemctlDisableAndStop() {
                     echo "stop:$1"
                     if [ "$1" = openibd ]; then
@@ -299,8 +295,8 @@ Describe 'cse_install_ubuntu.sh'
             }
 
             runBothGraceBlackwellTeardowns() {
-                runGraceBlackwellTeardown current || return 1
-                runGraceBlackwellTeardown legacy
+                runGraceBlackwellTeardown wave2 || return 1
+                runGraceBlackwellTeardown wave1
             }
 
             When call runBothGraceBlackwellTeardowns
@@ -333,13 +329,7 @@ Describe 'cse_install_ubuntu.sh'
         It 'retains GB provenance and retries after config removal or initramfs failure'
             runGraceBlackwellTeardownRetry() {
                 local status=0
-                createGraceBlackwellFixture
-                : > "$GPU_DKMS_MARKER_FILE"
-                gb_driver_version=580.159.04-1ubuntu1
-                gb_utils_version=580.159.04-1
-                gb_modprobe_version=1.10.0-1ubuntu1
-                gb_persistenced_version=1.10.0-1ubuntu1
-                gb_packages_purged=false
+                setupGraceBlackwellTeardown
                 gb_loaded_modules=""
                 mockGraceBlackwellPackageQueries
                 systemctlDisableAndStop() { :; }
@@ -421,17 +411,14 @@ Describe 'cse_install_ubuntu.sh'
         It 'retains the marker for mixed ownership and package-purge failures'
             runUnsafeGraceBlackwellTeardown() {
                 status=0
-                createGraceBlackwellFixture
-                : > "$GPU_DKMS_MARKER_FILE"
-                gb_packages_purged=false
+                setupGraceBlackwellTeardown
                 gb_extra_driver_package=nvidia-driver-535-server
                 mockGraceBlackwellPackageQueries
                 systemctlDisableAndStop() { echo "unexpected service stop"; }
                 apt_get_purge() { echo "unexpected package purge"; }
                 cleanUpGPUDrivers() { echo "unexpected marker cleanup"; }
                 runGraceBlackwellFailureAndRetainState || status=1
-                createGraceBlackwellFixture
-                : > "$GPU_DKMS_MARKER_FILE"
+                setupGraceBlackwellTeardown
                 gb_packages_purged=false
                 gb_extra_driver_package=""
                 mockGraceBlackwellPackageQueries
@@ -443,23 +430,18 @@ Describe 'cse_install_ubuntu.sh'
                 apt_get_purge() { echo "unexpected package purge"; }
                 cleanUpGPUDrivers() { echo "unexpected marker cleanup"; }
                 runGraceBlackwellFailureAndRetainState || status=1
-                createGraceBlackwellFixture
-                : > "$GPU_DKMS_MARKER_FILE"
+                setupGraceBlackwellTeardown
                 gb_extra_driver_package=""
-                gb_packages_purged=false
                 mockGraceBlackwellPackageQueries
                 lsmod() { echo "Module Size Used by"; }
                 systemctlDisableAndStop() { :; }
                 systemctl() { return 1; }
-                systemctlDisableAndStop() { :; }
-                systemctl() { return 1; }
                 update-initramfs() { :; }
                 apt_get_purge() { return 1; }
                 runGraceBlackwellFailureAndRetainState || status=1
-                createGraceBlackwellFixture
+                setupGraceBlackwellTeardown
                 gb_version_mismatch=true
                 mockGraceBlackwellPackageQueries
-                : > "$GPU_DKMS_MARKER_FILE"
                 apt_get_purge() { echo "unexpected package purge"; }
                 cleanUpGPUDrivers() { echo "unexpected marker cleanup"; }
                 cleanUpGPUDriversForBasePrep || status=1
@@ -491,7 +473,7 @@ Describe 'cse_install_ubuntu.sh'
     End
 
     Describe 'installPackageFromCache version matching'
-        deb_cache_root="/tmp/shellspec-deb-cache-$$"
+        deb_cache_root="${PWD}/.shellspec-deb-cache-$$"
 
         setup_deb_cache() {
             mkdir -p "$deb_cache_root"
@@ -613,7 +595,7 @@ EOF
     End
 
     Describe 'installContainerdWithAptGet revision comparison'
-        containerd_download_root="/tmp/cse-install-ubuntu-containerd-$$"
+        containerd_download_root="${PWD}/.shellspec-cse-install-ubuntu-containerd-$$"
 
         setup_containerd_revision() {
             mkdir -p "${containerd_download_root}"
@@ -664,7 +646,7 @@ EOF
     End
 
     Describe 'logResolvedPackageVersion'
-        resolved_version_log="/tmp/cse-install-ubuntu-resolved-version-$$"
+        resolved_version_log="${PWD}/.shellspec-cse-install-ubuntu-resolved-version-$$"
 
         cleanup_resolved_version_log() {
             rm -f "${resolved_version_log}"
@@ -694,7 +676,7 @@ EOF
     End
 
     Describe 'ensureRunc repository fallback'
-        runc_download_root="/tmp/cse-install-ubuntu-runc-$$"
+        runc_download_root="${PWD}/.shellspec-cse-install-ubuntu-runc-$$"
 
         setup_runc_fallback() {
             mkdir -p "${runc_download_root}"
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
index a6e2ead837c..eae39e54b9c 100644
--- 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
@@ -2,6 +2,17 @@
 
 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
+        BASE_PREP_RAN=false
+    }
 
     run_phases() {
         local dispatch cleanup
@@ -16,166 +27,87 @@ Describe 'GPU cleanup across PIS phases'
         nodePrep() { eval "$cleanup"; }
         logs_to_events() { shift; "$@"; }
         cleanUpGPUDrivers() {
-            rm -rf "$driver_root/usr/local/nvidia" "$driver_root/var/lib/dkms/nvidia" \
-                "$driver_root/lib/modules/test/updates/dkms/nvidia.ko" "$driver_root/usr/bin/nvidia-smi"
-            rm -f "$driver_root/dkms-marker"
+            rm -rf "${driver_root}/usr/local/nvidia"
+            touch "${fixture}/cleanup-called"
         }
         eval "$dispatch"
     }
 
-    driver_survives() {
-        run_phases
-        test -d "$driver_root/usr/local/nvidia" &&
-            test -d "$driver_root/var/lib/dkms/nvidia" &&
-            test -f "$driver_root/lib/modules/test/updates/dkms/nvidia.ko" &&
-            test -f "$driver_root/usr/bin/nvidia-smi" &&
-            test -f "$driver_root/dkms-marker"
+    run_base_prep_cleanup() {
+        local cleanup
+        cleanup=$(sed -n '/    if \[ "${PRE_PROVISION_ONLY}" = "true" \] && \[ "${GPU_NODE}" != "true" \] &&/,/    fi/p' "$CSE_MAIN")
+        logs_to_events() { shift; "$@"; }
+        cleanUpGPUDriversForBasePrep() { echo "basePrep driver cleanup"; return "${CLEANUP_STATUS:-0}"; }
+        isMarinerOrAzureLinux() { [ "$OS" = azurelinux ]; }
+        isAzureLinuxOSGuard() { [ "$OS_VARIANT" = osguard ]; }
+        UBUNTU_OS_NAME=ubuntu
+        ERR_NVIDIA_DRIVER_INSTALL=224
+        ( eval "$cleanup" )
     }
 
-    setup_driver() {
-        driver_root=$(mktemp -d)
-        marker="$driver_root/base_prep.complete"
-        mkdir -p "$driver_root/usr/local/nvidia" "$driver_root/var/lib/dkms/nvidia" \
-            "$driver_root/lib/modules/test/updates/dkms" "$driver_root/usr/bin"
-        touch "$driver_root/dkms-marker" \
-            "$driver_root/lib/modules/test/updates/dkms/nvidia.ko" "$driver_root/usr/bin/nvidia-smi"
-        GPU_NODE=false
-        PRE_PROVISION_ONLY=false
-        skip_nvidia_driver_install=false
+    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 "$driver_root"'
 
-    It 'preserves customer-installed driver artifacts on a derived opt-out node'
-        touch "$marker"
-        When call driver_survives
-        The status should be success
-        The output should include "Skipping basePrep"
-    End
+    BeforeEach 'setup_driver'
+    AfterEach 'rm -rf "$fixture"'
 
-    It 'preserves customer drivers on a derived GPU node with the legacy skip tag'
+    It 'preserves customer drivers on a PIS-derived opt-out node when basePrep is skipped'
         touch "$marker"
-        GPU_NODE=true
-        skip_nvidia_driver_install=true
-        When call driver_survives
-        The status should be success
-        The output should include "Skipping basePrep"
-    End
-
-    Describe 'PIS basePrep completion marker'
-        run_cse_marker() {
-            local marker_script
-            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/$marker_root}
-            marker_script=${marker_script//\/var\/log\/azure\/cluster-provision.log/$marker_root/provision.log}
-            eval "$marker_script"
-        }
-
-        setup_marker() { marker_root=$(mktemp -d); PRE_PROVISION_ONLY=true; }
-        BeforeEach 'setup_marker'
-        AfterEach 'rm -rf "$marker_root"'
-
-        It 'does not capture a failed GPU teardown as a completed image'
-            EXIT_CODE=224
-            When call run_cse_marker
-            The status should be success
-            The path "$marker_root/base_prep.complete" should not be exist
-        End
-
-        It 'marks a successful bake for derived nodes'
-            EXIT_CODE=0
-            When call run_cse_marker
-            The status should be success
-            The path "$marker_root/base_prep.complete" should be exist
-        End
-    End
-
-    It 'still cleans an opt-out node when 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 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 'still cleans a non-PIS node with the legacy skip-driver tag'
-        GPU_NODE=true
-        skip_nvidia_driver_install=true
+    It 'cleans an opt-out node when non-PIS 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 'does not clean a GPU bake VM'
-        PRE_PROVISION_ONLY=true
-        When call driver_survives
-        The status should be success
-        The output should include "basePrep"
-    End
-
-    base_cleanup() {
-        local cleanup
-        cleanup=$(sed -n '/    if \[ "${PRE_PROVISION_ONLY}" = "true" \] && \[ "${GPU_NODE}" != "true" \] &&/,/    fi/p' "$CSE_MAIN")
-        logs_to_events() { shift; "$@"; }
-        cleanUpGPUDriversForBasePrep() { echo "basePrep driver cleanup"; return "${CLEANUP_STATUS:-0}"; }
-        isMarinerOrAzureLinux() { [ "$OS" = azurelinux ]; }
-        isAzureLinuxOSGuard() { [ "$OS_VARIANT" = osguard ]; }
-        UBUNTU_OS_NAME=ubuntu
-        ERR_NVIDIA_DRIVER_INSTALL=224
-        ( eval "$cleanup" )
-    }
-
-    It 'tears down AKS drivers on a non-GPU bake VM before image capture'
+    It 'runs strict driver cleanup on an opt-out bake'
         PRE_PROVISION_ONLY=true
         OS=ubuntu
-        When call base_cleanup
+        When call run_base_prep_cleanup
         The status should be success
         The output should include "basePrep driver cleanup"
     End
 
-    It 'fails basePrep when bake-time teardown cannot be verified'
+    It 'fails basePrep when bake-time cleanup cannot be verified'
         PRE_PROVISION_ONLY=true
         OS=ubuntu
         CLEANUP_STATUS=1
-        When call base_cleanup
+        When call run_base_prep_cleanup
         The status should equal 224
         The output should include "basePrep driver cleanup"
     End
 
-    It 'does not tear down AKS drivers on a GPU bake VM'
-        PRE_PROVISION_ONLY=true
-        OS=ubuntu
-        GPU_NODE=true
-        When call base_cleanup
-        The status should be success
-        The output should not include "basePrep driver cleanup"
-    End
-
-    It 'does not tear down AKS drivers in non-PIS basePrep'
-        OS=ubuntu
-        When call base_cleanup
-        The status should be success
-        The output should not include "basePrep driver cleanup"
-    End
-
-    It 'tears down AKS caches on an Azure Linux bake VM'
-        PRE_PROVISION_ONLY=true
-        OS=azurelinux
-        When call base_cleanup
-        The status should be success
-        The output should include "basePrep driver cleanup"
-    End
+    Describe 'PIS basePrep completion marker'
+        It 'does not write the 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 'does not call unavailable cleanup on Azure Container Linux'
-        PRE_PROVISION_ONLY=true
-        OS=azurelinux
-        OS_VARIANT=osguard
-        When call base_cleanup
-        The status should be success
-        The output should not include "basePrep driver cleanup"
+        It 'writes the 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
 End

From bf12a514f36efe7fb7e32fc04a3dc691fc568a3b Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Thu, 8 Oct 2026 12:20:20 +1300
Subject: [PATCH 08/24] fix(gpu): make GB cleanup retries idempotent

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 .../artifacts/ubuntu/cse_install_ubuntu.sh          | 13 +++++++++++++
 .../cloud-init/artifacts/cse_install_ubuntu_spec.sh |  8 ++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

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 21766ac5a0c..d35f826e5aa 100755
--- a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
+++ b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
@@ -500,6 +500,19 @@ cleanUpGraceBlackwellGPUDriver() {
         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
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 91a87be7d79..47ede0bc9f5 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
@@ -238,7 +238,7 @@ Describe 'cse_install_ubuntu.sh'
 
         It 'tears down current and legacy BOM-owned Grace Blackwell drivers before marker cleanup'
             runGraceBlackwellTeardown() {
-                local wave="${1:-wave2}"
+                local wave="${1:-wave2}" retry="${2:-false}"
                 setupGraceBlackwellTeardown "$wave"
                 systemctlDisableAndStop() {
                     echo "stop:$1"
@@ -290,13 +290,17 @@ Describe 'cse_install_ubuntu.sh'
                 prebakedGPUDriverArtifactsRemain() { [ "$gb_packages_purged" != true ]; }
                 cleanUpGPUDriversForBasePrep
                 status=$?
+                if [ "$status" -eq 0 ] && [ "$retry" = true ]; then
+                    cleanUpGPUDriversForBasePrep
+                    status=$?
+                fi
                 cleanupGraceBlackwellFixture
                 return "$status"
             }
 
             runBothGraceBlackwellTeardowns() {
                 runGraceBlackwellTeardown wave2 || return 1
-                runGraceBlackwellTeardown wave1
+                runGraceBlackwellTeardown wave1 true
             }
 
             When call runBothGraceBlackwellTeardowns

From 1334403a5cc3c7b812157ab50c0c6eca8c8beb15 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 9 Oct 2026 12:53:37 +1300
Subject: [PATCH 09/24] fix(gpu): preserve RDMA service on PIS bakes

Stop openibd without disabling it, then refresh loaded modules before teardown. Cover stop failures, retry behavior and service enablement.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 .../artifacts/ubuntu/cse_install_ubuntu.sh    |  13 +-
 .../artifacts/cse_install_ubuntu_spec.sh      | 198 ++++++++++++------
 2 files changed, 149 insertions(+), 62 deletions(-)

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 d35f826e5aa..122b31407d7 100755
--- a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
+++ b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
@@ -538,7 +538,7 @@ cleanUpGraceBlackwellGPUDriver() {
         [ -n "${package}" ] && purge_packages+=("${package}")
     done <<< "${packages}"
 
-    for service in nvidia-device-plugin nvidia-dcgm-exporter nvidia-dcgm nvidia-imex openibd; do
+    for service in nvidia-device-plugin nvidia-dcgm-exporter nvidia-dcgm nvidia-imex; 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
@@ -547,6 +547,17 @@ cleanUpGraceBlackwellGPUDriver() {
     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; }
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 47ede0bc9f5..c2e06cc1ca6 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
@@ -89,6 +89,7 @@ Describe 'cse_install_ubuntu.sh'
         It 'keeps the marker if any prebaked userspace artifact remains'
             marker="$(mktemp)"
             GPU_DKMS_MARKER_FILE="$marker"
+            AfterRun 'command rm -f "$marker"'
             rm() { echo "mock rm $*"; }
             ldconfig() { :; }
             lsmod() { :; }
@@ -102,6 +103,7 @@ Describe 'cse_install_ubuntu.sh'
         It 'keeps the marker when loaded-module inspection fails'
             marker="$(mktemp)"
             GPU_DKMS_MARKER_FILE="$marker"
+            AfterRun 'command rm -f "$marker"'
             rm() { echo "mock rm $*"; }
             ldconfig() { :; }
             lsmod() { return 1; }
@@ -206,6 +208,102 @@ Describe 'cse_install_ubuntu.sh'
             mockGraceBlackwellPackageQueries
         }
 
+        runGraceBlackwellTeardown() {
+            local wave="${1:-wave2}" retry="${2:-false}" enabled="${3:-true}"
+            setupGraceBlackwellTeardown "$wave"
+            gb_openibd_enabled="$enabled"
+            gb_openibd_active=true
+            gb_openibd_stop_count=0
+            gb_openibd_stop_failure="${4:-false}"
+            gb_openibd_remain_active="${5:-false}"
+            gb_openibd_unloads_module="${6:-false}"
+            systemctlDisableAndStop() {
+                echo "stop:$1"
+                if [ "$1" = openibd ]; then
+                    gb_services_stopped=true
+                    gb_openibd_active=false
+                    gb_openibd_enabled=false
+                fi
+            }
+            systemctl() {
+                case "$1" in
+                    cat) [ "$2" = openibd ] ;;
+                    is-active) [ "$3" = openibd ] && [ "$gb_openibd_active" = true ] ;;
+                    is-enabled) [ "$3" = openibd ] && [ "$gb_openibd_enabled" = true ] ;;
+                    *) return 1 ;;
+                esac
+            }
+            systemctl_stop() {
+                [ "$4" = openibd ] || return 1
+                gb_openibd_stop_count=$((gb_openibd_stop_count + 1))
+                [ "$gb_openibd_stop_failure" != true ] || return 1
+                [ "$gb_openibd_remain_active" = true ] || gb_openibd_active=false
+                [ "$gb_openibd_active" = false ] && gb_services_stopped=true
+                echo "stop:openibd"
+                if [ "$gb_openibd_unloads_module" = true ]; then
+                    local unloaded_module="${gb_loaded_modules%% *}"
+                    gb_loaded_modules="${gb_loaded_modules#* }"
+                    gb_expected_modules="${gb_expected_modules#* }"
+                    echo "openibd_unloaded:$unloaded_module"
+                fi
+                return 0
+            }
+            lsmod() {
+                echo "Module Size Used by"
+                for gb_module in $gb_loaded_modules; do
+                    echo "$gb_module 123 0"
+                done
+            }
+            rmmod() {
+                [ "$gb_services_stopped" = true ] || return 1
+                [ "$gb_openibd_active" = false ] || return 1
+                [ "$gb_openibd_enabled" = "$enabled" ] || return 1
+                case " $gb_loaded_modules " in *" $1 "*) ;; *) return 1 ;; esac
+                [ "${gb_expected_modules%% *}" = "$1" ] || return 1
+                echo "rmmod:$1 openibd_enabled=$gb_openibd_enabled"
+                gb_expected_modules="${gb_expected_modules#* }"
+                gb_loaded_modules="${gb_loaded_modules#* }"
+            }
+            update-initramfs() {
+                [ "$gb_packages_purged" = true ] || return 1
+                [ ! -e "$GB_NVIDIA_PEERMEM_CONFIG_FILE" ] || return 1
+                [ ! -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || return 1
+                [ ! -e "$GB_NOUVEAU_MODPROBE_CONFIG_FILE" ] || return 1
+                [ -f "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
+                echo "update-initramfs"
+            }
+            apt_get_purge() {
+                [ "$1 $2 $3" = '10 5 300' ] || return 1
+                shift 3
+                local expected='libnvidia-common-580 nvidia-dkms-580-open nvidia-driver-580-open nvidia-utils-580 nvidia-modprobe nvidia-persistenced'
+                [ "$*" = "$expected" ] || return 1
+                [ -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || return 1
+                [ "$gb_services_stopped" = true ] || return 1
+                [ -z "$gb_loaded_modules" ] || return 1
+                [ -z "$gb_expected_modules" ] || return 1
+                echo "purge:$*"
+                gb_packages_purged=true
+            }
+            cleanUpGPUDrivers() {
+                [ "$gb_packages_purged" = true ] || return 1
+                [ ! -e "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
+                echo "aks-marker-cleanup"
+                rm -f "$GPU_DKMS_MARKER_FILE"
+            }
+            prebakedGPUDriverArtifactsRemain() { [ "$gb_packages_purged" != true ]; }
+            cleanUpGPUDriversForBasePrep
+            status=$?
+            if [ "$status" -eq 0 ] && [ "$retry" = true ]; then
+                cleanUpGPUDriversForBasePrep
+                status=$?
+                [ "$gb_openibd_stop_count" -eq 1 ] || status=1
+                echo "openibd_retry_stop_count=$gb_openibd_stop_count"
+            fi
+            [ "$status" -ne 0 ] || echo "openibd_enabled_after_unload=$gb_openibd_enabled"
+            cleanupGraceBlackwellFixture
+            return "$status"
+        }
+
         It 'fails when the prebake marker survives cleanup'
             GPU_DKMS_MARKER_FILE="$(mktemp)"
             cleanUpGPUDrivers() { :; }
@@ -237,67 +335,6 @@ Describe 'cse_install_ubuntu.sh'
         End
 
         It 'tears down current and legacy BOM-owned Grace Blackwell drivers before marker cleanup'
-            runGraceBlackwellTeardown() {
-                local wave="${1:-wave2}" retry="${2:-false}"
-                setupGraceBlackwellTeardown "$wave"
-                systemctlDisableAndStop() {
-                    echo "stop:$1"
-                    if [ "$1" = openibd ]; then
-                        gb_services_stopped=true
-                    fi
-                    return 0
-                }
-                systemctl() { return 1; }
-                lsmod() {
-                    echo "Module Size Used by"
-                    for gb_module in $gb_loaded_modules; do
-                        echo "$gb_module 123 0"
-                    done
-                }
-                rmmod() {
-                    [ "$gb_services_stopped" = true ] || return 1
-                    [ "${gb_expected_modules%% *}" = "$1" ] || return 1
-                    echo "rmmod:$1"
-                    gb_expected_modules="${gb_expected_modules#* }"
-                    gb_loaded_modules="${gb_loaded_modules#* }"
-                }
-                update-initramfs() {
-                    [ "$gb_packages_purged" = true ] || return 1
-                    [ ! -e "$GB_NVIDIA_PEERMEM_CONFIG_FILE" ] || return 1
-                    [ ! -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || return 1
-                    [ ! -e "$GB_NOUVEAU_MODPROBE_CONFIG_FILE" ] || return 1
-                    [ -f "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
-                    echo "update-initramfs"
-                }
-                apt_get_purge() {
-                    [ "$1 $2 $3" = '10 5 300' ] || return 1
-                    shift 3
-                    local expected='libnvidia-common-580 nvidia-dkms-580-open nvidia-driver-580-open nvidia-utils-580 nvidia-modprobe nvidia-persistenced'
-                    [ "$*" = "$expected" ] || return 1
-                    [ -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || return 1
-                    [ "$gb_services_stopped" = true ] || return 1
-                    [ -z "$gb_loaded_modules" ] || return 1
-                    [ -z "$gb_expected_modules" ] || return 1
-                    echo "purge:$*"
-                    gb_packages_purged=true
-                }
-                cleanUpGPUDrivers() {
-                    [ "$gb_packages_purged" = true ] || return 1
-                    [ ! -e "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
-                    echo "aks-marker-cleanup"
-                    rm -f "$GPU_DKMS_MARKER_FILE"
-                }
-                prebakedGPUDriverArtifactsRemain() { [ "$gb_packages_purged" != true ]; }
-                cleanUpGPUDriversForBasePrep
-                status=$?
-                if [ "$status" -eq 0 ] && [ "$retry" = true ]; then
-                    cleanUpGPUDriversForBasePrep
-                    status=$?
-                fi
-                cleanupGraceBlackwellFixture
-                return "$status"
-            }
-
             runBothGraceBlackwellTeardowns() {
                 runGraceBlackwellTeardown wave2 || return 1
                 runGraceBlackwellTeardown wave1 true
@@ -306,10 +343,49 @@ Describe 'cse_install_ubuntu.sh'
             When call runBothGraceBlackwellTeardowns
             The status should be success
             The output should include "rmmod:nvidia_peermem"
+            The output should include "rmmod:nvidia_peermem openibd_enabled=true"
             The output should include "stop:nvidia-imex"
             The output should include "stop:openibd"
             The output should include "purge:"
             The output should include "aks-marker-cleanup"
+            The output should include "openibd_retry_stop_count=1"
+        End
+
+        It 'preserves openibd enablement while stopping it for module unload and retries'
+            When call runGraceBlackwellTeardown wave2 true
+            The status should be success
+            The output should include "rmmod:nvidia_peermem openibd_enabled=true"
+            The output should include "openibd_enabled_after_unload=true"
+            The output should include "openibd_retry_stop_count=1"
+        End
+
+        It 'keeps openibd disabled when it was disabled before module unload'
+            When call runGraceBlackwellTeardown wave2 false false
+            The status should be success
+            The output should include "rmmod:nvidia_peermem openibd_enabled=false"
+            The output should include "openibd_enabled_after_unload=false"
+        End
+
+        It 'fails before module unload when openibd cannot be stopped'
+            When call runGraceBlackwellTeardown wave2 false true true
+            The status should be failure
+            The stderr should include "Grace Blackwell service openibd could not be stopped"
+            The output should not include "rmmod:"
+        End
+
+        It 'fails before module unload when openibd remains active'
+            When call runGraceBlackwellTeardown wave2 false true false true
+            The status should be failure
+            The stderr should include "Grace Blackwell service openibd remains active"
+            The output should not include "rmmod:"
+        End
+
+        It 'refreshes the loaded-module snapshot after openibd stops'
+            When call runGraceBlackwellTeardown wave2 false true false false true
+            The status should be success
+            The output should include "openibd_unloaded:nvidia_peermem"
+            The output should not include "rmmod:nvidia_peermem"
+            The output should include "rmmod:nvidia_uvm"
         End
 
         It 'uses a BOM-listed dependency version instead of comparing it to the driver version'

From 678a7dd075d80be2a09e7528054fa407139db8c6 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 9 Oct 2026 12:53:38 +1300
Subject: [PATCH 10/24] test(gpu): verify PIS installer preservation

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 e2e/scenario/validators.go      | 13 +++++++++---
 e2e/scenario/validators_test.go | 37 +++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/e2e/scenario/validators.go b/e2e/scenario/validators.go
index 8cdff2a2277..682ed135838 100644
--- a/e2e/scenario/validators.go
+++ b/e2e/scenario/validators.go
@@ -585,8 +585,16 @@ func ValidateCustomerNvidiaDriver(ctx context.Context, s *Scenario) error {
 	if err != nil {
 		return err
 	}
-	command := fmt.Sprintf(`set -euo pipefail
+	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}"
@@ -595,9 +603,8 @@ test "${gpu_version}" = "${expected_version}"
 sudo nvidia-smi
 `,
 		shellSingleQuote(expectedVersion),
+		shellSingleQuote(installerPath),
 	)
-	_, err = execScriptOnVMForScenarioValidateExitCode(ctx, s, command, 0, "customer-installed NVIDIA driver did not survive PIS image capture")
-	return err
 }
 
 func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) error {
diff --git a/e2e/scenario/validators_test.go b/e2e/scenario/validators_test.go
index 4d21a8b0e2f..f1f028746dd 100644
--- a/e2e/scenario/validators_test.go
+++ b/e2e/scenario/validators_test.go
@@ -3,11 +3,48 @@ package scenario
 import (
 	"crypto/sha256"
 	"encoding/hex"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
+func TestCustomerNvidiaDriverValidationRequiresStagedInstaller(t *testing.T) {
+	const expectedVersion = "580.159.04"
+	binDir := t.TempDir()
+	installerPath := filepath.Join(t.TempDir(), "nvidia-installer")
+	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 TestValidateSysctlOutput(t *testing.T) {
 	tests := []struct {
 		name     string

From 114abe70b191fea0bd66f1334b723bf437d27424 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 9 Oct 2026 15:57:56 +1300
Subject: [PATCH 11/24] fix(cloud-init): avoid nested script compression

Keep file contents in the outer compressed cloud-init document to reduce CustomData below the Azure Compute limit. Preserve gzip fallback for Ignition and scriptless paths, and recognize both template forms in hotfix generation.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 hotfix/hotfix_generate.py                 |  22 ++-
 hotfix/hotfix_generate_test.py            |  95 ++++++++++++-
 parts/linux/cloud-init/nodecustomdata.yml | 104 ++++----------
 pkg/agent/baker.go                        |  50 ++++++-
 pkg/agent/baker_test.go                   |  77 +++++++++-
 pkg/agent/bakerapi_test.go                | 162 ++++++++++++++++++++++
 pkg/agent/utils.go                        |  14 +-
 pkg/agent/variables.go                    |  63 +++++----
 8 files changed, 472 insertions(+), 115 deletions(-)

diff --git a/hotfix/hotfix_generate.py b/hotfix/hotfix_generate.py
index 0f927504853..7028b70730d 100644
--- a/hotfix/hotfix_generate.py
+++ b/hotfix/hotfix_generate.py
@@ -55,8 +55,8 @@
 SCRIPTS_BEGIN = "# ---- hotfix-scripts: auto-generated ----"
 SCRIPTS_END = "# ---- end hotfix-scripts ----"
 
-# Map from source file paths (relative to artifacts/) to the GetVariableProperty
-# keys used in nodecustomdata.yml. Only scripts that appear as write_files entries
+# Map from source file paths (relative to artifacts/) to the template variable
+# keys used in nodecustomdata.yml. Only files that appear as write_files entries
 # in the traditional section are included.
 SOURCE_TO_VARKEY = {
     # CSE helpers — base (non-distro)
@@ -379,7 +379,9 @@ def parse_write_files_blocks(traditional_lines):
     """Parse write_files blocks from the traditional section.
 
     Each block is either a simple '- path:' entry or an entire conditional
-    block (e.g., {{if IsAzlOSGuard}}...{{end}}) treated as a single unit.
+    block (e.g., {{if IsAzlOSGuard}}...{{end}}) treated as a single unit. File
+    keys may use either the legacy GetVariableProperty form or the literal-file
+    GetCloudInitFileProperties form.
 
     Returns a list of (varkeys_set, lines_list) tuples.
     """
@@ -420,9 +422,17 @@ def parse_write_files_blocks(traditional_lines):
 
         if in_block:
             current_block.append(line)
-            match = re.search(r'GetVariableProperty\s+"cloudInitData"\s+"(\w+)"', stripped)
-            if match:
-                current_varkeys.add(match.group(1))
+            matches = (
+                re.search(
+                    r'GetVariableProperty\s+"cloudInitData"\s+"(\w+)"',
+                    stripped,
+                ),
+                re.search(r'GetCloudInitFileProperties\s+"(\w+)"', stripped),
+            )
+            for match in matches:
+                if match:
+                    current_varkeys.add(match.group(1))
+                    break
 
     if current_block and current_varkeys:
         blocks.append((current_varkeys.copy(), list(current_block)))
diff --git a/hotfix/hotfix_generate_test.py b/hotfix/hotfix_generate_test.py
index 5f7eac72792..c52a6419e7a 100644
--- a/hotfix/hotfix_generate_test.py
+++ b/hotfix/hotfix_generate_test.py
@@ -54,6 +54,16 @@
 {{end}}
 """
 
+MIXED_ENCODING_TEMPLATE = """- path: /literal.sh
+  owner: root
+{{GetCloudInitFileProperties "literalFile"}}
+- path: /legacy.sh
+  owner: root
+  encoding: gzip
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "legacyFile"}}
+"""
+
 
 class HotfixGenerateTests(unittest.TestCase):
     def test_config_modules_are_independently_hotfixable(self):
@@ -109,7 +119,10 @@ def test_config_modules_are_independently_hotfixable(self):
                 self.assertEqual({variable}, blocks[0][0])
                 self.assertIn(f"- path: {{{{{path_function}}}}}", template)
                 self.assertIn('permissions: "0744"', template)
-                self.assertIn("encoding: gzip", template)
+                self.assertTrue(
+                    f'GetCloudInitFileProperties "{variable}"' in template
+                    or f'GetVariableProperty "cloudInitData" "{variable}"' in template
+                )
                 self.assertNotIn("{{if", template)
 
     def test_config_refactor_hotfix_selects_parent_and_new_modules(self):
@@ -171,6 +184,86 @@ def test_parse_write_files_blocks_keeps_distro_chain_together(self):
             blocks[1][0],
         )
 
+    def test_parse_write_files_blocks_accepts_literal_and_legacy_syntax(self):
+        blocks = hotfix_generate.parse_write_files_blocks(
+            MIXED_ENCODING_TEMPLATE.splitlines(keepends=True)
+        )
+
+        self.assertEqual(
+            [
+                ({"literalFile"}, [
+                    "- path: /literal.sh\n",
+                    "  owner: root\n",
+                    '{{GetCloudInitFileProperties "literalFile"}}\n',
+                ]),
+                ({"legacyFile"}, [
+                    "- path: /legacy.sh\n",
+                    "  owner: root\n",
+                    "  encoding: gzip\n",
+                    "  content: !!binary |\n",
+                    '    {{GetVariableProperty "cloudInitData" "legacyFile"}}\n',
+                ]),
+            ],
+            blocks,
+        )
+
+    def test_current_template_literal_entries_are_injected_and_rendered(self):
+        repository = Path(__file__).resolve().parents[1]
+        template_lines = (repository / hotfix_generate.TEMPLATE).read_text().splitlines(
+            keepends=True
+        )
+        _, outer_else, end = hotfix_generate.find_block_boundaries(template_lines)
+        self.assertIsNotNone(outer_else)
+        self.assertIsNotNone(end)
+        traditional = template_lines[outer_else + 1:end]
+        available = set().union(
+            *(keys for keys, _ in hotfix_generate.parse_write_files_blocks(traditional))
+        )
+        selected = {"provisionConfigsGPU"}
+        self.assertIn("provisionConfigsGPU", available)
+
+        hotfix_template = hotfix_generate.build_hotfix_template(selected, traditional)
+        self.assertIn(
+            '{{GetCloudInitFileProperties "provisionConfigsGPU"}}',
+            hotfix_template,
+        )
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            injected_template = Path(temp_dir) / "nodecustomdata.yml"
+            injected_template.write_text("".join(template_lines))
+            generated_dir = Path(temp_dir) / "generated"
+            with mock.patch.object(
+                hotfix_generate, "TEMPLATE", str(injected_template)
+            ), mock.patch.object(
+                hotfix_generate, "GENERATED_DIR", str(generated_dir)
+            ):
+                hotfix_generate.update_nodecustomdata(selected)
+                injected = injected_template.read_text()
+                marked_block = injected.split(hotfix_generate.SCRIPTS_BEGIN, 1)[1].split(
+                    hotfix_generate.SCRIPTS_END, 1
+                )[0]
+                self.assertIn(
+                    '{{GetCloudInitFileProperties "provisionConfigsGPU"}}',
+                    marked_block,
+                )
+
+                hotfix_generate.write_rendered_payload(selected, traditional)
+
+            for platform in ("ubuntu", "mariner"):
+                rendered = (
+                    generated_dir / f"rendered_nodecustomdata_{platform}.yml"
+                ).read_text()
+                self.assertIn(
+                    "- path: /opt/azure/containers/provision_configs_gpu.sh",
+                    rendered,
+                )
+                self.assertIn('permissions: "0744"', rendered)
+                self.assertIn(r'content: "#!/bin/bash\n', rendered)
+                self.assertIn("configGPUDrivers() {", rendered)
+                self.assertNotIn("GetCloudInitFileProperties", rendered)
+                self.assertNotIn("{{", rendered)
+                self.assertNotIn("encoding: gzip", rendered)
+
     def test_build_hotfix_template_selects_only_requested_blocks(self):
         rendered = hotfix_generate.build_hotfix_template(
             {"provisionSource"},
diff --git a/parts/linux/cloud-init/nodecustomdata.yml b/parts/linux/cloud-init/nodecustomdata.yml
index f2822c3034c..de9305ec558 100644
--- a/parts/linux/cloud-init/nodecustomdata.yml
+++ b/parts/linux/cloud-init/nodecustomdata.yml
@@ -28,194 +28,142 @@ write_files:
 {{if IsAKSCustomCloud}}
 - path: {{GetInitAKSCloudFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "initAKSCloud"}}
+{{GetCloudInitFileProperties "initAKSCloud"}}
 {{end}}
 
 {{- else }}
 - path: {{GetCSEHelpersScriptFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionSource"}}
+{{GetCloudInitFileProperties "provisionSource"}}
 
 
 {{if IsACL }}
 - path: {{GetCSEHelpersScriptDistroFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionSourceACL"}}
+{{GetCloudInitFileProperties "provisionSourceACL"}}
 {{- else if IsAzlOSGuard}}
 - path: {{GetCSEHelpersScriptDistroFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionSourceAzlOSGuard"}}
+{{GetCloudInitFileProperties "provisionSourceAzlOSGuard"}}
 {{- else if IsMariner}}
 - path: {{GetCSEHelpersScriptDistroFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionSourceMariner"}}
+{{GetCloudInitFileProperties "provisionSourceMariner"}}
 {{- else if IsFlatcar }}
 - path: {{GetCSEHelpersScriptDistroFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionSourceFlatcar"}}
+{{GetCloudInitFileProperties "provisionSourceFlatcar"}}
 {{- else }}
 - path: {{GetCSEHelpersScriptDistroFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionSourceUbuntu"}}
+{{GetCloudInitFileProperties "provisionSourceUbuntu"}}
 {{end}}
 
 {{ if not IsCustomImage -}}
 - path: /opt/azure/containers/provision_start.sh
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionStartScript"}}
+{{GetCloudInitFileProperties "provisionStartScript"}}
 {{- end }}
 
 - path: /opt/azure/containers/provision.sh
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionScript"}}
+{{GetCloudInitFileProperties "provisionScript"}}
 
 - path: {{GetCSEInstallScriptFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionInstalls"}}
+{{GetCloudInitFileProperties "provisionInstalls"}}
 
 {{if IsACL }}
 - path: {{GetCSEInstallScriptDistroFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionInstallsACL"}}
+{{GetCloudInitFileProperties "provisionInstallsACL"}}
 {{- else if IsAzlOSGuard}}
 - path: {{GetCSEInstallScriptDistroFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionInstallsAzlOSGuard"}}
+{{GetCloudInitFileProperties "provisionInstallsAzlOSGuard"}}
 {{- else if IsMariner}}
 - path: {{GetCSEInstallScriptDistroFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionInstallsMariner"}}
+{{GetCloudInitFileProperties "provisionInstallsMariner"}}
 {{- else if IsFlatcar }}
 - path: {{GetCSEInstallScriptDistroFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionInstallsFlatcar"}}
+{{GetCloudInitFileProperties "provisionInstallsFlatcar"}}
 {{- else }}
 - path: {{GetCSEInstallScriptDistroFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionInstallsUbuntu"}}
+{{GetCloudInitFileProperties "provisionInstallsUbuntu"}}
 {{end}}
 
 - path: {{GetCSEConfigScriptFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionConfigs"}}
+{{GetCloudInitFileProperties "provisionConfigs"}}
 
 - path: {{GetCSEConfigGPUScriptFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionConfigsGPU"}}
+{{GetCloudInitFileProperties "provisionConfigsGPU"}}
 
 - path: {{GetCSEConfigLocalDNSScriptFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionConfigsLocalDNS"}}
+{{GetCloudInitFileProperties "provisionConfigsLocalDNS"}}
 
 - path: {{GetCSEConfigKubeletScriptFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionConfigsKubelet"}}
+{{GetCloudInitFileProperties "provisionConfigsKubelet"}}
 
 - path: {{GetCSEConfigNetworkScriptFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionConfigsNetwork"}}
+{{GetCloudInitFileProperties "provisionConfigsNetwork"}}
 
 - path: {{GetCSEConfigAddonsScriptFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "provisionConfigsAddons"}}
+{{GetCloudInitFileProperties "provisionConfigsAddons"}}
 
 {{if IsAKSCustomCloud}}
 - path: {{GetInitAKSCloudFilepath}}
   permissions: "0744"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "initAKSCloud"}}
+{{GetCloudInitFileProperties "initAKSCloud"}}
 {{end}}
 
 - path: /etc/systemd/system/reconcile-private-hosts.service
   permissions: "0644"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "reconcilePrivateHostsService"}}
+{{GetCloudInitFileProperties "reconcilePrivateHostsService"}}
 
 - path: /etc/systemd/system/kubelet.service
   permissions: "0600"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "kubeletSystemdService"}}
+{{GetCloudInitFileProperties "kubeletSystemdService"}}
 
 - path: /opt/azure-network/configure-azure-network.sh
   permissions: "0755"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "configureAzureNetworkScript"}}
+{{GetCloudInitFileProperties "configureAzureNetworkScript"}}
 
 - path: /etc/udev/rules.d/99-azure-network.rules
   permissions: "0644"
-  encoding: gzip
   owner: root
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "azureNetworkUdevRule"}}
+{{GetCloudInitFileProperties "azureNetworkUdevRule"}}
 {{- end }}
diff --git a/pkg/agent/baker.go b/pkg/agent/baker.go
index b8812c4c60a..88f559a5639 100644
--- a/pkg/agent/baker.go
+++ b/pkg/agent/baker.go
@@ -312,13 +312,13 @@ func (t *TemplateGenerator) getLinuxNodeCustomDataJSONObject(config *datamodel.N
 	parameters := getParameters(config)
 	// get variable cloudInit
 	variables := getCustomDataVariables(config)
-	str, e := t.getSingleLineForTemplate(kubernetesNodeCustomDataYaml, config.AgentPoolProfile, getBakerFuncMap(config, parameters, variables), true)
+	str, e := t.getSingleLine(kubernetesNodeCustomDataYaml, config.AgentPoolProfile, getBakerFuncMap(config, parameters, variables), true)
 
 	if e != nil {
 		panic(e)
 	}
 
-	return fmt.Sprintf("{\"customData\": \"%s\"}", str)
+	return getCustomDataJSONObject(str)
 }
 
 const (
@@ -657,6 +657,20 @@ func RenderLinuxNodeCustomDataTemplate(templateContent []byte, config *datamodel
 func getBakerFuncMap(config *datamodel.NodeBootstrappingConfiguration, params paramsMap, variables paramsMap) template.FuncMap {
 	funcMap := getContainerServiceFuncMap(config)
 
+	funcMap["GetCloudInitFileProperties"] = func(name string) (string, error) {
+		if files, ok := variables["cloudInitFileData"].(paramsMap); ok {
+			if content, ok := files[name].(string); ok {
+				return renderCloudInitFileContent(content)
+			}
+		}
+		if cloudInitData, ok := variables["cloudInitData"].(paramsMap); ok {
+			if encoded, ok := cloudInitData[name].(string); ok {
+				return fmt.Sprintf("  encoding: gzip\n  content: !!binary |\n    %s\n", encoded), nil
+			}
+		}
+		return "", fmt.Errorf("cloud-init file %q is unavailable", name)
+	}
+
 	funcMap["GetParameter"] = func(s string) interface{} {
 		if v, ok := params[s].(paramsMap); ok && v != nil {
 			if v["value"] == nil {
@@ -704,6 +718,38 @@ func getBakerFuncMap(config *datamodel.NodeBootstrappingConfiguration, params pa
 	return funcMap
 }
 
+func renderCloudInitFileContent(content string) (string, error) {
+	style := yaml.LiteralStyle
+	for _, line := range strings.Split(content, "\n") {
+		if strings.HasPrefix(line, "\t") {
+			style = yaml.DoubleQuotedStyle
+			break
+		}
+	}
+	document := &yaml.Node{
+		Kind: yaml.MappingNode,
+		Tag:  "!!map",
+		Content: []*yaml.Node{
+			{Kind: yaml.ScalarNode, Tag: "!!str", Value: "content"},
+			{Kind: yaml.ScalarNode, Tag: "!!str", Value: content, Style: style},
+		},
+	}
+	var buffer bytes.Buffer
+	encoder := yaml.NewEncoder(&buffer)
+	encoder.SetIndent(2)
+	if err := encoder.Encode(document); err != nil {
+		return "", fmt.Errorf("encode cloud-init file content: %w", err)
+	}
+	if err := encoder.Close(); err != nil {
+		return "", fmt.Errorf("close cloud-init file content encoder: %w", err)
+	}
+	lines := strings.Split(strings.TrimSuffix(buffer.String(), "\n"), "\n")
+	for index := range lines {
+		lines[index] = "  " + lines[index]
+	}
+	return strings.Join(lines, "\n") + "\n", nil
+}
+
 /* normalizeResourceGroupNameForLabel normalizes resource group name to be used as a label,
 similar to what the ARM template used to do.
 
diff --git a/pkg/agent/baker_test.go b/pkg/agent/baker_test.go
index 764b253d47d..e740c8456b8 100644
--- a/pkg/agent/baker_test.go
+++ b/pkg/agent/baker_test.go
@@ -26,6 +26,7 @@ import (
 	. "github.com/onsi/gomega"
 	"github.com/stretchr/testify/require"
 	"github.com/vincent-petithory/dataurl"
+	"gopkg.in/yaml.v3"
 )
 
 // this regex looks for groups of the following forms, returning KEY and VALUE as submatches.
@@ -88,6 +89,71 @@ write_files:
 	}
 }
 
+func TestRenderCloudInitFileContentRoundTrips(t *testing.T) {
+	testCases := []struct {
+		name    string
+		content string
+	}{
+		{name: "empty", content: ""},
+		{name: "no trailing newline", content: "echo ready"},
+		{name: "one trailing newline", content: "echo ready\n"},
+		{name: "multiple trailing newlines", content: "echo ready\n\n"},
+		{name: "YAML-significant text", content: "---\n# shell text\nkey: value\n"},
+		{name: "tabs and indentation", content: "\tcommand\n  indented\n"},
+		{name: "Unicode", content: "printf 'héllo 🌍'\n"},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			content, err := renderCloudInitFileContent(testCase.content)
+			require.NoError(t, err)
+
+			customData := "#cloud-config\nwrite_files:\n- path: /opt/script.sh\n  owner: root\n" + content
+			var parsed cloudInit
+			require.NoError(t, yaml.Unmarshal([]byte(customData), &parsed))
+			require.Len(t, parsed.WriteFiles, 1)
+			require.Equal(t, testCase.content, parsed.WriteFiles[0].Content)
+		})
+	}
+}
+
+func TestGetCloudInitFilePropertiesKeepsGzipFallback(t *testing.T) {
+	testCases := []struct {
+		name                   string
+		distro                 datamodel.Distro
+		enableScriptlessCSECmd bool
+	}{
+		{name: "Flatcar", distro: datamodel.AKSFlatcarGen2},
+		{name: "ACL", distro: datamodel.AKSACLGen2TL},
+		{name: "scriptless", distro: datamodel.AKSUbuntuContainerd2204Gen2, enableScriptlessCSECmd: true},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(testCase.name, func(t *testing.T) {
+			config := newNodeCustomDataRenderConfig(testCase.distro)
+			config.EnableScriptlessCSECmd = testCase.enableScriptlessCSECmd
+			variables := getCustomDataVariables(config)
+			require.NotContains(t, variables, "cloudInitFileData")
+
+			funcMap := getBakerFuncMap(config, getParameters(config), variables)
+			getProperties := funcMap["GetCloudInitFileProperties"].(func(string) (string, error))
+			properties, err := getProperties("provisionSource")
+			require.NoError(t, err)
+			require.Contains(t, properties, "encoding: gzip")
+
+			customData := "#cloud-config\nwrite_files:\n- path: /opt/script.sh\n  owner: root\n" + properties
+			var parsed cloudInit
+			require.NoError(t, yaml.Unmarshal([]byte(customData), &parsed))
+			require.Len(t, parsed.WriteFiles, 1)
+			require.Equal(t, "gzip", parsed.WriteFiles[0].Encoding)
+
+			actual, err := getGzipDecodedValue([]byte(parsed.WriteFiles[0].Content))
+			require.NoError(t, err)
+			require.Equal(t, getCustomScriptContent(kubernetesCSEHelpersScript, config), string(actual))
+		})
+	}
+}
+
 func newNodeCustomDataRenderConfig(distro datamodel.Distro) *datamodel.NodeBootstrappingConfiguration {
 	profile := &datamodel.AgentPoolProfile{
 		Name:   "hotfix-render-test",
@@ -2181,8 +2247,15 @@ var _ = Describe("getLinuxNodeBootstrappingPayload", func() {
 		expectedCustomData := getCustomDataFromJSON(templateGenerator.getLinuxNodeCustomDataJSONObject(config))
 
 		Expect(string(decompressedPayload)).To(Equal(expectedCustomData))
-		Expect(string(decompressedPayload)).NotTo(ContainSubstring(aksNodeCustomDataFilepath))
-		Expect(string(decompressedPayload)).NotTo(ContainSubstring(aksNbcCmdFilepath))
+		var parsedCustomData cloudInit
+		Expect(yaml.Unmarshal(decompressedPayload, &parsedCustomData)).To(Succeed())
+		for _, file := range parsedCustomData.WriteFiles {
+			Expect(file.Path).NotTo(Equal(aksNodeCustomDataFilepath))
+			Expect(file.Path).NotTo(Equal(aksNbcCmdFilepath))
+		}
+		for _, command := range parsedCustomData.BootCommands {
+			Expect(command).NotTo(ContainSubstring(aksNbcCmdFilepath))
+		}
 	})
 })
 
diff --git a/pkg/agent/bakerapi_test.go b/pkg/agent/bakerapi_test.go
index 32998201e16..23c5a4a9292 100644
--- a/pkg/agent/bakerapi_test.go
+++ b/pkg/agent/bakerapi_test.go
@@ -1,13 +1,19 @@
 package agent
 
 import (
+	"bytes"
+	"compress/gzip"
 	"context"
+	"encoding/base64"
+	"fmt"
+	"io"
 
 	"github.com/Azure/agentbaker/pkg/agent/datamodel"
 	agenttoggles "github.com/Azure/agentbaker/pkg/agent/toggles"
 	"github.com/barkimedes/go-deepcopy"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
+	"gopkg.in/yaml.v3"
 )
 
 type testToggles struct {
@@ -179,6 +185,162 @@ var _ = Describe("AgentBaker API implementation tests", func() {
 			Expect(nodeBootStrapping.SigImageConfig.Definition).To(Equal("2204gen2containerd"))
 		})
 
+		It("should render Linux cloud-init files as literal write_files content without changing their bytes", func() {
+			testCases := []struct {
+				name   string
+				distro datamodel.Distro
+			}{
+				{name: "Ubuntu 22.04", distro: datamodel.AKSUbuntuContainerd2204Gen2},
+				{name: "Ubuntu 24.04", distro: datamodel.AKSUbuntuContainerd2404Gen2},
+				{name: "Azure Linux 3", distro: datamodel.AKSAzureLinuxV3Gen2},
+			}
+
+			for _, testCase := range testCases {
+				for _, preProvisionOnly := range []bool{false, true} {
+					name := testCase.name
+					if preProvisionOnly {
+						name += " PIS"
+					} else {
+						name += " ordinary"
+					}
+					By(name)
+
+					configCopy, err := deepcopy.Anything(config)
+					Expect(err).NotTo(HaveOccurred())
+					testConfig, ok := configCopy.(*datamodel.NodeBootstrappingConfiguration)
+					Expect(ok).To(BeTrue())
+					testConfig.AgentPoolProfile.Distro = testCase.distro
+					testConfig.AgentPoolProfile.VMSize = "Standard_NC4as_T4_v3"
+					testConfig.PreProvisionOnly = preProvisionOnly
+					testConfig.EnableScriptlessCSECmd = false
+					testConfig.ConfigGPUDriverIfNeeded = false
+					testConfig.EnableGPUDevicePluginIfNeeded = false
+					testConfig.EnableNvidia = false
+
+					payload := InitializeTemplateGenerator().getLinuxNodeBootstrappingPayload(testConfig)
+					Expect(len(payload)).To(BeNumerically("<", MaxCustomDataLength))
+					encoded, err := base64.StdEncoding.Strict().DecodeString(payload)
+					Expect(err).NotTo(HaveOccurred())
+					Expect(len(encoded)).To(BeNumerically("<=", 65535))
+					reader, err := gzip.NewReader(bytes.NewReader(encoded))
+					Expect(err).NotTo(HaveOccurred())
+					customDataBytes, err := io.ReadAll(reader)
+					Expect(err).NotTo(HaveOccurred())
+					Expect(reader.Close()).To(Succeed())
+					var customData cloudInit
+					Expect(yaml.Unmarshal(customDataBytes, &customData)).To(Succeed())
+					fmt.Fprintf(
+						GinkgoWriter,
+						"PIS_CUSTOMDATA distro=%s preProvisionOnly=%t base64=%d gzip=%d yaml=%d\n",
+						testCase.distro,
+						preProvisionOnly,
+						len(payload),
+						len(encoded),
+						len(customDataBytes),
+					)
+
+					customDataVariables := getCustomDataVariables(testConfig)
+					variables := customDataVariables["cloudInitFileData"].(paramsMap)
+					expectedScripts := []struct {
+						path string
+						key  string
+					}{
+						{cseHelpersScriptFilepath, "provisionSource"},
+						{cseHelpersScriptDistroFilepath, "provisionSourceUbuntu"},
+						{"/opt/azure/containers/provision_start.sh", "provisionStartScript"},
+						{"/opt/azure/containers/provision.sh", "provisionScript"},
+						{cseInstallScriptFilepath, "provisionInstalls"},
+						{cseInstallScriptDistroFilepath, "provisionInstallsUbuntu"},
+						{cseConfigScriptFilepath, "provisionConfigs"},
+						{cseConfigGPUScriptFilepath, "provisionConfigsGPU"},
+						{cseConfigLocalDNSScriptFilepath, "provisionConfigsLocalDNS"},
+						{cseConfigKubeletScriptFilepath, "provisionConfigsKubelet"},
+						{cseConfigNetworkScriptFilepath, "provisionConfigsNetwork"},
+						{cseConfigAddonsScriptFilepath, "provisionConfigsAddons"},
+					}
+					if testCase.distro == datamodel.AKSAzureLinuxV3Gen2 {
+						expectedScripts[1].key = "provisionSourceMariner"
+						expectedScripts[5].key = "provisionInstallsMariner"
+					}
+
+					expectedPaths := make([]string, 0, len(customData.WriteFiles))
+					for _, script := range expectedScripts {
+						expectedPaths = append(expectedPaths, script.path)
+					}
+					allVariables := customDataVariables["cloudInitData"].(paramsMap)
+					hasChronyScript := allVariables["provisionConfigsChrony"] != nil
+					if hasChronyScript {
+						expectedPaths = append(expectedPaths, "/opt/azure/containers/provision_configs_chrony.sh")
+					}
+					expectedPaths = append(expectedPaths,
+						"/etc/systemd/system/reconcile-private-hosts.service",
+						"/etc/systemd/system/kubelet.service",
+						"/opt/azure-network/configure-azure-network.sh",
+						"/etc/udev/rules.d/99-azure-network.rules",
+					)
+					actualPaths := make([]string, 0, len(customData.WriteFiles))
+					for _, file := range customData.WriteFiles {
+						actualPaths = append(actualPaths, file.Path)
+					}
+					Expect(actualPaths).To(Equal(expectedPaths))
+
+					filesByPath := make(map[string]cloudInitWriteFile, len(customData.WriteFiles))
+					for _, file := range customData.WriteFiles {
+						filesByPath[file.Path] = file
+						Expect(file.Owner).To(Equal("root"))
+					}
+					for _, script := range expectedScripts {
+						file := filesByPath[script.path]
+						Expect(file.Permissions).To(Equal("0744"))
+						Expect(file.Encoding).To(BeEmpty())
+						Expect(file.Content).To(Equal(variables[script.key].(string)), script.path)
+					}
+					expectedLiteralFiles := []struct {
+						path        string
+						key         string
+						permissions string
+					}{
+						{"/etc/systemd/system/reconcile-private-hosts.service", "reconcilePrivateHostsService", "0644"},
+						{"/etc/systemd/system/kubelet.service", "kubeletSystemdService", "0600"},
+						{"/opt/azure-network/configure-azure-network.sh", "configureAzureNetworkScript", "0755"},
+						{"/etc/udev/rules.d/99-azure-network.rules", "azureNetworkUdevRule", "0644"},
+					}
+					for _, expected := range expectedLiteralFiles {
+						file := filesByPath[expected.path]
+						Expect(file.Permissions).To(Equal(expected.permissions))
+						Expect(file.Encoding).To(BeEmpty())
+						Expect(file.Content).To(Equal(variables[expected.key].(string)), expected.path)
+					}
+					expectedGzippedFiles := []struct {
+						path        string
+						key         string
+						permissions string
+					}{}
+					if hasChronyScript {
+						expectedGzippedFiles = append(expectedGzippedFiles, struct {
+							path        string
+							key         string
+							permissions string
+						}{"/opt/azure/containers/provision_configs_chrony.sh", "provisionConfigsChrony", "0744"})
+					}
+					for _, expected := range expectedGzippedFiles {
+						file := filesByPath[expected.path]
+						Expect(file.Permissions).To(Equal(expected.permissions))
+						Expect(file.Encoding).To(Equal("gzip"))
+						encodedFile, err := base64.StdEncoding.DecodeString(allVariables[expected.key].(string))
+						Expect(err).NotTo(HaveOccurred())
+						expectedBytes, err := getGzipDecodedValue(encodedFile)
+						Expect(err).NotTo(HaveOccurred())
+						actualBytes, err := getGzipDecodedValue([]byte(file.Content))
+						Expect(err).NotTo(HaveOccurred())
+						Expect(actualBytes).To(Equal(expectedBytes), expected.path)
+					}
+
+					By(fmt.Sprintf("%s customData size: encoded=%d decoded=%d", name, len(payload), len(customDataBytes)))
+				}
+			}
+		})
+
 		It("should accept supported custom Linux transparent huge page values", func() {
 			config.AgentPoolProfile.CustomLinuxOSConfig = &datamodel.CustomLinuxOSConfig{
 				TransparentHugePageEnabled: "never",
diff --git a/pkg/agent/utils.go b/pkg/agent/utils.go
index 7b6d4072fa9..d67838fd8cb 100644
--- a/pkg/agent/utils.go
+++ b/pkg/agent/utils.go
@@ -181,6 +181,10 @@ func escapeSingleLine(escapedStr string) string {
 
 // getBase64EncodedGzippedCustomScript will return a base64 of the CSE.
 func getBase64EncodedGzippedCustomScript(csFilename string, config *datamodel.NodeBootstrappingConfiguration) string {
+	return getBase64EncodedGzippedCustomScriptFromStr(getCustomScriptContent(csFilename, config))
+}
+
+func getCustomScriptContent(csFilename string, config *datamodel.NodeBootstrappingConfiguration) string {
 	b, err := parts.Templates.ReadFile(csFilename)
 	if err != nil {
 		// this should never happen and this is a bug.
@@ -202,7 +206,7 @@ func getBase64EncodedGzippedCustomScript(csFilename string, config *datamodel.No
 	}
 	csStr := buffer.String()
 	csStr = strings.ReplaceAll(csStr, "\r\n", "\n")
-	return getBase64EncodedGzippedCustomScriptFromStr(csStr)
+	return csStr
 }
 
 // This is "best-effort" - removes MOST of the comments with obvious formats, to lower the space required by CustomData component.
@@ -392,6 +396,14 @@ func getCustomDataFromJSON(jsonStr string) string {
 	return customDataObj["customData"]
 }
 
+func getCustomDataJSONObject(customData string) string {
+	customDataJSON, err := json.Marshal(map[string]string{"customData": customData})
+	if err != nil {
+		panic(fmt.Errorf("failed to marshal customData: %w", err))
+	}
+	return string(customDataJSON)
+}
+
 // GetOrderedKubeletConfigFlagString returns an ordered string of key/val pairs.
 // copied from AKS-Engine and filter out flags that already translated to config file.
 func GetOrderedKubeletConfigFlagString(config *datamodel.NodeBootstrappingConfiguration) string {
diff --git a/pkg/agent/variables.go b/pkg/agent/variables.go
index 3e16b778874..32bab6210bd 100644
--- a/pkg/agent/variables.go
+++ b/pkg/agent/variables.go
@@ -13,36 +13,46 @@ import (
 // getCustomDataVariables returns cloudinit data used by Linux.
 func getCustomDataVariables(config *datamodel.NodeBootstrappingConfiguration) paramsMap {
 	cs := config.ContainerService
+	useCloudInitLiteralFiles := !config.IsFlatcar() && !config.IsACL() && !config.EnableScriptlessCSECmd
+	cloudInitFileData := paramsMap{}
+	getCloudInitFile := func(name, filename string) string {
+		content := getCustomScriptContent(filename, config)
+		if useCloudInitLiteralFiles {
+			cloudInitFileData[name] = content
+			return ""
+		}
+		return getBase64EncodedGzippedCustomScriptFromStr(content)
+	}
 	cloudInitFiles := map[string]interface{}{
 		"cloudInitData": paramsMap{
-			"provisionStartScript":                  getBase64EncodedGzippedCustomScript(kubernetesCSEStartScript, config),
-			"provisionScript":                       getBase64EncodedGzippedCustomScript(kubernetesCSEMainScript, config),
-			"provisionSource":                       getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScript, config),
-			"provisionSourceUbuntu":                 getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScriptUbuntu, config),
-			"provisionSourceMariner":                getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScriptMariner, config),
-			"provisionSourceAzlOSGuard":             getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScriptAzlOSGuard, config),
-			"provisionSourceFlatcar":                getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScriptFlatcar, config),
-			"provisionSourceACL":                    getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScriptACL, config),
-			"provisionInstalls":                     getBase64EncodedGzippedCustomScript(kubernetesCSEInstall, config),
-			"provisionInstallsUbuntu":               getBase64EncodedGzippedCustomScript(kubernetesCSEInstallUbuntu, config),
-			"provisionInstallsMariner":              getBase64EncodedGzippedCustomScript(kubernetesCSEInstallMariner, config),
-			"provisionInstallsAzlOSGuard":           getBase64EncodedGzippedCustomScript(kubernetesCSEInstallAzlOSGuard, config),
-			"provisionInstallsFlatcar":              getBase64EncodedGzippedCustomScript(kubernetesCSEInstallFlatcar, config),
-			"provisionInstallsACL":                  getBase64EncodedGzippedCustomScript(kubernetesCSEInstallACL, config),
-			"provisionConfigs":                      getBase64EncodedGzippedCustomScript(kubernetesCSEConfig, config),
-			"provisionConfigsGPU":                   getBase64EncodedGzippedCustomScript(kubernetesCSEConfigGPU, config),
-			"provisionConfigsLocalDNS":              getBase64EncodedGzippedCustomScript(kubernetesCSEConfigLocalDNS, config),
-			"provisionConfigsKubelet":               getBase64EncodedGzippedCustomScript(kubernetesCSEConfigKubelet, config),
-			"provisionConfigsNetwork":               getBase64EncodedGzippedCustomScript(kubernetesCSEConfigNetwork, config),
-			"provisionConfigsAddons":                getBase64EncodedGzippedCustomScript(kubernetesCSEConfigAddons, config),
+			"provisionStartScript":                  getCloudInitFile("provisionStartScript", kubernetesCSEStartScript),
+			"provisionScript":                       getCloudInitFile("provisionScript", kubernetesCSEMainScript),
+			"provisionSource":                       getCloudInitFile("provisionSource", kubernetesCSEHelpersScript),
+			"provisionSourceUbuntu":                 getCloudInitFile("provisionSourceUbuntu", kubernetesCSEHelpersScriptUbuntu),
+			"provisionSourceMariner":                getCloudInitFile("provisionSourceMariner", kubernetesCSEHelpersScriptMariner),
+			"provisionSourceAzlOSGuard":             getCloudInitFile("provisionSourceAzlOSGuard", kubernetesCSEHelpersScriptAzlOSGuard),
+			"provisionSourceFlatcar":                getCloudInitFile("provisionSourceFlatcar", kubernetesCSEHelpersScriptFlatcar),
+			"provisionSourceACL":                    getCloudInitFile("provisionSourceACL", kubernetesCSEHelpersScriptACL),
+			"provisionInstalls":                     getCloudInitFile("provisionInstalls", kubernetesCSEInstall),
+			"provisionInstallsUbuntu":               getCloudInitFile("provisionInstallsUbuntu", kubernetesCSEInstallUbuntu),
+			"provisionInstallsMariner":              getCloudInitFile("provisionInstallsMariner", kubernetesCSEInstallMariner),
+			"provisionInstallsAzlOSGuard":           getCloudInitFile("provisionInstallsAzlOSGuard", kubernetesCSEInstallAzlOSGuard),
+			"provisionInstallsFlatcar":              getCloudInitFile("provisionInstallsFlatcar", kubernetesCSEInstallFlatcar),
+			"provisionInstallsACL":                  getCloudInitFile("provisionInstallsACL", kubernetesCSEInstallACL),
+			"provisionConfigs":                      getCloudInitFile("provisionConfigs", kubernetesCSEConfig),
+			"provisionConfigsGPU":                   getCloudInitFile("provisionConfigsGPU", kubernetesCSEConfigGPU),
+			"provisionConfigsLocalDNS":              getCloudInitFile("provisionConfigsLocalDNS", kubernetesCSEConfigLocalDNS),
+			"provisionConfigsKubelet":               getCloudInitFile("provisionConfigsKubelet", kubernetesCSEConfigKubelet),
+			"provisionConfigsNetwork":               getCloudInitFile("provisionConfigsNetwork", kubernetesCSEConfigNetwork),
+			"provisionConfigsAddons":                getCloudInitFile("provisionConfigsAddons", kubernetesCSEConfigAddons),
 			"provisionSendLogs":                     getBase64EncodedGzippedCustomScript(kubernetesCSESendLogs, config),
 			"provisionRedactCloudConfig":            getBase64EncodedGzippedCustomScript(kubernetesCSERedactCloudConfig, config),
 			"customSearchDomainsScript":             getBase64EncodedGzippedCustomScript(kubernetesCustomSearchDomainsScript, config),
 			"dhcpv6SystemdService":                  getBase64EncodedGzippedCustomScript(dhcpv6SystemdService, config),
 			"dhcpv6ConfigurationScript":             getBase64EncodedGzippedCustomScript(dhcpv6ConfigurationScript, config),
-			"kubeletSystemdService":                 getBase64EncodedGzippedCustomScript(kubeletSystemdService, config),
+			"kubeletSystemdService":                 getCloudInitFile("kubeletSystemdService", kubeletSystemdService),
 			"reconcilePrivateHostsScript":           getBase64EncodedGzippedCustomScript(reconcilePrivateHostsScript, config),
-			"reconcilePrivateHostsService":          getBase64EncodedGzippedCustomScript(reconcilePrivateHostsService, config),
+			"reconcilePrivateHostsService":          getCloudInitFile("reconcilePrivateHostsService", reconcilePrivateHostsService),
 			"ensureNoDupEbtablesScript":             getBase64EncodedGzippedCustomScript(ensureNoDupEbtablesScript, config),
 			"ensureNoDupEbtablesService":            getBase64EncodedGzippedCustomScript(ensureNoDupEbtablesService, config),
 			"bindMountScript":                       getBase64EncodedGzippedCustomScript(bindMountScript, config),
@@ -62,13 +72,16 @@ func getCustomDataVariables(config *datamodel.NodeBootstrappingConfiguration) pa
 			"cloudInitStatusCheckScript":            getBase64EncodedGzippedCustomScript(cloudInitStatusCheckScript, config),
 			"measureTLSBootstrappingLatencyScript":  getBase64EncodedGzippedCustomScript(measureTLSBootstrappingLatencyScript, config),
 			"measureTLSBootstrappingLatencyService": getBase64EncodedGzippedCustomScript(measureTLSBootstrappingLatencyService, config),
-			"configureAzureNetworkScript":           getBase64EncodedGzippedCustomScript(configureAzureNetworkScript, config),
-			"azureNetworkUdevRule":                  getBase64EncodedGzippedCustomScript(azureNetworkUdevRule, config),
+			"configureAzureNetworkScript":           getCloudInitFile("configureAzureNetworkScript", configureAzureNetworkScript),
+			"azureNetworkUdevRule":                  getCloudInitFile("azureNetworkUdevRule", azureNetworkUdevRule),
 		},
 	}
 
 	cloudInitData := cloudInitFiles["cloudInitData"].(paramsMap) //nolint:errcheck // no error is actually here
-	cloudInitData["initAKSCloud"] = getBase64EncodedGzippedCustomScript(initAKSCloudScript, config)
+	cloudInitData["initAKSCloud"] = getCloudInitFile("initAKSCloud", initAKSCloudScript)
+	if useCloudInitLiteralFiles {
+		cloudInitFiles["cloudInitFileData"] = cloudInitFileData
+	}
 
 	if config.IsFlatcar() || config.IsACL() {
 		cloudInitData["provisionRedactCloudConfig"] = "" // Flatcar and ACL do not have cloud-init

From d6440ae9b2e64eaa9a5951c63220776bdfc22777 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 9 Oct 2026 18:05:10 +1300
Subject: [PATCH 12/24] ci: pin lint Go toolchain to 1.27.1

Go 1.27.2 changed export data format, which golangci-lint v2.13.2 cannot decode. Keep lint on the latest verified patch until its tooling supports the new format.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 .github/workflows/golangci-lint.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/golangci-lint.yml b/.github/workflows/golangci-lint.yml
index 6c3f17bbdb0..0cfc53f4d37 100644
--- a/.github/workflows/golangci-lint.yml
+++ b/.github/workflows/golangci-lint.yml
@@ -22,7 +22,7 @@ jobs:
     steps:
       - uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0
         with:
-          go-version: '1.27'
+          go-version: '1.27.1'
       - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
         with:
           persist-credentials: false

From c25cfd6574c1ab401a3edc006cda77881f90d481 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 9 Oct 2026 18:12:24 +1300
Subject: [PATCH 13/24] fix(cloud-init): satisfy lint for literal file
 rendering

Extract the file-property renderer and check test type assertions so the merged payload test passes the PR lint gate.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 pkg/agent/baker.go         | 26 +++++++++++++++-----------
 pkg/agent/baker_test.go    |  3 ++-
 pkg/agent/bakerapi_test.go | 11 ++++++++---
 3 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/pkg/agent/baker.go b/pkg/agent/baker.go
index 702103ba1e2..8d06ca51420 100644
--- a/pkg/agent/baker.go
+++ b/pkg/agent/baker.go
@@ -660,17 +660,7 @@ func getBakerFuncMap(config *datamodel.NodeBootstrappingConfiguration, params pa
 	funcMap := getContainerServiceFuncMap(config)
 
 	funcMap["GetCloudInitFileProperties"] = func(name string) (string, error) {
-		if files, ok := variables["cloudInitFileData"].(paramsMap); ok {
-			if content, ok := files[name].(string); ok {
-				return renderCloudInitFileContent(content)
-			}
-		}
-		if cloudInitData, ok := variables["cloudInitData"].(paramsMap); ok {
-			if encoded, ok := cloudInitData[name].(string); ok {
-				return fmt.Sprintf("  encoding: gzip\n  content: !!binary |\n    %s\n", encoded), nil
-			}
-		}
-		return "", fmt.Errorf("cloud-init file %q is unavailable", name)
+		return getCloudInitFileProperties(variables, name)
 	}
 
 	funcMap["GetParameter"] = func(s string) interface{} {
@@ -720,6 +710,20 @@ func getBakerFuncMap(config *datamodel.NodeBootstrappingConfiguration, params pa
 	return funcMap
 }
 
+func getCloudInitFileProperties(variables paramsMap, name string) (string, error) {
+	if files, ok := variables["cloudInitFileData"].(paramsMap); ok {
+		if content, ok := files[name].(string); ok {
+			return renderCloudInitFileContent(content)
+		}
+	}
+	if cloudInitData, ok := variables["cloudInitData"].(paramsMap); ok {
+		if encoded, ok := cloudInitData[name].(string); ok {
+			return fmt.Sprintf("  encoding: gzip\n  content: !!binary |\n    %s\n", encoded), nil
+		}
+	}
+	return "", fmt.Errorf("cloud-init file %q is unavailable", name)
+}
+
 func renderCloudInitFileContent(content string) (string, error) {
 	style := yaml.LiteralStyle
 	for _, line := range strings.Split(content, "\n") {
diff --git a/pkg/agent/baker_test.go b/pkg/agent/baker_test.go
index 7885aeee7a8..7a20ddd82d3 100644
--- a/pkg/agent/baker_test.go
+++ b/pkg/agent/baker_test.go
@@ -124,7 +124,8 @@ func TestGetCloudInitFilePropertiesKeepsGzipFallback(t *testing.T) {
 			require.NotContains(t, variables, "cloudInitFileData")
 
 			funcMap := getBakerFuncMap(config, getParameters(config), variables)
-			getProperties := funcMap["GetCloudInitFileProperties"].(func(string) (string, error))
+			getProperties, ok := funcMap["GetCloudInitFileProperties"].(func(string) (string, error))
+			require.True(t, ok)
 			properties, err := getProperties("provisionSource")
 			require.NoError(t, err)
 			require.Contains(t, properties, "encoding: gzip")
diff --git a/pkg/agent/bakerapi_test.go b/pkg/agent/bakerapi_test.go
index 5de35d09c4f..c155c824f65 100644
--- a/pkg/agent/bakerapi_test.go
+++ b/pkg/agent/bakerapi_test.go
@@ -240,7 +240,8 @@ var _ = Describe("AgentBaker API implementation tests", func() {
 					)
 
 					customDataVariables := getCustomDataVariables(testConfig)
-					variables := customDataVariables["cloudInitFileData"].(paramsMap)
+					variables, ok := customDataVariables["cloudInitFileData"].(paramsMap)
+					Expect(ok).To(BeTrue())
 					expectedScripts := []struct {
 						path string
 						key  string
@@ -289,7 +290,9 @@ var _ = Describe("AgentBaker API implementation tests", func() {
 						file := filesByPath[script.path]
 						Expect(file.Permissions).To(Equal("0744"))
 						Expect(file.Encoding).To(BeEmpty())
-						Expect(file.Content).To(Equal(variables[script.key].(string)), script.path)
+						expectedContent, ok := variables[script.key].(string)
+						Expect(ok).To(BeTrue())
+						Expect(file.Content).To(Equal(expectedContent), script.path)
 					}
 					expectedLiteralFiles := []struct {
 						path        string
@@ -305,7 +308,9 @@ var _ = Describe("AgentBaker API implementation tests", func() {
 						file := filesByPath[expected.path]
 						Expect(file.Permissions).To(Equal(expected.permissions))
 						Expect(file.Encoding).To(BeEmpty())
-						Expect(file.Content).To(Equal(variables[expected.key].(string)), expected.path)
+						expectedContent, ok := variables[expected.key].(string)
+						Expect(ok).To(BeTrue())
+						Expect(file.Content).To(Equal(expectedContent), expected.path)
 					}
 					By(fmt.Sprintf("%s customData size: encoded=%d decoded=%d", name, len(payload), len(customDataBytes)))
 				}

From 258161f77f6e39ff0a96b717ee96e5ab1291bc84 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 9 Oct 2026 23:40:46 +1300
Subject: [PATCH 14/24] fix(gpu): stop persistence before module unload

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 .../artifacts/ubuntu/cse_install_ubuntu.sh     |  2 +-
 .../artifacts/cse_install_ubuntu_spec.sh       | 18 ++++++++++++++++--
 2 files changed, 17 insertions(+), 3 deletions(-)

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 122b31407d7..6e9acaabb88 100755
--- a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
+++ b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
@@ -538,7 +538,7 @@ cleanUpGraceBlackwellGPUDriver() {
         [ -n "${package}" ] && purge_packages+=("${package}")
     done <<< "${packages}"
 
-    for service in nvidia-device-plugin nvidia-dcgm-exporter nvidia-dcgm nvidia-imex; do
+    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
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 c2e06cc1ca6..fb606cae318 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
@@ -204,6 +204,8 @@ Describe 'cse_install_ubuntu.sh'
             gb_loaded_modules='nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia '
             gb_expected_modules='nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia '
             gb_services_stopped=false
+            gb_persistenced_active=true
+            gb_persistenced_enabled=true
             gb_packages_purged=false
             mockGraceBlackwellPackageQueries
         }
@@ -219,6 +221,10 @@ Describe 'cse_install_ubuntu.sh'
             gb_openibd_unloads_module="${6:-false}"
             systemctlDisableAndStop() {
                 echo "stop:$1"
+                if [ "$1" = nvidia-persistenced ]; then
+                    gb_persistenced_active=false
+                    gb_persistenced_enabled=false
+                fi
                 if [ "$1" = openibd ]; then
                     gb_services_stopped=true
                     gb_openibd_active=false
@@ -228,8 +234,14 @@ Describe 'cse_install_ubuntu.sh'
             systemctl() {
                 case "$1" in
                     cat) [ "$2" = openibd ] ;;
-                    is-active) [ "$3" = openibd ] && [ "$gb_openibd_active" = true ] ;;
-                    is-enabled) [ "$3" = openibd ] && [ "$gb_openibd_enabled" = true ] ;;
+                    is-active)
+                        { [ "$3" = openibd ] && [ "$gb_openibd_active" = true ]; } ||
+                            { [ "$3" = nvidia-persistenced ] && [ "$gb_persistenced_active" = true ]; }
+                        ;;
+                    is-enabled)
+                        { [ "$3" = openibd ] && [ "$gb_openibd_enabled" = true ]; } ||
+                            { [ "$3" = nvidia-persistenced ] && [ "$gb_persistenced_enabled" = true ]; }
+                        ;;
                     *) return 1 ;;
                 esac
             }
@@ -256,6 +268,7 @@ Describe 'cse_install_ubuntu.sh'
             }
             rmmod() {
                 [ "$gb_services_stopped" = true ] || return 1
+                [ "$gb_persistenced_active" = false ] || return 1
                 [ "$gb_openibd_active" = false ] || return 1
                 [ "$gb_openibd_enabled" = "$enabled" ] || return 1
                 case " $gb_loaded_modules " in *" $1 "*) ;; *) return 1 ;; esac
@@ -344,6 +357,7 @@ Describe 'cse_install_ubuntu.sh'
             The status should be success
             The output should include "rmmod:nvidia_peermem"
             The output should include "rmmod:nvidia_peermem openibd_enabled=true"
+            The output should include "stop:nvidia-persistenced"
             The output should include "stop:nvidia-imex"
             The output should include "stop:openibd"
             The output should include "purge:"

From bb146ccc79aac5260c595d21a52898f4c637e6fa Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Sat, 10 Oct 2026 00:18:00 +1300
Subject: [PATCH 15/24] test(gpu): fail closed on persistence stop failure

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 .../artifacts/cse_install_ubuntu_spec.sh          | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

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 fb606cae318..6457a0936b2 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
@@ -211,7 +211,7 @@ Describe 'cse_install_ubuntu.sh'
         }
 
         runGraceBlackwellTeardown() {
-            local wave="${1:-wave2}" retry="${2:-false}" enabled="${3:-true}"
+            local wave="${1:-wave2}" retry="${2:-false}" enabled="${3:-true}" persistenced_stop_failure="${7:-false}"
             setupGraceBlackwellTeardown "$wave"
             gb_openibd_enabled="$enabled"
             gb_openibd_active=true
@@ -222,8 +222,10 @@ Describe 'cse_install_ubuntu.sh'
             systemctlDisableAndStop() {
                 echo "stop:$1"
                 if [ "$1" = nvidia-persistenced ]; then
-                    gb_persistenced_active=false
-                    gb_persistenced_enabled=false
+                    if [ "$persistenced_stop_failure" != true ]; then
+                        gb_persistenced_active=false
+                        gb_persistenced_enabled=false
+                    fi
                 fi
                 if [ "$1" = openibd ]; then
                     gb_services_stopped=true
@@ -380,6 +382,13 @@ Describe 'cse_install_ubuntu.sh'
             The output should include "openibd_enabled_after_unload=false"
         End
 
+        It 'fails before module unload when nvidia-persistenced cannot be stopped'
+            When call runGraceBlackwellTeardown wave2 false true false false false true
+            The status should be failure
+            The stderr should include "Grace Blackwell service nvidia-persistenced remains active or enabled"
+            The output should not include "rmmod:"
+        End
+
         It 'fails before module unload when openibd cannot be stopped'
             When call runGraceBlackwellTeardown wave2 false true true
             The status should be failure

From def56f1736e9eb83aedfb02943137bee57250aac Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 23 Oct 2026 10:00:43 +1300
Subject: [PATCH 16/24] refactor: separate payload encoding from GPU fix

Remove the bundled compression change and unrelated lint pin so the GPU PR can use a baked helper without depending on PR #9753.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 .github/workflows/golangci-lint.yml       |   2 +-
 hotfix/hotfix_generate.py                 |  22 +---
 hotfix/hotfix_generate_test.py            |  95 +--------------
 parts/linux/cloud-init/nodecustomdata.yml | 108 ++++++++++++-----
 pkg/agent/baker.go                        |  54 +--------
 pkg/agent/baker_test.go                   |  78 +------------
 pkg/agent/bakerapi_test.go                | 136 ----------------------
 pkg/agent/utils.go                        |  14 +--
 pkg/agent/variables.go                    |  65 +++++------
 9 files changed, 120 insertions(+), 454 deletions(-)

diff --git a/.github/workflows/golangci-lint.yml b/.github/workflows/golangci-lint.yml
index d9835c16691..af34812dc6a 100644
--- a/.github/workflows/golangci-lint.yml
+++ b/.github/workflows/golangci-lint.yml
@@ -22,7 +22,7 @@ jobs:
     steps:
       - uses: actions/setup-go@b7ad1dad31e06c5925ef5d2fc7ad053ef454303e # v7.0.0
         with:
-          go-version: '1.27.1'
+          go-version: '1.27'
       - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
         with:
           persist-credentials: false
diff --git a/hotfix/hotfix_generate.py b/hotfix/hotfix_generate.py
index 080e0f20434..2623604bac3 100644
--- a/hotfix/hotfix_generate.py
+++ b/hotfix/hotfix_generate.py
@@ -55,8 +55,8 @@
 SCRIPTS_BEGIN = "# ---- hotfix-scripts: auto-generated ----"
 SCRIPTS_END = "# ---- end hotfix-scripts ----"
 
-# Map from source file paths (relative to artifacts/) to the template variable
-# keys used in nodecustomdata.yml. Only files that appear as write_files entries
+# Map from source file paths (relative to artifacts/) to the GetVariableProperty
+# keys used in nodecustomdata.yml. Only scripts that appear as write_files entries
 # in the traditional section are included.
 SOURCE_TO_VARKEY = {
     # CSE helpers — base (non-distro)
@@ -380,9 +380,7 @@ def parse_write_files_blocks(traditional_lines):
     """Parse write_files blocks from the traditional section.
 
     Each block is either a simple '- path:' entry or an entire conditional
-    block (e.g., {{if IsAzlOSGuard}}...{{end}}) treated as a single unit. File
-    keys may use either the legacy GetVariableProperty form or the literal-file
-    GetCloudInitFileProperties form.
+    block (e.g., {{if IsAzlOSGuard}}...{{end}}) treated as a single unit.
 
     Returns a list of (varkeys_set, lines_list) tuples.
     """
@@ -423,17 +421,9 @@ def parse_write_files_blocks(traditional_lines):
 
         if in_block:
             current_block.append(line)
-            matches = (
-                re.search(
-                    r'GetVariableProperty\s+"cloudInitData"\s+"(\w+)"',
-                    stripped,
-                ),
-                re.search(r'GetCloudInitFileProperties\s+"(\w+)"', stripped),
-            )
-            for match in matches:
-                if match:
-                    current_varkeys.add(match.group(1))
-                    break
+            match = re.search(r'GetVariableProperty\s+"cloudInitData"\s+"(\w+)"', stripped)
+            if match:
+                current_varkeys.add(match.group(1))
 
     if current_block and current_varkeys:
         blocks.append((current_varkeys.copy(), list(current_block)))
diff --git a/hotfix/hotfix_generate_test.py b/hotfix/hotfix_generate_test.py
index 9acdab4f6c1..31ed80b03d3 100644
--- a/hotfix/hotfix_generate_test.py
+++ b/hotfix/hotfix_generate_test.py
@@ -54,16 +54,6 @@
 {{end}}
 """
 
-MIXED_ENCODING_TEMPLATE = """- path: /literal.sh
-  owner: root
-{{GetCloudInitFileProperties "literalFile"}}
-- path: /legacy.sh
-  owner: root
-  encoding: gzip
-  content: !!binary |
-    {{GetVariableProperty "cloudInitData" "legacyFile"}}
-"""
-
 
 class HotfixGenerateTests(unittest.TestCase):
     def test_production_renderer_omits_empty_platform_payloads(self):
@@ -171,10 +161,7 @@ def test_config_modules_are_independently_hotfixable(self):
                 self.assertEqual({variable}, blocks[0][0])
                 self.assertIn(f"- path: {{{{{path_function}}}}}", template)
                 self.assertIn('permissions: "0744"', template)
-                self.assertTrue(
-                    f'GetCloudInitFileProperties "{variable}"' in template
-                    or f'GetVariableProperty "cloudInitData" "{variable}"' in template
-                )
+                self.assertIn("encoding: gzip", template)
                 self.assertNotIn("{{if", template)
 
     def test_mariner_package_update_is_not_embedded_in_custom_data(self):
@@ -247,86 +234,6 @@ def test_parse_write_files_blocks_keeps_distro_chain_together(self):
             blocks[1][0],
         )
 
-    def test_parse_write_files_blocks_accepts_literal_and_legacy_syntax(self):
-        blocks = hotfix_generate.parse_write_files_blocks(
-            MIXED_ENCODING_TEMPLATE.splitlines(keepends=True)
-        )
-
-        self.assertEqual(
-            [
-                ({"literalFile"}, [
-                    "- path: /literal.sh\n",
-                    "  owner: root\n",
-                    '{{GetCloudInitFileProperties "literalFile"}}\n',
-                ]),
-                ({"legacyFile"}, [
-                    "- path: /legacy.sh\n",
-                    "  owner: root\n",
-                    "  encoding: gzip\n",
-                    "  content: !!binary |\n",
-                    '    {{GetVariableProperty "cloudInitData" "legacyFile"}}\n',
-                ]),
-            ],
-            blocks,
-        )
-
-    def test_current_template_literal_entries_are_injected_and_rendered(self):
-        repository = Path(__file__).resolve().parents[1]
-        template_lines = (repository / hotfix_generate.TEMPLATE).read_text().splitlines(
-            keepends=True
-        )
-        _, outer_else, end = hotfix_generate.find_block_boundaries(template_lines)
-        self.assertIsNotNone(outer_else)
-        self.assertIsNotNone(end)
-        traditional = template_lines[outer_else + 1:end]
-        available = set().union(
-            *(keys for keys, _ in hotfix_generate.parse_write_files_blocks(traditional))
-        )
-        selected = {"provisionConfigsGPU"}
-        self.assertIn("provisionConfigsGPU", available)
-
-        hotfix_template = hotfix_generate.build_hotfix_template(selected, traditional)
-        self.assertIn(
-            '{{GetCloudInitFileProperties "provisionConfigsGPU"}}',
-            hotfix_template,
-        )
-
-        with tempfile.TemporaryDirectory() as temp_dir:
-            injected_template = Path(temp_dir) / "nodecustomdata.yml"
-            injected_template.write_text("".join(template_lines))
-            generated_dir = Path(temp_dir) / "generated"
-            with mock.patch.object(
-                hotfix_generate, "TEMPLATE", str(injected_template)
-            ), mock.patch.object(
-                hotfix_generate, "GENERATED_DIR", str(generated_dir)
-            ):
-                hotfix_generate.update_nodecustomdata(selected)
-                injected = injected_template.read_text()
-                marked_block = injected.split(hotfix_generate.SCRIPTS_BEGIN, 1)[1].split(
-                    hotfix_generate.SCRIPTS_END, 1
-                )[0]
-                self.assertIn(
-                    '{{GetCloudInitFileProperties "provisionConfigsGPU"}}',
-                    marked_block,
-                )
-
-                hotfix_generate.write_rendered_payload(selected, traditional)
-
-            for platform in ("ubuntu", "mariner"):
-                rendered = (
-                    generated_dir / f"rendered_nodecustomdata_{platform}.yml"
-                ).read_text()
-                self.assertIn(
-                    "- path: /opt/azure/containers/provision_configs_gpu.sh",
-                    rendered,
-                )
-                self.assertIn('permissions: "0744"', rendered)
-                self.assertIn(r'content: "#!/bin/bash\n', rendered)
-                self.assertIn("configGPUDrivers() {", rendered)
-                self.assertNotIn("GetCloudInitFileProperties", rendered)
-                self.assertNotIn("{{", rendered)
-                self.assertNotIn("encoding: gzip", rendered)
-
     def test_build_hotfix_template_selects_only_requested_blocks(self):
         rendered = hotfix_generate.build_hotfix_template(
             {"provisionSource"},
diff --git a/parts/linux/cloud-init/nodecustomdata.yml b/parts/linux/cloud-init/nodecustomdata.yml
index 46ca391d300..200eeb0ce5f 100644
--- a/parts/linux/cloud-init/nodecustomdata.yml
+++ b/parts/linux/cloud-init/nodecustomdata.yml
@@ -28,147 +28,201 @@ write_files:
 {{if IsAKSCustomCloud}}
 - path: {{GetInitAKSCloudFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "initAKSCloud"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "initAKSCloud"}}
 {{end}}
 
 {{- else }}
 - path: {{GetCSEHelpersScriptFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionSource"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionSource"}}
 
 
 {{if IsACL }}
 - path: {{GetCSEHelpersScriptDistroFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionSourceACL"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionSourceACL"}}
 {{- else if IsAzlOSGuard}}
 - path: {{GetCSEHelpersScriptDistroFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionSourceAzlOSGuard"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionSourceAzlOSGuard"}}
 {{- else if IsMariner}}
 - path: {{GetCSEHelpersScriptDistroFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionSourceMariner"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionSourceMariner"}}
 {{- else if IsFlatcar }}
 - path: {{GetCSEHelpersScriptDistroFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionSourceFlatcar"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionSourceFlatcar"}}
 {{- else }}
 - path: {{GetCSEHelpersScriptDistroFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionSourceUbuntu"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionSourceUbuntu"}}
 {{end}}
 
 {{ if not IsCustomImage -}}
 - path: /opt/azure/containers/provision_start.sh
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionStartScript"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionStartScript"}}
 {{- end }}
 
 - path: /opt/azure/containers/provision.sh
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionScript"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionScript"}}
 
 - path: {{GetCSEInstallScriptFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionInstalls"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionInstalls"}}
 
 {{if IsACL }}
 - path: {{GetCSEInstallScriptDistroFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionInstallsACL"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionInstallsACL"}}
 {{- else if IsAzlOSGuard}}
 - path: {{GetCSEInstallScriptDistroFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionInstallsAzlOSGuard"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionInstallsAzlOSGuard"}}
 {{- else if IsMariner}}
 - path: {{GetCSEInstallScriptDistroFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionInstallsMariner"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionInstallsMariner"}}
 {{- else if IsFlatcar }}
 - path: {{GetCSEInstallScriptDistroFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionInstallsFlatcar"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionInstallsFlatcar"}}
 {{- else }}
 - path: {{GetCSEInstallScriptDistroFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionInstallsUbuntu"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionInstallsUbuntu"}}
 {{end}}
 
 - path: {{GetCSEConfigScriptFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionConfigs"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionConfigs"}}
 
 - path: {{GetCSEConfigGPUScriptFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionConfigsGPU"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionConfigsGPU"}}
 
 - path: {{GetCSEConfigLocalDNSScriptFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionConfigsLocalDNS"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionConfigsLocalDNS"}}
 
 - path: {{GetCSEConfigKubeletScriptFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionConfigsKubelet"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionConfigsKubelet"}}
 
 - path: {{GetCSEConfigNetworkScriptFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionConfigsNetwork"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionConfigsNetwork"}}
 
 - path: {{GetCSEConfigAddonsScriptFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionConfigsAddons"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionConfigsAddons"}}
 
 - path: {{GetCSEConfigChronyScriptFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "provisionConfigsChrony"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "provisionConfigsChrony"}}
 
 {{if IsAKSCustomCloud}}
 - path: {{GetInitAKSCloudFilepath}}
   permissions: "0744"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "initAKSCloud"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "initAKSCloud"}}
 {{end}}
 
 - path: /etc/systemd/system/reconcile-private-hosts.service
   permissions: "0644"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "reconcilePrivateHostsService"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "reconcilePrivateHostsService"}}
 
 - path: /etc/systemd/system/kubelet.service
   permissions: "0600"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "kubeletSystemdService"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "kubeletSystemdService"}}
 
 - path: /opt/azure-network/configure-azure-network.sh
   permissions: "0755"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "configureAzureNetworkScript"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "configureAzureNetworkScript"}}
 
 - path: /etc/udev/rules.d/99-azure-network.rules
   permissions: "0644"
+  encoding: gzip
   owner: root
-{{GetCloudInitFileProperties "azureNetworkUdevRule"}}
+  content: !!binary |
+    {{GetVariableProperty "cloudInitData" "azureNetworkUdevRule"}}
 {{- end }}
diff --git a/pkg/agent/baker.go b/pkg/agent/baker.go
index 31c4ca797bc..91eacc7f9fb 100644
--- a/pkg/agent/baker.go
+++ b/pkg/agent/baker.go
@@ -350,13 +350,13 @@ func (t *TemplateGenerator) getLinuxNodeCustomDataJSONObject(config *datamodel.N
 	parameters := getParameters(config)
 	// get variable cloudInit
 	variables := getCustomDataVariables(config)
-	str, e := t.getSingleLine(kubernetesNodeCustomDataYaml, config.AgentPoolProfile, getBakerFuncMap(config, parameters, variables), true)
+	str, e := t.getSingleLineForTemplate(kubernetesNodeCustomDataYaml, config.AgentPoolProfile, getBakerFuncMap(config, parameters, variables), true)
 
 	if e != nil {
 		panic(e)
 	}
 
-	return getCustomDataJSONObject(str)
+	return fmt.Sprintf("{\"customData\": \"%s\"}", str)
 }
 
 const (
@@ -698,10 +698,6 @@ func RenderLinuxNodeCustomDataTemplate(templateContent []byte, config *datamodel
 func getBakerFuncMap(config *datamodel.NodeBootstrappingConfiguration, params paramsMap, variables paramsMap) template.FuncMap {
 	funcMap := getContainerServiceFuncMap(config)
 
-	funcMap["GetCloudInitFileProperties"] = func(name string) (string, error) {
-		return getCloudInitFileProperties(variables, name)
-	}
-
 	funcMap["GetParameter"] = func(s string) interface{} {
 		if v, ok := params[s].(paramsMap); ok && v != nil {
 			if v["value"] == nil {
@@ -749,52 +745,6 @@ func getBakerFuncMap(config *datamodel.NodeBootstrappingConfiguration, params pa
 	return funcMap
 }
 
-func getCloudInitFileProperties(variables paramsMap, name string) (string, error) {
-	if files, ok := variables["cloudInitFileData"].(paramsMap); ok {
-		if content, ok := files[name].(string); ok {
-			return renderCloudInitFileContent(content)
-		}
-	}
-	if cloudInitData, ok := variables["cloudInitData"].(paramsMap); ok {
-		if encoded, ok := cloudInitData[name].(string); ok {
-			return fmt.Sprintf("  encoding: gzip\n  content: !!binary |\n    %s\n", encoded), nil
-		}
-	}
-	return "", fmt.Errorf("cloud-init file %q is unavailable", name)
-}
-
-func renderCloudInitFileContent(content string) (string, error) {
-	style := yaml.LiteralStyle
-	for _, line := range strings.Split(content, "\n") {
-		if strings.HasPrefix(line, "\t") {
-			style = yaml.DoubleQuotedStyle
-			break
-		}
-	}
-	document := &yaml.Node{
-		Kind: yaml.MappingNode,
-		Tag:  "!!map",
-		Content: []*yaml.Node{
-			{Kind: yaml.ScalarNode, Tag: "!!str", Value: "content"},
-			{Kind: yaml.ScalarNode, Tag: "!!str", Value: content, Style: style},
-		},
-	}
-	var buffer bytes.Buffer
-	encoder := yaml.NewEncoder(&buffer)
-	encoder.SetIndent(2)
-	if err := encoder.Encode(document); err != nil {
-		return "", fmt.Errorf("encode cloud-init file content: %w", err)
-	}
-	if err := encoder.Close(); err != nil {
-		return "", fmt.Errorf("close cloud-init file content encoder: %w", err)
-	}
-	lines := strings.Split(strings.TrimSuffix(buffer.String(), "\n"), "\n")
-	for index := range lines {
-		lines[index] = "  " + lines[index]
-	}
-	return strings.Join(lines, "\n") + "\n", nil
-}
-
 /* normalizeResourceGroupNameForLabel normalizes resource group name to be used as a label,
 similar to what the ARM template used to do.
 
diff --git a/pkg/agent/baker_test.go b/pkg/agent/baker_test.go
index 13e857067ce..d0d34b9e0cd 100644
--- a/pkg/agent/baker_test.go
+++ b/pkg/agent/baker_test.go
@@ -27,7 +27,6 @@ import (
 	. "github.com/onsi/gomega"
 	"github.com/stretchr/testify/require"
 	"github.com/vincent-petithory/dataurl"
-	"gopkg.in/yaml.v3"
 )
 
 // this regex looks for groups of the following forms, returning KEY and VALUE as submatches.
@@ -103,72 +102,6 @@ write_files:
 	}
 }
 
-func TestRenderCloudInitFileContentRoundTrips(t *testing.T) {
-	testCases := []struct {
-		name    string
-		content string
-	}{
-		{name: "empty", content: ""},
-		{name: "no trailing newline", content: "echo ready"},
-		{name: "one trailing newline", content: "echo ready\n"},
-		{name: "multiple trailing newlines", content: "echo ready\n\n"},
-		{name: "YAML-significant text", content: "---\n# shell text\nkey: value\n"},
-		{name: "tabs and indentation", content: "\tcommand\n  indented\n"},
-		{name: "Unicode", content: "printf 'héllo 🌍'\n"},
-	}
-
-	for _, testCase := range testCases {
-		t.Run(testCase.name, func(t *testing.T) {
-			content, err := renderCloudInitFileContent(testCase.content)
-			require.NoError(t, err)
-
-			customData := "#cloud-config\nwrite_files:\n- path: /opt/script.sh\n  owner: root\n" + content
-			var parsed cloudInit
-			require.NoError(t, yaml.Unmarshal([]byte(customData), &parsed))
-			require.Len(t, parsed.WriteFiles, 1)
-			require.Equal(t, testCase.content, parsed.WriteFiles[0].Content)
-		})
-	}
-}
-
-func TestGetCloudInitFilePropertiesKeepsGzipFallback(t *testing.T) {
-	testCases := []struct {
-		name                   string
-		distro                 datamodel.Distro
-		enableScriptlessCSECmd bool
-	}{
-		{name: "Flatcar", distro: datamodel.AKSFlatcarGen2},
-		{name: "ACL", distro: datamodel.AKSACLGen2TL},
-		{name: "scriptless", distro: datamodel.AKSUbuntuContainerd2204Gen2, enableScriptlessCSECmd: true},
-	}
-
-	for _, testCase := range testCases {
-		t.Run(testCase.name, func(t *testing.T) {
-			config := newNodeCustomDataRenderConfig(testCase.distro)
-			config.EnableScriptlessCSECmd = testCase.enableScriptlessCSECmd
-			variables := getCustomDataVariables(config)
-			require.NotContains(t, variables, "cloudInitFileData")
-
-			funcMap := getBakerFuncMap(config, getParameters(config), variables)
-			getProperties, ok := funcMap["GetCloudInitFileProperties"].(func(string) (string, error))
-			require.True(t, ok)
-			properties, err := getProperties("provisionSource")
-			require.NoError(t, err)
-			require.Contains(t, properties, "encoding: gzip")
-
-			customData := "#cloud-config\nwrite_files:\n- path: /opt/script.sh\n  owner: root\n" + properties
-			var parsed cloudInit
-			require.NoError(t, yaml.Unmarshal([]byte(customData), &parsed))
-			require.Len(t, parsed.WriteFiles, 1)
-			require.Equal(t, "gzip", parsed.WriteFiles[0].Encoding)
-
-			actual, err := getGzipDecodedValue([]byte(parsed.WriteFiles[0].Content))
-			require.NoError(t, err)
-			require.Equal(t, getCustomScriptContent(kubernetesCSEHelpersScript, config), string(actual))
-		})
-	}
-}
-
 func newNodeCustomDataRenderConfig(distro datamodel.Distro) *datamodel.NodeBootstrappingConfiguration {
 	profile := &datamodel.AgentPoolProfile{
 		Name:   "hotfix-render-test",
@@ -1618,15 +1551,8 @@ var _ = Describe("getLinuxNodeBootstrappingPayload", func() {
 		expectedCustomData := getCustomDataFromJSON(templateGenerator.getLinuxNodeCustomDataJSONObject(config))
 
 		Expect(string(decompressedPayload)).To(Equal(expectedCustomData))
-		var parsedCustomData cloudInit
-		Expect(yaml.Unmarshal(decompressedPayload, &parsedCustomData)).To(Succeed())
-		for _, file := range parsedCustomData.WriteFiles {
-			Expect(file.Path).NotTo(Equal(aksNodeCustomDataFilepath))
-			Expect(file.Path).NotTo(Equal(aksNbcCmdFilepath))
-		}
-		for _, command := range parsedCustomData.BootCommands {
-			Expect(command).NotTo(ContainSubstring(aksNbcCmdFilepath))
-		}
+		Expect(string(decompressedPayload)).NotTo(ContainSubstring(aksNodeCustomDataFilepath))
+		Expect(string(decompressedPayload)).NotTo(ContainSubstring(aksNbcCmdFilepath))
 	})
 })
 
diff --git a/pkg/agent/bakerapi_test.go b/pkg/agent/bakerapi_test.go
index 194d445e65d..d1bd44ff392 100644
--- a/pkg/agent/bakerapi_test.go
+++ b/pkg/agent/bakerapi_test.go
@@ -1,19 +1,15 @@
 package agent
 
 import (
-	"bytes"
-	"compress/gzip"
 	"context"
 	"encoding/base64"
 	"fmt"
-	"io"
 
 	"github.com/Azure/agentbaker/pkg/agent/datamodel"
 	agenttoggles "github.com/Azure/agentbaker/pkg/agent/toggles"
 	"github.com/barkimedes/go-deepcopy"
 	. "github.com/onsi/ginkgo"
 	. "github.com/onsi/gomega"
-	"gopkg.in/yaml.v3"
 )
 
 type testToggles struct {
@@ -246,138 +242,6 @@ var _ = Describe("AgentBaker API implementation tests", func() {
 			Expect(nodeBootStrapping.SigImageConfig.Definition).To(Equal("2204gen2containerd"))
 		})
 
-		It("should render Linux cloud-init files as literal write_files content without changing their bytes", func() {
-			testCases := []struct {
-				name   string
-				distro datamodel.Distro
-			}{
-				{name: "Ubuntu 22.04", distro: datamodel.AKSUbuntuContainerd2204Gen2},
-				{name: "Ubuntu 24.04", distro: datamodel.AKSUbuntuContainerd2404Gen2},
-				{name: "Azure Linux 3", distro: datamodel.AKSAzureLinuxV3Gen2},
-			}
-
-			for _, testCase := range testCases {
-				for _, preProvisionOnly := range []bool{false, true} {
-					name := testCase.name
-					if preProvisionOnly {
-						name += " PIS"
-					} else {
-						name += " ordinary"
-					}
-					By(name)
-
-					configCopy, err := deepcopy.Anything(config)
-					Expect(err).NotTo(HaveOccurred())
-					testConfig, ok := configCopy.(*datamodel.NodeBootstrappingConfiguration)
-					Expect(ok).To(BeTrue())
-					testConfig.AgentPoolProfile.Distro = testCase.distro
-					testConfig.AgentPoolProfile.VMSize = "Standard_NC4as_T4_v3"
-					testConfig.PreProvisionOnly = preProvisionOnly
-					testConfig.EnableScriptlessCSECmd = false
-					testConfig.ConfigGPUDriverIfNeeded = false
-					testConfig.EnableGPUDevicePluginIfNeeded = false
-					testConfig.EnableNvidia = false
-
-					payload := InitializeTemplateGenerator().getLinuxNodeBootstrappingPayload(testConfig)
-					Expect(len(payload)).To(BeNumerically("<", MaxCustomDataLength))
-					encoded, err := base64.StdEncoding.Strict().DecodeString(payload)
-					Expect(err).NotTo(HaveOccurred())
-					Expect(len(encoded)).To(BeNumerically("<=", 65535))
-					reader, err := gzip.NewReader(bytes.NewReader(encoded))
-					Expect(err).NotTo(HaveOccurred())
-					customDataBytes, err := io.ReadAll(reader)
-					Expect(err).NotTo(HaveOccurred())
-					Expect(reader.Close()).To(Succeed())
-					var customData cloudInit
-					Expect(yaml.Unmarshal(customDataBytes, &customData)).To(Succeed())
-					fmt.Fprintf(
-						GinkgoWriter,
-						"PIS_CUSTOMDATA distro=%s preProvisionOnly=%t base64=%d gzip=%d yaml=%d\n",
-						testCase.distro,
-						preProvisionOnly,
-						len(payload),
-						len(encoded),
-						len(customDataBytes),
-					)
-
-					customDataVariables := getCustomDataVariables(testConfig)
-					variables, ok := customDataVariables["cloudInitFileData"].(paramsMap)
-					Expect(ok).To(BeTrue())
-					expectedScripts := []struct {
-						path string
-						key  string
-					}{
-						{cseHelpersScriptFilepath, "provisionSource"},
-						{cseHelpersScriptDistroFilepath, "provisionSourceUbuntu"},
-						{"/opt/azure/containers/provision_start.sh", "provisionStartScript"},
-						{"/opt/azure/containers/provision.sh", "provisionScript"},
-						{cseInstallScriptFilepath, "provisionInstalls"},
-						{cseInstallScriptDistroFilepath, "provisionInstallsUbuntu"},
-						{cseConfigScriptFilepath, "provisionConfigs"},
-						{cseConfigGPUScriptFilepath, "provisionConfigsGPU"},
-						{cseConfigLocalDNSScriptFilepath, "provisionConfigsLocalDNS"},
-						{cseConfigKubeletScriptFilepath, "provisionConfigsKubelet"},
-						{cseConfigNetworkScriptFilepath, "provisionConfigsNetwork"},
-						{cseConfigAddonsScriptFilepath, "provisionConfigsAddons"},
-						{cseConfigChronyScriptFilepath, "provisionConfigsChrony"},
-					}
-					if testCase.distro == datamodel.AKSAzureLinuxV3Gen2 {
-						expectedScripts[1].key = "provisionSourceMariner"
-						expectedScripts[5].key = "provisionInstallsMariner"
-					}
-
-					expectedPaths := make([]string, 0, len(customData.WriteFiles))
-					for _, script := range expectedScripts {
-						expectedPaths = append(expectedPaths, script.path)
-					}
-					expectedPaths = append(expectedPaths,
-						"/etc/systemd/system/reconcile-private-hosts.service",
-						"/etc/systemd/system/kubelet.service",
-						"/opt/azure-network/configure-azure-network.sh",
-						"/etc/udev/rules.d/99-azure-network.rules",
-					)
-					actualPaths := make([]string, 0, len(customData.WriteFiles))
-					for _, file := range customData.WriteFiles {
-						actualPaths = append(actualPaths, file.Path)
-					}
-					Expect(actualPaths).To(Equal(expectedPaths))
-
-					filesByPath := make(map[string]cloudInitWriteFile, len(customData.WriteFiles))
-					for _, file := range customData.WriteFiles {
-						filesByPath[file.Path] = file
-						Expect(file.Owner).To(Equal("root"))
-					}
-					for _, script := range expectedScripts {
-						file := filesByPath[script.path]
-						Expect(file.Permissions).To(Equal("0744"))
-						Expect(file.Encoding).To(BeEmpty())
-						expectedContent, ok := variables[script.key].(string)
-						Expect(ok).To(BeTrue())
-						Expect(file.Content).To(Equal(expectedContent), script.path)
-					}
-					expectedLiteralFiles := []struct {
-						path        string
-						key         string
-						permissions string
-					}{
-						{"/etc/systemd/system/reconcile-private-hosts.service", "reconcilePrivateHostsService", "0644"},
-						{"/etc/systemd/system/kubelet.service", "kubeletSystemdService", "0600"},
-						{"/opt/azure-network/configure-azure-network.sh", "configureAzureNetworkScript", "0755"},
-						{"/etc/udev/rules.d/99-azure-network.rules", "azureNetworkUdevRule", "0644"},
-					}
-					for _, expected := range expectedLiteralFiles {
-						file := filesByPath[expected.path]
-						Expect(file.Permissions).To(Equal(expected.permissions))
-						Expect(file.Encoding).To(BeEmpty())
-						expectedContent, ok := variables[expected.key].(string)
-						Expect(ok).To(BeTrue())
-						Expect(file.Content).To(Equal(expectedContent), expected.path)
-					}
-					By(fmt.Sprintf("%s customData size: encoded=%d decoded=%d", name, len(payload), len(customDataBytes)))
-				}
-			}
-		})
-
 		It("should accept supported custom Linux transparent huge page values", func() {
 			config.AgentPoolProfile.CustomLinuxOSConfig = &datamodel.CustomLinuxOSConfig{
 				TransparentHugePageEnabled: "never",
diff --git a/pkg/agent/utils.go b/pkg/agent/utils.go
index d67838fd8cb..7b6d4072fa9 100644
--- a/pkg/agent/utils.go
+++ b/pkg/agent/utils.go
@@ -181,10 +181,6 @@ func escapeSingleLine(escapedStr string) string {
 
 // getBase64EncodedGzippedCustomScript will return a base64 of the CSE.
 func getBase64EncodedGzippedCustomScript(csFilename string, config *datamodel.NodeBootstrappingConfiguration) string {
-	return getBase64EncodedGzippedCustomScriptFromStr(getCustomScriptContent(csFilename, config))
-}
-
-func getCustomScriptContent(csFilename string, config *datamodel.NodeBootstrappingConfiguration) string {
 	b, err := parts.Templates.ReadFile(csFilename)
 	if err != nil {
 		// this should never happen and this is a bug.
@@ -206,7 +202,7 @@ func getCustomScriptContent(csFilename string, config *datamodel.NodeBootstrappi
 	}
 	csStr := buffer.String()
 	csStr = strings.ReplaceAll(csStr, "\r\n", "\n")
-	return csStr
+	return getBase64EncodedGzippedCustomScriptFromStr(csStr)
 }
 
 // This is "best-effort" - removes MOST of the comments with obvious formats, to lower the space required by CustomData component.
@@ -396,14 +392,6 @@ func getCustomDataFromJSON(jsonStr string) string {
 	return customDataObj["customData"]
 }
 
-func getCustomDataJSONObject(customData string) string {
-	customDataJSON, err := json.Marshal(map[string]string{"customData": customData})
-	if err != nil {
-		panic(fmt.Errorf("failed to marshal customData: %w", err))
-	}
-	return string(customDataJSON)
-}
-
 // GetOrderedKubeletConfigFlagString returns an ordered string of key/val pairs.
 // copied from AKS-Engine and filter out flags that already translated to config file.
 func GetOrderedKubeletConfigFlagString(config *datamodel.NodeBootstrappingConfiguration) string {
diff --git a/pkg/agent/variables.go b/pkg/agent/variables.go
index 41631d5faef..7efc595a410 100644
--- a/pkg/agent/variables.go
+++ b/pkg/agent/variables.go
@@ -13,47 +13,37 @@ import (
 // getCustomDataVariables returns cloudinit data used by Linux.
 func getCustomDataVariables(config *datamodel.NodeBootstrappingConfiguration) paramsMap {
 	cs := config.ContainerService
-	useCloudInitLiteralFiles := !config.IsFlatcar() && !config.IsACL() && !config.EnableScriptlessCSECmd
-	cloudInitFileData := paramsMap{}
-	getCloudInitFile := func(name, filename string) string {
-		content := getCustomScriptContent(filename, config)
-		if useCloudInitLiteralFiles {
-			cloudInitFileData[name] = content
-			return ""
-		}
-		return getBase64EncodedGzippedCustomScriptFromStr(content)
-	}
 	cloudInitFiles := map[string]interface{}{
 		"cloudInitData": paramsMap{
-			"provisionStartScript":                  getCloudInitFile("provisionStartScript", kubernetesCSEStartScript),
-			"provisionScript":                       getCloudInitFile("provisionScript", kubernetesCSEMainScript),
-			"provisionSource":                       getCloudInitFile("provisionSource", kubernetesCSEHelpersScript),
-			"provisionSourceUbuntu":                 getCloudInitFile("provisionSourceUbuntu", kubernetesCSEHelpersScriptUbuntu),
-			"provisionSourceMariner":                getCloudInitFile("provisionSourceMariner", kubernetesCSEHelpersScriptMariner),
-			"provisionSourceAzlOSGuard":             getCloudInitFile("provisionSourceAzlOSGuard", kubernetesCSEHelpersScriptAzlOSGuard),
-			"provisionSourceFlatcar":                getCloudInitFile("provisionSourceFlatcar", kubernetesCSEHelpersScriptFlatcar),
-			"provisionSourceACL":                    getCloudInitFile("provisionSourceACL", kubernetesCSEHelpersScriptACL),
-			"provisionInstalls":                     getCloudInitFile("provisionInstalls", kubernetesCSEInstall),
-			"provisionInstallsUbuntu":               getCloudInitFile("provisionInstallsUbuntu", kubernetesCSEInstallUbuntu),
-			"provisionInstallsMariner":              getCloudInitFile("provisionInstallsMariner", kubernetesCSEInstallMariner),
-			"provisionInstallsAzlOSGuard":           getCloudInitFile("provisionInstallsAzlOSGuard", kubernetesCSEInstallAzlOSGuard),
-			"provisionInstallsFlatcar":              getCloudInitFile("provisionInstallsFlatcar", kubernetesCSEInstallFlatcar),
-			"provisionInstallsACL":                  getCloudInitFile("provisionInstallsACL", kubernetesCSEInstallACL),
-			"provisionConfigs":                      getCloudInitFile("provisionConfigs", kubernetesCSEConfig),
-			"provisionConfigsGPU":                   getCloudInitFile("provisionConfigsGPU", kubernetesCSEConfigGPU),
-			"provisionConfigsLocalDNS":              getCloudInitFile("provisionConfigsLocalDNS", kubernetesCSEConfigLocalDNS),
-			"provisionConfigsKubelet":               getCloudInitFile("provisionConfigsKubelet", kubernetesCSEConfigKubelet),
-			"provisionConfigsNetwork":               getCloudInitFile("provisionConfigsNetwork", kubernetesCSEConfigNetwork),
-			"provisionConfigsAddons":                getCloudInitFile("provisionConfigsAddons", kubernetesCSEConfigAddons),
-			"provisionConfigsChrony":                getCloudInitFile("provisionConfigsChrony", kubernetesCSEConfigChrony),
+			"provisionStartScript":                  getBase64EncodedGzippedCustomScript(kubernetesCSEStartScript, config),
+			"provisionScript":                       getBase64EncodedGzippedCustomScript(kubernetesCSEMainScript, config),
+			"provisionSource":                       getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScript, config),
+			"provisionSourceUbuntu":                 getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScriptUbuntu, config),
+			"provisionSourceMariner":                getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScriptMariner, config),
+			"provisionSourceAzlOSGuard":             getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScriptAzlOSGuard, config),
+			"provisionSourceFlatcar":                getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScriptFlatcar, config),
+			"provisionSourceACL":                    getBase64EncodedGzippedCustomScript(kubernetesCSEHelpersScriptACL, config),
+			"provisionInstalls":                     getBase64EncodedGzippedCustomScript(kubernetesCSEInstall, config),
+			"provisionInstallsUbuntu":               getBase64EncodedGzippedCustomScript(kubernetesCSEInstallUbuntu, config),
+			"provisionInstallsMariner":              getBase64EncodedGzippedCustomScript(kubernetesCSEInstallMariner, config),
+			"provisionInstallsAzlOSGuard":           getBase64EncodedGzippedCustomScript(kubernetesCSEInstallAzlOSGuard, config),
+			"provisionInstallsFlatcar":              getBase64EncodedGzippedCustomScript(kubernetesCSEInstallFlatcar, config),
+			"provisionInstallsACL":                  getBase64EncodedGzippedCustomScript(kubernetesCSEInstallACL, config),
+			"provisionConfigs":                      getBase64EncodedGzippedCustomScript(kubernetesCSEConfig, config),
+			"provisionConfigsGPU":                   getBase64EncodedGzippedCustomScript(kubernetesCSEConfigGPU, config),
+			"provisionConfigsLocalDNS":              getBase64EncodedGzippedCustomScript(kubernetesCSEConfigLocalDNS, config),
+			"provisionConfigsKubelet":               getBase64EncodedGzippedCustomScript(kubernetesCSEConfigKubelet, config),
+			"provisionConfigsNetwork":               getBase64EncodedGzippedCustomScript(kubernetesCSEConfigNetwork, config),
+			"provisionConfigsAddons":                getBase64EncodedGzippedCustomScript(kubernetesCSEConfigAddons, config),
+			"provisionConfigsChrony":                getBase64EncodedGzippedCustomScript(kubernetesCSEConfigChrony, config),
 			"provisionSendLogs":                     getBase64EncodedGzippedCustomScript(kubernetesCSESendLogs, config),
 			"provisionRedactCloudConfig":            getBase64EncodedGzippedCustomScript(kubernetesCSERedactCloudConfig, config),
 			"customSearchDomainsScript":             getBase64EncodedGzippedCustomScript(kubernetesCustomSearchDomainsScript, config),
 			"dhcpv6SystemdService":                  getBase64EncodedGzippedCustomScript(dhcpv6SystemdService, config),
 			"dhcpv6ConfigurationScript":             getBase64EncodedGzippedCustomScript(dhcpv6ConfigurationScript, config),
-			"kubeletSystemdService":                 getCloudInitFile("kubeletSystemdService", kubeletSystemdService),
+			"kubeletSystemdService":                 getBase64EncodedGzippedCustomScript(kubeletSystemdService, config),
 			"reconcilePrivateHostsScript":           getBase64EncodedGzippedCustomScript(reconcilePrivateHostsScript, config),
-			"reconcilePrivateHostsService":          getCloudInitFile("reconcilePrivateHostsService", reconcilePrivateHostsService),
+			"reconcilePrivateHostsService":          getBase64EncodedGzippedCustomScript(reconcilePrivateHostsService, config),
 			"ensureNoDupEbtablesScript":             getBase64EncodedGzippedCustomScript(ensureNoDupEbtablesScript, config),
 			"ensureNoDupEbtablesService":            getBase64EncodedGzippedCustomScript(ensureNoDupEbtablesService, config),
 			"bindMountScript":                       getBase64EncodedGzippedCustomScript(bindMountScript, config),
@@ -73,16 +63,13 @@ func getCustomDataVariables(config *datamodel.NodeBootstrappingConfiguration) pa
 			"cloudInitStatusCheckScript":            getBase64EncodedGzippedCustomScript(cloudInitStatusCheckScript, config),
 			"measureTLSBootstrappingLatencyScript":  getBase64EncodedGzippedCustomScript(measureTLSBootstrappingLatencyScript, config),
 			"measureTLSBootstrappingLatencyService": getBase64EncodedGzippedCustomScript(measureTLSBootstrappingLatencyService, config),
-			"configureAzureNetworkScript":           getCloudInitFile("configureAzureNetworkScript", configureAzureNetworkScript),
-			"azureNetworkUdevRule":                  getCloudInitFile("azureNetworkUdevRule", azureNetworkUdevRule),
+			"configureAzureNetworkScript":           getBase64EncodedGzippedCustomScript(configureAzureNetworkScript, config),
+			"azureNetworkUdevRule":                  getBase64EncodedGzippedCustomScript(azureNetworkUdevRule, config),
 		},
 	}
 
 	cloudInitData := cloudInitFiles["cloudInitData"].(paramsMap) //nolint:errcheck // no error is actually here
-	cloudInitData["initAKSCloud"] = getCloudInitFile("initAKSCloud", initAKSCloudScript)
-	if useCloudInitLiteralFiles {
-		cloudInitFiles["cloudInitFileData"] = cloudInitFileData
-	}
+	cloudInitData["initAKSCloud"] = getBase64EncodedGzippedCustomScript(initAKSCloudScript, config)
 
 	if config.IsFlatcar() || config.IsACL() {
 		cloudInitData["provisionRedactCloudConfig"] = "" // Flatcar and ACL do not have cloud-init

From 6358f7f97444b2da07be67633365398d28ebd8f7 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 23 Oct 2026 11:09:58 +1300
Subject: [PATCH 17/24] fix(gpu): bake strict PIS cleanup helper

Keep bake-only cleanup out of CustomData. Opted-out Ubuntu PIS bakes, including CPU-only pools, require a new source image containing the helper and fail explicitly when it is missing. Ordinary and derived nodes do not require the helper.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 parts/linux/cloud-init/artifacts/cse_main.sh  |   7 +
 .../artifacts/ubuntu/cse_install_ubuntu.sh    | 289 +-----------------
 .../ubuntu/cse_install_ubuntu_gpu.sh          | 288 +++++++++++++++++
 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 +
 8 files changed, 323 insertions(+), 288 deletions(-)
 create mode 100755 parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh

diff --git a/parts/linux/cloud-init/artifacts/cse_main.sh b/parts/linux/cloud-init/artifacts/cse_main.sh
index a102dd4497f..fc8fa44727f 100755
--- a/parts/linux/cloud-init/artifacts/cse_main.sh
+++ b/parts/linux/cloud-init/artifacts/cse_main.sh
@@ -409,6 +409,13 @@ EOF
 
     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
 
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 6e9acaabb88..0d86742de72 100755
--- a/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
+++ b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu.sh
@@ -369,7 +369,7 @@ cleanUpPrebakedGPUDriver() {
     local dkms_after=false modprobe_after=false marker_after=true status=cleaned
     [ -d /var/lib/dkms/nvidia ] && dkms_after=true
     [ -e /usr/bin/nvidia-modprobe ] && modprobe_after=true
-    if ! prebakedGPUDriverArtifactsRemain; then
+    if [ "${dkms_after}" = false ] && [ "${modprobe_after}" = false ] && [ "${module_after}" = false ]; then
         rm -f "${marker}" || true
         [ -f "${marker}" ] || marker_after=false
     fi
@@ -379,222 +379,6 @@ cleanUpPrebakedGPUDriver() {
     echo "AKS_GPU_PREBAKE event=teardown gpu_node=${GPU_NODE:-} status=${status} dkms_before=${dkms_before} module_before=${module_before} module_after=${module_after} marker_after=${marker_after} dkms_after=${dkms_after} modprobe_after=${modprobe_after}"
 }
 
-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}"
-}
-
-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
-}
-
-cleanUpGraceBlackwellGPUDriver() {
-    local entries="${1}" bom="${2}" wave="${3}" pending="${4:-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
-}
-
 cleanUpGPUDrivers() {
     rm -Rf $GPU_DEST /opt/gpu
 
@@ -609,77 +393,6 @@ cleanUpGPUDrivers() {
     cleanUpPrebakedGPUDriver
 }
 
-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}" "${bom}" "${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
-}
-
 installCriCtlPackage() {
     local version="${1:-}"
     local fullPackageVersion
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..c45a6a41076
--- /dev/null
+++ b/parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh
@@ -0,0 +1,288 @@
+#!/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}"
+}
+
+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
+}
+
+cleanUpGraceBlackwellGPUDriver() {
+    local entries="${1}" bom="${2}" wave="${3}" pending="${4:-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
+}
+
+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}" "${bom}" "${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 44a32731653..07a7bb93ebf 100644
--- a/vhdbuilder/packer/packer_source.sh
+++ b/vhdbuilder/packer/packer_source.sh
@@ -189,6 +189,12 @@ configureServiceExecutionTelemetryForSystemdVersion() {
   done
 }
 
+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
@@ -487,6 +493,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 ff45bef598a1d02ce32806a857cf25f7f474eb44 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 23 Oct 2026 11:09:58 +1300
Subject: [PATCH 18/24] fix(gpu): scope AMD payload to Azure Linux

Omit Azure Linux-only AMD installation code from other known distro payloads. Preserve raw VHD scripts and unknown custom-image implementations.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 parts/linux/cloud-init/artifacts/cse_config_gpu.sh | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/parts/linux/cloud-init/artifacts/cse_config_gpu.sh b/parts/linux/cloud-init/artifacts/cse_config_gpu.sh
index 3ad27398952..a807499d2c4 100755
--- a/parts/linux/cloud-init/artifacts/cse_config_gpu.sh
+++ b/parts/linux/cloud-init/artifacts/cse_config_gpu.sh
@@ -231,6 +231,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.
@@ -251,6 +252,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
@@ -260,6 +262,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
@@ -325,6 +330,7 @@ setupAmdAma() {
         if [ "$(systemctl is-active kubelet)" = "active" ]; then
             systemctl restart kubelet
         fi
+        #{{end}}
     fi
 }
 

From 1976e57cbbdb4edd5c929411bc25455b6f109ee3 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 23 Oct 2026 11:09:58 +1300
Subject: [PATCH 19/24] test(gpu): cover baked helper delivery

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 pkg/agent/gpu_pis_delivery_test.go            | 183 ++++++++++++++++++
 .../artifacts/cse_install_ubuntu_spec.sh      |  27 +--
 .../artifacts/cse_main_gpu_pis_spec.sh        |  41 +++-
 spec/vhdbuilder/packer/pis_gpu_helper_spec.sh |  37 ++++
 4 files changed, 261 insertions(+), 27 deletions(-)
 create mode 100644 pkg/agent/gpu_pis_delivery_test.go
 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..a7e13ad065e
--- /dev/null
+++ b/pkg/agent/gpu_pis_delivery_test.go
@@ -0,0 +1,183 @@
+package agent
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/Azure/agentbaker/pkg/agent/datamodel"
+	"github.com/stretchr/testify/require"
+	"gopkg.in/yaml.v3"
+)
+
+func TestPISGPUBakedHelperDelivery(t *testing.T) {
+	for _, distro := range []datamodel.Distro{datamodel.AKSUbuntuContainerd2204Gen2, datamodel.AKSUbuntuContainerd2404Gen2} {
+		for _, mode := range []struct {
+			preProvisionOnly bool
+			phase1           bool
+		}{{true, true}, {true, false}, {false, false}} {
+			t.Run(fmt.Sprintf("%s/PreProvisionOnly=%t/Phase1=%t", distro, mode.preProvisionOnly, mode.phase1), func(t *testing.T) {
+				pool := &datamodel.AgentPoolProfile{
+					Name: "agent2", VMSize: "Standard_NC4as_T4_v3", StorageProfile: "ManagedDisks",
+					OSType: datamodel.Linux, AvailabilityProfile: datamodel.VirtualMachineScaleSets,
+					VnetSubnetID: "/subscriptions/test/resourceGroups/test/providers/Microsoft.Network/virtualNetworks/test/subnets/default",
+					Distro:       distro,
+				}
+				properties := &datamodel.Properties{
+					OrchestratorProfile:     &datamodel.OrchestratorProfile{OrchestratorType: datamodel.Kubernetes, OrchestratorVersion: "1.32.1", KubernetesConfig: &datamodel.KubernetesConfig{}},
+					HostedMasterProfile:     &datamodel.HostedMasterProfile{DNSPrefix: "fixture"},
+					LinuxProfile:            &datamodel.LinuxProfile{AdminUsername: "azureuser"},
+					ServicePrincipalProfile: &datamodel.ServicePrincipalProfile{ClientID: "fixture", Secret: "fixture"},
+					AgentPoolProfiles:       []*datamodel.AgentPoolProfile{pool},
+				}
+				properties.LinuxProfile.SSH.PublicKeys = []datamodel.PublicKey{{KeyData: "fixture-key"}}
+				config := &datamodel.NodeBootstrappingConfiguration{
+					ContainerService: &datamodel.ContainerService{Location: "southcentralus", Type: "Microsoft.ContainerService/ManagedClusters", Properties: properties},
+					CloudSpecConfig:  datamodel.AzurePublicCloudSpecForTest, K8sComponents: &datamodel.K8sComponents{},
+					AgentPoolProfile: pool, TenantID: "tenantID", SubscriptionID: "subID",
+					ResourceGroupName: "resourceGroupName", UserAssignedIdentityClientID: "userAssignedID",
+					KubeletConfig:       map[string]string{"--cluster-domain": "cluster.local", "--cluster-dns": "10.0.0.10", "--cloud-provider": "azure", "--max-pods": "110"},
+					PrimaryScaleSetName: "aks-agent2-36873793-vmss", PreProvisionOnly: mode.preProvisionOnly,
+					EnableNvidia: false, ConfigGPUDriverIfNeeded: false,
+					EnableGPUDevicePluginIfNeeded: false, EnableScriptlessCSECmd: mode.phase1,
+				}
+
+				require.False(t, supportsScriptlessProvisioning(config))
+				generator := InitializeTemplateGenerator()
+				command := generator.getNodeBootstrappingCmd(config)
+				require.Contains(t, command, fmt.Sprintf(`PRE_PROVISION_ONLY="%t"`, mode.preProvisionOnly))
+				require.Contains(t, command, "GPU_NODE=false")
+				require.Contains(t, command, `/bin/bash /opt/azure/containers/provision_start.sh`)
+				require.NotContains(t, command, "aks-node-controller provision-wait")
+
+				payload := generator.getLinuxNodeBootstrappingPayload(config)
+				compressed, err := base64.StdEncoding.DecodeString(payload)
+				require.NoError(t, err)
+				decoded, err := getGzipDecodedValue(compressed)
+				require.NoError(t, err)
+				t.Logf("PIS T4 payload: yaml=%d gzip=%d base64=%d gzip-margin=%d base64-margin=%d", len(decoded), len(compressed), len(payload), 65535-len(compressed), 87380-len(payload))
+				require.LessOrEqual(t, len(payload), 87380)
+				require.LessOrEqual(t, len(compressed), 65535)
+				var document cloudInit
+				require.NoError(t, yaml.Unmarshal(decoded, &document))
+				installerFound := false
+				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" {
+						installerFound = true
+						script, err := getGzipDecodedValue([]byte(file.Content))
+						require.NoError(t, err)
+						require.NotContains(t, string(script), "graceBlackwellConfigIsUnmodified()")
+						require.NotContains(t, string(script), "cleanUpGPUDriversForBasePrep()")
+						require.Contains(t, string(script), "cleanUpGPUDrivers()")
+					}
+				}
+				if mode.phase1 {
+					require.Len(t, document.WriteFiles, 1)
+					require.Equal(t, "/opt/azure/containers/scriptless-cse-overrides.txt", document.WriteFiles[0].Path)
+				} else {
+					require.Greater(t, len(document.WriteFiles), 1)
+					require.True(t, installerFound)
+				}
+			})
+		}
+	}
+}
+
+func TestPISGPUBakedHelperProvenance(t *testing.T) {
+	const helper = "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh"
+	root := filepath.Join("..", "..")
+	for _, variant := range []string{"base", "cvm", "arm64-gen2", "arm64-gb"} {
+		data, err := os.ReadFile(filepath.Join(root, "vhdbuilder", "packer", "vhd-image-builder-"+variant+".json"))
+		require.NoError(t, err)
+		var image struct {
+			Provisioners []struct{ Type, Source, Destination string }
+		}
+		require.NoError(t, json.Unmarshal(data, &image))
+		matches := 0
+		for _, provisioner := range image.Provisioners {
+			if provisioner.Type == "file" && provisioner.Source == helper &&
+				provisioner.Destination == "/home/packer/provision_installs_ubuntu_gpu.sh" {
+				matches++
+			}
+		}
+		require.Equal(t, 1, matches, variant)
+	}
+	// Execute the raw Packer source, not the comment-stripped CustomData copy.
+	config := filepath.Join(t.TempDir(), "nvidia.conf")
+	require.NoError(t, os.WriteFile(config, []byte("fixture"), 0600))
+	command := exec.Command("bash", "-c", `
+source "$1"
+graceBlackwellInstalledDriverPackages() { :; }
+graceBlackwellPackageManagedDriverRemains() { return 1; }
+graceBlackwellConfigIsUnmodified() {
+    printf '%d\n' "$#"
+    printf '%s\n' "$@"
+    return 1
+}
+GB_NVIDIA_MODPROBE_CONFIG_FILE="$2"
+cleanUpGraceBlackwellGPUDriver unused unused unused
+`, "gpu-helper-test", filepath.Join(root, helper), config)
+	// A present config reaches the comparison without executing teardown.
+	output, err := command.CombinedOutput()
+	require.Error(t, err)
+	require.True(t, strings.HasPrefix(string(output), "8\n"), string(output))
+	require.Contains(t, string(output), "# Explicitly load nvidia_peermem at boot for GPUDirect RDMA.")
+}
+
+func TestGPUAMADeliveryIsDistroScoped(t *testing.T) {
+	distros := append([]datamodel.Distro{
+		datamodel.AKSUbuntuContainerd2204Gen2,
+		datamodel.AKSUbuntuContainerd2404Gen2,
+		datamodel.AKSACLGen2TL,
+		datamodel.AKSFlatcarGen2,
+		datamodel.CustomizedImage,
+		datamodel.CustomizedImageKata,
+		datamodel.CustomizedImageLinuxGuard,
+		datamodel.CustomizedImageTrustedLaunch,
+	}, datamodel.AvailableAzureLinuxDistros...)
+	for _, distro := range distros {
+		t.Run(string(distro), func(t *testing.T) {
+			config := newNodeCustomDataRenderConfig(distro)
+			encoded := getBase64EncodedGzippedCustomScript(kubernetesCSEConfigGPU, config)
+			compressed, err := base64.StdEncoding.DecodeString(encoded)
+			require.NoError(t, err)
+			script, err := getGzipDecodedValue(compressed)
+			require.NoError(t, err)
+			custom := distro == datamodel.CustomizedImage || distro == datamodel.CustomizedImageKata ||
+				distro == datamodel.CustomizedImageLinuxGuard || distro == datamodel.CustomizedImageTrustedLaunch
+			keepImplementation := distro.IsAzureLinuxDistro() || custom
+			require.Equal(t, keepImplementation, strings.Contains(string(script), "dnf_install_amd_ama_core_package()"))
+			require.Equal(t, keepImplementation, strings.Contains(string(script), `dnf repoquery -y --available "amd-ama-driver-*"`))
+			require.Contains(t, string(script), "setupAmdAma()")
+			syntax := exec.Command("bash", "-n")
+			syntax.Stdin = strings.NewReader(string(script))
+			output, err := syntax.CombinedOutput()
+			require.NoError(t, err, string(output))
+			if keepImplementation {
+				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))
+				rendered := exec.Command("bash", "-c", string(script)+"\ndeclare -f setupAmdAma dnf_install_amd_ama_core_package")
+				renderedFunctions, err := rendered.CombinedOutput()
+				require.NoError(t, err, string(renderedFunctions))
+				require.Equal(t, string(rawFunctions), string(renderedFunctions))
+			} else {
+				command := exec.Command("bash", "-c", string(script)+`
+isARM64() { echo 0; }
+isAzureLinux() { return 1; }
+setupAmdAma
+`)
+				output, err := command.CombinedOutput()
+				require.NoError(t, err, string(output))
+				require.Empty(t, output)
+			}
+		})
+	}
+}
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 6457a0936b2..2dc65af836f 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
@@ -86,35 +86,10 @@ Describe 'cse_install_ubuntu.sh'
             The output should include "status=incomplete"
         End
 
-        It 'keeps the marker if any prebaked userspace artifact remains'
-            marker="$(mktemp)"
-            GPU_DKMS_MARKER_FILE="$marker"
-            AfterRun 'command rm -f "$marker"'
-            rm() { echo "mock rm $*"; }
-            ldconfig() { :; }
-            lsmod() { :; }
-            prebakedGPUDriverArtifactsRemain() { return 0; }
-            When call cleanUpPrebakedGPUDriver
-            The status should be success
-            The output should include "status=incomplete"
-            The output should not include "mock rm -f $marker"
-        End
-
-        It 'keeps the marker when loaded-module inspection fails'
-            marker="$(mktemp)"
-            GPU_DKMS_MARKER_FILE="$marker"
-            AfterRun 'command rm -f "$marker"'
-            rm() { echo "mock rm $*"; }
-            ldconfig() { :; }
-            lsmod() { return 1; }
-            When call cleanUpPrebakedGPUDriver
-            The status should be success
-            The output should include "status=incomplete"
-            The output should not include "mock rm -f $marker"
-        End
     End
 
     Describe 'cleanUpGPUDriversForBasePrep'
+        Include "./parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh"
         GPU_DEST="${PWD}/.shellspec-nonexistent-gpu-dest"
         managedGPUPackageList() { :; }
 
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
index eae39e54b9c..58a2a6b5bf8 100644
--- 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
@@ -12,6 +12,7 @@ Describe 'GPU cleanup across PIS phases'
         PRE_PROVISION_ONLY=false
         skip_nvidia_driver_install=false
         BASE_PREP_RAN=false
+        printf '%s\n' 'cleanUpGPUDriversForBasePrep() { echo "basePrep driver cleanup"; return "${CLEANUP_STATUS:-0}"; }' > "${fixture}/provision_installs_ubuntu_gpu.sh"
     }
 
     run_phases() {
@@ -35,7 +36,8 @@ Describe 'GPU cleanup across PIS phases'
 
     run_base_prep_cleanup() {
         local cleanup
-        cleanup=$(sed -n '/    if \[ "${PRE_PROVISION_ONLY}" = "true" \] && \[ "${GPU_NODE}" != "true" \] &&/,/    fi/p' "$CSE_MAIN")
+        cleanup=$(sed -n '/    if \[ "${PRE_PROVISION_ONLY}" = "true" \] && \[ "${GPU_NODE}" != "true" \] &&/,/^    fi$/p' "$CSE_MAIN")
+        cleanup=${cleanup//\/opt\/azure\/containers\/provision_installs_ubuntu_gpu.sh/$fixture/provision_installs_ubuntu_gpu.sh}
         logs_to_events() { shift; "$@"; }
         cleanUpGPUDriversForBasePrep() { echo "basePrep driver cleanup"; return "${CLEANUP_STATUS:-0}"; }
         isMarinerOrAzureLinux() { [ "$OS" = azurelinux ]; }
@@ -95,6 +97,43 @@ Describe 'GPU cleanup across PIS phases'
         The output should include "basePrep driver cleanup"
     End
 
+    It 'fails an Ubuntu bake without the required helper before customization'
+        PRE_PROVISION_ONLY=true
+        OS=ubuntu
+        rm "${fixture}/provision_installs_ubuntu_gpu.sh"
+        When call run_base_prep_cleanup
+        The status should equal 224
+        The output should include "image containing PIS GPU cleanup helper required"
+        The output should not include "basePrep driver cleanup"
+    End
+
+    It 'does not require the helper on an ordinary Ubuntu node'
+        OS=ubuntu
+        rm "${fixture}/provision_installs_ubuntu_gpu.sh"
+        When call run_base_prep_cleanup
+        The status should be success
+        The output should be blank
+    End
+
+    It 'does not require the helper for managed GPU image preparation'
+        PRE_PROVISION_ONLY=true
+        GPU_NODE=true
+        OS=ubuntu
+        rm "${fixture}/provision_installs_ubuntu_gpu.sh"
+        When call run_base_prep_cleanup
+        The status should be success
+        The output should be blank
+    End
+
+    It 'keeps Azure Linux bake cleanup independent of the Ubuntu helper'
+        PRE_PROVISION_ONLY=true
+        OS=azurelinux
+        rm "${fixture}/provision_installs_ubuntu_gpu.sh"
+        When call run_base_prep_cleanup
+        The status should be success
+        The output should include "basePrep driver cleanup"
+    End
+
     Describe 'PIS basePrep completion marker'
         It 'does not write the marker after failed basePrep'
             EXIT_CODE=224
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..615c003108e
--- /dev/null
+++ b/spec/vhdbuilder/packer/pis_gpu_helper_spec.sh
@@ -0,0 +1,37 @@
+#!/bin/bash
+# shellcheck disable=SC2329
+
+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 status should be success
+    The output should equal '/home/packer/provision_installs_ubuntu_gpu.sh /opt/azure/containers/provision_installs_ubuntu_gpu.sh 0744'
+  End
+
+  Parameters
+    MARINER
+    AZURELINUX
+    FLATCAR
+  End
+  It "does not require the Ubuntu helper on $1"
+    OS="$1"
+    When call copyUbuntuGPUCleanupHelper
+    The status should be success
+    The output should be blank
+  End
+
+  It 'installs the helper when copying Packer files'
+    When call sed -n '/^copyPackerFiles()/,/^}/p' vhdbuilder/packer/packer_source.sh
+    The output should include '  copyUbuntuGPUCleanupHelper'
+  End
+End

From ad208234577f1e332fdb72c8e9bbb1a09b546ca9 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 23 Oct 2026 11:09:58 +1300
Subject: [PATCH 20/24] test(e2e): exercise mixed PIS script delivery

Use baked scripts for the bake unless explicitly disabled, then traditional delivery on the derived node. Document the required new source image for affected Ubuntu PIS pools.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 e2e/README.md            | 20 ++++++++++++++++++++
 e2e/scenario/scenario.go |  3 +++
 2 files changed, 23 insertions(+)

diff --git a/e2e/README.md b/e2e/README.md
index 2d7c56c02ee..677b34e6ad5 100644
--- a/e2e/README.md
+++ b/e2e/README.md
@@ -94,6 +94,26 @@ 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` must bake from a newly built Ubuntu
+24.04 Gen2 VHD containing the GPU-preservation changes in the current branch;
+an older gallery image cannot verify those changes. Use the VHD builder
+pipeline's same-run E2E VHD metadata, or pass `--vhd-metadata-file`
+(`E2E_VHD_METADATA_FILE`) with `2404gen2containerd` mapped to the new image
+resource ID and replicated regions.
+
+The scenario uses Phase 1 (baked shell scripts) for the bake unless scriptless
+is explicitly disabled, and traditional script delivery for the derived node.
+The Ubuntu strict cleanup helper is baked into the image and is not transferred
+in CustomData, even when the other CSE scripts are transferred.
+
+Ubuntu PIS bakes with managed GPU installation disabled require a new source
+image containing this helper. This includes CPU-only Ubuntu PIS pools, because
+the opt-out condition also includes CPU nodes. A missing helper fails the bake
+with an explicit image requirement instead of silently skipping cleanup.
+Ordinary nodes and derived PIS nodes do not require the helper.
+
 ### LocalDNS fallback checks
 
 The existing LocalDNS scenarios use short-lived DNS pods on the scenario node
diff --git a/e2e/scenario/scenario.go b/e2e/scenario/scenario.go
index 33634d99e4b..1101b68c480 100755
--- a/e2e/scenario/scenario.go
+++ b/e2e/scenario/scenario.go
@@ -3118,6 +3118,9 @@ var _ = Register(&Scenario{
 		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"

From bb5b398e95165635a84b65e455997f2349f240ec Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 23 Oct 2026 11:16:52 +1300
Subject: [PATCH 21/24] test(gpu): wrap delivery fixture for lint

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 pkg/agent/gpu_pis_delivery_test.go | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/pkg/agent/gpu_pis_delivery_test.go b/pkg/agent/gpu_pis_delivery_test.go
index a7e13ad065e..4f791bdbe61 100644
--- a/pkg/agent/gpu_pis_delivery_test.go
+++ b/pkg/agent/gpu_pis_delivery_test.go
@@ -29,7 +29,11 @@ func TestPISGPUBakedHelperDelivery(t *testing.T) {
 					Distro:       distro,
 				}
 				properties := &datamodel.Properties{
-					OrchestratorProfile:     &datamodel.OrchestratorProfile{OrchestratorType: datamodel.Kubernetes, OrchestratorVersion: "1.32.1", KubernetesConfig: &datamodel.KubernetesConfig{}},
+					OrchestratorProfile: &datamodel.OrchestratorProfile{
+						OrchestratorType:    datamodel.Kubernetes,
+						OrchestratorVersion: "1.32.1",
+						KubernetesConfig:    &datamodel.KubernetesConfig{},
+					},
 					HostedMasterProfile:     &datamodel.HostedMasterProfile{DNSPrefix: "fixture"},
 					LinuxProfile:            &datamodel.LinuxProfile{AdminUsername: "azureuser"},
 					ServicePrincipalProfile: &datamodel.ServicePrincipalProfile{ClientID: "fixture", Secret: "fixture"},

From 89580d94158150b127fa292cce1d9521a998c258 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 23 Oct 2026 12:09:42 +1300
Subject: [PATCH 22/24] test(gpu): focus PIS regression coverage

Keep representative preservation, teardown failure, GB ownership and retry, payload, and E2E checks. Remove the expanded test matrices and repeated fixture setup without changing production or Packer behavior.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 e2e/README.md                                 |  22 +-
 e2e/scenario/validators_test.go               |   2 +-
 pkg/agent/gpu_pis_delivery_test.go            | 227 +++------
 .../artifacts/cse_install_mariner_spec.sh     |   2 +-
 .../cse_install_ubuntu_gpu_pis_spec.sh        |  75 +++
 .../artifacts/cse_install_ubuntu_spec.sh      | 471 +-----------------
 .../artifacts/cse_main_gpu_pis_spec.sh        |  90 +---
 spec/vhdbuilder/packer/pis_gpu_helper_spec.sh |  24 +-
 8 files changed, 161 insertions(+), 752 deletions(-)
 create mode 100644 spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_gpu_pis_spec.sh

diff --git a/e2e/README.md b/e2e/README.md
index 677b34e6ad5..2c578cf6130 100644
--- a/e2e/README.md
+++ b/e2e/README.md
@@ -96,23 +96,11 @@ VHD-caching behavior remain in effect.
 
 ### Ubuntu PIS GPU cleanup
 
-`Ubuntu2404_GPUDriverOptOut_VHDCaching` must bake from a newly built Ubuntu
-24.04 Gen2 VHD containing the GPU-preservation changes in the current branch;
-an older gallery image cannot verify those changes. Use the VHD builder
-pipeline's same-run E2E VHD metadata, or pass `--vhd-metadata-file`
-(`E2E_VHD_METADATA_FILE`) with `2404gen2containerd` mapped to the new image
-resource ID and replicated regions.
-
-The scenario uses Phase 1 (baked shell scripts) for the bake unless scriptless
-is explicitly disabled, and traditional script delivery for the derived node.
-The Ubuntu strict cleanup helper is baked into the image and is not transferred
-in CustomData, even when the other CSE scripts are transferred.
-
-Ubuntu PIS bakes with managed GPU installation disabled require a new source
-image containing this helper. This includes CPU-only Ubuntu PIS pools, because
-the opt-out condition also includes CPU nodes. A missing helper fails the bake
-with an explicit image requirement instead of silently skipping cleanup.
-Ordinary nodes and derived PIS nodes do not require the helper.
+`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.
 
 ### LocalDNS fallback checks
 
diff --git a/e2e/scenario/validators_test.go b/e2e/scenario/validators_test.go
index e9b7d20a169..197508dcb8a 100644
--- a/e2e/scenario/validators_test.go
+++ b/e2e/scenario/validators_test.go
@@ -17,7 +17,7 @@ import (
 func TestCustomerNvidiaDriverValidationRequiresStagedInstaller(t *testing.T) {
 	const expectedVersion = "580.159.04"
 	binDir := t.TempDir()
-	installerPath := filepath.Join(t.TempDir(), "nvidia-installer")
+	installerPath := filepath.Join(binDir, "nvidia-installer")
 	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
diff --git a/pkg/agent/gpu_pis_delivery_test.go b/pkg/agent/gpu_pis_delivery_test.go
index 4f791bdbe61..74f4a990a71 100644
--- a/pkg/agent/gpu_pis_delivery_test.go
+++ b/pkg/agent/gpu_pis_delivery_test.go
@@ -2,12 +2,8 @@ package agent
 
 import (
 	"encoding/base64"
-	"encoding/json"
-	"fmt"
-	"os"
 	"os/exec"
 	"path/filepath"
-	"strings"
 	"testing"
 
 	"github.com/Azure/agentbaker/pkg/agent/datamodel"
@@ -15,173 +11,74 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
-func TestPISGPUBakedHelperDelivery(t *testing.T) {
-	for _, distro := range []datamodel.Distro{datamodel.AKSUbuntuContainerd2204Gen2, datamodel.AKSUbuntuContainerd2404Gen2} {
-		for _, mode := range []struct {
-			preProvisionOnly bool
-			phase1           bool
-		}{{true, true}, {true, false}, {false, false}} {
-			t.Run(fmt.Sprintf("%s/PreProvisionOnly=%t/Phase1=%t", distro, mode.preProvisionOnly, mode.phase1), func(t *testing.T) {
-				pool := &datamodel.AgentPoolProfile{
-					Name: "agent2", VMSize: "Standard_NC4as_T4_v3", StorageProfile: "ManagedDisks",
-					OSType: datamodel.Linux, AvailabilityProfile: datamodel.VirtualMachineScaleSets,
-					VnetSubnetID: "/subscriptions/test/resourceGroups/test/providers/Microsoft.Network/virtualNetworks/test/subnets/default",
-					Distro:       distro,
-				}
-				properties := &datamodel.Properties{
-					OrchestratorProfile: &datamodel.OrchestratorProfile{
-						OrchestratorType:    datamodel.Kubernetes,
-						OrchestratorVersion: "1.32.1",
-						KubernetesConfig:    &datamodel.KubernetesConfig{},
-					},
-					HostedMasterProfile:     &datamodel.HostedMasterProfile{DNSPrefix: "fixture"},
-					LinuxProfile:            &datamodel.LinuxProfile{AdminUsername: "azureuser"},
-					ServicePrincipalProfile: &datamodel.ServicePrincipalProfile{ClientID: "fixture", Secret: "fixture"},
-					AgentPoolProfiles:       []*datamodel.AgentPoolProfile{pool},
-				}
-				properties.LinuxProfile.SSH.PublicKeys = []datamodel.PublicKey{{KeyData: "fixture-key"}}
-				config := &datamodel.NodeBootstrappingConfiguration{
-					ContainerService: &datamodel.ContainerService{Location: "southcentralus", Type: "Microsoft.ContainerService/ManagedClusters", Properties: properties},
-					CloudSpecConfig:  datamodel.AzurePublicCloudSpecForTest, K8sComponents: &datamodel.K8sComponents{},
-					AgentPoolProfile: pool, TenantID: "tenantID", SubscriptionID: "subID",
-					ResourceGroupName: "resourceGroupName", UserAssignedIdentityClientID: "userAssignedID",
-					KubeletConfig:       map[string]string{"--cluster-domain": "cluster.local", "--cluster-dns": "10.0.0.10", "--cloud-provider": "azure", "--max-pods": "110"},
-					PrimaryScaleSetName: "aks-agent2-36873793-vmss", PreProvisionOnly: mode.preProvisionOnly,
-					EnableNvidia: false, ConfigGPUDriverIfNeeded: false,
-					EnableGPUDevicePluginIfNeeded: false, EnableScriptlessCSECmd: mode.phase1,
-				}
+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
 
-				require.False(t, supportsScriptlessProvisioning(config))
-				generator := InitializeTemplateGenerator()
-				command := generator.getNodeBootstrappingCmd(config)
-				require.Contains(t, command, fmt.Sprintf(`PRE_PROVISION_ONLY="%t"`, mode.preProvisionOnly))
-				require.Contains(t, command, "GPU_NODE=false")
-				require.Contains(t, command, `/bin/bash /opt/azure/containers/provision_start.sh`)
-				require.NotContains(t, command, "aks-node-controller provision-wait")
+		payload := InitializeTemplateGenerator().getLinuxNodeBootstrappingPayload(config)
+		compressed, err := base64.StdEncoding.DecodeString(payload)
+		require.NoError(t, err)
+		decoded, err := getGzipDecodedValue(compressed)
+		require.NoError(t, err)
+		require.LessOrEqual(t, len(compressed), 65535)
+		require.LessOrEqual(t, len(payload), 87380)
 
-				payload := generator.getLinuxNodeBootstrappingPayload(config)
-				compressed, err := base64.StdEncoding.DecodeString(payload)
+		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)
-				decoded, err := getGzipDecodedValue(compressed)
-				require.NoError(t, err)
-				t.Logf("PIS T4 payload: yaml=%d gzip=%d base64=%d gzip-margin=%d base64-margin=%d", len(decoded), len(compressed), len(payload), 65535-len(compressed), 87380-len(payload))
-				require.LessOrEqual(t, len(payload), 87380)
-				require.LessOrEqual(t, len(compressed), 65535)
-				var document cloudInit
-				require.NoError(t, yaml.Unmarshal(decoded, &document))
-				installerFound := false
-				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" {
-						installerFound = true
-						script, err := getGzipDecodedValue([]byte(file.Content))
-						require.NoError(t, err)
-						require.NotContains(t, string(script), "graceBlackwellConfigIsUnmodified()")
-						require.NotContains(t, string(script), "cleanUpGPUDriversForBasePrep()")
-						require.Contains(t, string(script), "cleanUpGPUDrivers()")
-					}
-				}
-				if mode.phase1 {
-					require.Len(t, document.WriteFiles, 1)
-					require.Equal(t, "/opt/azure/containers/scriptless-cse-overrides.txt", document.WriteFiles[0].Path)
-				} else {
-					require.Greater(t, len(document.WriteFiles), 1)
-					require.True(t, installerFound)
-				}
-			})
-		}
-	}
-}
-
-func TestPISGPUBakedHelperProvenance(t *testing.T) {
-	const helper = "parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh"
-	root := filepath.Join("..", "..")
-	for _, variant := range []string{"base", "cvm", "arm64-gen2", "arm64-gb"} {
-		data, err := os.ReadFile(filepath.Join(root, "vhdbuilder", "packer", "vhd-image-builder-"+variant+".json"))
-		require.NoError(t, err)
-		var image struct {
-			Provisioners []struct{ Type, Source, Destination string }
-		}
-		require.NoError(t, json.Unmarshal(data, &image))
-		matches := 0
-		for _, provisioner := range image.Provisioners {
-			if provisioner.Type == "file" && provisioner.Source == helper &&
-				provisioner.Destination == "/home/packer/provision_installs_ubuntu_gpu.sh" {
-				matches++
+				require.Contains(t, string(script), "cleanUpGPUDrivers()")
 			}
 		}
-		require.Equal(t, 1, matches, variant)
+		if mode.phase1 {
+			require.Equal(t, "/opt/azure/containers/scriptless-cse-overrides.txt", document.WriteFiles[0].Path)
+		} else {
+			require.True(t, installer)
+		}
 	}
-	// Execute the raw Packer source, not the comment-stripped CustomData copy.
-	config := filepath.Join(t.TempDir(), "nvidia.conf")
-	require.NoError(t, os.WriteFile(config, []byte("fixture"), 0600))
-	command := exec.Command("bash", "-c", `
-source "$1"
-graceBlackwellInstalledDriverPackages() { :; }
-graceBlackwellPackageManagedDriverRemains() { return 1; }
-graceBlackwellConfigIsUnmodified() {
-    printf '%d\n' "$#"
-    printf '%s\n' "$@"
-    return 1
-}
-GB_NVIDIA_MODPROBE_CONFIG_FILE="$2"
-cleanUpGraceBlackwellGPUDriver unused unused unused
-`, "gpu-helper-test", filepath.Join(root, helper), config)
-	// A present config reaches the comparison without executing teardown.
-	output, err := command.CombinedOutput()
-	require.Error(t, err)
-	require.True(t, strings.HasPrefix(string(output), "8\n"), string(output))
-	require.Contains(t, string(output), "# Explicitly load nvidia_peermem at boot for GPUDirect RDMA.")
 }
 
-func TestGPUAMADeliveryIsDistroScoped(t *testing.T) {
-	distros := append([]datamodel.Distro{
-		datamodel.AKSUbuntuContainerd2204Gen2,
-		datamodel.AKSUbuntuContainerd2404Gen2,
-		datamodel.AKSACLGen2TL,
-		datamodel.AKSFlatcarGen2,
-		datamodel.CustomizedImage,
-		datamodel.CustomizedImageKata,
-		datamodel.CustomizedImageLinuxGuard,
-		datamodel.CustomizedImageTrustedLaunch,
-	}, datamodel.AvailableAzureLinuxDistros...)
-	for _, distro := range distros {
-		t.Run(string(distro), func(t *testing.T) {
-			config := newNodeCustomDataRenderConfig(distro)
-			encoded := getBase64EncodedGzippedCustomScript(kubernetesCSEConfigGPU, config)
-			compressed, err := base64.StdEncoding.DecodeString(encoded)
-			require.NoError(t, err)
-			script, err := getGzipDecodedValue(compressed)
-			require.NoError(t, err)
-			custom := distro == datamodel.CustomizedImage || distro == datamodel.CustomizedImageKata ||
-				distro == datamodel.CustomizedImageLinuxGuard || distro == datamodel.CustomizedImageTrustedLaunch
-			keepImplementation := distro.IsAzureLinuxDistro() || custom
-			require.Equal(t, keepImplementation, strings.Contains(string(script), "dnf_install_amd_ama_core_package()"))
-			require.Equal(t, keepImplementation, strings.Contains(string(script), `dnf repoquery -y --available "amd-ama-driver-*"`))
-			require.Contains(t, string(script), "setupAmdAma()")
-			syntax := exec.Command("bash", "-n")
-			syntax.Stdin = strings.NewReader(string(script))
-			output, err := syntax.CombinedOutput()
-			require.NoError(t, err, string(output))
-			if keepImplementation {
-				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))
-				rendered := exec.Command("bash", "-c", string(script)+"\ndeclare -f setupAmdAma dnf_install_amd_ama_core_package")
-				renderedFunctions, err := rendered.CombinedOutput()
-				require.NoError(t, err, string(renderedFunctions))
-				require.Equal(t, string(rawFunctions), string(renderedFunctions))
-			} else {
-				command := exec.Command("bash", "-c", string(script)+`
-isARM64() { echo 0; }
-isAzureLinux() { return 1; }
-setupAmdAma
-`)
-				output, err := command.CombinedOutput()
-				require.NoError(t, err, string(output))
-				require.Empty(t, output)
-			}
-		})
-	}
+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 c243d5e7d88..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
@@ -32,7 +32,7 @@ Describe 'cse_install_mariner.sh'
         GPU_DEST="${PWD}/.shellspec-nonexistent-gpu-dest"
         managedGPUPackageList() { echo "driver-cache"; }
         getPackageCacheDir() { echo "$gpu_test_root/driver-cache"; }
-        setup_gpu_cache() { gpu_test_root=$(mktemp -d); }
+        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'
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..de9fb823273
--- /dev/null
+++ b/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_gpu_pis_spec.sh
@@ -0,0 +1,75 @@
+#!/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"
+    }
+
+    BeforeEach 'setup_cleanup'
+    AfterEach 'rm -rf "$fixture"'
+
+    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' "${fixture}/bom" 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_install_ubuntu_spec.sh b/spec/parts/linux/cloud-init/artifacts/cse_install_ubuntu_spec.sh
index 2dc65af836f..8c18fb3badd 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
@@ -85,473 +85,10 @@ Describe 'cse_install_ubuntu.sh'
             The output should include "module_after=true"
             The output should include "status=incomplete"
         End
-
-    End
-
-    Describe 'cleanUpGPUDriversForBasePrep'
-        Include "./parts/linux/cloud-init/artifacts/ubuntu/cse_install_ubuntu_gpu.sh"
-        GPU_DEST="${PWD}/.shellspec-nonexistent-gpu-dest"
-        managedGPUPackageList() { :; }
-
-        createGraceBlackwellFixture() {
-            local wave="${1:-wave2}"
-            GB_MAI_BOM_FILE="${PWD}/.shellspec-gb-bom-$$.json"
-            GB200_MAI_BOM_FILE="${PWD}/.shellspec-gb200-bom-$$.json"
-            GPU_DKMS_MARKER_FILE="${PWD}/.shellspec-gpu-marker-$$"
-            GB_NVIDIA_PEERMEM_CONFIG_FILE="${PWD}/.shellspec-nvidia-peermem-$$.conf"
-            GB_NVIDIA_MODPROBE_CONFIG_FILE="${PWD}/.shellspec-nvidia-modprobe-$$.conf"
-            GB_NOUVEAU_MODPROBE_CONFIG_FILE="${PWD}/.shellspec-nouveau-$$.conf"
-            GB_DRIVER_CLEANUP_PENDING_FILE="${PWD}/.shellspec-gb-cleanup-pending-$$"
-            if [ "$wave" = wave1 ]; then
-                printf '%s\n' '{"versions-wave1":{"libnvidia-common-580":"580.105.08-0ubuntu1","nvidia-dkms-580-open":"580.105.08-0ubuntu1","nvidia-driver-580-open":"580.105.08-0ubuntu1"}}' > "$GB200_MAI_BOM_FILE"
-                printf '%s\n' 'options nvidia NVreg_RestrictProfilingToAdminUsers=0' \
-                    'options nvidia NVreg_CreateImexChannel0=1' \
-                    'options nvidia NVreg_CoherentGPUMemoryMode=driver' \
-                    'options nvidia NVreg_RegistryDwords="RMBug5172204War=4"' > "$GB_NVIDIA_MODPROBE_CONFIG_FILE"
-            else
-                printf '%s\n' '{"versions-wave2":{"libnvidia-common-580":"580.159.04-1ubuntu1","nvidia-dkms-580-open":"580.159.04-1ubuntu1","nvidia-driver-580-open":"580.159.04-1ubuntu1"}}' > "$GB_MAI_BOM_FILE"
-                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"
-            fi
-        }
-
-        cleanupGraceBlackwellFixture() {
-            rm -f "${GB_MAI_BOM_FILE:-}" "${GB200_MAI_BOM_FILE:-}" "${GB_DRIVER_CLEANUP_PENDING_FILE:-}" "${GPU_DKMS_MARKER_FILE:-}" \
-                "${GB_NVIDIA_PEERMEM_CONFIG_FILE:-}" "${GB_NVIDIA_MODPROBE_CONFIG_FILE:-}" \
-                "${GB_NOUVEAU_MODPROBE_CONFIG_FILE:-}"
-        }
-
-        runGraceBlackwellFailureAndRetainState() {
-            cleanUpGPUDriversForBasePrep
-            status=$?
-            [ -f "$GPU_DKMS_MARKER_FILE" ] && echo "aks marker retained" || status=0
-            [ -f "$GB_NVIDIA_PEERMEM_CONFIG_FILE" ] && echo "GB boot config retained"
-            cleanupGraceBlackwellFixture
-            return "$status"
-        }
-
-        mockGraceBlackwellPackageQueries() {
-            dpkg-query() {
-                if [ "$#" -eq 2 ]; then
-                    [ "${gb_packages_purged:-false}" = true ] && return 0
-                    printf 'libnvidia-common-580\tinstalled\nnvidia-dkms-580-open\tinstalled\nnvidia-driver-580-open\tinstalled\nnvidia-utils-580\tinstalled\nnvidia-modprobe\tinstalled\nnvidia-persistenced\tinstalled\n'
-                    [ -n "${gb_extra_driver_package:-}" ] && printf '%s\tinstalled\n' "$gb_extra_driver_package"
-                    return 0
-                fi
-                case "$3" in
-                    libnvidia-common-580|nvidia-dkms-580-open|nvidia-driver-580-open)
-                        if [ "${gb_packages_purged:-false}" = true ]; then
-                            printf 'not-installed\t\n'
-                        elif [ "${gb_version_mismatch:-false}" = true ] && [ "$3" = nvidia-driver-580-open ]; then
-                            printf 'installed\t580.200.01-1ubuntu1\n'
-                        else
-                            printf 'installed\t%s\n' "${gb_driver_version:-580.159.04-1ubuntu1}"
-                        fi
-                        ;;
-                    nvidia-utils-580)
-                        [ "${gb_packages_purged:-false}" = true ] && printf 'not-installed\t\n' || printf 'installed\t%s\n' "${gb_utils_version:-580.159.04-1}"
-                        ;;
-                    nvidia-modprobe)
-                        [ "${gb_packages_purged:-false}" = true ] && printf 'not-installed\t\n' || printf 'installed\t%s\n' "${gb_modprobe_version:-1.10.0-1ubuntu1}"
-                        ;;
-                    nvidia-persistenced)
-                        [ "${gb_packages_purged:-false}" = true ] && printf 'not-installed\t\n' || printf 'installed\t%s\n' "${gb_persistenced_version:-1.10.0-1ubuntu1}"
-                        ;;
-                    *) printf 'not-installed\t\n' ;;
-                esac
-            }
-            apt-mark() { printf 'nvidia-utils-580\nnvidia-modprobe\nnvidia-persistenced\n'; }
-        }
-
-        setupGraceBlackwellTeardown() {
-            local wave="${1:-wave2}"
-            createGraceBlackwellFixture "$wave"
-            : > "$GPU_DKMS_MARKER_FILE"
-            gb_driver_version=580.159.04-1ubuntu1
-            gb_utils_version=580.159.04-1
-            gb_modprobe_version=1.10.0-1ubuntu1
-            gb_persistenced_version=1.10.0-1ubuntu1
-            if [ "$wave" = wave1 ]; then
-                gb_driver_version=580.105.08-0ubuntu1
-                gb_utils_version=580.105.08-1
-            fi
-            gb_loaded_modules='nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia '
-            gb_expected_modules='nvidia_peermem nvidia_uvm nvidia_drm nvidia_modeset nvidia '
-            gb_services_stopped=false
-            gb_persistenced_active=true
-            gb_persistenced_enabled=true
-            gb_packages_purged=false
-            mockGraceBlackwellPackageQueries
-        }
-
-        runGraceBlackwellTeardown() {
-            local wave="${1:-wave2}" retry="${2:-false}" enabled="${3:-true}" persistenced_stop_failure="${7:-false}"
-            setupGraceBlackwellTeardown "$wave"
-            gb_openibd_enabled="$enabled"
-            gb_openibd_active=true
-            gb_openibd_stop_count=0
-            gb_openibd_stop_failure="${4:-false}"
-            gb_openibd_remain_active="${5:-false}"
-            gb_openibd_unloads_module="${6:-false}"
-            systemctlDisableAndStop() {
-                echo "stop:$1"
-                if [ "$1" = nvidia-persistenced ]; then
-                    if [ "$persistenced_stop_failure" != true ]; then
-                        gb_persistenced_active=false
-                        gb_persistenced_enabled=false
-                    fi
-                fi
-                if [ "$1" = openibd ]; then
-                    gb_services_stopped=true
-                    gb_openibd_active=false
-                    gb_openibd_enabled=false
-                fi
-            }
-            systemctl() {
-                case "$1" in
-                    cat) [ "$2" = openibd ] ;;
-                    is-active)
-                        { [ "$3" = openibd ] && [ "$gb_openibd_active" = true ]; } ||
-                            { [ "$3" = nvidia-persistenced ] && [ "$gb_persistenced_active" = true ]; }
-                        ;;
-                    is-enabled)
-                        { [ "$3" = openibd ] && [ "$gb_openibd_enabled" = true ]; } ||
-                            { [ "$3" = nvidia-persistenced ] && [ "$gb_persistenced_enabled" = true ]; }
-                        ;;
-                    *) return 1 ;;
-                esac
-            }
-            systemctl_stop() {
-                [ "$4" = openibd ] || return 1
-                gb_openibd_stop_count=$((gb_openibd_stop_count + 1))
-                [ "$gb_openibd_stop_failure" != true ] || return 1
-                [ "$gb_openibd_remain_active" = true ] || gb_openibd_active=false
-                [ "$gb_openibd_active" = false ] && gb_services_stopped=true
-                echo "stop:openibd"
-                if [ "$gb_openibd_unloads_module" = true ]; then
-                    local unloaded_module="${gb_loaded_modules%% *}"
-                    gb_loaded_modules="${gb_loaded_modules#* }"
-                    gb_expected_modules="${gb_expected_modules#* }"
-                    echo "openibd_unloaded:$unloaded_module"
-                fi
-                return 0
-            }
-            lsmod() {
-                echo "Module Size Used by"
-                for gb_module in $gb_loaded_modules; do
-                    echo "$gb_module 123 0"
-                done
-            }
-            rmmod() {
-                [ "$gb_services_stopped" = true ] || return 1
-                [ "$gb_persistenced_active" = false ] || return 1
-                [ "$gb_openibd_active" = false ] || return 1
-                [ "$gb_openibd_enabled" = "$enabled" ] || return 1
-                case " $gb_loaded_modules " in *" $1 "*) ;; *) return 1 ;; esac
-                [ "${gb_expected_modules%% *}" = "$1" ] || return 1
-                echo "rmmod:$1 openibd_enabled=$gb_openibd_enabled"
-                gb_expected_modules="${gb_expected_modules#* }"
-                gb_loaded_modules="${gb_loaded_modules#* }"
-            }
-            update-initramfs() {
-                [ "$gb_packages_purged" = true ] || return 1
-                [ ! -e "$GB_NVIDIA_PEERMEM_CONFIG_FILE" ] || return 1
-                [ ! -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || return 1
-                [ ! -e "$GB_NOUVEAU_MODPROBE_CONFIG_FILE" ] || return 1
-                [ -f "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
-                echo "update-initramfs"
-            }
-            apt_get_purge() {
-                [ "$1 $2 $3" = '10 5 300' ] || return 1
-                shift 3
-                local expected='libnvidia-common-580 nvidia-dkms-580-open nvidia-driver-580-open nvidia-utils-580 nvidia-modprobe nvidia-persistenced'
-                [ "$*" = "$expected" ] || return 1
-                [ -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || return 1
-                [ "$gb_services_stopped" = true ] || return 1
-                [ -z "$gb_loaded_modules" ] || return 1
-                [ -z "$gb_expected_modules" ] || return 1
-                echo "purge:$*"
-                gb_packages_purged=true
-            }
-            cleanUpGPUDrivers() {
-                [ "$gb_packages_purged" = true ] || return 1
-                [ ! -e "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
-                echo "aks-marker-cleanup"
-                rm -f "$GPU_DKMS_MARKER_FILE"
-            }
-            prebakedGPUDriverArtifactsRemain() { [ "$gb_packages_purged" != true ]; }
-            cleanUpGPUDriversForBasePrep
-            status=$?
-            if [ "$status" -eq 0 ] && [ "$retry" = true ]; then
-                cleanUpGPUDriversForBasePrep
-                status=$?
-                [ "$gb_openibd_stop_count" -eq 1 ] || status=1
-                echo "openibd_retry_stop_count=$gb_openibd_stop_count"
-            fi
-            [ "$status" -ne 0 ] || echo "openibd_enabled_after_unload=$gb_openibd_enabled"
-            cleanupGraceBlackwellFixture
-            return "$status"
-        }
-
-        It 'fails when the prebake marker survives cleanup'
-            GPU_DKMS_MARKER_FILE="$(mktemp)"
-            cleanUpGPUDrivers() { :; }
-            lsmod() { echo ""; }
-            When call cleanUpGPUDriversForBasePrep
-            The status should be failure
-            The stderr should include "GPU basePrep cleanup incomplete"
-            The path "$GPU_DKMS_MARKER_FILE" should be exist
-            AfterRun 'rm -f "$GPU_DKMS_MARKER_FILE"'
-        End
-
-        It 'fails when nvidia remains loaded without a prebake marker'
-            GPU_DKMS_MARKER_FILE="${PWD}/.shellspec-absent-gpu-marker-$$"
-            cleanUpGPUDrivers() { :; }
-            lsmod() { echo "nvidia 104165376 0"; }
-            When call cleanUpGPUDriversForBasePrep
-            The status should be failure
-            The stderr should include "GPU basePrep cleanup incomplete"
-        End
-
-        It 'fails when an unmarked prebake artifact remains'
-            GPU_DKMS_MARKER_FILE="${PWD}/.shellspec-absent-gpu-marker-$$"
-            cleanUpGPUDrivers() { :; }
-            compgen() { [ "$2" = '/lib/modules/*/updates/dkms/nvidia*.ko*' ]; }
-            lsmod() { echo ""; }
-            When call cleanUpGPUDriversForBasePrep
-            The status should be failure
-            The stderr should include "GPU basePrep cleanup incomplete"
-        End
-
-        It 'tears down current and legacy BOM-owned Grace Blackwell drivers before marker cleanup'
-            runBothGraceBlackwellTeardowns() {
-                runGraceBlackwellTeardown wave2 || return 1
-                runGraceBlackwellTeardown wave1 true
-            }
-
-            When call runBothGraceBlackwellTeardowns
-            The status should be success
-            The output should include "rmmod:nvidia_peermem"
-            The output should include "rmmod:nvidia_peermem openibd_enabled=true"
-            The output should include "stop:nvidia-persistenced"
-            The output should include "stop:nvidia-imex"
-            The output should include "stop:openibd"
-            The output should include "purge:"
-            The output should include "aks-marker-cleanup"
-            The output should include "openibd_retry_stop_count=1"
-        End
-
-        It 'preserves openibd enablement while stopping it for module unload and retries'
-            When call runGraceBlackwellTeardown wave2 true
-            The status should be success
-            The output should include "rmmod:nvidia_peermem openibd_enabled=true"
-            The output should include "openibd_enabled_after_unload=true"
-            The output should include "openibd_retry_stop_count=1"
-        End
-
-        It 'keeps openibd disabled when it was disabled before module unload'
-            When call runGraceBlackwellTeardown wave2 false false
-            The status should be success
-            The output should include "rmmod:nvidia_peermem openibd_enabled=false"
-            The output should include "openibd_enabled_after_unload=false"
-        End
-
-        It 'fails before module unload when nvidia-persistenced cannot be stopped'
-            When call runGraceBlackwellTeardown wave2 false true false false false true
-            The status should be failure
-            The stderr should include "Grace Blackwell service nvidia-persistenced remains active or enabled"
-            The output should not include "rmmod:"
-        End
-
-        It 'fails before module unload when openibd cannot be stopped'
-            When call runGraceBlackwellTeardown wave2 false true true
-            The status should be failure
-            The stderr should include "Grace Blackwell service openibd could not be stopped"
-            The output should not include "rmmod:"
-        End
-
-        It 'fails before module unload when openibd remains active'
-            When call runGraceBlackwellTeardown wave2 false true false true
-            The status should be failure
-            The stderr should include "Grace Blackwell service openibd remains active"
-            The output should not include "rmmod:"
-        End
-
-        It 'refreshes the loaded-module snapshot after openibd stops'
-            When call runGraceBlackwellTeardown wave2 false true false false true
-            The status should be success
-            The output should include "openibd_unloaded:nvidia_peermem"
-            The output should not include "rmmod:nvidia_peermem"
-            The output should include "rmmod:nvidia_uvm"
-        End
-
-        It 'uses a BOM-listed dependency version instead of comparing it to the driver version'
-            runGraceBlackwellBOMDependencyVersion() {
-                createGraceBlackwellFixture
-                gb_driver_version=580.159.04-1ubuntu1
-                gb_modprobe_version=1.10.0-1ubuntu1
-                printf '%s\n' '{"versions-wave2":{"libnvidia-common-580":"580.159.04-1ubuntu1","nvidia-dkms-580-open":"580.159.04-1ubuntu1","nvidia-driver-580-open":"580.159.04-1ubuntu1","nvidia-modprobe":"1.10.0-1ubuntu1"}}' > "$GB_MAI_BOM_FILE"
-                mockGraceBlackwellPackageQueries
-                entries=$(graceBlackwellDriverBOMEntries "$GB_MAI_BOM_FILE" versions-wave2) || return 1
-                graceBlackwellInstalledDriverPackages "$entries"
-                status=$?
-                cleanupGraceBlackwellFixture
-                return "$status"
-            }
-            When call runGraceBlackwellBOMDependencyVersion
-            The status should be success
-            The output should include "nvidia-modprobe=1.10.0-1ubuntu1"
-        End
-
-        It 'retains GB provenance and retries after config removal or initramfs failure'
-            runGraceBlackwellTeardownRetry() {
-                local status=0
-                setupGraceBlackwellTeardown
-                gb_loaded_modules=""
-                mockGraceBlackwellPackageQueries
-                systemctlDisableAndStop() { :; }
-                systemctl() { return 1; }
-                lsmod() { echo "Module Size Used by"; }
-                apt_get_purge() { gb_packages_purged=true; }
-                gb_fail_config_remove=true
-                rm() {
-                    if [ "${gb_fail_config_remove:-false}" = true ] &&
-                        [ "$*" = "-f $GB_NVIDIA_PEERMEM_CONFIG_FILE $GB_NVIDIA_MODPROBE_CONFIG_FILE $GB_NOUVEAU_MODPROBE_CONFIG_FILE" ]; then
-                        gb_fail_config_remove=false
-                        return 1
-                    fi
-                    command rm "$@"
-                }
-                gb_initramfs_failures=1
-                update-initramfs() {
-                    [ ! -e "$GB_NVIDIA_PEERMEM_CONFIG_FILE" ] || return 1
-                    [ ! -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || return 1
-                    [ ! -e "$GB_NOUVEAU_MODPROBE_CONFIG_FILE" ] || return 1
-                    [ -f "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
-                    if [ "$gb_initramfs_failures" -gt 0 ]; then
-                        gb_initramfs_failures=$((gb_initramfs_failures - 1))
-                        return 1
-                    fi
-                }
-                cleanUpGPUDrivers() {
-                    [ "$gb_packages_purged" = true ] || return 1
-                    [ ! -e "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || return 1
-                    rm -f "$GPU_DKMS_MARKER_FILE"
-                }
-                prebakedGPUDriverArtifactsRemain() { [ "$gb_packages_purged" != true ]; }
-
-                cleanUpGPUDriversForBasePrep && status=1
-                [ -f "$GPU_DKMS_MARKER_FILE" ] || status=1
-                [ -f "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || status=1
-                [ -f "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || status=1
-                cleanUpGPUDriversForBasePrep && status=1
-                [ -f "$GPU_DKMS_MARKER_FILE" ] || status=1
-                [ -f "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || status=1
-                [ ! -e "$GB_NVIDIA_MODPROBE_CONFIG_FILE" ] || status=1
-                cleanUpGPUDriversForBasePrep || status=1
-                [ ! -e "$GB_DRIVER_CLEANUP_PENDING_FILE" ] || status=1
-                [ ! -e "$GPU_DKMS_MARKER_FILE" ] || status=1
-                cleanupGraceBlackwellFixture
-                return "$status"
-            }
-            When call runGraceBlackwellTeardownRetry
-            The status should be success
-            The stderr should include "GPU basePrep cleanup incomplete: Grace Blackwell driver teardown failed"
-        End
-
-        It 'retains the marker on a legacy GB VHD without a BOM'
-            runGraceBlackwellWithoutBOM() {
-                status=0
-                createGraceBlackwellFixture
-                rm -f "$GB_MAI_BOM_FILE"
-                : > "$GPU_DKMS_MARKER_FILE"
-                cleanUpGPUDrivers() { echo "unexpected marker cleanup"; }
-                runGraceBlackwellFailureAndRetainState || status=1
-                createGraceBlackwellFixture
-                rm -f "$GB_MAI_BOM_FILE" "$GB_NVIDIA_PEERMEM_CONFIG_FILE"
-                : > "$GPU_DKMS_MARKER_FILE"
-                dpkg-query() {
-                    [ "$#" -eq 2 ] && printf 'nvidia-dkms-580-open\tinstalled\n'
-                }
-                runGraceBlackwellFailureAndRetainState || status=1
-                return "$status"
-            }
-            When call runGraceBlackwellWithoutBOM
-            The status should be failure
-            The stderr should include "Grace Blackwell BOM is unavailable"
-            The stderr should include "has no verifiable GB BOM ownership"
-            The output should include "aks marker retained"
-            The output should include "GB boot config retained"
-            The output should not include "unexpected marker cleanup"
-        End
-
-        It 'retains the marker for mixed ownership and package-purge failures'
-            runUnsafeGraceBlackwellTeardown() {
-                status=0
-                setupGraceBlackwellTeardown
-                gb_extra_driver_package=nvidia-driver-535-server
-                mockGraceBlackwellPackageQueries
-                systemctlDisableAndStop() { echo "unexpected service stop"; }
-                apt_get_purge() { echo "unexpected package purge"; }
-                cleanUpGPUDrivers() { echo "unexpected marker cleanup"; }
-                runGraceBlackwellFailureAndRetainState || status=1
-                setupGraceBlackwellTeardown
-                gb_packages_purged=false
-                gb_extra_driver_package=""
-                mockGraceBlackwellPackageQueries
-                printf '\n# customer change\n' >> "$GB_NVIDIA_MODPROBE_CONFIG_FILE"
-                systemctlDisableAndStop() { echo "unexpected service stop"; }
-                systemctl() { return 1; }
-                lsmod() { echo "nvidia 123 0"; }
-                rmmod() { echo "unexpected module unload"; }
-                apt_get_purge() { echo "unexpected package purge"; }
-                cleanUpGPUDrivers() { echo "unexpected marker cleanup"; }
-                runGraceBlackwellFailureAndRetainState || status=1
-                setupGraceBlackwellTeardown
-                gb_extra_driver_package=""
-                mockGraceBlackwellPackageQueries
-                lsmod() { echo "Module Size Used by"; }
-                systemctlDisableAndStop() { :; }
-                systemctl() { return 1; }
-                update-initramfs() { :; }
-                apt_get_purge() { return 1; }
-                runGraceBlackwellFailureAndRetainState || status=1
-                setupGraceBlackwellTeardown
-                gb_version_mismatch=true
-                mockGraceBlackwellPackageQueries
-                apt_get_purge() { echo "unexpected package purge"; }
-                cleanUpGPUDrivers() { echo "unexpected marker cleanup"; }
-                cleanUpGPUDriversForBasePrep || status=1
-                cleanupGraceBlackwellFixture
-                return "$status"
-            }
-            When call runUnsafeGraceBlackwellTeardown
-            The status should be failure
-            The stderr should include "is not owned by the Grace Blackwell BOM"
-            The stderr should include "module configuration"
-            The stderr should include "Grace Blackwell driver package purge failed"
-            The stderr should include "does not match the GB BOM"
-            The output should include "aks marker retained"
-            The output should include "GB boot config retained"
-            The output should not include "unexpected service stop"
-            The output should not include "unexpected module unload"
-            The output should not include "unexpected package purge"
-            The output should not include "unexpected marker cleanup"
-        End
-
-        It 'succeeds when the prebake marker and AKS paths are gone'
-            GPU_DKMS_MARKER_FILE="$(mktemp)"
-            cleanUpGPUDrivers() { rm -f "$GPU_DKMS_MARKER_FILE"; }
-            lsmod() { echo ""; }
-            When call cleanUpGPUDriversForBasePrep
-            The status should be success
-            The path "$GPU_DKMS_MARKER_FILE" should not be exist
-        End
     End
 
     Describe 'installPackageFromCache version matching'
-        deb_cache_root="${PWD}/.shellspec-deb-cache-$$"
+        deb_cache_root="/tmp/shellspec-deb-cache-$$"
 
         setup_deb_cache() {
             mkdir -p "$deb_cache_root"
@@ -673,7 +210,7 @@ EOF
     End
 
     Describe 'installContainerdWithAptGet revision comparison'
-        containerd_download_root="${PWD}/.shellspec-cse-install-ubuntu-containerd-$$"
+        containerd_download_root="/tmp/cse-install-ubuntu-containerd-$$"
 
         setup_containerd_revision() {
             mkdir -p "${containerd_download_root}"
@@ -724,7 +261,7 @@ EOF
     End
 
     Describe 'logResolvedPackageVersion'
-        resolved_version_log="${PWD}/.shellspec-cse-install-ubuntu-resolved-version-$$"
+        resolved_version_log="/tmp/cse-install-ubuntu-resolved-version-$$"
 
         cleanup_resolved_version_log() {
             rm -f "${resolved_version_log}"
@@ -754,7 +291,7 @@ EOF
     End
 
     Describe 'ensureRunc repository fallback'
-        runc_download_root="${PWD}/.shellspec-cse-install-ubuntu-runc-$$"
+        runc_download_root="/tmp/cse-install-ubuntu-runc-$$"
 
         setup_runc_fallback() {
             mkdir -p "${runc_download_root}"
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
index 58a2a6b5bf8..a828b7c0391 100644
--- 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
@@ -12,7 +12,6 @@ Describe 'GPU cleanup across PIS phases'
         PRE_PROVISION_ONLY=false
         skip_nvidia_driver_install=false
         BASE_PREP_RAN=false
-        printf '%s\n' 'cleanUpGPUDriversForBasePrep() { echo "basePrep driver cleanup"; return "${CLEANUP_STATUS:-0}"; }' > "${fixture}/provision_installs_ubuntu_gpu.sh"
     }
 
     run_phases() {
@@ -34,19 +33,6 @@ Describe 'GPU cleanup across PIS phases'
         eval "$dispatch"
     }
 
-    run_base_prep_cleanup() {
-        local cleanup
-        cleanup=$(sed -n '/    if \[ "${PRE_PROVISION_ONLY}" = "true" \] && \[ "${GPU_NODE}" != "true" \] &&/,/^    fi$/p' "$CSE_MAIN")
-        cleanup=${cleanup//\/opt\/azure\/containers\/provision_installs_ubuntu_gpu.sh/$fixture/provision_installs_ubuntu_gpu.sh}
-        logs_to_events() { shift; "$@"; }
-        cleanUpGPUDriversForBasePrep() { echo "basePrep driver cleanup"; return "${CLEANUP_STATUS:-0}"; }
-        isMarinerOrAzureLinux() { [ "$OS" = azurelinux ]; }
-        isAzureLinuxOSGuard() { [ "$OS_VARIANT" = osguard ]; }
-        UBUNTU_OS_NAME=ubuntu
-        ERR_NVIDIA_DRIVER_INSTALL=224
-        ( eval "$cleanup" )
-    }
-
     run_cse_marker() {
         local marker_script
         PRE_PROVISION_ONLY=true
@@ -63,7 +49,7 @@ Describe 'GPU cleanup across PIS phases'
     BeforeEach 'setup_driver'
     AfterEach 'rm -rf "$fixture"'
 
-    It 'preserves customer drivers on a PIS-derived opt-out node when basePrep is skipped'
+    It 'preserves customer drivers on a PIS-derived opt-out node'
         touch "$marker"
         When call run_phases
         The status should be success
@@ -72,7 +58,7 @@ Describe 'GPU cleanup across PIS phases'
         The path "$fixture/cleanup-called" should not be exist
     End
 
-    It 'cleans an opt-out node when non-PIS basePrep ran in this execution'
+    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"
@@ -80,73 +66,17 @@ Describe 'GPU cleanup across PIS phases'
         The path "$fixture/cleanup-called" should be exist
     End
 
-    It 'runs strict driver cleanup on an opt-out bake'
-        PRE_PROVISION_ONLY=true
-        OS=ubuntu
-        When call run_base_prep_cleanup
+    It 'does not write the bake marker after failed basePrep'
+        EXIT_CODE=224
+        When call run_cse_marker
         The status should be success
-        The output should include "basePrep driver cleanup"
+        The path "$fixture/base_prep.complete" should not be exist
     End
 
-    It 'fails basePrep when bake-time cleanup cannot be verified'
-        PRE_PROVISION_ONLY=true
-        OS=ubuntu
-        CLEANUP_STATUS=1
-        When call run_base_prep_cleanup
-        The status should equal 224
-        The output should include "basePrep driver cleanup"
-    End
-
-    It 'fails an Ubuntu bake without the required helper before customization'
-        PRE_PROVISION_ONLY=true
-        OS=ubuntu
-        rm "${fixture}/provision_installs_ubuntu_gpu.sh"
-        When call run_base_prep_cleanup
-        The status should equal 224
-        The output should include "image containing PIS GPU cleanup helper required"
-        The output should not include "basePrep driver cleanup"
-    End
-
-    It 'does not require the helper on an ordinary Ubuntu node'
-        OS=ubuntu
-        rm "${fixture}/provision_installs_ubuntu_gpu.sh"
-        When call run_base_prep_cleanup
+    It 'writes the bake marker after successful basePrep'
+        EXIT_CODE=0
+        When call run_cse_marker
         The status should be success
-        The output should be blank
-    End
-
-    It 'does not require the helper for managed GPU image preparation'
-        PRE_PROVISION_ONLY=true
-        GPU_NODE=true
-        OS=ubuntu
-        rm "${fixture}/provision_installs_ubuntu_gpu.sh"
-        When call run_base_prep_cleanup
-        The status should be success
-        The output should be blank
-    End
-
-    It 'keeps Azure Linux bake cleanup independent of the Ubuntu helper'
-        PRE_PROVISION_ONLY=true
-        OS=azurelinux
-        rm "${fixture}/provision_installs_ubuntu_gpu.sh"
-        When call run_base_prep_cleanup
-        The status should be success
-        The output should include "basePrep driver cleanup"
-    End
-
-    Describe 'PIS basePrep completion marker'
-        It 'does not write the 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 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
+        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
index 615c003108e..cb556b84757 100644
--- a/spec/vhdbuilder/packer/pis_gpu_helper_spec.sh
+++ b/spec/vhdbuilder/packer/pis_gpu_helper_spec.sh
@@ -1,36 +1,18 @@
 #!/bin/bash
-# shellcheck disable=SC2329
 
 Describe 'Ubuntu PIS GPU cleanup helper VHD installation'
   Include './vhdbuilder/packer/packer_source.sh'
 
-  isUbuntu() {
-    [ "$1" = UBUNTU ]
-  }
-  cpAndMode() {
-    printf '%s\n' "$*"
-  }
+  isUbuntu() { [ "$1" = UBUNTU ]; }
+  cpAndMode() { printf '%s\n' "$*"; }
 
   It 'copies the helper into Ubuntu images'
     OS=UBUNTU
     When call copyUbuntuGPUCleanupHelper
-    The status should be success
     The output should equal '/home/packer/provision_installs_ubuntu_gpu.sh /opt/azure/containers/provision_installs_ubuntu_gpu.sh 0744'
   End
 
-  Parameters
-    MARINER
-    AZURELINUX
-    FLATCAR
-  End
-  It "does not require the Ubuntu helper on $1"
-    OS="$1"
-    When call copyUbuntuGPUCleanupHelper
-    The status should be success
-    The output should be blank
-  End
-
-  It 'installs the helper when copying Packer files'
+  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

From a547f293ce3fba2265f297c92c786a46234a3867 Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 23 Oct 2026 12:19:52 +1300
Subject: [PATCH 23/24] fix(e2e): honor effective scriptless mode

PIS derived nodes deliberately override the global scriptless default.
Validate CustomData against the effective NBC flag instead.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 e2e/scenario/vmss.go      |  2 +-
 e2e/scenario/vmss_test.go | 61 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+), 1 deletion(-)

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

From 30226157252b762134fde7c712c06f40bfb46e9b Mon Sep 17 00:00:00 2001
From: Ben Brady <86134744+benjamin-brady@users.noreply.github.com>
Date: Fri, 23 Oct 2026 12:23:33 +1300
Subject: [PATCH 24/24] ci(gpu): allow targeted scenario runs

Keep the default GPU suite; allow a single fresh-image PIS scenario
without rerunning every GPU scenario.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
---
 .pipelines/e2e-gpu.yaml | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

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.