From aba255768f9c3b545823070a3b1c8c8ea727632f Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 13 Aug 2026 20:40:18 +0000 Subject: [PATCH 1/7] fix(advisorylock): set statement_timeout alongside lock_timeout WithLockCtx already overrides a deployment-wide lock_timeout on its dedicated connection so a blocking pg_advisory_lock() waits its turn instead of failing with 55P03. statement_timeout aborts that exact same statement independently, with SQLSTATE 57014, and was not overridden. Production roles commonly carry statement_timeout=60s. Any guarded section longer than that (a cold model load stages for tens of minutes) therefore killed every concurrent waiter: advisorylock: acquiring lock 9003261067483446873: ERROR: canceling statement due to statement timeout (SQLSTATE 57014) Derive it from the same context budget as lock_timeout, with a matching RESET so the pooled connection is returned clean. Assisted-by: Claude Opus 5 [claude-code] --- core/services/advisorylock/advisorylock.go | 10 ++++ .../advisorylock/advisorylock_test.go | 48 +++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/core/services/advisorylock/advisorylock.go b/core/services/advisorylock/advisorylock.go index a37403c31940..ccaea2d967c4 100644 --- a/core/services/advisorylock/advisorylock.go +++ b/core/services/advisorylock/advisorylock.go @@ -180,6 +180,16 @@ func WithLockCtx(ctx context.Context, db *gorm.DB, key int64, fn func() error) e // Restore the session default before this pooled connection is reused. defer func() { _, _ = conn.ExecContext(context.Background(), "RESET lock_timeout") }() + // statement_timeout aborts the same blocking pg_advisory_lock() call + // independently of lock_timeout, with SQLSTATE 57014. Deployments that set a + // short statement_timeout on the role (60s is a common default) would + // otherwise kill every waiter regardless of the lock_timeout override above. + if _, err := conn.ExecContext(ctx, + fmt.Sprintf("SET statement_timeout = %d", waitBudget.Milliseconds())); err != nil { + return fmt.Errorf("advisorylock: setting statement_timeout: %w", err) + } + defer func() { _, _ = conn.ExecContext(context.Background(), "RESET statement_timeout") }() + if _, err := conn.ExecContext(ctx, "SELECT pg_advisory_lock($1)", key); err != nil { return fmt.Errorf("advisorylock: acquiring lock %d: %w", key, err) } diff --git a/core/services/advisorylock/advisorylock_test.go b/core/services/advisorylock/advisorylock_test.go index a77df56c446a..f1bd3e75ed5c 100644 --- a/core/services/advisorylock/advisorylock_test.go +++ b/core/services/advisorylock/advisorylock_test.go @@ -205,6 +205,54 @@ var _ = Describe("AdvisoryLock", func() { <-released }) + It("waits out a short server-side statement_timeout instead of failing with 57014", func() { + const lockKey int64 = 705 + + // Same shape as the lock_timeout case above, but for the *other* + // server-side bound that aborts the very same blocking + // pg_advisory_lock() statement. Production roles routinely carry + // statement_timeout=60s; a cold model load holds the lock far longer, + // so every concurrent caller died with SQLSTATE 57014 ("canceling + // statement due to statement timeout") rather than waiting its turn. + Expect(db.Exec("ALTER DATABASE testdb SET statement_timeout = '300ms'").Error).ToNot(HaveOccurred()) + sqlDB, err := db.DB() + Expect(err).ToNot(HaveOccurred()) + // Drop pooled connections so subsequent ones reconnect and inherit + // the new database-level statement_timeout default. + sqlDB.SetMaxIdleConns(0) + + holding := make(chan struct{}) + released := make(chan struct{}) + go func() { + defer GinkgoRecover() + herr := WithLockCtx(context.Background(), db, lockKey, func() error { + close(holding) + // Hold well past the 300ms server statement_timeout. + time.Sleep(1 * time.Second) + return nil + }) + Expect(herr).ToNot(HaveOccurred()) + close(released) + }() + + <-holding // ensure the holder owns the lock before we contend + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + executed := false + start := time.Now() + werr := WithLockCtx(ctx, db, lockKey, func() error { + executed = true + return nil + }) + Expect(werr).ToNot(HaveOccurred(), + "waiter should wait out the in-progress hold, not fail with statement_timeout (57014)") + Expect(executed).To(BeTrue()) + Expect(time.Since(start)).To(BeNumerically(">=", 400*time.Millisecond), + "waiter should have actually waited for the holder to release") + <-released + }) + It("bounds a deadline-less waiter with the backstop instead of waiting forever", func() { const lockKey int64 = 704 From f75c1c81c1ae3a758919c35fec0750a9b751f8a3 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 13 Aug 2026 20:59:12 +0000 Subject: [PATCH 2/7] feat(distributed): add ModelLoadJob, the durable cold-load record MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A cold load in distributed mode is a long-running background job, but it was modelled as a synchronous side effect of an inference request: the whole of it (backend install, multi-GB staging, checkpoint load) ran inside the per-model advisory lock. Loading a 35.7 GB GGUF held that lock for ~20 minutes, so every concurrent request for the same model blocked on pg_advisory_lock and died at the role's 60s statement_timeout. Introduce the row that lets the lock shrink to a decision. Exactly one ModelLoadJob may be active per tracking key; that uniqueness — not the lifetime of a lock — is what de-duplicates concurrent loaders across replicas. ClaimLoadJob does its read-then-write under the advisory lock and nothing else: no network, file or gRPC I/O inside the guarded section, so a claim costs milliseconds no matter how long the resulting load takes. LastProgress is a heartbeat rather than a byte counter. A checkpoint load legitimately moves zero bytes for many minutes, so a reaper keyed on byte movement would reclaim a healthy job mid-load; byte progress stays the concern of load_deadline.go. A job whose heartbeat stops for longer than the orphan window is reclaimable, so a replica killed mid-load cannot wedge a model permanently. Failed jobs keep their row for a short grace so an immediately-following request reports the real cause instead of silently starting a fresh load of a model that just failed. No caller yet — the router moves onto this in the next commit. Assisted-by: Claude Opus 5 [claude-code] --- core/services/nodes/interfaces.go | 12 + core/services/nodes/model_load_job.go | 262 +++++++++++++++++++++ core/services/nodes/model_load_job_test.go | 214 +++++++++++++++++ core/services/nodes/model_router_test.go | 1 + core/services/nodes/registry.go | 42 +++- core/services/nodes/router_test.go | 82 +++++++ 6 files changed, 612 insertions(+), 1 deletion(-) create mode 100644 core/services/nodes/model_load_job.go create mode 100644 core/services/nodes/model_load_job_test.go diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index 399ad74f8456..1be5f20e84a9 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -39,6 +39,18 @@ type ModelRouter interface { GetNodeLabels(ctx context.Context, nodeID string) ([]NodeLabel, error) FindNodesWithModel(ctx context.Context, modelName string) ([]BackendNode, error) LoadedReplicaStats(ctx context.Context, modelName string, candidateNodeIDs []string) ([]ReplicaCandidate, error) + LoadJobStore +} + +// LoadJobStore is the durable cold-load job record SmartRouter uses to +// de-duplicate concurrent loaders across replicas without holding the per-model +// advisory lock for the whole load. See ModelLoadJob. +type LoadJobStore interface { + ClaimLoadJob(ctx context.Context, trackingKey, owner string) (*ModelLoadJob, bool, error) + GetLoadJob(ctx context.Context, trackingKey string) (*ModelLoadJob, error) + UpdateLoadJob(ctx context.Context, trackingKey string, u LoadJobUpdate) error + FailLoadJob(ctx context.Context, trackingKey, msg string) error + DeleteLoadJob(ctx context.Context, trackingKey string) error } // ConcurrencyConflictResolver returns the names of configured models that diff --git a/core/services/nodes/model_load_job.go b/core/services/nodes/model_load_job.go new file mode 100644 index 000000000000..eaf50c7899a8 --- /dev/null +++ b/core/services/nodes/model_load_job.go @@ -0,0 +1,262 @@ +package nodes + +import ( + "context" + "errors" + "fmt" + "sync" + "time" + + "github.com/google/uuid" + "github.com/mudler/LocalAI/core/services/advisorylock" + "gorm.io/gorm" +) + +// Cold-load job states. `pending` covers node selection and replica +// allocation, which report nothing a waiter could act on; the rest name the +// phase the load is actually in. There is no terminal `ready` state — a +// successful job deletes its row and leaves the NodeModel row as the record. +const ( + LoadJobStatePending = "pending" + LoadJobStateInstalling = "installing" + LoadJobStateStaging = "staging" + LoadJobStateLoading = "loading" + LoadJobStateFailed = "failed" +) + +const ( + // loadJobHeartbeatInterval is how often a running job touches LastProgress. + // It matches the staging broadcast debounce so a job writes at most one row + // per second regardless of how many 32 KB chunks land in it. + loadJobHeartbeatInterval = stagingBroadcastInterval + + // loadJobOrphanWindow is how long a job may go without a heartbeat before + // another replica may reclaim it. Generous relative to the 1s heartbeat: a + // frontend under GC pressure or a stalled DB write must not have its + // perfectly healthy multi-GB transfer stolen and restarted from zero. + loadJobOrphanWindow = 60 * time.Second + + // loadJobFailureGrace is how long a failed job row is kept before deletion. + // Without it a waiter polling just after the failure finds no row, concludes + // "not loading", and starts a duplicate load of a model that just failed — + // a retry storm dressed as recovery. + loadJobFailureGrace = 15 * time.Second + + // loadJobPollInterval is how often a waiter on a non-owning replica polls + // the job row. The DB is the authority: NATS staging broadcasts are + // fire-and-forget, so a missed terminal event must not strand a waiter. + loadJobPollInterval = 2 * time.Second +) + +// loadJobLockPrefix namespaces the per-model advisory lock key. It is the same +// key the whole cold load used to hold; only the guarded section changed. +const loadJobLockPrefix = "model-load:" + +var ( + replicaIDOnce sync.Once + replicaIDValue string +) + +// ReplicaID returns this process's identity, generated once at startup and held +// for the process lifetime. It is recorded on jobs for diagnostics only, never +// for correctness decisions: a replica cannot be assumed alive just because its +// ID is on a row, which is what the LastProgress heartbeat is for. +func ReplicaID() string { + replicaIDOnce.Do(func() { replicaIDValue = uuid.New().String() }) + return replicaIDValue +} + +// IsOrphaned reports whether the job's owner has stopped heartbeating and the +// job may be reclaimed by another replica. +func (j *ModelLoadJob) IsOrphaned(now time.Time) bool { + return now.Sub(j.LastProgress) > loadJobOrphanWindow +} + +// Progress returns overall completion as a percentage, or 0 when the job has +// not reported enough to compute one. +func (j *ModelLoadJob) Progress() float64 { + if j.TotalBytes <= 0 { + return 0 + } + filePct := float64(j.BytesSent) / float64(j.TotalBytes) * 100 + if j.TotalFiles <= 1 || j.FileIndex <= 0 { + return filePct + } + return (float64(j.FileIndex-1)*100 + filePct) / float64(j.TotalFiles) +} + +// ETA returns the estimated time remaining for the transfer, and false when the +// job has not moved enough bytes for the observed rate to mean anything. A +// confidently wrong ETA on a twenty-minute wait is worse than none, so this +// omits rather than guesses. +func (j *ModelLoadJob) ETA(now time.Time) (time.Duration, bool) { + if j.State != LoadJobStateStaging || j.BytesSent <= 0 || j.TotalBytes <= j.BytesSent { + return 0, false + } + if j.StartedAt.IsZero() { + return 0, false + } + elapsed := now.Sub(j.StartedAt) + if elapsed < loadJobHeartbeatInterval { + return 0, false + } + rate := float64(j.BytesSent) / elapsed.Seconds() + if rate <= 0 { + return 0, false + } + return time.Duration(float64(j.TotalBytes-j.BytesSent)/rate) * time.Second, true +} + +// LoadJobUpdate is a partial update to a running job. Empty node fields are +// left untouched so a heartbeat does not erase the placement the runner +// reported earlier. +type LoadJobUpdate struct { + State string + NodeID string + NodeName string + ReplicaIndex int + BytesSent int64 + TotalBytes int64 + FileIndex int + TotalFiles int + // StartedAt anchors the rate the ETA is derived from. Set by the runner the + // first time the transfer reports bytes; zero leaves the stored value alone. + StartedAt time.Time +} + +// ClaimLoadJob decides, under the per-model advisory lock, whether this replica +// owns the cold load of trackingKey. It returns the live job and claimed=false +// when another replica is already loading it (or it just failed and is inside +// its grace window), or a fresh `pending` job with claimed=true when this +// replica took the work. +// +// The lock is held only across these statements — no network, file, or gRPC I/O +// happens inside it, which is the entire point of the job row. The primary key +// on TrackingKey is the real guard: if the lock were somehow bypassed the +// INSERT fails rather than producing two loaders. +func (r *NodeRegistry) ClaimLoadJob(ctx context.Context, trackingKey, owner string) (*ModelLoadJob, bool, error) { + var ( + job *ModelLoadJob + claimed bool + ) + lockKey := advisorylock.KeyFromString(loadJobLockPrefix + trackingKey) + err := advisorylock.WithLockCtx(ctx, r.db, lockKey, func() error { + var existing ModelLoadJob + err := r.db.WithContext(ctx).First(&existing, "tracking_key = ?", trackingKey).Error + switch { + case err == nil: + if !existing.IsOrphaned(time.Now()) { + job, claimed = &existing, false + return nil + } + // The owning replica died mid-load. Without this a crashed frontend + // would wedge the model permanently: every later request would find + // a job row that nobody is running and wait for a load that will + // never progress. + if err := r.db.WithContext(ctx).Delete(&ModelLoadJob{}, "tracking_key = ?", trackingKey).Error; err != nil { + return fmt.Errorf("deleting orphaned model load job: %w", err) + } + case errors.Is(err, gorm.ErrRecordNotFound): + default: + return fmt.Errorf("reading model load job: %w", err) + } + + now := time.Now() + fresh := &ModelLoadJob{ + TrackingKey: trackingKey, + State: LoadJobStatePending, + OwnerReplica: owner, + CreatedAt: now, + UpdatedAt: now, + LastProgress: now, + } + if err := r.db.WithContext(ctx).Create(fresh).Error; err != nil { + return fmt.Errorf("creating model load job: %w", err) + } + job, claimed = fresh, true + return nil + }) + if err != nil { + return nil, false, err + } + return job, claimed, nil +} + +// GetLoadJob returns the active job for trackingKey, or (nil, nil) when none is +// active. Callers on a non-owning replica poll this; it is the authority for +// both readiness and failure. +func (r *NodeRegistry) GetLoadJob(ctx context.Context, trackingKey string) (*ModelLoadJob, error) { + var job ModelLoadJob + err := r.db.WithContext(ctx).First(&job, "tracking_key = ?", trackingKey).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("reading model load job: %w", err) + } + return &job, nil +} + +// UpdateLoadJob applies a phase transition or heartbeat. LastProgress is always +// touched: it is the liveness signal the orphan check reads, and it must tick +// even during phases that move no bytes at all. +func (r *NodeRegistry) UpdateLoadJob(ctx context.Context, trackingKey string, u LoadJobUpdate) error { + now := time.Now() + fields := map[string]any{ + "last_progress": now, + "updated_at": now, + "bytes_sent": u.BytesSent, + "total_bytes": u.TotalBytes, + "file_index": u.FileIndex, + "total_files": u.TotalFiles, + } + if u.State != "" { + fields["state"] = u.State + } + if u.NodeID != "" { + fields["node_id"] = u.NodeID + fields["replica_index"] = u.ReplicaIndex + } + if u.NodeName != "" { + fields["node_name"] = u.NodeName + } + if !u.StartedAt.IsZero() { + fields["started_at"] = u.StartedAt + } + res := r.db.WithContext(ctx).Model(&ModelLoadJob{}). + Where("tracking_key = ?", trackingKey).Updates(fields) + if res.Error != nil { + return fmt.Errorf("updating model load job: %w", res.Error) + } + return nil +} + +// FailLoadJob records the real failure on the job row so every waiter — local +// or on another replica — reports the same cause instead of an anonymous +// timeout. The row is deleted after loadJobFailureGrace by the runner. +func (r *NodeRegistry) FailLoadJob(ctx context.Context, trackingKey, msg string) error { + now := time.Now() + res := r.db.WithContext(ctx).Model(&ModelLoadJob{}). + Where("tracking_key = ?", trackingKey). + Updates(map[string]any{ + "state": LoadJobStateFailed, + "last_error": msg, + "last_progress": now, + "updated_at": now, + }) + if res.Error != nil { + return fmt.Errorf("failing model load job: %w", res.Error) + } + return nil +} + +// DeleteLoadJob removes a terminal job row. Success deletes immediately (the +// NodeModel row is the record of a loaded model); failures delete after their +// grace window. +func (r *NodeRegistry) DeleteLoadJob(ctx context.Context, trackingKey string) error { + if err := r.db.WithContext(ctx). + Delete(&ModelLoadJob{}, "tracking_key = ?", trackingKey).Error; err != nil { + return fmt.Errorf("deleting model load job: %w", err) + } + return nil +} diff --git a/core/services/nodes/model_load_job_test.go b/core/services/nodes/model_load_job_test.go new file mode 100644 index 000000000000..a304a6b3dd67 --- /dev/null +++ b/core/services/nodes/model_load_job_test.go @@ -0,0 +1,214 @@ +package nodes + +import ( + "context" + "runtime" + "sync" + "sync/atomic" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/testutil" + "gorm.io/gorm" +) + +var _ = Describe("ModelLoadJob", func() { + var ( + db *gorm.DB + registry *NodeRegistry + ctx context.Context + ) + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + db = testutil.SetupTestDB() + var err error + registry, err = NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + ctx = context.Background() + }) + + Describe("ClaimLoadJob", func() { + It("claims a model that has no job yet", func() { + job, claimed, err := registry.ClaimLoadJob(ctx, "qwen3", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeTrue()) + Expect(job.State).To(Equal(LoadJobStatePending)) + Expect(job.OwnerReplica).To(Equal("replica-a")) + }) + + It("hands a second caller the live job instead of a second claim", func() { + _, claimed, err := registry.ClaimLoadJob(ctx, "qwen3", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeTrue()) + + job, claimed, err := registry.ClaimLoadJob(ctx, "qwen3", "replica-b") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeFalse(), "the second caller must attach as a waiter, not start a duplicate load") + Expect(job.OwnerReplica).To(Equal("replica-a")) + }) + + It("gives exactly one of many concurrent claimers the job", func() { + const claimers = 8 + var claims int32 + var wg sync.WaitGroup + for range claimers { + wg.Go(func() { + defer GinkgoRecover() + _, claimed, err := registry.ClaimLoadJob(ctx, "contended", "replica") + Expect(err).ToNot(HaveOccurred()) + if claimed { + atomic.AddInt32(&claims, 1) + } + }) + } + wg.Wait() + Expect(claims).To(Equal(int32(1))) + }) + + It("returns promptly while another replica's job is running", func() { + // The whole point of the split: a claim is a decision that takes + // milliseconds, so a concurrent request never waits behind the + // minutes-long load the owner is running. + _, claimed, err := registry.ClaimLoadJob(ctx, "slow-model", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeTrue()) + + running := make(chan struct{}) + done := make(chan struct{}) + go func() { + defer GinkgoRecover() + close(running) + // Stand in for a multi-GB staging run owned by replica-a. + time.Sleep(2 * time.Second) + Expect(registry.DeleteLoadJob(ctx, "slow-model")).To(Succeed()) + close(done) + }() + <-running + + start := time.Now() + _, claimed, err = registry.ClaimLoadJob(ctx, "slow-model", "replica-b") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeFalse()) + Expect(time.Since(start)).To(BeNumerically("<", 100*time.Millisecond), + "claiming must not block behind the running job") + <-done + }) + + It("reclaims a job whose owner stopped heartbeating", func() { + _, claimed, err := registry.ClaimLoadJob(ctx, "orphan", "dead-replica") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeTrue()) + + // Backdate the heartbeat past the orphan window, as a replica killed + // mid-load would leave it. + stale := time.Now().Add(-2 * loadJobOrphanWindow) + Expect(db.Model(&ModelLoadJob{}).Where("tracking_key = ?", "orphan"). + Update("last_progress", stale).Error).ToNot(HaveOccurred()) + + job, claimed, err := registry.ClaimLoadJob(ctx, "orphan", "live-replica") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeTrue(), "a dead replica must not wedge the model permanently") + Expect(job.OwnerReplica).To(Equal("live-replica")) + }) + }) + + Describe("lifecycle", func() { + It("records progress and clears the row on completion", func() { + _, _, err := registry.ClaimLoadJob(ctx, "m1", "replica-a") + Expect(err).ToNot(HaveOccurred()) + + started := time.Now() + Expect(registry.UpdateLoadJob(ctx, "m1", LoadJobUpdate{ + State: LoadJobStateStaging, NodeID: "node-1", NodeName: "nvidia-thor", + ReplicaIndex: 2, BytesSent: 500, TotalBytes: 1000, FileIndex: 1, TotalFiles: 1, + StartedAt: started, + })).To(Succeed()) + + job, err := registry.GetLoadJob(ctx, "m1") + Expect(err).ToNot(HaveOccurred()) + Expect(job.State).To(Equal(LoadJobStateStaging)) + Expect(job.NodeName).To(Equal("nvidia-thor")) + Expect(job.ReplicaIndex).To(Equal(2)) + Expect(job.Progress()).To(BeNumerically("~", 50, 0.01)) + + Expect(registry.DeleteLoadJob(ctx, "m1")).To(Succeed()) + job, err = registry.GetLoadJob(ctx, "m1") + Expect(err).ToNot(HaveOccurred()) + Expect(job).To(BeNil()) + }) + + It("keeps the placement across a byte-less heartbeat", func() { + _, _, err := registry.ClaimLoadJob(ctx, "m2", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(registry.UpdateLoadJob(ctx, "m2", LoadJobUpdate{ + State: LoadJobStateStaging, NodeID: "node-1", NodeName: "nvidia-thor", + })).To(Succeed()) + + before, err := registry.GetLoadJob(ctx, "m2") + Expect(err).ToNot(HaveOccurred()) + + // A checkpoint load moves no bytes for minutes; the heartbeat must + // still tick, and must not erase where the model is loading. + time.Sleep(10 * time.Millisecond) + Expect(registry.UpdateLoadJob(ctx, "m2", LoadJobUpdate{State: LoadJobStateLoading})).To(Succeed()) + + after, err := registry.GetLoadJob(ctx, "m2") + Expect(err).ToNot(HaveOccurred()) + Expect(after.NodeName).To(Equal("nvidia-thor")) + Expect(after.State).To(Equal(LoadJobStateLoading)) + Expect(after.LastProgress.After(before.LastProgress)).To(BeTrue()) + }) + + It("records the failure cause for waiters to read", func() { + _, _, err := registry.ClaimLoadJob(ctx, "m3", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(registry.FailLoadJob(ctx, "m3", "no available nodes")).To(Succeed()) + + job, err := registry.GetLoadJob(ctx, "m3") + Expect(err).ToNot(HaveOccurred()) + Expect(job.State).To(Equal(LoadJobStateFailed)) + Expect(job.LastError).To(Equal("no available nodes")) + }) + + It("hands a request arriving inside the failure grace the real error", func() { + _, _, err := registry.ClaimLoadJob(ctx, "m4", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(registry.FailLoadJob(ctx, "m4", "worker out of VRAM")).To(Succeed()) + + job, claimed, err := registry.ClaimLoadJob(ctx, "m4", "replica-b") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeFalse(), "a fresh load must not silently start on top of a just-failed one") + Expect(job.LastError).To(Equal("worker out of VRAM")) + }) + }) + + Describe("ETA", func() { + It("is omitted until the observed rate means something", func() { + job := &ModelLoadJob{State: LoadJobStateStaging, BytesSent: 0, TotalBytes: 1000} + _, ok := job.ETA(time.Now()) + Expect(ok).To(BeFalse()) + + job = &ModelLoadJob{State: LoadJobStateStaging, BytesSent: 10, TotalBytes: 1000, StartedAt: time.Now()} + _, ok = job.ETA(time.Now()) + Expect(ok).To(BeFalse(), "less than one broadcast interval of data is not a rate") + }) + + It("derives the remaining time from the job's own rate", func() { + now := time.Now() + job := &ModelLoadJob{ + State: LoadJobStateStaging, + BytesSent: 1000, TotalBytes: 3000, + StartedAt: now.Add(-10 * time.Second), + } + eta, ok := job.ETA(now) + Expect(ok).To(BeTrue()) + // 100 B/s observed, 2000 B left. + Expect(eta).To(BeNumerically("~", 20*time.Second, time.Second)) + }) + }) +}) diff --git a/core/services/nodes/model_router_test.go b/core/services/nodes/model_router_test.go index f68caa143925..b31e9dc064bb 100644 --- a/core/services/nodes/model_router_test.go +++ b/core/services/nodes/model_router_test.go @@ -14,6 +14,7 @@ import ( // --- fakeModelRouterForSmartRouter implements ModelRouter --- type fakeModelRouterForSmartRouter struct { + fakeLoadJobStore mu sync.Mutex node *BackendNode nodeModel *NodeModel diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index 0b85a1f44164..1a59e361f2e5 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -241,6 +241,46 @@ type PendingBackendOp struct { NextRetryAt time.Time `gorm:"index" json:"next_retry_at"` } +// ModelLoadJob is one in-flight cold load of a model. Exactly one row per +// trackingKey may be active at a time; that uniqueness — not the lifetime of an +// advisory lock — is what de-duplicates concurrent loaders across replicas. +// +// Before this table the whole cold load (backend install, multi-GB staging, +// checkpoint load) ran inside the per-model advisory lock, so every other +// replica's request for the model blocked on pg_advisory_lock for tens of +// minutes and was killed by the role's statement_timeout. The job row lets the +// lock shrink to the claim while the work itself runs unlocked and observable. +// +// Terminal rows are deleted rather than retained: NodeModel is already the +// record of what is loaded, and keeping finished jobs would create a second +// source of truth about it. +type ModelLoadJob struct { + TrackingKey string `gorm:"primaryKey;size:255" json:"tracking_key"` + State string `gorm:"size:16;not null;index" json:"state"` + OwnerReplica string `gorm:"size:64" json:"owner_replica"` + NodeID string `gorm:"size:36" json:"node_id"` + NodeName string `gorm:"size:255" json:"node_name"` + ReplicaIndex int `json:"replica_index"` + BytesSent int64 `json:"bytes_sent"` + TotalBytes int64 `json:"total_bytes"` + FileIndex int `json:"file_index"` + TotalFiles int `json:"total_files"` + LastError string `gorm:"type:text" json:"last_error,omitempty"` + // StartedAt is when the job first reported bytes, and is what the ETA rate + // is measured from. Distinct from CreatedAt, which also covers node + // selection and backend install — phases that move no bytes and would skew + // the derived rate low. + StartedAt time.Time `json:"started_at"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` + // LastProgress is a heartbeat, not a byte counter: the runner touches it on + // a fixed interval for as long as it is alive, whether or not bytes are + // moving. A checkpoint load legitimately transfers zero bytes for many + // minutes, so a reaper keyed on byte movement would reclaim a healthy job + // mid-load. Byte progress is measured separately, by load_deadline.go. + LastProgress time.Time `gorm:"index" json:"last_progress_at"` +} + // Op constants mirror the operation names used by DistributedBackendManager // so callers don't repeat stringly-typed values. const ( @@ -343,7 +383,7 @@ func (r *NodeRegistry) nodeModelNames(ctx context.Context, db *gorm.DB, nodeID s // when multiple instances (frontend + workers) start at the same time. func NewNodeRegistry(db *gorm.DB) (*NodeRegistry, error) { if err := advisorylock.WithLockCtx(context.Background(), db, advisorylock.KeySchemaMigrate, func() error { - return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}) + return db.AutoMigrate(&BackendNode{}, &NodeModel{}, &NodeLabel{}, &ModelSchedulingConfig{}, &PendingBackendOp{}, &ModelLoadInfo{}, &ModelLoadJob{}) }); err != nil { return nil, fmt.Errorf("migrating node tables: %w", err) } diff --git a/core/services/nodes/router_test.go b/core/services/nodes/router_test.go index a3f63cc4b13a..6e1d86d47929 100644 --- a/core/services/nodes/router_test.go +++ b/core/services/nodes/router_test.go @@ -60,6 +60,8 @@ func (f *fakeFileStager) ListRemoteDir(_ context.Context, _, _ string) ([]string // fakeModelRouter implements ModelRouter with configurable return values. type fakeModelRouter struct { + fakeLoadJobStore + // FindAndLockNodeWithModel returns findAndLockNode *BackendNode findAndLockNM *NodeModel @@ -148,6 +150,86 @@ func (f *fakeModelRouter) LoadedReplicaStats(_ context.Context, modelName string return f.loadedReplicaStatsByName[modelName], nil } +// fakeLoadJobStore is an in-memory LoadJobStore so tests that build a +// SmartRouter over a fake registry get the real claim/wait semantics without a +// database. +type fakeLoadJobStore struct { + mu sync.Mutex + jobs map[string]*ModelLoadJob +} + +func (s *fakeLoadJobStore) ClaimLoadJob(_ context.Context, trackingKey, owner string) (*ModelLoadJob, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.jobs == nil { + s.jobs = map[string]*ModelLoadJob{} + } + if existing, ok := s.jobs[trackingKey]; ok && !existing.IsOrphaned(time.Now()) { + cp := *existing + return &cp, false, nil + } + now := time.Now() + job := &ModelLoadJob{TrackingKey: trackingKey, State: LoadJobStatePending, OwnerReplica: owner, CreatedAt: now, UpdatedAt: now, LastProgress: now} + s.jobs[trackingKey] = job + cp := *job + return &cp, true, nil +} + +func (s *fakeLoadJobStore) GetLoadJob(_ context.Context, trackingKey string) (*ModelLoadJob, error) { + s.mu.Lock() + defer s.mu.Unlock() + job, ok := s.jobs[trackingKey] + if !ok { + return nil, nil + } + cp := *job + return &cp, nil +} + +func (s *fakeLoadJobStore) UpdateLoadJob(_ context.Context, trackingKey string, u LoadJobUpdate) error { + s.mu.Lock() + defer s.mu.Unlock() + job, ok := s.jobs[trackingKey] + if !ok { + return nil + } + if u.State != "" { + job.State = u.State + } + if u.NodeID != "" { + job.NodeID = u.NodeID + job.ReplicaIndex = u.ReplicaIndex + } + if u.NodeName != "" { + job.NodeName = u.NodeName + } + if !u.StartedAt.IsZero() { + job.StartedAt = u.StartedAt + } + job.BytesSent, job.TotalBytes = u.BytesSent, u.TotalBytes + job.FileIndex, job.TotalFiles = u.FileIndex, u.TotalFiles + job.LastProgress = time.Now() + return nil +} + +func (s *fakeLoadJobStore) FailLoadJob(_ context.Context, trackingKey, msg string) error { + s.mu.Lock() + defer s.mu.Unlock() + if job, ok := s.jobs[trackingKey]; ok { + job.State = LoadJobStateFailed + job.LastError = msg + job.LastProgress = time.Now() + } + return nil +} + +func (s *fakeLoadJobStore) DeleteLoadJob(_ context.Context, trackingKey string) error { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.jobs, trackingKey) + return nil +} + func (f *fakeModelRouter) DecrementInFlight(_ context.Context, nodeID, modelName string, _ int) error { f.decrementCalls = append(f.decrementCalls, nodeID+":"+modelName) return nil From 37ad4f73a3cfb15fde2a314ca879f7b955844c52 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 13 Aug 2026 21:13:21 +0000 Subject: [PATCH 3/7] refactor(distributed): run cold loads as jobs, outside the advisory lock MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Route wrapped the entire cold load — node selection, backend install, multi-GB staging and the remote LoadModel — in the per-model advisory lock. The lock's job is to de-duplicate concurrent loaders, a decision that takes milliseconds; holding it for the tens of minutes the resulting work takes is what turned a dedup mechanism into a cluster-wide outage for that model. Split it into a claim and a run. The claim is the only thing left inside the lock. The run is a background job owned by the claiming replica and bounded by the same progress-extended deadline as before; every other request for that model — local or on another replica — attaches as a waiter and is served the moment the model is ready, with no duplicate load and no lock contention. Waiters share one broadcast rather than an ordered queue: they all want the identical outcome, so ordering them would add fairness machinery that changes no result. The local channel wakes same-replica waiters instantly and a 2s DB poll is the authority, because a waiter on another replica has no channel to close. On wake a waiter re-runs the warm path rather than trusting the signal — the model may have been evicted in between. A waiter whose client disconnects returns immediately and the job keeps running; it belongs to the job record, not to the request. A failure is recorded on the row so every waiter reports the real cause, and the row survives briefly so the next request does not read "no job" as "not loading" and start a duplicate load of a model that just failed. The runner heartbeats the row on a fixed interval whether or not bytes are moving, which is what keeps a legitimately silent checkpoint load from being reclaimed as an orphan. Phase (installing/staging/loading) and placement ride to the heartbeat on the context, the same seam load_deadline.go already uses, so single-host paths are untouched. Non-distributed mode (no DB) keeps the inline load exactly as it was. Assisted-by: Claude Opus 5 [claude-code] --- core/services/nodes/load_job_phase.go | 64 +++++ core/services/nodes/load_job_runner.go | 249 +++++++++++++++++++ core/services/nodes/router.go | 255 ++++++++++---------- core/services/nodes/router_load_job_test.go | 219 +++++++++++++++++ 4 files changed, 665 insertions(+), 122 deletions(-) create mode 100644 core/services/nodes/load_job_phase.go create mode 100644 core/services/nodes/load_job_runner.go create mode 100644 core/services/nodes/router_load_job_test.go diff --git a/core/services/nodes/load_job_phase.go b/core/services/nodes/load_job_phase.go new file mode 100644 index 000000000000..ece0d099d96a --- /dev/null +++ b/core/services/nodes/load_job_phase.go @@ -0,0 +1,64 @@ +package nodes + +import ( + "context" + "sync" +) + +// loadPhaseReporter carries the current phase of a cold load from the code that +// performs it back to the job runner's heartbeat, without threading a job +// handle through every scheduling function. +// +// It rides on the context the same way the cold-load deadline does (see +// load_deadline.go), so the single-host paths and every test that constructs a +// router directly stay untouched: with no reporter on the context, the report +// calls are no-ops. +type loadPhaseReporter struct { + mu sync.Mutex + state string + nodeID string + nodeName string + replicaIndex int +} + +type loadPhaseKey struct{} + +func newLoadPhaseReporter() *loadPhaseReporter { + return &loadPhaseReporter{state: LoadJobStatePending} +} + +func withLoadPhaseReporter(ctx context.Context, p *loadPhaseReporter) context.Context { + return context.WithValue(ctx, loadPhaseKey{}, p) +} + +// snapshot returns the phase as a job update. Byte counts are filled in by the +// caller from the staging tracker. +func (p *loadPhaseReporter) snapshot() LoadJobUpdate { + p.mu.Lock() + defer p.mu.Unlock() + return LoadJobUpdate{ + State: p.state, + NodeID: p.nodeID, + NodeName: p.nodeName, + ReplicaIndex: p.replicaIndex, + } +} + +func (p *loadPhaseReporter) set(state string, node *BackendNode, replicaIndex int) { + p.mu.Lock() + defer p.mu.Unlock() + p.state = state + if node != nil { + p.nodeID, p.nodeName, p.replicaIndex = node.ID, node.Name, replicaIndex + } +} + +// reportLoadPhase records which phase of a cold load is running, so a waiting +// request can be told "staging to nvidia-thor" rather than nothing at all. A +// context without a reporter (non-distributed loads, reconciler scale-ups, +// tests) is a no-op. +func reportLoadPhase(ctx context.Context, state string, node *BackendNode, replicaIndex int) { + if p, ok := ctx.Value(loadPhaseKey{}).(*loadPhaseReporter); ok { + p.set(state, node, replicaIndex) + } +} diff --git a/core/services/nodes/load_job_runner.go b/core/services/nodes/load_job_runner.go new file mode 100644 index 000000000000..ec806cd4265c --- /dev/null +++ b/core/services/nodes/load_job_runner.go @@ -0,0 +1,249 @@ +package nodes + +import ( + "context" + "fmt" + "time" + + "github.com/mudler/xlog" +) + +// maxColdLoadRounds bounds how many times a request may claim-or-wait before +// giving up. A round ends when the job reaches a terminal state; a second round +// only happens when the model was evicted between the job finishing and the +// waiter re-checking, which is rare and must not become a spin. +const maxColdLoadRounds = 3 + +// routeViaLoadJob serves a request whose model is not loaded, in distributed +// mode. The cold load itself becomes a durable job owned by whichever replica +// claims it; every other request for the same model — on this replica or any +// other — attaches as a waiter and is served the moment the model is ready. +// +// The per-model advisory lock still de-duplicates loaders, but it is held only +// for the claim. Before this split it wrapped the whole load, so a 35.7 GB +// staging run pinned it for ~20 minutes and every concurrent request died at +// the role's 60s statement_timeout with SQLSTATE 57014. +func (r *SmartRouter) routeViaLoadJob(ctx context.Context, att *routeAttempt) (*RouteResult, error) { + for range maxColdLoadRounds { + // Register interest BEFORE claiming, so a job that finishes immediately + // cannot close the channel before this waiter exists. + waiter := r.loadWaiterChan(att.trackingKey) + + job, claimed, err := r.registry.ClaimLoadJob(ctx, att.trackingKey, ReplicaID()) + if err != nil { + // A broken job table must not make the model unroutable: fall back + // to loading inline, which is what every release before this did. + xlog.Warn("Claiming the model load job failed; loading inline instead", + "model", att.trackingKey, "error", err) + loadCtx, cancelLoad := r.newColdLoadContext(context.WithoutCancel(ctx)) + defer cancelLoad() + return r.coldLoad(loadCtx, att, 1) + } + + switch { + case claimed: + // The model may have been loaded between this request's warm-path + // check and the claim — the check the old code did after acquiring + // the lock. Without it the claim would schedule a second copy of a + // model that is already up. + if result := r.tryWarmPath(ctx, att); result != nil { + r.finishLoadJob(ctx, att.trackingKey) + return result, nil + } + r.startLoadJob(ctx, att) + case job != nil && job.State == LoadJobStateFailed: + // Inside the failure grace window: report the real cause rather + // than silently starting a fresh load of a model that just failed. + return nil, fmt.Errorf("loading model %s: %s", att.trackingKey, job.LastError) + default: + xlog.Info("Model is already loading on another replica; waiting for it", + "model", att.trackingKey, "state", job.State, "node", job.NodeName, "owner", job.OwnerReplica) + } + + if err := r.waitForLoadJob(ctx, att.trackingKey, waiter); err != nil { + return nil, err + } + + // The signal is not the authority — the model may have been evicted + // between ready and wake, so re-run the warm path. + if result := r.tryWarmPath(ctx, att); result != nil { + return result, nil + } + } + return nil, fmt.Errorf("loading model %s: the load finished but the model is not available", att.trackingKey) +} + +// startLoadJob runs the claimed cold load in the background, detached from the +// request that triggered it. The job is owned by its record, not by that +// request: the client may disconnect, be retried onto another replica, or time +// out, and the transfer keeps going. +func (r *SmartRouter) startLoadJob(ctx context.Context, att *routeAttempt) { + trackingKey := att.trackingKey + // Keep the request's context VALUES (prefix chain and friends) but none of + // its cancellation — see newColdLoadContext. + parent := context.WithoutCancel(ctx) + + go func() { + loadCtx, cancelLoad := r.newColdLoadContext(parent) + defer cancelLoad() + + phase := newLoadPhaseReporter() + loadCtx = withLoadPhaseReporter(loadCtx, phase) + + stopHeartbeat := r.startLoadJobHeartbeat(parent, trackingKey, phase) + + _, err := r.coldLoad(loadCtx, att, 0) + + stopHeartbeat() + + // Bookkeeping must survive the load context, which may be exactly what + // just expired. + bookCtx, cancelBook := context.WithTimeout(context.WithoutCancel(parent), 30*time.Second) + defer cancelBook() + + if err != nil { + xlog.Error("Cold load job failed", "model", trackingKey, "error", err) + if ferr := r.registry.FailLoadJob(bookCtx, trackingKey, err.Error()); ferr != nil { + xlog.Warn("Failed to record cold load failure", "model", trackingKey, "error", ferr) + } + r.closeLoadWaiters(trackingKey) + // Keep the row briefly so a request arriving right now reports this + // failure instead of starting a duplicate load. Deleting it + // immediately turns a failure into a retry storm. + time.AfterFunc(loadJobFailureGrace, func() { + delCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if derr := r.registry.DeleteLoadJob(delCtx, trackingKey); derr != nil { + xlog.Warn("Failed to clear failed cold load job", "model", trackingKey, "error", derr) + } + }) + return + } + + r.finishLoadJob(bookCtx, trackingKey) + }() +} + +// finishLoadJob ends a job that succeeded. The NodeModel row (state `loaded`) +// is the record from here, so the job row is dropped BEFORE waiters are woken: +// they re-run the warm path and must not find a job that is really done. +func (r *SmartRouter) finishLoadJob(ctx context.Context, trackingKey string) { + if err := r.registry.DeleteLoadJob(ctx, trackingKey); err != nil { + xlog.Warn("Failed to clear completed cold load job", "model", trackingKey, "error", err) + } + r.closeLoadWaiters(trackingKey) +} + +// startLoadJobHeartbeat keeps the job row's liveness and progress fresh while +// the load runs, and returns a function that stops it. +// +// The heartbeat is deliberately time-driven rather than byte-driven: a +// checkpoint load moves no bytes for many minutes, and a job that only wrote a +// row when bytes moved would look orphaned and be reclaimed mid-load. Byte +// progress is copied in from the staging tracker, which already debounces the +// per-chunk callbacks, so the row is written at most once per interval. +func (r *SmartRouter) startLoadJobHeartbeat(parent context.Context, trackingKey string, phase *loadPhaseReporter) func() { + done := make(chan struct{}) + stopped := make(chan struct{}) + + go func() { + defer close(stopped) + ticker := time.NewTicker(loadJobHeartbeatInterval) + defer ticker.Stop() + var startedAt time.Time + for { + select { + case <-done: + return + case <-ticker.C: + u := phase.snapshot() + if st := r.stagingTracker.Get(trackingKey); st != nil { + u.BytesSent, u.TotalBytes = st.BytesSent, st.TotalBytes + u.FileIndex, u.TotalFiles = st.FileIndex, st.TotalFiles + if u.BytesSent > 0 && startedAt.IsZero() { + startedAt = time.Now() + } + u.StartedAt = startedAt + } + ctx, cancel := context.WithTimeout(context.WithoutCancel(parent), loadJobHeartbeatInterval*5) + if err := r.registry.UpdateLoadJob(ctx, trackingKey, u); err != nil { + xlog.Debug("Failed to heartbeat cold load job", "model", trackingKey, "error", err) + } + cancel() + } + } + }() + + return func() { + close(done) + <-stopped + } +} + +// waitForLoadJob blocks until the cold load of trackingKey reaches a terminal +// state, the job's failure is known, or the caller gives up. +// +// Waiters share one broadcast rather than an ordered queue: they all want the +// identical outcome — the model loaded — so ordering them would add fairness +// machinery that changes no result. The local channel wakes same-replica +// waiters instantly; the DB poll is the authority, because a waiter on another +// replica has no channel to close and NATS broadcasts are fire-and-forget, so a +// missed terminal event must not strand it. +func (r *SmartRouter) waitForLoadJob(ctx context.Context, trackingKey string, waiter <-chan struct{}) error { + ticker := time.NewTicker(loadJobPollInterval) + defer ticker.Stop() + + for { + select { + case <-waiter: + return nil + case <-ctx.Done(): + // The client gave up. The job is unaffected: it is owned by the job + // record, not by this request. + return ctx.Err() + case <-ticker.C: + job, err := r.registry.GetLoadJob(ctx, trackingKey) + if err != nil { + xlog.Debug("Polling the model load job failed", "model", trackingKey, "error", err) + continue + } + if job == nil { + // Terminal: either it succeeded, or it was reaped. Either way + // the caller re-checks the warm path. + return nil + } + if job.State == LoadJobStateFailed { + return fmt.Errorf("loading model %s: %s", trackingKey, job.LastError) + } + } + } +} + +// loadWaiterChan returns the broadcast channel for trackingKey, creating it on +// first use. Same shape as advisorylock.localLocks: N local requests share one +// wait and wake together. +func (r *SmartRouter) loadWaiterChan(trackingKey string) <-chan struct{} { + r.loadWaitersMu.Lock() + defer r.loadWaitersMu.Unlock() + if r.loadWaiters == nil { + r.loadWaiters = map[string]chan struct{}{} + } + ch, ok := r.loadWaiters[trackingKey] + if !ok { + ch = make(chan struct{}) + r.loadWaiters[trackingKey] = ch + } + return ch +} + +// closeLoadWaiters wakes every local waiter on trackingKey. A waiter that +// registers after this sees a fresh channel and falls back to the DB poll. +func (r *SmartRouter) closeLoadWaiters(trackingKey string) { + r.loadWaitersMu.Lock() + ch, ok := r.loadWaiters[trackingKey] + delete(r.loadWaiters, trackingKey) + r.loadWaitersMu.Unlock() + if ok { + close(ch) + } +} diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index fd051e9544b4..07816c5bf1bc 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -14,7 +14,6 @@ import ( "time" "github.com/mudler/LocalAI/core/config" - "github.com/mudler/LocalAI/core/services/advisorylock" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/pkg/distributedhdr" grpc "github.com/mudler/LocalAI/pkg/grpc" @@ -188,6 +187,12 @@ type SmartRouter struct { // hard countdown into a progress-extended hold (see load_deadline.go). stagingStallWindow time.Duration modelLoadAbsoluteMax time.Duration + // loadWaiters is one broadcast channel per model being cold-loaded, closed + // when the job reaches a terminal state. Same-model waiters all want the + // identical outcome, so they share one wait instead of queueing. See + // load_job_runner.go. + loadWaitersMu sync.Mutex + loadWaiters map[string]chan struct{} } // probeCacheTTL is how long a successful gRPC HealthCheck on a backend is @@ -239,6 +244,7 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter // ceiling, so nothing to normalize here. stagingStallWindow: opts.StagingStallWindow, modelLoadAbsoluteMax: opts.ModelLoadAbsoluteMax, + loadWaiters: map[string]chan struct{}{}, } } @@ -329,6 +335,7 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking if err := r.registry.SetNodeModel(ctx, node.ID, trackingKey, replicaIndex, "staging", backendAddr, 0); err != nil { xlog.Warn("Failed to record staging state", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", err) } + reportLoadPhase(ctx, LoadJobStateStaging, node, replicaIndex) lifecycleSettled := false defer func() { if lifecycleSettled { @@ -367,6 +374,7 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking if err := r.registry.SetNodeModel(ctx, node.ID, trackingKey, replicaIndex, "loading", backendAddr, 0); err != nil { xlog.Warn("Failed to record loading state", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", err) } + reportLoadPhase(ctx, LoadJobStateLoading, node, replicaIndex) // The cold-load hold above this call extends on STAGING progress, and // the remote LoadModel reports none — so once the last byte lands the @@ -555,138 +563,140 @@ func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType // below. Both are nil (no-op) when prefix-cache routing is disabled. pref, observeChain := r.buildPreference(ctx, trackingKey, candidateNodeIDs, sched) + att := &routeAttempt{ + trackingKey: trackingKey, + modelName: modelName, + backendType: backendType, + modelOpts: modelOpts, + parallel: parallel, + sched: sched, + candidateNodeIDs: candidateNodeIDs, + pref: pref, + observeChain: observeChain, + } + // Step 1: Find and atomically lock a node with this model loaded - node, nm, err := r.registry.FindAndLockNodeWithModel(ctx, trackingKey, candidateNodeIDs, pref) - if err == nil && node != nil { - modelAddr := node.Address - if nm.Address != "" { - modelAddr = nm.Address - } - replicaIdx := nm.ReplicaIndex - - // Verify the backend process is still alive via gRPC health check - if !r.probeHealth(ctx, node, modelAddr) { - // Stale — roll back the increment, remove the specific replica row, fall through - r.registry.DecrementInFlight(ctx, node.ID, trackingKey, replicaIdx) - r.registry.RemoveNodeModel(ctx, node.ID, trackingKey, replicaIdx) - xlog.Warn("Backend not reachable for cached model, falling through to reload", - "node", node.Name, "model", modelName, "replica", replicaIdx) - } else { - // Verify node still matches scheduling constraints - if !r.nodeMatchesScheduling(ctx, node, sched) { - r.registry.DecrementInFlight(ctx, node.ID, trackingKey, replicaIdx) - xlog.Info("Cached model on node that no longer matches selector, falling through", - "node", node.Name, "model", trackingKey, "replica", replicaIdx) - // Fall through to step 2 (scheduleNewModel) - } else { - // Node is alive — FindAndLockNodeWithModel already incremented in-flight as a - // reservation. InFlightTrackingClient handles per-inference tracking, and its - // onFirstComplete callback releases the reservation after the first inference - // call finishes, so in-flight returns to 0 when idle. - r.registry.TouchNodeModel(ctx, node.ID, trackingKey, replicaIdx) - r.observePrefix(trackingKey, observeChain, prefixcache.ReplicaKey{NodeID: node.ID, Replica: replicaIdx}) - grpcClient := r.buildClientForAddr(node, modelAddr, parallel) - tracked := NewInFlightTrackingClient(grpcClient, r.registry, node.ID, trackingKey, replicaIdx) - return r.newRouteResult(node, trackingKey, replicaIdx, grpcClient, tracked), nil - } - } + if result := r.tryWarmPath(ctx, att); result != nil { + return result, nil } - // Step 2: Model not loaded — schedule loading with distributed lock to prevent duplicates. - // - // Detach the cold-load from the caller's context. Staging a model can - // transfer multiple GB to a worker, which takes far longer than any client - // keeps its HTTP request open — a browser refresh, an ingress/LB idle - // timeout, or a round-robined retry landing on another replica all cancel - // the request context. If staging were bound to it, the multi-GB upload - // aborts with "context canceled" mid-transfer and large models can never - // finish staging (the model-load outage). WithoutCancel keeps the request's - // values (prefix chain, etc.) but drops its cancellation/deadline. + // Step 2: model not loaded — it has to be cold-loaded. // - // Detaching from the caller is necessary, but it must not be unbounded: the - // load runs while holding the per-model advisory lock, and a worker that - // dies mid-install (its backend.install never replies) would otherwise pin - // that lock (and every other replica's request for the same model) until - // the NATS install deadline alone expires. Re-impose a single hard ceiling - // over the whole sequence so the lock is always released in bounded time, - // even if a sub-step wedges. Each long step still has its own (tighter) - // bound; this only backstops them. The per-model advisory lock below - // de-dupes concurrent loaders across replicas. - // The backstop is progress-based, not wall-clock: staging time is bytes over - // bandwidth, so a fixed ceiling is a model-size cliff (a 70 GB checkpoint - // transferring healthily at 26 MB/s needs ~45m and was killed at exactly - // 25m00s). The hold instead extends while the transfer reports bytes and - // expires a stall window after they stop. See load_deadline.go. - loadCtx, cancelLoad := newLoadDeadlineContext(context.WithoutCancel(ctx), - r.modelLoadCeiling, r.stagingStallWindow, r.modelLoadAbsoluteMax) + // In distributed mode that runs as a durable job (see load_job_runner.go): + // the per-model advisory lock now guards only the claim, and the transfer + // itself runs unlocked so a concurrent request for the same model never + // blocks on pg_advisory_lock for the tens of minutes a multi-GB stage takes. + if r.db != nil { + return r.routeViaLoadJob(ctx, att) + } + + // No DB (non-distributed): there is no other replica to coordinate with, so + // the load stays inline on the request exactly as before. + loadCtx, cancelLoad := r.newColdLoadContext(context.WithoutCancel(ctx)) defer cancelLoad() - loadModel := func(ctx context.Context) (*RouteResult, error) { - // Re-check after acquiring lock — another request may have loaded it - node, nm, err := r.registry.FindAndLockNodeWithModel(ctx, trackingKey, candidateNodeIDs, pref) - if err == nil && node != nil { - modelAddr := node.Address - if nm.Address != "" { - modelAddr = nm.Address - } - replicaIdx := nm.ReplicaIndex - - // Verify the backend process is still alive via gRPC health check - if !r.probeHealth(ctx, node, modelAddr) { - // Stale — roll back the increment, remove the specific replica row, continue loading - r.registry.DecrementInFlight(ctx, node.ID, trackingKey, replicaIdx) - r.registry.RemoveNodeModel(ctx, node.ID, trackingKey, replicaIdx) - xlog.Warn("Backend not reachable for cached model inside lock, proceeding to load", - "node", node.Name, "model", modelName, "replica", replicaIdx) - } else { - // Verify node still matches scheduling constraints - if !r.nodeMatchesScheduling(ctx, node, sched) { - r.registry.DecrementInFlight(ctx, node.ID, trackingKey, replicaIdx) - xlog.Info("Cached model on node that no longer matches selector, falling through", - "node", node.Name, "model", trackingKey, "replica", replicaIdx) - // Fall through to scheduling below - } else { - // Model loaded while we waited — FindAndLockNodeWithModel already incremented - // in-flight as a reservation. Release it after the first inference completes. - r.registry.TouchNodeModel(ctx, node.ID, trackingKey, replicaIdx) - r.observePrefix(trackingKey, observeChain, prefixcache.ReplicaKey{NodeID: node.ID, Replica: replicaIdx}) - grpcClient := r.buildClientForAddr(node, modelAddr, parallel) - tracked := NewInFlightTrackingClient(grpcClient, r.registry, node.ID, trackingKey, replicaIdx) - return r.newRouteResult(node, trackingKey, replicaIdx, grpcClient, tracked), nil - } - } - } + if result := r.tryWarmPath(loadCtx, att); result != nil { + return result, nil + } + return r.coldLoad(loadCtx, att, 1) +} - // Still not loaded — use shared schedule-and-load logic, which picks - // both the node and the replica slot. - result, err := r.scheduleAndLoad(ctx, backendType, trackingKey, modelName, modelOpts, parallel, 1) - if err != nil { - return nil, err - } +// routeAttempt is the per-request routing state shared by the warm path, the +// cold-load job runner, and the waiter loop. Bundling it keeps those from +// drifting apart on which candidate set or preference they used. +type routeAttempt struct { + trackingKey string + modelName string + backendType string + modelOpts *pb.ModelOptions + parallel bool + sched *ModelSchedulingConfig + candidateNodeIDs []string + pref *RoutePreference + observeChain []uint64 +} - // Cold load landed on result.Node replica result.ReplicaIndex: record the - // assignment so subsequent requests with the same prefix prefer it. - r.observePrefix(trackingKey, observeChain, prefixcache.ReplicaKey{NodeID: result.Node.ID, Replica: result.ReplicaIndex}) +// tryWarmPath returns a route to an already-loaded, reachable replica, or nil +// when the model has to be cold-loaded. It is the authority on readiness: a +// waiter woken by a finished job re-runs it rather than trusting the signal, +// because the model may have been evicted between ready and wake. +func (r *SmartRouter) tryWarmPath(ctx context.Context, att *routeAttempt) *RouteResult { + node, nm, err := r.registry.FindAndLockNodeWithModel(ctx, att.trackingKey, att.candidateNodeIDs, att.pref) + if err != nil || node == nil { + return nil + } + modelAddr := node.Address + if nm.Address != "" { + modelAddr = nm.Address + } + replicaIdx := nm.ReplicaIndex - replicaIdx := result.ReplicaIndex - tracked := NewInFlightTrackingClient(result.Client, r.registry, result.Node.ID, trackingKey, replicaIdx) - return r.newRouteResult(result.Node, trackingKey, replicaIdx, result.Client, tracked), nil + // Verify the backend process is still alive via gRPC health check + if !r.probeHealth(ctx, node, modelAddr) { + // Stale — roll back the increment, remove the specific replica row, fall through + r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx) + r.registry.RemoveNodeModel(ctx, node.ID, att.trackingKey, replicaIdx) + xlog.Warn("Backend not reachable for cached model, falling through to reload", + "node", node.Name, "model", att.modelName, "replica", replicaIdx) + return nil } - if r.db != nil { - lockKey := advisorylock.KeyFromString("model-load:" + trackingKey) - var result *RouteResult - lockErr := advisorylock.WithLockCtx(loadCtx, r.db, lockKey, func() error { - var err error - result, err = loadModel(loadCtx) - return err - }) - if lockErr != nil { - return nil, fmt.Errorf("loading model %s: %w", trackingKey, lockErr) - } - return result, nil + // Verify node still matches scheduling constraints + if !r.nodeMatchesScheduling(ctx, node, att.sched) { + r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx) + xlog.Info("Cached model on node that no longer matches selector, falling through", + "node", node.Name, "model", att.trackingKey, "replica", replicaIdx) + return nil } - // No DB (non-distributed) — proceed without lock - return loadModel(loadCtx) + + // Node is alive — FindAndLockNodeWithModel already incremented in-flight as a + // reservation. InFlightTrackingClient handles per-inference tracking, and its + // onFirstComplete callback releases the reservation after the first inference + // call finishes, so in-flight returns to 0 when idle. + r.registry.TouchNodeModel(ctx, node.ID, att.trackingKey, replicaIdx) + r.observePrefix(att.trackingKey, att.observeChain, prefixcache.ReplicaKey{NodeID: node.ID, Replica: replicaIdx}) + grpcClient := r.buildClientForAddr(node, modelAddr, att.parallel) + tracked := NewInFlightTrackingClient(grpcClient, r.registry, node.ID, att.trackingKey, replicaIdx) + return r.newRouteResult(node, att.trackingKey, replicaIdx, grpcClient, tracked) +} + +// coldLoad schedules the model onto a node and loads it, returning a route to +// the replica it landed on. initialInFlight reserves the slot for the calling +// request; the job runner passes 0 because it is loading on nobody's behalf. +func (r *SmartRouter) coldLoad(ctx context.Context, att *routeAttempt, initialInFlight int) (*RouteResult, error) { + result, err := r.scheduleAndLoad(ctx, att.backendType, att.trackingKey, att.modelName, att.modelOpts, att.parallel, initialInFlight) + if err != nil { + return nil, err + } + + // Cold load landed on result.Node replica result.ReplicaIndex: record the + // assignment so subsequent requests with the same prefix prefer it. + r.observePrefix(att.trackingKey, att.observeChain, prefixcache.ReplicaKey{NodeID: result.Node.ID, Replica: result.ReplicaIndex}) + + tracked := NewInFlightTrackingClient(result.Client, r.registry, result.Node.ID, att.trackingKey, result.ReplicaIndex) + return r.newRouteResult(result.Node, att.trackingKey, result.ReplicaIndex, result.Client, tracked), nil +} + +// newColdLoadContext builds the detached, progress-extended context a cold load +// runs under. +// +// Detach the cold load from the caller's context. Staging a model can transfer +// multiple GB to a worker, which takes far longer than any client keeps its +// HTTP request open — a browser refresh, an ingress/LB idle timeout, or a +// round-robined retry landing on another replica all cancel the request +// context. If staging were bound to it, the multi-GB upload aborts with +// "context canceled" mid-transfer and large models can never finish staging +// (the model-load outage). The caller passes context.WithoutCancel, which keeps +// the request's values (prefix chain, etc.) but drops its cancellation. +// +// Detaching must not be unbounded either: a worker that dies mid-install (its +// backend.install never replies) would otherwise leave the job wedged until the +// NATS install deadline alone expires. The backstop is progress-based, not +// wall-clock: staging time is bytes over bandwidth, so a fixed ceiling is a +// model-size cliff (a 70 GB checkpoint transferring healthily at 26 MB/s needs +// ~45m and was killed at exactly 25m00s). The hold extends while the transfer +// reports bytes and expires a stall window after they stop. See load_deadline.go. +func (r *SmartRouter) newColdLoadContext(parent context.Context) (context.Context, context.CancelFunc) { + return newLoadDeadlineContext(parent, r.modelLoadCeiling, r.stagingStallWindow, r.modelLoadAbsoluteMax) } // parseSelectorJSON decodes a JSON node selector string into a map. @@ -1188,6 +1198,7 @@ func (r *SmartRouter) installBackendOnNode(ctx context.Context, node *BackendNod if r.unloader == nil { return "", fmt.Errorf("no NATS connection for backend installation") } + reportLoadPhase(ctx, LoadJobStateInstalling, node, replicaIndex) key := fmt.Sprintf("%s|%s|%s|%d", node.ID, backendType, modelID, replicaIndex) // DoChan rather than Do so this wait honors ctx cancellation. InstallBackend diff --git a/core/services/nodes/router_load_job_test.go b/core/services/nodes/router_load_job_test.go new file mode 100644 index 000000000000..3ba53a7493e4 --- /dev/null +++ b/core/services/nodes/router_load_job_test.go @@ -0,0 +1,219 @@ +package nodes + +import ( + "context" + "runtime" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/testutil" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + "gorm.io/gorm" +) + +// These specs cover the claim/run split: a cold load runs as a durable job +// outside the per-model advisory lock, and concurrent requests attach to it as +// waiters instead of blocking on pg_advisory_lock (which the production role's +// statement_timeout killed at 60s). +var _ = Describe("Route cold-load jobs", func() { + var ( + db *gorm.DB + registry *NodeRegistry + backend *stubBackend + factory *stubClientFactory + unloader *fakeUnloader + node *BackendNode + ) + + BeforeEach(func() { + if runtime.GOOS == "darwin" { + Skip("testcontainers requires Docker, not available on macOS CI") + } + db = testutil.SetupTestDB() + var err error + registry, err = NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + + node = &BackendNode{ + Name: "worker-1", + NodeType: NodeTypeBackend, + Address: "10.0.0.1:50051", + TotalVRAM: 64_000_000_000, + AvailableVRAM: 64_000_000_000, + } + Expect(registry.Register(context.Background(), node, true)).To(Succeed()) + + backend = &stubBackend{healthResult: true, loadResult: &pb.Result{Success: true}} + factory = &stubClientFactory{client: backend} + unloader = &fakeUnloader{ + installReply: &messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:9001"}, + } + }) + + newRouter := func() *SmartRouter { + return NewSmartRouter(registry, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + DB: db, + }) + } + + It("serves a concurrent request for a loading model without a duplicate load", func() { + // The load takes far longer than a request would tolerate holding a + // lock. Both callers must still be served, from ONE load. + release := make(chan struct{}) + unloader.installHook = func() { <-release } + + router := newRouter() + + first := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/big.gguf"}, false) + first <- err + }() + + // Give the first request time to claim and start its job. + Eventually(func() *ModelLoadJob { + job, _ := registry.GetLoadJob(context.Background(), "big-model") + return job + }, 5*time.Second, 50*time.Millisecond).ShouldNot(BeNil()) + + second := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/big.gguf"}, false) + second <- err + }() + + // Neither caller may be stuck behind a lock: the second request must + // still be waiting (not failed) while the load runs. + Consistently(second, 300*time.Millisecond).ShouldNot(Receive()) + + close(release) + + var firstErr, secondErr error + Eventually(first, 15*time.Second).Should(Receive(&firstErr)) + Eventually(second, 15*time.Second).Should(Receive(&secondErr)) + Expect(firstErr).ToNot(HaveOccurred()) + Expect(secondErr).ToNot(HaveOccurred()) + + unloader.mu.Lock() + installs := len(unloader.installCalls) + unloader.mu.Unlock() + Expect(installs).To(Equal(1), "the waiter must attach to the running job, not start a second load") + + // The job row is the record of an IN-FLIGHT load only; NodeModel is the + // record of a loaded model. + Eventually(func() *ModelLoadJob { + job, _ := registry.GetLoadJob(context.Background(), "big-model") + return job + }, 5*time.Second, 50*time.Millisecond).Should(BeNil()) + }) + + It("reports the load's real failure to every waiter", func() { + release := make(chan struct{}) + unloader.installHook = func() { <-release } + unloader.installReply = &messaging.BackendInstallReply{Success: false, Error: "worker out of disk"} + + router := newRouter() + + first := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(context.Background(), "doomed", "models/doomed.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/doomed.gguf"}, false) + first <- err + }() + Eventually(func() *ModelLoadJob { + job, _ := registry.GetLoadJob(context.Background(), "doomed") + return job + }, 5*time.Second, 50*time.Millisecond).ShouldNot(BeNil()) + + second := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(context.Background(), "doomed", "models/doomed.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/doomed.gguf"}, false) + second <- err + }() + + close(release) + + var firstErr, secondErr error + Eventually(first, 15*time.Second).Should(Receive(&firstErr)) + Eventually(second, 15*time.Second).Should(Receive(&secondErr)) + Expect(firstErr).To(HaveOccurred()) + Expect(secondErr).To(HaveOccurred()) + Expect(secondErr.Error()).To(ContainSubstring("worker out of disk"), + "a waiter must learn the real cause, not an anonymous timeout") + }) + + It("returns immediately when the client disconnects, leaving the job running", func() { + release := make(chan struct{}) + unloader.installHook = func() { <-release } + defer close(release) + + router := newRouter() + + go func() { + defer GinkgoRecover() + _, _ = router.Route(context.Background(), "detached", "models/detached.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/detached.gguf"}, false) + }() + Eventually(func() *ModelLoadJob { + job, _ := registry.GetLoadJob(context.Background(), "detached") + return job + }, 5*time.Second, 50*time.Millisecond).ShouldNot(BeNil()) + + ctx, cancel := context.WithCancel(context.Background()) + waiter := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(ctx, "detached", "models/detached.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/detached.gguf"}, false) + waiter <- err + }() + time.Sleep(100 * time.Millisecond) + cancel() + + var waitErr error + Eventually(waiter, 3*time.Second).Should(Receive(&waitErr)) + Expect(waitErr).To(HaveOccurred()) + + // The job is owned by its record, not by the request that triggered it. + job, err := registry.GetLoadJob(context.Background(), "detached") + Expect(err).ToNot(HaveOccurred()) + Expect(job).ToNot(BeNil(), "cancelling a waiter must not abort the load") + }) + + It("heartbeats the job row so a live load is never mistaken for an orphan", func() { + release := make(chan struct{}) + unloader.installHook = func() { <-release } + defer close(release) + + router := newRouter() + go func() { + defer GinkgoRecover() + _, _ = router.Route(context.Background(), "beating", "models/beating.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/beating.gguf"}, false) + }() + + var first *ModelLoadJob + Eventually(func() *ModelLoadJob { + first, _ = registry.GetLoadJob(context.Background(), "beating") + return first + }, 5*time.Second, 50*time.Millisecond).ShouldNot(BeNil()) + + Eventually(func() bool { + job, _ := registry.GetLoadJob(context.Background(), "beating") + return job != nil && job.LastProgress.After(first.LastProgress) + }, 5*time.Second, 200*time.Millisecond).Should(BeTrue(), + "the runner must heartbeat even while no bytes move") + }) +}) From e11d76dfbc4bc999a1fb2c29b3e25b5439fcac90 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 13 Aug 2026 21:41:54 +0000 Subject: [PATCH 4/7] feat(distributed): bound the wait for a loading model and answer with progress MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A request whose model is cold-loading now attaches to the running job and is served the moment the model is ready. That wait has to be bounded: a held HTTP request cannot survive real infrastructure, and an ingress or LB idle timeout kills a twenty-minute request regardless of what LocalAI does. New LOCALAI_MODEL_LOAD_WAIT (default 60s) bounds the CALLER, never the load — the job keeps running either way. On expiry the request gets 503 with Retry-After and a structured body naming the model, the node, the phase, byte progress and an ETA. The `error` envelope keeps OpenAI clients working; `loading` is additive so they ignore it. The ETA comes from the job's own observed rate and is omitted rather than guessed until enough bytes have moved for that rate to mean anything: a confidently wrong ETA on a twenty-minute wait is worse than none. Retry-After is that ETA when known, clamped to [5s, 300s], and the wait budget otherwise. LOCALAI_MODEL_LOAD_WAIT=0 waits unbounded, for deployments with no proxy in front. Zero in the config struct still means "unset, use the default", so the CLI records the operator's zero as ModelLoadWaitUnbounded rather than losing the distinction. The distributed branch of ModelLoader.loadModel wrapped the router's error with %s, which flattened it to a string. Use %w: the typed error is what the HTTP layer keys the 503 off. Assisted-by: Claude Opus 5 [claude-code] --- core/application/distributed.go | 3 + core/cli/run.go | 8 +++ core/config/distributed_config.go | 34 ++++++++++ core/http/app.go | 53 +++++++++++++++ core/schema/model_loading.go | 28 ++++++++ core/services/nodes/load_job_runner.go | 52 ++++++++++++++- core/services/nodes/model_loading_error.go | 72 +++++++++++++++++++++ core/services/nodes/router.go | 9 +++ core/services/nodes/router_load_job_test.go | 59 +++++++++++++++++ docs/content/features/distributed-mode.md | 43 ++++++++++++ pkg/model/loader.go | 5 +- 11 files changed, 364 insertions(+), 2 deletions(-) create mode 100644 core/schema/model_loading.go create mode 100644 core/services/nodes/model_loading_error.go diff --git a/core/application/distributed.go b/core/application/distributed.go index 32269a44a550..e22ac5534361 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -378,6 +378,9 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade cfg.Distributed.BackendInstallTimeoutOrDefault(), cfg.Distributed.ModelLoadTimeoutOrDefault(), ), + // Bounds the REQUEST, not the load: a caller out of budget gets 503 with + // live staging progress while the job keeps running underneath. + ModelLoadWait: cfg.Distributed.ModelLoadWait, }) // Wire staging-progress broadcasting so file-staging shows up on every diff --git a/core/cli/run.go b/core/cli/run.go index 7d35fb693352..53c8cbf4794d 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -179,6 +179,7 @@ type RunCMD struct { BackendInstallTimeout string `env:"LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT" help:"NATS round-trip timeout for backend.install requests sent to worker nodes (default 15m). Increase for slow links pulling multi-GB images." group:"distributed"` BackendUpgradeTimeout string `env:"LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT" help:"NATS round-trip timeout for backend.upgrade requests (default 15m)." group:"distributed"` ModelLoadTimeout string `env:"LOCALAI_NATS_MODEL_LOAD_TIMEOUT" help:"Fixed gRPC deadline for the remote LoadModel call sent to a worker node once its backend is installed and model files are staged. Unset (the default), the deadline is derived from the checkpoint size instead: 5m plus 20s per GiB, capped at 6h, so multi-tens-of-GB diffusion/video checkpoints get the minutes they need without a fixed cliff. Set this only to pin a specific budget; the value is used verbatim, including when it is shorter than the derived one." group:"distributed"` + ModelLoadWait string `env:"LOCALAI_MODEL_LOAD_WAIT" help:"How long an inference request waits for a model that is still cold-loading onto a worker before it is answered with 503, a Retry-After header and live staging progress (default 60s). The request is served the moment the model becomes ready, so a model already most of the way staged needs no client retry. Set to 0 to wait as long as the load takes — only safe when no ingress or load balancer with an idle timeout sits in front." group:"distributed"` NatsAccountSeed string `env:"LOCALAI_NATS_ACCOUNT_SEED" help:"NATS account signing seed (SU...) used to mint per-node worker JWTs at registration" group:"distributed"` NatsServiceJWT string `env:"LOCALAI_NATS_SERVICE_JWT" help:"NATS user JWT for the frontend (and agent workers) to publish control-plane messages" group:"distributed"` NatsServiceSeed string `env:"LOCALAI_NATS_SERVICE_SEED" help:"NATS user signing seed (SU...) paired with LOCALAI_NATS_SERVICE_JWT" group:"distributed"` @@ -386,6 +387,13 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { } opts = append(opts, config.WithModelLoadTimeout(d)) } + if r.ModelLoadWait != "" { + d, err := parseDistributedDuration("LOCALAI_MODEL_LOAD_WAIT", r.ModelLoadWait) + if err != nil { + return err + } + opts = append(opts, config.WithModelLoadWait(d)) + } if r.RegistrationToken != "" { opts = append(opts, config.WithRegistrationToken(r.RegistrationToken)) } diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index 3cbbcf77b790..5a48a84e9b44 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -82,6 +82,16 @@ type DistributedConfig struct { // pipeline init, which for a multi-tens-of-GB diffusion/video checkpoint on // unified memory can far exceed the 5m default. ModelLoadTimeout time.Duration // gRPC deadline for remote LoadModel (default 5m) + // ModelLoadWait bounds how long an inference request waits for a model that + // is being cold-loaded before it is answered with 503 plus live progress. A + // held HTTP request cannot survive real infrastructure — an ingress or LB + // idle timeout kills a 20-minute request regardless of what LocalAI does — + // so the wait is bounded by default. + // + // Zero means unset (DefaultModelLoadWait applies); ModelLoadWaitUnbounded + // records the operator asking for unbounded waiting with + // LOCALAI_MODEL_LOAD_WAIT=0. + ModelLoadWait time.Duration MaxUploadSize int64 // Maximum upload body size in bytes (default 50 GB) @@ -315,6 +325,18 @@ func WithModelLoadTimeout(d time.Duration) AppOption { } } +// WithModelLoadWait sets how long a request waits for a cold-loading model. A +// zero d records the operator asking for unbounded waiting: "set the knob to +// zero" cannot sensibly mean "use the default". +func WithModelLoadWait(d time.Duration) AppOption { + return func(o *ApplicationConfig) { + if d == 0 { + d = ModelLoadWaitUnbounded + } + o.Distributed.ModelLoadWait = d + } +} + var EnableAutoApproveNodes = func(o *ApplicationConfig) { o.Distributed.AutoApproveNodes = true } @@ -379,6 +401,7 @@ const ( FlagBackendInstallTimeout = "backend-install-timeout" FlagBackendUpgradeTimeout = "backend-upgrade-timeout" FlagModelLoadTimeout = "model-load-timeout" + FlagModelLoadWait = "model-load-wait" // FlagDiskHeadroomCheck names the disk-headroom toggle. It is quoted in // the warning the check emits while disabled, so the operator reading a // log line knows exactly which knob produced it. @@ -397,8 +420,19 @@ const ( DefaultBackendInstallTimeout = 15 * time.Minute DefaultBackendUpgradeTimeout = 15 * time.Minute DefaultModelLoadTimeout = 5 * time.Minute + // DefaultModelLoadWait is how long a request waits for a cold-loading model + // before it is answered with 503 and live progress. Chosen to sit under the + // idle timeout of typical ingress/LB defaults, so the answer comes from + // LocalAI (with progress the client can act on) rather than from a proxy + // dropping the connection. + DefaultModelLoadWait = 60 * time.Second ) +// ModelLoadWaitUnbounded records LOCALAI_MODEL_LOAD_WAIT=0 — "wait as long as +// it takes" — which a plain zero cannot express, since zero also means "unset, +// use the default". Only deployments with no proxy in front should use it. +const ModelLoadWaitUnbounded = -1 * time.Second + // DefaultMaxUploadSize is the default maximum upload body size (50 GB). const DefaultMaxUploadSize int64 = 50 << 30 diff --git a/core/http/app.go b/core/http/app.go index 8b08af4adffe..65e8ee917cac 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -78,6 +78,53 @@ func applyModelLoadCooldown(err error, code int, c echo.Context) int { return http.StatusServiceUnavailable } +// respondModelLoading answers a request whose model is still cold-loading with +// 503, a Retry-After header and the live `loading` object, reporting true when +// it handled the error. +// +// The distinction from applyModelLoadCooldown matters: a cooldown means "the +// last load FAILED, back off", this means "the load is progressing, here is how +// far it got". Both used to look like the same anonymous error, so an operator +// watching a 35 GB model stage normally onto a new worker saw only failures. +func respondModelLoading(err error, c echo.Context) bool { + var loadErr *nodes.ModelLoadingError + if !errors.As(err, &loadErr) { + return false + } + setModelLoadingRetryAfter(loadErr, c) + status := loadErr.Status + if jerr := c.JSON(http.StatusServiceUnavailable, schema.ModelLoadingResponse{ + Error: &schema.APIError{ + Message: loadErr.Error(), + Code: "model_loading", + Type: "model_loading", + }, + Loading: &status, + }); jerr != nil { + xlog.Debug("Failed to write model-loading response", "error", jerr) + } + return true +} + +// applyModelLoading is the body-less half of respondModelLoading, for the +// opaque-errors handler: status and Retry-After only. +func applyModelLoading(err error, code int, c echo.Context) int { + var loadErr *nodes.ModelLoadingError + if !errors.As(err, &loadErr) { + return code + } + setModelLoadingRetryAfter(loadErr, c) + return http.StatusServiceUnavailable +} + +func setModelLoadingRetryAfter(loadErr *nodes.ModelLoadingError, c echo.Context) { + secs := int(math.Ceil(loadErr.RetryAfter.Seconds())) + if secs < 1 { + secs = 1 + } + c.Response().Header().Set("Retry-After", strconv.Itoa(secs)) +} + // @title LocalAI API // @version 2.0.0 // @description The LocalAI Rest API. @@ -141,6 +188,9 @@ func API(application *application.Application) (*echo.Echo, error) { // Set error handler if !application.ApplicationConfig().OpaqueErrors { e.HTTPErrorHandler = func(err error, c echo.Context) { + if respondModelLoading(err, c) { + return + } code := http.StatusInternalServerError var he *echo.HTTPError if errors.As(err, &he) { @@ -175,6 +225,9 @@ func API(application *application.Application) (*echo.Echo, error) { code = he.Code } code = applyModelLoadCooldown(err, code, c) + // Opaque errors deliberately withhold the body, so a still-loading + // model gets the status and Retry-After but no progress detail. + code = applyModelLoading(err, code, c) c.NoContent(code) } } diff --git a/core/schema/model_loading.go b/core/schema/model_loading.go new file mode 100644 index 000000000000..3e2927508e25 --- /dev/null +++ b/core/schema/model_loading.go @@ -0,0 +1,28 @@ +package schema + +// ModelLoadingStatus describes a cold load that is still in progress. In +// distributed mode a model can take tens of minutes to stage onto a worker, +// which is far longer than a request may be held; a caller that runs out of +// wait budget gets this instead of an anonymous hang or a misleading error. +type ModelLoadingStatus struct { + Model string `json:"model"` + State string `json:"state"` + Node string `json:"node,omitempty"` + Progress float64 `json:"progress"` + BytesSent int64 `json:"bytes_sent"` + TotalBytes int64 `json:"total_bytes"` + FileIndex int `json:"file_index"` + TotalFiles int `json:"total_files"` + // ETASeconds is omitted rather than guessed until enough bytes have moved + // for the observed rate to mean anything. A confidently wrong ETA on a + // twenty-minute wait is worse than none. + ETASeconds int `json:"eta_seconds,omitempty"` +} + +// ModelLoadingResponse is the 503 body served while a model is still loading. +// The `error` envelope keeps OpenAI-client compatibility; `loading` is additive, +// so existing clients ignore it and load-aware ones can render real progress. +type ModelLoadingResponse struct { + Error *APIError `json:"error,omitempty"` + Loading *ModelLoadingStatus `json:"loading,omitempty"` +} diff --git a/core/services/nodes/load_job_runner.go b/core/services/nodes/load_job_runner.go index ec806cd4265c..bfb18b1e50c8 100644 --- a/core/services/nodes/load_job_runner.go +++ b/core/services/nodes/load_job_runner.go @@ -5,6 +5,7 @@ import ( "fmt" "time" + "github.com/mudler/LocalAI/core/config" "github.com/mudler/xlog" ) @@ -24,6 +25,18 @@ const maxColdLoadRounds = 3 // staging run pinned it for ~20 minutes and every concurrent request died at // the role's 60s statement_timeout with SQLSTATE 57014. func (r *SmartRouter) routeViaLoadJob(ctx context.Context, att *routeAttempt) (*RouteResult, error) { + // A held HTTP request cannot survive real infrastructure: an ingress or LB + // idle timeout kills a twenty-minute request regardless of what LocalAI + // does. So the wait is bounded, and expiry produces a structured answer + // carrying live progress rather than letting the connection die anonymously. + budget := r.loadWaitBudget() + waitCtx := ctx + if budget > 0 { + var cancelWait context.CancelFunc + waitCtx, cancelWait = context.WithTimeout(ctx, budget) + defer cancelWait() + } + for range maxColdLoadRounds { // Register interest BEFORE claiming, so a job that finishes immediately // cannot close the channel before this waiter exists. @@ -60,7 +73,12 @@ func (r *SmartRouter) routeViaLoadJob(ctx context.Context, att *routeAttempt) (* "model", att.trackingKey, "state", job.State, "node", job.NodeName, "owner", job.OwnerReplica) } - if err := r.waitForLoadJob(ctx, att.trackingKey, waiter); err != nil { + if err := r.waitForLoadJob(waitCtx, att.trackingKey, waiter); err != nil { + // The caller's own context is still live, so it was the wait budget + // that ran out, not the client giving up: answer with progress. + if ctx.Err() == nil && waitCtx.Err() != nil { + return nil, r.loadingAnswer(ctx, att.trackingKey, budget) + } return nil, err } @@ -73,6 +91,38 @@ func (r *SmartRouter) routeViaLoadJob(ctx context.Context, att *routeAttempt) (* return nil, fmt.Errorf("loading model %s: the load finished but the model is not available", att.trackingKey) } +// loadWaitBudget resolves the configured wait into a duration, where 0 means +// "no timer — wait as long as the load takes". +func (r *SmartRouter) loadWaitBudget() time.Duration { + switch { + case r.modelLoadWait < 0: // LOCALAI_MODEL_LOAD_WAIT=0 + return 0 + case r.modelLoadWait == 0: // unset + return config.DefaultModelLoadWait + default: + return r.modelLoadWait + } +} + +// loadingAnswer builds the 503 payload for a caller whose wait budget expired, +// reading the job row for live progress. A job that finished in the meantime +// leaves nothing to report, so the caller is told to retry against a plain +// deadline instead. +func (r *SmartRouter) loadingAnswer(ctx context.Context, trackingKey string, budget time.Duration) error { + // The wait context is spent; read on a fresh, short-lived one. + readCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + + job, err := r.registry.GetLoadJob(readCtx, trackingKey) + if err != nil || job == nil { + return fmt.Errorf("timed out waiting for model %s to load", trackingKey) + } + if job.State == LoadJobStateFailed { + return fmt.Errorf("loading model %s: %s", trackingKey, job.LastError) + } + return newModelLoadingError(job, budget) +} + // startLoadJob runs the claimed cold load in the background, detached from the // request that triggered it. The job is owned by its record, not by that // request: the client may disconnect, be retried onto another replica, or time diff --git a/core/services/nodes/model_loading_error.go b/core/services/nodes/model_loading_error.go new file mode 100644 index 000000000000..a7c8d7fb4e05 --- /dev/null +++ b/core/services/nodes/model_loading_error.go @@ -0,0 +1,72 @@ +package nodes + +import ( + "fmt" + "time" + + "github.com/mudler/LocalAI/core/schema" +) + +const ( + // retryAfterFloor and retryAfterCeiling clamp the Retry-After we hand a + // client. Below the floor a client hammers a load that cannot possibly be + // done yet; above the ceiling it stops polling long enough that a model + // which became ready in the meantime sits idle. + retryAfterFloor = 5 * time.Second + retryAfterCeiling = 300 * time.Second +) + +// ModelLoadingError reports that the request's model is still cold-loading and +// the caller's wait budget ran out. It carries live progress so the answer is +// actionable — "staging to nvidia-thor, 41%, ETA ~11m" — rather than an +// anonymous timeout, which is what every UI retry produced before. +type ModelLoadingError struct { + Status schema.ModelLoadingStatus + RetryAfter time.Duration +} + +func (e *ModelLoadingError) Error() string { + msg := fmt.Sprintf("model %s is %s", e.Status.Model, e.Status.State) + if e.Status.Node != "" { + msg += " on node " + e.Status.Node + } + if e.Status.Progress > 0 { + msg += fmt.Sprintf(" (%.0f%%", e.Status.Progress) + if e.Status.ETASeconds > 0 { + msg += fmt.Sprintf(", ETA ~%s", (time.Duration(e.Status.ETASeconds) * time.Second).Round(time.Minute)) + } + msg += ")" + } + return msg +} + +// LoadingStatus renders a job row as the API's `loading` object. +func LoadingStatus(job *ModelLoadJob) schema.ModelLoadingStatus { + status := schema.ModelLoadingStatus{ + Model: job.TrackingKey, + State: job.State, + Node: job.NodeName, + Progress: job.Progress(), + BytesSent: job.BytesSent, + TotalBytes: job.TotalBytes, + FileIndex: job.FileIndex, + TotalFiles: job.TotalFiles, + } + if eta, ok := job.ETA(time.Now()); ok { + status.ETASeconds = int(eta.Seconds()) + } + return status +} + +// newModelLoadingError builds the 503 answer for a caller whose wait budget +// expired. Retry-After is the ETA when the job has one, clamped so it stays a +// useful poll interval, and the caller's own budget otherwise. +func newModelLoadingError(job *ModelLoadJob, budget time.Duration) *ModelLoadingError { + status := LoadingStatus(job) + retryAfter := budget + if status.ETASeconds > 0 { + retryAfter = time.Duration(status.ETASeconds) * time.Second + } + retryAfter = min(max(retryAfter, retryAfterFloor), retryAfterCeiling) + return &ModelLoadingError{Status: status, RetryAfter: retryAfter} +} diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 07816c5bf1bc..9322ed740cc5 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -105,6 +105,11 @@ type SmartRouterOptions struct { // arriving, so a peer trickling bytes forever cannot pin the advisory lock // indefinitely. Zero selects modelLoadAbsoluteMax (24h). ModelLoadAbsoluteMax time.Duration + // ModelLoadWait bounds how long a REQUEST waits for a cold load that is + // already running before it is answered with live progress. It bounds the + // caller, never the load: the job keeps running either way. Zero selects + // config.DefaultModelLoadWait; config.ModelLoadWaitUnbounded waits forever. + ModelLoadWait time.Duration } // modelLoadStagingMargin is the slack ModelLoadCeilingFor adds on top of the @@ -187,6 +192,9 @@ type SmartRouter struct { // hard countdown into a progress-extended hold (see load_deadline.go). stagingStallWindow time.Duration modelLoadAbsoluteMax time.Duration + // modelLoadWait bounds the REQUEST's wait for a running cold load, not the + // load itself (see SmartRouterOptions.ModelLoadWait). + modelLoadWait time.Duration // loadWaiters is one broadcast channel per model being cold-loaded, closed // when the job reaches a terminal state. Same-model waiters all want the // identical outcome, so they share one wait instead of queueing. See @@ -244,6 +252,7 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter // ceiling, so nothing to normalize here. stagingStallWindow: opts.StagingStallWindow, modelLoadAbsoluteMax: opts.ModelLoadAbsoluteMax, + modelLoadWait: opts.ModelLoadWait, loadWaiters: map[string]chan struct{}{}, } } diff --git a/core/services/nodes/router_load_job_test.go b/core/services/nodes/router_load_job_test.go index 3ba53a7493e4..f67e1399b5fc 100644 --- a/core/services/nodes/router_load_job_test.go +++ b/core/services/nodes/router_load_job_test.go @@ -2,12 +2,14 @@ package nodes import ( "context" + "errors" "runtime" "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/testutil" pb "github.com/mudler/LocalAI/pkg/grpc/proto" @@ -192,6 +194,63 @@ var _ = Describe("Route cold-load jobs", func() { Expect(job).ToNot(BeNil(), "cancelling a waiter must not abort the load") }) + It("answers with live progress once the wait budget is spent", func() { + release := make(chan struct{}) + unloader.installHook = func() { <-release } + defer close(release) + + router := NewSmartRouter(registry, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + DB: db, + ModelLoadWait: 500 * time.Millisecond, + }) + + start := time.Now() + _, err := router.Route(context.Background(), "slow-model", "models/slow.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/slow.gguf"}, false) + Expect(err).To(HaveOccurred()) + Expect(time.Since(start)).To(BeNumerically("<", 10*time.Second)) + + var loadingErr *ModelLoadingError + Expect(errors.As(err, &loadingErr)).To(BeTrue(), + "a caller out of budget must get a structured answer, not an anonymous timeout") + Expect(loadingErr.Status.Model).To(Equal("slow-model")) + Expect(loadingErr.Status.State).ToNot(BeEmpty()) + Expect(loadingErr.RetryAfter).To(BeNumerically(">", 0)) + + // The job is untouched: the caller gave up, the load did not. + job, err := registry.GetLoadJob(context.Background(), "slow-model") + Expect(err).ToNot(HaveOccurred()) + Expect(job).ToNot(BeNil()) + }) + + It("waits unbounded when the budget is explicitly disabled", func() { + release := make(chan struct{}) + unloader.installHook = func() { <-release } + + router := NewSmartRouter(registry, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + DB: db, + ModelLoadWait: config.ModelLoadWaitUnbounded, + }) + + done := make(chan error, 1) + go func() { + defer GinkgoRecover() + _, err := router.Route(context.Background(), "patient", "models/patient.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/patient.gguf"}, false) + done <- err + }() + + // Well past any default budget shrunk for tests; the caller must still + // be waiting rather than 503-ing. + Consistently(done, 2*time.Second).ShouldNot(Receive()) + close(release) + Eventually(done, 15*time.Second).Should(Receive(BeNil())) + }) + It("heartbeats the job row so a live load is never mistaken for an orphan", func() { release := make(chan struct{}) unloader.installHook = func() { <-release } diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index f3a0f6534d1d..bf5a872619fc 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -75,6 +75,7 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These | `--backend-install-timeout` | `LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT` | `15m` | How long the frontend waits for a worker to acknowledge a backend install before considering the request stalled. Raise it when workers pull large backend images over slow links. If a worker takes longer than this, the operation shows as "still installing in background" in the admin UI and clears once the worker finishes. | | `--backend-upgrade-timeout` | `LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT` | `15m` | Same as the install timeout, applied to backend upgrades (force-reinstall). | | `--model-load-timeout` | `LOCALAI_NATS_MODEL_LOAD_TIMEOUT` | *(derived from checkpoint size)* | Pins the deadline for the `LoadModel` gRPC call the frontend issues to a worker. Leave it unset: by default the deadline is **derived from the checkpoint's on-disk size** (see below), which is what the worker actually spends its load time reading. Set it only to pin a specific budget — the value is then used verbatim, including when it is *shorter* than the derived one, so an operator who wants fast failure gets it. | +| *(env only)* | `LOCALAI_MODEL_LOAD_WAIT` | `60s` | How long an inference request waits for a model that is still cold-loading onto a worker before it is answered with `503`, a `Retry-After` header and live staging progress. The request is served the moment the model becomes ready, so a model already most of the way staged needs no client retry. Set to `0` to wait as long as the load takes — only safe when no ingress or load balancer with an idle timeout sits in front. See [Requests for a model that is still loading](#requests-for-a-model-that-is-still-loading). | | `--expose-node-header` | `LOCALAI_EXPOSE_NODE_HEADER` | `false` | When enabled, inference responses carry an `X-LocalAI-Node` header with the ID of the worker node that served the request. Coverage spans the OpenAI-compatible endpoints (chat completions, completions, embeddings, audio transcriptions, audio speech / TTS, image generations, image inpainting), the Jina rerank endpoint (`/v1/rerank`), the VAD endpoints (`/v1/vad`, `/vad`), and the Anthropic Messages (`/v1/messages`) and Ollama (`/api/chat`, `/api/generate`, `/api/embed`) shims. Useful for debugging, observability and load-balancer attribution. Off by default: the node ID reveals internal cluster topology and should not be exposed on a public endpoint. Best-effort: under heavy concurrency for the same model across multiple replicas, the header may reflect a recent routing decision rather than this exact request's. Acceptable for observability and debugging. | ### The model load deadline scales with the checkpoint @@ -111,6 +112,48 @@ While **model files are staging**, however, the deadline extends every time stag An absolute cap of 24h ends the hold even if progress keeps arriving, so a degenerate peer trickling a few bytes at a time cannot pin the lock forever. No configuration is needed for either value; both are sized well above any legitimate transfer. +### Requests for a model that is still loading + +A cold load in distributed mode is a long-running background job: install the backend, stage multi-GB model files to the worker, then load the checkpoint. Staging a 35.7 GB GGUF onto a fresh worker takes roughly twenty minutes on a fast LAN — far longer than any HTTP request can be held open. + +So the load does **not** run on the request. The first request for an unloaded model claims a durable **model load job** — that claim takes milliseconds and is the only part that holds the per-model advisory lock — and the job then runs in the background on the frontend replica that claimed it. Every other request for the same model, on any replica, attaches to that job as a waiter: + +- It is **served the moment the model is ready**, with no client-side retry. A model already 90% staged usually needs no second request. +- It never starts a duplicate load and never blocks on the database lock. (Before this split, concurrent requests blocked on `pg_advisory_lock` for the whole load and were killed by the PostgreSQL role's `statement_timeout` — `SQLSTATE 57014` — so from the operator's seat the model simply never loaded.) +- If the load fails, the waiter gets the *real* cause (`worker out of disk`), not an anonymous timeout. +- If the client disconnects, the load keeps going. It belongs to the job record, not to the request. + +When the wait budget (`LOCALAI_MODEL_LOAD_WAIT`, default `60s`) runs out, the request is answered with `503`, a `Retry-After` header, and a body that says exactly where the load is: + +```json +{ + "error": { + "message": "model Qwen3.6-27B-MTP-GGUF is staging on node nvidia-thor (41%, ETA ~11m)", + "type": "model_loading", + "code": "model_loading" + }, + "loading": { + "model": "Qwen3.6-27B-MTP-GGUF", + "state": "staging", + "node": "nvidia-thor", + "progress": 41.2, + "bytes_sent": 14730000000, + "total_bytes": 35776484480, + "file_index": 1, + "total_files": 2, + "eta_seconds": 660 + } +} +``` + +The `error` envelope keeps OpenAI clients working unchanged; `loading` is additive, so a client that understands it renders progress instead of an error. `eta_seconds` is derived from the job's own observed transfer rate and is **omitted rather than guessed** until enough bytes have moved for that rate to mean anything — a confidently wrong ETA on a twenty-minute wait is worse than none. `state` is one of `pending` (choosing a node), `installing`, `staging` (transferring files) or `loading` (the worker is reading the checkpoint). + +The chat UI renders this state inline and retries automatically once the model reports ready. Poll `GET /api/models/{id}/load-status` for the same `loading` object at any time. + +{{% notice note %}} +A frontend replica that dies mid-load does not wedge the model: the job row carries a heartbeat and another replica reclaims a job whose heartbeat has stopped. The heartbeat is time-based, not byte-based, because a checkpoint load legitimately transfers zero bytes for many minutes. +{{% /notice %}} + ### NATS JWT authentication (recommended for production) By default, NATS connections are anonymous: any client that can reach port `4222` may publish control-plane subjects such as `nodes..backend.install`. Enable JWT auth to scope workers to their own node subjects and give the frontend a dedicated service credential. diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 6e0abee6926c..fd54d358f591 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -469,7 +469,10 @@ func (ml *ModelLoader) loadModel(modelID, modelName, modelFileName string, loade modelFile := filepath.Join(ml.ModelPath, modelFileName) model, err := loader(modelID, modelName, modelFile) if err != nil { - return nil, fmt.Errorf("failed to route model with internal loader: %s", err) + // %w, not %s: the router reports a still-loading model as a typed + // error that the HTTP layer turns into 503 plus live progress, and + // that only survives an unbroken chain. + return nil, fmt.Errorf("failed to route model with internal loader: %w", err) } if model == nil { return nil, fmt.Errorf("loader didn't return a model") From e648f1b802768f2d56f05abc7d985b5705bca157 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 13 Aug 2026 21:49:48 +0000 Subject: [PATCH 5/7] feat(api): add GET /api/models/{id}/load-status MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A client that receives 503 while a model stages onto a worker needs somewhere to poll. This returns the same `loading` object the 503 carries — phase, node, byte progress and ETA — or 404 when no load is running. Read-only and observability-shaped, so it is deliberately neither admin-gated nor feature-gated: it explains a 503 the caller just received, and hiding that behind a per-modality feature would make the explanation for a failed image request depend on chat permissions. It also gets no MCP tool, since there is nothing here an admin would manage conversationally. Registered on the surfaces from .agents/api-endpoints-and-auth.md: the swagger block (existing `models` tag, so /api/instructions needs no new area), the endpoint discovery maps in RegisterLocalAIRoutes, regenerated swagger, and the distributed-mode docs page. No FLAG_* usecase is involved, so capabilities.js is unchanged. Assisted-by: Claude Opus 5 [claude-code] --- .../endpoints/localai/model_load_status.go | 66 ++++++++++++++++ .../localai/model_load_status_test.go | 75 +++++++++++++++++++ core/http/routes/localai.go | 15 ++++ swagger/docs.go | 68 +++++++++++++++++ swagger/swagger.json | 68 +++++++++++++++++ swagger/swagger.yaml | 51 +++++++++++++ 6 files changed, 343 insertions(+) create mode 100644 core/http/endpoints/localai/model_load_status.go create mode 100644 core/http/endpoints/localai/model_load_status_test.go diff --git a/core/http/endpoints/localai/model_load_status.go b/core/http/endpoints/localai/model_load_status.go new file mode 100644 index 000000000000..70cebdcaa9f9 --- /dev/null +++ b/core/http/endpoints/localai/model_load_status.go @@ -0,0 +1,66 @@ +package localai + +import ( + "net/http" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/nodes" +) + +// ModelLoadStatusEndpoint reports the progress of a cold load that is +// currently running for a model, so a client that got a 503 while the model +// stages onto a worker can poll rather than blind-retry. +// +// Read-only and observability-shaped: it is not admin-gated (a caller allowed +// to ask for inference on a model may see why it is not answering yet) and it +// is not feature-gated, since a per-capability gate would make the explanation +// for a 503 depend on which modality the model happens to be. +// +// @Summary Report the progress of an in-flight model load. +// @Description Returns the live state of a distributed cold load — phase, node, byte progress and ETA — or 404 when no load is running for the model. This is the same `loading` object the 503 response carries while a model is still staging. +// @Tags models +// @Produce json +// @Param id path string true "Model ID" +// @Success 200 {object} schema.ModelLoadingStatus "Live load progress" +// @Failure 404 {object} schema.ErrorResponse "No load is running for this model" +// @Router /api/models/{id}/load-status [get] +func ModelLoadStatusEndpoint(loadJobs func() nodes.LoadJobStore) echo.HandlerFunc { + return func(c echo.Context) error { + modelID := c.Param("id") + if modelID == "" { + return c.JSON(http.StatusBadRequest, schema.ErrorResponse{ + Error: &schema.APIError{Message: "model id is required", Code: http.StatusBadRequest, Type: "invalid_request_error"}, + }) + } + + notLoading := schema.ErrorResponse{ + Error: &schema.APIError{ + Message: "no load is running for model " + modelID, + Code: http.StatusNotFound, + Type: "not_found_error", + }, + } + + // Cold-load jobs are a distributed-mode concept: a single-host load is + // synchronous and has no job to report on. + var store nodes.LoadJobStore + if loadJobs != nil { + store = loadJobs() + } + if store == nil { + return c.JSON(http.StatusNotFound, notLoading) + } + + job, err := store.GetLoadJob(c.Request().Context(), modelID) + if err != nil { + return c.JSON(http.StatusInternalServerError, schema.ErrorResponse{ + Error: &schema.APIError{Message: err.Error(), Code: http.StatusInternalServerError, Type: "server_error"}, + }) + } + if job == nil { + return c.JSON(http.StatusNotFound, notLoading) + } + return c.JSON(http.StatusOK, nodes.LoadingStatus(job)) + } +} diff --git a/core/http/endpoints/localai/model_load_status_test.go b/core/http/endpoints/localai/model_load_status_test.go new file mode 100644 index 000000000000..9f989a54868a --- /dev/null +++ b/core/http/endpoints/localai/model_load_status_test.go @@ -0,0 +1,75 @@ +package localai_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + + "github.com/labstack/echo/v4" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/http/endpoints/localai" + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" +) + +var _ = Describe("ModelLoadStatusEndpoint", func() { + get := func(store func() nodes.LoadJobStore, modelID string) *httptest.ResponseRecorder { + e := echo.New() + e.GET("/api/models/:id/load-status", localai.ModelLoadStatusEndpoint(store)) + req := httptest.NewRequest(http.MethodGet, "/api/models/"+modelID+"/load-status", nil) + rec := httptest.NewRecorder() + e.ServeHTTP(rec, req) + return rec + } + + It("404s when the server is not running distributed", func() { + rec := get(nil, "some-model") + Expect(rec.Code).To(Equal(http.StatusNotFound)) + }) + + Context("with a registry", func() { + var registry *nodes.NodeRegistry + + BeforeEach(func() { + db := testutil.SetupTestDB() + var err error + registry, err = nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + }) + + store := func(r *nodes.NodeRegistry) func() nodes.LoadJobStore { + return func() nodes.LoadJobStore { return r } + } + + It("404s when no load is running for the model", func() { + rec := get(store(registry), "idle-model") + Expect(rec.Code).To(Equal(http.StatusNotFound)) + }) + + It("reports the live progress of a running load", func() { + ctx := context.Background() + _, claimed, err := registry.ClaimLoadJob(ctx, "big-model", "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(claimed).To(BeTrue()) + Expect(registry.UpdateLoadJob(ctx, "big-model", nodes.LoadJobUpdate{ + State: nodes.LoadJobStateStaging, NodeID: "n1", NodeName: "nvidia-thor", + BytesSent: 1000, TotalBytes: 4000, FileIndex: 1, TotalFiles: 1, + })).To(Succeed()) + + rec := get(store(registry), "big-model") + Expect(rec.Code).To(Equal(http.StatusOK)) + + var status schema.ModelLoadingStatus + Expect(json.Unmarshal(rec.Body.Bytes(), &status)).To(Succeed()) + Expect(status.Model).To(Equal("big-model")) + Expect(status.State).To(Equal(nodes.LoadJobStateStaging)) + Expect(status.Node).To(Equal("nvidia-thor")) + Expect(status.Progress).To(BeNumerically("~", 25, 0.01)) + Expect(status.TotalBytes).To(Equal(int64(4000))) + }) + }) +}) diff --git a/core/http/routes/localai.go b/core/http/routes/localai.go index bfed2d9edb94..30a2803f7ae2 100644 --- a/core/http/routes/localai.go +++ b/core/http/routes/localai.go @@ -11,6 +11,7 @@ import ( "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/monitoring" + "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/templates" "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/model" @@ -154,6 +155,18 @@ func RegisterLocalAIRoutes(router *echo.Echo, // Forget does not load a voice model — it only needs the registry. router.POST("/v1/voice/forget", localai.VoiceForgetEndpoint(app.VoiceRegistry())) + // Progress of an in-flight cold load. Standard auth only: it explains a 503 + // the caller just received, so gating it behind admin (or a per-modality + // feature) would hide the explanation from exactly the client that needs it. + // Resolved per request, not at registration: distributed services are wired + // during startup and a snapshot taken here could be nil forever. + router.GET("/api/models/:id/load-status", localai.ModelLoadStatusEndpoint(func() nodes.LoadJobStore { + if d := app.Distributed(); d != nil && d.Registry != nil { + return d.Registry + } + return nil + })) + voiceProfiles := app.VoiceProfileStore() router.GET("/api/voice-profiles", localai.ListVoiceProfilesEndpoint(voiceProfiles)) router.GET("/api/voice-profiles/:id/audio", localai.ServeVoiceProfileAudioEndpoint(voiceProfiles)) @@ -309,6 +322,7 @@ func RegisterLocalAIRoutes(router *echo.Echo, "config_patch": "/api/models/config-json/:name", "autocomplete": "/api/models/config-metadata/autocomplete/:provider", "vram_estimate": "/api/models/vram-estimate", + "model_load_status": "/api/models/:id/load-status", "tts": "/tts", "voice_profiles": "/api/voice-profiles", "transcription": "/v1/audio/transcriptions", @@ -344,6 +358,7 @@ func RegisterLocalAIRoutes(router *echo.Echo, "import": "/models/import", "reload": "/models/reload", "list_aliases": "/api/aliases", + "load_status": "/api/models/:id/load-status", }, "ai_functions": map[string]string{ "tts": "/tts", diff --git a/swagger/docs.go b/swagger/docs.go index 14b0f335232f..d785faa686d8 100644 --- a/swagger/docs.go +++ b/swagger/docs.go @@ -1097,6 +1097,41 @@ const docTemplate = `{ } } }, + "/api/models/{id}/load-status": { + "get": { + "description": "Returns the live state of a distributed cold load — phase, node, byte progress and ETA — or 404 when no load is running for the model. This is the same ` + "`" + `loading` + "`" + ` object the 503 response carries while a model is still staging.", + "produces": [ + "application/json" + ], + "tags": [ + "models" + ], + "summary": "Report the progress of an in-flight model load.", + "parameters": [ + { + "type": "string", + "description": "Model ID", + "name": "id", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "Live load progress", + "schema": { + "$ref": "#/definitions/schema.ModelLoadingStatus" + } + }, + "404": { + "description": "No load is running for this model", + "schema": { + "$ref": "#/definitions/schema.ErrorResponse" + } + } + } + } + }, "/api/models/{name}/{action}": { "put": { "description": "Enable or disable a model from being loaded on demand. Disabled models remain installed but cannot be loaded.", @@ -6064,6 +6099,39 @@ const docTemplate = `{ } } }, + "schema.ModelLoadingStatus": { + "type": "object", + "properties": { + "bytes_sent": { + "type": "integer" + }, + "eta_seconds": { + "description": "ETASeconds is omitted rather than guessed until enough bytes have moved\nfor the observed rate to mean anything. A confidently wrong ETA on a\ntwenty-minute wait is worse than none.", + "type": "integer" + }, + "file_index": { + "type": "integer" + }, + "model": { + "type": "string" + }, + "node": { + "type": "string" + }, + "progress": { + "type": "number" + }, + "state": { + "type": "string" + }, + "total_bytes": { + "type": "integer" + }, + "total_files": { + "type": "integer" + } + } + }, "schema.ModelsDataResponse": { "type": "object", "properties": { diff --git a/swagger/swagger.json b/swagger/swagger.json index 4181a0e82bd7..8b7f4803aa16 100644 --- a/swagger/swagger.json +++ b/swagger/swagger.json @@ -1094,6 +1094,41 @@ } } }, + "/api/models/{id}/load-status": { + "get": { + "description": "Returns the live state of a distributed cold load — phase, node, byte progress and ETA — or 404 when no load is running for the model. This is the same `loading` object the 503 response carries while a model is still staging.", + "produces": [ + "application/json" + ], + "tags": [ + "models" + ], + "summary": "Report the progress of an in-flight model load.", + "parameters": [ + { + "type": "string", + "description": "Model ID", + "name": "id", + "in": "path", + "required": true + } + ], + "responses": { + "200": { + "description": "Live load progress", + "schema": { + "$ref": "#/definitions/schema.ModelLoadingStatus" + } + }, + "404": { + "description": "No load is running for this model", + "schema": { + "$ref": "#/definitions/schema.ErrorResponse" + } + } + } + } + }, "/api/models/{name}/{action}": { "put": { "description": "Enable or disable a model from being loaded on demand. Disabled models remain installed but cannot be loaded.", @@ -6061,6 +6096,39 @@ } } }, + "schema.ModelLoadingStatus": { + "type": "object", + "properties": { + "bytes_sent": { + "type": "integer" + }, + "eta_seconds": { + "description": "ETASeconds is omitted rather than guessed until enough bytes have moved\nfor the observed rate to mean anything. A confidently wrong ETA on a\ntwenty-minute wait is worse than none.", + "type": "integer" + }, + "file_index": { + "type": "integer" + }, + "model": { + "type": "string" + }, + "node": { + "type": "string" + }, + "progress": { + "type": "number" + }, + "state": { + "type": "string" + }, + "total_bytes": { + "type": "integer" + }, + "total_files": { + "type": "integer" + } + } + }, "schema.ModelsDataResponse": { "type": "object", "properties": { diff --git a/swagger/swagger.yaml b/swagger/swagger.yaml index 8b66f6892117..f944d88f882a 100644 --- a/swagger/swagger.yaml +++ b/swagger/swagger.yaml @@ -1583,6 +1583,31 @@ definitions: an error). type: string type: object + schema.ModelLoadingStatus: + properties: + bytes_sent: + type: integer + eta_seconds: + description: |- + ETASeconds is omitted rather than guessed until enough bytes have moved + for the observed rate to mean anything. A confidently wrong ETA on a + twenty-minute wait is worse than none. + type: integer + file_index: + type: integer + model: + type: string + node: + type: string + progress: + type: number + state: + type: string + total_bytes: + type: integer + total_files: + type: integer + type: object schema.ModelsDataResponse: properties: data: @@ -3454,6 +3479,32 @@ paths: summary: Get an instruction's API guide or OpenAPI fragment tags: - instructions + /api/models/{id}/load-status: + get: + description: Returns the live state of a distributed cold load — phase, node, + byte progress and ETA — or 404 when no load is running for the model. This + is the same `loading` object the 503 response carries while a model is still + staging. + parameters: + - description: Model ID + in: path + name: id + required: true + type: string + produces: + - application/json + responses: + "200": + description: Live load progress + schema: + $ref: '#/definitions/schema.ModelLoadingStatus' + "404": + description: No load is running for this model + schema: + $ref: '#/definitions/schema.ErrorResponse' + summary: Report the progress of an in-flight model load. + tags: + - models /api/models/{name}/{action}: put: description: Enable or disable a model from being loaded on demand. Disabled From 9b4ac35d47fac8417c3de6c65d5e81ed48900882 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 13 Aug 2026 22:01:04 +0000 Subject: [PATCH 6/7] feat(ui): show cold-load progress in Chat and retry when the model is ready MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A chat request for a model that is still staging onto a worker now gets a 503 carrying live progress instead of an error. Render it: the composer shows the phase (installing / staging / loading), the node, the percent and the ETA, then polls load-status and re-sends the request the moment the model is ready. Reuses the staging progress idiom the page already had rather than inventing a second one — the two sources are folded into one loadProgress, with the load job winning because it is authoritative across frontend replicas and knows the phase, where the staging operation only knows about a byte transfer this replica happens to be performing. Waiting is bounded (three send attempts, ~30 min of polling each), so a load that never finishes still surfaces as an error rather than as a spinner nobody questions. An aborted generation stops the polling too. Assisted-by: Claude Opus 5 [claude-code] --- .../react-ui/e2e/chat-model-loading.spec.js | 87 +++++++++++++++++++ .../http/react-ui/public/locales/en/chat.json | 10 ++- core/http/react-ui/src/hooks/useChat.js | 76 +++++++++++++++- core/http/react-ui/src/pages/Chat.jsx | 55 ++++++++++-- core/http/react-ui/src/utils/config.js | 3 + 5 files changed, 218 insertions(+), 13 deletions(-) create mode 100644 core/http/react-ui/e2e/chat-model-loading.spec.js diff --git a/core/http/react-ui/e2e/chat-model-loading.spec.js b/core/http/react-ui/e2e/chat-model-loading.spec.js new file mode 100644 index 000000000000..e8e564d21ac5 --- /dev/null +++ b/core/http/react-ui/e2e/chat-model-loading.spec.js @@ -0,0 +1,87 @@ +import { test, expect } from './coverage-fixtures.js' + +// A model that is cold-loading onto a worker answers 503 with a `loading` +// object rather than an error. The chat must show that as progress and retry +// itself once the load finishes — the whole point of the change: an operator +// watching a 35 GB model stage normally used to see only repeated failures. +test.describe('Chat - model loading', () => { + test('renders staging progress on 503 and retries when the model is ready', async ({ page }) => { + await page.route('**/api/models/capabilities', (route) => { + route.fulfill({ + contentType: 'application/json', + body: JSON.stringify({ data: [{ id: 'test-model', capabilities: ['FLAG_CHAT'] }] }), + }) + }) + + let completions = 0 + await page.route('**/v1/chat/completions', (route) => { + completions++ + if (completions === 1) { + route.fulfill({ + status: 503, + contentType: 'application/json', + headers: { 'Retry-After': '5' }, + body: JSON.stringify({ + error: { + message: 'model test-model is staging on node nvidia-thor (41%)', + type: 'model_loading', + code: 'model_loading', + }, + loading: { + model: 'test-model', + state: 'staging', + node: 'nvidia-thor', + progress: 41.2, + bytes_sent: 14730000000, + total_bytes: 35776484480, + file_index: 1, + total_files: 2, + eta_seconds: 660, + }, + }), + }) + return + } + route.fulfill({ + status: 200, + contentType: 'text/event-stream', + body: + 'data: {"choices":[{"delta":{"content":"loaded at last"}}]}\n\n' + + 'data: [DONE]\n\n', + }) + }) + + // First poll still staging, then the job is gone: the model is ready. + let polls = 0 + await page.route('**/api/models/*/load-status', (route) => { + polls++ + if (polls === 1) { + route.fulfill({ + contentType: 'application/json', + body: JSON.stringify({ + model: 'test-model', state: 'staging', node: 'nvidia-thor', + progress: 62.5, bytes_sent: 22000000000, total_bytes: 35776484480, + file_index: 2, total_files: 2, eta_seconds: 300, + }), + }) + return + } + route.fulfill({ status: 404, contentType: 'application/json', body: '{}' }) + }) + + await page.goto('/app/chat') + await expect(page.getByRole('button', { name: 'test-model' })).toBeVisible({ timeout: 10_000 }) + + await page.locator('.chat-input').fill('Hello') + await page.locator('.chat-send-btn').click() + + // The 503 is shown as progress, not as an error. + await expect(page.locator('.chat-staging-progress')).toBeVisible({ timeout: 10_000 }) + await expect(page.locator('.chat-staging-label')).toContainText('nvidia-thor') + await expect(page.locator('.chat-staging-pct')).toContainText('41%') + + // ...and the request retries itself once the load reports ready. + await expect(page.getByText('loaded at last')).toBeVisible({ timeout: 25_000 }) + expect(completions).toBeGreaterThan(1) + }) +}) diff --git a/core/http/react-ui/public/locales/en/chat.json b/core/http/react-ui/public/locales/en/chat.json index 0e7af136b8d5..f004cee56526 100644 --- a/core/http/react-ui/public/locales/en/chat.json +++ b/core/http/react-ui/public/locales/en/chat.json @@ -81,7 +81,15 @@ }, "streaming": { "transferring": "Transferring model...", - "transferringTo": "Transferring model to {{node}}..." + "transferringTo": "Transferring model to {{node}}...", + "onNode": "on {{node}}", + "eta": "~{{value}} left", + "modelState": { + "pending": "Preparing to load the model...", + "installing": "Installing the backend...", + "staging": "Staging model files...", + "loading": "Loading the model..." + } }, "tokens": { "perSec": "{{count}} tok/s", diff --git a/core/http/react-ui/src/hooks/useChat.js b/core/http/react-ui/src/hooks/useChat.js index 2bcf955ce116..567cf7f55cfa 100644 --- a/core/http/react-ui/src/hooks/useChat.js +++ b/core/http/react-ui/src/hooks/useChat.js @@ -16,6 +16,68 @@ async function extractHttpError(response) { return errorMsg } +// How long the UI keeps waiting for a model that is staging onto a worker. +// Staging a multi-GB checkpoint runs for tens of minutes — far longer than the +// server's own per-request wait budget — so the poll has to outlive several +// 503s. Bounded, so a load that never finishes eventually surfaces as an error +// rather than as a spinner nobody questions. +const MODEL_LOAD_POLL_INTERVAL = 3000 +const MODEL_LOAD_MAX_ATTEMPTS = 3 +const MODEL_LOAD_MAX_POLLS = 600 // ~30 min per attempt + +// readModelLoading returns the `loading` object from the 503 the server sends +// while a model is still cold-loading, or null for any other failure. It reads +// a clone so the caller can still parse the body for its error message. +async function readModelLoading(response) { + if (response.status !== 503) return null + try { + const data = await response.clone().json() + if (data?.error?.type !== 'model_loading') return null + return { ...(data.loading || {}), message: data.error?.message } + } catch { + return null + } +} + +// waitForModelReady polls load-status until the model finishes loading (404 — +// no job left), the load fails, or the caller aborts. Returns true when it is +// worth re-sending the request. +async function waitForModelReady(modelID, onProgress, signal) { + for (let i = 0; i < MODEL_LOAD_MAX_POLLS; i++) { + await new Promise(resolve => setTimeout(resolve, MODEL_LOAD_POLL_INTERVAL)) + if (signal?.aborted) return false + try { + const res = await fetch(apiUrl(API_CONFIG.endpoints.modelLoadStatus(modelID)), { signal }) + if (res.status === 404) return true // job gone: loaded, or worth one retry + if (!res.ok) return false + const status = await res.json() + if (status?.state === 'failed') return false + onProgress(status) + } catch { + return false + } + } + return false +} + +// fetchWithModelLoadWait issues the request and, when the model is still +// staging onto a worker, waits for it rather than surfacing an error. The +// server answers 503 within its own wait budget so no connection is held for +// the whole load; the UI picks the wait back up here and retries once ready. +async function fetchWithModelLoadWait(url, init, modelID, onLoading, signal) { + let response = await fetch(url, init) + for (let attempt = 0; attempt < MODEL_LOAD_MAX_ATTEMPTS; attempt++) { + const loading = await readModelLoading(response) + if (!loading) return response + onLoading(loading) + const ready = await waitForModelReady(loading.model || modelID, onLoading, signal) + onLoading(null) + if (!ready) return response + response = await fetch(url, init) + } + return response +} + function extractThinking(text) { let regularContent = '' let thinkingContent = '' @@ -125,6 +187,10 @@ export function useChat(initialModel = '') { const [streamingToolCalls, setStreamingToolCalls] = useState([]) const [tokensPerSecond, setTokensPerSecond] = useState(null) const [maxTokensPerSecond, setMaxTokensPerSecond] = useState(null) + // Live progress of a cold load the request is waiting on, so the composer can + // say "staging to nvidia-thor, 41%" instead of showing an error for a model + // that is loading exactly as it should. + const [modelLoading, setModelLoading] = useState(null) const abortControllerRef = useRef(null) const startTimeRef = useRef(null) const tokenCountRef = useRef(0) @@ -360,12 +426,12 @@ export function useChat(initialModel = '') { // Legacy MCP SSE streaming (custom event types from /v1/mcp/chat/completions) try { const timeoutId = setTimeout(() => controller.abort(), 300000) // 5 min timeout - const response = await fetch(apiUrl(endpoint), { + const response = await fetchWithModelLoadWait(apiUrl(endpoint), { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(requestBody), signal: controller.signal, - }) + }, requestBody.model, setModelLoading, controller.signal) clearTimeout(timeoutId) if (!response.ok) { @@ -506,12 +572,12 @@ export function useChat(initialModel = '') { let fullToolCalls = [] // Tool calls with id for agentic loop try { - const response = await fetch(apiUrl(endpoint), { + const response = await fetchWithModelLoadWait(apiUrl(endpoint), { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(loopBody), signal: controller.signal, - }) + }, loopBody.model, setModelLoading, controller.signal) if (!response.ok) { throw new Error(await extractHttpError(response)) @@ -746,6 +812,7 @@ export function useChat(initialModel = '') { // Finalize setIsStreaming(false) setStreamingChatId(null) + setModelLoading(null) abortControllerRef.current = null setStreamingContent('') setStreamingReasoning('') @@ -821,6 +888,7 @@ export function useChat(initialModel = '') { streamingToolCalls: isActiveStreaming ? streamingToolCalls : [], tokensPerSecond, maxTokensPerSecond, + modelLoading: isActiveStreaming ? modelLoading : null, addChat, forkChat, switchChat, diff --git a/core/http/react-ui/src/pages/Chat.jsx b/core/http/react-ui/src/pages/Chat.jsx index 28e2c36d523a..e47ef9eab37b 100644 --- a/core/http/react-ui/src/pages/Chat.jsx +++ b/core/http/react-ui/src/pages/Chat.jsx @@ -294,6 +294,17 @@ function editableMessageText(message) { return typeof textBlock?.text === 'string' ? textBlock.text : null } +// formatLoadEta renders the server's remaining-seconds estimate. The server +// omits it entirely until its observed transfer rate is meaningful, so anything +// arriving here is worth showing. +function formatLoadEta(seconds) { + if (!Number.isFinite(seconds) || seconds <= 0) return '' + if (seconds < 60) return `${Math.round(seconds)}s` + const minutes = Math.round(seconds / 60) + if (minutes < 60) return `${minutes} min` + return `${Math.floor(minutes / 60)}h ${minutes % 60}m` +} + function withEditedMessageText(message, text) { if (typeof message.content === 'string') return { ...message, content: text } const textIndex = message.content.findIndex(block => block?.type === 'text') @@ -315,7 +326,7 @@ export default function Chat() { const { operations } = useOperations() const { chats, activeChat, activeChatId, isStreaming, streamingChatId, streamingContent, - streamingReasoning, streamingToolCalls, tokensPerSecond, maxTokensPerSecond, + streamingReasoning, streamingToolCalls, tokensPerSecond, maxTokensPerSecond, modelLoading, addChat, forkChat, switchChat, deleteChat, deleteAllChats, renameChat, updateChatSettings, sendMessage, stopGeneration, clearHistory, getContextUsagePercent, addMessage, } = useChat(urlModel || '') @@ -326,6 +337,34 @@ export default function Chat() { return operations.find(op => op.taskType === 'staging' && op.name === activeChat.model) || null }, [operations, isStreaming, activeChat?.model]) + // What to show instead of the thinking dots while the model is not up yet. + // The load job wins over the staging operation: it is authoritative across + // frontend replicas and names the phase (installing / staging / loading), + // where the operation only knows about a byte transfer this replica is + // performing. The operation stays as the fallback for a transfer with no job + // attached to this request (a reconciler scale-up, for instance). + const loadProgress = useMemo(() => { + if (modelLoading) { + const eta = formatLoadEta(modelLoading.eta_seconds) + return { + label: t(`streaming.modelState.${modelLoading.state}`, t('streaming.transferring')) + + (modelLoading.node ? ` ${t('streaming.onNode', { node: modelLoading.node })}` : ''), + progress: modelLoading.progress || 0, + detail: eta ? t('streaming.eta', { value: eta }) : '', + } + } + if (stagingOp) { + return { + label: stagingOp.nodeName + ? t('streaming.transferringTo', { node: stagingOp.nodeName }) + : t('streaming.transferring'), + progress: stagingOp.progress || 0, + detail: stagingOp.message || '', + } + } + return null + }, [modelLoading, stagingOp, t]) + const [input, setInput] = useState('') const [files, setFiles] = useState([]) const [showSettings, setShowSettings] = useState(false) @@ -1339,21 +1378,21 @@ export default function Chat() {
- {stagingOp ? ( + {loadProgress ? (
- {stagingOp.nodeName ? t('streaming.transferringTo', { node: stagingOp.nodeName }) : t('streaming.transferring')} + {loadProgress.label}
- {stagingOp.progress > 0 && ( + {loadProgress.progress > 0 && (
-
+
- {Math.round(stagingOp.progress)}% + {Math.round(loadProgress.progress)}%
)} - {stagingOp.message && ( -
{stagingOp.message}
+ {loadProgress.detail && ( +
{loadProgress.detail}
)}
) : ( diff --git a/core/http/react-ui/src/utils/config.js b/core/http/react-ui/src/utils/config.js index 80d437c912af..c72e7e91ba65 100644 --- a/core/http/react-ui/src/utils/config.js +++ b/core/http/react-ui/src/utils/config.js @@ -66,6 +66,9 @@ export const API_CONFIG = { cancelAgentJob: (id) => `/api/agent/jobs/${id}/cancel`, executeAgentJob: '/api/agent/jobs/execute', + // Progress of a cold load still staging a model onto a worker + modelLoadStatus: (id) => `/api/models/${encodeURIComponent(id)}/load-status`, + // OpenAI-compatible endpoints chatCompletions: '/v1/chat/completions', mcpChatCompletions: '/v1/mcp/chat/completions', From bbe8d7c8cd3932ee37260e7c15d882c631037565 Mon Sep 17 00:00:00 2001 From: localai-org-maint-bot <306269227+localai-org-maint-bot@users.noreply.github.com> Date: Fri, 14 Aug 2026 02:13:40 +0000 Subject: [PATCH 7/7] fix(distributed): check warm-path cleanup errors The router moved legacy cleanup calls onto newly linted lines. Report cleanup failures while preserving the fallback to a cold load. Assisted-by: Codex:gpt-5 [golangci-lint] --- core/services/nodes/router.go | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 9322ed740cc5..dbcd23bb8eb7 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -642,8 +642,14 @@ func (r *SmartRouter) tryWarmPath(ctx context.Context, att *routeAttempt) *Route // Verify the backend process is still alive via gRPC health check if !r.probeHealth(ctx, node, modelAddr) { // Stale — roll back the increment, remove the specific replica row, fall through - r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx) - r.registry.RemoveNodeModel(ctx, node.ID, att.trackingKey, replicaIdx) + if err := r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx); err != nil { + xlog.Warn("Failed to release stale routing reservation", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + } + if err := r.registry.RemoveNodeModel(ctx, node.ID, att.trackingKey, replicaIdx); err != nil { + xlog.Warn("Failed to remove stale model from registry", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + } xlog.Warn("Backend not reachable for cached model, falling through to reload", "node", node.Name, "model", att.modelName, "replica", replicaIdx) return nil @@ -651,7 +657,10 @@ func (r *SmartRouter) tryWarmPath(ctx context.Context, att *routeAttempt) *Route // Verify node still matches scheduling constraints if !r.nodeMatchesScheduling(ctx, node, att.sched) { - r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx) + if err := r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx); err != nil { + xlog.Warn("Failed to release unmatched routing reservation", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + } xlog.Info("Cached model on node that no longer matches selector, falling through", "node", node.Name, "model", att.trackingKey, "replica", replicaIdx) return nil