diff --git a/core/application/distributed.go b/core/application/distributed.go index 750b73b02..32269a44a 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -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 { diff --git a/core/application/startup.go b/core/application/startup.go index 7dd5a8d61..a28f0f847 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -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 diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index 11e3bfba9..5f09adda5 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -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 { diff --git a/core/services/nodes/local_stub_invalidator.go b/core/services/nodes/local_stub_invalidator.go new file mode 100644 index 000000000..86adb171b --- /dev/null +++ b/core/services/nodes/local_stub_invalidator.go @@ -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) + } +} diff --git a/core/services/nodes/local_stub_invalidator_test.go b/core/services/nodes/local_stub_invalidator_test.go new file mode 100644 index 000000000..a8f591417 --- /dev/null +++ b/core/services/nodes/local_stub_invalidator_test.go @@ -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()) + }) +}) diff --git a/core/services/nodes/model_router_test.go b/core/services/nodes/model_router_test.go index a0c14ab26..f68caa143 100644 --- a/core/services/nodes/model_router_test.go +++ b/core/services/nodes/model_router_test.go @@ -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)) }) }) }) diff --git a/core/services/nodes/reconciler.go b/core/services/nodes/reconciler.go index 1f86c1321..38fd97cc2 100644 --- a/core/services/nodes/reconciler.go +++ b/core/services/nodes/reconciler.go @@ -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) + } } } diff --git a/core/services/nodes/reconciler_busy_probe_test.go b/core/services/nodes/reconciler_busy_probe_test.go new file mode 100644 index 000000000..9cc7b07f9 --- /dev/null +++ b/core/services/nodes/reconciler_busy_probe_test.go @@ -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)) + }) +}) diff --git a/core/services/nodes/reconciler_inflight_leak_test.go b/core/services/nodes/reconciler_inflight_leak_test.go new file mode 100644 index 000000000..f3574fab7 --- /dev/null +++ b/core/services/nodes/reconciler_inflight_leak_test.go @@ -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")) + }) +}) diff --git a/core/services/nodes/reconciler_test.go b/core/services/nodes/reconciler_test.go index 6a0171dfd..049fb9441 100644 --- a/core/services/nodes/reconciler_test.go +++ b/core/services/nodes/reconciler_test.go @@ -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, diff --git a/core/services/nodes/reconciler_worker_processes_test.go b/core/services/nodes/reconciler_worker_processes_test.go new file mode 100644 index 000000000..8fd848b1d --- /dev/null +++ b/core/services/nodes/reconciler_worker_processes_test.go @@ -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") + }) +}) diff --git a/core/services/nodes/registry.go b/core/services/nodes/registry.go index b9c78ee03..0b85a1f44 100644 --- a/core/services/nodes/registry.go +++ b/core/services/nodes/registry.go @@ -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 diff --git a/core/services/nodes/registry_test.go b/core/services/nodes/registry_test.go index 589a62fa3..a9d4b059a 100644 --- a/core/services/nodes/registry_test.go +++ b/core/services/nodes/registry_test.go @@ -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{}{} diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 62bf302ac..fd051e954 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -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 ".sha256" for every file it accepts, so re-staging a +// directory that already held sidecars produced ".sha256.sha256", then +// ".sha256.sha256.sha256", multiplying the tree on every pass. +// +// The check is deliberately "a sidecar sitting next to a real file" rather than +// a blanket suffix ban, so a model that genuinely ships a .sha256 payload with +// no corresponding base file is still transferred. +func isHashSidecar(path string) bool { + for _, suffix := range []string{targetSidecarSuffix, hashSidecarSuffix} { + base, ok := strings.CutSuffix(path, suffix) + if !ok { + continue + } + if fi, err := os.Stat(base); err == nil && !fi.IsDir() { + return true + } + } + return false +} + func countStageableFiles(path string) int { fi, err := os.Stat(path) if err != nil { @@ -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() diff --git a/core/services/nodes/router_reservation_test.go b/core/services/nodes/router_reservation_test.go new file mode 100644 index 000000000..c9f6cf73b --- /dev/null +++ b/core/services/nodes/router_reservation_test.go @@ -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") + }) +}) diff --git a/core/services/nodes/router_sidecar_stage_test.go b/core/services/nodes/router_sidecar_stage_test.go new file mode 100644 index 000000000..8d29fcd45 --- /dev/null +++ b/core/services/nodes/router_sidecar_stage_test.go @@ -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 ".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)) + }) +}) diff --git a/core/services/nodes/router_test.go b/core/services/nodes/router_test.go index 422822857..a3f63cc4b 100644 --- a/core/services/nodes/router_test.go +++ b/core/services/nodes/router_test.go @@ -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 diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index 24a19e866..decfa2b12 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -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. diff --git a/core/services/worker/lifecycle.go b/core/services/worker/lifecycle.go index 904a7a7f5..c80e00ea0 100644 --- a/core/services/worker/lifecycle.go +++ b/core/services/worker/lifecycle.go @@ -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) } diff --git a/core/services/worker/models_running.go b/core/services/worker/models_running.go new file mode 100644 index 000000000..efc4700a0 --- /dev/null +++ b/core/services/worker/models_running.go @@ -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}) +} diff --git a/core/services/worker/models_running_test.go b/core/services/worker/models_running_test.go new file mode 100644 index 000000000..bd676ea45 --- /dev/null +++ b/core/services/worker/models_running_test.go @@ -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)) + }) +}) diff --git a/tests/e2e/distributed/prefix_cache_routing_test.go b/tests/e2e/distributed/prefix_cache_routing_test.go index ad0cb709d..542cd29e0 100644 --- a/tests/e2e/distributed/prefix_cache_routing_test.go +++ b/tests/e2e/distributed/prefix_cache_routing_test.go @@ -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 {