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
7 changes: 6 additions & 1 deletion core/application/distributed.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down
10 changes: 10 additions & 0 deletions core/services/messaging/subjects.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"`
}
7 changes: 7 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,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() {
Expand All @@ -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"))
})
})
30 changes: 30 additions & 0 deletions core/services/nodes/prefixcache/metrics.go
Original file line number Diff line number Diff line change
@@ -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)))
}
}
78 changes: 78 additions & 0 deletions core/services/nodes/prefixcache/pressure.go
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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 {
Expand All @@ -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)
}
51 changes: 51 additions & 0 deletions core/services/nodes/prefixcache/pressure_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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))
})
})
4 changes: 3 additions & 1 deletion docs/content/features/distributed-mode.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)):

Expand Down
1 change: 1 addition & 0 deletions pkg/natsauth/permissions_coverage_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"),
}
Expand Down
Loading