diff --git a/core/http/endpoints/localai/backend_monitor.go b/core/http/endpoints/localai/backend_monitor.go index d138c5bee3e0..c60e92a92626 100644 --- a/core/http/endpoints/localai/backend_monitor.go +++ b/core/http/endpoints/localai/backend_monitor.go @@ -1,9 +1,15 @@ package localai import ( + "errors" + "fmt" + "net/http" + "github.com/labstack/echo/v4" "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/core/services/monitoring" + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/xlog" ) // BackendMonitorEndpoint returns the status of the specified backend @@ -48,7 +54,25 @@ func BackendShutdownEndpoint(bm *monitoring.BackendMonitorService) echo.HandlerF if err := c.Bind(input); err != nil { return err } + if input.Model == "" { + return echo.NewHTTPError(http.StatusBadRequest, "model is required") + } - return bm.ShutdownModel(input.Model) + if err := bm.ShutdownModel(input.Model); err != nil { + // "Not loaded" is a client-side condition, not a server fault, so + // it should not surface as a 500. In distributed mode this branch + // used to be reached for models that were running on a worker — + // only this replica's local store was empty. The loader now + // consults the node registry first, so reaching it means the + // model is loaded neither here nor anywhere in the cluster. + if errors.Is(err, model.ErrModelNotFound) { + xlog.Info("Shutdown requested for a model that is not loaded", "model", input.Model) + return echo.NewHTTPError(http.StatusNotFound, + fmt.Sprintf("model %q is not loaded on this instance or on any worker node", input.Model)) + } + xlog.Error("Failed to shut down model", "model", input.Model, "error", err) + return err + } + return c.JSON(http.StatusOK, map[string]string{"message": "model stopped"}) } } diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index bbdd9f0f23bf..dd05fe2ede64 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -226,6 +226,14 @@ type BackendUpgradeRequest struct { type BackendUpgradeReply struct { Success bool `json:"success"` Error string `json:"error,omitempty"` + + // StoppedProcessKeys / ReportsStoppedProcesses carry the same + // stale-row-invalidation contract as on BackendDeleteReply; an upgrade + // force-stops every process using the binary and starts none back up, so it + // recycles ports exactly the way a delete does. See that type for why the + // boolean is not redundant with an empty list. + StoppedProcessKeys []string `json:"stopped_process_keys,omitempty"` + ReportsStoppedProcesses bool `json:"reports_stopped_processes,omitempty"` } // SubjectNodeBackendList queries a worker node for its installed backends. @@ -290,6 +298,23 @@ type BackendDeleteRequest struct { type BackendDeleteReply struct { Success bool `json:"success"` Error string `json:"error,omitempty"` + + // StoppedProcessKeys names every `modelID#replica` process the worker + // terminated while serving this delete. Stopping a process returns its gRPC + // port to the worker's allocator, so any NodeModel row still pointing at + // that address becomes a live misroute the moment an unrelated backend + // binds the recycled port: probeHealth verifies liveness, not identity, so + // the request is served by the wrong backend rather than failing. The + // controller uses these keys to drop the rows eagerly. + StoppedProcessKeys []string `json:"stopped_process_keys,omitempty"` + + // ReportsStoppedProcesses distinguishes "this worker enumerates what it + // stopped and stopped nothing" from "this worker predates the field". Both + // send an empty list, and only the first is authoritative. Without this + // flag a controller cannot tell them apart and would eventually be tempted + // to read silence as a completed cleanup, which is precisely the wrong + // conclusion against an older worker. + ReportsStoppedProcesses bool `json:"reports_stopped_processes,omitempty"` } // SubjectNodeModelUnload tells a worker node to unload a model (gRPC Free) without killing the backend. diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index 8699968178c7..24a19e866d27 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -12,6 +12,7 @@ import ( "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/xlog" ) @@ -71,6 +72,17 @@ func (a *RemoteUnloaderAdapter) InstallTimeout() time.Duration { return a.installTimeout } +// Compile-time proof that the adapter still satisfies the loader's optional +// extensions. Both are consumed via runtime type assertion in deleteProcess, so +// a signature drift here would silently downgrade behavior — losing force +// propagation, or making ShutdownModel unable to tell a cluster-wide miss from +// a completed unload — rather than failing the build. +var ( + _ model.RemoteModelUnloader = (*RemoteUnloaderAdapter)(nil) + _ model.RemoteModelContextUnloader = (*RemoteUnloaderAdapter)(nil) + _ model.RemoteModelPresenceChecker = (*RemoteUnloaderAdapter)(nil) +) + // 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 a bounded Free() followed by process termination; @@ -80,6 +92,22 @@ func (a *RemoteUnloaderAdapter) UnloadRemoteModel(modelName string) error { return a.UnloadRemoteModelContext(context.Background(), modelName, false) } +// HasRemoteModel reports whether any node currently holds the model. It exists +// because UnloadRemoteModel is idempotent and so cannot signal "there was +// nothing to stop"; ShutdownModel consults this first so it can answer 404 for +// a model loaded neither locally nor anywhere in the cluster, instead of the +// misleading 500 "model not found" that a local-store miss used to produce. +func (a *RemoteUnloaderAdapter) HasRemoteModel(ctx context.Context, modelName string) (bool, error) { + if ctx == nil { + ctx = context.Background() + } + nodes, err := a.registry.FindNodesWithModel(ctx, modelName) + if err != nil { + return false, fmt.Errorf("finding nodes with model %q: %w", modelName, err) + } + return len(nodes) > 0, nil +} + // 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 { @@ -91,6 +119,10 @@ func (a *RemoteUnloaderAdapter) UnloadRemoteModelContext(ctx context.Context, mo return fmt.Errorf("finding nodes with model %q: %w", modelName, err) } if len(nodes) == 0 { + // Unloading is idempotent by contract: cleanup paths (model deletion, + // config edits, watchdog eviction) legitimately run against an + // already-unloaded model and must not fail. Callers that need to tell + // this case apart use HasRemoteModel before unloading. xlog.Debug("No remote nodes found with model", "model", modelName) return nil } @@ -238,6 +270,9 @@ func (a *RemoteUnloaderAdapter) UpgradeBackend(nodeID, backendType, galleriesJSO return nil, fmt.Errorf("%w (subject=%s nodeID=%s backend=%s): %v", galleryop.ErrWorkerStillInstalling, subject, nodeID, backendType, err) } + if err == nil { + a.dropStoppedReplicaRows(nodeID, "backend.upgrade", backendType, reply.StoppedProcessKeys, reply.ReportsStoppedProcesses) + } return reply, err } @@ -305,7 +340,55 @@ func (a *RemoteUnloaderAdapter) DeleteBackend(nodeID, backendName string) (*mess subject := messaging.SubjectNodeBackendDelete(nodeID) xlog.Info("Sending NATS backend.delete", "nodeID", nodeID, "backend", backendName) - return messaging.RequestJSON[messaging.BackendDeleteRequest, messaging.BackendDeleteReply](a.nats, subject, messaging.BackendDeleteRequest{Backend: backendName}, 2*time.Minute) + reply, err := messaging.RequestJSON[messaging.BackendDeleteRequest, messaging.BackendDeleteReply](a.nats, subject, messaging.BackendDeleteRequest{Backend: backendName}, 2*time.Minute) + if err != nil { + return reply, err + } + a.dropStoppedReplicaRows(nodeID, "backend.delete", backendName, reply.StoppedProcessKeys, reply.ReportsStoppedProcesses) + return reply, nil +} + +// dropStoppedReplicaRows removes the NodeModel rows addressing processes a +// worker just terminated. +// +// Why eagerly, rather than leaving it to the existing health checks: stopping a +// process returns its gRPC port to the worker's allocator, and the next backend +// started there can be handed that same port. Until the row is gone it names a +// live address, so both SmartRouter.probeHealth and the HealthMonitor per-model +// probe — which verify liveness, not identity — pass against whatever now +// occupies the port, and the request is served by the wrong backend rather than +// failing. Nothing else on the delete/upgrade path tells the controller the +// address just became invalid, unlike model.unload which drops its rows itself. +// +// reported=false means the worker predates this reply field. Its empty list is +// then indistinguishable from "stopped nothing", so it must NOT be read as a +// completed cleanup: leave the rows alone and fall back to the probe-based +// staleness recovery that was the only mechanism before this change. +func (a *RemoteUnloaderAdapter) dropStoppedReplicaRows(nodeID, op, backendName string, processKeys []string, reported bool) { + if !reported { + xlog.Debug("Worker did not report stopped processes; relying on probe-based staleness recovery", + "nodeID", nodeID, "op", op, "backend", backendName) + return + } + ctx := context.Background() + for _, key := range processKeys { + modelName, replicaIndex, ok := model.ParseBackendProcessKey(key) + if !ok { + // Acting on a guess could evict the row of a healthy sibling replica. + xlog.Warn("Ignoring unparseable process key reported by worker", + "nodeID", nodeID, "op", op, "backend", backendName, "processKey", key) + continue + } + xlog.Info("Dropping replica row for a process the worker stopped", + "nodeID", nodeID, "op", op, "backend", backendName, "model", modelName, "replica", replicaIndex) + if err := a.registry.RemoveNodeModel(ctx, nodeID, modelName, replicaIndex); err != nil { + // Best-effort: probe-based recovery remains the backstop, and failing + // the operator's delete over a bookkeeping error would be worse than + // the stale row this prevents. + xlog.Warn("Failed to drop replica row for a stopped process", + "nodeID", nodeID, "op", op, "model", modelName, "replica", replicaIndex, "error", err) + } + } } // UnloadModelOnNode sends a model.unload request to a specific node. diff --git a/core/services/nodes/unloader_stale_rows_test.go b/core/services/nodes/unloader_stale_rows_test.go new file mode 100644 index 000000000000..0fff87bafe20 --- /dev/null +++ b/core/services/nodes/unloader_stale_rows_test.go @@ -0,0 +1,186 @@ +package nodes + +import ( + "encoding/json" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/messaging" +) + +// Replies are handed to the adapter as raw JSON rather than as marshalled +// structs on purpose: these specs pin the behavior a controller must exhibit +// when a worker on a *different* build answers it. Marshalling the current +// struct would only ever produce the current field set and could not express +// an old worker's reply at all. +var _ = Describe("RemoteUnloaderAdapter stale replica rows", func() { + var ( + locator *fakeModelLocator + mc *fakeMessagingClient + adapter *RemoteUnloaderAdapter + ) + + BeforeEach(func() { + locator = &fakeModelLocator{} + mc = &fakeMessagingClient{} + adapter = NewRemoteUnloaderAdapter(locator, mc, 3*time.Minute, 15*time.Minute) + }) + + Describe("DeleteBackend", func() { + It("drops the replica rows for every process the worker stopped", func() { + // The worker has just terminated these processes and returned their + // gRPC ports to its allocator. Any row still pointing at those + // addresses will pass probeHealth as soon as an unrelated backend + // binds the recycled port, and the request is then silently served + // by the wrong backend. + mc.requestReply = []byte(`{ + "success": true, + "reports_stopped_processes": true, + "stopped_process_keys": ["qwen3-0.6b#0", "qwen3-0.6b#2"] + }`) + + reply, err := adapter.DeleteBackend("node-1", "llama-cpp") + Expect(err).NotTo(HaveOccurred()) + Expect(reply.Success).To(BeTrue()) + + Expect(locator.removedReplicas).To(ConsistOf( + modelReplicaRef{"node-1", "qwen3-0.6b", 0}, + modelReplicaRef{"node-1", "qwen3-0.6b", 2}, + )) + }) + + It("keeps model IDs that contain the replica separator intact", func() { + // Process keys are `modelID#replica` and model IDs are free to + // contain '#' themselves, so the split must be anchored at the last + // separator. Splitting at the first one addresses a row that does + // not exist and leaves the real stale row in place. + mc.requestReply = []byte(`{ + "success": true, + "reports_stopped_processes": true, + "stopped_process_keys": ["weird#name#3"] + }`) + + _, err := adapter.DeleteBackend("node-1", "llama-cpp") + Expect(err).NotTo(HaveOccurred()) + Expect(locator.removedReplicas).To(ConsistOf(modelReplicaRef{"node-1", "weird#name", 3})) + }) + + It("removes nothing when a worker that reports stopped processes stopped none", func() { + // Deleting a backend that was never loaded is routine. The reply is + // authoritative here, so "no keys" genuinely means "no rows". + mc.requestReply = []byte(`{"success": true, "reports_stopped_processes": true}`) + + _, err := adapter.DeleteBackend("node-1", "llama-cpp") + Expect(err).NotTo(HaveOccurred()) + Expect(locator.removedReplicas).To(BeEmpty()) + }) + + It("does not read an old worker's silence as a completed cleanup", func() { + // A worker predating this change never sets either field. Its empty + // list is indistinguishable from "stopped nothing", so the + // controller must not conclude the node is clean, and equally must + // not guess at rows to delete. It falls back to the probe-based + // self-heal in SmartRouter.probeHealth, which is exactly the + // pre-change behavior. + mc.requestReply = []byte(`{"success": true}`) + + reply, err := adapter.DeleteBackend("node-1", "llama-cpp") + Expect(err).NotTo(HaveOccurred()) + Expect(reply.Success).To(BeTrue()) + Expect(reply.ReportsStoppedProcesses).To(BeFalse()) + Expect(locator.removedReplicas).To(BeEmpty()) + Expect(locator.removedPairs).To(BeEmpty()) + }) + + It("leaves rows alone when the worker could not stop the process", func() { + // The worker aborts the delete without listing the key it failed to + // kill: that process survived, so its address is still correct and + // dropping the row would force a needless reload of a live replica. + mc.requestReply = []byte(`{ + "success": false, + "error": "could not stop running process qwen3-0.6b#0", + "reports_stopped_processes": true + }`) + + reply, err := adapter.DeleteBackend("node-1", "llama-cpp") + Expect(err).NotTo(HaveOccurred()) + Expect(reply.Success).To(BeFalse()) + Expect(locator.removedReplicas).To(BeEmpty()) + }) + + It("drops rows for processes that died before a later step failed", func() { + // The worker reports a key only once that process is confirmed gone, + // so a failure further along (removing files, re-registering) does + // not make the already-recycled ports any less dangerous. Gating + // removal on overall success would strand exactly those rows. + mc.requestReply = []byte(`{ + "success": false, + "error": "failed to delete backend files", + "reports_stopped_processes": true, + "stopped_process_keys": ["qwen3-0.6b#0"] + }`) + + _, err := adapter.DeleteBackend("node-1", "llama-cpp") + Expect(err).NotTo(HaveOccurred()) + Expect(locator.removedReplicas).To(ConsistOf(modelReplicaRef{"node-1", "qwen3-0.6b", 0})) + }) + + It("ignores malformed process keys instead of removing a wrong row", func() { + mc.requestReply = []byte(`{ + "success": true, + "reports_stopped_processes": true, + "stopped_process_keys": ["no-replica-suffix", "qwen3-0.6b#notanumber", "good#1"] + }`) + + _, err := adapter.DeleteBackend("node-1", "llama-cpp") + Expect(err).NotTo(HaveOccurred()) + Expect(locator.removedReplicas).To(ConsistOf(modelReplicaRef{"node-1", "good", 1})) + }) + }) + + Describe("UpgradeBackend", func() { + It("drops the replica rows for every process the worker stopped", func() { + // An upgrade force-stops every process using the binary and starts + // none of them back up, so it recycles ports exactly as delete does + // while leaving the same rows behind. + mc.requestReply = []byte(`{ + "success": true, + "reports_stopped_processes": true, + "stopped_process_keys": ["whisper#0", "whisper#1"] + }`) + + reply, err := adapter.UpgradeBackend("node-1", "whisper", "", "", "", "", 0, "", nil) + Expect(err).NotTo(HaveOccurred()) + Expect(reply.Success).To(BeTrue()) + + Expect(locator.removedReplicas).To(ConsistOf( + modelReplicaRef{"node-1", "whisper", 0}, + modelReplicaRef{"node-1", "whisper", 1}, + )) + }) + + It("does not read an old worker's silence as a completed cleanup", func() { + mc.requestReply = []byte(`{"success": true}`) + + reply, err := adapter.UpgradeBackend("node-1", "whisper", "", "", "", "", 0, "", nil) + Expect(err).NotTo(HaveOccurred()) + Expect(reply.ReportsStoppedProcesses).To(BeFalse()) + Expect(locator.removedReplicas).To(BeEmpty()) + }) + }) + + Describe("reply wire compatibility", func() { + It("decodes an old worker's delete reply without the new fields", func() { + // Guards the rolling-upgrade direction that matters: a new + // controller must keep working against every worker already + // deployed, not just ones rebuilt from this commit. + var reply messaging.BackendDeleteReply + Expect(json.Unmarshal([]byte(`{"success": true}`), &reply)).To(Succeed()) + Expect(reply.Success).To(BeTrue()) + Expect(reply.ReportsStoppedProcesses).To(BeFalse()) + Expect(reply.StoppedProcessKeys).To(BeEmpty()) + }) + }) +}) diff --git a/core/services/nodes/unloader_test.go b/core/services/nodes/unloader_test.go index dc2c7e8b1ff6..f95169eb83e7 100644 --- a/core/services/nodes/unloader_test.go +++ b/core/services/nodes/unloader_test.go @@ -20,9 +20,10 @@ import ( // fakeModelLocator implements ModelLocator with configurable node lists. type fakeModelLocator struct { - nodes []BackendNode - findErr error - removedPairs []modelNodePair // records RemoveNodeModel calls + nodes []BackendNode + findErr error + removedPairs []modelNodePair // records RemoveNodeModel calls + removedReplicas []modelReplicaRef // records RemoveNodeModel calls including the replica index } type modelNodePair struct { @@ -30,12 +31,24 @@ type modelNodePair struct { modelName string } +// modelReplicaRef records a row removal at full replica granularity. +// modelNodePair drops the index because RemoveAllNodeModelReplicas has none; +// the backend delete/upgrade paths address exactly one replica row, so an +// assertion that ignored the index could not tell a correct removal from one +// that wiped a sibling replica still serving traffic. +type modelReplicaRef struct { + nodeID string + modelName string + replicaIndex int +} + func (f *fakeModelLocator) FindNodesWithModel(_ context.Context, _ string) ([]BackendNode, error) { return f.nodes, f.findErr } -func (f *fakeModelLocator) RemoveNodeModel(_ context.Context, nodeID, modelName string, _ int) error { +func (f *fakeModelLocator) RemoveNodeModel(_ context.Context, nodeID, modelName string, replicaIndex int) error { f.removedPairs = append(f.removedPairs, modelNodePair{nodeID, modelName}) + f.removedReplicas = append(f.removedReplicas, modelReplicaRef{nodeID, modelName, replicaIndex}) return nil } @@ -126,8 +139,44 @@ var _ = Describe("RemoteUnloaderAdapter", func() { adapter = NewRemoteUnloaderAdapter(locator, mc, 3*time.Minute, 15*time.Minute) }) + // HasRemoteModel carries the distinction that UnloadRemoteModel + // deliberately does not, so ShutdownModel can answer 404 for a model that + // is loaded neither locally nor anywhere in the cluster without making the + // shared unload path fail for every idempotent cleanup caller. + Describe("HasRemoteModel", func() { + It("reports false when no node has the model", func() { + locator.nodes = nil + loaded, err := adapter.HasRemoteModel(context.Background(), "my-model") + Expect(err).ToNot(HaveOccurred()) + Expect(loaded).To(BeFalse()) + }) + + It("reports true when a node has the model", func() { + locator.nodes = []BackendNode{{ID: "node-1", Name: "worker-1"}} + loaded, err := adapter.HasRemoteModel(context.Background(), "my-model") + Expect(err).ToNot(HaveOccurred()) + Expect(loaded).To(BeTrue()) + }) + + It("surfaces a registry failure instead of reporting absence", func() { + // An unreachable registry is not evidence that the model is gone; + // reporting false would let ShutdownModel answer a confident 404 + // on the strength of a failed lookup. + locator.findErr = errors.New("registry unavailable") + _, err := adapter.HasRemoteModel(context.Background(), "my-model") + Expect(err).To(HaveOccurred()) + }) + }) + Describe("UnloadRemoteModel", func() { It("with no nodes returns nil", func() { + // Unloading is idempotent: cleanup paths (model deletion, config + // edits, watchdog eviction) legitimately run against an already + // unloaded model, and turning that into an error wedges the + // watchdog's LRU reclaimer, which only untracks a model when + // shutdown reports success. The same contract is pinned end to end + // by "should be no-op for models not on any node" in + // tests/e2e/distributed/node_lifecycle_test.go — keep them in step. locator.nodes = nil Expect(adapter.UnloadRemoteModel("my-model")).To(Succeed()) Expect(mc.published).To(BeEmpty()) diff --git a/core/services/worker/backend_delete_test.go b/core/services/worker/backend_delete_test.go new file mode 100644 index 000000000000..1fb0c52ecced --- /dev/null +++ b/core/services/worker/backend_delete_test.go @@ -0,0 +1,208 @@ +package worker + +import ( + "github.com/mudler/LocalAI/core/gallery" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// Production incident (Jetson/Thor worker): deleting a backend returned HTTP +// 200 but left its gRPC process alive for ~40 minutes with its directory +// removed from disk. A later model load was routed to that orphan and failed +// with a certifi path pointing into the deleted directory. +// +// Root cause: s.processes is keyed by `modelID#replicaIndex` +// (buildProcessKey), so a delete keyed on the *backend* name resolved to no +// keys at all and the stop silently no-op'd. These specs pin the two lookups +// that must work by backend name. +var _ = Describe("Backend deletion reaps the backend's processes", func() { + const ( + concrete = "cuda13-nvidia-l4t-arm64-longcat-video" + development = "cuda13-nvidia-l4t-arm64-longcat-video-development" + alias = "longcat-video" + ) + + Describe("resolveProcessKeysForBackend", func() { + It("finds processes started for a backend even though they are keyed by modelID", func() { + s := &backendSupervisor{ + processes: map[string]*backendProcess{ + // Started by a model load: key is modelID#replica, and the + // backend name survives only in backendName. + "LongCat-Video#0": {addr: "127.0.0.1:30232", backendName: concrete}, + "LongCat-Video#1": {addr: "127.0.0.1:30234", backendName: concrete}, + "other-model#0": {addr: "127.0.0.1:30235", backendName: "llama-cpp"}, + }, + } + + Expect(s.resolveProcessKeysForBackend(setOf(concrete))).To( + ConsistOf("LongCat-Video#0", "LongCat-Video#1"), + "a delete keyed on the backend name must reach every process started for it") + }) + + It("matches through the alias the model config used at install time", func() { + // The model config referenced `backend: longcat-video` (the alias), + // so the process recorded the alias; the delete request carries the + // concrete directory name. Without alias resolution the orphan + // survives exactly as it did in production. + s := &backendSupervisor{ + processes: map[string]*backendProcess{ + "LongCat-Video#0": {addr: "127.0.0.1:30232", backendName: alias}, + }, + } + + Expect(s.resolveProcessKeysForBackend(setOf(concrete, alias))).To( + ConsistOf("LongCat-Video#0"), + "alias and concrete name must resolve to the same process") + }) + + It("does not reap processes belonging to a different backend", func() { + s := &backendSupervisor{ + processes: map[string]*backendProcess{ + "LongCat-Video#0": {addr: "127.0.0.1:30233", backendName: development}, + }, + } + + Expect(s.resolveProcessKeysForBackend(setOf(concrete))).To(BeEmpty(), + "deleting the non-development backend must not kill the -development process") + }) + + It("still resolves legacy entries that predate backendName via the modelID prefix", func() { + // Installs with an empty modelID key the map by the backend name + // itself (buildProcessKey falls back). Those must keep working. + s := &backendSupervisor{ + processes: map[string]*backendProcess{ + concrete + "#0": {addr: "127.0.0.1:30232"}, + }, + } + + Expect(s.resolveProcessKeysForBackend(setOf(concrete))).To(ConsistOf(concrete + "#0")) + }) + }) + + Describe("resolveStopTargets", func() { + // backend.stop is ambiguous by design: the admin backends UI publishes + // a BACKEND name, UnloadRemoteModel publishes a MODEL name, and the + // router's abandoned-load reap (#10948) publishes an exact + // modelID#replica key. Resolving only one of the three silently + // strands the others. + newSupervisor := func() *backendSupervisor { + return &backendSupervisor{ + processes: map[string]*backendProcess{ + "LongCat-Video#0": {addr: "127.0.0.1:30232", backendName: concrete}, + }, + } + } + + It("stops a process addressed by model ID", func() { + Expect(newSupervisor().resolveStopTargets("LongCat-Video")).To(ConsistOf("LongCat-Video#0"), + "model-name stop (UnloadRemoteModel) must still reach the process") + }) + + It("stops a process addressed by backend name", func() { + Expect(newSupervisor().resolveStopTargets(concrete)).To(ConsistOf("LongCat-Video#0"), + "backend-name stop must reach the process too") + }) + + It("stops a process addressed by an exact replica key", func() { + Expect(newSupervisor().resolveStopTargets("LongCat-Video#0")).To(ConsistOf("LongCat-Video#0"), + "the router reaps an abandoned load by exact modelID#replica key") + }) + + It("does not stop unrelated processes", func() { + Expect(newSupervisor().resolveStopTargets("some-other-model")).To(BeEmpty()) + }) + }) + + Describe("processMatchesBackend", func() { + It("rejects reusing a process that was started for a different backend", func() { + // The install fast path returns the address of any live process + // under this processKey. After the concrete backend was deleted and + // the -development variant installed, the same model+replica got + // the deleted backend's port handed back to it. + s := &backendSupervisor{ + processes: map[string]*backendProcess{ + "LongCat-Video#0": {addr: "127.0.0.1:30232", backendName: concrete}, + }, + } + + Expect(s.processMatchesBackend("LongCat-Video#0", development)).To(BeFalse(), + "a process started for the deleted backend must not be reused for another backend") + }) + + It("accepts reusing a process started for the same backend", func() { + s := &backendSupervisor{ + processes: map[string]*backendProcess{ + "LongCat-Video#0": {addr: "127.0.0.1:30232", backendName: concrete}, + }, + } + + Expect(s.processMatchesBackend("LongCat-Video#0", concrete)).To(BeTrue()) + }) + + It("accepts a process recorded under the alias of the requested backend", func() { + s := &backendSupervisor{ + processes: map[string]*backendProcess{ + "LongCat-Video#0": {addr: "127.0.0.1:30232", backendName: alias}, + }, + } + + Expect(s.processMatchesBackend("LongCat-Video#0", alias)).To(BeTrue()) + }) + + It("accepts legacy entries with no recorded backend name", func() { + // Pre-upgrade processes carry no backendName. Treating them as a + // mismatch would restart every running backend once on rollout. + s := &backendSupervisor{ + processes: map[string]*backendProcess{ + "LongCat-Video#0": {addr: "127.0.0.1:30232"}, + }, + } + + Expect(s.processMatchesBackend("LongCat-Video#0", concrete)).To(BeTrue()) + }) + }) + + Describe("backendIdentitySet", func() { + It("maps an alias to its concrete backend and back", func() { + backends := gallery.SystemBackends{ + concrete: {Name: concrete, Metadata: &gallery.BackendMetadata{Name: concrete, Alias: alias}}, + alias: {Name: alias, Metadata: &gallery.BackendMetadata{Name: concrete, Alias: alias}}, + } + + Expect(backendIdentitySet(backends, concrete)).To(HaveKey(alias)) + Expect(backendIdentitySet(backends, alias)).To(HaveKey(concrete)) + }) + + It("falls back to the bare name when the backend is unknown", func() { + // A delete must never fail (or over-reach) because the gallery + // listing is unavailable or the entry is already gone from disk. + Expect(backendIdentitySet(gallery.SystemBackends{}, concrete)).To( + Equal(map[string]struct{}{concrete: {}})) + }) + + It("does not conflate two concrete backends sharing an alias", func() { + // Both variants declare alias "longcat-video"; ListSystemBackends + // picks one for the alias row. Deleting the non-chosen concrete + // must not pull in the chosen one's identity. + backends := gallery.SystemBackends{ + development: {Name: development, Metadata: &gallery.BackendMetadata{Name: development, Alias: alias}}, + alias: {Name: alias, Metadata: &gallery.BackendMetadata{Name: development, Alias: alias}}, + } + + set := backendIdentitySet(backends, development) + Expect(set).To(HaveKey(development)) + Expect(set).ToNot(HaveKey(concrete)) + }) + }) +}) + +// setOf builds the identity set the resolver consumes, keeping the specs +// readable without repeating the map literal. +func setOf(names ...string) map[string]struct{} { + set := make(map[string]struct{}, len(names)) + for _, n := range names { + set[n] = struct{}{} + } + return set +} diff --git a/core/services/worker/free_timeout_test.go b/core/services/worker/free_timeout_test.go new file mode 100644 index 000000000000..4f1b6346e749 --- /dev/null +++ b/core/services/worker/free_timeout_test.go @@ -0,0 +1,147 @@ +package worker + +import ( + "context" + "net" + "os" + "strconv" + "syscall" + + process "github.com/mudler/go-processmanager" + gogrpc "google.golang.org/grpc" + + pb "github.com/mudler/LocalAI/pkg/grpc/proto" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// pidAlive probes the OS directly for a process ID. The supervisor's own +// liveness helpers all go through go-processmanager's pidfile, which Stop +// deletes as part of releasing the handle, so they report "not alive" even if +// the signal never landed. Only asking the kernel proves termination. +func pidAlive(pid string) bool { + n, err := strconv.Atoi(pid) + if err != nil { + return false + } + proc, err := os.FindProcess(n) + if err != nil { + return false + } + return proc.Signal(syscall.Signal(0)) == nil +} + +// hangingBackend is a real gRPC backend server whose Free handler never +// returns. It models the production failure mode: 37 Python backends default to +// PYTHON_GRPC_MAX_WORKERS=1, so while a LoadModel handler is wedged (e.g. +// inside torch.load) the single gRPC worker thread is occupied and a Free RPC +// queues behind it forever. +// +// The server must be real rather than a socket that merely accepts: without a +// completed HTTP/2 handshake the call is bounded by gRPC's own ~20s connect +// timeout, which would mask an unbounded Free instead of exposing it. Here the +// connection reaches READY, so nothing but the caller's own deadline can end +// the call. +type hangingBackend struct { + pb.UnimplementedBackendServer + entered chan struct{} + release chan struct{} +} + +func (h *hangingBackend) Free(_ context.Context, _ *pb.HealthMessage) (*pb.Result, error) { + close(h.entered) + // Deliberately ignores the server-side context: a backend whose only + // worker thread is blocked cannot observe cancellation either. + <-h.release + return &pb.Result{Success: true}, nil +} + +var _ = Describe("Stopping a backend whose Free never returns", func() { + var ( + proc *process.Process + procPID string + backend *hangingBackend + server *gogrpc.Server + s *backendSupervisor + ) + + BeforeEach(func() { + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + + backend = &hangingBackend{ + entered: make(chan struct{}), + release: make(chan struct{}), + } + server = gogrpc.NewServer() + pb.RegisterBackendServer(server, backend) + go func() { _ = server.Serve(lis) }() + + // A real, long-running child so the spec can assert the process is + // actually dead afterwards, not merely that Stop() returned. It + // outlives every timeout below, so if it is gone at the end it is + // because the supervisor signalled it. + proc = process.New( + process.WithTemporaryStateDir(), + process.WithName("/bin/sleep"), + process.WithArgs("300"), + ) + Expect(proc.Run()).To(Succeed()) + + procPID = proc.CurrentPID() + Expect(procPID).ToNot(BeEmpty(), "the fixture process must report a PID once started") + Expect(pidAlive(procPID)).To(BeTrue(), "the fixture process must be running before the stop") + + s = &backendSupervisor{ + cfg: &Config{}, + processes: map[string]*backendProcess{ + "wedged-model#0": { + proc: proc, + addr: lis.Addr().String(), + port: lis.Addr().(*net.TCPAddr).Port, + backendName: "longcat-video", + }, + }, + } + }) + + AfterEach(func() { + close(backend.release) + server.Stop() + // A spec that fails before the stop would otherwise leave the sleep + // running for its full duration. + if pidAlive(procPID) { + _ = proc.Stop() + } + }) + + It("reaches the stop even though Free never returns", func() { + done := make(chan error, 1) + go func() { done <- s.stopBackendExact("wedged-model#0", false) }() + + Eventually(backend.entered, "20s", "100ms").Should(BeClosed(), + "Free must be attempted — it is a courtesy call we still want to make") + + // The whole point: a Free that never answers must not hold the stop + // hostage. Unbounded, this never receives, the process is never + // signalled, and the router's abandoned-load reap (#10948) is inert + // against a wedged single-worker Python backend. + Eventually(done, "60s", "200ms").Should(Receive(BeNil()), + "stopBackendExact must fall through to the process stop despite a hung Free") + + // Reaching the stop is not enough on its own: the signal has to land. + // Done closes only once go-processmanager has waited on the child, so + // this is the supervisor's kill being observed, not a pidfile that + // Stop deleted on its way out. + Eventually(proc.Done(), "30s", "100ms").Should(BeClosed(), + "the backend process must actually exit, not just be signalled") + Eventually(func() bool { return pidAlive(procPID) }, "30s", "100ms").Should(BeFalse(), + "the backend process must be gone from the OS after the stop") + + s.mu.Lock() + defer s.mu.Unlock() + Expect(s.processes).ToNot(HaveKey("wedged-model#0"), + "the supervisor must release the slot so the port can be recycled") + }) +}) diff --git a/core/services/worker/install.go b/core/services/worker/install.go index 19056ea7a20d..67c5bc2f0eb2 100644 --- a/core/services/worker/install.go +++ b/core/services/worker/install.go @@ -72,30 +72,46 @@ func (s *backendSupervisor) installBackend(req messaging.BackendInstallRequest, // replica failed, retries the install, the supervisor says "already // running" again, and the cluster loops on a dead replica forever. if addr := s.getAddr(processKey); addr != "" { - if s.isRunning(processKey) { + switch { + case !s.processMatchesBackend(processKey, req.Backend): + // The slot is held by a process started from a DIFFERENT + // backend — typically this model's previous backend was + // deleted (or superseded by a -development variant) while its + // process stayed up. Reusing that address would serve the load + // from a backend directory that may no longer exist on disk. + xlog.Warn("Process for this model replica belongs to another backend; restarting it", + "backend", req.Backend, "model", req.ModelID, "replica", req.ReplicaIndex, "addr", addr) + if err := s.stopBackendExact(processKey, false); err != nil { + return "", fmt.Errorf("stopping mismatched backend process before reinstall: %w", err) + } + case s.isRunning(processKey): xlog.Info("Backend already running for model replica", "backend", req.Backend, "model", req.ModelID, "replica", req.ReplicaIndex, "addr", addr) return addr, nil + default: + xlog.Warn("Stale process entry for backend (dead process); cleaning up before reinstall", + "backend", req.Backend, "model", req.ModelID, "replica", req.ReplicaIndex, "addr", addr) + if err := s.stopBackendExact(processKey, false); err != nil { + xlog.Warn("Failed to clean up stale process entry", "processKey", processKey, "error", err) + } } - 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, false) } } else { // Upgrade path: stop every live process that shares this backend so the // gallery install can overwrite the on-disk artifact and the restarted - // process picks up the new binary. resolveProcessKeys catches peer - // replicas of the same backend (whisper#0, whisper#1, ...) on workers - // configured with MaxReplicasPerModel>1. We also stop the exact - // processKey from the request tuple — keys created with an explicit - // modelID don't share the bare-name prefix the resolver matches, but - // they're still using the old binary and need to come down. Both calls - // are no-ops on missing keys. - toStop := s.resolveProcessKeys(req.Backend) + // process picks up the new binary. resolveProcessKeysForBackend finds + // them by recorded backend name (and alias), which also covers peer + // replicas (whisper#0, whisper#1, ...) on workers configured with + // MaxReplicasPerModel>1. We also stop the exact processKey from the + // request tuple, whose recorded name may be absent on legacy entries. + // Both are no-ops on missing keys. + toStop := s.resolveProcessKeysForBackend(s.backendIdentity(req.Backend)) toStop = append(toStop, processKey) for _, key := range toStop { xlog.Info("Force install: stopping running backend before reinstall", "backend", req.Backend, "processKey", key) - s.stopBackendExact(key, true) + if err := s.stopBackendExact(key, true); err != nil { + return "", fmt.Errorf("stopping running backend before reinstall: %w", err) + } } } @@ -160,8 +176,10 @@ func (s *backendSupervisor) installBackend(req messaging.BackendInstallRequest, xlog.Info("Found backend binary", "path", backendPath, "processKey", processKey) - // Start the gRPC process on a new port (keyed by model, not just backend) - return s.startBackend(processKey, backendPath) + // Start the gRPC process on a new port (keyed by model, not just backend). + // req.Backend is recorded on the process so a later backend.delete/stop can + // find it by backend name. + return s.startBackend(processKey, req.Backend, backendPath) } // upgradeBackend stops every running process for `backend`, force-reinstalls @@ -170,22 +188,38 @@ func (s *backendSupervisor) installBackend(req messaging.BackendInstallRequest, // backend.install will spawn a fresh process picking up the new binary. // // The caller is responsible for holding s.lockBackend(req.Backend). -func (s *backendSupervisor) upgradeBackend(req messaging.BackendUpgradeRequest) error { +// +// It returns the process keys it terminated so the controller can drop the +// NodeModel rows addressing them: an upgrade stops every process using the +// binary and starts none back up, recycling their gRPC ports while the rows +// still point at those addresses. +func (s *backendSupervisor) upgradeBackend(req messaging.BackendUpgradeRequest) ([]string, error) { // Stop every live process for this backend (peer replicas + the bare // processKey). Same logic as the force branch in installBackend. - toStop := s.resolveProcessKeys(req.Backend) + toStop := s.resolveProcessKeysForBackend(s.backendIdentity(req.Backend)) toStop = append(toStop, buildProcessKey("", req.Backend, int(req.ReplicaIndex))) + stopped := make([]string, 0, len(toStop)) for _, key := range toStop { + // The bare-backend key appended above is speculative — it exists only + // on legacy entries. Reporting a key we never had would ask the + // controller to drop a row that a different, still-running replica may + // own, so only keys backed by a tracked process are reported. + tracked := s.getAddr(key) != "" xlog.Info("Upgrade: stopping running backend before reinstall", "backend", req.Backend, "processKey", key) - s.stopBackendExact(key, true) + if err := s.stopBackendExact(key, true); err != nil { + return stopped, fmt.Errorf("stopping running backend before upgrade: %w", err) + } + if tracked { + stopped = append(stopped, key) + } } galleries := s.galleries if req.BackendGalleries != "" { var reqGalleries []config.Gallery if err := json.Unmarshal([]byte(req.BackendGalleries), &reqGalleries); err != nil { - return fmt.Errorf("decoding backend galleries: %w", err) + return stopped, fmt.Errorf("decoding backend galleries: %w", err) } galleries = reqGalleries } @@ -207,7 +241,7 @@ func (s *backendSupervisor) upgradeBackend(req messaging.BackendUpgradeRequest) if err := galleryop.InstallExternalBackend( context.Background(), galleries, s.systemState, s.ml, downloadCb, req.URI, req.Name, req.Alias, true, s.cfg.RequireBackendIntegrity, ); err != nil { - return fmt.Errorf("upgrading backend from external URI: %w", err) + return stopped, fmt.Errorf("upgrading backend from external URI: %w", err) } } else { xlog.Info("Upgrading backend from gallery", "backend", req.Backend) @@ -215,14 +249,14 @@ func (s *backendSupervisor) upgradeBackend(req messaging.BackendUpgradeRequest) context.Background(), galleries, s.systemState, s.ml, req.Backend, downloadCb, true, /* force */ s.cfg.RequireBackendIntegrity, ); err != nil { - return fmt.Errorf("upgrading backend from gallery: %w", err) + return stopped, fmt.Errorf("upgrading backend from gallery: %w", err) } } if err := gallery.RegisterBackends(s.systemState, s.ml); err != nil { - return fmt.Errorf("refreshing registered backends after upgrade: %w", err) + return stopped, fmt.Errorf("refreshing registered backends after upgrade: %w", err) } - return nil + return stopped, nil } // findBackend looks for the backend binary in the backends path and system path. diff --git a/core/services/worker/lifecycle.go b/core/services/worker/lifecycle.go index 0348e41ecf3c..904a7a7f5db9 100644 --- a/core/services/worker/lifecycle.go +++ b/core/services/worker/lifecycle.go @@ -4,7 +4,9 @@ import ( "context" "encoding/json" "fmt" + "maps" "net" + "slices" "syscall" "github.com/mudler/LocalAI/core/gallery" @@ -114,12 +116,25 @@ func (s *backendSupervisor) handleBackendUpgrade(data []byte, reply func([]byte) release := s.lockBackend(req.Backend) defer release() - if err := s.upgradeBackend(req); err != nil { + // stopped is meaningful even on the error paths: it lists processes + // already terminated (and ports already recycled) before the failure, so + // the controller must drop those rows regardless of the outcome. + stopped, err := s.upgradeBackend(req) + if err != nil { xlog.Error("Failed to upgrade backend via NATS", "error", err) - replyJSON(reply, messaging.BackendUpgradeReply{Success: false, Error: err.Error()}) + replyJSON(reply, messaging.BackendUpgradeReply{ + Success: false, + Error: err.Error(), + StoppedProcessKeys: stopped, + ReportsStoppedProcesses: true, + }) return } - replyJSON(reply, messaging.BackendUpgradeReply{Success: true}) + replyJSON(reply, messaging.BackendUpgradeReply{ + Success: true, + StoppedProcessKeys: stopped, + ReportsStoppedProcesses: true, + }) }() } @@ -137,7 +152,14 @@ func (s *backendSupervisor) handleBackendStop(data []byte) { return } xlog.Info("Received NATS backend.stop event", "backend", req.Backend, "force", req.Force) - s.stopBackend(req.Backend, req.Force) + // The identifier may be a backend name, a model name, or an exact + // modelID#replica key depending on the publisher; resolveStopTargets + // handles all three. stopBackend alone resolves only the model meanings. + for _, key := range s.resolveStopTargets(req.Backend) { + if err := s.stopBackendExact(key, req.Force); err != nil { + xlog.Error("Failed to stop backend process", "backend", req.Backend, "processKey", key, "error", err) + } + } } func decodeBackendStopRequest(data []byte) (messaging.BackendStopRequest, bool, error) { @@ -162,28 +184,67 @@ func (s *backendSupervisor) handleBackendDelete(data []byte, reply func([]byte)) } xlog.Info("Received NATS backend.delete event", "backend", req.Backend) - // Stop if running this backend - if s.isRunning(req.Backend) { - s.stopBackend(req.Backend, false) + // Resolve the backend's identity (concrete name + alias) BEFORE touching + // the filesystem: DeleteBackendFromSystem removes the metadata.json that + // carries the alias, and a model loaded via the alias records the alias as + // its process's backend name. + identity := s.backendIdentity(req.Backend) + + // Stop every process started for this backend. Processes are keyed by + // modelID#replica, so the lookup must match the recorded backend name — a + // lookup by backend name alone resolved to nothing and left the process + // running with its directory deleted underneath it. + keys := s.resolveProcessKeysForBackend(identity) + if len(keys) == 0 { + // Not an error: deleting a backend that was never loaded is routine. + // But log it — silence here is what made the orphan case invisible. + xlog.Info("Deleting backend with no matching running process", + "backend", req.Backend, "identity", slices.Sorted(maps.Keys(identity))) + } + // Accumulate the processes we actually terminate. Every stop hands a gRPC + // port back to this worker's allocator while the controller still holds a + // NodeModel row for that address, so the controller needs these keys to + // drop those rows before the port is re-bound by an unrelated backend. A + // key is appended only after its process is confirmed gone, which is what + // lets the controller trust the list on the partial-failure replies below. + stopped := make([]string, 0, len(keys)) + deleteReply := func(success bool, errMsg string) messaging.BackendDeleteReply { + return messaging.BackendDeleteReply{ + Success: success, + Error: errMsg, + StoppedProcessKeys: stopped, + ReportsStoppedProcesses: true, + } + } + + for _, key := range keys { + if err := s.stopBackendExact(key, false); err != nil { + // We knew about this process and could not kill it. Replying + // success would repeat the original defect: the operator is told + // "backend deleted" while the process keeps serving requests. + xlog.Error("Failed to stop backend process during delete; aborting delete", + "backend", req.Backend, "processKey", key, "error", err) + replyJSON(reply, deleteReply(false, fmt.Sprintf("could not stop running process %s: %v", key, err))) + return + } + stopped = append(stopped, key) } // Delete the backend files if err := gallery.DeleteBackendFromSystem(s.systemState, req.Backend); err != nil { xlog.Warn("Failed to delete backend files", "backend", req.Backend, "error", err) - resp := messaging.BackendDeleteReply{Success: false, Error: err.Error()} - replyJSON(reply, resp) + replyJSON(reply, deleteReply(false, err.Error())) return } // Re-register backends after deletion 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()}) + replyJSON(reply, deleteReply(false, err.Error())) return } - resp := messaging.BackendDeleteReply{Success: true} - replyJSON(reply, resp) + replyJSON(reply, deleteReply(true, "")) } // handleBackendList is the NATS callback for backend.list — reply with the diff --git a/core/services/worker/port_quarantine_test.go b/core/services/worker/port_quarantine_test.go new file mode 100644 index 000000000000..3dcd25d2d5da --- /dev/null +++ b/core/services/worker/port_quarantine_test.go @@ -0,0 +1,120 @@ +package worker + +import ( + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// quarantinedPortNumbers exposes the held ports for assertions without leaking +// the quarantine deadlines, which are wall-clock and not reproducible. +func quarantinedPortNumbers(s *backendSupervisor) []int { + ports := make([]int, 0, len(s.quarantinedPorts)) + for _, q := range s.quarantinedPorts { + ports = append(ports, q.port) + } + return ports +} + +var _ = Describe("Backend port quarantine", func() { + Describe("finishBackendStop", func() { + It("does not offer a just-stopped port to the next backend", func() { + // The controller still holds a NodeModel row pointing at this + // address. Handing the port straight to the next backend lets that + // row pass probeHealth against a different process and dispatch to + // it silently. The quarantine holds the port for the NATS + // round-trip the controller needs to drop the row. + bp := &backendProcess{port: 50051} + s := &backendSupervisor{ + processes: map[string]*backendProcess{"model#0": bp}, + nextPort: 50060, + portQuarantine: time.Hour, + } + + Expect(s.finishBackendStop("model#0", bp, nil)).To(Succeed()) + + Expect(s.freePorts).To(BeEmpty()) + Expect(s.allocatePort()).NotTo(Equal(50051)) + }) + + It("returns the port to the free pool once the quarantine expires", func() { + // The quarantine must not leak ports: nextPort only ever grows, so a + // port that never comes back is a permanent loss of allocator range. + bp := &backendProcess{port: 50051} + s := &backendSupervisor{ + processes: map[string]*backendProcess{"model#0": bp}, + nextPort: 50060, + portQuarantine: time.Millisecond, + } + + Expect(s.finishBackendStop("model#0", bp, nil)).To(Succeed()) + time.Sleep(5 * time.Millisecond) + + Expect(s.allocatePort()).To(Equal(50051)) + }) + }) + + Describe("releaseBackendStart", func() { + It("quarantines the port of a startup that failed after binding", func() { + // A backend that died during startup may have bound and served on + // the port for a moment, which is long enough for a load to have + // recorded the address. + bp := &backendProcess{port: 50051} + s := &backendSupervisor{ + processes: map[string]*backendProcess{"model#0": bp}, + nextPort: 50060, + portQuarantine: time.Hour, + } + + s.releaseBackendStart("model#0", bp) + + Expect(s.freePorts).To(BeEmpty()) + Expect(s.allocatePort()).To(Equal(50060)) + }) + }) + + Describe("allocatePort", func() { + It("prefers a released port over growing the range", func() { + s := &backendSupervisor{ + processes: map[string]*backendProcess{}, + nextPort: 50060, + portQuarantine: time.Millisecond, + } + + s.releasePort(50051) + time.Sleep(5 * time.Millisecond) + + Expect(s.allocatePort()).To(Equal(50051)) + Expect(s.allocatePort()).To(Equal(50060)) + Expect(s.nextPort).To(Equal(50061)) + }) + + It("does not sweep a port whose quarantine is still running", func() { + s := &backendSupervisor{ + processes: map[string]*backendProcess{}, + nextPort: 50060, + portQuarantine: time.Hour, + } + + s.releasePort(50051) + s.releasePort(50052) + + Expect(s.allocatePort()).To(Equal(50060)) + Expect(s.allocatePort()).To(Equal(50061)) + }) + + It("applies the default quarantine when none is configured", func() { + // The zero value must not degrade to "no quarantine": every + // supervisor built outside the tests leaves the field unset. + s := &backendSupervisor{ + processes: map[string]*backendProcess{}, + nextPort: 50060, + } + + s.releasePort(50051) + + Expect(s.allocatePort()).To(Equal(50060)) + }) + }) +}) diff --git a/core/services/worker/replica_test.go b/core/services/worker/replica_test.go index 1600b982ee1e..a8cbfdc6c5f0 100644 --- a/core/services/worker/replica_test.go +++ b/core/services/worker/replica_test.go @@ -136,10 +136,13 @@ var _ = Describe("Worker per-replica process keying", func() { Expect(bp.stopping).To(BeTrue()) Expect(s.processes).To(HaveKeyWithValue("model#0", bp)) Expect(s.freePorts).To(BeEmpty()) + Expect(quarantinedPortNumbers(s)).To(BeEmpty()) s.finishBackendStop("model#0", bp, nil) Expect(s.processes).NotTo(HaveKey("model#0")) - Expect(s.freePorts).To(ConsistOf(50051)) + // Released into quarantine rather than straight into the free pool: + // a controller row may still name this address. + Expect(quarantinedPortNumbers(s)).To(ConsistOf(50051)) }) It("does not report a starting backend as ready after stop begins", func() { @@ -163,7 +166,7 @@ var _ = Describe("Worker per-replica process keying", func() { s.releaseBackendStart("model#0", bp) Expect(s.processes).NotTo(HaveKey("model#0")) - Expect(s.freePorts).To(ConsistOf(50051)) + Expect(quarantinedPortNumbers(s)).To(ConsistOf(50051)) }) It("does not recycle the port owned by a replacement startup", func() { @@ -177,6 +180,7 @@ var _ = Describe("Worker per-replica process keying", func() { Expect(s.processes).To(HaveKeyWithValue("model#0", replacement)) Expect(s.freePorts).To(BeEmpty()) + Expect(quarantinedPortNumbers(s)).To(BeEmpty()) }) }) diff --git a/core/services/worker/supervisor.go b/core/services/worker/supervisor.go index 702071b74d1c..e065518d5982 100644 --- a/core/services/worker/supervisor.go +++ b/core/services/worker/supervisor.go @@ -11,6 +11,7 @@ import ( "time" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/gallery" "github.com/mudler/LocalAI/core/services/messaging" grpc "github.com/mudler/LocalAI/pkg/grpc" "github.com/mudler/LocalAI/pkg/model" @@ -25,10 +26,58 @@ type backendProcess struct { addr string // gRPC address (host:port) port int stopping bool + // backendName is the gallery backend this process was started for (e.g. + // "cuda13-nvidia-l4t-arm64-longcat-video"). It is NOT derivable from the + // map key: keys are `modelID#replicaIndex`, so without it a backend.delete + // keyed on the backend name resolves to nothing and the process is + // orphaned while its files are removed from disk. + backendName string } const workerBackendFreeTimeout = 5 * time.Second +// backendIdentitySet returns every name that refers to the same on-disk +// backend as `name`: the concrete directory name plus its alias. Deletes +// arrive with the concrete name while installs arrive with whatever the model +// config declared (often the alias), so both must land on one identity. +// +// An unknown name yields just itself. A delete must never over-reach (or fail) +// because the gallery listing is unavailable or the entry is already gone. +func backendIdentitySet(backends gallery.SystemBackends, name string) map[string]struct{} { + set := map[string]struct{}{name: {}} + b, ok := backends[name] + if !ok || b.Metadata == nil { + return set + } + // Only this entry's own metadata is consulted. Two concrete backends can + // share one alias (foo and foo-development); walking every candidate of + // that alias would let deleting one reap the other's process. + if b.Metadata.Name != "" { + set[b.Metadata.Name] = struct{}{} + } + if b.Metadata.Alias != "" { + set[b.Metadata.Alias] = struct{}{} + } + return set +} + +// backendIdentity resolves `name` against this worker's installed backends. +// Callers must invoke it BEFORE deleting files: DeleteBackendFromSystem +// removes the metadata.json that carries the alias mapping. +func (s *backendSupervisor) backendIdentity(name string) map[string]struct{} { + if s.systemState == nil { + return map[string]struct{}{name: {}} + } + backends, err := gallery.ListSystemBackends(s.systemState) + if err != nil { + // Degrade to name-only matching rather than failing the operation. + xlog.Warn("Could not list backends for alias resolution; matching on name only", + "backend", name, "error", err) + return map[string]struct{}{name: {}} + } + return backendIdentitySet(backends, name) +} + // 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 { @@ -43,7 +92,13 @@ type backendSupervisor struct { mu sync.Mutex processes map[string]*backendProcess // key: backend name nextPort int // next available port for new backends - freePorts []int // ports freed by stopBackend, reused before nextPort + freePorts []int // ports out of quarantine, reused before nextPort + + // quarantinedPorts holds ports whose process has terminated but which are + // not yet safe to re-bind, and portQuarantine is how long they wait. + // Overridden only by tests; zero means defaultPortQuarantine. + quarantinedPorts []quarantinedPort + portQuarantine time.Duration // backendLocks serializes gallery operations against the same on-disk // artifact. Two installs of different backends on the same worker run @@ -54,9 +109,92 @@ type backendSupervisor struct { backendLocks map[string]*sync.Mutex } +// defaultPortQuarantine is how long a released gRPC port waits before it can be +// re-bound by another backend. +// +// This is an interlock for one specific window and nothing more: between the +// moment this worker frees a port and the moment the controller processes our +// reply and drops the NodeModel rows naming that address, a row still resolves +// to a live listener. probeHealth verifies liveness, not identity, so a port +// re-bound inside that window is dispatched to as if it were the original +// backend. The window is a NATS round-trip plus a row delete, so seconds of +// slack are ample. +// +// Deliberately NOT derived from the controller's HealthCheckInterval or from +// the per-model miss threshold. Tying a worker-local constant to a +// controller-side cadence would be unsound: that cadence is operator-tunable +// and the per-model reaper can be switched off entirely with +// DisablePerModelHealthCheck, so the coupling would be silently wrong on some +// clusters. Eager row removal, not this delay, is what actually fixes stale +// rows; raising this value is not a substitute for it. +const defaultPortQuarantine = 15 * time.Second + +// quarantinedPort is a released port that must not be re-bound until `until`. +type quarantinedPort struct { + port int + until time.Time +} + +// quarantineWindow returns the configured quarantine, defaulting when unset so +// the zero-value supervisor never degrades to immediate reuse. +func (s *backendSupervisor) quarantineWindow() time.Duration { + if s.portQuarantine <= 0 { + return defaultPortQuarantine + } + return s.portQuarantine +} + +// allocatePort returns a gRPC port for a new backend process, preferring a +// previously released port over growing the range. Callers must hold s.mu. +func (s *backendSupervisor) allocatePort() int { + s.sweepQuarantine() + if len(s.freePorts) > 0 { + port := s.freePorts[len(s.freePorts)-1] + s.freePorts = s.freePorts[:len(s.freePorts)-1] + return port + } + port := s.nextPort + s.nextPort++ + return port +} + +// sweepQuarantine moves ports whose quarantine has elapsed into the free pool. +// Sweeping lazily on allocation avoids a timer goroutine per stopped backend; +// the only observer of the free pool is allocation itself. Callers must hold +// s.mu. +func (s *backendSupervisor) sweepQuarantine() { + if len(s.quarantinedPorts) == 0 { + return + } + now := time.Now() + held := s.quarantinedPorts[:0] + for _, q := range s.quarantinedPorts { + if now.Before(q.until) { + held = append(held, q) + continue + } + s.freePorts = append(s.freePorts, q.port) + } + s.quarantinedPorts = held +} + +// releasePort returns a port to the allocator after a quarantine period. +// Callers must hold s.mu. See defaultPortQuarantine for why the port is not +// immediately reusable. +func (s *backendSupervisor) releasePort(port int) { + s.quarantinedPorts = append(s.quarantinedPorts, quarantinedPort{ + port: port, + until: time.Now().Add(s.quarantineWindow()), + }) +} + // startBackend starts a gRPC backend process on a dynamically allocated port. // Returns the gRPC address. -func (s *backendSupervisor) startBackend(backend, backendPath string) (string, error) { +// +// `backend` is the process key (modelID#replica); `backendName` is the gallery +// backend it was started from, recorded so delete/stop can find this process +// by backend name later. +func (s *backendSupervisor) startBackend(backend, backendName, backendPath string) (string, error) { s.mu.Lock() // Already running? @@ -74,29 +212,22 @@ func (s *backendSupervisor) startBackend(backend, backendPath string) (string, e delete(s.processes, backend) } - // Allocate port — recycle freed ports first, then grow upward from basePort - var port int - if len(s.freePorts) > 0 { - port = s.freePorts[len(s.freePorts)-1] - s.freePorts = s.freePorts[:len(s.freePorts)-1] - } else { - port = s.nextPort - s.nextPort++ - } + port := s.allocatePort() bindAddr := fmt.Sprintf("0.0.0.0:%d", port) clientAddr := fmt.Sprintf("127.0.0.1:%d", port) proc, err := s.ml.StartProcess(backendPath, backend, bindAddr) if err != nil { - s.freePorts = append(s.freePorts, port) + s.releasePort(port) s.mu.Unlock() return "", fmt.Errorf("starting backend process: %w", err) } s.processes[backend] = &backendProcess{ - proc: proc, - addr: clientAddr, - port: port, + proc: proc, + addr: clientAddr, + port: port, + backendName: backendName, } xlog.Info("Backend process started", "backend", backend, "addr", clientAddr) @@ -185,7 +316,7 @@ func (s *backendSupervisor) releaseBackendStart(key string, bp *backendProcess) 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) + s.releasePort(bp.port) } // resolveProcessKeys turns a caller-supplied identifier into the set of @@ -219,12 +350,119 @@ func (s *backendSupervisor) resolveProcessKeys(id string) []string { return keys } +// resolveProcessKeysForBackend returns the process keys of every process +// started for any name in `names` (the backend's identity set). +// +// resolveProcessKeys alone is not enough here: it resolves a bare *modelID* +// via the `id#N` prefix, but backend.delete/backend.stop carry a *backend* +// name, which never prefixes a `modelID#replica` key. That mismatch is what +// let a deleted backend's process survive with its files removed. Matching on +// the recorded backendName is the primary path; the prefix resolution is +// unioned in so legacy entries keyed by the backend name itself (installs with +// an empty modelID) keep resolving. +func (s *backendSupervisor) resolveProcessKeysForBackend(names map[string]struct{}) []string { + seen := make(map[string]struct{}) + var keys []string + add := func(k string) { + if _, dup := seen[k]; dup { + return + } + seen[k] = struct{}{} + keys = append(keys, k) + } + + s.mu.Lock() + for key, bp := range s.processes { + if bp.backendName == "" { + continue + } + if _, ok := names[bp.backendName]; ok { + add(key) + } + } + s.mu.Unlock() + + // Legacy fallback for entries predating backendName: an install with an + // empty modelID keys the map by the backend name itself. Restricted to + // entries with no recorded name so it cannot reap a process that belongs + // to a different backend but whose modelID happens to equal this backend's + // name or alias. + for name := range names { + for _, k := range s.resolveProcessKeys(name) { + s.mu.Lock() + bp, ok := s.processes[k] + unnamed := ok && bp.backendName == "" + s.mu.Unlock() + if unnamed { + add(k) + } + } + } + return keys +} + +// resolveStopTargets resolves the ambiguous identifier carried on backend.stop. +// +// The payload field is named "backend", but it is published with both +// meanings: the admin backends UI sends a BACKEND name, while +// UnloadRemoteModel and the router's abandoned-load reap send a MODEL name (or +// an exact modelID#replica key). Resolving only one meaning silently strands +// the other — matching on backend name alone stops nothing for a model unload, +// and matching on the modelID prefix alone is the bug that orphaned deleted +// backends' processes. Stop is best-effort and idempotent, so accepting both is +// safe here; delete stays strict because its identifier is unambiguously a +// backend name. +func (s *backendSupervisor) resolveStopTargets(id string) []string { + keys := s.resolveProcessKeysForBackend(s.backendIdentity(id)) + seen := make(map[string]struct{}, len(keys)) + for _, k := range keys { + seen[k] = struct{}{} + } + for _, k := range s.resolveProcessKeys(id) { + if _, dup := seen[k]; !dup { + seen[k] = struct{}{} + keys = append(keys, k) + } + } + return keys +} + +// processMatchesBackend reports whether the process under `key` may be reused +// to serve `backend`. The install fast path returns any live process for a +// (model, replica) slot; without this check a slot whose backend was deleted +// and replaced handed the new install the deleted backend's port, so the load +// ran against a directory that no longer exists. +// +// A process with no recorded backendName predates this field and is accepted: +// treating it as a mismatch would restart every running backend once on +// rollout. +func (s *backendSupervisor) processMatchesBackend(key, backend string) bool { + s.mu.Lock() + bp, ok := s.processes[key] + if !ok { + s.mu.Unlock() + return false + } + recorded := bp.backendName + s.mu.Unlock() + + if recorded == "" || recorded == backend { + return true + } + // Names differ textually — they may still be alias and concrete for the + // same on-disk backend, which is a legitimate reuse. + _, match := s.backendIdentity(backend)[recorded] + return match +} + // 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, force bool) { for _, key := range s.resolveProcessKeys(id) { - s.stopBackendExact(key, force) + if err := s.stopBackendExact(key, force); err != nil { + xlog.Error("Failed to stop backend process", "processKey", key, "error", err) + } } } @@ -232,10 +470,15 @@ func (s *backendSupervisor) stopBackend(id string, force bool) { // 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) { +// +// Returns an error only when a process was found and is still alive after the +// stop attempt — a missing key is a no-op, not a failure. Callers that report +// success to an operator (backend.delete) must not claim success while the +// process keeps serving requests from a directory they just removed. +func (s *backendSupervisor) stopBackendExact(key string, force bool) error { bp := s.beginBackendStop(key) if bp == nil { - return + return nil } if !force { @@ -248,12 +491,12 @@ func (s *backendSupervisor) stopBackendExact(key string, force bool) { cancel() } - xlog.Info("Stopping backend process", "backend", key, "addr", bp.addr, "force", force) + xlog.Info("Stopping backend process", "backend", key, "addr", bp.addr, "force", force, "backendName", bp.backendName) stopErr := bp.proc.Stop() if stopErr != nil { xlog.Error("Error stopping backend process", "backend", key, "error", stopErr) } - s.finishBackendStop(key, bp, stopErr) + return s.finishBackendStop(key, bp, stopErr) } // beginBackendStop reserves both the process entry and its port while network @@ -269,25 +512,28 @@ func (s *backendSupervisor) beginBackendStop(key string) *backendProcess { return bp } -func (s *backendSupervisor) finishBackendStop(key string, bp *backendProcess, stopErr error) { +// finishBackendStop returns a non-nil error when the process survived the stop +// attempt, so callers can distinguish a real reap from a failed one. +func (s *backendSupervisor) finishBackendStop(key string, bp *backendProcess, stopErr error) 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 + return nil } if stopErr != nil && bp.proc.IsAlive() { bp.stopping = false - return + return fmt.Errorf("stopping backend process %s: %w", key, stopErr) } 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 + return nil } - s.freePorts = append(s.freePorts, bp.port) + s.releasePort(bp.port) + return nil } // stopAllBackends stops all running backend processes. @@ -303,8 +549,12 @@ func (s *backendSupervisor) stopAllBackends(force bool) { // isRunning returns whether at least one backend process matching the given // identifier is currently running. Accepts a bare modelID (matches any -// replica) or a full processKey (exact match). Callers like the -// backend.delete pre-check rely on the bare-name path. +// replica) or a full processKey (exact match). +// +// It does NOT accept a backend name: backend names never prefix a +// modelID#replica key. Callers holding a backend name must go through +// resolveProcessKeysForBackend — assuming otherwise here is what left deleted +// backends' processes running. func (s *backendSupervisor) isRunning(id string) bool { keys := s.resolveProcessKeys(id) if len(keys) == 0 { diff --git a/go.mod b/go.mod index fa3608c620e0..0b3037f9ccb2 100644 --- a/go.mod +++ b/go.mod @@ -38,7 +38,7 @@ require ( github.com/modelcontextprotocol/go-sdk v1.5.0 github.com/mudler/cogito v0.10.1-0.20260609212329-bf4010d31047 github.com/mudler/edgevpn v0.34.0 - github.com/mudler/go-processmanager v0.1.1 + github.com/mudler/go-processmanager v0.1.2-0.20260719202549-a94e2b7e5f4f github.com/mudler/memory v0.0.0-20260406210934-424c1ecf2cf8 github.com/mudler/xlog v0.0.6 github.com/nats-io/jwt/v2 v2.7.4 diff --git a/go.sum b/go.sum index 129976bbcd0f..b30f0e83510b 100644 --- a/go.sum +++ b/go.sum @@ -974,8 +974,8 @@ github.com/mudler/edgevpn v0.34.0 h1:qDrD/rCPFY/FdURbXudIZWihVKY4VOX3nMn3CcbeQEU github.com/mudler/edgevpn v0.34.0/go.mod h1:yki7uMi5LR9gSMrw8PdPieuxsrk8BLV2Ui7VBEmbbIA= github.com/mudler/go-piper v0.0.0-20241023091659-2494246fd9fc h1:RxwneJl1VgvikiX28EkpdAyL4yQVnJMrbquKospjHyA= github.com/mudler/go-piper v0.0.0-20241023091659-2494246fd9fc/go.mod h1:O7SwdSWMilAWhBZMK9N9Y/oBDyMMzshE3ju8Xkexwig= -github.com/mudler/go-processmanager v0.1.1 h1:c/1NRZOZpW8HuFv9RhBG57nQu1oDMRomEHedwBFMlrw= -github.com/mudler/go-processmanager v0.1.1/go.mod h1:h6kmHUZeafr+k5hRYpGLMzJFH4hItHffgpRo2QIkP+o= +github.com/mudler/go-processmanager v0.1.2-0.20260719202549-a94e2b7e5f4f h1:td22tuommnMukfgj92yAzKhkQOQ/V+uWbyzRtPWlLbw= +github.com/mudler/go-processmanager v0.1.2-0.20260719202549-a94e2b7e5f4f/go.mod h1:h6kmHUZeafr+k5hRYpGLMzJFH4hItHffgpRo2QIkP+o= github.com/mudler/localrecall v0.6.3 h1:uXOrP9JmetzxgVKzSrawviyBHZfAcvPBBIrvVUdZjDA= github.com/mudler/localrecall v0.6.3/go.mod h1:28k5n19raUrkuwXkacdNsBlj8yuSnGhpT16tu+2+4dU= github.com/mudler/memory v0.0.0-20260406210934-424c1ecf2cf8 h1:Ry8RiWy8fZ6Ff4E7dPmjRsBrnHOnPeOOj2LhCgyjQu0= diff --git a/pkg/model/backend_log_store.go b/pkg/model/backend_log_store.go index a0988790ff9e..c5b5253ddc40 100644 --- a/pkg/model/backend_log_store.go +++ b/pkg/model/backend_log_store.go @@ -27,6 +27,28 @@ func BackendProcessKey(modelID string, replicaIndex int) string { return modelID + replicaSeparator + strconv.Itoa(replicaIndex) } +// ParseBackendProcessKey is the inverse of BackendProcessKey. It exists so the +// controller can map process keys a worker reports back to the (model_name, +// replica_index) pair identifying a NodeModel row, without a fourth hand-rolled +// copy of the format. +// +// The split is anchored at the LAST separator because model IDs may contain +// '#' themselves and BackendProcessKey appends the index at the end. Splitting +// at the first separator would address a row that does not exist and silently +// leave the real one in place. ok is false for anything this function did not +// produce, so callers drop unparseable keys rather than acting on a guess. +func ParseBackendProcessKey(key string) (modelID string, replicaIndex int, ok bool) { + i := strings.LastIndex(key, replicaSeparator) + if i <= 0 || i == len(key)-1 { + return "", 0, false + } + idx, err := strconv.Atoi(key[i+1:]) + if err != nil || idx < 0 { + return "", 0, false + } + return key[:i], idx, true +} + // BackendLogLine represents a single line of output from a backend process. type BackendLogLine struct { Timestamp time.Time `json:"timestamp"` diff --git a/pkg/model/backend_log_store_test.go b/pkg/model/backend_log_store_test.go index 0b8cab4c2594..775e07cdb561 100644 --- a/pkg/model/backend_log_store_test.go +++ b/pkg/model/backend_log_store_test.go @@ -199,3 +199,40 @@ var _ = Describe("BackendLogStore", func() { }) }) }) + +var _ = Describe("ParseBackendProcessKey", func() { + It("round-trips every key BackendProcessKey produces", func() { + for _, tc := range []struct { + modelID string + replica int + }{ + {"qwen3-0.6b", 0}, + {"qwen3-0.6b", 7}, + {"weird#name", 3}, + {"a", 0}, + } { + modelID, replica, ok := model.ParseBackendProcessKey(model.BackendProcessKey(tc.modelID, tc.replica)) + Expect(ok).To(BeTrue(), "key for %s#%d", tc.modelID, tc.replica) + Expect(modelID).To(Equal(tc.modelID)) + Expect(replica).To(Equal(tc.replica)) + } + }) + + It("splits at the last separator so '#' in a model ID survives", func() { + // Splitting at the first separator would yield the model "weird" and + // address a row that does not exist, leaving the real stale row behind. + modelID, replica, ok := model.ParseBackendProcessKey("weird#name#3") + Expect(ok).To(BeTrue()) + Expect(modelID).To(Equal("weird#name")) + Expect(replica).To(Equal(3)) + }) + + It("rejects anything it did not produce", func() { + // Callers act on the result by deleting a routing row, so a guess is + // worse than a refusal. + for _, key := range []string{"", "#", "no-suffix", "model#", "#0", "model#x", "model#-1", "model#0.5"} { + _, _, ok := model.ParseBackendProcessKey(key) + Expect(ok).To(BeFalse(), "expected %q to be rejected", key) + } + }) +}) diff --git a/pkg/model/loader.go b/pkg/model/loader.go index 01fc9f8b51b2..827ce1bdde15 100644 --- a/pkg/model/loader.go +++ b/pkg/model/loader.go @@ -41,6 +41,20 @@ type RemoteModelContextUnloader interface { UnloadRemoteModelContext(ctx context.Context, modelName string, force bool) error } +// RemoteModelPresenceChecker reports whether any node in the cluster currently +// holds the model. +// +// Unloading is deliberately idempotent: UnloadRemoteModel succeeds when no node +// has the model, because cleanup paths (model deletion, config edits, watchdog +// eviction) legitimately run against an already-unloaded model and must not +// fail. That makes the unload result unable to distinguish "stopped it" from +// "there was nothing to stop", so the one caller that needs the distinction — +// ShutdownModel, which must report 404 rather than a misleading success for a +// model loaded nowhere — asks separately, before unloading. +type RemoteModelPresenceChecker interface { + HasRemoteModel(ctx context.Context, modelName string) (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. diff --git a/pkg/model/process.go b/pkg/model/process.go index 33eb6377e0ca..f6533e57827a 100644 --- a/pkg/model/process.go +++ b/pkg/model/process.go @@ -20,7 +20,9 @@ import ( var forceBackendShutdown bool = os.Getenv("LOCALAI_FORCE_BACKEND_SHUTDOWN") == "true" var ( - modelNotFoundErr = errors.New("model not found") + // ErrModelNotFound reports that a model is not loaded. Exported so HTTP + // handlers can map it to 404 instead of a blanket 500. + ErrModelNotFound = 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") @@ -33,6 +35,16 @@ const ( busyPollInterval = 100 * time.Millisecond ) +// unloadRemote asks the remote unloader to stop `s` on whichever node holds +// it, preferring the context-aware extension so a forced shutdown and the +// caller's deadline both survive the distributed boundary. +func unloadRemote(ctx context.Context, u RemoteModelUnloader, s string, force bool) error { + if contextUnloader, ok := u.(RemoteModelContextUnloader); ok { + return contextUnloader.UnloadRemoteModelContext(ctx, s, force) + } + return u.UnloadRemoteModel(s) +} + // 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 may be stuck in an @@ -58,8 +70,34 @@ func (ml *ModelLoader) deleteProcess(ctx context.Context, s string, force bool) model, ok := store.Get(s) if !ok { + // A local-store miss is not proof the model is unloaded. In + // distributed mode the model runs on a worker and the authoritative + // record is the shared node registry: any frontend replica that did + // not itself serve the model (load balancer picked a peer, or this + // replica restarted) has no local entry. Returning "model not found" + // here reported a model that was demonstrably running as absent, and + // left its backend process untouched. + if remoteUnloader != nil { + xlog.Debug("Model not in local store; asking the remote unloader", "model", s) + // Ask BEFORE unloading. Unloading is idempotent by contract, so it + // cannot afterwards tell us whether anything was actually stopped, + // and only a model absent locally AND cluster-wide may be reported + // as not found. + if checker, ok := remoteUnloader.(RemoteModelPresenceChecker); ok { + loaded, err := checker.HasRemoteModel(ctx, s) + if err != nil { + // An unreachable registry is not evidence of absence; + // saying "not found" here would be a guess presented as fact. + return fmt.Errorf("checking whether model %q is loaded on any node: %w", s, err) + } + if !loaded { + return ErrModelNotFound + } + } + return unloadRemote(ctx, remoteUnloader, s, force) + } xlog.Debug("Model not found", "model", s) - return modelNotFoundErr + return ErrModelNotFound } if !force { @@ -114,11 +152,7 @@ func (ml *ModelLoader) deleteProcess(ctx context.Context, s string, force bool) var unloadErr error if remoteUnloader != nil { xlog.Debug("Delegating model unload to remote unloader", "model", s) - if contextUnloader, ok := remoteUnloader.(RemoteModelContextUnloader); ok { - unloadErr = contextUnloader.UnloadRemoteModelContext(ctx, s, force) - } else { - unloadErr = remoteUnloader.UnloadRemoteModel(s) - } + unloadErr = unloadRemote(ctx, remoteUnloader, s, force) if unloadErr != nil { xlog.Warn("Remote model unload failed", "model", s, "error", unloadErr) } @@ -187,7 +221,7 @@ func (ml *ModelLoader) GetGRPCPID(id string) (int, error) { if p.Process() == nil { return -1, fmt.Errorf("no grpc backend found for %s", id) } - return strconv.Atoi(p.Process().PID) + return strconv.Atoi(p.Process().CurrentPID()) } // StartProcess starts a gRPC backend process and returns its process handle. diff --git a/pkg/model/remote_shutdown_test.go b/pkg/model/remote_shutdown_test.go new file mode 100644 index 000000000000..4433cd28d4cb --- /dev/null +++ b/pkg/model/remote_shutdown_test.go @@ -0,0 +1,124 @@ +package model_test + +import ( + "context" + "errors" + + "github.com/mudler/LocalAI/pkg/model" + "github.com/mudler/LocalAI/pkg/system" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// fakeRemoteUnloader records the models it was asked to unload so the specs can +// assert the remote path was actually taken (not merely that no error +// surfaced). It mirrors the real adapter's contract: unloading is idempotent +// and reports nil even when nothing was loaded, so presence is a separate +// question. +type fakeRemoteUnloader struct { + called []string + unloadErr error + + present bool + presenceErr error + asked []string +} + +func (f *fakeRemoteUnloader) UnloadRemoteModel(modelName string) error { + f.called = append(f.called, modelName) + return f.unloadErr +} + +func (f *fakeRemoteUnloader) HasRemoteModel(_ context.Context, modelName string) (bool, error) { + f.asked = append(f.asked, modelName) + return f.present, f.presenceErr +} + +// unloaderWithoutPresence is a RemoteModelUnloader that does NOT implement +// RemoteModelPresenceChecker, pinning the compatibility path for third-party +// implementations of the older interface. +type unloaderWithoutPresence struct{ called []string } + +func (u *unloaderWithoutPresence) UnloadRemoteModel(modelName string) error { + u.called = append(u.called, modelName) + return nil +} + +// In distributed mode the authoritative record of "is this model loaded" is +// the shared node registry, not this replica's in-memory store. A frontend +// replica that never served the model itself (load balancer picked another +// replica, or this one restarted) has no local entry, so ShutdownModel +// short-circuited on a local-store miss and reported "model not found" for a +// model that was demonstrably running on a worker — while the remote unload +// path it documents was never reached. +var _ = Describe("ShutdownModel in distributed mode", func() { + var ( + modelLoader *model.ModelLoader + unloader *fakeRemoteUnloader + ) + + BeforeEach(func() { + systemState, err := system.GetSystemState(system.WithModelPath(GinkgoT().TempDir())) + Expect(err).ToNot(HaveOccurred()) + modelLoader = model.NewModelLoader(systemState) + unloader = &fakeRemoteUnloader{} + }) + + It("delegates to the remote unloader when the model is not in the local store", func() { + unloader.present = true + modelLoader.SetRemoteUnloader(unloader) + + err := modelLoader.ShutdownModel("longcat-video-avatar-1.5") + + Expect(unloader.called).To(ConsistOf("longcat-video-avatar-1.5"), + "a model absent locally may still be loaded on a worker; the remote unloader must be consulted") + Expect(err).ToNot(HaveOccurred(), + "stopping a model that is running on a worker must succeed, not report 'model not found'") + }) + + It("reports not-found only after the registry confirms no node has it", func() { + unloader.present = false + modelLoader.SetRemoteUnloader(unloader) + + err := modelLoader.ShutdownModel("never-loaded") + + Expect(unloader.asked).To(ConsistOf("never-loaded"), + "the registry must be consulted before declaring a model not found") + Expect(err).To(MatchError(model.ErrModelNotFound), + "absent locally AND cluster-wide is the only case that may report not-found") + Expect(unloader.called).To(BeEmpty(), + "nothing to unload — no point publishing a stop for a model no node holds") + }) + + It("does not claim not-found when the registry lookup fails", func() { + // An unreachable registry is not evidence of absence. Reporting 404 + // here would tell an operator the model is gone on the strength of a + // failed lookup. + unloader.presenceErr = errors.New("registry unavailable") + modelLoader.SetRemoteUnloader(unloader) + + err := modelLoader.ShutdownModel("maybe-loaded") + + Expect(err).To(HaveOccurred()) + Expect(err).ToNot(MatchError(model.ErrModelNotFound)) + }) + + It("still unloads via an unloader that cannot answer presence", func() { + // Older RemoteModelUnloader implementations have no presence check. + // They must keep working: attempt the unload rather than refusing it. + legacy := &unloaderWithoutPresence{} + modelLoader.SetRemoteUnloader(legacy) + + err := modelLoader.ShutdownModel("some-model") + + Expect(legacy.called).To(ConsistOf("some-model")) + Expect(err).ToNot(HaveOccurred()) + }) + + It("still reports not-found when no remote unloader is configured", func() { + // Single-node behavior must be unchanged. + err := modelLoader.ShutdownModel("never-loaded") + Expect(err).To(MatchError(model.ErrModelNotFound)) + }) +}) diff --git a/pkg/model/watchdog.go b/pkg/model/watchdog.go index efe400e41f6f..3a604668a28a 100644 --- a/pkg/model/watchdog.go +++ b/pkg/model/watchdog.go @@ -1,6 +1,7 @@ package model import ( + "errors" "slices" "sync" "time" @@ -899,7 +900,7 @@ func (wd *WatchDog) evictLRUModel() { if wasBusy { shutdown = wd.pm.ShutdownModelForce } - if err := shutdown(lruModel.model); err != nil && err != modelNotFoundErr { + if err := shutdown(lruModel.model); err != nil && !errors.Is(err, ErrModelNotFound) { xlog.Error("[WatchDog] error shutting down model during memory reclamation", "error", err, "model", lruModel.model) } else { // Untrack the model diff --git a/tests/e2e/distributed/node_lifecycle_test.go b/tests/e2e/distributed/node_lifecycle_test.go index e0539e9120fc..04b7342e794b 100644 --- a/tests/e2e/distributed/node_lifecycle_test.go +++ b/tests/e2e/distributed/node_lifecycle_test.go @@ -141,6 +141,15 @@ var _ = Describe("Node Backend Lifecycle (NATS-driven)", Label("Distributed"), f }) It("should be no-op for models not on any node", func() { + // Unloading is idempotent by contract: cleanup paths (model + // deletion, config edits, watchdog eviction) legitimately run + // against an already-unloaded model, and the watchdog's LRU + // reclaimer only untracks a model when shutdown reports success. + // Callers needing to tell "stopped it" from "nothing to stop" — + // ShutdownModel, so it can answer 404 — use HasRemoteModel instead. + // The same contract is pinned at unit level by "with no nodes + // returns nil" in core/services/nodes/unloader_test.go; keep them + // in step. adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) Expect(adapter.UnloadRemoteModel("nonexistent-model")).To(Succeed()) })