From 81b1924b545ac9d81fc6bee03f09f726ac212f81 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Tue, 21 Jul 2026 20:51:50 +0000 Subject: [PATCH] fix(distributed): scale the remote model-load deadline with checkpoint size The gRPC deadline for the remote LoadModel call was a fixed 5m. It starts only after the backend install and file staging have completed, so it covers the worker's checkpoint read and pipeline init alone - work whose duration is proportional to the bytes on disk. A fixed value is therefore a model-size cliff, not a timeout. Measured in production: a 70 GB video checkpoint (longcat-video-avatar-1.5) on an NVIDIA Jetson Thor worker failed reproducibly with "rpc error: code = DeadlineExceeded" after 953.5s of wall clock. Backend install plus staging consumed ~11m, then LoadModel got its 5m and expired. The load never had a chance, and the operator saw only a generic DeadlineExceeded with no hint that a config value was the cause. Raising the constant does not fix this. It moves the cliff to the next larger model - the cluster has to support 600 GB checkpoints - and it makes a genuinely wedged SMALL model hang for the whole inflated duration before anyone notices, which is a real regression in failure latency. So derive the budget from the checkpoint size instead: budget = 5m + 20s/GiB, capped at 6h 2 GiB -> 5m40s, 70 GiB -> 28m20s, 600 GiB -> 3h25m. The per-GiB rate is deliberately pessimistic (~54 MB/s of weight read) because the errors are not symmetric: too long costs only failure latency on a load that was going to fail anyway, too short is a guaranteed false failure on a healthy load. The size is measured from the frontend's local model files, over the same path set stageModelFiles uploads. When those files are not present locally - a backend handed a bare HuggingFace repo id fetches its own weights on the worker - there is nothing to measure and the budget stays at today's 5m. An explicit LOCALAI_NATS_MODEL_LOAD_TIMEOUT still wins outright, in both directions: a shorter override is honoured, so an operator who wants fast failure is not silently extended by the heuristic. The cold-load hold needed widening to match. It extends on staging progress, but LoadModel reports none, so once the last byte lands the hold expires a stall window later and would cancel a load still well inside its own budget. scheduleAndLoad now extends the hold by the load budget plus the staging margin as it enters the load phase; ModelLoadCeilingFor stays the hold's starting budget rather than its maximum. Finally, a deadline that does expire now names the budget, the checkpoint size it was derived from, and the knob that overrides it, instead of surfacing a bare "context deadline exceeded". Signed-off-by: Ettore Di Giacinto Assisted-by: Claude Code:claude-opus-4-8[1m] [Read] [Edit] [Bash] --- core/application/distributed.go | 7 +- core/cli/run.go | 2 +- core/config/model_load_budget.go | 78 ++++++++ core/config/model_load_budget_test.go | 53 +++++ core/services/nodes/load_deadline.go | 38 ++++ core/services/nodes/router.go | 138 +++++++++++-- .../services/nodes/router_load_budget_test.go | 188 ++++++++++++++++++ docs/content/features/distributed-mode.md | 28 ++- 8 files changed, 514 insertions(+), 18 deletions(-) create mode 100644 core/config/model_load_budget.go create mode 100644 core/config/model_load_budget_test.go create mode 100644 core/services/nodes/router_load_budget_test.go diff --git a/core/application/distributed.go b/core/application/distributed.go index 0519d96e787e..3cef85bc9e2d 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -356,7 +356,12 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade PrefixConfig: prefixCfg, Pressure: pressure, SharedModels: cfg.Distributed.SharedModels, - ModelLoadTimeout: cfg.Distributed.ModelLoadTimeoutOrDefault(), + // RAW, not OrDefault: zero means "derive the budget per model from the + // checkpoint size" (config.ModelLoadTimeoutForSize), which is what makes + // a 70 GB video checkpoint work without the operator first hitting a + // DeadlineExceeded and going looking for a knob. A non-zero value here is + // an explicit override and is used verbatim. + ModelLoadTimeout: cfg.Distributed.ModelLoadTimeout, // Cap how long a cold load may hold the per-model advisory lock. Derived // from BOTH configured budgets it has to cover, so raising either the // install timeout (slow links pulling multi-GB images) or the model load diff --git a/core/cli/run.go b/core/cli/run.go index 659977cc1dd0..57689e792d6b 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -173,7 +173,7 @@ type RunCMD struct { DistributedPrefixCacheTTL string `env:"LOCALAI_DISTRIBUTED_PREFIX_CACHE_TTL" help:"Idle-timeout for prefix-cache index entries; also drives the background eviction cadence (every TTL/2). Default 5m." group:"distributed"` 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:"gRPC deadline for the remote LoadModel call sent to a worker node once its backend is installed and model files are staged (default 5m). Increase for very large checkpoints (multi-tens-of-GB diffusion/video models) whose load and pipeline init exceed 5 minutes. Raising it also widens the cold-load lock ceiling." 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"` 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"` diff --git a/core/config/model_load_budget.go b/core/config/model_load_budget.go new file mode 100644 index 000000000000..0b9a55017362 --- /dev/null +++ b/core/config/model_load_budget.go @@ -0,0 +1,78 @@ +package config + +import "time" + +// The remote LoadModel deadline used to be the fixed DefaultModelLoadTimeout. +// That is a model-size cliff, not a timeout: the deadline starts only after the +// backend install and file staging have finished, so it covers the worker's +// checkpoint read and pipeline init alone — work whose duration is proportional +// to the bytes on disk. A 70 GB video checkpoint on a Jetson Thor worker +// therefore failed reproducibly (953.5s wall clock, ~11m of it install and +// staging, then DeadlineExceeded on a load that never had a chance). +// +// Raising the constant does not fix that: it moves the cliff to the next larger +// model, and it makes a genuinely wedged SMALL model hang for the whole inflated +// duration before anyone notices. So the budget is derived from the size instead. +const ( + // ModelLoadTimeoutPerGiB is the budget granted per GiB of checkpoint on top + // of DefaultModelLoadTimeout. + // + // It is deliberately generous. This is a timeout: erring long costs only + // failure LATENCY on a load that was going to fail anyway, while erring + // short costs a guaranteed FALSE failure on a load that was healthy. 20s/GiB + // corresponds to reading weights at ~54 MB/s, which is below what any + // supported medium sustains (NVMe is orders of magnitude faster; the slow end + // is an eMMC or SD-backed Jetson, or a checkpoint faulted in over a network + // filesystem) and so leaves headroom for the dequantisation and pipeline init + // that follow the read. Hardware and quantisation both move the real figure + // by an order of magnitude, which is exactly why the constant sits at the + // pessimistic end rather than at a measured average. + // + // Worked examples: 2 GiB -> 5m40s (a wedged small model still fails fast), + // 70 GiB -> 28m20s (the production checkpoint that failed at 5m), + // 600 GiB -> 3h25m (the size the cluster has to support). + ModelLoadTimeoutPerGiB = 20 * time.Second + + // MaxModelLoadTimeout caps the derived budget so a nonsense size (a corrupted + // stat, a future 10 TB artifact) cannot hand out an effectively infinite + // deadline. At ModelLoadTimeoutPerGiB the cap binds only above ~1 TiB, + // comfortably past the 600 GB requirement. + MaxModelLoadTimeout = 6 * time.Hour +) + +// bytesPerGiB is the divisor for the per-GiB rate above. +const bytesPerGiB int64 = 1 << 30 + +// ModelLoadTimeoutForSize derives the gRPC deadline for the remote LoadModel +// call from the checkpoint's on-disk size. +// +// A non-positive size means the frontend could not measure the payload — a +// backend given a bare HuggingFace repo id fetches its own weights on the +// worker, so there is nothing local to stat. Rather than guess, that case keeps +// the historical DefaultModelLoadTimeout, which is also the floor of the derived +// range: size only ever adds budget. +// +// An explicit LOCALAI_NATS_MODEL_LOAD_TIMEOUT always wins over this; see +// SmartRouterOptions.ModelLoadTimeout. +func ModelLoadTimeoutForSize(bytes int64) time.Duration { + if bytes <= 0 { + return DefaultModelLoadTimeout + } + + // Split into whole GiB plus remainder before scaling: multiplying a 600 GB + // byte count by a 20e9-nanosecond rate overflows int64 long before the cap + // could clamp it. + gib := bytes / bytesPerGiB + remainder := bytes % bytesPerGiB + + extra := time.Duration(gib) * ModelLoadTimeoutPerGiB + extra += time.Duration(int64(ModelLoadTimeoutPerGiB) * remainder / bytesPerGiB) + + // A large enough gib still overflows the multiply above into a negative + // duration; treat any non-positive extra as "past the cap". + if extra <= 0 { + return MaxModelLoadTimeout + } + + return min(DefaultModelLoadTimeout+extra, MaxModelLoadTimeout) +} diff --git a/core/config/model_load_budget_test.go b/core/config/model_load_budget_test.go new file mode 100644 index 000000000000..f1a762a9cf9c --- /dev/null +++ b/core/config/model_load_budget_test.go @@ -0,0 +1,53 @@ +package config_test + +import ( + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" +) + +const gib int64 = 1 << 30 + +var _ = Describe("ModelLoadTimeoutForSize", func() { + // The remote LoadModel deadline used to be a fixed 5m. That is a model-size + // cliff: a 70 GB video checkpoint on a Jetson Thor worker failed + // reproducibly with DeadlineExceeded because its weight load and pipeline + // init alone exceed 5 minutes. Raising the constant only moves the cliff, so + // the budget is derived from the bytes the worker has to read. + + It("keeps a small checkpoint close to the historical 5m default", func() { + // A wedged 2 GB model must still fail fast: inflating every load's + // budget is a real regression in failure latency. + Expect(config.ModelLoadTimeoutForSize(2 * gib)).To(BeNumerically("<", 10*time.Minute)) + }) + + It("gives a 70 GB checkpoint materially more budget than a 2 GB one", func() { + small := config.ModelLoadTimeoutForSize(2 * gib) + big := config.ModelLoadTimeoutForSize(70 * gib) + Expect(big).To(BeNumerically(">", small*3)) + // The measured production failure had ~5m of load budget and needed + // more; anything under 20m would still be a cliff for this exact model. + Expect(big).To(BeNumerically(">=", 20*time.Minute)) + }) + + It("scales monotonically with size", func() { + Expect(config.ModelLoadTimeoutForSize(600 * gib)). + To(BeNumerically(">", config.ModelLoadTimeoutForSize(70*gib))) + }) + + It("still gives a 600 GB checkpoint hours, not minutes", func() { + Expect(config.ModelLoadTimeoutForSize(600 * gib)).To(BeNumerically(">=", 3*time.Hour)) + }) + + It("falls back to the plain default when the size is unknown", func() { + Expect(config.ModelLoadTimeoutForSize(0)).To(Equal(config.DefaultModelLoadTimeout)) + Expect(config.ModelLoadTimeoutForSize(-1)).To(Equal(config.DefaultModelLoadTimeout)) + }) + + It("never exceeds the absolute maximum, however absurd the size", func() { + Expect(config.ModelLoadTimeoutForSize(100_000 * gib)).To(Equal(config.MaxModelLoadTimeout)) + }) +}) diff --git a/core/services/nodes/load_deadline.go b/core/services/nodes/load_deadline.go index f5cd6c37ca99..958fd465b9c5 100644 --- a/core/services/nodes/load_deadline.go +++ b/core/services/nodes/load_deadline.go @@ -246,6 +246,44 @@ func (d *loadDeadline) Observe() { d.timer.Reset(time.Until(next)) } +// Extend pushes the expiry out to at least now+d, for a phase whose duration is +// known up front but which reports no progress of its own. +// +// Observe() cannot serve that case: it grants exactly one stall window, so a +// step that legitimately runs for half an hour in silence — the remote +// LoadModel reading a 70 GB checkpoint — is cancelled long before it finishes. +// Extend states the budget once instead of inferring it from a heartbeat that +// will never arrive. Like Observe it only ever moves the expiry forward and +// never past the absolute cap, so it cannot be used to escape the hard bound. +func (d *loadDeadline) Extend(dur time.Duration) { + if d == nil || dur <= 0 { + return + } + d.mu.Lock() + defer d.mu.Unlock() + if d.stopped { + return + } + next := time.Now().Add(dur) + if next.After(d.hardExpiry) { + next = d.hardExpiry + } + if !next.After(d.expiry) { + return + } + d.expiry = next + d.timer.Reset(time.Until(next)) +} + +// extendLoadDeadline widens the cold-load hold attached to ctx to cover a phase +// of known duration. A context without a load deadline (single-host paths, tests +// constructing routers directly) is a no-op. +func extendLoadDeadline(ctx context.Context, d time.Duration) { + if ld, ok := ctx.Value(loadDeadlineKey{}).(*loadDeadline); ok { + ld.Extend(d) + } +} + // observeLoadProgress pushes out the cold-load deadline attached to ctx, if any. // Callers report byte-level movement here; a context without a load deadline // (single-host paths, tests constructing stagers directly) is a no-op. diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 7fa7142b70a2..a463ef07aae8 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -79,13 +79,17 @@ type SmartRouterOptions struct { // it from the default install budget and ModelLoadTimeout via // ModelLoadCeilingFor. ModelLoadCeiling time.Duration - // ModelLoadTimeout is the gRPC deadline for the remote LoadModel call, which - // runs after the backend install and file staging have already completed and - // so covers only the worker backend's checkpoint load and pipeline init. - // Zero selects config.DefaultModelLoadTimeout. Operators raise it via - // LOCALAI_NATS_MODEL_LOAD_TIMEOUT for very large checkpoints; ModelLoadCeiling - // must be widened in step (see ModelLoadCeilingFor) or the ceiling clips the - // longer deadline before it can be used. + // ModelLoadTimeout pins the gRPC deadline for the remote LoadModel call, + // which runs after the backend install and file staging have already + // completed and so covers only the worker backend's checkpoint load and + // pipeline init. + // + // Zero means "derive it per model from the checkpoint size" via + // config.ModelLoadTimeoutForSize — the default, because load duration is + // proportional to the bytes the worker reads and any fixed value is a + // model-size cliff. A non-zero value is an explicit operator override + // (LOCALAI_NATS_MODEL_LOAD_TIMEOUT) and always wins over the derived budget, + // in BOTH directions: an operator who wants faster failure gets it. ModelLoadTimeout time.Duration // StagingStallWindow is how long file staging may report zero bytes before // the cold load is declared wedged and the advisory lock released. While @@ -116,6 +120,12 @@ const minModelLoadCeiling = 25 * time.Minute // raising LOCALAI_NATS_MODEL_LOAD_TIMEOUT for an 80 GB video checkpoint actually // takes effect, instead of being cut short by a constant that silently went // stale. Non-positive inputs fall back to their package defaults. +// +// This is the hold's STARTING budget, not its maximum. A per-model budget +// derived from checkpoint size (config.ModelLoadTimeoutForSize) can exceed any +// ceiling computed here, so scheduleAndLoad widens the hold as it enters the +// load phase — see extendLoadDeadline. Without that, a 70 GB checkpoint's ~28m +// load budget would be cancelled by a 25m ceiling that knew nothing about it. func ModelLoadCeilingFor(installTimeout, loadTimeout time.Duration) time.Duration { if installTimeout <= 0 { installTimeout = config.DefaultBackendInstallTimeout @@ -160,8 +170,9 @@ type SmartRouter struct { // modelLoadCeiling bounds how long a cold load may hold the per-model // advisory lock (see SmartRouterOptions.ModelLoadCeiling). modelLoadCeiling time.Duration - // modelLoadTimeout is the deadline for the remote LoadModel gRPC call - // (see SmartRouterOptions.ModelLoadTimeout). + // modelLoadTimeout is the operator's explicit override for the remote + // LoadModel deadline, or zero to derive it per model from the checkpoint + // size (see SmartRouterOptions.ModelLoadTimeout and loadTimeoutFor). modelLoadTimeout time.Duration // stagingStallWindow and modelLoadAbsoluteMax turn modelLoadCeiling from a // hard countdown into a progress-extended hold (see load_deadline.go). @@ -183,10 +194,10 @@ func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter if factory == nil { factory = &tokenClientFactory{token: opts.AuthToken} } - loadTimeout := opts.ModelLoadTimeout - if loadTimeout <= 0 { - loadTimeout = config.DefaultModelLoadTimeout - } + // Keep the override RAW: zero has to stay distinguishable from an explicit + // 5m, because zero now means "derive per model from the checkpoint size" + // while an explicit 5m means "hold every load to 5m". + loadTimeout := max(opts.ModelLoadTimeout, 0) ceiling := opts.ModelLoadCeiling if ceiling <= 0 { ceiling = ModelLoadCeilingFor(config.DefaultBackendInstallTimeout, loadTimeout) @@ -293,6 +304,12 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking // that built modelOpts may have no GPU at all in distributed mode. applyNodeHardwareDefaults(modelOpts, node, backendType) + // Size the remote load budget BEFORE staging: stageModelFiles rewrites the + // path fields to their remote equivalents on a clone, and only the local + // paths can be stat'ed here. + payloadBytes := modelPayloadBytes(modelOpts) + loadTimeout := r.loadTimeoutFor(payloadBytes) + // Pre-stage model files via FileStager before loading loadOpts := modelOpts if r.fileStager != nil && modelOpts != nil { @@ -307,13 +324,30 @@ func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, tracking // Load the model on the remote node if loadOpts != nil { - xlog.Info("Loading model on remote node", "node", node.Name, "model", modelName, "addr", backendAddr) + xlog.Info("Loading model on remote node", "node", node.Name, "model", modelName, "addr", backendAddr, + "payloadBytes", payloadBytes, "loadBudget", loadTimeout) + + // The cold-load hold above this call extends on STAGING progress, and + // the remote LoadModel reports none — so once the last byte lands the + // hold expires a stall window later and would cancel a load that is + // still well inside its own budget (#11026 in miniature). Widen the hold + // to cover the load phase before entering it. + extendLoadDeadline(ctx, loadTimeout+modelLoadStagingMargin) - loadCtx, cancel := context.WithTimeout(ctx, r.modelLoadTimeout) + loadCtx, cancel := context.WithTimeout(ctx, loadTimeout) defer cancel() res, err := client.LoadModel(loadCtx, loadOpts) if err != nil { + // A bare "context deadline exceeded" tells the operator nothing + // about which of the several budgets in a cold load ran out, and + // cost real debugging time in production. Name the budget, the + // payload it was derived from, and the knob that overrides it. + if errors.Is(err, context.DeadlineExceeded) { + err = fmt.Errorf("model load budget of %s for a %s checkpoint exceeded; "+ + "raise it with LOCALAI_NATS_MODEL_LOAD_TIMEOUT (%s): %w", + loadTimeout, humanFileSize(payloadBytes), config.FlagModelLoadTimeout, err) + } // 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, so it keeps loading with nobody @@ -1368,6 +1402,80 @@ func (r *SmartRouter) withStagingCallback(ctx context.Context, trackingKey, file }) } +// loadTimeoutFor resolves the gRPC deadline for one remote LoadModel call. +// +// An explicit LOCALAI_NATS_MODEL_LOAD_TIMEOUT wins outright — including when it +// is SHORTER than the derived value, because an operator who deliberately wants +// fast failure must not have a size heuristic silently extend their loads. +// Otherwise the budget scales with the checkpoint, which is what the worker +// actually spends its time reading. +func (r *SmartRouter) loadTimeoutFor(payloadBytes int64) time.Duration { + if r.modelLoadTimeout > 0 { + return r.modelLoadTimeout + } + return config.ModelLoadTimeoutForSize(payloadBytes) +} + +// modelPayloadBytes totals the on-disk size of everything the worker will have +// to read for this model, over the same field set stageModelFiles uploads. +// +// Paths that do not exist locally contribute nothing: a backend handed a bare +// HuggingFace repo id gets an optimistically constructed path that was never +// materialized and fetches its own weights on the worker. There is no way to +// size that from here, so it falls back to the plain default budget rather than +// to a guess (see config.ModelLoadTimeoutForSize). +func modelPayloadBytes(opts *pb.ModelOptions) int64 { + if opts == nil { + return 0 + } + paths := []string{ + opts.ModelFile, opts.MMProj, opts.LoraAdapter, opts.DraftModel, + opts.CLIPModel, opts.Tokenizer, opts.AudioPath, opts.LoraBase, + } + paths = append(paths, opts.LoraAdapters...) + + // The same file can legitimately appear in two fields (a GGUF that is both + // ModelFile and Tokenizer, say); counting it twice would inflate the budget. + seen := make(map[string]struct{}, len(paths)) + var total int64 + for _, p := range paths { + if p == "" { + continue + } + if _, dup := seen[p]; dup { + continue + } + seen[p] = struct{}{} + total += pathBytes(p) + } + return total +} + +// pathBytes returns the size of a regular file, the total size of a directory's +// contents, or 0 if the path cannot be stat'ed. +func pathBytes(path string) int64 { + fi, err := os.Stat(path) + if err != nil { + return 0 + } + if !fi.IsDir() { + return fi.Size() + } + var total int64 + _ = filepath.WalkDir(path, func(_ string, d fs.DirEntry, walkErr error) error { + if walkErr != nil || d.IsDir() { + return nil + } + info, infoErr := d.Info() + if infoErr != nil { + return nil + } + total += info.Size() + return nil + }) + return total +} + // countStageableFiles returns the number of regular files a model path expands // to for staging: 1 for a regular file, the contained file count for a // directory, and 0 if the path does not exist. diff --git a/core/services/nodes/router_load_budget_test.go b/core/services/nodes/router_load_budget_test.go new file mode 100644 index 000000000000..44779694a3b8 --- /dev/null +++ b/core/services/nodes/router_load_budget_test.go @@ -0,0 +1,188 @@ +package nodes + +import ( + "context" + "errors" + "os" + "path/filepath" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" + grpc "github.com/mudler/LocalAI/pkg/grpc" + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + ggrpc "google.golang.org/grpc" +) + +// sparseCheckpoint writes a file whose reported size is `size` but which +// occupies (almost) no blocks, so a spec can exercise the 70 GB and 600 GB +// paths on a host with a nearly full disk. +func sparseCheckpoint(dir, name string, size int64) string { + path := filepath.Join(dir, name) + f, err := os.Create(path) + Expect(err).ToNot(HaveOccurred()) + Expect(f.Truncate(size)).To(Succeed()) + Expect(f.Close()).To(Succeed()) + return path +} + +// holdBackend blocks inside LoadModel for `hold` and records whether the +// context it was handed survived. A budget that exists only as a +// context.WithTimeout value is not enough: the cold-load hold above it must +// also stay alive, or the load is killed by the ceiling regardless. +type holdBackend struct { + grpc.Backend // embedded: unused methods panic if called + + hold time.Duration + + mu sync.Mutex + budget time.Duration + hadOne bool + errAtEnd error +} + +func (b *holdBackend) HealthCheck(_ context.Context) (bool, error) { return true, nil } + +func (b *holdBackend) IsBusy() bool { return false } + +func (b *holdBackend) LoadModel(ctx context.Context, _ *pb.ModelOptions, _ ...ggrpc.CallOption) (*pb.Result, error) { + dl, ok := ctx.Deadline() + b.mu.Lock() + b.hadOne = ok + if ok { + b.budget = time.Until(dl).Round(time.Second) + } + b.mu.Unlock() + + if b.hold > 0 { + select { + case <-ctx.Done(): + case <-time.After(b.hold): + } + b.mu.Lock() + b.errAtEnd = ctx.Err() + b.mu.Unlock() + } + return &pb.Result{Success: true}, nil +} + +func (b *holdBackend) loadBudget() time.Duration { + b.mu.Lock() + defer b.mu.Unlock() + Expect(b.hadOne).To(BeTrue(), "LoadModel must be called with a deadline") + return b.budget +} + +func (b *holdBackend) ctxErrAtEnd() error { + b.mu.Lock() + defer b.mu.Unlock() + return b.errAtEnd +} + +type holdClientFactory struct{ client *holdBackend } + +func (f *holdClientFactory) NewClient(_ string, _ bool) grpc.Backend { return f.client } + +var _ = Describe("size-derived remote LoadModel budget", func() { + // Production, on an NVIDIA Jetson Thor worker: a 70 GB video checkpoint + // (longcat-video-avatar-1.5) failed reproducibly after 953.5s with + // "rpc error: code = DeadlineExceeded". Backend install plus staging ate + // ~11m of wall clock, then the fixed 5m LoadModel deadline expired before + // the worker had finished reading the weights. The failure is deterministic + // for any checkpoint whose init exceeds 5 minutes, and the cluster has to + // support 600 GB checkpoints, so the budget has to scale with the bytes. + var ( + reg *fakeModelRouter + backend *holdBackend + factory *holdClientFactory + unloader *fakeUnloader + dir string + ) + + BeforeEach(func() { + reg = &fakeModelRouter{findAndLockErr: errors.New("not found")} + reg.findIdleNode = &BackendNode{ID: "n1", Name: "worker", Address: "10.0.0.1:50051"} + backend = &holdBackend{} + factory = &holdClientFactory{client: backend} + unloader = &fakeUnloader{ + installReply: &messaging.BackendInstallReply{Success: true, Address: "10.0.0.1:9001"}, + } + dir = GinkgoT().TempDir() + }) + + routeFile := func(router *SmartRouter, modelFile string) { + _, err := router.Route(context.Background(), "big-model", "models/big.gguf", "llama-cpp", + &pb.ModelOptions{Model: "models/big.gguf", ModelFile: modelFile}, false) + Expect(err).ToNot(HaveOccurred()) + } + + It("gives a 70 GB checkpoint materially more than the 5m default", func() { + big := sparseCheckpoint(dir, "big.gguf", 70<<30) + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + }) + routeFile(router, big) + Expect(backend.loadBudget()).To(BeNumerically(">=", 20*time.Minute)) + }) + + It("keeps a small checkpoint on a short budget so a wedged load fails fast", func() { + small := sparseCheckpoint(dir, "small.gguf", 2<<30) + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + }) + routeFile(router, small) + Expect(backend.loadBudget()).To(BeNumerically("<", 10*time.Minute)) + }) + + It("lets an explicit LOCALAI_NATS_MODEL_LOAD_TIMEOUT beat the derived budget", func() { + big := sparseCheckpoint(dir, "big.gguf", 70<<30) + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + // Deliberately SHORTER than what 70 GB would derive: an operator + // asking for fast failure must get it, so the override cannot be a + // floor under the derived value. + ModelLoadTimeout: 7 * time.Minute, + ModelLoadCeiling: 2 * time.Hour, + }) + routeFile(router, big) + Expect(backend.loadBudget()).To(Equal(7 * time.Minute)) + }) + + It("does not let the cold-load ceiling clip the derived budget", func() { + // The hold that bounds the whole cold-load sequence extends on staging + // progress, but the remote LoadModel reports none — so once staging + // stops the hold expires a stall window later and cancels the load out + // from under a budget that was supposed to be much longer. Entering the + // load phase has to widen the hold in step with the load budget. + big := sparseCheckpoint(dir, "big.gguf", 70<<30) + backend.hold = 2 * time.Second + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + // A ceiling far shorter than the derived load budget. Without the + // widening, the hold fires while LoadModel is still working. + ModelLoadCeiling: time.Second, + }) + routeFile(router, big) + Expect(backend.ctxErrAtEnd()).ToNot(HaveOccurred(), + "the cold-load hold cancelled a LoadModel that was still inside its own budget") + }) + + It("falls back to the plain default when the checkpoint is not on the frontend disk", func() { + // Backends that take a bare HuggingFace repo id get an optimistically + // constructed path that was never materialized; there are no bytes to + // measure, so the budget stays at today's default rather than guessing. + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + }) + routeFile(router, filepath.Join(dir, "never-materialized.gguf")) + Expect(backend.loadBudget()).To(Equal(5 * time.Minute)) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 2330793878bb..5db0b9f21625 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -72,13 +72,39 @@ The frontend is a standard LocalAI instance with distributed mode enabled. These | `--auth-database-url` | `LOCALAI_AUTH_DATABASE_URL` | *(required)* | PostgreSQL connection URL | | `--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` | `5m` | Deadline for the `LoadModel` gRPC call the frontend issues to a worker. It starts *after* the backend is installed and the model files are staged, so it covers only the worker backend's own checkpoint load and pipeline init. Raise it for very large checkpoints: a multi-tens-of-GB diffusion or video model on unified memory routinely needs more than 5 minutes, and the load then fails with `rpc error: code = DeadlineExceeded` at exactly the timeout. Raising this also widens the cold-load lock ceiling below, so the two knobs never fight. | +| `--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. | | `--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 + +The `LoadModel` deadline starts *after* the backend is installed and the model files are staged, so it covers only the worker backend's own checkpoint read and pipeline init. That work is proportional to the bytes on disk, which makes any fixed deadline a model-size cliff rather than a timeout: a 70 GB video checkpoint on a Jetson Thor worker failed reproducibly against the old fixed 5m default (`rpc error: code = DeadlineExceeded` after 953.5s of wall clock, roughly 11m of which was backend install and staging), and simply raising the constant would only move the cliff to the next larger model while making a genuinely wedged *small* model hang for the whole inflated duration. + +So the deadline is derived per model: + +``` +budget = 5m + 20s per GiB of checkpoint, capped at 6h +``` + +| Checkpoint | Derived budget | +|---|---| +| 2 GB | 5m40s | +| 70 GB | 28m20s | +| 600 GB | 3h25m | + +The per-GiB rate is deliberately pessimistic — it corresponds to reading weights at about 54 MB/s, below what any supported storage sustains — because the two errors are not symmetric: a budget that is too long costs only *failure latency* on a load that was going to fail anyway, while a budget that is too short causes a guaranteed false failure on a load that was perfectly healthy. + +The size is measured from the model files on the frontend's disk, over the same set of paths that get staged to the worker. If those files are not present locally — a backend handed a bare HuggingFace repo id fetches its own weights on the worker — there is nothing to measure and the budget stays at the plain 5m default. Pin `LOCALAI_NATS_MODEL_LOAD_TIMEOUT` for those models if their load is slow. + +When the budget *is* exceeded, the error names the budget, the checkpoint size it was derived from, and the knob that overrides it, instead of surfacing a bare `context deadline exceeded`. + +### The cold-load lock ceiling + The router also bounds how long a single cold load may hold the per-model advisory lock, so a worker that dies mid-install cannot pin every other replica's request for that model. That bound is **derived**, not configured, and it is based on *progress* rather than on wall-clock time. The load starts with a base budget of `max(backend-install-timeout + model-load-timeout + 5m, 25m)` — with the defaults, `15m + 5m + 5m = 25m`. That budget covers the steps that report no progress: node selection, backend install, and the remote `LoadModel` call. Raising either timeout widens it in step, so a longer load deadline is never clipped. +That base is the hold's *starting* budget, not its maximum. Because the derived load budget above can exceed it — a 70 GB checkpoint's 28m20s against a 25m base — the hold is widened again as the router enters the load phase, by the derived budget plus the same 5m of slack. Without that step the ceiling would cancel a load that was still comfortably inside its own deadline. + While **model files are staging**, however, the deadline extends every time staging does real work, and expires only once staging has been silent for a 5-minute stall window. Real work means uploaded bytes, and also the resumable-upload verify phase: when a shard is already present on the worker from an earlier attempt, the frontend HEADs it and hashes the local copy to confirm it matches, then skips the transfer. That phase uploads nothing at all — on a 70 GB model resuming with 56 GB already staged it ran for six-plus consecutive minutes at ~45s per shard — so hashing counts as progress too. Otherwise a resumed transfer would be mistaken for a wedged one. Staging time is a function of checkpoint size and available bandwidth, not a constant: a 70 GB model at 26 MB/s needs about 45 minutes, and a 600 GB checkpoint needs hours. A fixed ceiling would therefore be a model-size cliff — every increase just moves the cliff to the next larger model. Extending on progress means a large model transfers for as long as it legitimately needs, while a worker that dies mid-transfer still releases the lock within the stall window. 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.