Ben Brady (benjamin-brady) · Oct 28, 2026 · Oct 28, 2026 · Oct 28, 2026
@@ -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.

@@ -94,6 +94,17 @@ coverage matching production. A Scriptless tag or an AKSNodeConfig mutator alone
 does not opt into Phase 3. Existing scriptless opt-outs, image restrictions, and
 VHD-caching behavior remain in effect.
 
+### Ubuntu PIS GPU cleanup
+
+`Ubuntu2404_GPUDriverOptOut_VHDCaching` needs a new Ubuntu 24.04 Gen2 image
+containing the baked cleanup helper; use same-run E2E VHD metadata or
+`--vhd-metadata-file` with `2404gen2containerd`. The helper is baked, not
+transferred in CustomData. An opt-out PIS bake without it fails explicitly;
+ordinary nodes and derived PIS nodes do not require it.
+Opt out of the managed driver in bake CustomData (`EnableNvidia=false`,
+rendering `GPU_NODE=false`) before capture; a tag only on the derived VM does
+not request bake cleanup.
+
 ### LocalDNS fallback checks
 
 The existing LocalDNS scenarios use short-lived DNS pods on the scenario node

@@ -70,7 +70,13 @@ func runVHDCachingScenario(ctx context.Context, name string, original *Scenario)
 				ValidateSystemdUnitIsNotRunning(ctx, scenario, "kubelet"),
 			)
 		}
-		return validationErr
+		if validationErr != nil {
+			return validationErr
+		}
+		if original.Config.VHDCachingPostBasePrep != nil {
+			return original.Config.VHDCachingPostBasePrep(ctx, scenario)
+		}
+		return nil
 	}
 
 	bakeScenario.Config.VMConfigMutator = func(vmss *armcompute.VirtualMachineScaleSet) {
@@ -832,6 +838,65 @@ func RunCommand(ctx context.Context, s *Scenario, command string) (armcompute.Vi
 	return view, runCommandScriptError(view)
 }
 
+const (
+	customerGPUDriverImage = "mcr.microsoft.com/aks/aks-gpu-cuda-lts"
+)
+
+func customerGPUDriverImageReference() (string, string, error) {
+	if datamodel.NvidiaCudaLTSDriverVersion == "" || datamodel.AKSGPUCudaLTSVersionSuffix == "" {
+		return "", "", fmt.Errorf("approved NVIDIA CUDA LTS driver image version is unavailable")
+	}
+	version := datamodel.NvidiaCudaLTSDriverVersion
+	tag := version + "-" + datamodel.AKSGPUCudaLTSVersionSuffix
+	return customerGPUDriverImage + ":" + tag, version, nil
+}
+
+// installCustomerNvidiaDriver manually runs the VHD-cached, components-approved installer after
+// basePrep. This models a customer RunCommand install; the scenario leaves managed GPU installation
+// opted out, so nodePrep does not invoke the installer.
+func installCustomerNvidiaDriver(ctx context.Context, s *Scenario) error {
+	imageRef, expectedVersion, err := customerGPUDriverImageReference()
+	if err != nil {
+		return err
+	}
+	command := fmt.Sprintf(`set -euo pipefail
+image_ref=%s
+expected_version=%s
+sudo mkdir -p /opt/gpu /opt/actions
+sudo ctr -n k8s.io images ls -q | grep -Fxq "${image_ref}"
+sudo ctr -n k8s.io run --privileged --rm --net-host \
+  --with-ns pid:/proc/1/ns/pid \
+  --mount type=bind,src=/opt/gpu,dst=/mnt/gpu,options=rbind \
+  --mount type=bind,src=/opt/actions,dst=/mnt/actions,options=rbind \
+  "${image_ref}" gpuinstall /entrypoint.sh copy
+sudo ctr -n k8s.io images rm "${image_ref}"
+
+installer_dir="/opt/gpu/NVIDIA-Linux-x86_64-${expected_version}"
+test -x "${installer_dir}/nvidia-installer"
+cd "${installer_dir}"
+sudo ./nvidia-installer -s -k="$(uname -r)" --no-drm --dkms
+
+sudo nvidia-modprobe -u -c0
+gpu_info="$(sudo nvidia-smi --query-gpu=driver_version,name --format=csv,noheader | head -n1)"
+driver_version="${gpu_info%%,*}"
+driver_version="$(printf '%%s' "${driver_version}" | tr -d '[:space:]')"
+gpu_name="${gpu_info#*,}"
+case "${gpu_name}" in *T4*) ;; *) echo "expected a T4 GPU, got: ${gpu_name}" >&2; exit 1 ;; esac
+test "${driver_version}" = "${expected_version}"
+sudo nvidia-smi -L | grep -qi 'T4'
+lsmod | grep -q '^nvidia[[:space:]]'
+test "$(sudo modinfo -F version nvidia)" = "${expected_version}"
+`,
+		shellSingleQuote(imageRef),
+		shellSingleQuote(expectedVersion),
+	)
+	_, err = RunCommand(ctx, s, command)
+	if err != nil {
+		return fmt.Errorf("install customer NVIDIA driver from cached approved artifact: %w", err)
+	}
+	return nil
+}
+
 func RestartVMSSVM(ctx context.Context, s *Scenario) error {
 	if s.Runtime == nil || s.Runtime.Cluster == nil || s.Runtime.Cluster.Model == nil ||
 		s.Runtime.Cluster.Model.Properties == nil || s.Runtime.Cluster.Model.Properties.NodeResourceGroup == nil ||

@@ -2981,6 +2981,7 @@ var _ = Register(&Scenario{
 			runcVersions := components.GetExpectedPackageVersions("runc", "ubuntu", "r2404")
 			return errors.Join(
 				ValidateNvidiaSMINotInstalled(ctx, s),
+				ValidateNvidiaModulesNotLoaded(ctx, s),
 				ValidateContainerd2Properties(ctx, s, containerdVersions),
 				ValidateRuncVersion(ctx, s, runcVersions),
 			)
@@ -3106,6 +3107,37 @@ var _ = Register(&Scenario{
 	},
 })
 
+var _ = Register(&Scenario{
+	Name:             "Ubuntu2404_GPUDriverOptOut_VHDCaching",
+	Description:      "Installs the VHD-cached, components-approved NVIDIA driver artifact with a post-basePrep customer RunCommand while managed driver installation remains opted out, then verifies the real driver survives PIS VHD-cached provisioning",
+	K8sSystemPoolSKU: "Standard_D2s_v3",
+	Tags: Tags{
+		GPU: true,
+	},
+	Config: Config{
+		Cluster:                ClusterKubenet,
+		VHD:                    config.VHDUbuntu2404Gen2Containerd,
+		VHDCaching:             true,
+		VHDCachingPostBasePrep: installCustomerNvidiaDriver,
+		PreProvisionBootstrapConfigMutator: func(_ *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) {
+			nbc.EnableScriptlessCSECmd = !config.Config.DisableScriptless
+		},
+		BootstrapConfigMutator: func(_ *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) {
+			nbc.EnableScriptlessCSECmd = false
+			nbc.AgentPoolProfile.VMSize = "Standard_NC4as_T4_v3"
+			nbc.ConfigGPUDriverIfNeeded = false
+			nbc.EnableGPUDevicePluginIfNeeded = false
+			nbc.EnableNvidia = false
+		},
+		VMConfigMutator: func(vmss *armcompute.VirtualMachineScaleSet) {
+			vmss.SKU.Name = to.Ptr("Standard_NC4as_T4_v3")
+		},
+		Validator: func(ctx context.Context, s *Scenario) error {
+			return ValidateCustomerNvidiaDriver(ctx, s)
+		},
+	},
+})
+
 var _ = Register(&Scenario{
 	Name:        "Ubuntu2404_AKSVMExtension_FilesystemCorruption",
 	Description: "Tests Ubuntu 24.04 NPD filesystem-corruption reporting with the AKS VM Extension and hotfix target resolution from the NBC command when ANC JSON is absent",

@@ -198,6 +198,10 @@ type Config struct {
 	// The main purpose is to validate VHD Caching logic and ensure a reboot step between basePrep and nodePrep doesn't break anything.
 	VHDCaching bool
 
+	// VHDCachingPostBasePrep runs on the bake VM after basePrep completes and before its disk is captured.
+	// Use it only for VHD-caching scenarios that need to seed state between the two provisioning phases.
+	VHDCachingPostBasePrep func(context.Context, *Scenario) error
+
 	// ExpectedError, when set, indicates that VMSS creation is expected to fail with an error containing this substring.
 	// The assertion is performed during the scenario run.
 	ExpectedError string

@@ -621,12 +621,49 @@ func ValidateNvidiaSMINotInstalled(ctx context.Context, s *Scenario) error {
 	return assert.Contains(execResult.stderr, "nvidia-smi: command not found", "expected stderr to contain 'nvidia-smi: command not found', but got %q", execResult.stderr)
 }
 
+func ValidateNvidiaModulesNotLoaded(ctx context.Context, s *Scenario) error {
+	_, err := execScriptOnVMForScenarioValidateExitCode(ctx, s, nvidiaModulesNotLoadedScript(), 0, "NVIDIA kernel modules remain loaded")
+	return err
+}
+
+func nvidiaModulesNotLoadedScript() string {
+	return `set -euo pipefail
+lsmod | awk '$1 == "nvidia" || $1 ~ /^nvidia_/ { print "unexpected loaded NVIDIA module: " $1 > "/dev/stderr"; exit 1 }'`
+}
+
 func ValidateNvidiaSMIInstalled(ctx context.Context, s *Scenario) error {
 	command := []string{"set -ex", "sudo nvidia-smi"}
 	_, err := execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "could not execute nvidia-smi command")
 	return err
 }
 
+func ValidateCustomerNvidiaDriver(ctx context.Context, s *Scenario) error {
+	_, expectedVersion, err := customerGPUDriverImageReference()
+	if err != nil {
+		return err
+	}
+	installerPath := fmt.Sprintf("/opt/gpu/NVIDIA-Linux-x86_64-%s/nvidia-installer", expectedVersion)
+	command := customerNvidiaDriverValidationScript(expectedVersion, installerPath)
+	_, err = execScriptOnVMForScenarioValidateExitCode(ctx, s, command, 0, "customer-installed NVIDIA driver did not survive PIS image capture")
+	return err
+}
+
+func customerNvidiaDriverValidationScript(expectedVersion, installerPath string) string {
+	return fmt.Sprintf(`set -euo pipefail
+expected_version=%s
+test -x %s
+lsmod | grep -q '^nvidia[[:space:]]'
+module_version="$(sudo modinfo -F version nvidia)"
+test "${module_version}" = "${expected_version}"
+gpu_version="$(sudo nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -n1 | tr -d '[:space:]')"
+test "${gpu_version}" = "${expected_version}"
+sudo nvidia-smi
+`,
+		shellSingleQuote(expectedVersion),
+		shellSingleQuote(installerPath),
+	)
+}
+
 func ValidateNvidiaModProbeInstalled(ctx context.Context, s *Scenario) error {
 	command := []string{
 		"set -ex",

@@ -65,3 +65,31 @@ func TestNvidiaDriverVersionValidationScriptRejectsInvalidVersion(t *testing.T)
 		})
 	}
 }
+
+func TestNvidiaModulesNotLoadedScript(t *testing.T) {
+	tests := []struct {
+		name      string
+		modules   string
+		lsmodExit string
+		wantExit  int
+	}{
+		{"unrelated modules", "Module Size Used by\nx_tables 123 0", "0", 0},
+		{"core module", "Module Size Used by\nnvidia 123 0", "0", 1},
+		{"auxiliary module", "Module Size Used by\nnvidia_uvm 123 0", "0", 1},
+		{"inspection failure", "Module Size Used by", "17", 17},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			command := exec.Command("bash", "-c", `lsmod() { printf '%s\n' "$TEST_LSMOD_OUTPUT"; return "$TEST_LSMOD_EXIT"; }
+`+nvidiaModulesNotLoadedScript())
+			command.Env = append(os.Environ(), "TEST_LSMOD_OUTPUT="+tt.modules, "TEST_LSMOD_EXIT="+tt.lsmodExit)
+			output, err := command.CombinedOutput()
+			if tt.wantExit == 0 {
+				require.NoError(t, err, string(output))
+			} else {
+				require.Error(t, err)
+			}
+			require.Equal(t, tt.wantExit, command.ProcessState.ExitCode(), string(output))
+		})
+	}
+}
@@ -3,13 +3,52 @@ package scenario
 import (
 	"crypto/sha256"
 	"encoding/hex"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
 	"testing"
 
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/stretchr/testify/require"
 	resourcev1 "k8s.io/api/resource/v1"
 )
 
+func TestCustomerNvidiaDriverValidationRequiresStagedInstaller(t *testing.T) {
+	const expectedVersion = "580.159.04"
+	binDir := t.TempDir()
+	installerPath := filepath.Join(binDir, "nvidia-installer")
+	// PATH-local lsmod and sudo stand in for GPU hardware while the real
+	// validation script checks the staged installer and driver state.
+	require.NoError(t, os.WriteFile(filepath.Join(binDir, "lsmod"), []byte("#!/bin/sh\nprintf 'nvidia 123 0\\n'\n"), 0o755))
+	sudo := fmt.Sprintf(`#!/bin/sh
+case "$1" in
+  modinfo) printf '%%s\n' '%s'; exit 0 ;;
+  nvidia-smi)
+    if [ "$2" = "--query-gpu=driver_version" ]; then
+      printf '%%s\n' '%s'
+    fi
+    exit 0
+    ;;
+esac
+exit 1
+`, expectedVersion, expectedVersion)
+	require.NoError(t, os.WriteFile(filepath.Join(binDir, "sudo"), []byte(sudo), 0o755))
+	t.Setenv("PATH", binDir+string(os.PathListSeparator)+os.Getenv("PATH"))
+
+	runValidation := func() error {
+		output, err := exec.Command("bash", "-x", "-c", customerNvidiaDriverValidationScript(expectedVersion, installerPath)).CombinedOutput()
+		if err != nil {
+			return fmt.Errorf("%w: %s", err, output)
+		}
+		return nil
+	}
+	require.Error(t, runValidation(), "validation must reject correct driver state when the staged installer was removed")
+
+	require.NoError(t, os.WriteFile(installerPath, []byte("#!/bin/sh\n"), 0o755))
+	require.NoError(t, runValidation(), "validation must pass when the staged installer and driver state are present")
+}
+
 func TestValidateDRAMIGAllocation(t *testing.T) {
 	tests := []struct {
 		name    string

@@ -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)

@@ -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