From d90f726345483e03d88f0cfd7f75cd8e11dda27b Mon Sep 17 00:00:00 2001 From: localai-org-maint-bot <306269227+localai-org-maint-bot@users.noreply.github.com> Date: Mon, 17 Aug 2026 18:06:02 +0000 Subject: [PATCH] fix(distributed): reload changed model context Distributed routing reused a warm worker by model ID even when the model context changed. The worker kept its original context and truncated requests at that old limit. Retire an idle replica when its stored context differs. The existing cold-load path then starts that replica with the current model options. Assisted-by: Codex:gpt-5 --- core/services/nodes/router.go | 37 +++++++++++++++++++++++++ core/services/nodes/router_test.go | 43 ++++++++++++++++++++++++++++-- 2 files changed, 78 insertions(+), 2 deletions(-) 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() {