From c3b47c6fcec8ae213efa7d4d3011b09414e91fcc Mon Sep 17 00:00:00 2001 From: Jonathan Meiri <33288957+Meiri28@users.noreply.github.com> Date: Wed, 22 Jul 2026 11:24:42 +0300 Subject: [PATCH 1/2] perf(rm): replace per-iteration sort in greedyAlloc with a min-heap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-up on top of #1621, which introduced the shared greedyAlloc loop with a pluggable replicaComparator (distributed vs packed). The loop still sorts the full candidate slice inside the allocation loop, paying O(n log n) per iteration for n iterations and giving O(n² log n) overall. Since all annotated replicas from the same underlying physical device share the same sort key, sorting at the replica granularity is wasted work — only m (the number of distinct physical devices contributing candidates) needs to be reordered. Refactor greedyAlloc to bucket candidates by their underlying physical device into a small gpuAllocState per device, holding a shared *replicaCount, the pickedFrom counter, and the remaining candidate IDs. A gpuPriorityQueue defers to the caller-supplied replicaComparator on allocated() for primary ordering and to pickedFrom for the tie-break (unchanged semantics). Each iteration pops the best device, takes one of its remaining replicas, updates counters, and pushes it back if any remain. Total cost drops to O(n log m). Both allocation policies (distributed and packed) benefit; no behavior change — the existing test suite (TestDistributedAlloc, TestPackedAlloc, TestPackedVsDistributedContrast, TestDistributedAlloc_PartiallyAllocated_DistributesAcrossDistinctGPUs, etc.) passes unchanged. 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> --- internal/rm/allocate.go | 104 +++++++++++++++++++++++++++++----------- 1 file changed, 75 insertions(+), 29 deletions(-) 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 From 498aba4613b1385a041261f283bdc7756cc6c12b Mon Sep 17 00:00:00 2001 From: Jonathan Meiri <33288957+Meiri28@users.noreply.github.com> Date: Wed, 2 Sep 2026 09:01:23 +0300 Subject: [PATCH 2/2] test(rm): add BenchmarkGreedyAlloc 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> --- internal/rm/allocate_bench_test.go | 58 ++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 internal/rm/allocate_bench_test.go 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) + } + } + }) + } +}