fix(distributed): route per request across loaded replicas + cache probeHealth (#9968)

* refactor(distributed): extract PickBestReplica from FindAndLockNodeWithModel

Lifts the replica-selection policy (in_flight ASC, last_used ASC,
available_vram DESC) out of the SQL ORDER BY into a pure Go function in
the new replicapicker.go. The SQL clause keeps its FOR UPDATE atomicity
and remains the production path used by SmartRouter; PickBestReplica is
the canonical implementation that the future per-frontend rotating
replica cache (TODO referenced from pkg/model) will call against an
in-memory snapshot without paying a DB round-trip per inference.

A new registry_test mirror spec seeds a multi-tier scenario and asserts
both layers pick the same replica, so any future tweak to either side
fails the test until the other side is updated.

No behavior change.

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-4-7 [Claude Code]

* fix(distributed): route per inference request and cache probeHealth

Two related fixes that together restore load balancing across loaded
replicas of the same model.

1. ModelLoader.Load and LoadModel bypass the local *Model cache when
   modelRouter is set. The cached *Model wraps an InFlightTrackingClient
   bound to a single (nodeID, replicaIndex) — reusing it pinned every
   subsequent request to whichever node won the very first pick, so
   FindAndLockNodeWithModel's round-robin never got a chance to run
   even after the reconciler scaled the model out to a second node. In
   distributed mode SmartRouter.Route now runs per request, and
   PickBestReplica picks the least-loaded replica each time.

   SmartRouter has its own coalescing (advisory DB lock for first-time
   loads + singleflight on backend.install RPC) so concurrent first
   requests for a not-yet-loaded model still produce a single worker
   side install.

2. SmartRouter.probeHealth memoizes successful gRPC HealthCheck results
   in a new probeCache (probe_cache.go) with a 30s TTL. With per-request
   routing every inference call hits probeHealth, and llama.cpp-style
   backends serialize HealthCheck behind active Predict — so a burst of
   incoming requests stalled on the probe to a node already mid-stream,
   tripping the 2s timeout and falling through to the install path.
   singleflight collapses N concurrent first-time probes for the same
   (node, addr) into one round-trip, failed probes invalidate the entry
   so the staleness-recovery path still triggers, and the TTL matches
   pkg/model/model.go's healthCheckTTL so the single-process and
   distributed paths share a staleness budget. The background
   HealthMonitor still reaps actually-dead backends within ~45s.

The bypass introduces one short FindAndLockNodeWithModel transaction per
inference. A TODO in pkg/model/loader.go documents the future per modelID
rotating-replica cache that would reuse PickBestReplica against an
in-memory snapshot and skip the DB round-trip for hot paths.

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Assisted-by: Claude:claude-opus-4-7 [Claude Code]

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Co-authored-by: Ettore Di Giacinto <mudler@localai.io>

pkg/model/initializers.go

@@ -276,6 +276,37 @@ func (ml *ModelLoader) updateModelLastUsed(m *Model) {
 func (ml *ModelLoader) Load(opts ...Option) (grpc.Backend, error) {
 	o := NewOptions(opts...)
 
+	ml.mu.Lock()
+	distributed := ml.modelRouter != nil
+	ml.mu.Unlock()
+
+	// In distributed mode, SmartRouter must run per inference request so
+	// PickBestReplica (core/services/nodes/replicapicker.go) picks the
+	// least-loaded replica each time. Bypass the local cache and the local
+	// LRU / concurrency-group watchdog enforcement: both are scoped to the
+	// in-process Model store, which in distributed mode only holds stubs for
+	// remote replicas. SmartRouter handles cluster-wide eviction
+	// (evictLRUAndFreeNode) and concurrency-group anti-affinity
+	// (narrowByGroupAntiAffinity) at the scheduler layer.
+	//
+	// TODO(distributed-cache): see LoadModel for the rotating-replica-cache
+	// integration point that would let hot paths skip the per-request DB
+	// round-trip without giving up the shared PickBestReplica policy.
+	if distributed {
+		client, err := ml.backendLoader(opts...)
+		if err != nil {
+			return nil, err
+		}
+		if m := ml.CheckIsLoaded(o.modelID); m != nil && m.Process() == nil {
+			client = newConnectionEvictingClient(client, o.modelID, func() {
+				if err := ml.ShutdownModel(o.modelID); err != nil {
+					xlog.Warn("Failed to shut down remote model after connection error", "model", o.modelID, "error", err)
+				}
+			})
+		}
+		return client, nil
+	}
+
 	// Return earlier if we have a model already loaded
 	// (avoid looping through all the backends)
 	if m := ml.CheckIsLoaded(o.modelID); m != nil {

pkg/model/loader.go

@@ -250,6 +250,49 @@ func (ml *ModelLoader) ListLoadedModels() []*Model {
 }
 
 func (ml *ModelLoader) LoadModel(modelID, modelName string, loader func(string, string, string) (*Model, error)) (*Model, error) {
+	ml.mu.Lock()
+	distributed := ml.modelRouter != nil
+	ml.mu.Unlock()
+
+	if distributed {
+		// Distributed mode: SmartRouter must run per inference request so
+		// PickBestReplica (core/services/nodes/replicapicker.go) picks the
+		// least-loaded replica each time. The cached *Model returned from a
+		// previous call holds a client wrapper bound to one (nodeID,
+		// replicaIndex), so reusing it pins every subsequent request to the
+		// node that won the very first pick — defeating per-replica load
+		// balancing. Bypass the cache and the loading-coalesce map; the
+		// router does its own coalescing for first-time loads (advisory DB
+		// lock + singleflight on backend.install RPC), so concurrent first
+		// requests still produce a single worker-side install.
+		//
+		// TODO(distributed-cache): if profiling shows the per-request
+		// FindAndLockNodeWithModel SELECT FOR UPDATE becomes a hot path
+		// under burst load, replace this branch with a per-modelID cache
+		// that holds a *list* of replicas (refreshed every ~5s in
+		// background) and picks per call via PickBestReplica against
+		// locally-tracked in-flight counters. Same policy, no DB round-trip
+		// per inference. Trade-off: cross-frontend in-flight visibility
+		// becomes eventually consistent, acceptable for 1-3 frontend
+		// deployments.
+		modelFile := filepath.Join(ml.ModelPath, modelName)
+		model, err := loader(modelID, modelName, modelFile)
+		if err != nil {
+			return nil, fmt.Errorf("failed to route model with internal loader: %s", err)
+		}
+		if model == nil {
+			return nil, fmt.Errorf("loader didn't return a model")
+		}
+		// Record the latest mapping so DistributedModelStore.Range, shutdown,
+		// and listing endpoints see a representative entry. The DB is the
+		// source of truth for cluster-wide state; the local store is just a
+		// stub for in-process callers.
+		ml.mu.Lock()
+		ml.store.Set(modelID, model)
+		ml.mu.Unlock()
+		return model, nil
+	}
+
 	ml.mu.Lock()
 
 	// Check if we already have a loaded model