diff --git a/internal/rm/allocate.go b/internal/rm/allocate.go index 64a686363..9588ee44b 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" ) @@ -94,6 +94,39 @@ func (r *resourceManager) prepareCandidates(available, required []string, size i return candidates, replicas, needed, nil } +// 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 +} + +// 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. +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 +} +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 +137,48 @@ 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. 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. + 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_bench_test.go b/internal/rm/allocate_bench_test.go new file mode 100644 index 000000000..e166e7baa --- /dev/null +++ b/internal/rm/allocate_bench_test.go @@ -0,0 +1,58 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package rm + +import ( + "fmt" + "testing" + + spec "github.com/NVIDIA/k8s-device-plugin/api/config/v1" +) + +// BenchmarkGreedyAlloc measures greedyAlloc across a range of node shapes. +// gpus = physical GPUs, replicas per GPU (candidates n = gpus*replicas), +// request = slots requested. +func BenchmarkGreedyAlloc(b *testing.B) { + scenarios := []struct { + gpus, replicas, request int + }{ + {gpus: 4, replicas: 4, request: 4}, // n=16 — small node + {gpus: 8, replicas: 8, request: 8}, // n=64 — typical dense (8-GPU) node + {gpus: 8, replicas: 16, request: 16}, // n=128 — 8-GPU node, aggressive sharing + } + + cmp := comparatorForPolicy(spec.AllocationPolicyDistributed) + for _, s := range scenarios { + ids := make([]string, s.gpus) + for i := 0; i < s.gpus; i++ { + ids[i] = fmt.Sprintf("gpu%d", i) + } + devices := newTestDevices(ids, s.replicas) + available := getDeviceIDs(devices) + r := &resourceManager{config: &spec.Config{}, devices: devices} + name := fmt.Sprintf("gpus=%d/replicas=%d/req=%d/n=%d", s.gpus, s.replicas, s.request, s.gpus*s.replicas) + + b.Run(name, func(b *testing.B) { + b.ReportAllocs() + for i := 0; i < b.N; i++ { + if _, err := r.greedyAlloc(available, nil, s.request, cmp); err != nil { + b.Fatal(err) + } + } + }) + } +}