Files
LocalAI/core/services/galleryop/service.go
T
mudler-agentandEttore Di Giacinto e4fa051ee3 refactor(distributed): put the NATS-only paths behind interfaces (#12395)
* feat(messaging): add shared subject rules

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* test(messaging): cover BroadcastRoots, ControlRoots and SubjectRoot

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat(messaging): add Broadcaster and enforce subject rules in every carrier

Broadcaster is the fan-out half of MessagingClient. The NATS client and
the in-memory FakeBus now refuse a subject outside the served roots and any
wildcard other than a whole single token, and FakeBus shares MatchSubject
instead of its own copy. FakeBus Unsubscribe now removes its own
subscription instead of the first one with the same subject.

A shared conformance suite in messagingtest runs against both carriers.
The distributed e2e specs that used invented test.* subjects, and the one
that subscribed with a > filter, now use subjects from subjects.go.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: depend on Broadcaster where only publish and subscribe are used

Narrowed to messaging.Broadcaster: nodes/staging_progress.go,
nodes/install_progress_publisher.go, galleryop/operation.go,
galleryop/service.go, agentpool/user_services.go, agentpool/agent_jobs.go,
openresponses/store.go, openresponses/sync.go, syncstate/syncstate.go,
finetune/service.go, quantization/service.go and
failover/distsync/distsync.go. SubscribeJSON now takes a Broadcaster
because it only calls Subscribe, which lets the narrowed consumers use it.

Stayed wide: worker/supervisor.go, because its client field also serves
the SubscribeReply handlers in worker/lifecycle.go. The request/reply,
queue and wiring files (nodes/unloader.go, nodes/file_stager_s3.go,
jobs/dispatcher.go, agents/dispatcher.go, agents/events.go,
worker/file_staging.go, cli/agent_worker.go, http/app.go) are unchanged by
design.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): name the no-route condition and confine the carrier error

Consumers matched nats.ErrNoResponders, which names an absence, to demote a node. They now match ErrNoRoute, the control path maps the carrier's failure onto it, and timeouts and worker refusals are pinned as not being no-route.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs(nodes): state which FileStager implementations return ErrNoRoute

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): build backend clients through one node-aware seam

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs: describe the distributed transport seams

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs: correct comments that overclaim after the seams refactor

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(agent-worker): refuse an unserved LOCALAI_AGENT_SUBJECT at startup

The messaging client now refuses a subject whose root no carrier serves.
An agent worker started with a custom LOCALAI_AGENT_SUBJECT such as
tenant-a.agent.execute used to start and then wait on a subject the
frontend never publishes to. After the subject rules landed it exited
at subscribe time with an error that did not name the setting.

Behaviour change: the worker now checks LOCALAI_AGENT_SUBJECT before it
registers or connects, and exits with an error that names the variable
and says to use a served subject under the agent root, for example
agent.execute. The served roots are not widened: a custom root was
never delivered by the frontend, and a wider set would reopen the
drift the subject rules exist to close. The flag help and the agent
worker docs state the constraint.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* test(nodes): pin the reactions to ErrNoRoute

Three callers react to ErrNoRoute and had no spec: the reconciler's
upgrade drain falls back to the legacy forced install, the reconciler
marks the node unhealthy when a pending op has no route, and the
backend-op fan-out marks the node unhealthy. Each spec drives the real
caller with a scripted no-responders reply and reads the result from
the registry or the recorded requests.

A fourth spec pins the other side: a pending op that times out leaves
the node healthy and only counts the attempt, so mapping timeouts onto
ErrNoRoute would fail here.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* test(messaging): pin client subject checks and fail the carrier suite in CI

Add specs that call Publish, Request, Subscribe, QueueSubscribe,
SubscribeReply and QueueSubscribeReply on a client with no connection.
Each call must return ErrUnservedSubject for bogus.thing and
ErrUnsupportedWildcard for jobs.>. This proves that the subject check
runs before the connection is used, and needs no server.

The NATS conformance suite is the only check that runs the subject rules
against a real carrier. Before this change it skipped without output
when Docker was missing. Now it fails when CI is set, so a Linux runner
without Docker cannot hide it. It still skips on local runs and on macOS
CI, which has no Docker.

Add SubjectNodeBackendInstallProgress to the list of constructors that
must build served subjects, and ask contributors to extend the list.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs: state what ErrNoRoute may change, and group the distributed guides

The seams note said MarkUnhealthy was the only state change allowed
on ErrNoRoute. A pending backend op still records the failed attempt,
counts toward the reconciler's retry limit and is dead-lettered after
the maximum attempts. The note now says that MarkUnhealthy is the only
change to the node's own state, and that the per-op accounting is not
a verdict about the node.

The note also documents that the NATS conformance run fails under CI
when Docker is missing. The distributed-seams row moves next to the
distributed-state row in the topics table. The liveness ping spec
header now says no route is a reason to skip the worker, not proof
that the worker is gone.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): give the backend client factory the node id

Mechanical: the method gains a nodeID parameter and the eight test fakes are updated. No behaviour change.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): drop the optional node-aware factory

The node id is now in the main method, so the optional interface and its helper had no behaviour of their own.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): dial backend probes through the client factory

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(nodes): dial workers' file servers through a per-node dialer

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(http): proxy backend logs through the per-node worker dialer

The admin backend-logs proxy (list, lines and the WebSocket stream) now reaches a worker through the same per-node dialer as the HTTP file stager, so every frontend-to-worker dial goes through one seam. The shared direct dialer keeps alive for 15s where the proxy used 30s. Harmless for requests bounded at 15s.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(http): keep the backend-logs proxy independent of the admin connection

The proxy request had no context before the dialer change and is bounded only by its 15s timeout. Keep it that way.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: move the worker control payloads to workerctl

Mechanical move of the request and reply structs, the install progress event and the file payloads out of messaging. The verbs no longer belong to one carrier. No alias is left behind.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(worker): serve the lifecycle verbs through a controlServer

The worker registers one handler per verb and a NATS server maps each verb to its subject. Registration errors now name the verb. node.stop is served with SubscribeReply, which is identical on the wire because the handler never replies.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(worker): report install progress through the control sink

Install and upgrade now emit download progress through the sink the control server hands them. The debounce and the terminal flush stay in the handler path, built over that sink by the new nodes.NewDebouncedInstallProgressSink, which replaces NewDebouncedInstallProgressPublisher. The subject and payload on the wire are unchanged. The supervisor no longer holds the bus, and installFn and upgradeFn let specs drive both verbs without a gallery.

The malformed-request log lines are restored for install, upgrade, backend.delete, model.unload, model.stop and model.delete, with the reply bytes unchanged. The signal adapter is renamed noReply, which also lets worker.go import os/signal without an alias again.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(worker): serve the file-staging verbs through a controlServer

An empty list-dir answer is now {} rather than {"files":null}, because the typed reply omits an empty Files slice. The frontend decodes both to a nil slice in nodes/file_stager_s3.go.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat(messaging): add WorkQueue and the NATS producer

This is the producer side of the competing-consumer seam. The work kinds map one to one to today's subjects and queue groups: task to jobs.new and mcp-ci to jobs.mcp-ci.new (both in group workers), agent-run to agent.execute (group agent-workers). Enqueue publishes the payload as Publish does today, with one JSON marshal. FakeBus now records queue groups and keeps reply handlers so later specs can pin and drive them.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* feat(messaging): add the NATS WorkConsumer

An in-flight limit of one runs the handler inline on the delivery goroutine, as the MCP CI consumer does today. Any other limit spawns per delivery, as the agent consumer does. Queue groups are unchanged.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: publish queued work through WorkQueue

The job dispatcher, the agent pool and the agent scheduler enqueue through messaging.WorkQueue; the NATS implementation publishes to the same subjects as before. DistributedServices builds the queue next to the NATS client and hands it to the dispatcher and the agent pool, whose distributed mode switch now reads a non-nil WorkQueue. The unused AgentPoolService.SetNATSClient is removed.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: consume queued work through WorkConsumer

The agent dispatcher and the MCP CI consumer register through messaging.WorkConsumer. The NATS implementation keeps the inline one-at-a-time model for MCP CI and the per-delivery model for agent runs. handleMCPCIJob reports on the events publisher the carrier hands it instead of a captured client.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: delete the consumers nothing in production reached

jobs.new has a producer and no production consumer, and the agent dispatcher's Dispatch was only called from tests. Publishing jobs.new is unchanged.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(mcp): send MCP requests to agent workers through AgentControl

Timeouts still honour only the deadline, not cancellation, exactly as today. The NATS no-responders error maps to ErrNoRoute and a timeout does not.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(agent-worker): serve MCP requests and backend.stop through agentRPCServer

The agent worker's MCP tool and discovery reply subscriptions and its backend stop listener move behind an unexported agentRPCServer interface, served on NATS by nodes.NATSAgentRPCServer. The handlers become typed mcp.ToolHandler and mcp.DiscoveryHandler values that answer every failure with a reply carrying Error.

Queue group (agent-workers), inline execution on the delivery goroutine, the background handler context, the unmarshal error reply texts and the reply-less backend stop subscription are unchanged. The backend stop handler takes the decoded backend name, so it can still close that backend's MCP sessions.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor(messaging): remove helpers that only tests used

BroadcastRoots, ControlRoots and SubjectRoot had no production caller. The roots spec now asserts every served root through ValidateSubject instead. MatchSubject moves back into the test support package, the only place that used it, with its table. NATSAgentRPCServer drops the subscription list it stored and never read, and NewNATSAgentRPCServer gets a doc comment.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* test(mcp): round trip the agent RPC server over a real NATS server

One spec sends a tool request and a discovery request through NATSAgentControl to NATSAgentRPCServer and checks that the handlers see the decoded requests and the replies come back. It also puts an undecodable body on the tool subject and checks the server answers with an unmarshal error instead of leaving the requester to time out.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs: describe the distributed transport seams

The developer note now lists the final seams: fan-out, queues, both halves of the control verbs and of agent RPC, and the dial. It records the open items a second carrier has to handle.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* test: pin the in-flight limit each queue consumer asks for

The work queue specs pin what Consume does for a given limit, but nothing
pinned which limit each production consumer passes. Changing the agent
worker's MCP CI limit from 1 to 0 would have let MCP CI jobs run
concurrently on each worker with every test green.

Move the MCP CI Consume call into startMCPCIConsumer with the same wiring
and pin that it asks for (WorkMCPCI, 1). Pin that NATSDispatcher.Start asks
for (WorkAgentRun, maxConcurrent) for several limits.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* refactor: remove helpers the branch left without a caller

SubjectJobCancelWildcard lost its last subscriber when the frontend
stopped listening on jobs.*.cancel; the NATS permissions and conformance
suite spell the subject out, so nothing reads the constant.

decodeBackendStopRequest returned a stopAll flag that production dropped
and only a test read. decodeBackendStop is now the single decoder with the
same semantics: an empty body is stop-all, an empty Backend is stop-all,
malformed JSON is an error. stopBackends still derives stop-all from
Backend, so no reply changes.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fix(messaging): keep an explicitly empty agent queue a plain subscription

Before the work queue seam the agent worker passed LOCALAI_AGENT_QUEUE
straight to QueueSubscribe, so an explicitly empty value made a plain
subscription and every agent worker ran every agent run.
WithAgentRunRoute replaced an empty queue with agent-workers, which
silently changed that.

Keep the queue as given once the option is applied. An empty subject still
falls back to agent.execute, since it never had a meaning of its own. The
flag default stays agent-workers, so only an explicitly empty value
reaches this.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* docs: correct comments and record the PR B notes

Fix the recordingFactory comment (it also records the parallel flag),
document that a negative maxInFlight is unbounded and that Unsubscribe from
a handler deadlocks, and say a permanently undecodable payload returns nil.

Record controlHandler's undecodable return as a kept exception, and add the
second carrier notes to the developer note: the reconciler has no
ClientFactory option, the logs proxy honours HTTP_PROXY, verbs one carrier
serves need an opt-out, terminal replies come from the result event, and
agent runs publish through the NATS-bound EventBridge, which is not an
additive change.

Assisted-by: Claude:claude-sonnet-5-5
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
2026-10-02 23:58:41 +02:00

839 lines
30 KiB
Go

package galleryop
import (
"context"
"errors"
"fmt"
"runtime/debug"
"sync"
"time"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/core/gallery"
"github.com/mudler/LocalAI/core/services/distributed"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/pkg/downloader"
"github.com/mudler/LocalAI/pkg/model"
"github.com/mudler/LocalAI/pkg/system"
"github.com/mudler/xlog"
)
type GalleryService struct {
appConfig *config.ApplicationConfig
sync.Mutex
ModelGalleryChannel chan ManagementOp[gallery.GalleryModel, gallery.ModelConfig]
BackendGalleryChannel chan ManagementOp[gallery.GalleryBackend, any]
modelLoader *model.ModelLoader
modelManager ModelManager
backendManager BackendManager
statuses map[string]*OpStatus
cancellations map[string]cancellationActions
// Distributed mode (nil when not in distributed mode).
// natsClient is a messaging.Broadcaster (Publisher + Subscribe)
// when wired by the distributed startup path; broadcastSubs holds the
// progress + cancel subscriptions opened by SubscribeBroadcasts.
natsClient messaging.Broadcaster
galleryStore *distributed.GalleryStore
broadcastSubs []messaging.Subscription
// OnBackendOpCompleted is fired after every successful install/upgrade/delete
// on the backend channel. The Application wires this to UpgradeChecker.TriggerCheck
// so `/api/backends/upgrades` stops surfacing a backend as upgradeable the moment
// the worker finishes — previously the cache only refreshed on the 6-hour tick,
// making manual upgrades look like they failed even when they hadn't.
//
// In distributed mode the same hook also fires on peer replicas via the
// SubjectCacheInvalidateBackends subscriber, so every replica's
// UpgradeChecker stays in sync.
OnBackendOpCompleted func()
// OnModelsChanged is fired on peer replicas when SOMEONE else publishes
// SubjectCacheInvalidateModels. The Application wires this to
// ModelConfigLoader.LoadModelConfigsFromPath so a chat completion that
// load-balances onto this replica can find the just-installed model.
// The originating replica reloads inline (models.go) so it does not need
// the hook.
OnModelsChanged func(messaging.CacheInvalidateEvent)
modelRevisionLifecycle interface {
ApplyConfigRevisions(context.Context, []config.ModelConfigRevisionTransition) (int, error)
}
}
// SetModelRevisionLifecycle wires the distributed config-generation boundary
// into gallery deletion without coupling gallery operations to node internals.
func (g *GalleryService) SetModelRevisionLifecycle(lifecycle interface {
ApplyConfigRevisions(context.Context, []config.ModelConfigRevisionTransition) (int, error)
}) {
g.Lock()
defer g.Unlock()
g.modelRevisionLifecycle = lifecycle
}
func NewGalleryService(appConfig *config.ApplicationConfig, ml *model.ModelLoader) *GalleryService {
return &GalleryService{
appConfig: appConfig,
ModelGalleryChannel: make(chan ManagementOp[gallery.GalleryModel, gallery.ModelConfig]),
BackendGalleryChannel: make(chan ManagementOp[gallery.GalleryBackend, any]),
modelLoader: ml,
modelManager: NewLocalModelManager(appConfig, ml),
backendManager: NewLocalBackendManager(appConfig, ml),
statuses: make(map[string]*OpStatus),
cancellations: make(map[string]cancellationActions),
}
}
// SetModelManager replaces the model manager (e.g. with a distributed implementation).
func (g *GalleryService) SetModelManager(m ModelManager) {
g.Lock()
defer g.Unlock()
g.modelManager = m
}
// SetBackendManager replaces the backend manager (e.g. with a distributed implementation).
func (g *GalleryService) SetBackendManager(b BackendManager) {
g.Lock()
defer g.Unlock()
g.backendManager = b
}
// BackendManager returns the current backend manager. Callers like the
// periodic upgrade checker need this so they run CheckUpgrades through the
// distributed implementation (which asks workers) instead of the frontend's
// local filesystem — the latter is always empty in distributed deployments.
func (g *GalleryService) BackendManager() BackendManager {
g.Lock()
defer g.Unlock()
return g.backendManager
}
// ModelArtifactMaterializer returns the controller-only acquisition capability
// used by startup paths that install gallery entries outside the operation loop.
func (g *GalleryService) ModelArtifactMaterializer() config.ArtifactMaterializer {
if g == nil || g.appConfig == nil {
return nil
}
return g.appConfig.ModelArtifactMaterializer
}
// SetNATSClient sets the NATS client for distributed progress publishing.
// Accepting a Broadcaster (vs. plain Publisher) lets
// SubscribeBroadcasts wire the wildcard subscriptions that keep peer
// replicas' statuses + cancellations in sync.
func (g *GalleryService) SetNATSClient(nc messaging.Broadcaster) {
g.Lock()
defer g.Unlock()
g.natsClient = nc
}
// SetGalleryStore sets the PostgreSQL gallery store for distributed persistence.
func (g *GalleryService) SetGalleryStore(s *distributed.GalleryStore) {
g.Lock()
defer g.Unlock()
g.galleryStore = s
}
// ListBackends returns installed backends via the backend manager.
// In standalone mode this checks the local filesystem; in distributed mode
// it aggregates from all healthy worker nodes.
func (g *GalleryService) ListBackends() (gallery.SystemBackends, error) {
g.Lock()
mgr := g.backendManager
g.Unlock()
return mgr.ListBackends()
}
// DeleteBackend delegates backend deletion to the backend manager, which in distributed
// mode fans out the deletion to worker nodes via NATS.
func (g *GalleryService) DeleteBackend(name string) error {
g.Lock()
mgr := g.backendManager
g.Unlock()
return mgr.DeleteBackend(name)
}
func (g *GalleryService) UpdateStatus(s string, op *OpStatus) {
g.Lock()
// Preserve any per-node entries already accumulated by UpdateNodeProgress:
// the legacy progressCb path (used by the Phase 2 install bridge) calls
// UpdateStatus with a fresh *OpStatus on every tick, which would otherwise
// wipe the Nodes slice and leave the UI flickering between one node and
// another. If the caller explicitly populates Nodes on the incoming op,
// that wins; an empty Nodes slice on the incoming op is treated as "no
// new per-node data" and the previous Nodes are carried forward.
if op != nil {
if prev := g.statuses[s]; prev != nil {
if len(op.Nodes) == 0 && len(prev.Nodes) > 0 {
op.Nodes = prev.Nodes
}
// A job is a delete or an install for its whole life. markQueued is
// the only writer that knows which; every later status omits the
// flag, so an unset value means "no new information", not "this is
// an install".
if !op.Deletion {
op.Deletion = prev.Deletion
}
}
}
g.statuses[s] = op
store := g.galleryStore
nc := g.natsClient
g.Unlock()
// I/O happens after Unlock. The NATS 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.
if store != nil && op != nil {
if op.Processed {
status, errMsg := "completed", ""
if op.Error != nil {
status = "failed"
errMsg = op.Error.Error()
}
if op.Cancelled {
status = "cancelled"
}
if err := store.UpdateStatus(s, status, errMsg); err != nil {
xlog.Warn("Failed to persist gallery operation status", "op_id", s, "error", err)
}
} else {
if err := store.UpdateProgress(s, op.Progress, op.Message, op.DownloadedFileSize, op.Cancellable,
distributed.OperationProgressDetails{
Phase: op.Phase, CurrentBytes: op.CurrentBytes, TotalBytes: op.TotalBytes,
}); err != nil {
xlog.Warn("Failed to persist gallery operation progress", "op_id", s, "error", err)
}
}
}
// 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.
if nc != nil {
if err := nc.Publish(messaging.SubjectGalleryProgress(s), GalleryProgressEvent{
JobID: s,
Status: op,
}); err != nil {
xlog.Warn("Failed to broadcast gallery progress", "op_id", s, "error", err)
}
}
}
// 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).
func (g *GalleryService) publishCacheInvalidate(subject string, evt messaging.CacheInvalidateEvent) {
g.Lock()
nc := g.natsClient
g.Unlock()
if nc == nil {
return
}
if err := nc.Publish(subject, evt); err != nil {
xlog.Warn("Failed to broadcast cache invalidation", "subject", subject, "error", err)
}
}
// BroadcastModelsChanged notifies peer replicas that a model config was
// created, edited, or removed out-of-band of the gallery install/delete
// channel (e.g. the admin /models/edit, /models/import and
// /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).
//
// 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,
// which a reload-from-path cannot do because the loader is additive).
func (g *GalleryService) BroadcastModelsChanged(element, op string) {
g.BroadcastModelsChangedRevision(element, op, "")
}
// BroadcastModelsChangedRevision includes the accepted semantic generation so
// peers can apply the same registry transition idempotently.
func (g *GalleryService) BroadcastModelsChangedRevision(element, op, configRevision string) {
g.publishCacheInvalidate(messaging.SubjectCacheInvalidateModels, messaging.CacheInvalidateEvent{
Element: element,
Op: op,
ConfigRevision: configRevision,
})
}
// 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
// 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
// silently re-applies → no second publish).
func (g *GalleryService) mergeStatus(opID string, op *OpStatus) {
if op == nil {
return
}
g.Lock()
defer g.Unlock()
if len(op.Nodes) == 0 {
if prev := g.statuses[opID]; prev != nil && len(prev.Nodes) > 0 {
op.Nodes = prev.Nodes
}
}
g.statuses[opID] = op
}
// UpdateNodeProgress merges a per-node progress tick into OpStatus.Nodes,
// keyed by nodeID, and mirrors the latest values into the aggregate
// Progress / FileName / DownloadedFileSize / TotalFileSize / Message
// fields so the legacy single-bar OperationsBar view keeps working
// unchanged alongside the new per-node breakdown.
//
// We deliberately do NOT delegate the aggregate mirror to UpdateStatus
// here: UpdateStatus overwrites the entire OpStatus, which would clobber
// the Nodes slice we just merged into. Doing the merge + mirror under a
// single lock keeps both views consistent and concurrent-safe.
func (g *GalleryService) UpdateNodeProgress(opID, nodeID string, np NodeProgress) {
g.Lock()
defer g.Unlock()
status := g.statuses[opID]
if status == nil {
status = &OpStatus{}
g.statuses[opID] = status
}
merged := false
for i := range status.Nodes {
if status.Nodes[i].NodeID == nodeID {
status.Nodes[i] = np
merged = true
break
}
}
if !merged {
status.Nodes = append(status.Nodes, np)
}
// Mirror the latest tick into the legacy aggregate fields so the
// existing single-bar UI keeps rendering meaningful progress.
status.FileName = np.FileName
status.Progress = np.Percentage
status.DownloadedFileSize = np.Current
status.TotalFileSize = np.Total
if np.Phase != "" {
status.Message = np.Phase
}
}
func (g *GalleryService) GetStatus(s string) *OpStatus {
g.Lock()
defer g.Unlock()
return g.statuses[s]
}
func (g *GalleryService) GetAllStatus() map[string]*OpStatus {
g.Lock()
defer g.Unlock()
return g.statuses
}
// ReapStaleOperations marks abandoned in-progress operations (pending/
// downloading/processing) older than `age` as failed, so an op orphaned by a
// replica that died mid-flight does not linger as "processing" forever. The
// store's CleanStale runs once on startup; this exposes it for periodic
// invocation (a post-startup orphan is otherwise not reaped until the next
// restart). No-op when no gallery store is wired. Returns rows reaped.
func (g *GalleryService) ReapStaleOperations(age time.Duration) (int64, error) {
g.Lock()
store := g.galleryStore
g.Unlock()
if store == nil {
return 0, nil
}
// Collect the IDs before the update: once CleanStale flips them to
// "failed" they no longer match the stale predicate.
staleIDs, err := store.ListStale(age)
if err != nil {
xlog.Warn("Failed to list stale gallery operations", "error", err)
}
n, err := store.CleanStale(age)
if err != nil {
return 0, err
}
if n > 0 {
xlog.Info("Reaped stale gallery operations", "count", n)
}
// The database row is only half the picture. GET /models/jobs/<id> and
// /api/operations read the in-memory statuses map, which is populated
// locally and via the NATS progress broadcast and never expires. An op
// 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
// on the row. Reconcile the in-memory copy with the reap.
for _, id := range staleIDs {
g.failStaleStatus(id)
}
return n, nil
}
// failStaleStatus flips a locally-cached in-progress status to a terminal
// failure after its store row was reaped. Statuses that already reached a
// terminal state are left alone so a genuine completion or cancellation that
// raced the reaper is not rewritten as a failure.
func (g *GalleryService) failStaleStatus(id string) {
g.Lock()
st, ok := g.statuses[id]
if !ok || st == nil || st.Processed {
g.Unlock()
return
}
elementName := st.GalleryElementName
g.Unlock()
xlog.Warn("Marking orphaned gallery operation as failed", "op_id", id, "element", elementName)
g.UpdateStatus(id, &OpStatus{
Processed: true,
Error: errors.New("stale operation reaped (abandoned by a crashed or restarted instance)"),
Message: "error: stale operation reaped (abandoned by a crashed or restarted instance)",
GalleryElementName: elementName,
Cancellable: false,
})
}
// CancelOperation cancels an in-progress operation by its ID.
//
// In distributed mode the UI's cancel click may land on a different replica
// than the one running the operation. We still publish the cancel event in
// that case — the peer holding the cancellation func picks it up via the
// SubjectGalleryCancelWildcard subscriber and runs it locally. The caller
// gets a non-error reply so the UI shows the cancel as accepted.
func (g *GalleryService) CancelOperation(id string) error {
return g.stopOperation(id, false)
}
// PauseOperation stops an in-progress download while preserving its partial
// file. Re-submitting the same install resumes it through the downloader's
// existing HTTP Range support.
func (g *GalleryService) PauseOperation(id string) error {
return g.stopOperation(id, true)
}
func (g *GalleryService) stopOperation(id string, pause bool) error {
g.Lock()
if status, ok := g.statuses[id]; ok && status.Cancelled {
g.Unlock()
return fmt.Errorf("operation %q is already cancelled", id)
}
actions, localExists := g.cancellations[id]
if localExists {
delete(g.cancellations, id)
}
nc := g.natsClient
store := g.galleryStore
if !localExists && nc == nil {
g.Unlock()
return fmt.Errorf("operation %q not found or already completed", id)
}
if status, ok := g.statuses[id]; ok {
status.Cancelled = true
status.Processed = true
status.Message = map[bool]string{true: "paused", false: "cancelled"}[pause]
} else {
g.statuses[id] = &OpStatus{
Cancelled: true,
Processed: true,
Message: map[bool]string{true: "paused", false: "cancelled"}[pause],
Cancellable: false,
}
}
g.Unlock()
// Persist the terminal status so the cancel survives a restart. Without
// this the row stays in its active state and re-hydrates straight back into
// processingBackends on the next replica boot — the UI spins again on an op
// the admin already cancelled. The peer that broadcasts wins the write; a
// no-op when standalone (store nil).
if store != nil {
if err := store.Cancel(id); err != nil {
xlog.Warn("Failed to persist gallery operation cancellation", "op_id", id, "error", err)
}
}
// I/O and user-provided callback after Unlock — the cancel-wildcard
// subscriber loops back into applyCancel on this same replica, which
// would otherwise deadlock on g.Mutex.
stopFunc := actions.cancel
if pause {
stopFunc = actions.pause
}
if stopFunc != nil {
stopFunc()
}
if nc != nil {
if err := nc.Publish(messaging.SubjectGalleryCancel(id), GalleryCancelEvent{JobID: id, Pause: pause}); err != nil {
xlog.Warn("Failed to broadcast gallery cancel", "op_id", id, "error", err)
}
}
return nil
}
// 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
// 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) {
g.Lock()
actions, hasCancel := g.cancellations[id]
if hasCancel {
delete(g.cancellations, id)
}
if status, ok := g.statuses[id]; ok {
if status.Cancelled {
g.Unlock()
return
}
status.Cancelled = true
status.Processed = true
status.Message = map[bool]string{true: "paused", false: "cancelled"}[pause]
} else {
g.statuses[id] = &OpStatus{
Cancelled: true,
Processed: true,
Message: map[bool]string{true: "paused", false: "cancelled"}[pause],
Cancellable: false,
}
}
g.Unlock()
// Invoke the cancel func after Unlock so a callback that touches
// GalleryService doesn't re-enter the mutex.
stopFunc := actions.cancel
if pause {
stopFunc = actions.pause
}
if hasCancel && stopFunc != nil {
stopFunc()
}
}
// NewUserCancellableContext returns a child context whose CancelFunc cancels
// with the downloader.ErrUserCancelled cause. This lets the download layer
// distinguish a deliberate user cancel (discard the half-downloaded .partial)
// from an incidental cancellation such as process shutdown (keep the .partial
// so the next run resumes via Range instead of restarting from zero).
// NewUserCancellableContext creates distinct callbacks for destructive cancel
// and resume-safe pause while sharing one operation context.
func NewUserCancellableContext(parent context.Context) (context.Context, context.CancelFunc, context.CancelFunc) {
ctx, cancelCause := context.WithCancelCause(parent)
return ctx,
func() { cancelCause(downloader.ErrUserCancelled) },
func() { cancelCause(context.Canceled) }
}
type cancellationActions struct {
cancel context.CancelFunc
pause context.CancelFunc
}
func (g *GalleryService) storeCancellation(id string, cancelFunc, pauseFunc context.CancelFunc) {
g.Lock()
defer g.Unlock()
if pauseFunc == nil {
pauseFunc = cancelFunc
}
g.cancellations[id] = cancellationActions{cancel: cancelFunc, pause: pauseFunc}
}
// StoreCancellation is a public method to store a cancellation function for an operation
// This allows cancellation functions to be stored immediately when operations are created,
// enabling cancellation of queued operations that haven't started processing yet.
func (g *GalleryService) StoreCancellation(id string, cancelFunc context.CancelFunc) {
g.storeCancellation(id, cancelFunc, cancelFunc)
}
// StoreCancellationActions registers distinct destructive-cancel and
// resume-safe pause callbacks for an operation.
func (g *GalleryService) StoreCancellationActions(id string, cancelFunc, pauseFunc context.CancelFunc) {
g.storeCancellation(id, cancelFunc, pauseFunc)
}
// removeCancellation removes a cancellation function when operation completes
func (g *GalleryService) removeCancellation(id string) {
g.Lock()
defer g.Unlock()
delete(g.cancellations, id)
}
// runOpHandler runs one operation handler and converts a panic into an error.
//
// The gallery worker is a single goroutine consuming both channels serially. A
// panic anywhere in an install handler (a malformed gallery entry, a nil
// dereference in a backend-specific path) took down the entire process with it,
// and every queued operation went with it. Containing the panic to the
// operation that caused it keeps the consumer alive so subsequent operations
// are still picked up, and surfaces the failure on the op itself instead of as
// an unexplained restart.
func runOpHandler(fn func() error) (err error) {
defer func() {
if r := recover(); r != nil {
xlog.Error("Gallery operation handler panicked", "panic", r, "stack", string(debug.Stack()))
err = fmt.Errorf("gallery operation handler panicked: %v", r)
}
}()
return fn()
}
func (g *GalleryService) Start(c context.Context, cl *config.ModelConfigLoader, systemState *system.SystemState) error {
// updates the status with an error
var updateError func(id string, e error)
if !g.appConfig.OpaqueErrors {
updateError = func(id string, e error) {
g.UpdateStatus(id, &OpStatus{Error: e, Processed: true, Message: "error: " + e.Error()})
}
} else {
updateError = func(id string, _ error) {
g.UpdateStatus(id, &OpStatus{Error: fmt.Errorf("an error occurred"), Processed: true})
}
}
go func() {
for {
select {
case <-c.Done():
return
case op := <-g.BackendGalleryChannel:
// Create context if not provided
if op.Context == nil {
op.Context, op.CancelFunc, op.PauseFunc = NewUserCancellableContext(c)
g.storeCancellation(op.ID, op.CancelFunc, op.PauseFunc)
} else if op.CancelFunc != nil {
g.storeCancellation(op.ID, op.CancelFunc, op.PauseFunc)
}
// Create DB record for distributed tracking
if g.galleryStore != nil {
opType := "backend_install"
if op.Delete {
opType = "backend_delete"
}
if err := g.galleryStore.Create(&distributed.GalleryOperationRecord{
ID: op.ID,
GalleryElementName: op.GalleryElementName,
OpType: opType,
Status: "pending",
// Create runs at dequeue, so this is the running-phase
// value: a running delete is not cancellable, an install
// is, matching the model channel. The queued phase before
// this point is cancellable either way (see markQueued).
Cancellable: !op.Delete,
}); err != nil {
// Not fatal: the install still runs and the in-memory
// status still updates. Logged because without the row
// the cross-replica dedup guard and hydration cannot
// see this operation at all.
xlog.Warn("Failed to create gallery operation record", "op_id", op.ID, "error", err)
}
}
err := runOpHandler(func() error { return g.backendHandler(&op, systemState) })
if err != nil {
updateError(op.ID, err)
} else if g.OnBackendOpCompleted != nil {
// Let listeners (e.g. UpgradeChecker) refresh their view of
// installed state. Run off the worker goroutine so a slow
// callback doesn't stall the next queued operation.
go g.OnBackendOpCompleted()
}
g.removeCancellation(op.ID)
case op := <-g.ModelGalleryChannel:
// Create context if not provided
if op.Context == nil {
op.Context, op.CancelFunc, op.PauseFunc = NewUserCancellableContext(c)
g.storeCancellation(op.ID, op.CancelFunc, op.PauseFunc)
} else if op.CancelFunc != nil {
g.storeCancellation(op.ID, op.CancelFunc, op.PauseFunc)
}
// Create DB record for distributed tracking
if g.galleryStore != nil {
opType := "model_install"
if op.Delete {
opType = "model_delete"
}
if err := g.galleryStore.Create(&distributed.GalleryOperationRecord{
ID: op.ID,
GalleryElementName: op.GalleryElementName,
OpType: opType,
Status: "pending",
// Create runs at dequeue, so this is the running-phase
// value: a running delete is not cancellable, an install
// is. The queued phase before this point is cancellable
// either way (see markQueued).
Cancellable: !op.Delete,
}); err != nil {
xlog.Warn("Failed to create gallery operation record", "op_id", op.ID, "error", err)
}
}
err := runOpHandler(func() error { return g.modelHandler(&op, cl, systemState) })
if err != nil {
updateError(op.ID, err)
}
g.removeCancellation(op.ID)
}
}
}()
return nil
}
// SubscribeBroadcasts opens the wildcard subscriptions that keep this
// replica's in-memory statuses + cancellation state in sync with peers.
// Returns an error if the progress subscription fails; cancel-sub failures
// are not fatal but are logged.
//
// Hydrate should be called before this so the freshly-started replica has
// the pre-existing operations before live updates start flowing.
func (g *GalleryService) SubscribeBroadcasts() error {
g.Lock()
nc := g.natsClient
g.Unlock()
if nc == nil {
return nil
}
progressSub, err := messaging.SubscribeJSON(nc, messaging.SubjectGalleryProgressWildcard, func(evt GalleryProgressEvent) {
if evt.JobID == "" || evt.Status == nil {
return
}
g.mergeStatus(evt.JobID, evt.Status)
})
if err != nil {
return fmt.Errorf("subscribing to gallery progress wildcard: %w", err)
}
cancelSub, err := messaging.SubscribeJSON(nc, messaging.SubjectGalleryCancelWildcard, func(evt GalleryCancelEvent) {
if evt.JobID == "" {
return
}
g.applyCancel(evt.JobID, evt.Pause)
})
if err != nil {
if uerr := progressSub.Unsubscribe(); uerr != nil {
xlog.Warn("failed to unsubscribe partial gallery progress sub", "error", uerr)
}
return fmt.Errorf("subscribing to gallery cancel wildcard: %w", err)
}
modelsSub, err := messaging.SubscribeJSON(nc, messaging.SubjectCacheInvalidateModels, func(evt messaging.CacheInvalidateEvent) {
g.Lock()
cb := g.OnModelsChanged
g.Unlock()
if cb != nil {
cb(evt)
}
})
if err != nil {
if uerr := progressSub.Unsubscribe(); uerr != nil {
xlog.Warn("failed to unsubscribe partial gallery progress sub", "error", uerr)
}
if uerr := cancelSub.Unsubscribe(); uerr != nil {
xlog.Warn("failed to unsubscribe partial gallery cancel sub", "error", uerr)
}
return fmt.Errorf("subscribing to models invalidation: %w", err)
}
backendsSub, err := messaging.SubscribeJSON(nc, messaging.SubjectCacheInvalidateBackends, func(_ messaging.CacheInvalidateEvent) {
g.Lock()
cb := g.OnBackendOpCompleted
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.
go cb()
}
})
if err != nil {
if uerr := progressSub.Unsubscribe(); uerr != nil {
xlog.Warn("failed to unsubscribe partial gallery progress sub", "error", uerr)
}
if uerr := cancelSub.Unsubscribe(); uerr != nil {
xlog.Warn("failed to unsubscribe partial gallery cancel sub", "error", uerr)
}
if uerr := modelsSub.Unsubscribe(); uerr != nil {
xlog.Warn("failed to unsubscribe partial models sub", "error", uerr)
}
return fmt.Errorf("subscribing to backends invalidation: %w", err)
}
g.Lock()
g.broadcastSubs = append(g.broadcastSubs, progressSub, cancelSub, modelsSub, backendsSub)
g.Unlock()
return nil
}
// CloseBroadcasts drops the wildcard subscriptions. Safe to call multiple times.
func (g *GalleryService) CloseBroadcasts() {
g.Lock()
subs := g.broadcastSubs
g.broadcastSubs = nil
g.Unlock()
for _, s := range subs {
if err := s.Unsubscribe(); err != nil {
xlog.Warn("GalleryService unsubscribe failed", "error", err)
}
}
}
// Hydrate loads still-active operations from the GalleryStore into the
// in-memory statuses map so a freshly-started replica does not return an
// empty /api/operations payload while a peer is mid-install. Idempotent.
// No-op when no store is wired.
//
// The reconstructed OpStatus carries no Error type — the DB stores the
// message as a string and Hydrate surfaces it via errors.New so the UI's
// "operation failed" banner survives a frontend restart.
func (g *GalleryService) Hydrate() error {
g.Lock()
store := g.galleryStore
g.Unlock()
if store == nil {
return nil
}
ops, err := store.ListActive()
if err != nil {
return fmt.Errorf("listing active gallery ops: %w", err)
}
g.Lock()
defer g.Unlock()
for _, op := range ops {
// Skip rows that already have an in-memory status — the live
// broadcast subscriber will fill any gaps with fresher data.
if _, ok := g.statuses[op.ID]; ok {
continue
}
st := &OpStatus{
Message: op.Message,
Progress: op.Progress,
Phase: op.Phase,
CurrentBytes: op.CurrentBytes,
TotalBytes: op.TotalBytes,
FileName: op.FileName,
TotalFileSize: op.TotalFileSize,
DownloadedFileSize: op.DownloadedFileSize,
GalleryElementName: op.GalleryElementName,
Cancellable: op.Cancellable,
Deletion: IsDeleteOpType(op.OpType),
}
if op.Error != "" {
st.Error = errors.New(op.Error)
}
g.statuses[op.ID] = st
}
xlog.Info("Hydrated gallery service statuses from store", "count", len(ops))
return nil
}