From a5fbd28b22c4122ba88948804c4fa6474371ab63 Mon Sep 17 00:00:00 2001 From: localai-org-maint-bot <306269227+localai-org-maint-bot@users.noreply.github.com> Date: Tue, 11 Aug 2026 12:08:47 +0000 Subject: [PATCH] feat(prefixcache): index reported KV residency Add a NATS event contract and exact-residency provider for backend KV cache reports. Keep guessed request observations as the default routing source while maintaining the reported index for future producers. Assisted-by: Codex:gpt-5 --- core/application/distributed.go | 8 ++ core/services/messaging/subjects.go | 20 +++ .../messaging/subjects_prefixcache_test.go | 9 ++ core/services/nodes/prefixcache/reported.go | 133 ++++++++++++++++++ .../nodes/prefixcache/reported_test.go | 53 +++++++ docs/content/features/distributed-mode.md | 22 +++ 6 files changed, 245 insertions(+) create mode 100644 core/services/nodes/prefixcache/reported.go create mode 100644 core/services/nodes/prefixcache/reported_test.go diff --git a/core/application/distributed.go b/core/application/distributed.go index 32269a44a550..0ae16cdf4ca2 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -319,6 +319,14 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheInvalidate, err) } + // Keep an exact-residency index current so backend producers can report + // their real KV state without coupling to router internals. Routing stays + // on the guessed provider until a backend producer is available. + reportedIndex := prefixcache.NewReportedIndex() + if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCacheResidency, reportedIndex.Apply); err != nil { + return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheResidency, err) + } + // Background eviction: sweep idle entries on the app context. Stopped // when the app context is cancelled (mirrors the reconciler loop which // also runs on options.Context). TTL/2 keeps stale entries from diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index 5f09adda5f50..e0648322b62e 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -478,8 +478,28 @@ const subjectSyncStatePrefix = "state." const ( SubjectPrefixCacheObserve = "prefixcache.observe" SubjectPrefixCacheInvalidate = "prefixcache.invalidate" + SubjectPrefixCacheResidency = "prefixcache.residency" ) +// PrefixCacheOperation describes a backend-reported KV-cache residency change. +type PrefixCacheOperation string + +const ( + PrefixCacheStore PrefixCacheOperation = "store" + PrefixCacheRemove PrefixCacheOperation = "remove" + PrefixCacheClear PrefixCacheOperation = "clear" +) + +// PrefixCacheResidencyEvent reports exact backend KV-cache residency. Chain is +// the compatible shallow-to-deep prefix hash chain used by the router. +type PrefixCacheResidencyEvent struct { + Operation PrefixCacheOperation `json:"operation"` + Model string `json:"model"` + NodeID string `json:"node_id"` + Replica int `json:"replica"` + Chain []uint64 `json:"chain,omitempty"` +} + // PrefixCacheObserveEvent announces that the replica (NodeID, Replica) served a // request whose prefix chain ends at the given hashes for model. Chain is the // full shallow-to-deep hash chain so peers can insert the same path. Affinity is diff --git a/core/services/messaging/subjects_prefixcache_test.go b/core/services/messaging/subjects_prefixcache_test.go index 2b8eb5771c6c..a59dff613de2 100644 --- a/core/services/messaging/subjects_prefixcache_test.go +++ b/core/services/messaging/subjects_prefixcache_test.go @@ -11,6 +11,15 @@ var _ = Describe("PrefixCache subjects", func() { It("exposes stable subject constants", func() { Expect(messaging.SubjectPrefixCacheObserve).To(Equal("prefixcache.observe")) Expect(messaging.SubjectPrefixCacheInvalidate).To(Equal("prefixcache.invalidate")) + Expect(messaging.SubjectPrefixCacheResidency).To(Equal("prefixcache.residency")) + }) + + It("defines the reported residency event operations", func() { + Expect([]messaging.PrefixCacheOperation{ + messaging.PrefixCacheStore, + messaging.PrefixCacheRemove, + messaging.PrefixCacheClear, + }).To(Equal([]messaging.PrefixCacheOperation{"store", "remove", "clear"})) }) It("carries a replica index on the observe event", func() { diff --git a/core/services/nodes/prefixcache/reported.go b/core/services/nodes/prefixcache/reported.go new file mode 100644 index 000000000000..016041e06435 --- /dev/null +++ b/core/services/nodes/prefixcache/reported.go @@ -0,0 +1,133 @@ +package prefixcache + +import ( + "sort" + "sync" + "time" + + "github.com/mudler/LocalAI/core/services/messaging" +) + +// ReportedIndex is an exact-residency Provider populated only by backend +// events. Request routing observations intentionally do not mutate it. +type ReportedIndex struct { + mu sync.RWMutex + residencies map[string]map[ReplicaKey][][]uint64 +} + +func NewReportedIndex() *ReportedIndex { + return &ReportedIndex{residencies: map[string]map[ReplicaKey][][]uint64{}} +} + +func (ix *ReportedIndex) Apply(event messaging.PrefixCacheResidencyEvent) { + key := ReplicaKey{NodeID: event.NodeID, Replica: event.Replica} + if event.Model == "" || key.NodeID == "" { + return + } + + ix.mu.Lock() + defer ix.mu.Unlock() + byReplica := ix.residencies[event.Model] + switch event.Operation { + case messaging.PrefixCacheStore: + if len(event.Chain) == 0 { + return + } + if byReplica == nil { + byReplica = map[ReplicaKey][][]uint64{} + ix.residencies[event.Model] = byReplica + } + for _, chain := range byReplica[key] { + if equalChain(chain, event.Chain) { + return + } + } + byReplica[key] = append(byReplica[key], append([]uint64(nil), event.Chain...)) + case messaging.PrefixCacheRemove: + if byReplica == nil || len(event.Chain) == 0 { + return + } + chains := byReplica[key] + for i, chain := range chains { + if equalChain(chain, event.Chain) { + chains = append(chains[:i], chains[i+1:]...) + break + } + } + if len(chains) == 0 { + delete(byReplica, key) + } else { + byReplica[key] = chains + } + case messaging.PrefixCacheClear: + delete(byReplica, key) + } +} + +func (ix *ReportedIndex) Decide(model string, chain []uint64, candidates []ReplicaKey, _ time.Time) PrefixDecision { + order := append([]ReplicaKey(nil), candidates...) + sort.Slice(order, func(i, j int) bool { return order[i].less(order[j]) }) + d := PrefixDecision{ColdOrder: order} + if len(chain) == 0 { + return d + } + + ix.mu.RLock() + defer ix.mu.RUnlock() + byReplica := ix.residencies[model] + bestDepth := 0 + for _, key := range order { + for _, reported := range byReplica[key] { + depth := commonDepth(chain, reported) + if depth > bestDepth { + bestDepth = depth + d.Hot = key + d.HasHot = true + } + } + } + if d.HasHot { + d.MatchRatio = float64(bestDepth) / float64(len(chain)) + } + return d +} + +func (ix *ReportedIndex) Observe(string, []uint64, ReplicaKey, time.Time) bool { return false } + +func (ix *ReportedIndex) Invalidate(model string, key ReplicaKey) { + ix.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheClear, Model: model, NodeID: key.NodeID, Replica: key.Replica}) +} + +func (ix *ReportedIndex) InvalidateNode(model, nodeID string) { + ix.mu.Lock() + defer ix.mu.Unlock() + for key := range ix.residencies[model] { + if key.NodeID == nodeID { + delete(ix.residencies[model], key) + } + } +} + +func (ix *ReportedIndex) Evict(time.Time) {} + +func equalChain(a, b []uint64) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +func commonDepth(a, b []uint64) int { + depth := min(len(a), len(b)) + for i := range depth { + if a[i] != b[i] { + return i + } + } + return depth +} diff --git a/core/services/nodes/prefixcache/reported_test.go b/core/services/nodes/prefixcache/reported_test.go new file mode 100644 index 000000000000..b3a70abc4120 --- /dev/null +++ b/core/services/nodes/prefixcache/reported_test.go @@ -0,0 +1,53 @@ +package prefixcache_test + +import ( + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes/prefixcache" +) + +var _ prefixcache.Provider = (*prefixcache.ReportedIndex)(nil) + +var _ = Describe("ReportedIndex", func() { + var idx *prefixcache.ReportedIndex + + BeforeEach(func() { idx = prefixcache.NewReportedIndex() }) + + It("routes to the replica with the longest reported prefix", func() { + idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheStore, Model: "m", NodeID: "A", Replica: 0, Chain: []uint64{1, 2}}) + idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheStore, Model: "m", NodeID: "B", Replica: 0, Chain: []uint64{1, 2, 3, 4}}) + d := idx.Decide("m", []uint64{1, 2, 3, 9}, []prefixcache.ReplicaKey{rk("A", 0), rk("B", 0)}, t0) + Expect(d.HasHot).To(BeTrue()) + Expect(d.Hot).To(Equal(rk("B", 0))) + Expect(d.MatchRatio).To(Equal(0.75)) + }) + + It("removes only the announced residency", func() { + for _, chain := range [][]uint64{{1, 2}, {7, 8}} { + idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheStore, Model: "m", NodeID: "A", Replica: 0, Chain: chain}) + } + idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheRemove, Model: "m", NodeID: "A", Replica: 0, Chain: []uint64{1, 2}}) + Expect(idx.Decide("m", []uint64{1, 2}, []prefixcache.ReplicaKey{rk("A", 0)}, t0).HasHot).To(BeFalse()) + Expect(idx.Decide("m", []uint64{7, 8}, []prefixcache.ReplicaKey{rk("A", 0)}, t0).HasHot).To(BeTrue()) + }) + + It("clears every residency for one replica only", func() { + idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheStore, Model: "m", NodeID: "A", Replica: 0, Chain: []uint64{1, 2}}) + idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheStore, Model: "m", NodeID: "B", Replica: 0, Chain: []uint64{3, 4}}) + idx.Apply(messaging.PrefixCacheResidencyEvent{Operation: messaging.PrefixCacheClear, Model: "m", NodeID: "A", Replica: 0}) + candidates := []prefixcache.ReplicaKey{rk("A", 0), rk("B", 0)} + Expect(idx.Decide("m", []uint64{1, 2}, candidates, t0).HasHot).To(BeFalse()) + Expect(idx.Decide("m", []uint64{3, 4}, candidates, t0).Hot).To(Equal(rk("B", 0))) + }) + + It("ignores guessed request observations and keeps cold ordering deterministic", func() { + Expect(idx.Observe("m", []uint64{1, 2}, rk("A", 0), t0)).To(BeFalse()) + d := idx.Decide("m", []uint64{1, 2}, []prefixcache.ReplicaKey{rk("B", 1), rk("A", 1), rk("A", 0)}, time.Now()) + Expect(d.HasHot).To(BeFalse()) + Expect(d.ColdOrder).To(Equal([]prefixcache.ReplicaKey{rk("A", 0), rk("A", 1), rk("B", 1)})) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index f3a0f6534d1d..d238dea3abc6 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -948,6 +948,28 @@ Notes: The scheduling algorithm above is load-based (least in-flight, then least-recently-used). Work is underway to make routing **prefix-cache-aware**: bias each request toward the replica that already holds the relevant KV/prefix cache (multi-turn conversations and shared system prompts), so backends reuse cache instead of recomputing it. The first step is a router-side radix tree of prompt-prefix hashes mapped to nodes, with longest-prefix match, a load guard that preserves round-robin behavior under imbalance, and NATS sync across frontends. It is purely a routing-layer hint (no backend changes) and never routes worse than today's round-robin. +Backends can report exact KV-cache residency on the `prefixcache.residency` +NATS subject. The JSON event contract is: + +```json +{ + "operation": "store", + "model": "model-name", + "node_id": "worker-id", + "replica": 0, + "chain": [1203053429005847826, 15485907386658061715] +} +``` + +`operation` is `store`, `remove`, or `clear`. `store` adds the announced +shallow-to-deep chain for one model replica, `remove` removes only that exact +announced chain, and `clear` removes all reported residency for that model +replica (and may omit `chain`). Producers must generate the chain with exactly +the same windowing and hashing algorithm as the router; hashes from a different +chain algorithm are not compatible and will never match requests correctly. +Reported events populate the exact-residency provider, but the guessed provider +remains the routing default until a backend producer is available. + Further enhancements, surfaced from a survey of SGLang, vLLM production-stack, Ray Serve, llm-d, AIBrix, and NVIDIA Dynamo, are tracked under the routing roadmap epic ([#10063](https://github.com/mudler/LocalAI/issues/10063)): - **Reported/precise KV-event mode** ([#10064](https://github.com/mudler/LocalAI/issues/10064)): subscribe to actual backend KV-cache events for exact residency instead of inferring it from routing history.