fix(distributed): reload changed model context

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.

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() {