feat: use MIG profile layout for partitioning

aks-node-controller/parser/helper.go

@@ -250,8 +250,8 @@ func isContainerdV2(version string) bool {
 	return IsKubernetesVersionGe(version, "2.0.0")
 }
 
-func getIsMIGNode(gpuInstanceProfile string) bool {
-	return gpuInstanceProfile != ""
+func getIsMIGNode(gpuInstanceProfile string, migProfileLayout []string) bool {
+	return gpuInstanceProfile != "" || len(migProfileLayout) > 0
 }
 
 func getCustomCACertsStatus(customCACerts []string) bool {

aks-node-controller/parser/parser.go

@@ -46,6 +46,7 @@ func getCSEEnv(ctx context.Context, config *aksnodeconfigv1.Configuration, gpuCo
 
 	containerdVersion, _ := detectContainerdVersion(ctx)
 	cloudProviderSettings := getCloudProviderSettings(config)
+	isMIGNode := getIsMIGNode(config.GetGpuConfig().GetGpuInstanceProfile(), config.GetGpuConfig().GetMigProfileLayout())
 	env := map[string]string{
 		"PROVISION_OUTPUT":                                     "/var/log/azure/cluster-provision-cse-output.log",
 		"MOBY_VERSION":                                         "",
@@ -96,7 +97,7 @@ func getCSEEnv(ctx context.Context, config *aksnodeconfigv1.Configuration, gpuCo
 		"IS_VHD":                                               fmt.Sprintf("%v", getIsVHD(config.IsVhd)),
 		"GPU_NODE":                                             fmt.Sprintf("%v", getEnableNvidia(config)),
 		"SGX_NODE":                                             fmt.Sprintf("%v", getIsSgxEnabledSKU(config.GetVmSize())),
-		"MIG_NODE":                                             fmt.Sprintf("%v", getIsMIGNode(config.GetGpuConfig().GetGpuInstanceProfile())),
+		"MIG_NODE":                                             fmt.Sprintf("%v", isMIGNode),
 		"CONFIG_GPU_DRIVER_IF_NEEDED":                          fmt.Sprintf("%v", config.GetGpuConfig().GetConfigGpuDriver()),
 		"ENABLE_GPU_DEVICE_PLUGIN_IF_NEEDED":                   fmt.Sprintf("%v", config.GetGpuConfig().GetGpuDevicePlugin()),
 		"MANAGED_GPU_EXPERIENCE_AFEC_ENABLED":                  fmt.Sprintf("%v", config.GetGpuConfig().GetManagedGpuExperienceAfecEnabled()),

aks-node-controller/parser/parser_test.go

@@ -32,6 +32,7 @@ func TestBuildCSECmd(t *testing.T) {
 			k8sVersion: "1.19.13",
 			aksNodeConfigUpdator: func(aksNodeConfig *aksnodeconfigv1.Configuration) {
 				aksNodeConfig.GpuConfig.GpuInstanceProfile = "MIG7g"
+				aksNodeConfig.GpuConfig.MigStrategy = "Single"
 				// Skip GPU driver install
 				aksNodeConfig.GpuConfig.EnableNvidia = to.Ptr(false)
 				aksNodeConfig.VmSize = "Standard_ND96asr_v4"
@@ -40,6 +41,10 @@ func TestBuildCSECmd(t *testing.T) {
 				vars := environToMap(cmd.Env)
 				assertHasKeyWithValue(t, vars, "LOCATION", "southcentralus")
 				assert.Equal(t, "false", vars["GPU_NODE"])
+				assertHasKeyWithValue(t, vars, "MIG_NODE", "true")
+				assertHasKeyWithValue(t, vars, "GPU_INSTANCE_PROFILE", "MIG7g")
+				assertHasKeyWithValue(t, vars, "NVIDIA_MIG_PROFILE_LAYOUT", "")
+				assertHasKeyWithValue(t, vars, "NVIDIA_MIG_STRATEGY", "Single")
 				assert.NotEmpty(t, vars["CONTAINERD_CONFIG_NO_GPU_CONTENT"])
 				// Ensure the containerd config does not use the
 				// nvidia container runtime when skipping the
@@ -78,9 +83,10 @@ oom_score = -999
 			},
 			validator: func(cmd *exec.Cmd) {
 				vars := environToMap(cmd.Env)
+				assertHasKeyWithValue(t, vars, "GPU_INSTANCE_PROFILE", "")
 				assertHasKeyWithValue(t, vars, "NVIDIA_MIG_PROFILE_LAYOUT", "MIG3g,MIG2g,MIG1g,MIG1g")
-				// TODO: Make MIG_NODE true if either NVIDIA_MIG_PROFILE_LAYOUT or GPU_INSTANCE_PROFILE is set.
-				assertHasKeyWithValue(t, vars, "MIG_NODE", "false")
+				assertHasKeyWithValue(t, vars, "NVIDIA_MIG_STRATEGY", "")
+				assertHasKeyWithValue(t, vars, "MIG_NODE", "true")
 			},
 		},
 		{
@@ -496,7 +502,9 @@ func TestAKSNodeConfigCompatibilityFromJsonToCSECommand(t *testing.T) {
 				assertHasKeyWithValue(t, vars, "VNET_CNI_PLUGINS_URL", "")
 				assertHasKeyWithValue(t, vars, "LOCATION", "")
 				assertHasKeyWithValue(t, vars, "GPU_NODE", "false")
+				assertHasKeyWithValue(t, vars, "MIG_NODE", "false")
 				assertHasKeyWithValue(t, vars, "GPU_INSTANCE_PROFILE", "")
+				assertHasKeyWithValue(t, vars, "NVIDIA_MIG_PROFILE_LAYOUT", "")
 				assertHasKeyWithValue(t, vars, "CUSTOM_CA_TRUST_COUNT", "0")
 				assertHasKeyWithValue(t, vars, "SHOULD_CONFIGURE_CUSTOM_CA_TRUST", "false")
 				assertHasKeyWithValue(t, vars, "KUBELET_FLAGS", "")

aks-node-controller/proto/aksnodeconfig/v1/gpu_config.proto

@@ -28,7 +28,7 @@ message GpuConfig {
 
   // mig_strategy specifies the MIG strategy for the nvidia-device-plugin.
   // Valid values are "None", "Single", "Mixed". Only meaningful when MIG is enabled
-  // (i.e., when gpu_instance_profile is set). Defaults to "Single" if not specified.
+  // (i.e., when gpu_instance_profile or mig_profile_layout is set). Defaults to "Single" if not specified.
   string mig_strategy = 8;
 
   // mig_profile_layout is the ordered list of GPU instance profiles used for MIG partitioning.

e2e/node_config.go

@@ -258,9 +258,11 @@ func nbcToAKSNodeConfigV1(nbc *datamodel.NodeBootstrappingConfiguration) (*aksno
 			VnetCniPluginsUrl: vnetCNIPluginURL,
 		},
 		GpuConfig: &aksnodeconfigv1.GpuConfig{
-			ConfigGpuDriver:  true,
-			GpuDevicePlugin:  false,
-			MigProfileLayout: nbc.MIGProfileLayout,
+			ConfigGpuDriver:    true,
+			GpuDevicePlugin:    false,
+			GpuInstanceProfile: nbc.GPUInstanceProfile,
+			MigStrategy:        nbc.MigStrategy,
+			MigProfileLayout:   nbc.MIGProfileLayout,
 		},
 		EnableUnattendedUpgrade: enableUnattendedUpgrade,
 		EnableArtifactStreaming: nbc.EnableArtifactStreaming,

e2e/scenario_gpu_managed_experience_test.go

@@ -534,8 +534,30 @@ func Test_AzureLinux3_NvidiaDevicePluginRunning(t *testing.T) {
 }
 
 func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG(t *testing.T) {
+	runUbuntu2404NvidiaDevicePluginMIGSingle(t,
+		"Tests that NVIDIA device plugin and DCGM Exporter work with the legacy GPUInstanceProfile field",
+		func(nbc *datamodel.NodeBootstrappingConfiguration) {
+			nbc.GPUInstanceProfile = "MIG2g"
+		},
+	)
+}
+
+func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIGProfileLayout_Single(t *testing.T) {
+	runUbuntu2404NvidiaDevicePluginMIGSingle(t,
+		"Tests that NVIDIA device plugin and DCGM Exporter work with MIGProfileLayout and the Single MIG strategy",
+		func(nbc *datamodel.NodeBootstrappingConfiguration) {
+			nbc.MIGProfileLayout = []string{"MIG2g", "MIG2g", "MIG2g"}
+		},
+	)
+}
+
+func runUbuntu2404NvidiaDevicePluginMIGSingle(
+	t *testing.T,
+	description string,
+	setMIGProfile func(*datamodel.NodeBootstrappingConfiguration),
+) {
 	RunScenario(t, &Scenario{
-		Description: "Tests that NVIDIA device plugin and DCGM Exporter work with MIG enabled on Ubuntu 24.04 GPU nodes",
+		Description: description,
 		Location:    "westus2",
 		Tags: Tags{
 			GPU: true,
@@ -549,7 +571,7 @@ func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG(t *testing.T) {
 				nbc.ConfigGPUDriverIfNeeded = true
 				nbc.EnableGPUDevicePluginIfNeeded = true
 				nbc.EnableNvidia = true
-				nbc.GPUInstanceProfile = "MIG2g"
+				setMIGProfile(nbc)
 				nbc.EnableManagedGPU = true
 				nbc.MigStrategy = "Single"
 			},
@@ -573,16 +595,14 @@ func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG(t *testing.T) {
 				// Validate that the NVIDIA device plugin systemd service is running
 				ValidateNvidiaDevicePluginServiceRunning(ctx, s)
 
-				// Validate that MIG mode is enabled via nvidia-smi
+				ValidateMIGInstanceProfileCounts(ctx, s, map[string]int{"MIG 2g.20gb": 3})
+				ValidateNvidiaDevicePluginMIGStrategy(ctx, s, "single")
 				ValidateMIGModeEnabled(ctx, s, 1)
 
-				// Validate that MIG instances are created
-				ValidateMIGInstancesCreated(ctx, s, "MIG 2g.20gb", 3)
-
-				// Validate that GPU resources are advertised by the device plugin
-				ValidateNodeAdvertisesGPUResources(ctx, s, 3, "nvidia.com/gpu")
+				// Single exposes all three uniform partitions through nvidia.com/gpu and no profile-specific resources.
+				ValidateNodeAdvertisesExactGPUResources(ctx, s, map[string]int64{"nvidia.com/gpu": 3})
 
-				// Validate that MIG workloads can be scheduled
+				// Exercise every advertised MIG partition through the generic resource.
 				ValidateGPUWorkloadSchedulable(ctx, s, 3, "nvidia.com/gpu")
 
 				// Validate that the NVIDIA DCGM packages were installed correctly
@@ -776,7 +796,7 @@ func Test_CreateVMExtensionLinuxAKSNode_Timing(t *testing.T) {
 
 func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG_Mixed(t *testing.T) {
 	RunScenario(t, &Scenario{
-		Description: "Tests that NVIDIA device plugin work with MIG Mixed mode on Ubuntu 24.04 GPU nodes",
+		Description: "Tests that NVIDIA device plugin provisions and advertises a heterogeneous Mixed MIG geometry on Ubuntu 24.04 GPU nodes",
 		Location:    "westus2",
 		Tags: Tags{
 			GPU: true,
@@ -790,7 +810,7 @@ func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG_Mixed(t *testing.T) {
 				nbc.ConfigGPUDriverIfNeeded = true
 				nbc.EnableGPUDevicePluginIfNeeded = true
 				nbc.EnableNvidia = true
-				nbc.GPUInstanceProfile = "MIG1g"
+				nbc.MIGProfileLayout = []string{"MIG3g", "MIG2g", "MIG1g", "MIG1g"}
 				nbc.EnableManagedGPU = true
 				nbc.MigStrategy = "Mixed"
 			},
@@ -814,18 +834,25 @@ func Test_Ubuntu2404_NvidiaDevicePluginRunning_MIG_Mixed(t *testing.T) {
 				// Validate that the NVIDIA device plugin systemd service is running
 				ValidateNvidiaDevicePluginServiceRunning(ctx, s)
 
-				// Validate that MIG mode is enabled via nvidia-smi
+				ValidateMIGInstanceProfileCounts(ctx, s, map[string]int{
+					"MIG 3g.40gb": 1,
+					"MIG 2g.20gb": 1,
+					"MIG 1g.10gb": 2,
+				})
+				ValidateNvidiaDevicePluginMIGStrategy(ctx, s, "mixed")
 				ValidateMIGModeEnabled(ctx, s, 1)
 
-				// Validate that MIG instances are created
-				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"
-				ValidateNodeAdvertisesGPUResources(ctx, s, 7, migResourceName)
-
-				// Validate that MIG workloads can be scheduled
-				ValidateGPUWorkloadSchedulable(ctx, s, 2, migResourceName)
+				// Mixed exposes every profile-specific resource and no generic nvidia.com/gpu resource.
+				ValidateNodeAdvertisesExactGPUResources(ctx, s, map[string]int64{
+					"nvidia.com/mig-3g.40gb": 1,
+					"nvidia.com/mig-2g.20gb": 1,
+					"nvidia.com/mig-1g.10gb": 2,
+				})
+
+				// Exercise every advertised resource type, including both duplicate 1g partitions.
+				ValidateGPUWorkloadSchedulable(ctx, s, 1, "nvidia.com/mig-3g.40gb")
+				ValidateGPUWorkloadSchedulable(ctx, s, 1, "nvidia.com/mig-2g.20gb")
+				ValidateGPUWorkloadSchedulable(ctx, s, 2, "nvidia.com/mig-1g.10gb")
 			},
 		},
 	})

e2e/validators.go

@@ -2631,6 +2631,12 @@ func ValidateNvidiaDevicePluginServiceRunning(ctx context.Context, s *Scenario)
 	execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "NVIDIA device plugin systemd service should be active and enabled")
 }
 
+func ValidateNvidiaDevicePluginMIGStrategy(ctx context.Context, s *Scenario, strategy string) {
+	s.T.Helper()
+	command := fmt.Sprintf("systemctl cat nvidia-device-plugin.service | grep -F -- '--mig-strategy %s'", strategy)
+	execScriptOnVMForScenarioValidateExitCode(ctx, s, command, 0, "NVIDIA device plugin is not configured with MIG strategy "+strategy)
+}
+
 func ValidateNodeAdvertisesGPUResources(ctx context.Context, s *Scenario, gpuCountExpected int64, resourceName string) {
 	s.T.Helper()
 	s.T.Logf("validating that node advertises GPU resources")
@@ -2652,6 +2658,33 @@ func ValidateNodeAdvertisesGPUResources(ctx context.Context, s *Scenario, gpuCou
 	s.T.Logf("node %s advertises %s=%d resources", nodeName, resourceName, gpuCount)
 }
 
+func ValidateNodeAdvertisesExactGPUResources(ctx context.Context, s *Scenario, expected map[string]int64) {
+	s.T.Helper()
+	s.T.Logf("validating that node advertises exactly the expected NVIDIA GPU resources")
+
+	for resourceName := range expected {
+		waitUntilResourceAvailable(ctx, s, resourceName)
+	}
+
+	nodeName := s.Runtime.VM.KubeName
+	node, err := s.Runtime.Kube.Typed.CoreV1().Nodes().Get(ctx, nodeName, metav1.GetOptions{})
+	require.NoError(s.T, err, "failed to get node %q", nodeName)
+
+	getNvidiaResources := func(resources corev1.ResourceList) map[string]int64 {
+		result := make(map[string]int64)
+		for resourceName, quantity := range resources {
+			name := string(resourceName)
+			if strings.HasPrefix(name, "nvidia.com/") && quantity.Value() > 0 {
+				result[name] = quantity.Value()
+			}
+		}
+		return result
+	}
+
+	require.Equal(s.T, expected, getNvidiaResources(node.Status.Capacity), "node %s advertises unexpected NVIDIA GPU capacity", nodeName)
+	require.Equal(s.T, expected, getNvidiaResources(node.Status.Allocatable), "node %s advertises unexpected allocatable NVIDIA GPU resources", nodeName)
+}
+
 func ValidateGPUWorkloadSchedulable(ctx context.Context, s *Scenario, gpuCount int, resourceName string) {
 	s.T.Helper()
 	s.T.Logf("validating that GPU workloads can be scheduled")
@@ -2661,9 +2694,10 @@ func ValidateGPUWorkloadSchedulable(ctx context.Context, s *Scenario, gpuCount i
 	time.Sleep(20 * time.Second) // Same delay as existing GPU tests
 
 	// Create a GPU test pod using the same pattern as podRunNvidiaWorkload
+	resourceID := strings.ReplaceAll(strings.TrimPrefix(resourceName, "nvidia.com/"), ".", "-")
 	pod := &corev1.Pod{
 		ObjectMeta: metav1.ObjectMeta{
-			Name:      fmt.Sprintf("%s-gpu-test", s.Runtime.VM.KubeName),
+			Name:      fmt.Sprintf("%s-%s-test", s.Runtime.VM.KubeName, resourceID),
 			Namespace: "default",
 		},
 		Spec: corev1.PodSpec{
@@ -2832,6 +2866,28 @@ func ValidateMIGModeEnabled(ctx context.Context, s *Scenario, gpuCountExpected i
 	s.T.Logf("MIG mode is enabled on %d GPUs", gpuCountExpected)
 }
 
+func ValidateMIGInstanceProfileCounts(ctx context.Context, s *Scenario, expected map[string]int) {
+	s.T.Helper()
+	s.T.Logf("validating exact MIG instance profile counts")
+
+	command := []string{
+		"set -ex",
+		"sudo nvidia-smi mig -lgi",
+	}
+	execResult := execScriptOnVMForScenarioValidateExitCode(ctx, s, strings.Join(command, "\n"), 0, "failed to list MIG instances")
+
+	stdout := execResult.stdout
+	require.NotContains(s.T, stdout, "No MIG-enabled devices found", "no MIG devices were created.\nOutput:\n%s", stdout)
+
+	profilePattern := regexp.MustCompile(`MIG [0-9]+g\.[0-9]+gb(\+me)?`)
+	actual := make(map[string]int)
+	for _, profile := range profilePattern.FindAllString(stdout, -1) {
+		actual[profile]++
+	}
+
+	require.Equal(s.T, expected, actual, "unexpected MIG instance geometry.\nOutput:\n%s", stdout)
+}
+
 func ValidateMIGInstancesCreated(ctx context.Context, s *Scenario, migProfile string, instanceCountExpected int) {
 	s.T.Helper()
 	s.T.Logf("validating that %d MIG instances are created with profile %s", instanceCountExpected, migProfile)