fix(mig): support RTX PRO 6000 Blackwell dynamic MIG

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 }}

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

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) {

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