Compare commits

...

1 Commits

Author SHA1 Message Date
localai-org-maint-bot
ace2a29a9c fix(distributed): aggregate prefix cache pressure
Broadcast forced-disturb events and successful scale-up resets so every frontend shares the same rolling autoscale signal. Deduplicate NATS echoes, expose an origin-only Prometheus counter, and document cluster behavior.

Closes #10083

Assisted-by: Codex:gpt-5
2026-07-31 12:02:53 +00:00
8 changed files with 186 additions and 2 deletions

View File

@@ -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

View File

@@ -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"`
}

View File

@@ -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"))
})
})

View File

@@ -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)))
}
}

View File

@@ -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,9 +38,24 @@ 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
// boundary itself is inclusive so an event exactly window-old still counts.
// Callers must hold p.mu.
@@ -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)
}

View File

@@ -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))
})
})

View File

@@ -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)):

View File

@@ -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"),
}