Skip to content

Commit 22ff86d

Browse files
localai-botmudler
andauthored
fix(distributed): round-robin replicas of the same model (#9695)
FindAndLockNodeWithModel previously ordered candidate replicas by in_flight ASC, available_vram DESC. The primary key is correct, but the tiebreaker meant that whenever in_flight tied — the common case at low to moderate concurrency where requests don't overlap — the node with the largest available VRAM won every pick. With autoscaling placing replicas of the same model on multiple nodes, the fattest GPU node ended up taking nearly all the load while the others sat idle. Insert last_used ASC between the two existing tiers. last_used is already refreshed inside the same transaction that increments in_flight (and by TouchNodeModel on cache hits in the router), so the "oldest-used" replica naturally rotates through the candidate set — strict round-robin without a schema change. available_vram DESC is demoted to a final tiebreaker for cold starts where last_used is identical across replicas. Placement queries (FindNodeWithVRAM, FindLeastLoadedNode, and the *FromSet variants) have the same fattest-GPU bias on tiebreakers but are higher-cost to fix consistently. Deferred to a follow-up so the routing fix can land first — for the user-observed symptom routing was the dominant cause anyway. Test: registry_test.go adds a focused spec that loads three replicas on three nodes with 24/16/8 GB VRAM and asserts each is picked at least twice across 9 in_flight-tied calls. Assisted-by: claude-code:claude-opus-4-7 [Read] [Edit] [Bash] [Grep] Signed-off-by: Ettore Di Giacinto <mudler@localai.io> Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
1 parent 4e154b5 commit 22ff86d

2 files changed

Lines changed: 49 additions & 3 deletions

File tree

core/services/nodes/registry.go

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -663,16 +663,24 @@ func (r *NodeRegistry) FindAndLockNodeWithModel(ctx context.Context, modelName s
663663
var node BackendNode
664664

665665
err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
666-
// Order by in_flight ASC (least busy replica), then by available_vram DESC
667-
// (prefer nodes with more free VRAM to spread load across the cluster).
666+
// Order by in_flight ASC (least busy replica), then by last_used ASC
667+
// (round-robin between equally-loaded replicas — oldest used wins, and
668+
// every successful pick refreshes last_used below, so the "oldest" naturally
669+
// rotates through the candidate set). available_vram DESC is the final
670+
// tiebreaker for cold starts where last_used is identical.
671+
//
672+
// Without the last_used tier, a tie on in_flight (the common case at low
673+
// to moderate concurrency where requests don't overlap) collapses to
674+
// "biggest GPU wins every time" and one node ends up taking nearly all
675+
// the load while replicas on other nodes sit idle.
668676
q := tx.Clauses(clause.Locking{Strength: "UPDATE"}).
669677
Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id").
670678
Where("node_models.model_name = ? AND node_models.state = ?", modelName, "loaded")
671679
if len(candidateNodeIDs) > 0 {
672680
q = q.Where("node_models.node_id IN ?", candidateNodeIDs)
673681
}
674682
if err := q.
675-
Order("node_models.in_flight ASC, backend_nodes.available_vram DESC").
683+
Order("node_models.in_flight ASC, node_models.last_used ASC, backend_nodes.available_vram DESC").
676684
First(&nm).Error; err != nil {
677685
return err
678686
}

core/services/nodes/registry_test.go

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -304,6 +304,44 @@ var _ = Describe("NodeRegistry", func() {
304304
Expect(foundNM.NodeID).To(Equal(included.ID))
305305
})
306306

307+
It("round-robins between replicas when in_flight ties (last_used tiebreaker)", func() {
308+
// Three replicas of the same model on three nodes, all with in_flight=0.
309+
// Without the last_used tiebreaker, the node with the largest available_vram
310+
// would win every pick and one node would take ~all the load. With it,
311+
// each successful pick refreshes last_used so the next pick rotates to
312+
// the oldest-used replica.
313+
fat := makeNode("rr-fat", "10.0.0.50:50051", 24_000_000_000)
314+
mid := makeNode("rr-mid", "10.0.0.51:50051", 16_000_000_000)
315+
small := makeNode("rr-small", "10.0.0.52:50051", 8_000_000_000)
316+
Expect(registry.Register(context.Background(), fat, true)).To(Succeed())
317+
Expect(registry.Register(context.Background(), mid, true)).To(Succeed())
318+
Expect(registry.Register(context.Background(), small, true)).To(Succeed())
319+
320+
Expect(registry.SetNodeModel(context.Background(), fat.ID, "rr-model", 0, "loaded", "", 0)).To(Succeed())
321+
Expect(registry.SetNodeModel(context.Background(), mid.ID, "rr-model", 0, "loaded", "", 0)).To(Succeed())
322+
Expect(registry.SetNodeModel(context.Background(), small.ID, "rr-model", 0, "loaded", "", 0)).To(Succeed())
323+
324+
// Decrement back to 0 after each pick so the next call sees a tie.
325+
// (FindAndLockNodeWithModel atomically increments to lock the row.)
326+
picks := make([]string, 0, 9)
327+
for i := 0; i < 9; i++ {
328+
n, nm, err := registry.FindAndLockNodeWithModel(context.Background(), "rr-model", nil)
329+
Expect(err).ToNot(HaveOccurred())
330+
picks = append(picks, n.Name)
331+
Expect(registry.DecrementInFlight(context.Background(), n.ID, "rr-model", nm.ReplicaIndex)).To(Succeed())
332+
}
333+
334+
// Each replica should have been picked at least twice across 9 ties —
335+
// proves we're rotating, not pinning to the largest-VRAM node.
336+
counts := map[string]int{}
337+
for _, p := range picks {
338+
counts[p]++
339+
}
340+
Expect(counts["rr-fat"]).To(BeNumerically(">=", 2), "fat node was picked %d times across 9 ties: %v", counts["rr-fat"], picks)
341+
Expect(counts["rr-mid"]).To(BeNumerically(">=", 2), "mid node was picked %d times across 9 ties: %v", counts["rr-mid"], picks)
342+
Expect(counts["rr-small"]).To(BeNumerically(">=", 2), "small node was picked %d times across 9 ties: %v", counts["rr-small"], picks)
343+
})
344+
307345
It("returns not-found when the model is loaded only on excluded nodes", func() {
308346
loadedExcluded := makeNode("excl-only-node", "10.0.0.45:50051", 8_000_000_000)
309347
emptyIncluded := makeNode("empty-included-node", "10.0.0.46:50051", 8_000_000_000)

0 commit comments

Comments
 (0)