mirror of
https://github.com/mudler/LocalAI.git
synced 2026-10-28 09:08:54 -04:00
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:
1 parent
f884d4b456
commit
783556bc93
6 files changed
+245
No files matched your search
@@ -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
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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)}))
|
||||
})
|
||||
})
|
||||
@@ -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.
|
||||
|
||||
Reference in new issue
Block a user