diff --git a/pkg/clusterrouting/affinity.go b/pkg/clusterrouting/affinity.go new file mode 100644 index 000000000..4067002ca --- /dev/null +++ b/pkg/clusterrouting/affinity.go @@ -0,0 +1,33 @@ +package clusterrouting + +// PickWithAffinity prefers the candidate whose NodeID equals preferredNodeID +// when that candidate's in-flight count is within slack of the least-loaded +// candidate; otherwise it falls back to PickBestReplica (least in-flight, then +// oldest last-used, then most free VRAM). This keeps a warm prefix-cache peer +// sticky without letting it become a hot-spot: once it is more than slack +// requests busier than the least-loaded peer, load wins. With an empty +// preferredNodeID, or a preferred node not in the set, it is exactly +// PickBestReplica. slack mirrors prefixcache's BalanceAbsThreshold. +func PickWithAffinity(candidates []ReplicaCandidate, preferredNodeID string, slack int) *ReplicaCandidate { + if len(candidates) == 0 { + return nil + } + if preferredNodeID == "" { + return PickBestReplica(candidates) + } + var preferred *ReplicaCandidate + minInFlight := candidates[0].InFlight + for i := range candidates { + c := &candidates[i] + if c.InFlight < minInFlight { + minInFlight = c.InFlight + } + if c.NodeID == preferredNodeID { + preferred = c + } + } + if preferred != nil && preferred.InFlight <= minInFlight+slack { + return preferred + } + return PickBestReplica(candidates) +} diff --git a/pkg/clusterrouting/affinity_test.go b/pkg/clusterrouting/affinity_test.go new file mode 100644 index 000000000..50ed7f3a5 --- /dev/null +++ b/pkg/clusterrouting/affinity_test.go @@ -0,0 +1,48 @@ +package clusterrouting + +import ( + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("PickWithAffinity", func() { + ref := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + + It("returns nil for an empty candidate list", func() { + Expect(PickWithAffinity(nil, "x", 2)).To(BeNil()) + }) + + It("falls back to PickBestReplica when no preferred node is given", func() { + cs := []ReplicaCandidate{ + {NodeID: "busy", InFlight: 3, LastUsed: ref, AvailableVRAM: 80}, + {NodeID: "idle", InFlight: 0, LastUsed: ref, AvailableVRAM: 8}, + } + Expect(PickWithAffinity(cs, "", 2).NodeID).To(Equal("idle")) + }) + + It("honors the preferred node when it is within the in-flight slack of the least-loaded", func() { + cs := []ReplicaCandidate{ + {NodeID: "cold", InFlight: 0, LastUsed: ref, AvailableVRAM: 80}, + {NodeID: "warm", InFlight: 2, LastUsed: ref, AvailableVRAM: 8}, + } + Expect(PickWithAffinity(cs, "warm", 2).NodeID).To(Equal("warm")) + }) + + It("ignores the preferred node when it is beyond the slack and falls back to load+VRAM", func() { + cs := []ReplicaCandidate{ + {NodeID: "cold", InFlight: 0, LastUsed: ref, AvailableVRAM: 80}, + {NodeID: "warm", InFlight: 5, LastUsed: ref, AvailableVRAM: 8}, + } + Expect(PickWithAffinity(cs, "warm", 2).NodeID).To(Equal("cold")) + }) + + It("falls back to load+VRAM when the preferred node is not among the candidates", func() { + cs := []ReplicaCandidate{ + {NodeID: "a", InFlight: 1, LastUsed: ref, AvailableVRAM: 8}, + {NodeID: "b", InFlight: 1, LastUsed: ref, AvailableVRAM: 24}, + } + Expect(PickWithAffinity(cs, "ghost", 2).NodeID).To(Equal("b")) + }) +})