diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index 14b19092e..08cf2fd78 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 000000000..8d0caaffb --- /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()) + }) +})