diff --git a/e2e/README.md b/e2e/README.md index 64c80a3b025..225289074c4 100644 --- a/e2e/README.md +++ b/e2e/README.md @@ -56,6 +56,15 @@ Use `_[_]`, 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 diff --git a/e2e/scenario/scenario_gpu_managed_experience.go b/e2e/scenario/scenario_gpu_managed_experience.go index c6d22a28b07..a9ca4ed7584 100644 --- a/e2e/scenario/scenario_gpu_managed_experience.go +++ b/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") + }, + }, +}) diff --git a/e2e/scenario/validation.go b/e2e/scenario/validation.go index dfa712ca21b..850eaf18699 100644 --- a/e2e/scenario/validation.go +++ b/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 } diff --git a/e2e/scenario/validators.go b/e2e/scenario/validators.go index 13bb2267da8..afaceac3039 100644 --- a/e2e/scenario/validators.go +++ b/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,7 +4651,27 @@ 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) } @@ -4636,6 +4679,45 @@ func ValidateDRAWorkloadSchedulable(ctx context.Context, s *Scenario) (err error 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 { diff --git a/e2e/scenario/validators_test.go b/e2e/scenario/validators_test.go index 4d21a8b0e2f..1b23beebaad 100644 --- a/e2e/scenario/validators_test.go +++ b/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 diff --git a/parts/linux/cloud-init/artifacts/cse_config_gpu.sh b/parts/linux/cloud-init/artifacts/cse_config_gpu.sh index 8dc03879ced..3ad27398952 100755 --- a/parts/linux/cloud-init/artifacts/cse_config_gpu.sh +++ b/parts/linux/cloud-init/artifacts/cse_config_gpu.sh @@ -447,7 +447,18 @@ EOF local NODE_NAME_DRA NODE_NAME_DRA=$(echo "${NODE_NAME}" | tr '[:upper:]' '[:lower:]') - tee "${DRA_DRIVER_NVIDIA_GPU_OVERRIDE_DIR}/10-dra-driver-nvidia-gpu.conf" > /dev/null < /dev/null < /dev/null <