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 <hugues.charbonnier@gmail.com>
This commit is contained in:
hcharbonnier committed 2026-10-23 17:29:13 +02:00
1 parent 7d528e90f4
commit 97390f7754
5 files changed
+605 -12

No files matched your search

+159
View File
@@ -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")
})
})
+13
View File
@@ -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 {
+159 -12
View File
@@ -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()
+72
View File
@@ -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()
+202
View File
@@ -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() {