From 9a35647ec0bf9e90c445c444945ad16d9c87feff Mon Sep 17 00:00:00 2001 From: hcharbonnier Date: Fri, 23 Oct 2026 18:19:31 +0200 Subject: [PATCH] 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 5e4852f67..6cc14f5d8 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 { + if wd.pendingEvictions[t.address] <= 1 { + delete(wd.pendingEvictions, t.address) + } else { + wd.pendingEvictions[t.address]-- + } + if _, taken := wd.addressModelMap[t.address]; taken { return } - wd.addressModelMap[address] = s.model - if s.process != nil { - wd.addressMap[address] = s.process + if !keep || t.detached == nil { + wd.clearRequestStateLocked(t.address) + return } - if s.hasBusy { - wd.busyTime[address] = s.busyTime + d := t.detached + wd.addressModelMap[t.address] = d.model + if d.process != nil { + wd.addressMap[t.address] = d.process } - 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 + 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 8de118187..60930ebe7 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