[pull] master from mudler:master

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
 }

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.

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

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())
+	})
+})