fix(mig): support RTX PRO 6000 Blackwell dynamic MIG - #2132
Closed
shkatara wants to merge 1 commit into
Closed
Conversation
…ver Edition The RTX PRO 6000 Blackwell Server Edition (96 GB, 188 SMs) supports up to 4 MIG instances but was not present in knownMigGeometries, so its model string matched no entry and dynamic MIG carving was silently skipped (zero geometries, no error). Add its allowed geometries, mirroring the A30's 4-slice layout scaled to 96 GB: - 1g.24gb x4 (core 25) - 2g.48gb x2 (core 50) - 4g.96gb x1 (core 100) Model matching is substring-based (strings.Contains), so the configured "RTX PRO 6000 Blackwell Server Edition" matches the full NVML name "NVIDIA RTX PRO 6000 Blackwell Server Edition". - charts: add entry to the default device-configmap - docs: document the new model in dynamic-mig.md - test: cover the new model in TestGenerateMigTemplate, asserting the full NVML name resolves to the 1g.24gb x4 geometry fix(mig): fall back to nvidia-smi when nvidia-mig-parted creates no instances On some newer cards (notably the RTX PRO 6000 Blackwell Server Edition on open-kernel driver 610), `nvidia-mig-parted apply` reports success but its NVML-based create path is a silent no-op: zero GPU instances are created, no error is returned, and the pod ends up with no GPU. The `nvidia-smi mig -cgi ... -C` CLI creates the same geometry correctly on that hardware. After `nvidia-mig-parted apply`, verify via `nvidia-smi -L` that each GPU actually has the requested number of MIG instances. For any GPU that is short, reset and recreate the geometry with nvidia-smi. This is a no-op on hardware where nvidia-mig-parted already carved the instances, so it is safe for all MIG-capable cards. Verified on 8× RTX PRO 6000 Blackwell Server Edition. Signed-off-by: shkatara <writeshubhamhere@gmail.com>
Contributor
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: shkatara The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
Contributor
|
Welcome @shkatara! It looks like this is your first PR to Project-HAMi/HAMi 🎉 |
📝 WalkthroughWalkthroughThe change adds RTX PRO 6000 Blackwell Server Edition MIG geometries, verifies instances after ChangesDynamic MIG fallback
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ApplyMigTemplate
participant nvidia-mig-parted
participant ensureMigInstancesViaSmi
participant nvidia-smi
ApplyMigTemplate->>nvidia-mig-parted: Apply MIG template
nvidia-mig-parted-->>ApplyMigTemplate: Return success
ApplyMigTemplate->>ensureMigInstancesViaSmi: Verify requested instances
ensureMigInstancesViaSmi->>nvidia-smi: Query MIG instances
nvidia-smi-->>ensureMigInstancesViaSmi: Return per-GPU counts
ensureMigInstancesViaSmi->>nvidia-smi: Recreate missing geometry
nvidia-smi-->>ensureMigInstancesViaSmi: Return creation result
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go`:
- Around line 363-399: Update ensureMigInstancesViaSmi to return an error
instead of logging and continuing when nvidia-smi -L or createMigDevicesViaSmi
fails, and propagate that error through ApplyMigTemplate. After recreating MIG
devices, rerun the nvidia-smi listing and verify the expected instance counts
before returning success; report any failed verification as an error.
- Around line 374-397: Compare exact per-GPU MIG profile counts instead of only
total counts in the validation logic around migInstanceCountsFromSmi and
createMigDevicesViaSmi; trigger recreation when any requested profile count in
migSpec.MigDevices differs from the discovered inventory. In
pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go lines 374-397, retain
per-profile counts and validate every requested geometry. In
pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go lines 190-210, add
a fixture with the same total instance count but different profiles and assert
it is treated as a mismatch.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 4227e104-d1a5-416e-9088-359c21c82bfe
📒 Files selected for processing (4)
charts/hami/templates/scheduler/device-configmap.yamldocs/develop/dynamic-mig.mdpkg/device-plugin/nvidiadevice/nvinternal/plugin/util.gopkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go
Comment on lines
+363
to
+399
| func (nv *NvidiaDevicePlugin) ensureMigInstancesViaSmi() { | ||
| current, ok := nv.migCurrent.MigConfigs["current"] | ||
| if !ok { | ||
| return | ||
| } | ||
|
|
||
| out, err := exec.Command("nvidia-smi", "-L").CombinedOutput() | ||
| if err != nil { | ||
| klog.Errorf("failed to list GPUs with nvidia-smi -L, skipping MIG fallback: %v, output: %s", err, string(out)) | ||
| return | ||
| } | ||
| counts := migInstanceCountsFromSmi(string(out)) | ||
|
|
||
| for _, migSpec := range current { | ||
| if !migSpec.MigEnabled { | ||
| continue | ||
| } | ||
| expected := 0 | ||
| for _, c := range migSpec.MigDevices { | ||
| expected += int(c) | ||
| } | ||
| if expected == 0 { | ||
| continue | ||
| } | ||
| for _, dev := range migSpec.Devices { | ||
| gpuIndex := int(dev) | ||
| if counts[gpuIndex] >= expected { | ||
| continue | ||
| } | ||
| klog.Warningf("GPU %d has %d MIG instance(s) but %d were requested; nvidia-mig-parted did not create them, falling back to nvidia-smi", gpuIndex, counts[gpuIndex], expected) | ||
| if err := createMigDevicesViaSmi(gpuIndex, migSpec.MigDevices); err != nil { | ||
| klog.Errorf("nvidia-smi MIG fallback failed for GPU %d: %v", gpuIndex, err) | ||
| continue | ||
| } | ||
| klog.Infof("nvidia-smi MIG fallback created geometry on GPU %d: %v", gpuIndex, migSpec.MigDevices) | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Propagate fallback failures instead of continuing startup.
A failed nvidia-smi -L query or creation command is only logged, so ApplyMigTemplate still returns successfully without a usable MIG geometry. Return and propagate an error, and perform a final verification after recreation before allowing the startup path to succeed.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go` around lines 363 -
399, Update ensureMigInstancesViaSmi to return an error instead of logging and
continuing when nvidia-smi -L or createMigDevicesViaSmi fails, and propagate
that error through ApplyMigTemplate. After recreating MIG devices, rerun the
nvidia-smi listing and verify the expected instance counts before returning
success; report any failed verification as an error.
Comment on lines
+374
to
+397
| counts := migInstanceCountsFromSmi(string(out)) | ||
|
|
||
| for _, migSpec := range current { | ||
| if !migSpec.MigEnabled { | ||
| continue | ||
| } | ||
| expected := 0 | ||
| for _, c := range migSpec.MigDevices { | ||
| expected += int(c) | ||
| } | ||
| if expected == 0 { | ||
| continue | ||
| } | ||
| for _, dev := range migSpec.Devices { | ||
| gpuIndex := int(dev) | ||
| if counts[gpuIndex] >= expected { | ||
| continue | ||
| } | ||
| klog.Warningf("GPU %d has %d MIG instance(s) but %d were requested; nvidia-mig-parted did not create them, falling back to nvidia-smi", gpuIndex, counts[gpuIndex], expected) | ||
| if err := createMigDevicesViaSmi(gpuIndex, migSpec.MigDevices); err != nil { | ||
| klog.Errorf("nvidia-smi MIG fallback failed for GPU %d: %v", gpuIndex, err) | ||
| continue | ||
| } | ||
| klog.Infof("nvidia-smi MIG fallback created geometry on GPU %d: %v", gpuIndex, migSpec.MigDevices) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Verify exact MIG geometry rather than total instance count. The runtime and test both reduce inventory to a per-GPU total, which accepts an incorrect profile mix with the requested number of instances.
pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go#L374-L397: retain and compare eachMIG <profile>count per GPU againstmigSpec.MigDevices; recreate on any mismatch.pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go#L190-L210: add a same-total/different-profile fixture and assert that it is treated as a mismatch.
📍 Affects 2 files
pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go#L374-L397(this comment)pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go#L190-L210
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go` around lines 374 -
397, Compare exact per-GPU MIG profile counts instead of only total counts in
the validation logic around migInstanceCountsFromSmi and createMigDevicesViaSmi;
trigger recreation when any requested profile count in migSpec.MigDevices
differs from the discovered inventory. In
pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go lines 374-397, retain
per-profile counts and validate every requested geometry. In
pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go lines 190-210, add
a fixture with the same total instance count but different profiles and assert
it is treated as a mismatch.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What type of PR is this?
/kind bug
What this PR does / why we need it:
nvidia-mig-parted applyactually created the requested number of MIG instances.nvidia-smi mig -cgi ... -Cwhennvidia-mig-partedexits successfully but creates no instances.Without this validation,
nvidia-mig-partedcan report success on PCI device2bb5while leaving the GPU without MIG instances, allowing allocation to continue with an empty MIG UUID.Which issue(s) this PR fixes:
Fixes #2131
Special notes for your reviewer:
Tested with:
A
linux/amd64test image was built from commit193e074asdocker.io/shkatara/master-hami:fallback-to-nvidia-smi.Cursor AI assisted with investigation, code generation, tests, and drafting the issue and pull request. The resulting changes were reviewed and the unit tests were run locally.
Does this PR introduce a user-facing change?:
Dynamic MIG allocation now supports NVIDIA RTX PRO 6000 Blackwell Server Edition GPUs and recovers when
nvidia-mig-partedsilently creates no MIG instances.Summary by CodeRabbit
New Features
Documentation
Bug Fixes