fix(gpu): preserve customer-installed PIS drivers

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

e2e/README.md

@@ -94,6 +94,14 @@ 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.
+
 ### LocalDNS fallback checks
 
 The existing LocalDNS scenarios use short-lived DNS pods on the scenario node

e2e/scenario/provision.go

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

e2e/scenario/scenario.go

@@ -3106,6 +3106,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",

e2e/scenario/types.go

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

e2e/scenario/validators.go

@@ -627,6 +627,33 @@ func ValidateNvidiaSMIInstalled(ctx context.Context, s *Scenario) error {
 	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",

e2e/scenario/validators_test.go

@@ -3,13 +3,50 @@ 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")
+	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

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)

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

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
 }
 

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

@@ -407,6 +407,18 @@ 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
+        if [ "${OS}" = "${UBUNTU_OS_NAME}" ]; then
+            if [ ! -r /opt/azure/containers/provision_installs_ubuntu_gpu.sh ]; then
+                echo "image containing PIS GPU cleanup helper required"
+                exit $ERR_NVIDIA_DRIVER_INSTALL
+            fi
+            source /opt/azure/containers/provision_installs_ubuntu_gpu.sh || exit $ERR_NVIDIA_DRIVER_INSTALL
+        fi
+        logs_to_events "AKS.CSE.cleanUpGPUDriversForBasePrep" cleanUpGPUDriversForBasePrep || exit $ERR_NVIDIA_DRIVER_INSTALL
+    fi
+
     # Enable localdns to handle node and pod DNS traffic via a local CoreDNS instance.
     # If hosts plugin is enabled, enableAKSLocalDNSHostsSetup() was already called at the
     # very start of basePrep (before disableSystemdResolved) to give the timer a head start
@@ -645,7 +657,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
 
@@ -753,6 +765,7 @@ logs_to_events "AKS.CSE.configureNodeTimeSync" configure_node_time_sync_or_repor
 
 if [ ! -f /opt/azure/containers/base_prep.complete ]; then
     basePrep
+    BASE_PREP_RAN=true
 else
     echo "Skipping basePrep - base_prep.complete file exists"
 fi