diff --git a/core/application/cache_fanout_wiring.go b/core/application/cache_fanout_wiring.go new file mode 100644 index 000000000..29330a49b --- /dev/null +++ b/core/application/cache_fanout_wiring.go @@ -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 +} diff --git a/core/application/cache_fanout_wiring_test.go b/core/application/cache_fanout_wiring_test.go new file mode 100644 index 000000000..0cc078230 --- /dev/null +++ b/core/application/cache_fanout_wiring_test.go @@ -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()) + }) +}) diff --git a/core/application/distributed.go b/core/application/distributed.go index 4fcb67a70..ef64417eb 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -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) } diff --git a/core/application/startup.go b/core/application/startup.go index b985d7e55..e2d995171 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -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 diff --git a/core/http/app.go b/core/http/app.go index e736d4cfd..530d149ef 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -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) } } diff --git a/core/http/endpoints/localai/config_meta_test.go b/core/http/endpoints/localai/config_meta_test.go index db7e0fd5f..c90ca9997 100644 --- a/core/http/endpoints/localai/config_meta_test.go +++ b/core/http/endpoints/localai/config_meta_test.go @@ -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)) diff --git a/core/http/endpoints/localai/edit_model_test.go b/core/http/endpoints/localai/edit_model_test.go index 223943e46..31f73e749 100644 --- a/core/http/endpoints/localai/edit_model_test.go +++ b/core/http/endpoints/localai/edit_model_test.go @@ -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)) diff --git a/core/services/galleryop/artifact_progress_coalescer_test.go b/core/services/galleryop/artifact_progress_coalescer_test.go index 9c02b3e9a..f26ed79d3 100644 --- a/core/services/galleryop/artifact_progress_coalescer_test.go +++ b/core/services/galleryop/artifact_progress_coalescer_test.go @@ -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", diff --git a/core/services/galleryop/distributed_pg_test.go b/core/services/galleryop/distributed_pg_test.go new file mode 100644 index 000000000..8c8affc42 --- /dev/null +++ b/core/services/galleryop/distributed_pg_test.go @@ -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") + }) +}) diff --git a/core/services/galleryop/distributed_sync_test.go b/core/services/galleryop/distributed_sync_test.go index ed88ab25b..3572c81c4 100644 --- a/core/services/galleryop/distributed_sync_test.go +++ b/core/services/galleryop/distributed_sync_test.go @@ -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()) }) }) diff --git a/core/services/galleryop/history.go b/core/services/galleryop/history.go index 5617c74aa..76f946607 100644 --- a/core/services/galleryop/history.go +++ b/core/services/galleryop/history.go @@ -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() diff --git a/core/services/galleryop/model_revision_delete_test.go b/core/services/galleryop/model_revision_delete_test.go index 7da910102..3daf9342f 100644 --- a/core/services/galleryop/model_revision_delete_test.go +++ b/core/services/galleryop/model_revision_delete_test.go @@ -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]{ diff --git a/core/services/galleryop/operation.go b/core/services/galleryop/operation.go index 4322b626d..c177e4cd5 100644 --- a/core/services/galleryop/operation.go +++ b/core/services/galleryop/operation.go @@ -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 { diff --git a/core/services/galleryop/service.go b/core/services/galleryop/service.go index 2a7410ab3..edd878885 100644 --- a/core/services/galleryop/service.go +++ b/core/services/galleryop/service.go @@ -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/ 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() } }) diff --git a/core/services/nodes/prefixcache/extractor.go b/core/services/nodes/prefixcache/extractor.go index 65531bafb..e17ed2720 100644 --- a/core/services/nodes/prefixcache/extractor.go +++ b/core/services/nodes/prefixcache/extractor.go @@ -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 diff --git a/core/services/nodes/prefixcache/sync.go b/core/services/nodes/prefixcache/sync.go index 421fea7e2..78d28f708 100644 --- a/core/services/nodes/prefixcache/sync.go +++ b/core/services/nodes/prefixcache/sync.go @@ -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) } } diff --git a/core/services/nodes/prefixcache/sync_test.go b/core/services/nodes/prefixcache/sync_test.go index 001454618..87f09a563 100644 --- a/core/services/nodes/prefixcache/sync_test.go +++ b/core/services/nodes/prefixcache/sync_test.go @@ -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()) + }) +}) diff --git a/core/services/nodes/staging_progress.go b/core/services/nodes/staging_progress.go index 0a6ddc50e..9b11924cd 100644 --- a/core/services/nodes/staging_progress.go +++ b/core/services/nodes/staging_progress.go @@ -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 { diff --git a/core/services/nodes/staging_progress_broadcast_test.go b/core/services/nodes/staging_progress_broadcast_test.go index 0f0f0db1e..77bfa70e5 100644 --- a/core/services/nodes/staging_progress_broadcast_test.go +++ b/core/services/nodes/staging_progress_broadcast_test.go @@ -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)) + }) }) }) diff --git a/core/services/nodes/staging_progress_pg_test.go b/core/services/nodes/staging_progress_pg_test.go new file mode 100644 index 000000000..a1718c290 --- /dev/null +++ b/core/services/nodes/staging_progress_pg_test.go @@ -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")) + }) +}) diff --git a/core/services/pgbus/bus.go b/core/services/pgbus/bus.go index a2c89faf0..1e0ff6dda 100644 --- a/core/services/pgbus/bus.go +++ b/core/services/pgbus/bus.go @@ -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 diff --git a/core/services/pgbus/spill_test.go b/core/services/pgbus/spill_test.go index 801b5f6e5..9bdc851e8 100644 --- a/core/services/pgbus/spill_test.go +++ b/core/services/pgbus/spill_test.go @@ -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) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 3acc30a68..2f50d78cf 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -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_`, 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..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..progress` | Merges the status so `/api/operations` answers the same on any replica | +| Gallery cancel | `gallery..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..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)): diff --git a/tests/e2e/distributed/gallery_distributed_test.go b/tests/e2e/distributed/gallery_distributed_test.go index 100063ed5..b6bf9a6e2 100644 --- a/tests/e2e/distributed/gallery_distributed_test.go +++ b/tests/e2e/distributed/gallery_distributed_test.go @@ -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"))