implement round-robin rotation for equal-ratio candidate selection in model queue

- Replace float64 ratio comparison with integer-only load ratio math
  to avoid floating-point precision issues.
- Collect all equal-ratio candidates and rotate via lastSelectedSlot so
  equally loaded providers do not pin to a single provider.
- Add loadRatioLess and loadRatioEqual helper functions.
- Track lastSelectedSlot in admit and tryDispatchLocked paths.
This commit is contained in:
toki 2026-06-24 22:43:37 +09:00
parent 395c0b4708
commit c341d622ea
2 changed files with 136 additions and 19 deletions

View file

@ -3,6 +3,7 @@ package service
import (
"context"
"fmt"
"sort"
"sync"
"time"
@ -70,12 +71,13 @@ type queueItem struct {
}
type modelQueueGroup struct {
key string
policy groupPolicy
queue []*queueItem
inflight map[string]int // nodeID → current in-flight count
adapter string
target string
key string
policy groupPolicy
queue []*queueItem
inflight map[string]int // slotKey → current in-flight count
lastSelectedSlot string
adapter string
target string
}
type modelQueueManager struct {
@ -177,32 +179,65 @@ func (m *modelQueueManager) updateGroupPolicyLocked(g *modelQueueGroup, policy g
}
// findAvailableNodeLocked returns the available candidate with the lowest
// in_flight/capacity load ratio. Ties are broken deterministically by
// providerID then nodeID to produce a stable ordering across calls.
// in_flight/capacity load ratio. Equal-ratio candidates rotate from the last
// selected slot so idle or equally loaded providers do not pin to one provider.
// Uses provider-aware slot keys so that multiple providers on the same node
// are tracked independently.
func (m *modelQueueManager) findAvailableNodeLocked(group *modelQueueGroup, candidates []candidateNode) *candidateNode {
var best *candidateNode
bestRatio := 2.0 // sentinel: any valid ratio is in [0, 1)
var best []candidateNode
bestInflight := 0
bestCapacity := 1
for i := range candidates {
c := &candidates[i]
c := candidates[i]
if c.capacity <= 0 {
continue
}
slot := c.slotKey()
inflight := group.inflight[slot]
if inflight >= c.capacity {
continue
}
ratio := float64(inflight) / float64(c.capacity)
if ratio < bestRatio || (ratio == bestRatio && candidateLess(c, best)) {
best = c
bestRatio = ratio
switch {
case len(best) == 0 || loadRatioLess(inflight, c.capacity, bestInflight, bestCapacity):
best = []candidateNode{c}
bestInflight = inflight
bestCapacity = c.capacity
case loadRatioEqual(inflight, c.capacity, bestInflight, bestCapacity):
best = append(best, c)
}
}
return best
if len(best) == 0 {
return nil
}
sort.Slice(best, func(i, j int) bool {
return candidateLess(&best[i], &best[j])
})
idx := 0
if group.lastSelectedSlot != "" {
for i := range best {
if best[i].slotKey() == group.lastSelectedSlot {
idx = (i + 1) % len(best)
break
}
}
}
return &best[idx]
}
// candidateLess provides a deterministic tie-break ordering for candidates
// with equal load ratios: providerID first, then nodeID.
func loadRatioLess(leftInflight, leftCapacity, rightInflight, rightCapacity int) bool {
return leftInflight*rightCapacity < rightInflight*leftCapacity
}
func loadRatioEqual(leftInflight, leftCapacity, rightInflight, rightCapacity int) bool {
return leftInflight*rightCapacity == rightInflight*leftCapacity
}
// candidateLess provides a deterministic ordering for equal-ratio rotation:
// providerID first, then nodeID.
func candidateLess(a, b *candidateNode) bool {
if b == nil {
return true
}
if a.providerID != b.providerID {
return a.providerID < b.providerID
}
@ -232,6 +267,7 @@ func (m *modelQueueManager) admit(ctx context.Context, groupKey, adapter, target
if candidate != nil {
slot := candidate.slotKey()
group.inflight[slot]++
group.lastSelectedSlot = slot
m.mu.Unlock()
return candidate, nil
}
@ -468,6 +504,7 @@ func (m *modelQueueManager) tryDispatchLocked(group *modelQueueGroup) {
select {
case head.waitCh <- admitResult{candidate: candidate}:
group.lastSelectedSlot = slot
return
default:
// Caller already timed out or cancelled; release the reserved slot and try next.

View file

@ -525,7 +525,7 @@ func TestModelQueueProviderLoadRatioSelection(t *testing.T) {
}
// B is now at inflight=1 (ratio=0.5). A is still at inflight=2 (ratio=0.5).
// With equal ratios, tie-break by providerID: "prov-a" < "prov-b" → A wins.
// Equal-ratio rotation advances after the last selected slot (B), so A wins.
sel2, err := m.admit(context.Background(), "g-lr", "", "", cands, groupPolicy{})
if err != nil {
t.Fatalf("admit2: %v", err)
@ -535,6 +535,86 @@ func TestModelQueueProviderLoadRatioSelection(t *testing.T) {
}
}
func TestModelQueueProviderLoadRatioFillsProportionally(t *testing.T) {
gx10 := &edgenode.NodeEntry{NodeID: "node-gx10"}
onex := &edgenode.NodeEntry{NodeID: "node-onex"}
cands := []candidateNode{
{entry: gx10, capacity: 4, providerID: "gx10-vllm"},
{entry: onex, capacity: 3, providerID: "onexplayer-lemonade"},
}
m := newModelQueueManager(nil)
var got []string
for i := 0; i < 7; i++ {
sel, err := m.admit(context.Background(), "g-proportional", "", "", cands, groupPolicy{})
if err != nil {
t.Fatalf("admit %d: %v", i+1, err)
}
got = append(got, sel.providerID)
}
want := []string{
"gx10-vllm",
"onexplayer-lemonade",
"gx10-vllm",
"onexplayer-lemonade",
"gx10-vllm",
"onexplayer-lemonade",
"gx10-vllm",
}
if len(got) != len(want) {
t.Fatalf("dispatch count: got %d want %d", len(got), len(want))
}
for i := range want {
if got[i] != want[i] {
t.Fatalf("dispatch[%d]: got %q want %q; full sequence=%v", i, got[i], want[i], got)
}
}
m.mu.Lock()
defer m.mu.Unlock()
g := m.groups["g-proportional"]
if g.inflight["node-gx10:gx10-vllm"] != 4 || g.inflight["node-onex:onexplayer-lemonade"] != 3 {
t.Fatalf("inflight: got gx10=%d onex=%d, want 4/3", g.inflight["node-gx10:gx10-vllm"], g.inflight["node-onex:onexplayer-lemonade"])
}
}
func TestModelQueueProviderEqualIdleTieRotates(t *testing.T) {
gx10 := &edgenode.NodeEntry{NodeID: "node-gx10-idle"}
onex := &edgenode.NodeEntry{NodeID: "node-onex-idle"}
cands := []candidateNode{
{entry: gx10, capacity: 4, providerID: "gx10-vllm"},
{entry: onex, capacity: 3, providerID: "onexplayer-lemonade"},
}
m := newModelQueueManager(nil)
sel1, err := m.admit(context.Background(), "g-idle-rotate", "", "", cands, groupPolicy{})
if err != nil {
t.Fatalf("first admit: %v", err)
}
if sel1.providerID != "gx10-vllm" {
t.Fatalf("first admit provider: got %q want gx10-vllm", sel1.providerID)
}
m.releaseSlot("g-idle-rotate", sel1.entry.NodeID, sel1.providerID)
sel2, err := m.admit(context.Background(), "g-idle-rotate", "", "", cands, groupPolicy{})
if err != nil {
t.Fatalf("second admit: %v", err)
}
if sel2.providerID != "onexplayer-lemonade" {
t.Fatalf("second admit provider: got %q want onexplayer-lemonade", sel2.providerID)
}
m.releaseSlot("g-idle-rotate", sel2.entry.NodeID, sel2.providerID)
sel3, err := m.admit(context.Background(), "g-idle-rotate", "", "", cands, groupPolicy{})
if err != nil {
t.Fatalf("third admit: %v", err)
}
if sel3.providerID != "gx10-vllm" {
t.Fatalf("third admit provider: got %q want gx10-vllm", sel3.providerID)
}
}
// TestModelQueueProviderServedTargetRewrite verifies that the selected
// candidateNode carries its servedTarget so callers can rewrite req.Target.
func TestModelQueueProviderServedTargetRewrite(t *testing.T) {