diff --git a/e2e/scenario_gpu_managed_experience_test.go b/e2e/scenario_gpu_managed_experience_test.go index 23305e15ddf..ebeafcb9d6e 100644 --- a/e2e/scenario_gpu_managed_experience_test.go +++ b/e2e/scenario_gpu_managed_experience_test.go @@ -574,10 +574,10 @@ func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG(t *testing.T) { ValidateNvidiaDevicePluginServiceRunning(ctx, s) // Validate that MIG mode is enabled via nvidia-smi - ValidateMIGModeEnabled(ctx, s) + ValidateMIGModeEnabled(ctx, s, 1) // Validate that MIG instances are created - ValidateMIGInstancesCreated(ctx, s, "MIG 2g.20gb") + ValidateMIGInstancesCreated(ctx, s, "MIG 2g.20gb", 3) // Validate that GPU resources are advertised by the device plugin ValidateNodeAdvertisesGPUResources(ctx, s, 3, "nvidia.com/gpu") @@ -611,6 +611,56 @@ func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG(t *testing.T) { }) } +func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG_MultiGPU(t *testing.T) { + const ( + gpuCount = 2 + migInstancesPerGPU = 3 + totalMIGInstances = gpuCount * migInstancesPerGPU + multiGPUA100VMSize = "Standard_NC48ads_A100_v4" + ) + + RunScenario(t, &Scenario{ + Description: "Tests that a MIG profile is applied to every GPU on an Ubuntu 24.04 multi-GPU VM", + Location: "westus3", + K8sSystemPoolSKU: "Standard_D2s_v3", + Tags: Tags{ + GPU: true, + }, + Config: Config{ + Cluster: ClusterKubenet, + VHD: config.VHDUbuntu2404Gen2Containerd, + WaitForSSHAfterReboot: 5 * time.Minute, + BootstrapConfigMutator: func(_ *Cluster, nbc *datamodel.NodeBootstrappingConfiguration) { + nbc.AgentPoolProfile.VMSize = multiGPUA100VMSize + nbc.ConfigGPUDriverIfNeeded = true + nbc.EnableGPUDevicePluginIfNeeded = true + nbc.EnableNvidia = true + nbc.GPUInstanceProfile = "MIG2g" + nbc.EnableManagedGPU = true + nbc.MigStrategy = "Single" + }, + VMConfigMutator: func(vmss *armcompute.VirtualMachineScaleSet) { + vmss.SKU.Name = to.Ptr(multiGPUA100VMSize) + + extension, err := createVMExtensionLinuxAKSNode(t.Context(), vmss.Location) + require.NoError(t, err, "creating AKS VM extension") + vmss.Properties = addVMExtensionToVMSS(vmss.Properties, extension) + }, + Validator: func(ctx context.Context, s *Scenario) { + versions := components.GetExpectedPackageVersions("nvidia-device-plugin", "ubuntu", "r2404") + require.Lenf(s.T, versions, 1, "Expected exactly one nvidia-device-plugin version for ubuntu r2404 but got %d", len(versions)) + ValidateInstalledPackageVersion(ctx, s, "nvidia-device-plugin", versions[0]) + + ValidateNvidiaDevicePluginServiceRunning(ctx, s) + ValidateMIGModeEnabled(ctx, s, gpuCount) + ValidateMIGInstancesCreated(ctx, s, "MIG 2g.20gb", totalMIGInstances) + ValidateNodeAdvertisesGPUResources(ctx, s, totalMIGInstances, "nvidia.com/gpu") + ValidateGPUWorkloadSchedulable(ctx, s, 1, "nvidia.com/gpu") + }, + }, + }) +} + func Test_Ubuntu2204_NvidiaDevicePluginRunning_WithoutVMSSTag(t *testing.T) { RunScenario(t, &Scenario{ Description: "Tests that NVIDIA device plugin and DCGM Exporter work via NBC EnableManagedGPU field without VMSS tag", @@ -766,10 +816,10 @@ func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG_Mixed(t *testing.T) { ValidateNvidiaDevicePluginServiceRunning(ctx, s) // Validate that MIG mode is enabled via nvidia-smi - ValidateMIGModeEnabled(ctx, s) + ValidateMIGModeEnabled(ctx, s, 1) // Validate that MIG instances are created - ValidateMIGInstancesCreated(ctx, s, "MIG 1g.10gb") + ValidateMIGInstancesCreated(ctx, s, "MIG 1g.10gb", 7) // Validate that MIG profile-specific GPU resources are advertised by the device plugin migResourceName := "nvidia.com/mig-1g.10gb" diff --git a/e2e/validators.go b/e2e/validators.go index abe8f2a3ab3..abe23b98a99 100644 --- a/e2e/validators.go +++ b/e2e/validators.go @@ -2589,26 +2589,29 @@ func ValidateNvidiaDCGMExporterScrapeCommonMetric(ctx context.Context, s *Scenar execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "Nvidia DCGM Exporter is not returning "+metric) } -func ValidateMIGModeEnabled(ctx context.Context, s *Scenario) { +func ValidateMIGModeEnabled(ctx context.Context, s *Scenario, gpuCountExpected int) { s.T.Helper() - s.T.Logf("validating that MIG mode is enabled") + s.T.Logf("validating that MIG mode is enabled on %d GPUs", gpuCountExpected) command := []string{ "set -ex", - // Grep to verify it contains 'Enabled' - this will fail if MIG is disabled - "sudo nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader | grep -i 'Enabled'", + "sudo nvidia-smi --query-gpu=mig.mode.current --format=csv,noheader", } execResult := execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "MIG mode is not enabled") stdout := strings.TrimSpace(execResult.stdout) s.T.Logf("MIG mode status: %s", stdout) - require.Contains(s.T, stdout, "Enabled", "expected MIG mode to be enabled, but got: %s", stdout) - s.T.Logf("MIG mode is enabled") + gpuStatuses := strings.Split(stdout, "\n") + require.Len(s.T, gpuStatuses, gpuCountExpected, "expected MIG status for %d GPUs, but got: %s", gpuCountExpected, stdout) + for gpuIndex, gpuStatus := range gpuStatuses { + require.Equalf(s.T, "Enabled", strings.TrimSpace(gpuStatus), "expected MIG mode to be enabled on GPU %d", gpuIndex) + } + s.T.Logf("MIG mode is enabled on %d GPUs", gpuCountExpected) } -func ValidateMIGInstancesCreated(ctx context.Context, s *Scenario, migProfile string) { +func ValidateMIGInstancesCreated(ctx context.Context, s *Scenario, migProfile string, instanceCountExpected int) { s.T.Helper() - s.T.Logf("validating that MIG instances are created with profile %s", migProfile) + s.T.Logf("validating that %d MIG instances are created with profile %s", instanceCountExpected, migProfile) command := []string{ "set -ex", @@ -2620,9 +2623,10 @@ func ValidateMIGInstancesCreated(ctx context.Context, s *Scenario, migProfile st execResult := execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "MIG instances with profile "+migProfile+" were not found") stdout := execResult.stdout - require.Contains(s.T, stdout, migProfile, "expected to find MIG profile %s in output, but did not.\nOutput:\n%s", migProfile, stdout) + instanceCount := strings.Count(stdout, migProfile) + require.Equal(s.T, instanceCountExpected, instanceCount, "expected %d MIG instances with profile %s, but found %d.\nOutput:\n%s", instanceCountExpected, migProfile, instanceCount, stdout) require.NotContains(s.T, stdout, "No MIG-enabled devices found", "no MIG devices were created.\nOutput:\n%s", stdout) - s.T.Logf("MIG instances with profile %s are created", migProfile) + s.T.Logf("%d MIG instances with profile %s are created", instanceCountExpected, migProfile) } // ValidateIPTablesCompatibleWithCiliumEBPF validates that all iptables rules in each table match the provided patterns which are accounted for