mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 22:33:54 -04:00
feat(distributed): move the last nine fan-out families onto PostgreSQL
Gallery progress and cancel, the operation cache's start and end, the model
and backend cache invalidations, staging progress, and the prefix cache's
observations and invalidations now travel on the LISTEN/NOTIFY carrier. No
subject is published or subscribed on messaging.Client anywhere in the tree,
which is what makes retiring that package a deletion rather than a migration:
$ grep -rn 'natsClient\.Publish\|nats\.Publish\|\.Nats\.Publish\|QueueSubscribe\|SubscribeReply\|\.Request(' \
--include='*.go' core/ pkg/ | grep -v _test \
| grep -v 'c\.Request()\|ctx\.Request()\|Request()\.Context' \
| grep -v 'core/services/testutil/fakebus.go'
core/services/messaging/client.go:168,170,172,227,234,236,250,252,254,268,269,287
core/services/messaging/interfaces.go:21,22,23
Every remaining hit is inside core/services/messaging itself. The production
reads of the NATS client are now three, all of them the documented agent-worker
exception: Close on shutdown, the agent pool's publisher, and the agent-cancel
carrier passed to newFanoutBridges.
Prefix-cache observations publish like every other family rather than through a
method that refuses a message too large for a notification. The plan proposed
such a refusal on the reasoning that a long prompt makes a chain of thousands of
entries; ExtractChain caps a chain at Config.MaxDepth blocks, MaxDepth is a
constant with no operator knob, and the chain reaching Sync.Observe has one
source, the router's own extraction hook. A worst-case observation is a few
kilobytes against an 8000-byte cap, so the hot-path spill the refusal was
designed to avoid cannot occur, and shipping it would have added the programme's
only deliberate message drop to guard a condition that cannot arise. pgbus gains
FitsInline instead, a predicate that shares one size decision with Publish and
decides nothing, and core/application refuses at startup to wire a prefix cache
whose configured depth would put every observation over the cap.
The carrier choice is no longer stated at four sites. StagingTracker.SetPublisher
and SubscribeBroadcasts become one SetBroadcaster, so a tracker that publishes
where its peers are not listening cannot be spelled; prefixcache.Sync gains
SubscribeBroadcasts, which reads the carrier it publishes on; and the gallery
service and the operation cache are wired by methods on DistributedServices that
name no carrier at all, so the NATS client beside it cannot be handed over.
OpCache.SetMessagingClient and GalleryService.SetNATSClient are renamed to
SetBroadcaster so a missed call site fails to compile.
Two pre-existing defects that the two-real-carrier specs surfaced are fixed. A
progress tick published before a cancel and delivered after it cleared Cancelled
and left the operation reading as still running on that replica; mergeStatus now
drops a stale tick rather than merging it. GetStatus and GetAllStatus handed out
the stored OpStatus pointer while the broadcast subscribers mutated it in place,
so an /api/operations response could be marshalled mid-write; both now copy.
Assisted-by: Claude Opus 5 [claude-code]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
1 parent
5a95bb3c0a
commit
44cb495169
24 files changed
+1510
-179
No files matched your search
@@ -0,0 +1,165 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/galleryop"
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/nodes/prefixcache"
|
||||
"github.com/mudler/LocalAI/core/services/pgbus"
|
||||
)
|
||||
|
||||
// The four process-lifetime caches, each wired onto the broadcast carrier by
|
||||
// one function here, and none of them by a call site naming a carrier.
|
||||
//
|
||||
// Two are methods on DistributedServices and take NO carrier at all, because
|
||||
// their call sites hold the whole struct and could otherwise reach for ds.Nats.
|
||||
// The other two run inside initDistributed before that struct exists, so they
|
||||
// take the CONCRETE *pgbus.Bus rather than messaging.Broadcaster.
|
||||
//
|
||||
// Concrete on purpose. *messaging.Client satisfies messaging.Broadcaster just
|
||||
// as well as the carrier does, so an interface parameter at these sites lets a
|
||||
// caller hand over the NATS client that is also in scope: it compiles, it
|
||||
// starts, it publishes, it is delivered, onto a carrier the deployment is being
|
||||
// taken off, and nothing fails until NATS goes away. With the concrete type
|
||||
// that mistake is a build error rather than a deployment that looks healthy.
|
||||
//
|
||||
// The adopters themselves still take the interface, so their own specs drive
|
||||
// them with an in-memory double. The narrowing is only here, where the wrong
|
||||
// carrier is in scope.
|
||||
|
||||
// wireGalleryBroadcasts puts the gallery service's progress, cancel and
|
||||
// cache-invalidation traffic on the carrier, and opens the wildcard
|
||||
// subscriptions that mirror a peer's.
|
||||
//
|
||||
// Set and subscribe in one call because they are one decision: a service that
|
||||
// published where nobody listened would show every operation it started and
|
||||
// none of its peers', which is what /api/operations looks like on a replica
|
||||
// that has been load-balanced away from.
|
||||
//
|
||||
// The caller must have hydrated from the store and bound OnModelsChanged
|
||||
// first; both are stated on the methods themselves.
|
||||
func (ds *DistributedServices) wireGallery(gs *galleryop.GalleryService) error {
|
||||
if ds == nil || ds.Bus == nil {
|
||||
return fmt.Errorf("wiring gallery broadcasts: no broadcast carrier, so gallery progress and cancels would reach no peer replica")
|
||||
}
|
||||
if gs == nil {
|
||||
return nil
|
||||
}
|
||||
gs.SetBroadcaster(ds.Bus)
|
||||
return gs.SubscribeBroadcasts()
|
||||
}
|
||||
|
||||
// WireOpCache puts the admin operation cache on the carrier and starts it,
|
||||
// which hydrates from PostgreSQL and subscribes.
|
||||
//
|
||||
// Exported, unlike its siblings, because the OpCache is built in the HTTP layer
|
||||
// rather than in initDistributed. It takes neither a carrier nor a store: both
|
||||
// come off this struct, so the HTTP layer cannot pass the NATS client that
|
||||
// hangs off it beside them.
|
||||
//
|
||||
// A hydrate failure is the OpCache's own business and is logged there; a
|
||||
// subscribe failure is returned, because a cache that hydrated and did not
|
||||
// subscribe reports the operations it found at boot and never learns of another.
|
||||
func (ds *DistributedServices) WireOpCache(ctx context.Context, cache *galleryop.OpCache) error {
|
||||
if cache == nil {
|
||||
return nil
|
||||
}
|
||||
if ds == nil || ds.Bus == nil {
|
||||
return fmt.Errorf("wiring the operation cache: no broadcast carrier, so /api/operations would answer with whatever this one replica admitted")
|
||||
}
|
||||
cache.SetBroadcaster(ds.Bus)
|
||||
if ds.DistStores != nil && ds.DistStores.Gallery != nil {
|
||||
cache.SetGalleryStore(ds.DistStores.Gallery)
|
||||
}
|
||||
return cache.Start(ctx)
|
||||
}
|
||||
|
||||
// wireStagingBroadcasts puts file-staging progress on the carrier in both
|
||||
// directions. The tracker's own SetBroadcaster is what makes those one call;
|
||||
// see the comment there.
|
||||
func wireStagingBroadcasts(bus *pgbus.Bus, tracker *nodes.StagingTracker) (messaging.Subscription, error) {
|
||||
if bus == nil {
|
||||
return nil, fmt.Errorf("wiring staging broadcasts: no broadcast carrier, so a staging transfer would show a progress bar only on the replica performing it")
|
||||
}
|
||||
if tracker == nil {
|
||||
return nil, nil
|
||||
}
|
||||
return tracker.SetBroadcaster(bus)
|
||||
}
|
||||
|
||||
// wirePrefixCacheBroadcasts builds the cross-frontend prefix-cache layer on the
|
||||
// carrier and subscribes it to peers, after refusing a configuration whose
|
||||
// observations could not travel in a notification.
|
||||
func wirePrefixCacheBroadcasts(bus *pgbus.Bus, cfg prefixcache.Config, idx prefixcache.Provider) (*prefixcache.Sync, error) {
|
||||
// The configuration first, and the carrier second. A depth this carrier
|
||||
// cannot hold is wrong whether or not a carrier was supplied, and naming
|
||||
// the more specific fault is what makes the startup message actionable.
|
||||
if err := requirePrefixCacheFitsInline(cfg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if bus == nil {
|
||||
return nil, fmt.Errorf("wiring the prefix cache: no broadcast carrier, so each frontend would route on nothing but its own history")
|
||||
}
|
||||
sync := prefixcache.NewSync(idx, bus)
|
||||
if _, err := sync.SubscribeBroadcasts(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return sync, nil
|
||||
}
|
||||
|
||||
// prefixCacheIdentifierAllowance is how many bytes of model id plus node id a
|
||||
// prefix-cache observation is budgeted for when its worst case is checked
|
||||
// against the notification cap.
|
||||
//
|
||||
// Both are operator-chosen strings with no enforced length, so no bound here is
|
||||
// a proof. It does not need to be one: an observation that does not fit is
|
||||
// SPILLED like any other broadcast, at the cost of a row and a SELECT, and is
|
||||
// never lost. What the check exists to catch is the other failure, the one that
|
||||
// has no symptom: a change to Config.MaxDepth that quietly puts every
|
||||
// observation over the cap and turns the inference path into a table write per
|
||||
// request. A generous allowance catches that and does not fire on a long model
|
||||
// name.
|
||||
const prefixCacheIdentifierAllowance = 512
|
||||
|
||||
// requirePrefixCacheFitsInline refuses a prefix-cache configuration whose
|
||||
// observations would spill.
|
||||
//
|
||||
// prefixcache.ExtractChain caps a chain at Config.MaxDepth blocks, so an
|
||||
// observation's size has a worst case that is known before the deployment
|
||||
// serves a request: MaxDepth hashes at their widest decimal encoding, plus the
|
||||
// identifiers. That bound is the reason Sync.Observe can publish like every
|
||||
// other family instead of being given a way to refuse.
|
||||
//
|
||||
// It is a startup error and not a warning because the alternative reading is
|
||||
// the one this programme exists to remove: a deployment that came up, spills a
|
||||
// row and reads it back on every replica for every request whose prefix
|
||||
// changed, and looks exactly like one that is merely slow.
|
||||
func requirePrefixCacheFitsInline(cfg prefixcache.Config) error {
|
||||
// The widest a uint64 encodes to in JSON, so the check does not depend on
|
||||
// which hashes a workload happens to produce.
|
||||
chain := make([]uint64, cfg.MaxDepth)
|
||||
for i := range chain {
|
||||
chain[i] = math.MaxUint64
|
||||
}
|
||||
worst := messaging.PrefixCacheObserveEvent{
|
||||
Model: strings.Repeat("m", prefixCacheIdentifierAllowance/2),
|
||||
Chain: chain,
|
||||
NodeID: strings.Repeat("n", prefixCacheIdentifierAllowance/2),
|
||||
Replica: math.MaxInt32,
|
||||
}
|
||||
fits, err := pgbus.FitsInline(messaging.SubjectPrefixCacheObserve, worst)
|
||||
if err != nil {
|
||||
return fmt.Errorf("sizing a prefix-cache observation: %w", err)
|
||||
}
|
||||
if !fits {
|
||||
return fmt.Errorf("the prefix-cache depth is too large to broadcast: an observation for %d blocks does not fit in one notification, so every request whose prefix changed would write a row and every replica would read it back on the inference path", cfg.MaxDepth)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,289 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/galleryop"
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/nodes/prefixcache"
|
||||
"github.com/mudler/LocalAI/core/services/pgbus"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// The four process-lifetime caches, each asserted from the OTHER replica's
|
||||
// carrier.
|
||||
//
|
||||
// Every case here wires the cache on busA and drives it from busB. A cache
|
||||
// talking to itself would pass with the wiring pointed at any carrier at all,
|
||||
// which is the defect these exist to catch: the NATS client is in scope at
|
||||
// three of the four call sites and satisfies the same interface, so a site left
|
||||
// holding it publishes successfully and is delivered, to nobody the deployment
|
||||
// will still be listening on.
|
||||
var _ = Describe("wiring the process-lifetime caches onto the broadcast carrier", func() {
|
||||
var (
|
||||
ctx context.Context
|
||||
db *gorm.DB
|
||||
busA, busB *pgbus.Bus
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
ctx = context.Background()
|
||||
|
||||
var dsn string
|
||||
db, dsn = testutil.SetupTestDBWithDSN()
|
||||
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
|
||||
|
||||
newBus := func() *pgbus.Bus {
|
||||
b, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(b.Close)
|
||||
return b
|
||||
}
|
||||
busA, busB = newBus(), newBus()
|
||||
})
|
||||
|
||||
// Each of the four refuses rather than coming up on nothing. A cache wired
|
||||
// to no carrier has no symptom of its own: it answers from whatever this
|
||||
// one replica happened to do, forever, and looks exactly like a fleet with
|
||||
// nothing going on elsewhere.
|
||||
DescribeTable("refuses to wire a cache with no carrier",
|
||||
func(wire func() error, want string) {
|
||||
err := wire()
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring(want))
|
||||
},
|
||||
Entry("gallery", func() error {
|
||||
return (&DistributedServices{}).wireGallery(galleryop.NewGalleryService(&config.ApplicationConfig{}, nil))
|
||||
}, "gallery progress and cancels"),
|
||||
Entry("operation cache", func() error {
|
||||
svc := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
return (&DistributedServices{}).WireOpCache(context.Background(), galleryop.NewOpCache(svc))
|
||||
}, "/api/operations"),
|
||||
Entry("staging", func() error {
|
||||
_, err := wireStagingBroadcasts(nil, nodes.NewStagingTracker())
|
||||
return err
|
||||
}, "progress bar only on the replica performing it"),
|
||||
Entry("prefix cache", func() error {
|
||||
_, err := wirePrefixCacheBroadcasts(nil, prefixcache.DefaultConfig(), prefixcache.NewIndex(prefixcache.DefaultConfig()))
|
||||
return err
|
||||
}, "its own history"),
|
||||
)
|
||||
|
||||
// S2. The gallery service applies a peer's progress, which it can only do
|
||||
// if the wildcard subscription wireGallery opened is on the carrier the
|
||||
// peer published to.
|
||||
It("subscribes the gallery service to progress a peer replica broadcasts", func() {
|
||||
svc := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
Expect((&DistributedServices{Bus: busA}).wireGallery(svc)).To(Succeed())
|
||||
DeferCleanup(svc.CloseBroadcasts)
|
||||
|
||||
Expect(busB.Publish(messaging.SubjectGalleryProgress("op-1"), galleryop.GalleryProgressEvent{
|
||||
JobID: "op-1",
|
||||
Status: &galleryop.OpStatus{Progress: 42, Message: "halfway"},
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() *galleryop.OpStatus { return svc.GetStatus("op-1") }, 20*time.Second).ShouldNot(BeNil())
|
||||
Expect(svc.GetStatus("op-1").Progress).To(Equal(42.0))
|
||||
})
|
||||
|
||||
// S2, the other direction. A service that only subscribed would pass the
|
||||
// row above and publish its own progress where no peer reads it.
|
||||
It("publishes the gallery service's progress onto the same carrier", func() {
|
||||
svc := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
Expect((&DistributedServices{Bus: busA}).wireGallery(svc)).To(Succeed())
|
||||
DeferCleanup(svc.CloseBroadcasts)
|
||||
|
||||
out := make(chan []byte, 4)
|
||||
_, err := busB.Subscribe(messaging.SubjectGalleryProgressWildcard, func(b []byte) { out <- b })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
svc.UpdateStatus("op-2", &galleryop.OpStatus{Progress: 7})
|
||||
|
||||
Eventually(out, 20*time.Second).Should(Receive())
|
||||
})
|
||||
|
||||
// S1. The OpCache is wired from the HTTP layer, and WireOpCache is what
|
||||
// keeps that call site from naming a carrier of its own.
|
||||
It("subscribes the operation cache to a peer replica's admissions", func() {
|
||||
svc := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
cache := galleryop.NewOpCache(svc)
|
||||
Expect((&DistributedServices{Bus: busA}).WireOpCache(ctx, cache)).To(Succeed())
|
||||
DeferCleanup(cache.Close)
|
||||
|
||||
Expect(busB.Publish(messaging.SubjectGalleryOpStart, galleryop.OpCacheEvent{
|
||||
JobID: "job-9", CacheKey: "official@vllm", IsBackend: true,
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() bool { return cache.Exists("official@vllm") }, 20*time.Second).Should(BeTrue())
|
||||
Expect(cache.IsBackendOp("official@vllm")).To(BeTrue())
|
||||
})
|
||||
|
||||
It("publishes the operation cache's admissions onto the same carrier", func() {
|
||||
svc := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
cache := galleryop.NewOpCache(svc)
|
||||
Expect((&DistributedServices{Bus: busA}).WireOpCache(ctx, cache)).To(Succeed())
|
||||
DeferCleanup(cache.Close)
|
||||
|
||||
out := make(chan []byte, 4)
|
||||
_, err := busB.Subscribe(messaging.SubjectGalleryOpStart, func(b []byte) { out <- b })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
cache.Set("llama-3-8b", "job-10")
|
||||
|
||||
Eventually(out, 20*time.Second).Should(Receive())
|
||||
})
|
||||
|
||||
// S3, in both directions and as two separate specs. It used to be two
|
||||
// calls, a publisher and a subscriber, and a tracker with one of them on
|
||||
// each carrier shows a staging progress bar on the originating replica and
|
||||
// nowhere else. SetBroadcaster is one method now, so that deployment cannot
|
||||
// be spelled, but each half still has to be held on its own: a mutation
|
||||
// that drops the subscribe leaves the publishing spec green and the reverse
|
||||
// leaves the mirroring spec green.
|
||||
It("mirrors a peer replica's staging progress into the tracker", func() {
|
||||
tracker := nodes.NewStagingTracker()
|
||||
sub, err := wireStagingBroadcasts(busA, tracker)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(sub).ToNot(BeNil())
|
||||
|
||||
Expect(busB.Publish(messaging.SubjectStagingProgress("model-x"), nodes.StagingProgressEvent{
|
||||
ModelID: "model-x",
|
||||
Status: &nodes.StagingStatus{ModelID: "model-x", NodeName: "worker-7"},
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() map[string]nodes.StagingStatus { return tracker.GetAll() }, 20*time.Second).
|
||||
Should(HaveKey("model-x"))
|
||||
})
|
||||
|
||||
It("publishes the tracker's own staging progress onto the same carrier", func() {
|
||||
tracker := nodes.NewStagingTracker()
|
||||
_, err := wireStagingBroadcasts(busA, tracker)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
out := make(chan []byte, 4)
|
||||
_, err = busB.Subscribe(messaging.SubjectStagingProgressWildcard, func(b []byte) { out <- b })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
tracker.Start("model-y", "worker-8", 1)
|
||||
|
||||
Eventually(out, 20*time.Second).Should(Receive())
|
||||
})
|
||||
|
||||
// S4, in both directions. The prefix cache is the family on the inference
|
||||
// path, and a Sync wired to a carrier its peers do not read leaves every
|
||||
// frontend routing on nothing but its own history while every publish
|
||||
// succeeds.
|
||||
It("applies a peer replica's observation into the prefix index", func() {
|
||||
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
sync, err := wirePrefixCacheBroadcasts(busA, prefixcache.DefaultConfig(), idx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
chain := []uint64{101, 202, 303}
|
||||
Expect(busB.Publish(messaging.SubjectPrefixCacheObserve, messaging.PrefixCacheObserveEvent{
|
||||
Model: "m", Chain: chain, NodeID: "A", Replica: 1,
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() bool {
|
||||
return sync.Decide("m", chain, []prefixcache.ReplicaKey{{NodeID: "A", Replica: 1}}, time.Now()).HasHot
|
||||
}, 20*time.Second).Should(BeTrue())
|
||||
})
|
||||
|
||||
It("applies a peer replica's invalidation, so a removed replica stops being routed to", func() {
|
||||
// The invalidation half separately: a missed one leaves this frontend
|
||||
// routing to a replica that is gone until the TTL, which is the reading
|
||||
// of a missed message this programme forbids.
|
||||
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
sync, err := wirePrefixCacheBroadcasts(busA, prefixcache.DefaultConfig(), idx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
chain := []uint64{404, 505}
|
||||
key := prefixcache.ReplicaKey{NodeID: "A", Replica: 0}
|
||||
sync.ApplyObserve(messaging.PrefixCacheObserveEvent{Model: "m", Chain: chain, NodeID: "A"}, time.Now())
|
||||
Expect(sync.Decide("m", chain, []prefixcache.ReplicaKey{key}, time.Now()).HasHot).To(BeTrue())
|
||||
|
||||
Expect(busB.Publish(messaging.SubjectPrefixCacheInvalidate, messaging.PrefixCacheInvalidateEvent{
|
||||
Model: "m", NodeID: "A", Replica: 0,
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() bool {
|
||||
return sync.Decide("m", chain, []prefixcache.ReplicaKey{key}, time.Now()).HasHot
|
||||
}, 20*time.Second).Should(BeFalse())
|
||||
})
|
||||
|
||||
It("publishes this replica's observations onto the same carrier", func() {
|
||||
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
sync, err := wirePrefixCacheBroadcasts(busA, prefixcache.DefaultConfig(), idx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
out := make(chan []byte, 4)
|
||||
_, err = busB.Subscribe(messaging.SubjectPrefixCacheObserve, func(b []byte) { out <- b })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
sync.Observe("m", []uint64{909}, prefixcache.ReplicaKey{NodeID: "B", Replica: 0}, time.Now())
|
||||
|
||||
Eventually(out, 20*time.Second).Should(Receive())
|
||||
})
|
||||
})
|
||||
|
||||
// The size bound that lets prefix-cache observations publish like every other
|
||||
// family instead of being given a way to refuse.
|
||||
//
|
||||
// The plan for this phase proposed a PublishNoSpill that would REFUSE an
|
||||
// observation too large for a notification, on the reasoning that a long prompt
|
||||
// makes a chain of thousands of entries. ExtractChain does not produce one, and
|
||||
// the refusal would have been the only deliberate message drop in the
|
||||
// programme, guarding a condition that cannot arise, with a counter nothing
|
||||
// alerts on as its only symptom. This is what took its place: the same
|
||||
// knowledge, asked at startup, where being wrong is a deployment that refuses
|
||||
// to come up and says why rather than one that runs with no cross-replica
|
||||
// affinity and looks healthy.
|
||||
//
|
||||
// It needs no bus. FitsInline is a pure function over the same encoder and the
|
||||
// same constant Publish measures against.
|
||||
var _ = Describe("the prefix-cache observation bound", func() {
|
||||
It("accepts the depth the extractor actually produces", func() {
|
||||
Expect(requirePrefixCacheFitsInline(prefixcache.DefaultConfig())).To(Succeed())
|
||||
})
|
||||
|
||||
It("refuses a depth whose observations would spill on every request", func() {
|
||||
// An absolute depth, not one derived from the carrier's cap, so this
|
||||
// row states a fact about this family rather than restating the
|
||||
// constant it is measured against.
|
||||
cfg := prefixcache.DefaultConfig()
|
||||
cfg.MaxDepth = 100000
|
||||
|
||||
err := requirePrefixCacheFitsInline(cfg)
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("write a row"))
|
||||
})
|
||||
|
||||
It("is checked before the prefix cache is wired at all", func() {
|
||||
// The check is worth nothing if the wiring runs anyway. Asserted
|
||||
// through the same function initDistributed calls, and on the MESSAGE
|
||||
// rather than on failure alone: this call has two things wrong with it,
|
||||
// and a spec that accepted any error would pass on the carrier
|
||||
// complaint with the bound check deleted.
|
||||
cfg := prefixcache.DefaultConfig()
|
||||
cfg.MaxDepth = 100000
|
||||
|
||||
sync, err := wirePrefixCacheBroadcasts(nil, cfg, prefixcache.NewIndex(cfg))
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("write a row"))
|
||||
Expect(sync).To(BeNil())
|
||||
})
|
||||
})
|
||||
@@ -495,8 +495,8 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
// with --distributed-prefix-cache=false, which leaves prefixProvider and
|
||||
// pressure nil so the SmartRouter and reconciler behave exactly as the
|
||||
// round-robin floor (true no-op). When enabled we build the local index,
|
||||
// wrap it in a NATS-backed Sync (publishes our observations, applies peers'
|
||||
// via the subscriptions below), install the extraction hook used by
|
||||
// wrap it in a Sync on the broadcast carrier (which both publishes our
|
||||
// observations and applies peers'), install the extraction hook used by
|
||||
// core/backend/llm.go, and run a background eviction ticker on the app ctx.
|
||||
var prefixProvider prefixcache.Provider
|
||||
var pressure *prefixcache.Pressure
|
||||
@@ -510,7 +510,13 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
return nil, fmt.Errorf("invalid prefix-cache configuration: %w", err)
|
||||
}
|
||||
idx := prefixcache.NewIndex(prefixCfg)
|
||||
prefixSync := prefixcache.NewSync(idx, natsClient)
|
||||
// S4. One call puts this replica's observations and its peers' on the
|
||||
// same carrier, and it takes the concrete carrier so the NATS client
|
||||
// still in scope here cannot be handed to it by accident.
|
||||
prefixSync, err := wirePrefixCacheBroadcasts(bus, prefixCfg, idx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pressure = prefixcache.NewPressure(prefixCfg.PressureWindow)
|
||||
prefixProvider = prefixSync
|
||||
|
||||
@@ -534,20 +540,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
return prefixcache.ExtractChain(model, prompt, prefixCfg)
|
||||
}
|
||||
|
||||
// Apply peers' observations/invalidations to the same Sync. ApplyObserve
|
||||
// and ApplyInvalidate update only the local index and do not re-publish,
|
||||
// so there is no broadcast loop.
|
||||
if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCacheObserve, func(ev messaging.PrefixCacheObserveEvent) {
|
||||
prefixSync.ApplyObserve(ev, time.Now())
|
||||
}); err != nil {
|
||||
return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheObserve, err)
|
||||
}
|
||||
if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCacheInvalidate, func(ev messaging.PrefixCacheInvalidateEvent) {
|
||||
prefixSync.ApplyInvalidate(ev)
|
||||
}); err != nil {
|
||||
return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheInvalidate, 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
|
||||
@@ -624,7 +616,8 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
// Wire staging-progress broadcasting so file-staging shows up on every
|
||||
// replica, not just the one performing the transfer. Without this, a
|
||||
// /api/operations poll that round-robins onto a peer sees no staging row and
|
||||
// the progress flickers. The origin publishes; peers mirror via the wildcard.
|
||||
// the progress flickers. The origin publishes; peers mirror via the
|
||||
// wildcard, on the same carrier.
|
||||
// A silently disabled safety check is how the original incident stayed
|
||||
// invisible for sixteen minutes. Say so once, loudly, at startup.
|
||||
if cfg.Distributed.DiskHeadroomDisabled {
|
||||
@@ -632,8 +625,10 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
"knob", config.FlagDiskHeadroomCheck, "env", "LOCALAI_DISTRIBUTED_DISK_HEADROOM_CHECK")
|
||||
}
|
||||
|
||||
router.StagingTracker().SetPublisher(natsClient)
|
||||
if _, err := router.StagingTracker().SubscribeBroadcasts(natsClient); err != nil {
|
||||
// S3, and it is ONE call rather than a publisher and a subscriber: see
|
||||
// StagingTracker.SetBroadcaster for why a tracker that could name two
|
||||
// carriers is a progress bar that only the originating replica shows.
|
||||
if _, err := wireStagingBroadcasts(bus, router.StagingTracker()); err != nil {
|
||||
xlog.Warn("Failed to subscribe to staging progress broadcasts", "error", err)
|
||||
}
|
||||
|
||||
|
||||
@@ -334,9 +334,10 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
// Wire skill store into AgentPoolService (wired at pool start time via closure)
|
||||
// The actual wiring happens in StartAgentPool since the pool doesn't exist yet.
|
||||
|
||||
// Wire NATS and gallery store into GalleryService for cross-instance progress/cancel
|
||||
// Wire the broadcast carrier and gallery store into GalleryService for
|
||||
// cross-instance progress/cancel. The carrier is wired below, next to
|
||||
// the subscriptions it feeds, so the two cannot name different carriers.
|
||||
if application.galleryService != nil {
|
||||
application.galleryService.SetNATSClient(distSvc.Nats)
|
||||
if distSvc.DistStores != nil && distSvc.DistStores.Gallery != nil {
|
||||
// Clean up stale in-progress operations from previous crashed instances
|
||||
if _, err := distSvc.DistStores.Gallery.CleanStale(30 * time.Minute); err != nil {
|
||||
@@ -396,7 +397,10 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
xlog.Warn("Failed to apply peer model config change", "error", err)
|
||||
}
|
||||
}
|
||||
if err := application.galleryService.SubscribeBroadcasts(); err != nil {
|
||||
// S2. One call sets the carrier and opens the wildcard
|
||||
// subscriptions, and it names no carrier, so the NATS client on
|
||||
// distSvc cannot be passed here by accident.
|
||||
if err := distSvc.wireGallery(application.galleryService); err != nil {
|
||||
xlog.Warn("Gallery service subscribe failed", "error", err)
|
||||
}
|
||||
// Wire distributed model/backend managers so delete propagates to workers
|
||||
|
||||
+7
-7
@@ -502,17 +502,17 @@ func API(application *application.Application) (*echo.Echo, error) {
|
||||
var opcache *galleryop.OpCache
|
||||
if !application.ApplicationConfig().DisableWebUI {
|
||||
opcache = galleryop.NewOpCache(application.GalleryService())
|
||||
// In distributed mode, wire the NATS client + gallery store so this
|
||||
// replica's OpCache stays in sync with peers — without this the
|
||||
// In distributed mode, wire the broadcast carrier + gallery store so
|
||||
// this replica's OpCache stays in sync with peers. Without this the
|
||||
// /api/operations endpoint returns whatever this single replica
|
||||
// happened to admit, and a load-balanced UI poll alternates between
|
||||
// "operation visible" and "operation gone" between replicas.
|
||||
//
|
||||
// S1. The carrier choice lives in core/application with the other three
|
||||
// caches, and this call names no carrier at all, so the NATS client
|
||||
// hanging off the same struct cannot be handed over here by accident.
|
||||
if d := application.Distributed(); d != nil {
|
||||
opcache.SetMessagingClient(d.Nats)
|
||||
if d.DistStores != nil && d.DistStores.Gallery != nil {
|
||||
opcache.SetGalleryStore(d.DistStores.Gallery)
|
||||
}
|
||||
if err := opcache.Start(application.ApplicationConfig().Context); err != nil {
|
||||
if err := d.WireOpCache(application.ApplicationConfig().Context, opcache); err != nil {
|
||||
xlog.Warn("OpCache distributed subscribe failed; running standalone", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -188,7 +188,7 @@ backend: llama-cpp
|
||||
lifecycle := &endpointLifecycleRecorder{}
|
||||
galleryService := galleryop.NewGalleryService(appConfig, nil)
|
||||
client := &endpointRecordingClient{}
|
||||
galleryService.SetNATSClient(client)
|
||||
galleryService.SetBroadcaster(client)
|
||||
endpointApp := echo.New()
|
||||
endpointApp.PATCH("/api/models/config-json/:name", PatchConfigEndpoint(configLoader, galleryService, appConfig, lifecycle))
|
||||
|
||||
|
||||
@@ -126,7 +126,7 @@ var _ = Describe("Edit Model test", func() {
|
||||
Expect(loader.LoadModelConfigsFromPath(tempDir)).To(Succeed())
|
||||
galleryService := galleryop.NewGalleryService(applicationConfig, nil)
|
||||
client := &endpointRecordingClient{}
|
||||
galleryService.SetNATSClient(client)
|
||||
galleryService.SetBroadcaster(client)
|
||||
|
||||
app := echo.New()
|
||||
app.POST("/models/edit/:name", EditModelEndpoint(loader, galleryService, applicationConfig))
|
||||
@@ -152,7 +152,7 @@ var _ = Describe("Edit Model test", func() {
|
||||
Expect(loader.LoadModelConfigsFromPath(tempDir)).To(Succeed())
|
||||
galleryService := galleryop.NewGalleryService(applicationConfig, nil)
|
||||
client := &endpointRecordingClient{}
|
||||
galleryService.SetNATSClient(client)
|
||||
galleryService.SetBroadcaster(client)
|
||||
|
||||
app := echo.New()
|
||||
app.POST("/models/edit/:name", EditModelEndpoint(loader, galleryService, applicationConfig))
|
||||
@@ -276,7 +276,7 @@ var _ = Describe("Edit Model test", func() {
|
||||
Expect(peerLoader.LoadModelConfigsFromPath(tempDir)).To(Succeed())
|
||||
galleryService := galleryop.NewGalleryService(applicationConfig, nil)
|
||||
client := &endpointRecordingClient{}
|
||||
galleryService.SetNATSClient(client)
|
||||
galleryService.SetBroadcaster(client)
|
||||
lifecycle := &endpointLifecycleRecorder{pendingCleanup: 2}
|
||||
app := echo.New()
|
||||
app.POST("/models/edit/:name", EditModelEndpoint(loader, galleryService, applicationConfig, lifecycle))
|
||||
|
||||
@@ -180,7 +180,7 @@ var _ = Describe("artifact progress coalescer", func() {
|
||||
progressClient := &recordingProgressClient{}
|
||||
service := NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
service.modelManager = &modelOperationProgressManager{err: installErr}
|
||||
service.natsClient = progressClient
|
||||
service.broadcaster = progressClient
|
||||
op := &ManagementOp[gallery.GalleryModel, gallery.ModelConfig]{
|
||||
ID: "model-operation",
|
||||
GalleryElementName: "model",
|
||||
@@ -205,7 +205,7 @@ var _ = Describe("artifact progress coalescer", func() {
|
||||
progressClient := &recordingProgressClient{}
|
||||
service := NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
service.modelManager = &legacyModelOperationProgressManager{err: installErr}
|
||||
service.natsClient = progressClient
|
||||
service.broadcaster = progressClient
|
||||
op := &ManagementOp[gallery.GalleryModel, gallery.ModelConfig]{
|
||||
ID: "legacy-model-operation",
|
||||
GalleryElementName: "model",
|
||||
|
||||
@@ -0,0 +1,237 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package galleryop_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/galleryop"
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/pgbus"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// The gallery families on the real carrier, between two services on two LISTEN
|
||||
// connections.
|
||||
//
|
||||
// The in-memory double in distributed_sync_test.go delivers synchronously and
|
||||
// cannot spill, so it proves the merge logic and nothing about the carrier.
|
||||
// Gallery progress is the family that crosses the 8000-byte notification cap in
|
||||
// ordinary operation, at a few tens of workers, so the path a real fleet takes
|
||||
// every tick is the spill path, and it is only reachable here.
|
||||
var _ = Describe("the gallery families on the broadcast carrier", func() {
|
||||
var (
|
||||
db *gorm.DB
|
||||
dsn string
|
||||
svcA, svcB *galleryop.GalleryService
|
||||
// peerBus stands in for a replica that has no GalleryService in this
|
||||
// process. It is how the backends family is driven: its publisher is
|
||||
// reached only from deep inside backendHandler, which needs a whole
|
||||
// gallery operation, while the shared publishCacheInvalidate it calls
|
||||
// is already pinned end to end by the models row below.
|
||||
peerBus *pgbus.Bus
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
ctx := context.Background()
|
||||
db, dsn = testutil.SetupTestDBWithDSN()
|
||||
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
|
||||
|
||||
newService := func() *galleryop.GalleryService {
|
||||
bus, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(bus.Close)
|
||||
svc := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
svc.SetBroadcaster(bus)
|
||||
Expect(svc.SubscribeBroadcasts()).To(Succeed())
|
||||
DeferCleanup(svc.CloseBroadcasts)
|
||||
return svc
|
||||
}
|
||||
svcA, svcB = newService(), newService()
|
||||
|
||||
var err error
|
||||
peerBus, err = pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(peerBus.Close)
|
||||
})
|
||||
|
||||
spilledRows := func(subject string) int64 {
|
||||
GinkgoHelper()
|
||||
var n int64
|
||||
Expect(db.Model(&pgbus.BusMessage{}).Where("subject = ?", subject).Count(&n).Error).To(Succeed())
|
||||
return n
|
||||
}
|
||||
|
||||
It("carries a 60-node progress event to a peer through the spill table, byte for byte", func() {
|
||||
// Sixty nodes is an ordinary fleet and roughly twice what the cap
|
||||
// holds, so this is the path a real deployment's progress ticks take
|
||||
// and not an edge case. Asserted as one row AND as identical content:
|
||||
// a spill that lost an entry would still deliver something.
|
||||
nodes := make([]galleryop.NodeProgress, 0, 60)
|
||||
for i := range 60 {
|
||||
nodes = append(nodes, galleryop.NodeProgress{
|
||||
NodeID: fmt.Sprintf("node-%02d", i),
|
||||
NodeName: fmt.Sprintf("worker-%02d.fleet.internal", i),
|
||||
Status: galleryop.NodeStatusDownloading,
|
||||
FileName: "backend-image.tar",
|
||||
Current: "512 MiB",
|
||||
Total: "1.5 GiB",
|
||||
Percentage: float64(i),
|
||||
Phase: "downloading",
|
||||
})
|
||||
}
|
||||
status := &galleryop.OpStatus{
|
||||
Progress: 33.0,
|
||||
Message: "installing on the fleet",
|
||||
GalleryElementName: "official@vllm",
|
||||
Nodes: nodes,
|
||||
}
|
||||
|
||||
svcA.UpdateStatus("op-fleet", status)
|
||||
|
||||
Eventually(func() *galleryop.OpStatus { return svcB.GetStatus("op-fleet") }, 20*time.Second).
|
||||
ShouldNot(BeNil())
|
||||
got := svcB.GetStatus("op-fleet")
|
||||
Expect(got.Nodes).To(HaveLen(60))
|
||||
Expect(got.Nodes).To(Equal(nodes), "a peer must see the same per-node breakdown, not a truncated one")
|
||||
Expect(got.GalleryElementName).To(Equal("official@vllm"))
|
||||
|
||||
Expect(spilledRows(messaging.SubjectGalleryProgress("op-fleet"))).To(Equal(int64(1)),
|
||||
"a progress event this size does not fit in a notification and must travel as a row")
|
||||
})
|
||||
|
||||
It("carries a cancel to the peer holding the operation", func() {
|
||||
// A cancel is a request and not a verdict, but it has to arrive: the
|
||||
// replica that admitted the job is rarely the one the admin's click
|
||||
// lands on.
|
||||
svcB.UpdateStatus("op-cancel", &galleryop.OpStatus{Progress: 10, Cancellable: true})
|
||||
|
||||
Expect(svcA.CancelOperation("op-cancel")).To(Succeed())
|
||||
|
||||
Eventually(func() bool {
|
||||
st := svcB.GetStatus("op-cancel")
|
||||
return st != nil && st.Cancelled
|
||||
}, 20*time.Second).Should(BeTrue())
|
||||
})
|
||||
|
||||
It("does not let a progress tick delivered after a cancel undo it", func() {
|
||||
// The carrier puts no order on two subjects, and the owning replica's
|
||||
// last tick is published BEFORE the admin's cancel and can arrive after
|
||||
// it, on that replica's own echo as readily as on a peer. Driven here
|
||||
// as an explicit late tick rather than by racing the two, so it states
|
||||
// the rule rather than reproducing a window.
|
||||
svcB.UpdateStatus("op-late", &galleryop.OpStatus{Progress: 10, Cancellable: true})
|
||||
Expect(svcA.CancelOperation("op-late")).To(Succeed())
|
||||
Eventually(func() bool {
|
||||
st := svcB.GetStatus("op-late")
|
||||
return st != nil && st.Cancelled
|
||||
}, 20*time.Second).Should(BeTrue())
|
||||
|
||||
Expect(peerBus.Publish(messaging.SubjectGalleryProgress("op-late"), galleryop.GalleryProgressEvent{
|
||||
JobID: "op-late",
|
||||
Status: &galleryop.OpStatus{Progress: 60, Message: "downloading"},
|
||||
})).To(Succeed())
|
||||
|
||||
Consistently(func() bool {
|
||||
st := svcB.GetStatus("op-late")
|
||||
return st != nil && st.Cancelled
|
||||
}, 3*time.Second).Should(BeTrue(),
|
||||
"a cancelled operation must not read as running again because a tick arrived late")
|
||||
})
|
||||
|
||||
// A missed invalidation must never read as a cache that is valid, which is
|
||||
// why this family is on Publish and spills rather than on anything that can
|
||||
// refuse. These pin that the peer's refresh hook actually fires.
|
||||
DescribeTable("carries a cache invalidation to the peer's refresh hook",
|
||||
func(broadcast func(*galleryop.GalleryService), bind func(*galleryop.GalleryService, chan<- string)) {
|
||||
fired := make(chan string, 4)
|
||||
bind(svcB, fired)
|
||||
|
||||
broadcast(svcA)
|
||||
|
||||
var got string
|
||||
Eventually(fired, 20*time.Second).Should(Receive(&got))
|
||||
Expect(got).ToNot(BeEmpty())
|
||||
},
|
||||
Entry("models",
|
||||
func(s *galleryop.GalleryService) { s.BroadcastModelsChanged("llama-3-8b", "install") },
|
||||
func(s *galleryop.GalleryService, fired chan<- string) {
|
||||
s.OnModelsChanged = func(evt messaging.CacheInvalidateEvent) { fired <- evt.Element }
|
||||
},
|
||||
),
|
||||
Entry("backends",
|
||||
func(*galleryop.GalleryService) {
|
||||
Expect(peerBus.Publish(messaging.SubjectCacheInvalidateBackends,
|
||||
messaging.CacheInvalidateEvent{Element: "official@vllm", Op: "upgrade"})).To(Succeed())
|
||||
},
|
||||
func(s *galleryop.GalleryService, fired chan<- string) {
|
||||
s.OnBackendOpCompleted = func() { fired <- "backends" }
|
||||
},
|
||||
),
|
||||
)
|
||||
})
|
||||
|
||||
// The OpCache's two subjects, likewise on the real carrier. The echo-loop guard
|
||||
// stays on the in-memory double, which is the only carrier that can count
|
||||
// publishes.
|
||||
var _ = Describe("the OpCache on the broadcast carrier", func() {
|
||||
It("propagates an admission and a dismissal to a peer replica", func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
ctx := context.Background()
|
||||
db, dsn := testutil.SetupTestDBWithDSN()
|
||||
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
|
||||
|
||||
newCache := func() *galleryop.OpCache {
|
||||
bus, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(bus.Close)
|
||||
cache := galleryop.NewOpCache(galleryop.NewGalleryService(&config.ApplicationConfig{}, nil))
|
||||
cache.SetBroadcaster(bus)
|
||||
Expect(cache.Start(ctx)).To(Succeed())
|
||||
DeferCleanup(cache.Close)
|
||||
return cache
|
||||
}
|
||||
cacheA, cacheB := newCache(), newCache()
|
||||
|
||||
cacheA.SetBackend("official@vllm", "job-1")
|
||||
|
||||
Eventually(func() bool { return cacheB.Exists("official@vllm") }, 20*time.Second).Should(BeTrue())
|
||||
Expect(cacheB.IsBackendOp("official@vllm")).To(BeTrue(),
|
||||
"the peer must learn this is a backend install, not a model install")
|
||||
|
||||
cacheA.DeleteUUID("job-1")
|
||||
|
||||
Eventually(func() bool { return cacheB.Exists("official@vllm") }, 20*time.Second).Should(BeFalse(),
|
||||
"a dismissed operation must clear from peer replicas too")
|
||||
})
|
||||
|
||||
It("does not re-broadcast an event it applied", func() {
|
||||
// The echo-loop guard, re-run on the new carrier's shape. applyStart
|
||||
// writes the local maps and must not publish: two replicas that
|
||||
// answered each other's broadcasts would never stop.
|
||||
bus := testutil.NewFakeBus()
|
||||
cache := galleryop.NewOpCache(galleryop.NewGalleryService(&config.ApplicationConfig{}, nil))
|
||||
cache.SetBroadcaster(bus)
|
||||
Expect(cache.Start(context.Background())).To(Succeed())
|
||||
DeferCleanup(cache.Close)
|
||||
|
||||
cache.Set("llama-3-8b", "job-2")
|
||||
|
||||
Expect(cache.Exists("llama-3-8b")).To(BeTrue())
|
||||
Expect(bus.PublishCount(messaging.SubjectGalleryOpStart)).To(Equal(1),
|
||||
"the replica's own broadcast came back to it and must not have produced a second")
|
||||
})
|
||||
})
|
||||
@@ -152,8 +152,8 @@ var _ = Describe("OpCache distributed sync", func() {
|
||||
svcB = galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
opA = galleryop.NewOpCache(svcA)
|
||||
opB = galleryop.NewOpCache(svcB)
|
||||
opA.SetMessagingClient(bus)
|
||||
opB.SetMessagingClient(bus)
|
||||
opA.SetBroadcaster(bus)
|
||||
opB.SetBroadcaster(bus)
|
||||
Expect(opA.Start(context.Background())).To(Succeed())
|
||||
Expect(opB.Start(context.Background())).To(Succeed())
|
||||
})
|
||||
@@ -241,8 +241,8 @@ var _ = Describe("GalleryService broadcast sync", func() {
|
||||
bus = newFakeBus()
|
||||
svcA = galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
svcB = galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
svcA.SetNATSClient(bus)
|
||||
svcB.SetNATSClient(bus)
|
||||
svcA.SetBroadcaster(bus)
|
||||
svcB.SetBroadcaster(bus)
|
||||
Expect(svcA.SubscribeBroadcasts()).To(Succeed())
|
||||
Expect(svcB.SubscribeBroadcasts()).To(Succeed())
|
||||
})
|
||||
@@ -329,8 +329,8 @@ var _ = Describe("GalleryService cache invalidation broadcasts", func() {
|
||||
bus = newFakeBus()
|
||||
svcA = galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
svcB = galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
svcA.SetNATSClient(bus)
|
||||
svcB.SetNATSClient(bus)
|
||||
svcA.SetBroadcaster(bus)
|
||||
svcB.SetBroadcaster(bus)
|
||||
})
|
||||
|
||||
AfterEach(func() {
|
||||
@@ -417,7 +417,7 @@ var _ = Describe("GalleryService cache invalidation broadcasts", func() {
|
||||
|
||||
It("BroadcastModelsChanged is a no-op when NATS is not wired (standalone)", func() {
|
||||
standalone := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
// No SetNATSClient: must not panic and must simply do nothing.
|
||||
// No SetBroadcaster: must not panic and must simply do nothing.
|
||||
Expect(func() { standalone.BroadcastModelsChanged("x", "delete") }).ToNot(Panic())
|
||||
})
|
||||
})
|
||||
|
||||
@@ -118,7 +118,7 @@ func newOpHistory(limit int) *opHistory {
|
||||
|
||||
// add appends rec unless its job ID was already recorded. Returns false when
|
||||
// the record was a duplicate. The originating replica both evicts locally and
|
||||
// receives its own NATS end broadcast, so without this every distributed
|
||||
// receives its own end broadcast, so without this every distributed
|
||||
// operation would be recorded twice.
|
||||
func (h *opHistory) add(rec OpRecord) bool {
|
||||
h.mu.Lock()
|
||||
|
||||
@@ -204,7 +204,7 @@ var _ = Describe("model deletion revision lifecycle", func() {
|
||||
|
||||
service := NewGalleryService(appConfig, nil)
|
||||
bus := &countingMessagingClient{}
|
||||
service.SetNATSClient(bus)
|
||||
service.SetBroadcaster(bus)
|
||||
service.SetModelManager(manager)
|
||||
service.SetModelRevisionLifecycle(lifecycle)
|
||||
op := &ManagementOp[gallery.GalleryModel, gallery.ModelConfig]{
|
||||
|
||||
@@ -91,7 +91,7 @@ type OpStatus struct {
|
||||
}
|
||||
|
||||
// opStatusWire is the JSON shape used when an OpStatus crosses a process
|
||||
// boundary (NATS broadcast). The Error field on OpStatus is an `error`
|
||||
// boundary (a cross-replica broadcast). The Error field on OpStatus is an `error`
|
||||
// interface, which json.Marshal flattens to `{}` because the concrete error
|
||||
// type usually has no exported fields — so a failed install replicated to a
|
||||
// peer frontend would arrive with a nil error and the UI would never surface
|
||||
@@ -165,7 +165,7 @@ func (o *OpStatus) UnmarshalJSON(data []byte) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// OpCacheEvent is the NATS payload broadcast by frontend replicas when an
|
||||
// OpCacheEvent is the payload broadcast by frontend replicas when an
|
||||
// admin operation is admitted (SubjectGalleryOpStart) or dismissed
|
||||
// (SubjectGalleryOpEnd). Peers merge these into their local OpCache so a
|
||||
// load-balanced /api/operations poll never returns an empty list while a
|
||||
@@ -176,15 +176,15 @@ type OpCacheEvent struct {
|
||||
IsBackend bool `json:"is_backend"`
|
||||
}
|
||||
|
||||
// GalleryProgressEvent is the NATS payload for an OpStatus broadcast. It
|
||||
// GalleryProgressEvent is the payload for an OpStatus broadcast. It
|
||||
// wraps OpStatus with the opID/JobID so subscribers reading the wildcard
|
||||
// subject don't need to parse it back out of the NATS subject string.
|
||||
// subject don't need to parse it back out of the subject string.
|
||||
type GalleryProgressEvent struct {
|
||||
JobID string `json:"job_id"`
|
||||
Status *OpStatus `json:"status"`
|
||||
}
|
||||
|
||||
// GalleryCancelEvent is the NATS payload for a gallery cancellation. The
|
||||
// GalleryCancelEvent is the payload for a gallery cancellation. The
|
||||
// local cancellation func may live on a different frontend replica than the
|
||||
// one that received the UI cancel button click; the broadcast subscriber
|
||||
// runs the cancel func on whichever replica registered it.
|
||||
@@ -236,10 +236,14 @@ type OpCache struct {
|
||||
started *xsync.SyncedMap[string, time.Time]
|
||||
|
||||
// Distributed sync (nil when standalone).
|
||||
mu sync.RWMutex
|
||||
nats messaging.MessagingClient
|
||||
store *distributed.GalleryStore
|
||||
subs []messaging.Subscription
|
||||
//
|
||||
// ONE carrier field, read by both the start/end publishes and by the
|
||||
// subscriptions Start opens, so this cache cannot publish where its peers
|
||||
// are not listening.
|
||||
mu sync.RWMutex
|
||||
broadcaster messaging.Broadcaster
|
||||
store *distributed.GalleryStore
|
||||
subs []messaging.Subscription
|
||||
}
|
||||
|
||||
func NewOpCache(galleryService *GalleryService) *OpCache {
|
||||
@@ -252,13 +256,17 @@ func NewOpCache(galleryService *GalleryService) *OpCache {
|
||||
}
|
||||
}
|
||||
|
||||
// SetMessagingClient enables cross-replica OpCache sync. Once set, Set/
|
||||
// SetBroadcaster enables cross-replica OpCache sync on the deployment's
|
||||
// fan-out carrier, which is PostgreSQL LISTEN/NOTIFY. Once set, Set/
|
||||
// SetBackend/DeleteUUID publish OpCacheEvent messages that peer OpCaches
|
||||
// merge into their local maps. Call Start after this to subscribe.
|
||||
func (m *OpCache) SetMessagingClient(nc messaging.MessagingClient) {
|
||||
//
|
||||
// messaging.Broadcaster and not the wider client: this cache needs Publish and
|
||||
// Subscribe and nothing else.
|
||||
func (m *OpCache) SetBroadcaster(b messaging.Broadcaster) {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
m.nats = nc
|
||||
m.broadcaster = b
|
||||
}
|
||||
|
||||
// SetGalleryStore enables PostgreSQL-backed OpCache persistence.
|
||||
@@ -273,17 +281,17 @@ func (m *OpCache) SetGalleryStore(s *distributed.GalleryStore) {
|
||||
}
|
||||
|
||||
// Start hydrates the in-memory maps from PostgreSQL (if a store was wired)
|
||||
// and subscribes to the broadcast subjects (if NATS was wired). It returns
|
||||
// and subscribes to the broadcast subjects (if a carrier was wired). It returns
|
||||
// the first subscribe error; hydration errors are logged but non-fatal so
|
||||
// the frontend still comes up.
|
||||
//
|
||||
// Safe to call exactly once after SetMessagingClient / SetGalleryStore. The
|
||||
// Safe to call exactly once after SetBroadcaster / SetGalleryStore. The
|
||||
// ctx parameter is reserved for future cancellation — current subscriptions
|
||||
// live for the lifetime of the OpCache and are released by Close.
|
||||
func (m *OpCache) Start(_ context.Context) error {
|
||||
m.mu.RLock()
|
||||
store := m.store
|
||||
nc := m.nats
|
||||
nc := m.broadcaster
|
||||
m.mu.RUnlock()
|
||||
|
||||
if store != nil {
|
||||
@@ -318,7 +326,7 @@ func (m *OpCache) Start(_ context.Context) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close drops all NATS subscriptions. Safe to call multiple times.
|
||||
// Close drops all broadcast subscriptions. Safe to call multiple times.
|
||||
func (m *OpCache) Close() {
|
||||
m.mu.Lock()
|
||||
subs := m.subs
|
||||
@@ -421,7 +429,7 @@ func (m *OpCache) dropReplacedStamp(key, jobID string) {
|
||||
func (m *OpCache) persistAndBroadcastStart(key, value string, isBackend bool) {
|
||||
m.mu.RLock()
|
||||
store := m.store
|
||||
nc := m.nats
|
||||
nc := m.broadcaster
|
||||
m.mu.RUnlock()
|
||||
|
||||
if store != nil {
|
||||
@@ -467,7 +475,7 @@ const (
|
||||
//
|
||||
// Safe to call for an unknown job ID (no key, no record) and safe to call
|
||||
// twice for the same job (the ring dedupes), which is what makes it usable
|
||||
// from both the local delete path and the NATS end event.
|
||||
// from both the local delete path and the broadcast end event.
|
||||
func (m *OpCache) recordTerminal(jobID string, src terminalSource) {
|
||||
if jobID == "" {
|
||||
return
|
||||
@@ -628,7 +636,7 @@ func (m *OpCache) DeleteUUID(uuid string) {
|
||||
return
|
||||
}
|
||||
m.mu.RLock()
|
||||
nc := m.nats
|
||||
nc := m.broadcaster
|
||||
m.mu.RUnlock()
|
||||
if nc != nil {
|
||||
if err := nc.Publish(messaging.SubjectGalleryOpEnd, OpCacheEvent{JobID: uuid}); err != nil {
|
||||
|
||||
@@ -31,10 +31,14 @@ type GalleryService struct {
|
||||
cancellations map[string]cancellationActions
|
||||
|
||||
// Distributed mode (nil when not in distributed mode).
|
||||
// natsClient is the wider MessagingClient (Publisher + subscribe methods)
|
||||
// when wired by the distributed startup path; broadcastSubs holds the
|
||||
// progress + cancel subscriptions opened by SubscribeBroadcasts.
|
||||
natsClient messaging.MessagingClient
|
||||
//
|
||||
// broadcaster is the deployment's fan-out carrier, and it is ONE field on
|
||||
// purpose: the progress and cancel publishes and the wildcard
|
||||
// subscriptions SubscribeBroadcasts opens all read it, so this service
|
||||
// cannot end up publishing on one carrier and listening on another. A
|
||||
// replica in that state shows every gallery operation it started and none
|
||||
// of its peers', with no error anywhere.
|
||||
broadcaster messaging.Broadcaster
|
||||
galleryStore *distributed.GalleryStore
|
||||
broadcastSubs []messaging.Subscription
|
||||
|
||||
@@ -117,14 +121,18 @@ func (g *GalleryService) ModelArtifactMaterializer() config.ArtifactMaterializer
|
||||
return g.appConfig.ModelArtifactMaterializer
|
||||
}
|
||||
|
||||
// SetNATSClient sets the NATS client for distributed progress publishing.
|
||||
// Accepting the wider MessagingClient (vs. plain Publisher) lets
|
||||
// SubscribeBroadcasts wire the wildcard subscriptions that keep peer
|
||||
// replicas' statuses + cancellations in sync.
|
||||
func (g *GalleryService) SetNATSClient(nc messaging.MessagingClient) {
|
||||
// SetBroadcaster wires the deployment's fan-out carrier, which is PostgreSQL
|
||||
// LISTEN/NOTIFY. Both halves of this service's cross-replica sync ride it: the
|
||||
// progress and cancel publishes, and the wildcard subscriptions
|
||||
// SubscribeBroadcasts opens.
|
||||
//
|
||||
// It takes messaging.Broadcaster, which is Publish plus Subscribe and nothing
|
||||
// else. Neither half needs request/reply or a queue group, and a wider
|
||||
// parameter here would let this service acquire one without anyone deciding to.
|
||||
func (g *GalleryService) SetBroadcaster(b messaging.Broadcaster) {
|
||||
g.Lock()
|
||||
defer g.Unlock()
|
||||
g.natsClient = nc
|
||||
g.broadcaster = b
|
||||
}
|
||||
|
||||
// SetGalleryStore sets the PostgreSQL gallery store for distributed persistence.
|
||||
@@ -178,10 +186,10 @@ func (g *GalleryService) UpdateStatus(s string, op *OpStatus) {
|
||||
}
|
||||
g.statuses[s] = op
|
||||
store := g.galleryStore
|
||||
nc := g.natsClient
|
||||
nc := g.broadcaster
|
||||
g.Unlock()
|
||||
|
||||
// I/O happens after Unlock. The NATS broadcast loops back into our own
|
||||
// I/O happens after Unlock. The broadcast loops back into our own
|
||||
// wildcard subscriber (mergeStatus), which would deadlock on this mutex
|
||||
// if we still held it. Holding the lock across a PostgreSQL round-trip
|
||||
// would also stall every concurrent reader on each progress tick.
|
||||
@@ -208,9 +216,14 @@ func (g *GalleryService) UpdateStatus(s string, op *OpStatus) {
|
||||
}
|
||||
}
|
||||
|
||||
// Publish progress to NATS in distributed mode. The payload wraps the
|
||||
// OpStatus with the opID so peer replicas reading the wildcard subject
|
||||
// don't need to parse it back out of the NATS subject string.
|
||||
// Broadcast progress in distributed mode. The payload wraps the OpStatus
|
||||
// with the opID so peer replicas reading the wildcard subject don't need
|
||||
// to parse it back out of the subject string.
|
||||
//
|
||||
// A progress event carries one entry per node, so on a fleet of a few tens
|
||||
// of workers it outgrows the notification cap and travels as a spilled row
|
||||
// instead. That is the carrier's ordinary path and not an error: the peer
|
||||
// receives the same bytes either way.
|
||||
if nc != nil {
|
||||
if err := nc.Publish(messaging.SubjectGalleryProgress(s), GalleryProgressEvent{
|
||||
JobID: s,
|
||||
@@ -222,11 +235,17 @@ func (g *GalleryService) UpdateStatus(s string, op *OpStatus) {
|
||||
}
|
||||
|
||||
// publishCacheInvalidate broadcasts a cache invalidation event so peer
|
||||
// replicas refresh whatever in-memory state mirrors disk. No-op when
|
||||
// natsClient is not wired (standalone mode).
|
||||
// replicas refresh whatever in-memory state mirrors disk. No-op when no
|
||||
// carrier is wired (standalone mode).
|
||||
//
|
||||
// An invalidation is not a hint and is never traded away for a cheaper
|
||||
// publish: a peer that misses one keeps serving from a cache it believes is
|
||||
// valid, which is the one reading of a missed message this programme forbids.
|
||||
// It goes out through Publish, which spills a message too large for a
|
||||
// notification rather than refusing it.
|
||||
func (g *GalleryService) publishCacheInvalidate(subject string, evt messaging.CacheInvalidateEvent) {
|
||||
g.Lock()
|
||||
nc := g.natsClient
|
||||
nc := g.broadcaster
|
||||
g.Unlock()
|
||||
if nc == nil {
|
||||
return
|
||||
@@ -242,7 +261,7 @@ func (g *GalleryService) publishCacheInvalidate(subject string, evt messaging.Ca
|
||||
// /models/toggle-state endpoints, which write the YAML and reload only the
|
||||
// local in-memory loader). Peers receive it via OnModelsChanged and refresh
|
||||
// their own ModelConfigLoader so a request load-balanced to any replica sees
|
||||
// the same config. No-op in standalone mode (no NATS client).
|
||||
// the same config. No-op in standalone mode (no carrier).
|
||||
//
|
||||
// op is "install" for a create/edit (the element must be (re)loaded from
|
||||
// disk) or "delete" for a removal (the element must be pruned from memory,
|
||||
@@ -262,7 +281,7 @@ func (g *GalleryService) BroadcastModelsChangedRevision(element, op, configRevis
|
||||
}
|
||||
|
||||
// mergeStatus is the broadcast-side merge: it updates the in-memory map from
|
||||
// a peer's GalleryProgressEvent without re-publishing to NATS or re-writing
|
||||
// a peer's GalleryProgressEvent without re-publishing to the carrier or re-writing
|
||||
// to PostgreSQL. UpdateStatus is the local-write entry point and does both;
|
||||
// mergeStatus is what the wildcard subscriber calls. Splitting them avoids
|
||||
// an echo loop (replica publishes → its own subscriber receives → mergeStatus
|
||||
@@ -273,8 +292,22 @@ func (g *GalleryService) mergeStatus(opID string, op *OpStatus) {
|
||||
}
|
||||
g.Lock()
|
||||
defer g.Unlock()
|
||||
prev := g.statuses[opID]
|
||||
// A cancellation is terminal and a progress tick is not, and the carrier
|
||||
// puts no order on the two. The owning replica's last tick is published
|
||||
// before the admin's cancel and can be DELIVERED after it, on the owner's
|
||||
// own echo as readily as on a peer; merging it wholesale would clear
|
||||
// Cancelled and leave the operation reading as still running on that
|
||||
// replica while every other one shows it stopped. A missed or late message
|
||||
// must never read as an operation that was not cancelled, so a stale tick
|
||||
// is dropped rather than merged. A terminal status that carries the
|
||||
// cancellation still merges, which is how the final "cancelled" message
|
||||
// arrives.
|
||||
if prev != nil && prev.Cancelled && !op.Cancelled {
|
||||
return
|
||||
}
|
||||
if len(op.Nodes) == 0 {
|
||||
if prev := g.statuses[opID]; prev != nil && len(prev.Nodes) > 0 {
|
||||
if prev != nil && len(prev.Nodes) > 0 {
|
||||
op.Nodes = prev.Nodes
|
||||
}
|
||||
}
|
||||
@@ -322,18 +355,48 @@ func (g *GalleryService) UpdateNodeProgress(opID, nodeID string, np NodeProgress
|
||||
}
|
||||
}
|
||||
|
||||
// GetStatus returns a COPY of the operation's status, not the stored pointer.
|
||||
//
|
||||
// The copy is what makes the lock mean anything. Every caller of this and of
|
||||
// GetAllStatus only reads what it gets back, but the broadcast subscribers
|
||||
// mutate the stored OpStatus IN PLACE - applyCancel sets Cancelled on the
|
||||
// struct a peer's event names, mergeStatus rewrites the fields a peer sent - so
|
||||
// handing out the pointer let an /api/operations response be marshalled while a
|
||||
// peer's cancel was being written into it. Returning the pointer under a mutex
|
||||
// serialises the map lookup and nothing else.
|
||||
//
|
||||
// Nodes is shared with the stored status and not deep-copied: UpdateNodeProgress
|
||||
// replaces that slice rather than writing through it, so a reader holding the
|
||||
// old header sees a consistent older breakdown rather than a torn one.
|
||||
func (g *GalleryService) GetStatus(s string) *OpStatus {
|
||||
g.Lock()
|
||||
defer g.Unlock()
|
||||
|
||||
return g.statuses[s]
|
||||
status, ok := g.statuses[s]
|
||||
if !ok || status == nil {
|
||||
return nil
|
||||
}
|
||||
copied := *status
|
||||
return &copied
|
||||
}
|
||||
|
||||
// GetAllStatus returns a snapshot of every operation's status. Same rule as
|
||||
// GetStatus, and for the same reason: the map and the statuses in it are both
|
||||
// copied, because this one is handed straight to a JSON encoder while peers'
|
||||
// broadcasts are still arriving.
|
||||
func (g *GalleryService) GetAllStatus() map[string]*OpStatus {
|
||||
g.Lock()
|
||||
defer g.Unlock()
|
||||
|
||||
return g.statuses
|
||||
snapshot := make(map[string]*OpStatus, len(g.statuses))
|
||||
for id, status := range g.statuses {
|
||||
if status == nil {
|
||||
continue
|
||||
}
|
||||
copied := *status
|
||||
snapshot[id] = &copied
|
||||
}
|
||||
return snapshot
|
||||
}
|
||||
|
||||
// ReapStaleOperations marks abandoned in-progress operations (pending/
|
||||
@@ -364,7 +427,7 @@ func (g *GalleryService) ReapStaleOperations(age time.Duration) (int64, error) {
|
||||
}
|
||||
// The database row is only half the picture. GET /models/jobs/<id> and
|
||||
// /api/operations read the in-memory statuses map, which is populated
|
||||
// locally and via the NATS progress broadcast and never expires. An op
|
||||
// locally and via the progress broadcast and never expires. An op
|
||||
// orphaned by a replica that died mid-download therefore kept serving its
|
||||
// last frozen tick (phase=downloading, processed=false, error=none) on
|
||||
// every replica indefinitely, long after the reaper had already given up
|
||||
@@ -430,7 +493,7 @@ func (g *GalleryService) stopOperation(id string, pause bool) error {
|
||||
delete(g.cancellations, id)
|
||||
}
|
||||
|
||||
nc := g.natsClient
|
||||
nc := g.broadcaster
|
||||
store := g.galleryStore
|
||||
|
||||
if !localExists && nc == nil {
|
||||
@@ -484,7 +547,7 @@ func (g *GalleryService) stopOperation(id string, pause bool) error {
|
||||
|
||||
// applyCancel is the broadcast-side counterpart to CancelOperation. The
|
||||
// wildcard subscriber calls it when a peer publishes a cancel event:
|
||||
// run the local cancel func if we have one (no echo via NATS), and reflect
|
||||
// run the local cancel func if we have one (no echo via the carrier), and reflect
|
||||
// the cancellation in the local statuses map. Idempotent: a replica that
|
||||
// already cancelled this op locally treats the inbound event as a no-op.
|
||||
func (g *GalleryService) applyCancel(id string, pause bool) {
|
||||
@@ -699,7 +762,7 @@ func (g *GalleryService) Start(c context.Context, cl *config.ModelConfigLoader,
|
||||
// the pre-existing operations before live updates start flowing.
|
||||
func (g *GalleryService) SubscribeBroadcasts() error {
|
||||
g.Lock()
|
||||
nc := g.natsClient
|
||||
nc := g.broadcaster
|
||||
g.Unlock()
|
||||
if nc == nil {
|
||||
return nil
|
||||
@@ -752,7 +815,7 @@ func (g *GalleryService) SubscribeBroadcasts() error {
|
||||
g.Unlock()
|
||||
if cb != nil {
|
||||
// Run off-goroutine so a slow UpgradeChecker doesn't stall the
|
||||
// NATS receive loop. Matches the local fire-after-install path.
|
||||
// carrier's delivery loop. Matches the local fire-after-install path.
|
||||
go cb()
|
||||
}
|
||||
})
|
||||
|
||||
@@ -24,7 +24,7 @@ import (
|
||||
// limitation, since the cap bounds the chain length for very long prompts.
|
||||
//
|
||||
// xxhash is used (not hash/maphash) because the hash MUST be identical across
|
||||
// frontend processes: peers exchange these hashes over NATS, and maphash uses a
|
||||
// frontend processes: peers exchange these hashes on the broadcast carrier, and maphash uses a
|
||||
// per-process random seed that would make peers disagree.
|
||||
func ExtractChain(model, prompt string, cfg Config) []uint64 {
|
||||
if prompt == "" {
|
||||
@@ -38,7 +38,7 @@ func ExtractChain(model, prompt string, cfg Config) []uint64 {
|
||||
// state, so Reset()+Write produces the byte-identical value to a fresh
|
||||
// New()+Write. xxhash seed 0 is stateless, so output is unchanged while we
|
||||
// avoid allocating a Digest per block. The output determinism across
|
||||
// processes (peers exchange these hashes over NATS) is preserved.
|
||||
// processes (peers exchange these hashes on the broadcast carrier) is preserved.
|
||||
h := xxhash.New()
|
||||
chain := make([]uint64, 0, depth)
|
||||
prev := salt
|
||||
|
||||
@@ -1,34 +1,99 @@
|
||||
package prefixcache
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// publisher is the minimal slice of messaging.Client that Sync needs.
|
||||
type publisher interface {
|
||||
Publish(subject string, v any) error
|
||||
}
|
||||
|
||||
// Sync wraps an Index, broadcasting new/extended observations to peers and
|
||||
// applying peers' broadcasts. It is the cross-frontend coherence layer.
|
||||
//
|
||||
// It holds ONE carrier, and both directions use it: what this replica observes
|
||||
// goes out on it, and what peers observe comes back in on it through
|
||||
// SubscribeBroadcasts. Splitting those across two carriers would leave each
|
||||
// frontend routing on nothing but its own history while every publish and every
|
||||
// subscribe succeeded, so there is no way to spell that here.
|
||||
type Sync struct {
|
||||
idx Provider
|
||||
pub publisher
|
||||
// bus is messaging.Broadcaster, which is Publish plus Subscribe and nothing
|
||||
// more. This package is reached from the routing hot path and a spec drives
|
||||
// it with a two-method double; a wider parameter would hand the hot path
|
||||
// request/reply it must never make.
|
||||
bus messaging.Broadcaster
|
||||
}
|
||||
|
||||
func NewSync(idx Provider, pub publisher) *Sync { return &Sync{idx: idx, pub: pub} }
|
||||
// NewSync wraps idx and puts its cross-frontend traffic on bus. A nil bus keeps
|
||||
// the Sync local: it records and answers, and broadcasts nothing.
|
||||
func NewSync(idx Provider, bus messaging.Broadcaster) *Sync { return &Sync{idx: idx, bus: bus} }
|
||||
|
||||
// SubscribeBroadcasts applies peers' observations and invalidations into the
|
||||
// wrapped index. ApplyObserve and ApplyInvalidate update only the local index
|
||||
// and never re-publish, so there is no broadcast loop.
|
||||
//
|
||||
// It takes no carrier argument on purpose. The carrier is the one this Sync
|
||||
// publishes on, read from the same field, which is what makes "this replica
|
||||
// hears what it would have said" a property of the type rather than of whoever
|
||||
// wired it.
|
||||
//
|
||||
// Returns every subscription it opened; on a partial failure it releases what
|
||||
// it already opened and returns nothing, so a caller cannot be left holding a
|
||||
// half-subscribed Sync it believes is whole.
|
||||
func (s *Sync) SubscribeBroadcasts() ([]messaging.Subscription, error) {
|
||||
if s.bus == nil {
|
||||
return nil, nil
|
||||
}
|
||||
var subs []messaging.Subscription
|
||||
release := func() {
|
||||
for _, sub := range subs {
|
||||
if err := sub.Unsubscribe(); err != nil {
|
||||
xlog.Warn("prefixcache: releasing a partial subscription", "error", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
observeSub, err := messaging.SubscribeJSON(s.bus, messaging.SubjectPrefixCacheObserve, func(ev messaging.PrefixCacheObserveEvent) {
|
||||
s.ApplyObserve(ev, time.Now())
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("prefixcache: subscribing to %s: %w", messaging.SubjectPrefixCacheObserve, err)
|
||||
}
|
||||
subs = append(subs, observeSub)
|
||||
|
||||
invalidateSub, err := messaging.SubscribeJSON(s.bus, messaging.SubjectPrefixCacheInvalidate, func(ev messaging.PrefixCacheInvalidateEvent) {
|
||||
s.ApplyInvalidate(ev)
|
||||
})
|
||||
if err != nil {
|
||||
release()
|
||||
return nil, fmt.Errorf("prefixcache: subscribing to %s: %w", messaging.SubjectPrefixCacheInvalidate, err)
|
||||
}
|
||||
subs = append(subs, invalidateSub)
|
||||
return subs, nil
|
||||
}
|
||||
|
||||
// Observe records locally and, if new/extended, broadcasts to peers. It returns
|
||||
// whether the local index treated the assignment as new or extended, so Sync
|
||||
// satisfies prefixcache.Provider.
|
||||
//
|
||||
// The LOCAL record happens first and is never conditional on the broadcast. A
|
||||
// hint this replica could not share is still a fact this replica learned, and
|
||||
// collapsing the two would cost the observing replica its own affinity for the
|
||||
// very prompts it is already serving.
|
||||
//
|
||||
// The broadcast is an ordinary Publish, like every other family on the carrier.
|
||||
// An observation is bounded by construction: ExtractChain caps a chain at
|
||||
// Config.MaxDepth blocks, so the event a frontend publishes has a known worst
|
||||
// case, and core/application refuses to start a deployment whose configured
|
||||
// depth would push that worst case past what a notification can carry. Nothing
|
||||
// here drops a message to stay under the cap: a drop would cost peers their
|
||||
// affinity silently, and the startup refusal says so instead.
|
||||
func (s *Sync) Observe(model string, chain []uint64, key ReplicaKey, now time.Time) bool {
|
||||
changed := s.idx.Observe(model, chain, key, now)
|
||||
if changed && s.pub != nil {
|
||||
if changed && s.bus != nil {
|
||||
ev := messaging.PrefixCacheObserveEvent{Model: model, Chain: chain, NodeID: key.NodeID, Replica: key.Replica}
|
||||
if err := s.pub.Publish(messaging.SubjectPrefixCacheObserve, ev); err != nil {
|
||||
if err := s.bus.Publish(messaging.SubjectPrefixCacheObserve, ev); err != nil {
|
||||
xlog.Debug("prefixcache: observe publish failed", "error", err)
|
||||
}
|
||||
}
|
||||
@@ -37,7 +102,7 @@ func (s *Sync) Observe(model string, chain []uint64, key ReplicaKey, now time.Ti
|
||||
|
||||
// Invalidate drops the local entry for one replica and broadcasts to peers. The
|
||||
// local drop is a no-op for models that were never cached (Index.Invalidate does
|
||||
// not intern a tree). The broadcast is UNCONDITIONAL (when a publisher is
|
||||
// not intern a tree). The broadcast is UNCONDITIONAL (when a carrier is
|
||||
// configured): the registry chokepoint fires for every replica removal, and a
|
||||
// peer frontend may hold a stale entry for the model even when THIS frontend
|
||||
// never cached it, so gating the broadcast on local-tree existence would drop
|
||||
@@ -45,9 +110,9 @@ func (s *Sync) Observe(model string, chain []uint64, key ReplicaKey, now time.Ti
|
||||
// until their TTL.
|
||||
func (s *Sync) Invalidate(model string, key ReplicaKey) {
|
||||
s.idx.Invalidate(model, key)
|
||||
if s.pub != nil {
|
||||
if s.bus != nil {
|
||||
ev := messaging.PrefixCacheInvalidateEvent{Model: model, NodeID: key.NodeID, Replica: key.Replica}
|
||||
if err := s.pub.Publish(messaging.SubjectPrefixCacheInvalidate, ev); err != nil {
|
||||
if err := s.bus.Publish(messaging.SubjectPrefixCacheInvalidate, ev); err != nil {
|
||||
xlog.Debug("prefixcache: invalidate publish failed", "error", err)
|
||||
}
|
||||
}
|
||||
@@ -58,9 +123,9 @@ func (s *Sync) Invalidate(model string, key ReplicaKey) {
|
||||
// coherence. A negative Replica on the wire means "all replicas of the node".
|
||||
func (s *Sync) InvalidateNode(model, node string) {
|
||||
s.idx.InvalidateNode(model, node)
|
||||
if s.pub != nil {
|
||||
if s.bus != nil {
|
||||
ev := messaging.PrefixCacheInvalidateEvent{Model: model, NodeID: node, Replica: -1}
|
||||
if err := s.pub.Publish(messaging.SubjectPrefixCacheInvalidate, ev); err != nil {
|
||||
if err := s.bus.Publish(messaging.SubjectPrefixCacheInvalidate, ev); err != nil {
|
||||
xlog.Debug("prefixcache: invalidate-node publish failed", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
package prefixcache_test
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"math"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
@@ -10,15 +14,71 @@ import (
|
||||
"github.com/mudler/LocalAI/core/services/nodes/prefixcache"
|
||||
)
|
||||
|
||||
type fakePub struct{ published []any }
|
||||
// fakePub is the two-method carrier Sync is given: Publish and Subscribe and
|
||||
// nothing else, which is the whole of messaging.Broadcaster. It delivers
|
||||
// synchronously to matching subscribers so a spec can drive one Sync and read
|
||||
// the other without polling.
|
||||
type fakePub struct {
|
||||
mu sync.Mutex
|
||||
published []any
|
||||
subjects []string
|
||||
subs []fakePubSub
|
||||
publishErr error
|
||||
}
|
||||
|
||||
type fakePubSub struct {
|
||||
subject string
|
||||
handler func([]byte)
|
||||
}
|
||||
|
||||
func (f *fakePub) Publish(subject string, v any) error {
|
||||
f.mu.Lock()
|
||||
f.published = append(f.published, v)
|
||||
f.subjects = append(f.subjects, subject)
|
||||
err := f.publishErr
|
||||
subs := append([]fakePubSub(nil), f.subs...)
|
||||
f.mu.Unlock()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
payload, merr := json.Marshal(v)
|
||||
if merr != nil {
|
||||
return merr
|
||||
}
|
||||
for _, sub := range subs {
|
||||
if messaging.SubjectMatches(sub.subject, subject) {
|
||||
sub.handler(payload)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakePub) Subscribe(subject string, handler func([]byte)) (messaging.Subscription, error) {
|
||||
if err := messaging.ValidFilter(subject); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.subs = append(f.subs, fakePubSub{subject: subject, handler: handler})
|
||||
return fakePubSubscription{}, nil
|
||||
}
|
||||
|
||||
func (f *fakePub) subscribed() []string {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
out := make([]string, 0, len(f.subs))
|
||||
for _, sub := range f.subs {
|
||||
out = append(out, sub.subject)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
type fakePubSubscription struct{}
|
||||
|
||||
func (fakePubSubscription) Unsubscribe() error { return nil }
|
||||
|
||||
// Sync must satisfy the Provider seam so SmartRouter can hold a single
|
||||
// prefixcache.Provider that broadcasts via NATS.
|
||||
// prefixcache.Provider that broadcasts to its peers.
|
||||
var _ prefixcache.Provider = (*prefixcache.Sync)(nil)
|
||||
|
||||
var _ = Describe("Sync", func() {
|
||||
@@ -116,3 +176,175 @@ var _ = Describe("Sync", func() {
|
||||
Expect(idx.Decide("m", []uint64{3, 4}, cands, t0).HasHot).To(BeFalse())
|
||||
})
|
||||
})
|
||||
|
||||
// The observation family's own rules, and the reason it needs none of its own
|
||||
// machinery.
|
||||
//
|
||||
// The plan for this phase proposed publishing observations through a method
|
||||
// that REFUSES a message too large for a notification instead of spilling it,
|
||||
// on the reasoning that a long prompt makes a chain of thousands of entries.
|
||||
// ExtractChain does not produce one: it caps a chain at Config.MaxDepth blocks,
|
||||
// and MaxDepth is a constant with no operator knob, so an observation's size
|
||||
// has a known worst case six times under the cap. A refusal would therefore
|
||||
// have been a deliberate, silent message drop guarding a condition that cannot
|
||||
// arise, and the first change to MaxDepth would have turned it into a
|
||||
// deployment that lost cross-replica affinity while reporting nothing.
|
||||
//
|
||||
// So Observe publishes like every other family, and the bound is the thing that
|
||||
// is pinned: here, that a full-depth chain is what the extractor can produce,
|
||||
// and in core/application, that a full-depth observation still fits in one
|
||||
// notification.
|
||||
var _ = Describe("Sync observation size", func() {
|
||||
It("caps the chain a full-length prompt produces at MaxDepth", func() {
|
||||
cfg := prefixcache.DefaultConfig()
|
||||
// Ten times more prompt than the extractor will hash, so the cap is
|
||||
// what decides the length and not the prompt.
|
||||
prompt := make([]byte, cfg.WindowBytes*cfg.MaxDepth*10)
|
||||
for i := range prompt {
|
||||
prompt[i] = byte('a' + i%26)
|
||||
}
|
||||
|
||||
chain := prefixcache.ExtractChain("m", string(prompt), cfg)
|
||||
|
||||
Expect(chain).To(HaveLen(cfg.MaxDepth),
|
||||
"an observation's worst case is Config.MaxDepth entries; if it is not, nothing bounds what Observe publishes")
|
||||
})
|
||||
|
||||
It("publishes a full-depth observation rather than trimming or dropping it", func() {
|
||||
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
pub := &fakePub{}
|
||||
s := prefixcache.NewSync(idx, pub)
|
||||
chain := make([]uint64, prefixcache.DefaultConfig().MaxDepth)
|
||||
for i := range chain {
|
||||
chain[i] = math.MaxUint64 - uint64(i)
|
||||
}
|
||||
|
||||
Expect(s.Observe("m", chain, rk("A", 0), t0)).To(BeTrue())
|
||||
|
||||
Expect(pub.published).To(HaveLen(1))
|
||||
ev := pub.published[0].(messaging.PrefixCacheObserveEvent)
|
||||
Expect(ev.Chain).To(Equal(chain),
|
||||
"a peer must be able to reconstruct the same prefix, so the chain travels whole")
|
||||
})
|
||||
|
||||
// The load-bearing half. A hint this replica could not share is still a
|
||||
// fact this replica learned, and an Observe that abandoned the local record
|
||||
// when the carrier said no would cost the observing frontend its own
|
||||
// affinity for the prompts it is already serving.
|
||||
It("records locally even when the broadcast fails", func() {
|
||||
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
pub := &fakePub{publishErr: errors.New("carrier refused it")}
|
||||
s := prefixcache.NewSync(idx, pub)
|
||||
chain := []uint64{7, 8, 9}
|
||||
|
||||
Expect(s.Observe("m", chain, rk("A", 0), t0)).To(BeTrue())
|
||||
|
||||
d := s.Decide("m", chain, []prefixcache.ReplicaKey{rk("A", 0)}, t0)
|
||||
Expect(d.HasHot).To(BeTrue(), "the observing replica must keep its own affinity")
|
||||
Expect(d.Hot).To(Equal(rk("A", 0)))
|
||||
})
|
||||
|
||||
// Invalidations are not hints. A peer that misses one routes to a replica
|
||||
// that is gone until its TTL, which is why Sync.Invalidate's own comment
|
||||
// calls the broadcast unconditional. Asserted as the SUBJECT each call
|
||||
// published on, so a family moved onto the wrong subject fails here rather
|
||||
// than in a deployment.
|
||||
DescribeTable("publishes every family on its own subject",
|
||||
func(act func(*prefixcache.Sync), wantSubject string) {
|
||||
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
pub := &fakePub{}
|
||||
s := prefixcache.NewSync(idx, pub)
|
||||
|
||||
act(s)
|
||||
|
||||
Expect(pub.subjects).To(ConsistOf(wantSubject))
|
||||
},
|
||||
Entry("observe", func(s *prefixcache.Sync) { s.Observe("m", []uint64{1}, rk("A", 0), t0) },
|
||||
messaging.SubjectPrefixCacheObserve),
|
||||
Entry("invalidate", func(s *prefixcache.Sync) { s.Invalidate("m", rk("A", 0)) },
|
||||
messaging.SubjectPrefixCacheInvalidate),
|
||||
Entry("invalidate-node", func(s *prefixcache.Sync) { s.InvalidateNode("m", "A") },
|
||||
messaging.SubjectPrefixCacheInvalidate),
|
||||
)
|
||||
})
|
||||
|
||||
// Both directions on ONE carrier. SubscribeBroadcasts takes no carrier
|
||||
// argument, so a Sync that publishes where its peers are not listening cannot
|
||||
// be built; these pin that the subscriptions it opens are on the carrier it was
|
||||
// given and that they carry both families.
|
||||
var _ = Describe("Sync.SubscribeBroadcasts", func() {
|
||||
It("subscribes both families on the carrier it publishes to", func() {
|
||||
pub := &fakePub{}
|
||||
s := prefixcache.NewSync(prefixcache.NewIndex(prefixcache.DefaultConfig()), pub)
|
||||
|
||||
subs, err := s.SubscribeBroadcasts()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(subs).To(HaveLen(2))
|
||||
Expect(pub.subscribed()).To(ConsistOf(
|
||||
messaging.SubjectPrefixCacheObserve,
|
||||
messaging.SubjectPrefixCacheInvalidate,
|
||||
))
|
||||
})
|
||||
|
||||
It("registers nothing at all when there is no carrier", func() {
|
||||
s := prefixcache.NewSync(prefixcache.NewIndex(prefixcache.DefaultConfig()), nil)
|
||||
|
||||
subs, err := s.SubscribeBroadcasts()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(subs).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("carries a peer's observation into this replica's index", func() {
|
||||
// Two Syncs, one carrier: exactly the production topology, with the
|
||||
// publish and the subscribe on the same bus because neither end can
|
||||
// name a different one.
|
||||
bus := &fakePub{}
|
||||
idxA := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
idxB := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
a := prefixcache.NewSync(idxA, bus)
|
||||
b := prefixcache.NewSync(idxB, bus)
|
||||
_, err := b.SubscribeBroadcasts()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
chain := []uint64{11, 22, 33}
|
||||
a.Observe("m", chain, rk("A", 3), t0)
|
||||
|
||||
d := b.Decide("m", chain, []prefixcache.ReplicaKey{rk("A", 3)}, t0)
|
||||
Expect(d.HasHot).To(BeTrue(), "the peer replica must learn where the prefix is warm")
|
||||
Expect(d.Hot).To(Equal(rk("A", 3)))
|
||||
})
|
||||
|
||||
It("carries a peer's invalidation, so a removed replica stops being routed to", func() {
|
||||
bus := &fakePub{}
|
||||
idxA := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
idxB := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
a := prefixcache.NewSync(idxA, bus)
|
||||
b := prefixcache.NewSync(idxB, bus)
|
||||
_, err := b.SubscribeBroadcasts()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
chain := []uint64{44, 55}
|
||||
a.Observe("m", chain, rk("A", 0), t0)
|
||||
Expect(b.Decide("m", chain, []prefixcache.ReplicaKey{rk("A", 0)}, t0).HasHot).To(BeTrue())
|
||||
|
||||
a.Invalidate("m", rk("A", 0))
|
||||
|
||||
Expect(b.Decide("m", chain, []prefixcache.ReplicaKey{rk("A", 0)}, t0).HasHot).To(BeFalse(),
|
||||
"a missed invalidation would leave this replica routing to a replica that is gone")
|
||||
})
|
||||
|
||||
It("does not re-broadcast what it applied", func() {
|
||||
// The echo-loop guard. ApplyObserve and ApplyInvalidate write only the
|
||||
// local index; a peer that answered a broadcast with a broadcast would
|
||||
// keep the carrier busy forever.
|
||||
bus := &fakePub{}
|
||||
b := prefixcache.NewSync(prefixcache.NewIndex(prefixcache.DefaultConfig()), bus)
|
||||
_, err := b.SubscribeBroadcasts()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
b.ApplyObserve(messaging.PrefixCacheObserveEvent{Model: "m", Chain: []uint64{1, 2}, NodeID: "A"}, t0)
|
||||
b.ApplyInvalidate(messaging.PrefixCacheInvalidateEvent{Model: "m", NodeID: "A", Replica: 0})
|
||||
|
||||
Expect(bus.published).To(BeEmpty())
|
||||
})
|
||||
})
|
||||
@@ -30,9 +30,14 @@ const (
|
||||
// FileComplete, Complete) always publish so peers never miss them.
|
||||
stagingBroadcastInterval = time.Second
|
||||
// stagingRemoteTTL drops a mirrored (remote) op whose last update is older
|
||||
// than this. NATS pub/sub is fire-and-forget, so a missed Done event would
|
||||
// otherwise leave a phantom staging row on a peer forever; a live op
|
||||
// than this. The broadcast carrier is at-most-once, so a missed Done event
|
||||
// would otherwise leave a phantom staging row on a peer forever; a live op
|
||||
// refreshes its mirror at least every stagingBroadcastInterval.
|
||||
//
|
||||
// This is the direction the TTL is allowed to be wrong in, and it is the
|
||||
// invariant this family has to keep: a missed event ages a mirror out, so
|
||||
// a staging op that never happened is never invented, and one that is
|
||||
// still running is re-asserted on the next tick.
|
||||
stagingRemoteTTL = 60 * time.Second
|
||||
)
|
||||
|
||||
@@ -51,13 +56,17 @@ type stagingEntry struct {
|
||||
// Used by SmartRouter to publish progress and by /api/operations to surface it.
|
||||
//
|
||||
// In distributed mode each frontend replica runs its own tracker. The replica
|
||||
// performing a transfer owns the op locally and broadcasts progress over NATS
|
||||
// (SetPublisher); peers mirror it via ApplyRemote (SubscribeBroadcasts) so a
|
||||
// performing a transfer owns the op locally and broadcasts progress on the
|
||||
// deployment's fan-out carrier; peers mirror it via ApplyRemote, so a
|
||||
// /api/operations poll that round-robins onto any replica surfaces the op.
|
||||
// SetBroadcaster wires both halves at once, and it is one method for a reason
|
||||
// written there.
|
||||
type StagingTracker struct {
|
||||
mu sync.RWMutex
|
||||
active map[string]*stagingEntry
|
||||
publisher messaging.Publisher
|
||||
mu sync.RWMutex
|
||||
active map[string]*stagingEntry
|
||||
// broadcaster is where this tracker publishes AND where it mirrors from.
|
||||
// One field, set by one method, so the two cannot name different carriers.
|
||||
broadcaster messaging.Broadcaster
|
||||
}
|
||||
|
||||
// StagingProgressEvent is the wire payload a frontend replica broadcasts on
|
||||
@@ -76,19 +85,24 @@ func NewStagingTracker() *StagingTracker {
|
||||
}
|
||||
}
|
||||
|
||||
// SetPublisher wires the NATS publisher used to broadcast staging progress to
|
||||
// peer replicas. No-op publisher (nil) keeps the tracker standalone.
|
||||
func (t *StagingTracker) SetPublisher(p messaging.Publisher) {
|
||||
// SetBroadcaster puts this tracker's publishing AND its mirroring of peers onto
|
||||
// ONE carrier, and returns the mirror subscription for cleanup. A nil
|
||||
// broadcaster keeps the tracker standalone and registers nothing.
|
||||
//
|
||||
// One method rather than a setter and a subscriber, because they were never two
|
||||
// decisions. A tracker that publishes on one carrier and listens on another
|
||||
// shows a staging progress bar on the replica performing the transfer and on no
|
||||
// other, which is the exact symptom SubjectStagingProgress exists to prevent,
|
||||
// and neither half reports an error while it happens. With one argument that
|
||||
// deployment cannot be spelled.
|
||||
func (t *StagingTracker) SetBroadcaster(b messaging.Broadcaster) (messaging.Subscription, error) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
t.publisher = p
|
||||
}
|
||||
|
||||
// SubscribeBroadcasts subscribes to peer replicas' staging-progress broadcasts
|
||||
// and mirrors them into this tracker, so /api/operations on any replica surfaces
|
||||
// staging ops it did not originate. Returns the subscription for cleanup.
|
||||
func (t *StagingTracker) SubscribeBroadcasts(nc messaging.MessagingClient) (messaging.Subscription, error) {
|
||||
return messaging.SubscribeJSON(nc, messaging.SubjectStagingProgressWildcard, func(evt StagingProgressEvent) {
|
||||
t.broadcaster = b
|
||||
t.mu.Unlock()
|
||||
if b == nil {
|
||||
return nil, nil
|
||||
}
|
||||
return messaging.SubscribeJSON(b, messaging.SubjectStagingProgressWildcard, func(evt StagingProgressEvent) {
|
||||
if evt.ModelID == "" {
|
||||
return
|
||||
}
|
||||
@@ -96,9 +110,9 @@ func (t *StagingTracker) SubscribeBroadcasts(nc messaging.MessagingClient) (mess
|
||||
})
|
||||
}
|
||||
|
||||
// publishStaging emits an event to the per-model staging subject. The publisher
|
||||
// publishStaging emits an event to the per-model staging subject. The carrier
|
||||
// is captured by the caller under the lock and passed in, so publishing happens
|
||||
// outside the lock (a slow NATS link must not stall the staging copy loop).
|
||||
// outside the lock (a slow carrier must not stall the staging copy loop).
|
||||
func publishStaging(p messaging.Publisher, evt StagingProgressEvent) {
|
||||
if p == nil {
|
||||
return
|
||||
@@ -121,7 +135,7 @@ func (t *StagingTracker) Start(modelID, nodeName string, totalFiles int) {
|
||||
// lastPub stays zero so the first UpdateFile tick always broadcasts.
|
||||
}
|
||||
t.active[modelID] = e
|
||||
pub := t.publisher
|
||||
pub := t.broadcaster
|
||||
snap := e.status
|
||||
t.mu.Unlock()
|
||||
|
||||
@@ -167,7 +181,7 @@ func (t *StagingTracker) UpdateFile(modelID, fileName string, fileIndex int, byt
|
||||
var snap StagingStatus
|
||||
if time.Since(e.lastPub) >= stagingBroadcastInterval {
|
||||
e.lastPub = time.Now()
|
||||
pub = t.publisher
|
||||
pub = t.broadcaster
|
||||
snap = e.status
|
||||
}
|
||||
t.mu.Unlock()
|
||||
@@ -194,7 +208,7 @@ func (t *StagingTracker) FileComplete(modelID string, fileIndex, totalFiles int)
|
||||
s.Speed = ""
|
||||
e.updatedAt = time.Now()
|
||||
e.lastPub = time.Now()
|
||||
pub := t.publisher
|
||||
pub := t.broadcaster
|
||||
snap := e.status
|
||||
t.mu.Unlock()
|
||||
|
||||
@@ -207,7 +221,7 @@ func (t *StagingTracker) Complete(modelID string) {
|
||||
t.mu.Lock()
|
||||
_, ok := t.active[modelID]
|
||||
delete(t.active, modelID)
|
||||
pub := t.publisher
|
||||
pub := t.broadcaster
|
||||
t.mu.Unlock()
|
||||
|
||||
if ok {
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// decodeStagingEvents extracts every StagingProgressEvent the fake messaging
|
||||
@@ -33,13 +34,14 @@ var _ = Describe("StagingTracker cross-replica broadcast", func() {
|
||||
It("broadcasts staging progress so a peer replica surfaces an op it did not originate", func() {
|
||||
mc := &fakeMessagingClient{}
|
||||
origin := NewStagingTracker()
|
||||
origin.SetPublisher(mc)
|
||||
_, err := origin.SetBroadcaster(mc)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
origin.Start("model-x", "worker-1", 1)
|
||||
origin.UpdateFile("model-x", "weights.gguf", 1, 5<<30, 10<<30, "100 MiB/s")
|
||||
|
||||
events := decodeStagingEvents(mc)
|
||||
Expect(events).ToNot(BeEmpty(), "writes must be broadcast over NATS")
|
||||
Expect(events).ToNot(BeEmpty(), "writes must be broadcast to peer replicas")
|
||||
Expect(mc.published[0].Subject).To(Equal(messaging.SubjectStagingProgress("model-x")))
|
||||
|
||||
// A peer replica that never ran the op merges the broadcast.
|
||||
@@ -58,7 +60,8 @@ var _ = Describe("StagingTracker cross-replica broadcast", func() {
|
||||
It("removes the op from the peer when the origin completes it", func() {
|
||||
mc := &fakeMessagingClient{}
|
||||
origin := NewStagingTracker()
|
||||
origin.SetPublisher(mc)
|
||||
_, err := origin.SetBroadcaster(mc)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
origin.Start("model-x", "worker-1", 1)
|
||||
origin.Complete("model-x")
|
||||
@@ -90,7 +93,7 @@ var _ = Describe("StagingTracker cross-replica broadcast", func() {
|
||||
})
|
||||
})
|
||||
|
||||
Context("when no publisher is wired (standalone mode)", func() {
|
||||
Context("when no carrier is wired (standalone mode)", func() {
|
||||
It("does not broadcast", func() {
|
||||
mc := &fakeMessagingClient{}
|
||||
t := NewStagingTracker()
|
||||
@@ -98,6 +101,48 @@ var _ = Describe("StagingTracker cross-replica broadcast", func() {
|
||||
t.UpdateFile("model-x", "weights.gguf", 1, 1<<30, 10<<30, "")
|
||||
Expect(mc.published).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("registers no mirror subscription for a nil carrier", func() {
|
||||
// The negative half of SetBroadcaster's one-carrier promise: a
|
||||
// standalone tracker must not subscribe either. "Nothing was
|
||||
// delivered" cannot tell that apart from a live subscription
|
||||
// nobody published to.
|
||||
sub, err := NewStagingTracker().SetBroadcaster(nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(sub).To(BeNil())
|
||||
})
|
||||
})
|
||||
|
||||
// SetBroadcaster is ONE method because publishing and mirroring are one
|
||||
// decision. These two pin each half separately: a mutation that drops the
|
||||
// subscribe leaves the publish spec green, and a mutation that drops the
|
||||
// field assignment leaves the subscribe green, so a single spec over both
|
||||
// would let either half go missing.
|
||||
Context("both halves ride the carrier it is given", func() {
|
||||
It("mirrors a peer's broadcast that arrives on the same carrier", func() {
|
||||
bus := testutil.NewFakeBus()
|
||||
peer := NewStagingTracker()
|
||||
_, err := peer.SetBroadcaster(bus)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(bus.Publish(messaging.SubjectStagingProgress("model-y"), StagingProgressEvent{
|
||||
ModelID: "model-y",
|
||||
Status: &StagingStatus{ModelID: "model-y", NodeName: "worker-2"},
|
||||
})).To(Succeed())
|
||||
|
||||
Expect(peer.GetAll()).To(HaveKey("model-y"))
|
||||
})
|
||||
|
||||
It("publishes onto the same carrier it mirrors from", func() {
|
||||
bus := testutil.NewFakeBus()
|
||||
origin := NewStagingTracker()
|
||||
_, err := origin.SetBroadcaster(bus)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
origin.Start("model-z", "worker-3", 1)
|
||||
|
||||
Expect(bus.PublishCount(messaging.SubjectStagingProgress("model-z"))).To(Equal(1))
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
|
||||
@@ -0,0 +1,86 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package nodes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/pgbus"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// Staging progress on the real carrier, between two trackers on two LISTEN
|
||||
// connections.
|
||||
//
|
||||
// One tracker talking to itself proves nothing here: the whole reason this
|
||||
// family exists is that a /api/operations poll round-robins onto a replica that
|
||||
// did not perform the transfer. An in-memory double cannot fail the way a
|
||||
// carrier fails, and it cannot tell a tracker that publishes and listens on two
|
||||
// different carriers apart from one that does not, which is the defect
|
||||
// SetBroadcaster is shaped to exclude.
|
||||
var _ = Describe("staging progress across two replicas on the broadcast carrier", func() {
|
||||
var trackerA, trackerB *StagingTracker
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
ctx := context.Background()
|
||||
db, dsn := testutil.SetupTestDBWithDSN()
|
||||
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
|
||||
|
||||
newTracker := func() *StagingTracker {
|
||||
bus, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(bus.Close)
|
||||
t := NewStagingTracker()
|
||||
sub, err := t.SetBroadcaster(bus)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(sub).ToNot(BeNil())
|
||||
return t
|
||||
}
|
||||
trackerA, trackerB = newTracker(), newTracker()
|
||||
})
|
||||
|
||||
It("mirrors a peer's start, progress and completion, in that order", func() {
|
||||
trackerA.Start("model-a", "worker-1", 2)
|
||||
|
||||
Eventually(func() map[string]StagingStatus { return trackerB.GetAll() }, 20*time.Second).
|
||||
Should(HaveKey("model-a"))
|
||||
Expect(trackerB.GetAll()["model-a"].NodeName).To(Equal("worker-1"))
|
||||
|
||||
trackerA.UpdateFile("model-a", "weights.gguf", 1, 5<<30, 10<<30, "100 MiB/s")
|
||||
|
||||
Eventually(func() string {
|
||||
if s := trackerB.Get("model-a"); s != nil {
|
||||
return s.FileName
|
||||
}
|
||||
return ""
|
||||
}, 20*time.Second).Should(Equal("weights.gguf"))
|
||||
|
||||
// The Done event is what stops a peer showing a transfer that
|
||||
// finished. A tracker that only ever received the start would pass
|
||||
// every assertion above and leave the row on screen until the TTL.
|
||||
trackerA.Complete("model-a")
|
||||
|
||||
Eventually(func() map[string]StagingStatus { return trackerB.GetAll() }, 20*time.Second).
|
||||
ShouldNot(HaveKey("model-a"))
|
||||
})
|
||||
|
||||
It("leaves the mirroring replica's own operations alone", func() {
|
||||
trackerB.Start("model-b", "worker-local", 1)
|
||||
trackerA.Start("model-a", "worker-1", 1)
|
||||
|
||||
Eventually(func() map[string]StagingStatus { return trackerB.GetAll() }, 20*time.Second).
|
||||
Should(HaveKey("model-a"))
|
||||
|
||||
own := trackerB.Get("model-b")
|
||||
Expect(own).ToNot(BeNil(), "a replica must keep the transfer it is itself performing")
|
||||
Expect(own.NodeName).To(Equal("worker-local"))
|
||||
})
|
||||
})
|
||||
@@ -331,6 +331,53 @@ func (b *Bus) Close() {
|
||||
})
|
||||
}
|
||||
|
||||
// inlineNotification encodes what Publish puts on the wire when a message fits
|
||||
// in the notification itself, and returns it together with the caller's encoded
|
||||
// payload so the spill path does not marshal a second time.
|
||||
//
|
||||
// It is the ONE place the wire form of an inline broadcast is built. Publish
|
||||
// measures what it returns and FitsInline asks about the same bytes, so a
|
||||
// caller that has proved its worst case fits cannot be proved wrong later by a
|
||||
// change to the envelope's keys.
|
||||
func inlineNotification(subject string, data any) ([]byte, json.RawMessage, error) {
|
||||
payload, err := json.Marshal(data)
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("pgbus: encoding a broadcast on %q: %w", subject, err)
|
||||
}
|
||||
encoded, err := json.Marshal(notification{Subject: subject, Data: payload})
|
||||
if err != nil {
|
||||
return nil, nil, fmt.Errorf("pgbus: encoding the notification for %q: %w", subject, err)
|
||||
}
|
||||
return encoded, payload, nil
|
||||
}
|
||||
|
||||
// FitsInline reports whether Publish would carry this message in the
|
||||
// notification itself rather than writing it to a row and notifying the id.
|
||||
//
|
||||
// It is a PREDICATE and never an action: nothing here refuses, drops or
|
||||
// truncates a message, and no caller may use it to decide not to publish. The
|
||||
// spill path is correct for every family on this carrier, and a message that
|
||||
// does not fit costs a row and a SELECT, never its contents.
|
||||
//
|
||||
// It exists for the one question that is worth asking BEFORE a message is ever
|
||||
// published: can a family whose payload has a known upper bound put that bound
|
||||
// on the wrong side of the cap? A family whose worst case fits is a family that
|
||||
// never spills, and a configuration change that moves the bound past the cap
|
||||
// can then be refused at startup instead of turning every request on the
|
||||
// inference path into a table write nobody notices. See
|
||||
// requirePrefixCacheFitsInline in core/application.
|
||||
//
|
||||
// It shares its size decision with Publish rather than restating one: the same
|
||||
// encoder, the same comparison against the same constant. A second comparison
|
||||
// would be a second constant in disguise.
|
||||
func FitsInline(subject string, data any) (bool, error) {
|
||||
encoded, _, err := inlineNotification(subject, data)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return len(encoded) < maxNotifyPayloadBytes, nil
|
||||
}
|
||||
|
||||
// Publish fans a message out to every subscriber of the subject on every
|
||||
// replica, including this one.
|
||||
func (b *Bus) Publish(subject string, data any) error {
|
||||
@@ -338,14 +385,9 @@ func (b *Bus) Publish(subject string, data any) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
payload, err := json.Marshal(data)
|
||||
encoded, payload, err := inlineNotification(subject, data)
|
||||
if err != nil {
|
||||
return fmt.Errorf("pgbus: encoding a broadcast on %q: %w", subject, err)
|
||||
}
|
||||
|
||||
encoded, err := json.Marshal(notification{Subject: subject, Data: payload})
|
||||
if err != nil {
|
||||
return fmt.Errorf("pgbus: encoding the notification for %q: %w", subject, err)
|
||||
return err
|
||||
}
|
||||
// The one size decision in this package. Several subjects on this carrier
|
||||
// exceed the cap in normal operation (a job result carries a whole LLM
|
||||
|
||||
@@ -127,6 +127,47 @@ var _ = Describe("broadcasts too large for a notification", func() {
|
||||
Expect(spilledRows()).To(Equal(int64(1)))
|
||||
})
|
||||
|
||||
// FitsInline and Publish, asserted TOGETHER on the same payload.
|
||||
//
|
||||
// FitsInline exists so a family with a known worst case can prove at
|
||||
// startup that it never spills, and that proof is worth nothing if the
|
||||
// predicate and the publisher can disagree: a family told "you fit" that
|
||||
// then spills on every request is exactly the invisible cost the check was
|
||||
// added to prevent. Both rows below state an ABSOLUTE size and assert both
|
||||
// halves, so a second size decision anywhere splits them.
|
||||
DescribeTable("answers the same size question Publish answers",
|
||||
func(total int, wantFits bool, wantRows int64) {
|
||||
data := payloadEncodingTo(subject, total)
|
||||
|
||||
fits, err := pgbus.FitsInline(subject, data)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(fits).To(Equal(wantFits))
|
||||
|
||||
_ = deliver(data)
|
||||
|
||||
Expect(spilledRows()).To(Equal(wantRows))
|
||||
},
|
||||
Entry("one byte under the cap", notifyCap-1, true, int64(0)),
|
||||
Entry("exactly at the cap", notifyCap, false, int64(1)),
|
||||
)
|
||||
|
||||
It("refuses to call a megabyte inline", func() {
|
||||
// Absolute, not derived from the cap, for the same reason the spill row
|
||||
// above is: this must stay red-for-the-right-reason under a mutated
|
||||
// constant.
|
||||
fits, err := pgbus.FitsInline(subject, map[string]string{"k": strings.Repeat("a", 1<<20)})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(fits).To(BeFalse())
|
||||
})
|
||||
|
||||
It("reports the encoding failure rather than guessing at a size", func() {
|
||||
// A payload that cannot be marshalled has no size, and answering
|
||||
// "false" would let a caller read an encoder bug as a payload that is
|
||||
// merely large.
|
||||
_, err := pgbus.FitsInline(subject, make(chan int))
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
|
||||
It("stores the caller's payload in the spill row, not the envelope", func() {
|
||||
data := payloadEncodingTo(subject, notifyCap)
|
||||
|
||||
|
||||
@@ -41,7 +41,7 @@ Each model gets its own gRPC backend process, so a single worker can serve multi
|
||||
- Each frontend replica holds **one extra PostgreSQL session** beyond its connection pool, pinned for the life of the process, and creates a `bus_messages` table. Both belong to the broadcast carrier that is replacing NATS for cross-replica fan-out; it already carries the four `state.*.delta` families (see [Cross-replica in-memory state](#cross-replica-in-memory-state)). Size `max_connections` for one additional session per frontend replica.
|
||||
- That session reports an `application_name` of `localai_pgbus_<id>`, so `SELECT count(*) FROM pg_stat_activity WHERE application_name LIKE 'localai_pgbus_%'` counts the replicas currently listening. If the carrier loses its session it redials and re-registers on its own; a broadcast published while it was down is not replayed, which is why nothing that must survive a gap is carried by a broadcast alone.
|
||||
- `bus_messages` holds only broadcasts too large for a PostgreSQL notification, and every replica retires rows older than ten minutes. The table is a spill buffer, not a log: it is not a place to read past events from.
|
||||
- **NATS** server - used for agent-worker coordination and the frontend's own cross-replica events. **Serve-backend workers do not connect to it at all**: every verb they take, and file staging with it, is an HTTP route on the worker's tunnel. Set no `LOCALAI_NATS_URL` on a `local-ai worker`. The frontend and any `local-ai agent-worker` still need one.
|
||||
- **NATS** server - used for agent-worker coordination. The frontend's own cross-replica events do not use it: they travel on the PostgreSQL the deployment already runs. **Serve-backend workers do not connect to it at all**: every verb they take, and file staging with it, is an HTTP route on the worker's tunnel. Set no `LOCALAI_NATS_URL` on a `local-ai worker`. The frontend and any `local-ai agent-worker` still need one.
|
||||
- All services must be on the same network (or reachable via configured URLs)
|
||||
|
||||
## Quick Start with Docker Compose
|
||||
@@ -134,6 +134,47 @@ Several features keep state in a frontend's process memory and surface it over t
|
||||
| Agent tasks | `state.agent-tasks.delta` and `state.agent-tasks.<user_id>.delta` |
|
||||
| Open Responses metadata | `state.responses-metadata.delta` |
|
||||
|
||||
### Cross-replica caches
|
||||
|
||||
A frontend also keeps caches that live for the life of the process rather than
|
||||
for the life of a request: which gallery operations are in flight and how far
|
||||
along they are, which admin operations have been admitted, which model files are
|
||||
being staged onto a worker, and which replica already holds the KV/prefix cache
|
||||
for a prompt. Each of those is kept current on every replica by a broadcast, and
|
||||
**every one of them is on PostgreSQL**. Nothing in this table uses NATS.
|
||||
|
||||
| Family | Subject | What a peer does with it |
|
||||
|--------|---------|--------------------------|
|
||||
| Gallery progress | `gallery.<op_id>.progress` | Merges the status so `/api/operations` answers the same on any replica |
|
||||
| Gallery cancel | `gallery.<op_id>.cancel` | Stops the install, on whichever replica is running it |
|
||||
| Operation cache admit | `gallery.opcache.start` | Learns that an operation was admitted, and whether it is a backend install |
|
||||
| Operation cache dismiss | `gallery.opcache.end` | Drops the operation from its own map |
|
||||
| Model cache invalidation | `cache.invalidate.models` | Reloads the model config from disk, or prunes a deleted one |
|
||||
| Backend cache invalidation | `cache.invalidate.backends` | Refreshes its upgrade-available cache |
|
||||
| Staging progress | `staging.<model_id>.progress` | Mirrors a transfer it did not perform, so the progress bar does not flicker |
|
||||
| Prefix-cache observation | `prefixcache.observe` | Learns which replica already holds the prefix for a prompt |
|
||||
| Prefix-cache invalidation | `prefixcache.invalidate` | Stops routing to a replica that is gone |
|
||||
|
||||
**An invalidation that does not arrive must never read as a cache that is
|
||||
valid.** The two `cache.invalidate.*` families and `prefixcache.invalidate` are
|
||||
therefore published like any other broadcast: one too large for a notification
|
||||
is written to `bus_messages` and read back by the peer, never dropped to save
|
||||
the write. A missed staging event ages the peer's mirrored row out after a
|
||||
minute rather than inventing a transfer, for the same reason.
|
||||
|
||||
**Gallery progress spills, routinely.** A progress event carries one entry per
|
||||
worker, so on a fleet of a few tens of nodes it is past the 8000-byte cap on
|
||||
every tick and travels as a row. That is the ordinary path, not an error.
|
||||
|
||||
**A prefix-cache observation never spills.** It carries one hash per prefix
|
||||
block, and the extractor caps a chain at 64 blocks, so the largest observation a
|
||||
frontend can publish is a few kilobytes and fits in the notification itself.
|
||||
That bound is checked at startup: a build that raised the cap past what a
|
||||
notification carries would put a table write and a read-back on the inference
|
||||
path for every request whose prefix changed, so the frontend refuses to start
|
||||
and says so rather than running slowly and quietly. Nothing here drops an
|
||||
observation to stay under the cap.
|
||||
|
||||
### Job and agent streams across replicas
|
||||
|
||||
The same carrier moves the traffic whose subscriber is an open HTTP response rather than a cache: a job's progress stream, its result, its cancel, an agent's SSE events, an agent cancel and an Open Responses cancel.
|
||||
@@ -1462,7 +1503,7 @@ Notes:
|
||||
|---|---|---|
|
||||
| **Discovery** | Automatic via libp2p token | Self-registration to frontend URL |
|
||||
| **State storage** | In-memory / ledger | PostgreSQL |
|
||||
| **Coordination** | Gossip protocol | The worker's own tunnel for serve-backend work; NATS for agent workers and cross-replica frontend events |
|
||||
| **Coordination** | Gossip protocol | The worker's own tunnel for serve-backend work; PostgreSQL `LISTEN`/`NOTIFY` for cross-replica frontend events; NATS for agent workers |
|
||||
| **Node management** | Automatic | REST API + WebUI |
|
||||
| **Health monitoring** | Peer heartbeats | Centralized HealthMonitor |
|
||||
| **Backend management** | Manual per node | Dynamic via the worker's `backend.install` control route |
|
||||
@@ -1538,7 +1579,7 @@ 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 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 cross-frontend sync on the PostgreSQL broadcast carrier. It is purely a routing-layer hint (no backend changes) and never routes worse than today's round-robin.
|
||||
|
||||
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)):
|
||||
|
||||
|
||||
@@ -68,8 +68,20 @@ var _ = Describe("Gallery Distributed", Label("Distributed"), func() {
|
||||
})
|
||||
})
|
||||
|
||||
Context("NATS progress updates", func() {
|
||||
It("should publish progress updates via NATS", func() {
|
||||
// The gallery families ride the broadcast carrier, not NATS.
|
||||
//
|
||||
// These used to publish and subscribe on infra.NC, which asserted that NATS
|
||||
// delivers to itself and nothing about this deployment: they would have
|
||||
// stayed green through the whole migration while the gallery service had
|
||||
// already moved. Two carriers on the deployment's own database is the shape
|
||||
// a fleet has, and it is the shape that fails when one end moves and the
|
||||
// other does not.
|
||||
//
|
||||
// No flush, unlike the NATS version: pgbus.Subscribe has already issued its
|
||||
// LISTEN by the time it returns, and Subscribers() counts only live
|
||||
// handlers, so there is no window to wait out.
|
||||
Context("gallery progress on the broadcast carrier", func() {
|
||||
It("delivers a peer replica's progress updates", func() {
|
||||
op := &distributed.GalleryOperationRecord{
|
||||
GalleryElementName: "whisper-large",
|
||||
OpType: "model_install",
|
||||
@@ -77,35 +89,28 @@ var _ = Describe("Gallery Distributed", Label("Distributed"), func() {
|
||||
}
|
||||
Expect(galleryStore.Create(op)).To(Succeed())
|
||||
|
||||
// Subscribe to gallery progress
|
||||
publisher, subscriber := infra.Bus(), infra.Bus()
|
||||
|
||||
var received atomic.Int32
|
||||
sub, err := infra.NC.Subscribe(messaging.SubjectGalleryProgress(op.ID), func(data []byte) {
|
||||
sub, err := subscriber.Subscribe(messaging.SubjectGalleryProgress(op.ID), func([]byte) {
|
||||
received.Add(1)
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer sub.Unsubscribe()
|
||||
defer func() { Expect(sub.Unsubscribe()).To(Succeed()) }()
|
||||
|
||||
FlushNATS(infra.NC)
|
||||
|
||||
// Publish progress events
|
||||
Expect(infra.NC.Publish(messaging.SubjectGalleryProgress(op.ID), map[string]any{
|
||||
"op_id": op.ID,
|
||||
"progress": 0.25,
|
||||
"message": "25%",
|
||||
Expect(publisher.Publish(messaging.SubjectGalleryProgress(op.ID), map[string]any{
|
||||
"op_id": op.ID, "progress": 0.25, "message": "25%",
|
||||
})).To(Succeed())
|
||||
Expect(publisher.Publish(messaging.SubjectGalleryProgress(op.ID), map[string]any{
|
||||
"op_id": op.ID, "progress": 0.50, "message": "50%",
|
||||
})).To(Succeed())
|
||||
|
||||
Expect(infra.NC.Publish(messaging.SubjectGalleryProgress(op.ID), map[string]any{
|
||||
"op_id": op.ID,
|
||||
"progress": 0.50,
|
||||
"message": "50%",
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() int32 { return received.Load() }, "5s").Should(Equal(int32(2)))
|
||||
Eventually(func() int32 { return received.Load() }, "20s").Should(Equal(int32(2)))
|
||||
})
|
||||
})
|
||||
|
||||
Context("NATS cancel across instances", func() {
|
||||
It("should cancel operation across instances via NATS", func() {
|
||||
Context("gallery cancel on the broadcast carrier", func() {
|
||||
It("delivers a cancel to the replica holding the operation", func() {
|
||||
op := &distributed.GalleryOperationRecord{
|
||||
GalleryElementName: "cancel-model",
|
||||
OpType: "model_install",
|
||||
@@ -114,24 +119,23 @@ var _ = Describe("Gallery Distributed", Label("Distributed"), func() {
|
||||
}
|
||||
Expect(galleryStore.Create(op)).To(Succeed())
|
||||
|
||||
// Simulate another instance listening for cancel
|
||||
publisher, subscriber := infra.Bus(), infra.Bus()
|
||||
|
||||
var cancelReceived atomic.Bool
|
||||
sub, err := infra.NC.Subscribe(messaging.SubjectGalleryCancel(op.ID), func(data []byte) {
|
||||
sub, err := subscriber.Subscribe(messaging.SubjectGalleryCancel(op.ID), func([]byte) {
|
||||
cancelReceived.Store(true)
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer sub.Unsubscribe()
|
||||
defer func() { Expect(sub.Unsubscribe()).To(Succeed()) }()
|
||||
|
||||
FlushNATS(infra.NC)
|
||||
|
||||
// Send cancel from this instance
|
||||
Expect(infra.NC.Publish(messaging.SubjectGalleryCancel(op.ID), map[string]string{
|
||||
Expect(publisher.Publish(messaging.SubjectGalleryCancel(op.ID), map[string]string{
|
||||
"op_id": op.ID,
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() bool { return cancelReceived.Load() }, "5s").Should(BeTrue())
|
||||
Eventually(func() bool { return cancelReceived.Load() }, "20s").Should(BeTrue())
|
||||
|
||||
// Mark cancelled in the store
|
||||
// The row is what survives a replica that was not listening. The
|
||||
// broadcast is the hint to go and look at it.
|
||||
Expect(galleryStore.Cancel(op.ID)).To(Succeed())
|
||||
updated, _ := galleryStore.Get(op.ID)
|
||||
Expect(updated.Status).To(Equal("cancelled"))
|
||||
|
||||
Reference in new issue
Block a user