diff --git a/core/application/distributed.go b/core/application/distributed.go index 1744126cf..af7bd5501 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -410,10 +410,17 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade }, cfg.Distributed.RegistrationToken, workerHTTPDialer) xlog.Info("File stager initialized (HTTP direct transfer)") } + // The frontend's control plane client. It reaches every worker over that + // worker's own tunnel, on the same `http` stream tag the file stager above + // uses, so a control RPC to a worker another replica holds is relayed the + // way an inference request is. + controlClient := nodes.NewControlClient(workerHTTPDialer, cfg.Distributed.RegistrationToken) + // Create RemoteUnloaderAdapter — needed by SmartRouter and startup.go remoteUnloader := nodes.NewRemoteUnloaderAdapter( registry, natsClient, + controlClient, cfg.Distributed.BackendInstallTimeoutOrDefault(), cfg.Distributed.BackendUpgradeTimeoutOrDefault(), ) diff --git a/core/services/messaging/backend_install_progress.go b/core/services/messaging/backend_install_progress.go index 268ef86b9..52065bf64 100644 --- a/core/services/messaging/backend_install_progress.go +++ b/core/services/messaging/backend_install_progress.go @@ -11,10 +11,11 @@ const ( PhaseStarting = "starting" // worker is spawning the gRPC backend process ) -// BackendInstallProgressEvent is the wire payload published by a worker to -// nodes..backend.install..progress while a long-running install -// is in flight. Transient: dropped events are acceptable, the master relies -// on BackendInstallReply for ground truth on success/failure. +// BackendInstallProgressEvent is the wire payload a worker writes as a progress +// line of its backend.install and backend.upgrade responses while a +// long-running install is in flight. Transient: a line the frontend cannot read +// is acceptable, and BackendInstallReply is the ground truth on +// success/failure. // // Phase holds one of the Phase* constants above. type BackendInstallProgressEvent struct { @@ -27,10 +28,3 @@ type BackendInstallProgressEvent struct { Percentage float64 `json:"percentage"` Phase string `json:"phase,omitempty"` } - -// SubjectNodeBackendInstallProgress returns the NATS subject for transient -// progress events emitted by a worker during a single backend.install run. -// Per-op so multiple concurrent installs on the same node never alias. -func SubjectNodeBackendInstallProgress(nodeID, opID string) string { - return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".backend.install." + sanitizeSubjectToken(opID) + ".progress" -} diff --git a/core/services/messaging/backend_install_progress_test.go b/core/services/messaging/backend_install_progress_test.go index ec45f4619..e57c5505f 100644 --- a/core/services/messaging/backend_install_progress_test.go +++ b/core/services/messaging/backend_install_progress_test.go @@ -2,7 +2,6 @@ package messaging_test import ( "encoding/json" - "strings" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -26,23 +25,6 @@ var _ = Describe("Phase constants", func() { }) var _ = Describe("BackendInstallProgress", func() { - Context("SubjectNodeBackendInstallProgress", func() { - It("composes the per-op progress subject", func() { - Expect(messaging.SubjectNodeBackendInstallProgress("node-abc", "op-123")). - To(Equal("nodes.node-abc.backend.install.op-123.progress")) - }) - - It("sanitizes NATS-reserved characters in node and op tokens", func() { - // '.' is the NATS hierarchy delimiter, '*' and '>' are wildcards, - // and whitespace must be stripped - sanitizeSubjectToken replaces - // all of them with '-'. The resulting subject must still parse as - // exactly six hierarchy segments: nodes//backend/install//progress. - subj := messaging.SubjectNodeBackendInstallProgress("a.b c", "x.y z") - Expect(subj).ToNot(ContainSubstring(" ")) - Expect(strings.Count(subj, ".")).To(Equal(5)) - }) - }) - Context("BackendInstallProgressEvent", func() { It("JSON round-trips with all known fields", func() { ev := messaging.BackendInstallProgressEvent{ diff --git a/core/services/messaging/subjects.go b/core/services/messaging/subjects.go index 44c223f9c..fe9c1691a 100644 --- a/core/services/messaging/subjects.go +++ b/core/services/messaging/subjects.go @@ -141,26 +141,29 @@ func SubjectResponseCancel(responseID string) string { return subjectResponseCancelPrefix + sanitizeSubjectToken(responseID) + ".cancel" } -// Node Backend Lifecycle (Pub/Sub — targeted to specific nodes) +// Node Backend Lifecycle // -// These subjects control the backend *process* lifecycle on a serve-backend node, -// mirroring how the local ModelLoader uses startProcess() / deleteProcess(). +// The frontend's control plane no longer travels on NATS. The ten verbs that +// drove a worker's backend and model lifecycle are HTTP routes under +// workerctl.Prefix, served on the worker's own loopback server and reached +// through its tunnel, so a subject builder for any of them would be a subject +// nothing publishes and nothing subscribes to. // -// Model loading (LoadModel gRPC) is done via direct gRPC calls to the node's -// address — no NATS needed for that, same as local mode. +// ONE survives: backend.stop, and only for AGENT workers. They hold no tunnel, +// so they have no control plane to serve, and they subscribe to it to drop the +// MCP sessions cached for a backend that is going away. See +// nodes.RemoteUnloaderAdapter.stopBackend for the split, and +// core/cli/agent_worker.go for the subscriber. +// +// The request and reply types below are UNCHANGED and still live here: they are +// the wire format of the control routes, byte for byte what the subjects +// carried, so a worker and a frontend from different releases still understand +// each other. const ( subjectNodePrefix = "nodes." ) -// SubjectNodeBackendInstall tells a worker node to install a backend and start its gRPC process. -// Uses NATS request-reply: the SmartRouter sends the request, the worker installs -// the backend from gallery (if not already installed), starts the gRPC process, -// and replies when ready. -func SubjectNodeBackendInstall(nodeID string) string { - return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".backend.install" -} - -// BackendInstallRequest is the payload for a backend.install NATS request. +// BackendInstallRequest is the payload for a backend.install control request. type BackendInstallRequest struct { Backend string `json:"backend"` ModelID string `json:"model_id,omitempty"` @@ -178,20 +181,20 @@ type BackendInstallRequest struct { ReplicaIndex int32 `json:"replica_index,omitempty"` // Force is retained on the wire only for backward compatibility with // pre-2026-05-08 masters that did not know about backend.upgrade. New - // callers MUST send to SubjectNodeBackendUpgrade instead. Workers continue + // callers MUST use workerctl.PathBackendUpgrade instead. Workers continue // to honor Force=true here so a rolling update with new master + old // worker still works (the master's install fallback path also uses this - // when backend.upgrade returns nats.ErrNoResponders). + // when the worker answers that it does not serve the upgrade verb). Force bool `json:"force,omitempty"` - // OpID identifies the admin-side operation. When non-empty the worker - // publishes BackendInstallProgressEvent values to - // SubjectNodeBackendInstallProgress(nodeID, OpID) while the install is - // running, debounced to roughly 250ms. Empty means the caller is a - // reconciler-driven retry that does not need progress streamed. + // OpID identifies the admin-side operation. It travels so the worker can + // name the operation on the BackendInstallProgressEvent values it writes + // into the install response ahead of the reply, debounced to roughly 250ms. + // Empty means the caller is a reconciler-driven retry that does not need + // progress streamed. OpID string `json:"op_id,omitempty"` } -// BackendInstallReply is the response from a backend.install NATS request. +// BackendInstallReply is the response from a backend.install control request. type BackendInstallReply struct { Success bool `json:"success"` // WorkerLocalAddress is where the backend process listens ON THE WORKER, @@ -207,17 +210,7 @@ type BackendInstallReply struct { Error string `json:"error,omitempty"` } -// SubjectNodeBackendUpgrade tells a worker node to force-reinstall a backend -// from the gallery, stop every running process for that backend, and restart. -// Uses NATS request-reply with a long deadline (gallery image pulls can take -// many minutes on slow links). Routine model loads use SubjectNodeBackendInstall -// instead — this subject exists so the slow path doesn't head-of-line-block -// the fast one through a shared subscription goroutine. -func SubjectNodeBackendUpgrade(nodeID string) string { - return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".backend.upgrade" -} - -// BackendUpgradeRequest is the payload for a backend.upgrade NATS request. +// BackendUpgradeRequest is the payload for a backend.upgrade control request. // It is intentionally a strict subset of BackendInstallRequest — there is no // Force field because the upgrade subject IS the force semantics; no ModelID // because upgrade is backend-scoped (it stops every replica using the binary @@ -232,13 +225,9 @@ type BackendUpgradeRequest struct { // but the field lets future per-replica metadata (e.g. progress reporting // scoped to a slot) ride the same wire without a v3 type. ReplicaIndex int32 `json:"replica_index,omitempty"` - // OpID identifies the admin-side operation. When non-empty the worker - // publishes BackendInstallProgressEvent values to - // SubjectNodeBackendInstallProgress(nodeID, OpID) while the force-reinstall - // runs, so the master can stream per-node progress for upgrades exactly as - // it already does for installs (an upgrade IS a force-reinstall, so the - // install-progress subject is reused rather than minting a new one — no new - // NATS permission or rolling-update compat surface). Empty on legacy callers. + // OpID identifies the admin-side operation, so an upgrade streams per-node + // progress in its own response exactly as an install does. Empty on legacy + // callers. OpID string `json:"op_id,omitempty"` } @@ -258,16 +247,10 @@ type BackendUpgradeReply struct { ReportsStoppedProcesses bool `json:"reports_stopped_processes,omitempty"` } -// SubjectNodeBackendList queries a worker node for its installed backends. -// Uses NATS request-reply. -func SubjectNodeBackendList(nodeID string) string { - return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".backend.list" -} - -// BackendListRequest is the payload for a backend.list NATS request. +// BackendListRequest is the payload for a backend.list control request. type BackendListRequest struct{} -// BackendListReply is the response from a backend.list NATS request. +// BackendListReply is the response from a backend.list control request. type BackendListReply struct { Backends []NodeBackendInfo `json:"backends"` Error string `json:"error,omitempty"` @@ -296,21 +279,17 @@ type BackendStopRequest struct { 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 bounded Free() via gRPC (unless Force is true) -// 2. Kill the backend process -// 3. Can be restarted via another backend.start event. +// SubjectNodeBackendStop tells an AGENT worker that a backend is going away, so +// it can close the MCP sessions it cached for that backend. +// +// It is the one node subject left, and it is addressed only to agent nodes. A +// BACKEND worker takes its stop on workerctl.PathBackendStop over its tunnel, +// where it also kills the process and recycles the port; an agent worker runs +// no backend processes and only needs to hear that one went. func SubjectNodeBackendStop(nodeID string) string { return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".backend.stop" } -// SubjectNodeModelStop targets one supervisor process and acknowledges only -// after that process has exited and its worker-side resources are released. -func SubjectNodeModelStop(nodeID string) string { - return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".model.stop" -} - type ModelStopRequest struct { ModelName string `json:"model_name"` ProcessKey string `json:"process_key"` @@ -328,18 +307,12 @@ type ModelStopReply struct { Error string `json:"error,omitempty"` } -// SubjectNodeBackendDelete tells a worker node to delete a backend (stop + remove files). -// Uses NATS request-reply. -func SubjectNodeBackendDelete(nodeID string) string { - return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".backend.delete" -} - -// BackendDeleteRequest is the payload for a backend.delete NATS request. +// BackendDeleteRequest is the payload for a backend.delete control request. type BackendDeleteRequest struct { Backend string `json:"backend"` } -// BackendDeleteReply is the response from a backend.delete NATS request. +// BackendDeleteReply is the response from a backend.delete control request. type BackendDeleteReply struct { Success bool `json:"success"` Error string `json:"error,omitempty"` @@ -362,57 +335,33 @@ type BackendDeleteReply struct { ReportsStoppedProcesses bool `json:"reports_stopped_processes,omitempty"` } -// SubjectNodeModelUnload tells a worker node to unload a model (gRPC Free) without killing the backend. -// Uses NATS request-reply. -func SubjectNodeModelUnload(nodeID string) string { - return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".model.unload" -} - -// ModelUnloadRequest is the payload for a model.unload NATS request. +// ModelUnloadRequest is the payload for a model.unload control request. type ModelUnloadRequest struct { ModelName string `json:"model_name"` Address string `json:"address,omitempty"` // gRPC address of the backend process to unload from } -// ModelUnloadReply is the response from a model.unload NATS request. +// ModelUnloadReply is the response from a model.unload control request. type ModelUnloadReply struct { Success bool `json:"success"` Error string `json:"error,omitempty"` } -// SubjectNodeModelDelete tells a worker node to delete model files from disk. -// Uses NATS request-reply. -func SubjectNodeModelDelete(nodeID string) string { - return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".model.delete" -} - -// ModelDeleteRequest is the payload for a model.delete NATS request. +// ModelDeleteRequest is the payload for a model.delete control request. type ModelDeleteRequest struct { ModelName string `json:"model_name"` } -// ModelDeleteReply is the response from a model.delete NATS request. +// ModelDeleteReply is the response from a model.delete control request. type ModelDeleteReply struct { Success bool `json:"success"` Error string `json:"error,omitempty"` } -// SubjectNodeModelsRunning asks a worker node which model backend processes it -// currently has running. Uses NATS request-reply. -// -// This is the authoritative answer to "is this replica still alive". The worker -// owns the process table, so unlike a health probe against the backend's own -// serving port, its reply does not depend on whether that backend happens to be -// busy: a model mid-generation cannot answer a gRPC health check for minutes at -// a time, but the worker answers immediately either way. -func SubjectNodeModelsRunning(nodeID string) string { - return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".models.running" -} - -// ModelsRunningRequest is the payload for a models.running NATS request. +// ModelsRunningRequest is the payload for a models.running control request. type ModelsRunningRequest struct{} -// ModelsRunningReply is the response from a models.running NATS request. +// ModelsRunningReply is the response from a models.running control request. type ModelsRunningReply struct { Models []RunningModelInfo `json:"models"` Error string `json:"error,omitempty"` @@ -427,12 +376,6 @@ type RunningModelInfo struct { Address string `json:"address,omitempty"` } -// SubjectNodeStop tells a serve-backend node to shut down entirely -// (deregister + exit). The node will not restart the backend process. -func SubjectNodeStop(nodeID string) string { - return subjectNodePrefix + sanitizeSubjectToken(nodeID) + ".stop" -} - // File Staging (Request-Reply — targeted to specific nodes) // These subjects use request-reply for synchronous file operations. diff --git a/core/services/messaging/subjects_upgrade_test.go b/core/services/messaging/subjects_upgrade_test.go index e60369cfc..e1a059fac 100644 --- a/core/services/messaging/subjects_upgrade_test.go +++ b/core/services/messaging/subjects_upgrade_test.go @@ -7,15 +7,19 @@ import ( "github.com/mudler/LocalAI/core/services/messaging" ) -var _ = Describe("SubjectNodeBackendUpgrade", func() { - It("returns the per-node upgrade subject", func() { - Expect(messaging.SubjectNodeBackendUpgrade("abc")). - To(Equal("nodes.abc.backend.upgrade")) +// The surviving node subject. It is written out BY HAND and not derived from +// the builder: an agent worker built from another commit subscribes to this +// literal, and a renamed subject is silence that looks exactly like a worker +// that never started. +var _ = Describe("SubjectNodeBackendStop", func() { + It("returns the per-node stop subject an agent worker subscribes to", func() { + Expect(messaging.SubjectNodeBackendStop("abc")). + To(Equal("nodes.abc.backend.stop")) }) It("sanitizes reserved NATS tokens in the node id", func() { - Expect(messaging.SubjectNodeBackendUpgrade("a.b*c")). - To(Equal("nodes.a-b-c.backend.upgrade")) + Expect(messaging.SubjectNodeBackendStop("a.b*c")). + To(Equal("nodes.a-b-c.backend.stop")) }) }) diff --git a/core/services/nodes/control_client.go b/core/services/nodes/control_client.go new file mode 100644 index 000000000..6c943892b --- /dev/null +++ b/core/services/nodes/control_client.go @@ -0,0 +1,307 @@ +package nodes + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "sync" + "time" + + "github.com/mudler/xlog" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" + "github.com/mudler/LocalAI/pkg/httpclient" +) + +// ErrWorkerControlUnsupported reports that the worker answered, and does not +// serve that control verb. +// +// It is a deployment fact and not a verdict about anything: the worker is +// running a build older than this frontend, so a path this frontend knows is a +// 404 there. The one caller that may act on it is the upgrade fallback, which +// re-issues the legacy force-install the older build does understand. +// +// It is a SPECIALISATION of ErrWorkerUnroutable for the same reason +// ErrNoWorkerDialer is: everything that DELETES a node_models row makes exactly +// one check, and a 404 must not pass it. The worker spoke, but it said nothing +// about any backend, and reaping on it would evict healthy work over a version +// skew. Callers that want the fallback match this sentinel specifically, which +// ErrWorkerUnroutable does not imply in that direction. +var ErrWorkerControlUnsupported = fmt.Errorf("%w: the worker does not serve that control verb", ErrWorkerUnroutable) + +// ControlClient issues the frontend's control RPCs to a worker over that +// worker's tunnel. +// +// It is the frontend half of core/services/workerctl: the worker mounts those +// paths on the loopback HTTP server it already runs, and this reaches them +// through the `http` stream tag, so a control RPC to a worker another replica +// holds is relayed exactly like an inference request. +// +// There is one of these per frontend rather than one per verb, because the +// mapping from a failed RPC onto the four conditions this system must never +// confuse belongs in ONE place. See controlFailure. +type ControlClient struct { + // dialFor supplies the transport for one worker. It is per node because a + // worker is reached over ITS OWN tunnel and an http.Transport carries one + // DialContext, so a single shared transport could only ever reach one + // worker. nil means no tunnel dialer is wired and every call is refused; + // see ErrNoWorkerDialer for why that is not a fallback to a direct dial. + dialFor WorkerNetDialerFor + token string + + // clients caches one *http.Client per node, so a verb issued twice in a + // row reuses the tunnel stream its transport already holds instead of + // opening a new one. + // + // Entries are never pruned, and that is judged rather than overlooked: the + // map is bounded by the number of distinct workers this frontend has ever + // commanded, which is bounded by the fleet, and a departed worker's entry + // holds a map slot plus a transport whose idle connections IdleConnTimeout + // reclaims. It is the same shape, and would need the same missing + // node-departure signal to fix, as HTTPFileStager.clients. + clientsMu sync.Mutex + clients map[string]*http.Client +} + +// NewControlClient returns the control client for the workers dialFor can +// reach, authenticating with the deployment's registration token. +func NewControlClient(dialFor WorkerNetDialerFor, token string) *ControlClient { + return &ControlClient{dialFor: dialFor, token: token, clients: map[string]*http.Client{}} +} + +// clientFor returns the HTTP client that reaches one worker, building it on +// first use. +// +// The transport mirrors HTTPFileStager.clientFor, which is the other consumer +// of a worker's own HTTP server, and differs only where control traffic +// differs. HTTP/2 stays OFF for the reason it always was: its flow control +// stalls large transfers, and the install verb streams for minutes. The +// stager's 256 KB socket buffers are dropped because a control body is a +// handful of kilobytes and two 256 KB buffers per worker would be paid for +// every node in the fleet. +// +// The net.Dialer's connect timeout and keepalive have nothing to act on here: +// there is no TCP connect to time out, and liveness on the link is the yamux +// session's keepalive rather than the socket's. No client.Timeout is set +// either, because it would bound the response BODY, and the install verb's +// body stays open for as long as the install runs. What bounds a call is the +// context its caller passes. +func (c *ControlClient) clientFor(nodeID string) (*http.Client, error) { + if c == nil || c.dialFor == nil { + return nil, fmt.Errorf("control rpc to node %q: %w", nodeID, ErrNoWorkerDialer) + } + c.clientsMu.Lock() + defer c.clientsMu.Unlock() + if cl, ok := c.clients[nodeID]; ok { + return cl, nil + } + dial := c.dialFor(nodeID) + if dial == nil { + return nil, fmt.Errorf("control rpc to node %q: %w", nodeID, ErrNoWorkerDialer) + } + transport := &http.Transport{ + DialContext: dial, + ForceAttemptHTTP2: false, // HTTP/2 flow control can stall a long streaming verb + MaxIdleConns: 10, + IdleConnTimeout: 90 * time.Second, + TLSHandshakeTimeout: 10 * time.Second, + ExpectContinueTimeout: 1 * time.Second, + } + cl := httpclient.New(httpclient.WithTransport(transport)) + c.clients[nodeID] = cl + return cl, nil +} + +// Call issues one control RPC and decodes its reply. reply may be nil for the +// verbs that answer 204. +// +// A reply that arrives with its Error field set is NOT an error here: the +// worker answered, and reading its answer is the caller's job. Only a failure +// to reach the worker, or an answer this frontend cannot read, comes back as an +// error, and every one of those is mapped by controlFailure. +func (c *ControlClient) Call(ctx context.Context, nodeID, path string, req, reply any) error { + resp, err := c.do(ctx, nodeID, path, req) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + + if resp.StatusCode == http.StatusNoContent || reply == nil { + // A verb that answers 204 sends no body at all; a caller that asked for + // no reply may still have been sent one, and draining it is what lets + // the transport keep the tunnel stream for the next verb instead of + // tearing it down. + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, maxControlErrorBodyBytes)) + return nil + } + if err := json.NewDecoder(resp.Body).Decode(reply); err != nil { + // A body this frontend cannot read is not the worker's verdict about + // anything, so it must not be reported as one. + return controlFailure(ctx, nodeID, fmt.Errorf("decoding the %s reply: %w", path, err)) + } + return nil +} + +// CallStreaming issues a control RPC whose response is an NDJSON envelope +// stream, invoking onProgress for each progress line and decoding the single +// terminal reply line into reply. onProgress may be nil. +// +// onProgress runs SYNCHRONOUSLY, on this goroutine. The NATS carrier ran each +// progress callback on a goroutine of its own because a slow callback there +// stalled the one reader thread every worker's events arrived on; here the only +// thing a slow callback holds up is this request's own body, which is the +// caller's business. Dropping the guard is also what makes the events arrive in +// the order the worker sent them. +func (c *ControlClient) CallStreaming(ctx context.Context, nodeID, path string, + req, reply any, onProgress func(messaging.BackendInstallProgressEvent)) error { + resp, err := c.do(ctx, nodeID, path, req) + if err != nil { + return err + } + defer func() { _ = resp.Body.Close() }() + + dec := json.NewDecoder(resp.Body) + var terminal json.RawMessage + for terminal == nil { + var env workerctl.Envelope + decErr := dec.Decode(&env) + if decErr != nil { + // EOF before a reply line is the tunnel dying mid-verb, and reading + // it as "the install failed" is precisely the collapse this phase + // exists to prevent: nothing was learned about the backend, so this + // is unroutable and a caller may not act on it. + if errors.Is(decErr, io.EOF) { + decErr = fmt.Errorf("the %s stream ended before its reply line: %w", path, io.ErrUnexpectedEOF) + } + return controlFailure(ctx, nodeID, decErr) + } + if env.Reply != nil { + terminal = env.Reply + break + } + if env.Progress == nil || onProgress == nil { + continue + } + var ev messaging.BackendInstallProgressEvent + if err := json.Unmarshal(env.Progress, &ev); err != nil { + // Progress is transient by contract, so a line this frontend cannot + // read costs a tick and never the operation. + xlog.Debug("unreadable control progress line", "node", nodeID, "path", path, "error", err) + continue + } + onProgress(ev) + } + if reply == nil { + return nil + } + if err := json.Unmarshal(terminal, reply); err != nil { + return controlFailure(ctx, nodeID, fmt.Errorf("decoding the %s reply line: %w", path, err)) + } + return nil +} + +// do issues the request and returns the response for any status this frontend +// can read a body from. Every other outcome is already mapped. +// +// The caller owns closing the body. +func (c *ControlClient) do(ctx context.Context, nodeID, path string, req any) (*http.Response, error) { + client, err := c.clientFor(nodeID) + if err != nil { + // ErrNoWorkerDialer already carries ErrWorkerUnroutable, so it must not + // go through controlFailure, which would wrap the umbrella twice. + return nil, err + } + body, err := json.Marshal(req) + if err != nil { + return nil, controlFailure(ctx, nodeID, fmt.Errorf("encoding the %s request: %w", path, err)) + } + + // The host is a name that resolves nowhere. What carries the request is the + // transport's DialContext, which opens a stream on this worker's tunnel; + // the host exists because an http.Request needs one, and it names the node + // so a log line is diagnosable. See WorkerHTTPHost. + url := "http://" + WorkerHTTPHost(nodeID, "") + path + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(body)) + if err != nil { + return nil, controlFailure(ctx, nodeID, fmt.Errorf("building the %s request: %w", path, err)) + } + httpReq.Header.Set("Content-Type", "application/json") + if c.token != "" { + httpReq.Header.Set("Authorization", "Bearer "+c.token) + } + + resp, err := client.Do(httpReq) + if err != nil { + return nil, controlFailure(ctx, nodeID, err) + } + + switch { + case resp.StatusCode == http.StatusOK, resp.StatusCode == http.StatusNoContent: + return resp, nil + case resp.StatusCode == http.StatusNotFound: + // The worker answered, about ITSELF rather than about a backend: it is + // older than this frontend and serves no such verb. The catch-all under + // the control prefix is what makes this distinguishable from a proxy's + // bare 404. + _ = resp.Body.Close() + return nil, fmt.Errorf("control rpc %s to node %q: %w", path, nodeID, ErrWorkerControlUnsupported) + default: + // A verb's own failure arrives as 200 with Error set, so a non-2xx is + // the worker failing to serve the request rather than answering it: a + // body it could not read, a method it refuses, a handler that panicked. + // None of those is evidence about a backend. + detail := readErrorBody(resp) + _ = resp.Body.Close() + return nil, controlFailure(ctx, nodeID, fmt.Errorf("the %s verb answered HTTP %d: %s", path, resp.StatusCode, detail)) + } +} + +// maxControlErrorBodyBytes bounds how much of a non-2xx body reaches a log +// line. The body is written by the worker and lands in this frontend's logs, so +// it is bounded for the same reason the worker bounds the path it echoes. +const maxControlErrorBodyBytes = 512 + +// readErrorBody reads the diagnostic text off a non-2xx control response. +func readErrorBody(resp *http.Response) string { + raw, err := io.ReadAll(io.LimitReader(resp.Body, maxControlErrorBodyBytes)) + if err != nil { + return "" + } + return string(bytes.TrimSpace(raw)) +} + +// controlFailure maps a control RPC's failure onto exactly one of the four +// conditions this system must never confuse, and it is the only place that +// mapping is made. +// +// The rule it enforces, stated as the code enforces it rather than as advice: a +// WORKER'S ANSWER passes through unwrapped so cluster.IsWorkerAnswer still sees +// it and a reap guard may act on it, and EVERYTHING ELSE is wrapped in +// ErrWorkerUnroutable so nothing can. There is no third branch, because a third +// branch is how the eight collapses on this branch happened: each was a place +// that decided for itself which errors were evidence. +// +// The caller's spent budget is checked FIRST, and that ordering is the same one +// cluster.WorkerDialer.handshake takes for the same reason. A worker's refusal +// that arrives in the instant the deadline expires would otherwise be reported +// as the worker's non-transient verdict and reap a row, and nothing orders the +// two timers: an expiry is never evidence about a backend, so it is answered as +// the caller's own timeout rather than as an answer. +func controlFailure(ctx context.Context, nodeID string, err error) error { + if err == nil { + return nil + } + if ctxErr := ctx.Err(); ctxErr != nil { + return fmt.Errorf("control rpc to node %q: %w: %w", nodeID, ErrWorkerUnroutable, ctxErr) + } + if cluster.IsWorkerAnswer(err) { + return err + } + return fmt.Errorf("control rpc to node %q: %w: %w", nodeID, ErrWorkerUnroutable, err) +} diff --git a/core/services/nodes/control_client_test.go b/core/services/nodes/control_client_test.go new file mode 100644 index 000000000..d8abc1197 --- /dev/null +++ b/core/services/nodes/control_client_test.go @@ -0,0 +1,317 @@ +package nodes + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// The whole of this task is the table controlFailure implements: every way a +// control RPC can fail lands in exactly one of four conditions, and none of +// them may be reported as another. The rows, and what each permits: +// +// - the WORKER'S OWN ANSWER, which a reap guard may act on; +// - an UNREACHABLE PEER or a lost route, which nobody may act on; +// - a SPENT BUDGET, which nobody may act on; +// - the worker answering that it is OLDER than this frontend, which only the +// legacy upgrade fallback may act on. +var _ = Describe("ControlClient", func() { + Context("error mapping", func() { + dialerReturning := func(err error) WorkerNetDialerFor { + return func(string) func(context.Context, string, string) (net.Conn, error) { + return func(context.Context, string, string) (net.Conn, error) { return nil, err } + } + } + + It("keeps a worker's own refusal matchable and does NOT wrap it as unroutable", func() { + c := NewControlClient(dialerReturning( + fmt.Errorf("%w: no such process", cluster.ErrStreamTargetUnavailable)), "tok") + err := c.Call(context.Background(), "n1", workerctl.PathModelsRunning, struct{}{}, &struct{}{}) + Expect(cluster.IsWorkerAnswer(err)).To(BeTrue()) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeFalse()) + }) + + DescribeTable("reports every non-answer as unroutable and never as a worker verdict", + func(cause error) { + c := NewControlClient(dialerReturning(cause), "tok") + err := c.Call(context.Background(), "n1", workerctl.PathModelsRunning, struct{}{}, &struct{}{}) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + }, + Entry("no route", fmt.Errorf("reaching node: %w", cluster.ErrNoRoute)), + Entry("no connection recorded", cluster.ErrNoConnection), + Entry("peer unreachable", cluster.ErrPeerUnreachable), + Entry("no relay path", cluster.ErrNoRelayPath), + // The fourth refusal code. The worker DID send it, and it still + // must not be a verdict: it is what a worker says when it learned + // nothing, and those clear on their own. + Entry("the worker learned nothing", fmt.Errorf("%w: late frame", cluster.ErrStreamNotServed)), + Entry("a reply code this frontend does not know", errors.New(`tunnel stream refused with unrecognised code "from-the-future": x`)), + ) + + It("refuses without a dialer rather than reaching for an address", func() { + c := NewControlClient(nil, "tok") + err := c.Call(context.Background(), "n1", workerctl.PathModelsRunning, struct{}{}, &struct{}{}) + Expect(err).To(MatchError(ErrNoWorkerDialer)) + Expect(err).To(MatchError(ErrWorkerUnroutable)) + }) + + It("refuses when the dialer has nothing for that node", func() { + c := NewControlClient(func(string) func(context.Context, string, string) (net.Conn, error) { + return nil + }, "tok") + err := c.Call(context.Background(), "n1", workerctl.PathModelsRunning, struct{}{}, &struct{}{}) + Expect(err).To(MatchError(ErrNoWorkerDialer)) + }) + + It("reports a spent budget as unroutable, so nothing acts on an expiry", func() { + expired, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second)) + defer cancel() + + c := NewControlClient(dialerReturning(errors.New("never reached")), "tok") + err := c.Call(expired, "n1", workerctl.PathModelsRunning, struct{}{}, &struct{}{}) + + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue()) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + }) + + // A timeout is not a verdict, and this is the ordering that enforces + // it. Phase 2's final defect was a worker refusal that arrived in the + // instant the budget expired and was reported as the worker's + // non-transient answer, which reaps a row. + // + // Asserted on the mapping function directly, and deliberately so. + // Nothing orders the two timers, so a spec that drove this through Call + // could only ever produce one of the two orderings by luck: with a + // context already spent, the HTTP client returns before the dialler is + // even asked, and the refusal this is about never happens. Reaching for + // the collapse through the public call was tried and left the suite + // green under the mutation that removes the guard. + DescribeTable("decides between a worker's refusal and a spent budget by the BUDGET first", + func(spent bool, wantAnswer bool) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + if spent { + var stop context.CancelFunc + ctx, stop = context.WithDeadline(ctx, time.Now().Add(-time.Second)) + defer stop() + } + refusal := fmt.Errorf("%w: no such process", cluster.ErrStreamTargetUnavailable) + err := controlFailure(ctx, "n1", refusal) + + Expect(cluster.IsWorkerAnswer(err)).To(Equal(wantAnswer)) + Expect(errors.Is(err, context.DeadlineExceeded)).To(Equal(spent)) + }, + // The negative control: with budget left, the very same refusal is + // the worker speaking and a reap guard may act on it. Without it + // this table would pass on a mapping that never reports an answer. + Entry("budget left: the worker spoke", false, true), + Entry("budget spent: the caller's own timeout", true, false), + ) + }) + + Context("against a worker's HTTP answers", func() { + var ( + srv *httptest.Server + handler http.HandlerFunc + client *ControlClient + ) + + BeforeEach(func() { + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + handler(w, r) + })) + DeferCleanup(srv.Close) + addr := srv.Listener.Addr().String() + client = NewControlClient(func(string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, _, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "tcp", addr) + } + }, "tok") + }) + + It("addresses the node in the URL host and carries the bearer token", func() { + var gotHost, gotAuth, gotPath, gotMethod string + handler = func(w http.ResponseWriter, r *http.Request) { + gotHost, gotAuth, gotPath, gotMethod = r.Host, r.Header.Get("Authorization"), r.URL.Path, r.Method + _, _ = w.Write([]byte(`{}`)) + } + Expect(client.Call(context.Background(), "n1", workerctl.PathModelsRunning, struct{}{}, &struct{}{})).To(Succeed()) + + Expect(gotHost).To(Equal("n1" + unroutableHostSuffix)) + Expect(gotAuth).To(Equal("Bearer tok")) + Expect(gotPath).To(Equal(workerctl.PathModelsRunning)) + // POST, because a control verb is a command: the worker refuses a + // GET so a probe cannot fire one, and a client that sent GET would + // be refused rather than served. + Expect(gotMethod).To(Equal(http.MethodPost)) + }) + + It("hands the worker's own reply back with its Error field intact", func() { + // A verb's own failure is a 200 with Error set. The CALLER reads + // it; this is not an error here, and reporting it as one would put + // the worker's verdict in the bucket reserved for a broken link. + handler = func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(`{"success":false,"error":"disk full"}`)) + } + var reply messaging.BackendDeleteReply + Expect(client.Call(context.Background(), "n1", workerctl.PathBackendDelete, struct{}{}, &reply)).To(Succeed()) + Expect(reply.Success).To(BeFalse()) + Expect(reply.Error).To(Equal("disk full")) + }) + + It("accepts a 204 for the verbs that answer nothing", func() { + handler = func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusNoContent) } + Expect(client.Call(context.Background(), "n1", workerctl.PathNodeStop, struct{}{}, nil)).To(Succeed()) + }) + + It("reports an unknown control verb as unsupported, and not as absence", func() { + handler = func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "unknown worker control path "+r.URL.Path, http.StatusNotFound) + } + err := client.Call(context.Background(), "n1", workerctl.Prefix+"invented", struct{}{}, &struct{}{}) + Expect(err).To(MatchError(ErrWorkerControlUnsupported)) + Expect(errors.Is(err, cluster.ErrNoRoute)).To(BeFalse()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + }) + + DescribeTable("reports a worker that failed to SERVE the request as unroutable, never as a verdict", + func(status int) { + handler = func(w http.ResponseWriter, _ *http.Request) { http.Error(w, "boom", status) } + err := client.Call(context.Background(), "n1", workerctl.PathBackendList, struct{}{}, &struct{}{}) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + Expect(errors.Is(err, ErrWorkerControlUnsupported)).To(BeFalse()) + }, + Entry("the handler failed", http.StatusInternalServerError), + Entry("the body could not be read", http.StatusBadRequest), + Entry("the method was refused", http.StatusMethodNotAllowed), + Entry("a proxy in the path", http.StatusBadGateway), + ) + + It("reports a reply it cannot decode as unroutable, not as an empty answer", func() { + // An empty ModelsRunningReply says "this worker is running nothing", + // which the reconciler acts on. It must never be manufactured from + // a body that would not parse. + handler = func(w http.ResponseWriter, _ *http.Request) { _, _ = w.Write([]byte(`not json`)) } + var reply messaging.ModelsRunningReply + err := client.Call(context.Background(), "n1", workerctl.PathModelsRunning, struct{}{}, &reply) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + }) + + It("streams install progress in order and returns the terminal reply", func() { + handler = func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", workerctl.ContentTypeStream) + enc := json.NewEncoder(w) + _ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":50}`)}) + _ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":100}`)}) + _ = enc.Encode(workerctl.Envelope{Reply: json.RawMessage(`{"success":true}`)}) + } + var seen []float64 + var reply messaging.BackendInstallReply + err := client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall, + messaging.BackendInstallRequest{Backend: "mock", OpID: "op-1"}, &reply, + func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) }) + Expect(err).NotTo(HaveOccurred()) + Expect(seen).To(Equal([]float64{50, 100})) + Expect(reply.Success).To(BeTrue()) + }) + + It("stops reading at the reply line, so nothing after it can be taken for progress", func() { + handler = func(w http.ResponseWriter, _ *http.Request) { + enc := json.NewEncoder(w) + _ = enc.Encode(workerctl.Envelope{Reply: json.RawMessage(`{"success":true}`)}) + _ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":10}`)}) + } + var seen []float64 + var reply messaging.BackendInstallReply + Expect(client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall, + struct{}{}, &reply, + func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) })).To(Succeed()) + Expect(seen).To(BeEmpty()) + Expect(reply.Success).To(BeTrue()) + }) + + It("reports a stream that ends before its reply line as unroutable, not as a failed install", func() { + // A tunnel that dies mid-install must not be read as the worker + // saying the install failed. This is the collapse the phase exists + // to prevent, one layer up from where phase 2 fixed it. + handler = func(w http.ResponseWriter, _ *http.Request) { + enc := json.NewEncoder(w) + _ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":50}`)}) + hijackAndClose(w) + } + var reply messaging.BackendInstallReply + err := client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall, + struct{}{}, &reply, nil) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + Expect(reply.Success).To(BeFalse()) + }) + + It("reports a stream with no lines at all as unroutable", func() { + handler = func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) } + var reply messaging.BackendInstallReply + err := client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall, + struct{}{}, &reply, nil) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue()) + Expect(errors.Is(err, io.ErrUnexpectedEOF)).To(BeTrue()) + }) + + It("keeps going past a progress line it cannot read, since progress is transient", func() { + handler = func(w http.ResponseWriter, _ *http.Request) { + enc := json.NewEncoder(w) + _ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":"not a number"}`)}) + _ = enc.Encode(workerctl.Envelope{Progress: json.RawMessage(`{"percentage":70}`)}) + _ = enc.Encode(workerctl.Envelope{Reply: json.RawMessage(`{"success":true}`)}) + } + var seen []float64 + var reply messaging.BackendInstallReply + Expect(client.CallStreaming(context.Background(), "n1", workerctl.PathBackendInstall, + struct{}{}, &reply, + func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) })).To(Succeed()) + Expect(seen).To(Equal([]float64{70})) + Expect(reply.Success).To(BeTrue()) + }) + + It("reports a streaming 404 as unsupported rather than as a truncated stream", func() { + handler = func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "unknown worker control path "+r.URL.Path, http.StatusNotFound) + } + var reply messaging.BackendUpgradeReply + err := client.CallStreaming(context.Background(), "n1", workerctl.PathBackendUpgrade, struct{}{}, &reply, nil) + Expect(err).To(MatchError(ErrWorkerControlUnsupported)) + }) + }) +}) + +// hijackAndClose ends a response mid-body without the chunked terminator, which +// is what a tunnel dying under an in-flight verb looks like to the reader. +func hijackAndClose(w http.ResponseWriter) { + GinkgoHelper() + if f, ok := w.(http.Flusher); ok { + f.Flush() + } + hj, ok := w.(http.Hijacker) + Expect(ok).To(BeTrue(), "the test server must support hijacking") + conn, buf, err := hj.Hijack() + Expect(err).NotTo(HaveOccurred()) + _ = buf.Flush() + _ = conn.Close() +} diff --git a/core/services/nodes/control_worker_fake_test.go b/core/services/nodes/control_worker_fake_test.go new file mode 100644 index 000000000..013634324 --- /dev/null +++ b/core/services/nodes/control_worker_fake_test.go @@ -0,0 +1,273 @@ +package nodes + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// controlKey names one control verb on one node. +// +// It replaces the NATS subject the scripted double used to key on, and it is +// the same pair the real client addresses: the node picks the tunnel and the +// path picks the verb. +func controlKey(nodeID, path string) string { return nodeID + " " + path } + +// scriptedControlWorkers is a fleet of fake workers reachable only the way a +// real one is: through a per-node dialer, over HTTP, on the control paths. +// +// One HTTP server serves every node, and which node a request is FOR is read +// off the Host header, because that is what ControlClient puts there. That is +// not a convenience: it means a client addressing the wrong node's URL through +// the right node's tunnel would be visible here, which a per-node server could +// not see. +type scriptedControlWorkers struct { + mu sync.Mutex + srv *httptest.Server + + replies map[string][]byte + unsupported map[string]bool + matched map[string][]matchedControlReply + progress map[string][]messaging.BackendInstallProgressEvent + + // unreachable and expired are keyed by NODE, not by verb, because they are + // failures of the ROUTE and a route belongs to a node. They are what the + // dialer answers with; see scriptUnroutable and scriptTimeout. + unreachable map[string]bool + expired map[string]bool + + // hangs names the verbs whose handler never answers, so the only thing + // that ends the call is the caller's own budget. + hangs map[string]bool + + // calls records every request that reached a worker, in order. + calls []requestCall +} + +// matchedControlReply is a canned reply that fires only for a request matching +// pred. It exists so a spec can tell "install with Force=true" (the legacy +// upgrade fallback) from an ordinary install on the same verb. +type matchedControlReply struct { + pred func(messaging.BackendInstallRequest) bool + reply []byte +} + +func newScriptedControlWorkers() *scriptedControlWorkers { + s := &scriptedControlWorkers{ + replies: map[string][]byte{}, + unsupported: map[string]bool{}, + matched: map[string][]matchedControlReply{}, + progress: map[string][]messaging.BackendInstallProgressEvent{}, + unreachable: map[string]bool{}, + expired: map[string]bool{}, + hangs: map[string]bool{}, + } + mux := http.NewServeMux() + mux.HandleFunc(workerctl.Prefix, s.serve) + s.srv = httptest.NewServer(mux) + DeferCleanup(s.srv.Close) + return s +} + +// dialer hands ControlClient the per-node transport it expects. +// +// A node scripted unreachable or expired never reaches the server, which is +// what a real route failure looks like: nothing is asked of the worker and +// nothing is learned about it. +func (s *scriptedControlWorkers) dialer() WorkerNetDialerFor { + return func(nodeID string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, _, _ string) (net.Conn, error) { + s.mu.Lock() + unreachable, expired := s.unreachable[nodeID], s.expired[nodeID] + s.mu.Unlock() + switch { + case expired: + return nil, context.DeadlineExceeded + case unreachable: + return nil, fmt.Errorf("reaching node %q: %w", nodeID, cluster.ErrNoRoute) + } + var d net.Dialer + return d.DialContext(ctx, "tcp", s.srv.Listener.Addr().String()) + } + } +} + +func (s *scriptedControlWorkers) controlClient() *ControlClient { + return NewControlClient(s.dialer(), "test-token") +} + +// nodeOf recovers the node a request was addressed to from its Host header. +func nodeOf(host string) string { + return strings.TrimSuffix(host, unroutableHostSuffix) +} + +func (s *scriptedControlWorkers) serve(w http.ResponseWriter, r *http.Request) { + body, readErr := io.ReadAll(r.Body) + Expect(readErr).ToNot(HaveOccurred()) + key := controlKey(nodeOf(r.Host), r.URL.Path) + + s.mu.Lock() + s.calls = append(s.calls, requestCall{Subject: key, Data: body}) + unsupported := s.unsupported[key] + hang := s.hangs[key] + reply := s.replies[key] + matchers := s.matched[key] + ticks := s.progress[key] + s.mu.Unlock() + + if hang { + // The worker took the request and never answered. Only the caller's + // own budget ends this, which is what makes the budget observable. + <-r.Context().Done() + return + } + if unsupported { + http.Error(w, "unknown worker control path "+r.URL.Path, http.StatusNotFound) + return + } + if len(matchers) > 0 { + var req messaging.BackendInstallRequest + _ = json.Unmarshal(body, &req) + for _, m := range matchers { + if m.pred(req) { + reply = m.reply + break + } + } + } + if reply == nil { + // A verb no spec scripted. Answered LOUDLY rather than plausibly: a + // forgotten script must be a red spec, never a worker that looks absent. + http.Error(w, "this spec scripted no answer for "+key, http.StatusInternalServerError) + return + } + + if r.URL.Path != workerctl.PathBackendInstall && r.URL.Path != workerctl.PathBackendUpgrade { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(reply) + return + } + + // The two streaming verbs: zero or more progress lines, then exactly one + // reply line, always last. + w.Header().Set("Content-Type", workerctl.ContentTypeStream) + w.WriteHeader(http.StatusOK) + enc := json.NewEncoder(w) + for _, ev := range ticks { + raw, err := json.Marshal(ev) + if err != nil { + continue + } + _ = enc.Encode(workerctl.Envelope{Progress: raw}) + } + _ = enc.Encode(workerctl.Envelope{Reply: reply}) +} + +func (s *scriptedControlWorkers) scriptReply(key string, reply any) { + raw, err := json.Marshal(reply) + Expect(err).ToNot(HaveOccurred()) + s.mu.Lock() + defer s.mu.Unlock() + s.replies[key] = raw +} + +// scriptRawReply scripts a reply byte for byte, so a spec can express what a +// worker on a DIFFERENT build would send rather than what today's struct +// marshals to. +func (s *scriptedControlWorkers) scriptRawReply(key string, raw []byte) { + s.mu.Lock() + defer s.mu.Unlock() + s.replies[key] = raw +} + +// scriptUnsupported makes the worker answer 404 for one verb, which is what a +// build older than this frontend does. +func (s *scriptedControlWorkers) scriptUnsupported(key string) { + s.mu.Lock() + defer s.mu.Unlock() + s.unsupported[key] = true +} + +// scriptUnroutable makes every route to a node fail without reaching it. The +// worker is asked nothing, so nothing is learned about it. +func (s *scriptedControlWorkers) scriptUnroutable(nodeID string) { + s.mu.Lock() + defer s.mu.Unlock() + s.unreachable[nodeID] = true +} + +// scriptTimeout makes a node's RPCs end with the caller's budget spent. +// +// Injected at the dial rather than by making a handler sleep, and the +// difference is only in how long the spec takes: a real worker that answers too +// late ends the same way, with the caller's own deadline, because +// cluster.WorkerDialer reports a handshake that outlives the budget as exactly +// this. Answering it instantly keeps the adapter's real install timeout, which +// the retry-scheduling assertions read. +func (s *scriptedControlWorkers) scriptTimeout(nodeID string) { + s.mu.Lock() + defer s.mu.Unlock() + s.expired[nodeID] = true +} + +// clearTimeout lets a node answer again, as a worker that finished a long +// install in the background does. +func (s *scriptedControlWorkers) clearTimeout(nodeID string) { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.expired, nodeID) +} + +func (s *scriptedControlWorkers) scriptReplyMatching(key string, pred func(messaging.BackendInstallRequest) bool, reply messaging.BackendInstallReply) { + raw, err := json.Marshal(reply) + Expect(err).ToNot(HaveOccurred()) + s.mu.Lock() + defer s.mu.Unlock() + s.matched[key] = append(s.matched[key], matchedControlReply{pred: pred, reply: raw}) +} + +// scriptProgress queues the progress lines a streaming verb writes ahead of its +// reply. There is no window to miss them in and nothing to subscribe to: they +// are part of the response the caller is already reading. +func (s *scriptedControlWorkers) scriptProgress(key string, events []messaging.BackendInstallProgressEvent) { + s.mu.Lock() + defer s.mu.Unlock() + s.progress[key] = events +} + +// scriptHang makes one verb on one node accept the request and never answer, +// so the call ends only when the caller's budget does. +// +// It is how a spec observes which budget a verb was given. HTTP carries no +// deadline and the transport dials on a context of its own, so the budget is +// invisible from the far side; how long the client is willing to wait is the +// only thing that shows it. +func (s *scriptedControlWorkers) scriptHang(key string) { + s.mu.Lock() + defer s.mu.Unlock() + s.hangs[key] = true +} + +// callSubjects reports the (node, verb) pairs that reached a worker, in order. +func (s *scriptedControlWorkers) callSubjects() []string { + s.mu.Lock() + defer s.mu.Unlock() + out := make([]string, 0, len(s.calls)) + for _, c := range s.calls { + out = append(out, c.Subject) + } + return out +} diff --git a/core/services/nodes/interfaces.go b/core/services/nodes/interfaces.go index c481113c3..6a0a1367c 100644 --- a/core/services/nodes/interfaces.go +++ b/core/services/nodes/interfaces.go @@ -101,10 +101,17 @@ type NodeHealthStore interface { } // ModelLocator is used by RemoteUnloaderAdapter for model discovery. +// +// Get is here for one reason: backend.stop is the only control verb whose +// carrier depends on what KIND of worker it is addressed to, because agent +// workers hold no tunnel and still take it over the bus. The callers that come +// through NodeCommandSender carry a node id and nothing else, so the type is +// read here rather than threaded through every one of them. type ModelLocator interface { FindNodesWithModel(ctx context.Context, modelName string) ([]BackendNode, error) RemoveNodeModel(ctx context.Context, nodeID, modelName string, replicaIndex int) error RemoveAllNodeModelReplicas(ctx context.Context, nodeID, modelName string) error + Get(ctx context.Context, nodeID string) (*BackendNode, error) } // ModelLookup is used by DistributedModelStore for model existence queries. diff --git a/core/services/nodes/managers_agent_node_test.go b/core/services/nodes/managers_agent_node_test.go index 8ee95c083..5666c093d 100644 --- a/core/services/nodes/managers_agent_node_test.go +++ b/core/services/nodes/managers_agent_node_test.go @@ -11,19 +11,20 @@ import ( "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/workerctl" ) -// Agent workers do not subscribe to the backend.* subjects, so asking one to -// list its backends can only answer "no responders". ListBackends read that as -// a node that had gone away and marked it unhealthy; the node's next heartbeat -// marked it healthy again. Every poll of the backends view therefore flapped -// every agent node in the cluster, and while it was unhealthy the router would -// not schedule onto it. +// Agent workers hold no tunnel and serve no control plane, so asking one to +// list its backends can only fail. ListBackends read that failure as a node +// that had gone away and marked it unhealthy; the node's next heartbeat marked +// it healthy again. Every poll of the backends view therefore flapped every +// agent node in the cluster, and while it was unhealthy the router would not +// schedule onto it. var _ = Describe("Backend listing across mixed node types", func() { var ( db *gorm.DB registry *NodeRegistry - mc *scriptedMessagingClient + mc *scriptedControlWorkers mgr *DistributedBackendManager ctx context.Context ) @@ -36,10 +37,10 @@ var _ = Describe("Backend listing across mixed node types", func() { var err error registry, err = NewNodeRegistry(db) Expect(err).ToNot(HaveOccurred()) - mc = newScriptedMessagingClient() + mc = newScriptedControlWorkers() mgr = &DistributedBackendManager{ local: stubLocalBackendManager{}, - adapter: NewRemoteUnloaderAdapter(nil, mc, 3*time.Minute, 15*time.Minute), + adapter: NewRemoteUnloaderAdapter(nil, nil, mc.controlClient(), 3*time.Minute, 15*time.Minute), registry: registry, } ctx = context.Background() @@ -62,23 +63,43 @@ var _ = Describe("Backend listing across mixed node types", func() { It("leaves an agent node healthy instead of flapping it", func() { agent := register("agent-worker-1", NodeTypeAgent) - mc.scriptNoResponders(messaging.SubjectNodeBackendList(agent.ID)) + mc.scriptUnroutable(agent.ID) _, err := mgr.ListBackends() Expect(err).ToNot(HaveOccurred()) Expect(statusOf(agent.ID)).To(Equal(StatusHealthy), - "an agent node cannot answer backend.list and must not be judged on it") + "an agent node serves no control plane and must not be judged on it") + Expect(mc.callSubjects()).To(BeEmpty(), + "an agent node must not be asked a backend-worker verb at all") }) - It("still marks a backend node unhealthy when it does not answer", func() { + // The direction changed with the carrier, and the change is the safety + // property rather than a regression. "No responders" meant the worker was + // not on the bus; a failed control RPC means THIS frontend could not route + // to it, which is equally what a healthy worker re-homing its tunnel + // between frontend replicas produces. Demoting on that is the fleet-wide + // eviction this phase exists to prevent, so a node that does not answer is + // skipped and left alone. cluster.Presence, read identically on every + // replica from the database, is what may say a worker is gone. + It("leaves a BACKEND node healthy when its control RPC could not be routed", func() { backendNode := register("worker-a", NodeTypeBackend) - mc.scriptNoResponders(messaging.SubjectNodeBackendList(backendNode.ID)) + mc.scriptUnroutable(backendNode.ID) _, err := mgr.ListBackends() Expect(err).ToNot(HaveOccurred()) - Expect(statusOf(backendNode.ID)).To(Equal(StatusUnhealthy), - "a backend worker that does not answer is genuinely gone") + Expect(statusOf(backendNode.ID)).To(Equal(StatusHealthy), + "a route this frontend could not open is not evidence the worker has gone") + }) + + It("still reports the backends of a node that does answer", func() { + backendNode := register("worker-b", NodeTypeBackend) + mc.scriptReply(controlKey(backendNode.ID, workerctl.PathBackendList), + messaging.BackendListReply{Backends: []messaging.NodeBackendInfo{{Name: "vllm"}}}) + + backends, err := mgr.ListBackends() + Expect(err).ToNot(HaveOccurred()) + Expect(backends).To(HaveKey("vllm")) }) }) diff --git a/core/services/nodes/managers_distributed.go b/core/services/nodes/managers_distributed.go index 4132eca79..e3b6ae5b5 100644 --- a/core/services/nodes/managers_distributed.go +++ b/core/services/nodes/managers_distributed.go @@ -14,7 +14,6 @@ import ( "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/system" "github.com/mudler/xlog" - "github.com/nats-io/nats.go" ) // DistributedModelManager wraps a local ModelManager and adds NATS fan-out @@ -213,8 +212,7 @@ func (d *DistributedBackendManager) enqueueAndDrainBackendOp(ctx context.Context continue } - // Record failure for backoff. If it's an ErrNoResponders, the node's - // gone AWOL - mark unhealthy so the router stops picking it too. + // Record failure for backoff. errMsg := applyErr.Error() // Worker-still-installing is a "soft" failure: the worker is most @@ -234,10 +232,15 @@ func (d *DistributedBackendManager) enqueueAndDrainBackendOp(ctx context.Context continue } - if errors.Is(applyErr, nats.ErrNoResponders) { - xlog.Warn("No NATS responders for node, marking unhealthy", "node", node.Name, "nodeID", node.ID) - d.registry.MarkUnhealthy(ctx, node.ID) - } + // A failed control RPC no longer demotes the node, and that is the + // point rather than an omission. It used to, on nats.ErrNoResponders, + // which meant "not on the bus"; the control plane's failures mean "this + // frontend could not route to it", which is equally true of a worker + // that is heartbeating, serving another replica and re-homing its + // tunnel. Demoting on that is the fleet-wide eviction this phase exists + // to prevent. Absence is a fact read from the database, identically on + // every replica, and the scheduler starts reading it in the task that + // replaces this signal with cluster.Presence. if id, err := d.findPendingRow(ctx, node.ID, backend, op); err == nil { _ = d.registry.RecordPendingBackendOpFailure(ctx, id, errMsg) } @@ -331,9 +334,9 @@ func (d *DistributedBackendManager) DeleteBackendDetailed(ctx context.Context, n // populated from the first node seen so single-node-minded callers still work. // // Pending/offline/draining nodes are skipped because they aren't expected to -// answer NATS requests, and so are non-backend workers, which do not subscribe -// to backend.list at all; unhealthy backend nodes are still queried — -// ErrNoResponders then marks them unhealthy and the loop continues. +// answer, and so are non-backend workers, which serve no control plane at all; +// unhealthy backend nodes are still queried, and a node that does not answer is +// skipped rather than demoted. func (d *DistributedBackendManager) ListBackends() (gallery.SystemBackends, error) { result := make(gallery.SystemBackends) allNodes, err := d.registry.List(context.Background()) @@ -345,9 +348,9 @@ func (d *DistributedBackendManager) ListBackends() (gallery.SystemBackends, erro if node.Status == StatusPending || node.Status == StatusOffline || node.Status == StatusDraining { continue } - // Only backend workers subscribe to backend.list. Asking an agent - // worker can only answer "no responders", which the error handling - // below reads as a node that has gone away, so every poll of this view + // Only backend workers serve backend.list. An agent worker holds no + // tunnel, so asking one can only fail, and the failure handling used to + // read that as a node that had gone away: every poll of this view // marked every agent node unhealthy and its next heartbeat marked it // healthy again. The backend-op fan-out skips them for the same reason. if node.NodeType != "" && node.NodeType != NodeTypeBackend { @@ -355,11 +358,9 @@ func (d *DistributedBackendManager) ListBackends() (gallery.SystemBackends, erro } reply, err := d.adapter.ListBackends(node.ID) if err != nil { - if errors.Is(err, nats.ErrNoResponders) { - xlog.Warn("No NATS responders for node, marking unhealthy", "node", node.Name, "nodeID", node.ID) - d.registry.MarkUnhealthy(context.Background(), node.ID) - continue - } + // Skipped, never demoted. Listing a node's backends is a read, and + // a read this frontend could not route says nothing about whether + // the worker is there; see the fan-out above for the same rule. xlog.Warn("Failed to list backends on worker", "node", node.Name, "error", err) continue } @@ -538,7 +539,7 @@ func (d *DistributedBackendManager) InstallBackend(ctx context.Context, op *gall // worker has no platform variant for a linux-only backend) and leaves a // forever-retrying pending_backend_ops row. // -// Rolling-update fallback: when a worker returns nats.ErrNoResponders on +// Rolling-update fallback: when a worker answers that it does not serve // backend.upgrade, we try the legacy backend.install Force=true path so a // new master + old worker still converges. Drop the fallback once every // worker in the fleet is on 2026-05-08 or newer. @@ -600,8 +601,11 @@ func (d *DistributedBackendManager) UpgradeBackend(ctx context.Context, op *gall reply, err := d.adapter.UpgradeBackend(node.ID, name, string(galleriesJSON), "", "", "", 0, opID, onProgressArg) if err != nil { // Rolling-update fallback: an older worker doesn't know - // backend.upgrade. Try the legacy install-with-force path. - if errors.Is(err, nats.ErrNoResponders) { + // backend.upgrade and answers 404 for it. ONLY that answer + // triggers the fallback: a worker this frontend merely could not + // route to has said nothing, and re-firing a force-reinstall at it + // would turn a lost route into a destructive retry. + if errors.Is(err, ErrWorkerControlUnsupported) { instReply, instErr := d.adapter.installWithForceFallback(node.ID, name, string(galleriesJSON), "", "", "", 0, opID, onProgressArg) if instErr != nil { return instErr diff --git a/core/services/nodes/managers_distributed_test.go b/core/services/nodes/managers_distributed_test.go index a72707ea2..62abd3fc9 100644 --- a/core/services/nodes/managers_distributed_test.go +++ b/core/services/nodes/managers_distributed_test.go @@ -2,13 +2,11 @@ package nodes import ( "context" - "encoding/json" "errors" "runtime" "sync" "time" - "github.com/nats-io/nats.go" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "gorm.io/gorm" @@ -18,245 +16,9 @@ import ( "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/workerctl" ) -// scriptedMessagingClient maps a NATS subject to a canned reply payload -// (or error). Used so each fan-out request can simulate a different worker -// outcome without spinning up real NATS. -type scriptedMessagingClient struct { - mu sync.Mutex - replies map[string][]byte - errs map[string]error - calls []requestCall - matchedReplies map[string][]matchedReply - publishes []progressPublishCall - scheduledProgressPublishes []scheduledProgressPublish - subscribes []string -} - -// progressPublishCall records a single Publish invocation. The progress -// publisher tests assert on the sequence of BackendInstallProgressEvent -// values written to a per-op subject, so we capture both subject and the -// decoded event. Named to avoid clashing with the simpler `publishCall` -// already defined in unloader_test.go (which stores raw JSON bytes for -// non-progress assertions). -type progressPublishCall struct { - Subject string - Event messaging.BackendInstallProgressEvent -} - -// scheduledProgressPublish queues a batch of BackendInstallProgressEvent -// values to be delivered the next time Subscribe is called with the matching -// subject. This lets master-side tests assert that the adapter installs its -// handler BEFORE publishing the install request, by scripting events to be -// delivered as soon as the subscription appears. -type scheduledProgressPublish struct { - subject string - events []messaging.BackendInstallProgressEvent -} - -// matchedReply lets a test script a canned reply that only fires when the -// inbound request matches a predicate. Used by scriptReplyMatching to -// distinguish "install Force=true" (the fallback) from "install Force=false" -// on the same subject. -type matchedReply struct { - pred func(messaging.BackendInstallRequest) bool - reply []byte - fallback []byte - fallbackErr error -} - -func newScriptedMessagingClient() *scriptedMessagingClient { - return &scriptedMessagingClient{ - replies: map[string][]byte{}, - errs: map[string]error{}, - } -} - -func (s *scriptedMessagingClient) scriptReply(subject string, reply any) { - raw, err := json.Marshal(reply) - Expect(err).ToNot(HaveOccurred()) - s.mu.Lock() - defer s.mu.Unlock() - s.replies[subject] = raw -} - -func (s *scriptedMessagingClient) scriptErr(subject string, err error) { - s.mu.Lock() - defer s.mu.Unlock() - s.errs[subject] = err -} - -// scriptNoResponders scripts a nats.ErrNoResponders error for `subject` so -// tests can simulate "old worker without backend.upgrade subscription" -// scenarios. Uses the real nats sentinel so errors.Is(...) works at the -// caller (the manager's NoResponders fallback path). -func (s *scriptedMessagingClient) scriptNoResponders(subject string) { - s.mu.Lock() - defer s.mu.Unlock() - s.errs[subject] = nats.ErrNoResponders -} - -// scriptReplyMatching is like scriptReply but the canned reply only fires -// when the inbound request payload matches `pred(req)`. Lets tests -// differentiate "install with Force=true" from "install Force=false" on -// the same subject — useful for asserting the rolling-update fallback -// path actually sets Force=true on its retry. -// -// If `pred` returns false (or the unmarshal of the payload into the -// predicate's expected type fails), the subject falls through to whatever -// was scripted before (or to the unscripted default ErrNoResponders). -func (s *scriptedMessagingClient) scriptReplyMatching(subject string, pred func(messaging.BackendInstallRequest) bool, reply messaging.BackendInstallReply) { - raw, err := json.Marshal(reply) - Expect(err).ToNot(HaveOccurred()) - s.mu.Lock() - defer s.mu.Unlock() - prev := s.replies[subject] // may be nil — that's fine - prevErr := s.errs[subject] // may be nil — that's fine - if s.matchedReplies == nil { - s.matchedReplies = map[string][]matchedReply{} - } - s.matchedReplies[subject] = append(s.matchedReplies[subject], matchedReply{ - pred: pred, - reply: raw, - fallback: prev, - fallbackErr: prevErr, - }) -} - -func (s *scriptedMessagingClient) Request(subject string, data []byte, timeout time.Duration) ([]byte, error) { - s.mu.Lock() - defer s.mu.Unlock() - s.calls = append(s.calls, requestCall{Subject: subject, Data: data, Timeout: timeout}) - - // Predicate-matched replies take precedence over flat scriptReply. - if matchers, ok := s.matchedReplies[subject]; ok { - var req messaging.BackendInstallRequest - _ = json.Unmarshal(data, &req) - for _, m := range matchers { - if m.pred(req) { - return m.reply, nil - } - } - // No predicate matched — fall through to the recorded fallback - // (whatever was scripted before scriptReplyMatching took over). - if matchers[0].fallback != nil { - return matchers[0].fallback, nil - } - if matchers[0].fallbackErr != nil { - return nil, matchers[0].fallbackErr - } - // No fallback either — default to ErrNoResponders. - return nil, nats.ErrNoResponders - } - - if err, ok := s.errs[subject]; ok && err != nil { - return nil, err - } - if reply, ok := s.replies[subject]; ok { - return reply, nil - } - // Simulate ErrNoResponders for any unscripted subject so tests fail - // loudly when they forget to script a node. - return nil, &fakeNoRespondersErr{} -} - -// Publish records each call so progress-publisher tests can assert on the -// stream of events written to a subject. The real messaging.Client JSON -// encodes the payload before sending, but our publisher hands a typed -// struct directly, so we handle both shapes. -func (s *scriptedMessagingClient) Publish(subject string, data any) error { - s.mu.Lock() - defer s.mu.Unlock() - switch ev := data.(type) { - case messaging.BackendInstallProgressEvent: - s.publishes = append(s.publishes, progressPublishCall{Subject: subject, Event: ev}) - case []byte: - var e messaging.BackendInstallProgressEvent - _ = json.Unmarshal(ev, &e) - s.publishes = append(s.publishes, progressPublishCall{Subject: subject, Event: e}) - } - return nil -} - -// publishCalls returns every BackendInstallProgressEvent that was published -// to `subject`, in order. Lets tests assert on debounce behavior without -// depending on internal Publish timing. -func (s *scriptedMessagingClient) publishCalls(subject string) []messaging.BackendInstallProgressEvent { - s.mu.Lock() - defer s.mu.Unlock() - out := make([]messaging.BackendInstallProgressEvent, 0) - for _, c := range s.publishes { - if c.Subject != subject { - continue - } - out = append(out, c.Event) - } - return out -} - -// scheduleProgressPublish queues a set of BackendInstallProgressEvent values -// to be delivered on the next Subscribe call matching the per-op progress -// subject. A short delay before delivery gives the subscriber time to install -// its message handler before the events arrive. -func (s *scriptedMessagingClient) scheduleProgressPublish(nodeID, opID string, events []messaging.BackendInstallProgressEvent) { - s.mu.Lock() - defer s.mu.Unlock() - s.scheduledProgressPublishes = append(s.scheduledProgressPublishes, scheduledProgressPublish{ - subject: messaging.SubjectNodeBackendInstallProgress(nodeID, opID), - events: events, - }) -} - -// subscribeCalls returns the subjects on which Subscribe was invoked. -// Used to confirm the master skipped subscription when onProgress was nil. -func (s *scriptedMessagingClient) subscribeCalls() []string { - s.mu.Lock() - defer s.mu.Unlock() - out := make([]string, len(s.subscribes)) - copy(out, s.subscribes) - return out -} - -func (s *scriptedMessagingClient) Subscribe(subject string, handler func([]byte)) (messaging.Subscription, error) { - s.mu.Lock() - s.subscribes = append(s.subscribes, subject) - matched := []scheduledProgressPublish{} - remaining := s.scheduledProgressPublishes[:0] - for _, sp := range s.scheduledProgressPublishes { - if sp.subject == subject { - matched = append(matched, sp) - } else { - remaining = append(remaining, sp) - } - } - s.scheduledProgressPublishes = remaining - s.mu.Unlock() - - go func() { - time.Sleep(20 * time.Millisecond) - for _, sp := range matched { - for _, ev := range sp.events { - raw, _ := json.Marshal(ev) - handler(raw) - } - } - }() - - return &fakeSubscription{}, nil -} -func (s *scriptedMessagingClient) QueueSubscribe(_ string, _ string, _ func([]byte)) (messaging.Subscription, error) { - return &fakeSubscription{}, nil -} -func (s *scriptedMessagingClient) QueueSubscribeReply(_ string, _ string, _ func([]byte, func([]byte))) (messaging.Subscription, error) { - return &fakeSubscription{}, nil -} -func (s *scriptedMessagingClient) SubscribeReply(_ string, _ func([]byte, func([]byte))) (messaging.Subscription, error) { - return &fakeSubscription{}, nil -} -func (s *scriptedMessagingClient) IsConnected() bool { return true } -func (s *scriptedMessagingClient) Close() {} - // recordingNodeCall captures a single UpdateNodeProgress invocation so // per-node OpStatus tests can assert on the sequence of writes the // DistributedBackendManager fans out into the sink. @@ -292,16 +54,6 @@ func (r *recordingProgressSink) callsFor(opID, nodeID string) []galleryop.NodePr return out } -// fakeNoRespondersErr is the unscripted-subject default. It matches -// nats.ErrNoResponders by string only - used when a test forgets to script -// a node so the failure is loud but doesn't tickle errors.Is(...) sentinel -// paths the test wasn't deliberately exercising. Tests that DO want the -// real sentinel (e.g. to drive the manager's NoResponders fallback) call -// scriptNoResponders instead, which scripts nats.ErrNoResponders directly. -type fakeNoRespondersErr struct{} - -func (e *fakeNoRespondersErr) Error() string { return "no responders" } - // stubLocalBackendManager satisfies galleryop.BackendManager for the // distributed manager's `local` field. The DeleteBackend path expects to // call into local first; in distributed mode the frontend rarely has @@ -329,7 +81,7 @@ var _ = Describe("DistributedBackendManager", func() { var ( db *gorm.DB registry *NodeRegistry - mc *scriptedMessagingClient + mc *scriptedControlWorkers adapter *RemoteUnloaderAdapter mgr *DistributedBackendManager ctx context.Context @@ -344,8 +96,8 @@ var _ = Describe("DistributedBackendManager", func() { registry, err = NewNodeRegistry(db) Expect(err).ToNot(HaveOccurred()) - mc = newScriptedMessagingClient() - adapter = NewRemoteUnloaderAdapter(nil, mc, 3*time.Minute, 15*time.Minute) + mc = newScriptedControlWorkers() + adapter = NewRemoteUnloaderAdapter(nil, nil, mc.controlClient(), 3*time.Minute, 15*time.Minute) mgr = &DistributedBackendManager{ local: stubLocalBackendManager{}, adapter: adapter, @@ -385,9 +137,9 @@ var _ = Describe("DistributedBackendManager", func() { n1 := registerHealthyBackend("worker-a", "10.0.0.1:50051") n2 := registerHealthyBackend("worker-b", "10.0.0.2:50051") - mc.scriptReply(messaging.SubjectNodeBackendInstall(n1.ID), + mc.scriptReply(controlKey(n1.ID, workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "10.0.0.1:50100"}) - mc.scriptReply(messaging.SubjectNodeBackendInstall(n2.ID), + mc.scriptReply(controlKey(n2.ID, workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "10.0.0.2:50100"}) Expect(mgr.InstallBackend(ctx, op("vllm-development"), nil)).To(Succeed()) @@ -399,9 +151,9 @@ var _ = Describe("DistributedBackendManager", func() { n1 := registerHealthyBackend("dgx-casa", "10.0.0.1:50051") n2 := registerHealthyBackend("nvidia-thor", "10.0.0.2:50051") - mc.scriptReply(messaging.SubjectNodeBackendInstall(n1.ID), + mc.scriptReply(controlKey(n1.ID, workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: false, Error: "no child with platform linux/arm64 in index quay.io/...master-cpu-vllm"}) - mc.scriptReply(messaging.SubjectNodeBackendInstall(n2.ID), + mc.scriptReply(controlKey(n2.ID, workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: false, Error: "disk full"}) err := mgr.InstallBackend(ctx, op("vllm-development"), nil) @@ -419,9 +171,9 @@ var _ = Describe("DistributedBackendManager", func() { ok := registerHealthyBackend("worker-ok", "10.0.0.1:50051") bad := registerHealthyBackend("worker-bad", "10.0.0.2:50051") - mc.scriptReply(messaging.SubjectNodeBackendInstall(ok.ID), + mc.scriptReply(controlKey(ok.ID, workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "10.0.0.1:50100"}) - mc.scriptReply(messaging.SubjectNodeBackendInstall(bad.ID), + mc.scriptReply(controlKey(bad.ID, workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: false, Error: "out of memory"}) err := mgr.InstallBackend(ctx, op("vllm-development"), nil) @@ -437,8 +189,8 @@ var _ = Describe("DistributedBackendManager", func() { registerUnhealthyBackend("worker-a", "10.0.0.1:50051") registerUnhealthyBackend("worker-b", "10.0.0.2:50051") - // No replies scripted: if the manager tried to call Request, - // it would hit the "no responders" default and we'd see it. + // No replies scripted: if the manager issued a control RPC at + // all, the unscripted-verb default answers 500 and we'd see it. Expect(mgr.InstallBackend(ctx, op("vllm-development"), nil)).To(Succeed()) mc.mu.Lock() calls := len(mc.calls) @@ -458,11 +210,10 @@ var _ = Describe("DistributedBackendManager", func() { target := registerHealthyBackend("worker-target", "10.0.0.1:50051") other := registerHealthyBackend("worker-other", "10.0.0.2:50051") - mc.scriptReply(messaging.SubjectNodeBackendInstall(target.ID), + mc.scriptReply(controlKey(target.ID, workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "10.0.0.1:50100"}) - // No reply scripted for `other`: if InstallBackend fans out - // to it, the fakeNoRespondersErr default would surface and - // the test would fail. + // No reply scripted for `other`: if InstallBackend fans out to + // it, the unscripted-verb default surfaces and the test fails. targetedOp := &galleryop.ManagementOp[gallery.GalleryBackend, any]{ GalleryElementName: "llama-cpp", @@ -473,13 +224,13 @@ var _ = Describe("DistributedBackendManager", func() { mc.mu.Lock() defer mc.mu.Unlock() Expect(mc.calls).To(HaveLen(1)) - Expect(mc.calls[0].Subject).To(Equal(messaging.SubjectNodeBackendInstall(target.ID))) - Expect(mc.calls[0].Subject).ToNot(Equal(messaging.SubjectNodeBackendInstall(other.ID))) + Expect(mc.calls[0].Subject).To(Equal(controlKey(target.ID, workerctl.PathBackendInstall))) + Expect(mc.calls[0].Subject).ToNot(Equal(controlKey(other.ID, workerctl.PathBackendInstall))) }) }) Context("when op.TargetNodeID is set to a node that does not exist", func() { - It("returns nil without sending any NATS request", func() { + It("returns nil without sending any control request", func() { registerHealthyBackend("worker-a", "10.0.0.1:50051") ghostOp := &galleryop.ManagementOp[gallery.GalleryBackend, any]{ @@ -498,10 +249,10 @@ var _ = Describe("DistributedBackendManager", func() { It("returns galleryop.ErrWorkerStillInstalling and keeps the queue row with NextRetryAt pushed out", func() { n := registerHealthyBackend("slow", "10.0.0.1:50051") - // Script a NATS timeout on the install subject. The adapter - // wraps this into galleryop.ErrWorkerStillInstalling, which - // the manager should treat as a soft failure. - mc.scriptErr(messaging.SubjectNodeBackendInstall(n.ID), nats.ErrTimeout) + // The install RPC ends with the caller's budget spent. The + // adapter wraps that into galleryop.ErrWorkerStillInstalling, + // which the manager treats as a soft failure. + mc.scriptTimeout(n.ID) err := mgr.InstallBackend(ctx, op("vllm"), nil) Expect(err).To(HaveOccurred()) @@ -513,7 +264,8 @@ var _ = Describe("DistributedBackendManager", func() { Expect(rows).To(HaveLen(1)) Expect(rows[0].Backend).To(Equal("vllm")) // The adapter is configured with a 3m install timeout in this - // suite (NewRemoteUnloaderAdapter above). NextRetryAt should + // suite (NewRemoteUnloaderAdapter above), and the RPC failed + // instantly rather than waiting it out. NextRetryAt should // be ~now+3m; a > now+2m bound is safe-but-tight enough to // catch the buggy short default (30s exponential backoff). Expect(rows[0].NextRetryAt).To(BeTemporally(">", time.Now().Add(2*time.Minute)), @@ -521,17 +273,17 @@ var _ = Describe("DistributedBackendManager", func() { }) }) - Context("end-to-end: timeout then successful reconcile via backend.list", func() { + Context("end-to-end: budget spent, then successful reconcile via backend.list", func() { It("surfaces the install in ListBackends after the worker finishes", func() { // Use the same node-registration helper the Task 5 test uses // so the test fixture is identical to the prior context. node := registerHealthyBackend("jetson", "10.0.0.2:50051") - // First install attempt: NATS times out. The adapter wraps + // First install attempt: the budget runs out. The adapter wraps // this as galleryop.ErrWorkerStillInstalling and the manager // keeps the pending_backend_ops row alive with NextRetryAt // pushed out (asserted in the previous context). - mc.scriptErr(messaging.SubjectNodeBackendInstall(node.ID), nats.ErrTimeout) + mc.scriptTimeout(node.ID) err := mgr.InstallBackend(ctx, op("vllm"), nil) Expect(err).To(HaveOccurred()) @@ -542,10 +294,11 @@ var _ = Describe("DistributedBackendManager", func() { Expect(listErr).ToNot(HaveOccurred()) Expect(rows).To(HaveLen(1)) - // The worker finished installing in the background. Script - // backend.list on the same scriptedMessagingClient so the - // manager's ListBackends fan-out reports the backend. - mc.scriptReply(messaging.SubjectNodeBackendList(node.ID), messaging.BackendListReply{ + // The worker finished installing in the background and answers + // again, so the manager's ListBackends fan-out reports the + // backend. + mc.clearTimeout(node.ID) + mc.scriptReply(controlKey(node.ID, workerctl.PathBackendList), messaging.BackendListReply{ Backends: []messaging.NodeBackendInfo{{Name: "vllm"}}, }) @@ -569,9 +322,9 @@ var _ = Describe("DistributedBackendManager", func() { It("deletes the pending_backend_ops install row when the backend is reported installed on its target node", func() { node := registerHealthyBackend("worker-a", "10.0.0.5:50051") - // Pre-stage: simulate an admin install that timed out at the NATS - // round-trip, leaving an install row in the queue. - mc.scriptErr(messaging.SubjectNodeBackendInstall(node.ID), nats.ErrTimeout) + // Pre-stage: simulate an admin install whose control RPC ran out + // of budget, leaving an install row in the queue. + mc.scriptTimeout(node.ID) err := mgr.InstallBackend(ctx, op("vllm"), nil) Expect(err).To(HaveOccurred()) Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeTrue()) @@ -581,7 +334,8 @@ var _ = Describe("DistributedBackendManager", func() { // Worker finishes installing in the background. backend.list now // confirms presence; ListBackends should proactively clear the row. - mc.scriptReply(messaging.SubjectNodeBackendList(node.ID), messaging.BackendListReply{ + mc.clearTimeout(node.ID) + mc.scriptReply(controlKey(node.ID, workerctl.PathBackendList), messaging.BackendListReply{ Backends: []messaging.NodeBackendInfo{{Name: "vllm"}}, }) @@ -599,7 +353,7 @@ var _ = Describe("DistributedBackendManager", func() { Expect(registry.UpsertPendingBackendOp(ctx, node.ID, "vllm", OpBackendUpgrade, []byte("[]"))).To(Succeed()) - mc.scriptReply(messaging.SubjectNodeBackendList(node.ID), messaging.BackendListReply{ + mc.scriptReply(controlKey(node.ID, workerctl.PathBackendList), messaging.BackendListReply{ Backends: []messaging.NodeBackendInfo{{Name: "vllm"}}, }) @@ -615,8 +369,8 @@ var _ = Describe("DistributedBackendManager", func() { It("invokes progressCb once per worker-published progress event", func() { node := registerHealthyBackend("worker-prog", "10.0.0.7:50051") - mc.scriptReply(messaging.SubjectNodeBackendInstall(node.ID), messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "10.0.0.7:50051"}) - mc.scheduleProgressPublish(node.ID, "op-prog-1", []messaging.BackendInstallProgressEvent{ + mc.scriptReply(controlKey(node.ID, workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "10.0.0.7:50051"}) + mc.scriptProgress(controlKey(node.ID, workerctl.PathBackendInstall), []messaging.BackendInstallProgressEvent{ {OpID: "op-prog-1", NodeID: node.ID, Backend: "vllm", FileName: "vllm.tar", Current: "100 MB", Total: "1 GB", Percentage: 10}, {OpID: "op-prog-1", NodeID: node.ID, Backend: "vllm", FileName: "vllm.tar", Current: "1 GB", Total: "1 GB", Percentage: 100}, }) @@ -646,21 +400,20 @@ var _ = Describe("DistributedBackendManager", func() { }, "1s").Should(Equal(2)) mu.Lock() defer mu.Unlock() - // The adapter dispatches each progress event to its own goroutine - // (see unloader.go: `go onProgress(ev)`) so two events emitted back - // to back can land at the bridge in either order. Assert the set of - // percentages observed contains both ticks, rather than depending - // on goroutine scheduling for ordering. - pcts := []float64{pcCalls[0].Percentage, pcCalls[1].Percentage} - Expect(pcts).To(ConsistOf(10.0, 100.0)) + // ORDER, not a set. Progress lines are read off the install + // response on the caller's own goroutine, so the order the + // worker wrote them in is the order the bridge sees; the + // goroutine-per-event dispatch that made this best-effort is + // gone with the subscription it existed for. + Expect([]float64{pcCalls[0].Percentage, pcCalls[1].Percentage}).To(Equal([]float64{10.0, 100.0})) }) }) Context("InstallBackend tolerates silent (pre-Phase-2) workers", func() { It("completes successfully even when no progress events are ever published", func() { node := registerHealthyBackend("worker-silent", "10.0.0.8:50051") - mc.scriptReply(messaging.SubjectNodeBackendInstall(node.ID), messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "10.0.0.8:50051"}) - // NO scheduleProgressPublish call - silent worker. + mc.scriptReply(controlKey(node.ID, workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "10.0.0.8:50051"}) + // NO scriptProgress call - silent worker. var ticks int var mu sync.Mutex @@ -695,13 +448,13 @@ var _ = Describe("DistributedBackendManager", func() { mgr = NewDistributedBackendManager(appCfg, nil, adapter, registry, sink) // stubLocalBackendManager mirrors the production behaviour // where the frontend node rarely has the backend installed - // locally - the NATS fan-out is what these specs verify. + // locally - the control-plane fan-out is what these specs verify. mgr.local = stubLocalBackendManager{} }) It("emits a success entry for each healthy node visited", func() { node := registerHealthyBackend("worker-ok", "10.0.0.9:50051") - mc.scriptReply(messaging.SubjectNodeBackendInstall(node.ID), + mc.scriptReply(controlKey(node.ID, workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "10.0.0.9:50051"}) opVal := op("vllm") @@ -714,9 +467,9 @@ var _ = Describe("DistributedBackendManager", func() { Expect(calls[len(calls)-1].NodeName).To(Equal("worker-ok")) }) - It("emits a running_on_worker entry when NATS times out", func() { + It("emits a running_on_worker entry when the install RPC runs out of budget", func() { node := registerHealthyBackend("worker-slow", "10.0.0.10:50051") - mc.scriptErr(messaging.SubjectNodeBackendInstall(node.ID), nats.ErrTimeout) + mc.scriptTimeout(node.ID) opVal := op("vllm") opVal.ID = "op-node-slow" @@ -730,9 +483,9 @@ var _ = Describe("DistributedBackendManager", func() { It("emits downloading entries from progress events", func() { node := registerHealthyBackend("worker-dl", "10.0.0.11:50051") - mc.scriptReply(messaging.SubjectNodeBackendInstall(node.ID), + mc.scriptReply(controlKey(node.ID, workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: true}) - mc.scheduleProgressPublish(node.ID, "op-node-dl", []messaging.BackendInstallProgressEvent{ + mc.scriptProgress(controlKey(node.ID, workerctl.PathBackendInstall), []messaging.BackendInstallProgressEvent{ {OpID: "op-node-dl", NodeID: node.ID, Backend: "vllm", FileName: "vllm.tar", Current: "1 GB", Total: "1 GB", Percentage: 100, Phase: messaging.PhaseDownloading}, }) @@ -765,13 +518,13 @@ var _ = Describe("DistributedBackendManager", func() { // upgrade should skip it. scriptInstalled := func(backend string, nodeIDs ...string) { for _, id := range nodeIDs { - mc.scriptReply(messaging.SubjectNodeBackendList(id), + mc.scriptReply(controlKey(id, workerctl.PathBackendList), messaging.BackendListReply{Backends: []messaging.NodeBackendInfo{{Name: backend}}}) } } scriptNoBackends := func(nodeIDs ...string) { for _, id := range nodeIDs { - mc.scriptReply(messaging.SubjectNodeBackendList(id), + mc.scriptReply(controlKey(id, workerctl.PathBackendList), messaging.BackendListReply{Backends: nil}) } } @@ -782,9 +535,9 @@ var _ = Describe("DistributedBackendManager", func() { n2 := registerHealthyBackend("worker-b", "10.0.0.2:50051") scriptInstalled("vllm-development", n1.ID, n2.ID) - mc.scriptReply(messaging.SubjectNodeBackendUpgrade(n1.ID), + mc.scriptReply(controlKey(n1.ID, workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: false, Error: "image manifest not found"}) - mc.scriptReply(messaging.SubjectNodeBackendUpgrade(n2.ID), + mc.scriptReply(controlKey(n2.ID, workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: false, Error: "registry unauthorized"}) err := mgr.UpgradeBackend(ctx, upgradeOp("vllm-development"), nil) @@ -800,7 +553,7 @@ var _ = Describe("DistributedBackendManager", func() { It("returns nil", func() { n1 := registerHealthyBackend("worker-a", "10.0.0.1:50051") scriptInstalled("vllm-development", n1.ID) - mc.scriptReply(messaging.SubjectNodeBackendUpgrade(n1.ID), + mc.scriptReply(controlKey(n1.ID, workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: true}) Expect(mgr.UpgradeBackend(ctx, upgradeOp("vllm-development"), nil)).To(Succeed()) }) @@ -818,7 +571,7 @@ var _ = Describe("DistributedBackendManager", func() { scriptInstalled("cpu-insightface-development", has.ID) scriptNoBackends(lacks.ID) - mc.scriptReply(messaging.SubjectNodeBackendUpgrade(has.ID), + mc.scriptReply(controlKey(has.ID, workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: true}) // Deliberately don't script SubjectNodeBackendUpgrade for `lacks`: // if the manager attempts it, the scripted-client default returns @@ -829,7 +582,7 @@ var _ = Describe("DistributedBackendManager", func() { mc.mu.Lock() defer mc.mu.Unlock() for _, call := range mc.calls { - Expect(call.Subject).ToNot(Equal(messaging.SubjectNodeBackendUpgrade(lacks.ID)), + Expect(call.Subject).ToNot(Equal(controlKey(lacks.ID, workerctl.PathBackendUpgrade)), "upgrade leaked to %s which does not have the backend installed", lacks.Name) } }) @@ -846,9 +599,9 @@ var _ = Describe("DistributedBackendManager", func() { n2 := registerHealthyBackend("worker-b", "10.0.0.2:50051") scriptInstalled("vllm-development", n1.ID, n2.ID) - mc.scriptReply(messaging.SubjectNodeBackendUpgrade(n1.ID), + mc.scriptReply(controlKey(n1.ID, workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: true}) - mc.scriptReply(messaging.SubjectNodeBackendUpgrade(n2.ID), + mc.scriptReply(controlKey(n2.ID, workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: true}) op := upgradeOp("vllm-development") @@ -859,10 +612,10 @@ var _ = Describe("DistributedBackendManager", func() { defer mc.mu.Unlock() upgraded := map[string]bool{} for _, call := range mc.calls { - if call.Subject == messaging.SubjectNodeBackendUpgrade(n1.ID) { + if call.Subject == controlKey(n1.ID, workerctl.PathBackendUpgrade) { upgraded[n1.ID] = true } - if call.Subject == messaging.SubjectNodeBackendUpgrade(n2.ID) { + if call.Subject == controlKey(n2.ID, workerctl.PathBackendUpgrade) { upgraded[n2.ID] = true } } @@ -876,7 +629,7 @@ var _ = Describe("DistributedBackendManager", func() { scriptInstalled("vllm-development", has.ID) scriptNoBackends(lacks.ID) - mc.scriptReply(messaging.SubjectNodeBackendUpgrade(has.ID), + mc.scriptReply(controlKey(has.ID, workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: true}) op := upgradeOp("vllm-development") @@ -888,9 +641,9 @@ var _ = Describe("DistributedBackendManager", func() { mc.mu.Lock() defer mc.mu.Unlock() for _, call := range mc.calls { - Expect(call.Subject).ToNot(Equal(messaging.SubjectNodeBackendUpgrade(has.ID)), + Expect(call.Subject).ToNot(Equal(controlKey(has.ID, workerctl.PathBackendUpgrade)), "a node-scoped upgrade for %s must not touch other nodes", lacks.Name) - Expect(call.Subject).ToNot(Equal(messaging.SubjectNodeBackendUpgrade(lacks.ID)), + Expect(call.Subject).ToNot(Equal(controlKey(lacks.ID, workerctl.PathBackendUpgrade)), "the target node lacks the backend; nothing should be sent") } }) @@ -908,37 +661,57 @@ var _ = Describe("DistributedBackendManager", func() { mc.mu.Lock() defer mc.mu.Unlock() for _, call := range mc.calls { - Expect(call.Subject).ToNot(Equal(messaging.SubjectNodeBackendUpgrade(n1.ID))) - Expect(call.Subject).ToNot(Equal(messaging.SubjectNodeBackendInstall(n1.ID))) + Expect(call.Subject).ToNot(Equal(controlKey(n1.ID, workerctl.PathBackendUpgrade))) + Expect(call.Subject).ToNot(Equal(controlKey(n1.ID, workerctl.PathBackendInstall))) } }) }) - // Rolling-update fallback: pre-2026-05-08 workers don't subscribe to - // backend.upgrade, so the manager catches nats.ErrNoResponders and + // Rolling-update fallback: pre-2026-05-08 workers do not serve + // backend.upgrade, so the manager catches the worker's own 404 and // re-fires the legacy backend.install Force=true on the same node. // Drop these specs once the fallback path itself is removed (see // managers_distributed.go UpgradeBackend godoc for the deprecation). Context("rolling-update fallback", func() { - It("falls back to backend.install Force=true when upgrade returns ErrNoResponders", func() { + It("falls back to backend.install Force=true when the worker does not serve the upgrade verb", func() { n := registerHealthyBackend("worker-old", "10.0.0.1:50051") scriptInstalled("vllm-development", n.ID) - // Old worker: no subscriber on backend.upgrade. - mc.scriptNoResponders(messaging.SubjectNodeBackendUpgrade(n.ID)) + // Old worker: it answers 404 for a verb it does not serve. + mc.scriptUnsupported(controlKey(n.ID, workerctl.PathBackendUpgrade)) // Fallback re-fires legacy backend.install with Force=true. - mc.scriptReplyMatching(messaging.SubjectNodeBackendInstall(n.ID), + mc.scriptReplyMatching(controlKey(n.ID, workerctl.PathBackendInstall), func(req messaging.BackendInstallRequest) bool { return req.Force }, messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "10.0.0.1:50100"}) Expect(mgr.UpgradeBackend(ctx, upgradeOp("vllm-development"), nil)).To(Succeed()) }) - It("returns the upgrade error when it is not ErrNoResponders", func() { + // The negative direction, and it is the one that matters: the + // fallback re-fires a DESTRUCTIVE force-reinstall, so it may run + // only on the worker's own "I do not serve that verb". A worker + // this frontend merely failed to reach has said nothing, and + // retrying a force-reinstall on silence is how a lost route becomes + // a reinstall of every backend on the fleet. + It("does NOT fall back to the legacy force install when the upgrade could not be routed", func() { + n := registerHealthyBackend("worker-unroutable", "10.0.0.1:50051") + scriptInstalled("vllm-development", n.ID) + // backend.upgrade is deliberately not scripted, so the worker + // fails to SERVE it rather than answering that it lacks it. + // backend.install is not scripted either, so a fallback would + // be visible in the calls below. + + err := mgr.UpgradeBackend(ctx, upgradeOp("vllm-development"), nil) + Expect(err).To(HaveOccurred()) + Expect(mc.callSubjects()).ToNot(ContainElement(controlKey(n.ID, workerctl.PathBackendInstall)), + "an unroutable upgrade must not re-fire a force-reinstall") + }) + + It("returns the upgrade error when the worker served the verb and refused", func() { n := registerHealthyBackend("worker-bad", "10.0.0.1:50051") scriptInstalled("vllm-development", n.ID) - mc.scriptReply(messaging.SubjectNodeBackendUpgrade(n.ID), + mc.scriptReply(controlKey(n.ID, workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: false, Error: "disk full"}) err := mgr.UpgradeBackend(ctx, upgradeOp("vllm-development"), nil) @@ -954,9 +727,9 @@ var _ = Describe("DistributedBackendManager", func() { n1 := registerHealthyBackend("worker-a", "10.0.0.1:50051") n2 := registerHealthyBackend("worker-b", "10.0.0.2:50051") - mc.scriptReply(messaging.SubjectNodeBackendDelete(n1.ID), + mc.scriptReply(controlKey(n1.ID, workerctl.PathBackendDelete), messaging.BackendDeleteReply{Success: false, Error: "backend not installed"}) - mc.scriptReply(messaging.SubjectNodeBackendDelete(n2.ID), + mc.scriptReply(controlKey(n2.ID, workerctl.PathBackendDelete), messaging.BackendDeleteReply{Success: false, Error: "permission denied"}) err := mgr.DeleteBackend("vllm-development") @@ -971,7 +744,7 @@ var _ = Describe("DistributedBackendManager", func() { Context("when every node succeeds", func() { It("returns nil", func() { n1 := registerHealthyBackend("worker-a", "10.0.0.1:50051") - mc.scriptReply(messaging.SubjectNodeBackendDelete(n1.ID), + mc.scriptReply(controlKey(n1.ID, workerctl.PathBackendDelete), messaging.BackendDeleteReply{Success: true}) Expect(mgr.DeleteBackend("vllm-development")).To(Succeed()) }) diff --git a/core/services/nodes/reconciler.go b/core/services/nodes/reconciler.go index 44c3d875f..ed231d93e 100644 --- a/core/services/nodes/reconciler.go +++ b/core/services/nodes/reconciler.go @@ -12,7 +12,6 @@ import ( "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/xlog" - "github.com/nats-io/nats.go" "google.golang.org/grpc/codes" "google.golang.org/grpc/status" "gorm.io/gorm" @@ -394,14 +393,17 @@ func (rc *ReplicaReconciler) drainPendingBackendOps(ctx context.Context) { // Pending-op drain for admin upgrade — fires backend.upgrade so // the slow re-pull doesn't head-of-line-block install traffic on // the same worker. Falls back to the legacy backend.install - // Force=true path on nats.ErrNoResponders for old workers that - // don't subscribe to backend.upgrade yet (rolling-update window). - // Reconciler retries are background reconciliation with no live - // admin watching a progress bar, so opID/onProgress are empty — - // the adapter skips the progress subscription entirely. + // Force=true path when the worker answers that it does not serve + // backend.upgrade (rolling-update window). Reconciler retries are + // background reconciliation with no live admin watching a progress + // bar, so opID/onProgress are empty and no progress is streamed. reply, err := rc.adapter.UpgradeBackend(op.NodeID, op.Backend, string(op.Galleries), "", "", "", 0, "", nil) if err != nil { - if errors.Is(err, nats.ErrNoResponders) { + // Only the worker's own "I do not serve that verb" may + // re-fire a force-reinstall. An unroutable worker has said + // nothing, and retrying a destructive verb on silence is how a + // lost route becomes a reinstall. + if errors.Is(err, ErrWorkerControlUnsupported) { instReply, instErr := rc.adapter.installWithForceFallback(op.NodeID, op.Backend, string(op.Galleries), "", "", "", 0, "", nil) if instErr != nil { applyErr = instErr @@ -429,17 +431,14 @@ func (rc *ReplicaReconciler) drainPendingBackendOps(ctx context.Context) { continue } - // ErrNoResponders means the node has no active NATS subscription for - // this subject. Either its connection dropped, or it's the wrong - // node type entirely. Mark unhealthy so the health monitor's - // heartbeat-only pass doesn't immediately flip it back — and so - // ListDuePendingBackendOps (which filters by status=healthy) stops - // picking the row until the node genuinely recovers. - if errors.Is(applyErr, nats.ErrNoResponders) { - xlog.Warn("Reconciler: no NATS responders — marking node unhealthy", - "op", op.Op, "backend", op.Backend, "node", op.NodeID) - _ = rc.registry.MarkUnhealthy(ctx, op.NodeID) - } + // A failed op no longer demotes the node. It used to, on + // nats.ErrNoResponders, which said the worker was not on the bus; a + // control RPC fails when THIS frontend cannot route to the worker, + // which a worker re-homing its tunnel does while it is heartbeating and + // serving. Demoting on it would take the node out of + // ListDuePendingBackendOps and out of scheduling for a reason that has + // nothing to do with the node. The row keeps its backoff and its + // dead-letter cap, and cluster.Presence becomes the absence signal. // Dead-letter cap: after maxAttempts the row is the reconciler // equivalent of a poison message. Delete it loudly so the queue diff --git a/core/services/nodes/reconciler_test.go b/core/services/nodes/reconciler_test.go index b91cacc29..075fee25f 100644 --- a/core/services/nodes/reconciler_test.go +++ b/core/services/nodes/reconciler_test.go @@ -9,8 +9,10 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" + "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/workerctl" "gorm.io/gorm" ) @@ -872,6 +874,95 @@ var _ = Describe("ReplicaReconciler — state reconciliation", func() { }) }) + // The pending-op drain runs the SAME rolling-update fallback the admin + // path does, and that fallback re-fires a DESTRUCTIVE force-reinstall. + // Only the worker's own "I do not serve that verb" may trigger it. A + // worker this frontend could not reach has said nothing about the verb, + // and a BACKGROUND drain that force-reinstalls on silence is worse than + // the admin path doing it: nobody is watching, it retries on every tick, + // and a frontend replica that has just lost its tunnels would reinstall + // every queued backend on the fleet. + // + // The admin path's negative direction is pinned in + // managers_distributed_test.go; this one is the second call site of the + // same rule and was unpinned, so the condition here could be widened to + // any error with the whole suite still green. + Describe("draining a pending backend upgrade", func() { + var ( + workers *scriptedControlWorkers + node *BackendNode + rc *ReplicaReconciler + ) + + BeforeEach(func() { + workers = newScriptedControlWorkers() + node = &BackendNode{Name: "worker-drain", NodeType: NodeTypeBackend, Address: "10.0.0.9:50051"} + Expect(registry.Register(context.Background(), node, true)).To(Succeed()) + Expect(registry.UpsertPendingBackendOp(context.Background(), node.ID, "vllm", OpBackendUpgrade, []byte("[]"))).To(Succeed()) + rc = NewReplicaReconciler(ReplicaReconcilerOptions{ + Registry: registry, + Scheduler: &fakeScheduler{}, + DB: db, + Adapter: NewRemoteUnloaderAdapter(registry, nil, workers.controlClient(), time.Minute, time.Minute), + }) + }) + + queuedOps := func() []PendingBackendOp { + var rows []PendingBackendOp + Expect(db.Find(&rows).Error).To(Succeed()) + return rows + } + + It("falls back to the legacy force install when the worker does not serve the upgrade verb", func() { + workers.scriptUnsupported(controlKey(node.ID, workerctl.PathBackendUpgrade)) + workers.scriptReplyMatching(controlKey(node.ID, workerctl.PathBackendInstall), + func(req messaging.BackendInstallRequest) bool { return req.Force }, + messaging.BackendInstallReply{Success: true}) + + rc.drainPendingBackendOps(context.Background()) + + Expect(workers.callSubjects()).To(Equal([]string{ + controlKey(node.ID, workerctl.PathBackendUpgrade), + controlKey(node.ID, workerctl.PathBackendInstall), + })) + Expect(queuedOps()).To(BeEmpty(), "an op the fallback converged is drained") + }) + + // The negative direction, arranged so it cannot pass vacuously: the + // force install IS scripted and IS reachable here, so a fallback that + // fired would show up as a call AND as a drained row. What the worker + // fails to do is SERVE the upgrade verb, which is a 5xx and not the + // 404 that means it lacks it. + It("does NOT fall back to the legacy force install when the worker failed to serve the upgrade", func() { + workers.scriptReplyMatching(controlKey(node.ID, workerctl.PathBackendInstall), + func(req messaging.BackendInstallRequest) bool { return req.Force }, + messaging.BackendInstallReply{Success: true}) + + rc.drainPendingBackendOps(context.Background()) + + Expect(workers.callSubjects()).ToNot(ContainElement(controlKey(node.ID, workerctl.PathBackendInstall)), + "an upgrade the worker failed to serve must not re-fire a force-reinstall") + rows := queuedOps() + Expect(rows).To(HaveLen(1), "the op keeps its backoff and is retried, not converged") + Expect(rows[0].Attempts).To(Equal(1)) + }) + + // The same rule for the other non-answer: no route at all. The install + // is unreachable too here, so the witness is the error the drain + // RECORDED, which names the verb that actually failed. + It("records the upgrade's own failure when the worker could not be routed to", func() { + workers.scriptUnroutable(node.ID) + + rc.drainPendingBackendOps(context.Background()) + + rows := queuedOps() + Expect(rows).To(HaveLen(1)) + Expect(rows[0].LastError).To(ContainSubstring(workerctl.PathBackendUpgrade)) + Expect(rows[0].LastError).ToNot(ContainSubstring(workerctl.PathBackendInstall), + "a fallback that fired would have replaced the upgrade's error with the install's") + }) + }) + Describe("NewNodeRegistry malformed-row pruning", func() { It("drops queue rows for agent nodes and non-existent nodes on startup", func() { agent := &BackendNode{Name: "agent-1", NodeType: NodeTypeAgent, Address: "x"} diff --git a/core/services/nodes/router_liveness.go b/core/services/nodes/router_liveness.go index 88646162f..9809d1b2d 100644 --- a/core/services/nodes/router_liveness.go +++ b/core/services/nodes/router_liveness.go @@ -14,14 +14,25 @@ import ( // whole fleet. const maxNodeLivenessRetries = 3 -// nodeAnswersOnBus reports whether a node still has a live subscription. +// nodeAnswersOnBus reports whether a node may still be given work. // -// Only nats.ErrNoResponders means "absent". Any other outcome, a timeout or a -// transport hiccup, leaves the node eligible: wrongly excluding a node that is -// merely slow costs real capacity, while the install that follows already -// reports its own failure. When no command sender is configured there is no bus -// to consult and every node is treated as reachable, which preserves the -// behaviour of deployments that do not run one. +// It probes the node's CONTROL PLANE, over that node's tunnel, and nothing it +// can learn there is proof of absence. A worker that answers is present by +// demonstration; every other outcome is this frontend failing to route, which +// is equally what a healthy worker re-homing its tunnel between frontend +// replicas produces, so excluding on it would demote a node that is +// heartbeating and serving. Only nats.ErrNoResponders excludes, and no control +// RPC produces it: the exclusion is therefore inert here, and stays only until +// the scheduler reads cluster.Presence, which is the one fact in the deployment +// that can say a worker is gone and says it identically on every replica. +// +// Before that carrier change this asked two NATS subjects the worker no longer +// subscribes to, so EVERY healthy worker answered "no responders" and was +// marked unhealthy on the scheduling path. +// +// When no command sender is configured there is nothing to consult and every +// node is treated as reachable, which preserves the behaviour of deployments +// that do not run one. func (r *SmartRouter) nodeAnswersOnBus(node *BackendNode) bool { if r.unloader == nil || node == nil { return true diff --git a/core/services/nodes/unloader.go b/core/services/nodes/unloader.go index caaebe5f1..542e7d8bd 100644 --- a/core/services/nodes/unloader.go +++ b/core/services/nodes/unloader.go @@ -2,22 +2,22 @@ package nodes import ( "context" - "encoding/json" "errors" "fmt" - "strings" "time" "github.com/nats-io/nats.go" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/xlog" ) -// NodeCommandSender abstracts NATS-based commands to worker nodes. -// Used by HTTP endpoint handlers to avoid coupling to the concrete RemoteUnloaderAdapter. +// NodeCommandSender abstracts the control commands a frontend issues to a +// worker node. They travel over the worker's tunnel as HTTP under +// workerctl.Prefix; see RemoteUnloaderAdapter. // // InstallBackend is idempotent: the worker short-circuits if the backend is // already running for the requested (modelID, replica) slot. Routine model @@ -27,8 +27,8 @@ import ( // every live process for the backend, re-pulls the gallery artifact, and // replies. Caller (DistributedBackendManager.UpgradeBackend) handles // rolling-update fallback to the legacy install Force=true path on -// nats.ErrNoResponders for old workers that don't subscribe to the new -// backend.upgrade subject. +// ErrWorkerControlUnsupported, which is what a worker older than the +// backend.upgrade verb answers. type NodeCommandSender interface { InstallBackend(nodeID, backendType, modelID, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(messaging.BackendInstallProgressEvent)) (*messaging.BackendInstallReply, error) UpgradeBackend(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(messaging.BackendInstallProgressEvent)) (*messaging.BackendUpgradeReply, error) @@ -36,34 +36,46 @@ type NodeCommandSender interface { ListBackends(nodeID string) (*messaging.BackendListReply, error) StopBackend(nodeID, backend string) error UnloadModelOnNode(nodeID, modelName string) error - // PingNode reports whether the node is still subscribed on the bus. It - // returns nats.ErrNoResponders when nothing answers for the node, which is - // the only condition callers may read as "this node cannot be given work". + // PingNode reports whether the worker answered a control RPC. Only an + // answer is evidence; every other outcome is this frontend failing to + // route, which callers must not read as the node having gone. Task 6 + // replaces it with a presence read on the database clock, which is the + // only thing in this deployment that CAN say a worker is gone. PingNode(nodeID string) error } -// RemoteUnloaderAdapter implements NodeCommandSender and model.RemoteModelUnloader -// by publishing NATS events for backend process lifecycle. The worker process -// subscribes and handles the actual process start/stop. +// RemoteUnloaderAdapter implements NodeCommandSender and +// model.RemoteModelUnloader by issuing control RPCs to the worker over its +// tunnel. The worker serves them on the loopback HTTP server it already runs +// (see core/services/worker/control_routes.go) and handles the actual process +// start/stop. // -// This mirrors the local ModelLoader's startProcess()/deleteProcess() but -// over NATS for remote nodes. +// This mirrors the local ModelLoader's startProcess()/deleteProcess() but for +// remote nodes. +// +// One verb is still carried by the bus, and only for one KIND of node: +// backend.stop to an AGENT node. Agent workers hold no tunnel yet, so they have +// nothing to serve a control route on, and they subscribe to +// nodes..backend.stop to drop cached MCP sessions. Removing that publish +// would strand them; see stopBackend. type RemoteUnloaderAdapter struct { registry ModelLocator nats messaging.MessagingClient + control *ControlClient installTimeout time.Duration upgradeTimeout time.Duration } -// NewRemoteUnloaderAdapter creates a new adapter. installTimeout and -// upgradeTimeout govern the NATS request-reply deadlines for backend.install -// and backend.upgrade respectively. Use -// DistributedConfig.BackendInstallTimeoutOrDefault() / +// NewRemoteUnloaderAdapter creates a new adapter. control carries every verb +// except backend.stop to an agent node, which stays on nats. installTimeout and +// upgradeTimeout bound the backend.install and backend.upgrade RPCs +// respectively; use DistributedConfig.BackendInstallTimeoutOrDefault() / // BackendUpgradeTimeoutOrDefault() at construction. -func NewRemoteUnloaderAdapter(registry ModelLocator, nats messaging.MessagingClient, installTimeout, upgradeTimeout time.Duration) *RemoteUnloaderAdapter { +func NewRemoteUnloaderAdapter(registry ModelLocator, nats messaging.MessagingClient, control *ControlClient, installTimeout, upgradeTimeout time.Duration) *RemoteUnloaderAdapter { return &RemoteUnloaderAdapter{ registry: registry, nats: nats, + control: control, installTimeout: installTimeout, upgradeTimeout: upgradeTimeout, } @@ -93,6 +105,13 @@ const exactModelStopTimeout = 10 * time.Second // StopModelReplica stops only the process represented by replica. Configuration // cleanup intentionally has no backend.stop fallback: an old worker that does // not understand this request leaves the quarantine row for a later retry. +// +// The caller's context is carried into the RPC rather than run alongside it on +// a goroutine, which is what the request/reply carrier needed because it took a +// timeout and not a context. Abandoning the request is safe: the worker's +// model.stop handler deliberately drops the caller's context and runs the stop +// to completion, so a frontend that gives up cannot leave a half-stopped +// process or an unreturned port behind. func (a *RemoteUnloaderAdapter) StopModelReplica(ctx context.Context, nodeID string, replica NodeModel, force bool) (messaging.ModelStopReply, error) { if ctx == nil { ctx = context.Background() @@ -100,31 +119,18 @@ func (a *RemoteUnloaderAdapter) StopModelReplica(ctx context.Context, nodeID str ctx, cancel := context.WithTimeout(ctx, exactModelStopTimeout) defer cancel() - type result struct { - reply *messaging.ModelStopReply - err error - } - done := make(chan result, 1) - go func() { - reply, err := messaging.RequestJSON[messaging.ModelStopRequest, messaging.ModelStopReply](a.nats, messaging.SubjectNodeModelStop(nodeID), messaging.ModelStopRequest{ - ModelName: replica.ModelName, - ProcessKey: model.BackendProcessKey(replica.ModelName, replica.ReplicaIndex), - ExpectedAddress: replica.WorkerLocalAddress, - Force: force, - ConfigRevision: replica.ConfigRevision, - }, exactModelStopTimeout) - done <- result{reply: reply, err: err} - }() - - select { - case <-ctx.Done(): - return messaging.ModelStopReply{}, ctx.Err() - case result := <-done: - if result.err != nil { - return messaging.ModelStopReply{}, result.err - } - return *result.reply, nil + var reply messaging.ModelStopReply + err := a.control.Call(ctx, nodeID, workerctl.PathModelStop, messaging.ModelStopRequest{ + ModelName: replica.ModelName, + ProcessKey: model.BackendProcessKey(replica.ModelName, replica.ReplicaIndex), + ExpectedAddress: replica.WorkerLocalAddress, + Force: force, + ConfigRevision: replica.ConfigRevision, + }, &reply) + if err != nil { + return messaging.ModelStopReply{}, err } + return reply, nil } // UnloadRemoteModel finds the node(s) hosting the given model and tells them @@ -173,8 +179,8 @@ func (a *RemoteUnloaderAdapter) UnloadRemoteModelContext(ctx context.Context, mo var unloadErr error for _, node := range nodes { - 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.Info("Sending backend.stop to node", "model", modelName, "node", node.Name, "nodeID", node.ID, "force", force) + if err := a.stopBackend(ctx, node.ID, node.NodeType, 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 @@ -189,7 +195,7 @@ func (a *RemoteUnloaderAdapter) UnloadRemoteModelContext(ctx context.Context, mo return unloadErr } -// InstallBackend sends a backend.install request-reply to a worker node. +// InstallBackend asks a worker node to install a backend and start its process. // Idempotent on the worker: if the (modelID, replica) process is already // running, the worker short-circuits and returns its address; if the binary // is on disk, the worker just spawns a process; only a missing binary @@ -201,23 +207,24 @@ func (a *RemoteUnloaderAdapter) UnloadRemoteModelContext(ctx context.Context, mo // case on slow links (Jetson Wi-Fi, multi-GB CUDA images) while still // failing fast enough to surface real worker hangs. // -// For force-reinstall (admin-driven Upgrade), use UpgradeBackend instead - -// it lives on a different NATS subject so it cannot head-of-line-block -// routine load traffic on the same worker. +// Progress needs no subscription and no window to miss events in: the worker +// writes its download ticks into THIS response ahead of the terminal reply, so +// there is nothing to arrange before the request is sent. +// +// For force-reinstall (admin-driven Upgrade), use UpgradeBackend instead. func (a *RemoteUnloaderAdapter) InstallBackend( nodeID, backendType, modelID, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(messaging.BackendInstallProgressEvent), ) (*messaging.BackendInstallReply, error) { - subject := messaging.SubjectNodeBackendInstall(nodeID) - xlog.Info("Sending NATS backend.install", "nodeID", nodeID, "backend", backendType, "modelID", modelID, "replica", replicaIndex, "opID", opID) + xlog.Info("Sending backend.install", "nodeID", nodeID, "backend", backendType, "modelID", modelID, "replica", replicaIndex, "opID", opID) - // Subscribe to the per-op progress subject BEFORE publishing the install - // request so we don't miss early events. - sub := a.subscribeProgress(nodeID, opID, onProgress) + ctx, cancel := context.WithTimeout(context.Background(), a.installTimeout) + defer cancel() - reply, err := messaging.RequestJSON[messaging.BackendInstallRequest, messaging.BackendInstallReply](a.nats, subject, messaging.BackendInstallRequest{ + var reply messaging.BackendInstallReply + err := a.control.CallStreaming(ctx, nodeID, workerctl.PathBackendInstall, messaging.BackendInstallRequest{ Backend: backendType, ModelID: modelID, BackendGalleries: galleriesJSON, @@ -226,75 +233,34 @@ func (a *RemoteUnloaderAdapter) InstallBackend( Alias: alias, ReplicaIndex: int32(replicaIndex), OpID: opID, - }, a.installTimeout) - - if sub != nil { - if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil { - xlog.Warn("Failed to unsubscribe from backend install progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr) + }, &reply, onProgress) + if err != nil { + if isRequestTimeout(err) { + return nil, fmt.Errorf("%w (nodeID=%s backend=%s): %v", + galleryop.ErrWorkerStillInstalling, nodeID, backendType, err) } + return nil, err } - - if err != nil && isNATSTimeout(err) { - return nil, fmt.Errorf("%w (subject=%s nodeID=%s backend=%s): %v", - galleryop.ErrWorkerStillInstalling, subject, nodeID, backendType, err) - } - return reply, err + return &reply, nil } -// subscribeProgress subscribes to the per-op backend-install progress subject -// so the master can stream per-node download ticks while a worker installs or -// upgrades. Returns nil (and subscribes to nothing) when onProgress is nil or -// opID is empty — the reconciler-driven retry path and legacy callers stay -// silent at no cost. Shared by InstallBackend, UpgradeBackend, and the legacy -// force-install fallback: an upgrade is a force-reinstall, so it reuses the -// install-progress subject rather than minting a new one (no new NATS -// permission, no new rolling-update compat surface). Caller must Unsubscribe -// the returned subscription after the request completes. -func (a *RemoteUnloaderAdapter) subscribeProgress(nodeID, opID string, onProgress func(messaging.BackendInstallProgressEvent)) messaging.Subscription { - if onProgress == nil || opID == "" { - return nil - } - progressSubject := messaging.SubjectNodeBackendInstallProgress(nodeID, opID) - s, subErr := a.nats.Subscribe(progressSubject, func(raw []byte) { - var ev messaging.BackendInstallProgressEvent - if err := json.Unmarshal(raw, &ev); err != nil { - xlog.Debug("malformed backend progress event", "subject", progressSubject, "error", err) - return - } - // Goroutine guard: a slow onProgress callback must not stall the NATS - // reader thread. Events spawn one goroutine each, so ordering at the - // consumer is best-effort; the worker debounces to ~250ms which dwarfs - // goroutine scheduling jitter, and its final Flush() is the terminal tick. - go onProgress(ev) - }) - if subErr != nil { - xlog.Warn("Failed to subscribe to backend progress subject; proceeding without progress streaming", - "subject", progressSubject, "error", subErr) - return nil - } - return s -} - -// UpgradeBackend sends a backend.upgrade request-reply to a worker node. +// UpgradeBackend asks a worker node to force-reinstall a backend. // The worker stops every live process for this backend, force-reinstalls // from the gallery (overwriting the on-disk artifact), and replies. The // next routine InstallBackend call spawns a fresh process with the new // binary - upgrade itself does not start a process. // -// When opID is non-empty and onProgress is set, the master subscribes to the -// per-op progress subject before firing the request so a long force-reinstall -// streams per-node download ticks instead of blocking opaque at progress 0. -// // Timeout: configured via DistributedConfig.BackendUpgradeTimeoutOrDefault // (default 15m). Real-world worst case observed: 8-10 minutes for large // CUDA-l4t backend images on Jetson over WiFi. func (a *RemoteUnloaderAdapter) UpgradeBackend(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(messaging.BackendInstallProgressEvent)) (*messaging.BackendUpgradeReply, error) { - subject := messaging.SubjectNodeBackendUpgrade(nodeID) - xlog.Info("Sending NATS backend.upgrade", "nodeID", nodeID, "backend", backendType, "replica", replicaIndex, "opID", opID) + xlog.Info("Sending backend.upgrade", "nodeID", nodeID, "backend", backendType, "replica", replicaIndex, "opID", opID) - sub := a.subscribeProgress(nodeID, opID, onProgress) + ctx, cancel := context.WithTimeout(context.Background(), a.upgradeTimeout) + defer cancel() - reply, err := messaging.RequestJSON[messaging.BackendUpgradeRequest, messaging.BackendUpgradeReply](a.nats, subject, messaging.BackendUpgradeRequest{ + var reply messaging.BackendUpgradeReply + err := a.control.CallStreaming(ctx, nodeID, workerctl.PathBackendUpgrade, messaging.BackendUpgradeRequest{ Backend: backendType, BackendGalleries: galleriesJSON, URI: uri, @@ -302,37 +268,31 @@ func (a *RemoteUnloaderAdapter) UpgradeBackend(nodeID, backendType, galleriesJSO Alias: alias, ReplicaIndex: int32(replicaIndex), OpID: opID, - }, a.upgradeTimeout) - - if sub != nil { - if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil { - xlog.Warn("Failed to unsubscribe from backend upgrade progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr) + }, &reply, onProgress) + if err != nil { + if isRequestTimeout(err) { + return nil, fmt.Errorf("%w (nodeID=%s backend=%s): %v", + galleryop.ErrWorkerStillInstalling, nodeID, backendType, err) } + return nil, err } - - if err != nil && isNATSTimeout(err) { - 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 + a.dropStoppedReplicaRows(nodeID, "backend.upgrade", backendType, reply.StoppedProcessKeys, reply.ReportsStoppedProcesses) + return &reply, nil } // installWithForceFallback is the rolling-update fallback used by -// DistributedBackendManager.UpgradeBackend when backend.upgrade returns -// nats.ErrNoResponders (the worker is on a pre-2026-05-08 build that -// doesn't subscribe to the new subject). It re-fires the legacy -// backend.install with Force=true. Drop this once every worker is on -// 2026-05-08 or newer. +// DistributedBackendManager.UpgradeBackend when backend.upgrade reports that +// the worker does not serve that verb (a pre-2026-05-08 build). It re-fires +// the legacy backend.install with Force=true. Drop this once every worker is +// on 2026-05-08 or newer. func (a *RemoteUnloaderAdapter) installWithForceFallback(nodeID, backendType, galleriesJSON, uri, name, alias string, replicaIndex int, opID string, onProgress func(messaging.BackendInstallProgressEvent)) (*messaging.BackendInstallReply, error) { - subject := messaging.SubjectNodeBackendInstall(nodeID) xlog.Warn("Falling back to legacy backend.install Force=true (old worker)", "nodeID", nodeID, "backend", backendType) - sub := a.subscribeProgress(nodeID, opID, onProgress) + ctx, cancel := context.WithTimeout(context.Background(), a.upgradeTimeout) + defer cancel() - reply, err := messaging.RequestJSON[messaging.BackendInstallRequest, messaging.BackendInstallReply](a.nats, subject, messaging.BackendInstallRequest{ + var reply messaging.BackendInstallReply + err := a.control.CallStreaming(ctx, nodeID, workerctl.PathBackendInstall, messaging.BackendInstallRequest{ Backend: backendType, BackendGalleries: galleriesJSON, URI: uri, @@ -341,108 +301,164 @@ func (a *RemoteUnloaderAdapter) installWithForceFallback(nodeID, backendType, ga ReplicaIndex: int32(replicaIndex), Force: true, OpID: opID, - }, a.upgradeTimeout) - - if sub != nil { - if unsubscribeErr := sub.Unsubscribe(); unsubscribeErr != nil { - xlog.Warn("Failed to unsubscribe from legacy backend install progress", "nodeID", nodeID, "backend", backendType, "opID", opID, "error", unsubscribeErr) + }, &reply, onProgress) + if err != nil { + if isRequestTimeout(err) { + return nil, fmt.Errorf("%w (nodeID=%s backend=%s): %v", + galleryop.ErrWorkerStillInstalling, nodeID, backendType, err) } + return nil, err } - - if err != nil && isNATSTimeout(err) { - return nil, fmt.Errorf("%w (subject=%s nodeID=%s backend=%s): %v", - galleryop.ErrWorkerStillInstalling, subject, nodeID, backendType, err) - } - return reply, err + return &reply, nil } -// ListBackends queries a worker node for its installed backends via NATS request-reply. +// Control-RPC budgets. Each is the deadline the corresponding NATS +// request/reply carried, kept unchanged so this cutover changes the carrier and +// not how long the frontend waits. +const ( + backendListTimeout = 30 * time.Second + modelsRunningTimeout = 10 * time.Second + backendStopTimeout = 30 * time.Second + backendDeleteTimeout = 2 * time.Minute + modelUnloadTimeout = 30 * time.Second + modelDeleteTimeout = 30 * time.Second + nodeStopTimeout = 30 * time.Second + nodePingTimeout = 5 * time.Second +) + +// ListBackends queries a worker node for its installed backends. func (a *RemoteUnloaderAdapter) ListBackends(nodeID string) (*messaging.BackendListReply, error) { - subject := messaging.SubjectNodeBackendList(nodeID) - xlog.Debug("Sending NATS backend.list", "nodeID", nodeID) + xlog.Debug("Sending backend.list", "nodeID", nodeID) - return messaging.RequestJSON[messaging.BackendListRequest, messaging.BackendListReply](a.nats, subject, messaging.BackendListRequest{}, 30*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), backendListTimeout) + defer cancel() + + var reply messaging.BackendListReply + if err := a.control.Call(ctx, nodeID, workerctl.PathBackendList, messaging.BackendListRequest{}, &reply); err != nil { + return nil, err + } + return &reply, nil } -// PingNode checks that a worker still has a live subscription on the bus. +// PingNode asks a worker one cheap control verb and reports what came back. // // A node's status in the database comes from its HTTP heartbeat, which is a -// separate channel from NATS. A worker that has died stops answering on NATS -// at once but keeps its healthy status until the heartbeat ages out, so the -// scheduler could pick a node that could not be given work and the request -// failed with "no responders available". +// separate channel, so a worker that has died keeps its healthy status until +// the heartbeat ages out and the scheduler could commit to a node it cannot +// reach. // -// The subject asked has to be one every worker in the fleet subscribes to, or -// this check condemns the workers that do not. models.running was the obvious -// choice and the wrong one: it arrived in 4.6, so a 4.5 worker that is alive -// and serving never answers it, and a model pinned to that node could never be -// scheduled. backend.list has been part of the worker protocol far longer, so -// it is the safer question to ask. +// What this can and cannot conclude changed with the carrier, and the +// difference is the whole point. A worker that ANSWERS is present by +// demonstration, whatever it answered. Every other outcome is this frontend +// failing to route to it, which is true of a worker that is heartbeating and +// serving another replica's requests while re-homing its tunnel; reported as +// absence it would have the scheduler demote a healthy node. So the error is +// ErrWorkerUnroutable, callers must not read it as the node having gone, and +// nothing here returns an absence sentinel at all. Task 6 replaces this with +// cluster.Registry.Presence, which is the one fact in the deployment that can +// say a worker is gone, and says it identically on every replica. // -// A worker that answers anything is alive. Only when every subject reports no -// responders is the node treated as absent, so adding a newer subject here can -// never condemn an older worker. +// backend.list rather than models.running for the reason it always was: it has +// been part of the worker protocol far longer, so it is the safer question to +// ask of a worker that may be older than this frontend. func (a *RemoteUnloaderAdapter) PingNode(nodeID string) error { - subjects := []string{ - messaging.SubjectNodeBackendList(nodeID), - messaging.SubjectNodeModelsRunning(nodeID), + ctx, cancel := context.WithTimeout(context.Background(), nodePingTimeout) + defer cancel() + + var reply messaging.BackendListReply + err := a.control.Call(ctx, nodeID, workerctl.PathBackendList, messaging.BackendListRequest{}, &reply) + if errors.Is(err, ErrWorkerControlUnsupported) { + // The worker answered, from a build that predates this verb. That is a + // worker that is very much there. + return nil } - var lastErr error - for _, subject := range subjects { - _, err := messaging.RequestJSON[messaging.BackendListRequest, messaging.BackendListReply]( - a.nats, subject, messaging.BackendListRequest{}, 5*time.Second) - if err == nil { - return nil - } - if !errors.Is(err, nats.ErrNoResponders) { - // Reached someone, or failed for a reason that is not absence. - // Either way the node is not proven gone. - return nil - } - lastErr = err - } - return lastErr + return err } // ListRunningModels asks a worker node which model backend processes it -// currently has running, via NATS request-reply. +// currently has running. // // The timeout is short on purpose: the worker answers straight out of its // in-memory process table, so a slow reply means the worker itself is in // trouble, and the caller treats no-answer as "don't know" rather than as // "nothing running". func (a *RemoteUnloaderAdapter) ListRunningModels(nodeID string) (*messaging.ModelsRunningReply, error) { - subject := messaging.SubjectNodeModelsRunning(nodeID) - return messaging.RequestJSON[messaging.ModelsRunningRequest, messaging.ModelsRunningReply]( - a.nats, subject, messaging.ModelsRunningRequest{}, 10*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), modelsRunningTimeout) + defer cancel() + + var reply messaging.ModelsRunningReply + if err := a.control.Call(ctx, nodeID, workerctl.PathModelsRunning, messaging.ModelsRunningRequest{}, &reply); err != nil { + return nil, err + } + return &reply, nil } // StopBackend tells a worker node to stop a specific gRPC backend process. // 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) + ctx, cancel := context.WithTimeout(context.Background(), backendStopTimeout) + defer cancel() + return a.stopBackend(ctx, nodeID, a.nodeTypeOf(ctx, nodeID), backend, false) } -func (a *RemoteUnloaderAdapter) stopBackend(nodeID, backend string, force bool) error { - subject := messaging.SubjectNodeBackendStop(nodeID) - if backend == "" && !force { - return a.nats.Publish(subject, nil) +// nodeTypeOf answers which KIND of worker a node id names, so stopBackend can +// pick the carrier that node actually listens on. +// +// A lookup that fails answers NodeTypeBackend, matching the column's own +// default and the "empty means backend" reading every other node-type branch in +// this package takes. The failure directions are not symmetric: sending a +// backend node's stop over the bus is silently lost, because nothing subscribes +// to it any more, while sending an agent node's stop over the tunnel returns an +// error the caller sees. +func (a *RemoteUnloaderAdapter) nodeTypeOf(ctx context.Context, nodeID string) string { + if a.registry == nil { + return NodeTypeBackend } - return a.nats.Publish(subject, messaging.BackendStopRequest{Backend: backend, Force: force}) + node, err := a.registry.Get(ctx, nodeID) + if err != nil || node == nil { + xlog.Debug("Could not resolve node type for a backend stop; assuming a backend worker", + "nodeID", nodeID, "error", err) + return NodeTypeBackend + } + return node.NodeType +} + +// stopBackend sends one backend.stop, over the carrier that kind of worker +// listens on. +// +// An AGENT node keeps the bus. It holds no tunnel, so it has no control route +// to serve, and it subscribes to nodes..backend.stop to drop the MCP +// sessions cached for a backend that is going away. This is the one verb of the +// ten that is split rather than moved, and the split is the honest intermediate +// state until agent workers hold tunnels too. +func (a *RemoteUnloaderAdapter) stopBackend(ctx context.Context, nodeID, nodeType, backend string, force bool) error { + if nodeType == NodeTypeAgent { + subject := messaging.SubjectNodeBackendStop(nodeID) + if backend == "" && !force { + return a.nats.Publish(subject, nil) + } + return a.nats.Publish(subject, messaging.BackendStopRequest{Backend: backend, Force: force}) + } + // An empty Backend is what the worker reads as "stop everything", the same + // meaning the bus carried as an empty payload; see decodeBackendStopRequest. + return a.control.Call(ctx, nodeID, workerctl.PathBackendStop, + messaging.BackendStopRequest{Backend: backend, Force: force}, nil) } // DeleteBackend tells a worker node to delete a backend (stop + remove files). func (a *RemoteUnloaderAdapter) DeleteBackend(nodeID, backendName string) (*messaging.BackendDeleteReply, error) { - subject := messaging.SubjectNodeBackendDelete(nodeID) - xlog.Info("Sending NATS backend.delete", "nodeID", nodeID, "backend", backendName) + xlog.Info("Sending backend.delete", "nodeID", nodeID, "backend", backendName) - reply, err := messaging.RequestJSON[messaging.BackendDeleteRequest, messaging.BackendDeleteReply](a.nats, subject, messaging.BackendDeleteRequest{Backend: backendName}, 2*time.Minute) - if err != nil { - return reply, err + ctx, cancel := context.WithTimeout(context.Background(), backendDeleteTimeout) + defer cancel() + + var reply messaging.BackendDeleteReply + if err := a.control.Call(ctx, nodeID, workerctl.PathBackendDelete, messaging.BackendDeleteRequest{Backend: backendName}, &reply); err != nil { + return nil, err } a.dropStoppedReplicaRows(nodeID, "backend.delete", backendName, reply.StoppedProcessKeys, reply.ReportsStoppedProcesses) - return reply, nil + return &reply, nil } // dropStoppedReplicaRows removes the NodeModel rows addressing processes a @@ -491,11 +507,13 @@ func (a *RemoteUnloaderAdapter) dropStoppedReplicaRows(nodeID, op, backendName s // UnloadModelOnNode sends a model.unload request to a specific node. // The worker calls gRPC Free() to release GPU memory. func (a *RemoteUnloaderAdapter) UnloadModelOnNode(nodeID, modelName string) error { - subject := messaging.SubjectNodeModelUnload(nodeID) - xlog.Info("Sending NATS model.unload", "nodeID", nodeID, "model", modelName) + xlog.Info("Sending model.unload", "nodeID", nodeID, "model", modelName) - reply, err := messaging.RequestJSON[messaging.ModelUnloadRequest, messaging.ModelUnloadReply](a.nats, subject, messaging.ModelUnloadRequest{ModelName: modelName}, 30*time.Second) - if err != nil { + ctx, cancel := context.WithTimeout(context.Background(), modelUnloadTimeout) + defer cancel() + + var reply messaging.ModelUnloadReply + if err := a.control.Call(ctx, nodeID, workerctl.PathModelUnload, messaging.ModelUnloadRequest{ModelName: modelName}, &reply); err != nil { return err } if !reply.Success { @@ -507,18 +525,20 @@ func (a *RemoteUnloaderAdapter) UnloadModelOnNode(nodeID, modelName string) erro // DeleteModelFiles sends model.delete to all nodes that have the model cached. // This removes model files from worker disks. func (a *RemoteUnloaderAdapter) DeleteModelFiles(modelName string) error { - nodes, err := a.registry.FindNodesWithModel(context.Background(), modelName) + ctx, cancel := context.WithTimeout(context.Background(), modelDeleteTimeout) + defer cancel() + + nodes, err := a.registry.FindNodesWithModel(ctx, modelName) if err != nil || len(nodes) == 0 { xlog.Debug("No nodes with model for file deletion", "model", modelName) return nil } for _, node := range nodes { - subject := messaging.SubjectNodeModelDelete(node.ID) - xlog.Info("Sending NATS model.delete", "nodeID", node.ID, "model", modelName) + xlog.Info("Sending model.delete", "nodeID", node.ID, "model", modelName) - reply, err := messaging.RequestJSON[messaging.ModelDeleteRequest, messaging.ModelDeleteReply](a.nats, subject, messaging.ModelDeleteRequest{ModelName: modelName}, 30*time.Second) - if err != nil { + var reply messaging.ModelDeleteReply + if err := a.control.Call(ctx, node.ID, workerctl.PathModelDelete, messaging.ModelDeleteRequest{ModelName: modelName}, &reply); err != nil { xlog.Warn("model.delete failed on node", "node", node.Name, "error", err) continue } @@ -531,17 +551,19 @@ func (a *RemoteUnloaderAdapter) DeleteModelFiles(modelName string) error { // StopNode tells a worker node to shut down entirely (deregister + exit). func (a *RemoteUnloaderAdapter) StopNode(nodeID string) error { - subject := messaging.SubjectNodeStop(nodeID) - return a.nats.Publish(subject, nil) + ctx, cancel := context.WithTimeout(context.Background(), nodeStopTimeout) + defer cancel() + return a.control.Call(ctx, nodeID, workerctl.PathNodeStop, struct{}{}, nil) } -// isNATSTimeout returns true if err looks like a NATS request-reply timeout. -// nats.ErrTimeout is the canonical sentinel; context.DeadlineExceeded can -// also surface depending on the client's path; we accept both, plus a -// string-match fallback for clients that return a bare error. -func isNATSTimeout(err error) bool { - if errors.Is(err, nats.ErrTimeout) || errors.Is(err, context.DeadlineExceeded) { - return true - } - return err != nil && strings.Contains(err.Error(), "nats: timeout") +// isRequestTimeout reports whether a control RPC ended because its budget ran +// out rather than because the worker said anything. +// +// context.DeadlineExceeded is the one signal, and it is matchable because +// controlFailure wraps the caller's own ctx.Err(). nats.ErrTimeout is still +// accepted for as long as any verb reaches a worker over the bus. The string +// match the NATS carrier needed is deliberately NOT reproduced: a message that +// merely quotes a timeout is not one. +func isRequestTimeout(err error) bool { + return errors.Is(err, nats.ErrTimeout) || errors.Is(err, context.DeadlineExceeded) } diff --git a/core/services/nodes/unloader_ping_test.go b/core/services/nodes/unloader_ping_test.go index a9b3a5889..ae2f594a0 100644 --- a/core/services/nodes/unloader_ping_test.go +++ b/core/services/nodes/unloader_ping_test.go @@ -1,51 +1,134 @@ package nodes import ( + "context" "errors" "time" + "github.com/nats-io/nats.go" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" - "github.com/nats-io/nats.go" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" ) -// The scheduler's liveness probe asks a worker a question over NATS and treats -// "no responders" as proof the worker is gone. That is only sound if every -// worker in the fleet subscribes to the subject asked. +// The scheduler asks a worker one cheap question before committing work to it, +// and then acts on the answer by DEMOTING the node. That is only sound if the +// question is one every worker in the fleet answers. // -// It originally asked models.running, which arrived in 4.6. A 4.5 worker is -// perfectly alive and serving, answers backend.list, and never subscribes to -// models.running, so the probe condemned it on every scheduling attempt. A -// model pinned to such a node could then never be placed at all. -var _ = Describe("Node liveness probe subject", func() { +// It has been the wrong question twice. It first asked models.running, which +// arrived in 4.6, so a 4.5 worker that was alive and serving never answered and +// a model pinned to it could never be placed. It then asked two NATS subjects +// the worker had stopped subscribing to when its control plane moved onto the +// tunnel, so EVERY healthy worker answered "no responders" and was marked +// unhealthy on the scheduling path. +// +// These specs pin both halves of the answer: what counts as the worker +// speaking, and what must never be read as the worker being gone. +var _ = Describe("Node liveness probe over the control plane", func() { var ( - mc *scriptedMessagingClient + workers *scriptedControlWorkers adapter *RemoteUnloaderAdapter ) const nodeID = "11111111-2222-3333-4444-555555555555" BeforeEach(func() { - mc = newScriptedMessagingClient() - adapter = NewRemoteUnloaderAdapter(nil, mc, 3*time.Minute, 15*time.Minute) + workers = newScriptedControlWorkers() + adapter = NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), 3*time.Minute, 15*time.Minute) }) It("treats a worker that answers backend.list as alive", func() { - // A worker old enough to predate models.running: it answers the - // long-standing backend.list subject and nothing else. - mc.scriptReply(messaging.SubjectNodeBackendList(nodeID), messaging.BackendListReply{}) - mc.scriptNoResponders(messaging.SubjectNodeModelsRunning(nodeID)) + workers.scriptReply(controlKey(nodeID, workerctl.PathBackendList), messaging.BackendListReply{}) - Expect(errors.Is(adapter.PingNode(nodeID), nats.ErrNoResponders)).To(BeFalse(), - "a worker answering backend.list is alive regardless of newer subjects") + Expect(adapter.PingNode(nodeID)).To(Succeed()) }) - It("still reports a worker that answers nothing as absent", func() { - mc.scriptNoResponders(messaging.SubjectNodeBackendList(nodeID)) - mc.scriptNoResponders(messaging.SubjectNodeModelsRunning(nodeID)) + It("treats a worker too old to serve the verb as alive, since it answered", func() { + // A 404 under the control prefix is the worker saying it is older than + // this frontend. It is a deployment fact, and a worker that can state + // one is a worker that is there. + workers.scriptUnsupported(controlKey(nodeID, workerctl.PathBackendList)) - Expect(errors.Is(adapter.PingNode(nodeID), nats.ErrNoResponders)).To(BeTrue()) + Expect(adapter.PingNode(nodeID)).To(Succeed()) + }) + + It("reports a worker it cannot route to as unroutable, and never as absent", func() { + workers.scriptUnroutable(nodeID) + + err := adapter.PingNode(nodeID) + Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue()) + // The two things a caller must not be able to conclude from it: that + // the worker spoke, and that it is gone. + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + Expect(errors.Is(err, nats.ErrNoResponders)).To(BeFalse()) + }) + + It("refuses without a tunnel dialer rather than reaching for an address", func() { + plain := NewRemoteUnloaderAdapter(nil, nil, NewControlClient(nil, "tok"), time.Minute, time.Minute) + + err := plain.PingNode(nodeID) + Expect(err).To(MatchError(ErrNoWorkerDialer)) + Expect(err).To(MatchError(ErrWorkerUnroutable)) + }) +}) + +// The merge gate this task exists to clear, asserted through the SCHEDULING +// path with the real adapter rather than through the router's own double. +// +// router_liveness_test.go drives a test-owned map of dead nodes and never +// touches a transport, so it stayed green for the whole window in which every +// healthy worker read as absent. These specs run pickReachableNode against a +// real RemoteUnloaderAdapter and a worker that answers over its control plane, +// which is the only arrangement that can see the difference. +var _ = Describe("Scheduling onto a worker reached over its control plane", func() { + var ( + workers *scriptedControlWorkers + reg *fakeModelRouter + router *SmartRouter + ) + + const nodeID = "22222222-3333-4444-5555-666666666666" + + BeforeEach(func() { + workers = newScriptedControlWorkers() + reg = &fakeModelRouter{} + router = NewSmartRouter(reg, SmartRouterOptions{ + Unloader: NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), time.Minute, time.Minute), + }) + }) + + selectOnce := func(node *BackendNode) func() *BackendNode { + handed := false + return func() *BackendNode { + if handed { + return nil + } + handed = true + return node + } + } + + It("schedules onto a healthy worker instead of demoting it", func() { + node := &BackendNode{ID: nodeID, Name: "healthy", Address: "10.0.0.1:50051"} + workers.scriptReply(controlKey(nodeID, workerctl.PathBackendList), messaging.BackendListReply{}) + + Expect(router.pickReachableNode(context.Background(), selectOnce(node))).To(Equal(node)) + Expect(reg.markedUnhealthy).To(BeEmpty()) + }) + + It("schedules onto a worker this frontend could not route to, rather than demoting it", func() { + // A worker re-homing its tunnel between frontend replicas, or one whose + // owning replica is restarting, is unroutable from here while it is + // heartbeating and serving. Excluding it costs capacity; demoting it + // tells every other scheduler the same wrong thing. + node := &BackendNode{ID: nodeID, Name: "re-homing", Address: "10.0.0.2:50051"} + workers.scriptUnroutable(nodeID) + + Expect(router.pickReachableNode(context.Background(), selectOnce(node))).To(Equal(node)) + Expect(reg.markedUnhealthy).To(BeEmpty()) }) }) diff --git a/core/services/nodes/unloader_stale_rows_test.go b/core/services/nodes/unloader_stale_rows_test.go index 0fff87baf..9be312917 100644 --- a/core/services/nodes/unloader_stale_rows_test.go +++ b/core/services/nodes/unloader_stale_rows_test.go @@ -8,6 +8,7 @@ import ( . "github.com/onsi/gomega" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" ) // Replies are handed to the adapter as raw JSON rather than as marshalled @@ -18,14 +19,16 @@ import ( var _ = Describe("RemoteUnloaderAdapter stale replica rows", func() { var ( locator *fakeModelLocator - mc *fakeMessagingClient + workers *scriptedControlWorkers adapter *RemoteUnloaderAdapter ) + const nodeID = "node-1" + BeforeEach(func() { locator = &fakeModelLocator{} - mc = &fakeMessagingClient{} - adapter = NewRemoteUnloaderAdapter(locator, mc, 3*time.Minute, 15*time.Minute) + workers = newScriptedControlWorkers() + adapter = NewRemoteUnloaderAdapter(locator, nil, workers.controlClient(), 3*time.Minute, 15*time.Minute) }) Describe("DeleteBackend", func() { @@ -35,11 +38,11 @@ var _ = Describe("RemoteUnloaderAdapter stale replica rows", func() { // 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(`{ + workers.scriptRawReply(controlKey(nodeID, workerctl.PathBackendDelete), []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()) @@ -56,11 +59,11 @@ var _ = Describe("RemoteUnloaderAdapter stale replica rows", func() { // 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(`{ + workers.scriptRawReply(controlKey(nodeID, workerctl.PathBackendDelete), []byte(`{ "success": true, "reports_stopped_processes": true, "stopped_process_keys": ["weird#name#3"] - }`) + }`)) _, err := adapter.DeleteBackend("node-1", "llama-cpp") Expect(err).NotTo(HaveOccurred()) @@ -70,7 +73,7 @@ var _ = Describe("RemoteUnloaderAdapter stale replica rows", func() { 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}`) + workers.scriptRawReply(controlKey(nodeID, workerctl.PathBackendDelete), []byte(`{"success": true, "reports_stopped_processes": true}`)) _, err := adapter.DeleteBackend("node-1", "llama-cpp") Expect(err).NotTo(HaveOccurred()) @@ -84,7 +87,7 @@ var _ = Describe("RemoteUnloaderAdapter stale replica rows", func() { // 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}`) + workers.scriptRawReply(controlKey(nodeID, workerctl.PathBackendDelete), []byte(`{"success": true}`)) reply, err := adapter.DeleteBackend("node-1", "llama-cpp") Expect(err).NotTo(HaveOccurred()) @@ -98,11 +101,11 @@ var _ = Describe("RemoteUnloaderAdapter stale replica rows", 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(`{ + workers.scriptRawReply(controlKey(nodeID, workerctl.PathBackendDelete), []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()) @@ -115,12 +118,12 @@ var _ = Describe("RemoteUnloaderAdapter stale replica rows", func() { // 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(`{ + workers.scriptRawReply(controlKey(nodeID, workerctl.PathBackendDelete), []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()) @@ -128,11 +131,11 @@ var _ = Describe("RemoteUnloaderAdapter stale replica rows", func() { }) It("ignores malformed process keys instead of removing a wrong row", func() { - mc.requestReply = []byte(`{ + workers.scriptRawReply(controlKey(nodeID, workerctl.PathBackendDelete), []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()) @@ -145,11 +148,11 @@ var _ = Describe("RemoteUnloaderAdapter stale replica rows", 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(`{ + workers.scriptRawReply(controlKey(nodeID, workerctl.PathBackendUpgrade), []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()) @@ -162,7 +165,7 @@ var _ = Describe("RemoteUnloaderAdapter stale replica rows", func() { }) It("does not read an old worker's silence as a completed cleanup", func() { - mc.requestReply = []byte(`{"success": true}`) + workers.scriptRawReply(controlKey(nodeID, workerctl.PathBackendUpgrade), []byte(`{"success": true}`)) reply, err := adapter.UpgradeBackend("node-1", "whisper", "", "", "", "", 0, "", nil) Expect(err).NotTo(HaveOccurred()) diff --git a/core/services/nodes/unloader_test.go b/core/services/nodes/unloader_test.go index 5a10370bc..9ea61e670 100644 --- a/core/services/nodes/unloader_test.go +++ b/core/services/nodes/unloader_test.go @@ -4,16 +4,15 @@ import ( "context" "encoding/json" "errors" - "fmt" "sync" "time" - "github.com/nats-io/nats.go" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" ) // --- Fakes --- @@ -22,6 +21,7 @@ import ( type fakeModelLocator struct { nodes []BackendNode findErr error + getErr error removedPairs []modelNodePair // records RemoveNodeModel calls removedReplicas []modelReplicaRef // records RemoveNodeModel calls including the replica index } @@ -46,6 +46,21 @@ func (f *fakeModelLocator) FindNodesWithModel(_ context.Context, _ string) ([]Ba return f.nodes, f.findErr } +// Get answers out of the same node list the locator hands out, so a spec that +// registers an AGENT node gets an agent node back and the carrier split is +// driven by the fixture rather than by a second thing to keep in step. +func (f *fakeModelLocator) Get(_ context.Context, nodeID string) (*BackendNode, error) { + if f.getErr != nil { + return nil, f.getErr + } + for i := range f.nodes { + if f.nodes[i].ID == nodeID { + return &f.nodes[i], nil + } + } + return nil, errors.New("no such node") +} + 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}) @@ -59,6 +74,10 @@ func (f *fakeModelLocator) RemoveAllNodeModelReplicas(_ context.Context, nodeID, // fakeMessagingClient implements messaging.MessagingClient, recording Publish // and Request calls so we can assert on subjects and payloads. +// +// Only ONE verb still reaches it: backend.stop to an agent node. Every other +// control verb travels over the tunnel, so a publish recorded here for a +// backend node is a bug, and several specs below assert exactly that. type fakeMessagingClient struct { mu sync.Mutex published []publishCall @@ -120,6 +139,17 @@ func (f *fakeMessagingClient) Request(subject string, data []byte, timeout time. func (f *fakeMessagingClient) IsConnected() bool { return true } func (f *fakeMessagingClient) Close() {} +// publishedSubjects reports what actually reached the bus. +func (f *fakeMessagingClient) publishedSubjects() []string { + f.mu.Lock() + defer f.mu.Unlock() + out := make([]string, 0, len(f.published)) + for _, c := range f.published { + out = append(out, c.Subject) + } + return out +} + type fakeSubscription struct{} func (f *fakeSubscription) Unsubscribe() error { return nil } @@ -129,16 +159,24 @@ func (f *fakeSubscription) Unsubscribe() error { return nil } var _ = Describe("RemoteUnloaderAdapter", func() { var ( locator *fakeModelLocator - mc *fakeMessagingClient + bus *fakeMessagingClient + workers *scriptedControlWorkers adapter *RemoteUnloaderAdapter ) BeforeEach(func() { locator = &fakeModelLocator{} - mc = &fakeMessagingClient{} - adapter = NewRemoteUnloaderAdapter(locator, mc, 3*time.Minute, 15*time.Minute) + bus = &fakeMessagingClient{} + workers = newScriptedControlWorkers() + adapter = NewRemoteUnloaderAdapter(locator, bus, workers.controlClient(), 3*time.Minute, 15*time.Minute) }) + // scriptStop lets a backend node accept the tunnelled backend.stop, which + // answers 204 and therefore carries no body. + scriptStop := func(nodeID string) { + workers.scriptRawReply(controlKey(nodeID, workerctl.PathBackendStop), []byte(`{}`)) + } + // 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 @@ -179,20 +217,25 @@ var _ = Describe("RemoteUnloaderAdapter", func() { // 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()) + Expect(workers.callSubjects()).To(BeEmpty()) + Expect(bus.publishedSubjects()).To(BeEmpty()) }) - It("broadcasts to all nodes with model", func() { + It("stops the backend on every node holding the model, over their tunnels", func() { locator.nodes = []BackendNode{ - {ID: "node-1", Name: "worker-1"}, - {ID: "node-2", Name: "worker-2"}, + {ID: "node-1", Name: "worker-1", NodeType: NodeTypeBackend}, + {ID: "node-2", Name: "worker-2", NodeType: NodeTypeBackend}, } + scriptStop("node-1") + scriptStop("node-2") + Expect(adapter.UnloadRemoteModel("llama")).To(Succeed()) - // Should have published a StopBackend for each node. - Expect(mc.published).To(HaveLen(2)) - Expect(mc.published[0].Subject).To(Equal(messaging.SubjectNodeBackendStop("node-1"))) - Expect(mc.published[1].Subject).To(Equal(messaging.SubjectNodeBackendStop("node-2"))) + Expect(workers.callSubjects()).To(Equal([]string{ + controlKey("node-1", workerctl.PathBackendStop), + controlKey("node-2", workerctl.PathBackendStop), + })) + Expect(bus.publishedSubjects()).To(BeEmpty()) // Should have removed the model from each node in the registry. Expect(locator.removedPairs).To(HaveLen(2)) @@ -202,46 +245,101 @@ var _ = Describe("RemoteUnloaderAdapter", func() { It("continues when one node fails", func() { locator.nodes = []BackendNode{ - {ID: "node-fail", Name: "worker-fail"}, - {ID: "node-ok", Name: "worker-ok"}, + {ID: "node-fail", Name: "worker-fail", NodeType: NodeTypeBackend}, + {ID: "node-ok", Name: "worker-ok", NodeType: NodeTypeBackend}, } - // Use a messaging client that fails the first Publish call only. - failOnce := &failOnceMessagingClient{inner: mc, failOn: 0} - adapter = NewRemoteUnloaderAdapter(locator, failOnce, 3*time.Minute, 15*time.Minute) + workers.scriptUnroutable("node-fail") + scriptStop("node-ok") 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. - // The second node's StopBackend succeeded, so RemoveNodeModel WAS called. + // The second node should still have been processed. The first + // node's stop errored, so its row was NOT dropped: a row deleted + // on a route this frontend could not open is a model reclaimed + // while it is still loaded. 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"}} + locator.nodes = []BackendNode{{ID: "node-1", Name: "worker-1", NodeType: NodeTypeBackend}} + scriptStop("node-1") + Expect(adapter.UnloadRemoteModelContext(context.Background(), "llama", true)).To(Succeed()) + workers.mu.Lock() + defer workers.mu.Unlock() + Expect(workers.calls).To(HaveLen(1)) var payload messaging.BackendStopRequest - Expect(json.Unmarshal(mc.published[0].Data, &payload)).To(Succeed()) + Expect(json.Unmarshal(workers.calls[0].Data, &payload)).To(Succeed()) Expect(payload).To(Equal(messaging.BackendStopRequest{Backend: "llama", Force: true})) }) }) - Describe("StopBackend", func() { - It("with empty backend publishes nil payload", func() { - Expect(adapter.StopBackend("node-1", "")).To(Succeed()) - Expect(mc.published).To(HaveLen(1)) - Expect(mc.published[0].Subject).To(Equal(messaging.SubjectNodeBackendStop("node-1"))) - Expect(mc.published[0].Data).To(BeNil()) + // The carrier split. It is the one verb of the ten that is decided by the + // KIND of worker, and getting it wrong is silent in both directions: a + // backend node's stop published on the bus reaches nothing, and an agent + // node's stop sent over a tunnel it does not hold reaches nothing either. + Describe("StopBackend and node type", func() { + It("sends a backend stop to a BACKEND node over the tunnel and not over the bus", func() { + locator.nodes = []BackendNode{{ID: "backend-1", Name: "gpu-1", NodeType: NodeTypeBackend}} + scriptStop("backend-1") + + Expect(adapter.StopBackend("backend-1", "llama-backend")).To(Succeed()) + + Expect(workers.callSubjects()).To(ContainElement(controlKey("backend-1", workerctl.PathBackendStop))) + Expect(bus.publishedSubjects()).To(BeEmpty()) }) - It("with backend name publishes JSON", func() { - Expect(adapter.StopBackend("node-1", "llama-backend")).To(Succeed()) - Expect(mc.published).To(HaveLen(1)) + It("sends a backend stop to an AGENT node over the bus, because agent workers hold no tunnel", func() { + locator.nodes = []BackendNode{{ID: "agent-1", Name: "agent", NodeType: NodeTypeAgent}} + Expect(adapter.StopBackend("agent-1", "llama-backend")).To(Succeed()) + + Expect(bus.publishedSubjects()).To(ContainElement(messaging.SubjectNodeBackendStop("agent-1"))) + Expect(workers.callSubjects()).ToNot(ContainElement(controlKey("agent-1", workerctl.PathBackendStop))) + }) + + It("treats a node whose type cannot be read as a backend worker", func() { + // The column defaults to backend and every other node-type branch + // in this package reads an empty value the same way. The lookup + // failing must not silently move a stop onto a carrier nothing is + // listening on. + locator.getErr = errors.New("database is down") + scriptStop("unknown-1") + + Expect(adapter.StopBackend("unknown-1", "llama-backend")).To(Succeed()) + + Expect(workers.callSubjects()).To(ContainElement(controlKey("unknown-1", workerctl.PathBackendStop))) + Expect(bus.publishedSubjects()).To(BeEmpty()) + }) + + It("with an empty backend asks the worker to stop everything", func() { + locator.nodes = []BackendNode{{ID: "backend-1", Name: "gpu-1", NodeType: NodeTypeBackend}} + scriptStop("backend-1") + + Expect(adapter.StopBackend("backend-1", "")).To(Succeed()) + + workers.mu.Lock() + defer workers.mu.Unlock() var payload messaging.BackendStopRequest - Expect(json.Unmarshal(mc.published[0].Data, &payload)).To(Succeed()) + Expect(json.Unmarshal(workers.calls[0].Data, &payload)).To(Succeed()) + // An empty Backend is what the worker reads as "stop everything"; + // see decodeBackendStopRequest. + Expect(payload.Backend).To(BeEmpty()) + Expect(payload.Force).To(BeFalse()) + }) + + It("names the backend when one is given", func() { + locator.nodes = []BackendNode{{ID: "backend-1", Name: "gpu-1", NodeType: NodeTypeBackend}} + scriptStop("backend-1") + + Expect(adapter.StopBackend("backend-1", "llama-backend")).To(Succeed()) + + workers.mu.Lock() + defer workers.mu.Unlock() + var payload messaging.BackendStopRequest + Expect(json.Unmarshal(workers.calls[0].Data, &payload)).To(Succeed()) Expect(payload.Backend).To(Equal("llama-backend")) Expect(payload.Force).To(BeFalse()) }) @@ -249,30 +347,46 @@ var _ = Describe("RemoteUnloaderAdapter", func() { Describe("StopModelReplica", func() { It("requests an acknowledged stop for the exact process", func() { - mc.requestReply, _ = json.Marshal(messaging.ModelStopReply{Matched: true, Terminated: true, ProcessKey: "llama#2"}) + workers.scriptReply(controlKey("node-1", workerctl.PathModelStop), + messaging.ModelStopReply{Matched: true, Terminated: true, ProcessKey: "llama#2"}) replica := NodeModel{ModelName: "llama", ReplicaIndex: 2, WorkerLocalAddress: "127.0.0.1:5002", ConfigRevision: "rev-1"} reply, err := adapter.StopModelReplica(context.Background(), "node-1", replica, true) Expect(err).NotTo(HaveOccurred()) Expect(reply.Terminated).To(BeTrue()) - Expect(mc.requestCalls).To(HaveLen(1)) - Expect(mc.requestCalls[0].Subject).To(Equal(messaging.SubjectNodeModelStop("node-1"))) - Expect(mc.requestCalls[0].Timeout).To(BeNumerically(">", 0)) + + workers.mu.Lock() + defer workers.mu.Unlock() + Expect(workers.calls).To(HaveLen(1)) + Expect(workers.calls[0].Subject).To(Equal(controlKey("node-1", workerctl.PathModelStop))) var request messaging.ModelStopRequest - Expect(json.Unmarshal(mc.requestCalls[0].Data, &request)).To(Succeed()) + Expect(json.Unmarshal(workers.calls[0].Data, &request)).To(Succeed()) Expect(request).To(Equal(messaging.ModelStopRequest{ ModelName: "llama", ProcessKey: "llama#2", ExpectedAddress: "127.0.0.1:5002", Force: true, ConfigRevision: "rev-1", })) }) + + It("reports an unroutable worker without inventing a stop reply", func() { + // A zero ModelStopReply reads as Matched=false, which the cleanup + // path treats as "there was nothing to stop" and drops the row. It + // must only ever be paired with an error. + workers.scriptUnroutable("node-gone") + + reply, err := adapter.StopModelReplica(context.Background(), "node-gone", NodeModel{ModelName: "llama"}, false) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue()) + Expect(reply).To(Equal(messaging.ModelStopReply{})) + }) }) Describe("StopNode", func() { - It("publishes to correct subject", func() { + It("asks the worker to shut down over its tunnel", func() { + workers.scriptRawReply(controlKey("node-abc", workerctl.PathNodeStop), []byte(`{}`)) + Expect(adapter.StopNode("node-abc")).To(Succeed()) - Expect(mc.published).To(HaveLen(1)) - Expect(mc.published[0].Subject).To(Equal(messaging.SubjectNodeStop("node-abc"))) - Expect(mc.published[0].Data).To(BeNil()) + + Expect(workers.callSubjects()).To(Equal([]string{controlKey("node-abc", workerctl.PathNodeStop)})) + Expect(bus.publishedSubjects()).To(BeEmpty()) }) }) @@ -284,94 +398,87 @@ var _ = Describe("RemoteUnloaderAdapter", func() { It("continues on failure", func() { locator.nodes = []BackendNode{ - {ID: "node-1", Name: "w1"}, - {ID: "node-2", Name: "w2"}, + {ID: "node-1", Name: "w1", NodeType: NodeTypeBackend}, + {ID: "node-2", Name: "w2", NodeType: NodeTypeBackend}, } - // Request will fail for all calls. - mc.requestErr = fmt.Errorf("timeout") + // Neither node is scripted, so both answer the loud + // unscripted-verb default. Both must still be attempted. Expect(adapter.DeleteModelFiles("my-model")).To(Succeed()) - // Both nodes attempted. - Expect(mc.requestCalls).To(HaveLen(2)) - Expect(mc.requestCalls[0].Subject).To(Equal(messaging.SubjectNodeModelDelete("node-1"))) - Expect(mc.requestCalls[1].Subject).To(Equal(messaging.SubjectNodeModelDelete("node-2"))) + Expect(workers.callSubjects()).To(Equal([]string{ + controlKey("node-1", workerctl.PathModelDelete), + controlKey("node-2", workerctl.PathModelDelete), + })) + }) + }) + + Describe("UnloadModelOnNode", func() { + It("succeeds when the worker reports the model freed", func() { + workers.scriptReply(controlKey("node-1", workerctl.PathModelUnload), messaging.ModelUnloadReply{Success: true}) + Expect(adapter.UnloadModelOnNode("node-1", "llama")).To(Succeed()) + }) + + It("surfaces the worker's own refusal, which is an answer and not a lost route", func() { + workers.scriptReply(controlKey("node-1", workerctl.PathModelUnload), + messaging.ModelUnloadReply{Success: false, Error: "Free failed"}) + + err := adapter.UnloadModelOnNode("node-1", "llama") + Expect(err).To(MatchError(ContainSubstring("Free failed"))) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeFalse()) }) }) }) -// failOnceMessagingClient wraps fakeMessagingClient but fails the Publish call -// at index failOn (0-based) and succeeds all others. -type failOnceMessagingClient struct { - inner *fakeMessagingClient - failOn int - callIdx int - mu sync.Mutex -} - -func (f *failOnceMessagingClient) Publish(subject string, data any) error { - f.mu.Lock() - idx := f.callIdx - f.callIdx++ - f.mu.Unlock() - if idx == f.failOn { - return fmt.Errorf("simulated failure") - } - return f.inner.Publish(subject, data) -} - -func (f *failOnceMessagingClient) Subscribe(subject string, handler func([]byte)) (messaging.Subscription, error) { - return f.inner.Subscribe(subject, handler) -} - -func (f *failOnceMessagingClient) QueueSubscribe(subject, queue string, handler func([]byte)) (messaging.Subscription, error) { - return f.inner.QueueSubscribe(subject, queue, handler) -} - -func (f *failOnceMessagingClient) QueueSubscribeReply(subject, queue string, handler func(data []byte, reply func([]byte))) (messaging.Subscription, error) { - return f.inner.QueueSubscribeReply(subject, queue, handler) -} - -func (f *failOnceMessagingClient) SubscribeReply(subject string, handler func(data []byte, reply func([]byte))) (messaging.Subscription, error) { - return f.inner.SubscribeReply(subject, handler) -} - -func (f *failOnceMessagingClient) Request(subject string, data []byte, timeout time.Duration) ([]byte, error) { - return f.inner.Request(subject, data, timeout) -} - -func (f *failOnceMessagingClient) IsConnected() bool { return true } -func (f *failOnceMessagingClient) Close() {} - var _ = Describe("RemoteUnloaderAdapter timeout configuration", func() { - It("passes the configured install timeout to the messaging client", func() { - mc := newScriptedMessagingClient() - mc.scriptReply(messaging.SubjectNodeBackendInstall("n1"), messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "127.0.0.1:0"}) - adapter := NewRemoteUnloaderAdapter(nil, mc, 7*time.Minute, 11*time.Minute) + // Each verb must be given ITS OWN configured budget, and the budget is + // observed the only way it can be: by how long the client waits on a worker + // that took the request and never answered. HTTP carries no deadline, and + // the transport dials on a context of its own, so nothing on the far side + // can see it. + // + // The two budgets are far apart so a swap cannot pass: install would wait + // the upgrade budget and upgrade would wait the install one, and each + // assertion below is on the wrong side of the divide for the other. + const ( + installBudget = 150 * time.Millisecond + upgradeBudget = 600 * time.Millisecond + divide = 400 * time.Millisecond + ) + It("gives backend.install the configured install timeout", func() { + workers := newScriptedControlWorkers() + workers.scriptHang(controlKey("n1", workerctl.PathBackendInstall)) + adapter := NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), installBudget, upgradeBudget) + + started := time.Now() _, err := adapter.InstallBackend("n1", "llama-cpp", "", "[]", "", "", "", 0, "", nil) - Expect(err).ToNot(HaveOccurred()) + elapsed := time.Since(started) - Expect(mc.calls).To(HaveLen(1)) - Expect(mc.calls[0].Timeout).To(Equal(7 * time.Minute)) + Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeTrue(), + "a spent budget is reported as still-installing, got %v", err) + Expect(elapsed).To(BeNumerically(">=", installBudget)) + Expect(elapsed).To(BeNumerically("<", divide)) }) - It("passes the configured upgrade timeout to the messaging client", func() { - mc := newScriptedMessagingClient() - mc.scriptReply(messaging.SubjectNodeBackendUpgrade("n1"), messaging.BackendUpgradeReply{Success: true}) - adapter := NewRemoteUnloaderAdapter(nil, mc, 7*time.Minute, 11*time.Minute) + It("gives backend.upgrade the configured upgrade timeout", func() { + workers := newScriptedControlWorkers() + workers.scriptHang(controlKey("n1", workerctl.PathBackendUpgrade)) + adapter := NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), installBudget, upgradeBudget) + started := time.Now() _, err := adapter.UpgradeBackend("n1", "llama-cpp", "[]", "", "", "", 0, "", nil) - Expect(err).ToNot(HaveOccurred()) + elapsed := time.Since(started) - Expect(mc.calls).To(HaveLen(1)) - Expect(mc.calls[0].Timeout).To(Equal(11 * time.Minute)) + Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeTrue(), + "a spent budget is reported as still-installing, got %v", err) + Expect(elapsed).To(BeNumerically(">", divide)) }) }) -var _ = Describe("RemoteUnloaderAdapter NATS timeout handling", func() { - It("wraps nats.ErrTimeout from InstallBackend in galleryop.ErrWorkerStillInstalling", func() { - mc := newScriptedMessagingClient() - mc.scriptErr(messaging.SubjectNodeBackendInstall("n1"), nats.ErrTimeout) - adapter := NewRemoteUnloaderAdapter(nil, mc, 100*time.Millisecond, 1*time.Second) +var _ = Describe("RemoteUnloaderAdapter timeout handling", func() { + It("reports a spent budget as still-installing, so the operation shows as running on the worker", func() { + workers := newScriptedControlWorkers() + workers.scriptTimeout("n1") + adapter := NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), 100*time.Millisecond, 1*time.Second) _, err := adapter.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "", nil) Expect(err).To(HaveOccurred()) @@ -379,57 +486,71 @@ var _ = Describe("RemoteUnloaderAdapter NATS timeout handling", func() { "expected wrapped ErrWorkerStillInstalling, got %v", err) }) - It("does NOT wrap non-timeout errors", func() { - mc := newScriptedMessagingClient() - mc.scriptErr(messaging.SubjectNodeBackendInstall("n1"), nats.ErrNoResponders) - adapter := NewRemoteUnloaderAdapter(nil, mc, 100*time.Millisecond, 1*time.Second) + It("does NOT report an unroutable worker as still installing", func() { + // The two are different facts and the operator UI shows them + // differently: one keeps the queue row and pushes the retry out, the + // other is a plain failure. Both are non-verdicts, so neither may reap. + workers := newScriptedControlWorkers() + workers.scriptUnroutable("n1") + adapter := NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), 100*time.Millisecond, 1*time.Second) _, err := adapter.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "", nil) Expect(err).To(HaveOccurred()) Expect(errors.Is(err, galleryop.ErrWorkerStillInstalling)).To(BeFalse()) - Expect(errors.Is(err, nats.ErrNoResponders)).To(BeTrue()) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue()) + }) + + It("does not read a message merely containing the words nats timeout as a timeout", func() { + // The string match the bus carrier needed would have matched a worker + // error that quoted the phrase, and would have turned a real failure + // into "still installing" forever. + workers := newScriptedControlWorkers() + workers.scriptReply(controlKey("n1", workerctl.PathBackendInstall), + messaging.BackendInstallReply{Success: false, Error: `the worker said "nats: timeout" in its log`}) + adapter := NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), time.Minute, time.Minute) + + reply, err := adapter.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "", nil) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Success).To(BeFalse()) }) }) var _ = Describe("RemoteUnloaderAdapter install progress streaming", func() { - It("forwards BackendInstallProgressEvent values into the onProgress callback when the worker publishes them", func() { - mc := newScriptedMessagingClient() - mc.scriptReply(messaging.SubjectNodeBackendInstall("n1"), messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "127.0.0.1:0"}) - mc.scheduleProgressPublish("n1", "op-abc", []messaging.BackendInstallProgressEvent{ + It("forwards the worker's progress lines to onProgress, in the order it wrote them", func() { + workers := newScriptedControlWorkers() + workers.scriptReply(controlKey("n1", workerctl.PathBackendInstall), + messaging.BackendInstallReply{Success: true, WorkerLocalAddress: "127.0.0.1:0"}) + workers.scriptProgress(controlKey("n1", workerctl.PathBackendInstall), []messaging.BackendInstallProgressEvent{ {OpID: "op-abc", NodeID: "n1", Backend: "vllm", FileName: "vllm.tar.zst", Current: "100 MB", Total: "1 GB", Percentage: 10}, {OpID: "op-abc", NodeID: "n1", Backend: "vllm", FileName: "vllm.tar.zst", Current: "500 MB", Total: "1 GB", Percentage: 50}, }) - adapter := NewRemoteUnloaderAdapter(nil, mc, 1*time.Second, 1*time.Second) - var ( - received []messaging.BackendInstallProgressEvent - mu sync.Mutex - ) + adapter := NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), time.Second, time.Second) + var received []messaging.BackendInstallProgressEvent onProgress := func(ev messaging.BackendInstallProgressEvent) { - mu.Lock() - defer mu.Unlock() + // No lock, and that is an assertion in itself: the callback runs + // synchronously on the caller's own goroutine, so a race detector + // run would fail here if it did not. received = append(received, ev) } _, err := adapter.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "op-abc", onProgress) Expect(err).ToNot(HaveOccurred()) - Eventually(func() int { - mu.Lock() - defer mu.Unlock() - return len(received) - }, "1s").Should(Equal(2)) + Expect(received).To(HaveLen(2)) + Expect([]float64{received[0].Percentage, received[1].Percentage}).To(Equal([]float64{10, 50})) }) - It("does NOT subscribe when onProgress is nil (reconciler retry path)", func() { - mc := newScriptedMessagingClient() - mc.scriptReply(messaging.SubjectNodeBackendInstall("n1"), messaging.BackendInstallReply{Success: true}) + It("completes when the caller wants no progress at all (reconciler retry path)", func() { + workers := newScriptedControlWorkers() + workers.scriptReply(controlKey("n1", workerctl.PathBackendInstall), messaging.BackendInstallReply{Success: true}) + workers.scriptProgress(controlKey("n1", workerctl.PathBackendInstall), []messaging.BackendInstallProgressEvent{ + {OpID: "", NodeID: "n1", Percentage: 42}, + }) - adapter := NewRemoteUnloaderAdapter(nil, mc, 1*time.Second, 1*time.Second) - _, err := adapter.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "", nil) + adapter := NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), time.Second, time.Second) + reply, err := adapter.InstallBackend("n1", "vllm", "", "[]", "", "", "", 0, "", nil) Expect(err).ToNot(HaveOccurred()) - - Expect(mc.subscribeCalls()).To(BeEmpty(), - "reconciler-driven retries must not subscribe to the per-op progress subject") + Expect(reply.Success).To(BeTrue()) }) }) diff --git a/core/services/nodes/unloader_upgrade_test.go b/core/services/nodes/unloader_upgrade_test.go index bad8f9ed5..41fe7316f 100644 --- a/core/services/nodes/unloader_upgrade_test.go +++ b/core/services/nodes/unloader_upgrade_test.go @@ -1,82 +1,78 @@ package nodes import ( - "sync" + "errors" "time" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/workerctl" ) var _ = Describe("RemoteUnloaderAdapter.UpgradeBackend", func() { - It("fires a NATS request to the backend.upgrade subject and returns the reply", func() { - mc := newScriptedMessagingClient() + It("calls the backend.upgrade verb on the worker and returns the reply", func() { + workers := newScriptedControlWorkers() nodeID := "node-x" - mc.scriptReply(messaging.SubjectNodeBackendUpgrade(nodeID), + workers.scriptReply(controlKey(nodeID, workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: true}) - adapter := NewRemoteUnloaderAdapter(nil, mc, 3*time.Minute, 15*time.Minute) + adapter := NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), 3*time.Minute, 15*time.Minute) reply, err := adapter.UpgradeBackend(nodeID, "llama-cpp", `[{"name":"x"}]`, "", "", "", 0, "", nil) Expect(err).ToNot(HaveOccurred()) Expect(reply.Success).To(BeTrue()) + Expect(workers.callSubjects()).To(Equal([]string{controlKey(nodeID, workerctl.PathBackendUpgrade)})) }) - It("returns the underlying error when the subject has no responders", func() { - mc := newScriptedMessagingClient() // unscripted subject => fakeNoRespondersErr by harness convention + It("reports a worker it cannot reach as unroutable rather than as a failed upgrade", func() { + workers := newScriptedControlWorkers() + workers.scriptUnroutable("missing-node") - adapter := NewRemoteUnloaderAdapter(nil, mc, 3*time.Minute, 15*time.Minute) + adapter := NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), 3*time.Minute, 15*time.Minute) _, err := adapter.UpgradeBackend("missing-node", "llama-cpp", "", "", "", "", 0, "", nil) - Expect(err).To(HaveOccurred()) + Expect(errors.Is(err, ErrWorkerUnroutable)).To(BeTrue()) + // Not the worker's 404 either: nothing was asked of it, so it did not + // say it lacks the verb, and the legacy force-install fallback must not + // fire on this. + Expect(errors.Is(err, ErrWorkerControlUnsupported)).To(BeFalse()) }) // Reproducer for "upgrade reports progress:0 the whole time" (Bug B). The // install path streamed per-node download ticks; the upgrade path did a bare - // request→single-reply with no progress subscription, so a long force-reinstall - // blocked opaque. The adapter must subscribe to the per-op progress subject - // (reused from install) BEFORE the request and deliver each tick to onProgress. + // request→single-reply with no progress at all, so a long force-reinstall + // blocked opaque. Both verbs now stream on the same envelope shape. It("streams per-node progress ticks during the upgrade", func() { - mc := newScriptedMessagingClient() + workers := newScriptedControlWorkers() nodeID := "node-slow" opID := "op-upgrade-1" - mc.scriptReply(messaging.SubjectNodeBackendUpgrade(nodeID), + workers.scriptReply(controlKey(nodeID, workerctl.PathBackendUpgrade), messaging.BackendUpgradeReply{Success: true}) - // The worker would publish these while force-reinstalling. The harness - // replays them as soon as the adapter subscribes to the per-op subject. - mc.scheduleProgressPublish(nodeID, opID, []messaging.BackendInstallProgressEvent{ + workers.scriptProgress(controlKey(nodeID, workerctl.PathBackendUpgrade), []messaging.BackendInstallProgressEvent{ {NodeID: nodeID, FileName: "llama-cpp.tar", Current: "10 MB", Total: "100 MB", Percentage: 10}, {NodeID: nodeID, FileName: "llama-cpp.tar", Current: "100 MB", Total: "100 MB", Percentage: 100}, }) - var mu sync.Mutex var got []messaging.BackendInstallProgressEvent onProgress := func(ev messaging.BackendInstallProgressEvent) { - mu.Lock() got = append(got, ev) - mu.Unlock() } - adapter := NewRemoteUnloaderAdapter(nil, mc, 3*time.Minute, 15*time.Minute) + adapter := NewRemoteUnloaderAdapter(nil, nil, workers.controlClient(), 3*time.Minute, 15*time.Minute) reply, err := adapter.UpgradeBackend(nodeID, "llama-cpp", `[{"name":"x"}]`, "", "", "", 0, opID, onProgress) Expect(err).ToNot(HaveOccurred()) Expect(reply.Success).To(BeTrue()) - // Confirm it subscribed to the (reused) install-progress subject for this op. - Expect(mc.subscribeCalls()).To(ContainElement(messaging.SubjectNodeBackendInstallProgress(nodeID, opID))) - - // Progress events are delivered asynchronously (goroutine-per-event), so - // poll for both and assert on the set — ordering is best-effort by design. - Eventually(func() []float64 { - mu.Lock() - defer mu.Unlock() - pcts := make([]float64, 0, len(got)) - for _, e := range got { - pcts = append(pcts, e.Percentage) - } - return pcts - }, 2*time.Second, 20*time.Millisecond).Should(ConsistOf(float64(10), float64(100))) + // Every tick, in the worker's order, and all of them BEFORE the reply + // the call returned: the reply line is the last thing on the body, so a + // tick arriving late is structurally impossible rather than merely + // unlikely. + pcts := make([]float64, 0, len(got)) + for _, e := range got { + pcts = append(pcts, e.Percentage) + } + Expect(pcts).To(Equal([]float64{10, 100})) }) }) diff --git a/core/services/worker/control_client_roundtrip_test.go b/core/services/worker/control_client_roundtrip_test.go new file mode 100644 index 000000000..52046f320 --- /dev/null +++ b/core/services/worker/control_client_roundtrip_test.go @@ -0,0 +1,151 @@ +package worker + +import ( + "context" + "errors" + "net" + "os" + "path/filepath" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// The frontend's control client against the REAL worker control plane: the real +// supervisor's handlers, mounted through the real nodes.AuthenticatedRoutes so +// the real bearer check runs first, reached by a real nodes.ControlClient over +// a real HTTP transport. The only thing standing in for production is the +// transport's dial, which is the seam the tunnel occupies. +// +// It lives in this package because the dependency runs worker -> nodes: the +// frontend's client cannot be exercised against the real handler from the other +// side without an import cycle. Both halves of the contract are written by +// different packages, so a spec on either side alone can only prove that side +// agrees with itself; the paths, the envelope shape and the status codes are +// only pinned together here. +var _ = Describe("the frontend's control client against the real worker", func() { + const ( + token = "s3cret-registration-token" + nodeID = "worker-under-test" + ) + + var ( + sup *backendSupervisor + client *nodes.ControlClient + srvAddr string + sigCh chan os.Signal + ) + + // newClient builds a control client whose transport dials srvAddr, + // authenticating with the token given. + newClient := func(srvAddr, tok string) *nodes.ControlClient { + return nodes.NewControlClient(func(string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, _, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "tcp", srvAddr) + } + }, tok) + } + + BeforeEach(func() { + sigCh = make(chan os.Signal, 1) + sup = &backendSupervisor{ + cfg: &Config{}, + nodeID: nodeID, + sigCh: sigCh, + processes: map[string]*backendProcess{}, + } + + dir := GinkgoT().TempDir() + lis, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).NotTo(HaveOccurred()) + srv, err := nodes.StartFileTransferServerWithRoutes(lis, + filepath.Join(dir, "staging"), filepath.Join(dir, "models"), filepath.Join(dir, "data"), + token, config.DefaultMaxUploadSize, nil, + &nodes.AuthenticatedRoutes{Prefix: workerctl.Prefix, Register: sup.RegisterControlRoutes}) + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { _ = srv.Close() }) + + srvAddr = lis.Addr().String() + client = newClient(srvAddr, token) + }) + + It("round-trips models.running end to end", func() { + var reply messaging.ModelsRunningReply + Expect(client.Call(context.Background(), nodeID, workerctl.PathModelsRunning, + messaging.ModelsRunningRequest{}, &reply)).To(Succeed()) + // Nothing is running, and an empty list is the worker's real answer + // rather than a decode that quietly produced nothing: the client + // reports a body it cannot read as unroutable instead. + Expect(reply.Models).To(BeEmpty()) + Expect(reply.Error).To(BeEmpty()) + }) + + It("streams install progress in order and returns the terminal reply", func() { + sup.installFn = func(_ context.Context, req messaging.BackendInstallRequest, _ bool, + onProgress func(messaging.BackendInstallProgressEvent)) (string, error) { + onProgress(messaging.BackendInstallProgressEvent{OpID: req.OpID, Percentage: 50}) + onProgress(messaging.BackendInstallProgressEvent{OpID: req.OpID, Percentage: 100}) + return "127.0.0.1:41234", nil + } + + var seen []float64 + var reply messaging.BackendInstallReply + err := client.CallStreaming(context.Background(), nodeID, workerctl.PathBackendInstall, + messaging.BackendInstallRequest{Backend: "mock", OpID: "op-1"}, &reply, + func(ev messaging.BackendInstallProgressEvent) { seen = append(seen, ev.Percentage) }) + Expect(err).NotTo(HaveOccurred()) + Expect(seen).To(Equal([]float64{50, 100})) + Expect(reply.Success).To(BeTrue()) + Expect(reply.WorkerLocalAddress).To(Equal("127.0.0.1:41234")) + }) + + It("reports a FAILED install as the worker's own answer, not as a transport failure", func() { + // The distinction the two sides exist to preserve: the worker answers + // 200 with Error set, so the frontend reads a verdict a caller may act + // on rather than a route it must not act on. + sup.installFn = func(context.Context, messaging.BackendInstallRequest, bool, + func(messaging.BackendInstallProgressEvent)) (string, error) { + return "", errors.New("no child with platform linux/arm64") + } + + var reply messaging.BackendInstallReply + err := client.CallStreaming(context.Background(), nodeID, workerctl.PathBackendInstall, + messaging.BackendInstallRequest{Backend: "mock"}, &reply, nil) + Expect(err).NotTo(HaveOccurred()) + Expect(reply.Success).To(BeFalse()) + Expect(reply.Error).To(ContainSubstring("linux/arm64")) + }) + + It("round-trips a 204 verb, which carries no body at all", func() { + Expect(client.Call(context.Background(), nodeID, workerctl.PathBackendStop, + messaging.BackendStopRequest{Backend: "no-such-backend"}, nil)).To(Succeed()) + }) + + It("reports an unknown control verb as unsupported, and not as absence", func() { + err := client.Call(context.Background(), nodeID, workerctl.Prefix+"invented", struct{}{}, &struct{}{}) + Expect(err).To(MatchError(nodes.ErrWorkerControlUnsupported)) + Expect(errors.Is(err, cluster.ErrNoRoute)).To(BeFalse()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + }) + + It("reports a rejected token as unroutable, never as a worker verdict about a backend", func() { + // The bearer check runs before routing, so a wrong token is a 401 for + // every verb. It says nothing about any backend and nothing may reap on + // it; it also must not be read as the verb being unsupported, which + // would send the upgrade path into its destructive legacy fallback. + wrong := newClient(srvAddr, "not-the-token") + var reply messaging.BackendListReply + err := wrong.Call(context.Background(), nodeID, workerctl.PathBackendList, + messaging.BackendListRequest{}, &reply) + Expect(errors.Is(err, nodes.ErrWorkerUnroutable)).To(BeTrue()) + Expect(errors.Is(err, nodes.ErrWorkerControlUnsupported)).To(BeFalse()) + Expect(cluster.IsWorkerAnswer(err)).To(BeFalse()) + }) +}) diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index 09b485c65..865b2b7af 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -17,7 +17,7 @@ Distributed mode requires authentication enabled with a **PostgreSQL** database **Frontends** are stateless LocalAI instances that receive API requests and route them to worker nodes via the **SmartRouter**. All frontends share state through PostgreSQL and coordinate via NATS. -**Workers** are generic processes that self-register with a frontend. They don't have a fixed backend type - the SmartRouter dynamically installs the required backend via NATS `backend.install` events when a model request arrives. +**Workers** are generic processes that self-register with a frontend. They don't have a fixed backend type - the SmartRouter dynamically installs the required backend by calling the worker's `backend.install` control route through its tunnel when a model request arrives. ### Scheduling Algorithm @@ -163,14 +163,14 @@ Reconnects use exponential backoff with jitter: the interval doubles from 500ms #### What the frontend sends through it -Every connection the frontend makes to a worker now goes through that worker's tunnel. There are three, and all three are the same path underneath: +Every connection the frontend makes to a worker now goes through that worker's tunnel. There are four, and all four are the same path underneath: | What | Protocol | Stream tag | |------|----------|-----------| | Inference, model load, health checks | gRPC to a backend process | `grpc` | | Model file staging, backend-log listing | HTTP to the worker's own server | `http` | | Live backend-log streaming | WebSocket to the same server | `http` | -| Backend install/upgrade/list/delete, model stop/unload/delete, node stop | HTTP to the worker's own server | `http` | +| The control plane: backend install/upgrade/list/stop/delete, model stop/unload/delete, models running, node stop | HTTP to the worker's own server | `http` | #### The worker control plane @@ -203,6 +203,32 @@ does not impose a timeout of its own on an install, and a caller that gives up cancels the download rather than leaving the worker pulling gigabytes for a response nobody will read. +On the frontend's side the ten verbs are ordinary HTTP calls on the `http` +stream tag, so a control RPC to a worker **another replica holds is relayed +exactly like an inference request** - same lookup, same one hop, same budget +arithmetic as [Reaching a worker another replica holds](#reaching-a-worker-another-replica-holds). +There is nothing to subscribe to before an install: its progress lines share the +install's own response, so no event can arrive before the caller is listening +and there is no per-op subject to grant a permission for. + +How a control RPC can FAIL is where absence is decided for the whole control +plane, so the frontend maps every outcome onto exactly one row of the table +below and never onto another. Only the two rows in which the WORKER spoke may +be acted on; everything else is this frontend failing to reach it, which says +nothing about the worker at all: + +| What happened | How the frontend reports it | May anything reap on it? | +|---|---|---| +| The worker refused the stream with one of its three evidence codes | the refusal itself, unwrapped | **Yes.** The worker spoke. | +| No route: no live owner, an unreachable peer, no relay path, a refusal code this frontend does not recognise, or the worker saying it learned nothing | "this frontend has no route to that worker" | No | +| The call ran out of budget | a deadline, which install and upgrade report as *still installing on the worker* | No | +| `404` under `/v1/control/` | "the worker does not serve that control verb" - it is older than this frontend, and only the upgrade path acts on it, by re-issuing the legacy force-install | No | +| `200` with a reply whose `error` is set | the worker's own answer, handed to the caller as-is | **Yes**, by the caller | + +A `5xx`, or a body the frontend cannot decode, is in the second row and not the +last: a worker's verdict arrives as a `200`, so a `5xx` is the server failing +rather than answering. + The address the frontend holds for a backend (the per-replica port a worker reports after an install) is still what identifies it, and it is still what appears in logs and errors. What it no longer is, is somewhere the frontend connects to: it travels inside the tunnel as the stream's target, and the worker decides what to do with it. A frontend with no way to reach a worker says so and fails. It does **not** fall back to connecting to the worker's advertised address. That fallback is what the tunnel exists to remove, and it is the kind of defect that works on a one-replica developer box and fails in production, so it is an error everywhere. The consequences are deliberately narrow: a model whose worker cannot be reached is not reaped, and its row is left alone, because a frontend that cannot reach a worker has learned nothing about whether that worker is still running the model. @@ -651,7 +677,7 @@ The system automatically applies hardware-detected labels on registration: ### How Workers Operate -Workers start as generic processes with no backend installed. When the SmartRouter needs to load a model on a worker, it sends a NATS `backend.install` event with the backend name and model ID. The worker: +Workers start as generic processes with no backend installed. When the SmartRouter needs to load a model on a worker, it calls `POST /v1/control/backend/install` through that worker's tunnel with the backend name and model ID. The worker: 1. Installs the backend from the gallery (if not already installed) 2. Starts a **new gRPC backend process on a dynamic port** (each model gets its own process) @@ -1276,7 +1302,7 @@ Notes: | **Coordination** | Gossip protocol | NATS messaging | | **Node management** | Automatic | REST API + WebUI | | **Health monitoring** | Peer heartbeats | Centralized HealthMonitor | -| **Backend management** | Manual per node | Dynamic via NATS backend.install | +| **Backend management** | Manual per node | Dynamic via the worker's `backend.install` control route | | **Best for** | Ad-hoc clusters, community sharing | Production, Kubernetes, managed infrastructure | | **Setup complexity** | Minimal (share a token) | Requires PostgreSQL + NATS | @@ -1319,11 +1345,11 @@ Notes: - Liveness is decided by the load job's progress heartbeat, not by elapsed time. Staging a large checkpoint legitimately runs for a long time without touching the replica row, so a transfer that is still progressing is never reclaimed however long it takes. - `Reconciler: reclaimed a replica slot held by a load nobody is driving` names each row reclaimed this way. -**A request fails with `nats: no responders available for request`:** -- The chosen worker was not subscribed on the bus when the frontend tried to install the backend on it. A node's status comes from its HTTP heartbeat, which is a separate channel: a worker that stops stays `healthy` until that heartbeat ages out. -- The scheduler now checks that a node still answers on the bus before it commits to it, marks one that does not as unhealthy, and picks another. A request should therefore see this only when no reachable node is left. -- Only a no-responders answer counts as absent. A worker that answers slowly stays eligible, because excluding it would cost capacity that is really there. -- Check the worker process is running and its NATS connection is up. `Scheduled node is not answering on the bus` in the frontend log names each node demoted this way. +**A request fails with `this frontend has no route to that worker`:** +- The chosen worker's tunnel was not reachable from the replica that handled the request. A node's status comes from its HTTP heartbeat, which is a separate channel: a worker that stops stays `healthy` until that heartbeat ages out, and a worker that is very much alive can be unroutable for a moment while its tunnel re-homes between frontend replicas. +- It is **not** the same as the worker being gone, and nothing acts on it as if it were. A model on an unroutable worker is not reaped, its rows are left alone, and the node is not demoted: doing any of those on a lost route is how a rolling frontend restart turns into a fleet-wide eviction. +- Check the worker process is running and that it has an open tunnel (`opened a tunnelled stream to a worker` in the frontend log, and the worker's own dial/reconnect lines). A worker behind a load balancer that keeps reconnecting is usually an idle-timeout or WebSocket-upgrade problem at the proxy; see the tunnel section above. +- A `nats: no responders available for request` in a modern deployment concerns only the subjects still on the bus: agent-worker jobs, MCP, and file staging. It is no longer how a serve-backend worker is reached. **A worker fills its own disk over time:** - A request that carries a file (an image, an audio clip, a video) stages that file to the worker under `/../staging/ephemeral/`. The worker deletes these 6 hours after the request that needed them, and sweeps every 30 minutes plus once at startup, so a worker that crashed mid-request still reclaims the space. diff --git a/pkg/natsauth/mint_test.go b/pkg/natsauth/mint_test.go index 4d5bb77ce..e33992ec1 100644 --- a/pkg/natsauth/mint_test.go +++ b/pkg/natsauth/mint_test.go @@ -37,7 +37,13 @@ var _ = Describe("MintWorkerJWT", func() { uc, err := jwt.DecodeUserClaims(token) Expect(err).NotTo(HaveOccurred()) Expect(uc.Permissions.Sub.Allow).To(ContainElement("nodes.550e8400-e29b-41d4-a716-446655440000.>")) - Expect(uc.Permissions.Pub.Allow).To(ContainElement("nodes.550e8400-e29b-41d4-a716-446655440000.backend.install.*.progress")) + Expect(uc.Permissions.Pub.Allow).To(ContainElement("nodes.550e8400-e29b-41d4-a716-446655440000.files.>")) + // The install-progress subject is gone with the carrier: progress is a + // line in the install response now, so a minted worker JWT must not + // still be granted a publish right for it. + for _, subj := range uc.Permissions.Pub.Allow { + Expect(subj).NotTo(ContainSubstring("backend.install")) + } }) It("mints agent permissions without backend install subscribe", func() { diff --git a/pkg/natsauth/permissions.go b/pkg/natsauth/permissions.go index d44e51990..a411e05fd 100644 --- a/pkg/natsauth/permissions.go +++ b/pkg/natsauth/permissions.go @@ -38,13 +38,21 @@ func WorkerPermissions(nodeID, nodeType string) (pubAllow, subAllow []string) { "_INBOX.>", } default: - // Backend worker: lifecycle + file staging on this node only. + // Backend worker: file staging on this node only. + // + // The backend and model lifecycle verbs left the bus: they are HTTP + // routes under workerctl.Prefix, served on the worker's own server and + // reached through its tunnel, so no subject is minted for them and none + // is allowed here. The wildcard stays because the file-staging subjects + // are still under this node's prefix; narrowing it to nodes..files.> + // is a separate change that would break a worker mid-upgrade. subAllow = []string{ prefix + ".>", "_INBOX.>", } + // backend.install.*.progress is gone with the subject: install progress + // is written into the install response the frontend is already reading. pubAllow = []string{ - prefix + ".backend.install.*.progress", prefix + ".files.>", "_INBOX.>", } diff --git a/pkg/natsauth/permissions_coverage_test.go b/pkg/natsauth/permissions_coverage_test.go index 05d2fbf0b..b02541e9a 100644 --- a/pkg/natsauth/permissions_coverage_test.go +++ b/pkg/natsauth/permissions_coverage_test.go @@ -33,6 +33,13 @@ func subjectMatches(pattern, subject string) bool { return len(p) == len(s) } +// workerSubjectTokenForTest mirrors the sanitizer both packages implement, so +// the negative assertion below names the exact prefix without reaching into +// either package's unexported copy. +func workerSubjectTokenForTest(nodeID string) string { + return strings.NewReplacer(".", "-", "*", "-", ">", "-", " ", "-", "\t", "-", "\n", "-").Replace(nodeID) +} + func anyAllows(allow []string, subject string) bool { for _, p := range allow { if subjectMatches(p, subject) { @@ -52,16 +59,11 @@ var _ = Describe("WorkerPermissions subject coverage", func() { Context("backend worker", func() { pub, sub := natsauth.WorkerPermissions(nodeID, "backend") - // Every subject core/services/worker/{lifecycle,file_staging}.go subscribes to. + // Every subject core/services/worker/file_staging.go subscribes to. + // The backend and model lifecycle verbs are NOT here: they left the bus + // for the worker's tunnelled control plane, so there is no subject to + // cover. See core/services/workerctl. subscribed := []string{ - messaging.SubjectNodeBackendInstall(nodeID), - messaging.SubjectNodeBackendUpgrade(nodeID), - messaging.SubjectNodeBackendStop(nodeID), - messaging.SubjectNodeBackendDelete(nodeID), - messaging.SubjectNodeBackendList(nodeID), - messaging.SubjectNodeModelUnload(nodeID), - messaging.SubjectNodeModelDelete(nodeID), - messaging.SubjectNodeStop(nodeID), messaging.SubjectNodeFilesEnsure(nodeID), messaging.SubjectNodeFilesStage(nodeID), messaging.SubjectNodeFilesTemp(nodeID), @@ -74,11 +76,21 @@ var _ = Describe("WorkerPermissions subject coverage", func() { }) } - It("allows publishing backend.install progress", func() { - subject := messaging.SubjectNodeBackendInstallProgress(nodeID, "op-123") + It("allows publishing file staging replies", func() { + subject := messaging.SubjectNodeFilesStage(nodeID) Expect(anyAllows(pub, subject)).To(BeTrue(), "backend JWT pub allow-list %v does not cover %s", pub, subject) }) + + // The negative half, and it is the one that would catch a verb quietly + // coming back to the bus: a backend worker is granted nothing to + // publish outside its own file-staging subtree and its inbox. + It("grants a backend worker no publish rights outside file staging and its inbox", func() { + Expect(pub).To(ConsistOf( + "nodes."+workerSubjectTokenForTest(nodeID)+".files.>", + "_INBOX.>", + )) + }) }) Context("agent worker", func() { @@ -115,7 +127,7 @@ var _ = Describe("Documented NATS service-user permissions", func() { frontendPublishes := []string{ messaging.SubjectPrefixCacheObserve, messaging.SubjectPrefixCacheInvalidate, - messaging.SubjectNodeBackendInstall("node-1"), + messaging.SubjectNodeBackendStop("node-1"), messaging.SubjectGalleryProgress("op-1"), } diff --git a/tests/e2e/distributed/control_workers_test.go b/tests/e2e/distributed/control_workers_test.go new file mode 100644 index 000000000..15ef583b6 --- /dev/null +++ b/tests/e2e/distributed/control_workers_test.go @@ -0,0 +1,112 @@ +package distributed_test + +import ( + "context" + "encoding/json" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/workerctl" +) + +// ControlWorkers is a fleet of fake workers serving the tunnelled control +// plane, which is where the frontend's backend and model lifecycle verbs go +// now that they have left the bus. +// +// It replaces the NATS SubscribeReply fakes these suites used to stand up. One +// HTTP server answers for every node and reads which node a request is for off +// the Host header, because that is what nodes.ControlClient puts there; a +// client addressing the wrong node would be visible rather than silently +// served. +// +// It is deliberately NOT a real worker: these suites are about what the +// frontend does with a worker's answer. The real handlers are exercised against +// the real client in core/services/worker. +type ControlWorkers struct { + mu sync.Mutex + srv *httptest.Server + handlers map[string]func(nodeID string, body []byte) any +} + +// controlHandlerKey is the (node, verb) a handler is registered for. AnyNode +// registers one handler for every node, which is how a suite fakes a fleet +// whose ids it does not know up front. +const AnyNode = "*" + +func controlHandlerKey(nodeID, path string) string { return nodeID + " " + path } + +// NewControlWorkers starts the fleet and stops it when the spec ends. +func NewControlWorkers() *ControlWorkers { + c := &ControlWorkers{handlers: map[string]func(string, []byte) any{}} + mux := http.NewServeMux() + mux.HandleFunc(workerctl.Prefix, c.serve) + c.srv = httptest.NewServer(mux) + DeferCleanup(c.srv.Close) + return c +} + +// Client returns a control client that reaches this fleet. +func (c *ControlWorkers) Client() *nodes.ControlClient { + addr := c.srv.Listener.Addr().String() + return nodes.NewControlClient(func(string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, _, _ string) (net.Conn, error) { + var d net.Dialer + return d.DialContext(ctx, "tcp", addr) + } + }, "") +} + +// On registers what one node answers for one verb. Returning nil answers 204, +// which is the shape the fire-and-forget verbs take. +func (c *ControlWorkers) On(nodeID, path string, fn func(nodeID string, body []byte) any) { + c.mu.Lock() + defer c.mu.Unlock() + c.handlers[controlHandlerKey(nodeID, path)] = fn +} + +func (c *ControlWorkers) serve(w http.ResponseWriter, r *http.Request) { + body, err := io.ReadAll(r.Body) + Expect(err).ToNot(HaveOccurred()) + nodeID := strings.TrimSuffix(r.Host, ".worker.invalid:80") + + c.mu.Lock() + fn, ok := c.handlers[controlHandlerKey(nodeID, r.URL.Path)] + if !ok { + fn, ok = c.handlers[controlHandlerKey(AnyNode, r.URL.Path)] + } + c.mu.Unlock() + + if !ok { + // Loud, never plausible: a verb no spec scripted must be a red spec + // rather than a worker that looks absent. + http.Error(w, "no control handler registered for "+r.URL.Path+" on "+nodeID, http.StatusInternalServerError) + return + } + + reply := fn(nodeID, body) + streaming := r.URL.Path == workerctl.PathBackendInstall || r.URL.Path == workerctl.PathBackendUpgrade + if !streaming { + if reply == nil { + w.WriteHeader(http.StatusNoContent) + return + } + w.Header().Set("Content-Type", "application/json") + Expect(json.NewEncoder(w).Encode(reply)).To(Succeed()) + return + } + + raw, err := json.Marshal(reply) + Expect(err).ToNot(HaveOccurred()) + w.Header().Set("Content-Type", workerctl.ContentTypeStream) + w.WriteHeader(http.StatusOK) + // The reply line, and it is the last thing on the body by contract. + Expect(json.NewEncoder(w).Encode(workerctl.Envelope{Reply: raw})).To(Succeed()) +} diff --git a/tests/e2e/distributed/distributed_full_flow_test.go b/tests/e2e/distributed/distributed_full_flow_test.go index de2657379..abf0cbb72 100644 --- a/tests/e2e/distributed/distributed_full_flow_test.go +++ b/tests/e2e/distributed/distributed_full_flow_test.go @@ -2,7 +2,6 @@ package distributed_test import ( "context" - "encoding/json" "fmt" "io" "net" @@ -13,6 +12,7 @@ import ( "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/workerctl" "github.com/mudler/LocalAI/pkg/grpc/base" pb "github.com/mudler/LocalAI/pkg/grpc/proto" @@ -21,7 +21,6 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" - "github.com/nats-io/nats.go" "google.golang.org/grpc" pgdriver "gorm.io/driver/postgres" @@ -225,10 +224,21 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() cancel() }) - // newTestSmartRouter creates a SmartRouter with NATS wired up and a mock - // backend.install handler that always replies success for all registered nodes. + // newTestSmartRouter creates a SmartRouter reaching a fleet of fake workers + // over the tunnelled control plane, with a backend.install handler that + // always replies success for every registered node. newTestSmartRouter := func(reg *nodes.NodeRegistry, extraOpts ...nodes.SmartRouterOptions) *nodes.SmartRouter { - unloader := nodes.NewRemoteUnloaderAdapter(reg, infra.NC, 3*time.Minute, 15*time.Minute) + workers := NewControlWorkers() + workers.On(AnyNode, workerctl.PathBackendInstall, func(string, []byte) any { + return messaging.BackendInstallReply{Success: true} + }) + workers.On(AnyNode, workerctl.PathModelsRunning, func(string, []byte) any { + return messaging.ModelsRunningReply{} + }) + workers.On(AnyNode, workerctl.PathBackendList, func(string, []byte) any { + return messaging.BackendListReply{} + }) + unloader := nodes.NewRemoteUnloaderAdapter(reg, infra.NC, workers.Client(), 3*time.Minute, 15*time.Minute) opts := nodes.SmartRouterOptions{ Unloader: unloader, @@ -251,22 +261,6 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() router := nodes.NewSmartRouter(reg, opts) - // Subscribe a mock backend.install handler that replies success for any node. - // We use a wildcard-style approach: subscribe to all nodes' install subjects - // by registering after each node. In practice, we rely on the test registering - // nodes before calling Route, so we subscribe to a catch-all pattern. - infra.NC.Conn().Subscribe("nodes.*.backend.install", func(msg *nats.Msg) { - reply := messaging.BackendInstallReply{Success: true} - data, _ := json.Marshal(reply) - msg.Respond(data) - }) - _, err := infra.NC.Conn().Subscribe("nodes.*.models.running", func(msg *nats.Msg) { - data, _ := json.Marshal(messaging.ModelsRunningReply{}) - _ = msg.Respond(data) - }) - Expect(err).NotTo(HaveOccurred()) - FlushNATS(infra.NC) - return router } // suppress unused warning in case some tests don't call it @@ -385,30 +379,25 @@ var _ = Describe("Full Distributed Inference Flow", Label("Distributed"), func() result.Release() }) - It("should unload remote model via NATS", func() { + It("should unload a remote model over the worker's control plane", func() { // Register a node with a loaded model node := &nodes.BackendNode{Name: "gpu-unload", Address: "127.0.0.1:50099"} Expect(registry.Register(context.Background(), node, true)).To(Succeed()) Expect(registry.SetNodeModel(context.Background(), node.ID, "old-model", 0, "loaded", "", 0)).To(Succeed()) - // Subscribe to NATS backend.stop for this node - stopSubject := messaging.SubjectNodeBackendStop(node.ID) + // A worker serving backend.stop on its own control plane. received := make(chan struct{}, 1) - rawConn, err := nats.Connect(infra.NatsURL) - Expect(err).ToNot(HaveOccurred()) - defer rawConn.Close() - - _, err = rawConn.Subscribe(stopSubject, func(msg *nats.Msg) { + workers := NewControlWorkers() + workers.On(node.ID, workerctl.PathBackendStop, func(string, []byte) any { received <- struct{}{} + return nil }) - Expect(err).ToNot(HaveOccurred()) // Create RemoteUnloaderAdapter and unload model - unloader := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) - err = unloader.UnloadRemoteModel("old-model") - Expect(err).ToNot(HaveOccurred()) + unloader := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, workers.Client(), 3*time.Minute, 15*time.Minute) + Expect(unloader.UnloadRemoteModel("old-model")).To(Succeed()) - // Verify NATS event received + // The worker got the stop over its tunnel, not over the bus. Eventually(received, 5*time.Second).Should(Receive()) // Verify model removed from registry diff --git a/tests/e2e/distributed/managers_test.go b/tests/e2e/distributed/managers_test.go index b4f51ef95..83080f0d7 100644 --- a/tests/e2e/distributed/managers_test.go +++ b/tests/e2e/distributed/managers_test.go @@ -12,6 +12,7 @@ import ( "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/workerctl" "github.com/mudler/LocalAI/pkg/model" "github.com/mudler/LocalAI/pkg/system" @@ -136,30 +137,18 @@ var _ = Describe("Model and Backend Managers", Label("Distributed"), func() { Expect(registry.SetNodeModel(context.Background(), node1.ID, "big-model", 0, "loaded", "", 0)).To(Succeed()) Expect(registry.SetNodeModel(context.Background(), node2.ID, "big-model", 0, "loaded", "", 0)).To(Succeed()) - // Subscribe to model.delete on both node subjects, track receipt + // Both workers serve model.delete on their own control plane. var deleteCount atomic.Int32 - sub1, err := infra.NC.SubscribeReply(messaging.SubjectNodeModelDelete(node1.ID), func(data []byte, reply func([]byte)) { - var req messaging.ModelDeleteRequest - json.Unmarshal(data, &req) - Expect(req.ModelName).To(Equal("big-model")) - deleteCount.Add(1) - resp, _ := json.Marshal(messaging.ModelDeleteReply{Success: true}) - reply(resp) - }) - Expect(err).ToNot(HaveOccurred()) - defer sub1.Unsubscribe() - - sub2, err := infra.NC.SubscribeReply(messaging.SubjectNodeModelDelete(node2.ID), func(data []byte, reply func([]byte)) { - var req messaging.ModelDeleteRequest - json.Unmarshal(data, &req) - deleteCount.Add(1) - resp, _ := json.Marshal(messaging.ModelDeleteReply{Success: true}) - reply(resp) - }) - Expect(err).ToNot(HaveOccurred()) - defer sub2.Unsubscribe() - - FlushNATS(infra.NC) + workers := NewControlWorkers() + for _, id := range []string{node1.ID, node2.ID} { + workers.On(id, workerctl.PathModelDelete, func(_ string, data []byte) any { + var req messaging.ModelDeleteRequest + Expect(json.Unmarshal(data, &req)).To(Succeed()) + Expect(req.ModelName).To(Equal("big-model")) + deleteCount.Add(1) + return messaging.ModelDeleteReply{Success: true} + }) + } // Create temp dir for local model files tempDir, err := os.MkdirTemp("", "dist-model-test-*") @@ -176,7 +165,7 @@ var _ = Describe("Model and Backend Managers", Label("Distributed"), func() { appCfg := config.NewApplicationConfig() appCfg.SystemState = ss - adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) + adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, workers.Client(), 3*time.Minute, 15*time.Minute) distMgr := nodes.NewDistributedModelManager(appCfg, ml, adapter) err = distMgr.DeleteModel("big-model") @@ -202,39 +191,24 @@ var _ = Describe("Model and Backend Managers", Label("Distributed"), func() { Expect(registry.Register(context.Background(), node3, true)).To(Succeed()) Expect(registry.MarkUnhealthy(context.Background(), node3.ID)).To(Succeed()) - // Subscribe to backend.delete on all 3 nodes + // All 3 workers serve backend.delete on their control plane. var deleteCount atomic.Int32 - sub1, err := infra.NC.SubscribeReply(messaging.SubjectNodeBackendDelete(node1.ID), func(data []byte, reply func([]byte)) { - var req messaging.BackendDeleteRequest - json.Unmarshal(data, &req) - Expect(req.Backend).To(Equal("my-backend")) - deleteCount.Add(1) - resp, _ := json.Marshal(messaging.BackendDeleteReply{Success: true}) - reply(resp) - }) - Expect(err).ToNot(HaveOccurred()) - defer sub1.Unsubscribe() - - sub2, err := infra.NC.SubscribeReply(messaging.SubjectNodeBackendDelete(node2.ID), func(data []byte, reply func([]byte)) { - var req messaging.BackendDeleteRequest - json.Unmarshal(data, &req) - deleteCount.Add(1) - resp, _ := json.Marshal(messaging.BackendDeleteReply{Success: true}) - reply(resp) - }) - Expect(err).ToNot(HaveOccurred()) - defer sub2.Unsubscribe() + workers := NewControlWorkers() + for _, id := range []string{node1.ID, node2.ID} { + workers.On(id, workerctl.PathBackendDelete, func(_ string, data []byte) any { + var req messaging.BackendDeleteRequest + Expect(json.Unmarshal(data, &req)).To(Succeed()) + Expect(req.Backend).To(Equal("my-backend")) + deleteCount.Add(1) + return messaging.BackendDeleteReply{Success: true} + }) + } var unhealthyReceived atomic.Int32 - sub3, err := infra.NC.SubscribeReply(messaging.SubjectNodeBackendDelete(node3.ID), func(data []byte, reply func([]byte)) { + workers.On(node3.ID, workerctl.PathBackendDelete, func(string, []byte) any { unhealthyReceived.Add(1) - resp, _ := json.Marshal(messaging.BackendDeleteReply{Success: true}) - reply(resp) + return messaging.BackendDeleteReply{Success: true} }) - Expect(err).ToNot(HaveOccurred()) - defer sub3.Unsubscribe() - - FlushNATS(infra.NC) // Create temp dir for local backend files tempDir, err := os.MkdirTemp("", "dist-backend-test-*") @@ -252,7 +226,7 @@ var _ = Describe("Model and Backend Managers", Label("Distributed"), func() { appCfg := config.NewApplicationConfig() appCfg.SystemState = ss - adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) + adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, workers.Client(), 3*time.Minute, 15*time.Minute) distMgr := nodes.NewDistributedBackendManager(appCfg, ml, adapter, registry, nil) err = distMgr.DeleteBackend("my-backend") @@ -275,18 +249,14 @@ var _ = Describe("Model and Backend Managers", Label("Distributed"), func() { Expect(registry.Register(context.Background(), node1, true)).To(Succeed()) var deleteCount atomic.Int32 - sub1, err := infra.NC.SubscribeReply(messaging.SubjectNodeBackendDelete(node1.ID), func(data []byte, reply func([]byte)) { + workers := NewControlWorkers() + workers.On(node1.ID, workerctl.PathBackendDelete, func(_ string, data []byte) any { var req messaging.BackendDeleteRequest - json.Unmarshal(data, &req) + Expect(json.Unmarshal(data, &req)).To(Succeed()) Expect(req.Backend).To(Equal("remote-only-backend")) deleteCount.Add(1) - resp, _ := json.Marshal(messaging.BackendDeleteReply{Success: true}) - reply(resp) + return messaging.BackendDeleteReply{Success: true} }) - Expect(err).ToNot(HaveOccurred()) - defer sub1.Unsubscribe() - - FlushNATS(infra.NC) // Use a temp dir with NO local backend directory — simulates frontend node tempDir, err := os.MkdirTemp("", "dist-backend-remote-only-*") @@ -299,7 +269,7 @@ var _ = Describe("Model and Backend Managers", Label("Distributed"), func() { appCfg := config.NewApplicationConfig() appCfg.SystemState = ss - adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) + adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, workers.Client(), 3*time.Minute, 15*time.Minute) distMgr := nodes.NewDistributedBackendManager(appCfg, ml, adapter, registry, nil) // Should NOT return an error even though the backend doesn't exist locally diff --git a/tests/e2e/distributed/node_lifecycle_test.go b/tests/e2e/distributed/node_lifecycle_test.go index 04b7342e7..e3fdb1eef 100644 --- a/tests/e2e/distributed/node_lifecycle_test.go +++ b/tests/e2e/distributed/node_lifecycle_test.go @@ -8,6 +8,7 @@ import ( "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/workerctl" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" @@ -17,7 +18,7 @@ import ( "gorm.io/gorm/logger" ) -var _ = Describe("Node Backend Lifecycle (NATS-driven)", Label("Distributed"), func() { +var _ = Describe("Node Backend Lifecycle over the worker control plane", Label("Distributed"), func() { var ( infra *TestInfra db *gorm.DB @@ -37,27 +38,23 @@ var _ = Describe("Node Backend Lifecycle (NATS-driven)", Label("Distributed"), f Expect(err).ToNot(HaveOccurred()) }) - Context("NATS backend.install events", func() { - It("should send backend.install request-reply to a specific node", func() { + Context("backend.install", func() { + It("should send backend.install to a specific node", func() { node := &nodes.BackendNode{ Name: "gpu-node-1", Address: "h1:50051", } Expect(registry.Register(context.Background(), node, true)).To(Succeed()) - // Simulate worker subscribing to backend.install and replying success - infra.NC.SubscribeReply(messaging.SubjectNodeBackendInstall(node.ID), func(data []byte, reply func([]byte)) { + // The worker serves backend.install on its own control plane. + workers := NewControlWorkers() + workers.On(node.ID, workerctl.PathBackendInstall, func(_ string, data []byte) any { var req messaging.BackendInstallRequest - json.Unmarshal(data, &req) + Expect(json.Unmarshal(data, &req)).To(Succeed()) Expect(req.Backend).To(Equal("llama-cpp")) - - resp := messaging.BackendInstallReply{Success: true} - respData, _ := json.Marshal(resp) - reply(respData) + return messaging.BackendInstallReply{Success: true} }) - FlushNATS(infra.NC) - - adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) + adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, workers.Client(), 3*time.Minute, 15*time.Minute) installReply, err := adapter.InstallBackend(node.ID, "llama-cpp", "", "", "", "", "", 0, "", nil) Expect(err).ToNot(HaveOccurred()) Expect(installReply.Success).To(BeTrue()) @@ -69,16 +66,13 @@ var _ = Describe("Node Backend Lifecycle (NATS-driven)", Label("Distributed"), f } Expect(registry.Register(context.Background(), node, true)).To(Succeed()) - // Simulate worker replying with error - infra.NC.SubscribeReply(messaging.SubjectNodeBackendInstall(node.ID), func(data []byte, reply func([]byte)) { - resp := messaging.BackendInstallReply{Success: false, Error: "backend not found"} - respData, _ := json.Marshal(resp) - reply(respData) + // The worker's own verdict: an answer, not a transport failure. + workers := NewControlWorkers() + workers.On(node.ID, workerctl.PathBackendInstall, func(string, []byte) any { + return messaging.BackendInstallReply{Success: false, Error: "backend not found"} }) - FlushNATS(infra.NC) - - adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) + adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, workers.Client(), 3*time.Minute, 15*time.Minute) installReply, err := adapter.InstallBackend(node.ID, "nonexistent", "", "", "", "", "", 0, "", nil) Expect(err).ToNot(HaveOccurred()) Expect(installReply.Success).To(BeFalse()) @@ -86,7 +80,7 @@ var _ = Describe("Node Backend Lifecycle (NATS-driven)", Label("Distributed"), f }) }) - Context("NATS backend.stop events (model unload)", func() { + Context("backend.stop (model unload)", func() { It("should send backend.stop to nodes hosting the model", func() { node := &nodes.BackendNode{ Name: "gpu-node-2", Address: "h2:50051", @@ -95,16 +89,14 @@ var _ = Describe("Node Backend Lifecycle (NATS-driven)", Label("Distributed"), f Expect(registry.SetNodeModel(context.Background(), node.ID, "whisper-large", 0, "loaded", "", 0)).To(Succeed()) var stopReceived atomic.Int32 - sub, err := infra.NC.Subscribe(messaging.SubjectNodeBackendStop(node.ID), func(data []byte) { + workers := NewControlWorkers() + workers.On(node.ID, workerctl.PathBackendStop, func(string, []byte) any { stopReceived.Add(1) + return nil }) - Expect(err).ToNot(HaveOccurred()) - defer sub.Unsubscribe() - - FlushNATS(infra.NC) // Frontend calls UnloadRemoteModel (triggered by UI "Stop" or WatchDog) - adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) + adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, workers.Client(), 3*time.Minute, 15*time.Minute) Expect(adapter.UnloadRemoteModel("whisper-large")).To(Succeed()) Eventually(func() int32 { return stopReceived.Load() }, "5s").Should(Equal(int32(1))) @@ -123,18 +115,15 @@ var _ = Describe("Node Backend Lifecycle (NATS-driven)", Label("Distributed"), f registry.SetNodeModel(context.Background(), node2.ID, "shared-model", 0, "loaded", "", 0) var count atomic.Int32 - sub1, _ := infra.NC.Subscribe(messaging.SubjectNodeBackendStop(node1.ID), func(data []byte) { - count.Add(1) - }) - sub2, _ := infra.NC.Subscribe(messaging.SubjectNodeBackendStop(node2.ID), func(data []byte) { - count.Add(1) - }) - defer sub1.Unsubscribe() - defer sub2.Unsubscribe() + workers := NewControlWorkers() + for _, id := range []string{node1.ID, node2.ID} { + workers.On(id, workerctl.PathBackendStop, func(string, []byte) any { + count.Add(1) + return nil + }) + } - FlushNATS(infra.NC) - - adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) + adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, workers.Client(), 3*time.Minute, 15*time.Minute) adapter.UnloadRemoteModel("shared-model") Eventually(func() int32 { return count.Load() }, "5s").Should(Equal(int32(2))) @@ -150,49 +139,57 @@ var _ = Describe("Node Backend Lifecycle (NATS-driven)", Label("Distributed"), f // 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) + adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, NewControlWorkers().Client(), 3*time.Minute, 15*time.Minute) Expect(adapter.UnloadRemoteModel("nonexistent-model")).To(Succeed()) }) }) - Context("NATS node stop events (full shutdown)", func() { - It("should publish stop event to a node", func() { + Context("node.stop (full shutdown)", func() { + It("should ask the node to shut down over its control plane", func() { node := &nodes.BackendNode{ Name: "stop-me", Address: "h3:50051", } Expect(registry.Register(context.Background(), node, true)).To(Succeed()) var stopped atomic.Int32 - sub, err := infra.NC.Subscribe(messaging.SubjectNodeStop(node.ID), func(data []byte) { + workers := NewControlWorkers() + workers.On(node.ID, workerctl.PathNodeStop, func(string, []byte) any { stopped.Add(1) + return nil }) - Expect(err).ToNot(HaveOccurred()) - defer sub.Unsubscribe() - FlushNATS(infra.NC) - - adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) + adapter := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, workers.Client(), 3*time.Minute, 15*time.Minute) Expect(adapter.StopNode(node.ID)).To(Succeed()) Eventually(func() int32 { return stopped.Load() }, "5s").Should(Equal(int32(1))) }) }) - Context("NATS subject naming", func() { - It("should generate correct backend lifecycle subjects", func() { - Expect(messaging.SubjectNodeBackendInstall("node-abc")).To(Equal("nodes.node-abc.backend.install")) + Context("wire naming", func() { + // Written out BY HAND, and not derived from the constants: a frontend + // and a worker built from different commits reach each other over these + // literals, and a renamed path is a 404 that looks exactly like a + // broken tunnel. + It("should name the backend lifecycle control verbs", func() { + Expect(workerctl.PathBackendInstall).To(Equal("/v1/control/backend/install")) + Expect(workerctl.PathBackendStop).To(Equal("/v1/control/backend/stop")) + Expect(workerctl.PathNodeStop).To(Equal("/v1/control/node/stop")) + }) + + // The one node subject left, and it is addressed only to AGENT workers: + // they hold no tunnel and subscribe to it to drop cached MCP sessions. + It("should keep the agent worker's backend.stop subject", func() { Expect(messaging.SubjectNodeBackendStop("node-abc")).To(Equal("nodes.node-abc.backend.stop")) - Expect(messaging.SubjectNodeStop("node-abc")).To(Equal("nodes.node-abc.stop")) }) }) - // Design note: LoadModel is a direct gRPC call to node.Address, NOT a NATS event. - // NATS is used for backend.install (install + start process) and backend.stop. - // The SmartRouter calls grpc.NewClient(node.Address).LoadModel() directly. + // Design note: LoadModel is a gRPC call through the worker's tunnel, not a + // control verb. The control plane installs and stops the process; the model + // is loaded into it over the `grpc` stream tag. // // Flow: - // 1. NATS backend.install → worker installs backend + starts gRPC process - // 2. SmartRouter.Route() → gRPC LoadModel(node.Address) directly - // 3. [inference via gRPC] - // 4. NATS backend.stop → worker stops gRPC process + // 1. backend.install → worker installs backend + starts gRPC process + // 2. SmartRouter.Route() → LoadModel over the worker's tunnel + // 3. [inference over the tunnel] + // 4. backend.stop → worker stops gRPC process }) diff --git a/tests/e2e/distributed/router_tracking_test.go b/tests/e2e/distributed/router_tracking_test.go index 75895a372..6a4b545e0 100644 --- a/tests/e2e/distributed/router_tracking_test.go +++ b/tests/e2e/distributed/router_tracking_test.go @@ -2,11 +2,11 @@ package distributed_test import ( "context" - "encoding/json" "time" "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/workerctl" "github.com/mudler/LocalAI/pkg/grpc/base" pb "github.com/mudler/LocalAI/pkg/grpc/proto" @@ -15,8 +15,6 @@ import ( . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" - "github.com/nats-io/nats.go" - pgdriver "gorm.io/driver/postgres" gormDB "gorm.io/gorm" "gorm.io/gorm/logger" @@ -60,18 +58,17 @@ var _ = Describe("SmartRouter trackingKey", Label("Distributed"), func() { registry, err = nodes.NewNodeRegistry(db) Expect(err).ToNot(HaveOccurred()) - // Mock backend.install handler — always replies success - infra.NC.Conn().Subscribe("nodes.*.backend.install", func(msg *nats.Msg) { - reply := messaging.BackendInstallReply{Success: true} - data, _ := json.Marshal(reply) - msg.Respond(data) + // Mock control plane — backend.install always replies success. + workers := NewControlWorkers() + workers.On(AnyNode, workerctl.PathBackendInstall, func(string, []byte) any { + return messaging.BackendInstallReply{Success: true} }) - _, err = infra.NC.Conn().Subscribe("nodes.*.models.running", func(msg *nats.Msg) { - data, _ := json.Marshal(messaging.ModelsRunningReply{}) - _ = msg.Respond(data) + workers.On(AnyNode, workerctl.PathModelsRunning, func(string, []byte) any { + return messaging.ModelsRunningReply{} + }) + workers.On(AnyNode, workerctl.PathBackendList, func(string, []byte) any { + return messaging.BackendListReply{} }) - Expect(err).NotTo(HaveOccurred()) - FlushNATS(infra.NC) // Start a mock gRPC backend using the same helper as full flow tests llm := &trackingTestLLM{} @@ -85,7 +82,7 @@ var _ = Describe("SmartRouter trackingKey", Label("Distributed"), func() { Expect(registry.Register(context.Background(), node, true)).To(Succeed()) nodeID = node.ID - unloader := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, 3*time.Minute, 15*time.Minute) + unloader := nodes.NewRemoteUnloaderAdapter(registry, infra.NC, workers.Client(), 3*time.Minute, 15*time.Minute) router = nodes.NewSmartRouter(registry, nodes.SmartRouterOptions{ Unloader: unloader, })