Group MIG instances by parent GPU when applying allocation policies

internal/rm/allocate.go

@@ -92,14 +92,14 @@ func (r *resourceManager) prepareCandidates(available, required []string, size i
 
 	replicas := make(map[string]*replicaCount)
 	for _, c := range candidates {
-		id := AnnotatedID(c).GetID()
+		id := r.devices.PhysicalGPUKey(c)
 		if _, exists := replicas[id]; !exists {
 			replicas[id] = &replicaCount{}
 		}
 		replicas[id].available++
 	}
 	for d := range r.devices {
-		id := AnnotatedID(d).GetID()
+		id := r.devices.PhysicalGPUKey(d)
 		if _, exists := replicas[id]; !exists {
 			continue
 		}
@@ -160,7 +160,7 @@ func (r *resourceManager) greedyAlloc(available, required []string, size int, pr
 	// updates the map entry, keeping a single source of truth.
 	byGPU := make(map[string]*gpuAllocState)
 	for _, c := range candidates {
-		id := AnnotatedID(c).GetID()
+		id := r.devices.PhysicalGPUKey(c)
 		item, ok := byGPU[id]
 		if !ok {
 			item = &gpuAllocState{count: replicas[id]}
@@ -173,7 +173,7 @@ func (r *resourceManager) greedyAlloc(available, required []string, size int, pr
 	// touched()); it keeps spread from re-picking a GPU already pinned to the
 	// pod, while leaving distributed and packed unchanged.
 	for _, req := range required {
-		if item, ok := byGPU[AnnotatedID(req).GetID()]; ok {
+		if item, ok := byGPU[r.devices.PhysicalGPUKey(req)]; ok {
 			item.requiredReplicas++
 		}
 	}

internal/rm/allocate_test.go

@@ -105,6 +105,39 @@ func TestDistributedAlloc_PartiallyAllocated_DistributesAcrossDistinctGPUs(t *te
 		counts)
 }
 
+// countPerPhysicalGPU counts allocated device IDs per physical GPU, grouping
+// MIG instances that share a parent GPU together (see Devices.PhysicalGPUKey).
+func countPerPhysicalGPU(devices Devices, allocated []string) map[string]int {
+	counts := make(map[string]int)
+	for _, id := range allocated {
+		counts[devices.PhysicalGPUKey(id)]++
+	}
+	return counts
+}
+
+func TestDistributedAlloc_MIG_SpansDistinctPhysicalGPUs(t *testing.T) {
+	// GPU 0 has two idle MIG instances; GPU 1 has one time-sliced instance with
+	// a replica already in use. Keyed by parent, distributed avoids busier GPU 0.
+	devices := Devices{
+		"MIG-0-0::0": {Device: pluginapi.Device{ID: "MIG-0-0::0", Health: pluginapi.Healthy}, Index: "0:0", Replicas: 1},
+		"MIG-0-1::0": {Device: pluginapi.Device{ID: "MIG-0-1::0", Health: pluginapi.Healthy}, Index: "0:1", Replicas: 1},
+		"MIG-1-0::0": {Device: pluginapi.Device{ID: "MIG-1-0::0", Health: pluginapi.Healthy}, Index: "1:0", Replicas: 2},
+		"MIG-1-0::1": {Device: pluginapi.Device{ID: "MIG-1-0::1", Health: pluginapi.Healthy}, Index: "1:0", Replicas: 2},
+	}
+	r := &resourceManager{devices: devices}
+
+	// MIG-1-0::1 is omitted from available, i.e. already allocated.
+	available := []string{"MIG-0-0::0", "MIG-0-1::0", "MIG-1-0::0"}
+
+	allocated, err := r.greedyAlloc(available, nil, 2, comparatorForPolicy(spec.AllocationPolicyDistributed))
+	require.NoError(t, err)
+	require.Len(t, allocated, 2)
+
+	counts := countPerPhysicalGPU(devices, allocated)
+	require.Lenf(t, counts, 2,
+		"expected the 2 slots on distinct physical GPUs, not two MIG instances of one card; got %v", allocated)
+}
+
 func TestDistributedAlloc(t *testing.T) {
 	testCases := []struct {
 		description string

internal/rm/devices.go

@@ -234,6 +234,16 @@ func (d *Device) GetUUID() string {
 	return AnnotatedID(d.ID).GetID()
 }
 
+// PhysicalGPUKey returns a key for the physical GPU backing id, grouping MIG
+// instances by parent (Device.Index prefix "<gpu>:<mig>") and others by UUID.
+func (ds Devices) PhysicalGPUKey(id string) string {
+	if d := ds.GetByID(id); d != nil && d.IsMigDevice() {
+		parent, _, _ := strings.Cut(d.Index, ":")
+		return "mig-parent:" + parent
+	}
+	return AnnotatedID(id).GetID()
+}
+
 // NewAnnotatedID creates a new AnnotatedID from an ID and a replica number.
 func NewAnnotatedID(id string, replica int) AnnotatedID {
 	return AnnotatedID(fmt.Sprintf("%s::%d", id, replica))