fix(distributed): run cold model loads as durable jobs instead of holding the advisory lock

[localai-bot]

The incident

On a localai-org-development frontend (2 replicas), loading a 35.7 GB GGUF onto a newly added Jetson Thor worker made the model permanently unloadable from the operator's seat, while staging was in fact progressing normally underneath.

Replica A acquired the per-model advisory lock model-load:Qwen3.6-27B-MTP-GGUF and began staging — ~20 minutes of transfer. Replica B received a request for the same model, blocked on pg_advisory_lock, and was killed at 60s by the localai role's statement_timeout:

routing model llama-cpp/models/Qwen3.6-27B-MTP-GGUF/Qwen3.6-27B-UD-Q8_K_XL.gguf:
loading model Qwen3.6-27B-MTP-GGUF: advisorylock: acquiring lock 9003261067483446873:
ERROR: canceling statement due to statement timeout (SQLSTATE 57014)

Every UI retry reproduced it. Three defects sit behind that one symptom:

1. The lock's lifetime was the transfer's lifetime. Route wrapped the whole cold load — backend install, multi-GB staging, checkpoint load — in advisorylock.WithLockCtx. The lock's job is to de-duplicate concurrent loaders, a decision that takes milliseconds; holding it for tens of minutes turns a dedup mechanism into a cluster-wide outage for that model.
2. WithLockCtx defended against lock_timeout but not statement_timeout. Both abort the same blocking pg_advisory_lock($1); only the former was overridden. Latent for every blocking caller, not just model loads.
3. There was no caller contract for "this model is staging." The request either blocked or failed. StagingTracker already tracked per-file byte progress and broadcast it over NATS, but nothing on the inference path consumed any of it.

The change: claim / run split

The cold load becomes a durable job (model_load_jobs), with waiters attached by broadcast.

Request ─► Route
             ├─ warm path: FindAndLockNodeWithModel ──────────────► serve (unchanged)
             │
             └─ cold path:
                  ClaimLoadJob(trackingKey)          ◄── advisory lock held ~ms
                     │
                     ├─ claimed ─► run the job in the background
                     │               install ─► stage ─► LoadModel
                     │
                     └─ already active ────► attach as waiter (no lock held)
                                                │
                  ┌─────────────────────────────┘
                  ▼
             wait for: job ready ─► retry warm path ─► serve
                       job failed ─► return the job's real error
                       wait budget elapsed ─► 503 + progress + Retry-After
                       client cancelled ─► return, the job continues

The advisory lock is kept — it is the right primitive — but its guarded section shrinks to the claim: a SELECT, possibly a DELETE of an orphan, an INSERT. No network, file or gRPC I/O inside it. Uniqueness of the primary key on tracking_key is the real guard; the lock only makes the read-then-write non-racy.

Waiters share one broadcast, not an ordered queue: every waiter for a model wants the identical outcome, so ordering them would add fairness machinery that changes no result. Local waiters wake on a closed channel; a 2s DB poll is the authority, because a waiter on another replica has no channel to close and NATS broadcasts are fire-and-forget.

Liveness is a heartbeat, not a byte counter. The runner touches last_progress on a fixed 1s interval whether or not bytes move — a checkpoint load legitimately transfers zero bytes for many minutes, and a reaper keyed on byte movement would reclaim a healthy job mid-load. Byte progress stays the concern of load_deadline.go, which is untouched. A job whose heartbeat stops past the orphan window is reclaimable, so a replica killed mid-load cannot wedge a model permanently.

Caller contract

• Within the wait budget, the request is served the moment the model is ready — no client retry. This is the common case for a model already most of the way staged.
• A failure propagates the real cause (worker out of disk) to every waiter. The failed job row is retained briefly so a request arriving right after does not read "no job" as "not loading" and start a duplicate load of a model that just failed.
• A client that disconnects returns immediately; the load keeps running. It is owned by the job record, not by the request.
• On budget expiry: 503 with Retry-After and

{
  "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-client compatibility; loading is additive, so existing clients ignore it. eta_seconds comes from the job's own observed rate and is omitted, not 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 the ETA when known, clamped to [5s, 300s], and the wait budget otherwise.

New config key

LOCALAI_MODEL_LOAD_WAIT (default 60s) bounds how long a request waits for a running cold load. It bounds the caller, never the load — the job runs on either way. 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 answer should come from LocalAI with progress attached rather than from a proxy dropping the connection. LOCALAI_MODEL_LOAD_WAIT=0 restores unbounded waiting 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 explicit zero as ModelLoadWaitUnbounded rather than losing the distinction.)

New endpoint

GET /api/models/{id}/load-status → the same loading object, or 404 when no load is running. Deliberately neither admin- 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. No MCP tool — there is nothing here an admin would manage conversationally.

UI

Chat renders a 503 with type: "model_loading" as inline progress (phase, node, percent, ETA), polls load-status, and re-sends the request once the model is ready. It reuses the staging progress idiom the page already had rather than inventing a second one; the load job wins over the /api/operations staging row because it is authoritative across replicas and names the phase, where the operation only knows about a byte transfer this replica happens to be performing.

Commits

1. fix(advisorylock) — set statement_timeout alongside lock_timeout. Independently backportable; it fixes a latent bug for every blocking WithLockCtx caller (schema migration, reconciler, health check), not just model loads.
2. feat(distributed) — the ModelLoadJob row + ClaimLoadJob.
3. refactor(distributed) — move staging/load out of the lock into the job runner, with waiter fan-out.
4. feat(distributed) — the wait budget and the 503 contract.
5. feat(api) — load-status + registration surfaces + docs.
6. feat(ui) — Chat staging state and auto-retry.

Each builds and passes on its own.

Testing

New specs (Ginkgo + testcontainers PostgreSQL, and Playwright for the UI):

• a waiter survives a short server-side statement_timeout instead of failing 57014 (mirrors the existing lock_timeout spec);
• exactly one of eight concurrent claimers wins;
• a claim returns in <100ms while another replica's job is running;
• an orphaned job (heartbeat stopped) is reclaimable;
• a concurrent request for a loading model is served from one load, with no duplicate install;
• a load failure reaches every waiter with its real cause;
• cancelling a waiter returns immediately and leaves the job running;
• the runner heartbeats while no bytes move;
• budget expiry yields a structured ModelLoadingError with Retry-After, and LOCALAI_MODEL_LOAD_WAIT=0 waits unbounded;
• load-status 404s with no job and reports progress with one;
• Chat renders the staging state on a 503 and auto-retries on ready (verified failing before the change).

Notes

• pkg/model/loader.go wrapped the router's error with %s, flattening the chain; changed to %w so the typed loading error survives to the HTTP layer.
• Non-distributed (no-DB) mode keeps the inline load byte-for-byte as it was.
• load_deadline.go (the progress-extended cold-load hold) and file_stager_http.go are untouched.

🤖 Generated with Claude Code