diff --git a/api/config/v1/consts.go b/api/config/v1/consts.go index 925b9507a..5fa6c4fba 100644 --- a/api/config/v1/consts.go +++ b/api/config/v1/consts.go @@ -52,6 +52,7 @@ const ( const ( AllocationPolicyDistributed = "distributed" AllocationPolicyPacked = "packed" + AllocationPolicySpread = "spread" ) // Constants related to generating CDI specifications diff --git a/cmd/nvidia-device-plugin/main.go b/cmd/nvidia-device-plugin/main.go index ebacd3edd..422c41ffe 100644 --- a/cmd/nvidia-device-plugin/main.go +++ b/cmd/nvidia-device-plugin/main.go @@ -219,6 +219,7 @@ func validateFlags(infolib nvinfo.Interface, config *spec.Config) error { switch *config.Flags.Plugin.SharedDevicesAllocationPolicy { case spec.AllocationPolicyDistributed: case spec.AllocationPolicyPacked: + case spec.AllocationPolicySpread: default: return fmt.Errorf("invalid --shared-devices-allocation-policy option: %s", *config.Flags.Plugin.SharedDevicesAllocationPolicy) } diff --git a/internal/rm/allocate.go b/internal/rm/allocate.go index 64a686363..ecb6f8777 100644 --- a/internal/rm/allocate.go +++ b/internal/rm/allocate.go @@ -17,8 +17,8 @@ package rm import ( + "container/heap" "fmt" - "sort" spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1" ) @@ -34,24 +34,57 @@ func (rc *replicaCount) allocated() int { return rc.total - rc.available } +// gpuAllocState is the per-physical-GPU bookkeeping the greedy allocator +// tracks while it consumes candidates. A comparator sees both the shared +// cluster-wide replicaCount and the per-allocation pickedFrom counter so it +// can express policies that mix the two axes (e.g. spread, which orders +// primarily by pickedFrom and only tie-breaks by allocated()). +type gpuAllocState struct { + count *replicaCount // shared reference to this GPU's replicaCount + pickedFrom int // slots picked from this GPU during this allocation + replicas []string // remaining annotated-ID candidates for this GPU +} + // replicaComparator decides whether the physical GPU represented by i should // be preferred over the one represented by j when greedily selecting the next -// device to allocate. -type replicaComparator func(i, j *replicaCount) bool +// device to allocate. Comparators are complete Less functions — they own both +// the primary ordering and the tie-break — so a new policy can freely choose +// which axis of gpuAllocState matters most. +type replicaComparator func(i, j *gpuAllocState) bool // allocationComparators maps each allocation policy to the comparator that // implements it. All policies share the same greedy selection loop // (greedyAlloc) and differ only in how the next best candidate is chosen. var allocationComparators = map[string]replicaComparator{ // distributed prefers GPUs with the fewest allocated replicas to spread - // workload evenly across physical GPUs. - spec.AllocationPolicyDistributed: func(i, j *replicaCount) bool { - return i.allocated() < j.allocated() + // workload evenly across physical GPUs. Equal allocated counts fall back + // to pickedFrom so the pod's own picks also rotate. + spec.AllocationPolicyDistributed: func(i, j *gpuAllocState) bool { + if i.count.allocated() != j.count.allocated() { + return i.count.allocated() < j.count.allocated() + } + return i.pickedFrom < j.pickedFrom }, // packed prefers GPUs with the most allocated replicas to consolidate - // workloads onto fewer physical GPUs. - spec.AllocationPolicyPacked: func(i, j *replicaCount) bool { - return i.allocated() > j.allocated() + // workloads onto fewer physical GPUs. Equal allocated counts still + // rotate by pickedFrom so a single pod that overflows a GPU spreads its + // overflow rather than concentrating arbitrarily. + spec.AllocationPolicyPacked: func(i, j *gpuAllocState) bool { + if i.count.allocated() != j.count.allocated() { + return i.count.allocated() > j.count.allocated() + } + return i.pickedFrom < j.pickedFrom + }, + // spread maximizes the number of distinct physical GPUs the current + // allocation touches: it prefers GPUs the pod has not yet picked from, + // falling back to less-allocated when pickedFrom ties. Useful for + // multi-GPU workloads (e.g. distributed training / NCCL) that benefit + // from spanning physical hardware. + spec.AllocationPolicySpread: func(i, j *gpuAllocState) bool { + if i.pickedFrom != j.pickedFrom { + return i.pickedFrom < j.pickedFrom + } + return i.count.allocated() < j.count.allocated() }, } @@ -94,6 +127,24 @@ func (r *resourceManager) prepareCandidates(available, required []string, size i return candidates, replicas, needed, nil } +// gpuPriorityQueue is a heap of *gpuAllocState whose ordering is fully +// determined by the caller-supplied comparator. +type gpuPriorityQueue struct { + items []*gpuAllocState + preferred replicaComparator +} + +func (q *gpuPriorityQueue) Len() int { return len(q.items) } +func (q *gpuPriorityQueue) Less(i, j int) bool { return q.preferred(q.items[i], q.items[j]) } +func (q *gpuPriorityQueue) Swap(i, j int) { q.items[i], q.items[j] = q.items[j], q.items[i] } +func (q *gpuPriorityQueue) Push(x any) { q.items = append(q.items, x.(*gpuAllocState)) } +func (q *gpuPriorityQueue) Pop() any { + n := len(q.items) - 1 + x := q.items[n] + q.items = q.items[:n] + return x +} + // greedyAlloc returns a list of devices by repeatedly selecting the best // remaining candidate according to the supplied comparator. It takes into // account already allocated replicas so that consecutive allocations keep @@ -104,35 +155,45 @@ func (r *resourceManager) greedyAlloc(available, required []string, size int, pr return nil, err } - // Track how many slots have already been picked from each physical device - // during this allocation. Used as the tie-break sort key below so that, - // when the comparator ranks two physical GPUs equally, the allocator - // rotates to a sibling device it has touched the least this round. This - // keeps the distributed policy spreading replicas across physical GPUs - // even when their allocated counts tie. - pickedFrom := make(map[string]int) - - // Select devices one-by-one. The supplied comparator decides which - // physical GPU is preferred for the current policy. Comparators order - // solely by allocated() (see TestComparatorsOrderSolelyByAllocated), so - // equal allocated counts mean the comparator has no preference and the - // pickedFrom tie-break above applies. - var devices []string + // Bucket candidates by their underlying physical GPU. Each gpuAllocState + // holds a shared *replicaCount so decrementing its available count also + // updates the map entry, keeping a single source of truth. + byGPU := make(map[string]*gpuAllocState) + for _, c := range candidates { + id := AnnotatedID(c).GetID() + item, ok := byGPU[id] + if !ok { + item = &gpuAllocState{count: replicas[id]} + byGPU[id] = item + } + item.replicas = append(item.replicas, c) + } + + // Build the heap once and let the comparator drive ordering. + pq := &gpuPriorityQueue{ + items: make([]*gpuAllocState, 0, len(byGPU)), + preferred: preferred, + } + for _, item := range byGPU { + pq.items = append(pq.items, item) + } + heap.Init(pq) + + // Pop the best GPU, take one of its replicas, update counters, push back + // if any remain. Total cost is O(n log m) where n is `needed` and m is + // the number of distinct physical devices contributing candidates. + devices := make([]string, 0, needed) for i := 0; i < needed; i++ { - sort.Slice(candidates, func(i, j int) bool { - iid := AnnotatedID(candidates[i]).GetID() - jid := AnnotatedID(candidates[j]).GetID() - ri, rj := replicas[iid], replicas[jid] - if ri.allocated() != rj.allocated() { - return preferred(ri, rj) - } - return pickedFrom[iid] < pickedFrom[jid] - }) - id := AnnotatedID(candidates[0]).GetID() - pickedFrom[id]++ - replicas[id].available-- - devices = append(devices, candidates[0]) - candidates = candidates[1:] + top := heap.Pop(pq).(*gpuAllocState) + last := len(top.replicas) - 1 + pick := top.replicas[last] + top.replicas = top.replicas[:last] + top.count.available-- + top.pickedFrom++ + if len(top.replicas) > 0 { + heap.Push(pq, top) + } + devices = append(devices, pick) } return append(required, devices...), nil diff --git a/internal/rm/allocate_test.go b/internal/rm/allocate_test.go index 7344149f0..df20bb1d3 100644 --- a/internal/rm/allocate_test.go +++ b/internal/rm/allocate_test.go @@ -405,8 +405,8 @@ func TestPackedVsDistributedContrast(t *testing.T) { // the comparator implementing it, and that unknown or empty policies fall // back to the default distributed comparator. func TestComparatorForPolicy(t *testing.T) { - moreAllocated := &replicaCount{total: 4, available: 1} // 3 allocated - lessAllocated := &replicaCount{total: 4, available: 3} // 1 allocated + moreAllocated := &gpuAllocState{count: &replicaCount{total: 4, available: 1}} // 3 allocated + lessAllocated := &gpuAllocState{count: &replicaCount{total: 4, available: 3}} // 1 allocated testCases := []struct { description string @@ -425,6 +425,11 @@ func TestComparatorForPolicy(t *testing.T) { policy: spec.AllocationPolicyPacked, expectPrefersLessAllocated: false, }, + { + description: "spread with equal pickedFrom falls back to less allocated", + policy: spec.AllocationPolicySpread, + expectPrefersLessAllocated: true, + }, { description: "empty policy falls back to distributed", policy: "", @@ -447,24 +452,50 @@ func TestComparatorForPolicy(t *testing.T) { } } +// TestSpreadPrefersUntouchedGPU pins spread's defining behavior: given equal +// (or even unequal) allocated counts, it always prefers the GPU the current +// allocation has touched the least. This is what makes it maximize distinct +// physical GPUs per pod. +func TestSpreadPrefersUntouchedGPU(t *testing.T) { + spread := comparatorForPolicy(spec.AllocationPolicySpread) + + // Even when GPU A has strictly less allocated capacity, spread still + // prefers GPU B if the current allocation has picked from A already. + touched := &gpuAllocState{count: &replicaCount{total: 8, available: 6}, pickedFrom: 1} // 2 allocated + untouched := &gpuAllocState{count: &replicaCount{total: 8, available: 3}, pickedFrom: 0} // 5 allocated + require.True(t, spread(untouched, touched), "spread must prefer the untouched GPU even when it has more allocated replicas") + require.False(t, spread(touched, untouched)) +} + // TestComparatorsOrderSolelyByAllocated pins the invariant that every // allocation comparator orders physical GPUs solely by their allocated() // count. The tie-break in greedyAlloc depends on this: it treats equal // allocated counts as "the comparator has no preference" and falls back to // the pickedFrom rotation, so a comparator that distinguishes GPUs by // anything else would be silently ignored there. +// TestComparatorsOrderSolelyByAllocated pins the invariant for the +// allocated()-primary policies (distributed, packed): with matching +// pickedFrom, they order GPUs solely by allocated(). The spread policy is +// explicitly excluded — it primarily orders by pickedFrom by design; see +// TestSpreadPrefersUntouchedGPU. func TestComparatorsOrderSolelyByAllocated(t *testing.T) { - for policy, preferred := range allocationComparators { + allocatedPrimaryPolicies := []string{ + spec.AllocationPolicyDistributed, + spec.AllocationPolicyPacked, + } + for _, policy := range allocatedPrimaryPolicies { + preferred := allocationComparators[policy] t.Run(policy, func(t *testing.T) { // Equal allocated counts with different total/available shapes - // must rank equal so the tie-break applies. - a := &replicaCount{total: 8, available: 6} // 2 allocated - b := &replicaCount{total: 4, available: 2} // 2 allocated + // must rank equal (when pickedFrom is also equal) so the + // greedyAlloc tie-break applies. + a := &gpuAllocState{count: &replicaCount{total: 8, available: 6}} // 2 allocated + b := &gpuAllocState{count: &replicaCount{total: 4, available: 2}} // 2 allocated require.False(t, preferred(a, b), "GPUs with equal allocated counts must rank equal") require.False(t, preferred(b, a), "GPUs with equal allocated counts must rank equal") // Different allocated counts must be strictly ordered. - c := &replicaCount{total: 8, available: 5} // 3 allocated + c := &gpuAllocState{count: &replicaCount{total: 8, available: 5}} // 3 allocated require.NotEqual(t, preferred(a, c), preferred(c, a), "GPUs with different allocated counts must be strictly ordered") }) } @@ -538,3 +569,136 @@ func TestFullGPUNodeIgnoresAllocationPolicy(t *testing.T) { require.True(t, AnnotatedIDs(replicatedAvailable).AnyHasAnnotations(), "replicated device IDs should have annotations") }) } + +func TestSpreadAlloc(t *testing.T) { + testCases := []struct { + description string + gpuIDs []string + replicas int + available []string // if nil, use all devices + required []string + size int + expectError bool + validate func(t *testing.T, allocated []string, allDevices Devices) + }{ + { + description: "2 GPUs, 4 replicas each, allocate 2 — should spread across distinct GPUs", + gpuIDs: []string{"gpu0", "gpu1"}, + replicas: 4, + required: []string{}, + size: 2, + validate: func(t *testing.T, allocated []string, _ Devices) { + counts := countPerGPU(allocated) + require.Len(t, allocated, 2) + require.Equal(t, 1, counts["gpu0"], "spread should pick one from each GPU") + require.Equal(t, 1, counts["gpu1"], "spread should pick one from each GPU") + }, + }, + { + description: "3 GPUs, 4 replicas each, allocate 3 — should spread across all 3 GPUs", + gpuIDs: []string{"gpu0", "gpu1", "gpu2"}, + replicas: 4, + required: []string{}, + size: 3, + validate: func(t *testing.T, allocated []string, _ Devices) { + counts := countPerGPU(allocated) + require.Len(t, allocated, 3) + require.Equal(t, 1, counts["gpu0"]) + require.Equal(t, 1, counts["gpu1"]) + require.Equal(t, 1, counts["gpu2"]) + }, + }, + { + description: "3 GPUs, 4 replicas each, allocate 6 — should hit each GPU twice", + gpuIDs: []string{"gpu0", "gpu1", "gpu2"}, + replicas: 4, + required: []string{}, + size: 6, + validate: func(t *testing.T, allocated []string, _ Devices) { + counts := countPerGPU(allocated) + require.Len(t, allocated, 6) + require.Equal(t, 2, counts["gpu0"]) + require.Equal(t, 2, counts["gpu1"]) + require.Equal(t, 2, counts["gpu2"]) + }, + }, + { + description: "allocate 1 from single GPU — trivial case", + gpuIDs: []string{"gpu0"}, + replicas: 4, + required: []string{}, + size: 1, + validate: func(t *testing.T, allocated []string, _ Devices) { + require.Len(t, allocated, 1) + counts := countPerGPU(allocated) + require.Equal(t, 1, counts["gpu0"]) + }, + }, + { + description: "not enough devices — should return error", + gpuIDs: []string{"gpu0"}, + replicas: 2, + required: []string{}, + size: 5, + expectError: true, + }, + } + + for _, tc := range testCases { + t.Run(tc.description, func(t *testing.T) { + devices := newTestDevices(tc.gpuIDs, tc.replicas) + available := tc.available + if available == nil { + available = getDeviceIDs(devices) + } + + rm := resourceManager{ + config: &spec.Config{}, + devices: devices, + } + + allocated, err := rm.greedyAlloc(available, tc.required, tc.size, comparatorForPolicy(spec.AllocationPolicySpread)) + if tc.expectError { + require.Error(t, err) + return + } + require.NoError(t, err) + if tc.validate != nil { + tc.validate(t, allocated, devices) + } + }) + } +} + +// TestSpreadPrefersDistinctGPUsEvenWhenUnbalanced captures the scenario that +// motivates the policy: with GPU-0 having fewer free slots than GPU-1, the +// distributed policy consolidates a 2-slot request onto GPU-1 (whichever has +// less allocated cluster-wide). Spread instead touches each physical GPU. +func TestSpreadPrefersDistinctGPUsEvenWhenUnbalanced(t *testing.T) { + // GPU-0 has replicas=8, of which 5 are already allocated externally + // (only 3 free). GPU-1 has 8 replicas, 3 already allocated (5 free). + devices := newTestDevices([]string{"gpu0", "gpu1"}, 8) + // available = 3 slots on gpu0 + 5 slots on gpu1 + available := []string{ + "gpu0::5", "gpu0::6", "gpu0::7", + "gpu1::3", "gpu1::4", "gpu1::5", "gpu1::6", "gpu1::7", + } + + rm := resourceManager{config: &spec.Config{}, devices: devices} + + // A pod requesting 2 slots under spread must land 1 on each physical GPU + // even though gpu1 has more free capacity (distributed would concentrate + // both slots on gpu1). + allocated, err := rm.greedyAlloc(available, nil, 2, comparatorForPolicy(spec.AllocationPolicySpread)) + require.NoError(t, err) + require.Len(t, allocated, 2) + counts := countPerGPU(allocated) + require.Equalf(t, 1, counts["gpu0"], "spread must include the less-free GPU; got: %v", counts) + require.Equalf(t, 1, counts["gpu1"], "spread must include the more-free GPU; got: %v", counts) + + // Contrast: same setup under distributed concentrates on gpu1. + allocated, err = rm.greedyAlloc(available, nil, 2, comparatorForPolicy(spec.AllocationPolicyDistributed)) + require.NoError(t, err) + distCounts := countPerGPU(allocated) + require.Equalf(t, 2, distCounts["gpu1"], "distributed should pick both from the less-loaded GPU; got: %v", distCounts) +}