From 7062bd0ddb13aab54267b49c016b3539583302bc Mon Sep 17 00:00:00 2001 From: Jonathan Meiri <33288957+Meiri28@users.noreply.github.com> Date: Wed, 22 Jul 2026 12:13:37 +0300 Subject: [PATCH] Add spread allocation policy for maximum per-pod distinct-GPU coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `distributed` minimizes cluster-wide load imbalance — it prefers the GPU with the fewest replicas already allocated across all pods. `packed` does the opposite for bin-packing. Neither guarantees that a single pod's multi-slot request touches distinct physical GPUs: given GPU-0 with 5 allocated and GPU-1 with 3 allocated, a pod requesting 2 slots under `distributed` gets both slots on GPU-1, because GPU-1 has the lower cluster-wide allocated count both before and after the first pick. Add a third policy, `spread`, that primarily orders by pickedFrom (the per-allocation counter of how many slots the current pod has taken from each GPU) and only tie-breaks by allocated(). The pod's own picks therefore drive selection: after taking one slot from GPU-1, GPU-0 becomes preferred (pickedFrom=0 < 1) regardless of cluster-wide load. Result: the pod's slots span as many distinct physical GPUs as possible, which is what multi-GPU workloads (data-parallel training, NCCL, tensor-parallel) actually need. The `replicaComparator` signature is enriched from `func(i, j *replicaCount) bool` to `func(i, j *gpuAllocState) bool` so a comparator can freely mix cluster-wide state (allocated()) and per-allocation state (pickedFrom). Each policy now owns both its primary ordering and its tie-break; the queue's Less becomes a pure passthrough. distributed and packed are behavior-preserving — the primary+tie-break they used to get from greedyAlloc's wrapping is spelled out in the comparator body itself. Validate `spread` in main.go alongside distributed and packed. Add TestSpreadAlloc mirroring the existing policy suites plus TestSpreadPrefersUntouchedGPU and TestSpreadPrefersDistinctGPUsEvenWhenUnbalanced for the defining behavior. TestComparatorsOrderSolelyByAllocated narrows to distributed/packed since spread intentionally violates that invariant. Co-Authored-By: Claude Opus 4.7 (1M context) Co-Authored-By: runatom-ai <258621014+runatom-ai@users.noreply.github.com> Signed-off-by: Jonathan Meiri <33288957+Meiri28@users.noreply.github.com> --- api/config/v1/consts.go | 1 + cmd/nvidia-device-plugin/main.go | 3 +- internal/rm/allocate.go | 66 +++++++--- internal/rm/allocate_test.go | 212 +++++++++++++++++++++++++++++-- 4 files changed, 248 insertions(+), 34 deletions(-) diff --git a/api/config/v1/consts.go b/api/config/v1/consts.go index 925b9507a0..5fa6c4fba3 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 ebacd3eddc..fbe9ddffa3 100644 --- a/cmd/nvidia-device-plugin/main.go +++ b/cmd/nvidia-device-plugin/main.go @@ -145,7 +145,7 @@ func main() { &cli.StringFlag{ Name: "shared-devices-allocation-policy", Value: spec.AllocationPolicyDistributed, - Usage: "the allocation policy for replicated and MIG resources:\n\t\t[distributed | packed]", + Usage: "the allocation policy for replicated and MIG resources:\n\t\t[distributed | packed | spread]", EnvVars: []string{"SHARED_DEVICES_ALLOCATION_POLICY"}, }, &cli.StringFlag{ @@ -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 adf29703f5..206921c3cc 100644 --- a/internal/rm/allocate.go +++ b/internal/rm/allocate.go @@ -37,7 +37,7 @@ func (rc *replicaCount) allocated() int { // 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 +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 @@ -45,13 +45,28 @@ type replicaComparator func(i, j *replicaCount) bool 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() + 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() + 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 prefers GPUs the current allocation has touched least, to span + // distinct physical GPUs. touched() folds in required replicas so a GPU + // already pinned to this pod counts as picked from. + spec.AllocationPolicySpread: func(i, j *gpuAllocState) bool { + if i.touched() != j.touched() { + return i.touched() < j.touched() + } + return i.count.allocated() < j.count.allocated() }, } @@ -97,14 +112,21 @@ func (r *resourceManager) prepareCandidates(available, required []string, size i // gpuAllocState is the per-physical-GPU bookkeeping the greedy allocator // tracks while it consumes candidates. 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 + count *replicaCount // shared reference to this GPU's replicaCount + pickedFrom int // slots picked from this GPU during this allocation + requiredReplicas int // required replicas already pinned to this GPU + replicas []string // remaining annotated-ID candidates for this GPU +} + +// touched reports how many slots this allocation has committed to the GPU, +// counting both picks made so far and required replicas already pinned. Only +// spread orders by this; distributed and packed tie-break on pickedFrom alone. +func (s *gpuAllocState) touched() int { + return s.pickedFrom + s.requiredReplicas } -// gpuPriorityQueue is a heap of *gpuAllocState whose ordering defers to the -// policy comparator on allocated() and falls back to pickedFrom for the -// tie-break so equal-allocated GPUs rotate rather than concentrating on one. +// gpuPriorityQueue is a heap of *gpuAllocState whose ordering is determined by +// the policy comparator. type gpuPriorityQueue struct { items []*gpuAllocState preferred replicaComparator @@ -112,11 +134,7 @@ type gpuPriorityQueue struct { func (q *gpuPriorityQueue) Len() int { return len(q.items) } func (q *gpuPriorityQueue) Less(i, j int) bool { - a, b := q.items[i], q.items[j] - if a.count.allocated() != b.count.allocated() { - return q.preferred(a.count, b.count) - } - return a.pickedFrom < b.pickedFrom + 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)) } @@ -151,10 +169,18 @@ func (r *resourceManager) greedyAlloc(available, required []string, size int, pr item.replicas = append(item.replicas, c) } - // Build the heap once. The comparator ranks GPUs on allocated() and the - // pickedFrom tie-break rotates between equal-ranked ones so, e.g., the - // distributed policy keeps spreading replicas across physical GPUs even - // when their allocated counts tie. + // Record required replicas per physical GPU. Only spread consults this (via + // 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 { + item.requiredReplicas++ + } + } + + // Build the heap once. Ordering is the policy comparator's: distributed and + // packed rank on allocated() (pickedFrom breaks ties), spread ranks on + // touched() (allocated() breaks ties). pq := &gpuPriorityQueue{ items: make([]*gpuAllocState, 0, len(byGPU)), preferred: preferred, diff --git a/internal/rm/allocate_test.go b/internal/rm/allocate_test.go index d7ddf79e34..a9ee2e3fe8 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,38 @@ func TestComparatorForPolicy(t *testing.T) { } } -// 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. +// TestSpreadPrefersUntouchedGPU: spread prefers the GPU the current allocation +// has touched least, even when it has more allocated replicas. +func TestSpreadPrefersUntouchedGPU(t *testing.T) { + spread := comparatorForPolicy(spec.AllocationPolicySpread) + + 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: distributed and packed order solely by +// allocated() when pickedFrom is equal. spread is excluded (it orders by +// pickedFrom first — 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 +557,170 @@ 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: gpu0 has 3 free slots, gpu1 +// has 5. A 2-slot request under spread lands 1 on each; distributed puts both +// on gpu1. +func TestSpreadPrefersDistinctGPUsEvenWhenUnbalanced(t *testing.T) { + devices := newTestDevices([]string{"gpu0", "gpu1"}, 8) + 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} + + 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) +} + +// TestSpreadAccountsForRequired: a required replica already fixes gpu0 to this +// allocation, so spread must place the additional slot on gpu1 — even though +// gpu0 has more free capacity (fewer allocated) and would win the tie-break if +// required replicas were not counted by touched(). +func TestSpreadAccountsForRequired(t *testing.T) { + devices := newTestDevices([]string{"gpu0", "gpu1"}, 4) + // gpu0 has 3 free slots; gpu1 has 1 free (3 allocated elsewhere), so gpu0 + // has the lower allocated() count. + available := []string{"gpu0::1", "gpu0::2", "gpu0::3", "gpu1::3"} + required := []string{"gpu0::0"} + + rm := resourceManager{config: &spec.Config{}, devices: devices} + + allocated, err := rm.greedyAlloc(available, required, 2, comparatorForPolicy(spec.AllocationPolicySpread)) + require.NoError(t, err) + require.Len(t, allocated, 2) + counts := countPerGPU(allocated) + require.Equalf(t, 1, counts["gpu0"], "spread must not pick gpu0 again — it already holds the required replica; got: %v", counts) + require.Equalf(t, 1, counts["gpu1"], "spread must span to gpu1 given gpu0 is already required; got: %v", counts) +} + +// TestRequiredReplicasOnlyAffectSpread pins the required-replica accounting to +// spread: with allocated() and pickedFrom tied, distributed and packed must +// treat a required-pinned GPU no differently, while spread prefers the +// untouched one. +func TestRequiredReplicasOnlyAffectSpread(t *testing.T) { + rc := &replicaCount{total: 4, available: 2} // allocated() == 2, shared by both + pinned := &gpuAllocState{count: rc, requiredReplicas: 1} + free := &gpuAllocState{count: rc} + + for _, policy := range []string{spec.AllocationPolicyDistributed, spec.AllocationPolicyPacked} { + cmp := comparatorForPolicy(policy) + require.Falsef(t, cmp(pinned, free), "%s must not deprioritize a required-pinned GPU", policy) + require.Falsef(t, cmp(free, pinned), "%s must not prefer a required-pinned GPU", policy) + } + + spread := comparatorForPolicy(spec.AllocationPolicySpread) + require.True(t, spread(free, pinned), "spread must prefer the untouched GPU") + require.False(t, spread(pinned, free), "spread must not prefer the required-pinned GPU") +}