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..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, @@ -471,6 +478,79 @@ type evictionTarget struct { address string model string wasBusy bool + + // 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 +} + +// 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 +} + +// 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 + } + d := &detachedBackend{model: model, process: wd.addressMap[address]} + d.size, d.hasSize = wd.modelSizes[model] + wd.untrackModelLocked(address) + wd.pendingEvictions[address]++ + return d +} + +// 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 wd.pendingEvictions[t.address] <= 1 { + delete(wd.pendingEvictions, t.address) + } else { + wd.pendingEvictions[t.address]-- + } + if _, taken := wd.addressModelMap[t.address]; taken { + return + } + if !keep || t.detached == nil { + wd.clearRequestStateLocked(t.address) + return + } + d := t.detached + wd.addressModelMap[t.address] = d.model + if d.process != nil { + wd.addressMap[t.address] = d.process + } + 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[d.model]; !ok { + wd.modelSizes[d.model] = d.size + } + } + xlog.Warn("[WatchDog] model could not be stopped, tracking it again", "model", d.model, "address", t.address) } // EnforceLRULimitResult contains the result of LRU enforcement @@ -544,14 +624,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,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}) - wd.untrack(m.address) + evicted = append(evicted, evictionTarget{address: m.address, model: m.model, wasBusy: isBusy, detached: wd.detachForEvictionLocked(m.address)}) } return evicted, skippedBusy } @@ -586,13 +670,50 @@ 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.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) } + 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 +780,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, } } @@ -727,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 { @@ -735,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}) + 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() @@ -755,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}) + 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() @@ -891,9 +1025,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() @@ -904,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. @@ -918,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_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..60930ebe784a 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,65 @@ 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] +} + +// 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() @@ -759,6 +819,61 @@ 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() + + // 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)) + }) + }) + + 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() + + // 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 + Eventually(func() map[string]string { + return wd.GetState().AddressModelMap + }, "500ms", "10ms").ShouldNot(HaveKey("addr1")) + }) }) Context("Concurrency Groups", func() { @@ -935,6 +1050,158 @@ 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("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 + // 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() {