fix: when MIG is enabled, start the DRA after MIG service

e2e/README.md

@@ -56,6 +56,15 @@ Use `<OS/image>_<distinguishing configuration>[_<lifecycle>]`, for example
   `nbc.AgentPoolProfile.VMSize` equal. See [vmss.go](scenario/vmss.go).
 - Use `Validator` to check the resulting node state.
 
+### DRA MIG workload validation
+
+The existing `Ubuntu2404_NvidiaDRADriver_With_MIG_Mixed` scenario uses
+`ValidateDRAWithMIGWorkloadSchedulable` with the advertised profile `1g.10gb`.
+It selects NVIDIA MIG devices of that profile, runs a consuming pod on the scenario
+node, and checks the ResourceClaim's allocated driver, pool, and device against that
+node's ResourceSlices before deleting the pod. Missing or mismatched type/profile
+attributes fail validation. Non-MIG DRA scenarios retain their existing scheduling checks.
+
 ### Scriptless provisioning
 
 Phase 3 is explicitly selected with `EnableScriptlessAKSNodeConfig` in the

e2e/scenario/scenario_gpu_managed_experience.go

@@ -959,3 +959,59 @@ var _ = Register(&Scenario{
 		},
 	},
 })
+
+var _ = Register(&Scenario{
+	Name:        "Ubuntu2404_NvidiaDRADriver_With_MIG_Mixed",
+	Description: "Tests DRA driver and workload allocation of MIG 1g.10gb devices with MIG Mixed mode on Ubuntu 24.04 GPU nodes",
+	Location:    "westus2",
+	Tags: Tags{
+		GPU: true,
+	},
+	Config: Config{
+		Cluster:               ClusterKubenet,
+		VHD:                   config.VHDUbuntu2404Gen2Containerd,
+		WaitForSSHAfterReboot: 5 * time.Minute,
+		BootstrapConfigMutator: func(_ *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) {
+			nbc.AgentPoolProfile.VMSize = "Standard_NC24ads_A100_v4"
+			nbc.ConfigGPUDriverIfNeeded = true
+			nbc.EnableNvidia = true
+			nbc.EnableManagedGPUDRA = true
+			nbc.GPUInstanceProfile = "MIG1g"
+			nbc.MigStrategy = "Mixed"
+		},
+		VMConfigMutatorWithError: func(ctx context.Context, vmss *armcompute.VirtualMachineScaleSet) error {
+			vmss.SKU.Name = to.Ptr("Standard_NC24ads_A100_v4")
+
+			extension, err := createVMExtensionLinuxAKSNode(ctx, vmss.Location)
+			if err != nil {
+				return fmt.Errorf("create AKS VM extension: %w", err)
+			}
+			vmss.Properties = addVMExtensionToVMSS(vmss.Properties, extension)
+			return nil
+		},
+		Validator: func(ctx context.Context, s *Scenario) error {
+			containerdVersions := components.GetExpectedPackageVersions("containerd", "ubuntu", "r2404")
+			runcVersions := components.GetExpectedPackageVersions("runc", "ubuntu", "r2404")
+
+			if err := ValidateMIGModeEnabled(ctx, s, 1); err != nil {
+				return err
+			}
+			if err := ValidateMIGInstancesCreated(ctx, s, "MIG 1g.10gb", 7); err != nil {
+				return err
+			}
+			if err := ValidateDraDriverNvidiaGpuServiceRunning(ctx, s); err != nil {
+				return err
+			}
+			if err := errors.Join(
+				ValidateFileExists(ctx, s, "/opt/azure/containers/nvidia-gpu-driver-mode"),
+				ValidateFileHasContent(ctx, s, "/opt/azure/containers/nvidia-gpu-driver-mode", "DRA"),
+				ValidateContainerd2Properties(ctx, s, containerdVersions),
+				ValidateRuncVersion(ctx, s, runcVersions),
+				ValidateContainerRuntimePlugins(ctx, s),
+			); err != nil {
+				return err
+			}
+			return ValidateDRAWithMIGWorkloadSchedulable(ctx, s, "1g.10gb")
+		},
+	},
+})

e2e/scenario/validation.go

@@ -223,6 +223,10 @@ func ValidateServicePrincipalData(ctx context.Context, s *Scenario) error {
 }
 
 func startPodAndCheckItRuns(ctx context.Context, s *Scenario, pod *corev1.Pod) error {
+	return startPodAndCheckItRunsWithValidation(ctx, s, pod, nil)
+}
+
+func startPodAndCheckItRunsWithValidation(ctx context.Context, s *Scenario, pod *corev1.Pod, validate func() error) error {
 	kube := s.Runtime.Kube
 	pod = pod.DeepCopy()
 	pod.Name = uniqueKubernetesResourceName(pod.Name)
@@ -262,6 +266,9 @@ func startPodAndCheckItRuns(ctx context.Context, s *Scenario, pod *corev1.Pod) e
 			pod.Name, pod.Namespace, s.Runtime.VM.KubeName, timeForReady.Round(time.Millisecond), readinessWarningThreshold))
 	logging.Logf(ctx, "node health validation: test pod %q in namespace %q is running on node %q",
 		pod.Name, pod.Namespace, s.Runtime.VM.KubeName)
+	if validate != nil {
+		return validate()
+	}
 	return nil
 }
 

e2e/scenario/validators.go

@@ -4531,9 +4531,32 @@ func ValidateDraDriverNvidiaGpuServiceRunning(ctx context.Context, s *Scenario)
 	return nil
 }
 
-func ValidateDRAWorkloadSchedulable(ctx context.Context, s *Scenario) (err error) {
+func ValidateDRAWorkloadSchedulable(ctx context.Context, s *Scenario) error {
+	return validateDRAWorkloadSchedulable(ctx, s, "")
+}
+
+func ValidateDRAWithMIGWorkloadSchedulable(ctx context.Context, s *Scenario, expectedProfile string) error {
+	if expectedProfile == "" {
+		return fmt.Errorf("expected MIG profile must not be empty")
+	}
+	return validateDRAWorkloadSchedulable(ctx, s, expectedProfile)
+}
+
+func validateDRAWorkloadSchedulable(ctx context.Context, s *Scenario, expectedMIGProfile string) (err error) {
 	logging.Logf(ctx, "validating that DRA workloads can be scheduled")
 
+	var selectors []resourcev1.DeviceSelector
+	if expectedMIGProfile != "" {
+		selectors = []resourcev1.DeviceSelector{{
+			CEL: &resourcev1.CELDeviceSelector{
+				Expression: fmt.Sprintf(
+					`device.driver == "gpu.nvidia.com" && device.attributes["gpu.nvidia.com"].type == "mig" && device.attributes["gpu.nvidia.com"].profile == %q`,
+					expectedMIGProfile,
+				),
+			},
+		}}
+	}
+
 	time.Sleep(20 * time.Second) // Same delay as existing GPU tests
 
 	ownerReference, err := scenarioNodeOwnerReference(ctx, s)
@@ -4549,7 +4572,7 @@ func ValidateDRAWorkloadSchedulable(ctx context.Context, s *Scenario) (err error
 			Name:            deviceClassName,
 			OwnerReferences: []metav1.OwnerReference{ownerReference},
 		},
-		Spec: resourcev1.DeviceClassSpec{},
+		Spec: resourcev1.DeviceClassSpec{Selectors: selectors},
 	}, metav1.CreateOptions{})
 	if createErr != nil {
 		return fmt.Errorf("create DeviceClass %q: %w", deviceClassName, createErr)
@@ -4628,14 +4651,73 @@ func ValidateDRAWorkloadSchedulable(ctx context.Context, s *Scenario) (err error
 			},
 		},
 	}
-	if err := ValidatePodRunning(ctx, s, pod); err != nil {
+	var validateAllocation func() error
+	if expectedMIGProfile != "" {
+		validateAllocation = func() error {
+			allocatedClaim, err := s.Runtime.Kube.Typed.ResourceV1().ResourceClaims(claim.Namespace).Get(ctx, claim.Name, metav1.GetOptions{})
+			if err != nil {
+				return fmt.Errorf("get allocated ResourceClaim %q: %w", claim.Name, err)
+			}
+			resourceSlices, err := s.Runtime.Kube.Typed.ResourceV1().ResourceSlices().List(ctx, metav1.ListOptions{
+				FieldSelector: resourcev1.ResourceSliceSelectorNodeName + "=" + s.Runtime.VM.KubeName,
+			})
+			if err != nil {
+				return fmt.Errorf("list ResourceSlices for node %q: %w", s.Runtime.VM.KubeName, err)
+			}
+			if err := validateDRAMIGAllocation(allocatedClaim, resourceSlices, s.Runtime.VM.KubeName, expectedMIGProfile); err != nil {
+				return fmt.Errorf("validate MIG allocation for ResourceClaim %q: %w", claim.Name, err)
+			}
+			logging.Logf(ctx, "DRA workload allocated a MIG device with profile %s", expectedMIGProfile)
+			return nil
+		}
+	}
+	if err := startPodAndCheckItRunsWithValidation(ctx, s, pod, validateAllocation); err != nil {
 		return fmt.Errorf("run DRA workload pod: %w", err)
 	}
 
 	logging.Logf(ctx, "GPU workload is schedulable and runs successfully")
 	return nil
 }
 
+func validateDRAMIGAllocation(claim *resourcev1.ResourceClaim, resourceSlices *resourcev1.ResourceSliceList, nodeName, expectedProfile string) error {
+	if claim.Status.Allocation == nil {
+		return fmt.Errorf("claim has no allocation")
+	}
+	results := claim.Status.Allocation.Devices.Results
+	if len(results) != 1 {
+		return fmt.Errorf("expected exactly one allocated device, got %d", len(results))
+	}
+	allocated := results[0]
+	if allocated.Request != "gpu" || allocated.Driver != "gpu.nvidia.com" {
+		return fmt.Errorf("expected request gpu from driver gpu.nvidia.com, got request %q from driver %q", allocated.Request, allocated.Driver)
+	}
+	for _, resourceSlice := range resourceSlices.Items {
+		if resourceSlice.Spec.Driver != allocated.Driver || resourceSlice.Spec.Pool.Name != allocated.Pool ||
+			resourceSlice.Spec.NodeName == nil || *resourceSlice.Spec.NodeName != nodeName {
+			continue
+		}
+		for _, device := range resourceSlice.Spec.Devices {
+			if device.Name != allocated.Device {
+				continue
+			}
+			for attribute, expected := range map[resourcev1.QualifiedName]string{"type": "mig", "profile": expectedProfile} {
+				value, ok := device.Attributes[attribute]
+				if !ok {
+					value = device.Attributes[resourcev1.QualifiedName(allocated.Driver+"/"+string(attribute))]
+				}
+				if value.StringValue == nil {
+					return fmt.Errorf("allocated device %q has no string attribute %q", allocated.Device, attribute)
+				}
+				if *value.StringValue != expected {
+					return fmt.Errorf("allocated device %q has %s %q, expected %q", allocated.Device, attribute, *value.StringValue, expected)
+				}
+			}
+			return nil
+		}
+	}
+	return fmt.Errorf("allocated device %q from driver %q and pool %q not found on node %q", allocated.Device, allocated.Driver, allocated.Pool, nodeName)
+}
+
 // ValidateRCV1PCertMode validates that the rcv1p certificate endpoint mode was used during
 // Linux node provisioning, certificates were downloaded and installed, and a refresh task was scheduled.
 func ValidateRCV1PCertMode(ctx context.Context, s *Scenario) error {

e2e/scenario/validators_test.go

@@ -5,9 +5,184 @@ import (
 	"encoding/hex"
 	"testing"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/stretchr/testify/require"
+	resourcev1 "k8s.io/api/resource/v1"
 )
 
+func TestValidateDRAMIGAllocation(t *testing.T) {
+	tests := []struct {
+		name    string
+		mutate  func(*resourcev1.ResourceClaim, *resourcev1.ResourceSliceList)
+		wantErr string
+	}{
+		{name: "matching allocated MIG device"},
+		{
+			name: "fully qualified attributes",
+			mutate: func(_ *resourcev1.ResourceClaim, slices *resourcev1.ResourceSliceList) {
+				attrs := slices.Items[0].Spec.Devices[0].Attributes
+				attrs["gpu.nvidia.com/type"] = attrs["type"]
+				attrs["gpu.nvidia.com/profile"] = attrs["profile"]
+				delete(attrs, "type")
+				delete(attrs, "profile")
+			},
+		},
+		{
+			name: "missing allocation",
+			mutate: func(claim *resourcev1.ResourceClaim, _ *resourcev1.ResourceSliceList) {
+				claim.Status.Allocation = nil
+			},
+			wantErr: "no allocation",
+		},
+		{
+			name: "empty allocation",
+			mutate: func(claim *resourcev1.ResourceClaim, _ *resourcev1.ResourceSliceList) {
+				claim.Status.Allocation.Devices.Results = nil
+			},
+			wantErr: "exactly one allocated device",
+		},
+		{
+			name: "multiple allocated devices",
+			mutate: func(claim *resourcev1.ResourceClaim, _ *resourcev1.ResourceSliceList) {
+				results := &claim.Status.Allocation.Devices.Results
+				*results = append(*results, (*results)[0])
+			},
+			wantErr: "exactly one allocated device",
+		},
+		{
+			name: "wrong request",
+			mutate: func(claim *resourcev1.ResourceClaim, _ *resourcev1.ResourceSliceList) {
+				claim.Status.Allocation.Devices.Results[0].Request = "other"
+			},
+			wantErr: `got request "other"`,
+		},
+		{
+			name: "wrong allocated driver",
+			mutate: func(claim *resourcev1.ResourceClaim, _ *resourcev1.ResourceSliceList) {
+				claim.Status.Allocation.Devices.Results[0].Driver = "other.example.com"
+			},
+			wantErr: `from driver "other.example.com"`,
+		},
+		{
+			name: "missing resource slices",
+			mutate: func(_ *resourcev1.ResourceClaim, slices *resourcev1.ResourceSliceList) {
+				slices.Items = nil
+			},
+			wantErr: "not found on node",
+		},
+		{
+			name: "wrong slice driver",
+			mutate: func(_ *resourcev1.ResourceClaim, slices *resourcev1.ResourceSliceList) {
+				slices.Items[0].Spec.Driver = "other.example.com"
+			},
+			wantErr: "not found on node",
+		},
+		{
+			name: "wrong pool",
+			mutate: func(_ *resourcev1.ResourceClaim, slices *resourcev1.ResourceSliceList) {
+				slices.Items[0].Spec.Pool.Name = "other-pool"
+			},
+			wantErr: "not found on node",
+		},
+		{
+			name: "wrong node",
+			mutate: func(_ *resourcev1.ResourceClaim, slices *resourcev1.ResourceSliceList) {
+				slices.Items[0].Spec.NodeName = to.Ptr("other-node")
+			},
+			wantErr: "not found on node",
+		},
+		{
+			name: "missing node",
+			mutate: func(_ *resourcev1.ResourceClaim, slices *resourcev1.ResourceSliceList) {
+				slices.Items[0].Spec.NodeName = nil
+			},
+			wantErr: "not found on node",
+		},
+		{
+			name: "matching device exists but was not allocated",
+			mutate: func(claim *resourcev1.ResourceClaim, _ *resourcev1.ResourceSliceList) {
+				claim.Status.Allocation.Devices.Results[0].Device = "other-device"
+			},
+			wantErr: "not found on node",
+		},
+		{
+			name: "full GPU instead of MIG",
+			mutate: func(_ *resourcev1.ResourceClaim, slices *resourcev1.ResourceSliceList) {
+				slices.Items[0].Spec.Devices[0].Attributes["type"] = resourcev1.DeviceAttribute{StringValue: to.Ptr("gpu")}
+			},
+			wantErr: `has type "gpu", expected "mig"`,
+		},
+		{
+			name: "wrong MIG profile",
+			mutate: func(_ *resourcev1.ResourceClaim, slices *resourcev1.ResourceSliceList) {
+				slices.Items[0].Spec.Devices[0].Attributes["profile"] = resourcev1.DeviceAttribute{StringValue: to.Ptr("2g.20gb")}
+			},
+			wantErr: `has profile "2g.20gb", expected "1g.10gb"`,
+		},
+		{
+			name: "missing type",
+			mutate: func(_ *resourcev1.ResourceClaim, slices *resourcev1.ResourceSliceList) {
+				delete(slices.Items[0].Spec.Devices[0].Attributes, "type")
+			},
+			wantErr: `no string attribute "type"`,
+		},
+		{
+			name: "missing profile",
+			mutate: func(_ *resourcev1.ResourceClaim, slices *resourcev1.ResourceSliceList) {
+				delete(slices.Items[0].Spec.Devices[0].Attributes, "profile")
+			},
+			wantErr: `no string attribute "profile"`,
+		},
+		{
+			name: "non-string profile",
+			mutate: func(_ *resourcev1.ResourceClaim, slices *resourcev1.ResourceSliceList) {
+				slices.Items[0].Spec.Devices[0].Attributes["profile"] = resourcev1.DeviceAttribute{IntValue: to.Ptr(int64(1))}
+			},
+			wantErr: `no string attribute "profile"`,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			claim := &resourcev1.ResourceClaim{
+				Status: resourcev1.ResourceClaimStatus{
+					Allocation: &resourcev1.AllocationResult{
+						Devices: resourcev1.DeviceAllocationResult{
+							Results: []resourcev1.DeviceRequestAllocationResult{{
+								Request: "gpu", Driver: "gpu.nvidia.com", Pool: "scenario-node", Device: "mig-device",
+							}},
+						},
+					},
+				},
+			}
+			slices := &resourcev1.ResourceSliceList{
+				Items: []resourcev1.ResourceSlice{{
+					Spec: resourcev1.ResourceSliceSpec{
+						Driver:   "gpu.nvidia.com",
+						Pool:     resourcev1.ResourcePool{Name: "scenario-node"},
+						NodeName: to.Ptr("scenario-node"),
+						Devices: []resourcev1.Device{{
+							Name: "mig-device",
+							Attributes: map[resourcev1.QualifiedName]resourcev1.DeviceAttribute{
+								"type":    {StringValue: to.Ptr("mig")},
+								"profile": {StringValue: to.Ptr("1g.10gb")},
+							},
+						}},
+					},
+				}},
+			}
+			if tt.mutate != nil {
+				tt.mutate(claim, slices)
+			}
+			err := validateDRAMIGAllocation(claim, slices, "scenario-node", "1g.10gb")
+			if tt.wantErr != "" {
+				require.ErrorContains(t, err, tt.wantErr)
+			} else {
+				require.NoError(t, err)
+			}
+		})
+	}
+}
+
 func TestValidateSysctlOutput(t *testing.T) {
 	tests := []struct {
 		name     string