Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions core/application/distributed.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
20 changes: 20 additions & 0 deletions core/services/messaging/subjects.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 9 additions & 0 deletions core/services/messaging/subjects_prefixcache_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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() {
Expand Down
133 changes: 133 additions & 0 deletions core/services/nodes/prefixcache/reported.go
Original file line number Diff line number Diff line change
@@ -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
}
53 changes: 53 additions & 0 deletions core/services/nodes/prefixcache/reported_test.go
Original file line number Diff line number Diff line change
@@ -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 changes: 22 additions & 0 deletions docs/content/features/distributed-mode.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading