diff --git a/.agents/distributed-seams.md b/.agents/distributed-seams.md index 2b313ed1f8eb..5957062ac966 100644 --- a/.agents/distributed-seams.md +++ b/.agents/distributed-seams.md @@ -2,25 +2,236 @@ Distributed mode talks to other processes through a few interfaces. Code above them must not name NATS, and a new transport is a new implementation of them. -Today every seam has one carrier: NATS, or a direct dial for the gRPC and HTTP -connections to a worker. Only the carrier files import `nats.go` -(`core/services/messaging/client.go`, `tls.go` and -`core/services/nodes/control_nats.go`). +Today the fan-out seam has two carriers: NATS, and `core/services/pgbus`, which +uses PostgreSQL LISTEN and NOTIFY. Every other seam has one: NATS, or a direct +dial for the gRPC and HTTP connections to a worker. Only the carrier files import +the NATS libraries (`core/services/messaging/client.go`, `tls.go`, +`core/services/nodes/control_nats.go` and `pkg/natsauth`). Only the pgbus package +imports `github.com/jackc/pgx`, and only `core/services/carrier/pgbus.go` imports +the pgbus package. Only `core/services/tunnel` imports yamux. A spec in +`core/services/carrier` fails when any other +non-test file imports one of them; it reads the import lines of every Go file, so +a build constraint does not hide one. A new carrier adds its library and its +files to that spec. | Seam | Interface | Lives in | Today | |---|---|---|---| | Fan-out | `messaging.Broadcaster` | `core/services/messaging` | NATS client | -| Queues | `messaging.WorkQueue` (producer), `messaging.WorkConsumer` (worker), keyed by `messaging.WorkKind` | `core/services/messaging` (`workqueue.go`) | `NewNATSWorkQueue`, `NewNATSWorkConsumer` | -| Control verbs, frontend half | `nodes.NodeCommandSender`, `nodes.FileStager` | `core/services/nodes` | `RemoteUnloaderAdapter`, `S3NATSFileStager` over NATS request/reply; `HTTPFileStager` over HTTP | -| Control verbs, worker half | unexported `controlServer` (`handle`, `handleWithProgress`) | `core/services/worker` | `natsControlServer` | -| Agent RPC, frontend half | `AgentControl` (package `mcp`) | `core/http/endpoints/mcp` | `nodes.NATSAgentControl` | -| Agent RPC, worker half | unexported `agentRPCServer` | `core/cli/agent_worker.go` | `nodes.NATSAgentRPCServer` | +| Queues | `messaging.WorkQueue` (producer), `messaging.WorkConsumer` (worker), keyed by `messaging.WorkKind` | `core/services/messaging` (`workqueue.go`) | `NewNATSWorkQueue`, `NewNATSWorkConsumer`; `jobs.ClaimQueue` and `jobs.ClaimConsumer` (the claim table); `agentworker.Work` (an agent worker on the tunnel) | +| Control verbs, frontend half | `nodes.NodeCommandSender`, `nodes.FileStager` | `core/services/nodes` | `RemoteUnloaderAdapter` over NATS request/reply and `TunnelControl` over the HTTP control plane of a worker; `S3FileStager` over either (`NewS3NATSFileStager`, `NewS3TunnelFileStager`); `HTTPFileStager` over HTTP | +| Control verbs, worker half | unexported `controlServer` (`handle`, `handleWithProgress`) | `core/services/worker` | `natsControlServer`, `httpControlServer` | +| Agent RPC, frontend half | `AgentControl` (package `mcp`) | `core/http/endpoints/mcp` | `nodes.NATSAgentControl`, `nodes.AgentControlClient` (tunnel) | +| Agent RPC, worker half | unexported `agentRPCServer` on NATS, `agentworker.Handler` on the tunnel | `core/cli/agent_worker.go`, `core/services/agentworker` | `nodes.NATSAgentRPCServer`, `agentworker.Config` and `agentworker.Work` | | Dial to a worker | `nodes.BackendClientFactory`, `nodes.ModelProber`, `nodes.WorkerNetDialerFor` | `core/services/nodes` | direct dial | Request and reply payloads of the control verbs live in `core/services/workerctl`, so the frontend half, the worker half and any carrier share one wire format without importing each other. +## Holders and the carrier row + +Code above the seams does not hold a carrier. It holds a forwarding holder from +`core/services/carrier`, one per seam: `Broadcaster`, `WorkQueue`, `Commands` +(`nodes.NodeControl`), `Files`, `Clients`, `NewWorkerDialer` and `Agents`. Each +call loads an `atomic.Pointer[carrier.Set]` once and calls the same method on +that set. A `carrier.Set` is the implementation of every seam for one carrier. +Build a set completely and call `Validate` before you store it. The holders take +no lock and allocate nothing, and a call that is already running is never moved. +`carrier.NewNATSSet` and `carrier.NewTunnelSet` are the places that build a set. + +`NodeControl` is the whole control surface a carrier gives the frontend: +`NodeCommandSender`, the process lister, the three unload interfaces the model +loader reads by type assertion, and the install timeout, model file deletion and +forced install that the managers and the reconciler call. Take it, not a +concrete sender. The `Files` holder implements `RequestFileReleaser` because +`FileStagingClient` looks for it, and `carrier.FileCarrier` makes every set +provide it. + +`carrier.Broadcaster` also keeps a registry of every subscription, because a +subscription lives on one carrier. A swap is `Listen(next)` (each handler now +listens on both carriers), then a store of `next` in the pointer (publishes go +to `next`), then `Release(old)`. No message is published on two carriers. +`NotifyReconnect` runs the reconnect hooks that `syncstate` and the failover +sync register, because messages across the flip are not ordered. + +`carrier.NewPgbusFanout` builds the fan-out member of a set for the pgbus +carrier. It opens the LISTEN connection when it is called and not before, so a +deployment on NATS holds no LISTEN session until a change of carrier asks for +one, and `Close` frees it. The set still needs the other members from the other +parts of that carrier before `Validate` accepts it. + +The active carrier is one row of table `cluster_carrier` +(`cluster.CarrierStore`). Each transition is a compare-and-set on its epoch. At +startup a replica reads the row, or inserts it if it is absent, and builds the +set the row names (`carrierRuntime.build` in `core/application`: `NewNATSSet` or +`NewTunnelSet`). A replica whose row names a carrier it does not implement fails +to start. It never falls back to another carrier, and its own flags never +override the row. + +The row is seeded by the first replica that starts: NATS when it has a NATS URL +(the flag `--nats-url`, or a URL stored in the cluster settings), the tunnel +otherwise. `--nats-url` means what it always meant for a deployment on NATS, and +it is no longer required: a deployment with only PostgreSQL runs on the tunnel. + +The address is stored before the row, so a replica that finds a row for NATS also +finds the address. A replica with no flag and no stored address does not seed the +tunnel into a cluster that has NATS: when a worker or a replica is registered, or +the row names NATS, its start fails and says to set `--nats-url`. + +The flag is copied into the cluster setting `nats.url` when none is stored, without +its user name, password and query string. A flag that names another server than the +stored address stops the start and names both: the stored address is what every +replica uses, and the operator who moves a broker stores the new address +(`PUT /api/cluster/settings`) and restarts. A replica whose flag names the stored +server dials the flag, because that is where it keeps its own credential. A NATS URL +that points at a server which is not up does not stop the start: +`messaging.New` retries on a failed connect, so a frontend can start before its +broker. A URL it cannot parse does. + +The row of a change to NATS carries the address (`cluster_carrier.nats_url`), copied +from the setting when the admin requests the change. While the change is under +way every replica builds from the row, and a prepared set built from another +address is dropped and built again. The setting itself can be written only when no +change is under way and this replica reaches the server at the new address. + +## Changing the carrier + +An admin changes the carrier of a running cluster with `POST /api/cluster/carrier` +(`{"target": "nats"|"tunnel", "dry_run": bool, "force": bool}`, or `{"abort": true}`). +`GET /api/cluster/carrier` reports the state. Both routes, and +`GET`/`PUT /api/cluster/settings`, answer to an admin only; the prefix +`/api/cluster/` is not public, and only the exact connect and peer paths skip the +global authentication. Saving a NATS URL stores it and checks that the serving +replica reaches it. It does not change the carrier. + +The row has a state, and every move is a compare-and-set on the epoch: + +| State | Meaning | Who moves it on | +|---|---|---| +| `stable` | One carrier is active. A previous carrier may still be attached (`draining`, `draining_until`). | the admin starts a change | +| `prepare` | Every replica builds the target and listens on it. Publishing stays on the old carrier. | the leader, when every live replica reported ready for this epoch | +| `commit` | The target is active. Every replica publishes on it and confirms with the epoch of the commit. | the leader, when every live replica confirmed | +| `stable` | The old carrier drains for `max_drain`. | the leader, at `draining_until` | + +`cluster.Switch` holds the protocol and keeps no state outside the row and the +`instances` table, so a leader that dies between two moves is replaced by one that +reads the same row. The leader is whichever replica holds the advisory lock +`KeyCarrierSwitch` at a tick (`advisorylock.RunLeaderLoop`). Every timeout is +measured on the clock of the database (`CarrierStore.DBNow` and the change stamp), +never on the clock of a replica. A replica that cannot build the target reports +the reason, and the leader aborts naming it. A replica that is late for the prepare +timeout aborts the change, or is left behind by a forced one. `abort` is accepted +only in `prepare`: after the commit the target is the active carrier, and going +back is a change like any other. + +The preflight (also the dry run) lists: the live replicas with their versions (old +frontends write no row and cannot be detected, so the admin confirms that the list +is complete, and every frontend must be upgraded before the first change); whether +each replica could build the target (`instances.availability`, written by the +replica when a dry run asks for it); every worker that cannot follow and why +(`nodes.SwitchWorkers`: a worker that reports no capabilities predates carrier +switching, a worker that reports an error says it); and the work in flight. A +blocker can be forced. A forced change commits without a replica that is not ready, +and names it in the note of the row. + +`carrier.Swapper` is what a replica does. It polls the row every two seconds with a +jitter, and a hint on `state.carrier` makes it look at once; the poll decides. In +`prepare` it builds the target, attaches every subscription to it +(`Broadcaster.Listen`) and reports ready for the epoch. In `commit` it stores the +set, so that publishes, enqueues and new calls use it, starts the work of the set +(`Set.Start`: the claim loop of the tunnel) and confirms. In `stable` with a drain it +does nothing: the old set stays attached. When the leader ends the drain it stops +listening on the old set, cancels its work with `messaging.ErrCarrierReleased`, +runs `Set.Handoff` (the tunnel moves its pending claims to the queue of the new +carrier: update, then publish, so a job is lost and not run twice if the replica +dies in between; it runs again after `Settle` for the producers that were slow to +flip), closes the set, and runs the reconnect hooks once. A replica that polls +late, restarts or joins during a change acts on the row alone. + +`carrier.Window` is the only per-worker routing, and it exists between the commit +and the end of the drain. A control verb, a file transfer, a dial or a client for +one worker goes to the active set when the worker is attached to it, else to the +previous set when it is attached to that, else it fails with `nodes.ErrNoRoute`. +The choice is read from the state of the worker (`BackendNode.Attached`, or a +tunnel held when it reports nothing) and never from the error of an earlier +attempt. A call that is running is never moved, so an install, a stream or a load +renewal ends on the carrier it began on while the old carrier is attached. + +A run that the end of the drain cuts off is cancelled with +`messaging.ErrCarrierReleased`. The claim driver completes its claim and leaves the +job to the reaper (`ReapStuckJobs`), because the run did start and the carrier that +took over must not start it again. The model loads, installs and agent runs that +were started before the change finish where they began; whatever is still running +at `max_drain` is failed by the reaper, and a load whose lease is not renewed is +killed by the watchdog of the worker, as it is for a dead owner (`#12524`). + +A worker that holds a tunnel has no address and binds its backends to loopback. The +NATS set therefore builds its clients with `nodes.NewGuardedDirectClientFactory`: +for a worker that registered no address, a loopback address is not dialled, and the +client reports a failure of the transport, which no code that decides whether a +backend is dead reads as an answer. A forced change to NATS cannot make the health +monitor or the reconciler reap a model that runs. + +The cluster settings (`cluster.SettingsStore`, table `cluster_settings`) hold what +every replica must read the same way: `nats.url`, `nats.worker_url` and the waits +of a change (`switch.prepare_timeout`, `switch.transition_window`, +`switch.max_drain`; defaults 1m, 2m and 15m; the flags +`--carrier-prepare-timeout`, `--carrier-transition-window` and `--carrier-max-drain` +are the fallback). The runtime settings of the application are a file of one +process and cannot carry them. Credentials are never stored there: an address +with a user name, a password or a query string is refused on write, and every view +and log line drops them. + +## The worker tunnel + +A worker that has no address that others can reach holds one outbound websocket +to a frontend. `core/services/tunnel` has the wire format (a request frame names +a tag and a target, and a reply frame always follows), the session, the registry +of the sessions that a replica holds, and `Splice`. The connect endpoint is +`GET /api/cluster/connect` (`core/http/endpoints/cluster`) and the client is +`core/services/worker/tunnel.go`. + +- Every websocket of the tunnel comes from `tunnel.NewUpgrader` or + `tunnel.NewDialer`. They set 64 KiB buffers. The 4 KiB default of gorilla makes + a transfer from the worker five times slower than the other direction. A spec + reads the source files and fails on a literal `websocket.Upgrader{` or + `websocket.Dialer{`. +- A worker has two sessions, called lanes. The inference lane carries model + calls, control and health, with the yamux windows at their defaults. The bulk + lane carries file transfers, with windows of 4 MiB and 32 MiB. The caller + chooses the lane in `Registry.Open`. A large transfer on the inference lane + delays every small call on it, because yamux queues up to one window ahead of + them. `Registry.Open` falls back to the inference lane when a node has no bulk + session. +- The inference lane owns the claim in `node_connections`. The bulk lane is held + under that claim, only on the replica that holds the inference lane, and it is + closed with it. A bulk dial that reaches another replica gets 409 and the + worker dials again. +- `Splice` passes the end of one direction on with `CloseWrite` and goes on + copying the other direction. It closes both streams at once when a direction + fails or when a stream cannot half-close. +- A node has its own tunnel credential. Registration mints it, returns it once, + and stores only its SHA-256 (`BackendNode.TunnelTokenHash`). The registration + token of the deployment never opens a tunnel. Every registration mints a new + credential, so the client reads it at every dial. +- The frontend resets every stream that a worker opens (`MaxIncomingStreams` is + 0 on the server side of a session), because nothing reads such a stream and an + accepted one would hold its window of unread data for as long as the session + lives. The worker side keeps the default: it accepts the streams of the + frontend. +- The credential is checked when a worker dials, and never again. A node that is + deleted, or whose credential a registration replaces, loses the session that + this replica holds (`Registry.Disconnect`), and a replica that shuts down + closes its registry (`Registry.Close`) before it leaves the instances table. A + session that another replica holds ends when its worker dials again and is + refused. A worker whose dial is refused with 401 or 403 registers again + (`TunnelConfig.Reauthorize`), once for each wait of its backoff. +- A claim has its own bound (10 seconds), and a claim that was written for a + request that already ended is released. +- A refusal of a stream is one of four. `tunnel.IsWorkerAnswer` is true for the + three that are evidence about a backend. `ErrStreamNotServed` says that the + worker learned nothing, and it must never count as evidence. + ## Queues `WorkKind` names the work (`WorkTask`, `WorkMCPCI`, `WorkAgentRun`). @@ -41,12 +252,104 @@ MCP CI with 1 (`startMCPCIConsumer`) and for agent runs with the dispatcher's An empty subject keeps `agent.execute`; an empty queue is kept and makes a plain subscription, as an explicitly empty `LOCALAI_AGENT_QUEUE` always did. +`messagingtest.RunWorkQueueConformance` is the suite every carrier of the +queue passes: one consumer gets a payload and a publisher for its events, +competing consumers split the work once when nothing fails, the kinds stay apart, +`maxInFlight` holds, `Unsubscribe` waits for the running handler, and a payload +above `messaging.MaxWorkPayloadBytes` (1 MiB, the default `max_payload` of NATS) +is refused. It does not state what differs between carriers: the fate of work +whose consumer dies, and how often a failed handler is called again. A handler +must tolerate a repeat. + +### The claim queue + +The tunnel carrier has no broker, so its queue is a table. `jobs.ClaimQueue` +writes a row (`work_claims`) and sends a wake hint on `messaging.SubjectClaimWake`. +`jobs.ClaimConsumer` runs on a frontend replica. It polls every two seconds, +wakes at once on a hint for its kind, and claims into as many slots as it has with +`FOR UPDATE SKIP LOCKED`. The hint is at-most-once, so the poll is what guarantees +that work is found; with the poll alone a unit waits one second on average, and +with the hint the median wait is a few milliseconds. A row stays until the work +has an answer, so `Enqueue` succeeds while nobody consumes. + +A replica claims and not an agent worker, because an agent worker has no database. +`jobs.AgentDriver` is the handler: it picks an agent worker with +`nodes.AgentSelector`, sends the payload to the verb of the kind +(`workerctl.VerbAgentExecute` or `VerbMCPCIRun`) and reads the stream that comes +back. `jobs.DispatchLoop` is a consumer for the three kinds with the driver as its +handler. + +- The outcome of the handler settles the row. `nil` deletes it. An error, or a + panic, releases it with a wait that doubles from two seconds to a minute (the + wait keeps one poison row from being claimed ahead of every newer row on every + tick). An error that wraps `jobs.ErrKeepClaim` leaves it with its replica: the + work ran and its answer could not be recorded, and releasing would run it again. +- A release is for work that was not started. The driver releases a claim when no + worker could be picked, when a worker has no slot, no route or is too old to + serve the verb (the next worker is tried first, with no wait), when the carrier + was released before the run began, and when a worker was reached and the call + failed before one line of the run came back. Once a line of the run has come + back, the claim is never released: a broken link then closes the job as failed + and completes the claim, because a second run could repeat a CI pipeline or the + side effects of an agent. +- A row ends after `MaxFailures` failures (default 10, about five minutes of + trying; `DispatchConfig.MaxFailures`). Only the failures that say something about + the work count: a worker was reached and the run failed on it. A release for lack + of a worker, a route or a free slot does not count, so a fleet that is down for a + day keeps its queue. The consumer records the job as failed first and deletes the + row after. +- A claim is settled only by the replica and the attempt that hold it, so a + holder that was reaped and answers late cannot delete the claim of the replica + that took the work over. +- `ReapAbandoned` returns the claims of replicas that are not live. It reads + `cluster.LiveInstanceIDsSQL` and has no age in it: a live replica that holds a + claim for an hour has a slow job. A replica whose own row is not live claims + nothing. +- A replica stops the handler of a claim that is no longer its own. Every tick the + consumer (once for all kinds) checks that the replica is live and lists the + claims it holds. When the instances row is missing for as long as a peer needs to + reap, the handler's context is cancelled, the peer runs the work, and this + replica does not settle the row. When this replica registers again before a peer + reaped, it gives the row back to the queue itself, because a row held by a live + id is never reaped. +- A process that starts to claim first releases the rows its own id holds + (`ReleaseOwned`). `LOCALAI_INSTANCE_ID` is stable across restarts, the id is + live again at once, and the reap never frees rows of a live id. +- Delivery is at-least-once. A replica that dies in the middle of a job loses + nothing and the job runs again. NATS loses the job of a consumer that dies. A + stall of a replica for as long as the liveness window can also run a job twice. + A handler for work that is not safe to repeat must record its progress. +- `MigratePending` takes the pending rows out of the queue once, for a change to a + carrier that has its own queue. Held rows are driven to their end where they are. + A row the new carrier refuses goes back to pending (`RestorePending`). Moved rows + stay as `migrated` for `MigratedRetention` (a day) and are then deleted with + their payloads. +- Payload at rest: the payload of a row is the whole request of the work, in + plain text. For an MCP job it holds the job, the task and the model + configuration, including the configuration of its MCP servers. It stays in + `work_claims.payload` until the work is complete (the row is deleted) or, for a + moved row, until the purge. Anyone who can read the table can read it. Treat the + table as the place where those secrets live, and give the database role the + same care. +- The poll reads the rows of a kind through `idx_work_claims_pick` (kind, state, + created_at) and filters the rows that wait for a backoff, so it passes over them + at each poll. The failure cap keeps their number small. +- A plain task (`WorkTask`) has no worker on either carrier. The driver drops the + claim and the job is left to the reaper (`ReapStuckJobs`), as it was on NATS. +- `Dispatcher.Cancel` marks the job as cancelled in the database and publishes + `jobs..cancel`. The mark is what stops a job that is still queued: the driver + reads the status before it dispatches, and again before each worker it offers the + run to. On NATS nothing subscribes to the subject, as on master, and a spec pins + that. On this carrier the replica that holds the job ends its stream, the run on + the worker ends with its request, and the job is closed as cancelled. + ## Control verbs Frontend half: `NodeCommandSender` sends the lifecycle verbs; `FileStager` -moves files. `S3NATSFileStager` returns `nodes.ErrNoRoute` when nothing is -listening for the node. `HTTPFileStager` reports connection failures as -ordinary errors. +moves files. `S3FileStager` sends the file verbs over a control link, so the same +stager runs over NATS and over the HTTP control plane of the tunnel, and it +returns `nodes.ErrNoRoute` when no route to the node exists. `HTTPFileStager` +reports connection failures as ordinary errors. `NodeCommandSender` embeds `LoadOperationControl`: the load operation verbs (`InstallBackendOp`, `StopLoadOperation`, `OperationControl` for renewals and @@ -57,15 +360,41 @@ runs every carrier in `loadOperationCarriers` against it. Whether a worker names operations is the worker's own report (`BackendInstallReply.ReportsOperations`), not a property of the carrier. -Worker half: each verb is a `controlVerb`. A handler is typed with `unary`, -`withProgress` or `noReply` and registered with `handle` (one request of the -verb at a time on NATS, panic not recovered) or `handleWithProgress` (a -goroutine per request, progress published on the install-progress subject). -An undecodable body is still answered with the verb's typed refusal. The -`undecodable` error a `controlHandler` returns is read only by tests today: the -NATS server drops it. It is a recorded exception to the no-dead-code rule, kept -as the hook a carrier that signals a malformed request out of band (HTTP 400) -needs. +Worker half: each verb is a `controlVerb`, named by the constants of +`core/services/workerctl` (`workerctl.AllVerbs` lists them). A handler is typed +with `unary`, `withProgress` or `noReply` and registered with `handle` (one +request of the verb at a time on NATS, panic not recovered) or +`handleWithProgress` (a goroutine per request, progress published on the +install-progress subject). An undecodable body is still answered with the verb's +typed refusal on NATS. + +`httpControlServer` is the second server of the same handlers. It mounts each +verb at `workerctl.PathOf(verb)` below `/v1/control/`, on the HTTP server that the +worker already runs (`nodes.StartFileTransferServerWithControl`, behind the same +bearer token as the file routes). It delivers requests of a verb concurrently, +where NATS delivers them one by one, so a handler must be safe to call from +several goroutines. The answer of a verb is 200 with the reply, and the refusal +of the worker stays in the reply. A body that cannot be read, including a body +that does not decode (the `undecodable` return of a `controlHandler`), is 400. A +verb the worker does not serve is 404 with a body that says so, and a method other +than POST is 405. A verb with progress streams lines (`workerctl.Envelope`): any +number of progress lines and then exactly one reply line. A body that ends +without a reply line is a link that broke, and it says nothing about the work. +A verb that only one server can serve is simply not registered on the other +(`files.*` need object storage on both). + +The frontend half is one implementation of the verbs over a `controlLink` +(`nodeControl` in `unloader.go`). The link holds how one request travels and how +its failure is told apart: `natsLink` for NATS and `httpLink` for the control +plane of a worker (`ControlClient` is the HTTP client). Adding a carrier of the +verbs is a new link, not a copy of the verbs. The link decides four things: what +is `ErrNoRoute`, whether a wait that ran out means that the worker may still be +installing (`waitExpired`), whether a missing acknowledgement of `backend.stop` +may be read as an older worker (`acknowledgementMissing`: true only for NATS), and +how the carrier is named in a log line. An HTTP link never reads silence as a +stop that was done. A 404 is a worker that does not serve the verb. It becomes +`ErrNoRoute` only for `backend.upgrade`, which keeps the fallback to the older +install, and it is a plain error for every other verb. ## Agent RPC @@ -77,6 +406,27 @@ cancellation, and uses the default MCP timeouts when there is none. undecodable body with an `unmarshal error: ` reply. It also serves the node's backend stop as `func(backend string)` and never replies to it. +On the tunnel the frontend picks the worker. `nodes.AgentSelector` lists the agent +nodes that may take work (not pending, not draining; a failed health probe does +not exclude one) and asks which of them hold a tunnel, preferring one that this +replica holds because that call needs no relay. `nodes.AgentControlClient` sends +the request over the control client, with the same contract as NATS: a decoded +reply is returned with a nil error, no agent worker wraps `nodes.ErrNoRoute`, and +only the deadline of ctx applies. A request goes to a second worker (up to three) +only when the failure proves that the first did not start it: no route, a worker +too old to serve the verb, or a worker with no free slot. A broken stream or a +timeout is returned, because the tool may have run. + +`agentworker.Handler` is the worker half on the tunnel: one HTTP server that +binds loopback only and is reached through the tunnel (its tunnel offers the http +tag alone). `agentworker.Work` is its `WorkConsumer`. A run is a streaming +request: the events that the handler publishes are progress lines that name their +subject, the result it publishes becomes the reply line, and a body that ends +without a reply line is a link that broke. A worker that is full answers busy +(`workerctl.WriteBusy`, `workerctl.ErrWorkerBusy`) and the frontend offers the run +to another worker. The frontend publishes a line only if +`nodes.Rebroadcaster` allows the subject for the node type of the worker. + ## Dial `BackendClientFactory.NewClient(nodeID, address, parallel)` builds the gRPC @@ -84,25 +434,123 @@ client for SmartRouter, HealthMonitor and the reconciler's default `ModelProber`. The node id is there because a dialer that must know which node it reaches, as a tunnel does, cannot recover it from the address; the direct factory ignores it. `ModelProber.Probe(ctx, nodeID, address)` likewise. +`NewDialerClientFactory` builds clients through `grpc.NewClientWithDialer`: the +address is then the name of a backend process of the worker, and gRPC does not +resolve it. `WorkerNetDialerFor` returns the dial function for one worker's own HTTP server. `DirectWorkerNetDialer` dials the address it is handed. It serves -`HTTPFileStager` and the backend-logs proxy (HTTP and WebSocket). -`HTTPFileStager.clientFor` keeps one HTTP client per node, because the idle -pool is keyed by host and port and two workers can report the same address. +`HTTPFileStager`, `ControlClient` and the backend-logs proxy (HTTP and +WebSocket). `HTTPFileStager.clientFor` keeps one HTTP client per node, because +the idle pool is keyed by host and port and two workers can report the same +address, and `ForgetNode` drops one. A worker that holds a tunnel has no address. +`nodes.WorkerHTTPHost` gives it a name under `.invalid` for the URL, which is +never dialled, and the proxy of the environment is not used for such a host. + +A dial that fails in the transport says nothing about the backend. gRPC reports +it as `codes.Unavailable`, the code of a backend that died, so the client keeps +the error of its last dial and `grpc.TransportFailureOf(client)` returns it, +looking through decorators (every decorator implements `Unwrap`). It is nil when +the dial worked and when the host answered about the backend (an error that +implements `grpc.BackendAnswer`). Code that decides a backend is dead must ask +it before it acts: the probe of the reconciler (`ProbeUnknown`), the per-model +check of the health monitor, the warm path of the router (`probeUnknown`), the +eviction of a remote model and the check of a cached remote model. None of them +may reap a row, count a miss or shut a model down on a transport failure. + +## The peer link and the relay + +A worker holds one tunnel and it lands on one replica. `tunnel.WorkerDialer` is +the one door to a worker: the replica that holds the tunnel opens the stream, and +any other replica relays through the owner over a link that the pool of peers +(`tunnel.PeerPool`) holds and the owner accepts at `GET /api/cluster/peer` +(`tunnel.PeerSessions`, `tunnel.Relay`). A replica publishes its address and the +hash of its own peer credential in its instances row. The credential of the +replica is the only secret that opens a peer link; the registration token opens +none. The link is plain `ws` unless `--peer-tls` is set (`tunnel.WithPeerTLS`), +and a replica that publishes an address that is not on its host warns at +start-up while it is plain. A relayed stream carries two request frames, the relay frame (worker, lane, +remaining budget of the caller) and then the frame of the tunnel. The relay never +relays onward. + +The errors of a dial are kept apart. A route that does not exist carries +`tunnel.ErrNoRoute`, and the carrier maps it onto `nodes.ErrNoRoute` in one place +(`carrier.mapDialError`). Absence claims (`cluster.ErrNoConnection`, +`ErrInstanceNotFound`) never reach a caller. A refusal of the worker keeps its +identity and is no `ErrNoRoute`. A peer that does not answer is +`tunnel.ErrPeerUnreachable` and is not `ErrNoRoute`: a link between two frontends +says nothing about a worker, and the scheduler demotes a worker on `ErrNoRoute`. A +bulk lane that is down is `tunnel.ErrNoBulkSession` and is not `ErrNoRoute` +either. A path that is slow (a socket deadline, a worker that takes the stream and +says nothing) is `tunnel.ErrTransport`. A reply that breaks the protocol is +`tunnel.ErrProtocol`. A database that cannot say who holds the tunnel is +`tunnel.ErrInfrastructure`. An owner link that breaks before it answers is +`tunnel.ErrPeerUnreachable`. None of them is `ErrNoRoute`, and a failure that +`routeFailure` does not recognise is none of them either: only the listed routing +facts demote a worker. The budget of the caller is checked first and is never a +route failure. + +A dial can ask for the bulk lane (`tunnel.WithBulkLane`) and then never falls +back to the lane of model calls. The file stager asks for it. ## Rules a carrier must keep +- The control verbs of a worker are never served unauthenticated to the network. + `nodes.controlGate` checks the bearer token, and with no token it serves only a + caller on the loopback address, which is the stream that the tunnel opens on + the worker. The file routes keep their older behaviour on an empty token. An + agent worker binds loopback only. + +- Payload: a broadcast above `messaging.MaxBroadcastBytes` is refused with + `messaging.ErrPayloadTooLarge` by every carrier. - Subjects: one closed set of roots, the `broadcastRoots` and `controlRoots` maps in `core/services/messaging/subject_rules.go`. Add a root there, with a row in the roots table in `subject_rules_test.go`, not in a carrier. `messaging.ValidateSubject` refuses anything else with - `messaging.ErrUnservedSubject`. + `messaging.ErrUnservedSubject`. A carrier with no request and reply, such as + pgbus, uses `messaging.ValidateBroadcastSubject`, which refuses the control + roots with the same error. `messagingtest.BroadcastRootSubjects` has one + subject for each broadcast root, and a spec keeps it equal to the rules. +- The `carrier.Broadcaster` holder carries fan-out only. Its `Subscribe` + refuses a subject outside the broadcast roots with + `messaging.ErrUnservedSubject`, so a subscription that it accepted is + servable by every carrier and `Listen` never fails for the subject. Control + traffic (the `nodes` and `mcp` roots) does not go through the holder: it uses + the request and reply clients of the set. - Wildcards: only a whole single token `*`, never the root. `>` is refused with `messaging.ErrUnsupportedWildcard`. - Delivery is at-most-once. Anything that must survive a gap belongs in a table. - A subscriber that is reconnecting misses messages. Do not read silence as evidence about a node. +- Database connections of pgbus: one replica uses at most 1 pinned LISTEN + connection, `pgbus.spillFetchers` (8) while it reads spilled rows, and + `pgbus.Config.MaxPublishers` (16 by default) while it publishes. `Publish` + waits for a free slot and never opens another connection, so a burst of + publishers cannot exhaust `max_connections`. Add the 25 connections of each + replica to the budget of the database. +- A read of a spilled row ends after `pgbus.Config.FetchTimeout` (8 seconds by + default). The broadcast is dropped and counted with `stage=resolve`, so that + one stuck read does not hold back the delivery on every subject. + +## The cluster registry + +`core/services/cluster` holds the facts that all replicas share through the +database. Table `instances` has one row for each live frontend replica +(`Registry`, kept fresh by `Membership`, which every replica runs). The +database clock stamps `last_seen`, and `instanceIsLive` is the one predicate +for "alive". The row also holds the epoch of the carrier row for which the +replica has built a carrier (`ReadyEpoch`) and the reason when it could not +(`ReadyReason`). Table `node_connections` records which live replica holds the +connection of a worker. `Registry.Claim` takes a new epoch from a sequence for +every claim, and `Release` needs both the owner and the epoch, so a replica that +lost a worker cannot clear the claim of the one that won it. Epochs are unique, and +they are not ordered: a claim that inserts after a `Release` can draw a number +that is lower than one already issued. Compare epochs for equality (`==`, +`!=`) and never for order. A spec (`core/services/cluster/epoch_order_test.go`) +fails on a `<`, `>`, `<=` or `>=` with an operand named like an epoch in code +under `core/services`. `Registry.Presence` tells a worker that is +reconnecting from one that is gone only when the departure is older than the +grace, and nothing but `PresenceGone` may be read as absence. ## Four conditions that are never reported as each other @@ -137,42 +585,67 @@ runner, where it fails so a runner without Docker cannot hide the check. The end-to-end specs in `tests/e2e/distributed` run the NATS implementations of the other seams against a real server, also through Docker. +## Workers follow the carrier + +A worker does not choose its carrier. Every heartbeat goes to the frontend over +HTTP, which exists on both carriers, and the answer carries the active carrier, +its epoch, the state of a change, the target and the carrier that drains +(`carrier`, `carrier_epoch`, `carrier_state`, `carrier_target`, +`carrier_draining`). `worker.Follower` reads it. The backend worker +(`worker.BackendPlane`) and the agent worker (`agentworker.FollowConfig`) give it +one `Attacher` for each carrier, so the code that attaches at start is the code +that attaches later. + +- At start the worker registers, learns the active carrier, and attaches to it. + It needs no NATS setting: the registration answer hands over the address + (`nats.worker_url`, else `nats.url`), the CA as PEM and the per-node JWT. A local + `--nats-url`, CA or credential still wins. A frontend that predates carriers + names none, and the worker then needs `--nats-url` as before. +- While a change is prepared the worker registers again with `carrier` set to the + target, which returns the credential of that carrier only. A tunnel token + replaces the previous one and ends the session of the replica, so a worker that + holds the tunnel never asks for one. Only the active carrier, the target of a + change and the carrier that drains can be asked for. +- It waits a random time of up to `LOCALAI_FOLLOW_MAX_DELAY` (ten seconds), then + attaches with a wait that grows from one to thirty seconds when the carrier + refuses. It keeps the old carrier open, and reports `attached`, `attached_epoch`, + `follow_capabilities` and `follow_error` in the heartbeat. That report is the + only source of where a worker is attached (`nodes.SwitchWorkers`). A worker that + reports nothing predates carrier switching and is on NATS. +- It closes a carrier when the cluster released it, the carrier it moved to is + connected, and no request runs on it (`inflight` in the control servers). It + reports first, then closes. Backend processes are not touched. A change that is + aborted makes it close the target. +- A worker that cannot follow keeps the carrier it has, reports why, and tries + again. A backend worker cannot follow to NATS without an address that the + frontends can dial (`--addr` or `--advertise-addr`; a worker that booted on NATS + counts, because its backends listen on every interface), or with a NATS server + that asks for a client certificate it does not have. Such a worker is a blocker + of the change, and a forced change leaves it where it is. +- After the cluster has settled, `HealthMonitor.UseCarrier` demotes a worker that + reports it is not attached to the active carrier. The change is status only, + like the tunnel read, and the worker is promoted when it reports the active + carrier. + ## Open items for a second carrier -- `DistributedModelStore.Range` (`core/services/nodes/distributed_store.go`) - builds a tokenless `model.Model` from `node.Address`. It dials nothing today, - but it is the one direct construction site left outside the dial seam. -- The worker's `files.ensure` handler passes the first caller's ctx into a - shared singleflight closure. The NATS server hands it `context.Background`. - A carrier with a request ctx needs `context.WithoutCancel` there, or one - cancelled caller fails the others. -- `HTTPFileStager` caches a client per node and never forgets one. There is no - `ForgetNode`. -- `clientFor` returns no error; a carrier with no dialer for a node needs that - path. -- The agent worker still uses NATS directly for its connection and for agent - events (`agents.NewEventBridge`). -- `ReplicaReconcilerOptions` has no `ClientFactory` field. Without a - `Prober`, `NewReplicaReconciler` builds its own `tokenClientFactory`, so a - carrier that dials through another factory must add the field. -- The backend-logs proxy (`proxyHTTPToWorker`) starts from - `httpclient.HardenedTransport()`, which keeps - `Proxy: http.ProxyFromEnvironment`. With `HTTP_PROXY` set, a dialer that - routes by node would be handed the proxy address. Clear `Proxy` when the - dialer is not the direct one. -- `natsControlServer.subject` refuses a verb it has no subject for, so - registration fails. A verb that only one carrier serves needs a per-carrier - opt-out where the verbs are registered. -- `WorkHandler` returns only an error. A carrier whose stream handler must send - a terminal reply derives it from the `jobs..result` event the handler - publishes on `events` (`handleMCPCIJob` does this). -- Not additive: the agent-run consumer (`NATSDispatcher.runDelivery`) ignores - the per-delivery `events` publisher and publishes through the process-wide - `EventBridge`, which is built on the NATS client. A second carrier must - change `NATSDispatcher` and `EventBridge`, for example by binding - `handleJob`'s publishes to `events` through a bridge view that shares the - cancel registry. -- `controlHandler`'s `undecodable` return is the recorded exception described - under Control verbs. +- `WorkHandler` returns only an error. A carrier whose stream handler must send a + terminal reply derives it from the `jobs..result` event the handler + publishes on `events` (`handleMCPCIJob` does this, and `agentworker.Work` + builds the reply line from it). +- The agent-run consumer (`NATSDispatcher.runDelivery`) binds the events of a run + to the publisher of its delivery through `EventBridge.WithPublisher`, which + shares the cancel registry with the bridge. On NATS that publisher is the NATS + client, so the subjects and payloads are the ones as before. - Agent cancel has no production sender (`EventBridge.CancelExecution` has no - caller), so it is not part of the agent RPC seam. + caller), so it is not part of the agent RPC seam, and the tunnel has no verb for + it. The cancel of a run on the tunnel is its request. +- A worker that registers an address and holds a tunnel (a worker that can use + either carrier) is reached by the backend-logs proxy through the host that its + address names, and the proxy of the environment is used for it. Only a worker + without an address skips the proxy. +- The health monitor reads the presence of each node (`HealthMonitor.UsePresence`) + only while the tunnel is the active carrier: a node that heartbeats and whose + tunnel has been gone longer than `cluster.DefaultReconnectGrace` is demoted with + a status-only change, and is neither promoted nor probed until its tunnel is + back. diff --git a/.github/workflows/tests-e2e-distributed.yml b/.github/workflows/tests-e2e-distributed.yml new file mode 100644 index 000000000000..884ca4e0bc92 --- /dev/null +++ b/.github/workflows/tests-e2e-distributed.yml @@ -0,0 +1,97 @@ +--- +name: 'E2E Distributed Tests' + +on: + pull_request: + # The suite's dependency graph is large, so an allowlist of paths would + # silently stop guarding the moment code moves. It is cheap enough to run + # unless the diff is confined to paths it provably cannot reach. + # See .agents/ci-caching.md. + paths-ignore: + - 'gallery/**' + - 'docs/**' + - 'examples/**' + - '**/*.md' + push: + branches: + - master + +concurrency: + group: ci-tests-e2e-distributed-${{ github.event.pull_request.number || github.sha }}-${{ github.repository }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + tests-e2e-distributed: + runs-on: ubuntu-latest + # Advisory: a failure is a visible red X on the pull request, and the job is + # deliberately not in branch protection. Promoting it to a required check + # is a repository-settings change, to be made once it has a track record. + continue-on-error: true + timeout-minutes: 45 + steps: + - name: Clone + uses: actions/checkout@v7 + with: + submodules: true + - name: Configure apt mirror on runner + uses: ./.github/actions/configure-apt-mirror + - name: Setup Go + uses: actions/setup-go@v5 + with: + go-version: '1.26.x' + cache: false + - name: Display Go version + run: go version + - name: Dependencies + run: | + sudo apt-get update + sudo apt-get install -y build-essential libopus-dev + - name: Proto Dependencies + run: | + curl -L -s https://github.com/protocolbuffers/protobuf/releases/download/v26.1/protoc-26.1-linux-x86_64.zip -o protoc.zip && \ + unzip -j -d /usr/local/bin protoc.zip bin/protoc && \ + rm protoc.zip + go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.34.2 + go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@1958fcbe2ca8bd93af633f11e97d44e567e945af + PATH="$PATH:$HOME/go/bin" make protogen-go + - name: Pre-pull test images + # The suite starts its own PostgreSQL and NATS servers with + # testcontainers. Pulling the images here keeps the pull time out of the + # spec timeouts and makes a registry outage read as a setup failure + # instead of a test failure. + run: | + docker pull postgres:16-alpine + docker pull nats:2-alpine + - name: Distributed E2E + # TESTCONTAINERS_RYUK_DISABLED keeps the pre-pull above meaningful. The + # reaper exists to clean up leaked containers on a long-lived host, but + # this runner is ephemeral and every container dies with the VM. + # Leaving it enabled would pull one more unpinned image mid-suite. + env: + TESTCONTAINERS_RYUK_DISABLED: "true" + run: | + PATH="$PATH:$HOME/go/bin" make test-e2e-distributed + - name: Distributed E2E, switch of the carrier + # Real processes: two frontends and two workers change the carrier + # between NATS and the database tunnel while a model serves requests. + # It runs apart from the suite above so that a failure names the + # switch and the two logs stay short. Like the suite, it needs only the + # public PostgreSQL and NATS images and no secret, so a pull request + # from a fork runs it the same way. + if: ${{ !cancelled() }} + env: + TESTCONTAINERS_RYUK_DISABLED: "true" + run: | + PATH="$PATH:$HOME/go/bin" make test-e2e-distributed-switch + # tmate keeps the runner busy until the 6 hour job limit, so a single + # failure costs a whole runner slot. Only open a session when someone + # asked for one by labelling the pull request `ci-debug`, and cap the + # session so a forgotten label cannot idle a runner either. + - name: Setup tmate session if tests fail + if: ${{ failure() && contains(github.event.pull_request.labels.*.name, 'ci-debug') }} + timeout-minutes: 30 + uses: mxschmitt/action-tmate@v3.23 + with: + detached: true + connect-timeout-seconds: 180 + limit-access-to-actor: true diff --git a/Makefile b/Makefile index 7c05fda6dd3d..4c28035b01fb 100644 --- a/Makefile +++ b/Makefile @@ -349,6 +349,18 @@ test-e2e-distributed: protogen-go @echo 'Running distributed e2e tests (label Distributed, incl. NatsJWT)' $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --label-filter='Distributed && !VLLMMultinode' --flake-attempts $(TEST_FLAKES) -v -r ./tests/e2e/distributed +# Carrier switch with real processes. Two frontends and two workers (one +# dual-capable, one tunnel-only) run as child processes against PostgreSQL and +# NATS containers, and the carrier is changed from NATS to the tunnel and back. +# The spec carries the CarrierSwitch label and not Distributed, so +# test-e2e-distributed above does not run it: it starts many processes and takes +# a few minutes. It builds local-ai when LOCALAI_E2E_BIN is not set, and skips +# when docker is missing. +test-e2e-distributed-switch: protogen-go build-mock-backend + @echo 'Running the carrier switch e2e (label CarrierSwitch)' + LOCALAI_E2E_MOCK_BACKEND=$(abspath tests/e2e/mock-backend/mock-backend) \ + $(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --label-filter='CarrierSwitch' --flake-attempts $(TEST_FLAKES) -v -r ./tests/e2e/distributed + # vLLM multi-node DP smoke (CPU). Builds local-ai:tests and the # cpu-vllm backend from the current working tree, then drives a # head + headless follower via testcontainers-go and asserts a chat diff --git a/core/application/agent_jobs.go b/core/application/agent_jobs.go index f380b075002a..27c1275a7c03 100644 --- a/core/application/agent_jobs.go +++ b/core/application/agent_jobs.go @@ -38,7 +38,7 @@ func (a *Application) RestartAgentJobService() error { agentJobService.SetDistributedJobStore(d.JobStore) } // Keep agent tasks consistent across replicas (same client the dispatcher uses). - agentJobService.SetTaskSyncNATS(d.Nats) + agentJobService.SetTaskSyncNATS(d.Broadcaster) } // Start the service diff --git a/core/application/application.go b/core/application/application.go index 326885b3105b..b3ab5e15fb60 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -611,6 +611,9 @@ func (a *Application) start() error { if a.distributed != nil && a.distributed.Unloader != nil { assistantClient.LoadStopper = a.distributed.Unloader } + if a.distributed != nil && a.distributed.Switch != nil { + assistantClient.Carrier = a.distributed.Switch + } // PII filter — same nil-or-real wiring. assistantClient.PIIRedactor = a.piiRedactor assistantClient.PIIEvents = a.piiEvents @@ -670,8 +673,8 @@ func (a *Application) agentPoolOptions() agentpool.AgentPoolOptions { opts.WorkQueue = d.WorkQueue opts.EventBridge = d.AgentBridge opts.AgentStore = d.AgentStore - if d.Nats != nil { - opts.Bus = d.Nats + if d.Broadcaster != nil { + opts.Bus = d.Broadcaster } } return opts @@ -723,7 +726,7 @@ func (a *Application) StartAgentPool() { } // Keep per-user agent tasks consistent across replicas (nil in standalone). if d := a.Distributed(); d != nil { - usm.SetJobSyncNATS(d.Nats) + usm.SetJobSyncNATS(d.Broadcaster) } aps.SetUserServicesManager(usm) diff --git a/core/application/carrier_runtime.go b/core/application/carrier_runtime.go new file mode 100644 index 000000000000..daa7ffa33f28 --- /dev/null +++ b/core/application/carrier_runtime.go @@ -0,0 +1,430 @@ +package application + +import ( + "context" + "errors" + "fmt" + "os" + "sync/atomic" + "time" + + "github.com/mudler/xlog" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/jobs" + "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/storage" + "github.com/mudler/LocalAI/core/services/tunnel" +) + +const ( + // natsConnectWait is how long a replica waits for the NATS server when it + // prepares the carrier for a change. At start-up it does not wait, because a + // frontend may start before its broker. + natsConnectWait = 10 * time.Second + // natsProbeWait is how long the report of what a replica can build waits for + // the NATS server. + natsProbeWait = 3 * time.Second + // availabilityEvery is how often a replica writes what it can build without + // being asked. A dry run asks, so this only keeps the status view from + // showing nothing. + availabilityEvery = 5 * time.Minute + // availabilityMaxAge is how old a report of what a replica can build may be + // for a change of carrier to rely on it. A dry run asks the replicas to look + // again first, so a report older than this is from a replica that did not + // answer. + availabilityMaxAge = 30 * time.Second + // probeWait is how long a dry run waits for the replicas to answer. + probeWait = 5 * time.Second + // hintPublishTimeout bounds the publish of a hint. + hintPublishTimeout = 3 * time.Second +) + +// carrierRuntime builds the carrier sets of this replica and reports what it +// could build. It is the one place of the application that names the carriers: +// everything above it holds the holders. +type carrierRuntime struct { + cfg *config.ApplicationConfig + db *gorm.DB + settings *cluster.SettingsStore + registry *nodes.NodeRegistry + fileMgr *storage.FileManager + instanceID string + jobStore *jobs.JobStore + clusterReg *cluster.Registry + tunnels *tunnel.Registry + peerPool *tunnel.PeerPool + + // bus is the broadcaster holder. The work that a set starts uses it, so that + // its subscriptions follow the sets. It is set before any set starts. + bus *carrier.Broadcaster + // active is the pointer the holders read. + active *atomic.Pointer[carrier.Set] + + probeWake chan struct{} +} + +// natsURL returns the NATS server of the cluster: the stored setting, else the +// flag of this replica. +func (rt *carrierRuntime) natsURL(ctx context.Context) (string, error) { + v, ok, err := rt.settings.Get(ctx, cluster.SettingNATSURL) + if err != nil { + return "", err + } + if ok && v != "" { + return v, nil + } + return rt.cfg.Distributed.NatsURL, nil +} + +// WorkerURL is the address of the NATS server that workers are told to use: the +// stored address for workers, else the address of the frontends, else the flag of +// this replica. A worker may reach the server by another name than a frontend +// does. It is empty when the deployment has no NATS. +func (rt *carrierRuntime) WorkerURL(ctx context.Context) (string, error) { + v, ok, err := rt.settings.Get(ctx, cluster.SettingNATSWorkerURL) + if err != nil { + return "", err + } + if ok && v != "" { + return v, nil + } + return rt.natsURL(ctx) +} + +// CAPEM is the private CA of the NATS server as PEM, or empty when the server has +// none. A worker is handed it, so that it needs no file of its own to verify the +// server. +func (rt *carrierRuntime) CAPEM() (string, error) { + path := rt.cfg.Distributed.NatsTLSCA + if path == "" { + return "", nil + } + // #nosec G304 -- the operator set this path for the CA of the NATS server. + raw, err := os.ReadFile(path) + if err != nil { + return "", fmt.Errorf("reading the NATS CA to hand to workers: %w", err) + } + return string(raw), nil +} + +// ClientTLS says that this replica connects to NATS with a client certificate. A +// worker cannot be handed such a file, so it needs its own. +func (rt *carrierRuntime) ClientTLS() bool { + d := rt.cfg.Distributed + return d.NatsTLSCert != "" || d.NatsTLSKey != "" +} + +// Build builds the set of a carrier for a change. It waits for the carrier to be +// usable, because a replica that reports ready must be. +func (rt *carrierRuntime) Build(ctx context.Context, row cluster.CarrierRow, target cluster.Carrier) (*carrier.Set, error) { + return rt.build(ctx, row, target, true) +} + +func (rt *carrierRuntime) build(ctx context.Context, row cluster.CarrierRow, target cluster.Carrier, wait bool) (*carrier.Set, error) { + switch target { + case cluster.CarrierNATS: + return rt.buildNATS(ctx, row, wait) + case cluster.CarrierTunnel: + return rt.buildTunnel(ctx, row) + } + return nil, fmt.Errorf("%w: %q", cluster.ErrInvalidCarrier, target) +} + +// natsTarget is the address of the NATS server that a build uses. While a change +// is under way the row names it, so that every replica builds from the same +// address whatever happens to the setting; otherwise it is the setting, else the +// flag of this replica. +func (rt *carrierRuntime) natsTarget(ctx context.Context, row *cluster.CarrierRow) (string, error) { + if row != nil && row.State != cluster.StateStable && row.NATSURL != "" { + return row.NATSURL, nil + } + return rt.natsURL(ctx) +} + +// connectURL is the address to dial for target. The cluster stores the address +// without a credential. When the flag of this replica names the same server, the +// flag is dialled, because it is where this replica keeps its user name, +// password or token. +func (rt *carrierRuntime) connectURL(target string) string { + if flag := rt.cfg.Distributed.NatsURL; flag != "" && cluster.PublicNATSURL(flag) == target { + return flag + } + return target +} + +func (rt *carrierRuntime) natsClient(ctx context.Context, wait time.Duration, row *cluster.CarrierRow) (*messaging.Client, string, error) { + target, err := rt.natsTarget(ctx, row) + if err != nil { + return nil, "", err + } + if target == "" { + return nil, "", errors.New("no NATS URL is configured: store one in the cluster settings (nats.url) or start the frontend with --nats-url") + } + client, err := rt.natsClientAt(ctx, rt.connectURL(target), wait) + return client, target, err +} + +// natsClientAt connects to the NATS server at url with the credentials of this +// replica. With a wait it returns only once the server has answered. +func (rt *carrierRuntime) natsClientAt(ctx context.Context, url string, wait time.Duration) (*messaging.Client, error) { + d := rt.cfg.Distributed + auth := d.NatsAuthConfig() + if auth.RequireAuth && (auth.ServiceUserJWT == "" || auth.ServiceUserSeed == "") { + return nil, errors.New("LOCALAI_NATS_REQUIRE_AUTH requires LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED") + } + // A URL or a credential the client cannot use stops here. A server that is + // not up does not, unless the caller waits for it: the client keeps trying. + client, err := messaging.New(url, d.NatsMessagingOptions("", "")...) + if err != nil { + return nil, fmt.Errorf("connecting to NATS: %w", err) + } + if wait > 0 { + deadline := time.Now().Add(wait) + for !client.IsConnected() { + if time.Now().After(deadline) || ctx.Err() != nil { + client.Close() + return nil, fmt.Errorf("the NATS server at %s did not answer within %s: check the address and the credentials", cluster.PublicNATSURL(url), wait) + } + select { + case <-ctx.Done(): + case <-time.After(50 * time.Millisecond): + } + } + } + return client, nil +} + +// CheckNATS says whether this replica reaches the NATS server at url with its own +// credentials. It opens a connection and closes it. +func (rt *carrierRuntime) CheckNATS(ctx context.Context, url string) error { + client, err := rt.natsClientAt(ctx, rt.connectURL(url), natsProbeWait) + if err != nil { + return err + } + client.Close() + return nil +} + +func (rt *carrierRuntime) buildNATS(ctx context.Context, row cluster.CarrierRow, wait bool) (*carrier.Set, error) { + var waitFor time.Duration + if wait { + waitFor = natsConnectWait + } + client, url, err := rt.natsClient(ctx, waitFor, &row) + if err != nil { + return nil, err + } + d := rt.cfg.Distributed + set, err := carrier.NewNATSSet(carrier.NATSOptions{ + Client: client, + Epoch: row.Epoch, + Registry: rt.registry, + InstallTimeout: d.BackendInstallTimeoutOrDefault(), + UpgradeTimeout: d.BackendUpgradeTimeoutOrDefault(), + Token: d.RegistrationToken, + S3Staging: d.StorageURL != "", + FileManager: rt.fileMgr, + HTTPAddrFor: func(nodeID string) (string, error) { + node, err := rt.registry.Get(context.Background(), nodeID) + if err != nil { + return "", err + } + if node.HTTPAddress == "" { + return "", fmt.Errorf("node %s has no HTTP address for file transfer", nodeID) + } + return node.HTTPAddress, nil + }, + NodeHasAddress: func(nodeID string) (bool, error) { + node, err := rt.registry.Get(context.Background(), nodeID) + if err != nil { + return false, err + } + return node.Address != "", nil + }, + }) + if err != nil { + client.Close() + return nil, err + } + set.NATSURL = row.NATSURL + xlog.Info("Connected to NATS", "url", cluster.PublicNATSURL(url), "epoch", row.Epoch) + return set, nil +} + +func (rt *carrierRuntime) buildTunnel(ctx context.Context, row cluster.CarrierRow) (*carrier.Set, error) { + d := rt.cfg.Distributed + fanout, err := carrier.NewPgbusFanout(ctx, carrier.PgbusOptions{DB: rt.db, DSN: rt.cfg.Auth.DatabaseURL}) + if err != nil { + return nil, err + } + selector := nodes.NewAgentSelector(rt.registry, rt.clusterReg, rt.instanceID) + set, err := carrier.NewTunnelSet(carrier.TunnelOptions{ + Epoch: row.Epoch, + Dialer: tunnel.NewWorkerDialer(rt.tunnels, rt.peerPool), + Fanout: fanout, + WorkQueue: jobs.NewClaimQueue(rt.db, fanout.Broadcaster), + AgentSelector: selector, + Registry: rt.registry, + InstallTimeout: d.BackendInstallTimeoutOrDefault(), + UpgradeTimeout: d.BackendUpgradeTimeoutOrDefault(), + Token: d.RegistrationToken, + S3Staging: d.StorageURL != "", + FileManager: rt.fileMgr, + // While the tunnel is the carrier in use, and while it drains, this replica + // claims queued work and drives it on agent workers. + Claims: &carrier.ClaimWork{ + DB: rt.db, Owner: rt.instanceID, Store: rt.jobStore, + Bus: func() messaging.Broadcaster { return rt.bus }, + }, + }) + if err != nil { + fanout.Close() + return nil, err + } + + return set, nil +} + +// availability says, for each carrier, why this replica could not build it now. +func (rt *carrierRuntime) availability(ctx context.Context) map[cluster.Carrier]string { + out := map[cluster.Carrier]string{} + cur := rt.active.Load() + + switch { + case cur != nil && cur.Name == cluster.CarrierNATS: + out[cluster.CarrierNATS] = "" + default: + client, _, err := rt.natsClient(ctx, natsProbeWait, nil) + if err != nil { + out[cluster.CarrierNATS] = err.Error() + } else { + client.Close() + out[cluster.CarrierNATS] = "" + } + } + + if cur != nil && cur.Name == cluster.CarrierTunnel { + out[cluster.CarrierTunnel] = "" + } else if err := carrier.ProbeListen(ctx, rt.cfg.Auth.DatabaseURL); err != nil { + out[cluster.CarrierTunnel] = err.Error() + } else { + out[cluster.CarrierTunnel] = "" + } + return out +} + +// reportAvailability looks at what this replica can build and writes it. +func (rt *carrierRuntime) reportAvailability(ctx context.Context) { + probeCtx, cancel := context.WithTimeout(ctx, 20*time.Second) + defer cancel() + if err := rt.clusterReg.ReportAvailability(ctx, rt.instanceID, rt.availability(probeCtx)); err != nil && ctx.Err() == nil { + xlog.Warn("Could not report which carriers this replica can build", "error", err) + } +} + +// RunAvailability reports what this replica can build at start, then when asked +// and every availabilityEvery. +func (rt *carrierRuntime) RunAvailability(ctx context.Context) { + rt.reportAvailability(ctx) + ticker := time.NewTicker(availabilityEvery) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + case <-rt.probeWake: + } + rt.reportAvailability(ctx) + } +} + +// wakeProbe asks RunAvailability to look now. +func (rt *carrierRuntime) wakeProbe() { + select { + case rt.probeWake <- struct{}{}: + default: + } +} + +// AskReplicas asks every replica, this one included, to look at what it can +// build, and does not wait for the answer. +func (rt *carrierRuntime) AskReplicas(ctx context.Context) { + rt.wakeProbe() + pubCtx, cancel := context.WithTimeout(ctx, hintPublishTimeout) + defer cancel() + done := make(chan struct{}) + go func() { + defer close(done) + if err := rt.bus.Publish(messaging.SubjectCarrierProbe, map[string]string{"from": rt.instanceID}); err != nil { + xlog.Debug("Could not ask the other replicas to look at the carriers", "error", err) + } + }() + select { + case <-done: + case <-pubCtx.Done(): + } +} + +// ProbeReplicas asks every replica to look at what it can build, and waits until +// the live replicas have answered or probeWait has passed. A replica that did not +// answer shows up as stale in the preflight. +func (rt *carrierRuntime) ProbeReplicas(ctx context.Context) { + asked := time.Now() + rt.AskReplicas(ctx) + + deadline := asked.Add(probeWait) + for time.Now().Before(deadline) && ctx.Err() == nil { + live, err := rt.clusterReg.ListLive(ctx, cluster.InstanceLiveness) + if err == nil { + fresh := true + for _, in := range live { + if _, known := in.AvailabilityFor(cluster.CarrierNATS); !known || in.AvailabilityAge > time.Since(asked) { + fresh = false + break + } + } + if fresh { + return + } + } + select { + case <-ctx.Done(): + case <-time.After(150 * time.Millisecond): + } + } +} + +// workCounts counts the work that a change of carrier can touch. +type workCounts struct{ db *gorm.DB } + +func (w workCounts) InFlight(ctx context.Context) (cluster.InFlight, error) { + var out cluster.InFlight + count := func(dst *int, q *gorm.DB) error { + var n int64 + if err := q.Count(&n).Error; err != nil { + return err + } + *dst = int(n) + return nil + } + db := w.db.WithContext(ctx) + if err := count(&out.Loads, db.Model(&nodes.NodeModel{}).Where("state IN ?", []string{"loading", "staging"})); err != nil { + return out, fmt.Errorf("counting the models that load: %w", err) + } + if err := count(&out.Jobs, db.Model(&jobs.JobRecord{}).Where("status = ?", "running")); err != nil { + return out, fmt.Errorf("counting the jobs that run: %w", err) + } + if err := count(&out.PendingClaims, db.Model(&jobs.WorkClaim{}).Where("state = ?", jobs.ClaimPending)); err != nil { + return out, fmt.Errorf("counting the pending claims: %w", err) + } + if err := count(&out.ClaimedClaims, db.Model(&jobs.WorkClaim{}).Where("state = ?", jobs.ClaimClaimed)); err != nil { + return out, fmt.Errorf("counting the claimed claims: %w", err) + } + return out, nil +} diff --git a/core/application/distributed.go b/core/application/distributed.go index 91cb2d053356..e264f57cd726 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -2,17 +2,27 @@ package application import ( "context" + "crypto/tls" + "crypto/x509" "encoding/json" + "errors" "fmt" "io" + "net" + "os" + "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/google/uuid" "github.com/mudler/LocalAI/core/config" mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" + "github.com/mudler/LocalAI/core/services/advisorylock" "github.com/mudler/LocalAI/core/services/agents" + "github.com/mudler/LocalAI/core/services/carrier" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/distributed" "github.com/mudler/LocalAI/core/services/jobs" "github.com/mudler/LocalAI/core/services/messaging" @@ -20,6 +30,8 @@ import ( "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "github.com/mudler/LocalAI/core/services/storage" + "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/distributedhdr" "github.com/mudler/LocalAI/pkg/sanitize" "github.com/mudler/xlog" @@ -27,8 +39,15 @@ import ( ) // DistributedServices holds all services initialized for distributed mode. +// +// The seams that cross a process boundary are holders over the active carrier +// set: they take the carrier through an atomic pointer, so a later change of +// carrier does not touch the code that holds them. Only Shutdown reaches the +// set itself. type DistributedServices struct { - Nats *messaging.Client + // Broadcaster is the fan-out seam. Code that only publishes and subscribes + // takes this, never the connection of a carrier. + Broadcaster messaging.Broadcaster WorkQueue messaging.WorkQueue AgentControl mcpTools.AgentControl Store storage.ObjectStore @@ -44,16 +63,83 @@ type DistributedServices struct { FileMgr *storage.FileManager FileStager nodes.FileStager ModelAdapter *nodes.ModelRouterAdapter - Unloader *nodes.RemoteUnloaderAdapter + Unloader nodes.NodeControl ModelCleanup *nodes.ModelCleanupService + // Clients builds the gRPC client for a backend of a worker through the + // carrier that is active. + Clients nodes.BackendClientFactory // WorkerHTTPDial reaches a worker's own HTTP server for the admin // backend-logs proxy, the same way the HTTP file stager does. WorkerHTTPDial nodes.WorkerNetDialerFor + // Membership keeps the row of this replica in the instances table and + // sweeps the rows of replicas that stopped answering. + Membership *cluster.Membership + + // Carriers reads the row that names the active carrier. The registration + // of a worker answers from it. + Carriers *cluster.CarrierStore + + // Tunnels holds the sessions that workers opened to this replica. It is + // empty while the cluster runs on NATS: the connect route is registered on + // every replica and refuses every worker that has no tunnel credential. + Tunnels *tunnel.Registry + + // Instances reads the instances table. The peer route checks the + // credential of a dialling replica against it. + Instances *cluster.Registry + // PeerSessions holds the links that other replicas dialled into this one, + // and relays their streams onto the tunnels that this replica holds. + PeerSessions *tunnel.PeerSessions + // PeerPool holds the links that this replica dialled to its peers. + PeerPool *tunnel.PeerPool + // WorkerDialer opens a stream to a worker through the replica that holds its + // tunnel: directly, or through a peer link. The tunnel carrier builds the + // dial seams on it; nothing calls it while the cluster runs on NATS. + WorkerDialer *tunnel.WorkerDialer + + // Settings holds the settings that decide what the whole cluster does: the + // NATS addresses and the waits of a change of carrier. + Settings *cluster.SettingsStore + // Switch is the protocol of a change of carrier. The admin API calls it, and + // the leader drives it. + Switch *cluster.Switch + // Swapper is what this replica does when the cluster changes its carrier. + Swapper *carrier.Swapper + // Runtime builds the carrier sets of this replica and reports what it could + // build. + Runtime *carrierRuntime + + // active names the carrier set the holders forward to. + active *atomic.Pointer[carrier.Set] + // window routes the calls to workers while two carriers are attached. + window *carrier.Window + // workers answers the questions about workers that a change asks. + workers *nodes.SwitchWorkers + shutdownOnce sync.Once } +// Disconnect ends the tunnel that this replica holds for a node, and drops what +// the active carrier caches for it. The node was removed or its credential was +// replaced. It reports whether this replica held a tunnel of the node. +func (ds *DistributedServices) Disconnect(nodeID string) bool { + held := ds.Tunnels.Disconnect(nodeID) + if ds.active != nil { + if set := ds.active.Load(); set != nil && set.ForgetNode != nil { + set.ForgetNode(nodeID) + } + } + // The carrier that is draining caches the same per-node state. + if ds.window != nil { + if prev := ds.window.Previous(); prev != nil && prev.ForgetNode != nil { + prev.ForgetNode(nodeID) + } + } + return held +} + // Shutdown stops all distributed services in reverse initialization order. // It is safe to call on a nil receiver and is idempotent (uses sync.Once). func (ds *DistributedServices) Shutdown() { @@ -61,6 +147,25 @@ func (ds *DistributedServices) Shutdown() { return } ds.shutdownOnce.Do(func() { + // Stop following the cluster carrier and stop the work that a carrier started, + // such as the claims this replica drives, before anything is torn down. + if ds.Swapper != nil { + ds.Swapper.Close() + } + // The tunnels next: their claims name this replica, and a peer that + // asks who owns a worker must not be told a replica that is leaving. + if ds.Tunnels != nil { + ds.Tunnels.Close() + } + if ds.PeerSessions != nil { + ds.PeerSessions.Close() + } + if ds.PeerPool != nil { + ds.PeerPool.Close() + } + // Then the membership, so the peers see this replica leave before its + // services stop. + ds.Membership.Stop() if ds.Health != nil { ds.Health.Stop() } @@ -70,10 +175,12 @@ func (ds *DistributedServices) Shutdown() { if closer, ok := ds.Store.(io.Closer); ok { closer.Close() } - // AgentBridge has no Close method — its NATS subscriptions are cleaned up - // when the NATS client is closed below. - if ds.Nats != nil { - ds.Nats.Close() + // AgentBridge has no Close method: its subscriptions are cleaned up + // when the carrier is closed below. + if ds.active != nil { + if set := ds.active.Load(); set != nil && set.Close != nil { + set.Close() + } } xlog.Info("Distributed services shut down") }) @@ -112,23 +219,35 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("Distributed instance", "id", cfg.Distributed.InstanceID) - // Connect to NATS - natsAuth := cfg.Distributed.NatsAuthConfig() - if natsAuth.RequireAuth && (natsAuth.ServiceUserJWT == "" || natsAuth.ServiceUserSeed == "") { - return nil, fmt.Errorf("LOCALAI_NATS_REQUIRE_AUTH requires LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED") + if authDB == nil { + return nil, fmt.Errorf("distributed mode requires auth database to be initialized first") + } + + // The cluster carrier row says which transport the whole cluster uses. The + // first replica to start writes it; later replicas, and every restart, follow + // it and never their own flags. A deployment that has a NATS URL (the flag, or + // a URL stored in the cluster settings) is seeded on NATS, so an existing + // deployment keeps what it has. One that has only PostgreSQL is seeded on the + // tunnel. + carrierStore, err := cluster.NewCarrierStore(authDB) + if err != nil { + return nil, fmt.Errorf("initializing cluster carrier state: %w", err) } - natsOpts := cfg.Distributed.NatsMessagingOptions("", "") - natsClient, err := messaging.New(cfg.Distributed.NatsURL, natsOpts...) + settingsStore, err := cluster.NewSettingsStore(authDB) if err != nil { - return nil, fmt.Errorf("connecting to NATS: %w", err) + return nil, fmt.Errorf("initializing cluster settings: %w", err) + } + carrierRow, err := seedCarrier(context.Background(), authDB, carrierStore, settingsStore, cfg.Distributed.NatsURL, cfg.Distributed.InstanceID) + if err != nil { + return nil, err } - xlog.Info("Connected to NATS", "url", sanitize.URL(cfg.Distributed.NatsURL)) - // Ensure NATS is closed if any subsequent initialization step fails. + // Whatever fails from here on must give back what was built. success := false + var startSet *carrier.Set defer func() { - if !success { - natsClient.Close() + if !success && startSet != nil && startSet.Close != nil { + startSet.Close() } }() @@ -162,10 +281,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } // Initialize node registry (requires the auth DB which is PostgreSQL) - if authDB == nil { - return nil, fmt.Errorf("distributed mode requires auth database to be initialized first") - } - registry, err := nodes.NewNodeRegistry(authDB) if err != nil { return nil, fmt.Errorf("initializing node registry: %w", err) @@ -207,6 +322,102 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade xlog.Info("Applied declarative model scheduling config", "models", len(schedConfigs)) } + // Initialize file manager with local cache + cacheDir := cfg.DataPath + "/cache" + fileMgr, err := storage.NewFileManager(store, cacheDir) + if err != nil { + return nil, fmt.Errorf("initializing file manager: %w", err) + } + xlog.Info("File manager initialized", "cacheDir", cacheDir) + + // The pieces of the tunnel carrier that do not depend on a set. They exist on + // every replica, whatever the carrier in use: a change of carrier needs them, + // and the routes of the tunnel are registered on every replica. + jobStore, err := jobs.NewJobStore(authDB) + if err != nil { + return nil, fmt.Errorf("initializing job store: %w", err) + } + // The claim queue is the queue of the tunnel carrier. Its table exists on + // every carrier, so that a change of carrier finds it. + if err := jobs.MigrateClaims(context.Background(), authDB); err != nil { + return nil, err + } + xlog.Info("Distributed job store initialized") + + // A replica proves who it is to its peers with a credential that it mints + // itself. Only the hash is published, in the same row as the address that + // the peers dial. + peerCred := cluster.NewPeerCredential() + advertised := peerAddress(cfg) + membership, err := startMembership(cfg.Context, authDB, cfg.Distributed.InstanceID, advertised, peerCred) + if err != nil { + return nil, err + } + defer func() { + if !success { + membership.Stop() + } + }() + clusterReg := cluster.NewRegistry(authDB) + tunnels := tunnel.NewRegistry(clusterReg, cfg.Distributed.InstanceID) + // If a peer sweeps this replica while it stalls, the claims of the tunnels + // that it still holds are written again when it registers again. + membership.SetReclaimer(tunnels) + peerOpts, err := peerPoolOptions(cfg, advertised) + if err != nil { + return nil, err + } + peerPool := tunnel.NewPeerPool(cfg.Distributed.InstanceID, peerCred, clusterReg, peerOpts...) + peerSessions := tunnel.NewPeerSessions(tunnel.NewRelay(tunnels).Stream) + + // Build the set of the carrier that the row names, and put it behind the + // holders. Everything below takes a holder, not a carrier. A NATS server that + // is not up yet does not stop the start; a URL or a credential that cannot be + // used does. + active := &atomic.Pointer[carrier.Set]{} + runtime := &carrierRuntime{ + cfg: cfg, db: authDB, settings: settingsStore, registry: registry, fileMgr: fileMgr, + instanceID: cfg.Distributed.InstanceID, jobStore: jobStore, clusterReg: clusterReg, + tunnels: tunnels, peerPool: peerPool, active: active, + probeWake: make(chan struct{}, 1), + } + startSet, err = runtime.build(context.Background(), carrierRow, carrierRow.Active, false) + if err != nil { + return nil, fmt.Errorf("building the %s carrier that the cluster uses: %w", carrierRow.Active, err) + } + active.Store(startSet) + broadcaster := carrier.NewBroadcaster(active) + runtime.bus = broadcaster + + // While two carriers are attached, a call to a worker goes to the carrier the + // worker is attached to. + workers := nodes.NewSwitchWorkers(registry, cfg.Distributed.StaleNodeThresholdOrDefault()) + window := carrier.NewWindow( + func(ctx context.Context, nodeID string) (carrier.Attachment, error) { + attached, err := workers.AttachedCarriers(ctx, nodeID) + if err != nil { + return carrier.Attachment{}, err + } + return attachmentOf(attached), nil + }, + workers.AgentsAttached, + ) + workQueue := carrier.NewWorkQueue(active) + clientFactory := carrier.NewClients(active) + clientFactory.UseWindow(window) + fileStager := carrier.NewFiles(active) + fileStager.UseWindow(window) + remoteUnloader := carrier.NewCommands(active) + remoteUnloader.UseWindow(window) + workerHTTPDial := carrier.NewRoutedWorkerDialer(active, window) + agentControl := carrier.NewAgents(active) + agentControl.UseWindow(window) + if cfg.Distributed.StorageURL != "" { + xlog.Info("File stager initialized (S3)", "carrier", carrierRow.Active) + } else { + xlog.Info("File stager initialized (HTTP direct transfer)", "carrier", carrierRow.Active) + } + // Collect SmartRouter option values; the router itself is created after all // dependencies (including FileStager and Unloader) are ready. var routerAuthToken string @@ -223,19 +434,11 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade cfg.Distributed.StaleNodeThresholdOrDefault(), routerAuthToken, !cfg.Distributed.DisablePerModelHealthCheck, + clientFactory, ) - // Initialize job store - jobStore, err := jobs.NewJobStore(authDB) - if err != nil { - return nil, fmt.Errorf("initializing job store: %w", err) - } - xlog.Info("Distributed job store initialized") - - workQueue := messaging.NewNATSWorkQueue(natsClient) - // Initialize job dispatcher - dispatcher := jobs.NewDispatcher(jobStore, workQueue, natsClient, authDB, cfg.Distributed.InstanceID) + dispatcher := jobs.NewDispatcher(jobStore, workQueue, broadcaster, authDB, cfg.Distributed.InstanceID) // Initialize agent store agentStore, err := agents.NewAgentStore(authDB) @@ -245,7 +448,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade xlog.Info("Distributed agent store initialized") // Initialize agent event bridge - agentBridge := agents.NewEventBridge(natsClient, agentStore, cfg.Distributed.InstanceID) + agentBridge := agents.NewEventBridge(broadcaster, agentStore, cfg.Distributed.InstanceID) // Start observable persister — captures observable_update events from workers // (which have no DB access) and persists them to PostgreSQL. @@ -261,41 +464,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade return nil, fmt.Errorf("initializing distributed stores: %w", err) } - // Initialize file manager with local cache - cacheDir := cfg.DataPath + "/cache" - fileMgr, err := storage.NewFileManager(store, cacheDir) - if err != nil { - return nil, fmt.Errorf("initializing file manager: %w", err) - } - xlog.Info("File manager initialized", "cacheDir", cacheDir) - - // Create FileStager for distributed file transfer - workerHTTPDial := nodes.DirectWorkerNetDialer() - var fileStager nodes.FileStager - if cfg.Distributed.StorageURL != "" { - fileStager = nodes.NewS3NATSFileStager(fileMgr, natsClient) - xlog.Info("File stager initialized (S3+NATS)") - } else { - fileStager = nodes.NewHTTPFileStager(func(nodeID string) (string, error) { - node, err := registry.Get(context.Background(), nodeID) - if err != nil { - return "", err - } - if node.HTTPAddress == "" { - return "", fmt.Errorf("node %s has no HTTP address for file transfer", nodeID) - } - return node.HTTPAddress, nil - }, cfg.Distributed.RegistrationToken, workerHTTPDial) - xlog.Info("File stager initialized (HTTP direct transfer)") - } - // Create RemoteUnloaderAdapter — needed by SmartRouter and startup.go - remoteUnloader := nodes.NewRemoteUnloaderAdapter( - registry, - natsClient, - cfg.Distributed.BackendInstallTimeoutOrDefault(), - cfg.Distributed.BackendUpgradeTimeoutOrDefault(), - ) - // Prefix-cache-aware routing. Enabled by default; an operator can opt out // with --distributed-prefix-cache=false, which leaves prefixProvider and // pressure nil so the SmartRouter and reconciler behave exactly as the @@ -315,8 +483,8 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade return nil, fmt.Errorf("invalid prefix-cache configuration: %w", err) } idx := prefixcache.NewIndex(prefixCfg) - prefixSync := prefixcache.NewSync(idx, natsClient) - pressure = prefixcache.NewSyncedPressure(prefixCfg.PressureWindow, natsClient) + prefixSync := prefixcache.NewSync(idx, broadcaster) + pressure = prefixcache.NewSyncedPressure(prefixCfg.PressureWindow, broadcaster) prefixProvider = prefixSync // Invalidate the prefix-cache index whenever a replica row is removed. @@ -342,17 +510,17 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // Apply peers' observations/invalidations to the same Sync. ApplyObserve // and ApplyInvalidate update only the local index and do not re-publish, // so there is no broadcast loop. - if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCacheObserve, func(ev messaging.PrefixCacheObserveEvent) { + if _, err := messaging.SubscribeJSON(broadcaster, messaging.SubjectPrefixCacheObserve, func(ev messaging.PrefixCacheObserveEvent) { prefixSync.ApplyObserve(ev, time.Now()) }); err != nil { return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheObserve, err) } - if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCacheInvalidate, func(ev messaging.PrefixCacheInvalidateEvent) { + if _, err := messaging.SubscribeJSON(broadcaster, messaging.SubjectPrefixCacheInvalidate, func(ev messaging.PrefixCacheInvalidateEvent) { prefixSync.ApplyInvalidate(ev) }); err != nil { return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheInvalidate, err) } - if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCachePressure, func(ev messaging.PrefixCachePressureEvent) { + if _, err := messaging.SubscribeJSON(broadcaster, messaging.SubjectPrefixCachePressure, func(ev messaging.PrefixCachePressureEvent) { pressure.ApplyPressure(ev, time.Now()) }); err != nil { return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCachePressure, err) @@ -362,7 +530,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // their real KV state without coupling to router internals. Routing stays // on the guessed provider until a backend producer is available. reportedIndex := prefixcache.NewReportedIndex() - if _, err := messaging.SubscribeJSON(natsClient, messaging.SubjectPrefixCacheResidency, reportedIndex.Apply); err != nil { + if _, err := messaging.SubscribeJSON(broadcaster, messaging.SubjectPrefixCacheResidency, reportedIndex.Apply); err != nil { return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheResidency, err) } @@ -404,9 +572,18 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } modelCleanup := nodes.NewModelCleanupService(registry, remoteUnloader) router := nodes.NewSmartRouter(registry, nodes.SmartRouterOptions{ - Unloader: remoteUnloader, - ModelCleanup: modelCleanup, + Unloader: remoteUnloader, + ModelCleanup: modelCleanup, + // A worker that is attached to another carrier and cannot follow the + // active one is not asked for a replica: its address cannot be dialled. + ActiveCarrier: func() cluster.Carrier { + if set := active.Load(); set != nil { + return set.Name + } + return "" + }, FileStager: fileStager, + ClientFactory: clientFactory, GalleriesJSON: routerGalleriesJSON, AuthToken: routerAuthToken, DB: authDB, @@ -454,8 +631,8 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade "knob", config.FlagDiskHeadroomCheck, "env", "LOCALAI_DISTRIBUTED_DISK_HEADROOM_CHECK") } - router.StagingTracker().SetPublisher(natsClient) - if _, err := router.StagingTracker().SubscribeBroadcasts(natsClient); err != nil { + router.StagingTracker().SetPublisher(broadcaster) + if _, err := router.StagingTracker().SubscribeBroadcasts(broadcaster); err != nil { xlog.Warn("Failed to subscribe to staging progress broadcasts", "error", err) } @@ -468,6 +645,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade Scheduler: router, Unloader: remoteUnloader, Adapter: remoteUnloader, + ClientFactory: clientFactory, RegistrationToken: cfg.Distributed.RegistrationToken, DB: authDB, Interval: 30 * time.Second, @@ -481,11 +659,90 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade // Create ModelRouterAdapter to wire into ModelLoader modelAdapter := nodes.NewModelRouterAdapter(router) + // While the tunnel is the active carrier, a node that heartbeats and holds no + // tunnel is demoted once its departure is older than the grace. On NATS no + // node holds a tunnel, and the monitor reads nothing. + healthMon.UsePresence(clusterReg, cluster.DefaultReconnectGrace, func() bool { + set := active.Load() + return set != nil && set.Name == cluster.CarrierTunnel + }) + + // A worker that reports it is not attached to the active carrier is demoted + // once the cluster has settled: after a forced change it is the worker that + // could not follow. While two carriers are attached the window routes it, so + // nothing is demoted then. The check also runs when a carrier is released, so + // that the demotion does not wait for the next tick. + healthMon.UseCarrier(func() (cluster.Carrier, bool) { + set := active.Load() + if set == nil { + return "", false + } + return set.Name, window.Previous() == nil + }) + broadcaster.OnReconnect(func() { + go healthMon.CheckNow(context.Background()) + }) + + // The protocol of a change of carrier, and what this replica does in one. The + // switch tells the other replicas of each move with a hint, so that they look + // at the row at once. The hint is a courtesy; the poll is what counts. + fallbackTimings := cluster.Timings{ + PrepareTimeout: cfg.Distributed.CarrierPrepareTimeout, + TransitionWindow: cfg.Distributed.CarrierTransitionWindow, + MaxDrain: cfg.Distributed.CarrierMaxDrain, + } + carrierSwitch, err := cluster.NewSwitch(cluster.SwitchOptions{ + NATSURL: runtime.natsURL, + Store: carrierStore, + Registry: clusterReg, + Workers: workers, + Work: workCounts{db: authDB}, + Timings: func() cluster.Timings { + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + t, err := settingsStore.Timings(ctx, fallbackTimings) + if err != nil { + xlog.Warn("Could not read the waits of a change of carrier from the cluster settings; using the ones of this replica", "error", err) + return fallbackTimings + } + return t + }, + AvailabilityMaxAge: availabilityMaxAge, + OnChange: func(row cluster.CarrierRow) { + go func() { + if err := broadcaster.Publish(messaging.SubjectCarrierChanged, row); err != nil { + xlog.Debug("Could not send the hint that the cluster carrier row moved; the other replicas read it on their poll", "error", err) + } + }() + }, + }) + if err != nil { + return nil, err + } + swapper, err := carrier.NewSwapper(carrier.SwapperOptions{ + Cur: active, Bus: broadcaster, Window: window, + Rows: carrierStore, Ready: membership, Build: runtime.Build, + }) + if err != nil { + return nil, err + } + success = true return &DistributedServices{ - Nats: natsClient, + Membership: membership, + Carriers: carrierStore, + Settings: settingsStore, + Switch: carrierSwitch, + Swapper: swapper, + Runtime: runtime, + Tunnels: tunnels, + Instances: clusterReg, + PeerSessions: peerSessions, + PeerPool: peerPool, + WorkerDialer: tunnel.NewWorkerDialer(tunnels, peerPool), + Broadcaster: broadcaster, WorkQueue: workQueue, - AgentControl: nodes.NewNATSAgentControl(natsClient), + AgentControl: agentControl, Store: store, Registry: registry, Router: router, @@ -501,11 +758,292 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade ModelAdapter: modelAdapter, Unloader: remoteUnloader, ModelCleanup: modelCleanup, + Clients: clientFactory, WorkerHTTPDial: workerHTTPDial, + + active: active, + window: window, + workers: workers, }, nil } +// StartCarrierSwitch starts what this replica does about changes of carrier: it +// follows the cluster carrier row, drives the protocol when it holds the +// leadership, and reports which carriers it could build. Call it once, after the +// services that use the carrier have started. +func (ds *DistributedServices) StartCarrierSwitch(ctx context.Context, db *gorm.DB) { + // The hints. Both are a courtesy that shortens a wait. + if _, err := ds.Broadcaster.Subscribe(messaging.SubjectCarrierChanged, func([]byte) { ds.Swapper.Wake() }); err != nil { + xlog.Warn("Could not listen for the hint that the cluster carrier changed; the poll still follows the row", "error", err) + } + if _, err := ds.Broadcaster.Subscribe(messaging.SubjectCarrierProbe, func([]byte) { ds.Runtime.wakeProbe() }); err != nil { + xlog.Warn("Could not listen for the request to report which carriers this replica can build", "error", err) + } + go ds.Swapper.Run(ctx) + go ds.Runtime.RunAvailability(ctx) + // One replica leads the protocol at a time. If it dies, the next tick on + // another replica goes on from the row. + go advisorylock.RunLeaderLoop(ctx, db, advisorylock.KeyCarrierSwitch, leaderTick, func() { + if err := ds.Switch.Drive(ctx); err != nil && ctx.Err() == nil { + xlog.Warn("The change of carrier could not advance", "error", err) + } + }) +} + +// leaderTick is how often the replica that leads makes the move that is due. +const leaderTick = 2 * time.Second + +// seedCarrier reads the cluster carrier row, and writes it when no replica has +// yet. It returns the row as it is after the call. +// +// The NATS address of the cluster is the setting that is stored without a +// credential. The flag of this replica is copied there when none is stored. A +// flag that names another server than the stored address stops the start: the +// stored address is what every replica uses, so a replica that was given another +// one by mistake, or an operator who moved the broker and edited only the flag, +// would otherwise run on the old server and see a warning in a log. +// +// A deployment that had NATS before the row existed, and a cluster whose row says +// NATS, need an address. A replica that has none of it, and that sees no address +// stored, stops its start, because seeding the tunnel would move the whole +// cluster off the broker that its workers use. +// +// The address is stored before the row, so that a replica which reads the row and +// finds NATS also finds the address. +func seedCarrier(ctx context.Context, db *gorm.DB, store *cluster.CarrierStore, settings *cluster.SettingsStore, flagURL, instanceID string) (cluster.CarrierRow, error) { + // A flag that is not an address is not stored. The connection that is opened + // later says what is wrong with it. + flagPublic := "" + if flagURL != "" { + flagPublic = cluster.PublicNATSURL(flagURL) + if cluster.CheckNATSURL(flagPublic) != nil { + flagPublic = "" + } + } + stored, _, err := settings.Get(ctx, cluster.SettingNATSURL) + if err != nil { + return cluster.CarrierRow{}, err + } + if flagPublic != "" && stored != "" && cluster.PublicNATSURL(stored) != flagPublic { + return cluster.CarrierRow{}, fmt.Errorf( + "the NATS address of this replica (--nats-url, %s) differs from the one stored for the cluster (%s). "+ + "Every replica uses the stored address. To move the cluster to the new address, store it with PUT /api/cluster/settings "+ + "({\"nats_url\": ...}) and restart; otherwise set --nats-url to the stored address or remove the flag", + flagPublic, cluster.PublicNATSURL(stored)) + } + + row, err := store.Get(ctx) + switch { + case err == nil: + case errors.Is(err, cluster.ErrNotSeeded): + seedWith := cluster.CarrierTunnel + known := stored + if known == "" { + known = flagPublic + } + switch { + case known != "" || flagURL != "": + seedWith = cluster.CarrierNATS + default: + era, err := natsEraDeployment(ctx, db) + if err != nil { + return cluster.CarrierRow{}, err + } + if era { + return cluster.CarrierRow{}, errNoNATSURL("this cluster already has registered workers or replicas, which use NATS") + } + } + if stored == "" && flagPublic != "" { + if _, err := settings.SetIfAbsent(ctx, cluster.SettingNATSURL, flagPublic, instanceID); err != nil { + return cluster.CarrierRow{}, err + } + stored = flagPublic + } + var created bool + row, created, err = store.SeedWith(ctx, seedWith, instanceID, stored) + if err != nil { + return cluster.CarrierRow{}, fmt.Errorf("reading cluster carrier state: %w", err) + } + if created { + xlog.Info("Cluster carrier seeded", "carrier", row.Active, "epoch", row.Epoch) + } + default: + return cluster.CarrierRow{}, fmt.Errorf("reading cluster carrier state: %w", err) + } + + // The address may have been stored by the replica that won the seed. + if stored == "" { + if stored, _, err = settings.Get(ctx, cluster.SettingNATSURL); err != nil { + return cluster.CarrierRow{}, err + } + } + if stored == "" && flagPublic != "" { + // The URL of the flag is copied into the cluster settings when none is + // stored, so that every replica uses the same one. + made, err := settings.SetIfAbsent(ctx, cluster.SettingNATSURL, flagPublic, instanceID) + if err != nil { + return cluster.CarrierRow{}, err + } + if made { + xlog.Info("The NATS URL of this replica is now a setting of the cluster", "url", flagPublic) + stored = flagPublic + } + } + if stored == "" && flagURL == "" && (row.Active == cluster.CarrierNATS || row.Draining == cluster.CarrierNATS || row.Target == cluster.CarrierNATS) { + return cluster.CarrierRow{}, errNoNATSURL("the cluster carrier row says NATS") + } + if flagURL != "" && flagPublic != flagURL { + xlog.Warn("The NATS URL of this replica carries a credential or a query string. This replica keeps it for its own connection; the cluster stores the address without it", + "url", cluster.PublicNATSURL(flagURL)) + } + if row.Active == cluster.CarrierTunnel && flagURL != "" { + xlog.Info("The cluster runs on the tunnel carrier. The NATS URL of this replica is kept for a change back to NATS, and no NATS connection is opened") + } + return row, nil +} + +func errNoNATSURL(why string) error { + return fmt.Errorf("no NATS address is known: %s, so a replica without one would flip the cluster to the tunnel or sit disconnected. "+ + "Start this replica with --nats-url (LOCALAI_NATS_URL), or store the address with PUT /api/cluster/settings", why) +} + +// natsEraDeployment says whether the database shows a cluster that ran before it +// had a carrier row: a worker or a replica is registered. Such a cluster used +// NATS, because NATS was the only carrier. +func natsEraDeployment(ctx context.Context, db *gorm.DB) (bool, error) { + for _, table := range []string{"backend_nodes", "instances"} { + if !db.Migrator().HasTable(table) { + continue + } + var n int64 + if err := db.WithContext(ctx).Table(table).Limit(1).Count(&n).Error; err != nil { + return false, fmt.Errorf("looking for registered workers: %w", err) + } + if n > 0 { + return true, nil + } + } + return false, nil +} + +// attachmentOf turns the carriers a worker is attached to into the routing form. +func attachmentOf(cs []cluster.Carrier) carrier.Attachment { + var a carrier.Attachment + for _, c := range cs { + switch c { + case cluster.CarrierNATS: + a.NATS = true + case cluster.CarrierTunnel: + a.Tunnel = true + } + } + return a +} + +// startMembership creates the cluster tables, then records this replica in the +// instances table and keeps its row fresh until the Membership is stopped. +// Every replica does this on every carrier, because the change of carrier needs +// to know which replicas are alive and ready. +func startMembership(ctx context.Context, db *gorm.DB, instanceID, advertisedAddr string, peerCred cluster.PeerCredential) (*cluster.Membership, error) { + if ctx == nil { + ctx = context.Background() + } + if err := advisorylock.WithLockCtx(ctx, db, advisorylock.KeySchemaMigrate, func() error { + return cluster.Migrate(ctx, db) + }); err != nil { + return nil, fmt.Errorf("migrating cluster tables: %w", err) + } + membership := cluster.NewMembership(cluster.NewRegistry(db), instanceID, internal.PrintableVersion()) + membership.SetPeer(advertisedAddr, peerCred) + if err := membership.Start(ctx); err != nil { + return nil, fmt.Errorf("registering this replica: %w", err) + } + return membership, nil +} + func isPostgresURL(url string) bool { return strings.HasPrefix(url, "postgres://") || strings.HasPrefix(url, "postgresql://") } + +// peerPoolOptions reads the TLS settings of the peer link. It also warns when the +// link would cross a network in clear text: the credential of this replica and +// every relayed request, prompts included, go over it. +func peerPoolOptions(cfg *config.ApplicationConfig, advertised string) ([]tunnel.PeerPoolOption, error) { + d := cfg.Distributed + if !d.PeerTLS && d.PeerTLSCA == "" { + if addr := advertised; clearTextOnNetwork(addr) { + xlog.Warn("The peer link is clear text and this replica publishes an address that is not on this host, so the credential of this replica and the requests that other replicas relay cross the network unencrypted. "+ + "Put the replicas behind TLS and set LOCALAI_PEER_TLS (and LOCALAI_PEER_TLS_CA for a private CA)", "address", addr) + } + return nil, nil + } + tlsCfg := &tls.Config{MinVersion: tls.VersionTLS12} + if d.PeerTLSCA != "" { + pem, err := os.ReadFile(d.PeerTLSCA) + if err != nil { + return nil, fmt.Errorf("reading the CA file of the peer link: %w", err) + } + roots := x509.NewCertPool() + if !roots.AppendCertsFromPEM(pem) { + return nil, fmt.Errorf("the CA file of the peer link, %s, holds no certificate", d.PeerTLSCA) + } + tlsCfg.RootCAs = roots + } + return []tunnel.PeerPoolOption{tunnel.WithPeerTLS(tlsCfg)}, nil +} + +// clearTextOnNetwork reports whether an address published for the peer link +// names a host other than this one. An empty address is not published, so +// nothing is dialled. +func clearTextOnNetwork(addr string) bool { + if addr == "" { + return false + } + host, _, err := net.SplitHostPort(addr) + if err != nil { + host = addr + } + if strings.EqualFold(host, "localhost") { + return false + } + ip := net.ParseIP(host) + return ip == nil || !ip.IsLoopback() +} + +// peerAddress returns the address at which the other replicas dial this one. An +// address that the operator configured is checked and used as it is. Without +// one, it comes from the route to the database, which all replicas share. +// +// A replica that has no usable address is not refused. It publishes none, and +// the replicas that need to relay through it find it unreachable, which is +// reported when they try. The log line says how to fix it. +func peerAddress(cfg *config.ApplicationConfig) string { + if configured := cfg.Distributed.PeerAddress; configured != "" { + reason, err := cluster.CheckAdvertisedAddr(configured) + if err != nil { + xlog.Warn("The configured peer address is unusable, so this replica publishes none and the other replicas cannot reach it", "address", configured, "error", err) + return "" + } + if reason != "" { + xlog.Warn("The configured peer address may not work for the other replicas", "address", configured, "reason", reason) + } + return configured + } + _, portText, err := net.SplitHostPort(cfg.APIAddress) + if err != nil { + xlog.Warn("The listen address has no port, so this replica publishes no peer address; set LOCALAI_PEER_ADDRESS", "address", cfg.APIAddress, "error", err) + return "" + } + port, err := strconv.Atoi(portText) + if err != nil { + xlog.Warn("The listen port is not a number, so this replica publishes no peer address; set LOCALAI_PEER_ADDRESS", "port", portText, "error", err) + return "" + } + addr, err := cluster.DiscoverAdvertisedAddr(cfg.Auth.DatabaseURL, port) + if err != nil { + xlog.Warn("This replica publishes no peer address, so a worker tunnel that it holds cannot be reached through the other replicas; set LOCALAI_PEER_ADDRESS", "error", err) + return "" + } + return addr +} diff --git a/core/application/distributed_carrier_test.go b/core/application/distributed_carrier_test.go new file mode 100644 index 000000000000..95009a73a71a --- /dev/null +++ b/core/application/distributed_carrier_test.go @@ -0,0 +1,448 @@ +package application + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "os" + "runtime" + "strings" + "sync" + "time" + + "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/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/testcontainers/testcontainers-go" + "github.com/testcontainers/testcontainers-go/wait" + "gorm.io/gorm" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("distributed startup and the cluster carrier", func() { + var ( + db *gorm.DB + cfg *config.ApplicationConfig + ) + + BeforeEach(func() { + db = testutil.SetupTestDB() + ctx, cancel := context.WithCancel(context.Background()) + DeferCleanup(cancel) + cfg = &config.ApplicationConfig{ + Context: ctx, + DataPath: GinkgoT().TempDir(), + Auth: config.AuthConfig{Enabled: true, DatabaseURL: "postgres://unused/db"}, + Distributed: config.DistributedConfig{ + Enabled: true, + InstanceID: "replica-a", + NatsURL: "nats://127.0.0.1:1", + }, + } + }) + + carrierRow := func() cluster.CarrierRow { + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + row, err := store.Get(context.Background()) + Expect(err).ToNot(HaveOccurred()) + return row + } + + It("seeds NATS before it connects, and stops the start on a NATS URL it cannot use", func() { + cfg.Distributed.NatsURL = "nats://[::1" + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).To(MatchError(ContainSubstring("connecting to NATS"))) + Expect(svc).To(BeNil()) + + // The row was written with NATS, because this deployment has a NATS URL. + row := carrierRow() + Expect(row.Active).To(Equal(cluster.CarrierNATS)) + Expect(row.Epoch).To(Equal(int64(1))) + Expect(row.State).To(Equal(cluster.StateStable)) + Expect(row.ChangedBy).To(Equal("replica-a")) + }) + + It("starts a deployment with only PostgreSQL on the tunnel carrier, and opens no NATS connection", func() { + pg, dsn := testutil.SetupTestDBWithDSN() + cfg.Auth.DatabaseURL = dsn + cfg.Distributed.NatsURL = "" + + svc, err := initDistributed(cfg, pg, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + + store, err := cluster.NewCarrierStore(pg) + Expect(err).ToNot(HaveOccurred()) + row, err := store.Get(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(row.Active).To(Equal(cluster.CarrierTunnel)) + Expect(row.Epoch).To(Equal(int64(1))) + Expect(svc.active.Load().Name).To(Equal(cluster.CarrierTunnel)) + + // Fan-out goes through the database, and one LISTEN session exists for it. + var sessions int64 + Expect(pg.Raw("SELECT count(*) FROM pg_stat_activity WHERE application_name LIKE 'localai_pgbus_%'").Scan(&sessions).Error).To(Succeed()) + Expect(sessions).To(Equal(int64(1))) + got := make(chan []byte, 1) + _, err = svc.Broadcaster.Subscribe("jobs.carrier-spec", func(b []byte) { got <- b }) + Expect(err).ToNot(HaveOccurred()) + Expect(svc.Broadcaster.Publish("jobs.carrier-spec", "hello")).To(Succeed()) + Eventually(got, "10s").Should(Receive(Equal([]byte(`"hello"`)))) + }) + + It("starts on the tunnel carrier that the row names, whatever the flags of this replica say, and does not try NATS", func() { + pg, dsn := testutil.SetupTestDBWithDSN() + cfg.Auth.DatabaseURL = dsn + store, err := cluster.NewCarrierStore(pg) + Expect(err).ToNot(HaveOccurred()) + _, _, err = store.Seed(context.Background(), cluster.CarrierTunnel, "replica-b") + Expect(err).ToNot(HaveOccurred()) + + // The NATS URL of the flag points at nothing. It is kept for a change back to + // NATS, and it is not dialled. + svc, err := initDistributed(cfg, pg, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + Expect(svc.active.Load().Name).To(Equal(cluster.CarrierTunnel)) + row, err := store.Get(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(row.Active).To(Equal(cluster.CarrierTunnel), "the row is left as it was") + Expect(row.ChangedBy).To(Equal("replica-b")) + + url, ok, err := svc.Settings.Get(context.Background(), cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(url).To(Equal("nats://127.0.0.1:1"), "the URL of the flag is stored for a change back to NATS") + }) + + It("claims queued work on the tunnel carrier once the switch has started", func() { + pg, dsn := testutil.SetupTestDBWithDSN() + cfg.Auth.DatabaseURL = dsn + cfg.Distributed.NatsURL = "" + svc, err := initDistributed(cfg, pg, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + svc.StartCarrierSwitch(cfg.Context, pg) + + Expect(svc.WorkQueue.Enqueue(context.Background(), messaging.WorkMCPCI, map[string]string{"job_id": "j1"})).To(Succeed()) + // No agent worker exists, so the claim goes back to the pool with a wait. That + // the attempt count moves shows that this replica took the row and drove it. + Eventually(func() int { + var attempts int + _ = pg.Raw("SELECT coalesce(max(attempts), 0) FROM work_claims").Scan(&attempts).Error + return attempts + }, "30s", "200ms").Should(BeNumerically(">=", 1)) + }) + + It("seeds NATS for a deployment that has a NATS URL, and stores the URL for the whole cluster", func() { + // The server at that address is not there. A frontend may start before its + // broker, so the start goes on and the client keeps trying. + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + Expect(svc.active.Load().Name).To(Equal(cluster.CarrierNATS)) + + row := carrierRow() + Expect(row.Active).To(Equal(cluster.CarrierNATS)) + store, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + url, ok, err := store.Get(context.Background(), cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(url).To(Equal("nats://127.0.0.1:1")) + }) + + It("stops the start when the NATS URL of the flag differs from the one stored for the cluster, and names both", func() { + store, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + Expect(store.Set(context.Background(), cluster.SettingNATSURL, "nats://old-broker:4222", "admin")).To(Succeed()) + cfg.Distributed.NatsURL = "nats://new-broker:4222" + + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).To(HaveOccurred()) + Expect(svc).To(BeNil()) + Expect(err.Error()).To(ContainSubstring("nats://new-broker:4222")) + Expect(err.Error()).To(ContainSubstring("nats://old-broker:4222")) + Expect(err.Error()).To(ContainSubstring("/api/cluster/settings"), "it says how to resolve it") + url, _, err := store.Get(context.Background(), cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(url).To(Equal("nats://old-broker:4222"), "a flag does not overwrite a setting of the cluster") + }) + + It("starts on the stored NATS URL when the flag names the same server, and keeps the credential of the flag out of the database", func() { + cfg.Distributed.NatsURL = "nats://user:secret@127.0.0.1:1?token=abc" + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + + store, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + url, ok, err := store.Get(context.Background(), cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(url).To(Equal("nats://127.0.0.1:1")) + Expect(carrierRow().NATSURL).To(Equal("nats://127.0.0.1:1")) + + // A second start with the same flag is the same deployment. + cfg.Distributed.InstanceID = "replica-b" + second, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(second.Shutdown) + }) + + It("does not seed the tunnel in a cluster that has workers, when the replica has no NATS URL", func() { + Expect(db.AutoMigrate(&nodes.BackendNode{})).To(Succeed()) + Expect(db.Create(&nodes.BackendNode{ID: "w1", Name: "w1", NodeType: "backend", Address: "w1:50051"}).Error).To(Succeed()) + cfg.Distributed.NatsURL = "" + + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).To(MatchError(ContainSubstring("--nats-url"))) + Expect(svc).To(BeNil()) + + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + _, err = store.Get(context.Background()) + Expect(err).To(MatchError(cluster.ErrNotSeeded), "the row is not written, so the next start with the flag seeds NATS") + }) + + It("does not start without a NATS URL when the row says NATS", func() { + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + _, _, err = store.Seed(context.Background(), cluster.CarrierNATS, "replica-b") + Expect(err).ToNot(HaveOccurred()) + cfg.Distributed.NatsURL = "" + + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).To(MatchError(ContainSubstring("the cluster carrier row says NATS"))) + Expect(svc).To(BeNil()) + }) + + It("starts without a NATS flag on the URL that the cluster stored", func() { + settings, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + Expect(settings.Set(context.Background(), cluster.SettingNATSURL, "nats://127.0.0.1:1", "admin")).To(Succeed()) + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + _, _, err = store.Seed(context.Background(), cluster.CarrierNATS, "replica-b") + Expect(err).ToNot(HaveOccurred()) + cfg.Distributed.NatsURL = "" + + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + Expect(svc.active.Load().Name).To(Equal(cluster.CarrierNATS)) + }) + + Context("with a NATS server", func() { + BeforeEach(func() { + url, err := startNATS() + if err != nil { + // A CI runner that lost Docker must go red, not skip. Local runs + // without Docker, and macOS CI, skip. + if os.Getenv("CI") != "" && runtime.GOOS != "darwin" { + Fail("testcontainers requires Docker and CI is set: " + err.Error()) + } + Skip("testcontainers requires Docker: " + err.Error()) + } + cfg.Distributed.NatsURL = url + }) + + It("seeds NATS on first start and wires every seam to it", func() { + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + + Expect(carrierRow().Active).To(Equal(cluster.CarrierNATS)) + Expect(svc.Broadcaster).ToNot(BeNil()) + Expect(svc.WorkQueue).ToNot(BeNil()) + Expect(svc.AgentControl).ToNot(BeNil()) + Expect(svc.FileStager).ToNot(BeNil()) + Expect(svc.Unloader).ToNot(BeNil()) + Expect(svc.WorkerHTTPDial).ToNot(BeNil()) + + // A message published through the holder reaches a subscriber made + // through the same holder, over the real server. + got := make(chan []byte, 1) + sub, err := svc.Broadcaster.Subscribe("jobs.carrier-spec", func(b []byte) { got <- b }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = sub.Unsubscribe() }) + Expect(svc.Broadcaster.Publish("jobs.carrier-spec", "hello")).To(Succeed()) + Eventually(got).Should(Receive(Equal([]byte(`"hello"`)))) + }) + + It("registers every replica in the instances table and removes it on shutdown", func() { + live := func() []string { + GinkgoHelper() + var ids []string + Expect(db.Raw(cluster.LiveInstanceIDsSQL+" ORDER BY id", cluster.InstanceLiveness.Seconds()).Scan(&ids).Error).To(Succeed()) + return ids + } + + first, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(first.Shutdown) + cfg2 := *cfg + cfg2.Distributed.InstanceID = "replica-b" + cfg2.DataPath = GinkgoT().TempDir() + second, err := initDistributed(&cfg2, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(second.Shutdown) + Expect(live()).To(ConsistOf("replica-a", "replica-b")) + + first.Shutdown() + Expect(live()).To(ConsistOf("replica-b")) + }) + + It("publishes the address and the peer credential of the replica in its row", func() { + cfg.Distributed.PeerAddress = "10.0.0.7:8080" + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + + inst, err := cluster.NewRegistry(db).Get(context.Background(), "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(inst.AdvertisedAddr).To(Equal("10.0.0.7:8080")) + Expect(inst.PeerTokenHash).To(HaveLen(64), "only the hash of the credential is published") + Expect(svc.PeerPool).ToNot(BeNil()) + Expect(svc.PeerSessions).ToNot(BeNil()) + Expect(svc.WorkerDialer).ToNot(BeNil()) + }) + + It("starts without an address when the operator gave none and the database route cannot give one", func() { + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + inst, err := cluster.NewRegistry(db).Get(context.Background(), "replica-a") + Expect(err).ToNot(HaveOccurred()) + Expect(inst.AdvertisedAddr).To(BeEmpty()) + Expect(inst.PeerTokenHash).ToNot(BeEmpty()) + }) + + It("ends the tunnel that it holds for a node that was removed, and says whether it held one", func() { + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(svc.Shutdown) + + Expect(svc.Disconnect("nobody")).To(BeFalse()) + + front, back := tunnelSessionPair() + _, err = svc.Tunnels.Attach(context.Background(), "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + Expect(svc.Disconnect("w1")).To(BeTrue()) + Eventually(back.CloseChan()).Should(BeClosed()) + _, _, err = cluster.NewRegistry(db).Owner(context.Background(), "w1") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("closes the tunnels it holds before it leaves the instances table", func() { + svc, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + front, back := tunnelSessionPair() + _, err = svc.Tunnels.Attach(context.Background(), "w2", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + + svc.Shutdown() + + Eventually(back.CloseChan()).Should(BeClosed()) + _, _, err = cluster.NewRegistry(db).Owner(context.Background(), "w2") + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("starts a second replica on the same row without changing it", func() { + first, err := initDistributed(cfg, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(first.Shutdown) + + cfg2 := *cfg + cfg2.Distributed.InstanceID = "replica-b" + cfg2.DataPath = GinkgoT().TempDir() + second, err := initDistributed(&cfg2, db, nil, nil) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(second.Shutdown) + + row := carrierRow() + Expect(row.Epoch).To(Equal(int64(1))) + Expect(row.ChangedBy).To(Equal("replica-a")) + }) + }) +}) + +// tunnelSessionPair returns the two ends of a tunnel session over a real +// websocket. +func tunnelSessionPair() (front, back *tunnel.Session) { + GinkgoHelper() + up := tunnel.NewUpgrader() + got := make(chan *tunnel.Session, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ws, err := up.Upgrade(w, r, nil) + if err != nil { + return + } + sess, err := tunnel.ServerSession(ws, tunnel.LaneInference) + if err != nil { + _ = ws.Close() + return + } + got <- sess + })) + DeferCleanup(srv.Close) + ws, _, err := tunnel.NewDialer(5*time.Second).Dial("ws"+strings.TrimPrefix(srv.URL, "http"), nil) + Expect(err).ToNot(HaveOccurred()) + back, err = tunnel.ClientSession(ws, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = back.Close() }) + Eventually(got).Should(Receive(&front)) + return front, back +} + +var ( + natsOnce sync.Once + natsURL string + natsCtr testcontainers.Container + natsErr error +) + +// startNATS starts one server for the process and returns its URL. +func startNATS() (string, error) { + natsOnce.Do(func() { + ctx := context.Background() + var ctr testcontainers.Container + ctr, natsErr = testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{ + ContainerRequest: testcontainers.ContainerRequest{ + Image: "nats:2.10-alpine", + ExposedPorts: []string{"4222/tcp"}, + WaitingFor: wait.ForListeningPort("4222/tcp"), + }, + Started: true, + }) + if natsErr != nil { + return + } + natsCtr = ctr + host, err := ctr.Host(ctx) + if err != nil { + natsErr = err + return + } + port, err := ctr.MappedPort(ctx, "4222/tcp") + if err != nil { + natsErr = err + return + } + natsURL = fmt.Sprintf("nats://%s:%s", host, port.Port()) + }) + return natsURL, natsErr +} + +var _ = AfterSuite(func() { + if natsCtr != nil { + _ = natsCtr.Terminate(context.Background()) + } +}) diff --git a/core/application/failover_distributed.go b/core/application/failover_distributed.go index 87adbdeb51d7..128aa566a8de 100644 --- a/core/application/failover_distributed.go +++ b/core/application/failover_distributed.go @@ -63,7 +63,7 @@ func (a *Application) startFailoverDistributed(ctx context.Context) { xlog.Error("failover: pins will not persist, could not prepare the pin store", "error", err) pins = nil // distsync.New treats a nil store as "no durable pins" } - s, err := distsync.New(ctx, a.distributed.Nats, pins, a.failoverManager) + s, err := distsync.New(ctx, a.distributed.Broadcaster, pins, a.failoverManager) if err != nil { xlog.Error("failover: state will not be shared between frontends", "error", err) return diff --git a/core/application/nats_handover_test.go b/core/application/nats_handover_test.go new file mode 100644 index 000000000000..a4f27be56207 --- /dev/null +++ b/core/application/nats_handover_test.go @@ -0,0 +1,70 @@ +package application + +import ( + "context" + "os" + "path/filepath" + + "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("What workers are told about NATS", func() { + var ( + ctx context.Context + settings *cluster.SettingsStore + cfg *config.ApplicationConfig + rt *carrierRuntime + ) + + BeforeEach(func() { + ctx = context.Background() + var err error + settings, err = cluster.NewSettingsStore(testutil.SetupTestDB()) + Expect(err).ToNot(HaveOccurred()) + cfg = &config.ApplicationConfig{Distributed: config.DistributedConfig{NatsURL: "nats://frontend-flag:4222"}} + rt = &carrierRuntime{cfg: cfg, settings: settings} + }) + + It("is the address for workers when one is stored", func() { + Expect(settings.Set(ctx, cluster.SettingNATSURL, "nats://frontends:4222", "spec")).To(Succeed()) + Expect(settings.Set(ctx, cluster.SettingNATSWorkerURL, "nats://workers.example:4222", "spec")).To(Succeed()) + Expect(rt.WorkerURL(ctx)).To(Equal("nats://workers.example:4222")) + }) + + It("falls back to the address of the frontends, then to the flag of this replica", func() { + Expect(rt.WorkerURL(ctx)).To(Equal("nats://frontend-flag:4222")) + Expect(settings.Set(ctx, cluster.SettingNATSURL, "nats://frontends:4222", "spec")).To(Succeed()) + Expect(rt.WorkerURL(ctx)).To(Equal("nats://frontends:4222")) + }) + + It("is empty when the deployment has no NATS", func() { + cfg.Distributed.NatsURL = "" + Expect(rt.WorkerURL(ctx)).To(BeEmpty()) + }) + + It("hands over the CA of the server and says nothing when there is none", func() { + Expect(rt.CAPEM()).To(BeEmpty()) + path := filepath.Join(GinkgoT().TempDir(), "ca.pem") + Expect(os.WriteFile(path, []byte("-----BEGIN CERTIFICATE-----\nx\n-----END CERTIFICATE-----\n"), 0o600)).To(Succeed()) + cfg.Distributed.NatsTLSCA = path + Expect(rt.CAPEM()).To(ContainSubstring("BEGIN CERTIFICATE")) + }) + + It("fails to hand over a CA file it cannot read, and does not make up one", func() { + cfg.Distributed.NatsTLSCA = filepath.Join(GinkgoT().TempDir(), "missing.pem") + _, err := rt.CAPEM() + Expect(err).To(HaveOccurred()) + }) + + It("says that workers need a certificate when this replica holds one", func() { + Expect(rt.ClientTLS()).To(BeFalse()) + cfg.Distributed.NatsTLSCert = "/some/client.pem" + cfg.Distributed.NatsTLSKey = "/some/client.key" + Expect(rt.ClientTLS()).To(BeTrue()) + }) +}) diff --git a/core/application/peer_tls_test.go b/core/application/peer_tls_test.go new file mode 100644 index 000000000000..7185e332f788 --- /dev/null +++ b/core/application/peer_tls_test.go @@ -0,0 +1,53 @@ +package application + +import ( + "os" + "path/filepath" + + "github.com/mudler/LocalAI/core/config" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("The TLS settings of the peer link", func() { + DescribeTable("warns about clear text only for an address that is not on this host", + func(addr string, want bool) { + Expect(clearTextOnNetwork(addr)).To(Equal(want)) + }, + Entry("nothing published", "", false), + Entry("loopback", "127.0.0.1:8080", false), + Entry("loopback v6", "[::1]:8080", false), + Entry("localhost", "localhost:8080", false), + Entry("a private address", "10.1.2.3:8080", true), + Entry("a public address", "203.0.113.9:443", true), + Entry("a name", "frontend-2.internal:8080", true), + ) + + It("adds no option and keeps the link plain when TLS is off", func() { + opts, err := peerPoolOptions(&config.ApplicationConfig{}, "10.1.2.3:8080") + Expect(err).ToNot(HaveOccurred()) + Expect(opts).To(BeEmpty()) + }) + + It("adds the TLS option when it is on, with the system roots", func() { + cfg := &config.ApplicationConfig{} + cfg.Distributed.PeerTLS = true + opts, err := peerPoolOptions(cfg, "10.1.2.3:8080") + Expect(err).ToNot(HaveOccurred()) + Expect(opts).To(HaveLen(1)) + }) + + It("refuses a CA file that holds no certificate, and one that is missing", func() { + cfg := &config.ApplicationConfig{} + cfg.Distributed.PeerTLSCA = filepath.Join(GinkgoT().TempDir(), "missing.pem") + _, err := peerPoolOptions(cfg, "") + Expect(err).To(HaveOccurred()) + + empty := filepath.Join(GinkgoT().TempDir(), "empty.pem") + Expect(os.WriteFile(empty, []byte("not a certificate"), 0o600)).To(Succeed()) + cfg.Distributed.PeerTLSCA = empty + _, err = peerPoolOptions(cfg, "") + Expect(err).To(HaveOccurred()) + }) +}) diff --git a/core/application/seed_carrier_test.go b/core/application/seed_carrier_test.go new file mode 100644 index 000000000000..b0da81bb386f --- /dev/null +++ b/core/application/seed_carrier_test.go @@ -0,0 +1,75 @@ +package application + +import ( + "context" + "fmt" + "sync" + + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("seeding the carrier of the cluster", func() { + It("lets a replica without a NATS flag start while replicas with one seed, whatever the order", func() { + for round := 0; round < 5; round++ { + db := testutil.SetupTestDB() + ctx := context.Background() + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + settings, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + + var wg sync.WaitGroup + errs := make([]error, 12) + start := make(chan struct{}) + for i := range errs { + flag := "" + if i%2 == 0 { + flag = "nats://broker:4222" + } + wg.Add(1) + go func() { + defer wg.Done() + <-start + _, errs[i] = seedCarrier(ctx, db, store, settings, flag, fmt.Sprintf("replica-%d", i)) + }() + } + close(start) + wg.Wait() + for i, err := range errs { + Expect(err).ToNot(HaveOccurred(), "replica %d", i) + } + + // Whoever won, the row and the address agree: a row that says NATS + // always has an address next to it. + row, err := store.Get(ctx) + Expect(err).ToNot(HaveOccurred()) + if row.Active == cluster.CarrierNATS { + url, ok, err := settings.Get(ctx, cluster.SettingNATSURL) + Expect(err).ToNot(HaveOccurred()) + Expect(ok).To(BeTrue()) + Expect(url).To(Equal("nats://broker:4222")) + } + } + }) + + It("finds the address stored, and not the row alone, when a flag-less replica reads between the two writes of another", func() { + db := testutil.SetupTestDB() + ctx := context.Background() + store, err := cluster.NewCarrierStore(db) + Expect(err).ToNot(HaveOccurred()) + settings, err := cluster.NewSettingsStore(db) + Expect(err).ToNot(HaveOccurred()) + + // The first write of the replica with the flag has happened; the second has not. + _, err = settings.SetIfAbsent(ctx, cluster.SettingNATSURL, "nats://broker:4222", "replica-a") + Expect(err).ToNot(HaveOccurred()) + + row, err := seedCarrier(ctx, db, store, settings, "", "replica-b") + Expect(err).ToNot(HaveOccurred()) + Expect(row.Active).To(Equal(cluster.CarrierNATS), "an address is stored, so the cluster is on NATS and not on the tunnel") + Expect(row.NATSURL).To(Equal("nats://broker:4222")) + }) +}) diff --git a/core/application/startup.go b/core/application/startup.go index 78d0a5b99fe7..af3926456c59 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -310,6 +310,7 @@ func New(opts ...config.AppOption) (*Application, error) { distStore := nodes.NewDistributedModelStore( model.NewInMemoryModelStore(), distSvc.Registry, + nodes.WithClientFactory(distSvc.Clients), ) application.modelLoader.SetModelStore(distStore) // Drop the local stub when a model's last replica leaves the registry. @@ -338,6 +339,10 @@ func New(opts ...config.AppOption) (*Application, error) { if err := distSvc.Dispatcher.Start(options.Context); err != nil { return nil, fmt.Errorf("starting job dispatcher: %w", err) } + // Follow the cluster carrier row, lead the protocol of a change of carrier when + // this replica holds the leadership, and report which carriers it could build. + // It starts after the services that use the carrier. + distSvc.StartCarrierSwitch(options.Context, application.authDB) // Start ephemeral file cleanup storage.StartEphemeralCleanup(options.Context, distSvc.FileMgr, 0, 0) // Wire distributed backends into AgentJobService (before Start) @@ -346,14 +351,14 @@ func New(opts ...config.AppOption) (*Application, error) { application.agentJobService.SetDistributedJobStore(distSvc.JobStore) // Keep agent tasks consistent across replicas (jobs already sync via the // dispatcher + DB read-through). Same NATS client the dispatcher uses. - application.agentJobService.SetTaskSyncNATS(distSvc.Nats) + application.agentJobService.SetTaskSyncNATS(distSvc.Broadcaster) } // Wire skill store into AgentPoolService (wired at pool start time via closure) // The actual wiring happens in StartAgentPool since the pool doesn't exist yet. // Wire NATS and gallery store into GalleryService for cross-instance progress/cancel if application.galleryService != nil { - application.galleryService.SetNATSClient(distSvc.Nats) + application.galleryService.SetNATSClient(distSvc.Broadcaster) if distSvc.DistStores != nil && distSvc.DistStores.Gallery != nil { // Clean up stale in-progress operations from previous crashed instances if _, err := distSvc.DistStores.Gallery.CleanStale(30 * time.Minute); err != nil { @@ -520,7 +525,7 @@ func New(opts ...config.AppOption) (*Application, error) { // frontend runs its own pass against the shared models directory; the // pass is idempotent, so no leader is needed. if modelConfigResync != nil && distSvc != nil { - modelConfigResync.Start(options.Context, options.Distributed.ModelConfigResyncIntervalOrDefault(), distSvc.Nats) + modelConfigResync.Start(options.Context, options.Distributed.ModelConfigResyncIntervalOrDefault(), distSvc.Broadcaster) } if options.PreloadJSONModels != "" { diff --git a/core/cli/agent_worker.go b/core/cli/agent_worker.go index 9c5422a51f7b..fc4d533c0c8d 100644 --- a/core/cli/agent_worker.go +++ b/core/cli/agent_worker.go @@ -15,11 +15,12 @@ import ( "github.com/mudler/LocalAI/core/cli/workerregistry" "github.com/mudler/LocalAI/core/config" mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp" - "github.com/mudler/LocalAI/core/services/agents" + "github.com/mudler/LocalAI/core/services/agentworker" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/jobs" mcpRemote "github.com/mudler/LocalAI/core/services/mcp" "github.com/mudler/LocalAI/core/services/messaging" - "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/worker" "github.com/mudler/LocalAI/internal" "github.com/mudler/LocalAI/pkg/sanitize" "github.com/mudler/cogito" @@ -28,17 +29,19 @@ import ( ) // AgentWorkerCMD starts a dedicated agent worker process for distributed mode. -// It registers with the frontend, subscribes to the NATS agent execution queue, -// and executes agent chats using cogito. The worker is a pure executor — it -// receives the full agent config and skills in the NATS job payload, so it -// does not need direct database access. +// It registers with the frontend, attaches to the carrier that the frontend +// names, and executes agent chats using cogito. On NATS it subscribes to the +// agent execution queue. On the tunnel it holds one outbound connection to the +// frontend and serves the requests that arrive on it. The worker is a pure +// executor — it receives the full agent config and skills in the job payload, +// so it does not need direct database access. // // Usage: // // localai agent-worker --nats-url nats://... --register-to http://localai:8080 type AgentWorkerCMD struct { - // NATS (required) - NatsURL string `env:"LOCALAI_NATS_URL" required:"" help:"NATS server URL" group:"distributed"` + // NATS (optional: the frontend hands over the address) + NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL. Optional: the frontend hands over the address when the cluster runs on NATS, and this overrides it" group:"distributed"` // Registration (required) RegisterTo string `env:"LOCALAI_REGISTER_TO" required:"" help:"Frontend URL for registration" group:"registration"` @@ -88,6 +91,28 @@ func validateAgentSubject(subject string) error { return nil } +// agentCarrier maps the carrier that the frontend named onto the one this agent +// worker attaches to, and reports whether it is the tunnel. A frontend that names +// none predates carriers and runs on NATS. The flag of the NATS URL cannot be +// required, because the carrier is known only after the registration, so the +// check is here. +func agentCarrier(named, natsURL, handedURL string) (onTunnel bool, err error) { + switch named { + case "tunnel": + if natsURL != "" { + xlog.Info("The cluster runs on the tunnel; the NATS URL of this agent worker is kept for a change to NATS") + } + return true, nil + case "", "nats": + if natsURL == "" && handedURL == "" { + return false, fmt.Errorf("the cluster runs on NATS and this agent worker has no NATS URL, and the frontend handed over none: set LOCALAI_NATS_URL") + } + return false, nil + default: + return false, fmt.Errorf("the frontend names the carrier %q, which this agent worker does not know: upgrade the agent worker", named) + } +} + func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { if err := validateAgentSubject(cmd.Subject); err != nil { return err @@ -108,14 +133,17 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { hostname, _ := os.Hostname() nodeName = "agent-" + hostname } - registrationBody := map[string]any{ - "name": nodeName, - "node_type": "agent", - "version": internal.Version, - "commit": internal.Commit, - } - if cmd.RegistrationToken != "" { - registrationBody["token"] = cmd.RegistrationToken + registrationBody := func() map[string]any { + body := map[string]any{ + "name": nodeName, + "node_type": "agent", + "version": internal.Version, + "commit": internal.Commit, + } + if cmd.RegistrationToken != "" { + body["token"] = cmd.RegistrationToken + } + return body } // Context cancelled on shutdown — used by registration waits, heartbeat, and @@ -126,11 +154,12 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { // Acquire credentials via (re)registration. When the bus requires auth and no // static fallback is configured, wait through admin approval until the // frontend mints credentials rather than starting unauthenticated. - credMgr := workerregistry.NewNATSCredentialManager( + staticNATS := cmd.NatsJWT != "" || cmd.NatsServiceJWT != "" + credMgr := workerregistry.NewCredentialManager( func(ctx context.Context) (*workerregistry.RegisterResponse, error) { - return regClient.RegisterFull(ctx, registrationBody) + return regClient.RegisterFull(ctx, registrationBody()) }, - cmd.natsAuthRequired() && cmd.NatsJWT == "" && cmd.NatsServiceJWT == "", + cmd.natsAuthRequired() && !staticNATS, ) res, err := credMgr.Acquire(shutdownCtx) if err != nil { @@ -139,143 +168,117 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error { nodeID := res.ID xlog.Info("Registered with frontend", "nodeID", nodeID, "frontend", cmd.RegisterTo) + onTunnel, err := agentCarrier(res.Carrier, cmd.NatsURL, res.NatsURL) + if err != nil { + return err + } + bootCarrier := cluster.CarrierNATS + if onTunnel { + bootCarrier = cluster.CarrierTunnel + } + // Use provisioned API token if none was set if cmd.APIToken == "" { cmd.APIToken = res.APIToken } - // Start heartbeat heartbeatInterval, err := time.ParseDuration(cmd.HeartbeatInterval) if err != nil && cmd.HeartbeatInterval != "" { xlog.Warn("invalid heartbeat interval, using default 10s", "input", cmd.HeartbeatInterval, "error", err) } heartbeatInterval = cmp.Or(heartbeatInterval, 10*time.Second) - go regClient.HeartbeatLoop(shutdownCtx, nodeID, heartbeatInterval, func() map[string]any { return map[string]any{} }) - - // Resolve NATS credentials with precedence: explicit env override, then - // frontend-minted (auto-refreshed before expiry), then service fallback. - // Each static source must supply JWT and seed together. - natsTLS := messaging.TLSFiles{CA: cmd.NatsTLSCA, Cert: cmd.NatsTLSCert, Key: cmd.NatsTLSKey} - var natsOpts []messaging.Option - switch { - case cmd.NatsJWT != "" || cmd.NatsUserSeed != "": - if (cmd.NatsJWT == "") != (cmd.NatsUserSeed == "") { - return fmt.Errorf("LOCALAI_NATS_JWT and LOCALAI_NATS_USER_SEED must be set together") - } - natsOpts = append(natsOpts, messaging.WithUserJWT(cmd.NatsJWT, cmd.NatsUserSeed)) - case credMgr.HasCredentials(): - natsOpts = append(natsOpts, messaging.WithUserJWTProvider(credMgr.Provider())) - go func() { - if err := credMgr.RefreshLoop(shutdownCtx); err != nil { - xlog.Error("NATS credential refresh permanently failed; shutting down agent worker", "error", err) - shutdownCancel() - } - }() - case cmd.NatsServiceJWT != "" || cmd.NatsServiceSeed != "": - if (cmd.NatsServiceJWT == "") != (cmd.NatsServiceSeed == "") { - return fmt.Errorf("LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED must be set together") - } - natsOpts = append(natsOpts, messaging.WithUserJWT(cmd.NatsServiceJWT, cmd.NatsServiceSeed)) - case cmd.natsAuthRequired(): - return fmt.Errorf("NATS JWT+seed required: enable frontend minting or set LOCALAI_NATS_* env vars") - } - if natsTLS.Enabled() { - natsOpts = append(natsOpts, messaging.WithTLS(natsTLS)) - } - natsClient, err := messaging.New(cmd.NatsURL, natsOpts...) - if err != nil { - return fmt.Errorf("connecting to NATS: %w", err) - } - defer natsClient.Close() - - // Create event bridge for publishing results back via NATS - eventBridge := agents.NewEventBridge(natsClient, nil, "agent-worker-"+nodeID) - - // Start cancel listener - cancelSub, err := eventBridge.StartCancelListener() - if err != nil { - xlog.Warn("Failed to start cancel listener", "error", err) - } else { - defer cancelSub.Unsubscribe() - } - - // One consumer serves both queued kinds; the route option only moves the - // agent-run subject and group, which operators may set. - work := messaging.NewNATSWorkConsumer(natsClient, messaging.WithAgentRunRoute(cmd.Subject, cmd.Queue)) - - // Create and start the NATS dispatcher. - // No ConfigProvider or SkillStore needed — config and skills arrive in the job payload. - dispatcher := agents.NewNATSDispatcher( - work, - eventBridge, - nil, // no ConfigProvider: config comes in the enriched NATS payload - apiURL, cmd.APIToken, - 0, // no concurrency limit (CLI worker) - ) - - if err := dispatcher.Start(shutdownCtx); err != nil { - return fmt.Errorf("starting dispatcher: %w", err) - } - - var rpc agentRPCServer = nodes.NewNATSAgentRPCServer(natsClient, nodeID) - - // Serve MCP tool execution requests (load-balanced across workers). - // The frontend routes model-level MCP tool calls here. - if err := rpc.ServeMCPTool(handleMCPToolRequest); err != nil { - return err - } - - // Serve MCP discovery requests (load-balanced across workers). - if err := rpc.ServeMCPDiscovery(handleMCPDiscoveryRequest); err != nil { - return err - } - - // Subscribe to MCP CI job execution (load-balanced across agent workers). - // In distributed mode, MCP CI jobs are routed here because the frontend - // cannot create MCP sessions (e.g., stdio servers using docker). mcpCIJobTimeout, err := time.ParseDuration(cmd.MCPCIJobTimeout) if err != nil && cmd.MCPCIJobTimeout != "" { xlog.Warn("invalid MCP CI job timeout, using default 10m", "input", cmd.MCPCIJobTimeout, "error", err) } mcpCIJobTimeout = cmp.Or(mcpCIJobTimeout, config.DefaultMCPCIJobTimeout) - if _, err := startMCPCIConsumer(shutdownCtx, work, apiURL, cmd.APIToken, mcpCIJobTimeout); err != nil { - return err + // Resolve NATS credentials with precedence: explicit env override, then + // frontend-minted (auto-refreshed before expiry), then service fallback. + // Each static source must supply JWT and seed together. + if (cmd.NatsJWT == "") != (cmd.NatsUserSeed == "") { + return fmt.Errorf("LOCALAI_NATS_JWT and LOCALAI_NATS_USER_SEED must be set together") + } + if (cmd.NatsServiceJWT == "") != (cmd.NatsServiceSeed == "") { + return fmt.Errorf("LOCALAI_NATS_SERVICE_JWT and LOCALAI_NATS_SERVICE_SEED must be set together") + } + local := worker.NATSLocal{ + URL: cmd.NatsURL, JWT: cmd.NatsJWT, Seed: cmd.NatsUserSeed, Required: cmd.natsAuthRequired(), + TLS: messaging.TLSFiles{CA: cmd.NatsTLSCA, Cert: cmd.NatsTLSCert, Key: cmd.NatsTLSKey}, + FallbackJWT: cmd.NatsServiceJWT, + FallbackSeed: cmd.NatsServiceSeed, + } + serve := agentworker.FollowConfig{ + NodeID: nodeID, FrontendURL: cmd.RegisterTo, ControlToken: cmd.RegistrationToken, + Local: local, Subject: cmd.Subject, Queue: cmd.Queue, + APIURL: apiURL, APIToken: cmd.APIToken, + // Model-level MCP tool calls and discovery from the frontend. The worker + // creates and caches the MCP sessions from the serialized config. + Tool: handleMCPToolRequest, + Discovery: handleMCPDiscoveryRequest, + // The agent worker has no model loader, so it drops the sessions cached for + // a backend that goes away when it hears the stop of that backend. + BackendStop: func(backend string) { + if backend != "" { + mcpTools.CloseMCPSessions(backend) + } + }, + // MCP CI jobs run here because the frontend cannot create MCP sessions + // (for example stdio servers that use docker). + Jobs: func(ctx context.Context, _ cluster.Carrier, consumer messaging.WorkConsumer) error { + _, err := startMCPCIConsumer(ctx, consumer, apiURL, cmd.APIToken, mcpCIJobTimeout) + return err + }, } - // Listen for backend stop events to clean up cached MCP sessions. - // In the main application this is done via ml.OnModelUnload, but the agent - // worker has no model loader, so it listens for the stop event instead. - if err := rpc.ServeBackendStop(func(backend string) { - if backend != "" { - mcpTools.CloseMCPSessions(backend) - } - }); err != nil { - return err + carriers := worker.NewFollower(worker.FollowerOptions{ + Attachers: serve.Attachers(), + Credentials: func(c cluster.Carrier) *workerregistry.CredentialManager { + return workerregistry.NewCredentialManagerFor( + func(ctx context.Context) (*workerregistry.RegisterResponse, error) { + body := registrationBody() + body["carrier"] = string(c) + return regClient.RegisterFull(ctx, body) + }, c == cluster.CarrierNATS && cmd.natsAuthRequired() && !staticNATS, string(c)) + }, + Heartbeat: func(ctx context.Context, body map[string]any) (*workerregistry.HeartbeatReply, error) { + return regClient.HeartbeatFull(ctx, nodeID, body) + }, + Cannot: serve.Cannot(), + Interval: heartbeatInterval, + }) + defer carriers.Close() + if _, err := carriers.Attach(shutdownCtx, bootCarrier, credMgr, res); err != nil { + return fmt.Errorf("attaching to the %s carrier: %w", bootCarrier, err) } + go carriers.Run(shutdownCtx) + xlog.Info("Agent worker ready, waiting for jobs", "carrier", bootCarrier, "subject", cmd.Subject, "queue", cmd.Queue) - xlog.Info("Agent worker ready, waiting for jobs", "subject", cmd.Subject, "queue", cmd.Queue) + runErr := awaitShutdown(shutdownCtx, "the carrier stopped") + xlog.Info("Shutting down agent worker") + shutdownCancel() // stop heartbeat loop immediately + mcpTools.CloseAllMCPSessions() + regClient.GracefulDeregister(nodeID) + return runErr +} - // Wait for an OS signal or an internal fatal condition (e.g. NATS - // credentials became unrenewable), so the worker restarts and re-acquires - // rather than lingering unable to serve. +// awaitShutdown blocks until the process is told to stop or the context of the +// worker ends for a fatal reason, and returns an error in the second case so +// that the worker restarts and registers again instead of lingering unable to +// serve. +func awaitShutdown(shutdownCtx context.Context, fatal string) error { sigCh := make(chan os.Signal, 1) signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM) - var runErr error + defer signal.Stop(sigCh) select { case <-sigCh: + return nil case <-shutdownCtx.Done(): - runErr = fmt.Errorf("agent worker shutting down: NATS credentials unavailable") - xlog.Error("Internal shutdown requested", "error", runErr) + err := fmt.Errorf("agent worker shutting down: %s", fatal) + xlog.Error("Internal shutdown requested", "error", err) + return err } - - xlog.Info("Shutting down agent worker") - shutdownCancel() // stop heartbeat loop immediately - dispatcher.Stop() - mcpTools.CloseAllMCPSessions() - regClient.GracefulDeregister(nodeID) - return runErr } // startMCPCIConsumer serves MCP CI jobs with maxInFlight 1, which keeps them @@ -286,14 +289,6 @@ func startMCPCIConsumer(ctx context.Context, consumer messaging.WorkConsumer, ap }) } -// agentRPCServer is how the agent worker serves the frontend's MCP requests -// and hears the node's backend stop events. -type agentRPCServer interface { - ServeMCPTool(h mcpRemote.ToolHandler) error - ServeMCPDiscovery(h mcpRemote.DiscoveryHandler) error - ServeBackendStop(h func(backend string)) error -} - // handleMCPToolRequest executes one MCP tool call. The worker creates/caches // MCP sessions from the serialized config and executes the tool. func handleMCPToolRequest(parent context.Context, req mcpRemote.MCPToolRequest) mcpRemote.MCPToolResponse { diff --git a/core/cli/agent_worker_subject_test.go b/core/cli/agent_worker_subject_test.go index 3f5f14fa3344..012c8abbef94 100644 --- a/core/cli/agent_worker_subject_test.go +++ b/core/cli/agent_worker_subject_test.go @@ -29,3 +29,39 @@ var _ = Describe("validateAgentSubject", func() { Expect(err.Error()).To(ContainSubstring("LOCALAI_AGENT_SUBJECT")) }) }) + +var _ = Describe("agentCarrier", func() { + It("attaches to NATS when the frontend names it or names nothing, and the URL is there", func() { + for _, named := range []string{"nats", ""} { + onTunnel, err := agentCarrier(named, "nats://n:4222", "") + Expect(err).ToNot(HaveOccurred(), named) + Expect(onTunnel).To(BeFalse(), named) + } + }) + + It("asks for the URL when the cluster runs on NATS and there is none", func() { + _, err := agentCarrier("nats", "", "") + Expect(err).To(MatchError(ContainSubstring("LOCALAI_NATS_URL"))) + _, err = agentCarrier("", "", "") + Expect(err).To(MatchError(ContainSubstring("LOCALAI_NATS_URL"))) + }) + + It("takes the URL the frontend hands over when the worker has none", func() { + onTunnel, err := agentCarrier("nats", "", "nats://handed:4222") + Expect(err).ToNot(HaveOccurred()) + Expect(onTunnel).To(BeFalse()) + }) + + It("attaches to the tunnel when the frontend names it, with a URL or without one", func() { + for _, url := range []string{"", "nats://n:4222"} { + onTunnel, err := agentCarrier("tunnel", url, "") + Expect(err).ToNot(HaveOccurred()) + Expect(onTunnel).To(BeTrue()) + } + }) + + It("asks for an upgrade when the frontend names a carrier it does not know", func() { + _, err := agentCarrier("smoke-signals", "nats://n:4222", "") + Expect(err).To(MatchError(ContainSubstring("upgrade"))) + }) +}) diff --git a/core/cli/cli.go b/core/cli/cli.go index 2b17ff96d5d0..b1ded4bda482 100644 --- a/core/cli/cli.go +++ b/core/cli/cli.go @@ -2,6 +2,7 @@ package cli import ( "github.com/mudler/LocalAI/core/cli/benchmark" + "github.com/mudler/LocalAI/core/cli/cluster" cliContext "github.com/mudler/LocalAI/core/cli/context" "github.com/mudler/LocalAI/core/cli/worker" ) @@ -10,6 +11,7 @@ var CLI struct { cliContext.Context `embed:""` Benchmark benchmark.Command `cmd:"" help:"Benchmark configured text models against a running LocalAI server"` + Cluster cluster.Command `cmd:"" help:"Read or change the carrier and the settings of a distributed cluster through a running frontend"` Run RunCMD `cmd:"" help:"Run LocalAI, this the default command if no other command is specified. Run 'local-ai run --help' for more information" default:"withargs"` Chat ChatCMD `cmd:"" help:"Run the built-in terminal agent against a LocalAI server"` Federated FederatedCLI `cmd:"" help:"Run LocalAI in federated mode"` diff --git a/core/cli/cluster/carrier.go b/core/cli/cluster/carrier.go new file mode 100644 index 000000000000..565dc8b158d7 --- /dev/null +++ b/core/cli/cluster/carrier.go @@ -0,0 +1,300 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "bufio" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "os" + "strings" + "text/tabwriter" + "time" + + "github.com/mudler/LocalAI/core/services/cluster" +) + +const carrierPath = "/api/cluster/carrier" + +// CarrierCommand groups the commands about the carrier. +type CarrierCommand struct { + Status StatusCommand `cmd:"" default:"withargs" help:"Show the active carrier, the state of a change, the replicas and the workers"` + Switch SwitchCommand `cmd:"" help:"Change the carrier of the cluster, or check what would block the change (--dry-run)"` + Abort AbortCommand `cmd:"" help:"Abort a change that is being prepared"` +} + +// StatusCommand shows the state of the carrier. +type StatusCommand struct { + Remote `embed:""` +} + +// SwitchCommand changes the carrier. +type SwitchCommand struct { + To string `required:"" enum:"nats,tunnel" help:"Carrier to change to [${enum}]."` + DryRun bool `name:"dry-run" help:"Run the preflight and change nothing. It lists what would block the change."` + Force bool `help:"Go ahead although some workers cannot follow or some replicas are not ready. Those workers become unroutable until they are fixed."` + Yes bool `help:"Confirm that the replicas that are listed are all the frontends of the cluster, and that all run this release. Without it the command asks, and a command with no terminal refuses."` + Wait bool `help:"Wait until the change is committed, and print the progress."` + // WaitTimeout bounds --wait. + WaitTimeout time.Duration `name:"wait-timeout" default:"10m" help:"How long --wait waits."` + Remote `embed:""` +} + +// AbortCommand aborts a change in the prepare state. +type AbortCommand struct { + Remote `embed:""` +} + +// pollInterval is how often --wait reads the state. +var pollInterval = 2 * time.Second + +func (c *StatusCommand) Run() error { + ctx, stop := signalContext() + defer stop() + return c.run(ctx, os.Stdout) +} + +func (c *StatusCommand) run(ctx context.Context, out io.Writer) error { + var view cluster.Report + _, raw, err := c.do(ctx, c.client(), http.MethodGet, carrierPath, nil, &view) + if err != nil { + return err + } + if c.JSON { + return printJSON(out, raw) + } + printStatus(out, view) + return nil +} + +func (c *AbortCommand) Run() error { + ctx, stop := signalContext() + defer stop() + return c.run(ctx, os.Stdout) +} + +func (c *AbortCommand) run(ctx context.Context, out io.Writer) error { + var row cluster.CarrierRow + _, raw, err := c.do(ctx, c.client(), http.MethodPost, carrierPath, map[string]any{"abort": true}, &row) + if err != nil { + return err + } + if c.JSON { + return printJSON(out, raw) + } + _, _ = fmt.Fprintf(out, "Aborted. The cluster stays on %s (epoch %d).\n", row.Active, row.Epoch) + return nil +} + +func (c *SwitchCommand) Run() error { + ctx, stop := signalContext() + defer stop() + return c.run(ctx, os.Stdout, os.Stdin, stdinIsTerminal()) +} + +func stdinIsTerminal() bool { + fi, err := os.Stdin.Stat() + return err == nil && fi.Mode()&os.ModeCharDevice != 0 +} + +func (c *SwitchCommand) run(ctx context.Context, out io.Writer, in io.Reader, interactive bool) error { + client := c.client() + target := c.To + + // The dry run is the preflight. A change runs it again on the server, so + // this one is for the person who reads the replica list. + var report cluster.Report + _, raw, err := c.do(ctx, client, http.MethodPost, carrierPath, map[string]any{"target": target, "dry_run": true}, &report) + if err != nil { + return err + } + if c.DryRun { + if c.JSON { + return printJSON(out, raw) + } + printPreflight(out, report, target) + if !report.OK { + return errors.New("the change would be blocked; fix the blockers, or use --force for those that can be forced") + } + return nil + } + + if !c.JSON { + printPreflight(out, report, target) + } + if !report.OK && !c.Force { + return errors.New("the change is blocked; fix the blockers, or use --force for those that can be forced") + } + if err := c.confirm(out, in, interactive, report); err != nil { + return err + } + + var started struct { + State string `json:"state"` + Active string `json:"active"` + Target string `json:"target"` + Epoch int64 `json:"epoch"` + } + _, raw, err = c.do(ctx, client, http.MethodPost, carrierPath, map[string]any{"target": target, "force": c.Force}, &started) + if err != nil { + var ae *apiError + if errors.As(err, &ae) && ae.Status == http.StatusUnprocessableEntity { + var body struct { + Blockers []cluster.Blocker `json:"blockers"` + } + _ = json.Unmarshal(ae.Body, &body) + var b strings.Builder + for _, bl := range body.Blockers { + fmt.Fprintf(&b, "\n %s %s: %s", bl.Kind, bl.ID, bl.Reason) + } + return fmt.Errorf("the change is blocked:%s", b.String()) + } + return err + } + if c.JSON && !c.Wait { + return printJSON(out, raw) + } + _, _ = fmt.Fprintf(out, "Change to %s started (epoch %d, state %s).\n", started.Target, started.Epoch, started.State) + if !c.Wait { + _, _ = fmt.Fprintln(out, "The replicas carry it out. Run `local-ai cluster carrier status` to follow it.") + return nil + } + return c.wait(ctx, client, out, started.Epoch) +} + +// confirm asks the person to confirm the replica list, which the server cannot +// check: a frontend of an older release writes no row and is not in the list. +func (c *SwitchCommand) confirm(out io.Writer, in io.Reader, interactive bool, report cluster.Report) error { + if c.Yes { + return nil + } + if !interactive { + return errors.New("confirm that the replicas listed above are all the frontends of the cluster and that all run this release: pass --yes") + } + q := fmt.Sprintf("Are these %d replicas all the frontends of the cluster, and do all of them run this release?", len(report.Replicas)) + if c.Force { + q += " --force makes the workers that cannot follow unroutable." + } + _, _ = fmt.Fprintf(out, "%s [y/N] ", q) + line, _ := bufio.NewReader(in).ReadString('\n') + switch strings.ToLower(strings.TrimSpace(line)) { + case "y", "yes": + return nil + } + return errors.New("not confirmed; nothing was changed") +} + +// wait reads the state until the change is committed or gone. +func (c *SwitchCommand) wait(ctx context.Context, client *http.Client, out io.Writer, epoch int64) error { + ctx, cancel := context.WithTimeout(ctx, c.WaitTimeout) + defer cancel() + last := "" + for { + var st cluster.Report + if _, _, err := c.do(ctx, client, http.MethodGet, carrierPath, nil, &st); err != nil { + if ctx.Err() != nil { + return fmt.Errorf("gave up waiting; the change goes on, check it with `local-ai cluster carrier status`: %w", ctx.Err()) + } + return err + } + line := fmt.Sprintf("state %s, active %s, epoch %d", st.State, st.Active, st.Epoch) + if line != last { + _, _ = fmt.Fprintln(out, line) + last = line + } + switch { + case st.State == cluster.StateStable && string(st.Active) == c.To: + _, _ = fmt.Fprintf(out, "The cluster is on %s.\n", c.To) + if st.DrainRemaining > 0 { + _, _ = fmt.Fprintf(out, "The previous carrier drains for %s more.\n", st.DrainRemaining.Round(time.Second)) + } + return nil + case st.State == cluster.StateStable && st.Epoch > epoch: + note := st.Row.Note + if note != "" { + note = ": " + note + } + return fmt.Errorf("the change was aborted; the cluster stays on %s%s", st.Active, note) + } + select { + case <-ctx.Done(): + return fmt.Errorf("gave up waiting; the change goes on, check it with `local-ai cluster carrier status`: %w", ctx.Err()) + case <-time.After(pollInterval): + } + } +} + +func printStatus(out io.Writer, r cluster.Report) { + _, _ = fmt.Fprintf(out, "Active carrier: %s (epoch %d)\n", r.Active, r.Epoch) + _, _ = fmt.Fprintf(out, "State: %s", r.State) + if r.Target != "" { + _, _ = fmt.Fprintf(out, " (changing to %s)", r.Target) + } + _, _ = fmt.Fprintln(out) + if r.Row.Draining != "" { + _, _ = fmt.Fprintf(out, "Draining: %s", r.Row.Draining) + if r.DrainRemaining > 0 { + _, _ = fmt.Fprintf(out, ", %s left", r.DrainRemaining.Round(time.Second)) + } + _, _ = fmt.Fprintln(out) + } + printReplicas(out, r) + printWorkers(out, r) +} + +func printReplicas(out io.Writer, r cluster.Report) { + _, _ = fmt.Fprintf(out, "\nReplicas (%d):\n", len(r.Replicas)) + tw := tabwriter.NewWriter(out, 0, 4, 2, ' ', 0) + _, _ = fmt.Fprintln(tw, " ID\tVERSION\tREADY EPOCH\tNOTE") + for _, rep := range r.Replicas { + _, _ = fmt.Fprintf(tw, " %s\t%s\t%d\t%s\n", rep.ID, rep.Version, rep.ReadyEpoch, rep.ReadyReason) + } + _ = tw.Flush() +} + +func printWorkers(out io.Writer, r cluster.Report) { + _, _ = fmt.Fprintf(out, "\nWorkers (%d):\n", len(r.Workers)) + tw := tabwriter.NewWriter(out, 0, 4, 2, ' ', 0) + _, _ = fmt.Fprintln(tw, " NAME\tID\tATTACHED\tCAN FOLLOW\tNOTE") + for _, w := range r.Workers { + var att []string + for _, a := range w.Attached { + att = append(att, string(a)) + } + note := w.Reason + if w.FollowError != "" { + if note != "" { + note += "; " + } + note += "follow_error: " + w.FollowError + } + _, _ = fmt.Fprintf(tw, " %s\t%s\t%s\t%t\t%s\n", w.Name, w.ID, strings.Join(att, ","), w.CanFollow, note) + } + _ = tw.Flush() +} + +func printPreflight(out io.Writer, r cluster.Report, target string) { + _, _ = fmt.Fprintf(out, "Preflight for a change from %s to %s:\n", r.Active, target) + if r.OK { + _, _ = fmt.Fprintln(out, " Nothing blocks the change.") + } + for _, b := range r.Blockers { + force := "" + if b.Forceable { + force = " (can be forced)" + } + _, _ = fmt.Fprintf(out, " BLOCKER %s %s: %s%s\n", b.Kind, b.ID, b.Reason, force) + } + for _, w := range r.Warnings { + _, _ = fmt.Fprintf(out, " warning: %s\n", w) + } + _, _ = fmt.Fprintf(out, " In flight: %d loads, %d jobs, %d pending claims, %d claimed\n", + r.InFlight.Loads, r.InFlight.Jobs, r.InFlight.PendingClaims, r.InFlight.ClaimedClaims) + printReplicas(out, r) + _, _ = fmt.Fprintln(out, " Frontends of an older release write no row and are not listed. Check that this list is complete.") + printWorkers(out, r) +} diff --git a/core/cli/cluster/cluster.go b/core/cli/cluster/cluster.go new file mode 100644 index 000000000000..1c50e6c898af --- /dev/null +++ b/core/cli/cluster/cluster.go @@ -0,0 +1,148 @@ +// SPDX-License-Identifier: MIT + +// Package cluster is the "local-ai cluster" command. It calls the admin API of +// a running distributed frontend to read or change the carrier of the cluster +// and to read or store the cluster settings. +package cluster + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "fmt" + "io" + "net/http" + "net/url" + "os" + "os/signal" + "strings" + "syscall" + "time" + + "github.com/mudler/LocalAI/pkg/httpclient" +) + +// Command is the "local-ai cluster" command group. +type Command struct { + Carrier CarrierCommand `cmd:"" help:"Read the carrier of a distributed cluster, or change it"` + Settings SettingsCommand `cmd:"" help:"Read or store the settings of a distributed cluster"` +} + +// Remote says which server a command talks to. +type Remote struct { + Endpoint string `default:"http://127.0.0.1:8080" env:"LOCALAI_ENDPOINT" help:"URL of a LocalAI frontend of the cluster."` + APIKey string `name:"api-key" env:"LOCALAI_API_KEY,API_KEY" help:"API key of an admin user."` + Timeout time.Duration `default:"2m" help:"Timeout for each request. A dry run asks every replica, so it can take a few seconds."` + JSON bool `name:"json" help:"Write the answer of the server as JSON."` +} + +// apiError is an answer that is not a success. +type apiError struct { + Status int + Msg string + // Body is the decoded body, for the callers that read more than the message. + Body []byte +} + +func (e *apiError) Error() string { + switch e.Status { + case http.StatusUnauthorized, http.StatusForbidden: + return fmt.Sprintf("the server refused the request (HTTP %d): use --api-key with the key of an admin user", e.Status) + case http.StatusServiceUnavailable: + return "the server is not running in distributed mode, or the cluster is not ready: " + e.Msg + } + if e.Msg == "" { + return fmt.Sprintf("HTTP status %d", e.Status) + } + return fmt.Sprintf("%s (HTTP %d)", e.Msg, e.Status) +} + +// endpointURL joins the base URL and the path of an API route. +func endpointURL(base, path string) (string, error) { + u, err := url.Parse(base) + if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || u.Fragment != "" || strings.Contains(base, "#") { + return "", errors.New("endpoint must be an HTTP(S) URL without credentials, query, or fragment") + } + u.Path = strings.TrimRight(u.Path, "/") + path + u.RawPath = "" + return u.String(), nil +} + +// do sends one request and decodes a JSON answer into out. It returns the status +// code and the raw body as well. A status outside 2xx is an *apiError. +func (r *Remote) do(ctx context.Context, client *http.Client, method, path string, in, out any) (int, []byte, error) { + target, err := endpointURL(r.Endpoint, path) + if err != nil { + return 0, nil, err + } + var body io.Reader + if in != nil { + raw, err := json.Marshal(in) + if err != nil { + return 0, nil, err + } + body = bytes.NewReader(raw) + } + req, err := http.NewRequestWithContext(ctx, method, target, body) + if err != nil { + return 0, nil, errors.New("cannot create the request") + } + req.Header.Set("Accept", "application/json") + if in != nil { + req.Header.Set("Content-Type", "application/json") + } + if r.APIKey != "" { + req.Header.Set("Authorization", "Bearer "+r.APIKey) + } + resp, err := client.Do(req) + if err != nil { + if ctx.Err() != nil { + return 0, nil, ctx.Err() + } + // A transport error can echo the URL and the credentials of a proxy. + return 0, nil, errors.New("cannot reach the server at the endpoint") + } + defer func() { _ = resp.Body.Close() }() + raw, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20)) + if err != nil { + return resp.StatusCode, nil, errors.New("cannot read the answer") + } + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + var e struct { + Error string `json:"error"` + } + _ = json.Unmarshal(raw, &e) + return resp.StatusCode, raw, &apiError{Status: resp.StatusCode, Msg: e.Error, Body: raw} + } + if out != nil { + if err := json.Unmarshal(raw, out); err != nil { + return resp.StatusCode, raw, fmt.Errorf("the server answered with a body that is not valid: %w", err) + } + } + return resp.StatusCode, raw, nil +} + +func (r *Remote) client() *http.Client { + timeout := r.Timeout + if timeout <= 0 { + timeout = 2 * time.Minute + } + return httpclient.NewWithTimeout(timeout) +} + +// printJSON writes raw, indented. +func printJSON(w io.Writer, raw []byte) error { + var buf bytes.Buffer + if err := json.Indent(&buf, raw, "", " "); err != nil { + _, err = w.Write(raw) + return err + } + buf.WriteByte('\n') + _, err := w.Write(buf.Bytes()) + return err +} + +func signalContext() (context.Context, context.CancelFunc) { + return signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) +} diff --git a/core/cli/cluster/cluster_test.go b/core/cli/cluster/cluster_test.go new file mode 100644 index 000000000000..416d3f867ff3 --- /dev/null +++ b/core/cli/cluster/cluster_test.go @@ -0,0 +1,416 @@ +// SPDX-License-Identifier: MIT +package cluster + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/alecthomas/kong" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestClusterCLI(t *testing.T) { RegisterFailHandler(Fail); RunSpecs(t, "Cluster CLI") } + +// fakeServer records the requests of a command and answers from a table. +type fakeServer struct { + mu sync.Mutex + requests []recorded + handler func(r *http.Request, body map[string]any) (int, any) + srv *httptest.Server +} + +type recorded struct { + Method, Path, Auth string + Body map[string]any +} + +func newFakeServer(h func(r *http.Request, body map[string]any) (int, any)) *fakeServer { + f := &fakeServer{handler: h} + f.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var body map[string]any + _ = json.NewDecoder(r.Body).Decode(&body) + f.mu.Lock() + f.requests = append(f.requests, recorded{r.Method, r.URL.Path, r.Header.Get("Authorization"), body}) + f.mu.Unlock() + code, resp := f.handler(r, body) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(code) + _ = json.NewEncoder(w).Encode(resp) + })) + return f +} + +func (f *fakeServer) last() recorded { + f.mu.Lock() + defer f.mu.Unlock() + return f.requests[len(f.requests)-1] +} + +func (f *fakeServer) count() int { + f.mu.Lock() + defer f.mu.Unlock() + return len(f.requests) +} + +func remote(f *fakeServer) Remote { + return Remote{Endpoint: f.srv.URL, APIKey: "k", Timeout: 5 * time.Second} +} + +var preflightOK = map[string]any{ + "active": "nats", "epoch": 3, "state": "stable", "ok": true, + "replicas": []map[string]any{{"id": "r1", "version": "v1", "ready_epoch": 3}, {"id": "r2", "version": "v1", "ready_epoch": 3}}, + "workers": []map[string]any{{"id": "w1", "name": "gpu", "attached": []string{"nats"}, "can_follow": true}}, + "in_flight": map[string]any{"loads": 1}, +} + +var preflightBlocked = map[string]any{ + "active": "nats", "epoch": 3, "state": "stable", "ok": false, + "blockers": []map[string]any{{"kind": "worker", "id": "w2", "reason": "no routable address", "forceable": true}}, + "replicas": []map[string]any{{"id": "r1", "version": "v1", "ready_epoch": 3}}, + "workers": []map[string]any{{"id": "w2", "name": "edge", "attached": []string{"nats"}, "can_follow": false, "reason": "no routable address", "follow_error": "needs address"}}, +} + +var _ = Describe("local-ai cluster", func() { + var out bytes.Buffer + BeforeEach(func() { out.Reset() }) + + Describe("parsing", func() { + It("reads the subcommands and their flags", func() { + var c struct { + Cluster Command `cmd:""` + } + parser, err := kong.New(&c) + Expect(err).NotTo(HaveOccurred()) + + ctx, err := parser.Parse([]string{"cluster", "carrier", "switch", "--to", "tunnel", "--dry-run", "--force", "--yes", "--wait"}) + Expect(err).NotTo(HaveOccurred()) + Expect(ctx.Command()).To(Equal("cluster carrier switch")) + sw := c.Cluster.Carrier.Switch + Expect(sw.To).To(Equal("tunnel")) + Expect(sw.DryRun && sw.Force && sw.Yes && sw.Wait).To(BeTrue()) + Expect(sw.Endpoint).To(Equal("http://127.0.0.1:8080")) + + _, err = parser.Parse([]string{"cluster", "carrier", "switch", "--to", "kafka"}) + Expect(err).To(HaveOccurred()) + var bare struct { + Cluster Command `cmd:""` + } + bareParser, err := kong.New(&bare) + Expect(err).NotTo(HaveOccurred()) + _, err = bareParser.Parse([]string{"cluster", "carrier", "switch"}) + Expect(err).To(HaveOccurred()) + + ctx, err = parser.Parse([]string{"cluster", "carrier", "abort"}) + Expect(err).NotTo(HaveOccurred()) + Expect(ctx.Command()).To(Equal("cluster carrier abort")) + + ctx, err = parser.Parse([]string{"cluster", "carrier"}) + Expect(err).NotTo(HaveOccurred()) + Expect(ctx.Command()).To(Equal("cluster carrier status")) + }) + + It("tells an absent setting from an empty one", func() { + var c struct { + Cluster Command `cmd:""` + } + parser, err := kong.New(&c) + Expect(err).NotTo(HaveOccurred()) + _, err = parser.Parse([]string{"cluster", "settings", "set", "--nats-url=", "--max-drain", "20m"}) + Expect(err).NotTo(HaveOccurred()) + set := c.Cluster.Settings.Set + Expect(set.NATSURL).NotTo(BeNil()) + Expect(*set.NATSURL).To(BeEmpty()) + Expect(set.NATSWorkerURL).To(BeNil()) + Expect(*set.MaxDrain).To(Equal("20m")) + }) + }) + + Describe("carrier status", func() { + It("prints the carrier, the replicas and the workers, and sends the key", func() { + f := newFakeServer(func(r *http.Request, _ map[string]any) (int, any) { + return 200, map[string]any{ + "active": "tunnel", "epoch": 9, "state": "stable", + "row": map[string]any{"draining": "nats"}, + "drain_remaining_ns": int64(90 * time.Second), + "replicas": preflightOK["replicas"], "workers": preflightBlocked["workers"], + } + }) + defer f.srv.Close() + c := &StatusCommand{Remote: remote(f)} + Expect(c.run(context.Background(), &out)).To(Succeed()) + Expect(f.last().Method).To(Equal("GET")) + Expect(f.last().Path).To(Equal("/api/cluster/carrier")) + Expect(f.last().Auth).To(Equal("Bearer k")) + s := out.String() + Expect(s).To(ContainSubstring("Active carrier: tunnel (epoch 9)")) + Expect(s).To(ContainSubstring("Draining: nats, 1m30s left")) + Expect(s).To(ContainSubstring("Replicas (2)")) + Expect(s).To(ContainSubstring("follow_error: needs address")) + }) + + It("writes the JSON of the server with --json", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 200, preflightOK }) + defer f.srv.Close() + r := remote(f) + r.JSON = true + Expect((&StatusCommand{Remote: r}).run(context.Background(), &out)).To(Succeed()) + var got map[string]any + Expect(json.Unmarshal(out.Bytes(), &got)).To(Succeed()) + Expect(got["active"]).To(Equal("nats")) + }) + + It("explains a refused key and a frontend that is not distributed", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 401, map[string]any{"error": "no"} }) + err := (&StatusCommand{Remote: remote(f)}).run(context.Background(), &out) + f.srv.Close() + Expect(err).To(MatchError(ContainSubstring("--api-key"))) + + f = newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 503, map[string]any{"error": "distributed mode is not enabled on this frontend"} + }) + defer f.srv.Close() + err = (&StatusCommand{Remote: remote(f)}).run(context.Background(), &out) + Expect(err).To(MatchError(ContainSubstring("not running in distributed mode"))) + }) + + It("refuses an endpoint that is not a plain HTTP URL", func() { + err := (&StatusCommand{Remote: Remote{Endpoint: "http://u:p@host:1"}}).run(context.Background(), &out) + Expect(err).To(MatchError(ContainSubstring("without credentials"))) + }) + }) + + Describe("carrier switch", func() { + It("only runs the preflight with --dry-run", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 200, preflightOK }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", DryRun: true, Remote: remote(f)} + Expect(c.run(context.Background(), &out, strings.NewReader(""), false)).To(Succeed()) + Expect(f.count()).To(Equal(1)) + Expect(f.last().Body).To(HaveKeyWithValue("dry_run", true)) + Expect(f.last().Body).To(HaveKeyWithValue("target", "tunnel")) + Expect(out.String()).To(ContainSubstring("Nothing blocks the change.")) + Expect(out.String()).To(ContainSubstring("Check that this list is complete.")) + }) + + It("fails a dry run that has a blocker and lists it", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 200, preflightBlocked }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", DryRun: true, Remote: remote(f)} + err := c.run(context.Background(), &out, strings.NewReader(""), false) + Expect(err).To(MatchError(ContainSubstring("blocked"))) + Expect(out.String()).To(ContainSubstring("BLOCKER worker w2: no routable address (can be forced)")) + }) + + It("refuses to start a change without a confirmation when there is no terminal", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 200, preflightOK }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Remote: remote(f)} + err := c.run(context.Background(), &out, strings.NewReader(""), false) + Expect(err).To(MatchError(ContainSubstring("--yes"))) + Expect(f.count()).To(Equal(1), "only the dry run was sent") + }) + + It("does not start a change when the answer to the question is no", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { return 200, preflightOK }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Remote: remote(f)} + err := c.run(context.Background(), &out, strings.NewReader("n\n"), true) + Expect(err).To(MatchError(ContainSubstring("nothing was changed"))) + Expect(f.count()).To(Equal(1)) + Expect(out.String()).To(ContainSubstring("Are these 2 replicas all the frontends")) + }) + + It("starts the change after a yes at the question", func() { + f := newFakeServer(func(_ *http.Request, b map[string]any) (int, any) { + if b["dry_run"] == true { + return 200, preflightOK + } + return 202, map[string]any{"state": "prepare", "active": "nats", "target": "tunnel", "epoch": 4} + }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Remote: remote(f)} + Expect(c.run(context.Background(), &out, strings.NewReader("yes\n"), true)).To(Succeed()) + Expect(f.last().Body).To(HaveKeyWithValue("target", "tunnel")) + Expect(f.last().Body).To(HaveKeyWithValue("force", false)) + Expect(f.last().Body).NotTo(HaveKey("dry_run")) + Expect(out.String()).To(ContainSubstring("Change to tunnel started (epoch 4, state prepare).")) + }) + + It("stops at a blocker unless --force, and sends force when given", func() { + var forced bool + f := newFakeServer(func(_ *http.Request, b map[string]any) (int, any) { + if b["dry_run"] == true { + return 200, preflightBlocked + } + forced = b["force"] == true + return 202, map[string]any{"state": "prepare", "active": "nats", "target": "tunnel", "epoch": 4} + }) + defer f.srv.Close() + err := (&SwitchCommand{To: "tunnel", Yes: true, Remote: remote(f)}).run(context.Background(), &out, strings.NewReader(""), false) + Expect(err).To(MatchError(ContainSubstring("blocked"))) + Expect(f.count()).To(Equal(1)) + + Expect((&SwitchCommand{To: "tunnel", Yes: true, Force: true, Remote: remote(f)}).run(context.Background(), &out, strings.NewReader(""), false)).To(Succeed()) + Expect(forced).To(BeTrue()) + }) + + It("shows the blockers of a 422 and the busy state of a 409", func() { + f := newFakeServer(func(_ *http.Request, b map[string]any) (int, any) { + if b["dry_run"] == true { + return 200, preflightOK + } + return 422, map[string]any{"error": "blocked", "blockers": []map[string]any{{"kind": "replica", "id": "r9", "reason": "not ready"}}} + }) + err := (&SwitchCommand{To: "tunnel", Yes: true, Remote: remote(f)}).run(context.Background(), &out, strings.NewReader(""), false) + f.srv.Close() + Expect(err).To(MatchError(ContainSubstring("replica r9: not ready"))) + + f = newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 409, map[string]any{"error": "a change is under way"} + }) + defer f.srv.Close() + err = (&SwitchCommand{To: "tunnel", DryRun: true, Remote: remote(f)}).run(context.Background(), &out, strings.NewReader(""), false) + Expect(err).To(MatchError(ContainSubstring("a change is under way"))) + }) + + It("waits for the commit with --wait", func() { + old := pollInterval + pollInterval = time.Millisecond + DeferCleanup(func() { pollInterval = old }) + var polls int + f := newFakeServer(func(r *http.Request, b map[string]any) (int, any) { + switch { + case r.Method == http.MethodGet: + polls++ + if polls < 3 { + return 200, map[string]any{"active": "nats", "state": "prepare", "epoch": 4} + } + return 200, map[string]any{"active": "tunnel", "state": "stable", "epoch": 6, "drain_remaining_ns": int64(15 * time.Minute)} + case b["dry_run"] == true: + return 200, preflightOK + } + return 202, map[string]any{"state": "prepare", "active": "nats", "target": "tunnel", "epoch": 4} + }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Yes: true, Wait: true, WaitTimeout: 5 * time.Second, Remote: remote(f)} + Expect(c.run(context.Background(), &out, strings.NewReader(""), false)).To(Succeed()) + Expect(out.String()).To(ContainSubstring("The cluster is on tunnel.")) + Expect(out.String()).To(ContainSubstring("drains for 15m0s more")) + }) + + It("reports an abort that --wait sees", func() { + old := pollInterval + pollInterval = time.Millisecond + DeferCleanup(func() { pollInterval = old }) + f := newFakeServer(func(r *http.Request, b map[string]any) (int, any) { + switch { + case r.Method == http.MethodGet: + return 200, map[string]any{"active": "nats", "state": "stable", "epoch": 5, "row": map[string]any{"note": "replica r2 not ready"}} + case b["dry_run"] == true: + return 200, preflightOK + } + return 202, map[string]any{"state": "prepare", "active": "nats", "target": "tunnel", "epoch": 4} + }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Yes: true, Wait: true, WaitTimeout: 5 * time.Second, Remote: remote(f)} + err := c.run(context.Background(), &out, strings.NewReader(""), false) + Expect(err).To(MatchError(ContainSubstring("aborted"))) + Expect(err).To(MatchError(ContainSubstring("replica r2 not ready"))) + }) + + It("gives up waiting when the time is out", func() { + old := pollInterval + pollInterval = time.Millisecond + DeferCleanup(func() { pollInterval = old }) + f := newFakeServer(func(r *http.Request, b map[string]any) (int, any) { + switch { + case r.Method == http.MethodGet: + return 200, map[string]any{"active": "nats", "state": "prepare", "epoch": 4} + case b["dry_run"] == true: + return 200, preflightOK + } + return 202, map[string]any{"state": "prepare", "active": "nats", "target": "tunnel", "epoch": 4} + }) + defer f.srv.Close() + c := &SwitchCommand{To: "tunnel", Yes: true, Wait: true, WaitTimeout: 30 * time.Millisecond, Remote: remote(f)} + Expect(c.run(context.Background(), &out, strings.NewReader(""), false)).To(MatchError(ContainSubstring("gave up waiting"))) + }) + }) + + Describe("carrier abort", func() { + It("posts abort and says where the cluster stays", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 200, map[string]any{"active": "nats", "epoch": 8, "state": "stable"} + }) + defer f.srv.Close() + Expect((&AbortCommand{Remote: remote(f)}).run(context.Background(), &out)).To(Succeed()) + Expect(f.last().Body).To(HaveKeyWithValue("abort", true)) + Expect(out.String()).To(ContainSubstring("The cluster stays on nats (epoch 8).")) + }) + + It("passes on the 409 of a change that cannot be aborted", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 409, map[string]any{"error": "nothing to abort"} + }) + defer f.srv.Close() + Expect((&AbortCommand{Remote: remote(f)}).run(context.Background(), &out)).To(MatchError(ContainSubstring("nothing to abort"))) + }) + }) + + Describe("settings", func() { + It("shows the settings", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 200, map[string]any{"nats_url": "nats://n:4222", "max_drain": "15m"} + }) + defer f.srv.Close() + Expect((&GetSettingsCommand{Remote: remote(f)}).run(context.Background(), &out)).To(Succeed()) + Expect(f.last().Path).To(Equal("/api/cluster/settings")) + Expect(out.String()).To(ContainSubstring("nats_url: nats://n:4222")) + Expect(out.String()).To(ContainSubstring("nats_worker_url: (not set)")) + }) + + It("sends only the flags that are given, and says the carrier did not change", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 200, map[string]any{"saved": true, "nats": map[string]any{"reachable": true}, "settings": map[string]any{"nats_url": "nats://n:4222"}} + }) + defer f.srv.Close() + url := "nats://n:4222" + c := &SetSettingsCommand{NATSURL: &url, Remote: remote(f)} + Expect(c.run(context.Background(), &out)).To(Succeed()) + Expect(f.last().Method).To(Equal("PUT")) + Expect(f.last().Body).To(Equal(map[string]any{"nats_url": url})) + Expect(out.String()).To(ContainSubstring("reaches the NATS server")) + Expect(out.String()).To(ContainSubstring("The carrier did not change.")) + }) + + It("sends an empty value to clear a setting", func() { + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 200, map[string]any{"saved": true, "settings": map[string]any{}} + }) + defer f.srv.Close() + empty := "" + Expect((&SetSettingsCommand{NATSWorkerURL: &empty, Remote: remote(f)}).run(context.Background(), &out)).To(Succeed()) + Expect(f.last().Body).To(Equal(map[string]any{"nats_worker_url": ""})) + }) + + It("refuses a call with nothing to store, and shows a refused address", func() { + Expect((&SetSettingsCommand{}).run(context.Background(), io.Discard)).To(MatchError(ContainSubstring("at least one setting"))) + + f := newFakeServer(func(*http.Request, map[string]any) (int, any) { + return 422, map[string]any{"error": "nats_url: this replica cannot reach the NATS server at nats://x:1"} + }) + defer f.srv.Close() + u := "nats://x:1" + Expect((&SetSettingsCommand{NATSURL: &u, Remote: remote(f)}).run(context.Background(), &out)).To(MatchError(ContainSubstring("cannot reach the NATS server"))) + }) + }) +}) diff --git a/core/cli/cluster/settings.go b/core/cli/cluster/settings.go new file mode 100644 index 000000000000..6cfcd52807d6 --- /dev/null +++ b/core/cli/cluster/settings.go @@ -0,0 +1,124 @@ +// SPDX-License-Identifier: MIT + +package cluster + +import ( + "context" + "errors" + "fmt" + "io" + "net/http" + "os" +) + +const settingsPath = "/api/cluster/settings" + +// SettingsCommand groups the commands about the settings of the cluster. +type SettingsCommand struct { + Get GetSettingsCommand `cmd:"" default:"withargs" help:"Show the settings"` + Set SetSettingsCommand `cmd:"" help:"Store settings. Storing a NATS address does not change the carrier."` +} + +// GetSettingsCommand shows the settings. +type GetSettingsCommand struct { + Remote `embed:""` +} + +// SetSettingsCommand stores settings. A flag that is not given leaves the +// setting as it is, and a flag given an empty value clears it. +type SetSettingsCommand struct { + NATSURL *string `name:"nats-url" help:"NATS address that the frontends use. The serving replica checks that it reaches the server. Empty clears it."` + NATSWorkerURL *string `name:"nats-worker-url" help:"NATS address that workers are told to use, when it differs from --nats-url. Empty clears it."` + PrepareTimeout *string `name:"prepare-timeout" help:"How long a change waits for every replica to be ready, as a duration such as 60s."` + TransitionWindow *string `name:"transition-window" help:"How long a replica may take to confirm a commit, as a duration."` + MaxDrain *string `name:"max-drain" help:"How long the previous carrier stays attached after a commit, as a duration such as 15m."` + Remote `embed:""` +} + +type settingsView struct { + NATSURL string `json:"nats_url"` + NATSWorkerURL string `json:"nats_worker_url"` + PrepareTimeout string `json:"prepare_timeout"` + TransitionWindow string `json:"transition_window"` + MaxDrain string `json:"max_drain"` +} + +func (c *GetSettingsCommand) Run() error { + ctx, stop := signalContext() + defer stop() + return c.run(ctx, os.Stdout) +} + +func (c *GetSettingsCommand) run(ctx context.Context, out io.Writer) error { + var view settingsView + _, raw, err := c.do(ctx, c.client(), http.MethodGet, settingsPath, nil, &view) + if err != nil { + return err + } + if c.JSON { + return printJSON(out, raw) + } + printSettings(out, view) + return nil +} + +func (c *SetSettingsCommand) Run() error { + ctx, stop := signalContext() + defer stop() + return c.run(ctx, os.Stdout) +} + +func (c *SetSettingsCommand) run(ctx context.Context, out io.Writer) error { + body := map[string]string{} + for name, v := range map[string]*string{ + "nats_url": c.NATSURL, + "nats_worker_url": c.NATSWorkerURL, + "prepare_timeout": c.PrepareTimeout, + "transition_window": c.TransitionWindow, + "max_drain": c.MaxDrain, + } { + if v != nil { + body[name] = *v + } + } + if len(body) == 0 { + return errors.New("give at least one setting to store, for example --nats-url nats://host:4222") + } + var answer struct { + Saved bool `json:"saved"` + NATS *struct { + Reachable bool `json:"reachable"` + } `json:"nats"` + Settings settingsView `json:"settings"` + } + _, raw, err := c.do(ctx, c.client(), http.MethodPut, settingsPath, body, &answer) + if err != nil { + return err + } + if c.JSON { + return printJSON(out, raw) + } + _, _ = fmt.Fprintln(out, "Saved.") + if answer.NATS != nil && answer.NATS.Reachable { + _, _ = fmt.Fprintln(out, "The serving replica reaches the NATS server.") + } + if c.NATSURL != nil || c.NATSWorkerURL != nil { + _, _ = fmt.Fprintln(out, "The carrier did not change. To use NATS, run `local-ai cluster carrier switch --to nats --dry-run`, then the same command without --dry-run.") + } + printSettings(out, answer.Settings) + return nil +} + +func printSettings(out io.Writer, v settingsView) { + show := func(s string) string { + if s == "" { + return "(not set)" + } + return s + } + _, _ = fmt.Fprintf(out, "nats_url: %s\n", show(v.NATSURL)) + _, _ = fmt.Fprintf(out, "nats_worker_url: %s\n", show(v.NATSWorkerURL)) + _, _ = fmt.Fprintf(out, "prepare_timeout: %s\n", show(v.PrepareTimeout)) + _, _ = fmt.Fprintf(out, "transition_window: %s\n", show(v.TransitionWindow)) + _, _ = fmt.Fprintf(out, "max_drain: %s\n", show(v.MaxDrain)) +} diff --git a/core/cli/run.go b/core/cli/run.go index 522cc04991b5..cf09066e5785 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -163,9 +163,15 @@ type RunCMD struct { DefaultAPIKeyExpiry string `env:"LOCALAI_DEFAULT_API_KEY_EXPIRY" help:"Default expiry for API keys (e.g. 90d, 1y; empty = no expiry)" group:"auth"` // Distributed / Horizontal Scaling - Distributed bool `env:"LOCALAI_DISTRIBUTED" default:"false" help:"Enable distributed mode (requires PostgreSQL + NATS)" group:"distributed"` + Distributed bool `env:"LOCALAI_DISTRIBUTED" default:"false" help:"Enable distributed mode (requires PostgreSQL; NATS is optional)" group:"distributed"` InstanceID string `env:"LOCALAI_INSTANCE_ID" help:"Unique instance ID for distributed mode (auto-generated UUID if empty)" group:"distributed"` + PeerAddress string `env:"LOCALAI_PEER_ADDRESS" help:"host:port at which the other frontends dial this one to reach a worker tunnel it holds (found from the route to the database if empty)" group:"distributed"` + PeerTLS bool `env:"LOCALAI_PEER_TLS" default:"false" help:"dial the other frontends over wss; they must sit behind TLS at the address they publish. Without it the peer link, which carries the credential of this replica and relayed requests, is clear text" group:"distributed"` + PeerTLSCA string `name:"peer-tls-ca" env:"LOCALAI_PEER_TLS_CA" type:"existingfile" help:"PEM file with the CA that signs the certificate of the other frontends (system roots if empty); use with --peer-tls" group:"distributed"` NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL (e.g., nats://localhost:4222)" group:"distributed"` + CarrierPrepareTimeout string `env:"LOCALAI_CARRIER_PREPARE_TIMEOUT" help:"How long a change of carrier waits for every frontend to be ready before it is aborted (default 1m). The cluster setting switch.prepare_timeout, when stored, wins." group:"distributed"` + CarrierTransitionWindow string `env:"LOCALAI_CARRIER_TRANSITION_WINDOW" help:"How long a change of carrier waits for every frontend to confirm that it uses the new carrier (default 2m). The cluster setting switch.transition_window, when stored, wins." group:"distributed"` + CarrierMaxDrain string `env:"LOCALAI_CARRIER_MAX_DRAIN" help:"How long the previous carrier stays attached after a change of carrier, so that work that started on it can finish (default 15m). Work still running after that is failed by the reaper. The cluster setting switch.max_drain, when stored, wins." group:"distributed"` StorageURL string `env:"LOCALAI_STORAGE_URL" help:"S3-compatible storage endpoint URL (e.g., http://minio:9000)" group:"distributed"` StorageBucket string `env:"LOCALAI_STORAGE_BUCKET" default:"localai" help:"S3 bucket name for object storage" group:"distributed"` StorageRegion string `env:"LOCALAI_STORAGE_REGION" default:"us-east-1" help:"S3 region" group:"distributed"` @@ -364,6 +370,12 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { if r.InstanceID != "" { opts = append(opts, config.WithDistributedInstanceID(r.InstanceID)) } + if r.PeerAddress != "" { + opts = append(opts, config.WithDistributedPeerAddress(r.PeerAddress)) + } + if r.PeerTLS || r.PeerTLSCA != "" { + opts = append(opts, config.WithDistributedPeerTLS(r.PeerTLSCA)) + } if r.NatsURL != "" { opts = append(opts, config.WithNatsURL(r.NatsURL)) } @@ -424,6 +436,26 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { } opts = append(opts, config.WithNodeHeartbeatCheckpoint(d)) } + if r.CarrierPrepareTimeout != "" || r.CarrierTransitionWindow != "" || r.CarrierMaxDrain != "" { + var prepare, window, drain time.Duration + var err error + for _, f := range []struct { + env, value string + into *time.Duration + }{ + {"LOCALAI_CARRIER_PREPARE_TIMEOUT", r.CarrierPrepareTimeout, &prepare}, + {"LOCALAI_CARRIER_TRANSITION_WINDOW", r.CarrierTransitionWindow, &window}, + {"LOCALAI_CARRIER_MAX_DRAIN", r.CarrierMaxDrain, &drain}, + } { + if f.value == "" { + continue + } + if *f.into, err = parseDistributedDuration(f.env, f.value); err != nil { + return err + } + } + opts = append(opts, config.WithCarrierTimings(prepare, window, drain)) + } if r.ModelConfigResyncInterval != "" { d, err := parseDistributedDuration("LOCALAI_MODEL_CONFIG_RESYNC_INTERVAL", r.ModelConfigResyncInterval) if err != nil { diff --git a/core/cli/workerregistry/client.go b/core/cli/workerregistry/client.go index cf46455c95c0..a6712d5337e4 100644 --- a/core/cli/workerregistry/client.go +++ b/core/cli/workerregistry/client.go @@ -9,6 +9,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "strings" "sync" @@ -63,6 +64,39 @@ type RegisterResponse struct { APIToken string `json:"api_token,omitempty"` NatsJWT string `json:"nats_jwt,omitempty"` NatsUserSeed string `json:"nats_user_seed,omitempty"` + // Carrier is the carrier that is active in the cluster: "nats" or "tunnel". + // A frontend that predates carriers does not send it, and then the cluster + // runs on NATS. CarrierEpoch is the epoch of the carrier row at that time. + Carrier string `json:"carrier,omitempty"` + CarrierEpoch int64 `json:"carrier_epoch,omitempty"` + // TunnelToken is the own credential of the node for the tunnel. It is sent + // once and only when the tunnel is the active carrier. + TunnelToken string `json:"tunnel_token,omitempty"` + // CarrierState, CarrierTarget and CarrierDraining say whether a change of + // carrier is under way. CredentialFor is the carrier that the credential in + // the answer is for: the one the worker asked for, when the frontend allowed + // it, else the active one. + CarrierState string `json:"carrier_state,omitempty"` + CarrierTarget string `json:"carrier_target,omitempty"` + CarrierDraining string `json:"carrier_draining,omitempty"` + CredentialFor string `json:"credential_for,omitempty"` + // NatsURL is the address workers use for NATS, NatsCA the CA of the server as + // PEM, and NatsClientTLS says the server asks for a client certificate, which + // the frontend cannot hand over. + NatsURL string `json:"nats_url,omitempty"` + NatsCA string `json:"nats_ca,omitempty"` + NatsClientTLS bool `json:"nats_client_tls,omitempty"` +} + +// HeartbeatReply is the answer to a heartbeat. The carrier fields are empty when +// the frontend predates carriers. +type HeartbeatReply struct { + Carrier string `json:"carrier,omitempty"` + CarrierEpoch int64 `json:"carrier_epoch,omitempty"` + CarrierState string `json:"carrier_state,omitempty"` + CarrierTarget string `json:"carrier_target,omitempty"` + CarrierDraining string `json:"carrier_draining,omitempty"` + NatsClientTLS bool `json:"nats_client_tls,omitempty"` } // RegisterFull sends a single registration request and returns the full @@ -97,58 +131,84 @@ func (c *RegistrationClient) RegisterFull(ctx context.Context, body map[string]a return &result, nil } -// Register sends a single registration request and returns the node ID and -// optional credentials (API token for agent workers, NATS JWT when configured). -func (c *RegistrationClient) Register(ctx context.Context, body map[string]any) (nodeID, apiToken, natsJWT, natsSeed string, err error) { - res, err := c.RegisterFull(ctx, body) +// RegisterWithRetry retries registration with exponential backoff. +func (c *RegistrationClient) RegisterWithRetry(ctx context.Context, body map[string]any, maxRetries int) (nodeID, apiToken, natsJWT, natsSeed string, err error) { + res, err := c.RegisterFullWithRetry(ctx, body, maxRetries) if err != nil { return "", "", "", "", err } return res.ID, res.APIToken, res.NatsJWT, res.NatsUserSeed, nil } -// RegisterWithRetry retries registration with exponential backoff. -func (c *RegistrationClient) RegisterWithRetry(ctx context.Context, body map[string]any, maxRetries int) (nodeID, apiToken, natsJWT, natsSeed string, err error) { +// RegisterFullWithRetry is RegisterWithRetry that returns the whole response. +func (c *RegistrationClient) RegisterFullWithRetry(ctx context.Context, body map[string]any, maxRetries int) (*RegisterResponse, error) { + if maxRetries < 1 { + return nil, fmt.Errorf("registering: %d attempts allowed, want at least one", maxRetries) + } backoff := 2 * time.Second maxBackoff := 30 * time.Second + var err error for attempt := 1; attempt <= maxRetries; attempt++ { - nodeID, apiToken, natsJWT, natsSeed, err = c.Register(ctx, body) + var res *RegisterResponse + res, err = c.RegisterFull(ctx, body) if err == nil { - return nodeID, apiToken, natsJWT, natsSeed, nil + return res, nil } if attempt == maxRetries { - return "", "", "", "", fmt.Errorf("failed after %d attempts: %w", maxRetries, err) + return nil, fmt.Errorf("failed after %d attempts: %w", maxRetries, err) } xlog.Warn("Registration failed, retrying", "attempt", attempt, "next_retry", backoff, "error", err) select { case <-ctx.Done(): - return "", "", "", "", ctx.Err() + return nil, ctx.Err() case <-time.After(backoff): } backoff = min(backoff*2, maxBackoff) } - return nodeID, apiToken, natsJWT, natsSeed, err + return nil, err } -// Heartbeat sends a single heartbeat POST with the given body. +// Heartbeat sends a single heartbeat POST with the given body. It fails only when +// the request cannot be made. func (c *RegistrationClient) Heartbeat(ctx context.Context, nodeID string, body map[string]any) error { + resp, err := c.post(ctx, nodeID, body) + if err != nil { + return err + } + return resp.Body.Close() +} + +// HeartbeatFull sends a heartbeat and returns the answer, which tells the worker +// which carrier the cluster uses. It fails when the frontend refuses the +// heartbeat, for example for a node it does not know. +func (c *RegistrationClient) HeartbeatFull(ctx context.Context, nodeID string, body map[string]any) (*HeartbeatReply, error) { + resp, err := c.post(ctx, nodeID, body) + if err != nil { + return nil, err + } + defer func() { _ = resp.Body.Close() }() + if resp.StatusCode < 200 || resp.StatusCode >= 300 { + return nil, fmt.Errorf("heartbeat refused with status %d", resp.StatusCode) + } + var reply HeartbeatReply + // An answer that is not JSON is the answer of a frontend that says nothing + // about carriers. + _ = json.NewDecoder(io.LimitReader(resp.Body, 1<<20)).Decode(&reply) + return &reply, nil +} + +func (c *RegistrationClient) post(ctx context.Context, nodeID string, body map[string]any) (*http.Response, error) { jsonBody, _ := json.Marshal(body) url := c.baseURL() + "/api/node/" + nodeID + "/heartbeat" req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(jsonBody)) if err != nil { - return fmt.Errorf("creating heartbeat request: %w", err) + return nil, fmt.Errorf("creating heartbeat request: %w", err) } req.Header.Set("Content-Type", "application/json") c.setAuth(req) - - resp, err := c.httpClient().Do(req) - if err != nil { - return err - } - defer resp.Body.Close() - return nil + return c.httpClient().Do(req) } // HeartbeatLoop runs heartbeats at the given interval until ctx is cancelled. diff --git a/core/cli/workerregistry/credentials.go b/core/cli/workerregistry/credentials.go index 24dd6f3c8ed7..d39a6cd0e449 100644 --- a/core/cli/workerregistry/credentials.go +++ b/core/cli/workerregistry/credentials.go @@ -15,6 +15,11 @@ import ( // package (and its gorm/DB dependencies). const statusPending = "pending" +// carrierTunnel mirrors cluster.CarrierTunnel, for the reason statusPending is +// copied: this package does not import the cluster package and its database +// dependencies. +const carrierTunnel = "tunnel" + // defaultMaxAttempts bounds how many times Acquire registers (and how many // consecutive times RefreshLoop may fail) before giving up. It is high enough // to ride out a slow admin approval or a transient frontend outage, but finite @@ -25,20 +30,27 @@ const defaultMaxAttempts = 100 // RegisterFunc performs one idempotent registration round-trip. type RegisterFunc func(ctx context.Context) (*RegisterResponse, error) -// NATSCredentialManager acquires NATS credentials at startup — waiting through -// admin approval when required — and refreshes them before the minted JWT -// expires, by re-registering (which mints a fresh JWT). The live NATS -// connection adopts a refreshed JWT on its next reconnect via Provider. Safe -// for concurrent use. +// CredentialManager acquires the credentials of the carrier at startup. For +// NATS it waits through admin approval when required, and it refreshes the +// credentials before the minted JWT expires, by registering again (which mints a +// new JWT). The live NATS connection adopts a refreshed JWT on its next +// reconnect via Provider. For the tunnel it keeps the token of the node that the +// last registration returned. Safe for concurrent use. // // It addresses two failure modes: a worker that needs credentials but registers // while still pending approval (it would otherwise give up and never connect), // and a long-running worker whose 24h JWT expires with no way to renew it. -type NATSCredentialManager struct { +type CredentialManager struct { register RegisterFunc requireCreds bool // block until credentials are present (frontend minting in use) + // carrier is the carrier whose credential this manager waits for. Empty + // means the one the frontend names as active, which is what a worker wants + // when it starts. A worker that follows a change asks for the carrier it + // attaches to, and that is not the active one yet while the change is + // prepared. + carrier string - // Tunables; defaults set by NewNATSCredentialManager, overridable in tests. + // Tunables; defaults set by NewCredentialManager, overridable in tests. initialBackoff time.Duration maxBackoff time.Duration maxAttempts int // bound on Acquire attempts / consecutive refresh failures (<=0 = unlimited) @@ -46,16 +58,40 @@ type NATSCredentialManager struct { refreshRetry time.Duration expiryOf func(jwt string) (time.Time, bool) - mu sync.RWMutex - jwt string - seed string - nodeID string + mu sync.RWMutex + jwt string + seed string + nodeID string + tunnelToken string +} + +// NewCredentialManagerFor is NewCredentialManager for a manager that waits for the +// credential of one carrier, whichever the cluster has active. register must ask +// the frontend for that carrier. +func NewCredentialManagerFor(register RegisterFunc, requireCreds bool, carrier string) *CredentialManager { + m := NewCredentialManager(register, requireCreds) + m.carrier = carrier + return m +} + +// credentialFor is the carrier that an answer carries a credential for. +func (m *CredentialManager) credentialFor(res *RegisterResponse) string { + if m.carrier != "" { + return m.carrier + } + return res.Carrier } -// NewNATSCredentialManager builds a manager over register. When requireCreds is +// wrongCarrier is true when the manager asked for a carrier and the frontend +// answered with the credential of another. +func (m *CredentialManager) wrongCarrier(res *RegisterResponse) bool { + return m.carrier != "" && res.CredentialFor != "" && res.CredentialFor != m.carrier +} + +// NewCredentialManager builds a manager over register. When requireCreds is // true, Acquire blocks until the node is approved and credentials are minted. -func NewNATSCredentialManager(register RegisterFunc, requireCreds bool) *NATSCredentialManager { - return &NATSCredentialManager{ +func NewCredentialManager(register RegisterFunc, requireCreds bool) *CredentialManager { + return &CredentialManager{ register: register, requireCreds: requireCreds, initialBackoff: 2 * time.Second, @@ -80,24 +116,36 @@ func jwtExpiry(token string) (time.Time, bool) { return time.Unix(uc.Expires, 0), true } -func (m *NATSCredentialManager) store(res *RegisterResponse) { +func (m *CredentialManager) store(res *RegisterResponse) { m.mu.Lock() defer m.mu.Unlock() m.nodeID = res.ID if res.NatsJWT != "" && res.NatsUserSeed != "" { m.jwt, m.seed = res.NatsJWT, res.NatsUserSeed } + if res.TunnelToken != "" { + m.tunnelToken = res.TunnelToken + } +} + +// TunnelToken returns the token of the node for the tunnel, or an empty string +// before the first registration that returned one. The tunnel client calls it +// for every dial, because a registration mints a new token. +func (m *CredentialManager) TunnelToken() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.tunnelToken } // Current returns the latest NATS credentials (both empty until acquired). -func (m *NATSCredentialManager) Current() (jwt, seed string) { +func (m *CredentialManager) Current() (jwt, seed string) { m.mu.RLock() defer m.mu.RUnlock() return m.jwt, m.seed } // NodeID returns the node ID from the most recent registration. -func (m *NATSCredentialManager) NodeID() string { +func (m *CredentialManager) NodeID() string { m.mu.RLock() defer m.mu.RUnlock() return m.nodeID @@ -105,12 +153,12 @@ func (m *NATSCredentialManager) NodeID() string { // Provider returns a callback compatible with messaging.WithUserJWTProvider, // supplying the current credentials on each (re)connect. -func (m *NATSCredentialManager) Provider() func() (string, string) { +func (m *CredentialManager) Provider() func() (string, string) { return m.Current } // HasCredentials reports whether complete NATS credentials have been obtained. -func (m *NATSCredentialManager) HasCredentials() bool { +func (m *CredentialManager) HasCredentials() bool { jwt, seed := m.Current() return jwt != "" && seed != "" } @@ -119,7 +167,13 @@ func (m *NATSCredentialManager) HasCredentials() bool { // exponential backoff until the node is approved (status != pending) and // credentials are minted. Without requireCreds it returns the first successful // response (the historical one-shot behavior, preserved for anonymous NATS). -func (m *NATSCredentialManager) Acquire(ctx context.Context) (*RegisterResponse, error) { +// +// When the frontend names the tunnel as the carrier, the token for the tunnel is +// what Acquire waits for, and it does not matter whether the node is approved. +// The frontend mints the token at registration and the connect route refuses a +// node that waits for approval, so the tunnel client waits for the approval by +// itself. +func (m *CredentialManager) Acquire(ctx context.Context) (*RegisterResponse, error) { backoff := m.initialBackoff var lastReason error for attempt := 1; m.maxAttempts <= 0 || attempt <= m.maxAttempts; attempt++ { @@ -128,6 +182,15 @@ func (m *NATSCredentialManager) Acquire(ctx context.Context) (*RegisterResponse, case err != nil: lastReason = err xlog.Warn("Registration failed, retrying", "attempt", attempt, "next_retry", backoff, "error", err) + case m.wrongCarrier(res): + lastReason = fmt.Errorf("asked for the credential of the %s carrier and got the one of the %s carrier", m.carrier, res.CredentialFor) + xlog.Info("The frontend does not hand out the credential of that carrier yet; waiting", "asked", m.carrier, "got", res.CredentialFor, "attempt", attempt, "next_retry", backoff) + case m.credentialFor(res) == carrierTunnel && res.TunnelToken == "": + lastReason = fmt.Errorf("node %s registered but the tunnel credential was not minted", res.ID) + xlog.Info("Node registered but the tunnel credential is not minted yet; waiting", "node", res.ID, "attempt", attempt, "next_retry", backoff) + case m.credentialFor(res) == carrierTunnel: + m.store(res) + return res, nil case !m.requireCreds: m.store(res) return res, nil @@ -151,6 +214,27 @@ func (m *NATSCredentialManager) Acquire(ctx context.Context) (*RegisterResponse, return nil, fmt.Errorf("giving up acquiring NATS credentials after %d attempts: %w", m.maxAttempts, lastReason) } +// Reregister registers once more and keeps the credentials of the answer. It +// does not wait and does not retry: the caller owns the backoff. The tunnel +// client calls it after the frontend refused its credential, which happens when +// another process registered under the same node name, or when the frontend +// lost its record of the node. A registration mints a new credential, and the +// registration token of the deployment authorises it, as it does at startup. +func (m *CredentialManager) Reregister(ctx context.Context) error { + res, err := m.register(ctx) + if err != nil { + return err + } + if m.wrongCarrier(res) { + return fmt.Errorf("asked for the credential of the %s carrier and got the one of the %s carrier", m.carrier, res.CredentialFor) + } + if m.credentialFor(res) == carrierTunnel && res.TunnelToken == "" { + return fmt.Errorf("node %s registered but the tunnel credential was not minted", res.ID) + } + m.store(res) + return nil +} + // RefreshLoop re-registers to mint a fresh JWT before the current one expires, // updating the credentials returned by Current/Provider so the NATS connection // adopts them on its next reconnect. It returns nil when ctx is cancelled or @@ -158,7 +242,7 @@ func (m *NATSCredentialManager) Acquire(ctx context.Context) (*RegisterResponse, // error after maxAttempts consecutive refresh failures — letting the caller // exit the worker so it restarts and re-acquires (or surfaces the outage) // rather than silently drifting toward an expired, unrenewable JWT. -func (m *NATSCredentialManager) RefreshLoop(ctx context.Context) error { +func (m *CredentialManager) RefreshLoop(ctx context.Context) error { failures := 0 for { jwt, _ := m.Current() diff --git a/core/cli/workerregistry/credentials_test.go b/core/cli/workerregistry/credentials_test.go index 02ed413c4e8f..76de4fd749eb 100644 --- a/core/cli/workerregistry/credentials_test.go +++ b/core/cli/workerregistry/credentials_test.go @@ -2,6 +2,7 @@ package workerregistry import ( "context" + "errors" "sync" "testing" "time" @@ -50,7 +51,7 @@ func (f *fakeRegister) count() int { return f.calls } -var _ = Describe("NATSCredentialManager", func() { +var _ = Describe("CredentialManager", func() { approved := func(jwt, seed string) *RegisterResponse { return &RegisterResponse{ID: "node-1", Status: "healthy", NatsJWT: jwt, NatsUserSeed: seed} } @@ -63,7 +64,7 @@ var _ = Describe("NATSCredentialManager", func() { {res: approved("", "")}, // approved but JWT not minted yet {res: approved("jwt-1", "seed-1")}, // finally minted }} - m := NewNATSCredentialManager(f.fn(), true /* requireCreds */) + m := NewCredentialManager(f.fn(), true /* requireCreds */) m.initialBackoff = time.Millisecond m.maxBackoff = time.Millisecond @@ -81,7 +82,7 @@ var _ = Describe("NATSCredentialManager", func() { It("returns immediately on the first success when credentials are not required (anonymous NATS)", func() { f := &fakeRegister{steps: []step{{res: pending}}} - m := NewNATSCredentialManager(f.fn(), false /* requireCreds */) + m := NewCredentialManager(f.fn(), false /* requireCreds */) res, err := m.Acquire(context.Background()) Expect(err).ToNot(HaveOccurred()) @@ -92,7 +93,7 @@ var _ = Describe("NATSCredentialManager", func() { It("aborts when the context is cancelled while waiting for approval", func() { f := &fakeRegister{steps: []step{{res: pending}}} - m := NewNATSCredentialManager(f.fn(), true) + m := NewCredentialManager(f.fn(), true) m.initialBackoff = 10 * time.Millisecond ctx, cancel := context.WithCancel(context.Background()) @@ -103,7 +104,7 @@ var _ = Describe("NATSCredentialManager", func() { It("gives up after a bounded number of attempts so the worker exits and alerts", func() { f := &fakeRegister{steps: []step{{res: pending}}} // never approved - m := NewNATSCredentialManager(f.fn(), true) + m := NewCredentialManager(f.fn(), true) m.initialBackoff = time.Millisecond m.maxBackoff = time.Millisecond m.maxAttempts = 5 @@ -116,10 +117,109 @@ var _ = Describe("NATSCredentialManager", func() { }) }) + Describe("Reregister", func() { + tunnelRes := func(token string) *RegisterResponse { + return &RegisterResponse{ID: "node-1", Status: "healthy", Carrier: "tunnel", TunnelToken: token} + } + + It("replaces the tunnel token with the one that the new registration minted", func() { + f := &fakeRegister{steps: []step{{res: tunnelRes("token-1")}, {res: tunnelRes("token-2")}}} + m := NewCredentialManager(f.fn(), false) + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + + Expect(m.Reregister(context.Background())).To(Succeed()) + Expect(m.TunnelToken()).To(Equal("token-2")) + Expect(f.count()).To(Equal(2), "one registration and no retry") + }) + + It("keeps the old token and returns the error when registration fails", func() { + f := &fakeRegister{steps: []step{{res: tunnelRes("token-1")}, {err: errors.New("frontend down")}}} + m := NewCredentialManager(f.fn(), false) + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + + Expect(m.Reregister(context.Background())).To(MatchError(ContainSubstring("frontend down"))) + Expect(m.TunnelToken()).To(Equal("token-1")) + Expect(f.count()).To(Equal(2)) + }) + + It("keeps the old token when the answer names the tunnel and mints none", func() { + f := &fakeRegister{steps: []step{{res: tunnelRes("token-1")}, {res: tunnelRes("")}}} + m := NewCredentialManager(f.fn(), false) + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + + Expect(m.Reregister(context.Background())).ToNot(Succeed()) + Expect(m.TunnelToken()).To(Equal("token-1")) + }) + }) + + Describe("Acquire on the tunnel carrier", func() { + tunnelRes := func(status, token string) *RegisterResponse { + return &RegisterResponse{ID: "node-1", Status: status, Carrier: "tunnel", CarrierEpoch: 4, TunnelToken: token} + } + + It("returns at once with the token, also for a node that waits for approval", func() { + f := &fakeRegister{steps: []step{{res: tunnelRes("pending", "token-1")}}} + m := NewCredentialManager(f.fn(), true) + + res, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(res.Carrier).To(Equal("tunnel")) + Expect(f.count()).To(Equal(1)) + Expect(m.TunnelToken()).To(Equal("token-1")) + Expect(m.HasCredentials()).To(BeFalse(), "a tunnel worker holds no NATS credentials") + }) + + It("waits when the frontend named the tunnel but did not mint a token", func() { + f := &fakeRegister{steps: []step{ + {res: tunnelRes("healthy", "")}, + {res: tunnelRes("healthy", "token-2")}, + }} + m := NewCredentialManager(f.fn(), false) + m.initialBackoff = time.Millisecond + m.maxBackoff = time.Millisecond + + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(f.count()).To(Equal(2)) + Expect(m.TunnelToken()).To(Equal("token-2")) + }) + + It("gives the token of the latest registration, because every registration mints a new one", func() { + f := &fakeRegister{steps: []step{{res: tunnelRes("healthy", "token-1")}, {res: tunnelRes("healthy", "token-2")}}} + m := NewCredentialManager(f.fn(), true) + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(m.TunnelToken()).To(Equal("token-1")) + + m.store(&RegisterResponse{ID: "node-1", Carrier: "tunnel", TunnelToken: "token-2"}) + Expect(m.TunnelToken()).To(Equal("token-2")) + }) + + It("does not touch the NATS path: a frontend that names NATS, or nothing, still waits for the JWT", func() { + for _, carrier := range []string{"", "nats"} { + f := &fakeRegister{steps: []step{ + {res: &RegisterResponse{ID: "node-1", Status: "healthy", Carrier: carrier}}, + {res: &RegisterResponse{ID: "node-1", Status: "healthy", Carrier: carrier, NatsJWT: "jwt", NatsUserSeed: "seed"}}, + }} + m := NewCredentialManager(f.fn(), true) + m.initialBackoff = time.Millisecond + m.maxBackoff = time.Millisecond + + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(f.count()).To(Equal(2), "carrier %q", carrier) + Expect(m.TunnelToken()).To(BeEmpty()) + } + }) + }) + Describe("RefreshLoop (#5 — renew before the JWT expires)", func() { It("re-registers before expiry and updates the credentials served to new connections", func() { f := &fakeRegister{steps: []step{{res: approved("jwt-2", "seed-2")}}} - m := NewNATSCredentialManager(f.fn(), true) + m := NewCredentialManager(f.fn(), true) m.refreshLead = 0.5 m.refreshRetry = time.Millisecond // jwt-1 expires soon; jwt-2 is long-lived so the loop then idles. @@ -147,7 +247,7 @@ var _ = Describe("NATSCredentialManager", func() { It("returns an error after the bounded number of consecutive failures so the caller can exit", func() { f := &fakeRegister{steps: []step{{err: context.DeadlineExceeded}}} // refresh always fails - m := NewNATSCredentialManager(f.fn(), true) + m := NewCredentialManager(f.fn(), true) m.refreshLead = 0.5 m.refreshRetry = time.Millisecond m.maxAttempts = 3 @@ -161,7 +261,7 @@ var _ = Describe("NATSCredentialManager", func() { It("exits promptly when the current credential has no expiry (nothing to refresh)", func() { f := &fakeRegister{steps: []step{{res: approved("x", "y")}}} - m := NewNATSCredentialManager(f.fn(), true) + m := NewCredentialManager(f.fn(), true) m.expiryOf = func(string) (time.Time, bool) { return time.Time{}, false } m.store(approved("static", "seed")) diff --git a/core/cli/workerregistry/follow_test.go b/core/cli/workerregistry/follow_test.go new file mode 100644 index 000000000000..96674e1a0653 --- /dev/null +++ b/core/cli/workerregistry/follow_test.go @@ -0,0 +1,110 @@ +package workerregistry + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("A credential for the carrier a worker follows to", func() { + It("waits for the tunnel token when it asks for the tunnel, although the active carrier is NATS", func() { + f := &fakeRegister{steps: []step{ + {res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "nats", CredentialFor: "tunnel"}}, + {res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "nats", CredentialFor: "tunnel", TunnelToken: "tok"}}, + }} + m := NewCredentialManagerFor(f.fn(), false, "tunnel") + m.initialBackoff = 1 + res, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(res.TunnelToken).To(Equal("tok")) + Expect(m.TunnelToken()).To(Equal("tok")) + Expect(f.count()).To(Equal(2)) + }) + + It("takes the NATS credential while the tunnel is the active carrier", func() { + f := &fakeRegister{steps: []step{{res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "tunnel", CredentialFor: "nats", NatsJWT: "j", NatsUserSeed: "s", NatsURL: "nats://x:4222"}}}} + m := NewCredentialManagerFor(f.fn(), true, "nats") + res, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(res.NatsURL).To(Equal("nats://x:4222")) + jwt, seed := m.Current() + Expect(jwt).To(Equal("j")) + Expect(seed).To(Equal("s")) + }) + + It("does not take the credential of another carrier for the one it asked for", func() { + f := &fakeRegister{steps: []step{{res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "nats", CredentialFor: "nats", NatsJWT: "j", NatsUserSeed: "s"}}}} + m := NewCredentialManagerFor(f.fn(), false, "tunnel") + m.maxAttempts = 2 + m.initialBackoff = 1 + _, err := m.Acquire(context.Background()) + Expect(err).To(MatchError(ContainSubstring("tunnel"))) + Expect(m.TunnelToken()).To(BeEmpty()) + }) + + It("asks again for the same carrier when it registers again", func() { + f := &fakeRegister{steps: []step{ + {res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "nats", CredentialFor: "tunnel", TunnelToken: "a"}}, + {res: &RegisterResponse{ID: "n", Status: "healthy", Carrier: "nats", CredentialFor: "tunnel", TunnelToken: "b"}}, + }} + m := NewCredentialManagerFor(f.fn(), false, "tunnel") + _, err := m.Acquire(context.Background()) + Expect(err).ToNot(HaveOccurred()) + Expect(m.Reregister(context.Background())).To(Succeed()) + Expect(m.TunnelToken()).To(Equal("b")) + }) +}) + +var _ = Describe("The heartbeat client", func() { + It("returns what the frontend says about the carrier", func() { + var calls atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls.Add(1) + Expect(r.URL.Path).To(Equal("/api/node/n1/heartbeat")) + var body map[string]any + Expect(json.NewDecoder(r.Body).Decode(&body)).To(Succeed()) + Expect(body).To(HaveKey("follow_capabilities")) + _ = json.NewEncoder(w).Encode(map[string]any{"message": "heartbeat received", "carrier": "nats", "carrier_epoch": 4, "carrier_state": "prepare", "carrier_target": "tunnel", "nats_client_tls": true}) + })) + DeferCleanup(srv.Close) + c := &RegistrationClient{FrontendURL: srv.URL} + reply, err := c.HeartbeatFull(context.Background(), "n1", map[string]any{"follow_capabilities": []string{"nats"}}) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Carrier).To(Equal("nats")) + Expect(reply.CarrierEpoch).To(BeEquivalentTo(4)) + Expect(reply.CarrierState).To(Equal("prepare")) + Expect(reply.CarrierTarget).To(Equal("tunnel")) + Expect(reply.NatsClientTLS).To(BeTrue()) + }) + + It("returns an empty answer from a frontend that predates carriers", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _ = json.NewEncoder(w).Encode(map[string]string{"message": "heartbeat received"}) + })) + DeferCleanup(srv.Close) + c := &RegistrationClient{FrontendURL: srv.URL} + reply, err := c.HeartbeatFull(context.Background(), "n1", nil) + Expect(err).ToNot(HaveOccurred()) + Expect(reply.Carrier).To(BeEmpty()) + }) + + It("reports a refusal and a node the frontend does not know", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusNotFound) })) + DeferCleanup(srv.Close) + c := &RegistrationClient{FrontendURL: srv.URL} + _, err := c.HeartbeatFull(context.Background(), "n1", nil) + Expect(err).To(HaveOccurred()) + }) + + It("keeps the heartbeat that returns only an error working as before", func() { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })) + DeferCleanup(srv.Close) + c := &RegistrationClient{FrontendURL: srv.URL} + Expect(c.Heartbeat(context.Background(), "n1", nil)).To(Succeed()) + }) +}) diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index 2ed848d73efa..82e8e8202a3a 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -11,13 +11,35 @@ import ( ) // DistributedConfig holds configuration for horizontal scaling mode. -// When Enabled is true, PostgreSQL and NATS are required. +// When Enabled is true, PostgreSQL is required. NATS is needed only while the +// cluster runs on the NATS carrier. type DistributedConfig struct { - Enabled bool // --distributed / LOCALAI_DISTRIBUTED - InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty) - NatsURL string // --nats-url / LOCALAI_NATS_URL - StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint) - RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration) + Enabled bool // --distributed / LOCALAI_DISTRIBUTED + InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty) + // PeerAddress is the host and port at which the other frontends dial this + // one, to reach a worker tunnel that this one holds. Empty means that the + // address is found from the route to the database. LOCALAI_PEER_ADDRESS. + PeerAddress string + // PeerTLS makes this frontend dial the other frontends over wss. They must sit + // behind TLS at the address that they publish. LOCALAI_PEER_TLS. + PeerTLS bool + // PeerTLSCA is a PEM file with the certificate authority that signs the + // certificate of the other frontends. Empty uses the system roots. + // LOCALAI_PEER_TLS_CA. + PeerTLSCA string + // NatsURL is the NATS server of a deployment that uses NATS. A deployment with + // no URL and no row in the database runs on the tunnel carrier. The URL is + // copied into the cluster settings when none is stored there, and the cluster + // setting wins afterwards, so that every replica uses the same one. + NatsURL string // --nats-url / LOCALAI_NATS_URL + // CarrierPrepareTimeout, CarrierTransitionWindow and CarrierMaxDrain are the + // waits of a change of carrier on the replica that leads it, used when the + // cluster settings do not set them. Zero means the default. + CarrierPrepareTimeout time.Duration // --carrier-prepare-timeout + CarrierTransitionWindow time.Duration // --carrier-transition-window + CarrierMaxDrain time.Duration // --carrier-max-drain + StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint) + RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration) // RegistrationRequireAuth fails startup when distributed mode is enabled but // RegistrationToken is empty. The default (false) keeps the historical // fail-open behavior with a loud warning; production should set it so the @@ -142,9 +164,9 @@ func (c DistributedConfig) Validate() error { if !c.Enabled { return nil } - if c.NatsURL == "" { - return fmt.Errorf("distributed mode requires --nats-url / LOCALAI_NATS_URL") - } + // No NATS URL is required. The carrier is a setting of the cluster in the + // database: a new deployment with only PostgreSQL runs on the tunnel, and a + // deployment with a NATS URL stays on NATS. // S3 credentials must be paired if (c.StorageAccessKey != "" && c.StorageSecretKey == "") || (c.StorageAccessKey == "" && c.StorageSecretKey != "") { @@ -166,7 +188,11 @@ func (c DistributedConfig) Validate() error { if err := c.NatsTLSFiles().Validate(); err != nil { return err } - c.NatsAuthConfig().WarnIfInsecure(true) + // A deployment with no NATS URL runs on the tunnel and has no NATS bus to warn + // about. The warning comes back when a URL is configured. + if c.NatsURL != "" { + c.NatsAuthConfig().WarnIfInsecure(true) + } // Check for negative durations for name, d := range map[string]time.Duration{ FlagMCPToolTimeout: c.MCPToolTimeout, @@ -188,6 +214,15 @@ func (c DistributedConfig) Validate() error { if c.ModelConfigResyncInterval < 0 { return fmt.Errorf("%s must not be negative", FlagModelConfigResyncInterval) } + for name, d := range map[string]time.Duration{ + FlagCarrierPrepareTimeout: c.CarrierPrepareTimeout, + FlagCarrierTransitionWindow: c.CarrierTransitionWindow, + FlagCarrierMaxDrain: c.CarrierMaxDrain, + } { + if d < 0 { + return fmt.Errorf("%s must not be negative", name) + } + } return nil } @@ -203,6 +238,23 @@ func WithDistributedInstanceID(id string) AppOption { } } +// WithDistributedPeerAddress sets the address at which other frontends dial this +// one. +func WithDistributedPeerAddress(addr string) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.PeerAddress = addr + } +} + +// WithDistributedPeerTLS makes this frontend dial its peers over wss, trusting +// the CA in caFile, or the system roots when caFile is empty. +func WithDistributedPeerTLS(caFile string) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.PeerTLS = true + o.Distributed.PeerTLSCA = caFile + } +} + func WithNatsURL(url string) AppOption { return func(o *ApplicationConfig) { o.Distributed.NatsURL = url @@ -446,6 +498,24 @@ const ( FlagDiskHeadroomCheck = "distributed-disk-headroom-check" ) +// Names of the timings of a change of carrier. +const ( + FlagCarrierPrepareTimeout = "carrier-prepare-timeout" + FlagCarrierTransitionWindow = "carrier-transition-window" + FlagCarrierMaxDrain = "carrier-max-drain" +) + +// WithCarrierTimings sets the waits of a change of carrier: how long prepare +// waits for every replica, how long commit waits for every replica to confirm, +// and how long the previous carrier stays attached. A zero keeps the default. +func WithCarrierTimings(prepareTimeout, transitionWindow, maxDrain time.Duration) AppOption { + return func(o *ApplicationConfig) { + o.Distributed.CarrierPrepareTimeout = prepareTimeout + o.Distributed.CarrierTransitionWindow = transitionWindow + o.Distributed.CarrierMaxDrain = maxDrain + } +} + // FlagModelConfigResyncInterval names the model config resync interval. const FlagModelConfigResyncInterval = "model-config-resync-interval" diff --git a/core/config/distributed_config_test.go b/core/config/distributed_config_test.go index 4296054379a3..7192f1118f2e 100644 --- a/core/config/distributed_config_test.go +++ b/core/config/distributed_config_test.go @@ -89,9 +89,43 @@ var _ = Describe("DistributedConfig flag-name constants", func() { Entry("backend upgrade timeout", config.FlagBackendUpgradeTimeout, "backend-upgrade-timeout"), Entry("model load timeout", config.FlagModelLoadTimeout, "model-load-timeout"), Entry("model config resync interval", config.FlagModelConfigResyncInterval, "model-config-resync-interval"), + Entry("carrier prepare timeout", config.FlagCarrierPrepareTimeout, "carrier-prepare-timeout"), + Entry("carrier transition window", config.FlagCarrierTransitionWindow, "carrier-transition-window"), + Entry("carrier max drain", config.FlagCarrierMaxDrain, "carrier-max-drain"), ) }) +var _ = Describe("DistributedConfig without a NATS URL", func() { + It("is valid: a deployment with only PostgreSQL runs on the tunnel carrier", func() { + Expect(config.DistributedConfig{Enabled: true}.Validate()).To(Succeed()) + }) + + It("keeps the NATS URL when one is given", func() { + o := config.NewApplicationConfig(config.WithNatsURL("nats://localhost:4222")) + Expect(o.Distributed.NatsURL).To(Equal("nats://localhost:4222")) + }) + + DescribeTable("rejects a negative timing of a change of carrier", + func(set func(*config.DistributedConfig), flag string) { + c := config.DistributedConfig{Enabled: true} + set(&c) + err := c.Validate() + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring(flag)) + }, + Entry("prepare timeout", func(c *config.DistributedConfig) { c.CarrierPrepareTimeout = -time.Second }, config.FlagCarrierPrepareTimeout), + Entry("transition window", func(c *config.DistributedConfig) { c.CarrierTransitionWindow = -time.Second }, config.FlagCarrierTransitionWindow), + Entry("max drain", func(c *config.DistributedConfig) { c.CarrierMaxDrain = -time.Second }, config.FlagCarrierMaxDrain), + ) + + It("sets the timings with the options", func() { + o := config.NewApplicationConfig(config.WithCarrierTimings(10*time.Second, 20*time.Second, 30*time.Second)) + Expect(o.Distributed.CarrierPrepareTimeout).To(Equal(10 * time.Second)) + Expect(o.Distributed.CarrierTransitionWindow).To(Equal(20 * time.Second)) + Expect(o.Distributed.CarrierMaxDrain).To(Equal(30 * time.Second)) + }) +}) + var _ = Describe("DistributedConfig.Validate negative-duration errors", func() { It("rejects a negative BackendInstallTimeout with the flag name in the error", func() { c := config.DistributedConfig{ diff --git a/core/http/app.go b/core/http/app.go index b6211686baf4..e360a97dc1fd 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -28,12 +28,14 @@ import ( "github.com/mudler/LocalAI/core/application" "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/distributed" "github.com/mudler/LocalAI/core/services/finetune" "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/quantization" + "github.com/mudler/LocalAI/core/services/tunnel" "github.com/mudler/xlog" ) @@ -500,7 +502,7 @@ func API(application *application.Application) (*echo.Echo, error) { // happened to admit, and a load-balanced UI poll alternates between // "operation visible" and "operation gone" between replicas. if d := application.Distributed(); d != nil { - opcache.SetMessagingClient(d.Nats) + opcache.SetMessagingClient(d.Broadcaster) if d.DistStores != nil && d.DistStores.Gallery != nil { opcache.SetGalleryStore(d.DistStores.Gallery) } @@ -518,10 +520,10 @@ func API(application *application.Application) (*echo.Echo, error) { // In distributed mode pass the shared NATS client + PostgreSQL store so // fine-tune jobs stay consistent across replicas (the SyncedMap broadcasts // mutations and hydrates from the DB); standalone passes nil for both. - var ftNats messaging.MessagingClient + var ftNats messaging.Broadcaster var ftStore *distributed.FineTuneStore if d := application.Distributed(); d != nil { - ftNats = d.Nats + ftNats = d.Broadcaster if d.DistStores != nil && d.DistStores.FineTune != nil { ftStore = d.DistStores.FineTune } @@ -540,10 +542,10 @@ func API(application *application.Application) (*echo.Echo, error) { // In distributed mode pass the shared NATS client + PostgreSQL store so // quantization jobs stay consistent across replicas (the SyncedMap broadcasts // mutations and hydrates from the DB); standalone passes nil for both. - var quantNats messaging.MessagingClient + var quantNats messaging.Broadcaster var quantStore *distributed.QuantStore if d := application.Distributed(); d != nil { - quantNats = d.Nats + quantNats = d.Broadcaster if d.DistStores != nil && d.DistStores.Quant != nil { quantStore = d.DistStores.Quant } @@ -562,16 +564,36 @@ func API(application *application.Application) (*echo.Echo, error) { var registry *nodes.NodeRegistry var remoteUnloader nodes.NodeCommandSender var workerHTTPDial nodes.WorkerNetDialerFor + var tunnels *tunnel.Registry + var registerOpts []localai.RegisterOption if d := application.Distributed(); d != nil { registry = d.Registry workerHTTPDial = d.WorkerHTTPDial + tunnels = d.Tunnels + registerOpts = append(registerOpts, localai.WithCarrierReader(d.Carriers)) + registerOpts = append(registerOpts, localai.WithTunnelDisconnector(d)) + registerOpts = append(registerOpts, localai.WithNATSHandover(d.Runtime)) if d.Router != nil { remoteUnloader = d.Router.Unloader() } } natsCfg := distCfg.NatsAuthConfig() - routes.RegisterNodeSelfServiceRoutes(e, registry, distCfg.RegistrationToken, distCfg.AutoApproveNodes, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, natsCfg) - routes.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, natsCfg, workerHTTPDial) + routes.RegisterNodeSelfServiceRoutes(e, registry, distCfg.RegistrationToken, distCfg.AutoApproveNodes, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, natsCfg, registerOpts...) + var ( + instances *cluster.Registry + onPeer func(string, *tunnel.Session) + ) + if d := application.Distributed(); d != nil { + instances = d.Instances + onPeer = d.PeerSessions.Accept + } + routes.RegisterClusterRoutes(e, registry, tunnels, instances, onPeer) + var clusterAdmin routes.ClusterAdmin + if d := application.Distributed(); d != nil { + clusterAdmin = routes.ClusterAdmin{Switch: d.Switch, Settings: d.Settings, Prober: d.Runtime, NATS: d.Runtime} + } + routes.RegisterClusterAdminRoutes(e, adminMiddleware, clusterAdmin) + routes.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, natsCfg, workerHTTPDial, registerOpts...) // Distributed SSE routes (job progress + agent events via NATS) if d := application.Distributed(); d != nil { diff --git a/core/http/auth/helpers_test.go b/core/http/auth/helpers_test.go index 1b52d4500149..8490211f9111 100644 --- a/core/http/auth/helpers_test.go +++ b/core/http/auth/helpers_test.go @@ -122,6 +122,17 @@ func newNodeSelfServiceTestApp(db *gorm.DB, appConfig *config.ApplicationConfig, return e } +// newClusterTestApp uses the real cluster route registration, with no registry, +// so a request that the global auth middleware lets through reaches a handler +// that answers 503 and one that it stops gets 401. +func newClusterTestApp(db *gorm.DB, appConfig *config.ApplicationConfig) *echo.Echo { + e := echo.New() + e.Use(auth.Middleware(db, appConfig)) + routes.RegisterClusterRoutes(e, nil, nil, nil, nil) + e.GET("/api/cluster/other", ok) + return e +} + // newAdminTestApp creates an Echo app with admin-protected routes. func newAdminTestApp(db *gorm.DB, appConfig *config.ApplicationConfig) *echo.Echo { e := echo.New() diff --git a/core/http/auth/middleware_test.go b/core/http/auth/middleware_test.go index bcbba0096193..6f6e729c6e77 100644 --- a/core/http/auth/middleware_test.go +++ b/core/http/auth/middleware_test.go @@ -442,6 +442,37 @@ var _ = Describe("Auth Middleware", func() { } }) + It("delegates the tunnel dial to its own credential check, and only that path", func() { + app = newClusterTestApp(db, appConfig) + + // A bearer token that is no API key and no session. The middleware + // lets the dial through, and the handler (which has no registry here) + // answers 503. If the middleware stopped it, the answer would be 401. + dial := doRequest(app, http.MethodGet, "/api/cluster/connect?id=w1", withBearerToken("own-tunnel-token")) + Expect(dial.Code).To(Equal(http.StatusServiceUnavailable)) + + // Another route under the same prefix is not delegated. + other := doRequest(app, http.MethodGet, "/api/cluster/other", withBearerToken("own-tunnel-token")) + Expect(other.Code).To(Equal(http.StatusUnauthorized)) + }) + + It("delegates the peer link to its own credential check, and only that path", func() { + app = newClusterTestApp(db, appConfig) + + // The header of the peer credential, which is no API key and no + // session. The middleware lets the dial through, and the handler (which + // has no table of replicas here) answers 503. If the middleware stopped + // it, the answer would be 401. + peer := doRequest(app, http.MethodGet, "/api/cluster/peer?id=replica-a", func(r *http.Request) { + r.Header.Set("X-LocalAI-Peer-Token", "own-peer-credential") + }) + Expect(peer.Code).To(Equal(http.StatusServiceUnavailable)) + + // A dial with no credential is refused whoever answers. + anonymous := doRequest(app, http.MethodGet, "/api/cluster/peer?id=replica-a") + Expect(anonymous.Code).To(Equal(http.StatusUnauthorized)) + }) + It("lets downstream node auth reject missing and invalid registration tokens", func() { const registrationToken = "node-registration-secret" app = newNodeSelfServiceTestApp(db, appConfig, registrationToken) diff --git a/core/http/auth/public_routes.go b/core/http/auth/public_routes.go index c2509e9a9f6c..55ea8d532da9 100644 --- a/core/http/auth/public_routes.go +++ b/core/http/auth/public_routes.go @@ -80,8 +80,24 @@ func isPublicRoute(method, path string) bool { return false } +// ClusterConnectPath is the route that a worker dials to open its tunnel. The +// handler checks the own tunnel credential of the node, which is not a user, +// a session or an API key, so the global middleware lets the request through. +// +// It is one exact path and not the prefix /api/cluster/, so that a route that is +// added under that prefix later is covered by the global middleware unless its +// author decides otherwise. The path is written out here because this package +// must not import the tunnel package. The spec of the connect endpoint compares +// it with tunnel.ConnectPath. +const ClusterConnectPath = "/api/cluster/connect" + +// ClusterPeerPath is the route that a replica dials to reach a tunnel that +// another replica holds. The handler checks the own credential of the dialling +// replica. The spec of the peer endpoint compares it with tunnel.PeerPath. +const ClusterPeerPath = "/api/cluster/peer" + // usesAlternativeAuthentication identifies requests whose credentials are // validated by route-group middleware instead of the global auth middleware. func usesAlternativeAuthentication(path string) bool { - return strings.HasPrefix(path, "/api/node/") + return strings.HasPrefix(path, "/api/node/") || path == ClusterConnectPath || path == ClusterPeerPath } diff --git a/core/http/endpoints/cluster/cluster_suite_test.go b/core/http/endpoints/cluster/cluster_suite_test.go new file mode 100644 index 000000000000..6127ee8c9100 --- /dev/null +++ b/core/http/endpoints/cluster/cluster_suite_test.go @@ -0,0 +1,13 @@ +package cluster_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestCluster(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Cluster endpoints test suite") +} diff --git a/core/http/endpoints/cluster/connect.go b/core/http/endpoints/cluster/connect.go new file mode 100644 index 000000000000..b2bddb150ac6 --- /dev/null +++ b/core/http/endpoints/cluster/connect.go @@ -0,0 +1,228 @@ +// Package cluster serves the endpoints that the processes of a cluster call +// among themselves. +package cluster + +import ( + "crypto/sha256" + "crypto/subtle" + "encoding/hex" + "errors" + "net/http" + "strings" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/mudler/xlog" + "gorm.io/gorm" +) + +// ConnectHandler serves the door that a worker dials. It checks the dial against +// the own credential of the node, upgrades it to a websocket, starts the server +// side of a session on it and gives the session to the tunnel registry. +// +// The worker dials out and never listens, so a worker behind NAT, in another +// cluster or on a laptop needs no inbound port. The worker is the yamux client +// and this side is the server. This side opens the streams. It accepts none: the +// frontend asks and the worker answers, so a stream that the worker opened would +// wait in the accept queue for ever. +// +// A worker holds two sessions. The query parameter lane names the one that a +// dial is for: "inference", which is also the default, or "bulk". +// +// The route is registered on every replica of every deployment, so that the +// route coverage test walks it. A registry that is nil is therefore a real +// answer (503) and not a panic. +// +// It is not in auth.RouteFeatureRegistry. That registry gates a route on the +// features of a user, and here there is no user: the caller is a worker process +// with a machine credential. The global auth middleware runs on this path and +// does not reject it, because usesAlternativeAuthentication names the path as +// one that checks its own credentials. Nothing here reads what the middleware +// set. +func ConnectHandler(registry *nodes.NodeRegistry, tunnels *tunnel.Registry) echo.HandlerFunc { + // The default origin check of gorilla refuses a browser of another origin + // and admits a client with no Origin header, which every worker is. + upgrader := tunnel.NewUpgrader() + + return func(c echo.Context) error { + // All checks come before the upgrade, and their order matters. The + // credential is read first, because a dial with no Authorization header + // is the anonymous case and the route coverage test sends exactly that, + // with no query. It must get 401 and not a 400 for a missing node id. + token, ok := bearerToken(c.Request()) + if !ok { + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + // Not 401: a frontend that is not distributed cannot authenticate + // anyone, and "unauthorized" would send the operator to look for a + // problem with a token that does not exist. + if registry == nil || tunnels == nil { + return echo.NewHTTPError(http.StatusServiceUnavailable, "distributed mode not enabled") + } + + nodeID := c.QueryParam("id") + if nodeID == "" { + return echo.NewHTTPError(http.StatusBadRequest, "missing node id") + } + + node, err := registry.Get(c.Request().Context(), nodeID) + switch { + case errors.Is(err, gorm.ErrRecordNotFound): + // 401 and not 404, so that a caller cannot list node IDs by status. + xlog.Debug("worker tunnel dial named an unknown node", "node", nodeID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + case err != nil: + // A query that failed is neither a rejection nor an absence. To tell + // a worker that its credentials are wrong when the database could not + // be read sends it to register again instead of retrying, and a + // worker that registers again has thrown away the credential that + // its tunnel is keyed by. + xlog.Error("Looking up a worker for its tunnel dial failed", "node", nodeID, "error", err) + return echo.NewHTTPError(http.StatusInternalServerError, "node lookup failed") + } + + // An empty hash and a wrong token are two problems with two fixes. An + // empty hash means that the node has not registered since the tunnel + // credential existed, so it has no secret to check and must register. + // A wrong token is a stale credential, usually from before a rotation. + if node.TunnelTokenHash == "" { + // Debug, because every such worker fails in this way on every dial. + xlog.Debug("refusing a worker tunnel: this node has no tunnel credential, so it has not registered since they were introduced", + "node", nodeID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + if !authorizedWorker(token, node.TunnelTokenHash) { + xlog.Debug("worker tunnel dial presented the wrong token", "node", nodeID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + // The credential is right, but the node is not allowed yet: 403, because + // the fix is that an admin approves the node and a different credential + // would not help. A tunnel is a standing pipe into the worker that other + // replicas relay to, so a node that waits for approval does not get one. + // Draining and unhealthy nodes keep their tunnels: draining means finish + // what you have, and a node marked unhealthy for missed heartbeats needs + // the pipe to recover. + if node.Status == nodes.StatusPending { + xlog.Warn("Refusing a worker tunnel: this node is awaiting admin approval", "node", nodeID) + return echo.NewHTTPError(http.StatusForbidden, "node is pending approval") + } + + lane, err := tunnel.ParseLane(c.QueryParam("lane")) + if err != nil { + return echo.NewHTTPError(http.StatusBadRequest, err.Error()) + } + + // The bulk lane belongs to the replica that holds the inference lane of + // the node. A load balancer can send the dial to another replica. Only a + // status can say that, and the status is possible before the upgrade, so + // the worker is told to dial again at once. + if lane == tunnel.LaneBulk && !tunnels.Holds(nodeID) { + return echo.NewHTTPError(http.StatusConflict, "this replica does not hold the tunnel of the node") + } + + ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil) + if err != nil { + // Upgrade has written its own failure to the client. + xlog.Debug("worker tunnel upgrade failed", "node", nodeID, "error", err) + return nil + } + + sess, err := tunnel.ServerSession(ws, lane) + if err != nil { + xlog.Error("Worker tunnel session setup failed", "node", nodeID, "error", err) + _ = ws.Close() + return nil + } + + // net/http recovers a panic of this goroutine but does not close the + // connection that was hijacked, and the recovery middleware does not + // either. A panic below would leave the worker with a session that this + // replica has no entry for and will never detach. Attach does database + // work, so this is not far-fetched. The panic goes on, because it is a + // bug and the recovery middleware is what reports it. + defer func() { + if r := recover(); r != nil { + _ = sess.Close() + panic(r) + } + }() + + // The connection is hijacked, so no status can reach the worker any more. + // A failure is a closed socket, which is what the reconnect loop of the + // worker reads. + attachment, err := tunnels.Attach(c.Request().Context(), nodeID, lane, sess) + if err != nil { + if errors.Is(err, tunnel.ErrNotOwner) { + // The inference session ended between the check before the + // upgrade and here. The worker dials again, and this is not an + // error of this replica. + xlog.Debug("worker bulk tunnel dial lost the inference session of its node", "node", nodeID) + } else { + xlog.Error("Attaching a worker tunnel failed", "node", nodeID, "lane", lane, "error", err) + } + _ = sess.Close() + return nil + } + + xlog.Info("Worker tunnel established", "node", nodeID, "lane", lane, "remote", ws.RemoteAddr().String()) + // The session outlives this handler, so something else has to see it end. + // yamux closes its channel from the receive loop as soon as the + // connection fails, and its keepalive notices a worker that vanished + // without a FIN. + go func() { + <-sess.CloseChan() + // The token that Attach returned. Detach compares it by equality: it + // names this attachment, so a worker that has already dialled again + // is not removed by the end of its predecessor. + tunnels.Detach(nodeID, lane, attachment) + xlog.Debug("worker tunnel closed", "node", nodeID, "lane", lane) + }() + return nil + } +} + +// bearerToken returns the token of an Authorization: Bearer header, and whether +// there was one. A missing credential and a wrong credential are different +// answers on purpose: "no credential" decides the 401 before the upgrade, and it +// must be decidable before anything about the node is known. +func bearerToken(r *http.Request) (string, bool) { + // The scheme is not case-sensitive (RFC 7235). The token is. + const prefix = "Bearer " + header := r.Header.Get("Authorization") + if len(header) < len(prefix) || !strings.EqualFold(header[:len(prefix)], prefix) { + return "", false + } + token := header[len(prefix):] + if token == "" { + return "", false + } + return token, true +} + +// authorizedWorker compares a token with the hash on the row of the node, in +// constant time. +// +// The token is the own credential of the node and not the registration token of +// the deployment. A tunnel is a standing pipe into a worker. A credential that +// opens every worker would let one leak take over the traffic of any worker +// whose ID can be read, by claiming its tunnel. The secret is minted for each +// node at registration, returned to that worker once and stored as this hash. +// It is not TokenHash, which hashes the token that the worker registered with. +// On most deployments that is the shared registration token. +// +// The credential stops a caller from claiming a node by its ID alone. It does +// not stop a holder of the registration token: that token authorises a +// registration under any node name, and a registration mints a new credential +// and replaces the old one. This is the same trust as the NATS credentials that +// registration mints. A worker whose credential was replaced registers again +// by itself. +func authorizedWorker(token, storedHash string) bool { + if storedHash == "" { + return false + } + sum := sha256.Sum256([]byte(token)) + return subtle.ConstantTimeCompare([]byte(hex.EncodeToString(sum[:])), []byte(storedHash)) == 1 +} diff --git a/core/http/endpoints/cluster/connect_test.go b/core/http/endpoints/cluster/connect_test.go new file mode 100644 index 000000000000..dd54060923ea --- /dev/null +++ b/core/http/endpoints/cluster/connect_test.go @@ -0,0 +1,398 @@ +package cluster_test + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "io" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "time" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/http/auth" + clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + "github.com/mudler/LocalAI/core/services/worker" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +func hashOf(secret string) string { + sum := sha256.Sum256([]byte(secret)) + return hex.EncodeToString(sum[:]) +} + +var _ = Describe("Connect endpoint", func() { + const ( + registrationToken = "shared-registration-token" + nodeToken = "own-tunnel-token-of-w1" + ) + + var ( + ctx context.Context + db *gorm.DB + clusterR *cluster.Registry + nodeReg *nodes.NodeRegistry + tunnels *tunnel.Registry + srv *httptest.Server + nodeID string + ) + + // register adds a node the way the registration endpoint does: the + // registration token is stored as TokenHash, and the own credential of the + // node as TunnelTokenHash. + register := func(name string, autoApprove bool) *nodes.BackendNode { + GinkgoHelper() + node := &nodes.BackendNode{Name: name, NodeType: nodes.NodeTypeBackend, TokenHash: hashOf(registrationToken)} + Expect(nodeReg.Register(ctx, node, autoApprove)).To(Succeed()) + return node + } + + serve := func(registry *nodes.NodeRegistry, tunnels *tunnel.Registry) *httptest.Server { + GinkgoHelper() + e := echo.New() + e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(registry, tunnels)) + s := httptest.NewServer(e) + DeferCleanup(s.Close) + return s + } + + dial := func(query, token string, header ...http.Header) (*websocket.Conn, *http.Response, error) { + h := http.Header{} + if token != "" { + h.Set("Authorization", "Bearer "+token) + } + for _, extra := range header { + for k, v := range extra { + h[k] = v + } + } + url := "ws" + srv.URL[len("http"):] + tunnel.ConnectPath + query + return tunnel.NewDialer(5*time.Second).Dial(url, h) + } + + // plainGet is a request that is no websocket upgrade. + plainGet := func(query, token string) *http.Response { + GinkgoHelper() + req, err := http.NewRequest(http.MethodGet, srv.URL+tunnel.ConnectPath+query, nil) + Expect(err).ToNot(HaveOccurred()) + if token != "" { + req.Header.Set("Authorization", "Bearer "+token) + } + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + return resp + } + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + clusterR = cluster.NewRegistry(db) + Expect(clusterR.Register(ctx, "replica-a", "test", 0, "")).To(Succeed()) + var err error + nodeReg, err = nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + tunnels = tunnel.NewRegistry(clusterR, "replica-a") + + node := register("w1", true) + nodeID = node.ID + Expect(nodeReg.SetTunnelTokenHash(ctx, nodeID, hashOf(nodeToken))).To(Succeed()) + + srv = serve(nodeReg, tunnels) + }) + + Describe("the answers before the upgrade", func() { + It("answers 401 to a dial with no credential, before anything else is checked", func() { + // No query string at all: an anonymous dial must not get a 400. + Expect(plainGet("", "").StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(plainGet("?id="+nodeID, "").StatusCode).To(Equal(http.StatusUnauthorized)) + // Not even the 503 for a missing registry may come first. + bare := serve(nil, nil) + resp, err := http.Get(bare.URL + tunnel.ConnectPath) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 503 to a credential when the frontend has no registry", func() { + bare := serve(nil, nil) + req, _ := http.NewRequest(http.MethodGet, bare.URL+tunnel.ConnectPath+"?id=x", nil) + req.Header.Set("Authorization", "Bearer anything") + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable)) + }) + + It("answers 400 to a credential with no node id", func() { + Expect(plainGet("", nodeToken).StatusCode).To(Equal(http.StatusBadRequest)) + }) + + It("answers 401, not 404, for a node it does not know", func() { + Expect(plainGet("?id=nobody", nodeToken).StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 401 to a wrong token", func() { + Expect(plainGet("?id="+nodeID, "not-the-token").StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("never accepts the shared registration token", func() { + Expect(plainGet("?id="+nodeID, registrationToken).StatusCode).To(Equal(http.StatusUnauthorized)) + _, resp, err := dial("?id="+nodeID, registrationToken) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(tunnels.Holds(nodeID)).To(BeFalse()) + }) + + It("answers 401 to a node that has no tunnel credential, whatever it presents", func() { + other := register("w-old", true) + Expect(plainGet("?id="+other.ID, "").StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(plainGet("?id="+other.ID, hashOf("")).StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(plainGet("?id="+other.ID, registrationToken).StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 403 to a node that waits for approval, and takes no tunnel from it", func() { + pending := register("w-pending", false) + Expect(pending.Status).To(Equal(nodes.StatusPending)) + Expect(nodeReg.SetTunnelTokenHash(ctx, pending.ID, hashOf("pending-token"))).To(Succeed()) + + _, resp, err := dial("?id="+pending.ID, "pending-token") + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusForbidden)) + Expect(tunnels.Holds(nodeID)).To(BeFalse()) + + Expect(nodeReg.ApproveNode(ctx, pending.ID)).To(Succeed()) + ws, _, err := dial("?id="+pending.ID, "pending-token") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(ws.Close) + }) + + It("answers 500 and not 401 when the node cannot be read", func() { + sqlDB, err := db.DB() + Expect(err).ToNot(HaveOccurred()) + Expect(sqlDB.Close()).To(Succeed()) + Expect(plainGet("?id="+nodeID, nodeToken).StatusCode).To(Equal(http.StatusInternalServerError)) + }) + + It("answers 400 to a lane it does not know", func() { + Expect(plainGet("?id="+nodeID+"&lane=express", nodeToken).StatusCode).To(Equal(http.StatusBadRequest)) + }) + + It("leaves the table alone when the request is no upgrade", func() { + resp := plainGet("?id="+nodeID, nodeToken) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + _, _, err := clusterR.Owner(ctx, nodeID) + Expect(err).To(MatchError(cluster.ErrNoConnection)) + }) + + It("refuses a browser of another origin", func() { + _, resp, err := dial("?id="+nodeID, nodeToken, http.Header{"Origin": []string{"https://evil.example"}}) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusForbidden)) + Expect(tunnels.Holds(nodeID)).To(BeFalse()) + }) + }) + + Describe("a dial with the own credential of the node", func() { + It("claims the node for this replica and carries streams to the worker", func() { + ws, _, err := dial("?id="+nodeID, nodeToken) + Expect(err).ToNot(HaveOccurred()) + worker, err := tunnel.ClientSession(ws, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(worker.Close) + + Eventually(func() bool { return tunnels.Holds(nodeID) }).Should(BeTrue()) + owner, _, err := clusterR.Owner(ctx, nodeID) + Expect(err).ToNot(HaveOccurred()) + Expect(owner).To(Equal("replica-a")) + + go func() { + for { + st, err := worker.AcceptStream() + if err != nil { + return + } + go func() { defer func() { _ = st.Close() }(); _, _ = io.Copy(st, st) }() + } + }() + st, err := tunnels.Open(ctx, nodeID, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + _, err = st.Write([]byte("ping")) + Expect(err).ToNot(HaveOccurred()) + got := make([]byte, 4) + _, err = io.ReadFull(st, got) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal("ping")) + }) + + It("releases the claim when the worker goes away", func() { + ws, _, err := dial("?id="+nodeID, nodeToken) + Expect(err).ToNot(HaveOccurred()) + worker, err := tunnel.ClientSession(ws, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + Eventually(func() bool { return tunnels.Holds(nodeID) }).Should(BeTrue()) + + Expect(worker.Close()).To(Succeed()) + + Eventually(func() bool { return tunnels.Holds(nodeID) }, 5*time.Second).Should(BeFalse()) + // The registry drops the entry first and releases the claim after it. + Eventually(func() error { + _, _, err := clusterR.Owner(ctx, nodeID) + return err + }, 5*time.Second).Should(MatchError(cluster.ErrNoConnection)) + }) + + It("replaces the tunnel when the worker dials again", func() { + ws1, _, err := dial("?id="+nodeID, nodeToken) + Expect(err).ToNot(HaveOccurred()) + first, err := tunnel.ClientSession(ws1, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + Eventually(func() bool { return tunnels.Holds(nodeID) }).Should(BeTrue()) + + ws2, _, err := dial("?id="+nodeID, nodeToken) + Expect(err).ToNot(HaveOccurred()) + second, err := tunnel.ClientSession(ws2, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(second.Close) + + Eventually(first.CloseChan(), 5*time.Second).Should(BeClosed()) + Expect(second.IsClosed()).To(BeFalse()) + // The end of the first session must not remove the second. + Consistently(func() bool { return tunnels.Holds(nodeID) }, time.Second).Should(BeTrue()) + _, _, err = clusterR.Owner(ctx, nodeID) + Expect(err).ToNot(HaveOccurred()) + }) + + It("stops accepting the old credential after a new one is minted", func() { + Expect(nodeReg.SetTunnelTokenHash(ctx, nodeID, hashOf("rotated"))).To(Succeed()) + _, resp, err := dial("?id="+nodeID, nodeToken) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + ws, _, err := dial("?id="+nodeID, "rotated") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(ws.Close) + }) + }) + + Describe("the bulk lane", func() { + It("is refused with 409 before the inference lane is held, and accepted after", func() { + _, resp, err := dial("?id="+nodeID+"&lane=bulk", nodeToken) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusConflict)) + _, _, ownerErr := clusterR.Owner(ctx, nodeID) + Expect(ownerErr).To(MatchError(cluster.ErrNoConnection), "a bulk dial must not claim the node") + + wsI, _, err := dial("?id="+nodeID, nodeToken) + Expect(err).ToNot(HaveOccurred()) + inference, err := tunnel.ClientSession(wsI, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(inference.Close) + Eventually(func() bool { return tunnels.Holds(nodeID) }).Should(BeTrue()) + + wsB, _, err := dial("?id="+nodeID+"&lane=bulk", nodeToken) + Expect(err).ToNot(HaveOccurred()) + bulk, err := tunnel.ClientSession(wsB, tunnel.LaneBulk) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(bulk.Close) + + // A transfer stream reaches the worker on the bulk session. + reached := make(chan string, 1) + go func() { + st, err := bulk.AcceptStream() + if err != nil { + return + } + buf := make([]byte, 4) + _, _ = io.ReadFull(st, buf) + reached <- string(buf) + }() + Eventually(func() error { + st, err := tunnels.Open(ctx, nodeID, tunnel.LaneBulk) + if err != nil { + return err + } + _, err = st.Write([]byte("blob")) + return err + }, 5*time.Second).Should(Succeed()) + Eventually(reached).Should(Receive(Equal("blob"))) + }) + + It("authenticates the bulk lane like the inference lane", func() { + _, resp, err := dial("?id="+nodeID+"&lane=bulk", registrationToken) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + }) +}) + +var _ = Describe("A worker whose credential was rotated on the frontend", func() { + It("registers again and gets its tunnel back without a restart", func() { + ctx := context.Background() + db := testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + clusterR := cluster.NewRegistry(db) + Expect(clusterR.Register(ctx, "replica-a", "test", 0, "")).To(Succeed()) + nodeReg, err := nodes.NewNodeRegistry(db) + Expect(err).ToNot(HaveOccurred()) + tunnels := tunnel.NewRegistry(clusterR, "replica-a") + + node := &nodes.BackendNode{Name: "w1", NodeType: nodes.NodeTypeBackend, TokenHash: hashOf("registration")} + Expect(nodeReg.Register(ctx, node, true)).To(Succeed()) + var mu sync.Mutex + current := "credential-1" + Expect(nodeReg.SetTunnelTokenHash(ctx, node.ID, hashOf(current))).To(Succeed()) + held := func() string { mu.Lock(); defer mu.Unlock(); return current } + + e := echo.New() + e.GET(tunnel.ConnectPath, clusterapi.ConnectHandler(nodeReg, tunnels)) + srv := httptest.NewServer(e) + DeferCleanup(srv.Close) + + // Another process registers under the same name and replaces the + // credential on the frontend. This worker still holds credential-1. + mu.Lock() + current = "credential-2-of-another-process" + Expect(nodeReg.SetTunnelTokenHash(ctx, node.ID, hashOf(current))).To(Succeed()) + mu.Unlock() + workerToken := "credential-1" + var reregistered atomic.Int32 + w, err := worker.StartTunnel(ctx, worker.TunnelConfig{ + FrontendURL: srv.URL, + NodeID: node.ID, + Token: func() string { mu.Lock(); defer mu.Unlock(); return workerToken }, + Services: map[string]worker.LocalService{}, + Reauthorize: func(ctx context.Context) error { + // What registration does on the frontend: mint, store the hash, + // return the value once. + mu.Lock() + defer mu.Unlock() + reregistered.Add(1) + current = "credential-3" + workerToken = current + return nodeReg.SetTunnelTokenHash(ctx, node.ID, hashOf(current)) + }, + }) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(w.Close) + Eventually(reregistered.Load, 20*time.Second).Should(BeNumerically(">=", 1)) + Eventually(w.Connected, 20*time.Second).Should(BeTrue()) + Expect(held()).To(Equal("credential-3")) + }) +}) + +var _ = Describe("The path of the connect route", func() { + It("is the one that the global auth middleware lets through", func() { + Expect(auth.ClusterConnectPath).To(Equal(tunnel.ConnectPath)) + }) +}) diff --git a/core/http/endpoints/cluster/peer.go b/core/http/endpoints/cluster/peer.go new file mode 100644 index 000000000000..8dd8a55b3fef --- /dev/null +++ b/core/http/endpoints/cluster/peer.go @@ -0,0 +1,122 @@ +package cluster + +import ( + "errors" + "net/http" + + "github.com/labstack/echo/v4" + "github.com/mudler/xlog" + + clustersvc "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/tunnel" +) + +// PeerHandler serves the link that a replica dials to reach a worker tunnel that +// another replica holds. The dial becomes a websocket, the websocket becomes a +// multiplexed session, and onSession gets the session. +// +// One credential decides, and it is the credential of the dialling replica. The +// id in ?id= is a claim. The handler resolves it to a row of the instances table +// and checks the secret in the header against the hash that the row publishes. +// The registration token of the deployment opens no peer link: it is held by +// every worker, and a worker that could pose as a replica could relay to every +// tunnel that a replica owns. +// +// All checks come before the upgrade, and the credential comes first. A dial with +// no credential is the anonymous case and gets 401 whatever else is missing. +// +// onSession runs on the goroutine of the request, so it returns promptly. The +// session outlives the handler, because the upgrade hijacks the connection, and +// onSession owns closing it. +func PeerHandler(instances *clustersvc.Registry, onSession func(peerID string, sess *tunnel.Session)) echo.HandlerFunc { + // The default origin check of gorilla refuses a browser of another origin and + // admits a client with no Origin header, which every peer is. + upgrader := tunnel.NewUpgrader() + + return func(c echo.Context) error { + presented := c.Request().Header.Get(clustersvc.PeerIdentityHeader) + if presented == "" { + // A missing header is refused and not waved through. Accepting a dial + // with no credential would leave the hole as open as it was, because an + // attacker omits the header too. + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + // Not 401. A frontend with no registry cannot resolve any replica, and + // "unauthorized" would send an operator after a problem with a credential + // that does not exist. The route is registered in every deployment so that + // the route coverage test sees it, and a handler that treated a nil registry + // as "nothing to check" would publish an unauthenticated multiplexer. + if instances == nil { + return echo.NewHTTPError(http.StatusServiceUnavailable, "distributed mode not enabled") + } + + peerID := c.QueryParam("id") + if peerID == "" { + return echo.NewHTTPError(http.StatusBadRequest, "missing peer id") + } + + inst, err := instances.Get(c.Request().Context(), peerID) + switch { + case errors.Is(err, clustersvc.ErrInstanceNotFound): + // 401 and not 404, so that a caller cannot list replica ids by status. + xlog.Warn("Refusing a peer link: the dial named a replica this deployment has no row for", "peer", peerID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + case err != nil: + // A query that failed is neither a rejection nor an absence. 401 on an + // unreadable database would tell a healthy replica that its credential + // is wrong. + xlog.Error("Looking up a replica for its peer dial failed", "peer", peerID, "error", err) + return echo.NewHTTPError(http.StatusInternalServerError, "peer lookup failed") + } + + // Two operator problems with two cures. An empty stored hash is a replica + // that has not registered a credential, so there is nothing to check and it + // must register again, which a restart does. A mismatch is a dialler that + // presents the wrong secret. An empty hash never matches: it is not "no + // restriction". + if inst.PeerTokenHash == "" { + xlog.Warn("Refusing a peer link: the replica it claims to be has no peer credential published, so nothing here can verify the claim", + "peer", peerID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + if !clustersvc.PeerTokenMatches(presented, inst.PeerTokenHash) { + xlog.Warn("Refusing a peer link: the dial presented the wrong credential for the replica it claims to be", "peer", peerID) + return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized") + } + + ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil) + if err != nil { + // Upgrade has written its own failure to the client. + xlog.Debug("tunnel peer link upgrade failed", "peer", peerID, "error", err) + return nil + } + + sess, err := tunnel.PeerServerSession(ws) + if err != nil { + xlog.Error("Tunnel peer link session setup failed", "peer", peerID, "error", err) + _ = ws.Close() + return nil + } + + if onSession == nil { + // Nothing reads this session, so the peer must not believe it has a + // live link. + _ = sess.Close() + return nil + } + + xlog.Debug("tunnel peer link established", "peer", peerID, "remote", ws.RemoteAddr().String()) + // net/http recovers a panic of this goroutine but does not close the + // hijacked connection, so a panicking callback would leave the peer with a + // link that nobody accepts streams on. + defer func() { + if r := recover(); r != nil { + _ = sess.Close() + panic(r) + } + }() + onSession(peerID, sess) + return nil + } +} diff --git a/core/http/endpoints/cluster/peer_test.go b/core/http/endpoints/cluster/peer_test.go new file mode 100644 index 000000000000..d2d5ec58e63d --- /dev/null +++ b/core/http/endpoints/cluster/peer_test.go @@ -0,0 +1,305 @@ +package cluster_test + +import ( + "context" + "errors" + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "time" + + "github.com/gorilla/websocket" + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/http/auth" + clusterapi "github.com/mudler/LocalAI/core/http/endpoints/cluster" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/core/services/tunnel" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" +) + +// These specs drive the production handler over a real websocket and a real +// yamux session, because what they pin is a property of the place where the +// connection is hijacked: the id of a peer is a query parameter, and whether it +// is checked is decided there. +var _ = Describe("Peer link endpoint", func() { + const ( + ownerID = "replica-owner" + callerID = "replica-caller" + // The registration token is held by every worker. Holding it must open + // no peer link. + registrationToken = "shared-registration-token" + ) + + var ( + ctx context.Context + db *gorm.DB + reg *cluster.Registry + sessions *tunnel.PeerSessions + srv *httptest.Server + ownerCred cluster.PeerCredential + callerCred cluster.PeerCredential + ) + + BeforeEach(func() { + ctx = context.Background() + db = testutil.SetupTestDB() + Expect(cluster.Migrate(ctx, db)).To(Succeed()) + reg = cluster.NewRegistry(db) + + // The owner holds a worker, and accepts peer links. + ownerTunnels := tunnel.NewRegistry(reg, ownerID) + sessions = tunnel.NewPeerSessions(tunnel.NewRelay(ownerTunnels).Stream) + DeferCleanup(sessions.Close) + + e := echo.New() + e.GET(tunnel.PeerPath, clusterapi.PeerHandler(reg, sessions.Accept)) + srv = httptest.NewServer(e) + DeferCleanup(srv.Close) + + ownerCred, callerCred = cluster.NewPeerCredential(), cluster.NewPeerCredential() + addr := strings.TrimPrefix(srv.URL, "http://") + Expect(reg.RegisterPeer(ctx, ownerID, "test", 0, "", addr, ownerCred.Hash())).To(Succeed()) + Expect(reg.RegisterPeer(ctx, callerID, "test", 0, "", "127.0.0.1:9", callerCred.Hash())).To(Succeed()) + }) + + // raw dials the route with the headers a spec chooses. + raw := func(query string, header http.Header) (*http.Response, error) { + GinkgoHelper() + url := "ws" + strings.TrimPrefix(srv.URL, "http") + tunnel.PeerPath + query + ws, resp, err := tunnel.NewDialer(5*time.Second).Dial(url, header) + if ws != nil { + DeferCleanup(func() { _ = ws.Close() }) + } + return resp, err + } + + Describe("the answers before the upgrade", func() { + It("answers 401 to a dial with no credential, whatever else is missing", func() { + resp, err := raw("", nil) + Expect(err).To(MatchError(websocket.ErrBadHandshake)) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 401 to a dial that holds only the registration token", func() { + header := http.Header{"Authorization": []string{"Bearer " + registrationToken}} + resp, err := raw("?id="+callerID, header) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 503 to a credential when the frontend has no table of replicas", func() { + e := echo.New() + e.GET(tunnel.PeerPath, clusterapi.PeerHandler(nil, nil)) + bare := httptest.NewServer(e) + DeferCleanup(bare.Close) + req, _ := http.NewRequest(http.MethodGet, bare.URL+tunnel.PeerPath+"?id=x", nil) + req.Header.Set(cluster.PeerIdentityHeader, "anything") + resp, err := http.DefaultClient.Do(req) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(resp.Body.Close) + Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable)) + }) + + It("answers 400 to a credential with no replica id", func() { + resp, err := raw("", http.Header{cluster.PeerIdentityHeader: []string{callerCred.Token()}}) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusBadRequest)) + }) + + It("answers 401, not 404, for a replica it has no row for", func() { + resp, err := raw("?id=ghost", http.Header{cluster.PeerIdentityHeader: []string{callerCred.Token()}}) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + }) + + It("answers 401 for a replica whose row publishes no credential", func() { + Expect(db.Model(&cluster.Instance{}).Where("id = ?", callerID).Update("peer_token_hash", "").Error).To(Succeed()) + resp, err := raw("?id="+callerID, http.Header{cluster.PeerIdentityHeader: []string{callerCred.Token()}}) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized), "an empty hash authorises nobody") + }) + + It("answers 401 for the id of a replica with the credential of another", func() { + resp, err := raw("?id="+callerID, http.Header{cluster.PeerIdentityHeader: []string{cluster.NewPeerCredential().Token()}}) + Expect(err).To(HaveOccurred()) + Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized)) + Expect(sessions.Holds(callerID)).To(BeFalse()) + }) + }) + + Describe("a replica that proves its id", func() { + It("gets a link, and the owner holds exactly one session for it", func() { + pool := tunnel.NewPeerPool(callerID, callerCred, reg) + DeferCleanup(pool.Close) + st, err := pool.Open(ctx, ownerID) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = st.Close() }) + Eventually(func() bool { return sessions.Holds(callerID) }, "10s").Should(BeTrue()) + }) + + It("is refused when it declares another replica's id, and cannot evict the link of the real one", func() { + good := tunnel.NewPeerPool(callerID, callerCred, reg) + DeferCleanup(good.Close) + st, err := good.Open(ctx, ownerID) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = st.Close() }) + Eventually(func() bool { return sessions.Holds(callerID) }, "10s").Should(BeTrue()) + + impostor := tunnel.NewPeerPool(callerID, cluster.NewPeerCredential(), reg) + DeferCleanup(impostor.Close) + _, err = impostor.Open(ctx, ownerID) + Expect(err).To(MatchError(tunnel.ErrPeerRejected)) + Expect(err).To(MatchError(tunnel.ErrPeerUnreachable), "every caller that handles an unreachable peer handles a refusal") + Expect(errors.Is(err, cluster.ErrInstanceNotFound)).To(BeFalse(), "a refusal is not absence") + Expect(errors.Is(err, tunnel.ErrNoRoute)).To(BeFalse()) + Expect(sessions.Holds(callerID)).To(BeTrue(), "the link of the real replica stays") + }) + + It("reaches a worker that the owner holds, through the link, with the bytes intact", func() { + // A worker behind the owner. + ownerTunnels := tunnel.NewRegistry(reg, ownerID) + worker := &workerEcho{} + sessions.Close() + sessions = tunnel.NewPeerSessions(tunnel.NewRelay(ownerTunnels).Stream) + DeferCleanup(sessions.Close) + e := echo.New() + e.GET(tunnel.PeerPath, clusterapi.PeerHandler(reg, sessions.Accept)) + srv2 := httptest.NewServer(e) + DeferCleanup(srv2.Close) + Expect(reg.RegisterPeer(ctx, ownerID, "test", 0, "", strings.TrimPrefix(srv2.URL, "http://"), ownerCred.Hash())).To(Succeed()) + + front, back := workerSessions() + worker.serve(back) + _, err := ownerTunnels.Attach(ctx, "w1", tunnel.LaneInference, front) + Expect(err).ToNot(HaveOccurred()) + + callerTunnels := tunnel.NewRegistry(reg, callerID) + pool := tunnel.NewPeerPool(callerID, callerCred, reg) + DeferCleanup(pool.Close) + dialer := tunnel.NewWorkerDialer(callerTunnels, pool) + + conn, err := dialer.Dial(ctx, "w1", tunnel.StreamTagGRPC, "10.0.0.5:50051") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = conn.Close() }) + big := strings.Repeat("0123456789abcdef", 1<<14) // 256 KiB + go func() { _, _ = conn.Write([]byte(big)) }() + got := make([]byte, len(big)) + Expect(conn.SetReadDeadline(time.Now().Add(10 * time.Second))).To(Succeed()) + _, err = io.ReadFull(conn, got) + Expect(err).ToNot(HaveOccurred()) + Expect(string(got)).To(Equal(big)) + }) + + It("reports a peer with no listener as unreachable, and a peer without an address as unreachable too", func() { + Expect(reg.RegisterPeer(ctx, "silent", "test", 0, "", "127.0.0.1:1", cluster.NewPeerCredential().Hash())).To(Succeed()) + Expect(reg.Register(ctx, "addressless", "test", 0, "")).To(Succeed()) + pool := tunnel.NewPeerPool(callerID, callerCred, reg) + DeferCleanup(pool.Close) + for _, peer := range []string{"silent", "addressless"} { + _, err := pool.Open(ctx, peer) + Expect(err).To(MatchError(tunnel.ErrPeerUnreachable), peer) + Expect(errors.Is(err, cluster.ErrInstanceNotFound)).To(BeFalse(), peer) + } + _, err := pool.Open(ctx, "no-such-replica") + Expect(err).To(MatchError(cluster.ErrInstanceNotFound)) + Expect(errors.Is(err, tunnel.ErrPeerUnreachable)).To(BeFalse(), "a replica that is not registered is not an unreachable one") + }) + + It("reports a closed pool as itself", func() { + pool := tunnel.NewPeerPool(callerID, callerCred, reg) + pool.Close() + pool.Close() + _, err := pool.Open(ctx, ownerID) + Expect(err).To(MatchError(tunnel.ErrPoolClosed)) + }) + + It("does not blame the peer when the budget of the caller runs out", func() { + // A listener that accepts the TCP connection and never answers the upgrade. + ln, err := net.Listen("tcp", "127.0.0.1:0") + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = ln.Close() }) + go func() { + for { + c, err := ln.Accept() + if err != nil { + return + } + DeferCleanup(func() { _ = c.Close() }) + } + }() + Expect(reg.RegisterPeer(ctx, "slow", "test", 0, "", ln.Addr().String(), cluster.NewPeerCredential().Hash())).To(Succeed()) + pool := tunnel.NewPeerPool(callerID, callerCred, reg) + DeferCleanup(pool.Close) + short, cancel := context.WithTimeout(ctx, 200*time.Millisecond) + defer cancel() + _, err = pool.Open(short, "slow") + Expect(errors.Is(err, context.DeadlineExceeded)).To(BeTrue(), "%v", err) + Expect(errors.Is(err, tunnel.ErrPeerUnreachable)).To(BeFalse()) + }) + }) +}) + +var _ = Describe("The path of the peer route", func() { + It("is the one that the global auth middleware lets through", func() { + Expect(auth.ClusterPeerPath).To(Equal(tunnel.PeerPath)) + }) +}) + +// workerEcho is a worker that accepts any stream and echoes. +type workerEcho struct{} + +func (workerEcho) serve(sess *tunnel.Session) { + go func() { + for { + st, err := sess.AcceptStream() + if err != nil { + return + } + go func() { + defer func() { _ = st.Close() }() + if _, _, err := tunnel.ReadStreamRequest(st); err != nil { + return + } + if err := tunnel.WriteStreamAccepted(st); err != nil { + return + } + _, _ = io.Copy(st, st) + }() + } + }() +} + +// workerSessions opens a real websocket on a loopback socket and returns the +// frontend side and the worker side of an inference session on it. +func workerSessions() (front, back *tunnel.Session) { + GinkgoHelper() + up := tunnel.NewUpgrader() + got := make(chan *tunnel.Session, 1) + s := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ws, err := up.Upgrade(w, r, nil) + if err != nil { + return + } + sess, err := tunnel.ServerSession(ws, tunnel.LaneInference) + if err != nil { + _ = ws.Close() + return + } + got <- sess + })) + DeferCleanup(s.Close) + ws, _, err := tunnel.NewDialer(5*time.Second).Dial("ws"+strings.TrimPrefix(s.URL, "http"), nil) + Expect(err).ToNot(HaveOccurred()) + back, err = tunnel.ClientSession(ws, tunnel.LaneInference) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(func() { _ = back.Close() }) + Eventually(got).Should(Receive(&front)) + DeferCleanup(func() { _ = front.Close() }) + return front, back +} diff --git a/core/http/endpoints/localai/cluster_carrier.go b/core/http/endpoints/localai/cluster_carrier.go new file mode 100644 index 000000000000..bf4f847622a8 --- /dev/null +++ b/core/http/endpoints/localai/cluster_carrier.go @@ -0,0 +1,402 @@ +package localai + +import ( + "context" + "errors" + "fmt" + "net/http" + "sync/atomic" + "time" + + "github.com/labstack/echo/v4" + + "github.com/mudler/LocalAI/core/http/auth" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/xlog" +) + +// CarrierSwitchAdmin is the protocol of a change of carrier, as the admin API +// uses it. *cluster.Switch is one. +type CarrierSwitchAdmin interface { + Status(ctx context.Context) (cluster.Report, error) + Preflight(ctx context.Context, target cluster.Carrier) (cluster.Report, error) + // PreflightForStatus is Preflight for a read of the state: an old report of a + // replica is shown with its age and is not a blocker. + PreflightForStatus(ctx context.Context, target cluster.Carrier) (cluster.Report, error) + // PreflightWithin is Preflight where a report of what a replica can build + // counts only when it is younger than within. + PreflightWithin(ctx context.Context, target cluster.Carrier, within time.Duration) (cluster.Report, error) + Request(ctx context.Context, req cluster.Request) (cluster.CarrierRow, cluster.Report, error) + Abort(ctx context.Context, by string) (cluster.CarrierRow, error) +} + +// ReplicaProber asks the replicas to look at which carriers they could build. +type ReplicaProber interface { + // ProbeReplicas asks and waits for the answers, or for a time limit. + ProbeReplicas(ctx context.Context) + // AskReplicas asks and does not wait. + AskReplicas(ctx context.Context) +} + +// NATSChecker tells whether this replica can reach a NATS server with its own +// credentials. +type NATSChecker interface { + CheckNATS(ctx context.Context, url string) error +} + +// ClusterSettings is the store of the settings of the cluster. +// *cluster.SettingsStore is one. +type ClusterSettings interface { + Get(ctx context.Context, key string) (string, bool, error) + Set(ctx context.Context, key, value, by string) error + Clear(ctx context.Context, key string) error + All(ctx context.Context) (map[string]string, error) +} + +// statusAskEvery is the least time between two asks of the replicas that a read +// of the status makes. +const statusAskEvery = 15 * time.Second + +func actor(c echo.Context) string { + if u := auth.GetUser(c); u != nil { + if u.Email != "" { + return u.Email + } + return u.ID + } + return "admin" +} + +func carrierError(c echo.Context, status int, msg string, extra map[string]any) error { + body := map[string]any{"error": msg} + for k, v := range extra { + body[k] = v + } + return c.JSON(status, body) +} + +// settingsView is the settings as an admin reads them. A user name, a password or a +// query string in the address of a server is dropped. +func settingsView(all map[string]string) ClusterSettingsView { + view := ClusterSettingsView{ + PrepareTimeout: all[cluster.SettingPrepareTimeout], + TransitionWindow: all[cluster.SettingTransitionWindow], + MaxDrain: all[cluster.SettingMaxDrain], + } + if v := all[cluster.SettingNATSURL]; v != "" { + view.NATSURL = cluster.PublicNATSURL(v) + } + if v := all[cluster.SettingNATSWorkerURL]; v != "" { + view.NATSWorkerURL = cluster.PublicNATSURL(v) + } + return view +} + +// ClusterSettingsView is the settings of the cluster as an admin reads them. A +// value that is not stored is empty and the default applies. +type ClusterSettingsView struct { + // NATSURL is the NATS address that the frontends use. + NATSURL string `json:"nats_url"` + // NATSWorkerURL is the NATS address that workers are told to use. Empty means NATSURL. + NATSWorkerURL string `json:"nats_worker_url"` + // PrepareTimeout is how long a change waits for every replica to be ready (a Go duration). + PrepareTimeout string `json:"prepare_timeout"` + // TransitionWindow is how long a replica may take to confirm a commit (a Go duration). + TransitionWindow string `json:"transition_window"` + // MaxDrain is how long the previous carrier stays attached after a commit (a Go duration). + MaxDrain string `json:"max_drain"` +} + +// carrierView is the report with the settings next to it. +type CarrierStatusResponse struct { + cluster.Report + Settings ClusterSettingsView `json:"settings"` +} + +// GetCarrierEndpoint reports the state of the carrier of the cluster: the active +// carrier, the epoch, the state of a change, each replica and whether it is ready, +// the workers that could not follow a change to the other carrier and why, the +// work in flight, and the list of live replicas for the admin to confirm before a +// change. +// +// @Summary Report the transport of a distributed cluster +// @Description Returns the active carrier (nats or tunnel), the epoch, the state of a change (stable, prepare or commit), the live replicas with their readiness, the workers and whether each could follow a change to the other carrier, the work in flight, and the cluster settings. Admin only. Answers 503 when distributed mode is off. +// @Tags Nodes +// @Success 200 {object} localai.CarrierStatusResponse +// @Failure 503 {object} map[string]string "Distributed mode is not enabled" +// @Router /api/cluster/carrier [get] +func GetCarrierEndpoint(sw CarrierSwitchAdmin, settings ClusterSettings, prober ReplicaProber) echo.HandlerFunc { + var lastAsk atomic.Int64 + return func(c echo.Context) error { + ctx := c.Request().Context() + // The page polls this route every few seconds. The replicas are asked to + // look again at most once in statusAskEvery, so that what the next read + // shows is fresh, and the reads do not turn into a stream of probes. + if now, last := time.Now().UnixNano(), lastAsk.Load(); now-last >= int64(statusAskEvery) && lastAsk.CompareAndSwap(last, now) { + prober.AskReplicas(ctx) + } + status, err := sw.Status(ctx) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + // The report is for a change to the other carrier, so the workers that could + // not follow it are listed with their reasons. + other := cluster.CarrierTunnel + if status.Active == cluster.CarrierTunnel { + other = cluster.CarrierNATS + } + report, err := sw.PreflightForStatus(ctx, other) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + all, err := settings.All(ctx) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + return c.JSON(http.StatusOK, CarrierStatusResponse{Report: report, Settings: settingsView(all)}) + } +} + +// CarrierSwitchRequest is the body of POST /api/cluster/carrier. +type CarrierSwitchRequest struct { + Target string `json:"target"` + DryRun bool `json:"dry_run"` + Force bool `json:"force"` + Abort bool `json:"abort"` +} + +// SwitchCarrierEndpoint changes the carrier of the cluster. It takes +// {"target": "nats"|"tunnel", "dry_run": bool, "force": bool}, or {"abort": true} +// to abort a change that is being prepared. +// +// A dry run asks every replica to look at what it can build, runs the preflight, +// and changes nothing. A request runs the same preflight and starts the change +// when nothing blocks. The change itself is carried out by the replicas and led by +// one of them, so the answer is 202: it says what was started and the epoch to +// watch. +// +// @Summary Dry-run, start or abort a change of carrier +// @Description Takes a target (nats or tunnel) with dry_run and force, or abort. A dry run returns the preflight report and changes nothing. A request starts the change and returns 202 with the epoch to watch. A request that a blocker stops returns 422 with the blockers; force accepts the blockers that are forceable. A change already under way returns 409, and only abort is accepted then. Admin only. +// @Tags Nodes +// @Param request body localai.CarrierSwitchRequest true "target, dry_run, force or abort" +// @Success 200 {object} cluster.Report "Dry run report, or the row after an abort" +// @Success 202 {object} map[string]interface{} "Change started: state, active, target, epoch, force, warnings, row" +// @Failure 400 {object} map[string]string "Bad body or target" +// @Failure 409 {object} map[string]string "A change is under way, or there is nothing to abort" +// @Failure 422 {object} map[string]interface{} "Blocked: error, blockers, report" +// @Failure 503 {object} map[string]string "Distributed mode is not enabled" +// @Router /api/cluster/carrier [post] +func SwitchCarrierEndpoint(sw CarrierSwitchAdmin, prober ReplicaProber) echo.HandlerFunc { + return func(c echo.Context) error { + var req CarrierSwitchRequest + if err := c.Bind(&req); err != nil { + return carrierError(c, http.StatusBadRequest, "invalid request body", nil) + } + ctx := c.Request().Context() + by := actor(c) + + if req.Abort { + row, err := sw.Abort(ctx, by) + switch { + case errors.Is(err, cluster.ErrNotAbortable): + return carrierError(c, http.StatusConflict, err.Error(), nil) + case err != nil: + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + xlog.Warn("Change of carrier aborted by an admin", "by", by, "epoch", row.Epoch) + return c.JSON(http.StatusOK, row) + } + + target := cluster.Carrier(req.Target) + if target != cluster.CarrierNATS && target != cluster.CarrierTunnel { + return carrierError(c, http.StatusBadRequest, fmt.Sprintf("target must be %q or %q", cluster.CarrierNATS, cluster.CarrierTunnel), nil) + } + // The replicas write what they can build when they are asked, so the answer + // of the preflight is about now and not about the last time they looked. + // A report counts as the answer only if it was written after the question. + asked := time.Now() + prober.ProbeReplicas(ctx) + answeredWithin := time.Since(asked) + + if req.DryRun { + report, err := sw.PreflightWithin(ctx, target, answeredWithin) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + return c.JSON(http.StatusOK, report) + } + + row, report, err := sw.Request(ctx, cluster.Request{Target: target, By: by, Force: req.Force, ReportedWithin: answeredWithin}) + var blocked *cluster.BlockedError + switch { + case errors.As(err, &blocked): + return carrierError(c, http.StatusUnprocessableEntity, err.Error(), map[string]any{"blockers": blocked.Report.Blockers, "report": blocked.Report}) + case errors.Is(err, cluster.ErrBusy): + return carrierError(c, http.StatusConflict, err.Error(), nil) + case err != nil: + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + xlog.Info("Change of carrier started by an admin", "by", by, "target", target, "force", req.Force, "epoch", row.Epoch) + return c.JSON(http.StatusAccepted, map[string]any{ + "state": row.State, "active": row.Active, "target": row.Target, "epoch": row.Epoch, + "force": row.Force, "warnings": report.Warnings, "row": row, + }) + } +} + +// GetClusterSettingsEndpoint reads the settings of the cluster. +// +// @Summary Read the settings of the cluster +// @Description Returns the NATS address of the frontends and of the workers and the waits of a change. A user name, password or query string in an address is dropped. Admin only. +// @Tags Nodes +// @Success 200 {object} localai.ClusterSettingsView +// @Failure 503 {object} map[string]string "Distributed mode is not enabled" +// @Router /api/cluster/settings [get] +func GetClusterSettingsEndpoint(settings ClusterSettings) echo.HandlerFunc { + return func(c echo.Context) error { + all, err := settings.All(c.Request().Context()) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + return c.JSON(http.StatusOK, settingsView(all)) + } +} + +// ClusterSettingsRequest is the body of PUT /api/cluster/settings. A field that +// is absent is left as it is, and an empty string clears the setting. +type ClusterSettingsRequest struct { + NATSURL *string `json:"nats_url"` + NATSWorkerURL *string `json:"nats_worker_url"` + PrepareTimeout *string `json:"prepare_timeout"` + TransitionWindow *string `json:"transition_window"` + MaxDrain *string `json:"max_drain"` +} + +// natsCheckTimeout bounds the reachability check of a saved URL. +const natsCheckTimeout = 8 * time.Second + +// CarrierStatus reads the state of the cluster carrier. *cluster.Switch is one. +type CarrierStatus interface { + Status(ctx context.Context) (cluster.Report, error) +} + +// PutClusterSettingsEndpoint stores settings of the cluster. It stores nothing +// unless it can store everything: the request is checked as a whole first. +// +// A NATS URL is stored only when no change of carrier is under way and this +// replica reaches the server at that address. A change builds the NATS carrier +// from the address that the row of the change holds, so an edit during one +// would only confuse the operator; an address that does not answer would stay in +// the settings until a change to NATS failed on it. Saving an address makes NATS +// available; it does not switch the cluster to NATS. The change of carrier is a +// separate request with a dry run. Every replica checks the stored URL on its +// next report, and a dry run asks them to. +// +// No address may carry a user name, a password or a query string: the settings +// are shared, and the credentials of NATS stay on each replica. +// +// @Summary Store settings of the cluster +// @Description Stores the NATS address, the address workers use, or the waits of a change. Saving a NATS address checks that the serving replica reaches it and makes NATS available; it does not switch the cluster. Admin only. +// @Tags Nodes +// @Param request body localai.ClusterSettingsRequest true "settings to store" +// @Success 200 {object} map[string]interface{} "saved, nats.reachable when a NATS address was checked, settings" +// @Failure 400 {object} map[string]string "Unknown setting, bad value, or a credential in an address" +// @Failure 409 {object} map[string]string "A change of carrier is under way" +// @Failure 422 {object} map[string]string "This replica cannot reach the NATS server" +// @Failure 503 {object} map[string]string "Distributed mode is not enabled" +// @Router /api/cluster/settings [put] +func PutClusterSettingsEndpoint(settings ClusterSettings, checker NATSChecker, prober ReplicaProber, state CarrierStatus) echo.HandlerFunc { + return func(c echo.Context) error { + var req ClusterSettingsRequest + if err := c.Bind(&req); err != nil { + return carrierError(c, http.StatusBadRequest, "invalid request body", nil) + } + ctx := c.Request().Context() + by := actor(c) + + type change struct { + key string + val *string + } + changes := []change{ + {cluster.SettingNATSURL, req.NATSURL}, + {cluster.SettingNATSWorkerURL, req.NATSWorkerURL}, + {cluster.SettingPrepareTimeout, req.PrepareTimeout}, + {cluster.SettingTransitionWindow, req.TransitionWindow}, + {cluster.SettingMaxDrain, req.MaxDrain}, + } + // Validate everything before storing anything. The store checks again, and a + // request that it would refuse half way must not be accepted here. + for _, ch := range changes { + if ch.val == nil || *ch.val == "" { + continue + } + if err := cluster.CheckSetting(ch.key, *ch.val); err != nil { + return carrierError(c, http.StatusBadRequest, err.Error(), nil) + } + } + + resp := map[string]any{"saved": true} + if req.NATSURL != nil { + st, err := state.Status(ctx) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + if st.State != cluster.StateStable { + return carrierError(c, http.StatusConflict, fmt.Sprintf( + "the NATS URL cannot be changed while a change of carrier is under way (%s to %s): wait for it to settle or abort it", + st.State, st.Row.Target), nil) + } + if *req.NATSURL != "" { + checkCtx, cancel := context.WithTimeout(ctx, natsCheckTimeout) + err := checker.CheckNATS(checkCtx, *req.NATSURL) + cancel() + if err != nil { + xlog.Warn("The NATS URL was refused: this replica cannot reach it", "error", err) + return carrierError(c, http.StatusUnprocessableEntity, fmt.Sprintf( + "nats_url: this replica cannot reach the NATS server at %s: %v", cluster.PublicNATSURL(*req.NATSURL), err), nil) + } + resp["nats"] = map[string]any{"reachable": true} + } + } + + for _, ch := range changes { + if ch.val == nil { + continue + } + var err error + if *ch.val == "" { + err = settings.Clear(ctx, ch.key) + } else { + err = settings.Set(ctx, ch.key, *ch.val, by) + } + switch { + case errors.Is(err, cluster.ErrUnknownSetting), errors.Is(err, cluster.ErrCredentialInURL): + return carrierError(c, http.StatusBadRequest, err.Error(), nil) + case err != nil: + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + } + xlog.Info("Cluster settings changed", "by", by) + + if req.NATSURL != nil || req.NATSWorkerURL != nil { + // The other replicas look at the new address and write what they find. + prober.AskReplicas(ctx) + } + all, err := settings.All(ctx) + if err != nil { + return carrierError(c, http.StatusInternalServerError, err.Error(), nil) + } + resp["settings"] = settingsView(all) + return c.JSON(http.StatusOK, resp) + } +} + +// UnavailableCarrierEndpoint answers 503. The routes of the carrier are +// registered on every frontend, and a frontend that is not distributed has no +// cluster to change. +func UnavailableCarrierEndpoint() echo.HandlerFunc { + return func(c echo.Context) error { + return carrierError(c, http.StatusServiceUnavailable, "distributed mode is not enabled on this frontend", nil) + } +} diff --git a/core/http/endpoints/localai/cluster_carrier_test.go b/core/http/endpoints/localai/cluster_carrier_test.go new file mode 100644 index 000000000000..0d5d9a51d649 --- /dev/null +++ b/core/http/endpoints/localai/cluster_carrier_test.go @@ -0,0 +1,392 @@ +package localai + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/services/cluster" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +type fakeSwitch struct { + mu sync.Mutex + report cluster.Report + row cluster.CarrierRow + requests []cluster.Request + within []time.Duration + aborts []string + requestE error + abortE error + order *[]string +} + +func (f *fakeSwitch) note(s string) { + if f.order != nil { + *f.order = append(*f.order, s) + } +} + +func (f *fakeSwitch) Status(context.Context) (cluster.Report, error) { return f.report, nil } +func (f *fakeSwitch) Preflight(_ context.Context, target cluster.Carrier) (cluster.Report, error) { + f.note("preflight") + if target != cluster.CarrierNATS && target != cluster.CarrierTunnel { + return cluster.Report{}, cluster.ErrInvalidCarrier + } + r := f.report + r.Target = target + return r, nil +} + +func (f *fakeSwitch) PreflightForStatus(_ context.Context, target cluster.Carrier) (cluster.Report, error) { + f.note("status-preflight") + r := f.report + r.Target = target + return r, nil +} + +func (f *fakeSwitch) PreflightWithin(ctx context.Context, target cluster.Carrier, within time.Duration) (cluster.Report, error) { + f.mu.Lock() + f.within = append(f.within, within) + f.mu.Unlock() + return f.Preflight(ctx, target) +} + +func (f *fakeSwitch) Request(_ context.Context, req cluster.Request) (cluster.CarrierRow, cluster.Report, error) { + f.note("request") + f.mu.Lock() + f.requests = append(f.requests, req) + f.mu.Unlock() + return f.row, f.report, f.requestE +} + +func (f *fakeSwitch) Abort(_ context.Context, by string) (cluster.CarrierRow, error) { + f.mu.Lock() + f.aborts = append(f.aborts, by) + f.mu.Unlock() + return f.row, f.abortE +} + +type fakeProber struct { + order *[]string + delay time.Duration + asked *atomic.Int32 +} + +func (f fakeProber) ProbeReplicas(context.Context) { + time.Sleep(f.delay) + if f.order != nil { + *f.order = append(*f.order, "probe") + } +} + +func (f fakeProber) AskReplicas(context.Context) { + if f.asked != nil { + f.asked.Add(1) + } +} + +type fakeNATS struct { + err error + seen []string +} + +func (f *fakeNATS) CheckNATS(_ context.Context, url string) error { + f.seen = append(f.seen, url) + return f.err +} + +type memorySettings struct { + mu sync.Mutex + m map[string]string +} + +func (s *memorySettings) Get(_ context.Context, k string) (string, bool, error) { + s.mu.Lock() + defer s.mu.Unlock() + v, ok := s.m[k] + return v, ok, nil +} + +func (s *memorySettings) Set(_ context.Context, k, v, _ string) error { + if k == "nats.password" { + return cluster.ErrUnknownSetting + } + s.mu.Lock() + defer s.mu.Unlock() + s.m[k] = v + return nil +} + +func (s *memorySettings) Clear(_ context.Context, k string) error { + s.mu.Lock() + defer s.mu.Unlock() + delete(s.m, k) + return nil +} + +func (s *memorySettings) All(context.Context) (map[string]string, error) { + s.mu.Lock() + defer s.mu.Unlock() + out := map[string]string{} + for k, v := range s.m { + out[k] = v + } + return out, nil +} + +var _ = Describe("The admin API of the carrier", func() { + var ( + sw *fakeSwitch + settings *memorySettings + nats *fakeNATS + order []string + ) + + call := func(h echo.HandlerFunc, method, body string) (int, map[string]any) { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(method, "/", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + Expect(h(e.NewContext(req, rec))).To(Succeed()) + var out map[string]any + if rec.Body.Len() > 0 { + Expect(json.Unmarshal(rec.Body.Bytes(), &out)).To(Succeed(), rec.Body.String()) + } + return rec.Code, out + } + + BeforeEach(func() { + order = nil + sw = &fakeSwitch{ + order: &order, + report: cluster.Report{Active: cluster.CarrierNATS, Epoch: 4, State: cluster.StateStable, OK: true, Replicas: []cluster.ReplicaStatus{{ID: "a", Version: "v1"}}}, + row: cluster.CarrierRow{Active: cluster.CarrierNATS, Epoch: 5, State: cluster.StatePrepare, Target: cluster.CarrierTunnel}, + } + settings = &memorySettings{m: map[string]string{cluster.SettingNATSURL: "nats://user:secret@broker:4222?token=abc"}} + nats = &fakeNATS{} + }) + + Describe("the status", func() { + It("reports the row, the replicas, the opposite carrier as the target, and a NATS URL without its password", func() { + code, body := call(GetCarrierEndpoint(sw, settings, fakeProber{}), http.MethodGet, "") + Expect(code).To(Equal(http.StatusOK)) + Expect(body["active"]).To(Equal("nats")) + Expect(body["epoch"]).To(BeEquivalentTo(4)) + Expect(body["target"]).To(Equal("tunnel"), "the report answers for a change to the other carrier") + Expect(body["replicas"]).To(HaveLen(1)) + shown := body["settings"].(map[string]any)["nats_url"].(string) + Expect(shown).To(ContainSubstring("broker:4222")) + Expect(shown).ToNot(ContainSubstring("secret")) + Expect(shown).ToNot(ContainSubstring("user")) + Expect(shown).ToNot(ContainSubstring("abc"), "a token in the query is not shown either") + Expect(settings.m[cluster.SettingNATSURL]).To(ContainSubstring("secret"), "the stored value is not changed") + }) + }) + + Describe("the status and old reports", func() { + It("reads the status without counting an old report of a replica as a blocker", func() { + code, _ := call(GetCarrierEndpoint(sw, settings, fakeProber{}), http.MethodGet, "") + Expect(code).To(Equal(http.StatusOK)) + Expect(order).To(Equal([]string{"status-preflight"})) + }) + + It("asks the replicas to look again, but not at every read", func() { + var asked atomic.Int32 + h := GetCarrierEndpoint(sw, settings, fakeProber{asked: &asked}) + for range 5 { + code, _ := call(h, http.MethodGet, "") + Expect(code).To(Equal(http.StatusOK)) + } + Expect(asked.Load()).To(Equal(int32(1)), "the page polls every few seconds; the replicas are asked once per interval") + }) + }) + + Describe("a request", func() { + It("asks the replicas to look, runs the preflight on a dry run, and changes nothing", func() { + code, body := call(SwitchCarrierEndpoint(sw, fakeProber{order: &order}), http.MethodPost, `{"target":"tunnel","dry_run":true}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(body["target"]).To(Equal("tunnel")) + Expect(body["ok"]).To(BeTrue()) + Expect(order).To(Equal([]string{"probe", "preflight"})) + Expect(sw.requests).To(BeEmpty()) + }) + + It("counts only a report that is younger than the check it just asked for", func() { + slow := fakeProber{delay: 40 * time.Millisecond} + code, _ := call(SwitchCarrierEndpoint(sw, slow), http.MethodPost, `{"target":"tunnel","dry_run":true}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(sw.within).To(HaveLen(1)) + Expect(sw.within[0]).To(BeNumerically(">=", 40*time.Millisecond), "the age that counts starts at the question") + Expect(sw.within[0]).To(BeNumerically("<", 5*time.Second)) + + code, _ = call(SwitchCarrierEndpoint(sw, slow), http.MethodPost, `{"target":"tunnel"}`) + Expect(code).To(Equal(http.StatusAccepted)) + Expect(sw.requests[0].ReportedWithin).To(BeNumerically(">=", 40*time.Millisecond)) + }) + + It("starts the change, and passes force on", func() { + code, body := call(SwitchCarrierEndpoint(sw, fakeProber{order: &order}), http.MethodPost, `{"target":"tunnel","force":true}`) + Expect(code).To(Equal(http.StatusAccepted)) + Expect(body["state"]).To(Equal("prepare")) + Expect(order).To(Equal([]string{"probe", "request"})) + Expect(sw.requests).To(HaveLen(1)) + Expect(sw.requests[0].Target).To(Equal(cluster.CarrierTunnel)) + Expect(sw.requests[0].Force).To(BeTrue()) + Expect(sw.requests[0].By).ToNot(BeEmpty(), "who asked is recorded") + }) + + It("answers 422 with the report when the preflight blocks", func() { + sw.requestE = &cluster.BlockedError{Report: cluster.Report{Blockers: []cluster.Blocker{{Kind: cluster.BlockerWorker, ID: "w1", Reason: "the worker predates carrier switching", Forceable: true}}}} + code, body := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"target":"tunnel"}`) + Expect(code).To(Equal(http.StatusUnprocessableEntity)) + Expect(body["blockers"]).To(HaveLen(1)) + Expect(body["error"]).To(ContainSubstring("blocked")) + }) + + It("answers 409 when a change is already under way", func() { + sw.requestE = cluster.ErrBusy + code, body := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"target":"tunnel"}`) + Expect(code).To(Equal(http.StatusConflict)) + Expect(body["error"]).To(ContainSubstring("under way")) + }) + + It("answers 400 for a target that is not a carrier, and for a body that is not JSON", func() { + code, _ := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"target":"pigeon"}`) + Expect(code).To(Equal(http.StatusBadRequest)) + code, _ = call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `nope`) + Expect(code).To(Equal(http.StatusBadRequest)) + code, _ = call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"dry_run":true}`) + Expect(code).To(Equal(http.StatusBadRequest), "a request without a target") + }) + + It("aborts a change in prepare", func() { + sw.row = cluster.CarrierRow{Active: cluster.CarrierNATS, State: cluster.StateStable, Epoch: 6, Note: "aborted by admin"} + code, body := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"abort":true}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(body["state"]).To(Equal("stable")) + Expect(sw.aborts).To(HaveLen(1)) + Expect(sw.requests).To(BeEmpty()) + }) + + It("answers 409 to an abort when there is nothing to abort", func() { + sw.abortE = cluster.ErrNotAbortable + code, _ := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"abort":true}`) + Expect(code).To(Equal(http.StatusConflict)) + }) + + It("reports an error of the store as a server error", func() { + sw.requestE = errors.New("the database is down") + code, _ := call(SwitchCarrierEndpoint(sw, fakeProber{}), http.MethodPost, `{"target":"tunnel"}`) + Expect(code).To(Equal(http.StatusInternalServerError)) + }) + }) + + Describe("the settings", func() { + It("stores a NATS URL, checks that this replica reaches it, and does not switch", func() { + code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, `{"nats_url":"nats://broker2:4222"}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(settings.m[cluster.SettingNATSURL]).To(Equal("nats://broker2:4222")) + Expect(nats.seen).To(Equal([]string{"nats://broker2:4222"})) + Expect(body["nats"].(map[string]any)["reachable"]).To(BeTrue()) + Expect(sw.requests).To(BeEmpty()) + }) + + It("refuses a URL that this replica cannot reach, says why, and stores nothing", func() { + nats.err = errors.New("no answer within 3s") + code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, + `{"nats_url":"nats://broker2:4222","max_drain":"20m"}`) + Expect(code).To(Equal(http.StatusUnprocessableEntity)) + Expect(body["error"]).To(ContainSubstring("no answer")) + Expect(settings.m[cluster.SettingNATSURL]).To(Equal("nats://user:secret@broker:4222?token=abc")) + Expect(settings.m).ToNot(HaveKey(cluster.SettingMaxDrain), "a refused request stores none of its parts") + }) + + It("refuses a NATS URL while a change of carrier is under way, and stores nothing", func() { + for _, st := range []cluster.State{cluster.StatePrepare, cluster.StateCommit} { + sw.report.State = st + code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, + `{"nats_url":"nats://broker2:4222"}`) + Expect(code).To(Equal(http.StatusConflict), string(st)) + Expect(body["error"]).To(ContainSubstring("change of carrier")) + Expect(settings.m[cluster.SettingNATSURL]).To(ContainSubstring("broker:4222")) + Expect(nats.seen).To(BeEmpty(), "the server is not even asked") + } + }) + + It("refuses an address that carries a credential or a query string", func() { + for _, bad := range []string{ + `{"nats_url":"nats://user:pass@broker2:4222"}`, + `{"nats_url":"nats://broker2:4222?token=abc"}`, + `{"nats_worker_url":"nats://user@broker2:4222"}`, + } { + code, body := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, bad) + Expect(code).To(Equal(http.StatusBadRequest), bad) + Expect(body["error"]).To(ContainSubstring("credentials stay on each replica"), bad) + } + }) + + It("refuses a URL that is not a NATS address", func() { + for _, bad := range []string{`{"nats_url":"http://broker:4222"}`, `{"nats_url":"nats://"}`, `{"nats_url":"::"}`} { + code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, bad) + Expect(code).To(Equal(http.StatusBadRequest), bad) + } + Expect(settings.m[cluster.SettingNATSURL]).To(ContainSubstring("broker:4222")) + }) + + It("clears a URL with an empty string", func() { + code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, `{"nats_url":""}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(settings.m).ToNot(HaveKey(cluster.SettingNATSURL)) + }) + + It("stores the waits of a change, and refuses one that is not a positive duration", func() { + code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, `{"max_drain":"20m","prepare_timeout":"90s"}`) + Expect(code).To(Equal(http.StatusOK)) + Expect(settings.m[cluster.SettingMaxDrain]).To(Equal("20m")) + Expect(settings.m[cluster.SettingPrepareTimeout]).To(Equal("90s")) + code, _ = call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, `{"max_drain":"soon"}`) + Expect(code).To(Equal(http.StatusBadRequest)) + }) + + It("refuses a wait that is not positive or is too short, with a 400, and stores none of the request", func() { + for _, bad := range []string{"0s", "-5s", "1ns", "1ms", "4s"} { + body := `{"nats_url":"nats://broker2:4222","max_drain":"20m","prepare_timeout":"` + bad + `"}` + code, resp := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, body) + Expect(code).To(Equal(http.StatusBadRequest), bad) + Expect(resp["error"]).To(ContainSubstring("prepare_timeout"), bad) + Expect(settings.m).ToNot(HaveKey(cluster.SettingMaxDrain), "%s: nothing is stored when one part is refused", bad) + Expect(settings.m[cluster.SettingNATSURL]).To(ContainSubstring("broker:4222"), bad) + } + code, _ := call(PutClusterSettingsEndpoint(settings, nats, fakeProber{}, sw), http.MethodPut, `{"prepare_timeout":"5s"}`) + Expect(code).To(Equal(http.StatusOK)) + }) + + It("reads the settings with the password of the URL masked", func() { + code, body := call(GetClusterSettingsEndpoint(settings), http.MethodGet, "") + Expect(code).To(Equal(http.StatusOK)) + Expect(body["nats_url"]).To(ContainSubstring("broker:4222")) + Expect(body["nats_url"]).ToNot(ContainSubstring("secret")) + Expect(body["nats_url"]).ToNot(ContainSubstring("token")) + }) + }) + + Describe("a frontend without distributed mode", func() { + It("answers 503 on every route", func() { + for _, h := range []echo.HandlerFunc{UnavailableCarrierEndpoint()} { + code, _ := call(h, http.MethodGet, "") + Expect(code).To(Equal(http.StatusServiceUnavailable)) + } + }) + }) +}) diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index b38cb5a6033f..ce179b68ee84 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -2,6 +2,7 @@ package localai import ( "context" + "crypto/rand" "crypto/sha256" "crypto/subtle" "encoding/hex" @@ -24,6 +25,7 @@ import ( "github.com/mudler/LocalAI/core/gallery" "github.com/mudler/LocalAI/core/http/auth" "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/core/services/cluster" "github.com/mudler/LocalAI/core/services/galleryop" "github.com/mudler/LocalAI/core/services/nodes" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" @@ -114,12 +116,152 @@ type RegisterNodeRequest struct { Version string `json:"version,omitempty"` // Commit is the git commit hash the worker binary was built from. Commit string `json:"commit,omitempty"` + // Routable says if the address of the worker can be reached from the + // frontends. A worker started with an address sends true. A worker that has + // none sends false: it holds an outbound tunnel and nothing dials it. A + // worker that predates the field sends nothing, and that means true. + Routable *bool `json:"routable,omitempty"` + // Carrier is the carrier the worker asks the credential of. A worker that + // follows a change of carrier registers again to get the credential of the + // carrier it attaches to, which is not the active one yet while the change is + // being prepared. Only the active carrier, the target of a change under way + // and the carrier that drains can be asked for. Anything else, and a worker + // that sends nothing (it predates the field, or it boots), gets the credential + // of the active carrier. + Carrier string `json:"carrier,omitempty"` +} + +// CarrierReader reads the row that names the carrier of the cluster. +// *cluster.CarrierStore is one. +type CarrierReader interface { + Get(ctx context.Context) (cluster.CarrierRow, error) +} + +type registerOptions struct { + carriers CarrierReader + tunnels TunnelDisconnector + nats NATSHandover +} + +// NATSHandover says what a worker needs to reach NATS besides its own credential. +// The frontend hands it over in the answer to a registration, so that a worker +// that was never given the address of NATS, or that booted on the tunnel, can +// attach to NATS when the cluster changes to it. +type NATSHandover interface { + // WorkerURL is the address that workers use. It is empty when no NATS is + // configured. + WorkerURL(ctx context.Context) (string, error) + // CAPEM is the CA of the NATS server as a PEM bundle, or empty when the + // server does not use a private CA. + CAPEM() (string, error) + // ClientTLS says that the server asks for a client certificate. Files of that + // kind are not handed over, so a worker needs its own. + ClientTLS() bool +} + +// WithNATSHandover makes a registration that hands over a NATS credential also +// hand over the address of the server, its CA, and whether the worker needs a +// client certificate of its own. +func WithNATSHandover(h NATSHandover) RegisterOption { + return func(o *registerOptions) { o.nats = h } +} + +// TunnelDisconnector ends the tunnel that this replica holds for a node. +// *tunnel.Registry is one. +type TunnelDisconnector interface { + Disconnect(nodeID string) bool +} + +// WithTunnelDisconnector makes a registration that replaces the tunnel +// credential of a node, and the removal of a node, end the tunnel session that +// the replica holds for it. The credential is checked when a worker dials, so +// without this a session that was open stays open for a node that no longer has +// the right to it. +// +// It acts on this replica only. A session held by another replica ends when its +// worker dials again and is refused, or when the session ends for another +// reason. +func WithTunnelDisconnector(d TunnelDisconnector) RegisterOption { + return func(o *registerOptions) { o.tunnels = d } +} + +// RegisterOption changes how a registration or an approval answers. +type RegisterOption func(*registerOptions) + +// WithCarrierReader makes the answer name the active carrier and its epoch, and +// hand over the credential for that carrier. Without it the answer is the one of +// a deployment that has only NATS. +func WithCarrierReader(r CarrierReader) RegisterOption { + return func(o *registerOptions) { o.carriers = r } +} + +func newRegisterOptions(opts []RegisterOption) registerOptions { + var o registerOptions + for _, opt := range opts { + opt(&o) + } + return o +} + +// activeCarrier returns the row of the carrier. The second result is false when +// there is no reader or the read failed. The caller then answers as a +// deployment with only NATS does. A failed read is logged and does not fail the +// registration: registration is how a worker joins the cluster, and a worker +// that is refused for a problem it cannot fix would not join at all. +func (o registerOptions) activeCarrier(ctx context.Context) (cluster.CarrierRow, bool) { + if o.carriers == nil { + return cluster.CarrierRow{}, false + } + row, err := o.carriers.Get(ctx) + if err != nil { + xlog.Warn("Reading the cluster carrier for a registration failed; answering as for NATS", "error", err) + return cluster.CarrierRow{}, false + } + return row, true +} + +// credentialCarrier is the carrier whose credential an answer carries: the one +// the worker asked for when it may have it, else the active one. +func credentialCarrier(row cluster.CarrierRow, asked string) cluster.Carrier { + c := cluster.Carrier(asked) + switch { + case c == "" || c == row.Active: + return row.Active + case row.State != cluster.StateStable && c == row.Target: + return c + case c == row.Draining: + return c + } + return row.Active +} + +// carrierNews is what a worker is told about the carrier of the cluster, in the +// answer to a registration and to a heartbeat. The heartbeat is the channel that +// exists whichever carrier is active, so it is the one that carries the news. +func carrierNews(row cluster.CarrierRow) map[string]any { + view := map[string]any{ + "carrier": row.Active, + "carrier_epoch": row.Epoch, + "carrier_state": row.State, + } + if row.Target != "" { + view["carrier_target"] = row.Target + } + if row.Draining != "" { + view["carrier_draining"] = row.Draining + if row.DrainingUntil != nil { + view["carrier_draining_until"] = row.DrainingUntil.UTC().Format(time.RFC3339) + } + } + return view } // RegisterNodeEndpoint registers a new backend node. // expectedToken is the registration token configured on the frontend (may be empty to disable auth). // autoApprove controls whether new nodes go directly to "healthy" or require admin approval. -func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, autoApprove bool, authDB *gorm.DB, hmacSecret string, natsCfg natsauth.Config) echo.HandlerFunc { +func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, autoApprove bool, authDB *gorm.DB, hmacSecret string, natsCfg natsauth.Config, options ...RegisterOption) echo.HandlerFunc { + opts := newRegisterOptions(options) + var openRegistrationWarning sync.Once return func(c echo.Context) error { var req RegisterNodeRequest if err := c.Bind(&req); err != nil { @@ -148,11 +290,30 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au fmt.Sprintf("invalid node_type %q; must be %q or %q", nodeType, nodes.NodeTypeBackend, nodes.NodeTypeAgent))) } - // Backend workers require address; agent workers don't serve gRPC + ctx := c.Request().Context() + carrierRow, haveCarrier := opts.activeCarrier(ctx) + // The carrier whose credential this answer carries. It decides whether + // the address of the worker matters: a worker that asks for NATS while + // the tunnel is still active is about to be dialled. + credFor := cluster.CarrierNATS + if haveCarrier { + credFor = credentialCarrier(carrierRow, req.Carrier) + } + tunnelActive := haveCarrier && credFor == cluster.CarrierTunnel + // A worker that said that its address cannot be reached holds an + // outbound tunnel and nothing dials it. Its address is dropped and not + // stored: a dialable-looking endpoint that nothing may dial misleads. + tunnelOnly := tunnelActive && req.Routable != nil && !*req.Routable + if tunnelOnly { + req.Address, req.HTTPAddress = "", "" + } + + // Backend workers require address on NATS, which dials it. Agent workers + // don't serve gRPC. With the tunnel active the address is optional. if req.Name == "" { return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "name is required")) } - if nodeType == nodes.NodeTypeBackend && req.Address == "" { + if nodeType == nodes.NodeTypeBackend && req.Address == "" && !tunnelActive { return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "address is required for backend workers")) } if len(req.Name) > 255 { @@ -204,11 +365,17 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au Commit: req.Commit, } - ctx := c.Request().Context() if err := registry.Register(ctx, node, autoApprove); err != nil { xlog.Error("Failed to register node", "name", req.Name, "error", err) return c.JSON(http.StatusInternalServerError, nodeError(http.StatusInternalServerError, "failed to register node")) } + if tunnelOnly { + // Register keeps a column when the new value is empty, so a node that + // had addresses from an earlier registration still has them. + if err := registry.ClearAddresses(ctx, node.ID); err != nil { + xlog.Warn("Failed to clear the addresses of a tunnel-only node", "node", node.Name, "error", err) + } + } // Merge worker-supplied labels into the node's existing label set, // then apply auto-labels. Using SetNodeLabels here would have @@ -257,7 +424,14 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au } } - attachNatsJWT(response, node, natsCfg) + attachCarrierCredential(ctx, response, registry, node, natsCfg, carrierRow, haveCarrier, credFor, opts) + if tunnelActive && expectedToken == "" && autoApprove { + openRegistrationWarning.Do(func() { + xlog.Warn("Worker registration is open: there is no registration token and new nodes are approved at once, " + + "so anyone who can reach this frontend can register a node, get a tunnel to it, and receive the requests that are sent to workers. " + + "Set a registration token or turn off automatic approval.") + }) + } return c.JSON(http.StatusCreated, response) } @@ -265,7 +439,8 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au // ApproveNodeEndpoint approves a pending node, setting its status to healthy. // For agent workers, it also provisions an API key so they can call the inference API. -func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecret string, natsCfg natsauth.Config) echo.HandlerFunc { +func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecret string, natsCfg natsauth.Config, options ...RegisterOption) echo.HandlerFunc { + opts := newRegisterOptions(options) return func(c echo.Context) error { ctx := c.Request().Context() id := c.Param("id") @@ -295,12 +470,131 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr } } - attachNatsJWT(response, node, natsCfg) + // A node that waited for approval already holds its tunnel credential, + // which was minted when it registered and is inert until now. Only the + // NATS credential is minted at approval, and only when NATS is active. + if row, haveCarrier := opts.activeCarrier(ctx); !haveCarrier || row.Active != cluster.CarrierTunnel { + attachNatsJWT(response, node, natsCfg) + } return c.JSON(http.StatusOK, response) } } +// attachCarrierCredential names the active carrier in the answer and adds the +// credential for credFor: the tunnel token when that is the tunnel, the NATS user +// JWT and the address of the server otherwise. A deployment with no carrier row +// is a deployment with NATS and gets the answer that it always had. +// +// A credential is minted for one carrier only. A tunnel token replaces the +// previous one and ends the session that the replica holds for the node, so a +// worker that holds the tunnel and registers for NATS must not mint one. +func attachCarrierCredential(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode, natsCfg natsauth.Config, row cluster.CarrierRow, haveCarrier bool, credFor cluster.Carrier, opts registerOptions) { + if haveCarrier { + for k, v := range carrierNews(row) { + response[k] = v + } + response["credential_for"] = credFor + } + if haveCarrier && credFor == cluster.CarrierTunnel { + attachTunnelToken(ctx, response, registry, node, opts.tunnels) + return + } + attachNatsJWT(response, node, natsCfg) + attachNATSHandover(ctx, response, opts.nats) +} + +// attachNATSHandover adds the address of the NATS server and what is needed to +// verify it. A failure is logged and the answer goes out without it: the worker +// then uses a URL of its own, or says that it cannot attach. +func attachNATSHandover(ctx context.Context, response map[string]any, h NATSHandover) { + if h == nil { + return + } + url, err := h.WorkerURL(ctx) + if err != nil { + xlog.Warn("Reading the NATS address for a registration failed", "error", err) + return + } + if url == "" { + return + } + response["nats_url"] = url + ca, err := h.CAPEM() + if err != nil { + xlog.Warn("Reading the NATS CA for a registration failed", "error", err) + } else if ca != "" { + response["nats_ca"] = ca + } + if h.ClientTLS() { + response["nats_client_tls"] = true + } +} + +// attachTunnelToken mints a new tunnel credential for the node, stores only its +// hash and puts the secret in the response. +// +// Every registration mints a new one, so a second process that registers under +// the same node name makes the next dial of the first process fail. The worker +// reads the token each time it dials. +// +// A node that waits for approval gets one too. That differs from the API key of +// an agent worker, which is not given before approval because it works at once. +// The tunnel credential does not: the connect endpoint reads the status of the +// node at every dial and refuses a pending node with 403. A worker that +// registers once, and gets no credential until it is approved, could not +// connect when the approval came, because an approval does not make it register +// again. +// +// Only a backend or agent node holds a credential. For any other type the +// column is cleared and not left alone. The connect endpoint does not look at +// the type of the node. It refuses a node with an empty hash, so a cleared +// column is what keeps such a node out. Register skips an empty value, so a +// node that registered again as another type would keep a credential that its +// type is not entitled to. +// +// A failure to mint or to store is logged and the answer goes out without the +// token. Registration is how a worker joins the cluster, and a tunnel problem +// must not become a node that cannot join. The worker sees no tunnel_token, +// says so, and tries at its next registration. +func attachTunnelToken(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode, tunnels TunnelDisconnector) { + if node == nil { + return + } + if !tunnelEligible(node.NodeType) { + if err := registry.SetTunnelTokenHash(ctx, node.ID, ""); err != nil { + xlog.Error("Failed to clear the tunnel credential of a node whose type holds none", + "node", node.Name, "type", node.NodeType, "error", err) + return + } + if tunnels != nil { + tunnels.Disconnect(node.ID) + } + return + } + // crypto/rand.Text has at least 128 bits of randomness and no error. + secret := rand.Text() + sum := sha256.Sum256([]byte(secret)) + if err := registry.SetTunnelTokenHash(ctx, node.ID, hex.EncodeToString(sum[:])); err != nil { + xlog.Error("Failed to store a tunnel credential for node", "node", node.Name, "error", err) + return + } + response["tunnel_token"] = secret + // The old credential is gone from the table. A session that was opened with + // it is ended, so that the old holder has to present the new one. + if tunnels != nil && tunnels.Disconnect(node.ID) { + xlog.Info("Ended the tunnel session of a node whose credential was replaced", "node", node.Name) + } +} + +// tunnelEligible reports whether a node of this type holds a tunnel credential. +// The minting and the clearing both use it, so they cannot disagree: if only one +// were widened, a type would get a credential and never lose it, or lose it at +// every registration. +func tunnelEligible(nodeType string) bool { + return nodeType == nodes.NodeTypeBackend || nodeType == nodes.NodeTypeAgent +} + // attachNatsJWT adds a per-node NATS user JWT to a register/approve response when minting is enabled. func attachNatsJWT(response map[string]any, node *nodes.BackendNode, natsCfg natsauth.Config) { if !natsCfg.CanMintWorkers() || node == nil || node.Status == nodes.StatusPending { @@ -378,7 +672,8 @@ func cancelNodeLoads(ctx context.Context, registry *nodes.NodeRegistry, unloader } // DeregisterNodeEndpoint removes a backend node permanently (admin use). -func DeregisterNodeEndpoint(registry *nodes.NodeRegistry, unloader nodes.NodeCommandSender) echo.HandlerFunc { +func DeregisterNodeEndpoint(registry *nodes.NodeRegistry, unloader nodes.NodeCommandSender, options ...RegisterOption) echo.HandlerFunc { + opts := newRegisterOptions(options) return func(c echo.Context) error { ctx := c.Request().Context() id := c.Param("id") @@ -387,6 +682,11 @@ func DeregisterNodeEndpoint(registry *nodes.NodeRegistry, unloader nodes.NodeCom xlog.Error("Failed to deregister node", "id", id, "error", err) return c.JSON(http.StatusInternalServerError, nodeError(http.StatusInternalServerError, "failed to deregister node")) } + // The row is gone, and so is the right to a tunnel. Without this the + // session stays up, because the credential is checked at the dial only. + if opts.tunnels != nil { + opts.tunnels.Disconnect(id) + } return c.JSON(http.StatusOK, map[string]string{"message": "node deregistered"}) } } @@ -406,7 +706,14 @@ func DeactivateNodeEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { } // HeartbeatEndpoint updates the heartbeat for a node. -func HeartbeatEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { +// +// The answer tells the worker which carrier the cluster uses and whether a change +// is under way. A worker follows the carrier from this answer, because the +// heartbeat reaches the frontend whichever carrier is active. A worker that +// reports the carriers it is attached to and the ones it can follow has them +// stored, and the frontends route by them during a change. +func HeartbeatEndpoint(registry *nodes.NodeRegistry, options ...RegisterOption) echo.HandlerFunc { + opts := newRegisterOptions(options) return func(c echo.Context) error { id := c.Param("id") @@ -431,7 +738,21 @@ func HeartbeatEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { xlog.Warn("Heartbeat failed for node", "id", id, "error", err) return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) } - return c.JSON(http.StatusOK, map[string]string{"message": "heartbeat received"}) + if report, ok := update.CarrierReport(); ok { + if err := registry.SetCarrierReport(ctx, id, report); err != nil { + xlog.Warn("Failed to store the carrier report of a node", "id", id, "error", err) + } + } + response := map[string]any{"message": "heartbeat received"} + if row, ok := opts.activeCarrier(ctx); ok { + for k, v := range carrierNews(row) { + response[k] = v + } + if opts.nats != nil && opts.nats.ClientTLS() { + response["nats_client_tls"] = true + } + } + return c.JSON(http.StatusOK, response) } } @@ -786,11 +1107,12 @@ func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) } - if node.HTTPAddress == "" { + host, ok := workerLogsHost(node) + if !ok { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, "node has no HTTP address")) } - resp, err := proxyHTTPToWorker(ctx, dialFor, node.ID, node.HTTPAddress, "/v1/backend-logs", registrationToken) + resp, err := proxyHTTPToWorker(ctx, dialFor, node.ID, host, "/v1/backend-logs", registrationToken) if err != nil { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err))) } @@ -816,12 +1138,13 @@ func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToke return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found")) } - if node.HTTPAddress == "" { + host, ok := workerLogsHost(node) + if !ok { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, "node has no HTTP address")) } path := "/v1/backend-logs/" + url.PathEscape(modelID) - resp, err := proxyHTTPToWorker(ctx, dialFor, node.ID, node.HTTPAddress, path, registrationToken) + resp, err := proxyHTTPToWorker(ctx, dialFor, node.ID, host, path, registrationToken) if err != nil { return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err))) } @@ -869,7 +1192,7 @@ func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken s } // Dial the worker WebSocket - workerURL := fmt.Sprintf("ws://%s/v1/backend-logs/%s/ws", node.HTTPAddress, url.PathEscape(modelID)) + workerURL := fmt.Sprintf("ws://%s/v1/backend-logs/%s/ws", nodes.WorkerHTTPHost(node.ID, node.HTTPAddress), url.PathEscape(modelID)) workerHeaders := http.Header{} if registrationToken != "" { workerHeaders.Set("Authorization", "Bearer "+registrationToken) @@ -1345,6 +1668,35 @@ func DeleteSchedulingEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc { } } +// workerLogsHost returns the host of the URL of a request to the HTTP server of a +// worker, and false when the node has no HTTP server to ask. +// +// A worker that holds a tunnel registers no address at all: nothing dials it, and +// the dialer opens a stream on its tunnel whatever host the URL names. Its host +// is the reserved name of nodes.WorkerHTTPHost. A worker that registered a +// gRPC address and no HTTP address is an older worker with no log endpoint, and +// the answer for it has not changed. +func workerLogsHost(node *nodes.BackendNode) (string, bool) { + if node.HTTPAddress == "" && node.Address != "" { + return "", false + } + return nodes.WorkerHTTPHost(node.ID, node.HTTPAddress), true +} + +// workerTransport returns the transport that reaches the HTTP server of a worker +// through the dialer of the node. +func workerTransport(dialFor nodes.WorkerNetDialerFor, nodeID, host string) *http.Transport { + t := httpclient.HardenedTransport() + t.DialContext = dialFor(nodeID) + // A worker without an address is reached through a tunnel, and a proxy of the + // environment cannot carry that: the transport would dial the proxy through + // the dialer of the worker. The proxy stays for every other worker. + if nodes.IsTunnelOnlyHost(host) { + t.Proxy = nil + } + return t +} + // proxyHTTPToWorker makes a GET request to a worker's HTTP server with bearer token auth. // The connection goes through dialFor(nodeID) because the advertised address // alone does not say how this frontend reaches that worker. @@ -1361,8 +1713,7 @@ func proxyHTTPToWorker(ctx context.Context, dialFor nodes.WorkerNetDialerFor, no req.Header.Set("Authorization", "Bearer "+token) } - t := httpclient.HardenedTransport() - t.DialContext = dialFor(nodeID) + t := workerTransport(dialFor, nodeID, httpAddress) client := httpclient.NewWithTimeout(15*time.Second, httpclient.WithTransport(t)) return client.Do(req) } diff --git a/core/http/endpoints/localai/nodes_carrier_test.go b/core/http/endpoints/localai/nodes_carrier_test.go new file mode 100644 index 000000000000..e15d76feb33f --- /dev/null +++ b/core/http/endpoints/localai/nodes_carrier_test.go @@ -0,0 +1,348 @@ +package localai + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "slices" + "strings" + "sync" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/pkg/natsauth" + "github.com/nats-io/nkeys" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// fixedCarrier answers with a row that the spec chose. +type fixedCarrier struct { + row cluster.CarrierRow + err error +} + +func (f fixedCarrier) Get(context.Context) (cluster.CarrierRow, error) { return f.row, f.err } + +var _ = Describe("Registration and the active carrier", func() { + var ( + ctx context.Context + registry *nodes.NodeRegistry + natsCfg natsauth.Config + ) + + tunnelRow := cluster.CarrierRow{Active: cluster.CarrierTunnel, Epoch: 7, State: cluster.StateStable} + natsRow := cluster.CarrierRow{Active: cluster.CarrierNATS, Epoch: 3, State: cluster.StateStable} + + BeforeEach(func() { + ctx = context.Background() + var err error + registry, err = nodes.NewNodeRegistry(testutil.SetupTestDB()) + Expect(err).ToNot(HaveOccurred()) + + akp, err := nkeys.CreateAccount() + Expect(err).ToNot(HaveOccurred()) + seed, err := akp.Seed() + Expect(err).ToNot(HaveOccurred()) + natsCfg = natsauth.Config{AccountSeed: string(seed)} + }) + + register := func(body string, autoApprove bool, opts ...RegisterOption) (int, map[string]any) { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + Expect(RegisterNodeEndpoint(registry, "", autoApprove, nil, "", natsCfg, opts...)(e.NewContext(req, rec))).To(Succeed()) + var resp map[string]any + if rec.Code == http.StatusCreated { + Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + } + return rec.Code, resp + } + + storedHash := func(name string) string { + GinkgoHelper() + node, err := registry.GetByName(ctx, name) + Expect(err).ToNot(HaveOccurred()) + return node.TunnelTokenHash + } + + Describe("a deployment with no carrier reader", func() { + It("answers as it always did: no carrier, no tunnel token", func() { + code, resp := register(`{"name":"w","address":"10.0.0.1:50051"}`, true) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp).ToNot(HaveKey("carrier")) + Expect(resp).ToNot(HaveKey("carrier_epoch")) + Expect(resp).ToNot(HaveKey("tunnel_token")) + Expect(resp).To(HaveKey("nats_jwt")) + Expect(storedHash("w")).To(BeEmpty()) + }) + }) + + Describe("with NATS active", func() { + reader := WithCarrierReader(fixedCarrier{row: natsRow}) + + It("names the carrier and its epoch, and hands over the NATS credential only", func() { + code, resp := register(`{"name":"w","address":"10.0.0.1:50051"}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp["carrier"]).To(Equal("nats")) + Expect(resp["carrier_epoch"]).To(BeEquivalentTo(3)) + Expect(resp).To(HaveKey("nats_jwt")) + Expect(resp).ToNot(HaveKey("tunnel_token")) + Expect(storedHash("w")).To(BeEmpty()) + }) + + It("still requires the address of a backend worker", func() { + code, _ := register(`{"name":"w"}`, true, reader) + Expect(code).To(Equal(http.StatusBadRequest)) + }) + + It("keeps the address of a worker that sends routable false, because NATS dials it", func() { + code, _ := register(`{"name":"w","address":"10.0.0.1:50051","routable":false}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Address).To(Equal("10.0.0.1:50051")) + }) + }) + + Describe("with the tunnel active", func() { + reader := WithCarrierReader(fixedCarrier{row: tunnelRow}) + + It("names the carrier and hands over a tunnel token, and no NATS credential", func() { + code, resp := register(`{"name":"w"}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp["carrier"]).To(Equal("tunnel")) + Expect(resp["carrier_epoch"]).To(BeEquivalentTo(7)) + Expect(resp).ToNot(HaveKey("nats_jwt")) + Expect(resp).ToNot(HaveKey("nats_user_seed")) + Expect(resp["tunnel_token"]).To(BeAssignableToTypeOf("")) + }) + + It("stores only the SHA-256 of the token", func() { + _, resp := register(`{"name":"w"}`, true, reader) + token := resp["tunnel_token"].(string) + Expect(len(token)).To(BeNumerically(">=", 26), "128 bits of randomness need at least 26 base32 characters") + + sum := sha256.Sum256([]byte(token)) + Expect(storedHash("w")).To(Equal(hex.EncodeToString(sum[:]))) + Expect(storedHash("w")).ToNot(ContainSubstring(token)) + }) + + It("mints a new token at every registration and the old token stops matching", func() { + _, first := register(`{"name":"w"}`, true, reader) + firstHash := storedHash("w") + _, second := register(`{"name":"w"}`, true, reader) + + Expect(second["tunnel_token"]).ToNot(Equal(first["tunnel_token"])) + Expect(storedHash("w")).ToNot(Equal(firstHash)) + Expect(second["id"]).To(Equal(first["id"]), "the node keeps its identity") + }) + + It("gives a token to a node that waits for approval, and the token is inert until then", func() { + code, resp := register(`{"name":"w"}`, false, reader) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp["status"]).To(Equal(nodes.StatusPending)) + Expect(resp["tunnel_token"]).ToNot(BeEmpty()) + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Status).To(Equal(nodes.StatusPending), "the connect endpoint refuses a pending node") + }) + + It("gives a token to an agent worker too", func() { + code, resp := register(`{"name":"agent","node_type":"agent"}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp["tunnel_token"]).ToNot(BeEmpty()) + }) + + It("does not need an address", func() { + code, _ := register(`{"name":"w"}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + }) + + It("keeps the address of a worker that has a routable one", func() { + code, _ := register(`{"name":"w","address":"10.0.0.1:50051","http_address":"10.0.0.1:50050","routable":true}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Address).To(Equal("10.0.0.1:50051")) + Expect(node.HTTPAddress).To(Equal("10.0.0.1:50050")) + }) + + It("keeps the address of a worker that predates the routable field", func() { + code, _ := register(`{"name":"w","address":"10.0.0.1:50051"}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Address).To(Equal("10.0.0.1:50051")) + }) + + It("drops the addresses of a worker that says nothing can reach it, also from an earlier registration", func() { + code, _ := register(`{"name":"w","address":"10.0.0.1:50051","http_address":"10.0.0.1:50050"}`, true, WithCarrierReader(fixedCarrier{row: natsRow})) + Expect(code).To(Equal(http.StatusCreated)) + + code, _ = register(`{"name":"w","address":"host:50051","http_address":"host:50050","routable":false}`, true, reader) + Expect(code).To(Equal(http.StatusCreated)) + + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Address).To(BeEmpty()) + Expect(node.HTTPAddress).To(BeEmpty()) + }) + }) + + It("answers as for NATS when the carrier row cannot be read", func() { + code, resp := register(`{"name":"w","address":"10.0.0.1:50051"}`, true, + WithCarrierReader(fixedCarrier{err: errors.New("database is down")})) + Expect(code).To(Equal(http.StatusCreated)) + Expect(resp).ToNot(HaveKey("carrier")) + Expect(resp).To(HaveKey("nats_jwt")) + Expect(resp).ToNot(HaveKey("tunnel_token")) + }) + + Describe("approval", func() { + approve := func(id string, opts ...RegisterOption) map[string]any { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetParamNames("id") + c.SetParamValues(id) + Expect(ApproveNodeEndpoint(registry, nil, "", natsCfg, opts...)(c)).To(Succeed()) + Expect(rec.Code).To(Equal(http.StatusOK)) + var resp map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + return resp + } + + It("mints the NATS credential when NATS is active", func() { + _, reg := register(`{"name":"w","address":"10.0.0.1:50051"}`, false) + resp := approve(reg["id"].(string), WithCarrierReader(fixedCarrier{row: natsRow})) + Expect(resp).To(HaveKey("nats_jwt")) + }) + + It("does not mint a NATS credential when the tunnel is active, and leaves the tunnel token alone", func() { + reader := WithCarrierReader(fixedCarrier{row: tunnelRow}) + _, reg := register(`{"name":"w"}`, false, reader) + before := storedHash("w") + + resp := approve(reg["id"].(string), reader) + + Expect(resp).ToNot(HaveKey("nats_jwt")) + Expect(resp).ToNot(HaveKey("tunnel_token"), "the worker holds the token from its registration") + Expect(storedHash("w")).To(Equal(before)) + }) + }) + + Describe("a node type that holds no tunnel credential", func() { + It("has the credential cleared and gets no token", func() { + _, reg := register(`{"name":"w"}`, true, WithCarrierReader(fixedCarrier{row: tunnelRow})) + id := reg["id"].(string) + Expect(storedHash("w")).ToNot(BeEmpty()) + + response := map[string]any{} + attachTunnelToken(ctx, response, registry, &nodes.BackendNode{ID: id, Name: "w", NodeType: "gpu-farm"}, nil) + + Expect(response).ToNot(HaveKey("tunnel_token")) + Expect(storedHash("w")).To(BeEmpty()) + }) + }) +}) + +// recordingDisconnector records the nodes whose tunnel it was asked to end. +type recordingDisconnector struct { + mu sync.Mutex + nodes []string +} + +func (r *recordingDisconnector) Disconnect(nodeID string) bool { + r.mu.Lock() + defer r.mu.Unlock() + r.nodes = append(r.nodes, nodeID) + return true +} + +func (r *recordingDisconnector) seen() []string { + r.mu.Lock() + defer r.mu.Unlock() + return slices.Clone(r.nodes) +} + +var _ = Describe("Ending the tunnel session of a node that lost its right to it", func() { + var ( + ctx context.Context + registry *nodes.NodeRegistry + tunnels *recordingDisconnector + natsCfg natsauth.Config + ) + tunnelRow := cluster.CarrierRow{Active: cluster.CarrierTunnel, Epoch: 7, State: cluster.StateStable} + + BeforeEach(func() { + ctx = context.Background() + var err error + registry, err = nodes.NewNodeRegistry(testutil.SetupTestDB()) + Expect(err).ToNot(HaveOccurred()) + tunnels = &recordingDisconnector{} + }) + + register := func() string { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(`{"name":"w"}`)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + h := RegisterNodeEndpoint(registry, "", true, nil, "", natsCfg, + WithCarrierReader(fixedCarrier{row: tunnelRow}), WithTunnelDisconnector(tunnels)) + Expect(h(e.NewContext(req, rec))).To(Succeed()) + var resp map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + return resp["id"].(string) + } + + It("ends the session when a registration replaces the credential", func() { + id := register() + Expect(tunnels.seen()).To(Equal([]string{id})) + // A second registration replaces it again. + register() + Expect(tunnels.seen()).To(HaveLen(2)) + }) + + It("ends the session when an admin removes the node", func() { + id := register() + before := len(tunnels.seen()) + + e := echo.New() + req := httptest.NewRequest(http.MethodDelete, "/", nil) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetParamNames("id") + c.SetParamValues(id) + Expect(DeregisterNodeEndpoint(registry, nil, WithTunnelDisconnector(tunnels))(c)).To(Succeed()) + + Expect(rec.Code).To(Equal(http.StatusOK)) + Expect(tunnels.seen()).To(HaveLen(before + 1)) + _, err := registry.Get(ctx, id) + Expect(err).To(HaveOccurred()) + }) + + It("does not end a session when the registration mints nothing", func() { + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(`{"name":"w"}`)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + natsRow := cluster.CarrierRow{Active: cluster.CarrierNATS, Epoch: 3, State: cluster.StateStable} + h := RegisterNodeEndpoint(registry, "", true, nil, "", natsCfg, + WithCarrierReader(fixedCarrier{row: natsRow}), WithTunnelDisconnector(tunnels)) + Expect(h(e.NewContext(req, rec))).To(Succeed()) + Expect(tunnels.seen()).To(BeEmpty()) + }) +}) diff --git a/core/http/endpoints/localai/nodes_follow_test.go b/core/http/endpoints/localai/nodes_follow_test.go new file mode 100644 index 000000000000..4a25f83ac482 --- /dev/null +++ b/core/http/endpoints/localai/nodes_follow_test.go @@ -0,0 +1,280 @@ +package localai + +import ( + "context" + "encoding/json" + "errors" + "net/http" + "net/http/httptest" + "strings" + "time" + + "github.com/labstack/echo/v4" + "github.com/mudler/LocalAI/core/services/cluster" + "github.com/mudler/LocalAI/core/services/nodes" + "github.com/mudler/LocalAI/core/services/testutil" + "github.com/mudler/LocalAI/pkg/natsauth" + "github.com/nats-io/nkeys" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// fixedNATS is the handover of a deployment with a NATS server. +type fixedNATS struct { + url, ca string + clientTLS bool + err error +} + +func (f fixedNATS) WorkerURL(context.Context) (string, error) { return f.url, f.err } +func (f fixedNATS) CAPEM() (string, error) { return f.ca, nil } +func (f fixedNATS) ClientTLS() bool { return f.clientTLS } + +var _ = Describe("What a worker needs to follow a change of carrier", func() { + var ( + ctx context.Context + registry *nodes.NodeRegistry + natsCfg natsauth.Config + tunnels *recordingDisconnector + ) + + stable := func(c cluster.Carrier) cluster.CarrierRow { + return cluster.CarrierRow{Active: c, Epoch: 5, State: cluster.StateStable} + } + preparing := func(from, to cluster.Carrier) cluster.CarrierRow { + return cluster.CarrierRow{Active: from, Epoch: 6, State: cluster.StatePrepare, Target: to} + } + handover := WithNATSHandover(fixedNATS{url: "nats://nats.example:4222", ca: "-----BEGIN CERTIFICATE-----\nCA\n-----END CERTIFICATE-----\n"}) + + BeforeEach(func() { + ctx = context.Background() + var err error + registry, err = nodes.NewNodeRegistry(testutil.SetupTestDB()) + Expect(err).ToNot(HaveOccurred()) + tunnels = &recordingDisconnector{} + akp, err := nkeys.CreateAccount() + Expect(err).ToNot(HaveOccurred()) + seed, err := akp.Seed() + Expect(err).ToNot(HaveOccurred()) + natsCfg = natsauth.Config{AccountSeed: string(seed)} + }) + + register := func(body string, opts ...RegisterOption) map[string]any { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + opts = append(opts, WithTunnelDisconnector(tunnels)) + Expect(RegisterNodeEndpoint(registry, "", true, nil, "", natsCfg, opts...)(e.NewContext(req, rec))).To(Succeed()) + Expect(rec.Code).To(Equal(http.StatusCreated), rec.Body.String()) + var resp map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + return resp + } + storedHash := func(name string) string { + GinkgoHelper() + node, err := registry.GetByName(ctx, name) + Expect(err).ToNot(HaveOccurred()) + return node.TunnelTokenHash + } + + Describe("the registration", func() { + It("hands over the NATS address and the CA with the credential, when NATS is active", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051"}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), handover) + Expect(resp["nats_url"]).To(Equal("nats://nats.example:4222")) + Expect(resp["nats_ca"]).To(ContainSubstring("BEGIN CERTIFICATE")) + Expect(resp).To(HaveKey("nats_jwt")) + Expect(resp["credential_for"]).To(Equal("nats")) + Expect(resp).ToNot(HaveKey("nats_client_tls")) + }) + + It("says that the worker needs a certificate of its own when the server asks for one", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051"}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), + WithNATSHandover(fixedNATS{url: "tls://nats.example:4222", clientTLS: true})) + Expect(resp["nats_client_tls"]).To(BeTrue()) + }) + + It("hands over nothing of NATS when none is configured, and when the read fails", func() { + reader := WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}) + Expect(register(`{"name":"a","address":"10.0.0.1:50051"}`, reader, WithNATSHandover(fixedNATS{}))).ToNot(HaveKey("nats_url")) + Expect(register(`{"name":"b","address":"10.0.0.1:50051"}`, reader, WithNATSHandover(fixedNATS{err: errors.New("down")}))).ToNot(HaveKey("nats_url")) + }) + + It("does not hand over the address of NATS to a worker that asks for the tunnel", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"tunnel"}`, + WithCarrierReader(fixedCarrier{row: preparing(cluster.CarrierNATS, cluster.CarrierTunnel)}), handover) + Expect(resp).To(HaveKey("tunnel_token")) + Expect(resp).ToNot(HaveKey("nats_url")) + Expect(resp["credential_for"]).To(Equal("tunnel")) + }) + + It("gives the credential of the target to a worker that asks for it while the change is prepared", func() { + By("a tunnel worker that follows to NATS") + first := register(`{"name":"w","address":"10.0.0.1:50051"}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierTunnel)}), handover) + held := storedHash("w") + Expect(first).To(HaveKey("tunnel_token")) + before := len(tunnels.seen()) + + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"nats"}`, + WithCarrierReader(fixedCarrier{row: preparing(cluster.CarrierTunnel, cluster.CarrierNATS)}), handover) + + Expect(resp["carrier"]).To(Equal("tunnel"), "the active carrier is still the tunnel") + Expect(resp["carrier_state"]).To(Equal("prepare")) + Expect(resp["carrier_target"]).To(Equal("nats")) + Expect(resp["credential_for"]).To(Equal("nats")) + Expect(resp).To(HaveKey("nats_jwt")) + Expect(resp).To(HaveKey("nats_url")) + Expect(resp).ToNot(HaveKey("tunnel_token")) + Expect(storedHash("w")).To(Equal(held), "the tunnel credential that the worker holds stays valid") + Expect(tunnels.seen()).To(HaveLen(before), "and its session is not ended") + }) + + It("mints a tunnel credential for a worker on NATS that asks for the tunnel, and leaves its NATS credential alone", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"tunnel"}`, + WithCarrierReader(fixedCarrier{row: preparing(cluster.CarrierNATS, cluster.CarrierTunnel)}), handover) + Expect(resp["carrier"]).To(Equal("nats")) + Expect(resp["credential_for"]).To(Equal("tunnel")) + Expect(resp["tunnel_token"]).ToNot(BeEmpty()) + Expect(resp).ToNot(HaveKey("nats_jwt")) + Expect(storedHash("w")).ToNot(BeEmpty()) + }) + + It("ignores a request for a carrier that is neither active, nor the target, nor draining", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"tunnel"}`, + WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), handover) + Expect(resp["credential_for"]).To(Equal("nats")) + Expect(resp).ToNot(HaveKey("tunnel_token"), "a node cannot get a tunnel while the cluster has none") + Expect(storedHash("w")).To(BeEmpty()) + + resp = register(`{"name":"w","address":"10.0.0.1:50051","carrier":"nonsense"}`, + WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), handover) + Expect(resp["credential_for"]).To(Equal("nats")) + }) + + It("gives the credential of the carrier that drains to a worker that asks for it", func() { + until := time.Now().Add(time.Minute) + row := cluster.CarrierRow{Active: cluster.CarrierTunnel, Epoch: 8, State: cluster.StateStable, Draining: cluster.CarrierNATS, DrainingUntil: &until} + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"nats"}`, WithCarrierReader(fixedCarrier{row: row}), handover) + Expect(resp["credential_for"]).To(Equal("nats")) + Expect(resp["carrier_draining"]).To(Equal("nats")) + Expect(resp).To(HaveKey("carrier_draining_until")) + Expect(resp).To(HaveKey("nats_jwt")) + }) + + It("keeps the address of a worker that has none to offer the tunnel while it asks for NATS", func() { + register(`{"name":"w","address":"10.0.0.1:50051","routable":true}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierTunnel)}), handover) + register(`{"name":"w","address":"10.0.0.1:50051","routable":false,"carrier":"nats"}`, + WithCarrierReader(fixedCarrier{row: preparing(cluster.CarrierTunnel, cluster.CarrierNATS)}), handover) + node, err := registry.GetByName(ctx, "w") + Expect(err).ToNot(HaveOccurred()) + Expect(node.Address).To(Equal("10.0.0.1:50051"), "NATS dials the address, so it must not be dropped") + }) + + It("answers a deployment with no carrier reader as it always did, whatever the worker asks", func() { + resp := register(`{"name":"w","address":"10.0.0.1:50051","carrier":"tunnel"}`) + Expect(resp).ToNot(HaveKey("carrier")) + Expect(resp).ToNot(HaveKey("credential_for")) + Expect(resp).ToNot(HaveKey("tunnel_token")) + Expect(resp).To(HaveKey("nats_jwt")) + }) + }) + + Describe("the heartbeat", func() { + var id string + + BeforeEach(func() { + id = register(`{"name":"w","address":"10.0.0.1:50051"}`)["id"].(string) + }) + + beat := func(body string, opts ...RegisterOption) (int, map[string]any) { + GinkgoHelper() + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + c.SetParamNames("id") + c.SetParamValues(id) + Expect(HeartbeatEndpoint(registry, opts...)(c)).To(Succeed()) + var resp map[string]any + Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed()) + return rec.Code, resp + } + + It("tells the worker the carrier, the epoch, the state and the target of a change", func() { + code, resp := beat(`{}`, WithCarrierReader(fixedCarrier{row: preparing(cluster.CarrierNATS, cluster.CarrierTunnel)})) + Expect(code).To(Equal(http.StatusOK)) + Expect(resp["message"]).To(Equal("heartbeat received")) + Expect(resp["carrier"]).To(Equal("nats")) + Expect(resp["carrier_epoch"]).To(BeEquivalentTo(6)) + Expect(resp["carrier_state"]).To(Equal("prepare")) + Expect(resp["carrier_target"]).To(Equal("tunnel")) + }) + + It("tells the worker that NATS asks for a client certificate", func() { + _, resp := beat(`{}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), WithNATSHandover(fixedNATS{url: "tls://n:4222", clientTLS: true})) + Expect(resp["nats_client_tls"]).To(BeTrue()) + _, resp = beat(`{}`, WithCarrierReader(fixedCarrier{row: stable(cluster.CarrierNATS)}), handover) + Expect(resp).ToNot(HaveKey("nats_client_tls")) + }) + + It("answers as it always did when there is no carrier reader, or the row cannot be read", func() { + _, resp := beat(`{}`) + Expect(resp).To(Equal(map[string]any{"message": "heartbeat received"})) + _, resp = beat(`{}`, WithCarrierReader(fixedCarrier{err: errors.New("down")})) + Expect(resp).To(Equal(map[string]any{"message": "heartbeat received"})) + }) + + It("stores what the worker reports about its carriers", func() { + code, _ := beat(`{"attached":["nats","tunnel"],"attached_epoch":6,"follow_capabilities":["nats","tunnel"],"follow_error":""}`) + Expect(code).To(Equal(http.StatusOK)) + node, err := registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Attached).To(Equal("nats,tunnel")) + Expect(node.AttachedEpoch).To(BeEquivalentTo(6)) + Expect(node.Follow).To(Equal("nats,tunnel")) + + beat(`{"attached":["tunnel"],"attached_epoch":7,"follow_capabilities":["tunnel"],"follow_error":"no address that a frontend can dial"}`) + node, err = registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Attached).To(Equal("tunnel")) + Expect(node.Follow).To(Equal("tunnel")) + Expect(node.FollowError).To(Equal("no address that a frontend can dial")) + }) + + It("stores an empty attachment as a report, because the worker still reports what it can follow", func() { + beat(`{"attached":["nats"],"follow_capabilities":["nats"]}`) + beat(`{"attached":[],"follow_capabilities":["nats","tunnel"]}`) + node, err := registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Attached).To(BeEmpty()) + Expect(node.Follow).To(Equal("nats,tunnel")) + }) + + It("stores nothing for a worker that predates the report, and leaves an earlier report alone", func() { + beat(`{"available_ram":1024}`) + node, err := registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Follow).To(BeEmpty()) + Expect(node.Attached).To(BeEmpty()) + }) + + It("keeps only carriers that the cluster knows and cuts a reason that does not fit the column", func() { + beat(`{"attached":["nats","gopher","nats"],"follow_capabilities":["tunnel","nope"],"follow_error":"` + strings.Repeat("é", 200) + `"}`) + node, err := registry.Get(ctx, id) + Expect(err).ToNot(HaveOccurred()) + Expect(node.Attached).To(Equal("nats")) + Expect(node.Follow).To(Equal("tunnel")) + Expect(len(node.FollowError)).To(BeNumerically("<=", 255)) + Expect(node.FollowError).To(HavePrefix("é")) + }) + + It("answers 404 for a node that is not registered and stores nothing", func() { + id = "no-such-node" + code, _ := beat(`{"attached":["nats"],"follow_capabilities":["nats"]}`) + Expect(code).To(Equal(http.StatusNotFound)) + }) + }) +}) diff --git a/core/http/endpoints/localai/nodes_logs_tunnel_test.go b/core/http/endpoints/localai/nodes_logs_tunnel_test.go new file mode 100644 index 000000000000..cff1ba11204c --- /dev/null +++ b/core/http/endpoints/localai/nodes_logs_tunnel_test.go @@ -0,0 +1,78 @@ +package localai + +import ( + "context" + "io" + "net" + "net/http" + "net/http/httptest" + "sync" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + + "github.com/mudler/LocalAI/core/services/nodes" +) + +var _ = Describe("Backend logs of a worker that holds a tunnel", func() { + Describe("the host of the request", func() { + It("names the reserved host for a worker with no address at all", func() { + host, ok := workerLogsHost(&nodes.BackendNode{ID: "n1"}) + Expect(ok).To(BeTrue()) + Expect(host).To(Equal("n1.worker.invalid")) + Expect(nodes.IsTunnelOnlyHost(host)).To(BeTrue()) + }) + + It("keeps the HTTP address of a worker that has one", func() { + host, ok := workerLogsHost(&nodes.BackendNode{ID: "n1", Address: "10.0.0.1:50051", HTTPAddress: "10.0.0.1:50050"}) + Expect(ok).To(BeTrue()) + Expect(host).To(Equal("10.0.0.1:50050")) + }) + + It("still refuses a worker that has a gRPC address and no HTTP server, as before", func() { + _, ok := workerLogsHost(&nodes.BackendNode{ID: "n1", Address: "10.0.0.1:50051"}) + Expect(ok).To(BeFalse()) + }) + }) + + Describe("the proxy of the environment", func() { + dial := func(string) func(context.Context, string, string) (net.Conn, error) { return nil } + + It("is not used to reach a worker through its tunnel, so the dialer gets the stream", func() { + Expect(workerTransport(dial, "n1", "n1.worker.invalid").Proxy).To(BeNil()) + }) + + It("is still used for a worker that has an address, as before", func() { + Expect(workerTransport(dial, "n1", "10.0.0.1:50050").Proxy).ToNot(BeNil()) + }) + + It("sends the request of a worker without an address to the dialer even when the environment names a proxy", func() { + // The proxy of the environment is read once for the process, so this + // spec does not rely on it: it checks that the transport has none. + var ( + srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, "logs for "+r.Host) })) + mu sync.Mutex + dials []string + ) + DeferCleanup(srv.Close) + dialer := func(string) func(context.Context, string, string) (net.Conn, error) { + return func(ctx context.Context, _, addr string) (net.Conn, error) { + mu.Lock() + dials = append(dials, addr) + mu.Unlock() + var d net.Dialer + return d.DialContext(ctx, "tcp", srv.Listener.Addr().String()) + } + } + resp, err := proxyHTTPToWorker(context.Background(), dialer, "n1", "n1.worker.invalid", "/v1/backend-logs", "token") + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = resp.Body.Close() }() + body, _ := io.ReadAll(resp.Body) + Expect(string(body)).To(Equal("logs for n1.worker.invalid")) + mu.Lock() + defer mu.Unlock() + Expect(dials).To(HaveLen(1)) + Expect(dials[0]).To(HavePrefix("n1.worker.invalid"), "the dialer was asked for the worker, not for a proxy") + }) + }) +}) diff --git a/core/http/endpoints/mcp/localai_assistant_test.go b/core/http/endpoints/mcp/localai_assistant_test.go index 3e539ce7d4bd..f6409218e89f 100644 --- a/core/http/endpoints/mcp/localai_assistant_test.go +++ b/core/http/endpoints/mcp/localai_assistant_test.go @@ -218,6 +218,10 @@ func (stubClient) ClearRouterCorpus(_ context.Context, routerModel string) (*loc return &localaitools.RouterCorpusClearResult{Router: routerModel}, nil } +func (stubClient) GetClusterCarrier(_ context.Context) (*localaitools.ClusterCarrierStatus, error) { + return &localaitools.ClusterCarrierStatus{}, nil +} + func (stubClient) ListFailoverChains(_ context.Context) ([]localaitools.FailoverChainInfo, error) { return []localaitools.FailoverChainInfo{}, nil } diff --git a/core/http/react-ui/e2e/cluster-transport.spec.js b/core/http/react-ui/e2e/cluster-transport.spec.js new file mode 100644 index 000000000000..5562d891b1d6 --- /dev/null +++ b/core/http/react-ui/e2e/cluster-transport.spec.js @@ -0,0 +1,246 @@ +import { test, expect } from './coverage-fixtures.js' + +const node = { id: 'n1', name: 'atlas', node_type: 'backend', address: '10.0.0.1:50051', status: 'healthy', labels: {}, total_vram: 100, available_vram: 40, total_ram: 200, available_ram: 100, total_disk: 1000, available_disk: 600, model_count: 0, in_flight_count: 0 } + +const stableReport = { + row: { active: 'nats', epoch: 4, state: 'stable', changed_by: 'admin', changed_at: '2026-10-08T10:00:00Z', prev_epoch: 3 }, + active: 'nats', epoch: 4, state: 'stable', ok: true, + replicas: [ + { id: 'frontend-a', version: 'v4.12.0', ready_epoch: 4 }, + { id: 'frontend-b', version: 'v4.12.0', ready_epoch: 4 }, + ], + workers: [ + { id: 'w1', name: 'dual-worker', type: 'backend', attached: ['nats'], follow: ['nats', 'tunnel'], can_follow: true }, + { id: 'w2', name: 'tunnel-worker', type: 'backend', attached: ['tunnel'], follow: ['tunnel'], follow_error: 'no routable address: it cannot follow a change to NATS', can_follow: false }, + { id: 'w3', name: 'old-worker', type: 'backend', attached: [], can_follow: false }, + ], + in_flight: { loads: 0, jobs: 1, pending_claims: 2, claimed_claims: 0 }, + timings: {}, + settings: { nats_url: 'nats://nats:4222', nats_worker_url: '', prepare_timeout: '', transition_window: '', max_drain: '' }, +} + +const blockedDryRun = { + ...stableReport, target: 'tunnel', ok: false, + blockers: [ + { kind: 'worker', id: 'w3', reason: 'the worker predates carrier switching: it reports no capabilities and stays on NATS', forceable: true }, + ], +} +const cleanDryRun = { ...stableReport, target: 'tunnel', ok: true } + +async function mockCluster(page, { carrier = stableReport, carrierStatus = 200, nodes = [node] } = {}) { + await page.route('**/api/features', route => route.fulfill({ status: 200, contentType: 'application/json', body: JSON.stringify({ distributed: true }) })) + await page.route('**/api/auth/status', route => route.fulfill({ + status: 200, contentType: 'application/json', + body: JSON.stringify({ authEnabled: true, staticApiKeyRequired: false, providers: ['local'], user: { id: 'admin', name: 'Admin', role: 'admin', provider: 'local' } }), + })) + await page.route('**/api/nodes', route => route.fulfill({ status: 200, contentType: 'application/json', body: JSON.stringify(nodes) })) + const calls = { posts: [], puts: [] } + const state = { carrier } + await page.route('**/api/cluster/carrier', async route => { + const request = route.request() + if (request.method() === 'POST') { + const body = request.postDataJSON() + calls.posts.push(body) + const handler = state.onPost || (() => ({ status: 202, body: { state: 'prepare', active: 'nats', target: 'tunnel', epoch: 5 } })) + const answer = handler(body) + return route.fulfill({ status: answer.status, contentType: 'application/json', body: JSON.stringify(answer.body) }) + } + return route.fulfill({ status: carrierStatus, contentType: 'application/json', body: JSON.stringify(carrierStatus === 200 ? state.carrier : { error: 'distributed mode is not enabled on this frontend' }) }) + }) + await page.route('**/api/cluster/settings', async route => { + const body = route.request().postDataJSON() + calls.puts.push(body) + const answer = (state.onPut || (() => ({ status: 200, body: { saved: true, nats: { reachable: true } } })))(body) + return route.fulfill({ status: answer.status, contentType: 'application/json', body: JSON.stringify(answer.body) }) + }) + return { calls, state } +} + +test.describe('Cluster transport panel', () => { + test('is hidden when the cluster routes answer 503', async ({ page }) => { + await mockCluster(page, { carrierStatus: 503 }) + await page.goto('/app/nodes') + await expect(page.getByLabel('Fleet health summary')).toBeVisible({ timeout: 15_000 }) + await expect(page.getByTestId('cluster-transport')).toHaveCount(0) + }) + + test('shows the active carrier, epoch, replicas and workers in the stable state', async ({ page }) => { + await mockCluster(page) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await expect(panel).toBeVisible({ timeout: 15_000 }) + await expect(panel.getByTestId('transport-active')).toHaveText('NATS') + await expect(panel.getByTestId('transport-state')).toHaveText('stable') + await expect(panel.getByTestId('transport-epoch')).toHaveText('epoch 4') + await expect(panel.getByTestId('transport-replica')).toHaveCount(2) + await expect(panel.getByTestId('transport-replica').first()).toContainText('ready') + const workers = panel.getByTestId('transport-worker') + await expect(workers).toHaveCount(3) + await expect(workers.nth(0)).toContainText('Dual-capable') + await expect(workers.nth(1)).toContainText('Tunnel-only') + await expect(workers.nth(1)).toContainText('no routable address') + await expect(workers.nth(2)).toContainText('Legacy (NATS only)') + // The legacy worker cannot follow a change to the tunnel, and the table says so. + await expect(workers.nth(2)).toContainText('predates carrier switching') + await expect(workers.nth(2)).not.toContainText(/\bok\b/) + await expect(workers.nth(0)).toContainText('ok') + await expect(panel.getByRole('button', { name: /Abort change/ })).toHaveCount(0) + await expect(panel.getByLabel('NATS URL', { exact: true })).toHaveValue('nats://nats:4222') + }) + + test('dry run lists blockers and the confirmation gates the switch', async ({ page }) => { + const { calls, state } = await mockCluster(page) + state.onPost = body => body.dry_run ? { status: 200, body: blockedDryRun } : { status: 202, body: { state: 'prepare', active: 'nats', target: 'tunnel', epoch: 5 } } + await page.goto('/app/nodes') + await page.getByRole('button', { name: /Switch to Database tunnel/ }).click() + const dialog = page.getByRole('dialog', { name: 'Switch cluster transport' }) + const start = dialog.getByRole('button', { name: /^Switch to Database tunnel$/ }) + await expect(start).toBeDisabled() + await dialog.getByRole('button', { name: 'Dry run' }).click() + await expect(dialog.getByTestId('transport-blocker')).toHaveCount(1) + await expect(dialog.getByTestId('transport-blocker')).toContainText('predates carrier switching') + await expect(dialog.getByTestId('transport-dry-verdict')).toContainText('1 blocker') + await expect(dialog.getByLabel('Live replicas').locator('li')).toHaveCount(2) + expect(calls.posts[0]).toEqual({ target: 'tunnel', dry_run: true }) + // A blocker without force keeps the button off, even once confirmed. + const confirm = dialog.getByLabel('I confirm these are all the frontends and all are upgraded') + await confirm.check() + await expect(start).toBeDisabled() + }) + + test('a clean dry run needs the confirmation, then starts the change', async ({ page }) => { + const { calls, state } = await mockCluster(page) + state.onPost = body => body.dry_run ? { status: 200, body: cleanDryRun } : { status: 202, body: { state: 'prepare', active: 'nats', target: 'tunnel', epoch: 5 } } + await page.goto('/app/nodes') + await page.getByRole('button', { name: /Switch to Database tunnel/ }).click() + const dialog = page.getByRole('dialog', { name: 'Switch cluster transport' }) + await dialog.getByRole('button', { name: 'Dry run' }).click() + await expect(dialog.getByTestId('transport-dry-verdict')).toContainText('Nothing blocks') + const start = dialog.getByRole('button', { name: /^Switch to Database tunnel$/ }) + await expect(start).toBeDisabled() + await dialog.getByLabel('I confirm these are all the frontends and all are upgraded').check() + await expect(start).toBeEnabled() + await start.click() + await expect(dialog).toHaveCount(0) + expect(calls.posts.at(-1)).toEqual({ target: 'tunnel', force: false }) + }) + + test('force needs its own confirmation', async ({ page }) => { + const { calls, state } = await mockCluster(page) + state.onPost = body => body.dry_run ? { status: 200, body: blockedDryRun } : { status: 202, body: { state: 'prepare', active: 'nats', target: 'tunnel', epoch: 5, force: true } } + await page.goto('/app/nodes') + await page.getByRole('button', { name: /Switch to Database tunnel/ }).click() + const dialog = page.getByRole('dialog', { name: 'Switch cluster transport' }) + await dialog.getByRole('button', { name: 'Dry run' }).click() + await dialog.getByLabel('I confirm these are all the frontends and all are upgraded').check() + await dialog.getByLabel('Force the change past the blockers above').check() + await dialog.getByRole('button', { name: /^Force switch to Database tunnel$/ }).click() + const confirm = page.getByRole('alertdialog') + await expect(confirm).toContainText('Force the change?') + // Nothing is sent until the second confirmation. + expect(calls.posts.filter(p => !p.dry_run)).toHaveLength(0) + await confirm.getByRole('button', { name: 'Force the change' }).click() + await expect(dialog).toHaveCount(0) + expect(calls.posts.at(-1)).toEqual({ target: 'tunnel', force: true }) + }) + + test('a refused request shows the blockers of the server', async ({ page }) => { + const { state } = await mockCluster(page) + state.onPost = body => body.dry_run ? { status: 200, body: cleanDryRun } : { status: 422, body: { error: 'the change is blocked', blockers: [{ kind: 'replica', id: 'frontend-b', reason: 'frontend-b cannot build the tunnel carrier', forceable: true }] } } + await page.goto('/app/nodes') + await page.getByRole('button', { name: /Switch to Database tunnel/ }).click() + const dialog = page.getByRole('dialog', { name: 'Switch cluster transport' }) + await dialog.getByRole('button', { name: 'Dry run' }).click() + await dialog.getByLabel('I confirm these are all the frontends and all are upgraded').check() + await dialog.getByRole('button', { name: /^Switch to Database tunnel$/ }).click() + await expect(dialog.getByRole('alert')).toContainText('frontend-b cannot build the tunnel carrier') + }) + + test('shows replica readiness and an abort button during prepare', async ({ page }) => { + const prepare = { + ...stableReport, state: 'prepare', epoch: 5, target: 'tunnel', + row: { ...stableReport.row, epoch: 5, state: 'prepare', target: 'tunnel' }, + replicas: [{ id: 'frontend-a', version: 'v4.12.0', ready_epoch: 5 }, { id: 'frontend-b', version: 'v4.12.0', ready_epoch: 4 }], + } + const { calls, state } = await mockCluster(page, { carrier: prepare }) + state.onPost = () => ({ status: 200, body: { active: 'nats', state: 'stable', epoch: 6 } }) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await expect(panel.getByTestId('transport-state')).toHaveText('prepare') + await expect(panel.getByRole('button', { name: /^Switch to/ })).toBeDisabled() + const replicas = panel.getByTestId('transport-replica') + await expect(replicas.nth(0)).toContainText('ready') + await expect(replicas.nth(1)).toContainText('not ready') + await panel.getByRole('button', { name: 'Abort change' }).click() + await page.getByRole('alertdialog').getByRole('button', { name: 'Abort the change' }).click() + await expect.poll(() => calls.posts.length).toBe(1) + expect(calls.posts[0]).toEqual({ abort: true }) + }) + + test('counts down the drain after a commit', async ({ page }) => { + const draining = { + ...stableReport, active: 'tunnel', epoch: 6, + row: { ...stableReport.row, active: 'tunnel', epoch: 6, draining: 'nats' }, + drain_remaining_ns: 125e9, + } + await mockCluster(page, { carrier: draining }) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await expect(panel.getByTestId('transport-active')).toHaveText('Database tunnel') + await expect(panel.getByTestId('transport-state')).toHaveText('draining') + const drain = panel.getByTestId('transport-drain') + await expect(drain).toContainText('Draining NATS') + await expect(drain).toContainText(/2:0\d left|2:1\d left|2:2\d left/) + const first = await drain.textContent() + await expect.poll(async () => drain.textContent(), { timeout: 5000 }).not.toBe(first) + }) + + test('saving the NATS settings reports reachability and does not switch', async ({ page }) => { + const { calls } = await mockCluster(page) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await panel.getByLabel('NATS URL', { exact: true }).fill('nats://nats.example:4222') + await panel.getByLabel('Worker NATS URL').fill('nats://nats.public.example:4222') + await panel.getByRole('button', { name: 'Save' }).click() + await expect(panel.getByTestId('transport-settings-result')).toContainText('reaches the NATS server') + await expect(panel.getByTestId('transport-settings-result')).toContainText('did not change') + expect(calls.puts[0]).toEqual({ nats_url: 'nats://nats.example:4222', nats_worker_url: 'nats://nats.public.example:4222' }) + expect(calls.posts).toHaveLength(0) + }) + + test('a NATS URL that cannot be reached shows the error', async ({ page }) => { + const { state } = await mockCluster(page) + state.onPut = () => ({ status: 422, body: { error: 'nats_url: this replica cannot reach the NATS server at nats://bad:4222: connection refused' } }) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await panel.getByLabel('NATS URL', { exact: true }).fill('nats://bad:4222') + await panel.getByRole('button', { name: 'Save' }).click() + await expect(panel.getByTestId('transport-settings-result')).toContainText('cannot reach the NATS server') + }) + + test('renders on a narrow viewport without horizontal overflow', async ({ page }) => { + await page.setViewportSize({ width: 390, height: 844 }) + await mockCluster(page) + await page.goto('/app/nodes') + const panel = page.getByTestId('cluster-transport') + await expect(panel).toBeVisible({ timeout: 15_000 }) + const overflow = await page.evaluate(() => document.documentElement.scrollWidth - document.documentElement.clientWidth) + expect(overflow).toBeLessThanOrEqual(1) + }) + + test('the worker hint does not ask for a NATS address', async ({ page }) => { + await mockCluster(page, { nodes: [] }) + await page.goto('/app/nodes') + const hint = page.getByTestId('worker-hint-carrier') + await expect(hint).toBeVisible({ timeout: 15_000 }) + await expect(hint).toContainText('only to override') + const commands = page.locator('.p2p-cmd pre') + await expect(commands.first()).toContainText('--register-to') + // One flag on each line, as a shell reads it. + expect(await commands.first().innerText()).toMatch(/ \\\n --register-to .* \\\n --registration-token /) + for (const command of await commands.allTextContents()) { + expect(command).not.toContain('nats') + expect(command).not.toContain('NATS') + } + }) +}) diff --git a/core/http/react-ui/src/App.css b/core/http/react-ui/src/App.css index 3ac4da254a34..4bad41b7d796 100644 --- a/core/http/react-ui/src/App.css +++ b/core/http/react-ui/src/App.css @@ -13875,3 +13875,35 @@ button.lane:hover { touch-action: none; cursor: grab; } + +/* Cluster transport panel (Nodes page) */ +.transport-panel { display: grid; gap: var(--spacing-md); } +.transport-head { align-items: flex-start; display: flex; flex-wrap: wrap; gap: var(--spacing-md); justify-content: space-between; } +.transport-head__line { align-items: center; display: flex; flex-wrap: wrap; gap: 10px; } +.transport-head__carrier { font-size: 1.375rem; letter-spacing: -.02em; } +.transport-head__actions { display: flex; flex-wrap: wrap; gap: var(--spacing-sm); } +.transport-grid { display: grid; gap: var(--spacing-lg); grid-template-columns: minmax(0, .8fr) minmax(0, 1.8fr); } +.transport-list { display: grid; gap: 6px; list-style: none; margin: 0 0 var(--spacing-md); padding: 0; } +.transport-list__row { align-items: center; display: flex; flex-wrap: wrap; font-size: var(--text-sm); gap: 8px; } +.transport-list__name { font-family: var(--font-mono); word-break: break-all; } +.transport-list__meta, .transport-list__state { color: var(--color-text-muted); font-size: var(--text-xs); } +.transport-dot { border-radius: 50%; display: inline-block; flex: none; height: 8px; width: 8px; } +.transport-dot--ok { background: var(--color-success); } +.transport-dot--warn { background: var(--color-warning); } +.transport-table-wrap { overflow-x: auto; } +.transport-table { border-collapse: collapse; font-size: var(--text-sm); width: 100%; } +.transport-table th { color: var(--color-text-muted); font-size: var(--text-xs); font-weight: 600; padding: 4px 10px 4px 0; text-align: left; } +.transport-table td { border-top: 1px solid var(--color-border-subtle); padding: 6px 10px 6px 0; vertical-align: top; } +.transport-table td:first-child { white-space: nowrap; } +.transport-error { color: var(--color-error); font-size: var(--text-sm); margin: 0; } +.transport-ok { color: var(--color-success); font-size: var(--text-sm); margin: 0; } +.transport-blockers { display: grid; gap: 8px; list-style: none; margin: 0 0 var(--spacing-md); padding: 0; } +.transport-blockers li { align-items: baseline; display: flex; flex-wrap: wrap; font-size: var(--text-sm); gap: 8px; } +.transport-settings { border-top: 1px solid var(--color-border-subtle); display: grid; gap: var(--spacing-sm); padding-top: var(--spacing-md); } +.transport-field { display: grid; font-size: var(--text-sm); gap: 4px; } +.transport-dialog { display: grid; gap: var(--spacing-md); padding: var(--spacing-lg); } +.transport-dialog__targets { justify-self: start; } +.transport-dialog__actions { display: flex; gap: var(--spacing-sm); } +.transport-dialog__result { display: grid; gap: var(--spacing-sm); } +.transport-dialog__footer { display: flex; flex-wrap: wrap; gap: var(--spacing-sm); justify-content: flex-end; } +@media (max-width: 760px) { .transport-grid { grid-template-columns: minmax(0, 1fr); } } diff --git a/core/http/react-ui/src/components/nodes/ClusterTransport.jsx b/core/http/react-ui/src/components/nodes/ClusterTransport.jsx new file mode 100644 index 000000000000..89c3e2beeef4 --- /dev/null +++ b/core/http/react-ui/src/components/nodes/ClusterTransport.jsx @@ -0,0 +1,367 @@ +import { useCallback, useEffect, useRef, useState } from 'react' +import { clusterApi } from '../../utils/api' +import { + CARRIERS, PROFILE_LABELS, canForce, carrierLabel, formatCountdown, otherCarrier, + followStatus, replicaReady, stateLabel, workerProfile, +} from '../../utils/clusterTransport' +import ConfirmDialog from '../ConfirmDialog' +import Modal from '../Modal' + +const POLL_IDLE_MS = 5000 +const POLL_BUSY_MS = 2000 + +// The routes answer 404 on a build without them and 503 on a frontend that is +// not distributed. Either way there is no transport to show. +function isUnavailable(error) { + return error?.status === 404 || error?.status === 503 || error?.status === 401 || error?.status === 403 +} + +function stateTone(state) { + if (state === 'stable') return 'badge-success' + if (state === 'draining') return 'badge-info' + return 'badge-warning' +} + +function ReplicaList({ replicas, report, label }) { + if (!replicas || replicas.length === 0) return

No live frontend replica reported.

+ return ( +
    + {replicas.map(replica => { + const ready = replicaReady(replica, report) + return ( +
  • +
    • + ) + })} +
+ ) +} + +function WorkerTable({ workers, active }) { + if (!workers || workers.length === 0) return

No worker is registered.

+ return ( +
+ + + + {workers.map(worker => ( + + + + + + + ))} + +
WorkerProfileAttached toFollow status
{worker.name || worker.id}{PROFILE_LABELS[workerProfile(worker)]}{(worker.attached || []).length ? worker.attached.map(carrierLabel).join(' + ') : 'not reported'}{(() => { + const status = followStatus(worker, active) + return status.ok + ? {status.text} + : {status.text} + })()}
+
+ ) +} + +function BlockerList({ blockers }) { + return ( +
    + {blockers.map((blocker, index) => ( +
  • + {blocker.kind}{blocker.id ? ` ${blocker.id}` : ''} + {blocker.reason} + {blocker.forceable && can be forced} +
    • + ))} +
+ ) +} + +function SwitchDialog({ report, onClose, onStarted, addToast }) { + const [target, setTarget] = useState(otherCarrier(report.active)) + const [dry, setDry] = useState(null) + const [dryError, setDryError] = useState('') + const [running, setRunning] = useState(false) + const [confirmed, setConfirmed] = useState(false) + const [force, setForce] = useState(false) + const [confirmForce, setConfirmForce] = useState(false) + const [submitting, setSubmitting] = useState(false) + const [submitError, setSubmitError] = useState('') + + const resetResult = () => { setDry(null); setDryError(''); setConfirmed(false); setForce(false); setSubmitError('') } + const pickTarget = value => { setTarget(value); resetResult() } + + const runDry = async () => { + setRunning(true) + resetResult() + try { + setDry(await clusterApi.dryRun(target)) + } catch (error) { + setDryError(error.message || 'The dry run failed') + } finally { + setRunning(false) + } + } + + const blockers = dry?.blockers || [] + const blocked = blockers.length > 0 + const forceable = canForce(dry) + const sameCarrier = target === report.active + const mayStart = !!dry && confirmed && !sameCarrier && (!blocked || (force && forceable)) + + const start = async () => { + setSubmitting(true) + setSubmitError('') + try { + const answer = await clusterApi.switchCarrier(target, force) + addToast(`Change to ${carrierLabel(target)} started (epoch ${answer.epoch}).`, 'success') + onStarted() + onClose() + } catch (error) { + const refused = error?.body?.blockers + setSubmitError(Array.isArray(refused) && refused.length + ? `${error.message}: ${refused.map(b => b.reason).join('; ')}` + : (error.message || 'The request failed')) + } finally { + setSubmitting(false) + setConfirmForce(false) + } + } + + const onPrimary = () => { if (force) setConfirmForce(true); else start() } + const forcedNodes = blockers.filter(b => b.kind === 'worker').length + + return ( + <> + {} : onClose} maxWidth="680px" ariaLabel="Switch cluster transport"> +
+