feat(routing): add load-guarded PickWithAffinity for prefix-cache routing

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

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)
+}

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"))
+	})
+})