fix(distributed): reaper reaps live backends, ghost model stubs, in_flight leak, sidecar staging runaway (#11142)

* fix(distributed): stop the probe reaper from orphaning busy backends

The reconciler's liveness probe is a 1s gRPC HealthCheck, and a single
failed probe deleted the model's node_models row. A backend that is
merely busy cannot answer it: single-threaded Python backends (video and
avatar generation) block for minutes inside one request, so the reaper
was deleting registry rows for backends that were alive and mid-request.

The model then vanished from the nodes page while it was still
generating, and because the row was gone the in-flight decrement had
nothing to decrement ("DecrementInFlight: no matching row or already
zero"). Every subsequent request re-routed and re-staged the full model
from scratch.

Two guards:

  - Replicas with in-flight requests are excluded in SQL. A row that is
    actively serving is proof of life, and the running request is
    exactly what stops the backend from answering the probe.

  - Idle replicas must miss three CONSECUTIVE probes before removal, so
    a transient blip cannot orphan a live replica. A successful probe
    resets the streak.

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

* fix(distributed): drop the local model stub when its last replica goes

In distributed mode every routed model leaves an in-process stub in the
frontend's ModelLoader, and DistributedModelStore.Range reports local
stubs UNION the registry rows. Every registry removal path deletes only
the DB row, so the stub outlived the replica and the model was reported
as loaded forever.

That is the "loaded on the home page, absent from every node" ghost:
/system reads the union and still sees the stub, while /api/nodes/models
reads the registry and correctly sees nothing. It never self-healed,
and both frontend replicas showed it independently.

The replica-removed chokepoint could not fix this as it stood, because
it held a SINGLE hook that the prefix cache already owned, and it was
registered only when the prefix cache was enabled. Registering a second
listener would have silently displaced the first.

  - Turn replicaRemovedHook into a list (AddReplicaRemovedHook), so
    independent subsystems can each register without displacing others.
  - Add NewLocalStubInvalidator, which drops the local stub once no
    healthy replica of the model remains anywhere in the cluster, and
    wire it unconditionally in startup.

The stub is kept while another node still serves the model: the
frontend is right to consider it loaded, and each request re-routes
through SmartRouter to pick a live replica anyway.

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

* fix(distributed): stop staging checksum sidecars back to workers

The file transfer server writes a "<file>.sha256" sidecar next to every
file it accepts. The sender walked the model directory with no filter,
so it staged those sidecars too, and the receiver duly wrote a sidecar
for each sidecar. Every staging pass multiplied the tree:

  config.json -> config.json.sha256 -> config.json.sha256.sha256 -> ...

One LongCat snapshot had grown to 498 files, 466 of them chained, up to
29 levels deep, and the staged file count climbed on every pass. This
inflates each transfer and grows disk without bound on both ends.

Skip hash sidecars in stageDirectory, and mirror the skip in
countStageableFiles so the progress bar still reaches 100%. The check is
"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.

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

* fix(distributed): classify the liveness probe instead of gating on in_flight

The previous commit excluded replicas with in-flight requests from the probe
reaper. That was the wrong guard, and could invert the bug it fixed.

in_flight has no decrement guarantee: track() balances its increment with a
defer, but a frontend killed mid-request never runs it, and the load-time
reservation is released only when the first inference completes. Nothing
resets a leaked counter. Gating the reaper on it therefore meant a leaked
counter would shield a genuinely dead replica from ever being reaped.

Nor was patience alone a fix: three misses at the default interval is ~90s of
silence, while the generation that triggered this blocks for 15+ minutes.

The real conflation was in the probe itself. A gRPC HealthCheck against the
backend's serving port measures "is it idle enough to answer", not "does the
process exist", and probeLoadedModels discarded the error that tells them
apart. Because the gRPC client is lazy, the status code is decisive:

  - DeadlineExceeded: transport fine, nothing serviced the RPC. Busy.
  - Unavailable: nothing is listening. Gone.

ModelProber now returns a ProbeOutcome, and only ProbeUnreachable counts
toward the reap threshold. ProbeBusy clears the streak: it is evidence of
life. A blackholed network reads as busy too, deliberately, since whole-node
failure is the health monitor's job.

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

* feat(distributed): reconcile replicas against worker-reported processes

Probing a backend's own serving port cannot distinguish "busy" from "gone"
without inferring it from an error code. The worker can answer directly: it
spawned the process, holds the handle, and its reply is not blocked by
whatever that backend is doing.

Adds a models.running request-reply subject. The worker answers out of its
in-memory process table, reporting each live process as (modelID,
replicaIndex, address) — the supervisor's process keys are `modelID#replica`,
which is isomorphic to a NodeModel row, so the reconciler can diff the two
directly.

reconcileNodeProcesses runs before the port probe and reaps rows for models
the worker is not running. Models the worker vouches for get updated_at
bumped, which takes them out of the port prober's stale set entirely: that is
what keeps a backend deep in a long generation away from the probe in the
first place, rather than relying on classifying its silence after the fact.

A worker that does not answer is skipped, not assumed empty. A messaging
failure says nothing about the processes, and assuming the worst would delete
a node's rows on a transient NATS blip; the port probe stays as the fallback
for those nodes. Rows younger than probeStaleAfter are ignored so a freshly
created row is never judged against a process table that has not caught up.

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

* fix(distributed): stop in_flight leaking and pin replicas against eviction

A leaked in_flight counter is not cosmetic. FindLRUModel,
FindGlobalLRUModelWithZeroInFlight and the router's eviction query all require
in_flight = 0, so a replica whose counter never came back is pinned and its
VRAM is unreclaimable for the lifetime of the process.

Two halves.

The source: routing reserves in_flight = 1 at load time so a freshly loaded
replica is not evicted out from under the request that caused the load. That
reservation was released ONLY by the first inference completing, so a route
torn down before any inference ran (client disconnect, handler error, failure
between load and the backend call) stranded it. newRouteResult now wires the
reservation to a sync.Once fired by whichever comes first, the first inference
or route teardown, and replaces three copies of the old wiring.

The backstop: a sweeper for counters leaked by paths that cannot run a defer
at all, such as a frontend killed mid-request.

Identifying a leak by elapsed time alone is unsafe. IncrementInFlight stamps
last_used at request START and nothing moves it while the request runs, so a
long generation is indistinguishable from a leak by age, and resetting there
would expose a serving model to eviction. The probe supplies the missing bit:
a backend that answers a health check promptly is not inside a request,
because that is precisely what a busy one cannot do. Requiring the row to also
be idle for 30 minutes covers backends that serve in parallel and can answer
while working, since those keep last_used fresh through each new increment.

Two existing tests asserted the old behaviour ("No decrement on Release").
That assertion was the leak, so both now pin the release instead.

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

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
Co-authored-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
mudler's LocalAI [bot]
2026-07-27 18:36:31 +02:00
committed by GitHub
parent e9c2754fc2
commit 878a0d00a1
22 changed files with 1784 additions and 101 deletions

View File

@@ -286,13 +286,14 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
prefixProvider = prefixSync
// Invalidate the prefix-cache index whenever a replica row is removed.
// SetReplicaRemovedHook fires from the single chokepoint all removal paths
// AddReplicaRemovedHook fires from the single chokepoint all removal paths
// funnel through (RemoveNodeModel / RemoveAllNodeModelReplicas), so this
// one hook covers every path: reconciler scale-down, probe reaper,
// health-monitor reap, RemoteUnloaderAdapter, and the router. Registering
// it only inside this enabled block keeps the disabled path a true no-op
// (the registry stays hook-less).
registry.SetReplicaRemovedHook(func(model, node string, replica int) {
// for the prefix cache; other subsystems register their own hooks
// independently and are unaffected either way.
registry.AddReplicaRemovedHook(func(model, node string, replica int) {
if replica < 0 {
prefixSync.InvalidateNode(model, node)
} else {

View File

@@ -283,6 +283,14 @@ func New(opts ...config.AppOption) (*Application, error) {
distSvc.Registry,
)
application.modelLoader.SetModelStore(distStore)
// Drop the local stub when a model's last replica leaves the registry.
// The store reports local stubs UNION registry rows, and every removal
// path deletes the row only, so without this the frontend keeps
// reporting a model as loaded long after the replica is gone.
// Registered unconditionally: this is independent of the prefix cache.
distSvc.Registry.AddReplicaRemovedHook(
nodes.NewLocalStubInvalidator(distSvc.Registry, distStore),
)
// Start health monitor
distSvc.Health.Start(options.Context)
// Start replica reconciler for auto-scaling model replicas

View File

@@ -365,6 +365,36 @@ type ModelDeleteReply struct {
Error string `json:"error,omitempty"`
}
// SubjectNodeModelsRunning asks a worker node which model backend processes it
// currently has running. Uses NATS request-reply.
//
// This is the authoritative answer to "is this replica still alive". The worker
// owns the process table, so unlike a health probe against the backend's own
// serving port, its reply does not depend on whether that backend happens to be
// busy: a model mid-generation cannot answer a gRPC health check for minutes at
// a time, but the worker answers immediately either way.
func SubjectNodeModelsRunning(nodeID string) string {
return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".models.running"
}
// ModelsRunningRequest is the payload for a models.running NATS request.
type ModelsRunningRequest struct{}
// ModelsRunningReply is the response from a models.running NATS request.
type ModelsRunningReply struct {
Models []RunningModelInfo `json:"models"`
Error string `json:"error,omitempty"`
}
// RunningModelInfo identifies one live backend process on a worker. The triple
// is isomorphic to a controller NodeModel row's (model_name, replica_index,
// address), which is what lets the reconciler diff the two directly.
type RunningModelInfo struct {
ModelID string `json:"model_id"`
ReplicaIndex int `json:"replica_index"`
Address string `json:"address,omitempty"`
}
// SubjectNodeStop tells a serve-backend node to shut down entirely
// (deregister + exit). The node will not restart the backend process.
func SubjectNodeStop(nodeID string) string {

View File

@@ -0,0 +1,58 @@
package nodes
import (
"context"
"github.com/mudler/LocalAI/pkg/model"
"github.com/mudler/xlog"
)
// NewLocalStubInvalidator returns a replica-removed hook that drops the
// frontend's in-process stub for a model once no healthy replica of it remains
// anywhere in the cluster.
//
// Why this is needed: in distributed mode every routed model leaves a stub in
// the frontend's ModelLoader, and DistributedModelStore.Range reports local
// stubs UNION the registry rows. Every registry removal path deletes only the
// DB row, so without this hook the stub outlives the replica and the model is
// reported as loaded forever. That is the "loaded on the home page, absent from
// every node" ghost: /system reads the union and still sees the stub, while
// /api/nodes/models reads the registry and correctly sees nothing.
//
// The stub is dropped only when the LAST replica is gone. While another node
// still serves the model the frontend is right to consider it loaded, and each
// request re-routes through SmartRouter to pick a live replica anyway.
//
// This deletes the store entry directly rather than going through
// ShutdownModel: the replica is already gone from the registry, so there is
// nothing left to unload, and ShutdownModel would emit a pointless remote
// unload to a node that no longer hosts it.
func NewLocalStubInvalidator(registry *NodeRegistry, store model.ModelStore) func(modelID, nodeID string, replicaIndex int) {
return func(modelID, nodeID string, replicaIndex int) {
if registry == nil || store == nil || modelID == "" {
return
}
if _, ok := store.Get(modelID); !ok {
return // no local stub for this model, nothing to invalidate
}
// Registry rows are the source of truth for "is this model loaded
// anywhere". Only healthy nodes with a loaded row are returned.
remaining, err := registry.FindNodesWithModel(context.Background(), modelID)
if err != nil {
// Leave the stub alone rather than risk dropping a live model on a
// transient DB error. A later removal re-runs this check, and the
// reconciler keeps probing, so the ghost is not permanent.
xlog.Warn("Local stub invalidation skipped: failed to count remaining replicas",
"model", modelID, "node", nodeID, "replica", replicaIndex, "error", err)
return
}
if len(remaining) > 0 {
return
}
store.Delete(modelID)
xlog.Info("Dropped local model stub after its last replica was removed",
"model", modelID, "node", nodeID, "replica", replicaIndex)
}
}

View File

@@ -0,0 +1,109 @@
package nodes
import (
"context"
"runtime"
"time"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"gorm.io/gorm"
"github.com/mudler/LocalAI/pkg/model"
"github.com/mudler/LocalAI/core/services/testutil"
)
// In distributed mode the frontend keeps an in-process stub for every model it
// has routed, and DistributedModelStore.Range reports local stubs UNION the
// registry rows. Every registry removal path deletes the DB row only, so
// without an invalidator the stub outlives the replica and the model is
// reported as loaded forever (visible on the home page, absent from the nodes
// page). These tests pin the invalidator that closes that gap.
var _ = Describe("LocalStubInvalidator", func() {
var (
db *gorm.DB
registry *NodeRegistry
local *model.InMemoryModelStore
nodeA *BackendNode
nodeB *BackendNode
)
BeforeEach(func() {
if runtime.GOOS == "darwin" {
Skip("testcontainers requires Docker, not available on macOS CI")
}
db = testutil.SetupTestDB()
var err error
registry, err = NewNodeRegistry(db)
Expect(err).ToNot(HaveOccurred())
local = model.NewInMemoryModelStore()
nodeA = &BackendNode{Name: "node-a", NodeType: NodeTypeBackend, Address: "10.0.0.1:50051"}
nodeB = &BackendNode{Name: "node-b", NodeType: NodeTypeBackend, Address: "10.0.0.2:50051"}
Expect(registry.Register(context.Background(), nodeA, true)).To(Succeed())
Expect(registry.Register(context.Background(), nodeB, true)).To(Succeed())
})
It("drops the local stub once the last replica of the model is gone", func() {
store := NewDistributedModelStore(local, registry)
Expect(registry.SetNodeModel(context.Background(), nodeA.ID, "ghost-model", 0, "loaded", "10.0.0.1:12345", 0)).To(Succeed())
local.Set("ghost-model", model.NewModel("ghost-model", "10.0.0.1:12345", nil))
registry.AddReplicaRemovedHook(NewLocalStubInvalidator(registry, store))
Expect(registry.RemoveNodeModel(context.Background(), nodeA.ID, "ghost-model", 0)).To(Succeed())
_, ok := local.Get("ghost-model")
Expect(ok).To(BeFalse(), "the stub must not outlive the last replica of the model")
// And the model must disappear from the loaded listing that feeds /system.
listed := []string{}
store.Range(func(id string, _ *model.Model) bool {
listed = append(listed, id)
return true
})
Expect(listed).ToNot(ContainElement("ghost-model"))
})
It("keeps the local stub while another replica still serves the model", func() {
store := NewDistributedModelStore(local, registry)
Expect(registry.SetNodeModel(context.Background(), nodeA.ID, "shared-model", 0, "loaded", "10.0.0.1:12345", 0)).To(Succeed())
Expect(registry.SetNodeModel(context.Background(), nodeB.ID, "shared-model", 0, "loaded", "10.0.0.2:12345", 0)).To(Succeed())
local.Set("shared-model", model.NewModel("shared-model", "10.0.0.1:12345", nil))
registry.AddReplicaRemovedHook(NewLocalStubInvalidator(registry, store))
Expect(registry.RemoveNodeModel(context.Background(), nodeA.ID, "shared-model", 0)).To(Succeed())
_, ok := local.Get("shared-model")
Expect(ok).To(BeTrue(), "a model still loaded on another node must stay in the local store")
})
It("fires every registered replica-removed hook", func() {
Expect(registry.SetNodeModel(context.Background(), nodeA.ID, "m", 0, "loaded", "10.0.0.1:12345", 0)).To(Succeed())
var first, second int
registry.AddReplicaRemovedHook(func(_, _ string, _ int) { first++ })
registry.AddReplicaRemovedHook(func(_, _ string, _ int) { second++ })
Expect(registry.RemoveNodeModel(context.Background(), nodeA.ID, "m", 0)).To(Succeed())
// Prefix-cache invalidation and local-stub invalidation are independent
// subsystems; registering one must never silently displace the other.
Expect(first).To(Equal(1))
Expect(second).To(Equal(1))
})
It("drops the local stub when a whole node's replicas are removed", func() {
store := NewDistributedModelStore(local, registry)
Expect(registry.SetNodeModel(context.Background(), nodeA.ID, "node-model", 0, "loaded", "10.0.0.1:12345", 0)).To(Succeed())
local.Set("node-model", model.NewModel("node-model", "10.0.0.1:12345", nil))
registry.AddReplicaRemovedHook(NewLocalStubInvalidator(registry, store))
// Negative replica index signals "all replicas of this model on the node",
// the shape used by node deregistration and health reaping.
Expect(registry.RemoveAllNodeModelReplicas(context.Background(), nodeA.ID, "node-model")).To(Succeed())
Eventually(func() bool {
_, ok := local.Get("node-model")
return ok
}, time.Second, 10*time.Millisecond).Should(BeFalse())
})
})

View File

@@ -197,19 +197,22 @@ var _ = Describe("ModelRouterAdapter", func() {
adapter.mu.Unlock()
Expect(hasRelease).To(BeTrue())
// The initial in-flight reservation is released via OnFirstComplete after
// the first inference call, not during ReleaseModel. ReleaseModel only
// closes the client.
// The initial in-flight reservation is released by whichever comes
// first: the first inference completing, or the route being released.
// Nothing has happened yet, so it is still held.
fakeReg.mu.Lock()
countBeforeRelease := fakeReg.decrementCalled["node-1:test-model"]
fakeReg.mu.Unlock()
Expect(countBeforeRelease).To(Equal(0))
// Releasing the route without any inference must give it back, or the
// counter leaks and pins the replica against every eviction query.
adapter.ReleaseModel("test-model")
fakeReg.mu.Lock()
countAfterRelease := fakeReg.decrementCalled["node-1:test-model"]
fakeReg.mu.Unlock()
Expect(countAfterRelease).To(Equal(0))
Eventually(func() int {
fakeReg.mu.Lock()
defer fakeReg.mu.Unlock()
return fakeReg.decrementCalled["node-1:test-model"]
}).Should(Equal(1))
})
})
})

View File

@@ -5,33 +5,106 @@ import (
"encoding/json"
"errors"
"fmt"
"sync"
"time"
"github.com/mudler/LocalAI/core/services/advisorylock"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/nodes/prefixcache"
grpcclient "github.com/mudler/LocalAI/pkg/grpc"
"github.com/mudler/xlog"
"github.com/nats-io/nats.go"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"gorm.io/gorm"
)
// ModelProber checks whether a model's backend process is still reachable.
// ProbeOutcome is the result of a model liveness probe. The distinction that
// matters is between silence caused by a BUSY backend and silence caused by a
// MISSING one: only the latter justifies deleting the registry row.
type ProbeOutcome int
const (
// ProbeAlive: the backend answered and reported healthy.
ProbeAlive ProbeOutcome = iota
// ProbeBusy: the backend was reachable but did not answer within the
// probe deadline. Single-threaded backends block here for the whole
// duration of a long generation, so this is evidence of life, not death.
ProbeBusy
// ProbeUnreachable: nothing is listening (connection refused), or the
// backend answered and affirmatively reported itself unhealthy.
ProbeUnreachable
)
// ModelProber checks the state of a model's backend process.
// Defaulted to a gRPC health probe but overridable for tests so we don't
// need to stand up a real server. Returning false without an error means the
// process is reachable but unhealthy (same as a timeout for our purposes).
// need to stand up a real server.
type ModelProber interface {
IsAlive(ctx context.Context, address string) bool
Probe(ctx context.Context, address string) ProbeOutcome
}
// grpcModelProber does a 1s HealthCheck on the model's stored gRPC address.
// NodeProcessLister asks a worker which model backend processes it currently
// has running. Implemented by RemoteUnloaderAdapter over NATS.
//
// This is the sounder liveness signal: the worker owns the process table, so
// its answer does not depend on whether a backend is busy. A health probe
// against the backend's own serving port cannot make that distinction, which
// is why it is only the fallback for workers that do not answer.
type NodeProcessLister interface {
ListRunningModels(nodeID string) (*messaging.ModelsRunningReply, error)
}
// probeTimeout bounds a single liveness probe. Kept short because a healthy
// idle backend answers immediately; a backend that needs longer is by
// definition busy, which classifyProbeOutcome reports as ProbeBusy rather than
// as death.
const probeTimeout = 1 * time.Second
// grpcModelProber does a short HealthCheck on the model's stored gRPC address.
type grpcModelProber struct{ token string }
func (g grpcModelProber) IsAlive(ctx context.Context, address string) bool {
func (g grpcModelProber) Probe(ctx context.Context, address string) ProbeOutcome {
client := grpcclient.NewClientWithToken(address, false, nil, false, g.token)
probeCtx, cancel := context.WithTimeout(ctx, 1*time.Second)
probeCtx, cancel := context.WithTimeout(ctx, probeTimeout)
defer cancel()
ok, _ := client.HealthCheck(probeCtx)
return ok
ok, err := client.HealthCheck(probeCtx)
return classifyProbeOutcome(ok, err)
}
// classifyProbeOutcome maps a HealthCheck result onto a ProbeOutcome.
//
// The gRPC client is lazy, so connection failures surface on the RPC rather
// than at dial time, and the status code tells the two cases apart:
//
// - DeadlineExceeded: the transport was fine but nothing serviced the RPC in
// time. That is a backend stuck inside a long synchronous request.
// - Unavailable: nothing is listening. The process is gone.
//
// A blackholed network also yields DeadlineExceeded and is therefore treated as
// busy. That is deliberate: whole-node failures are the health monitor's job
// (it marks the node offline and reaps its models), and mistaking a network
// blip for a dead backend is the failure this classification exists to prevent.
func classifyProbeOutcome(ok bool, err error) ProbeOutcome {
if err == nil {
if ok {
return ProbeAlive
}
// Answered, but reported unhealthy. An affirmative signal, unlike silence.
return ProbeUnreachable
}
if errors.Is(err, context.DeadlineExceeded) {
return ProbeBusy
}
switch status.Code(err) {
case codes.DeadlineExceeded:
return ProbeBusy
case codes.Unavailable:
return ProbeUnreachable
default:
// Any other status means something on the other end produced a
// response, so the process exists. Not healthy, but not a ghost.
return ProbeBusy
}
}
// ReplicaReconciler periodically ensures model replica counts match their
@@ -64,6 +137,24 @@ type ReplicaReconciler struct {
// the other scale-up paths. nil disables this signal (a true no-op).
pressure *prefixcache.Pressure
pressureThreshold int
// processLister asks workers which model processes they are running. nil
// disables the worker-authoritative pass, leaving only the port probe.
processLister NodeProcessLister
// probeFailures counts CONSECUTIVE failed liveness probes per node_models
// row id, so a single miss against a busy-but-alive backend cannot delete
// its registry row. Reset on any successful probe and on reap.
probeFailuresMu sync.Mutex
probeFailures map[string]int
// workerMisses counts CONSECUTIVE passes in which a worker did not report
// a row's model as running. Kept separate from probeFailures so the two
// independent signals never add into one another's budget.
workerMissesMu sync.Mutex
workerMisses map[string]int
// inFlightIdle counts CONSECUTIVE passes in which a row with a positive
// in_flight counter was observed answering health probes promptly, which is
// what a backend inside a request cannot do.
inFlightIdleMu sync.Mutex
inFlightIdle map[string]int
}
// ModelScheduler abstracts the scheduling logic needed by the reconciler.
@@ -86,7 +177,11 @@ type ReplicaReconcilerOptions struct {
// addresses. Matches the worker's token so HealthCheck auth succeeds.
RegistrationToken string
// Prober overrides the default gRPC health probe (used by tests).
Prober ModelProber
Prober ModelProber
// ProcessLister overrides the default worker process query. When nil and
// no Adapter is set, the worker-authoritative pass is skipped entirely and
// only the port probe runs.
ProcessLister NodeProcessLister
DB *gorm.DB
Interval time.Duration // default 30s
ScaleDownDelay time.Duration // default 5m
@@ -121,12 +216,19 @@ func NewReplicaReconciler(opts ReplicaReconcilerOptions) *ReplicaReconciler {
if pressureThreshold == 0 {
pressureThreshold = prefixcache.DefaultConfig().PressureScaleThreshold
}
// The adapter already speaks the models.running request, so it is the
// natural default; an explicit ProcessLister (tests) wins over it.
processLister := opts.ProcessLister
if processLister == nil && opts.Adapter != nil {
processLister = opts.Adapter
}
return &ReplicaReconciler{
registry: opts.Registry,
scheduler: opts.Scheduler,
unloader: opts.Unloader,
adapter: opts.Adapter,
prober: prober,
processLister: processLister,
db: opts.DB,
interval: interval,
scaleDownDelay: scaleDownDelay,
@@ -175,14 +277,21 @@ func (rc *ReplicaReconciler) reconcileOnce(ctx context.Context) {
}
// reconcileState runs the state-reconciliation passes: drain pending backend
// ops for freshly-healthy nodes, then probe model gRPC addresses to orphan
// ghosts. Both passes are best-effort: a failure on one node doesn't stop
// the rest.
// ops for freshly-healthy nodes, reconcile registry rows against what workers
// report they are running, then port-probe whatever is left. All passes are
// best-effort: a failure on one node doesn't stop the rest.
//
// Order matters. The worker pass runs first and refreshes updated_at for every
// model a worker vouches for, which takes those rows out of the port prober's
// stale set. The port probe is therefore only the fallback for nodes whose
// worker did not answer.
func (rc *ReplicaReconciler) reconcileState(ctx context.Context) {
if rc.adapter != nil {
rc.drainPendingBackendOps(ctx)
}
rc.reconcileNodeProcesses(ctx)
rc.probeLoadedModels(ctx)
rc.sweepLeakedInFlight(ctx)
}
// drainPendingBackendOps retries queued backend ops whose next_retry_at has
@@ -310,11 +419,34 @@ const maxPendingBackendOpAttempts = 10
// without this sweep they would leak forever and keep the UI op spinning.
const stalePendingBackendOpGrace = 15 * time.Minute
// probeFailuresBeforeReap is how many CONSECUTIVE failed liveness probes a
// replica must accumulate before its row is deleted. The probe is a 1s gRPC
// HealthCheck, and a backend that is merely busy cannot answer it: single
// threaded Python backends (video/avatar generation, long diffusion loops)
// block for minutes at a time inside one request. Treating the first miss as
// death deleted rows for backends that were alive and working, which made the
// model vanish from the nodes page while it was still generating. Three misses
// across three ticks is ~1.5 minutes of silence at the default interval, well
// past any GC pause or scheduling hiccup, while a genuinely dead process is
// still reaped promptly.
const probeFailuresBeforeReap = 3
// probeLoadedModels gRPC-health-checks model addresses that the DB says are
// loaded. If a model's backend process is gone (OOM, crash, manual restart)
// we remove the row so ghosts don't linger. Only probes rows older than
// probeStaleAfter so we don't hammer every worker every tick for models we
// just heard from.
//
// Two guards keep this from reaping healthy backends. A probe that times out is
// classified as ProbeBusy rather than as death, because a backend blocked
// inside a long generation cannot answer while it works. And a replica must
// then look unreachable on probeFailuresBeforeReap CONSECUTIVE passes, so a
// transient blip cannot orphan a live replica.
//
// Deliberately NOT gated on in_flight. That counter has no decrement guarantee
// (a frontend that dies mid-request leaves the increment behind), so using it
// to shield rows from reaping would make a leaked counter produce a replica
// that can never be cleaned up.
func (rc *ReplicaReconciler) probeLoadedModels(ctx context.Context) {
var stale []NodeModel
cutoff := time.Now().Add(-rc.probeStaleAfter)
@@ -327,22 +459,322 @@ func (rc *ReplicaReconciler) probeLoadedModels(ctx context.Context) {
xlog.Warn("Reconciler: failed to list loaded models for probe", "error", err)
return
}
seen := make(map[string]struct{}, len(stale))
for _, m := range stale {
if err := ctx.Err(); err != nil {
return
}
if rc.prober.IsAlive(ctx, m.Address) {
seen[m.ID] = struct{}{}
switch rc.prober.Probe(ctx, m.Address) {
case ProbeAlive:
rc.clearProbeFailures(m.ID)
// Bump updated_at so we don't probe this row again immediately.
_ = rc.registry.db.WithContext(ctx).Model(&NodeModel{}).
Where("id = ?", m.ID).Update("updated_at", time.Now()).Error
continue
case ProbeBusy:
// Reachable but mid-request. Proof of life, so clear the streak.
rc.clearProbeFailures(m.ID)
xlog.Debug("Reconciler: model busy, skipping liveness reap",
"node", m.NodeID, "model", m.ModelName, "replica", m.ReplicaIndex, "address", m.Address)
continue
}
failures := rc.recordProbeFailure(m.ID)
if failures < probeFailuresBeforeReap {
xlog.Debug("Reconciler: model unreachable, waiting for more misses before reaping",
"node", m.NodeID, "model", m.ModelName, "replica", m.ReplicaIndex, "address", m.Address,
"failures", failures, "threshold", probeFailuresBeforeReap)
continue
}
if err := rc.registry.RemoveNodeModel(ctx, m.NodeID, m.ModelName, m.ReplicaIndex); err != nil {
xlog.Warn("Reconciler: failed to remove unreachable model", "node", m.NodeID, "model", m.ModelName, "replica", m.ReplicaIndex, "error", err)
continue
}
rc.clearProbeFailures(m.ID)
xlog.Warn("Reconciler: model unreachable, removed from registry",
"node", m.NodeID, "model", m.ModelName, "replica", m.ReplicaIndex, "address", m.Address)
"node", m.NodeID, "model", m.ModelName, "replica", m.ReplicaIndex, "address", m.Address,
"failures", failures)
}
rc.pruneProbeFailures(seen)
}
// inFlightLeakIdleAfter is how long a positive in_flight counter must sit
// without its last_used moving before the sweeper will even consider it a leak.
// Generous on purpose: it only bounds how quickly a genuine leak is corrected,
// whereas being too eager risks overruling a real request.
const inFlightLeakIdleAfter = 30 * time.Minute
// inFlightLeakConfirmations is how many CONSECUTIVE passes must observe the
// backend answering promptly before the counter is overruled.
const inFlightLeakConfirmations = 2
// sweepLeakedInFlight corrects in_flight counters that no longer reflect any
// real request.
//
// The counter has no decrement guarantee. `track()` balances its increment with
// a defer, but that only holds while the process lives: a frontend killed
// mid-request (rolling restart, OOM) never runs the defer, and the load-time
// reservation is released only when the FIRST inference completes, so a request
// that dies before reaching the backend strands it. The result is not cosmetic —
// FindLRUModel, FindGlobalLRUModelWithZeroInFlight and the router's eviction
// query all require in_flight = 0, so a leaked counter pins that replica's VRAM
// forever.
//
// Elapsed time alone cannot identify a leak: IncrementInFlight stamps last_used
// at request START and nothing moves it while the request runs, so a long
// generation is indistinguishable from a leak by age. The probe supplies the
// missing bit — a backend that answers a health check promptly is not inside a
// request, because that is exactly what a busy backend cannot do. Requiring the
// row to ALSO be idle for inFlightLeakIdleAfter covers backends that serve
// requests in parallel and can answer while working: those keep last_used fresh
// through each new request's increment.
func (rc *ReplicaReconciler) sweepLeakedInFlight(ctx context.Context) {
var suspects []NodeModel
cutoff := time.Now().Add(-inFlightLeakIdleAfter)
err := rc.registry.db.WithContext(ctx).
Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id").
Where("node_models.state = ? AND backend_nodes.status = ? AND node_models.in_flight > 0 AND node_models.last_used < ? AND node_models.address != ''",
"loaded", StatusHealthy, cutoff).
Find(&suspects).Error
if err != nil {
xlog.Warn("Reconciler: failed to list rows for in-flight leak sweep", "error", err)
return
}
seen := make(map[string]struct{}, len(suspects))
for _, m := range suspects {
if err := ctx.Err(); err != nil {
return
}
seen[m.ID] = struct{}{}
if rc.prober.Probe(ctx, m.Address) != ProbeAlive {
// Busy or unreachable. Busy means the counter may well be real;
// unreachable is the reaper's business, not the sweeper's.
rc.clearInFlightIdle(m.ID)
continue
}
confirmations := rc.recordInFlightIdle(m.ID)
if confirmations < inFlightLeakConfirmations {
continue
}
res := rc.registry.db.WithContext(ctx).Model(&NodeModel{}).
Where("id = ? AND in_flight > 0", m.ID).
UpdateColumn("in_flight", 0)
if res.Error != nil {
xlog.Warn("Reconciler: failed to reset leaked in-flight counter",
"node", m.NodeID, "model", m.ModelName, "replica", m.ReplicaIndex, "error", res.Error)
continue
}
rc.clearInFlightIdle(m.ID)
if res.RowsAffected > 0 {
xlog.Warn("Reconciler: reset leaked in-flight counter on an idle replica",
"node", m.NodeID, "model", m.ModelName, "replica", m.ReplicaIndex,
"was", m.InFlight, "idleFor", time.Since(m.LastUsed).String())
}
}
rc.pruneInFlightIdle(seen)
}
func (rc *ReplicaReconciler) recordInFlightIdle(rowID string) int {
rc.inFlightIdleMu.Lock()
defer rc.inFlightIdleMu.Unlock()
if rc.inFlightIdle == nil {
rc.inFlightIdle = make(map[string]int)
}
rc.inFlightIdle[rowID]++
return rc.inFlightIdle[rowID]
}
func (rc *ReplicaReconciler) clearInFlightIdle(rowID string) {
rc.inFlightIdleMu.Lock()
defer rc.inFlightIdleMu.Unlock()
delete(rc.inFlightIdle, rowID)
}
// pruneInFlightIdle bounds the streak map. Dropping a streak only ever resets
// it, so this can never cause a premature correction.
func (rc *ReplicaReconciler) pruneInFlightIdle(seen map[string]struct{}) {
rc.inFlightIdleMu.Lock()
defer rc.inFlightIdleMu.Unlock()
for id := range rc.inFlightIdle {
if _, ok := seen[id]; !ok {
delete(rc.inFlightIdle, id)
}
}
}
// workerMissesBeforeReap is how many CONSECUTIVE passes a worker must fail to
// report a model before its row is deleted. The worker's answer is
// authoritative, so this only needs to absorb the narrow window where a row
// exists but the process table has not caught up; two passes is ample given
// rows are only considered once they are probeStaleAfter old.
const workerMissesBeforeReap = 2
// reconcileNodeProcesses reconciles registry rows against what each worker says
// it is actually running.
//
// This is the primary liveness pass, and it is sounder than probing the
// backend's own serving port: the worker owns the process table and answers
// immediately whether or not the backend is mid-request. A model confirmed here
// also gets its updated_at bumped, which keeps it out of the port prober's
// stale set entirely — that is what stops a backend deep in a long generation
// from ever being mistaken for a dead one.
//
// A worker that cannot be reached is skipped rather than treated as empty. A
// messaging failure says nothing about the processes, and assuming the worst
// would delete a whole node's rows on a transient NATS blip; the port probe
// remains as the fallback for those nodes.
func (rc *ReplicaReconciler) reconcileNodeProcesses(ctx context.Context) {
if rc.processLister == nil {
return
}
var stale []NodeModel
cutoff := time.Now().Add(-rc.probeStaleAfter)
err := rc.registry.db.WithContext(ctx).
Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id").
Where("node_models.state = ? AND backend_nodes.status = ? AND backend_nodes.node_type = ? AND node_models.updated_at < ?",
"loaded", StatusHealthy, NodeTypeBackend, cutoff).
Find(&stale).Error
if err != nil {
xlog.Warn("Reconciler: failed to list loaded models for worker reconciliation", "error", err)
return
}
if len(stale) == 0 {
return
}
byNode := make(map[string][]NodeModel)
for _, m := range stale {
byNode[m.NodeID] = append(byNode[m.NodeID], m)
}
seen := make(map[string]struct{}, len(stale))
for nodeID, rows := range byNode {
if err := ctx.Err(); err != nil {
return
}
reply, err := rc.processLister.ListRunningModels(nodeID)
if err != nil {
xlog.Debug("Reconciler: worker did not answer models.running, leaving its replicas to the port probe",
"node", nodeID, "error", err)
continue
}
if reply == nil || reply.Error != "" {
xlog.Debug("Reconciler: worker reported an error for models.running, skipping node",
"node", nodeID, "error", replyError(reply))
continue
}
running := make(map[replicaKey]struct{}, len(reply.Models))
for _, m := range reply.Models {
running[replicaKey{model: m.ModelID, replica: m.ReplicaIndex}] = struct{}{}
}
for _, row := range rows {
seen[row.ID] = struct{}{}
if _, ok := running[replicaKey{model: row.ModelName, replica: row.ReplicaIndex}]; ok {
rc.clearWorkerMisses(row.ID)
// Confirmed alive by the process owner. Bump updated_at so the
// port prober does not also go poking at a busy backend.
_ = rc.registry.db.WithContext(ctx).Model(&NodeModel{}).
Where("id = ?", row.ID).Update("updated_at", time.Now()).Error
continue
}
misses := rc.recordWorkerMiss(row.ID)
if misses < workerMissesBeforeReap {
xlog.Debug("Reconciler: worker does not report model as running, waiting for confirmation",
"node", nodeID, "model", row.ModelName, "replica", row.ReplicaIndex,
"misses", misses, "threshold", workerMissesBeforeReap)
continue
}
if err := rc.registry.RemoveNodeModel(ctx, row.NodeID, row.ModelName, row.ReplicaIndex); err != nil {
xlog.Warn("Reconciler: failed to remove model the worker is not running",
"node", nodeID, "model", row.ModelName, "replica", row.ReplicaIndex, "error", err)
continue
}
rc.clearWorkerMisses(row.ID)
xlog.Warn("Reconciler: worker is not running this model, removed from registry",
"node", nodeID, "model", row.ModelName, "replica", row.ReplicaIndex, "misses", misses)
}
}
rc.pruneWorkerMisses(seen)
}
// replicaKey identifies one replica of one model on a node.
type replicaKey struct {
model string
replica int
}
// replyError safely extracts the error text from a possibly-nil reply.
func replyError(reply *messaging.ModelsRunningReply) string {
if reply == nil {
return "nil reply"
}
return reply.Error
}
func (rc *ReplicaReconciler) recordWorkerMiss(rowID string) int {
rc.workerMissesMu.Lock()
defer rc.workerMissesMu.Unlock()
if rc.workerMisses == nil {
rc.workerMisses = make(map[string]int)
}
rc.workerMisses[rowID]++
return rc.workerMisses[rowID]
}
func (rc *ReplicaReconciler) clearWorkerMisses(rowID string) {
rc.workerMissesMu.Lock()
defer rc.workerMissesMu.Unlock()
delete(rc.workerMisses, rowID)
}
// pruneWorkerMisses drops streaks for rows that were not candidates this pass,
// bounding the map. Dropping a streak only ever resets it, so this can never
// cause a premature reap.
func (rc *ReplicaReconciler) pruneWorkerMisses(seen map[string]struct{}) {
rc.workerMissesMu.Lock()
defer rc.workerMissesMu.Unlock()
for id := range rc.workerMisses {
if _, ok := seen[id]; !ok {
delete(rc.workerMisses, id)
}
}
}
// recordProbeFailure increments and returns the consecutive-failure count for a
// replica row.
func (rc *ReplicaReconciler) recordProbeFailure(rowID string) int {
rc.probeFailuresMu.Lock()
defer rc.probeFailuresMu.Unlock()
if rc.probeFailures == nil {
rc.probeFailures = make(map[string]int)
}
rc.probeFailures[rowID]++
return rc.probeFailures[rowID]
}
// clearProbeFailures resets a replica's streak after it answers or is reaped.
func (rc *ReplicaReconciler) clearProbeFailures(rowID string) {
rc.probeFailuresMu.Lock()
defer rc.probeFailuresMu.Unlock()
delete(rc.probeFailures, rowID)
}
// pruneProbeFailures drops streaks for rows that were not probe candidates this
// pass, so the map cannot grow without bound as replicas come and go. Dropping
// a streak only ever resets it, so this can never cause a premature reap.
func (rc *ReplicaReconciler) pruneProbeFailures(seen map[string]struct{}) {
rc.probeFailuresMu.Lock()
defer rc.probeFailuresMu.Unlock()
for id := range rc.probeFailures {
if _, ok := seen[id]; !ok {
delete(rc.probeFailures, id)
}
}
}

View File

@@ -0,0 +1,193 @@
package nodes
import (
"context"
"runtime"
"time"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"gorm.io/gorm"
"github.com/mudler/LocalAI/core/services/testutil"
)
// The liveness probe is a short gRPC HealthCheck, and a backend that blocks
// inside a long synchronous generation (video/avatar models spend 10+ minutes
// in a denoising loop) cannot answer it. Treating that silence as death
// deleted registry rows for backends that were alive and working.
//
// The discriminator is HOW the probe failed, not whether the replica looked
// busy in the database: a dead process refuses the connection, while a busy one
// accepts it and never services the RPC. Crucially this does NOT consult
// in_flight — that counter can leak high and would then shield a genuinely dead
// replica from ever being reaped.
var _ = Describe("ReplicaReconciler — probe reaper vs busy backends", func() {
var (
db *gorm.DB
registry *NodeRegistry
node *BackendNode
)
const addr = "10.0.0.1:12345"
BeforeEach(func() {
if runtime.GOOS == "darwin" {
Skip("testcontainers requires Docker, not available on macOS CI")
}
db = testutil.SetupTestDB()
var err error
registry, err = NewNodeRegistry(db)
Expect(err).ToNot(HaveOccurred())
node = &BackendNode{Name: "busy-node", NodeType: NodeTypeBackend, Address: "10.0.0.1:50051"}
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
})
// seed inserts a stale loaded row so the probe pass picks it up.
seed := func(id string, inFlight int) {
Expect(db.Create(&NodeModel{
ID: id,
NodeID: node.ID,
ModelName: id,
Address: addr,
State: "loaded",
InFlight: inFlight,
UpdatedAt: time.Now().Add(-5 * time.Minute),
}).Error).To(Succeed())
}
// makeStale rewinds updated_at so the row qualifies for the next pass.
makeStale := func(id string) {
Expect(db.Model(&NodeModel{}).Where("id = ?", id).
Update("updated_at", time.Now().Add(-5*time.Minute)).Error).To(Succeed())
}
newReconciler := func(prober ModelProber) *ReplicaReconciler {
return NewReplicaReconciler(ReplicaReconcilerOptions{
Registry: registry,
DB: db,
Prober: prober,
ProbeStaleAfter: 2 * time.Minute,
})
}
It("never reaps a replica that is reachable but too busy to answer", func() {
seed("busy-1", 0)
prober := &fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeBusy}}
rc := newReconciler(prober)
// However long the generation runs, a busy backend is not a dead one.
for range 10 {
rc.probeLoadedModels(context.Background())
makeStale("busy-1")
}
var after NodeModel
Expect(db.First(&after, "id = ?", "busy-1").Error).To(Succeed(),
"a backend that accepted the connection but was mid-request must never be reaped")
})
It("reaps an unreachable replica even when in_flight leaked high", func() {
// in_flight has no decrement guarantee: a frontend that dies mid-request
// leaves the increment behind forever. Gating the reaper on it would
// make such a row permanently unreapable, so the reaper must ignore it.
seed("leaked-1", 7)
prober := &fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeUnreachable}}
rc := newReconciler(prober)
for range probeFailuresBeforeReap {
rc.probeLoadedModels(context.Background())
makeStale("leaked-1")
}
var after NodeModel
Expect(db.First(&after, "id = ?", "leaked-1").Error).To(MatchError(gorm.ErrRecordNotFound),
"a stale in_flight counter must not shield a dead replica from reaping")
})
It("requires consecutive unreachable probes before reaping", func() {
seed("idle-1", 0)
prober := &fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeUnreachable}}
rc := newReconciler(prober)
for i := 1; i < probeFailuresBeforeReap; i++ {
rc.probeLoadedModels(context.Background())
makeStale("idle-1")
var after NodeModel
Expect(db.First(&after, "id = ?", "idle-1").Error).To(Succeed(),
"a replica must survive until the failure threshold is reached")
}
rc.probeLoadedModels(context.Background())
var after NodeModel
Expect(db.First(&after, "id = ?", "idle-1").Error).To(MatchError(gorm.ErrRecordNotFound))
})
It("resets the failure streak when a replica answers again", func() {
seed("flaky-1", 0)
prober := &fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeUnreachable}}
rc := newReconciler(prober)
rc.probeLoadedModels(context.Background())
makeStale("flaky-1")
rc.probeLoadedModels(context.Background())
makeStale("flaky-1")
prober.outcomes[addr] = ProbeAlive
rc.probeLoadedModels(context.Background())
makeStale("flaky-1")
// Streak restarts from zero, so the threshold is not reached again yet.
prober.outcomes[addr] = ProbeUnreachable
for i := 1; i < probeFailuresBeforeReap; i++ {
rc.probeLoadedModels(context.Background())
makeStale("flaky-1")
}
var after NodeModel
Expect(db.First(&after, "id = ?", "flaky-1").Error).To(Succeed(),
"a successful probe must reset the consecutive-failure streak")
})
It("does not let a busy probe count toward the failure streak", func() {
seed("mixed-1", 0)
prober := &fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeUnreachable}}
rc := newReconciler(prober)
// Two genuine misses...
rc.probeLoadedModels(context.Background())
makeStale("mixed-1")
rc.probeLoadedModels(context.Background())
makeStale("mixed-1")
// ...then the backend starts a long generation. That is evidence of
// life, so it must not be the miss that tips the row over the edge.
prober.outcomes[addr] = ProbeBusy
rc.probeLoadedModels(context.Background())
makeStale("mixed-1")
var after NodeModel
Expect(db.First(&after, "id = ?", "mixed-1").Error).To(Succeed(),
"a busy probe is proof of life and must clear the failure streak")
})
})
// errConnRefused builds the error gRPC surfaces when nothing is listening on
// the address, which is what a dead backend process looks like.
func errConnRefused() error {
return status.Error(codes.Unavailable, "connection error: desc = \"transport: Error while dialing: dial tcp 10.0.0.1:12345: connect: connection refused\"")
}
var _ = Describe("classifyProbeOutcome", func() {
It("treats a deadline as busy and a refused connection as unreachable", func() {
// These are the two cases that matter: they are what the reaper keys
// its delete decision off.
Expect(classifyProbeOutcome(true, nil)).To(Equal(ProbeAlive))
Expect(classifyProbeOutcome(false, context.DeadlineExceeded)).To(Equal(ProbeBusy))
Expect(classifyProbeOutcome(false, errConnRefused())).To(Equal(ProbeUnreachable))
// Answered, but reported unhealthy: an affirmative signal, unlike silence.
Expect(classifyProbeOutcome(false, nil)).To(Equal(ProbeUnreachable))
})
})

View File

@@ -0,0 +1,163 @@
package nodes
import (
"context"
"runtime"
"time"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"gorm.io/gorm"
"github.com/mudler/LocalAI/core/services/testutil"
)
// in_flight has no decrement guarantee: a frontend that dies mid-request leaves
// the increment behind, and the load-time reservation is only released when the
// first inference completes. A leaked counter is not cosmetic — it pins the
// replica against LRU eviction (FindLRUModel, FindGlobalLRUModelWithZeroInFlight
// and the router's eviction query all require in_flight = 0), so the VRAM is
// never reclaimed.
//
// Resetting on elapsed time alone is unsafe: IncrementInFlight stamps last_used
// at request START and nothing moves it while the request runs, so a long
// generation looks identical to a leak. The discriminator is the health probe —
// a backend that answers immediately is not inside a request.
var _ = Describe("ReplicaReconciler — leaked in_flight sweeper", func() {
var (
db *gorm.DB
registry *NodeRegistry
node *BackendNode
)
const addr = "10.0.0.1:12345"
BeforeEach(func() {
if runtime.GOOS == "darwin" {
Skip("testcontainers requires Docker, not available on macOS CI")
}
db = testutil.SetupTestDB()
var err error
registry, err = NewNodeRegistry(db)
Expect(err).ToNot(HaveOccurred())
node = &BackendNode{Name: "n1", NodeType: NodeTypeBackend, Address: "10.0.0.1:50051"}
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
})
seed := func(id string, inFlight int, idleFor time.Duration) {
Expect(db.Create(&NodeModel{
ID: id,
NodeID: node.ID,
ModelName: id,
Address: addr,
State: "loaded",
InFlight: inFlight,
LastUsed: time.Now().Add(-idleFor),
UpdatedAt: time.Now(),
}).Error).To(Succeed())
}
inFlightOf := func(id string) int {
var m NodeModel
Expect(db.First(&m, "id = ?", id).Error).To(Succeed())
return m.InFlight
}
newReconciler := func(prober ModelProber) *ReplicaReconciler {
return NewReplicaReconciler(ReplicaReconcilerOptions{
Registry: registry,
DB: db,
Prober: prober,
})
}
It("resets a counter whose backend has been answering promptly while idle", func() {
seed("leaked", 3, 2*inFlightLeakIdleAfter)
rc := newReconciler(&fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeAlive}})
for range inFlightLeakConfirmations {
rc.sweepLeakedInFlight(context.Background())
}
Expect(inFlightOf("leaked")).To(Equal(0),
"an idle backend answering health checks cannot be serving 3 requests")
})
It("never resets a counter while the backend is mid-request", func() {
// The exact case a time-based sweeper gets wrong: a video generation
// that has been running for far longer than the idle threshold.
seed("generating", 1, 2*inFlightLeakIdleAfter)
rc := newReconciler(&fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeBusy}})
for range inFlightLeakConfirmations + 3 {
rc.sweepLeakedInFlight(context.Background())
}
Expect(inFlightOf("generating")).To(Equal(1),
"resetting here would expose a serving model to LRU eviction")
})
It("never resets a counter that was used recently", func() {
// A busy parallel backend keeps last_used fresh through each new
// request's increment, so recency alone protects it.
seed("recent", 2, time.Minute)
rc := newReconciler(&fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeAlive}})
for range inFlightLeakConfirmations + 3 {
rc.sweepLeakedInFlight(context.Background())
}
Expect(inFlightOf("recent")).To(Equal(2))
})
It("requires consecutive idle confirmations before resetting", func() {
seed("flappy", 1, 2*inFlightLeakIdleAfter)
prober := &fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeAlive}}
rc := newReconciler(prober)
for i := 1; i < inFlightLeakConfirmations; i++ {
rc.sweepLeakedInFlight(context.Background())
Expect(inFlightOf("flappy")).To(Equal(1),
"one idle observation is not enough to overrule the counter")
}
// Backend turns out to be working after all: the streak resets.
prober.outcomes[addr] = ProbeBusy
rc.sweepLeakedInFlight(context.Background())
prober.outcomes[addr] = ProbeAlive
for i := 1; i < inFlightLeakConfirmations; i++ {
rc.sweepLeakedInFlight(context.Background())
}
Expect(inFlightOf("flappy")).To(Equal(1))
})
It("leaves rows with a zero counter alone", func() {
seed("idle", 0, 2*inFlightLeakIdleAfter)
prober := &fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeAlive}}
rc := newReconciler(prober)
rc.sweepLeakedInFlight(context.Background())
Expect(inFlightOf("idle")).To(Equal(0))
Expect(prober.calls).To(Equal(0), "nothing to correct means nothing to probe")
})
It("frees the replica for LRU eviction once the leak is cleared", func() {
// The point of the sweep: FindLRUModel only considers in_flight = 0,
// so until the counter is corrected this replica's VRAM is unreclaimable.
seed("pinned", 1, 2*inFlightLeakIdleAfter)
rc := newReconciler(&fakeProber{outcomes: map[string]ProbeOutcome{addr: ProbeAlive}})
_, err := registry.FindLRUModel(context.Background(), node.ID)
Expect(err).To(HaveOccurred(), "precondition: the leak hides the row from LRU")
for range inFlightLeakConfirmations {
rc.sweepLeakedInFlight(context.Background())
}
lru, err := registry.FindLRUModel(context.Background(), node.ID)
Expect(err).ToNot(HaveOccurred())
Expect(lru.ModelName).To(Equal("pinned"))
})
})

View File

@@ -733,18 +733,19 @@ var _ = Describe("ReplicaReconciler", func() {
})
})
// fakeProber lets tests control whether a model's gRPC address "responds".
// fakeProber lets tests control how a model's gRPC address "responds".
// Addresses with no entry default to ProbeUnreachable.
type fakeProber struct {
alive map[string]bool
calls int
outcomes map[string]ProbeOutcome
calls int
}
func (f *fakeProber) IsAlive(_ context.Context, address string) bool {
func (f *fakeProber) Probe(_ context.Context, address string) ProbeOutcome {
f.calls++
if f.alive == nil {
return false
if f.outcomes == nil {
return ProbeUnreachable
}
return f.alive[address]
return f.outcomes[address]
}
var _ = Describe("ReplicaReconciler — state reconciliation", func() {
@@ -787,7 +788,7 @@ var _ = Describe("ReplicaReconciler — state reconciliation", func() {
Expect(db.Create(stale).Error).To(Succeed())
Expect(db.Create(fresh).Error).To(Succeed())
prober := &fakeProber{alive: map[string]bool{"10.0.0.1:12345": false}}
prober := &fakeProber{outcomes: map[string]ProbeOutcome{"10.0.0.1:12345": ProbeUnreachable}}
rc := NewReplicaReconciler(ReplicaReconcilerOptions{
Registry: registry,
DB: db,
@@ -795,15 +796,19 @@ var _ = Describe("ReplicaReconciler — state reconciliation", func() {
ProbeStaleAfter: 2 * time.Minute,
})
rc.probeLoadedModels(context.Background())
// A replica is reaped only after probeFailuresBeforeReap consecutive
// misses, so a busy-but-alive backend is not mistaken for a dead one.
for range probeFailuresBeforeReap {
rc.probeLoadedModels(context.Background())
}
// Stale was unreachable — row removed.
var after []NodeModel
Expect(db.Find(&after).Error).To(Succeed())
Expect(after).To(HaveLen(1))
Expect(after[0].ModelName).To(Equal("fresh-model"))
// Prober was only called once (the fresh row was filtered out).
Expect(prober.calls).To(Equal(1))
// Only the stale row was ever probed (the fresh row was filtered out).
Expect(prober.calls).To(Equal(probeFailuresBeforeReap))
})
It("keeps reachable models and bumps their updated_at", func() {
@@ -819,7 +824,7 @@ var _ = Describe("ReplicaReconciler — state reconciliation", func() {
}
Expect(db.Create(stale).Error).To(Succeed())
prober := &fakeProber{alive: map[string]bool{"10.0.0.1:12345": true}}
prober := &fakeProber{outcomes: map[string]ProbeOutcome{"10.0.0.1:12345": ProbeAlive}}
rc := NewReplicaReconciler(ReplicaReconcilerOptions{
Registry: registry,
DB: db,

View File

@@ -0,0 +1,212 @@
package nodes
import (
"context"
"errors"
"runtime"
"time"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"gorm.io/gorm"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/LocalAI/core/services/testutil"
)
// fakeProcessLister stands in for the NATS round-trip to a worker.
type fakeProcessLister struct {
running map[string][]messaging.RunningModelInfo
err error
calls int
}
func (f *fakeProcessLister) ListRunningModels(nodeID string) (*messaging.ModelsRunningReply, error) {
f.calls++
if f.err != nil {
return nil, f.err
}
return &messaging.ModelsRunningReply{Models: f.running[nodeID]}, nil
}
// The worker owns the backend processes, so its answer is authoritative and,
// unlike a health probe against the backend's own port, is unaffected by how
// busy those backends are. These tests pin the reconciliation of registry rows
// against that answer.
var _ = Describe("ReplicaReconciler — reconcile against worker processes", func() {
var (
db *gorm.DB
registry *NodeRegistry
node *BackendNode
)
BeforeEach(func() {
if runtime.GOOS == "darwin" {
Skip("testcontainers requires Docker, not available on macOS CI")
}
db = testutil.SetupTestDB()
var err error
registry, err = NewNodeRegistry(db)
Expect(err).ToNot(HaveOccurred())
node = &BackendNode{Name: "w1", NodeType: NodeTypeBackend, Address: "10.0.0.1:50051"}
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
})
seed := func(id, modelName string, replica int, age time.Duration) {
Expect(db.Create(&NodeModel{
ID: id,
NodeID: node.ID,
ModelName: modelName,
ReplicaIndex: replica,
Address: "10.0.0.1:12345",
State: "loaded",
UpdatedAt: time.Now().Add(-age),
}).Error).To(Succeed())
}
makeStale := func(id string) {
Expect(db.Model(&NodeModel{}).Where("id = ?", id).
Update("updated_at", time.Now().Add(-5*time.Minute)).Error).To(Succeed())
}
newReconciler := func(lister NodeProcessLister) *ReplicaReconciler {
return NewReplicaReconciler(ReplicaReconcilerOptions{
Registry: registry,
DB: db,
ProcessLister: lister,
ProbeStaleAfter: 2 * time.Minute,
})
}
It("reaps a row for a model the worker is not running", func() {
seed("ghost-1", "ghost-model", 0, 5*time.Minute)
lister := &fakeProcessLister{running: map[string][]messaging.RunningModelInfo{}}
rc := newReconciler(lister)
for range workerMissesBeforeReap {
rc.reconcileNodeProcesses(context.Background())
makeStale("ghost-1")
}
var after NodeModel
Expect(db.First(&after, "id = ?", "ghost-1").Error).To(MatchError(gorm.ErrRecordNotFound),
"the worker is the authority on which processes exist")
})
It("keeps and refreshes a row the worker confirms is running", func() {
seed("live-1", "live-model", 0, 5*time.Minute)
lister := &fakeProcessLister{running: map[string][]messaging.RunningModelInfo{
node.ID: {{ModelID: "live-model", ReplicaIndex: 0, Address: "10.0.0.1:12345"}},
}}
rc := newReconciler(lister)
for range workerMissesBeforeReap + 2 {
rc.reconcileNodeProcesses(context.Background())
makeStale("live-1")
}
var after NodeModel
Expect(db.First(&after, "id = ?", "live-1").Error).To(Succeed())
})
It("refreshes updated_at so the port probe skips a worker-confirmed model", func() {
// This is what keeps a busy backend off the port prober entirely: the
// worker vouches for it, so it never looks stale enough to probe.
seed("busy-1", "busy-model", 0, 5*time.Minute)
lister := &fakeProcessLister{running: map[string][]messaging.RunningModelInfo{
node.ID: {{ModelID: "busy-model", ReplicaIndex: 0, Address: "10.0.0.1:12345"}},
}}
rc := newReconciler(lister)
rc.reconcileNodeProcesses(context.Background())
var after NodeModel
Expect(db.First(&after, "id = ?", "busy-1").Error).To(Succeed())
Expect(after.UpdatedAt).To(BeTemporally("~", time.Now(), 5*time.Second))
})
It("distinguishes replicas of the same model", func() {
seed("rep-0", "multi-model", 0, 5*time.Minute)
seed("rep-1", "multi-model", 1, 5*time.Minute)
lister := &fakeProcessLister{running: map[string][]messaging.RunningModelInfo{
node.ID: {{ModelID: "multi-model", ReplicaIndex: 0, Address: "10.0.0.1:12345"}},
}}
rc := newReconciler(lister)
for range workerMissesBeforeReap {
rc.reconcileNodeProcesses(context.Background())
makeStale("rep-0")
makeStale("rep-1")
}
var kept NodeModel
Expect(db.First(&kept, "id = ?", "rep-0").Error).To(Succeed(),
"replica 0 is running and must survive")
var reaped NodeModel
Expect(db.First(&reaped, "id = ?", "rep-1").Error).To(MatchError(gorm.ErrRecordNotFound),
"replica 1 is not running and must be reaped")
})
It("reaps nothing when the worker cannot be reached", func() {
// A NATS failure says nothing about the processes. Reaping here would
// delete the whole node's rows on a transient messaging blip.
seed("unknown-1", "unknown-model", 0, 5*time.Minute)
lister := &fakeProcessLister{err: errors.New("nats: timeout")}
rc := newReconciler(lister)
for range workerMissesBeforeReap + 3 {
rc.reconcileNodeProcesses(context.Background())
makeStale("unknown-1")
}
var after NodeModel
Expect(db.First(&after, "id = ?", "unknown-1").Error).To(Succeed(),
"an unreachable worker must never cause its replicas to be reaped")
})
It("ignores rows younger than the stale cutoff", func() {
// A row created moments ago may legitimately not be in the worker's
// table yet; judging it immediately would race every fresh load.
seed("fresh-1", "fresh-model", 0, 0)
lister := &fakeProcessLister{running: map[string][]messaging.RunningModelInfo{}}
rc := newReconciler(lister)
for range workerMissesBeforeReap + 2 {
rc.reconcileNodeProcesses(context.Background())
}
var after NodeModel
Expect(db.First(&after, "id = ?", "fresh-1").Error).To(Succeed())
Expect(lister.calls).To(Equal(0), "no stale rows means no reason to ask the worker")
})
It("requires consecutive misses before reaping", func() {
seed("flap-1", "flap-model", 0, 5*time.Minute)
lister := &fakeProcessLister{running: map[string][]messaging.RunningModelInfo{}}
rc := newReconciler(lister)
for i := 1; i < workerMissesBeforeReap; i++ {
rc.reconcileNodeProcesses(context.Background())
makeStale("flap-1")
var after NodeModel
Expect(db.First(&after, "id = ?", "flap-1").Error).To(Succeed())
}
// It shows up again: the streak resets.
lister.running[node.ID] = []messaging.RunningModelInfo{
{ModelID: "flap-model", ReplicaIndex: 0, Address: "10.0.0.1:12345"},
}
rc.reconcileNodeProcesses(context.Background())
makeStale("flap-1")
lister.running[node.ID] = nil
for i := 1; i < workerMissesBeforeReap; i++ {
rc.reconcileNodeProcesses(context.Background())
makeStale("flap-1")
}
var after NodeModel
Expect(db.First(&after, "id = ?", "flap-1").Error).To(Succeed(),
"a reappearance must reset the consecutive-miss streak")
})
})

View File

@@ -252,35 +252,67 @@ const (
// NodeRegistry manages backend node registration and lookup in PostgreSQL.
type NodeRegistry struct {
db *gorm.DB
// replicaRemovedHook is invoked after a replica row for (modelName, nodeID)
// is removed. It is the single chokepoint that lets the prefix-cache index
// be invalidated no matter which removal path (router eviction, reconciler
// replicaRemovedHooks are invoked after a replica row for (modelName, nodeID)
// is removed. This is the single chokepoint that lets dependent state be
// invalidated no matter which removal path (router eviction, reconciler
// scale-down, probe reaper, health-monitor reap, RemoteUnloaderAdapter) ran.
// The replicaIndex argument is the SPECIFIC replica removed, or negative to
// signal "all replicas of (modelName, nodeID)". Stored in an atomic.Pointer
// so the startup wiring (setter) and request / reconcile handling (fire) are
// race-free.
replicaRemovedHook atomic.Pointer[func(modelName, nodeID string, replicaIndex int)]
// signal "all replicas of (modelName, nodeID)".
//
// A LIST, not a single slot: the prefix-cache index and the frontend's local
// model store are independent subsystems that both need invalidating, and
// they are wired from different places. With one slot the second registration
// silently displaced the first.
//
// Stored in an atomic.Pointer to an immutable slice so the startup wiring
// (append) and request / reconcile handling (fire) are race-free.
replicaRemovedHooks atomic.Pointer[[]func(modelName, nodeID string, replicaIndex int)]
}
// SetReplicaRemovedHook registers a callback invoked after a replica row for
// AddReplicaRemovedHook registers a callback invoked after a replica row for
// (modelName, nodeID) is removed from the registry. replicaIndex is the
// specific replica removed, or negative to mean "all replicas of the node".
// Used to invalidate the prefix-cache index so it never points at a replica
// that no longer hosts the model. Set once at startup before serving. Safe to
// leave unset (no-op).
func (r *NodeRegistry) SetReplicaRemovedHook(fn func(modelName, nodeID string, replicaIndex int)) {
r.replicaRemovedHook.Store(&fn)
// Every registered hook fires; registering one never displaces another.
// Called at startup before serving. Safe to leave unregistered (no-op).
func (r *NodeRegistry) AddReplicaRemovedHook(fn func(modelName, nodeID string, replicaIndex int)) {
if fn == nil {
return
}
for {
current := r.replicaRemovedHooks.Load()
existing := r.loadReplicaRemovedHooks()
updated := make([]func(string, string, int), 0, len(existing)+1)
updated = append(updated, existing...)
updated = append(updated, fn)
if r.replicaRemovedHooks.CompareAndSwap(current, &updated) {
return
}
}
}
// fireReplicaRemoved invokes the replica-removed hook if one is set. A negative
// loadReplicaRemovedHooks returns the currently registered hooks, or nil when
// none are registered.
func (r *NodeRegistry) loadReplicaRemovedHooks() []func(modelName, nodeID string, replicaIndex int) {
if p := r.replicaRemovedHooks.Load(); p != nil {
return *p
}
return nil
}
// fireReplicaRemoved invokes every registered replica-removed hook. A negative
// replicaIndex means all replicas of (modelName, nodeID). Nil-safe.
func (r *NodeRegistry) fireReplicaRemoved(modelName, nodeID string, replicaIndex int) {
if fn := r.replicaRemovedHook.Load(); fn != nil && *fn != nil {
(*fn)(modelName, nodeID, replicaIndex)
for _, fn := range r.loadReplicaRemovedHooks() {
fn(modelName, nodeID, replicaIndex)
}
}
// hasReplicaRemovedHook reports whether any hook is registered, so bulk delete
// paths can skip the extra enumeration query when nothing is listening.
func (r *NodeRegistry) hasReplicaRemovedHook() bool {
return len(r.loadReplicaRemovedHooks()) > 0
}
// nodeModelNames returns the DISTINCT model names that have node_models rows
// for nodeID, using db (which may be a transaction handle). Used by the bulk
// node-scoped delete paths (Register re-register cleanup, MarkOffline,
@@ -290,7 +322,7 @@ func (r *NodeRegistry) fireReplicaRemoved(modelName, nodeID string, replicaIndex
// the model. Skips the query entirely when no hook is set (these are lifecycle
// ops, not the request hot path, but the query is pure overhead with no hook).
func (r *NodeRegistry) nodeModelNames(ctx context.Context, db *gorm.DB, nodeID string) []string {
if fn := r.replicaRemovedHook.Load(); fn == nil || *fn == nil {
if !r.hasReplicaRemovedHook() {
return nil
}
var names []string

View File

@@ -1077,7 +1077,7 @@ var _ = Describe("NodeRegistry", func() {
})
})
Describe("SetReplicaRemovedHook", func() {
Describe("AddReplicaRemovedHook", func() {
type removed struct {
model, node string
replica int
@@ -1089,7 +1089,7 @@ var _ = Describe("NodeRegistry", func() {
Expect(registry.SetNodeModel(context.Background(), node.ID, "hook-model", 1, "loaded", "a", 0)).To(Succeed())
var fired []removed
registry.SetReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
fired = append(fired, removed{model: modelName, node: nodeID, replica: replicaIndex})
})
@@ -1106,7 +1106,7 @@ var _ = Describe("NodeRegistry", func() {
Expect(registry.SetNodeModel(context.Background(), node.ID, "hook-all-model", 1, "loaded", "b", 0)).To(Succeed())
var fired []removed
registry.SetReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
fired = append(fired, removed{model: modelName, node: nodeID, replica: replicaIndex})
})
@@ -1148,7 +1148,7 @@ var _ = Describe("NodeRegistry", func() {
seedTwoModels(node)
fired := map[removed]int{}
registry.SetReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
// Bulk node-scoped deletes signal "all replicas" with replica<0.
Expect(replicaIndex).To(BeNumerically("<", 0))
fired[removed{model: modelName, node: nodeID}]++
@@ -1165,7 +1165,7 @@ var _ = Describe("NodeRegistry", func() {
seedTwoModels(node)
fired := map[removed]int{}
registry.SetReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
// Bulk node-scoped deletes signal "all replicas" with replica<0.
Expect(replicaIndex).To(BeNumerically("<", 0))
fired[removed{model: modelName, node: nodeID}]++
@@ -1182,7 +1182,7 @@ var _ = Describe("NodeRegistry", func() {
seedTwoModels(node)
fired := map[removed]int{}
registry.SetReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
// Bulk node-scoped deletes signal "all replicas" with replica<0.
Expect(replicaIndex).To(BeNumerically("<", 0))
fired[removed{model: modelName, node: nodeID}]++
@@ -1199,7 +1199,7 @@ var _ = Describe("NodeRegistry", func() {
seedTwoModels(node)
fired := map[removed]int{}
registry.SetReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
// Bulk node-scoped deletes signal "all replicas" with replica<0.
Expect(replicaIndex).To(BeNumerically("<", 0))
fired[removed{model: modelName, node: nodeID}]++
@@ -1238,7 +1238,7 @@ var _ = Describe("NodeRegistry", func() {
Expect(expectedModels).To(HaveLen(2), "seed should create two distinct models")
fired := map[string]struct{}{}
registry.SetReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replicaIndex int) {
Expect(nodeID).To(Equal(node.ID))
Expect(replicaIndex).To(BeNumerically("<", 0))
fired[modelName] = struct{}{}

View File

@@ -10,6 +10,7 @@ import (
"os"
"path/filepath"
"strings"
"sync"
"time"
"github.com/mudler/LocalAI/core/config"
@@ -586,16 +587,7 @@ func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType
r.observePrefix(trackingKey, observeChain, prefixcache.ReplicaKey{NodeID: node.ID, Replica: replicaIdx})
grpcClient := r.buildClientForAddr(node, modelAddr, parallel)
tracked := NewInFlightTrackingClient(grpcClient, r.registry, node.ID, trackingKey, replicaIdx)
tracked.OnFirstComplete(func() {
r.registry.DecrementInFlight(context.Background(), node.ID, trackingKey, replicaIdx)
})
return &RouteResult{
Node: node,
Client: tracked,
Release: func() {
closeClient(grpcClient)
},
}, nil
return r.newRouteResult(node, trackingKey, replicaIdx, grpcClient, tracked), nil
}
}
}
@@ -659,16 +651,7 @@ func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType
r.observePrefix(trackingKey, observeChain, prefixcache.ReplicaKey{NodeID: node.ID, Replica: replicaIdx})
grpcClient := r.buildClientForAddr(node, modelAddr, parallel)
tracked := NewInFlightTrackingClient(grpcClient, r.registry, node.ID, trackingKey, replicaIdx)
tracked.OnFirstComplete(func() {
r.registry.DecrementInFlight(context.Background(), node.ID, trackingKey, replicaIdx)
})
return &RouteResult{
Node: node,
Client: tracked,
Release: func() {
closeClient(grpcClient)
},
}, nil
return r.newRouteResult(node, trackingKey, replicaIdx, grpcClient, tracked), nil
}
}
}
@@ -686,16 +669,7 @@ func (r *SmartRouter) Route(ctx context.Context, modelID, modelName, backendType
replicaIdx := result.ReplicaIndex
tracked := NewInFlightTrackingClient(result.Client, r.registry, result.Node.ID, trackingKey, replicaIdx)
tracked.OnFirstComplete(func() {
r.registry.DecrementInFlight(context.Background(), result.Node.ID, trackingKey, replicaIdx)
})
return &RouteResult{
Node: result.Node,
Client: tracked,
Release: func() {
closeClient(result.Client)
},
}, nil
return r.newRouteResult(result.Node, trackingKey, replicaIdx, result.Client, tracked), nil
}
if r.db != nil {
@@ -1581,6 +1555,31 @@ func pathBytes(path string) int64 {
// 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 {
@@ -1590,11 +1589,13 @@ func countStageableFiles(path string) int {
return 1
}
n := 0
_ = filepath.WalkDir(path, func(_ string, d fs.DirEntry, walkErr error) error {
_ = filepath.WalkDir(path, func(p string, d fs.DirEntry, walkErr error) error {
if walkErr != nil {
return nil
}
if !d.IsDir() {
// 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
@@ -1617,6 +1618,11 @@ func (r *SmartRouter) stageDirectory(ctx context.Context, node *BackendNode, tra
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)
@@ -1779,6 +1785,39 @@ func (r *SmartRouter) probeHealth(ctx context.Context, node *BackendNode, addr s
}
// 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()

View File

@@ -0,0 +1,92 @@
package nodes
import (
"context"
"runtime"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
"gorm.io/gorm"
"github.com/mudler/LocalAI/core/services/testutil"
)
// Routing reserves in_flight = 1 at load time so a freshly loaded replica is not
// immediately evicted out from under the request that caused the load. That
// reservation used to be released ONLY by the first inference completing, so a
// route torn down before any inference ran (client disconnect, handler error
// before the backend call) stranded the counter. A stranded counter is not
// cosmetic: every eviction query requires in_flight = 0, so the replica's VRAM
// becomes unreclaimable.
var _ = Describe("SmartRouter routing reservation", func() {
var (
db *gorm.DB
registry *NodeRegistry
router *SmartRouter
node *BackendNode
)
BeforeEach(func() {
if runtime.GOOS == "darwin" {
Skip("testcontainers requires Docker, not available on macOS CI")
}
db = testutil.SetupTestDB()
var err error
registry, err = NewNodeRegistry(db)
Expect(err).ToNot(HaveOccurred())
router = &SmartRouter{registry: registry}
node = &BackendNode{Name: "n1", NodeType: NodeTypeBackend, Address: "10.0.0.1:50051"}
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
// A loaded replica holding the load-time reservation.
Expect(registry.SetNodeModel(context.Background(), node.ID, "m", 0, "loaded", "10.0.0.1:12345", 1)).To(Succeed())
})
inFlight := func() int {
var m NodeModel
Expect(db.First(&m, "node_id = ? AND model_name = ?", node.ID, "m").Error).To(Succeed())
return m.InFlight
}
newResult := func() *RouteResult {
raw := &stubBackend{}
tracked := NewInFlightTrackingClient(raw, registry, node.ID, "m", 0)
return router.newRouteResult(node, "m", 0, raw, tracked)
}
It("releases the reservation when the route is torn down without any inference", func() {
result := newResult()
Expect(inFlight()).To(Equal(1))
result.Release()
Expect(inFlight()).To(Equal(0),
"a route that never ran an inference must still give the reservation back")
})
It("releases the reservation exactly once across both paths", func() {
result := newResult()
tracked, ok := result.Client.(*InFlightTrackingClient)
Expect(ok).To(BeTrue())
// Simulate the first inference completing, then the route being torn
// down. Only one of these may consume the single reservation.
tracked.firstOnce.Do(tracked.onFirstComplete)
result.Release()
result.Release()
Expect(inFlight()).To(Equal(0),
"double release would under-count and let a busy replica look idle")
})
It("does not disturb concurrent per-request tracking", func() {
result := newResult()
// A second request arrives and is tracked normally.
Expect(registry.IncrementInFlight(context.Background(), node.ID, "m", 0)).To(Succeed())
Expect(inFlight()).To(Equal(2))
result.Release()
Expect(inFlight()).To(Equal(1),
"releasing the reservation must leave a genuinely in-flight request counted")
})
})

View File

@@ -0,0 +1,107 @@
package nodes
import (
"context"
"os"
"path/filepath"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
pb "github.com/mudler/LocalAI/pkg/grpc/proto"
)
// The receiving side writes a "<file>.sha256" sidecar next to every file it
// accepts (see file_transfer_server.go). When the sender walks a model
// directory it must skip those sidecars: staging them makes the receiver write
// a sidecar for the sidecar, so each pass multiplies the tree
// (config.json -> config.json.sha256 -> config.json.sha256.sha256 -> ...) and
// grows disk without bound on both ends.
var _ = Describe("stageModelFiles hash sidecars", func() {
var (
stager *fakeFileStager
router *SmartRouter
node *BackendNode
tmp string
modelID = "sidecar-model"
)
BeforeEach(func() {
stager = &fakeFileStager{}
router = &SmartRouter{
fileStager: stager,
stagingTracker: NewStagingTracker(),
}
node = &BackendNode{ID: "node-1", Name: "node-1", Address: "10.0.0.1:50051"}
tmp = GinkgoT().TempDir()
})
It("does not stage hash sidecars left behind by a previous transfer", func() {
modelDir := filepath.Join(tmp, "models", modelID)
Expect(os.MkdirAll(modelDir, 0o755)).To(Succeed())
weights := filepath.Join(modelDir, "model.safetensors")
config := filepath.Join(modelDir, "config.json")
Expect(os.WriteFile(weights, []byte("weights"), 0o644)).To(Succeed())
Expect(os.WriteFile(config, []byte("{}"), 0o644)).To(Succeed())
// Sidecars a previous staging round deposited, including the chained
// ones produced by the runaway this guards against.
for _, sidecar := range []string{
"model.safetensors.sha256",
"config.json.sha256",
"config.json.sha256.sha256",
"config.json.sha256.sha256.sha256",
"config.json.sha256.target",
} {
Expect(os.WriteFile(filepath.Join(modelDir, sidecar), []byte("deadbeef"), 0o644)).To(Succeed())
}
opts := &pb.ModelOptions{Model: modelID, ModelFile: modelDir}
_, err := router.stageModelFiles(context.Background(), node, opts, "track-key")
Expect(err).ToNot(HaveOccurred())
staged := make([]string, 0, len(stager.ensureCalls))
for _, c := range stager.ensureCalls {
staged = append(staged, c.localPath)
}
Expect(staged).To(ConsistOf(weights, config),
"only real model files may be staged; hash sidecars are regenerated by the receiver")
})
It("excludes hash sidecars from the staging file count", func() {
modelDir := filepath.Join(tmp, "models", "counted")
Expect(os.MkdirAll(modelDir, 0o755)).To(Succeed())
Expect(os.WriteFile(filepath.Join(modelDir, "a.gguf"), []byte("a"), 0o644)).To(Succeed())
Expect(os.WriteFile(filepath.Join(modelDir, "b.gguf"), []byte("b"), 0o644)).To(Succeed())
for _, sidecar := range []string{"a.gguf.sha256", "b.gguf.sha256", "b.gguf.sha256.sha256"} {
Expect(os.WriteFile(filepath.Join(modelDir, sidecar), []byte("x"), 0o644)).To(Succeed())
}
// The progress bar counts what will actually be uploaded; counting
// sidecars strands it short of 100%.
Expect(countStageableFiles(modelDir)).To(Equal(2))
})
It("still stages a model file whose own name ends in .sha256", func() {
// Defensive: the skip must key off "sidecar sitting next to a real
// file", not a blanket suffix ban, so a model that genuinely ships a
// .sha256 payload with no base file is still transferred.
modelDir := filepath.Join(tmp, "models", "checksums-only")
Expect(os.MkdirAll(modelDir, 0o755)).To(Succeed())
orphan := filepath.Join(modelDir, "manifest.sha256")
Expect(os.WriteFile(orphan, []byte("abc"), 0o644)).To(Succeed())
opts := &pb.ModelOptions{Model: "checksums-only", ModelFile: modelDir}
_, err := router.stageModelFiles(context.Background(), node, opts, "track-key")
Expect(err).ToNot(HaveOccurred())
staged := make([]string, 0, len(stager.ensureCalls))
for _, c := range stager.ensureCalls {
staged = append(staged, c.localPath)
}
Expect(staged).To(ConsistOf(orphan))
})
})

View File

@@ -457,12 +457,17 @@ var _ = Describe("SmartRouter", func() {
// TouchNodeModel should have been called
Expect(reg.touchCalls).To(ContainElement("n1:my-model"))
// The initial in-flight reservation from FindAndLockNodeWithModel is released
// after the first inference call completes via OnFirstComplete callback.
// Release only closes the client.
// The initial in-flight reservation from FindAndLockNodeWithModel is
// released by whichever comes first: the first inference completing
// (OnFirstComplete) or the route being torn down. Teardown must
// release it too, or a route that never reached the backend leaks the
// counter and pins the replica against every eviction query.
result.Release()
// No decrement on Release — it happens via OnFirstComplete after first Predict
Expect(reg.decrementCalls).To(BeEmpty())
Expect(reg.decrementCalls).To(ContainElement("n1:my-model"))
// Exactly once, however many times teardown runs.
result.Release()
Expect(reg.decrementCalls).To(HaveLen(1))
})
})
@@ -1475,7 +1480,7 @@ var _ = Describe("SmartRouter prefix-cache routing", func() {
// UnloadModel must route the eviction through the registry removal
// chokepoint (RemoveAllNodeModelReplicas). The registry's
// SetReplicaRemovedHook is what invalidates the prefix index in
// AddReplicaRemovedHook is what invalidates the prefix index in
// production; the router no longer invalidates directly. Here the
// fake registry records the removal but fires no hook, so we assert
// the chokepoint is exercised rather than the downstream

View File

@@ -320,6 +320,19 @@ func (a *RemoteUnloaderAdapter) ListBackends(nodeID string) (*messaging.BackendL
return messaging.RequestJSON[messaging.BackendListRequest, messaging.BackendListReply](a.nats, subject, messaging.BackendListRequest{}, 30*time.Second)
}
// ListRunningModels asks a worker node which model backend processes it
// currently has running, via NATS request-reply.
//
// The timeout is short on purpose: the worker answers straight out of its
// in-memory process table, so a slow reply means the worker itself is in
// trouble, and the caller treats no-answer as "don't know" rather than as
// "nothing running".
func (a *RemoteUnloaderAdapter) ListRunningModels(nodeID string) (*messaging.ModelsRunningReply, error) {
subject := messaging.SubjectNodeModelsRunning(nodeID)
return messaging.RequestJSON[messaging.ModelsRunningRequest, messaging.ModelsRunningReply](
a.nats, subject, messaging.ModelsRunningRequest{}, 10*time.Second)
}
// StopBackend tells a worker node to stop a specific gRPC backend process.
// If backend is empty, the worker stops ALL backends.
// The node stays registered and can receive another InstallBackend later.

View File

@@ -36,6 +36,9 @@ func (s *backendSupervisor) subscribeLifecycleEvents() error {
if _, err := s.nats.SubscribeReply(messaging.SubjectNodeBackendList(s.nodeID), s.handleBackendList); err != nil {
return fmt.Errorf("subscribing to backend list events: %w", err)
}
if _, err := s.nats.SubscribeReply(messaging.SubjectNodeModelsRunning(s.nodeID), s.handleModelsRunning); err != nil {
return fmt.Errorf("subscribing to models running events: %w", err)
}
if _, err := s.nats.SubscribeReply(messaging.SubjectNodeModelUnload(s.nodeID), s.handleModelUnload); err != nil {
return fmt.Errorf("subscribing to model unload events: %w", err)
}

View File

@@ -0,0 +1,65 @@
package worker
import (
"strconv"
"strings"
"github.com/mudler/LocalAI/core/services/messaging"
"github.com/mudler/xlog"
)
// parseProcessKey is the inverse of buildProcessKey: it splits a
// `modelID#replicaIndex` process key back into its parts.
//
// The split is on the LAST '#' because model ids are user-supplied and may
// themselves contain one; only the trailing "#N" is the supervisor's suffix.
// Returns ok=false for anything that does not carry a numeric replica suffix,
// so a malformed key is skipped rather than reported under a wrong identity.
func parseProcessKey(key string) (modelID string, replicaIndex int, ok bool) {
hash := strings.LastIndex(key, "#")
if hash < 0 {
return "", 0, false
}
replica, err := strconv.Atoi(key[hash+1:])
if err != nil {
return "", 0, false
}
return key[:hash], replica, true
}
// runningModels returns the model backend processes this worker currently has
// alive, in the (modelID, replicaIndex, address) shape the controller's
// registry rows are keyed by.
//
// Processes being stopped are excluded: they are alive but on their way out,
// and reporting them would resurrect a replica the controller just released.
func (s *backendSupervisor) runningModels() []messaging.RunningModelInfo {
s.mu.Lock()
defer s.mu.Unlock()
running := make([]messaging.RunningModelInfo, 0, len(s.processes))
for key, bp := range s.processes {
if bp == nil || bp.stopping || bp.proc == nil || !bp.proc.IsAlive() {
continue
}
modelID, replicaIndex, ok := parseProcessKey(key)
if !ok {
xlog.Warn("Skipping unparseable process key when reporting running models", "key", key)
continue
}
running = append(running, messaging.RunningModelInfo{
ModelID: modelID,
ReplicaIndex: replicaIndex,
Address: bp.addr,
})
}
return running
}
// handleModelsRunning answers a models.running request with this worker's live
// process set.
func (s *backendSupervisor) handleModelsRunning(_ []byte, reply func([]byte)) {
running := s.runningModels()
xlog.Debug("Answering models.running", "nodeID", s.nodeID, "count", len(running))
replyJSON(reply, messaging.ModelsRunningReply{Models: running})
}

View File

@@ -0,0 +1,113 @@
package worker
import (
"os/exec"
"time"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
process "github.com/mudler/go-processmanager"
)
// The worker owns the backend processes and its answer is not affected by how
// busy those processes are, which is what makes it a sounder liveness source
// than probing the backend's own serving port. These tests pin the reply it
// gives the reconciler.
var _ = Describe("backendSupervisor running models", func() {
// liveProcess starts a real long-running child so IsAlive() is true.
liveProcess := func() *process.Process {
sleepBin, err := exec.LookPath("sleep")
Expect(err).ToNot(HaveOccurred())
p := process.New(
process.WithTemporaryStateDir(),
process.WithName(sleepBin),
process.WithArgs("60"),
)
Expect(p.Run()).To(Succeed())
DeferCleanup(func() { _ = p.Stop() })
Eventually(p.IsAlive, 5*time.Second, 20*time.Millisecond).Should(BeTrue())
return p
}
It("reports each live process as its model and replica index", func() {
s := &backendSupervisor{
processes: map[string]*backendProcess{
"qwen3.6-35B#0": {proc: liveProcess(), addr: "127.0.0.1:50051"},
"qwen3.6-35B#1": {proc: liveProcess(), addr: "127.0.0.1:50052"},
"other-model#0": {proc: liveProcess(), addr: "127.0.0.1:50053"},
},
}
running := s.runningModels()
Expect(running).To(ConsistOf(
HaveField("ModelID", "qwen3.6-35B"),
HaveField("ModelID", "qwen3.6-35B"),
HaveField("ModelID", "other-model"),
))
byAddr := map[string]int{}
for _, r := range running {
byAddr[r.Address] = r.ReplicaIndex
}
Expect(byAddr).To(Equal(map[string]int{
"127.0.0.1:50051": 0,
"127.0.0.1:50052": 1,
"127.0.0.1:50053": 0,
}))
})
It("omits processes that are not alive", func() {
// A recorded slot whose process died must not be reported as running,
// otherwise the reconciler would keep a dead replica's row forever.
s := &backendSupervisor{
processes: map[string]*backendProcess{
"dead-model#0": {proc: nil, addr: "127.0.0.1:50051"},
"live-model#0": {proc: liveProcess(), addr: "127.0.0.1:50052"},
},
}
Expect(s.runningModels()).To(ConsistOf(HaveField("ModelID", "live-model")))
})
It("omits processes that are being stopped", func() {
// Mid-stop the process is still alive, but its row is on its way out;
// reporting it would resurrect a replica the controller just released.
s := &backendSupervisor{
processes: map[string]*backendProcess{
"going-away#0": {proc: liveProcess(), addr: "127.0.0.1:50051", stopping: true},
},
}
Expect(s.runningModels()).To(BeEmpty())
})
})
var _ = Describe("parseProcessKey", func() {
It("splits a modelID#replica key, tolerating '#' inside the model id", func() {
modelID, replica, ok := parseProcessKey("qwen3.6-35B#2")
Expect(ok).To(BeTrue())
Expect(modelID).To(Equal("qwen3.6-35B"))
Expect(replica).To(Equal(2))
// Model ids are user-supplied, so the split must be on the LAST '#'.
modelID, replica, ok = parseProcessKey("weird#name#3")
Expect(ok).To(BeTrue())
Expect(modelID).To(Equal("weird#name"))
Expect(replica).To(Equal(3))
})
It("rejects keys without a valid replica suffix", func() {
_, _, ok := parseProcessKey("no-suffix")
Expect(ok).To(BeFalse())
_, _, ok = parseProcessKey("bad#suffix")
Expect(ok).To(BeFalse())
})
It("round-trips with buildProcessKey", func() {
modelID, replica, ok := parseProcessKey(buildProcessKey("my-model", "some-backend", 4))
Expect(ok).To(BeTrue())
Expect(modelID).To(Equal("my-model"))
Expect(replica).To(Equal(4))
})
})

View File

@@ -110,7 +110,7 @@ var _ = Describe("Prefix-cache aware routing", Label("Distributed"), func() {
// invalidation is exercised end to end. A negative replica index means
// "all replicas of the node" (InvalidateNode); otherwise drop the exact
// replica.
registry.SetReplicaRemovedHook(func(modelName, nodeID string, replica int) {
registry.AddReplicaRemovedHook(func(modelName, nodeID string, replica int) {
if replica < 0 {
idx.InvalidateNode(modelName, nodeID)
} else {