From 97390f7754f3812d39bf22fa017ab70e21e48825 Mon Sep 17 00:00:00 2001 From: hcharbonnier Date: Fri, 23 Oct 2026 16:09:48 +0200 Subject: [PATCH 1/2] fix(watchdog): keep tracking a model that an eviction could not stop An eviction untracks the backend before it stops it. When the stop failed, the error was only logged: the backend kept running and kept its VRAM, but the watchdog had already forgotten it. Group and LRU enforcement then saw no conflict, and the next load started next to it and ran out of memory. Stopping a model takes its exclusive lifecycle lock, which a model that is still loading holds for the whole load. The stop then times out ("waiting to shut down model: context deadline exceeded"), and the watchdog no longer knows the model is there. The tracking state is now captured before the backend is untracked, and put back when the stop fails. This applies to every eviction path that untracks first: group, LRU, idle timeout and busy timeout. A model that is already gone (ErrModelNotFound) still counts as stopped, and a replacement registered on the same address in the meantime is left alone. EvictedCount only counts backends that really stopped, and NeedMore is set when one did not, so the loader keeps retrying instead of treating the room as made. The idle and busy checks retry the stop on their next tick. The state is only put back while the loader still holds the model at that address (ModelLoader.HasModelAtAddress). deleteProcess drops a remote-only model, or the local half of a local+remote placement, before it reports a failed remote unload; tracking that model again would leave a phantom entry counted as loaded and as a group conflict. The memory reclaimer applies the same rule and untracks a model the loader already dropped. The untrack-before-stop ordering dates from #7535, so this is not a regression of 4.12.0. Assisted-by: Claude:claude-sonnet-5.5 go vet Signed-off-by: Hugues CHARBONNIER --- pkg/model/initializers_group_load_test.go | 159 +++++++++++++++++ pkg/model/loader.go | 13 ++ pkg/model/watchdog.go | 171 ++++++++++++++++-- pkg/model/watchdog_lifecycle_test.go | 72 ++++++++ pkg/model/watchdog_test.go | 202 ++++++++++++++++++++++ 5 files changed, 605 insertions(+), 12 deletions(-) create mode 100644 pkg/model/initializers_group_load_test.go diff --git a/pkg/model/initializers_group_load_test.go b/pkg/model/initializers_group_load_test.go new file mode 100644 index 000000000000..98a0b9aee5e0 --- /dev/null +++ b/pkg/model/initializers_group_load_test.go @@ -0,0 +1,159 @@ +package model + +import ( + "context" + "errors" + "sync" + "time" + + "github.com/mudler/LocalAI/pkg/system" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// storeEvicter drops a model from the loader store the way a real shutdown +// does, and can be told to fail like a stop that cannot get the model's +// lifecycle lock. +type storeEvicter struct { + ml *ModelLoader + mu sync.Mutex + evicted []string + failures map[string]error +} + +func (s *storeEvicter) setFailure(name string, err error) { + s.mu.Lock() + defer s.mu.Unlock() + if s.failures == nil { + s.failures = map[string]error{} + } + if err == nil { + delete(s.failures, name) + return + } + s.failures[name] = err +} + +func (s *storeEvicter) ShutdownModel(name string) error { + s.mu.Lock() + s.evicted = append(s.evicted, name) + failure := s.failures[name] + s.mu.Unlock() + if failure != nil { + return failure + } + s.ml.mu.Lock() + s.ml.store.Delete(name) + s.ml.mu.Unlock() + return nil +} + +func (s *storeEvicter) ShutdownModelForce(name string) error { return s.ShutdownModel(name) } + +func (s *storeEvicter) names() []string { + s.mu.Lock() + defer s.mu.Unlock() + return append([]string(nil), s.evicted...) +} + +// Load must evict a resident model of the same concurrency group BEFORE the +// requested model is started, and must not forget a resident model it failed +// to stop. This drives the real Load path. +var _ = Describe("Load with concurrency groups", func() { + const ( + gemma = "gemma-4-26b-a4b-it" + nemotron = "NVIDIA-Nemotron-3.5-Lightning-30B-A3B-UD-Q4_K_XL" + address = "127.0.0.1:65001" + ) + + var ( + ml *ModelLoader + wd *WatchDog + evic *storeEvicter + ) + + load := func() error { + _, err := ml.Load( + WithContext(context.Background()), + WithModelID(nemotron), + WithModel("nemotron.gguf"), + WithBackendString("fakebackend"), + WithGRPCAttempts(1), + WithGRPCAttemptsDelay(0), + ) + return err + } + + BeforeEach(func() { + systemState, err := system.GetSystemState(system.WithModelPath(GinkgoT().TempDir())) + Expect(err).ToNot(HaveOccurred()) + ml = NewModelLoader(systemState) + ml.SetLRUEvictionRetrySettings(3, 5*time.Millisecond) + ml.SetExternalBackend("fakebackend", "127.0.0.1:1") + + evic = &storeEvicter{ml: ml} + wd = NewWatchDog(WithProcessManager(evic), WithLRULimit(2)) + wd.ReplaceModelGroups(map[string][]string{ + gemma: {"vram-heavy"}, + nemotron: {"vram-heavy"}, + }) + ml.SetWatchDog(wd) + + // gemma is resident and idle: in the loader store and in the watchdog. + resident := NewModel(gemma, address, nil) + resident.MarkHealthy() + ml.mu.Lock() + ml.store.Set(gemma, resident) + ml.mu.Unlock() + wd.AddAddressModelMap(address, gemma) + wd.TrackRequest(address)() + }) + + It("evicts the resident model of the group before loading the requested one", func() { + // The load itself cannot succeed (no real backend); what matters is that + // the conflicting model was evicted first. + _ = load() + + Expect(evic.names()).To(ConsistOf(gemma), + "gemma shares the vram-heavy group and must be evicted before nemotron loads") + }) + + It("leaves a resident model alone when it is not in the group", func() { + wd.ReplaceModelGroups(map[string][]string{nemotron: {"vram-heavy"}}) + + _ = load() + + Expect(evic.names()).To(BeEmpty()) + }) + + It("keeps tracking a resident model it failed to stop", func() { + // The incident: stopping gemma timed out, gemma kept its VRAM, and it + // had been forgotten, so the next load saw no conflict and ran out of + // memory. + evic.setFailure(gemma, errors.New("waiting to shut down model: context deadline exceeded")) + + _ = load() + + Expect(wd.GetState().AddressModelMap).To(HaveKeyWithValue(address, gemma), + "gemma still runs, so the watchdog must still know it") + }) + + It("stops a resident model that first refused to stop on a later attempt", func() { + evic.setFailure(gemma, errors.New("waiting to shut down model: context deadline exceeded")) + ml.SetLRUEvictionRetrySettings(200, 10*time.Millisecond) + + go func() { + // Release the lifecycle lock once a first attempt has failed, rather + // than after a wall-clock delay that a slow CI could outrun. + defer GinkgoRecover() + Eventually(evic.names, "5s", "1ms").ShouldNot(BeEmpty()) + evic.setFailure(gemma, nil) + }() + + _ = load() + + Expect(wd.GetState().AddressModelMap).ToNot(HaveKey(address), + "gemma was finally stopped and untracked") + }) +}) diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 1db9a9cd74a8..f823f9de89e2 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -666,6 +666,19 @@ func (ml *ModelLoader) ShutdownModelAtAddress(modelName, address string, force b return err } +// HasModelAtAddress reports whether the loader store still holds modelName at +// address. The watchdog uses it after a failed shutdown to tell a backend that +// is still running (store entry intact, safe to track again) from one that +// deleteProcess already tore down before it returned an error, such as a failed +// remote unload or the remote half of a local+remote placement. +func (ml *ModelLoader) HasModelAtAddress(modelName, address string) bool { + ml.mu.Lock() + store := ml.store + ml.mu.Unlock() + m, ok := store.Get(modelName) + return ok && m.address == address +} + // ShutdownModelContext is the cancellation-aware lifecycle primitive used by // both graceful and forced shutdown wrappers. func (ml *ModelLoader) ShutdownModelContext(ctx context.Context, modelName string, force bool) error { diff --git a/pkg/model/watchdog.go b/pkg/model/watchdog.go index 977eca1ff565..5e4852f6767e 100644 --- a/pkg/model/watchdog.go +++ b/pkg/model/watchdog.go @@ -471,6 +471,98 @@ type evictionTarget struct { address string model string wasBusy bool + + // tracked is what untrack removed for this address, kept so a shutdown that + // fails can put it back. Every eviction path captures it before untrack; + // nil only when the address was not tracked. + tracked *trackedState +} + +// trackedState is a snapshot of everything the watchdog knows about one backend +// address. An eviction untracks the backend before it is stopped, so a stop that +// fails (typically a lifecycle lock held by a long model load) would otherwise +// leave a backend that is still running and still holding its VRAM invisible to +// the group and LRU checks: the next load sees no conflict and runs out of memory. +type trackedState struct { + model string + process *process.Process + busyTime time.Time + hasBusy bool + inFlight int + requestStarts map[uint64]time.Time + legacyRequests []uint64 + idleTime time.Time + hasIdle bool + lastUsed time.Time + hasLastUsed bool + size int64 + hasSize bool +} + +// snapshotLocked captures the tracking state of an address. Must be called with +// wd.Lock() held, before untrack. +func (wd *WatchDog) snapshotLocked(address string) *trackedState { + model, ok := wd.addressModelMap[address] + if !ok { + return nil + } + s := &trackedState{model: model, process: wd.addressMap[address], inFlight: wd.inFlight[address]} + s.busyTime, s.hasBusy = wd.busyTime[address] + s.idleTime, s.hasIdle = wd.idleTime[address] + s.lastUsed, s.hasLastUsed = wd.lastUsed[address] + s.size, s.hasSize = wd.modelSizes[model] + if starts := wd.requestStarts[address]; starts != nil { + s.requestStarts = make(map[uint64]time.Time, len(starts)) + for id, t := range starts { + s.requestStarts[id] = t + } + } + s.legacyRequests = slices.Clone(wd.legacyRequests[address]) + return s +} + +// restoreTracked puts a snapshot back after the backend it describes could not +// be stopped. It does nothing when the address was re-registered in the +// meantime: a replacement backend already owns the address and must win. +func (wd *WatchDog) restoreTracked(address string, s *trackedState) { + if s == nil { + return + } + wd.Lock() + defer wd.Unlock() + if _, taken := wd.addressModelMap[address]; taken { + return + } + wd.addressModelMap[address] = s.model + if s.process != nil { + wd.addressMap[address] = s.process + } + if s.hasBusy { + wd.busyTime[address] = s.busyTime + } + if s.inFlight > 0 { + wd.inFlight[address] = s.inFlight + } + if len(s.requestStarts) > 0 { + wd.requestStarts[address] = s.requestStarts + } + if len(s.legacyRequests) > 0 { + wd.legacyRequests[address] = s.legacyRequests + } + if s.hasIdle { + wd.idleTime[address] = s.idleTime + } + if s.hasLastUsed { + wd.lastUsed[address] = s.lastUsed + } + if s.hasSize { + // untrack dropped the estimate only when no other address served this + // model; never overwrite one a replacement has registered since. + if _, ok := wd.modelSizes[s.model]; !ok { + wd.modelSizes[s.model] = s.size + } + } + xlog.Warn("[WatchDog] model could not be stopped, tracking it again", "model", s.model, "address", address) } // EnforceLRULimitResult contains the result of LRU enforcement @@ -544,14 +636,19 @@ func (wd *WatchDog) EnforceLRULimit(pendingLoads int) EnforceLRULimitResult { // Now shutdown models without holding the watchdog lock to prevent // deadlock. Busy targets go through the force path so the loader doesn't // wait on their stuck in-flight call (which would re-deadlock here). - wd.shutdownEvicted(modelsToShutdown, "LRU eviction") + stopped := wd.shutdownEvictedStopped(modelsToShutdown, "LRU eviction") + if stopped < len(modelsToShutdown) { + // A backend that would not stop is still running and still holds its + // memory, so the room asked for does not exist yet. + needMore = true + } if needMore { - xlog.Warn("[WatchDog] LRU eviction incomplete", "evicted", len(modelsToShutdown), "needed", modelsToEvict, "skippedBusy", skippedBusyCount, "reason", "some models are busy with active API calls") + xlog.Warn("[WatchDog] LRU eviction incomplete", "evicted", stopped, "needed", modelsToEvict, "skippedBusy", skippedBusyCount, "reason", "some models are busy with active API calls or could not be stopped") } return EnforceLRULimitResult{ - EvictedCount: len(modelsToShutdown), + EvictedCount: stopped, NeedMore: needMore, } } @@ -576,7 +673,7 @@ func (wd *WatchDog) collectEvictionsLocked(candidates []modelUsageInfo, maxToEvi continue } xlog.Info("[WatchDog] evicting model", "model", m.model, "busy", isBusy) - evicted = append(evicted, evictionTarget{address: m.address, model: m.model, wasBusy: isBusy}) + evicted = append(evicted, evictionTarget{address: m.address, model: m.model, wasBusy: isBusy, tracked: wd.snapshotLocked(m.address)}) wd.untrack(m.address) } return evicted, skippedBusy @@ -586,13 +683,48 @@ func (wd *WatchDog) collectEvictionsLocked(candidates []modelUsageInfo, maxToEvi // that were busy at selection time so a stuck in-flight call does not consume // the graceful shutdown deadline. func (wd *WatchDog) shutdownEvicted(targets []evictionTarget, label string) { + wd.shutdownEvictedStopped(targets, label) +} + +// shutdownEvictedStopped is shutdownEvicted that also reports how many models +// were really stopped. A model that failed to stop is still running, so it is +// tracked again and not counted: callers use the count to decide whether the +// room they asked for exists, and counting a backend that still holds its VRAM +// would let the next load start next to it. +func (wd *WatchDog) shutdownEvictedStopped(targets []evictionTarget, label string) (stopped int) { for _, t := range targets { err := wd.shutdownTarget(t) - if err != nil { + if err != nil && !errors.Is(err, ErrModelNotFound) { xlog.Error("[WatchDog] error shutting down model during "+label, "error", err, "model", t.model, "busy", t.wasBusy) + if wd.stillResident(t) { + wd.restoreTracked(t.address, t.tracked) + } else { + // The loader already dropped the model before the error was + // returned (a failed remote unload, or the remote half of a + // local+remote placement). Tracking it again would leave a + // phantom entry that inflates LRU counts and group conflicts. + xlog.Warn("[WatchDog] model was already removed by the loader despite the shutdown error, not tracking it again", "model", t.model, "address", t.address) + } + continue } + stopped++ xlog.Debug("[WatchDog] "+label+" complete", "model", t.model, "busy", t.wasBusy) } + return stopped +} + +// stillResident reports whether the backend an eviction target describes is +// still held by the loader. A process manager that cannot answer is assumed to +// still hold it: after a failed stop, tracking a backend that may still run is +// the safe default, since forgetting it is what lets the next load run out of +// memory. +func (wd *WatchDog) stillResident(t evictionTarget) bool { + if pm, ok := wd.pm.(interface { + HasModelAtAddress(string, string) bool + }); ok { + return pm.HasModelAtAddress(t.model, t.address) + } + return true } // shutdownTarget retains the address selected under the watchdog lock. The @@ -659,14 +791,19 @@ func (wd *WatchDog) EnforceGroupExclusivity(requestedModel string) EnforceLRULim needMore := len(modelsToShutdown) < len(conflicts) wd.Unlock() - wd.shutdownEvicted(modelsToShutdown, "group eviction") + stopped := wd.shutdownEvictedStopped(modelsToShutdown, "group eviction") + if stopped < len(modelsToShutdown) { + // A conflict that would not stop is still running and still holds its + // memory: report it as unresolved so the loader keeps waiting. + needMore = true + } if needMore { - xlog.Warn("[WatchDog] Group eviction incomplete", "requested", requestedModel, "evicted", len(modelsToShutdown), "needed", len(conflicts), "skippedBusy", skippedBusyCount, "reason", "some conflicts are busy or pinned") + xlog.Warn("[WatchDog] Group eviction incomplete", "requested", requestedModel, "evicted", stopped, "needed", len(conflicts), "skippedBusy", skippedBusyCount, "reason", "some conflicts are busy, pinned or could not be stopped") } return EnforceLRULimitResult{ - EvictedCount: len(modelsToShutdown), + EvictedCount: stopped, NeedMore: needMore, } } @@ -735,7 +872,7 @@ func (wd *WatchDog) checkIdle() { continue } xlog.Warn("[WatchDog] Address is idle for too long, killing it", "address", address, "model", model) - modelsToShutdown = append(modelsToShutdown, evictionTarget{model: model, address: address}) + modelsToShutdown = append(modelsToShutdown, evictionTarget{model: model, address: address, tracked: wd.snapshotLocked(address)}) } else { xlog.Warn("[WatchDog] Address unresolvable", "address", address) } @@ -760,7 +897,7 @@ func (wd *WatchDog) checkBusy() { model, ok := wd.addressModelMap[address] if ok { xlog.Warn("[WatchDog] Model is busy for too long, killing it", "model", model) - modelsToShutdown = append(modelsToShutdown, evictionTarget{model: model, address: address, wasBusy: true}) + modelsToShutdown = append(modelsToShutdown, evictionTarget{model: model, address: address, wasBusy: true, tracked: wd.snapshotLocked(address)}) } else { xlog.Warn("[WatchDog] Address unresolvable", "address", address) } @@ -891,9 +1028,19 @@ func (wd *WatchDog) evictLRUModel() { wd.Unlock() - // Shutdown the model - if err := wd.shutdownTarget(evictionTarget{model: lruModel.model, address: lruModel.address, wasBusy: wasBusy}); err != nil && !errors.Is(err, ErrModelNotFound) { + // Shutdown the model. This path never untracks before the stop, so a + // failed stop of a backend that still runs keeps it tracked on its own. + target := evictionTarget{model: lruModel.model, address: lruModel.address, wasBusy: wasBusy} + if err := wd.shutdownTarget(target); err != nil && !errors.Is(err, ErrModelNotFound) { xlog.Error("[WatchDog] error shutting down model during memory reclamation", "error", err, "model", lruModel.model) + if !wd.stillResident(target) { + // Same rule as shutdownEvictedStopped: the loader dropped the model + // before returning the error, so keeping the entry would leave a + // phantom that is counted as loaded. + wd.Lock() + wd.untrack(lruModel.address) + wd.Unlock() + } } else { // Untrack the model wd.Lock() diff --git a/pkg/model/watchdog_lifecycle_test.go b/pkg/model/watchdog_lifecycle_test.go index 50e4e84d3030..cf3141c2d96a 100644 --- a/pkg/model/watchdog_lifecycle_test.go +++ b/pkg/model/watchdog_lifecycle_test.go @@ -4,6 +4,7 @@ package model import ( "context" + "errors" "os" "path/filepath" "time" @@ -20,6 +21,24 @@ type watchdogLifecycleBackend struct{ grpc.Backend } func (*watchdogLifecycleBackend) IsBusy() bool { return false } func (*watchdogLifecycleBackend) Free(context.Context) error { return nil } +// lockTimeoutManager behaves like a ModelLoader whose stop cannot get the +// model's lifecycle lock: it fails without touching the store. It embeds the +// loader so HasModelAtAddress still reports real residency. +type lockTimeoutManager struct{ *ModelLoader } + +func (lockTimeoutManager) ShutdownModel(string) error { return errLockTimeout } +func (lockTimeoutManager) ShutdownModelForce(string) error { return errLockTimeout } +func (lockTimeoutManager) ShutdownModelAtAddress(string, string, bool) error { + return errLockTimeout +} + +var errLockTimeout = errors.New("waiting to shut down model: context deadline exceeded") + +// failingUnloader fails every remote unload, like an unreachable worker. +type failingUnloader struct{ err error } + +func (u failingUnloader) UnloadRemoteModel(string) error { return u.err } + var _ = Describe("Watchdog backend lifecycle", func() { var loader *ModelLoader var wd *WatchDog @@ -99,6 +118,59 @@ var _ = Describe("Watchdog backend lifecycle", func() { Entry("forced", true), ) + Describe("a failed stop", func() { + // deleteProcess drops a remote-only model from the store and the + // watchdog before it reports that the remote unload failed, so the + // watchdog must not track it again: nothing local holds it anymore. + It("does not track a remote-only model the loader already dropped", func() { + loader.SetRemoteUnloader(failingUnloader{err: errors.New("worker unreachable")}) + loader.store.Set("model", NewModelWithClient("model", "remote", &watchdogLifecycleBackend{})) + wd.AddAddressModelMap("remote", "model") + wd.TrackRequest("remote")() + + targets, _ := wd.collectEvictionsLocked([]modelUsageInfo{{model: "model", address: "remote"}}, 1, false) + stopped := wd.shutdownEvictedStopped(targets, "test") + + Expect(stopped).To(Equal(0), "the unload failed, it is not counted as stopped") + _, ok := loader.store.Get("model") + Expect(ok).To(BeFalse()) + Expect(wd.GetState().AddressModelMap).To(BeEmpty(), "no phantom entry for a model the loader dropped") + }) + + It("tracks a model again when the loader still holds it", func() { + loader.store.Set("model", NewModelWithClient("model", "current", &watchdogLifecycleBackend{})) + wd.AddAddressModelMap("current", "model") + wd.TrackRequest("current")() + + // A stop that cannot get the model's lifecycle lock leaves the store + // entry intact. The real loader would wait out its 30s deadline, so + // the wrapper returns that failure immediately. + wd.pm = lockTimeoutManager{ModelLoader: loader} + + targets, _ := wd.collectEvictionsLocked([]modelUsageInfo{{model: "model", address: "current"}}, 1, false) + Expect(wd.GetState().AddressModelMap).To(BeEmpty(), "untracked before the stop") + + stopped := wd.shutdownEvictedStopped(targets, "test") + + Expect(stopped).To(Equal(0)) + Expect(wd.GetState().AddressModelMap).To(HaveKeyWithValue("current", "model"), + "the loader still holds the backend, so it may still run") + }) + + It("memory reclaimer does not keep a remote-only model the loader dropped", func() { + loader.SetRemoteUnloader(failingUnloader{err: errors.New("worker unreachable")}) + loader.store.Set("model", NewModelWithClient("model", "remote", &watchdogLifecycleBackend{})) + wd.AddAddressModelMap("remote", "model") + wd.TrackRequest("remote")() + + wd.evictLRUModel() + + _, ok := loader.store.Get("model") + Expect(ok).To(BeFalse()) + Expect(wd.GetState().AddressModelMap).To(BeEmpty()) + }) + }) + It("keeps replacement tracking when an old process cleanup arrives late", func() { oldProcess := process.New() newProcess := process.New() diff --git a/pkg/model/watchdog_test.go b/pkg/model/watchdog_test.go index 36d9ceed2709..8de1181873d4 100644 --- a/pkg/model/watchdog_test.go +++ b/pkg/model/watchdog_test.go @@ -1,6 +1,7 @@ package model_test import ( + "errors" "sync" "time" @@ -47,6 +48,48 @@ func (m *mockProcessManager) ShutdownModelForce(modelName string) error { return nil } +// setShutdownError and clearShutdownError are safe to call while a running +// watchdog is shutting models down. +func (m *mockProcessManager) setShutdownError(modelName string, err error) { + m.mu.Lock() + defer m.mu.Unlock() + m.shutdownErrors[modelName] = err +} + +func (m *mockProcessManager) clearShutdownError(modelName string) { + m.mu.Lock() + defer m.mu.Unlock() + delete(m.shutdownErrors, modelName) +} + +// residencyProcessManager also answers whether the loader still holds a model, +// like ModelLoader.HasModelAtAddress. It lets a test model a shutdown that +// returns an error after the loader already dropped the model. +type residencyProcessManager struct { + *mockProcessManager + resMu sync.Mutex + resident map[string]bool // keyed by model+"@"+address +} + +func newResidencyProcessManager() *residencyProcessManager { + return &residencyProcessManager{ + mockProcessManager: newMockProcessManager(), + resident: map[string]bool{}, + } +} + +func (r *residencyProcessManager) setResident(modelName, address string, resident bool) { + r.resMu.Lock() + defer r.resMu.Unlock() + r.resident[modelName+"@"+address] = resident +} + +func (r *residencyProcessManager) HasModelAtAddress(modelName, address string) bool { + r.resMu.Lock() + defer r.resMu.Unlock() + return r.resident[modelName+"@"+address] +} + func (m *mockProcessManager) getShutdownCalls() []string { m.mu.Lock() defer m.mu.Unlock() @@ -759,6 +802,60 @@ var _ = Describe("WatchDog", func() { Expect(pm.getShutdownCalls()).To(ContainElement("stuckModel")) Expect(pm.getGracefulShutdownCalls()).To(BeEmpty()) }) + + It("keeps tracking a stuck model whose force shutdown failed", func() { + pm.setShutdownError("stuckModel", errors.New("waiting to shut down model: context deadline exceeded")) + wd = model.NewWatchDog( + model.WithProcessManager(pm), + model.WithBusyTimeout(10*time.Millisecond), + model.WithBusyCheck(true), + model.WithWatchdogInterval(20*time.Millisecond), + ) + + wd.AddAddressModelMap("addr1", "stuckModel") + wd.Mark("addr1") + + go wd.Run() + defer wd.Shutdown() + + Eventually(func() []string { + return pm.getForceShutdownCalls() + }, "300ms", "10ms").Should(ContainElement("stuckModel")) + // The backend still runs: the watchdog must not forget it. + Consistently(func() map[string]string { + return wd.GetState().AddressModelMap + }, "100ms", "10ms").Should(HaveKeyWithValue("addr1", "stuckModel")) + }) + }) + + Context("Idle shutdown failure", func() { + It("keeps tracking an idle model that could not be stopped, then stops it once it can", func() { + pm.setShutdownError("idleModel", errors.New("waiting to shut down model: context deadline exceeded")) + wd = model.NewWatchDog( + model.WithProcessManager(pm), + model.WithIdleTimeout(10*time.Millisecond), + model.WithIdleCheck(true), + model.WithWatchdogInterval(20*time.Millisecond), + ) + + wd.AddAddressModelMap("addr1", "idleModel") + wd.Mark("addr1") + wd.UnMark("addr1") + + go wd.Run() + defer wd.Shutdown() + + Eventually(func() []string { + return pm.getShutdownCalls() + }, "300ms", "10ms").Should(ContainElement("idleModel")) + Expect(wd.GetState().AddressModelMap).To(HaveKeyWithValue("addr1", "idleModel"), + "the backend still runs, the watchdog must still know it") + + pm.clearShutdownError("idleModel") // the lifecycle lock is released + Eventually(func() map[string]string { + return wd.GetState().AddressModelMap + }, "500ms", "10ms").ShouldNot(HaveKey("addr1")) + }) }) Context("Concurrency Groups", func() { @@ -935,6 +1032,111 @@ var _ = Describe("WatchDog", func() { Expect(pm.getShutdownCalls()).To(ConsistOf("a")) }) }) + + Describe("a failed shutdown", func() { + // Reproduces a production incident: the eviction of a model timed out + // waiting for its lifecycle lock, the model kept running and kept its + // VRAM, but it had already been dropped from the watchdog. The next + // load saw no conflict and ran out of memory. + BeforeEach(func() { + wd = model.NewWatchDog(model.WithProcessManager(pm), model.WithLRULimit(1)) + wd.ReplaceModelGroups(map[string][]string{ + "a": {"heavy"}, + "b": {"heavy"}, + }) + pm.shutdownErrors["a"] = errors.New("waiting to shut down model: context deadline exceeded") + wd.AddAddressModelMap("addrA", "a") + wd.Mark("addrA") + wd.UnMark("addrA") + }) + + It("keeps tracking a model a group eviction could not stop", func() { + result := wd.EnforceGroupExclusivity("b") + + Expect(result.EvictedCount).To(Equal(0), "nothing was actually stopped") + Expect(result.NeedMore).To(BeTrue(), "the conflict is unresolved, the caller must keep waiting") + Expect(wd.GetLoadedModelCount()).To(Equal(1), "a is still running, it must stay tracked") + }) + + It("still sees the model as a conflict on the next attempt", func() { + Expect(wd.EnforceGroupExclusivity("b").NeedMore).To(BeTrue()) + + delete(pm.shutdownErrors, "a") // the lock is free this time + result := wd.EnforceGroupExclusivity("b") + Expect(result.EvictedCount).To(Equal(1)) + Expect(result.NeedMore).To(BeFalse()) + Expect(wd.GetLoadedModelCount()).To(Equal(0)) + }) + + It("keeps tracking a model an LRU eviction could not stop", func() { + wd.AddAddressModelMap("addrB", "b") + wd.Mark("addrB") + wd.UnMark("addrB") + + // Limit 1 with two models loaded: both must go to make room. + // a refuses to stop, b stops normally. + result := wd.EnforceLRULimit(0) + + Expect(result.EvictedCount).To(Equal(1), "only b was really stopped") + Expect(result.NeedMore).To(BeTrue(), "a still holds its memory, the room does not exist yet") + Expect(wd.GetLoadedModelCount()).To(Equal(1), "the model that would not stop is still tracked") + Expect(wd.GetState().AddressModelMap).To(HaveKeyWithValue("addrA", "a")) + }) + + It("untracks a model that really was stopped", func() { + delete(pm.shutdownErrors, "a") + Expect(wd.EnforceGroupExclusivity("b").EvictedCount).To(Equal(1)) + Expect(wd.GetLoadedModelCount()).To(Equal(0)) + }) + + It("treats a model that is already gone as stopped", func() { + pm.shutdownErrors["a"] = model.ErrModelNotFound + Expect(wd.EnforceGroupExclusivity("b").EvictedCount).To(Equal(1)) + Expect(wd.GetLoadedModelCount()).To(Equal(0)) + }) + }) + + Describe("a shutdown error after the loader dropped the model", func() { + // deleteProcess removes the store entry before it returns an error + // for a failed remote unload (or the remote half of a local+remote + // placement). Tracking such a model again would leave a phantom + // entry that is counted as loaded and as a group conflict. + var rpm *residencyProcessManager + + BeforeEach(func() { + rpm = newResidencyProcessManager() + wd = model.NewWatchDog(model.WithProcessManager(rpm), model.WithLRULimit(1)) + wd.ReplaceModelGroups(map[string][]string{ + "a": {"heavy"}, + "b": {"heavy"}, + }) + rpm.shutdownErrors["a"] = errors.New("unloading remote placements for model \"a\": node unreachable") + wd.AddAddressModelMap("addrA", "a") + wd.Mark("addrA") + wd.UnMark("addrA") + }) + + It("does not track it again when the loader no longer holds it", func() { + rpm.setResident("a", "addrA", false) + + result := wd.EnforceGroupExclusivity("b") + + Expect(result.NeedMore).To(BeTrue(), "the stop did report an error") + Expect(wd.GetLoadedModelCount()).To(Equal(0), "the loader dropped a, no phantom entry may remain") + }) + + It("still tracks it again when the loader still holds it", func() { + rpm.setResident("a", "addrA", true) + + result := wd.EnforceGroupExclusivity("b") + + Expect(result.EvictedCount).To(Equal(0)) + Expect(result.NeedMore).To(BeTrue()) + Expect(wd.GetState().AddressModelMap).To(HaveKeyWithValue("addrA", "a"), + "a is still held by the loader and may still run") + }) + }) + }) Context("Functional Options", func() { From 9a35647ec0bf9e90c445c444945ad16d9c87feff Mon Sep 17 00:00:00 2001 From: hcharbonnier Date: Fri, 23 Oct 2026 18:19:31 +0200 Subject: [PATCH 2/2] fix(watchdog): keep request state live while an eviction is pending Restoring a snapshot after a failed stop lost the request lifecycle events that happened while the stop was in progress. A request completing during the stop found no tracked request, and the restore then resurrected it, leaving the backend marked busy. A request starting during the stop was overwritten by the restore, so completing an older request marked the backend idle while the new one still ran, allowing an unsafe eviction. An eviction now detaches only the backend identity (model, process, size estimate) from the registry and marks the address as pending. Request state stays live, so callbacks keep updating it during the stop. Idle and busy checks skip pending addresses so they neither re-evict nor wipe that state. When the stop settles, a backend that still runs is registered again; otherwise its leftover request state is cleared. A replacement registered on the address meanwhile is left alone. Regression tests cover a request completing and a request starting inside a failing force stop. The idle and busy failure tests now prove re-tracking through a second stop attempt instead of sampling the registry, which raced with the pending window. Assisted-by: Claude:claude-opus-5.5 go vet golangci-lint Signed-off-by: Hugues CHARBONNIER --- pkg/model/watchdog.go | 169 +++++++++++++++++++------------------ pkg/model/watchdog_test.go | 87 ++++++++++++++++--- 2 files changed, 164 insertions(+), 92 deletions(-) diff --git a/pkg/model/watchdog.go b/pkg/model/watchdog.go index 5e4852f6767e..6cc14f5d854e 100644 --- a/pkg/model/watchdog.go +++ b/pkg/model/watchdog.go @@ -34,6 +34,12 @@ type WatchDog struct { idleTime map[string]time.Time lastUsed map[string]time.Time // LRU tracking: when each model was last used + // pendingEvictions counts, per address, the evictions detached from the + // model registry whose stop has not returned yet. Idle and busy checks + // skip these addresses so they neither re-evict nor wipe live request + // state while the stop is in progress. + pendingEvictions map[string]int + timeout, idletimeout time.Duration addressMap map[string]*process.Process addressModelMap map[string]string @@ -100,6 +106,7 @@ func NewWatchDog(opts ...WatchDogOption) *WatchDog { legacyRequests: make(map[string][]uint64), idleTime: make(map[string]time.Time), lastUsed: make(map[string]time.Time), + pendingEvictions: make(map[string]int), addressMap: make(map[string]*process.Process), busyCheck: o.busyCheck, idleCheck: o.idleCheck, @@ -472,97 +479,78 @@ type evictionTarget struct { model string wasBusy bool - // tracked is what untrack removed for this address, kept so a shutdown that - // fails can put it back. Every eviction path captures it before untrack; - // nil only when the address was not tracked. - tracked *trackedState + // detached is what the eviction removed from the model registry for this + // address, kept so a shutdown that fails can register it again. nil only + // when the address was not tracked. + detached *detachedBackend } -// trackedState is a snapshot of everything the watchdog knows about one backend -// address. An eviction untracks the backend before it is stopped, so a stop that -// fails (typically a lifecycle lock held by a long model load) would otherwise -// leave a backend that is still running and still holding its VRAM invisible to -// the group and LRU checks: the next load sees no conflict and runs out of memory. -type trackedState struct { - model string - process *process.Process - busyTime time.Time - hasBusy bool - inFlight int - requestStarts map[uint64]time.Time - legacyRequests []uint64 - idleTime time.Time - hasIdle bool - lastUsed time.Time - hasLastUsed bool - size int64 - hasSize bool +// detachedBackend is the identity of a backend whose eviction is pending. An +// eviction removes the backend from the model registry before it is stopped, +// so that concurrent LRU/group checks do not select it twice. A stop that fails +// (typically a lifecycle lock held by a long model load) leaves a backend that +// still runs and still holds its VRAM, so it must be registered again, or the +// next load sees no conflict and runs out of memory. +// +// Only the identity is detached. Request lifecycle state (in-flight requests, +// busy/idle/last-used times) stays live during the stop, so a request that +// completes or starts meanwhile updates it normally. Replaying a snapshot of +// that state after a failed stop would resurrect completed requests or erase +// requests started during the stop. +type detachedBackend struct { + model string + process *process.Process + size int64 + hasSize bool } -// snapshotLocked captures the tracking state of an address. Must be called with -// wd.Lock() held, before untrack. -func (wd *WatchDog) snapshotLocked(address string) *trackedState { +// detachForEvictionLocked removes the backend at address from the model +// registry and marks its eviction as pending. Must be called with wd.Lock() +// held. +func (wd *WatchDog) detachForEvictionLocked(address string) *detachedBackend { model, ok := wd.addressModelMap[address] if !ok { return nil } - s := &trackedState{model: model, process: wd.addressMap[address], inFlight: wd.inFlight[address]} - s.busyTime, s.hasBusy = wd.busyTime[address] - s.idleTime, s.hasIdle = wd.idleTime[address] - s.lastUsed, s.hasLastUsed = wd.lastUsed[address] - s.size, s.hasSize = wd.modelSizes[model] - if starts := wd.requestStarts[address]; starts != nil { - s.requestStarts = make(map[uint64]time.Time, len(starts)) - for id, t := range starts { - s.requestStarts[id] = t - } - } - s.legacyRequests = slices.Clone(wd.legacyRequests[address]) - return s + d := &detachedBackend{model: model, process: wd.addressMap[address]} + d.size, d.hasSize = wd.modelSizes[model] + wd.untrackModelLocked(address) + wd.pendingEvictions[address]++ + return d } -// restoreTracked puts a snapshot back after the backend it describes could not -// be stopped. It does nothing when the address was re-registered in the -// meantime: a replacement backend already owns the address and must win. -func (wd *WatchDog) restoreTracked(address string, s *trackedState) { - if s == nil { - return - } +// finishEviction settles a pending eviction. keep registers the backend again +// because it is still running; otherwise the request state left for the +// address is cleared. Nothing is touched when a replacement backend registered +// the address in the meantime: it owns the address and its request state. +func (wd *WatchDog) finishEviction(t evictionTarget, keep bool) { wd.Lock() defer wd.Unlock() - if _, taken := wd.addressModelMap[address]; taken { - return - } - wd.addressModelMap[address] = s.model - if s.process != nil { - wd.addressMap[address] = s.process - } - if s.hasBusy { - wd.busyTime[address] = s.busyTime - } - if s.inFlight > 0 { - wd.inFlight[address] = s.inFlight - } - if len(s.requestStarts) > 0 { - wd.requestStarts[address] = s.requestStarts + if wd.pendingEvictions[t.address] <= 1 { + delete(wd.pendingEvictions, t.address) + } else { + wd.pendingEvictions[t.address]-- } - if len(s.legacyRequests) > 0 { - wd.legacyRequests[address] = s.legacyRequests + if _, taken := wd.addressModelMap[t.address]; taken { + return } - if s.hasIdle { - wd.idleTime[address] = s.idleTime + if !keep || t.detached == nil { + wd.clearRequestStateLocked(t.address) + return } - if s.hasLastUsed { - wd.lastUsed[address] = s.lastUsed + d := t.detached + wd.addressModelMap[t.address] = d.model + if d.process != nil { + wd.addressMap[t.address] = d.process } - if s.hasSize { - // untrack dropped the estimate only when no other address served this + if d.hasSize { + // The estimate was dropped only when no other address served this // model; never overwrite one a replacement has registered since. - if _, ok := wd.modelSizes[s.model]; !ok { - wd.modelSizes[s.model] = s.size + if _, ok := wd.modelSizes[d.model]; !ok { + wd.modelSizes[d.model] = d.size } } - xlog.Warn("[WatchDog] model could not be stopped, tracking it again", "model", s.model, "address", address) + xlog.Warn("[WatchDog] model could not be stopped, tracking it again", "model", d.model, "address", t.address) } // EnforceLRULimitResult contains the result of LRU enforcement @@ -673,8 +661,7 @@ func (wd *WatchDog) collectEvictionsLocked(candidates []modelUsageInfo, maxToEvi continue } xlog.Info("[WatchDog] evicting model", "model", m.model, "busy", isBusy) - evicted = append(evicted, evictionTarget{address: m.address, model: m.model, wasBusy: isBusy, tracked: wd.snapshotLocked(m.address)}) - wd.untrack(m.address) + evicted = append(evicted, evictionTarget{address: m.address, model: m.model, wasBusy: isBusy, detached: wd.detachForEvictionLocked(m.address)}) } return evicted, skippedBusy } @@ -697,16 +684,18 @@ func (wd *WatchDog) shutdownEvictedStopped(targets []evictionTarget, label strin if err != nil && !errors.Is(err, ErrModelNotFound) { xlog.Error("[WatchDog] error shutting down model during "+label, "error", err, "model", t.model, "busy", t.wasBusy) if wd.stillResident(t) { - wd.restoreTracked(t.address, t.tracked) + wd.finishEviction(t, true) } else { // The loader already dropped the model before the error was // returned (a failed remote unload, or the remote half of a // local+remote placement). Tracking it again would leave a // phantom entry that inflates LRU counts and group conflicts. xlog.Warn("[WatchDog] model was already removed by the loader despite the shutdown error, not tracking it again", "model", t.model, "address", t.address) + wd.finishEviction(t, false) } continue } + wd.finishEviction(t, false) stopped++ xlog.Debug("[WatchDog] "+label+" complete", "model", t.model, "busy", t.wasBusy) } @@ -864,6 +853,10 @@ func (wd *WatchDog) checkIdle() { var modelsToShutdown []evictionTarget for address, t := range wd.idleTime { xlog.Debug("[WatchDog] idle connection", "address", address) + if wd.pendingEvictions[address] > 0 { + // Its stop is in progress; the eviction settles its state. + continue + } if time.Since(t) > wd.idletimeout { model, ok := wd.addressModelMap[address] if ok { @@ -872,11 +865,11 @@ func (wd *WatchDog) checkIdle() { continue } xlog.Warn("[WatchDog] Address is idle for too long, killing it", "address", address, "model", model) - modelsToShutdown = append(modelsToShutdown, evictionTarget{model: model, address: address, tracked: wd.snapshotLocked(address)}) + modelsToShutdown = append(modelsToShutdown, evictionTarget{model: model, address: address, detached: wd.detachForEvictionLocked(address)}) } else { xlog.Warn("[WatchDog] Address unresolvable", "address", address) + wd.untrack(address) } - wd.untrack(address) } } wd.Unlock() @@ -892,16 +885,20 @@ func (wd *WatchDog) checkBusy() { var modelsToShutdown []evictionTarget for address, t := range wd.busyTime { xlog.Debug("[WatchDog] active connection", "address", address) + if wd.pendingEvictions[address] > 0 { + // Its stop is in progress; the eviction settles its state. + continue + } if time.Since(t) > wd.timeout { model, ok := wd.addressModelMap[address] if ok { xlog.Warn("[WatchDog] Model is busy for too long, killing it", "model", model) - modelsToShutdown = append(modelsToShutdown, evictionTarget{model: model, address: address, wasBusy: true, tracked: wd.snapshotLocked(address)}) + modelsToShutdown = append(modelsToShutdown, evictionTarget{model: model, address: address, wasBusy: true, detached: wd.detachForEvictionLocked(address)}) } else { xlog.Warn("[WatchDog] Address unresolvable", "address", address) + wd.untrack(address) } - wd.untrack(address) } } wd.Unlock() @@ -1051,6 +1048,12 @@ func (wd *WatchDog) evictLRUModel() { } func (wd *WatchDog) untrack(address string) { + wd.untrackModelLocked(address) + wd.clearRequestStateLocked(address) +} + +// untrackModelLocked removes the backend identity registered at address. +func (wd *WatchDog) untrackModelLocked(address string) { if modelID, ok := wd.addressModelMap[address]; ok { // A stale address can coexist with the replacement's registration. // Its cleanup must not erase the replacement's size estimate. @@ -1065,14 +1068,18 @@ func (wd *WatchDog) untrack(address string) { delete(wd.modelSizes, modelID) } } + delete(wd.addressModelMap, address) + delete(wd.addressMap, address) +} + +// clearRequestStateLocked removes the request lifecycle state of address. +func (wd *WatchDog) clearRequestStateLocked(address string) { delete(wd.busyTime, address) delete(wd.inFlight, address) delete(wd.requestStarts, address) delete(wd.legacyRequests, address) delete(wd.idleTime, address) delete(wd.lastUsed, address) - delete(wd.addressModelMap, address) - delete(wd.addressMap, address) } // Untrack removes request and eviction state after a backend is removed. diff --git a/pkg/model/watchdog_test.go b/pkg/model/watchdog_test.go index 8de1181873d4..60930ebe784a 100644 --- a/pkg/model/watchdog_test.go +++ b/pkg/model/watchdog_test.go @@ -90,6 +90,23 @@ func (r *residencyProcessManager) HasModelAtAddress(modelName, address string) b return r.resident[modelName+"@"+address] } +// hookProcessManager runs a hook inside the stop, then fails it, so a test can +// drive request lifecycle events while an eviction is pending. +type hookProcessManager struct { + hook func() + err error +} + +func (h *hookProcessManager) ShutdownModel(string) error { + h.hook() + return h.err +} + +func (h *hookProcessManager) ShutdownModelForce(string) error { + h.hook() + return h.err +} + func (m *mockProcessManager) getShutdownCalls() []string { m.mu.Lock() defer m.mu.Unlock() @@ -818,13 +835,13 @@ var _ = Describe("WatchDog", func() { go wd.Run() defer wd.Shutdown() - Eventually(func() []string { - return pm.getForceShutdownCalls() - }, "300ms", "10ms").Should(ContainElement("stuckModel")) - // The backend still runs: the watchdog must not forget it. - Consistently(func() map[string]string { - return wd.GetState().AddressModelMap - }, "100ms", "10ms").Should(HaveKeyWithValue("addr1", "stuckModel")) + // The backend still runs: the watchdog must not forget it. The busy + // check only stops a tracked model, so a second stop attempt proves + // it was tracked again. Sampling the map instead would race with + // the short window in which the eviction is pending. + Eventually(func() int { + return len(pm.getForceShutdownCalls()) + }, "1s", "10ms").Should(BeNumerically(">=", 2)) }) }) @@ -845,10 +862,11 @@ var _ = Describe("WatchDog", func() { go wd.Run() defer wd.Shutdown() - Eventually(func() []string { - return pm.getShutdownCalls() - }, "300ms", "10ms").Should(ContainElement("idleModel")) - Expect(wd.GetState().AddressModelMap).To(HaveKeyWithValue("addr1", "idleModel"), + // The idle check only stops a tracked model: a second attempt proves + // the backend that still runs was tracked again. + Eventually(func() int { + return len(pm.getShutdownCalls()) + }, "1s", "10ms").Should(BeNumerically(">=", 2), "the backend still runs, the watchdog must still know it") pm.clearShutdownError("idleModel") // the lifecycle lock is released @@ -1096,6 +1114,53 @@ var _ = Describe("WatchDog", func() { }) }) + Describe("request lifecycle during a stop that fails", func() { + // Requests keep completing and starting while a busy backend is + // being force-stopped. When the stop fails, the watchdog must keep + // what happened meanwhile instead of replaying pre-stop state. + var hpm *hookProcessManager + + BeforeEach(func() { + hpm = &hookProcessManager{err: errors.New("waiting to shut down model: context deadline exceeded")} + wd = model.NewWatchDog( + model.WithProcessManager(hpm), + model.WithLRULimit(1), + model.WithForceEvictionWhenBusy(true), + ) + wd.AddAddressModelMap("addrA", "a") + }) + + It("does not resurrect a request that completed during the stop", func() { + done := wd.TrackRequest("addrA") + hpm.hook = done + + result := wd.EnforceLRULimit(0) + + Expect(result.EvictedCount).To(Equal(0)) + state := wd.GetState() + Expect(state.AddressModelMap).To(HaveKeyWithValue("addrA", "a")) + Expect(state.InFlight).ToNot(HaveKey("addrA"), "the request completed during the stop") + Expect(state.BusyTime).ToNot(HaveKey("addrA"), "the backend is idle, not busy") + }) + + It("keeps a request that started during the stop", func() { + first := wd.TrackRequest("addrA") + var second func() + hpm.hook = func() { second = wd.TrackRequest("addrA") } + + wd.EnforceLRULimit(0) + first() + + state := wd.GetState() + Expect(state.AddressModelMap).To(HaveKeyWithValue("addrA", "a")) + Expect(state.InFlight).To(HaveKeyWithValue("addrA", 1), "the request started during the stop is still running") + Expect(state.BusyTime).To(HaveKey("addrA"), "the backend must stay busy so it is not evicted unsafely") + + second() + Expect(wd.GetState().BusyTime).ToNot(HaveKey("addrA")) + }) + }) + Describe("a shutdown error after the loader dropped the model", func() { // deleteProcess removes the store entry before it returns an error // for a failed remote unload (or the remote half of a local+remote