fix(distributed): configurable remote model-load timeout, and reap the load when it times out

[localai-bot]

Two changes to the distributed cold-load path. They ship together because the first makes large cold loads possible and the second stops a timeout from leaking a multi-GB process every time one still happens.

The incident

An ARM64 Thor worker loading meituan-longcat/LongCat-Video-Avatar-1.5 (~83 GB):

• the cold load failed at exactly 302s with DeadlineExceeded, against a hardcoded 5m deadline
• the abandoned backend kept running — PID 31954 at elapsed 05:17 / cpu 01:06 / rss 3.87 GB, still alive at elapsed 30:05 / cpu 05:32, having pulled ~57 GB from HuggingFace with nobody waiting on the result
• it had to be reaped by hand via POST /api/nodes/:id/models/unload; we watched it happen twice in one session, each retry stacking another loader on the worker

Once the weights were cached locally the same model loaded in 108 seconds, comfortably inside the 5m budget. The default is not unreasonable for a warm load — the pathology is weight acquisition happening inside the load RPC, which is addressed separately.

1. LOCALAI_NATS_MODEL_LOAD_TIMEOUT / --model-load-timeout

Replaces the hardcoded context.WithTimeout(ctx, 5*time.Minute) in scheduleAndLoad. Default stays 5m, so unset clusters are unchanged. Mirrors the existing LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT precedent, whose help text already says "Increase for slow links pulling multi-GB images" — the install path got this treatment; the load path never did.

The advisory-lock ceiling is now derived rather than constant: ModelLoadCeilingFor(install, load) = max(install + load + 5m, 25m). That matters — the old defaultModelLoadCeiling = 25 * time.Minute was explicitly documented as "install (15m) plus staging and the remote LoadModel (5m)", so raising the load budget past it would have been silently clipped and the fix would have appeared not to work. 15m + 5m + 5m = 25m at defaults, identical to before, with a 25m floor so shrinking either budget can never tighten it.

2. Reap the abandoned replica on timeout

A gRPC deadline only cancels the client side of the call. A backend blocked in a synchronous weight load never observes its cancelled handler context, and scheduleAndLoad previously just returned the error: no StopBackend, no unload. The worker kept loading forever.

On a deadline or cancellation the router now sends backend.stop for the exact modelID#replicaIndex process key it just abandoned. The exact key matters: a bare model ID stops every replica on that node, including healthy ones serving traffic.

Any other LoadModel failure does not reap. An unsupported model or bad option means the backend answered, its handler returned and the process is idle; stopping it would discard a warm process and its downloaded weights before the next attempt. Classification uses errors.Is for context errors plus errors.As on the GRPCStatus() interface (not status.Code, which does not unwrap in all grpc-go versions).

The reap is best-effort: a failed stop is logged via xlog.Warn and never replaces the load error the caller is waiting on, with a spec asserting the stop error cannot appear in the returned error.

The modelID#replicaIndex format was already hand-rolled in two places, so rather than add a third this exports model.BackendProcessKey from pkg/model — the lowest common dependency of the router and the worker, with no new import edges.

Merge-order dependency

This PR should merge after the worker-side fix that bounds client.Free(context.Background()) in stopBackendExact.

proc.Stop() itself is sound (go-processmanager defaults: SIGTERM to the process group, 15s grace, then SIGKILL) and will kill a Python loader wedged in torch.load. But stopBackendExact calls client.Free() with an unbounded context first. 37 Python backends default to PYTHON_GRPC_MAX_WORKERS=1 — including backend/python/longcat-video/backend.py:46, the exact backend from this incident — so a backend wedged in LoadModel has no thread to service Free, and the stop blocks before ever reaching the kill.

So the reap here dispatches the right stop for the right replica, but against a fully wedged single-worker Python backend it is inert until that worker-side timeout lands. This matches the field evidence: the manual unload worked, but only once the loader had gone idle enough to answer Free.

Testing

New Ginkgo specs in core/services/nodes/router_reap_load_test.go drive Route → scheduleAndLoad with a stub grpc.Backend and assert: the stop is dispatched for the exact modelID#replicaIndex (pinned to replica 2, so a hardcoded #0 cannot pass), the original DeadlineExceeded still propagates when the stop itself fails, a wrapped context.DeadlineExceeded reaps too, and a clean InvalidArgument does not.

ok  core/services/nodes             193.025s
ok  core/services/nodes/prefixcache   0.009s
ok  core/config                       0.425s
ok  core/cli                          0.042s
ok  pkg/model                         1.343s

make lint → 0 issues. Coverage rose to 53.0% (baseline 48.5%).

---

🤖 Generated with Claude Code

[localai-bot]

Correction to the "Merge-order dependency" section of the description above: it was wrong, and this PR had no unmet dependency.

That section claimed stopBackendExact still called client.Free(context.Background()) unbounded, and that the reap added here would therefore be inert against a backend wedged in LoadModel with PYTHON_GRPC_MAX_WORKERS=1.

That was already fixed on master before this merged, by #10865, which added:

const workerBackendFreeTimeout = 5 * time.Second

and applied it at both relevant call sites — stopBackendExact (core/services/worker/supervisor.go:243) and handleModelUnload (core/services/worker/lifecycle.go:261). Master's implementation is also better than the one this PR's description anticipated: it releases the supervisor mutex across the Free and keeps the port reserved until termination.

So the reap in this PR is effective as merged. No follow-up is required to make it work.

The error came from checking a worktree branched off an older master rather than master itself. Leaving this as a comment rather than editing the description, so the original claim and its correction both stay visible.