diff --git a/internal/rm/allocate.go b/internal/rm/allocate.go index 206921c3c..60c93e0f0 100644 --- a/internal/rm/allocate.go +++ b/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++ } } diff --git a/internal/rm/allocate_test.go b/internal/rm/allocate_test.go index a9ee2e3fe..72edd70e2 100644 --- a/internal/rm/allocate_test.go +++ b/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 diff --git a/internal/rm/devices.go b/internal/rm/devices.go index 8a2b635b9..2d6bd1b4a 100644 --- a/internal/rm/devices.go +++ b/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 ":") 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))