feat: use MIG profile layout for partitioning - #9129
feat: use MIG profile layout for partitioning#9129Karen Chen (karenychen) wants to merge 15 commits into
Conversation
|
The latest Buf updates on your PR. Results from workflow Buf CI / buf (pull_request).
|
Windows Unit Test Results 3 files 13 suites 39s ⏱️ Results for commit 133c697. ♻️ This comment has been updated with latest results. |
There was a problem hiding this comment.
Pull request overview
This PR completes the runtime support for the ordered MIGProfileLayout contract (introduced in #9081) across classic CSE and aks-node-controller, enabling MIG partitioning based on an explicit ordered layout while preserving legacy GPUInstanceProfile uniform expansion behavior.
Changes:
- Treats nodes as MIG-enabled when either
MIGProfileLayoutor legacyGPUInstanceProfileis populated, withMIGProfileLayouttaking precedence when non-empty. - Updates
mig-partition.shto map the ordered layout directly to NVIDIA profile IDs (preserving order/duplicates) and to propagate failures fromnvidia-smi mig -cgiandnvidia-smi mig -cci. - Adds/updates focused ShellSpec coverage and GPU e2e coverage for legacy, Single, and Mixed MIG scenarios, plus additional validators for exact advertised resources and MIG strategy.
Reviewed changes
Copilot reviewed 13 out of 14 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| spec/parts/linux/cloud-init/artifacts/mig_partition_spec.sh | Adds ShellSpec coverage for legacy and layout-driven MIG partitioning behavior and failure propagation. |
| pkg/agent/variables.go | Sets MIG_NODE based on either legacy scalar or layout being populated. |
| pkg/agent/datamodel/helper.go | Extends IsMIGNode to consider both legacy and layout inputs. |
| pkg/agent/baker.go | Uses the shared IsMIGNode logic when templating MIG enablement. |
| pkg/agent/baker_test.go | Updates tests to assert MIG env vars for legacy and layout-driven cases. |
| parts/linux/cloud-init/artifacts/mig-partition.sh | Implements authoritative ordered layout parsing and legacy uniform expansion; propagates nvidia-smi failures. |
| e2e/validators.go | Adds validators for MIG strategy, exact GPU resource advertisement, and exact MIG instance profile counts. |
| e2e/scenario_gpu_managed_experience_test.go | Adds/updates Ubuntu 24.04 MIG scenarios for legacy field, layout+Single, and layout+Mixed with stricter validation. |
| e2e/node_config.go | Propagates MIG-related fields (instance profile, strategy, layout) into AKS node config v1 for e2e. |
| aks-node-controller/proto/aksnodeconfig/v1/gpu_config.proto | Clarifies MIG strategy semantics when MIG is enabled via either legacy scalar or layout. |
| aks-node-controller/pkg/gen/aksnodeconfig/v1/gpu_config.pb.go | Regenerates Go bindings reflecting the updated proto comment/metadata. |
| aks-node-controller/parser/parser.go | Sets MIG_NODE based on either legacy scalar or layout being populated. |
| aks-node-controller/parser/parser_test.go | Updates parser tests to assert MIG env vars and MIG enablement for layout-only inputs. |
| aks-node-controller/parser/helper.go | Extends MIG enablement helper to consider both legacy and layout inputs. |
Files not reviewed (1)
- aks-node-controller/pkg/gen/aksnodeconfig/v1/gpu_config.pb.go: Generated file
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 13 out of 14 changed files in this pull request and generated no new comments.
Files not reviewed (1)
- aks-node-controller/pkg/gen/aksnodeconfig/v1/gpu_config.pb.go: Generated file
Suppressed comments (2)
pkg/agent/datamodel/helper.go:131
- IsMIGNode treats any non-empty migProfileLayout slice as enabling MIG, but the layout is serialized via strings.Join. A slice like []string{""} makes MIG_NODE true while NVIDIA_MIG_PROFILE_LAYOUT becomes an empty string, causing mig-partition.sh to treat the layout as unset and fail with the misleading "neither ... is set" path. Consider only returning true when the layout contains at least one non-empty element (or trim-space non-empty) so MIG_NODE aligns with the serialized env var.
// IsMIGNode check if the node should be partitioned.
func IsMIGNode(gpuInstanceProfile string, migProfileLayout []string) bool {
return gpuInstanceProfile != "" || len(migProfileLayout) > 0
}
aks-node-controller/parser/helper.go:255
- getIsMIGNode uses len(migProfileLayout)>0, which can make MIG_NODE true even when the serialized NVIDIA_MIG_PROFILE_LAYOUT ends up empty (e.g., []string{""}). That leads to inconsistent behavior (MIG_NODE true, but mig-partition.sh sees no layout). Consider requiring at least one non-empty profile entry before enabling MIG based on migProfileLayout.
func getIsMIGNode(gpuInstanceProfile string, migProfileLayout []string) bool {
return gpuInstanceProfile != "" || len(migProfileLayout) > 0
}
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 13 out of 14 changed files in this pull request and generated no new comments.
Files not reviewed (1)
- aks-node-controller/pkg/gen/aksnodeconfig/v1/gpu_config.pb.go: Generated file
Suppressed comments (1)
parts/linux/cloud-init/artifacts/mig-partition.sh:60
parse_mig_profile_layoutuseslocal, which will fail the repo’s POSIX ShellCheck pass (shellcheck --shell=shruns on all *.sh except the small BASH_ONLY_LIST in .pipelines/scripts/verify_shell.sh). Sincemig-partition.shisn’t in that allowlist, please avoid bash-only syntax here so CI doesn’t break.
local profiles_csv="$1"
local layout=""
local is_last_profile
local profile
local profile_id
|
the agentbaker GPU E2E fails because this change requires rebuilding VHDs. Here is the passed GPU E2E run with VHD built from this branch: https://dev.azure.com/msazure/CloudNativeCompute/_build/results?buildId=177161624&view=results |
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 13 out of 14 changed files in this pull request and generated no new comments.
Files not reviewed (1)
- aks-node-controller/pkg/gen/aksnodeconfig/v1/gpu_config.pb.go: Generated file
Suppressed comments (1)
parts/linux/cloud-init/artifacts/mig-partition.sh:60
verify_shell.shruns an additional ShellCheck pass with--shell=shfor all scripts not inBASH_ONLY_LIST.mig-partition.shis not in that list, so usinglocalhere will likely fail the POSIX ShellCheck gate (e.g., SC3043). Either keep this script POSIX-compatible or explicitly exclude it from the POSIX pass.
A minimal fix is to remove local and use plain assignments (the rest of the script is already POSIX-friendly).
local profiles_csv="$1"
local layout=""
local is_last_profile
local profile
local profile_id
Karen Chen (karenychen)
left a comment
Karen Chen (karenychen)
left a comment
There was a problem hiding this comment.
Code review
Found 1 issue:
- High Risk - Backward Compatibility:
datamodel.IsMIGNodeis an exported API in current published AgentBaker module tags, and changing it fromIsMIGNode(string)toIsMIGNode(string, []string)breaks existing downstream callers at compile time. Preserve the one-argument function and add a new layout-aware helper (or another backward-compatible API). This is also tracked by the unresolved inline review thread on this line.
AgentBaker/pkg/agent/datamodel/helper.go
Lines 128 to 131 in 7131e37
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 13 out of 14 changed files in this pull request and generated no new comments.
Files not reviewed (1)
- aks-node-controller/pkg/gen/aksnodeconfig/v1/gpu_config.pb.go: Generated file
Suppressed comments (2)
parts/linux/cloud-init/artifacts/mig-partition.sh:61
parse_mig_profile_layoutuseslocal, but the repo’s validate-shell gate runs an additional ShellCheck pass with--shell=shfor all*.shnot in the smallBASH_ONLY_LIST(see.pipelines/scripts/verify_shell.sh:28-35,99-103). Sinceparts/linux/cloud-init/artifacts/mig-partition.shis not listed there, theselocaldeclarations will be flagged as non-POSIX (e.g., SC2039) and can break CI. Please either make the function POSIX-clean or explicitly exempt this script.
local profiles_csv="$1"
local layout=""
local is_last_profile
local profile
local profile_id
e2e/validators.go:3164
- Pod names must be a DNS-1123 label (<= 63 chars).
s.Runtime.VM.KubeNameis a Kubernetes node name derived from a VMSS name that is already capped at 57 chars (e2e/vmss.go:1140-1142), so appending-<resourceID>-testcan push the pod name over the 63-char limit and cause pod creation to fail in e2e. Consider truncating the generated name (similar to the existing truncation helper used in kata validators).
// 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-%s-test", s.Runtime.VM.KubeName, resourceID),
Namespace: "default",
| @@ -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 | |||
There was a problem hiding this comment.
still the case?
There was a problem hiding this comment.
I took a look at this before but not sure if it's the path we want to go: this would require us introducing a new dependency for the nvidia-mig-parted tool. Also I am not sure if it is as helpful with our mixed MIG model since it would require us to deliberately define all possible MIG profile combinations. I am okay with replacing this TODO with a comment
| esac | ||
| } | ||
|
|
||
| # TODO: Support GPU models with fewer than seven total partitions. |
There was a problem hiding this comment.
Which GPU instances are those?
There was a problem hiding this comment.
Oops I meant to remove this comment since the new GPUProfileLayout should already support it (we will compute the exact layout on RP side).
There are a few GPUs that support != 7 partitions documented on the official MIG user guide https://docs.nvidia.com/datacenter/tesla/mig-user-guide/latest/supported-mig-profiles.html
In Azure, we currently don't support any of these. However, I have spoken with Sachi and we want to maintain the flexibility to support them
What this PR does
Completes the runtime half of the MIG profile layout contract introduced by #9081.
MIGProfileLayoutor legacyGPUInstanceProfileas a MIG node in classic CSE and aks-node-controllerMIGProfileLayoutauthoritative when both fields are populatedMigStrategyinterpretationGPUInstanceProfileis populatednvidia-smi mig -cgiandnvidia-smi mig -cciMigStrategysolely for nvidia-device-plugin advertisement behaviorThis supersedes the runtime/test portion of closed #8869. AKS RP layout calculation remains out of scope.
Validation
make generateregenerated Go test data and the manifest without changes; its repository-wide ShellCheck tail stops on the existing POSIXSC3014baselinego test ./pkg/agent/... -count=1go test ./parser/... -count=1fromaks-node-controllergo test . -run '^$' -count=1frome2emig-partition.shand its ShellSpecAzure/AgentBakermaingit diff --check