diff --git a/core/services/nodes/reconciler.go b/core/services/nodes/reconciler.go index f7fefc835..88ca28bbe 100644 --- a/core/services/nodes/reconciler.go +++ b/core/services/nodes/reconciler.go @@ -278,8 +278,9 @@ func (rc *ReplicaReconciler) reconcileOnce(ctx context.Context) { // reconcileState runs the state-reconciliation passes: drain pending backend // 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. +// report they are running, port-probe whatever is left, then reclaim replica +// slots held by loads nobody is driving. 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 @@ -292,6 +293,9 @@ func (rc *ReplicaReconciler) reconcileState(ctx context.Context) { rc.reconcileNodeProcesses(ctx) rc.probeLoadedModels(ctx) rc.sweepLeakedInFlight(ctx) + // Runs last: the passes above can move a row into a serving state, and a + // row that just became loaded is no longer this sweeper's business. + rc.reclaimAbandonedLoads(ctx) } // drainPendingBackendOps retries queued backend ops whose next_retry_at has diff --git a/core/services/nodes/reconciler_abandoned_load.go b/core/services/nodes/reconciler_abandoned_load.go new file mode 100644 index 000000000..9b6b7b654 --- /dev/null +++ b/core/services/nodes/reconciler_abandoned_load.go @@ -0,0 +1,98 @@ +package nodes + +import ( + "context" + "errors" + "time" + + "github.com/mudler/xlog" + "gorm.io/gorm" +) + +const ( + // abandonedLoadGrace is how long a replica row may sit in a pre-serving + // state before the sweeper will consider it at all. + // + // It exists to cover the window between creating the replica row and + // writing the load job that vouches for it. Without it a load could be + // reclaimed in the moment before its own job row exists. It is not the + // thing that protects a long transfer: the job heartbeat does that. + abandonedLoadGrace = 5 * time.Minute +) + +// preServingStates are the replica states that hold a slot without being able +// to serve a request. NextFreeReplicaIndex counts every state except +// "unloading", so a row parked in one of these occupies capacity while +// answering nothing. +var preServingStates = []string{"loading", "staging"} + +// reclaimAbandonedLoads removes replica rows whose load will never finish. +// +// The other reconciler passes and the router's eviction query all filter +// state = "loaded", and the per-model probe skips rows without an address, so +// nothing reclaimed a row that never got that far. On a node with one replica +// slot per model, a single interrupted transfer made the model unschedulable +// there until an operator intervened: scheduling saw no free slot, and eviction +// found nothing it was allowed to evict. +// +// A row is abandoned when no live load job vouches for it. Ownership is decided +// by the job's LastProgress heartbeat rather than elapsed time, because staging +// a large checkpoint legitimately runs for a long while without touching the +// replica row. That is the same signal job takeover already trusts, so a +// transfer this sweeper reclaims is one no replica is still driving. +func (rc *ReplicaReconciler) reclaimAbandonedLoads(ctx context.Context) { + if rc.db == nil { + return + } + + cutoff := time.Now().Add(-abandonedLoadGrace) + var stuck []NodeModel + if err := rc.db.WithContext(ctx). + Where("state IN ? AND updated_at < ?", preServingStates, cutoff). + Find(&stuck).Error; err != nil { + xlog.Warn("Reconciler: failed to list replicas stuck before serving", "error", err) + return + } + + now := time.Now() + for _, row := range stuck { + if rc.loadStillRunning(ctx, row.ModelName, now) { + continue + } + if err := rc.registry.RemoveNodeModel(ctx, row.NodeID, row.ModelName, row.ReplicaIndex); err != nil { + xlog.Warn("Reconciler: failed to reclaim abandoned load", + "node", row.NodeID, "model", row.ModelName, "replica", row.ReplicaIndex, + "state", row.State, "error", err) + continue + } + xlog.Warn("Reconciler: reclaimed a replica slot held by a load nobody is driving", + "node", row.NodeID, "model", row.ModelName, "replica", row.ReplicaIndex, "state", row.State) + } +} + +// loadStillRunning reports whether a load job is actively driving this model. +// +// A missing job means nobody is loading it. A failed job has already given up. +// An orphaned job stopped heartbeating, which is the condition another replica +// uses to take it over, so the transfer behind it is not progressing either. +// Any error reading the job is treated as "still running": leaving a slot held +// for one more pass costs a scheduling opportunity, while removing a row out +// from under a live transfer would restart a multi-gigabyte load. +func (rc *ReplicaReconciler) loadStillRunning(ctx context.Context, modelName string, now time.Time) bool { + job, err := rc.registry.GetLoadJob(ctx, modelName) + if errors.Is(err, gorm.ErrRecordNotFound) { + return false + } + if err != nil { + xlog.Warn("Reconciler: cannot read load job, leaving the replica slot held", + "model", modelName, "error", err) + return true + } + if job == nil { + return false + } + if job.State == LoadJobStateFailed { + return false + } + return !job.IsOrphaned(now) +} diff --git a/core/services/nodes/reconciler_abandoned_load_test.go b/core/services/nodes/reconciler_abandoned_load_test.go new file mode 100644 index 000000000..6740b13aa --- /dev/null +++ b/core/services/nodes/reconciler_abandoned_load_test.go @@ -0,0 +1,133 @@ +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" +) + +// A replica row in loading or staging holds its slot: NextFreeReplicaIndex +// counts every state except unloading. Nothing reclaimed such a row. Every +// reconciler sweep and the router's eviction query filter state = "loaded", and +// the per-model health probe skips rows with no address, which is exactly what a +// row that never finished loading has. So a worker that dropped out mid-transfer +// left a row that pinned the only replica slot on that node for that model, and +// the next request failed with "no replica slot ... all models busy". +// +// Elapsed time alone cannot decide this: staging a large checkpoint legitimately +// runs for tens of minutes. The load job's LastProgress heartbeat is the +// discriminator, the same signal job takeover already trusts. +var _ = Describe("ReplicaReconciler — abandoned load sweeper", func() { + var ( + db *gorm.DB + registry *NodeRegistry + node *BackendNode + rc *ReplicaReconciler + ) + + 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()) + rc = NewReplicaReconciler(ReplicaReconcilerOptions{Registry: registry, DB: db}) + }) + + // seedReplica creates a replica row in the given state, aged so it is past + // the sweeper's grace period unless stated otherwise. + seedReplica := func(model, state string, age time.Duration) { + Expect(db.Create(&NodeModel{ + ID: model + "-row", + NodeID: node.ID, + ModelName: model, + State: state, + UpdatedAt: time.Now().Add(-age), + }).Error).To(Succeed()) + } + + seedJob := func(model, state string, sinceProgress time.Duration) { + Expect(db.Create(&ModelLoadJob{ + TrackingKey: model, + State: state, + OwnerReplica: "someone", + LastProgress: time.Now().Add(-sinceProgress), + CreatedAt: time.Now().Add(-sinceProgress), + UpdatedAt: time.Now().Add(-sinceProgress), + }).Error).To(Succeed()) + } + + rowExists := func(model string) bool { + var count int64 + Expect(db.Model(&NodeModel{}).Where("model_name = ?", model).Count(&count).Error).To(Succeed()) + return count > 0 + } + + It("reclaims a staging row whose load job has stopped heartbeating", func() { + seedReplica("abandoned", "staging", time.Hour) + seedJob("abandoned", LoadJobStateStaging, 30*time.Minute) + + rc.reclaimAbandonedLoads(context.Background()) + + Expect(rowExists("abandoned")).To(BeFalse()) + }) + + It("reclaims a loading row that has no load job at all", func() { + seedReplica("orphan", "loading", time.Hour) + + rc.reclaimAbandonedLoads(context.Background()) + + Expect(rowExists("orphan")).To(BeFalse()) + }) + + It("keeps a long transfer whose job is still heartbeating", func() { + // The row itself is old, because staging does not touch it. Only the + // job proves the transfer is alive. + seedReplica("big-model", "staging", time.Hour) + seedJob("big-model", LoadJobStateStaging, time.Second) + + rc.reclaimAbandonedLoads(context.Background()) + + Expect(rowExists("big-model")).To(BeTrue(), "a live transfer must never be reclaimed") + }) + + It("leaves a freshly created row alone while its job row is still being written", func() { + seedReplica("just-started", "loading", time.Second) + + rc.reclaimAbandonedLoads(context.Background()) + + Expect(rowExists("just-started")).To(BeTrue()) + }) + + It("does not touch loaded replicas, which the other sweeps own", func() { + seedReplica("serving", "loaded", time.Hour) + + rc.reclaimAbandonedLoads(context.Background()) + + Expect(rowExists("serving")).To(BeTrue()) + }) + + It("frees the slot so the model can be scheduled on that node again", func() { + seedReplica("wedged", "staging", time.Hour) + seedJob("wedged", LoadJobStateFailed, time.Minute) + + _, err := registry.NextFreeReplicaIndex(context.Background(), node.ID, "wedged", 1) + Expect(err).To(MatchError(ErrNoFreeSlot), "precondition: the stuck row holds the only slot") + + rc.reclaimAbandonedLoads(context.Background()) + + idx, err := registry.NextFreeReplicaIndex(context.Background(), node.ID, "wedged", 1) + Expect(err).ToNot(HaveOccurred()) + Expect(idx).To(Equal(0)) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 667408fbb..d82616024 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -1020,6 +1020,12 @@ Notes: - Upgrade the worker when it does not support the exact model-stop request. - Stop and restart the stale backend only as an operational recovery action. LocalAI keeps it non-routable while durable cleanup is pending. +**A model cannot be scheduled on a node that looks free (`no replica slot ... all models busy, cannot evict`):** +- A replica row in `staging` or `loading` holds its slot: slot allocation counts every state except `unloading`. If a worker drops out mid-transfer, that row never reaches `loaded`, and eviction only ever considers `loaded` replicas, so on a node with one replica slot per model the model became unschedulable there. +- The reconciler now reclaims a replica row stuck before serving when no load job is still driving it, and the freed slot is immediately reusable. +- Liveness is decided by the load job's progress heartbeat, not by elapsed time. Staging a large checkpoint legitimately runs for a long time without touching the replica row, so a transfer that is still progressing is never reclaimed however long it takes. +- `Reconciler: reclaimed a replica slot held by a load nobody is driving` names each row reclaimed this way. + **A request fails with `nats: no responders available for request`:** - The chosen worker was not subscribed on the bus when the frontend tried to install the backend on it. A node's status comes from its HTTP heartbeat, which is a separate channel: a worker that stops stays `healthy` until that heartbeat ages out. - The scheduler now checks that a node still answers on the bus before it commits to it, marks one that does not as unhealthy, and picks another. A request should therefore see this only when no reachable node is left.