mirror of
https://github.com/mudler/LocalAI.git
synced 2026-08-04 12:22:22 -04:00
Compare commits
1 Commits
bot/issue-
...
bot/issue-
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ace2a29a9c |
@@ -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
|
||||
|
||||
@@ -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"`
|
||||
}
|
||||
|
||||
@@ -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"))
|
||||
})
|
||||
})
|
||||
|
||||
30
core/services/nodes/prefixcache/metrics.go
Normal file
30
core/services/nodes/prefixcache/metrics.go
Normal 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)))
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
})
|
||||
})
|
||||
|
||||
@@ -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)):
|
||||
|
||||
|
||||
@@ -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"),
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user