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
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto committed 2026-09-11 21:53:33 +00:00
1 parent f884d4b456
commit 783556bc93
6 files changed
+245

No files matched your search

+8
View File
@@ -346,6 +346,14 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCachePressure, 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
+20
View File
@@ -509,8 +509,28 @@ const (
SubjectPrefixCacheObserve = "prefixcache.observe"
SubjectPrefixCacheInvalidate = "prefixcache.invalidate"
SubjectPrefixCachePressure = "prefixcache.pressure"
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
@@ -12,6 +12,15 @@ var _ = Describe("PrefixCache subjects", func() {
Expect(messaging.SubjectPrefixCacheObserve).To(Equal("prefixcache.observe"))
Expect(messaging.SubjectPrefixCacheInvalidate).To(Equal("prefixcache.invalidate"))
Expect(messaging.SubjectPrefixCachePressure).To(Equal("prefixcache.pressure"))
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() {
+133
View File
@@ -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
}
@@ -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)}))
})
})
+22
View File
@@ -1280,6 +1280,28 @@ The scheduling algorithm supports **prefix-cache-aware** routing: bias each requ
When the load guard must route away from a warm replica, the frontend emits a forced-disturb event. These events and the reset sent after a successful pressure-triggered scale-up are broadcast over NATS, so the rolling autoscale threshold is cluster-wide rather than per frontend. The Prometheus counter `localai_prefix_cache_forced_disturb_total{model="..."}` records events at their originating frontend; sum it across frontend replicas to inspect cluster pressure without counting the NATS copies.
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.