diff --git a/core/services/nodes/router.go b/core/services/nodes/router.go index dbcd23bb8eb7..a8bf913128ca 100644 --- a/core/services/nodes/router.go +++ b/core/services/nodes/router.go @@ -639,6 +639,31 @@ func (r *SmartRouter) tryWarmPath(ctx context.Context, att *routeAttempt) *Route } replicaIdx := nm.ReplicaIndex + // Load-time options cannot be changed on a running backend. Retire an idle + // replica whose context no longer matches the model config so the cold path + // recreates it with the options supplied by this request. Older registry + // rows without stored options remain usable for rolling-upgrade safety. + if nm.InFlight == 0 && !sameLoadContext(nm.ModelOptsBlob, att.modelOpts) && r.unloader != nil { + processKey := model.BackendProcessKey(att.trackingKey, replicaIdx) + if err := r.unloader.StopBackend(node.ID, processKey); err != nil { + xlog.Warn("Failed to retire model replica with stale context", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + } else { + if err := r.registry.DecrementInFlight(ctx, node.ID, att.trackingKey, replicaIdx); err != nil { + xlog.Warn("Failed to release stale-context routing reservation", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + } + if err := r.registry.RemoveNodeModel(ctx, node.ID, att.trackingKey, replicaIdx); err != nil { + xlog.Warn("Failed to remove model replica with stale context", + "node", node.ID, "model", att.trackingKey, "replica", replicaIdx, "error", err) + } + xlog.Info("Retired model replica after context size changed", + "node", node.Name, "model", att.trackingKey, "replica", replicaIdx, + "context", att.modelOpts.ContextSize) + return nil + } + } + // Verify the backend process is still alive via gRPC health check if !r.probeHealth(ctx, node, modelAddr) { // Stale — roll back the increment, remove the specific replica row, fall through @@ -677,6 +702,18 @@ func (r *SmartRouter) tryWarmPath(ctx context.Context, att *routeAttempt) *Route return r.newRouteResult(node, att.trackingKey, replicaIdx, grpcClient, tracked) } +func sameLoadContext(stored []byte, current *pb.ModelOptions) bool { + if len(stored) == 0 || current == nil { + return true + } + var loaded pb.ModelOptions + if err := proto.Unmarshal(stored, &loaded); err != nil { + xlog.Warn("Failed to read stored model options while checking context", "error", err) + return true + } + return loaded.ContextSize == current.ContextSize +} + // coldLoad schedules the model onto a node and loads it, returning a route to // the replica it landed on. initialInFlight reserves the slot for the calling // request; the job runner passes 0 because it is loading on nobody's behalf. diff --git a/core/services/nodes/router_test.go b/core/services/nodes/router_test.go index 6e1d86d47929..27f0e63928a1 100644 --- a/core/services/nodes/router_test.go +++ b/core/services/nodes/router_test.go @@ -18,6 +18,7 @@ import ( grpc "github.com/mudler/LocalAI/pkg/grpc" pb "github.com/mudler/LocalAI/pkg/grpc/proto" ggrpc "google.golang.org/grpc" + "google.golang.org/protobuf/proto" "gorm.io/gorm" ) @@ -520,8 +521,10 @@ var _ = Describe("SmartRouter", func() { Context("model already loaded on a healthy node", func() { It("returns the client and a release function", func() { + loadedOpts, err := proto.Marshal(&pb.ModelOptions{ContextSize: 32768}) + Expect(err).ToNot(HaveOccurred()) node := &BackendNode{ID: "n1", Name: "node-1", Address: "10.0.0.1:50051"} - nm := &NodeModel{NodeID: "n1", ModelName: "my-model", Address: "10.0.0.1:9001"} + nm := &NodeModel{NodeID: "n1", ModelName: "my-model", Address: "10.0.0.1:9001", ModelOptsBlob: loadedOpts} reg.findAndLockNode = node reg.findAndLockNM = nm backend.healthResult = true @@ -531,7 +534,8 @@ var _ = Describe("SmartRouter", func() { ClientFactory: factory, }) - result, err := router.Route(context.Background(), "my-model", "models/my-model.gguf", "llama-cpp", nil, false) + result, err := router.Route(context.Background(), "my-model", "models/my-model.gguf", "llama-cpp", + &pb.ModelOptions{ContextSize: 32768}, false) Expect(err).ToNot(HaveOccurred()) Expect(result).ToNot(BeNil()) Expect(result.Node.ID).To(Equal("n1")) @@ -551,6 +555,41 @@ var _ = Describe("SmartRouter", func() { result.Release() Expect(reg.decrementCalls).To(HaveLen(1)) }) + + It("retires a replica loaded with a different context size", func() { + loadedOpts, err := proto.Marshal(&pb.ModelOptions{ContextSize: 8192}) + Expect(err).ToNot(HaveOccurred()) + + node := &BackendNode{ID: "n1", Name: "node-1", Address: "10.0.0.1:50051"} + nm := &NodeModel{ + NodeID: "n1", + ModelName: "my-model", + ReplicaIndex: 2, + Address: "10.0.0.1:9001", + ModelOptsBlob: loadedOpts, + } + reg.findAndLockNode = node + reg.findAndLockNM = nm + backend.healthResult = true + + router := NewSmartRouter(reg, SmartRouterOptions{ + Unloader: unloader, + ClientFactory: factory, + }) + + result := router.tryWarmPath(context.Background(), &routeAttempt{ + trackingKey: "my-model", + modelName: "models/my-model.gguf", + backendType: "llama-cpp", + modelOpts: &pb.ModelOptions{ContextSize: 32768}, + }) + + Expect(result).To(BeNil()) + Expect(reg.decrementCalls).To(ConsistOf("n1:my-model")) + Expect(reg.removeCalls).To(ConsistOf("n1:my-model")) + Expect(unloader.stopCalls).To(ConsistOf("n1:my-model#2")) + Expect(reg.touchCalls).To(BeEmpty()) + }) }) Context("model not loaded, falls through to scheduling", func() {