fix: install absent or incomplete managed Ubuntu GPU drivers

e2e/scenario/scenario.go

@@ -1469,6 +1469,28 @@ func customNodeConfig(vhd *config.Image, cluster func(context.Context, ClusterRe
 
 var _ = Register(ubuntu2204GPUScenario("Ubuntu2204_GPUNC", "Standard_NC4as_T4_v3", "westus2"))
 var _ = Register(ubuntu2204GPUScenario("Ubuntu2204_GPUA100", "Standard_NC24ads_A100_v4", "westus2"))
+var _ = Register(ubuntuGPUValidationScenario("Ubuntu2204_GPUNC_ValidateDriver", config.VHDUbuntu2204Gen2Containerd))
+var _ = Register(ubuntuGPUValidationScenario("Ubuntu2404_GPUNC_ValidateDriver", config.VHDUbuntu2404Gen2Containerd))
+
+// Validation-only bootstrap is a distinct input from the existing unconditional-install GPU nodes.
+func ubuntuGPUValidationScenario(name string, vhd *config.Image) *Scenario {
+	s := ubuntu2204GPUScenario(name, "Standard_NC4as_T4_v3", "westus2")
+	s.Description = "Tests managed Ubuntu GPU validation and its full-install fallback without requesting unconditional driver installation"
+	s.Config.VHD = vhd
+	configureGPU := s.Config.BootstrapConfigMutator
+	s.Config.BootstrapConfigMutator = func(cluster *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) {
+		configureGPU(cluster, nbc)
+		nbc.ConfigGPUDriverIfNeeded = false
+		// Exercise this checkout's validation code, not the VHD's baked-in CSE.
+		nbc.EnableScriptlessCSECmd = false
+		nbc.EnableScriptlessNBCCSECmd = false
+	}
+	validateGPU := s.Config.Validator
+	s.Config.Validator = func(ctx context.Context, s *Scenario) error {
+		return errors.Join(validateGPU(ctx, s), ValidateNvidiaSMIInstalled(ctx, s))
+	}
+	return s
+}
 
 var _ = Register(&Scenario{
 	Name:        "Ubuntu2204_GPUA10",

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 s.Runtime.NBC.EnableScriptlessCSECmd && !scriptlessNBCCSECmdEnabled && 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

@@ -1,6 +1,8 @@
 package scenario
 
 import (
+	"bytes"
+	"compress/gzip"
 	"context"
 	"encoding/base64"
 	"errors"
@@ -441,6 +443,85 @@ func TestGetBaseVMSSModelUsesScenarioVMSize(t *testing.T) {
 	assert.Equal(t, s.Runtime.VMSize, scenarioVMSize(s))
 }
 
+func TestCreateVMSSModelUsesEffectiveCSEMode(t *testing.T) {
+	oldDisableScriptless := config.Config.DisableScriptless
+	oldDisableCompilation := config.Config.DisableScriptLessCompilation
+	oldPrepareVHD := CachedPrepareVHD
+	t.Cleanup(func() {
+		config.Config.DisableScriptless = oldDisableScriptless
+		config.Config.DisableScriptLessCompilation = oldDisableCompilation
+		CachedPrepareVHD = oldPrepareVHD
+	})
+	config.Config.DisableScriptLessCompilation = true
+	CachedPrepareVHD = func(context.Context, GetVHDRequest) (config.VHDResourceID, error) {
+		return "/subscriptions/test/resourceGroups/test/providers/Microsoft.Compute/images/test", nil
+	}
+
+	for _, image := range []*config.Image{config.VHDUbuntu2204Gen2Containerd, config.VHDUbuntu2404Gen2Containerd} {
+		for _, tc := range []struct {
+			name              string
+			disableScriptless bool
+			cseCmd            bool
+			nbcCmd            bool
+			wantFile          string
+		}{
+			{name: "legacy override", wantFile: "/opt/azure/containers/provision_configs_gpu.sh"},
+			{name: "legacy global default", disableScriptless: true, wantFile: "/opt/azure/containers/provision_configs_gpu.sh"},
+			{name: "scriptless CSE", cseCmd: true, wantFile: "/opt/azure/containers/scriptless-cse-overrides.txt"},
+			{name: "scriptless NBC", nbcCmd: true, wantFile: "/opt/azure/containers/aks-node-controller-nbc-cmd.sh"},
+		} {
+			t.Run(string(image.Distro)+"/"+tc.name, func(t *testing.T) {
+				config.Config.DisableScriptless = tc.disableScriptless
+				cluster := &Cluster{
+					Model: &armcontainerservice.ManagedCluster{
+						Name:     to.Ptr("test-cluster"),
+						Location: to.Ptr("eastus"),
+						Properties: &armcontainerservice.ManagedClusterProperties{
+							NodeResourceGroup:        to.Ptr("test-rg"),
+							CurrentKubernetesVersion: to.Ptr("1.33.0"),
+							NetworkProfile: &armcontainerservice.NetworkProfile{
+								NetworkPlugin: to.Ptr(armcontainerservice.NetworkPluginKubenet),
+							},
+						},
+					},
+					KubeletIdentity: &armcontainerservice.UserAssignedIdentity{
+						ClientID:   to.Ptr("test-client"),
+						ResourceID: to.Ptr("/subscriptions/test/resourceGroups/test/providers/Microsoft.ManagedIdentity/userAssignedIdentities/test"),
+					},
+					ClusterParams: &ClusterParams{BootstrapToken: "test-token", FQDN: "cluster.example.invalid"},
+				}
+				nbc, err := getBaseNBC(t.Context(), cluster, image)
+				require.NoError(t, err)
+				nbc.EnableScriptlessCSECmd = tc.cseCmd
+				nbc.EnableScriptlessNBCCSECmd = tc.nbcCmd
+				s := &Scenario{
+					Config: Config{VHD: image},
+					Runtime: &ScenarioRuntime{
+						Cluster:                   cluster,
+						NBC:                       nbc,
+						VMSSName:                  "test-vmss",
+						VMSize:                    config.DEFAULT_VMSKU,
+						EnableScriptlessNBCCSECmd: true,
+					},
+				}
+				s.Location = "eastus"
+				model, err := createVMSSModel(t.Context(), s)
+				require.NoError(t, err)
+				data, err := base64.StdEncoding.DecodeString(*model.Properties.VirtualMachineProfile.OSProfile.CustomData)
+				require.NoError(t, err)
+				if !tc.nbcCmd {
+					reader, err := gzip.NewReader(bytes.NewReader(data))
+					require.NoError(t, err)
+					data, err = io.ReadAll(reader)
+					require.NoError(t, reader.Close())
+					require.NoError(t, err)
+				}
+				assert.Contains(t, string(data), tc.wantFile)
+			})
+		}
+	}
+}
+
 func TestCSEExitCodeOutboundConnFail(t *testing.T) {
 	require.Equal(t, "50", cseExitCodeOutboundConnFail)
 }

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

@@ -108,7 +108,28 @@ configGPUDrivers() {
 }
 
 validateGPUDrivers() {
-    retrycmd_if_failure 24 5 25 nvidia-modprobe -u -c0 && echo "gpu driver loaded" || configGPUDrivers || exit $ERR_GPU_DRIVERS_START_FAIL
+    local driver_payload_missing=false
+    if [ "$OS" = "$UBUNTU_OS_NAME" ] && [ "${IS_VHD,,}" = "true" ] &&
+        [ "$(getCPUArch)" = "amd64" ]; then
+        if ! command -v nvidia-modprobe >/dev/null 2>&1; then
+            driver_payload_missing=true
+        elif command -v nvidia-smi >/dev/null 2>&1; then
+            local smi_status=0
+            timeout 30 nvidia-smi >/dev/null 2>&1 || smi_status=$?
+            # NVIDIA exit 12 means NVML cannot be loaded, even if a prebaked kernel module loads.
+            if [ "$smi_status" -eq 12 ]; then
+                driver_payload_missing=true
+            fi
+        fi
+    fi
+
+    # Retrying absent payload cannot make it ready and consumes 115 seconds of the CSE budget.
+    if [ "$driver_payload_missing" = true ]; then
+        echo "NVIDIA driver payload is absent or incomplete in this Ubuntu VHD; installing before validation"
+        configGPUDrivers || exit $ERR_GPU_DRIVERS_START_FAIL
+    else
+        retrycmd_if_failure 24 5 25 nvidia-modprobe -u -c0 && echo "gpu driver loaded" || configGPUDrivers || exit $ERR_GPU_DRIVERS_START_FAIL
+    fi
 
     if which nvidia-smi; then
         SMI_RESULT=$(retrycmd_if_failure 24 5 30 nvidia-smi)

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

@@ -74,6 +74,204 @@ Describe 'cse_config_gpu.sh'
         End
     End
 
+    Describe 'validation on managed Ubuntu VHDs'
+        OS="$UBUNTU_OS_NAME"
+        IS_VHD=true
+        MOCK_CPU_ARCH=amd64
+        MOCK_MODPROBE_PRESENT=1
+        MOCK_MODPROBE_RESULT=0
+        MOCK_SMI_PRESENT=0
+        MOCK_SMI_PROBE_RESULT=0
+        MOCK_SMI_VALIDATION_RESULT=0
+        MOCK_SMI_OUTPUT=""
+        MOCK_INSTALL_RESULT=0
+        ERR_GPU_DRIVERS_START_FAIL=84
+
+        getCPUArch() { echo "$MOCK_CPU_ARCH"; }
+        command() {
+            if [ "$*" = "-v nvidia-modprobe" ]; then
+                return "$MOCK_MODPROBE_PRESENT"
+            fi
+            if [ "$*" = "-v nvidia-smi" ]; then
+                return "$MOCK_SMI_PRESENT"
+            fi
+            builtin command "$@"
+        }
+        timeout() {
+            [ "$*" = "30 nvidia-smi" ] || return 99
+            return "$MOCK_SMI_PROBE_RESULT"
+        }
+        retrycmd_if_failure() {
+            echo "retry $*" >&2
+            if [ "$4" = "nvidia-modprobe" ]; then
+                return "$MOCK_MODPROBE_RESULT"
+            fi
+            printf '%s' "$MOCK_SMI_OUTPUT"
+            return "$MOCK_SMI_VALIDATION_RESULT"
+        }
+        configGPUDrivers() {
+            echo "normal driver installation"
+            return "$MOCK_INSTALL_RESULT"
+        }
+        which() { return 0; }
+
+        It 'installs an absent driver without retrying the missing executable'
+            When run validateGPUDrivers
+            The status should be success
+            The output should include "installing before validation"
+            The output should include "normal driver installation"
+            The output should include "gpu driver working fine"
+            The stderr should not include "nvidia-modprobe"
+            The stderr should include "retry 24 5 30 nvidia-smi"
+        End
+
+        It 'accepts the existing uppercase VHD flag'
+            IS_VHD=TRUE
+            When run validateGPUDrivers
+            The status should be success
+            The output should include "normal driver installation"
+            The stderr should not include "nvidia-modprobe"
+            The stderr should include "nvidia-smi"
+        End
+
+        It 'propagates installation failure without reporting successful validation'
+            MOCK_INSTALL_RESULT=1
+            When run validateGPUDrivers
+            The status should equal 84
+            The output should include "normal driver installation"
+            The output should not include "gpu driver working fine"
+            The stderr should equal ""
+        End
+
+        It 'keeps readiness retries for a present driver'
+            MOCK_MODPROBE_PRESENT=0
+            When run validateGPUDrivers
+            The status should be success
+            The output should include "gpu driver loaded"
+            The output should not include "normal driver installation"
+            The stderr should include "retry 24 5 25 nvidia-modprobe -u -c0"
+            The stderr should include "nvidia-smi"
+        End
+
+        It 'preserves the full-install fallback after an installed driver fails to load'
+            MOCK_MODPROBE_PRESENT=0
+            MOCK_MODPROBE_RESULT=1
+            When run validateGPUDrivers
+            The status should be success
+            The output should include "normal driver installation"
+            The output should not include "installing before validation"
+            The stderr should include "retry 24 5 25 nvidia-modprobe -u -c0"
+            The stderr should include "nvidia-smi"
+        End
+
+        It 'installs missing NVML even when nvidia-modprobe is present'
+            MOCK_MODPROBE_PRESENT=0
+            MOCK_SMI_PROBE_RESULT=12
+            When run validateGPUDrivers
+            The status should be success
+            The output should include "installing before validation"
+            The output should include "normal driver installation"
+            The output should include "gpu driver working fine"
+            The stderr should not include "nvidia-modprobe"
+            The stderr should include "retry 24 5 30 nvidia-smi"
+        End
+
+        It 'propagates failure to install missing NVML'
+            MOCK_MODPROBE_PRESENT=0
+            MOCK_SMI_PROBE_RESULT=12
+            MOCK_INSTALL_RESULT=1
+            When run validateGPUDrivers
+            The status should equal 84
+            The output should include "normal driver installation"
+            The output should not include "gpu driver working fine"
+            The stderr should equal ""
+        End
+
+        It 'requires successful validation after installing missing NVML'
+            MOCK_MODPROBE_PRESENT=0
+            MOCK_SMI_PROBE_RESULT=12
+            MOCK_SMI_VALIDATION_RESULT=1
+            When run validateGPUDrivers
+            The status should equal 84
+            The output should include "normal driver installation"
+            The output should not include "gpu driver working fine"
+            The stderr should include "retry 24 5 30 nvidia-smi"
+        End
+
+        Describe 'other nvidia-smi probe outcomes'
+            Parameters
+                0
+                9
+                15
+                124
+            End
+
+            It "keeps readiness retries rather than forcing installation for status=$1"
+                MOCK_MODPROBE_PRESENT=0
+                MOCK_SMI_PROBE_RESULT=$1
+                When run validateGPUDrivers
+                The status should be success
+                The output should not include "normal driver installation"
+                The output should include "gpu driver working fine"
+                The stderr should include "retry 24 5 25 nvidia-modprobe -u -c0"
+                The stderr should include "retry 24 5 30 nvidia-smi"
+            End
+        End
+
+        It 'does not mask a persistent GPU error as a missing library'
+            MOCK_MODPROBE_PRESENT=0
+            MOCK_SMI_PROBE_RESULT=15
+            MOCK_SMI_VALIDATION_RESULT=1
+            When run validateGPUDrivers
+            The status should equal 84
+            The output should not include "normal driver installation"
+            The output should not include "gpu driver working fine"
+            The stderr should include "retry 24 5 30 nvidia-smi"
+        End
+
+        It 'preserves the infoROM corruption error'
+            MOCK_MODPROBE_PRESENT=0
+            MOCK_SMI_PROBE_RESULT=14
+            MOCK_SMI_VALIDATION_RESULT=1
+            MOCK_SMI_OUTPUT="infoROM is corrupted"
+            When run validateGPUDrivers
+            The status should equal 87
+            The output should not include "normal driver installation"
+            The output should not include "gpu driver working fine"
+            The stderr should include "retry 24 5 30 nvidia-smi"
+        End
+
+        It 'preserves the existing alternate nvidia-smi path'
+            MOCK_MODPROBE_PRESENT=0
+            MOCK_SMI_PRESENT=1
+            GPU_DEST=/opt/gpu
+            which() { return 1; }
+            When run validateGPUDrivers
+            The status should be success
+            The output should not include "normal driver installation"
+            The stderr should include "retry 24 5 30 /opt/gpu/bin/nvidia-smi"
+        End
+
+        Parameters
+            UBUNTU false amd64
+            UBUNTU "" amd64
+            UBUNTU true arm64
+            AZURELINUX true amd64
+        End
+
+        It "preserves legacy validation for OS=$1 VHD=$2 arch=$3"
+            OS=$1
+            IS_VHD=$2
+            MOCK_CPU_ARCH=$3
+            When run validateGPUDrivers
+            The status should be success
+            The output should include "gpu driver loaded"
+            The output should not include "normal driver installation"
+            The stderr should include "retry 24 5 25 nvidia-modprobe -u -c0"
+            The stderr should include "nvidia-smi"
+        End
+    End
+
     Describe 'logGPUDriverPrebakeReadiness'
         It 'reports marker_present=false when no prebake marker exists'
             GPU_DKMS_MARKER_FILE="$(mktemp)"; rm -f "${GPU_DKMS_MARKER_FILE}"

vhdbuilder/packer/README.md

@@ -28,6 +28,17 @@ Goal1: remove mariner workflow so things will be simplified.
 - `cse_config_network.sh` -> `provision_configs_network.sh`
 - `cse_config_addons.sh` -> `provision_configs_addons.sh` (autoscaler, ACI connector, Azure Policy)
 
+On managed Ubuntu amd64 VHD nodes, GPU validation selects the normal installer
+immediately if `nvidia-modprobe` is absent or a bounded `nvidia-smi` probe returns
+12 (NVML cannot be loaded). Neither a cached installer image nor a loadable
+prebaked kernel module proves the host userspace driver is installed. Other
+probe outcomes retain the existing readiness retries and full-install fallback;
+ARM64, other OSes and non-VHD images retain their existing validation path.
+GPU eligibility and opt-out decisions still happen in `nodePrep`. The Ubuntu
+validation-only E2E scenarios explicitly deliver their checkout's CSE rather
+than using the tested VHD's baked-in scripts. VMSS artifact validation follows
+the effective NBC delivery flags, not the run-wide scriptless default.
+
 `cse_cmd.sh` and the ANC parser provide their paths through
 `CSE_CONFIG_GPU_FILEPATH`, `CSE_CONFIG_LOCALDNS_FILEPATH`,
 `CSE_CONFIG_KUBELET_FILEPATH`, `CSE_CONFIG_NETWORK_FILEPATH`, and