LocalAI/core/services/nodes/router.go

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/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
	// 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
	// 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]chan struct{}
}

// 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,
		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]chan struct{}{},
	}
}

// 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) {

	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 := modelPayloadBytes(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)
		}
	}

	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.
			if 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)
	}

	// initialInFlight=0: reconciler is pre-loading, not serving a request.
	// scheduleAndLoad picks both the node and the replica slot internally.
	result, err := r.scheduleAndLoad(ctx, backendType, modelName, modelName, revision, &modelOpts, false, 0)
	if err != nil {
		return nil, err
	}
	return result.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 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", modelID, sched.NodeSelector)
	}
	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.
	selectNode := func() *BackendNode {
		var candidate *BackendNode
		var selErr error
		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, _ = r.registry.FindLeastLoadedNodeFromSet(ctx, candidateNodeIDs)
				}
			} else {
				candidate, selErr = r.registry.FindIdleNode(ctx)
				if selErr != nil {
					candidate, _ = r.registry.FindLeastLoadedNode(ctx)
				}
			}
		}
		return candidate
	}

	node := r.pickReachableNode(ctx, selectNode)

	// 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", evictErr)
			}
			return nil, "", 0, fmt.Errorf("no healthy nodes available and eviction failed: %w", evictErr)
		}
		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 {
		// 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(modelPayloadBytes(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.
	resCh := r.installFlight.DoChan(key, func() (any, error) {
		reply, err := r.unloader.InstallBackend(node.ID, backendType, modelID, r.galleriesJSON, "", "", "", replicaIndex, "", nil)
		if err != nil {
			return "", err
		}
		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(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)
	}

	// 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++
		}
	}

	// 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 = ""
			}
		}
	}

	// 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.
	r.stageGenericOptions(ctx, node, opts.Options, frontendModelsDir, localModelDir, keyMapper.Key)
	r.stageGenericOptions(ctx, node, opts.Overrides, frontendModelsDir, localModelDir, keyMapper.Key)

	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 "<file>.sha256" for every file it accepts, so re-staging a
// directory that already held sidecars produced "<file>.sha256.sha256", then
// "<file>.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.
func (r *SmartRouter) stageGenericOptions(ctx context.Context, node *BackendNode, options []string, frontendModelsDir, modelDir string, keyFn func(string) string) {
	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() {
			r.stageOptionDir(ctx, node, absPath, keyFn)
			xlog.Debug("Staged option directory", "option", optKey, "localPath", absPath)
			continue
		}

		key := keyFn(absPath)
		if _, err := r.fileStager.EnsureRemote(ctx, node.ID, absPath, key); err != nil {
			xlog.Warn("Failed to stage option file, skipping", "option", opt, "path", absPath, "error", err)
			continue
		}
		xlog.Debug("Staged option file", "option", optKey, "localPath", absPath)
	}
}

// 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.
func (r *SmartRouter) stageOptionDir(ctx context.Context, node *BackendNode, dir string, keyFn func(string) string) {
	_ = 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 "<file>.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 "<file>.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
		}
		if _, err := r.fileStager.EnsureRemote(ctx, node.ID, path, keyFn(path)); err != nil {
			xlog.Warn("Failed to stage option directory file, skipping", "path", path, "error", err)
		}
		return nil
	})
}

// 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
}

// 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.
func (r *SmartRouter) EvictLRU(ctx context.Context, nodeID string) (string, error) {
	lru, err := r.registry.FindLRUModel(ctx, nodeID)
	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")

// 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)
			}
			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.UnloadModelOnNode(lru.NodeID, lru.ModelName); 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 all models have in-flight requests
		if attempt == 0 {
			xlog.Info("All models have in-flight requests, 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
}