From 193e074b0052b9d3a9a6de891e649f8cae796cc6 Mon Sep 17 00:00:00 2001 From: Shubham Katara Date: Fri, 24 Jul 2026 14:09:52 +0200 Subject: [PATCH] feat(nvidia): add MIG geometry support for RTX PRO 6000 Blackwell Server Edition MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .../templates/scheduler/device-configmap.yaml | 17 ++ docs/develop/dynamic-mig.md | 14 ++ .../nvidiadevice/nvinternal/plugin/util.go | 157 ++++++++++++++++++ .../nvinternal/plugin/util_test.go | 100 +++++++++++ 4 files changed, 288 insertions(+) diff --git a/charts/hami/templates/scheduler/device-configmap.yaml b/charts/hami/templates/scheduler/device-configmap.yaml index 7f0b2b254d..1335115bc6 100644 --- a/charts/hami/templates/scheduler/device-configmap.yaml +++ b/charts/hami/templates/scheduler/device-configmap.yaml @@ -230,6 +230,23 @@ data: core: 100 memory: 184320 count: 1 + - models: [ "RTX PRO 6000 Blackwell Server Edition" ] + allowedGeometries: + - + - name: 1g.24gb + core: 25 + memory: 24576 + count: 4 + - + - name: 2g.48gb + core: 50 + memory: 49152 + count: 2 + - + - name: 4g.96gb + core: 100 + memory: 98304 + count: 1 cambricon: resourceCountName: {{ .Values.mluResourceName }} resourceMemoryName: {{ .Values.mluResourceMem }} diff --git a/docs/develop/dynamic-mig.md b/docs/develop/dynamic-mig.md index b5805083c6..d95d7ad310 100644 --- a/docs/develop/dynamic-mig.md +++ b/docs/develop/dynamic-mig.md @@ -87,6 +87,20 @@ data: - name: 7g.79gb memory: 80896 count: 1 + - models: [ "RTX PRO 6000 Blackwell Server Edition" ] + allowedGeometries: + - + - name: 1g.24gb + memory: 24576 + count: 4 + - + - name: 2g.48gb + memory: 49152 + count: 2 + - + - name: 4g.96gb + memory: 98304 + count: 1 nodeconfig: - name: nodeA operatingmode: hami-core diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go index b622aa64bd..612f60dfc6 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util.go @@ -25,6 +25,7 @@ import ( "os" "os/exec" "path/filepath" + "sort" "strconv" "strings" "time" @@ -346,6 +347,162 @@ func (nv *NvidiaDevicePlugin) ApplyMigTemplate() { } outStr := stdout.String() klog.Infoln("Mig apply", outStr) + + // nvidia-mig-parted can report success while creating zero MIG instances on + // some newer cards (e.g. RTX PRO 6000 Blackwell Server Edition), where its + // NVML-based create path is a silent no-op. Verify the instances actually + // exist and fall back to the nvidia-smi CLI for any GPU that is short. + nv.ensureMigInstancesViaSmi() +} + +// ensureMigInstancesViaSmi verifies that the MIG instances requested in +// nv.migCurrent actually exist and, for any GPU that is short, recreates the +// geometry with the nvidia-smi CLI. It is a no-op on hardware where +// nvidia-mig-parted already carved the instances correctly, so it is safe for +// all MIG-capable cards. +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) + } + } +} + +// migInstanceCountsFromSmi parses `nvidia-smi -L` output and returns the number +// of already-created MIG devices per physical GPU index. GPUs with no MIG +// devices are still recorded with a count of zero. +func migInstanceCountsFromSmi(output string) map[int]int { + counts := make(map[int]int) + currentGPU := -1 + for _, line := range strings.Split(output, "\n") { + trimmed := strings.TrimSpace(line) + if strings.HasPrefix(trimmed, "GPU ") { + // e.g. "GPU 0: NVIDIA RTX PRO 6000 Blackwell Server Edition (UUID: GPU-...)" + idxStr := strings.TrimSpace(strings.SplitN(strings.TrimPrefix(trimmed, "GPU "), ":", 2)[0]) + if idx, err := strconv.Atoi(idxStr); err == nil { + currentGPU = idx + if _, seen := counts[currentGPU]; !seen { + counts[currentGPU] = 0 + } + } + continue + } + if strings.HasPrefix(trimmed, "MIG ") && currentGPU >= 0 { + counts[currentGPU]++ + } + } + return counts +} + +// createMigDevicesViaSmi resets and recreates the desired MIG geometry on a +// single physical GPU using the nvidia-smi CLI. It is the fallback for GPUs +// where nvidia-mig-parted reports success but the NVML create path produces no +// MIG instances. +func createMigDevicesViaSmi(gpuIndex int, migDevices map[string]int32) error { + gpu := strconv.Itoa(gpuIndex) + + cgi := buildCreateGpuInstancesArg(migDevices) + if cgi == "" { + return fmt.Errorf("no MIG devices requested for GPU %d", gpuIndex) + } + + // Best-effort: ensure MIG mode is enabled (a no-op if already enabled). + runNvidiaSmi("-i", gpu, "-mig", "1") + + // Destroy any pre-existing compute/GPU instances so the geometry is clean. + // Compute instances must be destroyed before their parent GPU instances. + // These fail harmlessly when nothing exists yet, so errors are only logged. + runNvidiaSmi("mig", "-i", gpu, "-dci") + runNvidiaSmi("mig", "-i", gpu, "-dgi") + + // -C also creates the matching compute instance for each GPU instance. + if out, err := runNvidiaSmi("mig", "-i", gpu, "-cgi", cgi, "-C"); err != nil { + return fmt.Errorf("nvidia-smi mig create failed on GPU %d (cgi=%s): %v, output: %s", gpuIndex, cgi, err, out) + } + return nil +} + +// buildCreateGpuInstancesArg expands a MigDevices map (profile name -> count) +// into a comma-separated argument for `nvidia-smi mig -cgi`, ordering larger +// GPU-instance slices first so placement succeeds for mixed geometries. +func buildCreateGpuInstancesArg(migDevices map[string]int32) string { + type profile struct { + name string + slices int + } + var profiles []profile + for name, count := range migDevices { + for i := int32(0); i < count; i++ { + profiles = append(profiles, profile{name: name, slices: migProfileSlices(name)}) + } + } + sort.SliceStable(profiles, func(i, j int) bool { + if profiles[i].slices != profiles[j].slices { + return profiles[i].slices > profiles[j].slices + } + return profiles[i].name < profiles[j].name + }) + names := make([]string, 0, len(profiles)) + for _, p := range profiles { + names = append(names, p.name) + } + return strings.Join(names, ",") +} + +// migProfileSlices returns the leading GPU-instance slice count of a MIG +// profile name (e.g. "2g.10gb" -> 2). It returns 1 when the name cannot be +// parsed so the profile is still created, just placed last. +func migProfileSlices(name string) int { + idx := strings.Index(name, "g") + if idx <= 0 { + return 1 + } + if n, err := strconv.Atoi(name[:idx]); err == nil && n > 0 { + return n + } + return 1 +} + +// runNvidiaSmi runs nvidia-smi with the given args and returns its combined +// output. Failures are logged at a high verbosity because some calls (instance +// teardown) are expected to fail when there is nothing to tear down. +func runNvidiaSmi(args ...string) (string, error) { + out, err := exec.Command("nvidia-smi", args...).CombinedOutput() + if err != nil { + klog.V(4).Infof("nvidia-smi %s: %v, output: %s", strings.Join(args, " "), err, string(out)) + } + return string(out), err } func (nv *NvidiaDevicePlugin) GenerateMigTemplate(devtype string, devindex int, val device.ContainerDevice) (int, bool) { diff --git a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go index e6c2d4f358..ceacaabd12 100644 --- a/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go +++ b/pkg/device-plugin/nvidiadevice/nvinternal/plugin/util_test.go @@ -60,6 +60,14 @@ func TestGenerateMigTemplate(t *testing.T) { {device.MigTemplate{Name: "7g.80gb", Core: 100, Memory: 81920, Count: 1}}, }, }, + { + Models: []string{"RTX PRO 6000 Blackwell Server Edition"}, + Geometries: []device.Geometry{ + {device.MigTemplate{Name: "1g.24gb", Core: 25, Memory: 24576, Count: 4}}, + {device.MigTemplate{Name: "2g.48gb", Core: 50, Memory: 49152, Count: 2}}, + {device.MigTemplate{Name: "4g.96gb", Core: 100, Memory: 98304, Count: 1}}, + }, + }, }, } @@ -134,6 +142,23 @@ func TestGenerateMigTemplate(t *testing.T) { "1g.5gb": 7, }, }, + { + // The full NVML model string must match the shorter configured + // model via substring matching (RTX PRO 6000 Blackwell Server Edition). + name: "rtx pro 6000 blackwell 1g.24gb template", + model: "NVIDIA RTX PRO 6000 Blackwell Server Edition", + deviceIdx: 0, + containerDev: device.ContainerDevice{ + Idx: 0, + UUID: "ccccdddd[0-3]", + Usedmem: 20000, + }, + expectedPos: 3, + expectedReset: true, + expectedMig: map[string]int32{ + "1g.24gb": 4, + }, + }, } for _, tc := range testCases { @@ -162,6 +187,81 @@ func TestGenerateMigTemplate(t *testing.T) { } } +func TestMigInstanceCountsFromSmi(t *testing.T) { + output := `GPU 0: NVIDIA RTX PRO 6000 Blackwell Server Edition (UUID: GPU-aaaa) + MIG 1g.24gb Device 0: (UUID: MIG-1111) + MIG 1g.24gb Device 1: (UUID: MIG-2222) + MIG 1g.24gb Device 2: (UUID: MIG-3333) + MIG 1g.24gb Device 3: (UUID: MIG-4444) +GPU 1: NVIDIA RTX PRO 6000 Blackwell Server Edition (UUID: GPU-bbbb) +GPU 2: NVIDIA A100-SXM4-40GB (UUID: GPU-cccc) + MIG 2g.10gb Device 0: (UUID: MIG-5555)` + + counts := migInstanceCountsFromSmi(output) + expected := map[int]int{0: 4, 1: 0, 2: 1} + for gpu, want := range expected { + if counts[gpu] != want { + t.Errorf("GPU %d: expected %d MIG devices, got %d", gpu, want, counts[gpu]) + } + } + if len(counts) != len(expected) { + t.Errorf("expected %d GPUs, got %d: %v", len(expected), len(counts), counts) + } +} + +func TestBuildCreateGpuInstancesArg(t *testing.T) { + testCases := []struct { + name string + migDevices map[string]int32 + want string + }{ + { + name: "homogeneous rtx pro 6000", + migDevices: map[string]int32{"1g.24gb": 4}, + want: "1g.24gb,1g.24gb,1g.24gb,1g.24gb", + }, + { + name: "mixed places larger slices first", + migDevices: map[string]int32{"1g.5gb": 1, "2g.10gb": 3}, + want: "2g.10gb,2g.10gb,2g.10gb,1g.5gb", + }, + { + name: "empty", + migDevices: map[string]int32{}, + want: "", + }, + { + name: "zero count skipped", + migDevices: map[string]int32{"1g.24gb": 0}, + want: "", + }, + } + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + if got := buildCreateGpuInstancesArg(tc.migDevices); got != tc.want { + t.Errorf("expected %q, got %q", tc.want, got) + } + }) + } +} + +func TestMigProfileSlices(t *testing.T) { + testCases := map[string]int{ + "1g.24gb": 1, + "2g.48gb": 2, + "4g.96gb": 4, + "7g.40gb": 7, + "garbage": 1, + "g.10gb": 1, + "0g.10gb": 1, + } + for name, want := range testCases { + if got := migProfileSlices(name); got != want { + t.Errorf("migProfileSlices(%q): expected %d, got %d", name, want, got) + } + } +} + func TestGetNextDeviceRequest_DeviceInRegularContainer(t *testing.T) { // Save and restore InRequestDevices oldInRequestDevices := device.InRequestDevices