diff --git a/aks-node-controller/parser/helper.go b/aks-node-controller/parser/helper.go index 9cc9e7eecca..14e0be88b7f 100644 --- a/aks-node-controller/parser/helper.go +++ b/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 { diff --git a/aks-node-controller/parser/parser.go b/aks-node-controller/parser/parser.go index 02003776a59..13e4d61e4a0 100644 --- a/aks-node-controller/parser/parser.go +++ b/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()), diff --git a/aks-node-controller/parser/parser_test.go b/aks-node-controller/parser/parser_test.go index bcb3362ee80..41207403f36 100644 --- a/aks-node-controller/parser/parser_test.go +++ b/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", "") diff --git a/aks-node-controller/pkg/gen/aksnodeconfig/v1/gpu_config.pb.go b/aks-node-controller/pkg/gen/aksnodeconfig/v1/gpu_config.pb.go index 1194703ef1b..e8fd356fd03 100644 --- a/aks-node-controller/pkg/gen/aksnodeconfig/v1/gpu_config.pb.go +++ b/aks-node-controller/pkg/gen/aksnodeconfig/v1/gpu_config.pb.go @@ -42,7 +42,7 @@ type GpuConfig struct { EnableManagedGpu bool `protobuf:"varint,7,opt,name=enable_managed_gpu,json=enableManagedGpu,proto3" json:"enable_managed_gpu,omitempty"` // 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. MigStrategy string `protobuf:"bytes,8,opt,name=mig_strategy,json=migStrategy,proto3" json:"mig_strategy,omitempty"` // mig_profile_layout is the ordered list of GPU instance profiles used for MIG partitioning. // It will become the unified MIG profile field after gpu_instance_profile is deprecated. diff --git a/aks-node-controller/proto/aksnodeconfig/v1/gpu_config.proto b/aks-node-controller/proto/aksnodeconfig/v1/gpu_config.proto index 2cbaa0b6697..927cef4a16a 100644 --- a/aks-node-controller/proto/aksnodeconfig/v1/gpu_config.proto +++ b/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. diff --git a/e2e/node_config.go b/e2e/node_config.go index 1bd82d24244..9be108c1e86 100644 --- a/e2e/node_config.go +++ b/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, diff --git a/e2e/scenario_gpu_managed_experience_test.go b/e2e/scenario_gpu_managed_experience_test.go index 371cd37b917..b73b9fbd4d4 100644 --- a/e2e/scenario_gpu_managed_experience_test.go +++ b/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") }, }, }) diff --git a/e2e/validators.go b/e2e/validators.go index cc5d5c2d629..9dc753adcaf 100644 --- a/e2e/validators.go +++ b/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) diff --git a/parts/linux/cloud-init/artifacts/mig-partition.sh b/parts/linux/cloud-init/artifacts/mig-partition.sh index d60b92d68b2..6c53d9030bf 100644 --- a/parts/linux/cloud-init/artifacts/mig-partition.sh +++ b/parts/linux/cloud-init/artifacts/mig-partition.sh @@ -2,26 +2,104 @@ #NOTE: Currently, Nvidia library mig-parted (https://github.com/NVIDIA/mig-parted) cannot work properly because of the outdated GPU driver version #TODO: Use mig-parted library to do the partition after the above issue is fixed -MIG_PROFILE=${1} -case ${MIG_PROFILE} in - "MIG1g") - nvidia-smi mig -cgi 19,19,19,19,19,19,19 - ;; - "MIG2g") - nvidia-smi mig -cgi 14,14,14 - ;; - "MIG3g") - nvidia-smi mig -cgi 9,9 - ;; - "MIG4g") - nvidia-smi mig -cgi 5 - ;; - "MIG7g") - nvidia-smi mig -cgi 0 - ;; - *) - echo "not a valid GPU instance profile" - exit 1 - ;; -esac -nvidia-smi mig -cci \ No newline at end of file + +mig_profile_id() { + case "$1" in + "MIG1g") + echo "19" + ;; + "MIG2g") + echo "14" + ;; + "MIG3g") + echo "9" + ;; + "MIG4g") + echo "5" + ;; + "MIG7g") + echo "0" + ;; + *) + echo "not a valid MIG profile: $1" >&2 + return 1 + ;; + esac +} + +# TODO: Support GPU models with fewer than seven total partitions. +uniform_mig_profile_layout() { + case "$1" in + "MIG1g") + echo "19,19,19,19,19,19,19" + ;; + "MIG2g") + echo "14,14,14" + ;; + "MIG3g") + echo "9,9" + ;; + "MIG4g") + echo "5" + ;; + "MIG7g") + echo "0" + ;; + *) + echo "not a valid MIG profile: $1" >&2 + return 1 + ;; + esac +} + +parse_mig_profile_layout() { + local profiles_csv="$1" + local layout="" + local is_last_profile + local profile + local profile_id + + while true; do + case "$profiles_csv" in + *,*) + profile="${profiles_csv%%,*}" + profiles_csv="${profiles_csv#*,}" + is_last_profile=false + ;; + *) + profile="$profiles_csv" + is_last_profile=true + ;; + esac + + profile_id="$(mig_profile_id "${profile}")" || return 1 + if [ -n "${layout}" ]; then + layout="${layout},${profile_id}" + else + layout="${profile_id}" + fi + + if $is_last_profile; then + break + fi + done + + echo "${layout}" +} + +# Backward-compatible: allow legacy positional argument usage (see mig-partition.service ExecStart). +if [ -z "${GPU_INSTANCE_PROFILE:-}" ] && [ -n "${1:-}" ]; then + GPU_INSTANCE_PROFILE="$1" +fi + +if [ -n "${NVIDIA_MIG_PROFILE_LAYOUT:-}" ]; then + MIG_LAYOUT="$(parse_mig_profile_layout "${NVIDIA_MIG_PROFILE_LAYOUT}")" || exit 1 +elif [ -n "${GPU_INSTANCE_PROFILE:-}" ]; then + MIG_LAYOUT="$(uniform_mig_profile_layout "${GPU_INSTANCE_PROFILE}")" || exit 1 +else + echo "neither NVIDIA_MIG_PROFILE_LAYOUT nor GPU_INSTANCE_PROFILE is set" >&2 + exit 1 +fi + +nvidia-smi mig -cgi "${MIG_LAYOUT}" || exit $? +nvidia-smi mig -cci diff --git a/pkg/agent/baker.go b/pkg/agent/baker.go index 873790f8eff..b89479f5b5e 100644 --- a/pkg/agent/baker.go +++ b/pkg/agent/baker.go @@ -894,7 +894,7 @@ func getContainerServiceFuncMap(config *datamodel.NodeBootstrappingConfiguration return strings.Join(config.MIGProfileLayout, ",") }, "IsMIGEnabledNode": func() bool { - return config.GPUInstanceProfile != "" + return datamodel.IsMIGNode(config.GPUInstanceProfile, config.MIGProfileLayout) }, "GetMigStrategy": func() string { return config.MigStrategy diff --git a/pkg/agent/baker_test.go b/pkg/agent/baker_test.go index 9ea69ecdc4e..2e816db8dc8 100644 --- a/pkg/agent/baker_test.go +++ b/pkg/agent/baker_test.go @@ -1481,6 +1481,7 @@ var _ = Describe("getLinuxNodeCSECommand", func() { It("should handle MIG GPU configuration", func() { baseConfig.GPUInstanceProfile = "MIG7g" + baseConfig.MigStrategy = "Single" baseConfig.ConfigGPUDriverIfNeeded = true baseConfig.EnableNvidia = true baseConfig.AgentPoolProfile.VMSize = "Standard_ND96asr_v4" @@ -1493,10 +1494,13 @@ var _ = Describe("getLinuxNodeCSECommand", func() { vars := decodeCSEVars(cseCmd) Expect(vars).To(HaveKeyWithValue("GPU_NODE", "true")) Expect(vars).To(HaveKeyWithValue("CONFIG_GPU_DRIVER_IF_NEEDED", "true")) + Expect(vars).To(HaveKeyWithValue("MIG_NODE", "true")) Expect(vars).To(HaveKeyWithValue("GPU_INSTANCE_PROFILE", "MIG7g")) + Expect(vars).To(HaveKeyWithValue("NVIDIA_MIG_PROFILE_LAYOUT", "")) + Expect(vars).To(HaveKeyWithValue("NVIDIA_MIG_STRATEGY", "Single")) }) - It("should populate a MIG profile layout without enabling partitioning", func() { + It("should enable partitioning for a MIG profile layout", func() { baseConfig.MIGProfileLayout = []string{"MIG3g", "MIG2g", "MIG1g", "MIG1g"} cseCmd := templateGenerator.getLinuxNodeCSECommand(baseConfig) @@ -1506,8 +1510,7 @@ var _ = Describe("getLinuxNodeCSECommand", func() { vars := decodeCSEVars(cseCmd) Expect(vars).To(HaveKeyWithValue("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. - Expect(vars).To(HaveKeyWithValue("MIG_NODE", "false")) + Expect(vars).To(HaveKeyWithValue("MIG_NODE", "true")) }) It("should handle disable unattended upgrades", func() { diff --git a/pkg/agent/datamodel/helper.go b/pkg/agent/datamodel/helper.go index c96137ff60f..6b21e0c07ef 100644 --- a/pkg/agent/datamodel/helper.go +++ b/pkg/agent/datamodel/helper.go @@ -126,8 +126,8 @@ func getComponentNameFromURL(downloadURL string) (string, error) { } // IsMIGNode check if the node should be partitioned. -func IsMIGNode(gpuInstanceProfile string) bool { - return gpuInstanceProfile != "" +func IsMIGNode(gpuInstanceProfile string, migProfileLayout []string) bool { + return gpuInstanceProfile != "" || len(migProfileLayout) > 0 } // returns the CSE timeout value in seconds. diff --git a/pkg/agent/variables.go b/pkg/agent/variables.go index 545861b8198..ca307c0649c 100644 --- a/pkg/agent/variables.go +++ b/pkg/agent/variables.go @@ -124,7 +124,7 @@ func getCSECommandVariables(config *datamodel.NodeBootstrappingConfiguration) pa "sgxNode": strconv.FormatBool(datamodel.IsSgxEnabledSKU(profile.VMSize)), "configGPUDriverIfNeeded": config.ConfigGPUDriverIfNeeded, "enableGPUDevicePluginIfNeeded": config.EnableGPUDevicePluginIfNeeded, - "migNode": strconv.FormatBool(datamodel.IsMIGNode(config.GPUInstanceProfile)), + "migNode": strconv.FormatBool(datamodel.IsMIGNode(config.GPUInstanceProfile, config.MIGProfileLayout)), "gpuInstanceProfile": config.GPUInstanceProfile, "migProfileLayout": strings.Join(config.MIGProfileLayout, ","), "windowsEnableCSIProxy": windowsProfile.IsCSIProxyEnabled(), diff --git a/spec/parts/linux/cloud-init/artifacts/mig_partition_spec.sh b/spec/parts/linux/cloud-init/artifacts/mig_partition_spec.sh new file mode 100644 index 00000000000..f0ddef8ed39 --- /dev/null +++ b/spec/parts/linux/cloud-init/artifacts/mig_partition_spec.sh @@ -0,0 +1,121 @@ +#!/bin/bash +# shellcheck disable=SC2329 # mock functions are invoked by the script under test. + +Describe 'mig-partition.sh' + SCRIPT_PATH="./parts/linux/cloud-init/artifacts/mig-partition.sh" + + nvidia-smi() { + echo "nvidia-smi $*" + } + export -f nvidia-smi + + Describe 'legacy GPU_INSTANCE_PROFILE' + Parameters + "MIG1g" "19,19,19,19,19,19,19" + "MIG2g" "14,14,14" + "MIG3g" "9,9" + "MIG4g" "5" + "MIG7g" "0" + End + + Example "maps $1 to its uniform layout" + When run env GPU_INSTANCE_PROFILE="$1" bash "$SCRIPT_PATH" + + The status should be success + The output should equal "nvidia-smi mig -cgi $2 +nvidia-smi mig -cci" + End + End + + It 'accepts the legacy positional profile used by older service units' + When run bash "$SCRIPT_PATH" MIG2g + + The status should be success + The output should equal "nvidia-smi mig -cgi 14,14,14 +nvidia-smi mig -cci" + End + + It 'maps the ordered layout without sorting or expanding profiles' + When run env NVIDIA_MIG_PROFILE_LAYOUT="MIG3g,MIG2g,MIG1g,MIG1g" bash "$SCRIPT_PATH" + + The status should be success + The output should equal "nvidia-smi mig -cgi 9,14,19,19 +nvidia-smi mig -cci" + End + + It 'prioritizes the layout when the legacy scalar is also populated' + When run env GPU_INSTANCE_PROFILE="MIG7g" NVIDIA_MIG_PROFILE_LAYOUT="MIG2g,MIG1g" bash "$SCRIPT_PATH" + + The status should be success + The output should include "nvidia-smi mig -cgi 14,19" + End + + It 'does not use MIG strategy to calculate the layout since it is already computed on caller site' + When run env NVIDIA_MIG_PROFILE_LAYOUT="MIG2g,MIG1g" NVIDIA_MIG_STRATEGY="Single" bash "$SCRIPT_PATH" + + The status should be success + The output should include "nvidia-smi mig -cgi 14,19" + End + + It 'rejects an invalid legacy profile' + When run env GPU_INSTANCE_PROFILE="MIG6g" bash "$SCRIPT_PATH" + + The status should be failure + The error should include "not a valid MIG profile: MIG6g" + End + + It 'rejects an invalid layout profile' + When run env NVIDIA_MIG_PROFILE_LAYOUT="MIG1g,MIG6g" bash "$SCRIPT_PATH" + + The status should be failure + The error should include "not a valid MIG profile: MIG6g" + End + + It 'rejects empty layout elements instead of repairing the plan' + When run env NVIDIA_MIG_PROFILE_LAYOUT="MIG1g,,MIG2g" bash "$SCRIPT_PATH" + + The status should be failure + The error should include "not a valid MIG profile: " + End + + It 'rejects whitespace around a layout element' + When run env NVIDIA_MIG_PROFILE_LAYOUT="MIG1g, MIG2g" bash "$SCRIPT_PATH" + + The status should be failure + The error should include "not a valid MIG profile: MIG2g" + End + + It 'rejects missing layout inputs' + When run bash "$SCRIPT_PATH" + + The status should be failure + The error should include "neither NVIDIA_MIG_PROFILE_LAYOUT nor GPU_INSTANCE_PROFILE is set" + End + + It 'propagates failure from creating GPU instances without running -cci' + nvidia-smi() { + echo "nvidia-smi $*" + [ "$2" != "-cgi" ] + } + export -f nvidia-smi + + When run env NVIDIA_MIG_PROFILE_LAYOUT="MIG2g" bash "$SCRIPT_PATH" + + The status should be failure + The output should equal "nvidia-smi mig -cgi 14" + End + + It 'propagates failure from creating compute instances' + nvidia-smi() { + echo "nvidia-smi $*" + [ "$2" != "-cci" ] + } + export -f nvidia-smi + + When run env NVIDIA_MIG_PROFILE_LAYOUT="MIG2g" bash "$SCRIPT_PATH" + + The status should be failure + The output should equal "nvidia-smi mig -cgi 14 +nvidia-smi mig -cci" + End +End