diff --git a/core/application/distributed.go b/core/application/distributed.go index 32269a44a550..36783ff39f48 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -282,7 +282,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } idx := prefixcache.NewIndex(prefixCfg) prefixSync := prefixcache.NewSync(idx, natsClient) - pressure = prefixcache.NewPressure(prefixCfg.PressureWindow) + pressure = prefixcache.NewSyncedPressure(prefixCfg.PressureWindow, natsClient) prefixProvider = prefixSync // Invalidate the prefix-cache index whenever a replica row is removed. @@ -318,6 +318,11 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade }); err != nil { return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheInvalidate, err) } + if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCachePressure, func(ev messaging.PrefixCachePressureEvent) { + pressure.ApplyPressure(ev, time.Now()) + }); err != nil { + return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCachePressure, err) + } // Background eviction: sweep idle entries on the app context. Stopped // when the app context is cancelled (mirrors the reconciler loop which diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index 5f09adda5f50..f0ea5bd0b678 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -478,6 +478,7 @@ const subjectSyncStatePrefix = "state." const ( SubjectPrefixCacheObserve = "prefixcache.observe" SubjectPrefixCacheInvalidate = "prefixcache.invalidate" + SubjectPrefixCachePressure = "prefixcache.pressure" ) // PrefixCacheObserveEvent announces that the replica (NodeID, Replica) served a @@ -501,3 +502,12 @@ type PrefixCacheInvalidateEvent struct { NodeID string `json:"node_id"` Replica int `json:"replica"` } + +// PrefixCachePressureEvent announces one forced-disturb observed by a frontend. +// ID lets the publisher ignore its own NATS echo and all frontends ignore +// redelivery without inflating the autoscale signal. +type PrefixCachePressureEvent struct { + ID string `json:"id"` + Model string `json:"model"` + Reset bool `json:"reset,omitempty"` +} diff --git a/core/services/messaging/subjects_prefixcache_test.go b/core/services/messaging/subjects_prefixcache_test.go index 2b8eb5771c6c..851e882389e8 100644 --- a/core/services/messaging/subjects_prefixcache_test.go +++ b/core/services/messaging/subjects_prefixcache_test.go @@ -11,6 +11,7 @@ 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.SubjectPrefixCachePressure).To(Equal("prefixcache.pressure")) }) It("carries a replica index on the observe event", func() { @@ -24,4 +25,10 @@ var _ = Describe("PrefixCache subjects", func() { one := messaging.PrefixCacheInvalidateEvent{Model: "m", NodeID: "A", Replica: 0} Expect(one.Replica).To(Equal(0)) }) + + It("carries a unique pressure event identity", func() { + ev := messaging.PrefixCachePressureEvent{ID: "frontend-a:1", Model: "m"} + Expect(ev.ID).To(Equal("frontend-a:1")) + Expect(ev.Model).To(Equal("m")) + }) }) diff --git a/core/services/nodes/prefixcache/metrics.go b/core/services/nodes/prefixcache/metrics.go new file mode 100644 index 000000000000..1267825bd15c --- /dev/null +++ b/core/services/nodes/prefixcache/metrics.go @@ -0,0 +1,30 @@ +package prefixcache + +import ( + "context" + "sync" + + "go.opentelemetry.io/otel" + "go.opentelemetry.io/otel/attribute" + "go.opentelemetry.io/otel/metric" +) + +var ( + pressureMetricsOnce sync.Once + pressureCounter metric.Int64Counter +) + +func recordPressureMetric(model string) { + pressureMetricsOnce.Do(func() { + counter, err := otel.Meter("github.com/mudler/LocalAI").Int64Counter( + "localai_prefix_cache_forced_disturb_total", + metric.WithDescription("Prefix-cache forced-disturb events originated by this frontend"), + ) + if err == nil { + pressureCounter = counter + } + }) + if pressureCounter != nil { + pressureCounter.Add(context.Background(), 1, metric.WithAttributes(attribute.String("model", model))) + } +} diff --git a/core/services/nodes/prefixcache/pressure.go b/core/services/nodes/prefixcache/pressure.go index 29e3e0f15a33..bc78ff20712d 100644 --- a/core/services/nodes/prefixcache/pressure.go +++ b/core/services/nodes/prefixcache/pressure.go @@ -1,8 +1,15 @@ package prefixcache import ( + "crypto/rand" + "encoding/hex" + "fmt" "sync" + "sync/atomic" "time" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/xlog" ) // Pressure is a concurrency-safe rolling per-model counter of forced-disturb @@ -19,6 +26,10 @@ type Pressure struct { mu sync.Mutex window time.Duration events map[string][]time.Time + seen map[string]time.Time + pub publisher + origin string + seq atomic.Uint64 } // NewPressure creates a Pressure counter that remembers events for the given @@ -27,7 +38,22 @@ func NewPressure(window time.Duration) *Pressure { return &Pressure{ window: window, events: make(map[string][]time.Time), + seen: make(map[string]time.Time), + } +} + +// NewSyncedPressure creates a Pressure counter that broadcasts locally +// originated events so every frontend sees the same cluster-wide signal. +func NewSyncedPressure(window time.Duration, pub publisher) *Pressure { + p := NewPressure(window) + p.pub = pub + var id [8]byte + if _, err := rand.Read(id[:]); err != nil { + p.origin = fmt.Sprintf("%d", time.Now().UnixNano()) + } else { + p.origin = hex.EncodeToString(id[:]) } + return p } // pruneLocked drops entries older than cutoff, compacting in place. The cutoff @@ -47,13 +73,50 @@ func pruneLocked(ts []time.Time, cutoff time.Time) []time.Time { // older than the window, so the per-model slice stays bounded regardless of how // often Count runs. func (p *Pressure) Record(model string, now time.Time) { + id := fmt.Sprintf("%s:%d", p.origin, p.seq.Add(1)) + p.record(model, id, now) + recordPressureMetric(model) + if p.pub != nil { + ev := messaging.PrefixCachePressureEvent{ID: id, Model: model} + if err := p.pub.Publish(messaging.SubjectPrefixCachePressure, ev); err != nil { + xlog.Debug("prefixcache: pressure publish failed", "error", err) + } + } +} + +func (p *Pressure) record(model, id string, now time.Time) { p.mu.Lock() defer p.mu.Unlock() cutoff := now.Add(-p.window) + for eventID, recordedAt := range p.seen { + if recordedAt.Before(cutoff) { + delete(p.seen, eventID) + } + } + if id != "" { + if _, exists := p.seen[id]; exists { + return + } + p.seen[id] = now + } kept := append(pruneLocked(p.events[model], cutoff), now) p.events[model] = kept } +// ApplyPressure records a pressure event received from NATS without +// re-broadcasting it. Duplicate deliveries, including the publisher's own echo, +// are ignored by event ID. +func (p *Pressure) ApplyPressure(ev messaging.PrefixCachePressureEvent, now time.Time) { + if ev.ID == "" || ev.Model == "" { + return + } + if ev.Reset { + p.reset(ev.Model, ev.ID, now) + return + } + p.record(ev.Model, ev.ID, now) +} + // Count returns the number of records for the model within [now-window, now], // dropping any entries older than the window so the backing slice stays bounded. func (p *Pressure) Count(model string, now time.Time) int { @@ -76,7 +139,22 @@ func (p *Pressure) Count(model string, now time.Time) int { // (a pressure-triggered scale-up) so a single burst does not trigger repeated // scale-ups across consecutive ticks. func (p *Pressure) Reset(model string) { + id := fmt.Sprintf("%s:%d", p.origin, p.seq.Add(1)) + p.reset(model, id, time.Now()) + if p.pub != nil { + ev := messaging.PrefixCachePressureEvent{ID: id, Model: model, Reset: true} + if err := p.pub.Publish(messaging.SubjectPrefixCachePressure, ev); err != nil { + xlog.Debug("prefixcache: pressure reset publish failed", "error", err) + } + } +} + +func (p *Pressure) reset(model, id string, now time.Time) { p.mu.Lock() defer p.mu.Unlock() + if _, exists := p.seen[id]; exists { + return + } + p.seen[id] = now delete(p.events, model) } diff --git a/core/services/nodes/prefixcache/pressure_test.go b/core/services/nodes/prefixcache/pressure_test.go index e6b741cb06c3..df6a728dbcd6 100644 --- a/core/services/nodes/prefixcache/pressure_test.go +++ b/core/services/nodes/prefixcache/pressure_test.go @@ -3,6 +3,7 @@ package prefixcache_test import ( "time" + "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -95,4 +96,54 @@ var _ = Describe("Pressure counter", func() { Expect(p.LenForTest("m")).To(Equal(1)) Expect(p.Count("m", t0.Add(198*time.Minute))).To(Equal(1)) }) + + It("broadcasts locally recorded pressure", func() { + pub := &fakePub{} + p := prefixcache.NewSyncedPressure(time.Minute, pub) + + p.Record("m", t0) + + Expect(p.Count("m", t0)).To(Equal(1)) + Expect(pub.published).To(HaveLen(1)) + ev := pub.published[0].(messaging.PrefixCachePressureEvent) + Expect(ev.ID).ToNot(BeEmpty()) + Expect(ev.Model).To(Equal("m")) + }) + + It("aggregates a peer pressure event and ignores redelivery", func() { + p := prefixcache.NewSyncedPressure(time.Minute, &fakePub{}) + ev := messaging.PrefixCachePressureEvent{ID: "frontend-a:1", Model: "m"} + + p.ApplyPressure(ev, t0) + p.ApplyPressure(ev, t0.Add(time.Second)) + + Expect(p.Count("m", t0.Add(time.Second))).To(Equal(1)) + }) + + It("does not double-count its own broadcast echo", func() { + pub := &fakePub{} + p := prefixcache.NewSyncedPressure(time.Minute, pub) + p.Record("m", t0) + ev := pub.published[0].(messaging.PrefixCachePressureEvent) + + p.ApplyPressure(ev, t0.Add(time.Second)) + + Expect(p.Count("m", t0.Add(time.Second))).To(Equal(1)) + }) + + It("broadcasts and applies a pressure reset after scaling", func() { + pub := &fakePub{} + origin := prefixcache.NewSyncedPressure(time.Minute, pub) + peer := prefixcache.NewSyncedPressure(time.Minute, &fakePub{}) + origin.Record("m", t0) + record := pub.published[0].(messaging.PrefixCachePressureEvent) + peer.ApplyPressure(record, t0) + + origin.Reset("m") + reset := pub.published[1].(messaging.PrefixCachePressureEvent) + peer.ApplyPressure(reset, t0.Add(time.Second)) + + Expect(reset.Reset).To(BeTrue()) + Expect(peer.Count("m", t0.Add(time.Second))).To(Equal(0)) + }) }) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index f3a0f6534d1d..493242a328b5 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -946,7 +946,9 @@ Notes: ## Roadmap: Routing and Caching Enhancements -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. +The scheduling algorithm supports **prefix-cache-aware** routing: 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. A router-side radix tree maps prompt-prefix hashes 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 round-robin. + +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. 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)): diff --git a/pkg/natsauth/permissions_coverage_test.go b/pkg/natsauth/permissions_coverage_test.go index 05d2fbf0bbee..51a60a86fa42 100644 --- a/pkg/natsauth/permissions_coverage_test.go +++ b/pkg/natsauth/permissions_coverage_test.go @@ -115,6 +115,7 @@ var _ = Describe("Documented NATS service-user permissions", func() { frontendPublishes := []string{ messaging.SubjectPrefixCacheObserve, messaging.SubjectPrefixCacheInvalidate, + messaging.SubjectPrefixCachePressure, messaging.SubjectNodeBackendInstall("node-1"), messaging.SubjectGalleryProgress("op-1"), }