package nodes import ( "context" "encoding/json" "errors" "fmt" "io" "io/fs" "os" "path/filepath" "strings" "sync" "time" "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/core/services/workerctl" "github.com/mudler/LocalAI/pkg/distributedhdr" grpc "github.com/mudler/LocalAI/pkg/grpc" pb "github.com/mudler/LocalAI/pkg/grpc/proto" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/vram" "github.com/mudler/xlog" "golang.org/x/sync/singleflight" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" "google.golang.org/protobuf/proto" "gorm.io/gorm" "gorm.io/gorm/clause" ) // companionSuffixes maps a file extension to additional suffixes that should // be staged alongside it. For example, piper TTS loads ".onnx.json" implicitly // when given an ".onnx" model file. var companionSuffixes = map[string][]string{ ".onnx": {".onnx.json"}, } // SmartRouterOptions holds all dependencies for constructing a SmartRouter. // Passing them at construction time eliminates data races from post-creation setters. type SmartRouterOptions struct { Unloader NodeCommandSender // ModelCleanup performs acknowledged exact-process cleanup when a load // finishes after its configuration revision became stale. ModelCleanup *ModelCleanupService FileStager FileStager GalleriesJSON string AuthToken string ClientFactory BackendClientFactory // optional; defaults to tokenClientFactory DB *gorm.DB // for advisory locks during routing // ConflictResolver, when set, lets the scheduler narrow placement // candidates by per-model concurrency_groups (#9659). When nil, group // anti-affinity is disabled at the scheduler layer; the per-node // watchdog still enforces the rule on arrival. ConflictResolver ConcurrencyConflictResolver // PinnedResolver, when set, excludes `pinned: true` models from the // automatic eviction paths (EvictLRU, evictLRUAndFreeNode) so the pin // contract holds cluster-wide, mirroring the per-node watchdog (#11101). // nil disables the exclusion. Deliberate teardown (UnloadModel, admin // endpoints, node drain) is unaffected. PinnedResolver PinnedModelResolver // ModelFiles, when set, returns the absolute local paths of every file a // model's install declared (gallery `files:`, config `download_files`). // The path fields of a load request name only what the backend opens // first; this is how staging learns about the rest, such as the other // shards of a split GGUF. nil stages the path fields alone. ModelFiles func(modelName string) []string // PrefixProvider, when set, enables prefix-cache-aware routing: requests // carrying a prompt prefix chain (distributedhdr.PrefixChain) are biased // toward the node that already holds the longest matching prefix, subject // to the load guard in prefixcache.Select. nil disables it entirely and // routing is byte-for-byte the round-robin floor. At runtime this is the // *prefixcache.Sync so Observe/Invalidate broadcast to peers. PrefixProvider prefixcache.Provider // PrefixConfig holds the global policy + thresholds. Per-model overrides on // ModelSchedulingConfig refine it per request. Unused when PrefixProvider // is nil. PrefixConfig prefixcache.Config // Pressure, when set, records a forced-disturb each time a request had a // usable hot prefix match but the load guard forced it off the warm node. // The reconciler reads the same instance to autoscale a saturated cache-warm // replica. nil disables recording (the disabled path stays a no-op). Pressure *prefixcache.Pressure // DiskHeadroomEnabled is read LIVE on every scheduling decision (not // snapshotted at construction) so the operator's runtime toggle takes // effect without a restart. nil means enabled — the safe default, and what // every embedder/test that never wires the knob gets. DiskHeadroomEnabled func() bool // SharedModels asserts that every node mounts the same models directory at // the same path. When true, stageModelFiles skips all uploading and leaves // the absolute model paths untouched so the worker loads them directly from // the shared volume (#10556). See config.DistributedConfig.SharedModels. SharedModels bool // ModelLoadCeiling is the hard upper bound on how long a single cold-load // attempt (node selection -> backend install -> file staging -> LoadModel) // may run while holding the per-model advisory lock. It backstops every // sub-step's own timeout so a wedged worker can never pin the lock - and // every other replica's request for that model - indefinitely. Zero derives // it from the default install budget and ModelLoadTimeout via // ModelLoadCeilingFor. ModelLoadCeiling time.Duration // 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 // bytes keep moving the hold extends, so transfer size no longer bounds the // load. Zero selects stagingStallWindow. It is clamped to ModelLoadCeiling. StagingStallWindow time.Duration // ModelLoadAbsoluteMax bounds the cold-load hold even while progress keeps // 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 // install + remote-load budgets to cover the steps that have no deadline of // their own: node selection, replica allocation and model file staging. const modelLoadStagingMargin = 5 * time.Minute // minModelLoadCeiling floors the derived ceiling at the historical 25m constant // so shrinking the install or load budgets can never tighten the hold ceiling // below what clusters relied on before it became derived. const minModelLoadCeiling = 25 * time.Minute // ModelLoadCeilingFor derives the cold-load hold ceiling from the budgets it has // to cover. The ceiling bounds how long one cold load may hold the per-model // advisory lock; it must comfortably exceed the slowest legitimate load (backend // install + staging + remote LoadModel) so it only ever fires when a step is // genuinely wedged - e.g. a worker that died mid-install. Deriving it means // 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 } if loadTimeout <= 0 { loadTimeout = config.DefaultModelLoadTimeout } return max(installTimeout+loadTimeout+modelLoadStagingMargin, minModelLoadCeiling) } // SmartRouter routes inference requests to the best available backend node. // It uses the ModelRouter interface (backed by NodeRegistry in production) for routing decisions. type SmartRouter struct { registry ModelRouter unloader NodeCommandSender // optional, for NATS-driven load/unload modelCleanup *ModelCleanupService fileStager FileStager // optional, for distributed file transfer galleriesJSON string // backend gallery config for dynamic installation clientFactory BackendClientFactory // creates gRPC backend clients db *gorm.DB // for advisory locks during routing stagingTracker *StagingTracker // tracks file staging progress for UI visibility conflictResolver ConcurrencyConflictResolver // pinnedResolver feeds the eviction paths the set of pinned model names // (see SmartRouterOptions.PinnedResolver). nil disables the exclusion. pinnedResolver PinnedModelResolver // modelFiles resolves a model's declared files (see // SmartRouterOptions.ModelFiles). nil stages the path fields alone. modelFiles func(modelName string) []string // prefixProvider is the prefix-cache routing seam (nil disables it; see // SmartRouterOptions.PrefixProvider). prefixConfig holds the global policy // and thresholds. prefixProvider prefixcache.Provider prefixConfig prefixcache.Config // pressure records forced-disturb events (hot match forced off the warm // node by the load guard). nil disables recording. See SmartRouterOptions. pressure *prefixcache.Pressure // installFlight coalesces concurrent identical NATS install requests // (same nodeID + backend + modelID + replica) so 6 simultaneous chat // completions for one not-yet-loaded model produce ONE round-trip, not // six. Avoids amplifying head-of-line blocking on the worker side. installFlight singleflight.Group // probeCache memoizes recent successful gRPC HealthCheck results so // per-request routing doesn't stall behind a busy backend's serialized // HealthCheck/Predict. See probe_cache.go for the rationale. probeCache *probeCache // sharedModels skips file staging when all nodes mount the same models // directory at the same path (see SmartRouterOptions.SharedModels). sharedModels bool // diskHeadroomEnabled is the live read of the operator's disk-headroom // toggle (see SmartRouterOptions.DiskHeadroomEnabled). Never nil after // NewSmartRouter. diskHeadroomEnabled func() bool // modelLoadCeiling bounds how long a cold load may hold the per-model // advisory lock (see SmartRouterOptions.ModelLoadCeiling). modelLoadCeiling time.Duration // 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). 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 // load_job_runner.go. loadWaitersMu sync.Mutex loadWaiters map[string]*loadWaiter // leaseTTL overrides loadJobLeaseTTL for the owner's own deadline (tests). leaseTTL time.Duration // opRenewEvery overrides loadOpRenewEvery, in heartbeat ticks (tests). opRenewEvery int } // probeCacheTTL is how long a successful gRPC HealthCheck on a backend is // trusted before the next request re-probes. Matches healthCheckTTL in // pkg/model/model.go so the single-process and distributed paths share a // staleness budget. The background HealthMonitor still reaps dead backends // independently within ~45s (see perModelMissThreshold). const probeCacheTTL = 30 * time.Second // NewSmartRouter creates a new SmartRouter backed by the given ModelRouter. // All optional dependencies are passed via SmartRouterOptions to avoid post-creation races. func NewSmartRouter(registry ModelRouter, opts SmartRouterOptions) *SmartRouter { factory := opts.ClientFactory if factory == nil { factory = &tokenClientFactory{token: opts.AuthToken} } // 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) } // Default ON: a nil provider must not silently disable a safety check. diskHeadroom := opts.DiskHeadroomEnabled if diskHeadroom == nil { diskHeadroom = func() bool { return true } } return &SmartRouter{ registry: registry, unloader: opts.Unloader, modelCleanup: opts.ModelCleanup, fileStager: opts.FileStager, galleriesJSON: opts.GalleriesJSON, clientFactory: factory, db: opts.DB, stagingTracker: NewStagingTracker(), conflictResolver: opts.ConflictResolver, pinnedResolver: opts.PinnedResolver, modelFiles: opts.ModelFiles, probeCache: newProbeCache(probeCacheTTL), prefixProvider: opts.PrefixProvider, prefixConfig: opts.PrefixConfig, pressure: opts.Pressure, sharedModels: opts.SharedModels, diskHeadroomEnabled: diskHeadroom, modelLoadCeiling: ceiling, modelLoadTimeout: loadTimeout, // Zero values are resolved to their defaults inside // newLoadDeadlineContext, which also clamps the stall window to the // ceiling, so nothing to normalize here. stagingStallWindow: opts.StagingStallWindow, modelLoadAbsoluteMax: opts.ModelLoadAbsoluteMax, modelLoadWait: opts.ModelLoadWait, loadWaiters: map[string]*loadWaiter{}, } } // Unloader returns the remote unloader adapter for external use. func (r *SmartRouter) Unloader() NodeCommandSender { return r.unloader } // StagingTracker returns the staging progress tracker for UI visibility. func (r *SmartRouter) StagingTracker() *StagingTracker { return r.stagingTracker } // scheduleLoadResult holds the result of scheduling and loading a model on a node. type scheduleLoadResult struct { Node *BackendNode Client grpc.Backend BackendAddr string ReplicaIndex int } // applyNodeHardwareDefaults tunes node-agnostic ModelOptions to the GPU of the // node that was actually selected to run the model, reusing the same hardware // heuristics as single-host config loading (core/config). On Blackwell it // raises the physical batch; on non-Blackwell it resets a hardware-default that // an upstream host (the GPU-less frontend in distributed mode) guessed higher. // Only values the heuristics themselves manage are touched, so an explicit user // batch (e.g. 1024) is never overridden. func applyNodeHardwareDefaults(opts *pb.ModelOptions, node *BackendNode, backend string) { if opts == nil || node == nil || config.HardwareDefaultsDisabled() { return } // Gate the throughput heuristics on the node's BUDGETED ceiling, not its // physical card. node.TotalVRAM is stored raw (so a percentage budget can be // recomputed if capacity changes), but the batch/parallel boosts allocate a // per-device compute buffer against real capacity: an operator who budgeted // only a slice of the GPU to LocalAI must not get defaults sized for the full // device, or the raised batch/slots would overflow the allocation (#10485). usableVRAM := node.TotalVRAM if node.VRAMBudgetBytes > 0 && node.VRAMBudgetBytes < usableVRAM { usableVRAM = node.VRAMBudgetBytes } gpu := config.GPU{ Vendor: node.GPUVendor, ComputeCapability: node.GPUComputeCapability, VRAM: usableVRAM, } if config.IsManagedPhysicalBatch(int(opts.NBatch)) { // Gate the raised batch on the selected node's per-device VRAM at this // model's context, so a large context can't overflow the node's compute // buffer (issue #10485). node.TotalVRAM is the node's reported ceiling. opts.NBatch = int32(config.PhysicalBatchForContext(gpu, int(opts.ContextSize))) } // Default concurrent serving for the selected node (the frontend that built // the options may have no GPU). Gated on the node's per-device VRAM at this // model's context, so a large context that already fills the device can't // tip it into OOM by adding slot scratch (issue #10485). Only adds when no // parallel option is set. parallel is a llama.cpp option string, so it is // also gated by backend: a strict backend (e.g. longcat-video) rejects an // unknown option at LoadModel. The typed NBatch above needs no such gate. if config.UsesLlamaCppServingOptions(backend) { opts.Options = config.EnsureParallelOptionForContext(opts.Options, gpu, int(opts.ContextSize)) } } // scheduleAndLoad is the shared core for loading a model on a new node. // Used by both Route() (for first-time loads) and ScheduleAndLoadModel() (for reconciler scale-ups). // // Steps: pick node + replica slot → install backend → stage files → LoadModel → SetNodeModel. // // scheduleNewModel allocates the replica index internally so the worker's // processKey, port, and the registry row all agree. func (r *SmartRouter) scheduleAndLoad(ctx context.Context, backendType, trackingKey, modelName string, configRevision string, modelOpts *pb.ModelOptions, parallel bool, initialInFlight int) (*scheduleLoadResult, error) { // With a database every cold load runs under a job row. Mark the path so a // registry write that lost its ownership value is refused, not waved through. if r.db != nil { ctx = withLoadPath(ctx) } node, backendAddr, replicaIndex, err := r.scheduleNewModel(ctx, backendType, trackingKey, modelOpts) if err != nil { return nil, fmt.Errorf("no available nodes: %w", err) } // Tune node-agnostic options to the SELECTED node's GPU. Only now do we know // which node (and its compute capability) will run the model — the frontend // that built modelOpts may have no GPU at all in distributed mode. applyNodeHardwareDefaults(modelOpts, node, backendType) // Publish the load lifecycle from the moment the node is chosen: staging a // large model takes minutes, and without a registry row the whole phase is // invisible to /api/nodes and the UI — a cold load looks exactly like // nothing happening. The row also reserves the replica slot against // concurrent schedulers. Removed on any failure below so a dead load does // not leave a phantom replica. if err := r.setNodeModelState(ctx, node.ID, trackingKey, replicaIndex, "staging", backendAddr, 0, configRevision, ""); err != nil { return nil, fmt.Errorf("recording staging state: %w", err) } reportLoadPhase(ctx, LoadJobStateStaging, node, replicaIndex) lifecycleSettled := false defer func() { if lifecycleSettled { return } cleanupCtx := context.WithoutCancel(ctx) // An edit may have quarantined this row while staging/loading was in // flight. Its cleanup intent is durable and must not be erased by the // ordinary failed-load cleanup path. if configRevision != "" { if current, err := r.registry.GetModelConfigRevision(cleanupCtx, trackingKey); err == nil && current != configRevision { return } } if err := r.registry.RemoveNodeModel(cleanupCtx, node.ID, trackingKey, replicaIndex); err != nil { xlog.Warn("Failed to clear lifecycle row after failed load", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", err) } }() // 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 := r.stagingPayloadBytes(trackingKey, modelOpts) loadTimeout := r.loadTimeoutFor(payloadBytes) // Pre-stage model files via FileStager before loading loadOpts := modelOpts if r.fileStager != nil && modelOpts != nil { staged, err := r.stageModelFiles(ctx, node, modelOpts, trackingKey) if err != nil { return nil, fmt.Errorf("staging model files for node %s: %w", node.Name, err) } loadOpts = staged } effectiveOptionsHash := "" if loadOpts != nil { var err error effectiveOptionsHash, err = config.EffectiveModelOptionsHash(loadOpts) if err != nil { return nil, fmt.Errorf("hashing effective model options: %w", err) } } reportLoadAddress(ctx, backendAddr) client := r.buildClientForAddr(node, backendAddr, parallel) // Load the model on the remote node if loadOpts != nil { xlog.Info("Loading model on remote node", "node", node.Name, "model", modelName, "addr", backendAddr, "payloadBytes", payloadBytes, "loadBudget", loadTimeout) // Staging is done; the checkpoint load on the worker begins. if err := r.setNodeModelState(ctx, node.ID, trackingKey, replicaIndex, "loading", backendAddr, 0, configRevision, effectiveOptionsHash); err != nil { return nil, fmt.Errorf("recording loading state: %w", 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 // 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, 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 // waiting: an 83GB checkpoint was seen still downloading 30 // minutes past the client timeout, and each retry stacked another // multi-GB loader process on the worker. Reap the replica we just // abandoned before handing the failure back. // A load owner stops its own work through the operation stop path, // so it does not reap here as well. if _, owned := ctx.Value(loadOwnershipKey{}).(LoadJobRef); !owned && loadAbandonedOnWorker(err) { r.reapAbandonedLoad(node, trackingKey, replicaIndex) } return nil, fmt.Errorf("loading model %s on node %s: %w", modelName, node.Name, err) } if !res.Success { return nil, fmt.Errorf("loading model %s on node %s: %s", modelName, node.Name, res.Message) } } // Record the model as loaded on this node (specific replica slot). From // here the row is authoritative; the failure-cleanup defer must not touch it. if err := r.setNodeModelState(ctx, node.ID, trackingKey, replicaIndex, "loaded", backendAddr, initialInFlight, configRevision, effectiveOptionsHash); err != nil { if errors.Is(err, ErrStaleModelConfigRevision) { lifecycleSettled = true r.cleanupStaleLoad(ctx, node, trackingKey, replicaIndex, backendAddr, configRevision, effectiveOptionsHash) } return nil, fmt.Errorf("publishing loaded model: %w", err) } lifecycleSettled = true // Store load metadata for future replica scale-ups by the reconciler. // Writes both per-replica (NodeModel.model_opts_blob) for backward compat // and per-model (ModelLoadInfo table) so the reconciler can recover after // every replica row has been removed (Bug-1). if modelOpts != nil { if optsBlob, marshalErr := proto.Marshal(modelOpts); marshalErr == nil { if storeErr := r.setNodeModelLoadInfo(ctx, node.ID, trackingKey, replicaIndex, backendType, configRevision, optsBlob); storeErr != nil { xlog.Warn("Failed to store model load info", "node", node.Name, "model", trackingKey, "replica", replicaIndex, "error", storeErr) } if storeErr := r.upsertModelLoadInfo(ctx, trackingKey, backendType, configRevision, optsBlob); storeErr != nil { xlog.Warn("Failed to upsert per-model load info", "model", trackingKey, "error", storeErr) } } } return &scheduleLoadResult{Node: node, Client: client, BackendAddr: backendAddr, ReplicaIndex: replicaIndex}, nil } func (r *SmartRouter) cleanupStaleLoad(ctx context.Context, node *BackendNode, modelName string, replicaIndex int, address, revision, hash string) { if r.modelCleanup == nil { xlog.Warn("Stale model load requires exact cleanup", "node", node.Name, "model", modelName, "replica", replicaIndex) return } replica, err := r.registry.GetNodeModel(context.WithoutCancel(ctx), node.ID, modelName, replicaIndex) if err != nil { replica = &NodeModel{NodeID: node.ID, ModelName: modelName, ReplicaIndex: replicaIndex, Address: address, State: "unloading", ConfigRevision: revision, EffectiveOptionsHash: hash} } r.modelCleanup.Cleanup(context.WithoutCancel(ctx), []NodeModel{*replica}, false) } func (r *SmartRouter) setNodeModelState(ctx context.Context, nodeID, modelName string, replicaIndex int, state, address string, initialInFlight int, revision, hash string) error { if revision == "" { if _, err := r.registry.GetModelConfigRevision(ctx, modelName); err == nil { return ErrStaleModelConfigRevision } else if !errors.Is(err, gorm.ErrRecordNotFound) { return err } return r.registry.SetNodeModel(ctx, nodeID, modelName, replicaIndex, state, address, initialInFlight) } return r.registry.SetNodeModelRevision(ctx, nodeID, modelName, replicaIndex, state, address, initialInFlight, revision, hash) } func (r *SmartRouter) setNodeModelLoadInfo(ctx context.Context, nodeID, modelName string, replicaIndex int, backendType, revision string, blob []byte) error { if revision == "" { return r.registry.SetNodeModelLoadInfo(ctx, nodeID, modelName, replicaIndex, backendType, blob) } return r.registry.SetNodeModelLoadInfoRevision(ctx, nodeID, modelName, replicaIndex, backendType, revision, blob) } func (r *SmartRouter) upsertModelLoadInfo(ctx context.Context, modelName, backendType, revision string, blob []byte) error { if revision == "" { return r.registry.UpsertModelLoadInfo(ctx, modelName, backendType, blob) } return r.registry.UpsertModelLoadInfoRevision(ctx, modelName, backendType, revision, blob) } // loadAbandonedOnWorker reports whether a failed remote LoadModel left the // worker process still running the load. // // Only a deadline or a cancellation qualifies: those are OUR side giving up // while the backend keeps going. Every other failure (an unsupported model, a // bad option, an OOM) is the backend answering, which means its handler // returned and the process is idle — stopping it there would throw away a // warm process and its downloaded weights for the next attempt. func loadAbandonedOnWorker(err error) bool { if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) { return true } // The deadline usually reaches us as a gRPC status rather than a wrapped // context error. errors.As (not status.Code) so a wrapped status is still // recognised. var st interface{ GRPCStatus() *status.Status } if errors.As(err, &st) { switch st.GRPCStatus().Code() { case codes.DeadlineExceeded, codes.Canceled: return true } } return false } // reapAbandonedLoad tells the worker to stop the one replica whose load we // just abandoned. The exact `modelID#replicaIndex` process key matters: the // bare model ID would stop every replica of the model on that node, including // healthy ones serving traffic. // // Best-effort by design — the caller is waiting on the load failure, and a // failed reap must be visible in the logs, never substituted for it. func (r *SmartRouter) reapAbandonedLoad(node *BackendNode, trackingKey string, replicaIndex int) { if r.unloader == nil { return } processKey := model.BackendProcessKey(trackingKey, replicaIndex) xlog.Warn("Reaping abandoned model load on worker", "node", node.Name, "model", trackingKey, "replica", replicaIndex) if err := r.unloader.StopBackend(node.ID, processKey); err != nil { xlog.Warn("Failed to reap abandoned model load; the worker may still be loading", "node", node.Name, "backend", processKey, "error", err) } } // ScheduleAndLoadModel implements ModelScheduler for the reconciler. // It retrieves stored model options from an existing replica and performs the // full load sequence (stage files, LoadModel, SetNodeModel) on a new node. func (r *SmartRouter) ScheduleAndLoadModel(ctx context.Context, modelName string, candidateNodeIDs []string) (*BackendNode, error) { // Get load info from an existing replica (stored when Route() first loaded the model) backendType, revision, optsBlob, err := r.registry.GetModelLoadInfoRevision(ctx, modelName) if err != nil { // No replica has ever been loaded for this model, so we have no // backend type or model options to replicate. The previous fallback // fired backend.install with backend="" every reconciler tick, which // the worker rejected ("backend name is empty"). Skip cleanly: the // model needs to be served at least once via Route() so its load // info is stored — then the reconciler can replicate it. return nil, fmt.Errorf("no load info for model %s: serve at least one request for it before the reconciler can replicate (cause: %w)", modelName, err) } // Deserialize the stored model options var modelOpts pb.ModelOptions if err := proto.Unmarshal(optsBlob, &modelOpts); err != nil { return nil, fmt.Errorf("unmarshalling stored model options for %s: %w", modelName, err) } // The reconciler is a load owner like a request is. It claims the job // before any remote work, so a request for the same model waits for this // load instead of scheduling a second copy, and a stale owner is stopped // by the same loop. job, claimed, err := r.registry.ClaimLoadJob(ctx, modelName, ReplicaID()) if err != nil { return nil, fmt.Errorf("claiming the load of model %s: %w", modelName, err) } if !claimed { return nil, fmt.Errorf("model %s is already being loaded by another owner", modelName) } var node *BackendNode err = r.runLoadOwner(ctx, job.Ref(), func(ownerCtx context.Context) error { // initialInFlight=0: reconciler is pre-loading, not serving a request. // scheduleAndLoad picks both the node and the replica slot internally. result, err := r.scheduleAndLoad(ownerCtx, backendType, modelName, modelName, revision, &modelOpts, false, 0) if err != nil { return err } node = result.Node return nil }) if err != nil { return nil, err } return node, nil } // RouteResult contains the routing decision. type RouteResult struct { Node *BackendNode Client grpc.Backend Release func() // Must be called when the request is done (decrements in-flight) } // Route finds the best node for the given model and backend type. // It tries: // 1. Nodes that already have the model loaded (least loaded first) — verified via gRPC health check // 2. Idle-first scheduling: pick an idle node, then fall back to least-loaded. // Sends backend.install via NATS to ensure the right backend is running. // // Returns a RouteResult with a release function that must be called when done. // // modelID is the logical model identifier used for DB tracking (e.g. "qwen_qwen3.5-0.8b"). // modelName is the model file path used for gRPC LoadModel (e.g. "llama-cpp/models/Qwen_...gguf"). // When modelID is empty, modelName is used for both purposes (backward compat). func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType, configRevision string, modelOpts *pb.ModelOptions, parallel bool) (*RouteResult, error) { // Use modelID for DB tracking; fall back to modelName if empty trackingKey := modelID if trackingKey == "" { trackingKey = modelName } if configRevision != "" { if err := r.registry.EstablishModelConfigRevision(ctx, trackingKey, configRevision); err != nil { return nil, fmt.Errorf("establishing config revision for %s: %w", trackingKey, err) } } else if _, err := r.registry.GetModelConfigRevision(ctx, trackingKey); err == nil { return nil, fmt.Errorf("routing %s without a config revision: %w", trackingKey, ErrStaleModelConfigRevision) } else if !errors.Is(err, gorm.ErrRecordNotFound) { return nil, fmt.Errorf("reading config revision for %s: %w", trackingKey, err) } if modelOpts != nil { modelOpts = proto.Clone(modelOpts).(*pb.ModelOptions) } // Fetch the model's scheduling config once: it is immutable for the life of // this request, and resolveSelectorCandidates, buildPreference, and // nodeMatchesScheduling all read it. Fetching once gives a consistent // snapshot and avoids three DB round-trips for one row. nil sched means // "no scheduling constraints", same as before. // GetGoverningScheduling, not GetModelScheduling: a rule may be keyed by an // alias of this model. Request middleware resolves an alias to its target // long before routing, so by here trackingKey is always the target's name // and the alias's rule can only be found by resolving the other way. sched, _ := r.registry.GetGoverningScheduling(ctx, trackingKey) // Resolve the model's NodeSelector once so cached-replica lookup and the // new-load scheduler agree on the candidate set. Without this, a cached // replica on a node the selector now excludes was picked over a matching // replica elsewhere, and the fall-through then tried to load on the // matching node where the model was already at capacity (eviction-busy). candidateNodeIDs, err := r.resolveSelectorCandidates(ctx, trackingKey, sched) if err != nil { return nil, err } // Compute the prefix-cache preference once for this request. pref biases // FindAndLockNodeWithModel toward the warm-cache node; observeChain is // non-nil only when this model uses prefix_cache, gating the Observe calls // 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, configRevision: configRevision, modelOpts: modelOpts, parallel: parallel, sched: sched, candidateNodeIDs: candidateNodeIDs, pref: pref, observeChain: observeChain, } // Step 1: Find and atomically lock a node with this model loaded if result := r.tryWarmPath(ctx, att); result != nil { return result, nil } // Step 2: model not loaded — it has to be cold-loaded. // // 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() if result := r.tryWarmPath(loadCtx, att); result != nil { return result, nil } return r.coldLoad(loadCtx, att, 1) } // 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 configRevision string modelOpts *pb.ModelOptions parallel bool sched *ModelSchedulingConfig candidateNodeIDs []string pref *RoutePreference observeChain []uint64 } // 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 // 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 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 } // Verify node still matches scheduling constraints if !r.nodeMatchesScheduling(ctx, node, att.sched) { 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 } // 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.configRevision, 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. func parseSelectorJSON(selectorJSON string) map[string]string { if selectorJSON == "" { return nil } var selector map[string]string if err := json.Unmarshal([]byte(selectorJSON), &selector); err != nil { xlog.Warn("Failed to parse node selector", "selector", selectorJSON, "error", err) return nil } return selector } func extractNodeIDs(nodes []BackendNode) []string { ids := make([]string, len(nodes)) for i, n := range nodes { ids[i] = n.ID } return ids } // buildPreference computes the per-request route preference from the prefix // chain on ctx and the model's resolved policy. The returned observeChain is // non-nil only when the resolved policy is prefix_cache, signalling Route to // record the assignment after a successful pick; for round-robin models it is // nil so the tree is never polluted. The *RoutePreference is non-nil only when // a load-eligible preferred node was chosen. // // When prefix-cache routing is disabled (nil provider), no chain is present, // or the policy resolves to round-robin, both returns are nil and routing is // the unchanged round-robin floor. func (r *SmartRouter) buildPreference(ctx context.Context, modelID string, candidateNodeIDs []string, sched *ModelSchedulingConfig) (*RoutePreference, []uint64) { if r.prefixProvider == nil { return nil, nil } chain := distributedhdr.PrefixChain(ctx) if len(chain) == 0 { return nil, nil } // Resolve per-model policy + thresholds over the global config. policy := r.prefixConfig.GlobalPolicy cfg := r.prefixConfig if sched != nil { policy = prefixcache.ParsePolicy(sched.RoutePolicy).Resolve(r.prefixConfig.GlobalPolicy) if sched.BalanceAbsThreshold > 0 { cfg.BalanceAbsThreshold = sched.BalanceAbsThreshold } if sched.BalanceRelThreshold > 0 { cfg.BalanceRelThreshold = sched.BalanceRelThreshold } if sched.MinPrefixMatch > 0 { cfg.MinPrefixMatch = sched.MinPrefixMatch } if sched.ScorerWeights != nil { cfg.ScorerWeights = sched.ScorerWeights } } if policy != prefixcache.RoutePolicyPrefixCache { return nil, nil } // Load the candidate replicas PER REPLICA. Affinity is tracked per replica // (each replica is a separate process with its own KV cache), so two // replicas of the same model on the same node are two distinct candidates. // FindAndLockNodeWithModel then locks the EXACT (node, replica) the policy // chose. stats, err := r.registry.LoadedReplicaStats(ctx, modelID, candidateNodeIDs) if err != nil { xlog.Debug("prefixcache: loading replica stats failed, skipping preference", "model", modelID, "error", err) return nil, chain } if len(stats) == 0 { return nil, chain } cands := make([]prefixcache.Candidate, 0, len(stats)) keys := make([]prefixcache.ReplicaKey, 0, len(stats)) for _, s := range stats { key := prefixcache.ReplicaKey{NodeID: s.NodeID, Replica: s.ReplicaIndex} cands = append(cands, prefixcache.Candidate{Key: key, InFlight: s.InFlight}) keys = append(keys, key) } d := r.prefixProvider.Decide(modelID, chain, keys, time.Now()) chosen, ok := prefixcache.Select(cands, d, cfg) // Observability for the prefix-cache routing decision. One line per request // at Debug: enable with DEBUG=true on the frontend to assess cache-aware // routing. hotMatchHonored=true means we routed to the cache-warm replica; // false with HasHot means the load guard forced a cold pick. xlog.Debug("prefix-cache routing decision", "model", modelID, "chainDepth", len(chain), "candidates", len(cands), "hotNode", d.Hot.NodeID, "hotReplica", d.Hot.Replica, "hasHot", d.HasHot, "matchRatio", d.MatchRatio, "minMatch", cfg.MinPrefixMatch, "chosen", fmt.Sprintf("%s/%d", chosen.NodeID, chosen.Replica), "hotMatchHonored", d.HasHot && chosen == d.Hot) // Forced-disturb: a usable hot prefix match existed but the load guard // forced us off the warm replica (Select picked a different replica). This // is the scale-worthy signal - the cache-warm replica is saturated. It // deliberately does not fire for all-unique workloads (no hot match), // avoiding false-positive scale-ups. nil pressure is a no-op. if r.pressure != nil && d.HasHot && d.MatchRatio >= cfg.MinPrefixMatch && chosen != d.Hot { r.pressure.Record(modelID, time.Now()) } if !ok { return nil, chain } return &RoutePreference{PreferredNodeID: chosen.NodeID, PreferredReplica: chosen.Replica}, chain } // observePrefix records that the replica `key` served the request whose prompt // prefix is chain. It is a no-op when prefix-cache routing is disabled or the // chain is empty (round-robin models pass a nil chain so the tree is never // polluted). func (r *SmartRouter) observePrefix(modelID string, chain []uint64, key prefixcache.ReplicaKey) { if r.prefixProvider == nil || len(chain) == 0 { return } r.prefixProvider.Observe(modelID, chain, key, time.Now()) xlog.Debug("prefix-cache observed assignment", "model", modelID, "node", key.NodeID, "replica", key.Replica, "chainDepth", len(chain)) } // resolveSelectorCandidates returns the node IDs that match the model's // NodeSelector. Returns nil when no selector is configured ("any healthy node" // — registry helpers treat nil as no filter). Returns an error when a // non-empty selector matches zero healthy nodes, since there is nothing to // route or schedule on. func (r *SmartRouter) resolveSelectorCandidates(ctx context.Context, modelID string, sched *ModelSchedulingConfig) ([]string, error) { if sched == nil || sched.NodeSelector == "" { return nil, nil } selector := parseSelectorJSON(sched.NodeSelector) if len(selector) == 0 { return nil, nil } candidates, err := r.registry.FindNodesBySelector(ctx, selector) if err != nil { return nil, fmt.Errorf("looking up nodes for selector %s: %w", sched.NodeSelector, err) } if len(candidates) == 0 { return nil, fmt.Errorf("no healthy nodes match selector for model %s: %s: %w", modelID, sched.NodeSelector, ErrNoAvailableNodes) } return extractNodeIDs(candidates), nil } // narrowByGroupAntiAffinity removes candidate nodes that already host a model // declared as concurrent-conflicting with modelID via concurrency_groups // (#9659). This is a soft filter: when *every* candidate would be excluded, // the original set is returned and the per-node watchdog evicts on arrival. // // candidates may be nil ("any healthy node" — registry helpers treat nil as // no filter). nil is returned unchanged: hard-narrowing the implicit "all // nodes" set would silently exclude every node we know nothing about. func (r *SmartRouter) narrowByGroupAntiAffinity(ctx context.Context, modelID string, candidates []string) ([]string, error) { if r.conflictResolver == nil || candidates == nil { return candidates, nil } conflicts := r.conflictResolver.GetModelsConflictingWith(modelID) if len(conflicts) == 0 { return candidates, nil } excluded := make(map[string]struct{}) for _, name := range conflicts { nodes, err := r.registry.FindNodesWithModel(ctx, name) if err != nil { // Best-effort: a single lookup failure shouldn't fail placement. // Log and move on — the watchdog still enforces the rule on arrival. xlog.Warn("Group anti-affinity: lookup failed, skipping", "model", name, "error", err) continue } for _, n := range nodes { excluded[n.ID] = struct{}{} } } if len(excluded) == 0 { return candidates, nil } narrowed := candidates[:0:0] for _, id := range candidates { if _, bad := excluded[id]; bad { continue } narrowed = append(narrowed, id) } if len(narrowed) == 0 { // Soft fallback: every candidate has a conflict. Return the original // set and let the per-node watchdog evict on arrival rather than // failing the request. xlog.Debug("Group anti-affinity: all candidates conflict, falling back to original set", "model", modelID, "conflicts", conflicts) return candidates, nil } xlog.Debug("Group anti-affinity narrowed candidates", "model", modelID, "before", len(candidates), "after", len(narrowed)) return narrowed, nil } // nodeMatchesScheduling checks if a node satisfies the scheduling constraints for a model. // Returns true if no constraints exist or the node matches all selector labels. func (r *SmartRouter) nodeMatchesScheduling(ctx context.Context, node *BackendNode, sched *ModelSchedulingConfig) bool { if sched == nil || sched.NodeSelector == "" { return true // no constraints } selector := parseSelectorJSON(sched.NodeSelector) if len(selector) == 0 { return true } labels, err := r.registry.GetNodeLabels(ctx, node.ID) if err != nil { xlog.Warn("Failed to get node labels for selector check", "node", node.ID, "error", err) return true // fail open } labelMap := make(map[string]string) for _, l := range labels { labelMap[l.Key] = l.Value } for k, v := range selector { if labelMap[k] != v { return false } } return true } // scheduleNewModel picks the best node for loading a new model and allocates // the replica slot. // Strategy: filter to nodes with a free slot for this model → VRAM-aware → // idle-first → least-loaded → eviction. // Sends backend.install via NATS so the chosen node has the right backend running. // // Returns (node, gRPC address, replicaIndex, err). replicaIndex is the slot // the worker has been told to use; the caller must pass the same index into // SetNodeModel so the registry row matches the live process. func (r *SmartRouter) scheduleNewModel(ctx context.Context, backendType, modelID string, modelOpts *pb.ModelOptions) (*BackendNode, string, int, error) { // Estimate VRAM required for the model var estimatedVRAM uint64 if modelOpts != nil { estimatedVRAM = r.estimateModelVRAM(ctx, modelOpts) } // Check for scheduling constraints (node selector). If a selector is set, // we restrict the candidate pool to matching nodes; otherwise nil means // "any healthy node". sched, _ := r.registry.GetGoverningScheduling(ctx, modelID) candidateNodeIDs, err := r.resolveSelectorCandidates(ctx, modelID, sched) if err != nil { return nil, "", 0, err } // Apply concurrency-group anti-affinity (#9659): prefer nodes that don't // already host a model declared exclusive with this one. Soft filter — if // every candidate has a conflict, the original set is returned and the // per-node watchdog evicts on arrival. candidateNodeIDs, err = r.narrowByGroupAntiAffinity(ctx, modelID, candidateNodeIDs) if err != nil { return nil, "", 0, err } // Narrow candidates to nodes that can physically STORE the model. // // Staging writes the checkpoint into the worker's models directory before // the backend ever sees it, so a node without room there cannot succeed no // matter how much VRAM it advertises. Doing this here, before the replica // slot and VRAM passes, is the whole point: the alternative is discovering // it sixteen minutes into a transfer, from the worker, as a 500. candidateNodeIDs, err = r.narrowByDiskHeadroom(ctx, modelID, modelOpts, candidateNodeIDs) if err != nil { return nil, "", 0, err } // Narrow candidates to nodes that still have a free replica slot for this // model. Without this filter, the scheduler would happily pick a node // already at capacity for this model (e.g. when MinReplicas > free // cluster capacity), which is what caused the original 30s flap loop. freeSlotNodes, err := r.registry.FindNodesWithFreeSlot(ctx, modelID, candidateNodeIDs) if err != nil { xlog.Warn("Failed to query nodes with free slot; falling back to selector-only filtering", "model", modelID, "error", err) } else if len(freeSlotNodes) > 0 { // Replace the candidate set with only those that have capacity. candidateNodeIDs = extractNodeIDs(freeSlotNodes) } // If freeSlotNodes is empty (everyone full), candidateNodeIDs is whatever // it was — we'll fall through to eviction below. // Node choice is wrapped in a liveness check: a node's stored status comes // from its HTTP heartbeat, which is a different channel from the bus that // carries the install. A worker that has died stops answering on the bus at // once but stays healthy in the database until its heartbeat ages out, so // without this the scheduler could commit to a node it cannot reach. // // The last selection error is kept because eviction below fires on a nil // node, and a lookup that failed is not the same answer as a cluster with // no room: a control-plane database slow enough to time out these queries // read as "everybody is full" and cost a healthy model its place. var selectErr error selectNode := func() *BackendNode { var candidate *BackendNode var selErr error selectErr = nil if estimatedVRAM > 0 { if candidateNodeIDs != nil { candidate, selErr = r.registry.FindNodeWithVRAMFromSet(ctx, estimatedVRAM, candidateNodeIDs) } else { candidate, selErr = r.registry.FindNodeWithVRAM(ctx, estimatedVRAM) } if selErr != nil { xlog.Warn("No nodes with enough VRAM, falling back to standard scheduling", "required_vram", vram.FormatBytes(estimatedVRAM), "error", selErr) } } if candidate == nil { if candidateNodeIDs != nil { candidate, selErr = r.registry.FindIdleNodeFromSet(ctx, candidateNodeIDs) if selErr != nil { candidate, selErr = r.registry.FindLeastLoadedNodeFromSet(ctx, candidateNodeIDs) } } else { candidate, selErr = r.registry.FindIdleNode(ctx) if selErr != nil { candidate, selErr = r.registry.FindLeastLoadedNode(ctx) } } if candidate == nil { selectErr = selErr } } return candidate } node := r.pickReachableNode(ctx, selectNode) // Same reasoning as the replica-slot guard further down: only // gorm.ErrRecordNotFound is a verdict that the cluster has no node to give. // Any other error left the question unanswered, and evicting on it costs a // healthy model its place for no evidence. if node == nil && selectErr != nil && !errors.Is(selectErr, gorm.ErrRecordNotFound) { return nil, "", 0, fmt.Errorf("selecting a node for %s: %w", modelID, selectErr) } // 4. Preemptive eviction: if no suitable node found, evict the LRU model with zero in-flight if node == nil { evictedNode, evictErr := r.evictLRUAndFreeNodeFrom(ctx, candidateNodeIDs) if evictErr != nil { if errors.Is(evictErr, ErrEvictionBusy) { return nil, "", 0, fmt.Errorf("no healthy nodes available: %w", errors.Join(evictErr, ErrNoAvailableNodes)) } return nil, "", 0, fmt.Errorf("no healthy nodes available and eviction failed: %w", errors.Join(evictErr, ErrNoAvailableNodes)) } node = evictedNode } // Allocate the replica slot before sending backend.install so the worker // uses the same slot for its processKey + port. Default to 0 when the // node's MaxReplicasPerModel is 1 (preserves single-replica behavior). maxSlots := node.MaxReplicasPerModel if maxSlots < 1 { maxSlots = 1 } replicaIdx, slotErr := r.registry.NextFreeReplicaIndex(ctx, node.ID, modelID, maxSlots) if slotErr != nil { // Only ErrNoFreeSlot means "this node is full". Any other error means // we could not find out, and evicting on a guess costs a healthy model // its place: a control-plane database slow enough to time out this // lookup made the scheduler evict loaded models it had no evidence to // evict, and they thrashed. if !errors.Is(slotErr, ErrNoFreeSlot) { return nil, "", 0, fmt.Errorf("determining free replica slot on %s: %w", node.Name, slotErr) } // All slots on this node are taken — fall back to eviction. This is // rare in practice because FindNodesWithFreeSlot already filtered; // it can race with another concurrent scheduler. xlog.Warn("Chosen node has no free replica slot, evicting LRU", "node", node.Name, "model", modelID, "max_slots", maxSlots) evictedNode, evictErr := r.evictLRUAndFreeNodeFrom(ctx, candidateNodeIDs) if evictErr != nil { return nil, "", 0, fmt.Errorf("no replica slot on %s and eviction failed: %w", node.Name, evictErr) } node = evictedNode replicaIdx, slotErr = r.registry.NextFreeReplicaIndex(ctx, node.ID, modelID, node.MaxReplicasPerModel) if slotErr != nil { return nil, "", 0, fmt.Errorf("no replica slot on %s after eviction: %w", node.Name, slotErr) } } // Soft-reserve VRAM up front so a second scheduling tick within the same // heartbeat window can't pick this node based on stale free-VRAM // numbers. The worker's next heartbeat resets reserved_vram to the // authoritative reading; explicit rollback below covers the failure // window between reservation and a successful install. reserved := false if estimatedVRAM > 0 { reserveErr := r.registry.ReserveVRAM(ctx, node.ID, estimatedVRAM) if reserveErr != nil { // ErrInsufficientVRAM races with another scheduler — log and // proceed without a reservation rather than failing the load. // FindNodeWithVRAM already accounted for reserved_vram, so this // is a tight race window; the worker will reconcile via heartbeat. xlog.Warn("Failed to reserve VRAM, proceeding without reservation", "node", node.Name, "bytes", estimatedVRAM, "error", reserveErr) } else { reserved = true } } // Send backend.install — the worker installs the backend if needed and // starts the gRPC process bound to a port for this (model, replica) slot. addr, installErr := r.installBackendOnNode(ctx, node, backendType, modelID, replicaIdx) if installErr != nil { // Roll back the reservation explicitly so the column is accurate // before the next heartbeat. Best-effort. if reserved { _ = r.registry.ReleaseVRAM(ctx, node.ID, estimatedVRAM) } return nil, "", 0, fmt.Errorf("installing backend on node %s: %w", node.Name, installErr) } return node, addr, replicaIdx, nil } // narrowByDiskHeadroom drops candidate nodes that cannot store the model and // returns the surviving set. // // When nothing fits it returns an error wrapping ErrInsufficientDisk so the // caller fails at scheduling time — unless the operator disabled the check, in // which case it warns and hands back the original candidates unchanged. // // modelPayloadBytes stats the same local paths stageModelFiles uploads, which // is why the requirement can be answered before a node is chosen at all. func (r *SmartRouter) narrowByDiskHeadroom(ctx context.Context, modelID string, modelOpts *pb.ModelOptions, candidateNodeIDs []string) ([]string, error) { // Shared-models mode uploads nothing: every node already mounts this exact // models directory at this exact path (stageModelFiles returns early). // Demanding the full checkpoint size of free space per node would reject a // cluster that needs no new bytes at all. if r.sharedModels { return candidateNodeIDs, nil } requiredDisk := DiskRequirementFor(r.stagingPayloadBytes(modelID, modelOpts)) diskCandidates, diskErr := r.registry.NarrowByDiskHeadroom(ctx, candidateNodeIDs, requiredDisk) // The check runs even when disabled. "Disabled" means do not BLOCK, not do // not LOOK: a safety check that goes quiet when switched off is how the // original incident stayed invisible for sixteen minutes. The operator // surrenders the veto and keeps the diagnosis. if !r.diskHeadroomEnabled() { if errors.Is(diskErr, ErrInsufficientDisk) { xlog.Warn("No node has room to store this model, but the disk-headroom check is DISABLED; scheduling anyway — staging will most likely fail with ENOSPC", "model", modelID, "knob", config.FlagDiskHeadroomCheck, "detail", diskErr) } return candidateNodeIDs, nil } switch { case errors.Is(diskErr, ErrInsufficientDisk): // Fail here rather than picking a node and letting staging discover it. return nil, fmt.Errorf("scheduling %s: %w", modelID, diskErr) case diskErr != nil: // A registry read failure is not a capacity verdict. Log and schedule // with the unnarrowed set — the old (worse) behaviour, but never a // wedged cluster because a query hiccuped. xlog.Warn("Failed to check node disk headroom; scheduling without the disk filter", "model", modelID, "required", vram.FormatBytes(requiredDisk), "error", diskErr) return candidateNodeIDs, nil } return diskCandidates, nil } // estimateModelVRAM estimates the VRAM required for a model using the unified estimator. func (r *SmartRouter) estimateModelVRAM(ctx context.Context, opts *pb.ModelOptions) uint64 { estCtx, cancel := context.WithTimeout(ctx, 10*time.Second) defer cancel() ctxSize := uint32(opts.ContextSize) if ctxSize == 0 { ctxSize = 8192 } input := vram.ModelEstimateInput{ Options: vram.EstimateOptions{ GPULayers: int(opts.NGPULayers), }, } // Try model file as a local file for GGUF metadata estimation if opts.ModelFile != "" { if _, err := os.Stat(opts.ModelFile); err == nil { input.Files = append(input.Files, vram.FileInput{URI: opts.ModelFile, Size: 0}) } } // Try HF repo from model name (e.g. "org/model") if opts.Model != "" { if repoID, ok := vram.ExtractHFRepoID(opts.Model); ok { input.HFRepo = repoID } } if len(input.Files) == 0 && input.HFRepo == "" && input.Size == "" { return 0 } result, err := vram.EstimateModelMultiContext(estCtx, input, []uint32{ctxSize}) if err != nil { return 0 } return result.VRAMForContext(ctxSize) } // installBackendOnNode sends a NATS backend.install request-reply to the node // and returns the gRPC address. Concurrent identical calls (same nodeID + // backend + modelID + replica) coalesce via singleflight: 6 chat completions // for the same not-yet-loaded model produce 1 NATS round-trip and 6 callers // share the result. This kills the load-amplification we saw in the live // cluster where 6× simultaneous BackendLoader logs sat behind one slow // install in the worker's NATS callback queue. // // Routine load: the worker's fast-path "already running → return current // address" is correct here. Upgrades go through // DistributedBackendManager.UpgradeBackend on the backend.upgrade subject. func (r *SmartRouter) installBackendOnNode(ctx context.Context, node *BackendNode, backendType, modelID string, replicaIndex int) (string, error) { 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 // blocks for its full NATS deadline (15m by default) when a worker accepts // the request but never replies (e.g. it died mid-install). Without ctx // awareness the caller (holding the per-model advisory lock) would sit there // the whole time; here a cancelled ctx (typically the model-load ceiling) // frees the caller promptly. The shared install keeps running in the // background and still coalesces other callers via singleflight. // A load owner starts the backend as an operation the worker bounds. The // operation id is the load job generation, so a stop can name exactly this // attempt. ref, owned := ctx.Value(loadOwnershipKey{}).(LoadJobRef) resCh := r.installFlight.DoChan(key, func() (any, error) { var reply *workerctl.BackendInstallReply var err error if owned { reply, err = r.unloader.InstallBackendOp(node.ID, backendType, modelID, r.galleriesJSON, replicaIndex, "", ref.Generation, r.loadOperationDeadline(), nil) } else { reply, err = r.unloader.InstallBackend(node.ID, backendType, modelID, r.galleriesJSON, "", "", "", replicaIndex, "", nil) } if err != nil { return "", err } if owned && reply.Success && !reply.ReportsOperations { // A worker that predates operations: it cannot confirm a stop. markLegacyWorker(ctx) } if !reply.Success { return "", fmt.Errorf("worker replied with error: %s", reply.Error) } // Return the backend's gRPC address (per-replica port from worker) addr := reply.Address if addr == "" { addr = node.Address // fallback to node base address } return addr, nil }) select { case <-ctx.Done(): return "", ctx.Err() case res := <-resCh: if res.Err != nil { return "", res.Err } return res.Val.(string), nil } } func (r *SmartRouter) buildClientForAddr(node *BackendNode, addr string, parallel bool) grpc.Backend { client := r.clientFactory.NewClient(node.ID, addr, parallel) // Wrap with file staging if configured if r.fileStager != nil { return NewFileStagingClient(client, r.fileStager, node.ID) } return client } // stageModelFiles uploads model files to the backend node via the FileStager. // Returns the ModelOptions with ModelFile and similar direct-path fields rewritten // to absolute remote paths. Generic options (vae_path, etc.) are left as relative // paths — backends resolve them via ModelPath. // // All files are namespaced under trackingKey so that worker-side deletion can // simply remove the {ModelsPath}/{trackingKey}/ directory. func (r *SmartRouter) stageModelFiles(ctx context.Context, node *BackendNode, opts *pb.ModelOptions, trackingKey string) (*pb.ModelOptions, error) { opts = proto.Clone(opts).(*pb.ModelOptions) // Shared-models mode: every node mounts the same models directory at the // same path, so the frontend's absolute model paths are already valid on the // worker. Staging would only re-upload files that already exist on the shared // volume (under a tracking-key subdir the probe never reuses), re-downloading // the model on every load (#10556). Return the clone untouched: no upload, no // path rewrite, no staging tracker. if r.sharedModels { xlog.Info("Skipping model file staging: shared-models mode is on (LOCALAI_DISTRIBUTED_SHARED_MODELS); worker loads directly from the shared volume", "node", node.Name, "modelFile", opts.ModelFile, "trackingKey", trackingKey) return opts, nil } xlog.Info("Staging model files for remote node", "node", node.Name, "modelFile", opts.ModelFile, "trackingKey", trackingKey) // Derive the frontend models directory from ModelFile and Model. // Example: ModelFile="/models/sd-cpp/models/flux.gguf", Model="sd-cpp/models/flux.gguf" // → frontendModelsDir="/models" // A managed artifact anchors on the models root that holds the whole // .artifacts tree instead, so sibling companion snapshots stay in scope. frontendModelsDir := ModelsRootForModelFile(opts.ModelFile, opts.Model) // Local model directory, captured before the ModelFile field is rewritten to // its remote path below. Companion assets declared as option paths (e.g. // sherpa-onnx's tokens.txt / espeak-ng-data) live beside the model, so option // values are resolved relative to this dir as well as frontendModelsDir — // letting a shared config declare them with bare names regardless of whether // Model includes a subdirectory. localModelDir := "" if opts.ModelFile != "" { localModelDir = filepath.Dir(opts.ModelFile) } // Resolved before the path fields are rewritten to remote paths below, // since that is what tells which declared files the fields already cover. declared := existingFiles(r.declaredExtraFiles(trackingKey, opts), node.Name, trackingKey) // keyMapper generates storage keys namespaced under trackingKey, preserving // subdirectory structure relative to frontendModelsDir. This ensures: // 1. All files for a model land in one directory on the worker for clean deletion // 2. Relative option paths (vae_path, etc.) resolve correctly via ModelPath keyMapper := &StagingKeyMapper{ TrackingKey: trackingKey, FrontendModelsDir: frontendModelsDir, } // Stage each model file path field. These fields are used directly by the // gRPC LoadModel call, so they must be rewritten to the absolute remote path. type pathField struct { name string val *string } fields := []pathField{ {"ModelFile", &opts.ModelFile}, {"MMProj", &opts.MMProj}, {"LoraAdapter", &opts.LoraAdapter}, {"DraftModel", &opts.DraftModel}, {"CLIPModel", &opts.CLIPModel}, {"Tokenizer", &opts.Tokenizer}, {"AudioPath", &opts.AudioPath}, } // Count stageable files for progress tracking. Directory models expand to // the number of files they contain, matching what stageDirectory uploads. totalFiles := 0 for _, f := range fields { if *f.val != "" { totalFiles += countStageableFiles(*f.val) } } for _, adapter := range opts.LoraAdapters { if adapter != "" { if _, err := os.Stat(adapter); err == nil { totalFiles++ } } } if opts.LoraBase != "" { if _, err := os.Stat(opts.LoraBase); err == nil { totalFiles++ } } totalFiles += len(declared) // Start tracking staging progress r.stagingTracker.Start(trackingKey, node.Name, totalFiles) defer r.stagingTracker.Complete(trackingKey) fileIdx := 0 for _, f := range fields { if *f.val == "" { continue } // Skip non-existent files. This is legitimate — a backend that takes a // bare HuggingFace repo id gets an optimistically constructed path that // was never materialized, and fetches its own weights on the worker — so // it must not fail the load. But it is warn-level because the same skip // is what a genuine controller-side acquisition gap looks like, and at // debug it left the operator with a reassuring "Staging model files" // line for work that never happened. if _, err := os.Stat(*f.val); os.IsNotExist(err) { xlog.Warn("Skipping staging for non-existent path; the worker will have to source this itself", "field", f.name, "path", *f.val, "node", node.Name, "trackingKey", trackingKey) *f.val = "" continue } localPath := *f.val // Directory models (e.g. qwen3-tts-cpp ships its weights and tokenizer // ggufs under one directory) can't be uploaded as a single file — the // stager would open the directory and read its fd, failing with // "is a directory" (EISDIR). Expand the directory and stage each // contained file, then rewrite the field to the remote directory. if fi, statErr := os.Stat(localPath); statErr == nil && fi.IsDir() { remoteDir, dirErr := r.stageDirectory(ctx, node, trackingKey, localPath, keyMapper, &fileIdx, totalFiles) if dirErr != nil { if f.name == "ModelFile" { xlog.Error("Failed to stage model directory for remote node", "node", node.Name, "field", f.name, "path", localPath, "error", dirErr) return nil, fmt.Errorf("staging model file: %w", dirErr) } xlog.Warn("Failed to stage model directory, clearing field", "field", f.name, "path", localPath, "error", dirErr) *f.val = "" continue } *f.val = remoteDir if f.name == "ModelFile" { opts.ModelPath = DeriveRemoteModelPath(remoteDir, relativeToModelsDir(frontendModelsDir, localPath, opts.Model)) xlog.Debug("Derived remote ModelPath", "modelPath", opts.ModelPath) } continue } fileIdx++ key := keyMapper.Key(localPath) // Attach progress callback to context for byte-level tracking fileName := filepath.Base(localPath) stageCtx := r.withStagingCallback(ctx, trackingKey, fileName, fileIdx, totalFiles) xlog.Info("Staging file", "model", trackingKey, "node", node.Name, "field", f.name, "file", fileName, "fileIndex", fileIdx, "totalFiles", totalFiles) remotePath, err := r.fileStager.EnsureRemote(stageCtx, node.ID, localPath, key) if err != nil { // ModelFile is required — fail the whole operation if f.name == "ModelFile" { xlog.Error("Failed to stage model file for remote node", "node", node.Name, "field", f.name, "path", localPath, "error", err) return nil, fmt.Errorf("staging model file: %w", err) } // Optional files: clear the path so the backend doesn't try a non-existent frontend path xlog.Warn("Failed to stage model file, clearing field", "field", f.name, "path", localPath, "error", err) *f.val = "" continue } r.stagingTracker.FileComplete(trackingKey, fileIdx, totalFiles) xlog.Debug("Staged model field", "field", f.name, "remotePath", remotePath) *f.val = remotePath // Derive ModelPath from the first staged file (ModelFile). // With tracking key namespacing: // remotePath = "/worker/models/{trackingKey}/sd-cpp/models/flux.gguf" // Model = "sd-cpp/models/flux.gguf" // → ModelPath = "/worker/models/{trackingKey}" if f.name == "ModelFile" { opts.ModelPath = DeriveRemoteModelPath(remotePath, relativeToModelsDir(frontendModelsDir, localPath, opts.Model)) xlog.Debug("Derived remote ModelPath", "modelPath", opts.ModelPath) } r.stageCompanionFiles(ctx, node, localPath, keyMapper.Key) } // Handle LoraAdapters (array) — rewritten to absolute remote paths stagedAdapters := make([]string, 0, len(opts.LoraAdapters)) for _, adapter := range opts.LoraAdapters { if adapter == "" { continue } if _, err := os.Stat(adapter); os.IsNotExist(err) { xlog.Debug("Skipping staging for non-existent lora adapter", "path", adapter) continue } fileIdx++ fileName := filepath.Base(adapter) stageCtx := r.withStagingCallback(ctx, trackingKey, fileName, fileIdx, totalFiles) key := keyMapper.Key(adapter) remotePath, err := r.fileStager.EnsureRemote(stageCtx, node.ID, adapter, key) if err != nil { xlog.Warn("Failed to stage lora adapter, skipping", "path", adapter, "error", err) continue } r.stagingTracker.FileComplete(trackingKey, fileIdx, totalFiles) stagedAdapters = append(stagedAdapters, remotePath) } opts.LoraAdapters = stagedAdapters // Handle LoraBase field — rewritten to absolute remote path if opts.LoraBase != "" { if _, err := os.Stat(opts.LoraBase); err == nil { fileIdx++ fileName := filepath.Base(opts.LoraBase) stageCtx := r.withStagingCallback(ctx, trackingKey, fileName, fileIdx, totalFiles) key := keyMapper.Key(opts.LoraBase) if remotePath, err := r.fileStager.EnsureRemote(stageCtx, node.ID, opts.LoraBase, key); err == nil { r.stagingTracker.FileComplete(trackingKey, fileIdx, totalFiles) opts.LoraBase = remotePath } else { xlog.Warn("Failed to stage LoraBase, clearing field", "path", opts.LoraBase, "error", err) opts.LoraBase = "" } } } for _, localPath := range declared { fileIdx++ fileName := filepath.Base(localPath) stageCtx := r.withStagingCallback(ctx, trackingKey, fileName, fileIdx, totalFiles) xlog.Info("Staging declared model file", "model", trackingKey, "node", node.Name, "file", fileName, "fileIndex", fileIdx, "totalFiles", totalFiles) if _, err := r.fileStager.EnsureRemote(stageCtx, node.ID, localPath, keyMapper.Key(localPath)); err != nil { // The install declared it, so the backend may read it: loading // without it fails later with a less useful error. xlog.Error("Failed to stage declared model file for remote node", "node", node.Name, "path", localPath, "error", err) return nil, fmt.Errorf("staging declared model file %s: %w", localPath, err) } r.stagingTracker.FileComplete(trackingKey, fileIdx, totalFiles) } // Stage file paths referenced in generic Options (key:value pairs where values // are file paths). Options stay as relative paths — backends resolve them via ModelPath. for _, options := range [][]string{opts.Options, opts.Overrides} { remoteRoot := r.stageGenericOptions(ctx, node, options, frontendModelsDir, localModelDir, keyMapper.Key) if opts.ModelFile == "" && remoteRoot != "" { // Virtual models have no primary file from which to derive the // worker root. Their relative options must resolve against the // companion assets we actually staged, not the frontend's root. opts.ModelPath = remoteRoot } } return opts, nil } // withStagingCallback creates a context with a progress callback that updates the staging tracker. func (r *SmartRouter) withStagingCallback(ctx context.Context, trackingKey, fileName string, fileIdx, totalFiles int) context.Context { start := time.Now() return WithStagingProgress(ctx, func(fn string, bytesSent, totalBytes int64) { // Byte-level movement is what keeps the cold-load hold alive. Observing // here (rather than at per-file completion) is what makes a single // 600 GB shard distinguishable from a wedged worker: file-granular // progress would look identical to a stall for hours. observeLoadProgress(ctx) var speed string elapsed := time.Since(start) if elapsed > 0 { bytesPerSec := float64(bytesSent) / elapsed.Seconds() speed = humanFileSize(int64(bytesPerSec)) + "/s" } r.stagingTracker.UpdateFile(trackingKey, fn, fileIdx, bytesSent, totalBytes, speed) }) } // 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. // isHashSidecar reports whether path is a checksum sidecar that the receiving // side generated for a neighbouring file (see hashSidecarSuffix in // file_transfer_server.go), rather than a file belonging to the model. // // Staging these is what made model directories grow without bound: the // receiver writes ".sha256" for every file it accepts, so re-staging a // directory that already held sidecars produced ".sha256.sha256", then // ".sha256.sha256.sha256", multiplying the tree on every pass. // // The check is deliberately "a sidecar sitting next to a real file" rather than // a blanket suffix ban, so a model that genuinely ships a .sha256 payload with // no corresponding base file is still transferred. func isHashSidecar(path string) bool { for _, suffix := range []string{targetSidecarSuffix, hashSidecarSuffix} { base, ok := strings.CutSuffix(path, suffix) if !ok { continue } if fi, err := os.Stat(base); err == nil && !fi.IsDir() { return true } } return false } func countStageableFiles(path string) int { fi, err := os.Stat(path) if err != nil { return 0 } if !fi.IsDir() { return 1 } n := 0 _ = filepath.WalkDir(path, func(p string, d fs.DirEntry, walkErr error) error { if walkErr != nil { return nil } // Must mirror stageDirectory's skip list, or the progress bar counts // files that are never uploaded and never reaches 100%. if !d.IsDir() && !isHashSidecar(p) { n++ } return nil }) return n } // stageDirectory stages every file under a directory-based model (e.g. // qwen3-tts-cpp, whose weights and tokenizer ggufs live in one directory). // Each file is uploaded individually with a structure-preserving key; the // returned path is the remote directory that contained them, suitable for the // backend's ModelFile/ModelPath. fileIdx is advanced per staged file so the // staging progress tracker stays accurate. func (r *SmartRouter) stageDirectory(ctx context.Context, node *BackendNode, trackingKey, dir string, keyMapper *StagingKeyMapper, fileIdx *int, totalFiles int) (string, error) { var remoteDir string err := filepath.WalkDir(dir, func(path string, d fs.DirEntry, walkErr error) error { if walkErr != nil { return walkErr } if d.IsDir() { return nil } // Checksum sidecars are regenerated by the receiver for every file it // accepts; re-uploading them makes it write sidecars for the sidecars. if isHashSidecar(path) { return nil } *fileIdx++ fileName := filepath.Base(path) stageCtx := r.withStagingCallback(ctx, trackingKey, fileName, *fileIdx, totalFiles) xlog.Info("Staging file", "model", trackingKey, "node", node.Name, "field", "ModelDir", "file", fileName, "fileIndex", *fileIdx, "totalFiles", totalFiles) remoteFile, err := r.fileStager.EnsureRemote(stageCtx, node.ID, path, keyMapper.Key(path)) if err != nil { return fmt.Errorf("staging %s: %w", path, err) } r.stagingTracker.FileComplete(trackingKey, *fileIdx, totalFiles) // Every file under dir shares the same remote parent directory; derive // it from this file's staged path and its path relative to dir. rel, relErr := filepath.Rel(dir, path) if relErr != nil { return relErr } remoteDir = DeriveRemoteModelPath(remoteFile, rel) r.stageCompanionFiles(ctx, node, path, keyMapper.Key) return nil }) if err != nil { return "", err } if remoteDir == "" { return "", fmt.Errorf("model directory %s contains no files", dir) } return remoteDir, nil } // stageCompanionFiles stages known companion files that exist alongside // localPath. For example, piper TTS implicitly loads ".onnx.json" next to // the ".onnx" model file. Errors are logged but not propagated. // keyFn generates the namespaced storage key for each file path. func (r *SmartRouter) stageCompanionFiles(ctx context.Context, node *BackendNode, localPath string, keyFn func(string) string) { ext := filepath.Ext(localPath) suffixes, ok := companionSuffixes[ext] if !ok { return } base := strings.TrimSuffix(localPath, ext) for _, suffix := range suffixes { companion := base + suffix if _, err := os.Stat(companion); err != nil { continue } key := keyFn(companion) if _, err := r.fileStager.EnsureRemote(ctx, node.ID, companion, key); err != nil { xlog.Warn("Failed to stage companion file", "path", companion, "error", err) } else { xlog.Debug("Staged companion file", "path", companion) } } } // stageGenericOptions iterates key:value option strings and stages any values // that resolve to existing files relative to the frontend models directory or // the model's own directory. Option values are NOT rewritten — backends resolve // them via ModelPath. keyFn generates the namespaced storage key for each file. // Returns the staged models root, or empty when no asset was staged. func (r *SmartRouter) stageGenericOptions(ctx context.Context, node *BackendNode, options []string, frontendModelsDir, modelDir string, keyFn func(string) string) string { remoteRoot := "" for _, opt := range options { optKey, val, ok := strings.Cut(opt, ":") if !ok || val == "" { continue } // Resolve the value to an existing path: absolute as-is, otherwise // relative to frontendModelsDir first, then the model's own directory // (where backends like sherpa-onnx keep companion assets such as // tokens.txt and espeak-ng-data). absPath, ok := resolveOptionPath(val, frontendModelsDir, modelDir) if !ok { continue } info, err := os.Stat(absPath) if err != nil { continue } // A directory option value (e.g. sherpa-onnx's espeak-ng-data) is staged // file-by-file so the whole tree is recreated beside the model on the // worker; a single file is staged directly. Values are never rewritten — // backends resolve relative paths via ModelPath. if err == nil && info.IsDir() { if remoteDir := r.stageOptionDir(ctx, node, absPath, keyFn); remoteDir != "" { remoteRoot = DeriveRemoteModelPath(remoteDir, relativeToModelsDir(frontendModelsDir, absPath, filepath.Base(absPath))) } xlog.Debug("Staged option directory", "option", optKey, "localPath", absPath) continue } key := keyFn(absPath) remotePath, err := r.fileStager.EnsureRemote(ctx, node.ID, absPath, key) if err != nil { xlog.Warn("Failed to stage option file, skipping", "option", opt, "path", absPath, "error", err) continue } remoteRoot = DeriveRemoteModelPath(remotePath, relativeToModelsDir(frontendModelsDir, absPath, filepath.Base(absPath))) xlog.Debug("Staged option file", "option", optKey, "localPath", absPath) } return remoteRoot } // resolveOptionPath finds an existing local path for an option value: an // absolute path as-is, otherwise relative to frontendModelsDir, then to the // model's own directory. Returns false when none exists. func resolveOptionPath(val, frontendModelsDir, modelDir string) (string, bool) { if filepath.IsAbs(val) { if _, err := os.Stat(val); err == nil { return val, true } return "", false } for _, base := range []string{frontendModelsDir, modelDir} { if base == "" { continue } p := filepath.Join(base, val) if _, err := os.Stat(p); err == nil { return p, true } } return "", false } // stageOptionDir stages every regular file under an option-declared directory // (e.g. sherpa-onnx's espeak-ng-data) using the structure-preserving key, so the // tree is recreated beside the model on the worker. Per-file errors are logged // and skipped; the option value itself is not rewritten. Returns the remote // directory derived from a successfully staged file, or empty when none succeeds. func (r *SmartRouter) stageOptionDir(ctx context.Context, node *BackendNode, dir string, keyFn func(string) string) string { remoteDir := "" _ = filepath.WalkDir(dir, func(path string, d fs.DirEntry, walkErr error) error { if walkErr != nil || d.IsDir() { return nil } // Same reason as stageDirectory: the receiver writes ".sha256" for // every file it accepts, so staging the sidecars makes it write sidecars // for those in turn. Option dirs are walked on every load, so each pass // added a level - an espeak-ng-data tree observed in the wild had grown // to ".sha256" repeated eleven times and 5077 junk files, which is // enough to keep a sherpa-onnx voice permanently "staging" and fail // every realtime warmup that needs it. if isHashSidecar(path) { return nil } remotePath, err := r.fileStager.EnsureRemote(ctx, node.ID, path, keyFn(path)) if err != nil { xlog.Warn("Failed to stage option directory file, skipping", "path", path, "error", err) return nil } if rel, err := filepath.Rel(dir, path); err == nil { remoteDir = DeriveRemoteModelPath(remotePath, rel) } return nil }) return remoteDir } // probeHealth checks whether a backend process on the given node/addr is alive // via a gRPC health check with a 2-second timeout. The client is closed after // the check. // // The result is memoized in r.probeCache for probeCacheTTL. With per-request // routing every inference call lands here, and unbounded re-probing can stall // behind a busy backend that serializes HealthCheck against active Predict. // Concurrent probes for the same (node, addr) coalesce via singleflight so a // burst of N requests for a cold cache costs at most one round-trip, not N. // Failed probes invalidate the cache so the staleness recovery path // (DecrementInFlight + RemoveNodeModel) still triggers on the next request. func (r *SmartRouter) probeHealth(ctx context.Context, node *BackendNode, addr string) bool { key := node.ID + "|" + addr return r.probeCache.DoOrCached(key, func() bool { client := r.buildClientForAddr(node, addr, false) defer closeClient(client) checkCtx, cancel := context.WithTimeout(ctx, 2*time.Second) defer cancel() ok, _ := client.HealthCheck(checkCtx) return ok }) } // closeClient closes a gRPC backend client if it implements io.Closer. // newRouteResult builds the RouteResult for a routed replica, wiring the // load-time in_flight reservation to a release that fires exactly once, on // whichever happens first: the triggering inference completing, or the route // being torn down without one ever running. // // Tying it to teardown as well as to the first inference is what stops the // reservation leaking. A route whose caller never reaches the backend (client // disconnect, handler error, validation failure after load) previously left // in_flight pinned at 1 forever, and every eviction query requires // in_flight = 0, so that replica's VRAM could never be reclaimed. func (r *SmartRouter) newRouteResult(node *BackendNode, trackingKey string, replicaIdx int, raw grpc.Backend, tracked *InFlightTrackingClient) *RouteResult { var once sync.Once release := func() { once.Do(func() { if err := r.registry.DecrementInFlight(context.Background(), node.ID, trackingKey, replicaIdx); err != nil { // Worth surfacing: a reservation that fails to come back is // exactly what the leak sweeper later has to clean up. xlog.Warn("Failed to release routing in-flight reservation", "node", node.ID, "model", trackingKey, "replica", replicaIdx, "error", err) } }) } tracked.OnFirstComplete(release) return &RouteResult{ Node: node, Client: tracked, Release: func() { release() closeClient(raw) }, } } func closeClient(client grpc.Backend) { if closer, ok := client.(io.Closer); ok { closer.Close() } } // UnloadModel sends a NATS unload event to a specific node for the given model // and removes every replica row for (nodeID, modelName). // The worker process handles Free() + kill + deregister. func (r *SmartRouter) UnloadModel(ctx context.Context, nodeID, modelName string) error { if r.unloader == nil { return fmt.Errorf("no remote unloader configured") } // Target the specific node, not all nodes hosting this model if err := r.unloader.StopBackend(nodeID, modelName); err != nil { return fmt.Errorf("failed to stop backend on node %s: %w", nodeID, err) } r.registry.RemoveAllNodeModelReplicas(ctx, nodeID, modelName) return nil } // pinnedModelNames returns the pinned set for eviction exclusion, or nil when // no resolver is wired (embedders, tests, deployments without a config loader). func (r *SmartRouter) pinnedModelNames() []string { if r.pinnedResolver == nil { return nil } return r.pinnedResolver.GetPinnedModelNames() } // EvictLRU evicts the least-recently-used model from a node to make room. // Returns the name of the evicted model, or empty string if nothing could be // evicted. Pinned models are never candidates (#11101). func (r *SmartRouter) EvictLRU(ctx context.Context, nodeID string) (string, error) { lru, err := r.registry.FindLRUModel(ctx, nodeID, r.pinnedModelNames()) if err != nil { return "", fmt.Errorf("finding LRU model on node %s: %w", nodeID, err) } if err := r.UnloadModel(ctx, nodeID, lru.ModelName); err != nil { return "", err } return lru.ModelName, nil } // ErrEvictionBusy is returned when all loaded models have in-flight requests // and none can be evicted to make room. var ErrEvictionBusy = errors.New("all models busy, cannot evict") // ErrNoAvailableNodes is returned when the scheduler cannot find any healthy // node to serve a model — all nodes are full and eviction cannot free a slot, // or a node selector excludes every candidate. The HTTP layer maps this to // 503 so clients treat it as a transient condition rather than a server bug // (which is what 500 would imply). var ErrNoAvailableNodes = errors.New("no available nodes") // evictLRUAndFreeNode finds the globally least-recently-used model with zero in-flight, // unloads it, and returns its node for reuse. If all models are busy, retries briefly. // // Uses SELECT FOR UPDATE inside a transaction to prevent two frontends from // simultaneously picking the same eviction target. The NodeModel row is deleted // inside the transaction; the NATS unload command is sent after commit. // evictLRUAndFreeNode evicts across every healthy node. Callers that hold a // candidate set must use evictLRUAndFreeNodeFrom instead. func (r *SmartRouter) evictLRUAndFreeNode(ctx context.Context) (*BackendNode, error) { return r.evictLRUAndFreeNodeFrom(ctx, nil) } // evictLRUAndFreeNodeFrom evicts the least-recently-used idle model from one of // candidateNodeIDs, or from any healthy node when the set is nil. // // Restricting eviction to the candidate set matters whenever the model being // scheduled has a node selector. Evicting globally freed a slot on a node the // selector forbids, so the model was then placed there anyway, on hardware it // was explicitly pinned away from, and an unrelated model was dropped to make // the room. On a cluster where the selector-matching node was momentarily // unavailable this repeated, and the evicted model appeared to bounce between // nodes. func (r *SmartRouter) evictLRUAndFreeNodeFrom(ctx context.Context, candidateNodeIDs []string) (*BackendNode, error) { const maxEvictionRetries = 5 const evictionRetryInterval = 500 * time.Millisecond if r.db == nil { return nil, ErrEvictionBusy // no DB means no row-level locking for safe eviction } for attempt := range maxEvictionRetries { var lru NodeModel err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { // Lock the row so no other frontend can evict the same model. // // The replica-floor guard matches a rule to a replica through // sc.target_model, so a rule keyed by an alias protects the model // the alias points at. It falls back to sc.model_name when the // stored target is empty, which keeps a rule inserted by some path // that never resolved it protecting itself rather than nothing. // // target_model is not unique (two names can resolve to one model), // so the floor is MAX over the matching rules: only one of them // governs, but over-protecting costs a retry while under-protecting // evicts below a floor the reconciler then has to rebuild. q := currentModelRevision(tx.Clauses(clause.Locking{Strength: "UPDATE"})). Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id"). Where(`node_models.in_flight = 0 AND node_models.state = ? AND backend_nodes.status = ? AND ( NOT EXISTS (SELECT 1 FROM model_scheduling_configs sc WHERE COALESCE(NULLIF(sc.target_model, ''), sc.model_name) = node_models.model_name AND (sc.min_replicas > 0 OR sc.max_replicas > 0)) OR (SELECT COUNT(*) FROM node_models nm2 WHERE nm2.model_name = node_models.model_name AND nm2.state = 'loaded' AND (NOT EXISTS (SELECT 1 FROM model_config_states mcs2 WHERE mcs2.model_name = nm2.model_name) OR nm2.config_revision = (SELECT mcs3.config_revision FROM model_config_states mcs3 WHERE mcs3.model_name = nm2.model_name))) > COALESCE((SELECT MAX(sc2.min_replicas) FROM model_scheduling_configs sc2 WHERE COALESCE(NULLIF(sc2.target_model, ''), sc2.model_name) = node_models.model_name), 1) )`, "loaded", StatusHealthy) if len(candidateNodeIDs) > 0 { q = q.Where("node_models.node_id IN ?", candidateNodeIDs) } // Pinned models are protected from automatic eviction (#11101). // Filtered in the query, not after selection, so the next-oldest // unpinned model is chosen instead of the attempt being wasted. if pinned := r.pinnedModelNames(); len(pinned) > 0 { q = q.Where("node_models.model_name NOT IN ?", pinned) } if err := q. Order("node_models.last_used ASC"). First(&lru).Error; err != nil { return err } // Remove inside the same transaction. Target the specific replica row // by ID so we don't accidentally delete sibling replicas of the same // model on the same node. return tx.Where("id = ?", lru.ID).Delete(&NodeModel{}).Error }) if err == nil { xlog.Info("Evicted LRU model to free capacity", "node", lru.NodeID, "model", lru.ModelName, "lastUsed", lru.LastUsed) // Unload outside the transaction (NATS call) if r.unloader != nil { if uerr := r.unloader.UnloadReplica(lru.NodeID, lru); uerr != nil { xlog.Warn("eviction unload failed (model already removed from registry)", "error", uerr) } } node, nodeErr := r.registry.Get(ctx, lru.NodeID) if nodeErr != nil { return nil, fmt.Errorf("node %s not found after eviction: %w", lru.NodeID, nodeErr) } return node, nil } // gorm.ErrRecordNotFound means every candidate is either mid-request // or excluded as pinned if attempt == 0 { xlog.Info("No evictable model (all busy or pinned), waiting for capacity") } select { case <-ctx.Done(): return nil, fmt.Errorf("context cancelled while waiting for eviction") case <-time.After(evictionRetryInterval): // retry } } return nil, ErrEvictionBusy } // loadOperationDeadline is the longest a worker may keep one load running: the // controller's own absolute cap. Renewals are what normally end a load sooner. func (r *SmartRouter) loadOperationDeadline() time.Duration { if r.modelLoadAbsoluteMax > 0 { return r.modelLoadAbsoluteMax } return modelLoadAbsoluteMax }