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