From 94e6cce8333cfa3cf74192b1988457faa043987b Mon Sep 17 00:00:00 2001 From: Richard Palethorpe Date: Thu, 16 Jul 2026 14:28:55 +0100 Subject: [PATCH] fix(model): make backend shutdown model-scoped Avoid holding the global loader lock across backend lifecycle waits and propagate forced shutdown through distributed workers. Track parallel requests with in-flight counters and reserve worker ports until process termination. Add focused race tests and an authoritative FizzBee lifecycle model with a fail-closed conformance target. Assisted-by: Codex:GPT-5 [FizzBee] [Ginkgo] Signed-off-by: Richard Palethorpe --- Makefile | 7 +- core/services/messaging/subjects.go | 10 +- core/services/nodes/unloader.go | 58 ++++-- core/services/nodes/unloader_test.go | 16 +- core/services/worker/addr_test.go | 2 + core/services/worker/file_staging.go | 44 ++-- core/services/worker/install.go | 24 ++- core/services/worker/lifecycle.go | 91 +++++--- core/services/worker/registration.go | 58 ++++-- core/services/worker/replica_test.go | 86 ++++++++ core/services/worker/supervisor.go | 160 +++++++++----- core/services/worker/worker.go | 19 +- docs/content/advanced/vram-management.md | 18 +- formal-verification/README.md | 41 +++- .../model_loader_shutdown.fizz | 195 ++++++++++++++++++ pkg/grpc/client.go | 172 ++++++--------- pkg/grpc/client_busy_test.go | 27 +++ pkg/model/loader.go | 94 ++++++--- pkg/model/loader_test.go | 192 +++++++++++++++++ pkg/model/model.go | 1 + pkg/model/operation_locks.go | 71 +++++++ pkg/model/process.go | 113 ++++++---- pkg/model/watchdog.go | 157 +++++++++++--- pkg/model/watchdog_test.go | 54 ++++- scripts/model-lifecycle-conformance.sh | 55 +++++ 25 files changed, 1417 insertions(+), 348 deletions(-) create mode 100644 formal-verification/model_loader_shutdown.fizz create mode 100644 pkg/grpc/client_busy_test.go create mode 100644 pkg/model/operation_locks.go create mode 100755 scripts/model-lifecycle-conformance.sh diff --git a/Makefile b/Makefile index 9e282e9273ab..c608de326f01 100644 --- a/Makefile +++ b/Makefile @@ -412,8 +412,13 @@ test-realtime: build-mock-backend test-realtime-conformance: GOCMD=$(GOCMD) ./scripts/realtime-conformance.sh +# Verify the shared model-loader shutdown behavior independently of any API +# modality (focused loader/gRPC/distributed/worker tests under -race + FizzBee). +test-model-lifecycle-conformance: + GOCMD=$(GOCMD) ./scripts/model-lifecycle-conformance.sh + # Install the pinned, checksum-verified FizzBee model checker (into .tools/, -# gitignored) used by test-realtime-conformance. Idempotent; no-op if present. +# gitignored) used by the conformance targets. Idempotent; no-op if present. install-fizzbee: ./scripts/install-fizzbee.sh diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index d2b11f53583e..bbdd9f0f23bf 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -258,9 +258,17 @@ type NodeBackendInfo struct { Digest string `json:"digest,omitempty"` } +// BackendStopRequest controls worker-side process shutdown. Force skips the +// best-effort Free RPC so a backend stuck serving a request can still be +// terminated by the watchdog. +type BackendStopRequest struct { + Backend string `json:"backend"` + Force bool `json:"force,omitempty"` +} + // SubjectNodeBackendStop tells a worker node to stop its gRPC backend process. // Equivalent to the local deleteProcess(). The node will: -// 1. Best-effort Free() via gRPC +// 1. Best-effort bounded Free() via gRPC (unless Force is true) // 2. Kill the backend process // 3. Can be restarted via another backend.start event. func SubjectNodeBackendStop(nodeID string) string { diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index bc1f7388a06c..8699968178c7 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -15,11 +15,6 @@ import ( "github.com/mudler/xlog" ) -// backendStopRequest is the request payload for backend.stop (fire-and-forget). -type backendStopRequest struct { - Backend string `json:"backend"` -} - // NodeCommandSender abstracts NATS-based commands to worker nodes. // Used by HTTP endpoint handlers to avoid coupling to the concrete RemoteUnloaderAdapter. // @@ -78,28 +73,44 @@ func (a *RemoteUnloaderAdapter) InstallTimeout() time.Duration { // UnloadRemoteModel finds the node(s) hosting the given model and tells them // to stop their backend process via NATS backend.stop event. -// The worker process handles: Free() → kill process. +// The worker process handles a bounded Free() followed by process termination; +// forced shutdown skips Free(). // This is called by ModelLoader.deleteProcess() when process == nil (remote model). func (a *RemoteUnloaderAdapter) UnloadRemoteModel(modelName string) error { - ctx := context.Background() + return a.UnloadRemoteModelContext(context.Background(), modelName, false) +} + +// UnloadRemoteModelContext is the cancellation-aware extension used by the +// model loader to preserve forced shutdown across the distributed boundary. +func (a *RemoteUnloaderAdapter) UnloadRemoteModelContext(ctx context.Context, modelName string, force bool) error { + if ctx == nil { + ctx = context.Background() + } nodes, err := a.registry.FindNodesWithModel(ctx, modelName) - if err != nil || len(nodes) == 0 { + if err != nil { + return fmt.Errorf("finding nodes with model %q: %w", modelName, err) + } + if len(nodes) == 0 { xlog.Debug("No remote nodes found with model", "model", modelName) return nil } + var unloadErr error for _, node := range nodes { - xlog.Info("Sending NATS backend.stop to node", "model", modelName, "node", node.Name, "nodeID", node.ID) - if err := a.StopBackend(node.ID, modelName); err != nil { + xlog.Info("Sending NATS backend.stop to node", "model", modelName, "node", node.Name, "nodeID", node.ID, "force", force) + if err := a.stopBackend(node.ID, modelName, force); err != nil { xlog.Warn("Failed to send backend.stop", "node", node.Name, "error", err) + unloadErr = errors.Join(unloadErr, fmt.Errorf("stopping model on node %s: %w", node.ID, err)) continue } // Remove every replica of this model on the node — the worker will // handle the actual process cleanup. - a.registry.RemoveAllNodeModelReplicas(ctx, node.ID, modelName) + if err := a.registry.RemoveAllNodeModelReplicas(ctx, node.ID, modelName); err != nil { + unloadErr = errors.Join(unloadErr, fmt.Errorf("removing model replicas from node %s: %w", node.ID, err)) + } } - return nil + return unloadErr } // InstallBackend sends a backend.install request-reply to a worker node. @@ -142,7 +153,9 @@ func (a *RemoteUnloaderAdapter) InstallBackend( }, a.installTimeout) if sub != nil { - _ = sub.Unsubscribe() + if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil { + xlog.Warn("Failed to unsubscribe from backend install progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr) + } } if err != nil && isNATSTimeout(err) { @@ -216,7 +229,9 @@ func (a *RemoteUnloaderAdapter) UpgradeBackend(nodeID, backendType, galleriesJSO }, a.upgradeTimeout) if sub != nil { - _ = sub.Unsubscribe() + if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil { + xlog.Warn("Failed to unsubscribe from backend upgrade progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr) + } } if err != nil && isNATSTimeout(err) { @@ -250,7 +265,9 @@ func (a *RemoteUnloaderAdapter) installWithForceFallback(nodeID, backendType, ga }, a.upgradeTimeout) if sub != nil { - _ = sub.Unsubscribe() + if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil { + xlog.Warn("Failed to unsubscribe from legacy backend install progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr) + } } if err != nil && isNATSTimeout(err) { @@ -272,14 +289,15 @@ func (a *RemoteUnloaderAdapter) ListBackends(nodeID string) (*messaging.BackendL // If backend is empty, the worker stops ALL backends. // The node stays registered and can receive another InstallBackend later. func (a *RemoteUnloaderAdapter) StopBackend(nodeID, backend string) error { + return a.stopBackend(nodeID, backend, false) +} + +func (a *RemoteUnloaderAdapter) stopBackend(nodeID, backend string, force bool) error { subject := messaging.SubjectNodeBackendStop(nodeID) - if backend == "" { + if backend == "" && !force { return a.nats.Publish(subject, nil) } - req := struct { - Backend string `json:"backend"` - }{Backend: backend} - return a.nats.Publish(subject, req) + return a.nats.Publish(subject, messaging.BackendStopRequest{Backend: backend, Force: force}) } // DeleteBackend tells a worker node to delete a backend (stop + remove files). diff --git a/core/services/nodes/unloader_test.go b/core/services/nodes/unloader_test.go index 41fdae147285..dc2c7e8b1ff6 100644 --- a/core/services/nodes/unloader_test.go +++ b/core/services/nodes/unloader_test.go @@ -160,7 +160,7 @@ var _ = Describe("RemoteUnloaderAdapter", func() { failOnce := &failOnceMessagingClient{inner: mc, failOn: 0} adapter = NewRemoteUnloaderAdapter(locator, failOnce, 3*time.Minute, 15*time.Minute) - Expect(adapter.UnloadRemoteModel("llama")).To(Succeed()) + Expect(adapter.UnloadRemoteModel("llama")).To(HaveOccurred()) // The second node should still have been processed. // The first node's StopBackend errored, so RemoveNodeModel was NOT called for it. @@ -168,6 +168,15 @@ var _ = Describe("RemoteUnloaderAdapter", func() { Expect(locator.removedPairs).To(HaveLen(1)) Expect(locator.removedPairs[0].nodeID).To(Equal("node-ok")) }) + + It("propagates forced shutdown to every worker", func() { + locator.nodes = []BackendNode{{ID: "node-1", Name: "worker-1"}} + Expect(adapter.UnloadRemoteModelContext(context.Background(), "llama", true)).To(Succeed()) + + var payload messaging.BackendStopRequest + Expect(json.Unmarshal(mc.published[0].Data, &payload)).To(Succeed()) + Expect(payload).To(Equal(messaging.BackendStopRequest{Backend: "llama", Force: true})) + }) }) Describe("StopBackend", func() { @@ -182,11 +191,10 @@ var _ = Describe("RemoteUnloaderAdapter", func() { Expect(adapter.StopBackend("node-1", "llama-backend")).To(Succeed()) Expect(mc.published).To(HaveLen(1)) - var payload struct { - Backend string `json:"backend"` - } + var payload messaging.BackendStopRequest Expect(json.Unmarshal(mc.published[0].Data, &payload)).To(Succeed()) Expect(payload.Backend).To(Equal("llama-backend")) + Expect(payload.Force).To(BeFalse()) }) }) diff --git a/core/services/worker/addr_test.go b/core/services/worker/addr_test.go index 5265a287a7fd..4f5b1ba67f4f 100644 --- a/core/services/worker/addr_test.go +++ b/core/services/worker/addr_test.go @@ -19,7 +19,9 @@ var _ = Describe("Worker address resolution", func() { Entry("falls back to ServeAddr", "", "0.0.0.0:50051", 50051), Entry("returns 50051 when neither set", "", "", 50051), Entry("Addr with custom port", "10.0.0.5:7000", "", 7000), + Entry("supports bracketed IPv6", "[2001:db8::1]:7001", "", 7001), Entry("invalid port in Addr falls through to ServeAddr", "host:notanumber", "0.0.0.0:9999", 9999), + Entry("out-of-range port in Addr falls through to ServeAddr", "host:70000", "0.0.0.0:9998", 9998), ) }) diff --git a/core/services/worker/file_staging.go b/core/services/worker/file_staging.go index 9d834a7127de..019afcba9ace 100644 --- a/core/services/worker/file_staging.go +++ b/core/services/worker/file_staging.go @@ -59,7 +59,7 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no } // Subscribe: files.ensure — download S3 key to local, reply with local path - natsClient.SubscribeReply(messaging.SubjectNodeFilesEnsure(nodeID), func(data []byte, reply func([]byte)) { + if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesEnsure(nodeID), func(data []byte, reply func([]byte)) { var req struct { Key string `json:"key"` } @@ -77,10 +77,12 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no xlog.Debug("File ensured locally", "key", req.Key, "path", localPath) replyJSON(reply, map[string]string{"local_path": localPath}) - }) + }); err != nil { + return fmt.Errorf("subscribing to files.ensure events: %w", err) + } // Subscribe: files.stage — upload local path to S3, reply with key - natsClient.SubscribeReply(messaging.SubjectNodeFilesStage(nodeID), func(data []byte, reply func([]byte)) { + if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesStage(nodeID), func(data []byte, reply func([]byte)) { var req struct { LocalPath string `json:"local_path"` Key string `json:"key"` @@ -107,10 +109,12 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no xlog.Debug("File staged to S3", "path", req.LocalPath, "key", req.Key) replyJSON(reply, map[string]string{"key": req.Key}) - }) + }); err != nil { + return fmt.Errorf("subscribing to files.stage events: %w", err) + } // Subscribe: files.temp — allocate temp file, reply with local path - natsClient.SubscribeReply(messaging.SubjectNodeFilesTemp(nodeID), func(data []byte, reply func([]byte)) { + if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesTemp(nodeID), func(data []byte, reply func([]byte)) { tmpDir := filepath.Join(cacheDir, "staging-tmp") if err := os.MkdirAll(tmpDir, 0750); err != nil { replyJSON(reply, map[string]string{"error": fmt.Sprintf("creating temp dir: %v", err)}) @@ -123,14 +127,19 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no return } localPath := f.Name() - f.Close() + if err := f.Close(); err != nil { + replyJSON(reply, map[string]string{"error": fmt.Sprintf("closing temp file: %v", err)}) + return + } xlog.Debug("Allocated temp file", "path", localPath) replyJSON(reply, map[string]string{"local_path": localPath}) - }) + }); err != nil { + return fmt.Errorf("subscribing to files.temp events: %w", err) + } // Subscribe: files.listdir — list files in a local directory, reply with relative paths - natsClient.SubscribeReply(messaging.SubjectNodeFilesListDir(nodeID), func(data []byte, reply func([]byte)) { + if _, err := natsClient.SubscribeReply(messaging.SubjectNodeFilesListDir(nodeID), func(data []byte, reply func([]byte)) { var req struct { KeyPrefix string `json:"key_prefix"` } @@ -163,22 +172,29 @@ func (cfg *Config) subscribeFileStaging(natsClient messaging.MessagingClient, no } var files []string - filepath.WalkDir(dirPath, func(path string, d os.DirEntry, err error) error { + if err := filepath.WalkDir(dirPath, func(path string, d os.DirEntry, err error) error { if err != nil { - return nil + return err } if !d.IsDir() { rel, err := filepath.Rel(dirPath, path) - if err == nil { - files = append(files, rel) + if err != nil { + return err } + files = append(files, rel) } return nil - }) + }); err != nil { + xlog.Error("Failed to list staged files", "keyPrefix", req.KeyPrefix, "dirPath", dirPath, "error", err) + replyJSON(reply, map[string]any{"error": err.Error()}) + return + } xlog.Debug("Listed remote dir", "keyPrefix", req.KeyPrefix, "dirPath", dirPath, "fileCount", len(files)) replyJSON(reply, map[string]any{"files": files}) - }) + }); err != nil { + return fmt.Errorf("subscribing to files.listdir events: %w", err) + } xlog.Info("Subscribed to file staging NATS subjects", "nodeID", nodeID) return nil diff --git a/core/services/worker/install.go b/core/services/worker/install.go index bcc7cc9b04ee..19056ea7a20d 100644 --- a/core/services/worker/install.go +++ b/core/services/worker/install.go @@ -78,7 +78,7 @@ func (s *backendSupervisor) installBackend(req messaging.BackendInstallRequest, } xlog.Warn("Stale process entry for backend (dead process); cleaning up before reinstall", "backend", req.Backend, "model", req.ModelID, "replica", req.ReplicaIndex, "addr", addr) - s.stopBackendExact(processKey) + s.stopBackendExact(processKey, false) } } else { // Upgrade path: stop every live process that shares this backend so the @@ -95,7 +95,7 @@ func (s *backendSupervisor) installBackend(req messaging.BackendInstallRequest, for _, key := range toStop { xlog.Info("Force install: stopping running backend before reinstall", "backend", req.Backend, "processKey", key) - s.stopBackendExact(key) + s.stopBackendExact(key, true) } } @@ -103,9 +103,10 @@ func (s *backendSupervisor) installBackend(req messaging.BackendInstallRequest, galleries := s.galleries if req.BackendGalleries != "" { var reqGalleries []config.Gallery - if err := json.Unmarshal([]byte(req.BackendGalleries), &reqGalleries); err == nil { - galleries = reqGalleries + if err := json.Unmarshal([]byte(req.BackendGalleries), &reqGalleries); err != nil { + return "", fmt.Errorf("decoding backend galleries: %w", err) } + galleries = reqGalleries } // When the master tagged this install with an OpID, stream the @@ -147,7 +148,9 @@ func (s *backendSupervisor) installBackend(req messaging.BackendInstallRequest, } } // Re-register after install and retry - gallery.RegisterBackends(s.systemState, s.ml) + if err := gallery.RegisterBackends(s.systemState, s.ml); err != nil { + return "", fmt.Errorf("refreshing registered backends after install: %w", err) + } backendPath = s.findBackend(req.Backend) } @@ -175,15 +178,16 @@ func (s *backendSupervisor) upgradeBackend(req messaging.BackendUpgradeRequest) for _, key := range toStop { xlog.Info("Upgrade: stopping running backend before reinstall", "backend", req.Backend, "processKey", key) - s.stopBackendExact(key) + s.stopBackendExact(key, true) } galleries := s.galleries if req.BackendGalleries != "" { var reqGalleries []config.Gallery - if err := json.Unmarshal([]byte(req.BackendGalleries), &reqGalleries); err == nil { - galleries = reqGalleries + if err := json.Unmarshal([]byte(req.BackendGalleries), &reqGalleries); err != nil { + return fmt.Errorf("decoding backend galleries: %w", err) } + galleries = reqGalleries } // When the master tagged this upgrade with an OpID, stream gallery download @@ -215,7 +219,9 @@ func (s *backendSupervisor) upgradeBackend(req messaging.BackendUpgradeRequest) } } - gallery.RegisterBackends(s.systemState, s.ml) + if err := gallery.RegisterBackends(s.systemState, s.ml); err != nil { + return fmt.Errorf("refreshing registered backends after upgrade: %w", err) + } return nil } diff --git a/core/services/worker/lifecycle.go b/core/services/worker/lifecycle.go index a3af45cfcf8d..0348e41ecf3c 100644 --- a/core/services/worker/lifecycle.go +++ b/core/services/worker/lifecycle.go @@ -18,15 +18,32 @@ import ( // keeping the dispatcher to a single line per subject makes adding a new // subject a 2-line patch (one line here, one new method) instead of grafting // onto a monolith. -func (s *backendSupervisor) subscribeLifecycleEvents() { - s.nats.SubscribeReply(messaging.SubjectNodeBackendInstall(s.nodeID), s.handleBackendInstall) - s.nats.SubscribeReply(messaging.SubjectNodeBackendUpgrade(s.nodeID), s.handleBackendUpgrade) - s.nats.Subscribe(messaging.SubjectNodeBackendStop(s.nodeID), s.handleBackendStop) - s.nats.SubscribeReply(messaging.SubjectNodeBackendDelete(s.nodeID), s.handleBackendDelete) - s.nats.SubscribeReply(messaging.SubjectNodeBackendList(s.nodeID), s.handleBackendList) - s.nats.SubscribeReply(messaging.SubjectNodeModelUnload(s.nodeID), s.handleModelUnload) - s.nats.SubscribeReply(messaging.SubjectNodeModelDelete(s.nodeID), s.handleModelDelete) - s.nats.Subscribe(messaging.SubjectNodeStop(s.nodeID), s.handleNodeStop) +func (s *backendSupervisor) subscribeLifecycleEvents() error { + if _, err := s.nats.SubscribeReply(messaging.SubjectNodeBackendInstall(s.nodeID), s.handleBackendInstall); err != nil { + return fmt.Errorf("subscribing to backend install events: %w", err) + } + if _, err := s.nats.SubscribeReply(messaging.SubjectNodeBackendUpgrade(s.nodeID), s.handleBackendUpgrade); err != nil { + return fmt.Errorf("subscribing to backend upgrade events: %w", err) + } + if _, err := s.nats.Subscribe(messaging.SubjectNodeBackendStop(s.nodeID), s.handleBackendStop); err != nil { + return fmt.Errorf("subscribing to backend stop events: %w", err) + } + if _, err := s.nats.SubscribeReply(messaging.SubjectNodeBackendDelete(s.nodeID), s.handleBackendDelete); err != nil { + return fmt.Errorf("subscribing to backend delete events: %w", err) + } + if _, err := s.nats.SubscribeReply(messaging.SubjectNodeBackendList(s.nodeID), s.handleBackendList); err != nil { + return fmt.Errorf("subscribing to backend list events: %w", err) + } + if _, err := s.nats.SubscribeReply(messaging.SubjectNodeModelUnload(s.nodeID), s.handleModelUnload); err != nil { + return fmt.Errorf("subscribing to model unload events: %w", err) + } + if _, err := s.nats.SubscribeReply(messaging.SubjectNodeModelDelete(s.nodeID), s.handleModelDelete); err != nil { + return fmt.Errorf("subscribing to model delete events: %w", err) + } + if _, err := s.nats.Subscribe(messaging.SubjectNodeStop(s.nodeID), s.handleNodeStop); err != nil { + return fmt.Errorf("subscribing to node stop events: %w", err) + } + return nil } // handleBackendInstall is the NATS callback for backend.install — install @@ -66,9 +83,14 @@ func (s *backendSupervisor) handleBackendInstall(data []byte, reply func([]byte) advertiseAddr := addr advAddr := s.cfg.advertiseAddr() if advAddr != addr { - _, port, _ := net.SplitHostPort(addr) - advertiseHost, _, _ := net.SplitHostPort(advAddr) - advertiseAddr = net.JoinHostPort(advertiseHost, port) + _, port, err := net.SplitHostPort(addr) + if err != nil { + xlog.Error("Failed to parse backend listen address; using it unchanged", "addr", addr, "error", err) + } else if advertiseHost, _, err := net.SplitHostPort(advAddr); err != nil { + xlog.Error("Failed to parse worker advertise address; using backend listen address", "addr", advAddr, "error", err) + } else { + advertiseAddr = net.JoinHostPort(advertiseHost, port) + } } resp := messaging.BackendInstallReply{Success: true, Address: advertiseAddr} replyJSON(reply, resp) @@ -104,17 +126,29 @@ func (s *backendSupervisor) handleBackendUpgrade(data []byte, reply func([]byte) // handleBackendStop is the NATS callback for backend.stop — stop a specific // backend process (fire-and-forget, no reply expected). func (s *backendSupervisor) handleBackendStop(data []byte) { - // Try to parse backend name from payload; if empty, stop all - var req struct { - Backend string `json:"backend"` + req, stopAll, err := decodeBackendStopRequest(data) + if err != nil { + xlog.Error("Ignoring malformed NATS backend.stop event", "error", err) + return } - if json.Unmarshal(data, &req) == nil && req.Backend != "" { - xlog.Info("Received NATS backend.stop event", "backend", req.Backend) - s.stopBackend(req.Backend) - } else { - xlog.Info("Received NATS backend.stop event (all)") - s.stopAllBackends() + if stopAll { + xlog.Info("Received NATS backend.stop event (all)", "force", req.Force) + s.stopAllBackends(req.Force) + return } + xlog.Info("Received NATS backend.stop event", "backend", req.Backend, "force", req.Force) + s.stopBackend(req.Backend, req.Force) +} + +func decodeBackendStopRequest(data []byte) (messaging.BackendStopRequest, bool, error) { + if len(data) == 0 { + return messaging.BackendStopRequest{}, true, nil + } + var req messaging.BackendStopRequest + if err := json.Unmarshal(data, &req); err != nil { + return messaging.BackendStopRequest{}, false, fmt.Errorf("decoding backend stop request: %w", err) + } + return req, req.Backend == "", nil } // handleBackendDelete is the NATS callback for backend.delete — stop the @@ -130,7 +164,7 @@ func (s *backendSupervisor) handleBackendDelete(data []byte, reply func([]byte)) // Stop if running this backend if s.isRunning(req.Backend) { - s.stopBackend(req.Backend) + s.stopBackend(req.Backend, false) } // Delete the backend files @@ -142,7 +176,11 @@ func (s *backendSupervisor) handleBackendDelete(data []byte, reply func([]byte)) } // Re-register backends after deletion - gallery.RegisterBackends(s.systemState, s.ml) + if err := gallery.RegisterBackends(s.systemState, s.ml); err != nil { + xlog.Error("Failed to refresh registered backends after deletion", "backend", req.Backend, "error", err) + replyJSON(reply, messaging.BackendDeleteReply{Success: false, Error: err.Error()}) + return + } resp := messaging.BackendDeleteReply{Success: true} replyJSON(reply, resp) @@ -217,11 +255,14 @@ func (s *backendSupervisor) handleModelUnload(data []byte, reply func([]byte)) { } if targetAddr != "" { - // Best-effort gRPC Free() + // Best-effort bounded gRPC Free(). A model.unload request must not + // occupy the NATS reply handler forever when a backend is wedged. client := grpc.NewClientWithToken(targetAddr, false, nil, false, s.cfg.RegistrationToken) - if err := client.Free(context.Background()); err != nil { + freeCtx, cancel := context.WithTimeout(context.Background(), workerBackendFreeTimeout) + if err := client.Free(freeCtx); err != nil { xlog.Warn("Free() failed during model.unload", "error", err, "addr", targetAddr) } + cancel() } resp := messaging.ModelUnloadReply{Success: true} diff --git a/core/services/worker/registration.go b/core/services/worker/registration.go index 7d3bbf2442b1..6be19c9a8ebe 100644 --- a/core/services/worker/registration.go +++ b/core/services/worker/registration.go @@ -3,24 +3,36 @@ package worker import ( "cmp" "fmt" + "net" "os" "strconv" "strings" "github.com/mudler/LocalAI/pkg/xsysinfo" + "github.com/mudler/xlog" ) // effectiveBasePort returns the port used as base for gRPC backend processes. // Priority: Addr port → ServeAddr port → 50051 func (cfg *Config) effectiveBasePort() int { for _, addr := range []string{cfg.Addr, cfg.ServeAddr} { - if addr != "" { - if _, portStr, ok := strings.Cut(addr, ":"); ok { - if p, _ := strconv.Atoi(portStr); p > 0 { - return p - } - } + if addr == "" { + continue + } + _, portStr, err := net.SplitHostPort(addr) + if err != nil { + xlog.Warn("Invalid worker address; trying the next base-port source", "addr", addr, "error", err) + continue } + port, err := strconv.Atoi(portStr) + if err != nil { + xlog.Warn("Invalid worker port; trying the next base-port source", "addr", addr, "port", portStr, "error", err) + continue + } + if port > 0 && port <= 65535 { + return port + } + xlog.Warn("Worker port is outside the valid range; trying the next base-port source", "addr", addr, "port", port) } return 50051 } @@ -33,7 +45,10 @@ func (cfg *Config) advertiseAddr() string { if cfg.Addr != "" { return cfg.Addr } - hostname, _ := os.Hostname() + hostname, err := os.Hostname() + if err != nil { + xlog.Warn("Failed to determine worker hostname; advertising localhost", "error", err) + } return fmt.Sprintf("%s:%d", cmp.Or(hostname, "localhost"), cfg.effectiveBasePort()) } @@ -53,9 +68,14 @@ func (cfg *Config) advertiseHTTPAddr() string { if cfg.AdvertiseHTTPAddr != "" { return cfg.AdvertiseHTTPAddr } - advHost, _, _ := strings.Cut(cfg.advertiseAddr(), ":") + advertiseAddr := cfg.advertiseAddr() + advHost, _, err := net.SplitHostPort(advertiseAddr) + if err != nil { + xlog.Warn("Invalid worker advertise address; advertising file transfer on localhost", "addr", advertiseAddr, "error", err) + advHost = "localhost" + } httpPort := cfg.effectiveBasePort() - 1 - return fmt.Sprintf("%s:%d", advHost, httpPort) + return net.JoinHostPort(advHost, strconv.Itoa(httpPort)) } // registrationBody builds the JSON body for node registration. @@ -71,8 +91,14 @@ func (cfg *Config) registrationBody() map[string]any { } // Detect GPU info for VRAM-aware scheduling - totalVRAM, _ := xsysinfo.TotalAvailableVRAM() - gpuVendor, _ := xsysinfo.DetectGPUVendor() + totalVRAM, err := xsysinfo.TotalAvailableVRAM() + if err != nil { + xlog.Debug("Failed to detect worker VRAM; registering without GPU capacity", "error", err) + } + gpuVendor, err := xsysinfo.DetectGPUVendor() + if err != nil { + xlog.Debug("Failed to detect worker GPU vendor; registering without vendor metadata", "error", err) + } // Compute capability (e.g. "12.1" for GB10) lets the router pick per-arch // options (e.g. larger physical batch on Blackwell). Detected on the worker // because only the worker sees the GPU in distributed mode. @@ -103,7 +129,10 @@ func (cfg *Config) registrationBody() map[string]any { // If no GPU detected, report system RAM so the scheduler/UI has capacity info if totalVRAM == 0 { - if ramInfo, err := xsysinfo.GetSystemRAMInfo(); err == nil { + ramInfo, err := xsysinfo.GetSystemRAMInfo() + if err != nil { + xlog.Debug("Failed to detect worker RAM for registration", "error", err) + } else { body["total_ram"] = ramInfo.Total body["available_ram"] = ramInfo.Available } @@ -147,7 +176,10 @@ func (cfg *Config) heartbeatBody() map[string]any { // CPU-only workers (or workers that lost GPU visibility momentarily): // report system RAM so the scheduler still has capacity info. if aggregate.TotalVRAM == 0 { - if ramInfo, err := xsysinfo.GetSystemRAMInfo(); err == nil { + ramInfo, err := xsysinfo.GetSystemRAMInfo() + if err != nil { + xlog.Debug("Failed to detect worker RAM for heartbeat", "error", err) + } else { body["available_ram"] = ramInfo.Available } } diff --git a/core/services/worker/replica_test.go b/core/services/worker/replica_test.go index a8c0cbb0557b..1600b982ee1e 100644 --- a/core/services/worker/replica_test.go +++ b/core/services/worker/replica_test.go @@ -1,6 +1,10 @@ package worker import ( + "encoding/json" + + "github.com/mudler/LocalAI/core/services/messaging" + process "github.com/mudler/go-processmanager" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) @@ -118,4 +122,86 @@ var _ = Describe("Worker per-replica process keying", func() { Expect(s.resolveProcessKeys("qwen3.6-35B")).To(ConsistOf("qwen3.6-35B#0")) }) }) + + Describe("Backend stop sequencing", func() { + It("keeps the process and port reserved until termination completes", func() { + proc := process.New() + bp := &backendProcess{proc: proc, addr: "not-reparsed-during-cleanup", port: 50051} + s := &backendSupervisor{ + processes: map[string]*backendProcess{"model#0": bp}, + } + + claimed := s.beginBackendStop("model#0") + Expect(claimed).To(BeIdenticalTo(bp)) + Expect(bp.stopping).To(BeTrue()) + Expect(s.processes).To(HaveKeyWithValue("model#0", bp)) + Expect(s.freePorts).To(BeEmpty()) + + s.finishBackendStop("model#0", bp, nil) + Expect(s.processes).NotTo(HaveKey("model#0")) + Expect(s.freePorts).To(ConsistOf(50051)) + }) + + It("does not report a starting backend as ready after stop begins", func() { + bp := &backendProcess{port: 50051} + s := &backendSupervisor{ + processes: map[string]*backendProcess{"model#0": bp}, + } + + Expect(s.backendStartStillValid("model#0", bp)).To(BeTrue()) + bp.stopping = true + Expect(s.backendStartStillValid("model#0", bp)).To(BeFalse()) + }) + + It("recycles a failed startup port at most once", func() { + bp := &backendProcess{port: 50051} + s := &backendSupervisor{ + processes: map[string]*backendProcess{"model#0": bp}, + } + + s.releaseBackendStart("model#0", bp) + s.releaseBackendStart("model#0", bp) + + Expect(s.processes).NotTo(HaveKey("model#0")) + Expect(s.freePorts).To(ConsistOf(50051)) + }) + + It("does not recycle the port owned by a replacement startup", func() { + failed := &backendProcess{port: 50051} + replacement := &backendProcess{port: 50052} + s := &backendSupervisor{ + processes: map[string]*backendProcess{"model#0": replacement}, + } + + s.releaseBackendStart("model#0", failed) + + Expect(s.processes).To(HaveKeyWithValue("model#0", replacement)) + Expect(s.freePorts).To(BeEmpty()) + }) + }) + + Describe("backend.stop request decoding", func() { + It("preserves the legacy empty-payload stop-all command", func() { + req, stopAll, err := decodeBackendStopRequest(nil) + Expect(err).NotTo(HaveOccurred()) + Expect(stopAll).To(BeTrue()) + Expect(req).To(Equal(messaging.BackendStopRequest{})) + }) + + It("preserves force for a structured stop-all command", func() { + data, err := json.Marshal(messaging.BackendStopRequest{Force: true}) + Expect(err).NotTo(HaveOccurred()) + + req, stopAll, err := decodeBackendStopRequest(data) + Expect(err).NotTo(HaveOccurred()) + Expect(stopAll).To(BeTrue()) + Expect(req.Force).To(BeTrue()) + }) + + It("rejects malformed JSON instead of treating it as stop-all", func() { + _, stopAll, err := decodeBackendStopRequest([]byte(`{"backend":`)) + Expect(err).To(MatchError(ContainSubstring("decoding backend stop request"))) + Expect(stopAll).To(BeFalse()) + }) + }) }) diff --git a/core/services/worker/supervisor.go b/core/services/worker/supervisor.go index 5fda9bae8af2..702071b74d1c 100644 --- a/core/services/worker/supervisor.go +++ b/core/services/worker/supervisor.go @@ -4,10 +4,8 @@ import ( "context" "fmt" "maps" - "net" "os" "slices" - "strconv" "strings" "sync" "time" @@ -23,11 +21,14 @@ import ( // backendProcess represents a single gRPC backend process. type backendProcess struct { - proc *process.Process - backend string - addr string // gRPC address (host:port) + proc *process.Process + addr string // gRPC address (host:port) + port int + stopping bool } +const workerBackendFreeTimeout = 5 * time.Second + // backendSupervisor manages multiple backend gRPC processes on different ports. // Each backend type (e.g., llama-cpp, bert-embeddings) gets its own process and port. type backendSupervisor struct { @@ -60,6 +61,10 @@ func (s *backendSupervisor) startBackend(backend, backendPath string) (string, e // Already running? if bp, ok := s.processes[backend]; ok { + if bp.stopping { + s.mu.Unlock() + return "", fmt.Errorf("backend %s is stopping", backend) + } if bp.proc != nil && bp.proc.IsAlive() { s.mu.Unlock() return bp.addr, nil @@ -83,14 +88,15 @@ func (s *backendSupervisor) startBackend(backend, backendPath string) (string, e proc, err := s.ml.StartProcess(backendPath, backend, bindAddr) if err != nil { + s.freePorts = append(s.freePorts, port) s.mu.Unlock() return "", fmt.Errorf("starting backend process: %w", err) } s.processes[backend] = &backendProcess{ - proc: proc, - backend: backend, - addr: clientAddr, + proc: proc, + addr: clientAddr, + port: port, } xlog.Info("Backend process started", "backend", backend, "addr", clientAddr) @@ -111,31 +117,32 @@ func (s *backendSupervisor) startBackend(backend, backendPath string) (string, e readinessTimeout = 30 * time.Second ) deadline := time.Now().Add(readinessTimeout) + var lastHealthErr error for time.Now().Before(deadline) { time.Sleep(readinessPollInterval) ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) - if ok, _ := client.HealthCheck(ctx); ok { + ok, healthErr := client.HealthCheck(ctx) + if ok { cancel() - // Verify the process wasn't stopped/replaced while health-checking - s.mu.Lock() - currentBP, exists := s.processes[backend] - s.mu.Unlock() - if !exists || currentBP != bp { + // Verify the process wasn't stopped/replaced while health-checking. + // A stopping entry remains in the map until process termination so its + // port stays reserved, but it must not be advertised as ready. + if !s.backendStartStillValid(backend, bp) { return "", fmt.Errorf("backend %s was stopped during startup", backend) } xlog.Debug("Backend gRPC server is ready", "backend", backend, "addr", clientAddr) return clientAddr, nil } + if healthErr != nil { + lastHealthErr = healthErr + } cancel() // Check if the process died (e.g. OOM, CUDA error, missing libs) if !proc.IsAlive() { stderrTail := readLastLinesFromFile(proc.StderrPath(), 20) - xlog.Warn("Backend process died during startup", "backend", backend, "stderr", stderrTail) - s.mu.Lock() - delete(s.processes, backend) - s.freePorts = append(s.freePorts, port) - s.mu.Unlock() + xlog.Warn("Backend process died during startup", "backend", backend, "healthError", lastHealthErr, "stderr", stderrTail) + s.releaseBackendStart(backend, bp) return "", fmt.Errorf("backend process %s died during startup. Last stderr:\n%s", backend, stderrTail) } } @@ -146,17 +153,39 @@ func (s *backendSupervisor) startBackend(backend, backendPath string) (string, e // real cause). Stop the half-started process, recycle the port, and // surface the failure to the caller with the backend's stderr tail. stderrTail := readLastLinesFromFile(proc.StderrPath(), 20) - xlog.Error("Backend gRPC server not ready before deadline; aborting install", "backend", backend, "addr", clientAddr, "timeout", readinessTimeout, "stderr", stderrTail) + xlog.Error("Backend gRPC server not ready before deadline; aborting install", "backend", backend, "addr", clientAddr, "timeout", readinessTimeout, "healthError", lastHealthErr, "stderr", stderrTail) if killErr := proc.Stop(); killErr != nil { xlog.Warn("Failed to stop unready backend process", "backend", backend, "error", killErr) } + s.releaseBackendStart(backend, bp) + return "", fmt.Errorf("backend %s did not become ready within %s. Last stderr:\n%s", backend, readinessTimeout, stderrTail) +} + +// backendStartStillValid verifies that a successful readiness probe still +// belongs to the active startup attempt. Stop keeps an entry tracked while it +// terminates, so pointer identity alone is not enough. +func (s *backendSupervisor) backendStartStillValid(key string, bp *backendProcess) bool { s.mu.Lock() - if cur, ok := s.processes[backend]; ok && cur == bp { - delete(s.processes, backend) - s.freePorts = append(s.freePorts, port) + defer s.mu.Unlock() + current, exists := s.processes[key] + return exists && current == bp && !current.stopping +} + +// releaseBackendStart removes a failed startup and recycles its port only when +// the map still owns that exact attempt. A concurrent stop or replacement may +// already have removed it and recycled the port. +func (s *backendSupervisor) releaseBackendStart(key string, bp *backendProcess) { + s.mu.Lock() + defer s.mu.Unlock() + if current, exists := s.processes[key]; !exists || current != bp { + return } - s.mu.Unlock() - return "", fmt.Errorf("backend %s did not become ready within %s. Last stderr:\n%s", backend, readinessTimeout, stderrTail) + delete(s.processes, key) + if bp.port <= 0 { + xlog.Error("Cannot recycle backend port: startup has invalid recorded port", "backend", key, "addr", bp.addr, "port", bp.port) + return + } + s.freePorts = append(s.freePorts, bp.port) } // resolveProcessKeys turns a caller-supplied identifier into the set of @@ -193,51 +222,82 @@ func (s *backendSupervisor) resolveProcessKeys(id string) []string { // stopBackend stops the backend process(es) matching the given identifier. // Accepts a bare modelID (stops every replica) or a full processKey // (stops just that replica). -func (s *backendSupervisor) stopBackend(id string) { +func (s *backendSupervisor) stopBackend(id string, force bool) { for _, key := range s.resolveProcessKeys(id) { - s.stopBackendExact(key) + s.stopBackendExact(key, force) } } -// stopBackendExact stops the process under exactly this key. Locking and -// network I/O are split: the map mutation runs under the lock, the gRPC -// Free() and proc.Stop() calls run after release so they don't block -// other supervisor operations. -func (s *backendSupervisor) stopBackendExact(key string) { - s.mu.Lock() - bp, ok := s.processes[key] - if !ok || bp.proc == nil { - s.mu.Unlock() +// stopBackendExact stops the process under exactly this key. It marks the +// entry as stopping under the lock, then runs Free() and proc.Stop() after +// release so network I/O cannot block other supervisor operations. The entry +// and port remain reserved until process termination completes. +func (s *backendSupervisor) stopBackendExact(key string, force bool) { + bp := s.beginBackendStop(key) + if bp == nil { return } - delete(s.processes, key) - if _, portStr, err := net.SplitHostPort(bp.addr); err == nil { - if p, err := strconv.Atoi(portStr); err == nil { - s.freePorts = append(s.freePorts, p) + + if !force { + client := grpc.NewClientWithToken(bp.addr, false, nil, false, s.cfg.RegistrationToken) + freeCtx, cancel := context.WithTimeout(context.Background(), workerBackendFreeTimeout) + xlog.Debug("Calling bounded Free() before stopping backend", "backend", key, "timeout", workerBackendFreeTimeout) + if err := client.Free(freeCtx); err != nil { + xlog.Warn("Free() failed (best-effort)", "backend", key, "error", err) } + cancel() } - s.mu.Unlock() - client := grpc.NewClientWithToken(bp.addr, false, nil, false, s.cfg.RegistrationToken) - xlog.Debug("Calling Free() before stopping backend", "backend", key) - if err := client.Free(context.Background()); err != nil { - xlog.Warn("Free() failed (best-effort)", "backend", key, "error", err) + xlog.Info("Stopping backend process", "backend", key, "addr", bp.addr, "force", force) + stopErr := bp.proc.Stop() + if stopErr != nil { + xlog.Error("Error stopping backend process", "backend", key, "error", stopErr) } + s.finishBackendStop(key, bp, stopErr) +} - xlog.Info("Stopping backend process", "backend", key, "addr", bp.addr) - if err := bp.proc.Stop(); err != nil { - xlog.Error("Error stopping backend process", "backend", key, "error", err) +// beginBackendStop reserves both the process entry and its port while network +// cleanup and process termination run without the supervisor mutex. +func (s *backendSupervisor) beginBackendStop(key string) *backendProcess { + s.mu.Lock() + defer s.mu.Unlock() + bp, ok := s.processes[key] + if !ok || bp.proc == nil || bp.stopping { + return nil + } + bp.stopping = true + return bp +} + +func (s *backendSupervisor) finishBackendStop(key string, bp *backendProcess, stopErr error) { + // Keep the process and port reserved until termination completes. Recycling + // the port before this point can start a second backend on an address still + // owned by the stuck process. + s.mu.Lock() + defer s.mu.Unlock() + if current, exists := s.processes[key]; !exists || current != bp { + return + } + if stopErr != nil && bp.proc.IsAlive() { + bp.stopping = false + return + } + delete(s.processes, key) + if bp.port <= 0 { + xlog.Error("Cannot recycle backend port: process has invalid recorded port", "backend", key, "addr", bp.addr, "port", bp.port) + return } + s.freePorts = append(s.freePorts, bp.port) } // stopAllBackends stops all running backend processes. -func (s *backendSupervisor) stopAllBackends() { +func (s *backendSupervisor) stopAllBackends(force bool) { s.mu.Lock() backends := slices.Collect(maps.Keys(s.processes)) s.mu.Unlock() for _, b := range backends { - s.stopBackend(b) + s.stopBackend(b, force) } } diff --git a/core/services/worker/worker.go b/core/services/worker/worker.go index 9ed5de8f11a6..5cdafb7c0f76 100644 --- a/core/services/worker/worker.go +++ b/core/services/worker/worker.go @@ -51,7 +51,9 @@ func Run(ctx *cliContext.Context, cfg *Config) error { ml.SetBackendLoggingEnabled(true) // Register already-installed backends - gallery.RegisterBackends(systemState, ml) + if err := gallery.RegisterBackends(systemState, ml); err != nil { + return fmt.Errorf("registering installed backends: %w", err) + } // Parse galleries config var galleries []config.Gallery @@ -190,7 +192,10 @@ func Run(ctx *cliContext.Context, cfg *Config) error { // Set the registration token once before any backends are started if cfg.RegistrationToken != "" { - os.Setenv(grpc.AuthTokenEnvVar, cfg.RegistrationToken) + if err := os.Setenv(grpc.AuthTokenEnvVar, cfg.RegistrationToken); err != nil { + nodes.ShutdownFileTransferServer(httpServer) + return fmt.Errorf("setting backend authentication token: %w", err) + } } supervisor := &backendSupervisor{ @@ -204,12 +209,16 @@ func Run(ctx *cliContext.Context, cfg *Config) error { processes: make(map[string]*backendProcess), nextPort: basePort, } - supervisor.subscribeLifecycleEvents() + if err := supervisor.subscribeLifecycleEvents(); err != nil { + nodes.ShutdownFileTransferServer(httpServer) + return fmt.Errorf("subscribing to worker lifecycle events: %w", err) + } // Subscribe to file staging NATS subjects if S3 is configured if cfg.StorageURL != "" { if err := cfg.subscribeFileStaging(natsClient, nodeID); err != nil { - xlog.Error("Failed to subscribe to file staging subjects", "error", err) + nodes.ShutdownFileTransferServer(httpServer) + return fmt.Errorf("subscribing to file staging subjects: %w", err) } } @@ -228,7 +237,7 @@ func Run(ctx *cliContext.Context, cfg *Config) error { xlog.Info("Shutting down worker") shutdownCancel() // stop heartbeat loop immediately regClient.GracefulDeregister(nodeID) - supervisor.stopAllBackends() + supervisor.stopAllBackends(false) nodes.ShutdownFileTransferServer(httpServer) return runErr } diff --git a/docs/content/advanced/vram-management.md b/docs/content/advanced/vram-management.md index bb7583a07b2b..6b9b47569c24 100644 --- a/docs/content/advanced/vram-management.md +++ b/docs/content/advanced/vram-management.md @@ -246,6 +246,11 @@ Via web UI: Navigate to Settings → Watchdog Settings and enable "Watchdog Idle The busy watchdog monitors models that have been processing requests for an unusually long time and terminates them if they exceed a threshold. This is useful for detecting and recovering from stuck or hung backends. +For parallel backends, the timeout is measured from the start of the oldest +request that is still in flight. Sustained overlapping traffic does not by +itself make a backend stale: when the oldest request completes, the watchdog +continues from the start time of the next-oldest request. + #### Configuration Via environment variables or CLI: @@ -259,6 +264,18 @@ LOCALAI_WATCHDOG_BUSY=true LOCALAI_WATCHDOG_BUSY_TIMEOUT=10m ./local-ai Via web UI: Navigate to Settings → Watchdog Settings and enable "Watchdog Busy Enabled" with your desired timeout. +#### Backend shutdown behavior + +Administrative and API-triggered model shutdown is graceful by default. It +waits up to 30 seconds for in-flight requests to finish and returns a +`model is still busy` error if the deadline expires. This bounded wait keeps a +busy backend from blocking lifecycle operations for other models. + +Set `LOCALAI_FORCE_BACKEND_SHUTDOWN=true` when starting LocalAI to escalate a +timed-out graceful shutdown to forced process termination. Forced shutdown can +interrupt active requests and skips the backend `Free()` RPC; terminating the +process releases its resources instead. + ### Combined Configuration You can enable both watchdogs simultaneously for comprehensive memory management: @@ -440,4 +457,3 @@ Conversely, you can pre-load a model into memory ahead of its first request with - See [Advanced Usage]({{%relref "advanced/advanced-usage" %}}) for other configuration options - See [GPU Acceleration]({{%relref "features/GPU-acceleration" %}}) for GPU setup and configuration - See [Backend Flags]({{%relref "advanced/advanced-usage#backend-flags" %}}) for all available backend configuration options - diff --git a/formal-verification/README.md b/formal-verification/README.md index 7f38e1462088..a13ab1e3cd9c 100644 --- a/formal-verification/README.md +++ b/formal-verification/README.md @@ -1,12 +1,14 @@ -# Formal verification — realtime state machines +# Formal verification -Formal designs (FizzBee specs) for the realtime API state machines and the harness -that keeps the Go implementation provably in step with them. Background and -rationale: [../docs/design/realtime-state-machines.md](../docs/design/realtime-state-machines.md) (Part 6). +Formal designs expressed as FizzBee specs. +Most specs cover the realtime API state machines; `model_loader_shutdown.fizz` +covers the shared model-loader lifecycle that all modalities use. Realtime +background and rationale: +[../docs/design/realtime-state-machines.md](../docs/design/realtime-state-machines.md) +(Part 6). -The design is **authoritative**: behaviour changes go through the spec first, then -the implementation is checked against it. The `realtime-conformance` gate makes -that the path of least resistance — you cannot land a non-conforming change green. +The designs are **authoritative**: behaviour changes go through the spec first, +then the implementation is checked against them. ## What's here @@ -18,6 +20,7 @@ that the path of least resistance — you cannot land a non-conforming change gr | `compaction.fizz` | **Authoritative** FizzBee model of machine M4 (conversation compaction): single-flight. | | `tts_pipeline.fizz` | **Authoritative** FizzBee model of machine M5 (TTS pipeline): open->closing->closed, idempotent close. | | `session_lifecycle.fizz` | **Composition** spec: the M1–M5 hierarchy — conn (M1) is the parent; when it is torn down, every child (vad/M2, resp/M3, compaction/M4) is terminal. Models the relationship the per-machine specs can't express. | +| `model_loader_shutdown.fizz` | **Authoritative** shared lifecycle design: bounded busy/Free waits, distributed force propagation/port reservation, and parallel in-flight accounting. Checked via `make test-model-lifecycle-conformance`. | | `fizzbee.sha256` | Pinned checksum(s) of the FizzBee release the gate uses (created on first `install-fizzbee.sh` run). | The implementations under test live in @@ -27,7 +30,7 @@ The implementations under test live in [`core/http/endpoints/openai/compactcoord`](../../../core/http/endpoints/openai/compactcoord) (M4), and [`core/http/endpoints/openai/ttscoord`](../../../core/http/endpoints/openai/ttscoord) (M5). -## Running the gate +## Running the realtime gate ```sh make test-realtime-conformance @@ -56,6 +59,23 @@ local work on unrelated code. CI never sets it, and `pre-commit` runs the full gate whenever `respcoord/**`, `turncoord/**`, `conncoord/**`, `compactcoord/**`, `ttscoord/**`, or `formal-verification/**` is staged (so a pure `.fizz` edit still re-verifies). +## Running the model-loader lifecycle gate + +```sh +make test-model-lifecycle-conformance +# or directly: +./scripts/model-lifecycle-conformance.sh +``` + +This focused, fail-closed gate runs the loader, gRPC client, distributed +unloader, and worker lifecycle specs under the race detector and then checks +`model_loader_shutdown.fizz`. The model proves that local force and bounded +graceful busy waiting cannot wedge unrelated loads, distributed force skips +`Free()` and reserves the worker port until termination, and parallel request +tracking stays busy until the final request completes. Process termination is +an explicit fairness assumption; the concrete Go tests connect the abstract +model to the implementation paths. + ## Installing FizzBee (pinned) FizzBee is pre-1.0 and single-maintainer, so we pin a version + sha256 and use the @@ -127,6 +147,11 @@ reproduce the legacy bug it guards against: teardown cancelling + joining each conversation's compaction (and respcoord/ compactcoord gained an absorbing `Terminated` state so no child can start after teardown). +- `model_loader_shutdown.fizz`: in `ForceShutdown`, delete `self.backend = 2` to + violate local progress; move `self.port_recycled = 1` from `ProcessStops` to + `WorkerReceivesStop` to violate distributed port safety; or set `self.busy = 0` + in `FinishOne` to violate parallel busy accounting. The checker rejects each + mutation. ## Adding another machine diff --git a/formal-verification/model_loader_shutdown.fizz b/formal-verification/model_loader_shutdown.fizz new file mode 100644 index 000000000000..d60793420e80 --- /dev/null +++ b/formal-verification/model_loader_shutdown.fizz @@ -0,0 +1,195 @@ +--- +# Formal design of the model-loader shutdown paths exercised by pkg/model and +# distributed workers. This is deliberately not a realtime model: realtime +# merely made a loader-wide stall visible while loading the shared opus backend. +# +# The model distinguishes four cases: +# 1. Local forced shutdown makes progress. +# 2. Graceful busy waiting is bounded without holding the global loader lock. +# 3. Distributed force reaches the worker, skips Free(), and reserves its port +# until process termination completes. +# 4. Parallel request tracking remains busy until every request completes. +deadlock_detection: false +options: + max_actions: 30 + max_concurrent_actions: 2 +--- + +role LocalForce: + action Init: + self.backend = 1 # 1 = stuck busy, 2 = stopped + self.timed_out = 0 + self.loader = 0 # 0 = free, 1 = held by shutdown + self.other = 0 # 0 = absent, 1 = pending, 2 = loaded + + atomic action RequestOtherLoad: + if self.other == 0: + self.other = 1 + + atomic fair action BusyTimeout: + if self.timed_out == 0: + self.timed_out = 1 + + # ShutdownModelForce skips both the IsBusy loop and Free RPC. Stopping the + # local process and releasing the loader are one abstract transition; the + # process manager's eventual termination is the fairness assumption. + atomic fair action ForceShutdown: + if self.timed_out == 1 and self.backend == 1: + self.loader = 1 + self.backend = 2 + self.loader = 0 + + atomic fair action CompleteOtherLoad: + if self.other == 1 and self.loader == 0: + self.other = 2 + +role GracefulShutdown: + action Init: + self.backend = 1 # permanently stuck busy + self.waiting = 0 + self.done = 0 # bounded wait returned ErrModelBusy + self.global_loader = 0 + self.other = 0 + + atomic action RequestOtherLoad: + if self.other == 0: + self.other = 1 + + atomic action BeginGracefulShutdown: + if self.waiting == 0 and self.done == 0: + self.waiting = 1 + + # The backend is allowed to remain busy forever. The shutdown context is + # the fairness boundary: expiry returns ErrModelBusy and releases only the + # per-model operation lock; the global loader lock was never acquired. + atomic fair action BusyDeadline: + if self.waiting == 1: + self.waiting = 0 + self.done = 1 + + atomic fair action CompleteOtherLoad: + if self.other == 1 and self.global_loader == 0: + self.other = 2 + +role DistributedForce: + action Init: + self.timed_out = 0 + self.stop_sent = 0 + self.process = 1 # 1 = running + self.supervisor_tracked = 1 + self.port_recycled = 0 + self.stopping = 0 + self.free_called = 0 + self.reinstall = 0 + + atomic fair action BusyTimeout: + if self.timed_out == 0: + self.timed_out = 1 + + # The controller publishes force=true in BackendStopRequest. + atomic fair action SendRemoteStop: + if self.timed_out == 1 and self.stop_sent == 0: + self.stop_sent = 1 + + # Force marks the process as stopping but keeps both supervisor ownership + # and the port reservation. The Free RPC is skipped. + atomic fair action WorkerReceivesStop: + if self.stop_sent == 1 and self.supervisor_tracked == 1: + self.stopping = 1 + + # Only completed process termination releases supervisor ownership and the + # port. Eventual termination is the process-manager fairness assumption. + atomic fair action ProcessStops: + if self.stopping == 1 and self.process == 1: + self.process = 0 + self.supervisor_tracked = 0 + self.port_recycled = 1 + + atomic action RequestReinstall: + if self.port_recycled == 1 and self.reinstall == 0: + self.reinstall = 1 + +role ParallelBusyTracker: + action Init: + self.inflight = 0 + self.busy = 0 + + atomic action StartFirst: + if self.inflight == 0: + self.inflight = 1 + self.busy = 1 + + atomic action StartSecond: + if self.inflight == 1 and self.busy == 1: + self.inflight = 2 + self.busy = 1 + + atomic action FinishOne: + if self.inflight == 2: + self.inflight = 1 + self.busy = 1 + + atomic action FinishLast: + if self.inflight == 1: + self.inflight = 0 + self.busy = 0 + +action Init: + local = LocalForce() + graceful = GracefulShutdown() + remote = DistributedForce() + tracker = ParallelBusyTracker() + +# Proven behavior of the implementation's local force path. Once selected by the watchdog, +# the stuck backend is eventually stopped and an unrelated pending load can +# eventually acquire the loader. +always eventually assertion LocalTimedOutBackendStops: + return local.timed_out == 0 or local.backend == 2 + +always eventually assertion LocalUnrelatedLoadProgresses: + return local.other != 1 + +always assertion LocalNeverWaitsBusyWithLoaderHeld: + return not (local.loader == 1 and local.backend == 1) + +exists assertion LocalForcePathExercised: + return local.timed_out == 1 and local.backend == 2 + +exists assertion LocalOtherLoadCompletes: + return local.other == 2 + +# The permanently busy backend cannot wedge the global loader, and its bounded +# busy-wait attempt eventually returns. +always assertion GracefulNeverHoldsGlobalLoader: + return graceful.global_loader == 0 + +always eventually assertion GracefulShutdownIsBounded: + return graceful.waiting == 0 + +always eventually assertion GracefulUnrelatedLoadProgresses: + return graceful.other != 1 + +exists assertion GracefulDeadlineExercised: + return graceful.done == 1 and graceful.other == 2 + +# Distributed force skips Free, terminates the process, and never recycles its +# port while the old process can still own it. +always assertion DistributedForceSkipsFree: + return remote.free_called == 0 + +always assertion DistributedPortReservedUntilStop: + return not (remote.port_recycled == 1 and remote.process == 1) + +always eventually assertion DistributedForcedStopProgresses: + return remote.stop_sent == 0 or remote.process == 0 + +exists assertion DistributedForcedStopExercised: + return remote.process == 0 and remote.port_recycled == 1 and remote.free_called == 0 + +# Busy is derived from the in-flight count, so a partial completion cannot +# expose a false-idle state. +always assertion ParallelBusyMatchesInflight: + return (tracker.inflight > 0 and tracker.busy == 1) or (tracker.inflight == 0 and tracker.busy == 0) + +exists assertion ParallelOverlapPreserved: + return tracker.inflight == 1 and tracker.busy == 1 diff --git a/pkg/grpc/client.go b/pkg/grpc/client.go index 62865d53eabc..5ea6d8c625aa 100644 --- a/pkg/grpc/client.go +++ b/pkg/grpc/client.go @@ -29,7 +29,7 @@ func (b bearerToken) RequireTransportSecurity() bool { return false } type Client struct { address string - busy bool + inFlight int parallel bool token string sync.Mutex @@ -38,32 +38,33 @@ type Client struct { } type WatchDog interface { - Mark(address string) - UnMark(address string) + TrackRequest(address string) func() } func (c *Client) IsBusy() bool { c.Lock() defer c.Unlock() - return c.busy + return c.inFlight > 0 } +// setBusy preserves the existing call-site shape while maintaining a count. +// Parallel requests can finish in any order, so a boolean would let the first +// completion report the backend idle while other calls were still running. func (c *Client) setBusy(v bool) { c.Lock() - c.busy = v - c.Unlock() -} - -func (c *Client) wdMark() { - if c.wd != nil { - c.wd.Mark(c.address) + if v { + c.inFlight++ + } else if c.inFlight > 0 { + c.inFlight-- } + c.Unlock() } -func (c *Client) wdUnMark() { +func (c *Client) wdMark() func() { if c.wd != nil { - c.wd.UnMark(c.address) + return c.wd.TrackRequest(c.address) } + return func() {} } // dial creates a gRPC client connection with common options. @@ -119,8 +120,7 @@ func (c *Client) Embeddings(ctx context.Context, in *pb.PredictOptions, opts ... } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -138,8 +138,7 @@ func (c *Client) Predict(ctx context.Context, in *pb.PredictOptions, opts ...grp } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -157,8 +156,7 @@ func (c *Client) LoadModel(ctx context.Context, in *pb.ModelOptions, opts ...grp } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -175,8 +173,7 @@ func (c *Client) PredictStream(ctx context.Context, in *pb.PredictOptions, f fun } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return err @@ -223,8 +220,7 @@ func (c *Client) GenerateImage(ctx context.Context, in *pb.GenerateImageRequest, } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -241,8 +237,7 @@ func (c *Client) GenerateVideo(ctx context.Context, in *pb.GenerateVideoRequest, } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -259,8 +254,7 @@ func (c *Client) TTS(ctx context.Context, in *pb.TTSRequest, opts ...grpc.CallOp } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -277,8 +271,7 @@ func (c *Client) TTSStream(ctx context.Context, in *pb.TTSRequest, f func(reply } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return err @@ -323,8 +316,7 @@ func (c *Client) SoundGeneration(ctx context.Context, in *pb.SoundGenerationRequ } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -341,8 +333,7 @@ func (c *Client) AudioTranscription(ctx context.Context, in *pb.TranscriptReques } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -359,8 +350,7 @@ func (c *Client) AudioTranscriptionStream(ctx context.Context, in *pb.Transcript } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return err @@ -403,8 +393,7 @@ func (c *Client) TokenizeString(ctx context.Context, in *pb.PredictOptions, opts } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -443,8 +432,7 @@ func (c *Client) StoresSet(ctx context.Context, in *pb.StoresSetOptions, opts .. } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -459,8 +447,7 @@ func (c *Client) StoresDelete(ctx context.Context, in *pb.StoresDeleteOptions, o c.opMutex.Lock() defer c.opMutex.Unlock() } - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() c.setBusy(true) defer c.setBusy(false) conn, err := c.dial() @@ -479,8 +466,7 @@ func (c *Client) StoresGet(ctx context.Context, in *pb.StoresGetOptions, opts .. } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -497,8 +483,7 @@ func (c *Client) StoresFind(ctx context.Context, in *pb.StoresFindOptions, opts } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -515,8 +500,7 @@ func (c *Client) Rerank(ctx context.Context, in *pb.RerankRequest, opts ...grpc. } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -533,8 +517,7 @@ func (c *Client) TokenClassify(ctx context.Context, in *pb.TokenClassifyRequest, } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -551,8 +534,7 @@ func (c *Client) Score(ctx context.Context, in *pb.ScoreRequest, opts ...grpc.Ca } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -569,8 +551,7 @@ func (c *Client) GetTokenMetrics(ctx context.Context, in *pb.MetricsRequest, opt } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -587,8 +568,7 @@ func (c *Client) VAD(ctx context.Context, in *pb.VADRequest, opts ...grpc.CallOp } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -605,8 +585,7 @@ func (c *Client) Diarize(ctx context.Context, in *pb.DiarizeRequest, opts ...grp } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -623,8 +602,7 @@ func (c *Client) SoundDetection(ctx context.Context, in *pb.SoundDetectionReques } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -641,8 +619,7 @@ func (c *Client) Detect(ctx context.Context, in *pb.DetectOptions, opts ...grpc. } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -659,8 +636,7 @@ func (c *Client) Depth(ctx context.Context, in *pb.DepthRequest, opts ...grpc.Ca } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -677,8 +653,7 @@ func (c *Client) FaceVerify(ctx context.Context, in *pb.FaceVerifyRequest, opts } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -695,8 +670,7 @@ func (c *Client) FaceAnalyze(ctx context.Context, in *pb.FaceAnalyzeRequest, opt } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -713,8 +687,7 @@ func (c *Client) VoiceVerify(ctx context.Context, in *pb.VoiceVerifyRequest, opt } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -731,8 +704,7 @@ func (c *Client) VoiceAnalyze(ctx context.Context, in *pb.VoiceAnalyzeRequest, o } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -749,8 +721,7 @@ func (c *Client) VoiceEmbed(ctx context.Context, in *pb.VoiceEmbedRequest, opts } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -767,8 +738,7 @@ func (c *Client) AudioEncode(ctx context.Context, in *pb.AudioEncodeRequest, opt } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -785,8 +755,7 @@ func (c *Client) AudioDecode(ctx context.Context, in *pb.AudioDecodeRequest, opt } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -803,8 +772,7 @@ func (c *Client) AudioTransform(ctx context.Context, in *pb.AudioTransformReques } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -857,10 +825,10 @@ func (c *Client) Forward(ctx context.Context, opts ...grpc.CallOption) (ForwardC c.opMutex.Lock() } c.setBusy(true) - c.wdMark() + completeRequest := c.wdMark() cleanup := func() { - c.wdUnMark() + completeRequest() c.setBusy(false) if !c.parallel { c.opMutex.Unlock() @@ -923,10 +891,10 @@ func (c *Client) AudioTransformStream(ctx context.Context, opts ...grpc.CallOpti c.opMutex.Lock() } c.setBusy(true) - c.wdMark() + completeRequest := c.wdMark() cleanup := func() { - c.wdUnMark() + completeRequest() c.setBusy(false) if !c.parallel { c.opMutex.Unlock() @@ -1002,10 +970,10 @@ func (c *Client) AudioTranscriptionLive(ctx context.Context, opts ...grpc.CallOp c.opMutex.Lock() } c.setBusy(true) - c.wdMark() + completeRequest := c.wdMark() cleanup := func() { - c.wdUnMark() + completeRequest() c.setBusy(false) if !c.parallel { c.opMutex.Unlock() @@ -1066,10 +1034,10 @@ func (c *Client) AudioToAudioStream(ctx context.Context, opts ...grpc.CallOption c.opMutex.Lock() } c.setBusy(true) - c.wdMark() + completeRequest := c.wdMark() cleanup := func() { - c.wdUnMark() + completeRequest() c.setBusy(false) if !c.parallel { c.opMutex.Unlock() @@ -1104,8 +1072,7 @@ func (c *Client) StartFineTune(ctx context.Context, in *pb.FineTuneRequest, opts } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -1122,8 +1089,7 @@ func (c *Client) FineTuneProgress(ctx context.Context, in *pb.FineTuneProgressRe } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return err @@ -1166,8 +1132,7 @@ func (c *Client) StopFineTune(ctx context.Context, in *pb.FineTuneStopRequest, o } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -1184,8 +1149,7 @@ func (c *Client) ListCheckpoints(ctx context.Context, in *pb.ListCheckpointsRequ } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -1202,8 +1166,7 @@ func (c *Client) ExportModel(ctx context.Context, in *pb.ExportModelRequest, opt } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -1220,8 +1183,7 @@ func (c *Client) StartQuantization(ctx context.Context, in *pb.QuantizationReque } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -1238,8 +1200,7 @@ func (c *Client) QuantizationProgress(ctx context.Context, in *pb.QuantizationPr } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return err @@ -1282,8 +1243,7 @@ func (c *Client) StopQuantization(ctx context.Context, in *pb.QuantizationStopRe } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err @@ -1300,8 +1260,7 @@ func (c *Client) Free(ctx context.Context) error { } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { @@ -1321,8 +1280,7 @@ func (c *Client) ModelMetadata(ctx context.Context, in *pb.ModelOptions, opts .. } c.setBusy(true) defer c.setBusy(false) - c.wdMark() - defer c.wdUnMark() + defer c.wdMark()() conn, err := c.dial() if err != nil { return nil, err diff --git a/pkg/grpc/client_busy_test.go b/pkg/grpc/client_busy_test.go new file mode 100644 index 000000000000..b989b87a9519 --- /dev/null +++ b/pkg/grpc/client_busy_test.go @@ -0,0 +1,27 @@ +package grpc + +import ( + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Client busy accounting", func() { + It("remains busy until every parallel operation completes", func() { + client := &Client{parallel: true} + + client.setBusy(true) + client.setBusy(true) + client.setBusy(false) + Expect(client.IsBusy()).To(BeTrue()) + + client.setBusy(false) + Expect(client.IsBusy()).To(BeFalse()) + }) + + It("does not underflow on a redundant completion", func() { + client := &Client{parallel: true} + client.setBusy(false) + client.setBusy(true) + Expect(client.IsBusy()).To(BeTrue()) + }) +}) diff --git a/pkg/model/loader.go b/pkg/model/loader.go index b64331d3c295..01fc9f8b51b2 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -2,6 +2,7 @@ package model import ( "context" + "errors" "fmt" "maps" "os" @@ -33,6 +34,13 @@ type RemoteModelUnloader interface { UnloadRemoteModel(modelName string) error } +// RemoteModelContextUnloader is the optional cancellation/force extension used +// by distributed workers. Keeping it separate preserves compatibility with +// existing RemoteModelUnloader implementations. +type RemoteModelContextUnloader interface { + UnloadRemoteModelContext(ctx context.Context, modelName string, force bool) error +} + // ModelRouter is a callback that routes model loading to a remote node // instead of starting a local process. When set on the ModelLoader, // grpcModel() will delegate to this function before attempting local loading. @@ -65,6 +73,7 @@ type ModelLoader struct { mu sync.Mutex store ModelStore loading map[string]chan struct{} // tracks models currently being loaded + operations *modelOperationLocks wd *WatchDog externalBackends map[string]string lruEvictionMaxRetries int // Maximum number of retries when waiting for busy models @@ -118,6 +127,7 @@ func NewModelLoader(system *system.SystemState) *ModelLoader { ModelPath: system.Model.ModelsPath, store: NewInMemoryModelStore(), loading: make(map[string]chan struct{}), + operations: newModelOperationLocks(), externalBackends: make(map[string]string), lruEvictionMaxRetries: 30, // Default: 30 retries lruEvictionRetryInterval: 1 * time.Second, // Default: 1 second @@ -257,6 +267,8 @@ func (ml *ModelLoader) SetModelStore(s ModelStore) { } func (ml *ModelLoader) GetWatchDog() *WatchDog { + ml.mu.Lock() + defer ml.mu.Unlock() return ml.wd } @@ -303,6 +315,8 @@ func (ml *ModelLoader) GetExternalBackend(name string) string { } func (ml *ModelLoader) GetAllExternalBackends(o *Options) map[string]string { + ml.mu.Lock() + defer ml.mu.Unlock() backends := make(map[string]string) maps.Copy(backends, ml.externalBackends) if o != nil { @@ -341,8 +355,6 @@ var knownModelsNameSuffixToSkip []string = []string{ ".tar.gz", } -const retryTimeout = time.Duration(2 * time.Minute) - func (ml *ModelLoader) ListFilesInModelPath() ([]string, error) { files, err := os.ReadDir(ml.ModelPath) if err != nil { @@ -385,10 +397,11 @@ FILE: func (ml *ModelLoader) ListLoadedModels() []*Model { ml.mu.Lock() - defer ml.mu.Unlock() + store := ml.store + ml.mu.Unlock() models := []*Model{} - ml.store.Range(func(_ string, m *Model) bool { + store.Range(func(_ string, m *Model) bool { models = append(models, m) return true }) @@ -404,6 +417,8 @@ func (ml *ModelLoader) LoadModelWithFile(modelID, modelName, modelFileName strin if modelFileName == "" { modelFileName = modelName } + release := ml.operations.acquire(modelID, false) + defer release() return ml.loadModel(modelID, modelName, modelFileName, loader, true) } @@ -455,10 +470,16 @@ func (ml *ModelLoader) loadModel(modelID, modelName, modelFileName string, loade return model, nil } - ml.mu.Lock() - // Check if we already have a loaded model if model := ml.checkIsLoaded(modelID); model != nil { + return model, nil + } + + ml.mu.Lock() + // A concurrent leader can publish its model between the health check above + // and this lock acquisition. It is fresh by definition, so avoid starting a + // duplicate load without performing another network probe under ml.mu. + if model, ok := ml.store.Get(modelID); ok { ml.mu.Unlock() return model, nil } @@ -483,9 +504,7 @@ func (ml *ModelLoader) loadModel(modelID, modelName, modelFileName string, loade xlog.Debug("Waiting for model to be loaded by another request", "modelID", modelID) <-loadingChan // Now check if the model is loaded - ml.mu.Lock() model := ml.checkIsLoaded(modelID) - ml.mu.Unlock() if model != nil { return model, nil } @@ -532,36 +551,65 @@ func (ml *ModelLoader) loadModel(modelID, modelName, modelFileName string, loade } func (ml *ModelLoader) ShutdownModel(modelName string) error { - ml.mu.Lock() - defer ml.mu.Unlock() - - return ml.deleteProcess(modelName, false) + ctx, cancel := context.WithTimeout(context.Background(), gracefulShutdownTimeout) + defer cancel() + err := ml.shutdownModel(ctx, modelName, false) + if errors.Is(err, ErrModelBusy) && forceBackendShutdown { + xlog.Warn("Model remained busy until the graceful shutdown deadline; forcing shutdown", "model", modelName) + return ml.ShutdownModelForce(modelName) + } + return err } // ShutdownModelForce stops a backend without waiting for an in-flight gRPC // call to finish first. It is used by the watchdog's busy-killer, which only // fires once a backend has been stuck on a call past the busy timeout — the -// graceful ShutdownModel would block forever on that stuck call (while -// holding ml.mu), preventing every other model load. See deleteProcess. +// graceful ShutdownModel intentionally waits for active calls until its +// bounded deadline. See deleteProcess. func (ml *ModelLoader) ShutdownModelForce(modelName string) error { - ml.mu.Lock() - defer ml.mu.Unlock() + ctx, cancel := context.WithTimeout(context.Background(), forcedShutdownTimeout) + defer cancel() + return ml.shutdownModel(ctx, modelName, true) +} + +// 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 { + return ml.shutdownModel(ctx, modelName, force) +} - return ml.deleteProcess(modelName, true) +func (ml *ModelLoader) shutdownModel(ctx context.Context, modelName string, force bool) error { + release, err := ml.operations.acquireContext(ctx, modelName, true) + if err != nil { + return fmt.Errorf("waiting to shut down model %q: %w", modelName, err) + } + defer release() + return ml.deleteProcess(ctx, modelName, force) } func (ml *ModelLoader) CheckIsLoaded(s string) *Model { - ml.mu.Lock() - defer ml.mu.Unlock() + release := ml.operations.acquire(s, false) + defer release() return ml.checkIsLoaded(s) } func (ml *ModelLoader) checkIsLoaded(s string) *Model { - m, ok := ml.store.Get(s) + ml.mu.Lock() + store := ml.store + wd := ml.wd + ml.mu.Unlock() + + m, ok := store.Get(s) if !ok { return nil } + // Health checks are per model. Serializing them here prevents a cold health + // cache from generating a burst of identical probes without blocking any + // other model's lifecycle. + m.healthMu.Lock() + defer m.healthMu.Unlock() + xlog.Debug("Model already loaded in memory", "model", s) // Skip the gRPC health check if the model was recently verified. @@ -572,7 +620,7 @@ func (ml *ModelLoader) checkIsLoaded(s string) *Model { return m } - client := m.GRPC(false, ml.wd) + client := m.GRPC(false, wd) xlog.Debug("Checking model availability", "model", s) cTimeout, cancel := context.WithTimeout(context.Background(), 2*time.Minute) @@ -589,7 +637,7 @@ func (ml *ModelLoader) checkIsLoaded(s string) *Model { // Timeouts may mean the node is busy, so keep the model cached. if isConnectionError(err) { xlog.Warn("Remote model unreachable (connection error), removing from cache", "model", s, "error", err) - if delErr := ml.deleteProcess(s, false); delErr != nil { + if delErr := ml.deleteProcess(cTimeout, s, false); delErr != nil { xlog.Error("error cleaning up remote model", "error", delErr, "model", s) } return nil @@ -600,7 +648,7 @@ func (ml *ModelLoader) checkIsLoaded(s string) *Model { if !process.IsAlive() { xlog.Debug("GRPC Process is not responding", "model", s) // stop and delete the process, this forces to re-load the model and re-create again the service - err := ml.deleteProcess(s, false) + err := ml.deleteProcess(cTimeout, s, true) if err != nil { xlog.Error("error stopping process", "error", err, "process", s) } diff --git a/pkg/model/loader_test.go b/pkg/model/loader_test.go index 1c8892bc3592..1a882943127f 100644 --- a/pkg/model/loader_test.go +++ b/pkg/model/loader_test.go @@ -1,6 +1,7 @@ package model_test import ( + "context" "errors" "os" "path/filepath" @@ -8,12 +9,53 @@ import ( "sync/atomic" "time" + grpcPkg "github.com/mudler/LocalAI/pkg/grpc" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/system" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" ) +// lifecycleBackend embeds the full backend interface so these tests only need +// to override the lifecycle methods they exercise. A nil embedded backend is +// safe because no inference method is called. +type lifecycleBackend struct { + grpcPkg.Backend + busy atomic.Bool + freeOnce sync.Once + freeStarted chan struct{} + freeRelease chan struct{} +} + +type failingRemoteUnloader struct { + err error +} + +func (u failingRemoteUnloader) UnloadRemoteModel(string) error { + return u.err +} + +func newLifecycleBackend() *lifecycleBackend { + return &lifecycleBackend{ + freeStarted: make(chan struct{}), + freeRelease: make(chan struct{}), + } +} + +func (b *lifecycleBackend) IsBusy() bool { + return b.busy.Load() +} + +func (b *lifecycleBackend) Free(ctx context.Context) error { + b.freeOnce.Do(func() { close(b.freeStarted) }) + select { + case <-b.freeRelease: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + var _ = Describe("ModelLoader", func() { var ( modelLoader *model.ModelLoader @@ -161,6 +203,156 @@ var _ = Describe("ModelLoader", func() { Expect(err).To(BeNil()) Expect(modelLoader.CheckIsLoaded("foo")).To(BeNil()) }) + + It("evicts the local remote-model entry when remote unload fails", func() { + remote := model.NewModel("remote", "worker.example:50051", nil) + remote.MarkHealthy() + _, err := modelLoader.LoadModel("remote", "remote", func(_, _, _ string) (*model.Model, error) { + return remote, nil + }) + Expect(err).NotTo(HaveOccurred()) + + unloadErr := errors.New("worker unreachable") + modelLoader.SetRemoteUnloader(failingRemoteUnloader{err: unloadErr}) + Expect(modelLoader.ShutdownModel("remote")).To(MatchError(unloadErr)) + Expect(modelLoader.ListLoadedModels()).To(BeEmpty()) + + replacement := model.NewModel("remote", "replacement.example:50051", nil) + replacement.MarkHealthy() + var reloads atomic.Int32 + loaded, err := modelLoader.LoadModel("remote", "remote", func(_, _, _ string) (*model.Model, error) { + reloads.Add(1) + return replacement, nil + }) + Expect(err).NotTo(HaveOccurred()) + Expect(loaded).To(BeIdenticalTo(replacement)) + Expect(reloads.Load()).To(Equal(int32(1))) + }) + }) + + Context("Shutdown lifecycle conformance", func() { + loadBackend := func(id string, backend grpcPkg.Backend) { + _, err := modelLoader.LoadModel(id, id, func(_, _, _ string) (*model.Model, error) { + return model.NewModelWithClient(id, id, backend), nil + }) + Expect(err).NotTo(HaveOccurred()) + } + + It("force shutdown bypasses a permanently busy backend and leaves the loader available", func() { + stuck := newLifecycleBackend() + stuck.busy.Store(true) + loadBackend("stuck", stuck) + + shutdownDone := make(chan error, 1) + go func() { shutdownDone <- modelLoader.ShutdownModelForce("stuck") }() + + var shutdownErr error + Eventually(shutdownDone, "500ms").Should(Receive(&shutdownErr)) + Expect(shutdownErr).NotTo(HaveOccurred()) + Consistently(stuck.freeStarted, "50ms").ShouldNot(Receive(), "force shutdown must skip Free on a stuck backend") + + other := newLifecycleBackend() + loadBackend("unrelated", other) + Expect(modelLoader.ListLoadedModels()).To(ConsistOf(HaveField("ID", "unrelated"))) + }) + + It("keeps unrelated loads available while graceful Free is blocked", func() { + blocked := newLifecycleBackend() + loadBackend("blocked", blocked) + + shutdownDone := make(chan error, 1) + go func() { shutdownDone <- modelLoader.ShutdownModel("blocked") }() + Eventually(blocked.freeStarted, "500ms").Should(BeClosed()) + + otherDone := make(chan error, 1) + go func() { + _, err := modelLoader.LoadModel("unrelated", "unrelated", func(_, _, _ string) (*model.Model, error) { + return model.NewModelWithClient("unrelated", "unrelated", newLifecycleBackend()), nil + }) + otherDone <- err + }() + + var otherErr error + Eventually(otherDone, "500ms").Should(Receive(&otherErr)) + Expect(otherErr).NotTo(HaveOccurred()) + Consistently(shutdownDone, "50ms").ShouldNot(Receive(), "shutdown must still be waiting for Free") + close(blocked.freeRelease) + + var shutdownErr error + Eventually(shutdownDone, "500ms").Should(Receive(&shutdownErr)) + Expect(shutdownErr).NotTo(HaveOccurred()) + }) + + It("serializes a same-model reload behind shutdown", func() { + blocked := newLifecycleBackend() + loadBackend("replace-me", blocked) + + shutdownDone := make(chan error, 1) + go func() { shutdownDone <- modelLoader.ShutdownModel("replace-me") }() + Eventually(blocked.freeStarted, "500ms").Should(BeClosed()) + + reloadStarted := make(chan struct{}) + reloadDone := make(chan error, 1) + go func() { + _, err := modelLoader.LoadModel("replace-me", "replace-me", func(_, _, _ string) (*model.Model, error) { + close(reloadStarted) + return model.NewModelWithClient("replace-me", "replace-me", newLifecycleBackend()), nil + }) + reloadDone <- err + }() + + Consistently(reloadStarted, "50ms").ShouldNot(BeClosed()) + close(blocked.freeRelease) + + Eventually(shutdownDone, "500ms").Should(Receive(Succeed())) + Eventually(reloadStarted, "500ms").Should(BeClosed()) + Eventually(reloadDone, "500ms").Should(Receive(Succeed())) + }) + + It("bounds a graceful wait for a permanently busy backend without blocking other models", func() { + stuck := newLifecycleBackend() + stuck.busy.Store(true) + loadBackend("stuck", stuck) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + shutdownDone := make(chan error, 1) + go func() { shutdownDone <- modelLoader.ShutdownModelContext(ctx, "stuck", false) }() + + other := newLifecycleBackend() + loadBackend("unrelated", other) + + var shutdownErr error + Eventually(shutdownDone, "500ms").Should(Receive(&shutdownErr)) + Expect(errors.Is(shutdownErr, model.ErrModelBusy)).To(BeTrue()) + Expect(modelLoader.ListLoadedModels()).To(ConsistOf( + HaveField("ID", "stuck"), + HaveField("ID", "unrelated"), + )) + }) + + It("honors cancellation while waiting for an in-progress load of the same model", func() { + loadStarted := make(chan struct{}) + loadRelease := make(chan struct{}) + loadDone := make(chan error, 1) + go func() { + _, err := modelLoader.LoadModel("loading", "loading", func(_, _, _ string) (*model.Model, error) { + close(loadStarted) + <-loadRelease + return model.NewModelWithClient("loading", "loading", newLifecycleBackend()), nil + }) + loadDone <- err + }() + Eventually(loadStarted, "500ms").Should(BeClosed()) + + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + err := modelLoader.ShutdownModelContext(ctx, "loading", false) + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue()) + + close(loadRelease) + Eventually(loadDone, "500ms").Should(Receive(Succeed())) + }) }) Context("Concurrent Loading", func() { diff --git a/pkg/model/model.go b/pkg/model/model.go index 4a0e6af2aa76..3af79abca4b4 100644 --- a/pkg/model/model.go +++ b/pkg/model/model.go @@ -19,6 +19,7 @@ type Model struct { client grpc.Backend process *process.Process lastHealthCheck time.Time + healthMu sync.Mutex sync.Mutex } diff --git a/pkg/model/operation_locks.go b/pkg/model/operation_locks.go new file mode 100644 index 000000000000..ce64312ebba4 --- /dev/null +++ b/pkg/model/operation_locks.go @@ -0,0 +1,71 @@ +package model + +import ( + "context" + "math" + "sync" + + "golang.org/x/sync/semaphore" +) + +// modelOperationLocks serializes destructive lifecycle operations per model +// without coupling unrelated models. Entries are reference-counted so model IDs +// observed from requests do not accumulate forever. +type modelOperationLocks struct { + mu sync.Mutex + entries map[string]*modelOperationLock +} + +type modelOperationLock struct { + sem *semaphore.Weighted + refs int +} + +func newModelOperationLocks() *modelOperationLocks { + return &modelOperationLocks{entries: make(map[string]*modelOperationLock)} +} + +func (l *modelOperationLocks) acquire(modelID string, exclusive bool) func() { + release, err := l.acquireContext(context.Background(), modelID, exclusive) + if err != nil { + panic(err) + } + return release +} + +func (l *modelOperationLocks) acquireContext(ctx context.Context, modelID string, exclusive bool) (func(), error) { + if ctx == nil { + ctx = context.Background() + } + l.mu.Lock() + entry := l.entries[modelID] + if entry == nil { + entry = &modelOperationLock{sem: semaphore.NewWeighted(math.MaxInt64)} + l.entries[modelID] = entry + } + entry.refs++ + l.mu.Unlock() + + weight := int64(1) + if exclusive { + weight = math.MaxInt64 + } + if err := entry.sem.Acquire(ctx, weight); err != nil { + l.releaseReference(modelID, entry) + return nil, err + } + + return func() { + entry.sem.Release(weight) + l.releaseReference(modelID, entry) + }, nil +} + +func (l *modelOperationLocks) releaseReference(modelID string, entry *modelOperationLock) { + l.mu.Lock() + entry.refs-- + if entry.refs == 0 { + delete(l.entries, modelID) + } + l.mu.Unlock() +} diff --git a/pkg/model/process.go b/pkg/model/process.go index 200f8ae30db3..33eb6377e0ca 100644 --- a/pkg/model/process.go +++ b/pkg/model/process.go @@ -21,39 +21,55 @@ var forceBackendShutdown bool = os.Getenv("LOCALAI_FORCE_BACKEND_SHUTDOWN") == " var ( modelNotFoundErr = errors.New("model not found") + // ErrModelBusy indicates that a graceful shutdown context ended while + // requests were still in flight. + ErrModelBusy = errors.New("model is still busy") +) + +const ( + gracefulShutdownTimeout = 30 * time.Second + forcedShutdownTimeout = 30 * time.Second + backendFreeTimeout = 5 * time.Second + busyPollInterval = 100 * time.Millisecond ) // deleteProcess stops and removes a backend. The force flag trades a graceful // shutdown for a prompt one and is meant for the watchdog's busy-killer: a -// backend that has been busy past the watchdog timeout is, by definition, -// stuck in an in-flight gRPC call. Waiting for that call to finish before -// stopping the process (the graceful path) would block forever — and since -// deleteProcess runs under ml.mu, it would stall every other model load, -// including the shared opus backend load at the start of every realtime -// (WebRTC) session, hanging new connections at "Connected, waiting for -// session...". The force path stops the process first, which drops the -// in-flight call's gRPC connection and unblocks it, then cleans up. -func (ml *ModelLoader) deleteProcess(s string, force bool) error { - model, ok := ml.store.Get(s) +// backend that has been busy past the watchdog timeout may be stuck in an +// in-flight gRPC call. The graceful path waits only until ctx expires; the +// force path skips that wait and Free(), stops the process, then cleans up. +// Callers serialize this operation per model, never with the global loader +// mutex, so a faulty backend cannot stall unrelated model lifecycle work. +func (ml *ModelLoader) deleteProcess(ctx context.Context, s string, force bool) error { + if ctx == nil { + ctx = context.Background() + } + + // Snapshot mutable loader configuration while holding ml.mu, then perform + // every wait, callback, RPC, and process operation without the global lock. + // Same-model ordering is provided by modelOperationLocks at the public + // lifecycle boundary. + ml.mu.Lock() + store := ml.store + wd := ml.wd + hooks := append([]ModelUnloadHook(nil), ml.onUnloadHooks...) + remoteUnloader := ml.remoteUnloader + ml.mu.Unlock() + + model, ok := store.Get(s) if !ok { xlog.Debug("Model not found", "model", s) return modelNotFoundErr } if !force { - retries := 1 - for model.GRPC(false, ml.wd).IsBusy() { + client := model.GRPC(false, wd) + for client.IsBusy() { xlog.Debug("Model busy. Waiting.", "model", s) - dur := time.Duration(retries*2) * time.Second - if dur > retryTimeout { - dur = retryTimeout - } - time.Sleep(dur) - retries++ - - if retries > 10 && forceBackendShutdown { - xlog.Warn("Model is still busy after retries. Forcing shutdown.", "model", s, "retries", retries) - break + select { + case <-ctx.Done(): + return fmt.Errorf("%w: %s: %w", ErrModelBusy, s, ctx.Err()) + case <-time.After(busyPollInterval): } } } @@ -61,7 +77,7 @@ func (ml *ModelLoader) deleteProcess(s string, force bool) error { xlog.Debug("Deleting process", "model", s, "force", force) // Run unload hooks (e.g. close MCP sessions) - for _, hook := range ml.onUnloadHooks { + for _, hook := range hooks { hook(s) } @@ -76,7 +92,10 @@ func (ml *ModelLoader) deleteProcess(s string, force bool) error { // don't surface it as an error. if !force { xlog.Debug("Calling Free() to release GPU resources", "model", s) - if err := model.GRPC(false, ml.wd).Free(context.Background()); err != nil { + freeCtx, cancel := context.WithTimeout(ctx, backendFreeTimeout) + err := model.GRPC(false, wd).Free(freeCtx) + cancel() + if err != nil { if grpcerrors.IsUnimplemented(err) { xlog.Debug("Backend does not implement Free(); GPU release handled on process stop", "model", s) } else { @@ -92,16 +111,26 @@ func (ml *ModelLoader) deleteProcess(s string, force bool) error { // No local process — this is a remote/external backend. // In distributed mode, delegate to the remote unloader to tell // the backend node to free the model (GPU resources, etc.). - if ml.remoteUnloader != nil { + var unloadErr error + if remoteUnloader != nil { xlog.Debug("Delegating model unload to remote unloader", "model", s) - if err := ml.remoteUnloader.UnloadRemoteModel(s); err != nil { - xlog.Warn("Remote model unload failed", "model", s, "error", err) + if contextUnloader, ok := remoteUnloader.(RemoteModelContextUnloader); ok { + unloadErr = contextUnloader.UnloadRemoteModelContext(ctx, s, force) + } else { + unloadErr = remoteUnloader.UnloadRemoteModel(s) + } + if unloadErr != nil { + xlog.Warn("Remote model unload failed", "model", s, "error", unloadErr) } } else { xlog.Debug("No local process and no remote unloader", "model", s) } - ml.store.Delete(s) - return nil + // The store is only the frontend's local representative of a remote + // model. Never retain it after an unload attempt: on failure it may point + // at a known-unreachable worker, while the distributed registry remains + // the source of truth for anything that is still running remotely. + store.Delete(s) + return unloadErr } // Mark the stop as intentional so the exit-watcher logs it as an @@ -110,29 +139,34 @@ func (ml *ModelLoader) deleteProcess(s string, force bool) error { err := process.Stop() if err != nil { xlog.Error("(deleteProcess) error while deleting process", "error", err, "model", s) + if !process.IsAlive() { + // A concurrently crashed/already-reaped process can no longer own + // resources even if Stop could not read or signal its PID. + store.Delete(s) + return nil + } + return err } - if err == nil { - ml.store.Delete(s) - } - - return err + store.Delete(s) + return nil } func (ml *ModelLoader) StopGRPC(filter GRPCProcessFilter) error { var err error = nil ml.mu.Lock() - defer ml.mu.Unlock() + store := ml.store + ml.mu.Unlock() // Collect matching keys first — can't mutate store during Range var toDelete []string - ml.store.Range(func(k string, m *Model) bool { + store.Range(func(k string, m *Model) bool { if filter(k, m.Process()) { toDelete = append(toDelete, k) } return true }) for _, k := range toDelete { - e := ml.deleteProcess(k, false) + e := ml.ShutdownModel(k) err = errors.Join(err, e) } return err @@ -144,8 +178,9 @@ func (ml *ModelLoader) StopAllGRPC() error { func (ml *ModelLoader) GetGRPCPID(id string) (int, error) { ml.mu.Lock() - defer ml.mu.Unlock() - p, exists := ml.store.Get(id) + store := ml.store + ml.mu.Unlock() + p, exists := store.Get(id) if !exists { return -1, fmt.Errorf("no grpc backend found for %s", id) } diff --git a/pkg/model/watchdog.go b/pkg/model/watchdog.go index a961333477d8..efe400e41f6f 100644 --- a/pkg/model/watchdog.go +++ b/pkg/model/watchdog.go @@ -25,9 +25,14 @@ import ( // and the GRPC client talks to it via a channel to send status updates type WatchDog struct { sync.Mutex - busyTime map[string]time.Time - idleTime map[string]time.Time - lastUsed map[string]time.Time // LRU tracking: when each model was last used + busyTime map[string]time.Time + inFlight map[string]int + requestStarts map[string]map[uint64]time.Time + legacyRequests map[string][]uint64 + nextRequestID uint64 + idleTime map[string]time.Time + lastUsed map[string]time.Time // LRU tracking: when each model was last used + timeout, idletimeout time.Duration addressMap map[string]*process.Process addressModelMap map[string]string @@ -64,8 +69,8 @@ type ProcessManager interface { ShutdownModel(modelName string) error // ShutdownModelForce stops the backend without waiting for an in-flight // gRPC call to finish. Used when the watchdog evicts a backend that is - // stuck busy: the graceful ShutdownModel would block on that stuck call - // (while holding the loader's mutex), stalling every other model load. + // stuck busy: the graceful ShutdownModel intentionally waits for active + // calls until its bounded deadline. ShutdownModelForce(modelName string) error } @@ -89,6 +94,9 @@ func NewWatchDog(opts ...WatchDogOption) *WatchDog { idletimeout: o.idleTimeout, pm: o.processManager, busyTime: make(map[string]time.Time), + inFlight: make(map[string]int), + requestStarts: make(map[string]map[uint64]time.Time), + legacyRequests: make(map[string][]uint64), idleTime: make(map[string]time.Time), lastUsed: make(map[string]time.Time), addressMap: make(map[string]*process.Process), @@ -241,22 +249,96 @@ func (wd *WatchDog) Add(address string, p *process.Process) { wd.addressMap[address] = p } +// TrackRequest records one request and returns an idempotent completion +// callback. The per-request identity lets the busy watchdog follow the oldest +// request that is still active instead of treating uninterrupted parallel +// traffic as one indefinitely old request. +func (wd *WatchDog) TrackRequest(address string) func() { + wd.Lock() + requestID := wd.beginRequestLocked(address, false) + wd.Unlock() + + var once sync.Once + return func() { + once.Do(func() { + wd.Lock() + defer wd.Unlock() + wd.finishRequestLocked(address, requestID) + }) + } +} + +// Mark and UnMark retain the legacy public API for non-gRPC callers. Callers +// that can keep the completion callback should use TrackRequest so overlapping +// requests can complete in any order without losing their individual ages. func (wd *WatchDog) Mark(address string) { wd.Lock() defer wd.Unlock() - now := time.Now() - wd.busyTime[address] = now - wd.lastUsed[address] = now // Update LRU tracking - delete(wd.idleTime, address) + wd.beginRequestLocked(address, true) } -func (wd *WatchDog) UnMark(ModelAddress string) { +func (wd *WatchDog) UnMark(address string) { wd.Lock() defer wd.Unlock() + requests := wd.legacyRequests[address] + if len(requests) == 0 { + return + } + requestID := requests[len(requests)-1] + requests = requests[:len(requests)-1] + if len(requests) == 0 { + delete(wd.legacyRequests, address) + } else { + wd.legacyRequests[address] = requests + } + wd.finishRequestLocked(address, requestID) +} + +func (wd *WatchDog) beginRequestLocked(address string, legacy bool) uint64 { now := time.Now() - delete(wd.busyTime, ModelAddress) - wd.idleTime[ModelAddress] = now - wd.lastUsed[ModelAddress] = now // Update LRU tracking + wd.nextRequestID++ + requestID := wd.nextRequestID + starts := wd.requestStarts[address] + if starts == nil { + starts = make(map[uint64]time.Time) + wd.requestStarts[address] = starts + } + starts[requestID] = now + if legacy { + wd.legacyRequests[address] = append(wd.legacyRequests[address], requestID) + } + wd.inFlight[address] = len(starts) + if len(starts) == 1 { + wd.busyTime[address] = now + } + wd.lastUsed[address] = now + delete(wd.idleTime, address) + return requestID +} + +func (wd *WatchDog) finishRequestLocked(address string, requestID uint64) { + starts := wd.requestStarts[address] + if _, exists := starts[requestID]; !exists { + return + } + delete(starts, requestID) + now := time.Now() + wd.lastUsed[address] = now // Update LRU tracking + if len(starts) > 0 { + wd.inFlight[address] = len(starts) + var oldest time.Time + for _, started := range starts { + if oldest.IsZero() || started.Before(oldest) { + oldest = started + } + } + wd.busyTime[address] = oldest + return + } + delete(wd.requestStarts, address) + delete(wd.inFlight, address) + delete(wd.busyTime, address) + wd.idleTime[address] = now } // UpdateLastUsed updates the last used time for a model address (for LRU tracking) @@ -278,6 +360,7 @@ func (wd *WatchDog) GetLoadedModelCount() int { type WatchDogState struct { AddressModelMap map[string]string BusyTime map[string]time.Time + InFlight map[string]int IdleTime map[string]time.Time LastUsed map[string]time.Time AddressMap map[string]*process.Process @@ -300,6 +383,11 @@ func (wd *WatchDog) GetState() WatchDogState { busyTime[k] = v } + inFlight := make(map[string]int, len(wd.inFlight)) + for k, v := range wd.inFlight { + inFlight[k] = v + } + idleTime := make(map[string]time.Time, len(wd.idleTime)) for k, v := range wd.idleTime { idleTime[k] = v @@ -318,6 +406,7 @@ func (wd *WatchDog) GetState() WatchDogState { return WatchDogState{ AddressModelMap: addressModelMap, BusyTime: busyTime, + InFlight: inFlight, IdleTime: idleTime, LastUsed: lastUsed, AddressMap: addressMap, @@ -332,6 +421,30 @@ func (wd *WatchDog) RestoreState(state WatchDogState) { wd.addressModelMap = state.AddressModelMap wd.busyTime = state.BusyTime + wd.inFlight = state.InFlight + if wd.inFlight == nil { + wd.inFlight = make(map[string]int, len(wd.busyTime)) + for address := range wd.busyTime { + wd.inFlight[address] = 1 + } + } + wd.requestStarts = make(map[string]map[uint64]time.Time, len(wd.inFlight)) + wd.legacyRequests = make(map[string][]uint64, len(wd.inFlight)) + for address, count := range wd.inFlight { + started := wd.busyTime[address] + if started.IsZero() { + started = time.Now() + wd.busyTime[address] = started + } + requests := make(map[uint64]time.Time, count) + for range count { + wd.nextRequestID++ + requestID := wd.nextRequestID + requests[requestID] = started + wd.legacyRequests[address] = append(wd.legacyRequests[address], requestID) + } + wd.requestStarts[address] = requests + } wd.idleTime = state.IdleTime wd.lastUsed = state.LastUsed wd.addressMap = state.AddressMap @@ -351,8 +464,8 @@ type modelUsageInfo struct { // at selection time. A busy target must be shut down via ShutdownModelForce: // the graceful path waits for its in-flight gRPC call to finish, but a busy // eviction only happens when that call is stuck (busy-killer) or the operator -// opted into forceEvictionWhenBusy — either way, waiting would deadlock while -// holding the loader mutex. +// opted into forceEvictionWhenBusy — either way, waiting for the graceful +// deadline would defeat prompt eviction. type evictionTarget struct { model string wasBusy bool @@ -468,9 +581,8 @@ func (wd *WatchDog) collectEvictionsLocked(candidates []modelUsageInfo, maxToEvi } // shutdownEvicted shuts down each evicted model, using the force path for any -// that were busy at selection time so a stuck in-flight call doesn't deadlock -// the graceful ShutdownModel (which waits for that call while holding the -// loader's mutex). +// 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) { for _, t := range targets { var err error @@ -654,11 +766,7 @@ func (wd *WatchDog) checkBusy() { // The busy-killer targets backends whose in-flight gRPC call has been // stuck past the busy timeout. Use the force path so the loader stops // the process FIRST (dropping the stuck call's gRPC connection) instead - // of waiting for that call to finish — the graceful path would block on - // that stuck call while holding the loader mutex and stall every other - // model load (notably the opus backend load at the start of every - // realtime WebRTC session, which hangs new sessions at "Connected, - // waiting for session..."). + // of waiting for the graceful shutdown deadline. for _, model := range modelsToShutdown { if err := wd.pm.ShutdownModelForce(model); err != nil { xlog.Error("[watchdog] error shutting down model", "error", err, "model", model) @@ -807,6 +915,9 @@ func (wd *WatchDog) untrack(address string) { delete(wd.modelSizes, modelID) } 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) diff --git a/pkg/model/watchdog_test.go b/pkg/model/watchdog_test.go index e158b84c005e..36d9ceed2709 100644 --- a/pkg/model/watchdog_test.go +++ b/pkg/model/watchdog_test.go @@ -63,6 +63,14 @@ func (m *mockProcessManager) getForceShutdownCalls() []string { return result } +func (m *mockProcessManager) getGracefulShutdownCalls() []string { + m.mu.Lock() + defer m.mu.Unlock() + result := make([]string, len(m.gracefulCalls)) + copy(result, m.gracefulCalls) + return result +} + var _ = Describe("WatchDog", func() { var ( wd *model.WatchDog @@ -183,6 +191,44 @@ var _ = Describe("WatchDog", func() { wd.UpdateLastUsed("addr1") // Verify the time was updated }) + + It("remains busy until every parallel request completes", func() { + wd.AddAddressModelMap("addr1", "model1") + wd.Mark("addr1") + busySince := wd.GetState().BusyTime["addr1"] + wd.Mark("addr1") + wd.UnMark("addr1") + + state := wd.GetState() + Expect(state.InFlight).To(HaveKeyWithValue("addr1", 1)) + Expect(state.BusyTime).To(HaveKeyWithValue("addr1", busySince)) + Expect(state.IdleTime).NotTo(HaveKey("addr1")) + + wd.UnMark("addr1") + state = wd.GetState() + Expect(state.InFlight).NotTo(HaveKey("addr1")) + Expect(state.BusyTime).NotTo(HaveKey("addr1")) + Expect(state.IdleTime).To(HaveKey("addr1")) + }) + + It("tracks the oldest active request instead of continuous parallel traffic", func() { + wd.AddAddressModelMap("addr1", "model1") + finishFirst := wd.TrackRequest("addr1") + firstStarted := wd.GetState().BusyTime["addr1"] + time.Sleep(10 * time.Millisecond) + finishSecond := wd.TrackRequest("addr1") + + finishFirst() + state := wd.GetState() + Expect(state.InFlight).To(HaveKeyWithValue("addr1", 1)) + Expect(state.BusyTime["addr1"]).To(BeTemporally(">", firstStarted)) + + finishSecond() + state = wd.GetState() + Expect(state.InFlight).NotTo(HaveKey("addr1")) + Expect(state.BusyTime).NotTo(HaveKey("addr1")) + Expect(state.IdleTime).To(HaveKey("addr1")) + }) }) Context("EnforceLRULimit", func() { @@ -692,11 +738,8 @@ var _ = Describe("WatchDog", func() { It("force-shuts down a model that is stuck busy past the busy timeout", func() { // Regression: a backend stuck on an in-flight gRPC call must be // killed via the force path (stop the process first), not the - // graceful one (wait for the stuck call to finish, which would - // deadlock while holding the loader mutex and stall every other - // model load — e.g. the opus backend load at the start of every - // realtime WebRTC session, hanging new "Connected, waiting for - // session..." connections). + // graceful one, which intentionally waits for active calls until + // its bounded deadline. wd = model.NewWatchDog( model.WithProcessManager(pm), model.WithBusyTimeout(10*time.Millisecond), @@ -714,6 +757,7 @@ var _ = Describe("WatchDog", func() { return pm.getForceShutdownCalls() }, "300ms", "10ms").Should(ContainElement("stuckModel")) Expect(pm.getShutdownCalls()).To(ContainElement("stuckModel")) + Expect(pm.getGracefulShutdownCalls()).To(BeEmpty()) }) }) diff --git a/scripts/model-lifecycle-conformance.sh b/scripts/model-lifecycle-conformance.sh new file mode 100755 index 000000000000..bec2fa552fa1 --- /dev/null +++ b/scripts/model-lifecycle-conformance.sh @@ -0,0 +1,55 @@ +#!/bin/sh +# model-lifecycle-conformance: check the model-loader shutdown behavior +# specified by formal-verification/model_loader_shutdown.fizz against the +# focused loader, gRPC client, distributed unloader, and worker tests. +# +# Both layers are required. The Go tests pin the concrete local/distributed +# paths; FizzBee exhaustively checks their progress and safety properties. +set -eu + +ROOT=$(CDPATH= cd "$(dirname "$0")/.." && pwd) +GOCMD=${GOCMD:-go} +SPEC="$ROOT/formal-verification/model_loader_shutdown.fizz" + +cd "$ROOT" + +echo "==> [1/3] Go lifecycle and busy-accounting conformance with -race" +"$GOCMD" test -race -count=1 ./pkg/model ./pkg/grpc ./core/services/worker + +echo "==> [2/3] Go distributed force-propagation conformance with -race" +"$GOCMD" test -race -count=1 ./core/services/nodes -ginkgo.focus=RemoteUnloaderAdapter + +# install-fizzbee.sh publishes a stable wrapper at .tools/fizzbee/fizz. The +# wrapper, rather than the raw fizzbee binary, runs both parser and checker. +FIZZBEE_BIN=${FIZZBEE_BIN:-} +if [ -z "$FIZZBEE_BIN" ]; then + if [ -x "$ROOT/.tools/fizzbee/fizz" ]; then + FIZZBEE_BIN="$ROOT/.tools/fizzbee/fizz" + elif command -v fizz >/dev/null 2>&1; then + FIZZBEE_BIN=fizz + fi +fi + +echo "==> [3/3] FizzBee model-loader lifecycle check" +if [ -z "$FIZZBEE_BIN" ] || { [ ! -x "$FIZZBEE_BIN" ] && ! command -v "$FIZZBEE_BIN" >/dev/null 2>&1; }; then + echo "ERROR: FizzBee not found; lifecycle verification is fail-closed." >&2 + echo " Install the pinned binary with: make install-fizzbee" >&2 + exit 1 +fi + +set +e +out=$("$FIZZBEE_BIN" "$SPEC" 2>&1) +rc=$? +set -e +printf '%s\n' "$out" + +if [ "$rc" -ne 0 ]; then + echo "ERROR: FizzBee exited $rc on $SPEC" >&2 + exit 1 +fi +if printf '%s\n' "$out" | grep -qE '^(FAILED|DEADLOCK)'; then + echo "ERROR: FizzBee reported a lifecycle property violation in $SPEC" >&2 + exit 1 +fi + +echo "==> model-lifecycle-conformance OK (Go + FizzBee)"