diff --git a/pkg/model/initializers_group_load_test.go b/pkg/model/initializers_group_load_test.go new file mode 100644 index 000000000..98a0b9aee --- /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 1db9a9cd7..f823f9de8 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 977eca1ff..5e4852f67 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 50e4e84d3..cf3141c2d 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 36d9ceed2..8de118187 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() {