diff --git a/core/services/nodes/reconciler_abandoned_load.go b/core/services/nodes/reconciler_abandoned_load.go index 9b6b7b65430b..f8d5dc702ef3 100644 --- a/core/services/nodes/reconciler_abandoned_load.go +++ b/core/services/nodes/reconciler_abandoned_load.go @@ -35,11 +35,17 @@ var preServingStates = []string{"loading", "staging"} // 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. +// A row is only reclaimed when something proves the load is not progressing: +// either a load job that has failed or stopped heartbeating, or, for a row with +// no job at all, a node that is no longer healthy. +// +// The no-job case has to be conservative. Only the request path creates load +// jobs; the reconciler's own scale-up loads a replica without one. Treating a +// missing job as proof of abandonment would let this sweeper delete a healthy +// reconciler-driven transfer the moment it ran past the grace period, which for +// a multi-gigabyte checkpoint is every time. A healthy node with no job is +// therefore left alone; when the node is gone, nothing can be progressing and +// the row is safe to reclaim. func (rc *ReplicaReconciler) reclaimAbandonedLoads(ctx context.Context) { if rc.db == nil { return @@ -56,7 +62,7 @@ func (rc *ReplicaReconciler) reclaimAbandonedLoads(ctx context.Context) { now := time.Now() for _, row := range stuck { - if rc.loadStillRunning(ctx, row.ModelName, now) { + if !rc.loadAbandoned(ctx, row, now) { continue } if err := rc.registry.RemoveNodeModel(ctx, row.NodeID, row.ModelName, row.ReplicaIndex); err != nil { @@ -70,29 +76,38 @@ func (rc *ReplicaReconciler) reclaimAbandonedLoads(ctx context.Context) { } } -// loadStillRunning reports whether a load job is actively driving this model. +// loadAbandoned reports whether this row's load has demonstrably stopped. // -// 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 { +// Every uncertain case answers false. Leaving a slot held for another pass +// costs one scheduling opportunity; reclaiming a row out from under a live +// transfer restarts a multi-gigabyte load and, on a single-slot node, makes the +// model unschedulable there for as long as the retry loop runs. +func (rc *ReplicaReconciler) loadAbandoned(ctx context.Context, row NodeModel, now time.Time) bool { + job, err := rc.registry.GetLoadJob(ctx, row.ModelName) + switch { + case errors.Is(err, gorm.ErrRecordNotFound), err == nil && job == nil: + // No job: only the request path creates them, so this may be a healthy + // reconciler-driven load. Reclaim only once its node is gone. + return !rc.nodeHealthy(ctx, row.NodeID) + case err != nil: xlog.Warn("Reconciler: cannot read load job, leaving the replica slot held", - "model", modelName, "error", err) - return true - } - if job == nil { + "model", row.ModelName, "error", err) return false + case job.State == LoadJobStateFailed: + return true + default: + return job.IsOrphaned(now) } - if job.State == LoadJobStateFailed { - return false +} + +// nodeHealthy reports whether the row's node is still healthy. An unreadable +// node counts as healthy so a database blip cannot trigger a reclaim. +func (rc *ReplicaReconciler) nodeHealthy(ctx context.Context, nodeID string) bool { + node, err := rc.registry.Get(ctx, nodeID) + if err != nil || node == nil { + xlog.Warn("Reconciler: cannot read node for a stuck replica, leaving the slot held", + "node", nodeID, "error", err) + return true } - return !job.IsOrphaned(now) + return node.Status == StatusHealthy } diff --git a/core/services/nodes/reconciler_abandoned_load_test.go b/core/services/nodes/reconciler_abandoned_load_test.go index 6740b13aafa7..3289fc31bb92 100644 --- a/core/services/nodes/reconciler_abandoned_load_test.go +++ b/core/services/nodes/reconciler_abandoned_load_test.go @@ -82,14 +82,29 @@ var _ = Describe("ReplicaReconciler — abandoned load sweeper", func() { Expect(rowExists("abandoned")).To(BeFalse()) }) - It("reclaims a loading row that has no load job at all", func() { + It("reclaims a jobless row once its node is gone", func() { seedReplica("orphan", "loading", time.Hour) + Expect(registry.MarkUnhealthy(context.Background(), node.ID)).To(Succeed()) rc.reclaimAbandonedLoads(context.Background()) Expect(rowExists("orphan")).To(BeFalse()) }) + // Only the request path creates load jobs. The reconciler's own scale-up + // loads a replica without one, so treating a missing job as abandonment + // deleted healthy transfers the moment they outran the grace period, which + // for a multi-gigabyte checkpoint is every time. That is what made a replica + // appear to hop between nodes instead of finishing anywhere. + It("keeps a jobless row while its node is still healthy", func() { + seedReplica("scaling-up", "staging", time.Hour) + + rc.reclaimAbandonedLoads(context.Background()) + + Expect(rowExists("scaling-up")).To(BeTrue(), + "a reconciler-driven load has no job row and must not be reclaimed for it") + }) + 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. diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 14b19092e5f7..08cf2fd78659 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -1129,7 +1129,7 @@ func (r *SmartRouter) scheduleNewModel(ctx context.Context, backendType, modelID // 4. Preemptive eviction: if no suitable node found, evict the LRU model with zero in-flight if node == nil { - evictedNode, evictErr := r.evictLRUAndFreeNode(ctx) + evictedNode, evictErr := r.evictLRUAndFreeNodeFrom(ctx, candidateNodeIDs) if evictErr != nil { if errors.Is(evictErr, ErrEvictionBusy) { return nil, "", 0, fmt.Errorf("no healthy nodes available: %w", evictErr) @@ -1153,7 +1153,7 @@ func (r *SmartRouter) scheduleNewModel(ctx context.Context, backendType, modelID // it can race with another concurrent scheduler. xlog.Warn("Chosen node has no free replica slot, evicting LRU", "node", node.Name, "model", modelID, "max_slots", maxSlots) - evictedNode, evictErr := r.evictLRUAndFreeNode(ctx) + evictedNode, evictErr := r.evictLRUAndFreeNodeFrom(ctx, candidateNodeIDs) if evictErr != nil { return nil, "", 0, fmt.Errorf("no replica slot on %s and eviction failed: %w", node.Name, evictErr) } @@ -1979,7 +1979,23 @@ var ErrEvictionBusy = errors.New("all models busy, cannot evict") // Uses SELECT FOR UPDATE inside a transaction to prevent two frontends from // simultaneously picking the same eviction target. The NodeModel row is deleted // inside the transaction; the NATS unload command is sent after commit. +// evictLRUAndFreeNode evicts across every healthy node. Callers that hold a +// candidate set must use evictLRUAndFreeNodeFrom instead. func (r *SmartRouter) evictLRUAndFreeNode(ctx context.Context) (*BackendNode, error) { + return r.evictLRUAndFreeNodeFrom(ctx, nil) +} + +// evictLRUAndFreeNodeFrom evicts the least-recently-used idle model from one of +// candidateNodeIDs, or from any healthy node when the set is nil. +// +// Restricting eviction to the candidate set matters whenever the model being +// scheduled has a node selector. Evicting globally freed a slot on a node the +// selector forbids, so the model was then placed there anyway, on hardware it +// was explicitly pinned away from, and an unrelated model was dropped to make +// the room. On a cluster where the selector-matching node was momentarily +// unavailable this repeated, and the evicted model appeared to bounce between +// nodes. +func (r *SmartRouter) evictLRUAndFreeNodeFrom(ctx context.Context, candidateNodeIDs []string) (*BackendNode, error) { const maxEvictionRetries = 5 const evictionRetryInterval = 500 * time.Millisecond @@ -1991,7 +2007,7 @@ func (r *SmartRouter) evictLRUAndFreeNode(ctx context.Context) (*BackendNode, er var lru NodeModel err := r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { // Lock the row so no other frontend can evict the same model - if err := currentModelRevision(tx.Clauses(clause.Locking{Strength: "UPDATE"})). + q := currentModelRevision(tx.Clauses(clause.Locking{Strength: "UPDATE"})). Joins("JOIN backend_nodes ON backend_nodes.id = node_models.node_id"). Where(`node_models.in_flight = 0 AND node_models.state = ? AND backend_nodes.status = ? AND ( @@ -2000,7 +2016,11 @@ func (r *SmartRouter) evictLRUAndFreeNode(ctx context.Context) (*BackendNode, er AND (NOT EXISTS (SELECT 1 FROM model_config_states mcs2 WHERE mcs2.model_name = nm2.model_name) OR nm2.config_revision = (SELECT mcs3.config_revision FROM model_config_states mcs3 WHERE mcs3.model_name = nm2.model_name))) > COALESCE((SELECT sc2.min_replicas FROM model_scheduling_configs sc2 WHERE sc2.model_name = node_models.model_name), 1) - )`, "loaded", StatusHealthy). + )`, "loaded", StatusHealthy) + if len(candidateNodeIDs) > 0 { + q = q.Where("node_models.node_id IN ?", candidateNodeIDs) + } + if err := q. Order("node_models.last_used ASC"). First(&lru).Error; err != nil { return err diff --git a/core/services/nodes/router_eviction_selector_test.go b/core/services/nodes/router_eviction_selector_test.go new file mode 100644 index 000000000000..8d0caaffbef5 --- /dev/null +++ b/core/services/nodes/router_eviction_selector_test.go @@ -0,0 +1,110 @@ +package nodes + +import ( + "context" + "fmt" + "runtime" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/testutil" +) + +// When no node the selector allows has a free slot, scheduling falls back to +// evicting the globally least-recently-used model. That eviction knew nothing +// about the selector, so a model pinned to one class of hardware would evict an +// unrelated model from a node it is not allowed to run on, and then be placed +// there. Two models lose: the pinned one runs on the wrong hardware, and the +// evicted one is dropped for nothing and has to reload elsewhere. +var _ = Describe("Eviction under a node selector", func() { + var ( + db *gorm.DB + registry *NodeRegistry + router *SmartRouter + ctx context.Context + ) + + 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 = NewSmartRouter(registry, SmartRouterOptions{DB: db}) + ctx = context.Background() + }) + + register := func(name string) *BackendNode { + node := &BackendNode{Name: name, NodeType: NodeTypeBackend, Address: name + ":50051"} + Expect(registry.Register(ctx, node, true)).To(Succeed()) + fetched, err := registry.GetByName(ctx, name) + Expect(err).ToNot(HaveOccurred()) + return fetched + } + + rowID := 0 + seed := func(node *BackendNode, model string, idleFor time.Duration, inFlight int) { + rowID++ + Expect(db.Create(&NodeModel{ + ID: fmt.Sprintf("row-%d", rowID), NodeID: node.ID, ModelName: model, + Address: node.Address, State: "loaded", InFlight: inFlight, + LastUsed: time.Now().Add(-idleFor), UpdatedAt: time.Now(), + }).Error).To(Succeed()) + } + seedLoaded := func(node *BackendNode, model string, idleFor time.Duration) { + seed(node, model, idleFor, 0) + } + + rowExists := func(model string) bool { + var n int64 + Expect(db.Model(&NodeModel{}).Where("model_name = ?", model).Count(&n).Error).To(Succeed()) + return n > 0 + } + + It("does not evict from a node the selector excludes", func() { + allowed := register("allowed-node") + excluded := register("excluded-node") + // The only eviction candidate sits on the excluded node and is the + // global LRU, so an unconstrained eviction would take it. + seedLoaded(excluded, "innocent-bystander", time.Hour) + // In-flight, so it is not an eviction candidate: the allowed node has + // nothing that can be freed. + seed(allowed, "busy-here", time.Minute, 1) + + _, err := router.evictLRUAndFreeNodeFrom(ctx, []string{allowed.ID}) + + Expect(err).To(HaveOccurred(), "no eviction candidate exists on an allowed node") + Expect(rowExists("innocent-bystander")).To(BeTrue(), + "a model on a node the selector excludes must not be evicted to make room") + }) + + It("evicts the LRU among the allowed nodes only", func() { + allowed := register("allowed-node") + excluded := register("excluded-node") + seedLoaded(excluded, "older-elsewhere", 2*time.Hour) + seedLoaded(allowed, "newer-but-allowed", time.Hour) + + node, err := router.evictLRUAndFreeNodeFrom(ctx, []string{allowed.ID}) + + Expect(err).ToNot(HaveOccurred()) + Expect(node.ID).To(Equal(allowed.ID)) + Expect(rowExists("newer-but-allowed")).To(BeFalse()) + Expect(rowExists("older-elsewhere")).To(BeTrue()) + }) + + It("keeps evicting globally when the model has no selector", func() { + a := register("node-a") + seedLoaded(a, "anything", time.Hour) + + node, err := router.evictLRUAndFreeNodeFrom(ctx, nil) + + Expect(err).ToNot(HaveOccurred()) + Expect(node.ID).To(Equal(a.ID)) + Expect(rowExists("anything")).To(BeFalse()) + }) +})