Add spread allocation policy for maximum per-pod distinct-GPU coverage

api/config/v1/consts.go

@@ -52,6 +52,7 @@ const (
 const (
 	AllocationPolicyDistributed = "distributed"
 	AllocationPolicyPacked      = "packed"
+	AllocationPolicySpread      = "spread"
 )
 
 // Constants related to generating CDI specifications

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

internal/rm/allocate.go

@@ -37,21 +37,36 @@ 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
 // (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()
+	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,26 +112,29 @@ 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
 }
 
 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,

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")
+}