mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-08 12:14:06 -04:00
Compare commits
119
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
490e4ba778 | ||
|
|
4e4597dfc2 | ||
|
|
2756eaeab0 | ||
|
|
0291189e82 | ||
|
|
4e32249227 | ||
|
|
994ec9dde0 | ||
|
|
d397393022 | ||
|
|
8403d9da05 | ||
|
|
e7b2e1ee55 | ||
|
|
d170de2095 | ||
|
|
7b66df6651 | ||
|
|
3462e57a3b | ||
|
|
21a63c8edf | ||
|
|
80f4da42ad | ||
|
|
4b9cd31dd3 | ||
|
|
f207976281 | ||
|
|
a1b5c177bc | ||
|
|
348b0860dc | ||
|
|
afe7741773 | ||
|
|
730d259976 | ||
|
|
d3dfad90b9 | ||
|
|
a6b2d7c0ec | ||
|
|
3b4858851a | ||
|
|
b45076c5f7 | ||
|
|
44cb495169 | ||
|
|
5a95bb3c0a | ||
|
|
245010f2f6 | ||
|
|
09acb3823a | ||
|
|
43f7a5d108 | ||
|
|
78f4ff7934 | ||
|
|
fcb93b128e | ||
|
|
e5ad2e18c4 | ||
|
|
9d7a2457d9 | ||
|
|
1a8384a8e1 | ||
|
|
5effa47527 | ||
|
|
64059cd7d7 | ||
|
|
49128cf486 | ||
|
|
8f71c08d94 | ||
|
|
cf619fc91b | ||
|
|
cbdd495850 | ||
|
|
6717856212 | ||
|
|
dd9aff58ff | ||
|
|
749cc7ad81 | ||
|
|
93af91419b | ||
|
|
4c3e0deb19 | ||
|
|
7fc617c8ed | ||
|
|
5880f4e6dd | ||
|
|
44f12b2adb | ||
|
|
f9d0d4c5c6 | ||
|
|
49b3d22974 | ||
|
|
6e50f060a3 | ||
|
|
6b4ce58207 | ||
|
|
d99f7ff7ab | ||
|
|
49e77447fb | ||
|
|
6b712e76db | ||
|
|
c19ed5ab32 | ||
|
|
d26263f9c0 | ||
|
|
0dc6ebd525 | ||
|
|
0683fb1579 | ||
|
|
3338d7bc56 | ||
|
|
1cf847f29e | ||
|
|
ed9a4b6b52 | ||
|
|
b8d47cc29b | ||
|
|
b4d8e23abb | ||
|
|
a8ac2af167 | ||
|
|
75953d9f63 | ||
|
|
1036f5664e | ||
|
|
cce914b3fd | ||
|
|
3b6d32c1c4 | ||
|
|
5108be222d | ||
|
|
29a2020f3d | ||
|
|
5b7d65e679 | ||
|
|
a816bf9b84 | ||
|
|
48ece89c63 | ||
|
|
6e55092a4b | ||
|
|
62476e553e | ||
|
|
63ed55a929 | ||
|
|
f3ba1f692b | ||
|
|
e4777915f9 | ||
|
|
5e2938ebf0 | ||
|
|
0d13056d53 | ||
|
|
e26d556594 | ||
|
|
aca383d263 | ||
|
|
667ce1d9a9 | ||
|
|
7f1599e83d | ||
|
|
d9288513e6 | ||
|
|
d73730b545 | ||
|
|
e7aac9b52b | ||
|
|
78958beef1 | ||
|
|
5847f6ee80 | ||
|
|
e957ff1ca2 | ||
|
|
8ebc24194c | ||
|
|
77e2041249 | ||
|
|
6ce3edaa43 | ||
|
|
6ef642f60d | ||
|
|
a3054442e2 | ||
|
|
b13ebeaa1b | ||
|
|
c5796d407f | ||
|
|
2c314d66d6 | ||
|
|
b0e01e9814 | ||
|
|
73e7ef85b6 | ||
|
|
6620aa0937 | ||
|
|
58232a3f04 | ||
|
|
875ff339ba | ||
|
|
737eb6c34c | ||
|
|
2e4c731ea1 | ||
|
|
1f241fb310 | ||
|
|
fc0fce8b7d | ||
|
|
ce3f360219 | ||
|
|
23a2bd5f1b | ||
|
|
53cd640a89 | ||
|
|
a7847b8a37 | ||
|
|
c0af66a7eb | ||
|
|
2d37ee10e6 | ||
|
|
3257fc5cd8 | ||
|
|
53639c4df3 | ||
|
|
f0fa4a7b1f | ||
|
|
1974bc1ea0 | ||
|
|
115aeb47c0 |
No files matched your search
@@ -45,3 +45,26 @@ Rules (both gates):
|
||||
- **Don't weaken the gate:** never hand-lower a baseline or widen a tolerance to turn a red gate green. The ratchet only moves up.
|
||||
- If a change drops coverage, **add tests** (sort `coverage-summary.json` by line% ascending to find untested code) rather than editing the baseline. When coverage legitimately rises, commit the regenerated baseline (`make test-coverage-baseline` / `test-ui-coverage-baseline`).
|
||||
- The Go gate is **strict — no tolerance**; `covermode=atomic` keeps it deterministic. The UI gate keeps a small tolerance only because its e2e coverage isn't.
|
||||
|
||||
## Distributed-mode test suites
|
||||
|
||||
Two suites cover distributed mode (frontend replicas, worker nodes, PostgreSQL), split by a Ginkgo label. Neither starts a message broker: a distributed deployment needs PostgreSQL and the frontends' own HTTP listener, and the `nats-io` modules are not in the build.
|
||||
|
||||
- `make test-e2e-distributed` runs `Distributed && !VLLMMultinode && !Cluster` over `./tests/e2e/distributed` recursively. Services are wired directly into the test binary. 231 specs in about 2m20s (measured 133.3s for the main suite plus the 8 argument-validation specs in the subpackage).
|
||||
- `make test-e2e-cluster` runs `Cluster` and spawns real `local-ai` child processes through the `tests/e2e/distributed/cluster` helper package. It **builds `./local-ai` itself** (the `e2e-binary` target), so you do not run `make build` first and cannot forget to. 26 specs, 933.8s of specs in a 15m37s target run (four measurements; the suite grew from 6 specs and roughly 509s, then 21 specs and 802.5s, then 24 specs and 897 to 907s, so budget on the newest number).
|
||||
|
||||
Both jobs live in `.github/workflows/tests-e2e-distributed.yml`, with `timeout-minutes: 45` each. They trigger on pull requests *and* on every push to `master`; the `paths-ignore` filter (see [.agents/ci-caching.md](ci-caching.md)) sits on the pull-request trigger only, so a master push always runs both. They are advisory only because `master` carries no branch protection, which is a repository setting and not a YAML key: `continue-on-error: true` would flip the run's *conclusion* to success and hide the failure, so it is not used.
|
||||
|
||||
- **The container is suite-scoped, not spec-scoped.** `SetupInfra` used to start a PostgreSQL (~10s) and a broker (~3.5s) per spec. Across the 213 specs behind it that was roughly **48 minutes of pure container startup per run**, which is why this suite was never in CI. (213 rather than the ~240 above: the larger number is everything the label filter selects, the smaller one is just the specs that call `SetupInfra`.) The broker container is gone entirely, and the PostgreSQL one now starts once in `BeforeSuite` with each spec getting its own database via `CREATE DATABASE` (~67ms), which is what the `dbName` argument was always describing. Adding a spec needs no change: call `SetupInfra("some-name")` as before, the name is a prefix and a counter keeps it unique.
|
||||
- **Consequence for new specs:** a fan-out carrier is *shared* within a Ginkgo process, so a wildcard subscriber can observe another spec's traffic. Filter assertions on an identifier your spec owns (a node ID, a job ID) instead of counting everything on `jobs.*.progress`, and verify the spec with `--randomize-all`.
|
||||
- **`BeforeSuite`, not `SynchronizedBeforeSuite`.** Under `ginkgo -p` each process then gets its own container, keeping each process's carrier isolated. A single shared database across parallel processes would let specs on different processes see each other's messages on the same subject.
|
||||
- **The label split.** The 8 argument-validation specs under `tests/e2e/distributed/cluster/` carry `Label("Distributed")` only, on purpose: they need no binary and no PostgreSQL, so they belong in the fast job. That is why `test-e2e-distributed` keeps `-r` (it must reach the subpackage) and `test-e2e-cluster` deliberately does **not** (the subpackage is out of its scope).
|
||||
- **The binary is built, and a stale one is refused.** `test-e2e-cluster` used to take `./local-ai` as given and only check that a file by that name existed. An edit that was never rebuilt therefore ran the entire suite against an older build while reporting on the working tree, which surfaced once as a spec failing on caller line numbers that matched no line in any source file. Two things close it. The target depends on `e2e-binary`, which is `go build -o local-ai ./cmd/local-ai` and not `make build` (that one also builds the React UI, which this suite never touches and CI stubs). And `localAIBinary` in the harness **fails**, on a developer box as well as under CI, when the binary is older than the newest non-test Go source in the tree, which covers a run started with `LOCALAI_E2E_BINARY` or by invoking ginkgo directly. `_test.go` files are excluded from that scan on purpose: they compile into the ginkgo suite and never into `local-ai`, so counting them would fire on every edit to the spec being run. Set `LOCALAI_E2E_BINARY` to opt out of the build; there is no opt-out of the staleness check, because a rebuild is the remedy.
|
||||
- **`--fail-on-empty` is load-bearing on both targets.** Ginkgo exits 0 when a label filter selects nothing, so without it a refactor that renames or drops `Label("Cluster")` leaves the target reporting "Test Suite Passed" having started no cluster at all. `LOCALAI_E2E_REQUIRE_BINARIES` does not cover this case: it only fires inside a spec that is actually running.
|
||||
- **The binary gate.** `localAIBinary()` and `mockBackendBinary()` **fail** rather than skip when `CI` is set, or when `LOCALAI_E2E_REQUIRE_BINARIES` is truthy; `LOCALAI_E2E_REQUIRE_BINARIES=0` (also `off`, `no`, `n`, `disabled`, and anything `strconv.ParseBool` reads as false) forces skipping even under CI. **Any value that parses as neither reads as ON**, not off: setting the variable to something meaningless means someone meant to turn the gate on, and reading it as false would quietly restore the silent skip the flag exists to remove. The whole polarity is deliberate, because in CI a skipped cluster spec is indistinguishable from a passing one: Ginkgo exits 0 on skips. Locally a missing binary still just skips, since `CI` is unset in an ordinary shell.
|
||||
- **Flake budget: no retries at all.** `--flake-attempts` is *total attempts*, not retries (ginkgo v2.29.0 `internal/group.go` sets `maxAttempts = FlakeAttempts` and loops `attempt < maxAttempts`; the flag's own usage string reads "0 - failed tests are not retried"). `DISTRIBUTED_TEST_FLAKES` defaults to **1**, so each spec runs once and a failure is a failure, and `test-e2e-cluster` pins `--flake-attempts 1` outright rather than reading the variable. The repo-wide `TEST_FLAKES=5` means up to five attempts, so up to four retries. These suites exist to surface nondeterminism, and a retry converts exactly that signal into a green run. Raise it locally when bisecting something unrelated, not in the Makefile.
|
||||
- **Coverage:** `tests/e2e/distributed` is excluded from the coverage roots (`COVERAGE_E2E_ROOTS = ./tests/e2e`, run non-recursively), and so is the `cluster` helper package beneath it. Neither suite moves the baseline, so production code that these suites are the only cover for reads as **uncovered**. Unit tests for such code belong under `./core/...` with `testutil.SetupTestDB()`.
|
||||
- **The cluster job builds against a stubbed React UI.** `core/http/react-ui/dist` is gitignored and built by Node, so the workflow writes a one-line `index.html` there to satisfy the `//go:embed react-ui/dist/*` in `core/http/app.go` and skips a full Node and Vite install. That holds only while the suite drives the HTTP API and never the UI, which has its own e2e suite. A spec that ever asserts on a UI asset would pass locally, where a real `dist/` exists, and be served the stub in CI: if you write one, the stub step has to go and the real build come back.
|
||||
- **The topology that matters is two frontends and two workers.** `cluster_busless_test.go` runs the whole programme on it, and it is the only shape in which the OWNER path and the RELAY path are live at the same time: with the two workers' tunnels landing on different replicas, every request either goes to the replica holding that worker's tunnel or is relayed through the one that does. A one-worker cluster cannot tell those apart, and cannot state that killing a replica re-homes only ITS worker. Which worker served a request is read back from `node_models` rather than assumed, and the two models are pinned to one worker each through the admin API (`PUT /api/nodes/:id/labels` plus `POST /api/nodes/scheduling`), so a selector that failed open reddens the spec instead of hiding in it.
|
||||
- **The fan-out carrier is proven under real binaries by `cluster_fanout_test.go`.** Both specs run two frontends and no workers against one PostgreSQL, publish at frontend 0 and read at frontend 1, and cost 7.0s together. They drive a gallery operation on purpose: `GET /models/jobs/<id>` answers from `galleryop`'s in-memory statuses map, which on a peer replica is filled by the `gallery.*.progress` subscriber and by nothing else (`Hydrate` runs once at startup, long before these operations exist), so a peer that answers was told over the carrier. Every other broadcast family has a durable table behind it that a peer would converge through anyway, and a spec on one of those cannot separate "the broadcast arrived" from "the row was read". The specs make that checkable rather than argued: the `gallery_operations` row is written when the gallery worker DEQUEUES an operation, so an operation still queued has none, and both specs assert zero rows while the peer is already answering with the operation's own bytes. Holding the queue is what `cluster.Options.Galleries` is for: it points every frontend at one gated index and turns the startup estimate warmer off, so the operation the spec submits is parked inside a fetch only the spec can complete. The second spec adds the oversized path: an element name of 9000 bytes puts the encoded notification past PostgreSQL's 8000-byte cap (exclusive), so the broadcast becomes a `bus_messages` row plus a notification carrying its id, and the spec asserts the row exists, decodes it, and compares its element name byte for byte with what frontend 1 answers. A one-byte control operation in the same run is asserted to leave NO row, which is what keeps "a row exists" a statement about size rather than one that is true of every operation.
|
||||
- **Do not shorten the cluster suite's waits.** Three of its six specs sit at ~167s each because they wait out a 60s staleness threshold plus a 15s health-check tick. That wait is what stops the assertions from passing before the system could have reacted, which was a real false green earlier on. If the job has to get faster, the levers are CI concurrency or making the thresholds configurable, not shorter waits.
|
||||
@@ -153,7 +153,7 @@ This is worth more than it looks. Measured over the week to 2026-07-30, **97% of
|
||||
|
||||
The volume is real: 13 gallery-only PRs merged that week with 10 open at once, and 78 of the 137 PRs opened were bot-generated.
|
||||
|
||||
`paths-ignore` on the PR trigger of `image-pr.yml` (7 jobs), `build-test.yaml` (3), `lint.yml` (2) and `tests-e2e.yml` (1) drops 13 of those 20. The excluded set:
|
||||
`paths-ignore` on the PR trigger of `image-pr.yml` (7 jobs), `build-test.yaml` (3), `lint.yml` (2) and `tests-e2e.yml` (1) drops 13 of those 20, measured before `tests-e2e-distributed.yml` (2 jobs) landed. That workflow carries the same exclusion set for the same reason: its dependency graph is 99 packages, so an allowlist of paths would silently stop guarding the moment code moved, while a diff confined to the paths below provably cannot reach it. The excluded set:
|
||||
|
||||
| Path | Why no image or Go build can see it |
|
||||
|---|---|
|
||||
@@ -192,7 +192,7 @@ What still runs, and why it has to:
|
||||
Two properties this relies on:
|
||||
|
||||
- `paths-ignore` skips a run only when **every** changed file matches, so a PR touching the gallery *and* Go code still runs everything. That is what makes the exclusion safe rather than a hole.
|
||||
- `master` carries no branch protection and no rulesets, so a skipped workflow reports no status and nothing waits on it. If required status checks are ever introduced, these four entries must be excluded from the required set or PRs will hang on "Expected — Waiting for status to be reported".
|
||||
- `master` carries no branch protection and no rulesets, so a skipped workflow reports no status and nothing waits on it. If required status checks are ever introduced, these five entries must be excluded from the required set or PRs will hang on "Expected — Waiting for status to be reported".
|
||||
|
||||
### `image.yml` on master push is gated too, by a job rather than a path filter
|
||||
|
||||
|
||||
@@ -94,4 +94,4 @@ The user chose KISS. Every mutating tool has a safety rule (`prompts/10_safety.m
|
||||
|
||||
## Distributed mode
|
||||
|
||||
The in-memory MCP server runs only on the head node (where the chat handler runs). `inproc.Client` wraps services that are already distributed-aware (`GalleryService` coordinates with workers; `ListNodes` reads the NATS-populated registry). No NATS routing of MCP tools — the admin surface lives on the head, period.
|
||||
The in-memory MCP server runs only on the head node (where the chat handler runs). `inproc.Client` wraps services that are already distributed-aware (`GalleryService` coordinates with workers; `ListNodes` reads the PostgreSQL-backed node registry). MCP tools are not routed anywhere: the admin surface lives on the head, period.
|
||||
@@ -0,0 +1,206 @@
|
||||
---
|
||||
name: 'E2E Distributed Tests'
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
# The suite's dependency graph is 99 packages, so an allowlist of paths
|
||||
# silently stops guarding the moment code moves. At ~75s the job 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 because it is deliberately not in branch protection, so a failure
|
||||
# is a visible red X rather than a blocked merge. Promoting it to a required
|
||||
# check is a repository-settings change, to be made once it has a track
|
||||
# record; a heavy suite made required on day one gets disabled instead of
|
||||
# fixed.
|
||||
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.0'
|
||||
cache: false
|
||||
- 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
|
||||
# Pulling here rather than inside the suite keeps container-start timing
|
||||
# out of the spec timeouts and makes a registry outage read as a
|
||||
# setup failure instead of a test failure. This is the only image the
|
||||
# suite needs once the testcontainers reaper is disabled below: the
|
||||
# suite stands up no message broker, because nothing under test dials
|
||||
# one.
|
||||
run: |
|
||||
docker pull postgres:16-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 a third, unpinned image (testcontainers/ryuk)
|
||||
# from Docker Hub mid-suite: exactly the registry dependency the
|
||||
# pre-pull step exists to remove.
|
||||
env:
|
||||
TESTCONTAINERS_RYUK_DISABLED: "true"
|
||||
run: |
|
||||
PATH="$PATH:$HOME/go/bin" make test-e2e-distributed
|
||||
- name: Setup tmate session if tests fail
|
||||
if: ${{ failure() }}
|
||||
uses: mxschmitt/action-tmate@v3.23
|
||||
with:
|
||||
detached: true
|
||||
connect-timeout-seconds: 180
|
||||
limit-access-to-actor: true
|
||||
|
||||
tests-e2e-cluster:
|
||||
runs-on: ubuntu-latest
|
||||
# Advisory for the same reason as the job above: master has no branch
|
||||
# protection, so a failure here is a visible red X rather than a blocked
|
||||
# merge. That is a repository-settings property, not a YAML key. The key
|
||||
# that looks like it says "advisory" instead flips the run's conclusion to
|
||||
# success, which hides the failure rather than flagging it, so it appears in
|
||||
# none of this repo's workflows and must not be added here.
|
||||
#
|
||||
# Separate job from tests-e2e-distributed so the fast in-process suite is
|
||||
# not held behind a Go build of local-ai. Serial on purpose: each Ginkgo
|
||||
# process would get its own PostgreSQL container and each spec spawns two or
|
||||
# three local-ai children, so --procs on an unmeasured runner is a change to
|
||||
# make with numbers, not by default.
|
||||
#
|
||||
# The two timeouts bound different things and are not alternatives. Ginkgo's
|
||||
# --timeout=20m bounds the SUITE only; this job timeout must additionally
|
||||
# cover setup, which here is the larger and more variable half: submodule
|
||||
# checkout, apt, protoc plus two go installs plus protogen-go, a cold-cache
|
||||
# module download (cache: false), a full go build of ./cmd/local-ai, and a
|
||||
# separate ginkgo test compile. That build alone is ~316s of CPU, so on a
|
||||
# 4-vCPU runner setup is realistically 8-12 minutes.
|
||||
#
|
||||
# 45 minutes therefore, matching the sibling job. A tighter number does not
|
||||
# make a hang fail faster, it just moves the kill from Ginkgo, which prints
|
||||
# which spec hung, to the runner, which prints nothing: a red job with no
|
||||
# evidence, which is how a suite gets disabled rather than fixed.
|
||||
#
|
||||
# The suite itself is about 8m30s over three consecutive runs (509.1s /
|
||||
# 509.8s / 512.3s, so 8m29s to 8m32s) on a developer box, and will be slower
|
||||
# here. Three specs sit at ~167s each because they wait out a 60s staleness
|
||||
# threshold plus a 15s health-check tick (HealthCheckInterval, in
|
||||
# core/config/distributed_config.go; core/services/nodes/health.go runs the
|
||||
# ticker on the unexported checkInterval, not one of the reconcilers). Do
|
||||
# not shorten those windows to make this job faster: the wait is what stops
|
||||
# the assertions from passing before the system could have reacted, which
|
||||
# was a real false green earlier on.
|
||||
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.0'
|
||||
cache: false
|
||||
- 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: Stub the embedded React UI
|
||||
# core/http/react-ui/dist is gitignored and built by Node, but this
|
||||
# suite drives the HTTP API and never the UI, which has its own e2e
|
||||
# suite. A single index.html satisfies the //go:embed react-ui/dist/*
|
||||
# in core/http/app.go, so the job skips a full Node and Vite install.
|
||||
# If a cluster spec ever asserts on a UI asset, this step must go and
|
||||
# the real build come back: a developer box has a real dist/, so such a
|
||||
# spec would pass locally and fail only here, or worse be served the
|
||||
# stub and pass in both places.
|
||||
run: |
|
||||
mkdir -p core/http/react-ui/dist
|
||||
printf '<!doctype html><title>stub</title>\n' > core/http/react-ui/dist/index.html
|
||||
- name: Pre-pull test images
|
||||
# Same reasoning as the job above: pulling here keeps container-start
|
||||
# timing out of the spec timeouts and makes a registry outage read as a
|
||||
# setup failure rather than a test failure.
|
||||
run: |
|
||||
docker pull postgres:16-alpine
|
||||
- name: Cluster E2E
|
||||
env:
|
||||
# No LOCALAI_E2E_BINARY and no separate build step: make test-e2e-cluster
|
||||
# builds ./local-ai itself, AFTER the protogen-go it also depends on.
|
||||
# Building it in a step of its own put the generated .pb.go files newer
|
||||
# than the binary, which the harness's staleness check reads, correctly,
|
||||
# as a binary that does not contain the tree it is about to be measured
|
||||
# against. One owner for the build is also what makes a local run and
|
||||
# this job exercise the same bytes.
|
||||
#
|
||||
# Must live under the workspace so the upload step below can reach it.
|
||||
# The harness defaults to GinkgoT().TempDir(), which lands under
|
||||
# TMPDIR and would leave the artifact glob matching nothing.
|
||||
LOCALAI_E2E_LOG_DIR: ${{ github.workspace }}/cluster-logs
|
||||
# Belt and braces: the harness already fails rather than skips when CI
|
||||
# is set, and GitHub Actions always sets CI. Stating it here means a
|
||||
# future edit to that default cannot silently turn this job into one
|
||||
# that passes without ever starting a cluster, since a skipped cluster
|
||||
# spec is indistinguishable from a passing one.
|
||||
LOCALAI_E2E_REQUIRE_BINARIES: "true"
|
||||
# See the job above: the runner is ephemeral, so the reaper buys
|
||||
# nothing and would pull a third, unpinned Docker Hub image mid-suite.
|
||||
TESTCONTAINERS_RYUK_DISABLED: "true"
|
||||
run: |
|
||||
PATH="$PATH:$HOME/go/bin" make test-e2e-cluster
|
||||
- name: Upload process logs
|
||||
# The per-process logs are the only way to read a cluster failure: the
|
||||
# Ginkgo output says which assertion failed, not what the four child
|
||||
# processes were doing. Without this a red job is undebuggable.
|
||||
if: ${{ failure() }}
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: cluster-process-logs
|
||||
path: cluster-logs/**/*.log
|
||||
if-no-files-found: ignore
|
||||
retention-days: 7
|
||||
- name: Setup tmate session if tests fail
|
||||
if: ${{ failure() }}
|
||||
uses: mxschmitt/action-tmate@v3.23
|
||||
with:
|
||||
detached: true
|
||||
connect-timeout-seconds: 180
|
||||
limit-access-to-actor: true
|
||||
@@ -19,7 +19,22 @@ linters:
|
||||
- staticcheck
|
||||
enable:
|
||||
- forbidigo
|
||||
# gocritic is enabled for ONE checker: ruleguard, which runs the rules in
|
||||
# hack/lint/. Every other gocritic check is off (disable-all below), so
|
||||
# this adds no style noise; it is here purely as the gate that catches a
|
||||
# gRPC backend wrapper written without Unwrap. See
|
||||
# hack/lint/backend_wrappers.go for why that cannot be a compile-time
|
||||
# assertion.
|
||||
- gocritic
|
||||
settings:
|
||||
gocritic:
|
||||
disable-all: true
|
||||
enabled-checks:
|
||||
- ruleguard
|
||||
settings:
|
||||
ruleguard:
|
||||
failOn: all
|
||||
rules: '${base-path}/hack/lint/backend_wrappers.go'
|
||||
forbidigo:
|
||||
forbid:
|
||||
- pattern: '^t\.Errorf$'
|
||||
@@ -126,3 +141,10 @@ linters:
|
||||
- path: ^backend/go/whisper/sources/
|
||||
text: 'http\.(DefaultClient|Get|Post|PostForm|Head)'
|
||||
linters: [forbidigo]
|
||||
# Test doubles embed grpc.Backend to inherit the interface's method set
|
||||
# over a NIL value; they decorate nothing, hold no inner client, and have
|
||||
# no transport answer to forward. The rule targets production wrappers,
|
||||
# which is where swallowing that answer deletes replica rows.
|
||||
# gocritic here is only the backend-wrapper ruleguard rule.
|
||||
- path: _test\.go$
|
||||
linters: [gocritic]
|
||||
@@ -265,6 +265,37 @@ The e2e tests run LocalAI in a Docker container and exercise the API:
|
||||
make test-e2e
|
||||
```
|
||||
|
||||
### Running distributed-mode tests
|
||||
|
||||
Distributed mode (several frontend replicas, worker nodes and PostgreSQL) has two suites. Both bring up their PostgreSQL with testcontainers, so Docker has to be available. There is no message broker to bring up: a distributed deployment needs PostgreSQL and the frontends' own HTTP listener, and nothing else.
|
||||
|
||||
```bash
|
||||
make test-e2e-distributed # in-process: services wired directly into the test binary
|
||||
make test-e2e-cluster # process-level: real local-ai child processes
|
||||
```
|
||||
|
||||
`make test-e2e-distributed` is the fast one (231 specs in roughly 2m20s). It starts one PostgreSQL for the whole run and gives each spec its own database. It runs each spec exactly once, with no retry: `DISTRIBUTED_TEST_FLAKES` defaults to 1 and feeds ginkgo's `--flake-attempts`, which counts *total attempts*, not retries. That is deliberately below the repo-wide `TEST_FLAKES=5`, because this suite exists to catch nondeterministic cluster behaviour and a retry hides exactly the failure it is meant to catch. Raise it locally when bisecting something unrelated.
|
||||
|
||||
`make test-e2e-cluster` runs `local-ai` as real child processes, one per frontend replica and one per worker, so a spec can kill a replica and assert what the survivors do. Budget about 13m30s (one measurement: 802.5s of specs in a 13m26s target run, which includes building the binary): several of its 21 specs wait out real staleness and health-check windows. An earlier record of 6 specs and roughly 509s is superseded. It needs a built binary and the mock backend:
|
||||
|
||||
```bash
|
||||
make build build-mock-backend
|
||||
make test-e2e-cluster
|
||||
```
|
||||
|
||||
Two environment variables steer it:
|
||||
|
||||
| Variable | Purpose |
|
||||
|---|---|
|
||||
| `LOCALAI_E2E_BINARY` | path to the `local-ai` binary (default: `local-ai` in the repository root) |
|
||||
| `LOCALAI_E2E_LOG_DIR` | directory for the per-process logs (default: a Ginkgo temp dir) |
|
||||
|
||||
Set `LOCALAI_E2E_LOG_DIR` when debugging. A cluster failure is unreadable without the individual frontend and worker logs, and Ginkgo only tells you which assertion failed.
|
||||
|
||||
A missing binary skips the cluster specs locally but fails them whenever `CI` is set, so a build problem cannot turn the CI job green without ever starting a cluster. `LOCALAI_E2E_REQUIRE_BINARIES=1` forces that failing behaviour anywhere; `LOCALAI_E2E_REQUIRE_BINARIES=0` forces the skip back on even under CI.
|
||||
|
||||
Both suites run in `.github/workflows/tests-e2e-distributed.yml`, on pull requests and on every push to `master`. The `paths-ignore` filter is on the pull-request trigger only, so a master push always runs them.
|
||||
|
||||
### React UI tests and coverage
|
||||
|
||||
The React UI (`core/http/react-ui/`) is covered by Playwright e2e specs, gated by a **monotonic line-coverage ratchet** (`make test-ui-coverage-check`, run in CI). The metric is non-deterministic — a fast local box reads higher than a slow CI runner for the same code — so a small tolerance is unavoidable.
|
||||
|
||||
@@ -340,12 +340,108 @@ run-e2e-aio: protogen-go
|
||||
@echo 'Running e2e AIO tests'
|
||||
$(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts $(TEST_FLAKES) -v -r ./tests/e2e-aio
|
||||
|
||||
# Distributed architecture e2e (PostgreSQL + NATS via testcontainers).
|
||||
# Includes NatsJWT specs (JWT-enabled NATS). Requires Docker.
|
||||
# Total ginkgo attempts per spec for the distributed suite: --flake-attempts counts
|
||||
# attempts, not retries. Defaults to 1, so each spec runs once and is never retried,
|
||||
# unlike TEST_FLAKES=5. This suite exists to catch nondeterministic cluster behaviour,
|
||||
# and a retry hides exactly the failures it is meant to surface. Raise it locally if
|
||||
# you are bisecting something unrelated.
|
||||
DISTRIBUTED_TEST_FLAKES?=1
|
||||
|
||||
# Distributed architecture e2e (PostgreSQL via testcontainers). Requires Docker.
|
||||
# There is no broker container and no NatsJWT label any more: a distributed
|
||||
# deployment needs PostgreSQL and the frontends' own HTTP listener, and this
|
||||
# suite stands up exactly that.
|
||||
# VLLMMultinode is excluded here; use test-e2e-vllm-multinode for that.
|
||||
# Cluster is excluded too and runs in test-e2e-cluster below, which needs a
|
||||
# built binary. The argument-validation specs under tests/e2e/distributed/cluster
|
||||
# carry Label("Distributed") only, so they run here and not there, on purpose.
|
||||
# -r stays because of those: they are in a subpackage this target must reach.
|
||||
# --fail-on-empty because ginkgo exits 0 when a label filter matches nothing, so
|
||||
# without it a rename of the label would turn this target into a silent no-op
|
||||
# that still reports "Test Suite Passed".
|
||||
#
|
||||
# This target needs NO built binary and therefore has no stale-binary hole:
|
||||
# ginkgo compiles the suite from the working tree on every run, and everything
|
||||
# these specs exercise runs in that process against containers. The one target
|
||||
# in this file that does exec a prebuilt binary is test-e2e-cluster below, and
|
||||
# the reason it now BUILDS it is written there.
|
||||
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
|
||||
@echo 'Running distributed e2e tests (label Distributed)'
|
||||
$(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --label-filter='Distributed && !VLLMMultinode && !Cluster' --fail-on-empty --flake-attempts $(DISTRIBUTED_TEST_FLAKES) --timeout=40m -v -r ./tests/e2e/distributed
|
||||
|
||||
# The local-ai binary the cluster e2e execs as its frontend and worker
|
||||
# processes.
|
||||
#
|
||||
# Deliberately not `make build`: that target also builds the embedded React UI,
|
||||
# which this suite never touches, and CI stubs core/http/react-ui/dist for
|
||||
# exactly that reason. What is left is the go build, which is incremental and
|
||||
# costs seconds when nothing changed, so making it a prerequisite of the suite
|
||||
# is cheap enough that nobody has a reason to route around it.
|
||||
#
|
||||
# LOCALAI_E2E_BINARY means the caller has its own binary and owns keeping it
|
||||
# current, so this builds nothing. The staleness check in the harness still
|
||||
# applies to it, which is the point: a binary this target did not produce is
|
||||
# precisely the one nothing else vouches for.
|
||||
.PHONY: e2e-binary
|
||||
e2e-binary: protogen-go
|
||||
@if [ -n "$$LOCALAI_E2E_BINARY" ]; then \
|
||||
echo "LOCALAI_E2E_BINARY=$$LOCALAI_E2E_BINARY is set; not building $(BINARY_NAME)"; \
|
||||
else \
|
||||
echo "Building $(BINARY_NAME) for the cluster e2e"; \
|
||||
CGO_LDFLAGS="$(CGO_LDFLAGS)" $(GOCMD) build -ldflags "$(LD_FLAGS)" -tags "$(GO_TAGS)" -o $(BINARY_NAME) ./cmd/local-ai; \
|
||||
fi
|
||||
|
||||
# Cluster e2e: runs local-ai as real child processes (frontend replicas +
|
||||
# workers) against PostgreSQL, and kills them to assert failover.
|
||||
# It BUILDS that binary rather than checking that a file by that name exists,
|
||||
# and that is a correctness fix rather than a convenience. This target used to
|
||||
# take ./local-ai as given, so an edit to core/ that was never rebuilt left the
|
||||
# whole suite exercising an older build while reporting on the working tree. The
|
||||
# failure mode is not theoretical and not loud: it surfaced as a spec failing on
|
||||
# caller line numbers that matched no line in any source file, and it would
|
||||
# otherwise have surfaced as a green run that proved nothing. build-mock-backend
|
||||
# already rebuilds unconditionally, so ./local-ai was the only gap.
|
||||
#
|
||||
# The harness carries the other half. Building here cannot help a run started
|
||||
# with LOCALAI_E2E_BINARY or by invoking ginkgo directly, so localAIBinary also
|
||||
# FAILS when the binary it is about to exec is older than the newest non-test Go
|
||||
# source in the tree. Together: this target makes the common path correct, and
|
||||
# the harness makes every other path loud.
|
||||
#
|
||||
# The argument-validation specs in tests/e2e/distributed/cluster deliberately
|
||||
# stay in test-e2e-distributed above: they need no binary and no PostgreSQL, so
|
||||
# no -r here and that package is simply out of scope.
|
||||
#
|
||||
# --fail-on-empty is load-bearing, not tidiness. Ginkgo exits 0 when a label
|
||||
# filter selects nothing, so without it a refactor that renames or drops
|
||||
# Label("Cluster") leaves this target reporting "Test Suite Passed" having
|
||||
# started no cluster at all. LOCALAI_E2E_REQUIRE_BINARIES does not cover this:
|
||||
# it only fires inside a spec that is actually running.
|
||||
#
|
||||
# --flake-attempts is pinned to 1 rather than $(DISTRIBUTED_TEST_FLAKES), and
|
||||
# should stay there: this suite exists to catch nondeterministic cluster
|
||||
# behaviour, and a retry turns exactly that signal into a green run.
|
||||
#
|
||||
# Budget: 26 specs, measured at 933.8 seconds of Ginkgo time (15m37s wall
|
||||
# including the compile) on a fast developer box. It was 591 to 612 seconds
|
||||
# before the phase 3 control-plane specs and 800 to 830 after them; the three
|
||||
# two-frontend two-worker specs in cluster_busless_test.go added 118 to 127
|
||||
# seconds (3s, 70s and 46 to 53s), nearly all of it in the churn spec, which
|
||||
# cannot be shortened: it waits for a killed replica to leave the live set
|
||||
# (cluster.InstanceLiveness is 30s, measured at 27s) before it may assert
|
||||
# anything, and then holds a window inside the reconnect grace. The two
|
||||
# cross-replica fan-out specs added 7.0 seconds of spec time (5.0s and 2.0s):
|
||||
# they run two frontends and no workers, so they pay for no registration, and
|
||||
# what they wait on is a broadcast rather than a threshold.
|
||||
#
|
||||
# --timeout is 30m rather than 20m because of that. The margin is not slack: a
|
||||
# Ginkgo timeout kills the suite mid-spec and reports a spec name rather than a
|
||||
# cause, and 20m on a loaded CI runner was one slow health tick away from
|
||||
# turning a green suite into an unreadable red one. At 15 minutes measured, 30m
|
||||
# is still twice the budget.
|
||||
test-e2e-cluster: protogen-go build-mock-backend e2e-binary
|
||||
@echo 'Running cluster e2e tests (label Cluster, real local-ai processes)'
|
||||
$(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --label-filter='Cluster' --fail-on-empty --flake-attempts 1 --timeout=30m -v ./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
|
||||
@@ -383,11 +479,13 @@ prepare-e2e:
|
||||
run-e2e-image:
|
||||
docker run -p 5390:8080 -e MODELS_PATH=/models -e THREADS=1 -e DEBUG=true -d --rm -v $(TEST_DIR):/models --name e2e-tests-$(RANDOM) localai-tests
|
||||
|
||||
# Distributed subpackages have separate targets that provision PostgreSQL and
|
||||
# build the host binaries; this target exercises the container-backed API only.
|
||||
test-e2e: build-mock-backend build-cloud-proxy-backend prepare-e2e run-e2e-image
|
||||
@echo 'Running e2e tests'
|
||||
BUILD_TYPE=$(BUILD_TYPE) \
|
||||
LOCALAI_API=http://$(E2E_BRIDGE_IP):5390 \
|
||||
$(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts $(TEST_FLAKES) -v -r ./tests/e2e
|
||||
$(GOCMD) run github.com/onsi/ginkgo/v2/ginkgo --flake-attempts $(TEST_FLAKES) -v ./tests/e2e
|
||||
$(MAKE) clean-mock-backend
|
||||
$(MAKE) clean-cloud-proxy-backend
|
||||
$(MAKE) teardown-e2e
|
||||
|
||||
@@ -209,7 +209,7 @@ For older news and full release notes, see [GitHub Releases](https://github.com/
|
||||
- [Object Detection](https://localai.io/features/object-detection/)
|
||||
- [Reranker API](https://localai.io/features/reranker/)
|
||||
- [P2P Inferencing](https://localai.io/features/distribute/)
|
||||
- [Distributed Mode](https://localai.io/features/distributed-mode/) — Horizontal scaling with PostgreSQL + NATS
|
||||
- [Distributed Mode](https://localai.io/features/distributed-mode/): horizontal scaling on PostgreSQL, with no message broker to run
|
||||
- [Model Context Protocol (MCP)](https://localai.io/docs/features/mcp/)
|
||||
- [Built-in Agents](https://localai.io/features/agents/) — Autonomous AI agents with tool use, RAG, skills, SSE streaming, and [Agent Hub](https://agenthub.localai.io)
|
||||
- [Backend Gallery](https://localai.io/backends/) — Install/remove backends on the fly via OCI images
|
||||
|
||||
@@ -9,7 +9,7 @@
|
||||
# recipe is a make target (not a prepare.sh) so 'make purge && make' is a clean
|
||||
# rebuild and so the bump bot can see the pin.
|
||||
|
||||
AUDIO_CPP_VERSION?=9c6a282337cc83f227cc10428867a478947706ad
|
||||
AUDIO_CPP_VERSION?=f6277c1695a83cf388a8282c1c1a8757cf626f18
|
||||
AUDIO_CPP_REPO?=https://github.com/0xShug0/audio.cpp
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
# ds4 backend Makefile.
|
||||
#
|
||||
# Upstream pin lives below as DS4_VERSION?=f62ca29a308724cde5bc99134ede19104b2a3260
|
||||
# Upstream pin lives below as DS4_VERSION?=9ab705347c1775e7599ede7eb81a6255ec7dccb5
|
||||
# (.github/bump_deps.sh) can find and update it - matches the
|
||||
# llama-cpp / ik-llama-cpp / turboquant convention.
|
||||
|
||||
DS4_VERSION?=f62ca29a308724cde5bc99134ede19104b2a3260
|
||||
DS4_VERSION?=9ab705347c1775e7599ede7eb81a6255ec7dccb5
|
||||
DS4_REPO?=https://github.com/antirez/ds4
|
||||
|
||||
CURRENT_MAKEFILE_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST))))
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
|
||||
LLAMA_VERSION?=67672dc5b76f8bc17785a19d3dc6d1463fc2902c
|
||||
LLAMA_VERSION?=465e49b9cea78a68b9c244ffb48d0ee24a82873d
|
||||
LLAMA_REPO?=https://github.com/ggerganov/llama.cpp
|
||||
|
||||
CMAKE_ARGS?=
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# CrispASR version (release tag)
|
||||
CRISPASR_REPO?=https://github.com/CrispStrobe/CrispASR
|
||||
CRISPASR_VERSION?=301acd87b036764973b8bfba71e0a21818036d33
|
||||
CRISPASR_VERSION?=c6d4265ae2ee2b8931b09d7d25d5c65c75c36a41
|
||||
SO_TARGET?=libgocrispasr.so
|
||||
|
||||
CMAKE_ARGS+=-DBUILD_SHARED_LIBS=OFF
|
||||
|
||||
@@ -14,7 +14,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
# It is kept alive by the upstream tag da2-support (survives a squash-merge);
|
||||
# repoint to the master merge commit once mudler/depth-anything.cpp PR #1 lands.
|
||||
DEPTHANYTHING_REPO?=https://github.com/mudler/depth-anything.cpp.git
|
||||
DEPTHANYTHING_VERSION?=14f7461d1f704761a038ac9f50dbde8fdb7275e2
|
||||
DEPTHANYTHING_VERSION?=02ba082274e001a63e50de5a1eb0ccc50c6af4b1
|
||||
|
||||
ifeq ($(NATIVE),false)
|
||||
CMAKE_ARGS+=-DGGML_NATIVE=OFF
|
||||
|
||||
@@ -8,7 +8,7 @@ JOBS?=$(shell nproc --ignore=1)
|
||||
|
||||
# stablediffusion.cpp (ggml)
|
||||
STABLEDIFFUSION_GGML_REPO?=https://github.com/leejet/stable-diffusion.cpp
|
||||
STABLEDIFFUSION_GGML_VERSION?=d04e8950c1ec8d30248cbe996682b3182fb1adf6
|
||||
STABLEDIFFUSION_GGML_VERSION?=d8fb10c02977c8ca999f3fb4e02df9ecf10f7ba6
|
||||
|
||||
CMAKE_ARGS+=-DGGML_MAX_NAME=128
|
||||
|
||||
|
||||
@@ -122,21 +122,6 @@ from diffusers.schedulers import (
|
||||
UniPCMultistepScheduler,
|
||||
)
|
||||
|
||||
def select_device(request_cuda, device_option, cuda_available, xpu, mps_available):
|
||||
"""Pick the pipeline device. An explicit `device:` model option wins;
|
||||
otherwise CUDA is used whenever torch reports it available (ROCm
|
||||
builds included) or the model config forces it with `cuda: true`,
|
||||
keeping the pre-existing XPU/MPS overrides. CPU is the fallback, not
|
||||
the default."""
|
||||
if device_option:
|
||||
return device_option
|
||||
device = "cuda" if (request_cuda or cuda_available) else "cpu"
|
||||
if xpu:
|
||||
device = "xpu"
|
||||
if mps_available:
|
||||
device = "mps"
|
||||
return device
|
||||
|
||||
def is_float(s):
|
||||
"""Check if a string can be converted to float."""
|
||||
try:
|
||||
@@ -642,13 +627,12 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
|
||||
# modify LoraAdapter to be relative to modelFileBase
|
||||
request.LoraAdapter = os.path.join(request.ModelPath, request.LoraAdapter)
|
||||
|
||||
device = select_device(
|
||||
request.CUDA,
|
||||
self.options.pop("device", None),
|
||||
torch.cuda.is_available(),
|
||||
XPU,
|
||||
hasattr(torch.backends, "mps") and torch.backends.mps.is_available(),
|
||||
)
|
||||
device = "cpu" if not request.CUDA else "cuda"
|
||||
if XPU:
|
||||
device = "xpu"
|
||||
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
||||
if mps_available:
|
||||
device = "mps"
|
||||
self.device = device
|
||||
if request.LoraAdapter:
|
||||
# Check if its a local file and not a directory ( we load lora differently for a safetensor file )
|
||||
|
||||
@@ -7,7 +7,6 @@ import time
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
# Import dynamic loader for testing (these don't need gRPC)
|
||||
import backend
|
||||
import diffusers_dynamic_loader as loader
|
||||
from diffusers import DiffusionPipeline, StableDiffusionPipeline
|
||||
|
||||
@@ -426,22 +425,3 @@ class TestGenerateImageOptionsKwargsMerge(unittest.TestCase):
|
||||
self.assertEqual(pipeline.kwargs["num_inference_steps"], 4)
|
||||
finally:
|
||||
os.unlink(dst_path)
|
||||
|
||||
|
||||
class TestDeviceSelection(unittest.TestCase):
|
||||
"""Unit tests for backend.select_device (no GPU required)."""
|
||||
|
||||
def test_autodetect_cuda(self):
|
||||
self.assertEqual(backend.select_device(False, None, True, False, False), "cuda")
|
||||
|
||||
def test_cpu_fallback(self):
|
||||
self.assertEqual(backend.select_device(False, None, False, False, False), "cpu")
|
||||
|
||||
def test_forced_cuda(self):
|
||||
self.assertEqual(backend.select_device(True, None, False, False, False), "cuda")
|
||||
|
||||
def test_device_option_wins(self):
|
||||
self.assertEqual(backend.select_device(True, "cpu", True, True, True), "cpu")
|
||||
|
||||
def test_mps_overrides(self):
|
||||
self.assertEqual(backend.select_device(False, None, True, False, True), "mps")
|
||||
@@ -0,0 +1,65 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
)
|
||||
|
||||
// distributedSchedulerOptions stamps the absence wiring onto the scheduler's
|
||||
// options and returns them.
|
||||
//
|
||||
// Two assignments in a named function rather than two more fields in the
|
||||
// twenty-field literal they used to live in. The literal cannot be reached by a
|
||||
// unit spec, because the function that builds it also opens a NATS connection
|
||||
// and a database; these two lines can, and they are the two lines this whole
|
||||
// change comes down to. Losing them in the literal was silent and green.
|
||||
//
|
||||
// The grace comes from the same expression the membership loop is given
|
||||
// (Membership.SetReconnectGrace), so the window a departure is measured against
|
||||
// and the window a departure is RETAINED for cannot drift apart.
|
||||
func distributedSchedulerOptions(cfg config.DistributedConfig, presence nodes.NodePresenceReader, opts nodes.SmartRouterOptions) nodes.SmartRouterOptions {
|
||||
opts.Presence = presence
|
||||
opts.ReconnectGrace = cfg.ReconnectGraceOrDefault()
|
||||
return opts
|
||||
}
|
||||
|
||||
// requireAbsenceWiring refuses to start a distributed deployment in which
|
||||
// nothing can decide that a worker has gone away.
|
||||
//
|
||||
// Two components read absence, from one source and against one window: the
|
||||
// scheduler, which stops placing work on a departed worker, and the health
|
||||
// monitor, which stops reporting one as healthy. Each reads it through a field
|
||||
// assigned in a large construction literal in initDistributed.
|
||||
//
|
||||
// It is checked rather than assumed because losing either assignment is
|
||||
// SILENT. A scheduler with no absence source places work on workers that are
|
||||
// gone and demotes none; a health monitor with none reports a worker whose
|
||||
// tunnel died an hour ago as healthy, forever, with every request for a model
|
||||
// loaded on it failing "no route to that worker". Neither logs anything,
|
||||
// neither fails a request that would not have failed anyway, and both look
|
||||
// exactly like a fleet that is fine. Refusing to boot is the only symptom
|
||||
// either failure has, and it is the reason this is a startup error and not a
|
||||
// warning: a deployment that came up and quietly decided absence by nothing is
|
||||
// the state the tunnel work exists to remove.
|
||||
//
|
||||
// What this guard itself rests on, stated because it is a real limit: the two
|
||||
// helper functions below and above are pinned by unit specs, but the CALL to
|
||||
// this one lives in initDistributed, which opens NATS and a database and so has
|
||||
// no unit spec at all. Deleting the call, or writing a literal nil where
|
||||
// initDistributed passes the cluster registry, compiles and leaves every suite
|
||||
// in this repository green. Only tests/e2e/distributed/cluster catches it, by
|
||||
// booting the real binary: the error is returned from initDistributed and
|
||||
// aborts application startup, so a frontend so wired never comes up.
|
||||
func requireAbsenceWiring(router *nodes.SmartRouter, health *nodes.HealthMonitor) error {
|
||||
if !router.ReadsAbsence() {
|
||||
return fmt.Errorf("the distributed scheduler was built with no source of worker absence: it would place work on workers that have gone away and never demote one")
|
||||
}
|
||||
if !health.ReadsAbsence() {
|
||||
return fmt.Errorf("the node health monitor was built with no source of worker absence: it would report a worker whose tunnel is gone as healthy indefinitely")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,121 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/cluster"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
)
|
||||
|
||||
// presenceStub is any source of absence. What it answers does not matter here:
|
||||
// these specs are about whether a source was wired at all, which is the one
|
||||
// property that has no other symptom.
|
||||
type presenceStub struct{}
|
||||
|
||||
func (presenceStub) Presence(context.Context, string, time.Duration) (cluster.Presence, error) {
|
||||
return cluster.PresenceConnected, nil
|
||||
}
|
||||
|
||||
// The guard on the two lines that connect the absence decision to production.
|
||||
//
|
||||
// Absence is read in exactly two places, and each reads it through one field
|
||||
// assigned in a twenty-field construction literal in initDistributed. Deleting
|
||||
// either assignment compiles, passes every suite in this repository, and
|
||||
// returns the deployment to "absence is decided by nothing" without a log line.
|
||||
// That is the failure this guard exists for, and these specs are what keep the
|
||||
// guard honest: an assertion that never fails is not one.
|
||||
var _ = Describe("stamping the absence wiring onto the scheduler's options", func() {
|
||||
It("gives the scheduler the deployment's source of absence", func() {
|
||||
reg := presenceStub{}
|
||||
|
||||
opts := distributedSchedulerOptions(config.DistributedConfig{}, reg, nodes.SmartRouterOptions{})
|
||||
|
||||
Expect(opts.Presence).To(Equal(nodes.NodePresenceReader(reg)))
|
||||
})
|
||||
|
||||
It("gives it the operator's reconnect grace", func() {
|
||||
opts := distributedSchedulerOptions(
|
||||
config.DistributedConfig{WorkerReconnectGrace: 4 * time.Minute}, presenceStub{}, nodes.SmartRouterOptions{})
|
||||
|
||||
Expect(opts.ReconnectGrace).To(Equal(4 * time.Minute))
|
||||
})
|
||||
|
||||
It("falls back to the documented default when the operator set no grace", func() {
|
||||
opts := distributedSchedulerOptions(config.DistributedConfig{}, presenceStub{}, nodes.SmartRouterOptions{})
|
||||
|
||||
Expect(opts.ReconnectGrace).To(Equal(config.DefaultWorkerReconnectGrace))
|
||||
})
|
||||
|
||||
It("leaves every other option the caller built untouched", func() {
|
||||
// The negative control: a stamp that rebuilt the options would drop the
|
||||
// twenty fields the caller assembled, and the two assertions above
|
||||
// would still pass.
|
||||
opts := distributedSchedulerOptions(config.DistributedConfig{}, presenceStub{},
|
||||
nodes.SmartRouterOptions{GalleriesJSON: "[]", SharedModels: true})
|
||||
|
||||
Expect(opts.GalleriesJSON).To(Equal("[]"))
|
||||
Expect(opts.SharedModels).To(BeTrue())
|
||||
})
|
||||
|
||||
It("produces a scheduler that reads absence", func() {
|
||||
// The property the boot guard checks, asserted through the same call
|
||||
// initDistributed makes.
|
||||
router := nodes.NewSmartRouter(nil,
|
||||
distributedSchedulerOptions(config.DistributedConfig{}, presenceStub{}, nodes.SmartRouterOptions{}))
|
||||
|
||||
Expect(router.ReadsAbsence()).To(BeTrue())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("the absence wiring a distributed deployment refuses to start without", func() {
|
||||
present := func() (*nodes.SmartRouter, *nodes.HealthMonitor) {
|
||||
router := nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{Presence: presenceStub{}})
|
||||
health := nodes.NewHealthMonitor(nil, nil, time.Second, time.Minute, "", false, presenceStub{}, time.Minute, nil)
|
||||
return router, health
|
||||
}
|
||||
|
||||
It("accepts a deployment where both readers have a source", func() {
|
||||
router, health := present()
|
||||
Expect(requireAbsenceWiring(router, health)).To(Succeed())
|
||||
})
|
||||
|
||||
It("refuses a scheduler built without one, and says what it would do", func() {
|
||||
_, health := present()
|
||||
blind := nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{})
|
||||
|
||||
err := requireAbsenceWiring(blind, health)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("scheduler"))
|
||||
Expect(err.Error()).To(ContainSubstring("never demote"))
|
||||
})
|
||||
|
||||
It("refuses a health monitor built without one, and says what it would do", func() {
|
||||
router, _ := present()
|
||||
blind := nodes.NewHealthMonitor(nil, nil, time.Second, time.Minute, "", false, nil, 0, nil)
|
||||
|
||||
err := requireAbsenceWiring(router, blind)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("health monitor"))
|
||||
Expect(err.Error()).To(ContainSubstring("healthy indefinitely"))
|
||||
})
|
||||
|
||||
// Each reader is named separately on purpose. One guard covering "at least
|
||||
// one of them" would accept a deployment that had lost the other, and the
|
||||
// two failures are different: the scheduler's places work on a dead worker,
|
||||
// the monitor's leaves it listed healthy while its models are unreachable.
|
||||
It("names the scheduler and the health monitor as separate requirements", func() {
|
||||
blindRouter := nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{})
|
||||
blindHealth := nodes.NewHealthMonitor(nil, nil, time.Second, time.Minute, "", false, nil, 0, nil)
|
||||
router, health := present()
|
||||
|
||||
Expect(requireAbsenceWiring(blindRouter, health)).ToNot(Succeed())
|
||||
Expect(requireAbsenceWiring(router, blindHealth)).ToNot(Succeed())
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,46 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
)
|
||||
|
||||
// newAgentControl builds the frontend's agent control client: the SELECTION
|
||||
// that decides which agent worker answers a verb, and the control client that
|
||||
// carries the verb to it.
|
||||
//
|
||||
// A named function rather than two more lines in initDistributed, for the
|
||||
// reason distributedSchedulerOptions is one: initDistributed opens a database
|
||||
// and a bus, so no unit spec reaches it, and the argument that matters here has
|
||||
// no symptom when it is wrong. An empty instance id makes every connection read
|
||||
// report nothing as held by this replica, so every MCP call takes a relay hop
|
||||
// through a peer even when this replica holds the worker's tunnel itself, and
|
||||
// nothing anywhere says so: the calls all succeed, just through one more
|
||||
// process than they need. It is refused here instead of shipped as latency.
|
||||
//
|
||||
// The other two refusals are the ordinary kind. A selector with no registry has
|
||||
// nothing to select from and a client with no transport reaches nobody, and
|
||||
// both would present as MCP being quietly unavailable in a deployment that
|
||||
// looks healthy.
|
||||
//
|
||||
// This client is also the deployment's agent CANCELLER, which is why the agent
|
||||
// event bridge takes it: a cancel is a control RPC on the tunnels the workers
|
||||
// hold, and the reconnect grace it is built with is what decides whether a
|
||||
// worker that is not connected makes a cancel undelivered or is simply gone.
|
||||
func newAgentControl(cfg config.DistributedConfig, registry *nodes.NodeRegistry,
|
||||
conns nodes.AgentConnectionReader, control *nodes.ControlClient) (*nodes.AgentControlClient, error) {
|
||||
if cfg.InstanceID == "" {
|
||||
return nil, fmt.Errorf("the agent control client was built with no instance id: every MCP call would relay through a peer even for a worker whose tunnel this replica holds")
|
||||
}
|
||||
if registry == nil || conns == nil {
|
||||
return nil, fmt.Errorf("the agent control client was built with no way to find a connected agent worker")
|
||||
}
|
||||
if control == nil {
|
||||
return nil, fmt.Errorf("the agent control client was built with no control transport to reach an agent worker over")
|
||||
}
|
||||
return nodes.NewAgentControlClient(nodes.NewAgentSelector(registry, conns, cfg.InstanceID, cfg.WorkerReconnectGrace), control), nil
|
||||
}
|
||||
@@ -0,0 +1,148 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/cluster"
|
||||
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/testutil"
|
||||
)
|
||||
|
||||
// recordingConnections captures the owner id the selection was made with. It is
|
||||
// how these specs see the one argument whose loss has no other symptom.
|
||||
//
|
||||
// The channel is what lets a spec observe a selection made on ANOTHER
|
||||
// goroutine, which is what the dispatch loop's own tick is. A slice read from
|
||||
// the spec goroutine would be a data race, and waiting on it would be a sleep.
|
||||
type recordingConnections struct {
|
||||
owners []string
|
||||
seen chan string
|
||||
// graces records the reconnect window each presence read was made with. It
|
||||
// is the other argument with no other symptom: with a zero one every worker
|
||||
// that lost its tunnel a moment ago reads as GONE, and a cancel addressed to
|
||||
// it is reported as a run no worker is running rather than as one this
|
||||
// deployment could not deliver.
|
||||
graces []time.Duration
|
||||
}
|
||||
|
||||
func (r *recordingConnections) ConnectedAmong(_ context.Context, _ []string, owner string) ([]string, []string, error) {
|
||||
r.owners = append(r.owners, owner)
|
||||
if r.seen != nil {
|
||||
select {
|
||||
case r.seen <- owner:
|
||||
default:
|
||||
}
|
||||
}
|
||||
return nil, nil, nil
|
||||
}
|
||||
|
||||
func (r *recordingConnections) Presence(_ context.Context, _ string, grace time.Duration) (cluster.Presence, error) {
|
||||
r.graces = append(r.graces, grace)
|
||||
return cluster.PresenceReconnecting, nil
|
||||
}
|
||||
|
||||
// newRecordingConnections returns a reader whose channel is ready BEFORE any
|
||||
// loop can be started against it. Creating it lazily from the spec goroutine
|
||||
// would race the loop's own goroutine reading it.
|
||||
func newRecordingConnections() *recordingConnections {
|
||||
return &recordingConnections{seen: make(chan string, 8)}
|
||||
}
|
||||
|
||||
// calledBy delivers the owner id of each selection this reader answers.
|
||||
func (r *recordingConnections) calledBy() chan string { return r.seen }
|
||||
|
||||
// The wiring that connects MCP to the agent workers, guarded the way the
|
||||
// absence wiring is and for the same reason: initDistributed opens a database
|
||||
// and a bus, so no unit spec reaches the construction literal, and one of these
|
||||
// arguments is silent when it is wrong.
|
||||
var _ = Describe("building the frontend's agent control client", func() {
|
||||
var registry *nodes.NodeRegistry
|
||||
var ctx context.Context
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
ctx = context.Background()
|
||||
var err error
|
||||
registry, err = nodes.NewNodeRegistry(testutil.SetupTestDB())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
})
|
||||
|
||||
// Every refusal below is given a valid value for everything except the one
|
||||
// argument it is about, so no assertion can be satisfied by a guard that
|
||||
// fires for the wrong reason.
|
||||
It("makes the selection with THIS replica's instance id", func() {
|
||||
// The silent one. With an empty id every connection read reports
|
||||
// nothing as held here, so every MCP call relays through a peer even
|
||||
// for a worker whose tunnel this replica holds: correct answers, one
|
||||
// extra hop, and no log line anywhere.
|
||||
Expect(registry.Register(ctx, &nodes.BackendNode{
|
||||
Name: "agent-1", NodeType: nodes.NodeTypeAgent, Address: "a:50051",
|
||||
}, true)).To(Succeed())
|
||||
|
||||
conns := newRecordingConnections()
|
||||
client, err := newAgentControl(
|
||||
config.DistributedConfig{InstanceID: "replica-7"}, registry, conns,
|
||||
nodes.NewControlClient(nil, "token"))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// Driven through a real call rather than read off a field: what has to
|
||||
// be true is that the id reaches the SELECTION, not that it was stored.
|
||||
_, _ = client.ExecuteMCPTool(ctx, mcpremote.MCPToolRequest{ModelName: "m"})
|
||||
Expect(conns.owners).To(ConsistOf("replica-7"))
|
||||
})
|
||||
|
||||
It("makes the reconnect grace the one a cancel measures a lost tunnel against", func() {
|
||||
// The other silent argument. With a zero grace every registered agent
|
||||
// worker that is not connected right now reads as GONE, so a cancel
|
||||
// addressed to one that is merely re-homing is reported as "no worker
|
||||
// is running that execution" rather than as one that was not delivered.
|
||||
// The only symptom is a wrong answer to a caller.
|
||||
Expect(registry.Register(ctx, &nodes.BackendNode{
|
||||
Name: "agent-grace", NodeType: nodes.NodeTypeAgent, Address: "a:50051",
|
||||
}, true)).To(Succeed())
|
||||
|
||||
conns := newRecordingConnections()
|
||||
client, err := newAgentControl(
|
||||
config.DistributedConfig{InstanceID: "replica-7", WorkerReconnectGrace: 7 * time.Minute},
|
||||
registry, conns, nodes.NewControlClient(nil, "token"))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// Driven through a real cancel, so what is asserted is the value that
|
||||
// reached the presence read rather than one stored on a field.
|
||||
_ = client.CancelAgentRun(ctx, messaging.AgentCancelRequest{AgentName: "a", MessageID: "m"})
|
||||
Expect(conns.graces).To(ConsistOf(7 * time.Minute))
|
||||
})
|
||||
|
||||
It("refuses to build with no instance id", func() {
|
||||
_, err := newAgentControl(config.DistributedConfig{}, registry, newRecordingConnections(),
|
||||
nodes.NewControlClient(nil, "token"))
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("instance id"))
|
||||
})
|
||||
|
||||
It("refuses to build with nothing to read connections through", func() {
|
||||
_, err := newAgentControl(config.DistributedConfig{InstanceID: "replica-7"}, registry, nil,
|
||||
nodes.NewControlClient(nil, "token"))
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("connected agent worker"))
|
||||
})
|
||||
|
||||
It("refuses to build with no control transport", func() {
|
||||
_, err := newAgentControl(config.DistributedConfig{InstanceID: "replica-7"}, registry,
|
||||
newRecordingConnections(), nil)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("control transport"))
|
||||
})
|
||||
})
|
||||
@@ -37,8 +37,11 @@ func (a *Application) RestartAgentJobService() error {
|
||||
if d.JobStore != nil {
|
||||
agentJobService.SetDistributedJobStore(d.JobStore)
|
||||
}
|
||||
// Keep agent tasks consistent across replicas (same client the dispatcher uses).
|
||||
agentJobService.SetTaskSyncNATS(d.Nats)
|
||||
// Keep agent tasks consistent across replicas, on the deployment's
|
||||
// broadcast carrier. This is the restart path and it is a second site
|
||||
// for the same rule: a fix applied only in startup.go leaves every
|
||||
// service the settings UI restarts on whatever carrier it picked here.
|
||||
agentJobService.SetTaskSyncBus(d.Broadcast())
|
||||
}
|
||||
|
||||
// Start the service
|
||||
|
||||
@@ -636,7 +636,6 @@ func (a *Application) StartAgentPool() {
|
||||
if d.DistStores != nil && d.DistStores.Skills != nil {
|
||||
opts.SkillStore = d.DistStores.Skills
|
||||
}
|
||||
opts.NATSClient = d.Nats
|
||||
opts.EventBridge = d.AgentBridge
|
||||
opts.AgentStore = d.AgentStore
|
||||
}
|
||||
@@ -680,7 +679,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.SetJobSyncBus(d.Broadcast())
|
||||
}
|
||||
aps.SetUserServicesManager(usm)
|
||||
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"math"
|
||||
"strings"
|
||||
|
||||
"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/nodes/prefixcache"
|
||||
"github.com/mudler/LocalAI/core/services/pgbus"
|
||||
)
|
||||
|
||||
// The four process-lifetime caches, each wired onto the broadcast carrier by
|
||||
// one function here, and none of them by a call site naming a carrier.
|
||||
//
|
||||
// Two are methods on DistributedServices and take NO carrier at all. The other
|
||||
// two run inside initDistributed before that struct exists, so they take the
|
||||
// CONCRETE *pgbus.Bus rather than messaging.Broadcaster.
|
||||
//
|
||||
// Concrete on purpose, and the purpose has outlived what prompted it. It was
|
||||
// written against a SECOND carrier that was in scope at these call sites and
|
||||
// satisfied the same interface, so an interface parameter let a caller hand
|
||||
// over the wrong one: it compiled, started, published and was delivered, to
|
||||
// nobody the deployment would still be listening on. That carrier went with the
|
||||
// message broker and there is one left. The narrowing stays because a second
|
||||
// carrier is exactly what it guards against, and widening these parameters back
|
||||
// would give the guard away silently on the day another one arrives.
|
||||
//
|
||||
// The adopters themselves still take the interface, so their own specs drive
|
||||
// them with an in-memory double. The narrowing is only here.
|
||||
|
||||
// wireGalleryBroadcasts puts the gallery service's progress, cancel and
|
||||
// cache-invalidation traffic on the carrier, and opens the wildcard
|
||||
// subscriptions that mirror a peer's.
|
||||
//
|
||||
// Set and subscribe in one call because they are one decision: a service that
|
||||
// published where nobody listened would show every operation it started and
|
||||
// none of its peers', which is what /api/operations looks like on a replica
|
||||
// that has been load-balanced away from.
|
||||
//
|
||||
// The caller must have hydrated from the store and bound OnModelsChanged
|
||||
// first; both are stated on the methods themselves.
|
||||
func (ds *DistributedServices) wireGallery(gs *galleryop.GalleryService) error {
|
||||
if ds == nil || ds.Bus == nil {
|
||||
return fmt.Errorf("wiring gallery broadcasts: no broadcast carrier, so gallery progress and cancels would reach no peer replica")
|
||||
}
|
||||
if gs == nil {
|
||||
return nil
|
||||
}
|
||||
gs.SetBroadcaster(ds.Bus)
|
||||
return gs.SubscribeBroadcasts()
|
||||
}
|
||||
|
||||
// WireOpCache puts the admin operation cache on the carrier and starts it,
|
||||
// which hydrates from PostgreSQL and subscribes.
|
||||
//
|
||||
// Exported, unlike its siblings, because the OpCache is built in the HTTP layer
|
||||
// rather than in initDistributed. It takes neither a carrier nor a store: both
|
||||
// come off this struct, so the HTTP layer names no carrier and cannot pass one
|
||||
// that is merely in scope beside them.
|
||||
//
|
||||
// A hydrate failure is the OpCache's own business and is logged there; a
|
||||
// subscribe failure is returned, because a cache that hydrated and did not
|
||||
// subscribe reports the operations it found at boot and never learns of another.
|
||||
func (ds *DistributedServices) WireOpCache(ctx context.Context, cache *galleryop.OpCache) error {
|
||||
if cache == nil {
|
||||
return nil
|
||||
}
|
||||
if ds == nil || ds.Bus == nil {
|
||||
return fmt.Errorf("wiring the operation cache: no broadcast carrier, so /api/operations would answer with whatever this one replica admitted")
|
||||
}
|
||||
cache.SetBroadcaster(ds.Bus)
|
||||
if ds.DistStores != nil && ds.DistStores.Gallery != nil {
|
||||
cache.SetGalleryStore(ds.DistStores.Gallery)
|
||||
}
|
||||
return cache.Start(ctx)
|
||||
}
|
||||
|
||||
// wireStagingBroadcasts puts file-staging progress on the carrier in both
|
||||
// directions. The tracker's own SetBroadcaster is what makes those one call;
|
||||
// see the comment there.
|
||||
func wireStagingBroadcasts(bus *pgbus.Bus, tracker *nodes.StagingTracker) (messaging.Subscription, error) {
|
||||
if bus == nil {
|
||||
return nil, fmt.Errorf("wiring staging broadcasts: no broadcast carrier, so a staging transfer would show a progress bar only on the replica performing it")
|
||||
}
|
||||
if tracker == nil {
|
||||
return nil, nil
|
||||
}
|
||||
return tracker.SetBroadcaster(bus)
|
||||
}
|
||||
|
||||
// wirePrefixCacheBroadcasts builds the cross-frontend prefix-cache layer on the
|
||||
// carrier and subscribes it to peers, after refusing a configuration whose
|
||||
// observations could not travel in a notification.
|
||||
func wirePrefixCacheBroadcasts(bus *pgbus.Bus, cfg prefixcache.Config, idx prefixcache.Provider) (*prefixcache.Sync, error) {
|
||||
// The configuration first, and the carrier second. A depth this carrier
|
||||
// cannot hold is wrong whether or not a carrier was supplied, and naming
|
||||
// the more specific fault is what makes the startup message actionable.
|
||||
if err := requirePrefixCacheFitsInline(cfg); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if bus == nil {
|
||||
return nil, fmt.Errorf("wiring the prefix cache: no broadcast carrier, so each frontend would route on nothing but its own history")
|
||||
}
|
||||
sync := prefixcache.NewSync(idx, bus)
|
||||
if _, err := sync.SubscribeBroadcasts(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return sync, nil
|
||||
}
|
||||
|
||||
// prefixCacheIdentifierAllowance is how many bytes of model id plus node id a
|
||||
// prefix-cache observation is budgeted for when its worst case is checked
|
||||
// against the notification cap.
|
||||
//
|
||||
// Both are operator-chosen strings with no enforced length, so no bound here is
|
||||
// a proof. It does not need to be one: an observation that does not fit is
|
||||
// SPILLED like any other broadcast, at the cost of a row and a SELECT, and is
|
||||
// never lost. What the check exists to catch is the other failure, the one that
|
||||
// has no symptom: a change to Config.MaxDepth that quietly puts every
|
||||
// observation over the cap and turns the inference path into a table write per
|
||||
// request. A generous allowance catches that and does not fire on a long model
|
||||
// name.
|
||||
const prefixCacheIdentifierAllowance = 512
|
||||
|
||||
// requirePrefixCacheFitsInline refuses a prefix-cache configuration whose
|
||||
// observations would spill.
|
||||
//
|
||||
// prefixcache.ExtractChain caps a chain at Config.MaxDepth blocks, so an
|
||||
// observation's size has a worst case that is known before the deployment
|
||||
// serves a request: MaxDepth hashes at their widest decimal encoding, plus the
|
||||
// identifiers. That bound is the reason Sync.Observe can publish like every
|
||||
// other family instead of being given a way to refuse.
|
||||
//
|
||||
// It is a startup error and not a warning because the alternative reading is
|
||||
// the one this programme exists to remove: a deployment that came up, spills a
|
||||
// row and reads it back on every replica for every request whose prefix
|
||||
// changed, and looks exactly like one that is merely slow.
|
||||
func requirePrefixCacheFitsInline(cfg prefixcache.Config) error {
|
||||
// The widest a uint64 encodes to in JSON, so the check does not depend on
|
||||
// which hashes a workload happens to produce.
|
||||
chain := make([]uint64, cfg.MaxDepth)
|
||||
for i := range chain {
|
||||
chain[i] = math.MaxUint64
|
||||
}
|
||||
worst := messaging.PrefixCacheObserveEvent{
|
||||
Model: strings.Repeat("m", prefixCacheIdentifierAllowance/2),
|
||||
Chain: chain,
|
||||
NodeID: strings.Repeat("n", prefixCacheIdentifierAllowance/2),
|
||||
Replica: math.MaxInt32,
|
||||
}
|
||||
fits, err := pgbus.FitsInline(messaging.SubjectPrefixCacheObserve, worst)
|
||||
if err != nil {
|
||||
return fmt.Errorf("sizing a prefix-cache observation: %w", err)
|
||||
}
|
||||
if !fits {
|
||||
return fmt.Errorf("the prefix-cache depth is too large to broadcast: an observation for %d blocks does not fit in one notification, so every request whose prefix changed would write a row and every replica would read it back on the inference path", cfg.MaxDepth)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"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/nodes/prefixcache"
|
||||
"github.com/mudler/LocalAI/core/services/pgbus"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// The four process-lifetime caches, each asserted from the OTHER replica's
|
||||
// carrier.
|
||||
//
|
||||
// Every case here wires the cache on busA and drives it from busB. A cache
|
||||
// talking to itself would pass with the wiring pointed at any carrier at all,
|
||||
// which is the defect these exist to catch: a site holding a carrier other than
|
||||
// the deployment's publishes successfully and is delivered, to nobody the
|
||||
// deployment will still be listening on. The second carrier that made that
|
||||
// concrete went with the message broker, and these stay because they are what
|
||||
// would catch it on the day a third arrives: busB IS the other carrier, so the
|
||||
// defect is exercised rather than argued about.
|
||||
var _ = Describe("wiring the process-lifetime caches onto the broadcast carrier", func() {
|
||||
var (
|
||||
ctx context.Context
|
||||
db *gorm.DB
|
||||
busA, busB *pgbus.Bus
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
ctx = context.Background()
|
||||
|
||||
var dsn string
|
||||
db, dsn = testutil.SetupTestDBWithDSN()
|
||||
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
|
||||
|
||||
newBus := func() *pgbus.Bus {
|
||||
b, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(b.Close)
|
||||
return b
|
||||
}
|
||||
busA, busB = newBus(), newBus()
|
||||
})
|
||||
|
||||
// Each of the four refuses rather than coming up on nothing. A cache wired
|
||||
// to no carrier has no symptom of its own: it answers from whatever this
|
||||
// one replica happened to do, forever, and looks exactly like a fleet with
|
||||
// nothing going on elsewhere.
|
||||
DescribeTable("refuses to wire a cache with no carrier",
|
||||
func(wire func() error, want string) {
|
||||
err := wire()
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring(want))
|
||||
},
|
||||
Entry("gallery", func() error {
|
||||
return (&DistributedServices{}).wireGallery(galleryop.NewGalleryService(&config.ApplicationConfig{}, nil))
|
||||
}, "gallery progress and cancels"),
|
||||
Entry("operation cache", func() error {
|
||||
svc := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
return (&DistributedServices{}).WireOpCache(context.Background(), galleryop.NewOpCache(svc))
|
||||
}, "/api/operations"),
|
||||
Entry("staging", func() error {
|
||||
_, err := wireStagingBroadcasts(nil, nodes.NewStagingTracker())
|
||||
return err
|
||||
}, "progress bar only on the replica performing it"),
|
||||
Entry("prefix cache", func() error {
|
||||
_, err := wirePrefixCacheBroadcasts(nil, prefixcache.DefaultConfig(), prefixcache.NewIndex(prefixcache.DefaultConfig()))
|
||||
return err
|
||||
}, "its own history"),
|
||||
)
|
||||
|
||||
// S2. The gallery service applies a peer's progress, which it can only do
|
||||
// if the wildcard subscription wireGallery opened is on the carrier the
|
||||
// peer published to.
|
||||
It("subscribes the gallery service to progress a peer replica broadcasts", func() {
|
||||
svc := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
Expect((&DistributedServices{Bus: busA}).wireGallery(svc)).To(Succeed())
|
||||
DeferCleanup(svc.CloseBroadcasts)
|
||||
|
||||
Expect(busB.Publish(messaging.SubjectGalleryProgress("op-1"), galleryop.GalleryProgressEvent{
|
||||
JobID: "op-1",
|
||||
Status: &galleryop.OpStatus{Progress: 42, Message: "halfway"},
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() *galleryop.OpStatus { return svc.GetStatus("op-1") }, 20*time.Second).ShouldNot(BeNil())
|
||||
Expect(svc.GetStatus("op-1").Progress).To(Equal(42.0))
|
||||
})
|
||||
|
||||
// S2, the other direction. A service that only subscribed would pass the
|
||||
// row above and publish its own progress where no peer reads it.
|
||||
It("publishes the gallery service's progress onto the same carrier", func() {
|
||||
svc := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
Expect((&DistributedServices{Bus: busA}).wireGallery(svc)).To(Succeed())
|
||||
DeferCleanup(svc.CloseBroadcasts)
|
||||
|
||||
out := make(chan []byte, 4)
|
||||
_, err := busB.Subscribe(messaging.SubjectGalleryProgressWildcard, func(b []byte) { out <- b })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
svc.UpdateStatus("op-2", &galleryop.OpStatus{Progress: 7})
|
||||
|
||||
Eventually(out, 20*time.Second).Should(Receive())
|
||||
})
|
||||
|
||||
// S1. The OpCache is wired from the HTTP layer, and WireOpCache is what
|
||||
// keeps that call site from naming a carrier of its own.
|
||||
It("subscribes the operation cache to a peer replica's admissions", func() {
|
||||
svc := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
cache := galleryop.NewOpCache(svc)
|
||||
Expect((&DistributedServices{Bus: busA}).WireOpCache(ctx, cache)).To(Succeed())
|
||||
DeferCleanup(cache.Close)
|
||||
|
||||
Expect(busB.Publish(messaging.SubjectGalleryOpStart, galleryop.OpCacheEvent{
|
||||
JobID: "job-9", CacheKey: "official@vllm", IsBackend: true,
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() bool { return cache.Exists("official@vllm") }, 20*time.Second).Should(BeTrue())
|
||||
Expect(cache.IsBackendOp("official@vllm")).To(BeTrue())
|
||||
})
|
||||
|
||||
It("publishes the operation cache's admissions onto the same carrier", func() {
|
||||
svc := galleryop.NewGalleryService(&config.ApplicationConfig{}, nil)
|
||||
cache := galleryop.NewOpCache(svc)
|
||||
Expect((&DistributedServices{Bus: busA}).WireOpCache(ctx, cache)).To(Succeed())
|
||||
DeferCleanup(cache.Close)
|
||||
|
||||
out := make(chan []byte, 4)
|
||||
_, err := busB.Subscribe(messaging.SubjectGalleryOpStart, func(b []byte) { out <- b })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
cache.Set("llama-3-8b", "job-10")
|
||||
|
||||
Eventually(out, 20*time.Second).Should(Receive())
|
||||
})
|
||||
|
||||
// S3, in both directions and as two separate specs. It used to be two
|
||||
// calls, a publisher and a subscriber, and a tracker with one of them on
|
||||
// each carrier shows a staging progress bar on the originating replica and
|
||||
// nowhere else. SetBroadcaster is one method now, so that deployment cannot
|
||||
// be spelled, but each half still has to be held on its own: a mutation
|
||||
// that drops the subscribe leaves the publishing spec green and the reverse
|
||||
// leaves the mirroring spec green.
|
||||
It("mirrors a peer replica's staging progress into the tracker", func() {
|
||||
tracker := nodes.NewStagingTracker()
|
||||
sub, err := wireStagingBroadcasts(busA, tracker)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(sub).ToNot(BeNil())
|
||||
|
||||
Expect(busB.Publish(messaging.SubjectStagingProgress("model-x"), nodes.StagingProgressEvent{
|
||||
ModelID: "model-x",
|
||||
Status: &nodes.StagingStatus{ModelID: "model-x", NodeName: "worker-7"},
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() map[string]nodes.StagingStatus { return tracker.GetAll() }, 20*time.Second).
|
||||
Should(HaveKey("model-x"))
|
||||
})
|
||||
|
||||
It("publishes the tracker's own staging progress onto the same carrier", func() {
|
||||
tracker := nodes.NewStagingTracker()
|
||||
_, err := wireStagingBroadcasts(busA, tracker)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
out := make(chan []byte, 4)
|
||||
_, err = busB.Subscribe(messaging.SubjectStagingProgressWildcard, func(b []byte) { out <- b })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
tracker.Start("model-y", "worker-8", 1)
|
||||
|
||||
Eventually(out, 20*time.Second).Should(Receive())
|
||||
})
|
||||
|
||||
// S4, in both directions. The prefix cache is the family on the inference
|
||||
// path, and a Sync wired to a carrier its peers do not read leaves every
|
||||
// frontend routing on nothing but its own history while every publish
|
||||
// succeeds.
|
||||
It("applies a peer replica's observation into the prefix index", func() {
|
||||
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
sync, err := wirePrefixCacheBroadcasts(busA, prefixcache.DefaultConfig(), idx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
chain := []uint64{101, 202, 303}
|
||||
Expect(busB.Publish(messaging.SubjectPrefixCacheObserve, messaging.PrefixCacheObserveEvent{
|
||||
Model: "m", Chain: chain, NodeID: "A", Replica: 1,
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() bool {
|
||||
return sync.Decide("m", chain, []prefixcache.ReplicaKey{{NodeID: "A", Replica: 1}}, time.Now()).HasHot
|
||||
}, 20*time.Second).Should(BeTrue())
|
||||
})
|
||||
|
||||
It("applies a peer replica's invalidation, so a removed replica stops being routed to", func() {
|
||||
// The invalidation half separately: a missed one leaves this frontend
|
||||
// routing to a replica that is gone until the TTL, which is the reading
|
||||
// of a missed message this programme forbids.
|
||||
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
sync, err := wirePrefixCacheBroadcasts(busA, prefixcache.DefaultConfig(), idx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
chain := []uint64{404, 505}
|
||||
key := prefixcache.ReplicaKey{NodeID: "A", Replica: 0}
|
||||
sync.ApplyObserve(messaging.PrefixCacheObserveEvent{Model: "m", Chain: chain, NodeID: "A"}, time.Now())
|
||||
Expect(sync.Decide("m", chain, []prefixcache.ReplicaKey{key}, time.Now()).HasHot).To(BeTrue())
|
||||
|
||||
Expect(busB.Publish(messaging.SubjectPrefixCacheInvalidate, messaging.PrefixCacheInvalidateEvent{
|
||||
Model: "m", NodeID: "A", Replica: 0,
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() bool {
|
||||
return sync.Decide("m", chain, []prefixcache.ReplicaKey{key}, time.Now()).HasHot
|
||||
}, 20*time.Second).Should(BeFalse())
|
||||
})
|
||||
|
||||
It("publishes this replica's observations onto the same carrier", func() {
|
||||
idx := prefixcache.NewIndex(prefixcache.DefaultConfig())
|
||||
sync, err := wirePrefixCacheBroadcasts(busA, prefixcache.DefaultConfig(), idx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
out := make(chan []byte, 4)
|
||||
_, err = busB.Subscribe(messaging.SubjectPrefixCacheObserve, func(b []byte) { out <- b })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
sync.Observe("m", []uint64{909}, prefixcache.ReplicaKey{NodeID: "B", Replica: 0}, time.Now())
|
||||
|
||||
Eventually(out, 20*time.Second).Should(Receive())
|
||||
})
|
||||
})
|
||||
|
||||
// The size bound that lets prefix-cache observations publish like every other
|
||||
// family instead of being given a way to refuse.
|
||||
//
|
||||
// The plan for this phase proposed a PublishNoSpill that would REFUSE an
|
||||
// observation too large for a notification, on the reasoning that a long prompt
|
||||
// makes a chain of thousands of entries. ExtractChain does not produce one, and
|
||||
// the refusal would have been the only deliberate message drop in the
|
||||
// programme, guarding a condition that cannot arise, with a counter nothing
|
||||
// alerts on as its only symptom. This is what took its place: the same
|
||||
// knowledge, asked at startup, where being wrong is a deployment that refuses
|
||||
// to come up and says why rather than one that runs with no cross-replica
|
||||
// affinity and looks healthy.
|
||||
//
|
||||
// It needs no bus. FitsInline is a pure function over the same encoder and the
|
||||
// same constant Publish measures against.
|
||||
var _ = Describe("the prefix-cache observation bound", func() {
|
||||
It("accepts the depth the extractor actually produces", func() {
|
||||
Expect(requirePrefixCacheFitsInline(prefixcache.DefaultConfig())).To(Succeed())
|
||||
})
|
||||
|
||||
It("refuses a depth whose observations would spill on every request", func() {
|
||||
// An absolute depth, not one derived from the carrier's cap, so this
|
||||
// row states a fact about this family rather than restating the
|
||||
// constant it is measured against.
|
||||
cfg := prefixcache.DefaultConfig()
|
||||
cfg.MaxDepth = 100000
|
||||
|
||||
err := requirePrefixCacheFitsInline(cfg)
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("write a row"))
|
||||
})
|
||||
|
||||
It("is checked before the prefix cache is wired at all", func() {
|
||||
// The check is worth nothing if the wiring runs anyway. Asserted
|
||||
// through the same function initDistributed calls, and on the MESSAGE
|
||||
// rather than on failure alone: this call has two things wrong with it,
|
||||
// and a spec that accepted any error would pass on the carrier
|
||||
// complaint with the bound check deleted.
|
||||
cfg := prefixcache.DefaultConfig()
|
||||
cfg.MaxDepth = 100000
|
||||
|
||||
sync, err := wirePrefixCacheBroadcasts(nil, cfg, prefixcache.NewIndex(cfg))
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("write a row"))
|
||||
Expect(sync).To(BeNil())
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,87 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/jobs"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// startJobDispatchLoop builds AND STARTS the loop that takes queued work off the
|
||||
// job store and drives it on an agent worker.
|
||||
//
|
||||
// One function rather than a construction here and a Start somewhere else, and
|
||||
// that is the point rather than tidiness. A loop that is built and never
|
||||
// started is a replica that writes claim rows and takes none, so every job in
|
||||
// the deployment is accepted and none is ever run, and nothing anywhere says
|
||||
// so. As a separate statement in the start-up path that line's loss has no
|
||||
// symptom and no spec reaches it: initDistributed opens a database and a bus.
|
||||
// Fused here, the loop cannot exist without running.
|
||||
//
|
||||
// The rest is a named function for the reason newAgentControl is one: two of
|
||||
// these arguments are silent when they are wrong.
|
||||
//
|
||||
// The BROADCASTER is the one worth naming. It is what re-publishes the progress
|
||||
// and result lines a worker asks for, and it is checked against the allow list
|
||||
// for that worker's node type. A loop built without one dispatches work
|
||||
// perfectly well and every SSE stream in the deployment goes quiet: the job
|
||||
// runs, the answer is persisted, and the user watching it sees nothing until
|
||||
// they reload. That is a whole feature lost to a nil field, with no error
|
||||
// anywhere, so it is refused here.
|
||||
//
|
||||
// The re-broadcaster is taken already built, from newFanoutBridges, and is a
|
||||
// *nodes.Rebroadcaster rather than the jobs.ProgressBroadcaster interface the
|
||||
// loop stores it as. Both of those are deliberate. Taking it built leaves ONE
|
||||
// expression in the tree that decides which carrier job and agent fan-out goes
|
||||
// on, next to the dispatcher and the bridge that must read the same one, so
|
||||
// there is no separate line here to point at a carrier nobody subscribes to:
|
||||
// that mis-wiring publishes successfully, returns true, reddens no spec in any
|
||||
// package, and shows up only as an SSE stream with no progress in it. Naming
|
||||
// the concrete type is what makes the refusal below fire, too: widened to the
|
||||
// interface, a nil re-broadcaster is a non-nil value holding a nil pointer.
|
||||
//
|
||||
// The SELECTOR is built here rather than borrowed from newAgentControl, and
|
||||
// deliberately: nodes.AgentSelector holds no per-caller state, and sharing one
|
||||
// would couple the dispatch loop's lifetime to MCP's for nothing.
|
||||
//
|
||||
// The reconnect grace it is built with is INERT on this path and is passed
|
||||
// correctly anyway. The selector reads it only in Reachable, which is what a
|
||||
// fan-out verb (an agent cancel) asks; this loop asks PickConnected, which
|
||||
// never needs it, because a worker that is not connected cannot be picked
|
||||
// whatever the reason. Passing a value this loop cannot observe is still
|
||||
// cheaper than a second constructor.
|
||||
func startJobDispatchLoop(ctx context.Context, cfg config.DistributedConfig, db *gorm.DB, store *jobs.JobStore,
|
||||
registry *nodes.NodeRegistry, conns nodes.AgentConnectionReader,
|
||||
control *nodes.ControlClient, broadcast *nodes.Rebroadcaster) (*jobs.DispatchLoop, error) {
|
||||
if cfg.InstanceID == "" {
|
||||
return nil, fmt.Errorf("the job dispatch loop was built with no instance id: its claims could not be told from ones a dead replica left")
|
||||
}
|
||||
if registry == nil || conns == nil {
|
||||
return nil, fmt.Errorf("the job dispatch loop was built with no way to find a connected agent worker")
|
||||
}
|
||||
if broadcast == nil {
|
||||
return nil, fmt.Errorf("the job dispatch loop was built with no broadcaster: every job would run with its progress and its result reaching no SSE stream in the deployment")
|
||||
}
|
||||
loop, err := jobs.NewDispatchLoop(jobs.DispatchConfig{
|
||||
DB: db,
|
||||
Owner: cfg.InstanceID,
|
||||
Selector: nodes.NewAgentSelector(registry, conns, cfg.InstanceID, cfg.WorkerReconnectGrace),
|
||||
Control: control,
|
||||
// The allow list lives in nodes and is keyed on the worker's node type;
|
||||
// nothing here decides what a worker may broadcast on.
|
||||
Broadcast: broadcast,
|
||||
Store: store,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := loop.Start(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return loop, nil
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"runtime"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/cluster"
|
||||
"github.com/mudler/LocalAI/core/services/jobs"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// The wiring that turns queued claim rows into work on an agent worker.
|
||||
//
|
||||
// Guarded the way newAgentControl is and for the same reason: initDistributed
|
||||
// opens a database and a bus, so no unit spec reaches the construction literal,
|
||||
// and two of these arguments are silent when they are wrong.
|
||||
var _ = Describe("building the job dispatch loop", func() {
|
||||
var registry *nodes.NodeRegistry
|
||||
var conns *recordingConnections
|
||||
var ctx context.Context
|
||||
var broadcast *nodes.Rebroadcaster
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
ctx = context.Background()
|
||||
var err error
|
||||
registry, err = nodes.NewNodeRegistry(testutil.SetupTestDB())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
conns = newRecordingConnections()
|
||||
|
||||
// A double is enough HERE, and only here. Which carrier this
|
||||
// re-broadcaster publishes on is not this function's decision any more:
|
||||
// it is handed one already built by newFanoutBridges, and that is where
|
||||
// the carrier is pinned, by receipt on a second connection. What is
|
||||
// left for these to say is that the loop refuses to be built without
|
||||
// one and starts when it is.
|
||||
broadcast = nodes.NewRebroadcaster(testutil.NewFakeBus())
|
||||
})
|
||||
|
||||
// The silent one. A loop with no broadcaster dispatches work perfectly
|
||||
// well: the job runs, the answer is persisted, and every SSE stream in the
|
||||
// deployment goes quiet, with no error anywhere.
|
||||
It("refuses to build with no broadcaster to re-publish a worker's progress on", func() {
|
||||
_, err := startJobDispatchLoop(ctx, config.DistributedConfig{InstanceID: "replica-7"},
|
||||
testutil.SetupTestDB(), nil, registry, conns, nodes.NewControlClient(nil, "token"), nil)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("broadcaster"))
|
||||
})
|
||||
|
||||
// The nil that the interface would have hidden. The loop stores its
|
||||
// re-broadcaster as the jobs.ProgressBroadcaster interface, and widened to
|
||||
// that here a nil *nodes.Rebroadcaster is a NON-nil value holding a nil
|
||||
// pointer, so the refusal above would never fire for the way one is
|
||||
// actually absent: newFanoutBridges returns a typed nil alongside its
|
||||
// error. This drives that exact value, which is why the parameter is the
|
||||
// concrete type.
|
||||
It("refuses a typed-nil broadcaster, which an interface parameter would have accepted", func() {
|
||||
var absent *nodes.Rebroadcaster
|
||||
_, err := startJobDispatchLoop(ctx, config.DistributedConfig{InstanceID: "replica-7"},
|
||||
testutil.SetupTestDB(), nil, registry, conns, nodes.NewControlClient(nil, "token"), absent)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("broadcaster"))
|
||||
})
|
||||
|
||||
It("refuses to build with no instance id", func() {
|
||||
_, err := startJobDispatchLoop(ctx, config.DistributedConfig{},
|
||||
testutil.SetupTestDB(), nil, registry, conns, nodes.NewControlClient(nil, "token"),
|
||||
broadcast)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("instance id"))
|
||||
})
|
||||
|
||||
It("refuses to build with nothing to read connections through", func() {
|
||||
_, err := startJobDispatchLoop(ctx, config.DistributedConfig{InstanceID: "replica-7"},
|
||||
testutil.SetupTestDB(), nil, registry, nil, nodes.NewControlClient(nil, "token"),
|
||||
broadcast)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("connected agent worker"))
|
||||
})
|
||||
|
||||
// Driven through the loop's OWN tick, not through a call this spec makes,
|
||||
// and that is deliberate. Two things are pinned here at once: that this
|
||||
// replica's id reaches the SELECTION (a loop built with the wrong id relays
|
||||
// every RPC through a peer and says so nowhere), and that building the loop
|
||||
// STARTED it (a loop that is never started writes claim rows and takes
|
||||
// none, so every job in the deployment is accepted and never run).
|
||||
It("starts on construction, and selects as THIS replica", func() {
|
||||
db := testutil.SetupTestDB()
|
||||
Expect(cluster.Migrate(ctx, db)).To(Succeed())
|
||||
Expect(jobs.MigrateClaims(ctx, db)).To(Succeed())
|
||||
// Registered, because a replica that is not in the instances table
|
||||
// refuses to claim: its claims could not be told from a dead one's.
|
||||
Expect(cluster.NewRegistry(db).Register(ctx, "replica-7", "127.0.0.1:8080", "v1", "")).To(Succeed())
|
||||
_, err := jobs.EnqueueClaim(ctx, db, jobs.ClaimKindAgentRun, json.RawMessage(`{}`))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
loop, err := startJobDispatchLoop(ctx, config.DistributedConfig{InstanceID: "replica-7"},
|
||||
db, nil, registry, conns, nodes.NewControlClient(nil, "token"), broadcast)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(loop.Stop)
|
||||
|
||||
// Nothing prods it. The only thing that can make this happen is the
|
||||
// loop's own goroutine.
|
||||
Eventually(conns.calledBy, "20s").Should(Receive(Equal("replica-7")))
|
||||
})
|
||||
})
|
||||
+639
-83
@@ -5,6 +5,8 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -12,13 +14,16 @@ import (
|
||||
"github.com/google/uuid"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/services/agents"
|
||||
"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"
|
||||
"github.com/mudler/LocalAI/core/services/monitoring"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/nodes/prefixcache"
|
||||
"github.com/mudler/LocalAI/core/services/pgbus"
|
||||
"github.com/mudler/LocalAI/core/services/storage"
|
||||
"github.com/mudler/LocalAI/internal"
|
||||
"github.com/mudler/LocalAI/pkg/distributedhdr"
|
||||
"github.com/mudler/LocalAI/pkg/sanitize"
|
||||
"github.com/mudler/xlog"
|
||||
@@ -27,7 +32,6 @@ import (
|
||||
|
||||
// DistributedServices holds all services initialized for distributed mode.
|
||||
type DistributedServices struct {
|
||||
Nats *messaging.Client
|
||||
Store storage.ObjectStore
|
||||
Registry *nodes.NodeRegistry
|
||||
Router *nodes.SmartRouter
|
||||
@@ -44,6 +48,55 @@ type DistributedServices struct {
|
||||
Unloader *nodes.RemoteUnloaderAdapter
|
||||
ModelCleanup *nodes.ModelCleanupService
|
||||
|
||||
// Bus is the deployment's fan-out carrier, riding the auth database's
|
||||
// PostgreSQL rather than a message broker. Every cross-replica family the
|
||||
// frontend broadcasts is on it now, with the single exception named above.
|
||||
// Adopters reach it through Broadcast() rather than through this field, so
|
||||
// that "which carrier does this family travel on" is decided once instead
|
||||
// of at every adopter; its DSN likewise has exactly one legitimate source,
|
||||
// settled in newBroadcastBus rather than invented per call site.
|
||||
Bus *pgbus.Bus
|
||||
|
||||
// Cluster is the replica-membership registry: which frontend replicas are
|
||||
// alive, at which address, and which of them holds a given worker's tunnel.
|
||||
Cluster *cluster.Registry
|
||||
// Membership publishes this replica's row and reaps the dead. Nil when no
|
||||
// peer-reachable address could be determined, which leaves this replica
|
||||
// invisible to its peers but otherwise fully functional.
|
||||
Membership *cluster.Membership
|
||||
// PeerSessions owns the peer links other replicas dialled into this one,
|
||||
// and relays the streams that arrive on them onto the worker tunnels this
|
||||
// replica holds.
|
||||
PeerSessions *cluster.SessionStore
|
||||
// Peers owns the peer links this replica dialled OUT, the mirror of
|
||||
// PeerSessions. It is what the relaying dialer opens a stream on when a
|
||||
// request arrives here for a worker another replica holds.
|
||||
Peers *cluster.PeerPool
|
||||
// Tunnels holds the worker tunnels this replica has accepted and keeps the
|
||||
// node_connections table agreeing with them. It is handed to the membership
|
||||
// loop, which re-claims what it holds after this replica has been reaped,
|
||||
// and to the route that accepts a worker's dial.
|
||||
Tunnels *cluster.TunnelRegistry
|
||||
// WorkerDialer is how anything in this process reaches a worker: locally
|
||||
// when this replica holds the tunnel, and through the owning replica when
|
||||
// it does not. The HTTP layer takes its WebSocket log proxy from here.
|
||||
WorkerDialer *cluster.WorkerDialer
|
||||
// BackendClients builds the gRPC clients for worker backend processes, over
|
||||
// WorkerDialer. Exposed so the model store built in startup.go reaches
|
||||
// remote models the same way every other caller does.
|
||||
BackendClients nodes.BackendClientFactory
|
||||
// AgentControl carries the frontend's MCP verbs to whichever agent worker
|
||||
// holds a tunnel this deployment can reach. It is what the chat, responses,
|
||||
// messages and MCP endpoints reach an agent worker through; a nil one means
|
||||
// this frontend cannot run MCP at all, which is why initDistributed refuses
|
||||
// to come up without it rather than leaving the endpoints to discover it
|
||||
// one request at a time.
|
||||
AgentControl *nodes.AgentControlClient
|
||||
// JobDispatch takes queued work off the job store and drives it on an agent
|
||||
// worker over that worker's tunnel. It is what replaces the three NATS
|
||||
// queue groups: dispatch is a claim, and a claim is a row and a lock.
|
||||
JobDispatch *jobs.DispatchLoop
|
||||
|
||||
shutdownOnce sync.Once
|
||||
}
|
||||
|
||||
@@ -54,19 +107,43 @@ func (ds *DistributedServices) Shutdown() {
|
||||
return
|
||||
}
|
||||
ds.shutdownOnce.Do(func() {
|
||||
// Peer state first: a replica that is going away should stop claiming
|
||||
// to be alive before it stops answering, so peers re-home rather than
|
||||
// dial a process in teardown.
|
||||
if ds.Membership != nil {
|
||||
ds.Membership.Stop()
|
||||
}
|
||||
if ds.PeerSessions != nil {
|
||||
ds.PeerSessions.CloseAll()
|
||||
}
|
||||
// Both halves of the peer mesh go down together. A pool left open
|
||||
// holds a WebSocket and two yamux loop goroutines per peer for as long
|
||||
// as the process lives, and an Open after this reports ErrPoolClosed,
|
||||
// which is a fact about this process and never node absence.
|
||||
if ds.Peers != nil {
|
||||
ds.Peers.Close()
|
||||
}
|
||||
if ds.Health != nil {
|
||||
ds.Health.Stop()
|
||||
}
|
||||
if ds.JobDispatch != nil {
|
||||
ds.JobDispatch.Stop()
|
||||
}
|
||||
if ds.Dispatcher != nil {
|
||||
ds.Dispatcher.Stop()
|
||||
}
|
||||
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 and needs none: it holds no
|
||||
// process-lifetime subscription of its own beyond the observable
|
||||
// persister, whose carrier is closed below.
|
||||
// The broadcast carrier holds a PostgreSQL session pinned for the life
|
||||
// of the process, plus the goroutine parked on it. A replica that
|
||||
// leaves one behind on every restart runs the server out of
|
||||
// connections, and the symptom lands on whatever connects next.
|
||||
if ds.Bus != nil {
|
||||
ds.Bus.Close()
|
||||
}
|
||||
xlog.Info("Distributed services shut down")
|
||||
})
|
||||
@@ -77,7 +154,13 @@ func (ds *DistributedServices) Shutdown() {
|
||||
// Returns nil if distributed mode is not enabled.
|
||||
// configLoader is used by the SmartRouter to compute concurrency-group
|
||||
// anti-affinity at placement time (#9659); it may be nil in tests.
|
||||
func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoader *config.ModelConfigLoader) (*DistributedServices, error) {
|
||||
//
|
||||
// galleryProgress is the gallery service, narrowed to the one method a node
|
||||
// departure needs. It is a PARAMETER and not a later setter because the
|
||||
// registration of every per-node cache a departure evicts happens here, in one
|
||||
// place, and a cache registered somewhere else is a cache a reader cannot find
|
||||
// by reading this function.
|
||||
func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoader *config.ModelConfigLoader, galleryProgress nodeProgressDropper) (*DistributedServices, error) {
|
||||
if !cfg.Distributed.Enabled {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -103,25 +186,16 @@ 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")
|
||||
}
|
||||
natsOpts := cfg.Distributed.NatsMessagingOptions("", "")
|
||||
natsClient, err := messaging.New(cfg.Distributed.NatsURL, natsOpts...)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("connecting to NATS: %w", err)
|
||||
}
|
||||
xlog.Info("Connected to NATS", "url", sanitize.URL(cfg.Distributed.NatsURL))
|
||||
|
||||
// Ensure NATS is closed if any subsequent initialization step fails.
|
||||
// No message bus is dialled here, and there is none left to dial. The last
|
||||
// family that needed one was agent.<name>.cancel, whose subscriber is an
|
||||
// agent worker that has no database and so could not join the broadcast
|
||||
// carrier below; it is now a control RPC on the tunnel that worker holds.
|
||||
// A distributed deployment needs PostgreSQL and the frontends' own HTTP
|
||||
// listener, and nothing else.
|
||||
//
|
||||
// success guards the carriers opened below, which must not be left pinned
|
||||
// when a later initialization step fails.
|
||||
success := false
|
||||
defer func() {
|
||||
if !success {
|
||||
natsClient.Close()
|
||||
}
|
||||
}()
|
||||
|
||||
// Initialize object storage
|
||||
var store storage.ObjectStore
|
||||
@@ -157,6 +231,21 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
return nil, fmt.Errorf("distributed mode requires auth database to be initialized first")
|
||||
}
|
||||
|
||||
// The fan-out carrier, opened before anything that might want it. It is
|
||||
// built here and not by its first adopter because its DSN has one
|
||||
// legitimate source, and a setting that decides whether every broadcast in
|
||||
// the deployment is delivered should not be settled under the time pressure
|
||||
// of a migration.
|
||||
bus, err := newBroadcastBus(cfg.Context, cfg, authDB)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer func() {
|
||||
if !success {
|
||||
bus.Close()
|
||||
}
|
||||
}()
|
||||
|
||||
registry, err := nodes.NewNodeRegistry(authDB)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("initializing node registry: %w", err)
|
||||
@@ -174,6 +263,122 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
xlog.Warn("Control-plane database metrics unavailable", "error", err)
|
||||
}
|
||||
|
||||
// Replica membership. NewNodeRegistry has just migrated the tables this
|
||||
// reads, so it has to come after it.
|
||||
clusterRegistry := cluster.NewRegistry(authDB)
|
||||
// This replica's proof of which replica it is, minted ONCE per process and
|
||||
// handed to both halves of the peer mesh: the membership loop publishes its
|
||||
// hash in the instances row, and the peer pool presents its plaintext on
|
||||
// every outbound dial. Two mints would leave a replica whose published hash
|
||||
// and presented secret disagree, which every peer refuses and which reads
|
||||
// from the logs like a peer running an older release.
|
||||
//
|
||||
// The plaintext never leaves this process except in a peer dial's header.
|
||||
// There is nothing to configure and nothing to rotate: a restart mints a
|
||||
// new one, and the same registration that republishes this replica's
|
||||
// address republishes the hash beside it.
|
||||
peerCredential := cluster.NewPeerCredential()
|
||||
var membership *cluster.Membership
|
||||
if advertised, err := advertisedPeerAddr(cfg); err != nil {
|
||||
// Not fatal, and the cost is worth stating exactly rather than as
|
||||
// "peers cannot reach it", because it is larger than that now.
|
||||
//
|
||||
// Without a row in the instances table this replica is not a live
|
||||
// owner as far as Registry.Owner is concerned: that read joins a
|
||||
// connection against a live instance, so a worker whose tunnel lands
|
||||
// HERE is answered as unroutable at every OTHER replica, for as long
|
||||
// as it stays here. This replica serves that worker perfectly well
|
||||
// itself; nobody else can. On N replicas behind round robin that is
|
||||
// (N-1)/N of the traffic for that worker.
|
||||
//
|
||||
// It does not refuse to START. Refusing would take out every existing
|
||||
// single-host deployment, whose route to a local database is loopback
|
||||
// and which has no peers to be unreachable by; the deployments this
|
||||
// hurts are multi-replica ones, and telling those two apart at startup
|
||||
// is a change with its own design and its own specs rather than a line
|
||||
// here.
|
||||
//
|
||||
// What it does not get to do is stay quiet. One startup line scrolls
|
||||
// away in seconds and the cost is paid for the whole life of the
|
||||
// process, on a symptom (workers that 5xx from most of the fleet) whose
|
||||
// obvious reading is "the worker is broken". So this is an ERROR, not a
|
||||
// warning, and nagUnadvertisedReplica below repeats it for as long as
|
||||
// the state lasts, naming the workers it is currently costing.
|
||||
//
|
||||
// It now costs the OTHER direction too, which is why the line says so.
|
||||
// A peer link is authenticated by the dialling replica's own credential,
|
||||
// published in the instances row this replica never writes, so this
|
||||
// replica cannot dial a peer either: its own relayed requests are
|
||||
// refused as an unproven identity rather than merely arriving nowhere.
|
||||
xlog.Error("This replica is not registered in the cluster: no advertised address. Peers cannot reach it, any worker whose tunnel lands here will be unroutable from every other replica, and this replica cannot relay OUT either, because a peer link is authenticated by the credential an instance row publishes and this replica has no row",
|
||||
"error", err, "knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR")
|
||||
} else {
|
||||
membership = cluster.NewMembership(clusterRegistry, cfg.Distributed.InstanceID, advertised, internal.PrintableVersion(), peerCredential)
|
||||
// Before Start, so the first sweep already purges on the retention this
|
||||
// deployment's grace requires rather than on the floor.
|
||||
membership.SetReconnectGrace(cfg.Distributed.ReconnectGraceOrDefault())
|
||||
if err := membership.Start(cfg.Context); err != nil {
|
||||
return nil, fmt.Errorf("registering this replica in the cluster: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The worker tunnels this replica accepts. It claims as the SAME instance
|
||||
// ID membership registers under, because that is the ID a peer's Owner
|
||||
// lookup joins a claim against to decide the owner is alive; two IDs here
|
||||
// would make every claim this replica writes look like it belongs to a
|
||||
// replica that does not exist.
|
||||
tunnels := cluster.NewTunnelRegistry(clusterRegistry, cfg.Distributed.InstanceID)
|
||||
// Without this the re-claim in the heartbeat loop is dead code: a replica
|
||||
// stalled long enough to be swept loses the connection rows it owned, and
|
||||
// nothing would ever write them back, so every other replica would answer
|
||||
// "not connected" for workers that are connected right here.
|
||||
//
|
||||
// Nil when no peer-reachable address could be determined above. There is no
|
||||
// heartbeat loop to hand it to in that case, and no other replica can reach
|
||||
// this one anyway; the registry is still built, because it is what the
|
||||
// tunnel endpoint attaches to and what this replica opens its own streams
|
||||
// through.
|
||||
if membership != nil {
|
||||
membership.SetTunnels(tunnels)
|
||||
} else {
|
||||
// The runtime symptom the startup line cannot be. See
|
||||
// nagUnadvertisedReplica.
|
||||
go nagUnadvertisedReplica(cfg.Context, tunnels.Held, unadvertisedNagInterval, logUnroutableWorkers)
|
||||
}
|
||||
|
||||
// The links peers dial IN, with the relay installed on them. This is what
|
||||
// makes more than one replica work: a worker holds one tunnel, it lands on
|
||||
// one replica, and every request that arrives anywhere else reaches the
|
||||
// worker through this handler. Passing nil here would leave every such
|
||||
// request refused, promptly and only at debug level, which presents as a
|
||||
// worker that is connected and unusable from most of the deployment.
|
||||
peerSessions := cluster.NewSessionStore(cluster.NewRelay(tunnels).Stream)
|
||||
// The links this replica dials OUT, the other half of the same mesh. It
|
||||
// authenticates with the registration token because that is the token the
|
||||
// peer route checks (see RegisterClusterRoutes); two different tokens here
|
||||
// would make every peer dial 401 with nothing naming the mismatch.
|
||||
//
|
||||
// And with this replica's own credential, which is the half that says WHICH
|
||||
// replica is dialling. It is the same value membership published the hash
|
||||
// of, by construction: there is one mint above and both call sites read it.
|
||||
peers := cluster.NewPeerPool(cfg.Distributed.InstanceID, cfg.Distributed.RegistrationToken, peerCredential, clusterRegistry)
|
||||
// The one door to every worker. Nothing in the frontend may dial a worker's
|
||||
// advertised address any more: a worker holds ONE tunnel, it lands on ONE
|
||||
// replica, and this resolves which replica that is and relays through it
|
||||
// when it is not this one. The three transports the frontend speaks to a
|
||||
// worker (gRPC to backend processes, HTTP for file staging and logs, a
|
||||
// WebSocket for live log streaming) are all pointed at it below.
|
||||
workerDialer := cluster.NewWorkerDialer(tunnels, peers)
|
||||
backendClients, err := nodes.NewTunnelClientFactory(cfg.Distributed.RegistrationToken, workerDialer.GRPCDialerFor)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("wiring the worker backend client factory: %w", err)
|
||||
}
|
||||
// Bound to the http tag: the worker ignores the target for it and routes to
|
||||
// its own file-transfer and log server, wherever that bound.
|
||||
workerHTTPDialer := nodes.WorkerNetDialerFor(func(nodeID string) func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
return workerDialer.DialerFor(nodeID, cluster.StreamTagHTTP)
|
||||
})
|
||||
|
||||
// Let scheduling rules be keyed by a model alias. The registry resolves a
|
||||
// rule's name through the config loader to find the model it governs, so an
|
||||
// operator can pin placement to a stable name like "production" and have it
|
||||
@@ -209,11 +414,26 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
routerGalleriesJSON = string(galleriesJSON)
|
||||
}
|
||||
|
||||
// The health monitor is the SECOND reader of absence, and it reads it from
|
||||
// the same place and against the same window as the scheduler: a heartbeat
|
||||
// says the worker's supervisor is alive, presence says whether anything
|
||||
// here can still reach its backends, and a worker can be the first without
|
||||
// being the second indefinitely.
|
||||
//
|
||||
// The departure notifier is built HERE, before its only caller, and its
|
||||
// subscribers are registered further down once the caches they drop exist.
|
||||
// One object, one caller, so "what does a departure evict" is answered by
|
||||
// reading registerDepartureEvictions and nothing else.
|
||||
departures := nodes.NewDepartureNotifier()
|
||||
healthMon := nodes.NewHealthMonitor(registry, authDB,
|
||||
cfg.Distributed.HealthCheckIntervalOrDefault(),
|
||||
cfg.Distributed.StaleNodeThresholdOrDefault(),
|
||||
routerAuthToken,
|
||||
!cfg.Distributed.DisablePerModelHealthCheck,
|
||||
clusterRegistry,
|
||||
cfg.Distributed.ReconnectGraceOrDefault(),
|
||||
departures,
|
||||
backendClients,
|
||||
)
|
||||
|
||||
// Initialize job store
|
||||
@@ -223,9 +443,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
}
|
||||
xlog.Info("Distributed job store initialized")
|
||||
|
||||
// Initialize job dispatcher
|
||||
dispatcher := jobs.NewDispatcher(jobStore, natsClient, authDB, cfg.Distributed.InstanceID, cfg.Distributed.JobWorkerConcurrency)
|
||||
|
||||
// Initialize agent store
|
||||
agentStore, err := agents.NewAgentStore(authDB)
|
||||
if err != nil {
|
||||
@@ -233,15 +450,38 @@ 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)
|
||||
// The frontend's control plane client. It reaches every worker over that
|
||||
// worker's own tunnel, on the same `http` stream tag the file stager below
|
||||
// uses, so a control RPC to a worker another replica holds is relayed the
|
||||
// way an inference request is.
|
||||
//
|
||||
// ONE of these for the whole frontend, and the S3 file stager takes this
|
||||
// one rather than minting a second. The client caches an http.Client per
|
||||
// node, which is what keeps a worker's tunnel stream warm between verbs; a
|
||||
// second client would open its own and the two would never share one.
|
||||
controlClient := nodes.NewControlClient(workerHTTPDialer, cfg.Distributed.RegistrationToken)
|
||||
|
||||
// Start observable persister — captures observable_update events from workers
|
||||
// (which have no DB access) and persists them to PostgreSQL.
|
||||
if err := agentBridge.StartObservablePersister(); err != nil {
|
||||
xlog.Warn("Failed to start observable persister", "error", err)
|
||||
} else {
|
||||
xlog.Info("Observable persister started")
|
||||
// The caller the agent worker's control plane has been waiting for. MCP
|
||||
// execution and discovery used to be a NATS request onto a queue group,
|
||||
// where the bus chose the worker and neither side could say which one had
|
||||
// answered; they are now a query against the connection rows plus an
|
||||
// ordinary control RPC over the chosen worker's tunnel.
|
||||
agentControl, err := newAgentControl(cfg.Distributed, registry, clusterRegistry, controlClient)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("wiring the agent control client: %w", err)
|
||||
}
|
||||
|
||||
// The job dispatcher and the agent event bridge, both on the broadcast
|
||||
// carrier. See newFanoutBridges for why the two constructors are reached
|
||||
// through one function that names *pgbus.Bus.
|
||||
//
|
||||
// The bridge takes agentControl and not a carrier: a cancel is the one
|
||||
// family whose far end is an agent worker, and it now rides that worker's
|
||||
// tunnel as a control RPC. It is built ABOVE for that reason, rather than
|
||||
// with the rest of the control plane below.
|
||||
dispatcher, agentBridge, rebroadcast, err := newFanoutBridges(bus, agentControl, jobStore, agentStore, authDB, cfg.Distributed.InstanceID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Initialize Phase 4 stores (MCP, Gallery, FineTune, Skills)
|
||||
@@ -258,28 +498,37 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
}
|
||||
xlog.Info("File manager initialized", "cacheDir", cacheDir)
|
||||
|
||||
// The consumer side of the claim queue, built and started in one act: see
|
||||
// startJobDispatchLoop for why those are not two lines.
|
||||
jobDispatch, err := startJobDispatchLoop(cfg.Context, cfg.Distributed, authDB, jobStore, registry, clusterRegistry, controlClient, rebroadcast)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("wiring the job dispatch loop: %w", err)
|
||||
}
|
||||
|
||||
// Create FileStager for distributed file transfer
|
||||
var fileStager nodes.FileStager
|
||||
if cfg.Distributed.StorageURL != "" {
|
||||
fileStager = nodes.NewS3NATSFileStager(fileMgr, natsClient)
|
||||
xlog.Info("File stager initialized (S3+NATS)")
|
||||
fileStager = nodes.NewS3FileStager(fileMgr, controlClient)
|
||||
xlog.Info("File stager initialized (object store + worker tunnel)")
|
||||
} 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)
|
||||
// An empty HTTPAddress is no longer a refusal. A tunnel-only worker
|
||||
// reports none and does not need one: the http stream tag ignores
|
||||
// the target and the worker routes to its own server. The host is
|
||||
// only ever the URL's host component here, and WorkerHTTPHost
|
||||
// supplies one that resolves nowhere so it cannot become a dial.
|
||||
return nodes.WorkerHTTPHost(nodeID, node.HTTPAddress), nil
|
||||
}, cfg.Distributed.RegistrationToken, workerHTTPDialer)
|
||||
xlog.Info("File stager initialized (HTTP direct transfer)")
|
||||
}
|
||||
// Create RemoteUnloaderAdapter — needed by SmartRouter and startup.go
|
||||
remoteUnloader := nodes.NewRemoteUnloaderAdapter(
|
||||
registry,
|
||||
natsClient,
|
||||
controlClient,
|
||||
cfg.Distributed.BackendInstallTimeoutOrDefault(),
|
||||
cfg.Distributed.BackendUpgradeTimeoutOrDefault(),
|
||||
)
|
||||
@@ -288,12 +537,17 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
// with --distributed-prefix-cache=false, which leaves prefixProvider and
|
||||
// pressure nil so the SmartRouter and reconciler behave exactly as the
|
||||
// round-robin floor (true no-op). When enabled we build the local index,
|
||||
// wrap it in a NATS-backed Sync (publishes our observations, applies peers'
|
||||
// via the subscriptions below), install the extraction hook used by
|
||||
// wrap it in a Sync on the broadcast carrier (which both publishes our
|
||||
// observations and applies peers'), install the extraction hook used by
|
||||
// core/backend/llm.go, and run a background eviction ticker on the app ctx.
|
||||
var prefixProvider prefixcache.Provider
|
||||
var pressure *prefixcache.Pressure
|
||||
var prefixCfg prefixcache.Config
|
||||
// The CONCRETE Sync, declared out here so the departure wiring below can be
|
||||
// handed it. Concrete and not prefixProvider, because a nil interface value
|
||||
// carrying a nil *Sync is not nil, and the disabled deployment would then
|
||||
// register an eviction that dereferences it.
|
||||
var prefixDrop *prefixcache.Sync
|
||||
if !cfg.Distributed.PrefixCacheDisabled {
|
||||
prefixCfg = prefixcache.DefaultConfig()
|
||||
if cfg.Distributed.PrefixCacheTTL > 0 {
|
||||
@@ -303,9 +557,19 @@ 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)
|
||||
// S4. One call puts this replica's observations and its peers' on the
|
||||
// same carrier, and it takes the CONCRETE carrier so that no other
|
||||
// thing satisfying messaging.Broadcaster can be handed to it by
|
||||
// accident. There is no second carrier in this scope to hand over any
|
||||
// more; the type stays narrow so there is still none on the day one is
|
||||
// added. See cache_fanout_wiring.go for the whole argument.
|
||||
prefixSync, err := wirePrefixCacheBroadcasts(bus, prefixCfg, idx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pressure = prefixcache.NewPressure(prefixCfg.PressureWindow)
|
||||
prefixProvider = prefixSync
|
||||
prefixDrop = prefixSync
|
||||
|
||||
// Invalidate the prefix-cache index whenever a replica row is removed.
|
||||
// AddReplicaRemovedHook fires from the single chokepoint all removal paths
|
||||
@@ -327,20 +591,6 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
return prefixcache.ExtractChain(model, prompt, prefixCfg)
|
||||
}
|
||||
|
||||
// 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) {
|
||||
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) {
|
||||
prefixSync.ApplyInvalidate(ev)
|
||||
}); err != nil {
|
||||
return nil, fmt.Errorf("subscribing to %s: %w", messaging.SubjectPrefixCacheInvalidate, err)
|
||||
}
|
||||
|
||||
// Background eviction: sweep idle entries on the app context. Stopped
|
||||
// when the app context is cancelled (mirrors the reconciler loop which
|
||||
// also runs on options.Context). TTL/2 keeps stale entries from
|
||||
@@ -369,12 +619,19 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
conflictResolver = configLoader
|
||||
}
|
||||
modelCleanup := nodes.NewModelCleanupService(registry, remoteUnloader)
|
||||
router := nodes.NewSmartRouter(registry, nodes.SmartRouterOptions{
|
||||
// Absence is stamped on by distributedSchedulerOptions rather than written
|
||||
// here. It is the only source of absence the scheduler has -- a fact read
|
||||
// from the database, so every replica answers it identically, where the bus
|
||||
// sentinel it replaces was one frontend's observation that nobody answered
|
||||
// IT within a budget -- and a field carrying that in a literal this size is
|
||||
// the easiest thing in this file to lose without a symptom.
|
||||
router := nodes.NewSmartRouter(registry, distributedSchedulerOptions(cfg.Distributed, clusterRegistry, nodes.SmartRouterOptions{
|
||||
Unloader: remoteUnloader,
|
||||
ModelCleanup: modelCleanup,
|
||||
FileStager: fileStager,
|
||||
GalleriesJSON: routerGalleriesJSON,
|
||||
AuthToken: routerAuthToken,
|
||||
ClientFactory: backendClients,
|
||||
DB: authDB,
|
||||
ConflictResolver: conflictResolver,
|
||||
PrefixProvider: prefixProvider,
|
||||
@@ -405,12 +662,13 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
// Bounds the REQUEST, not the load: a caller out of budget gets 503 with
|
||||
// live staging progress while the job keeps running underneath.
|
||||
ModelLoadWait: cfg.Distributed.ModelLoadWait,
|
||||
})
|
||||
}))
|
||||
|
||||
// Wire staging-progress broadcasting so file-staging shows up on every
|
||||
// replica, not just the one performing the transfer. Without this, a
|
||||
// /api/operations poll that round-robins onto a peer sees no staging row and
|
||||
// the progress flickers. The origin publishes; peers mirror via the wildcard.
|
||||
// the progress flickers. The origin publishes; peers mirror via the
|
||||
// wildcard, on the same carrier.
|
||||
// A silently disabled safety check is how the original incident stayed
|
||||
// invisible for sixteen minutes. Say so once, loudly, at startup.
|
||||
if cfg.Distributed.DiskHeadroomDisabled {
|
||||
@@ -418,11 +676,19 @@ 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 {
|
||||
// S3, and it is ONE call rather than a publisher and a subscriber: see
|
||||
// StagingTracker.SetBroadcaster for why a tracker that could name two
|
||||
// carriers is a progress bar that only the originating replica shows.
|
||||
if _, err := wireStagingBroadcasts(bus, router.StagingTracker()); err != nil {
|
||||
xlog.Warn("Failed to subscribe to staging progress broadcasts", "error", err)
|
||||
}
|
||||
|
||||
// Every per-node cache a departure leaves stale, onto the one notification
|
||||
// point, after the router that owns two of them exists.
|
||||
if err := registerDepartureEvictions(departures, prefixDrop, router, galleryProgress, controlClient, fileStager); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Create ReplicaReconciler for auto-scaling model replicas. Adapter +
|
||||
// RegistrationToken feed the state-reconciliation passes: pending op
|
||||
// drain uses the adapter, and model health probes use the token to auth
|
||||
@@ -433,6 +699,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
Unloader: remoteUnloader,
|
||||
Adapter: remoteUnloader,
|
||||
RegistrationToken: cfg.Distributed.RegistrationToken,
|
||||
ClientFactory: backendClients,
|
||||
DB: authDB,
|
||||
Interval: 30 * time.Second,
|
||||
ScaleDownDelay: 5 * time.Minute,
|
||||
@@ -441,28 +708,317 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
PressureThreshold: prefixCfg.PressureScaleThreshold,
|
||||
})
|
||||
|
||||
// Both readers of absence, checked once, here. See requireAbsenceWiring for
|
||||
// why a missing assignment has no other symptom.
|
||||
if err := requireAbsenceWiring(router, healthMon); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Create ModelRouterAdapter to wire into ModelLoader
|
||||
modelAdapter := nodes.NewModelRouterAdapter(router)
|
||||
|
||||
success = true
|
||||
return &DistributedServices{
|
||||
Nats: natsClient,
|
||||
Store: store,
|
||||
Registry: registry,
|
||||
Router: router,
|
||||
Health: healthMon,
|
||||
Reconciler: reconciler,
|
||||
JobStore: jobStore,
|
||||
Dispatcher: dispatcher,
|
||||
AgentStore: agentStore,
|
||||
AgentBridge: agentBridge,
|
||||
DistStores: distStores,
|
||||
FileMgr: fileMgr,
|
||||
FileStager: fileStager,
|
||||
ModelAdapter: modelAdapter,
|
||||
Unloader: remoteUnloader,
|
||||
ModelCleanup: modelCleanup,
|
||||
}, nil
|
||||
ds := &DistributedServices{
|
||||
Store: store,
|
||||
Registry: registry,
|
||||
Router: router,
|
||||
Health: healthMon,
|
||||
Reconciler: reconciler,
|
||||
JobStore: jobStore,
|
||||
Dispatcher: dispatcher,
|
||||
AgentStore: agentStore,
|
||||
AgentBridge: agentBridge,
|
||||
DistStores: distStores,
|
||||
FileMgr: fileMgr,
|
||||
FileStager: fileStager,
|
||||
ModelAdapter: modelAdapter,
|
||||
Unloader: remoteUnloader,
|
||||
ModelCleanup: modelCleanup,
|
||||
Cluster: clusterRegistry,
|
||||
Membership: membership,
|
||||
PeerSessions: peerSessions,
|
||||
Peers: peers,
|
||||
Tunnels: tunnels,
|
||||
WorkerDialer: workerDialer,
|
||||
BackendClients: backendClients,
|
||||
AgentControl: agentControl,
|
||||
JobDispatch: jobDispatch,
|
||||
Bus: bus,
|
||||
}
|
||||
// Checked once, here, on the assembled struct. See requireBroadcastCarrier.
|
||||
if err := requireBroadcastCarrier(ds); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ds, nil
|
||||
}
|
||||
|
||||
// nodeProgressDropper is the gallery service narrowed to the one method a node
|
||||
// departure needs. An interface so the wiring below cannot reach for anything
|
||||
// else on the gallery service from inside an eviction hook.
|
||||
type nodeProgressDropper interface {
|
||||
DropNodeProgress(nodeID string)
|
||||
}
|
||||
|
||||
// The names each per-node cache is registered under. They are constants because
|
||||
// they are asserted: a wiring spec that spelled them itself would keep passing
|
||||
// after a subscriber was renamed and lost.
|
||||
const (
|
||||
departurePrefixCache = "prefix-cache"
|
||||
departureProbeCache = "probe-cache"
|
||||
departureStagingTracker = "staging-tracker"
|
||||
departureGalleryNodes = "gallery-node-progress"
|
||||
departureControlClients = "control-http-clients"
|
||||
departureStagerClients = "file-stager-http-clients"
|
||||
)
|
||||
|
||||
// registerDepartureEvictions registers every per-node cache that a node's
|
||||
// departure leaves stale on the deployment's one departure notification point.
|
||||
//
|
||||
// One function, and every subscriber in it, because the reason the notifier
|
||||
// exists is that a reader cannot otherwise enumerate what a demotion
|
||||
// invalidates: before it, one node type could depart and each stale cache was
|
||||
// dropped from wherever its owner happened to notice. Adding a per-node cache
|
||||
// without adding a line here is the failure this shape exists to make visible,
|
||||
// which is why the subscribers are NAMED and the names are asserted.
|
||||
//
|
||||
// It REFUSES rather than skipping when the router or the gallery service is
|
||||
// missing. A deployment whose departed nodes keep their probe entries, staging
|
||||
// rows and per-node operation progress does not fail, log or slow down: it
|
||||
// answers with state for a node that left, indefinitely.
|
||||
//
|
||||
// prefix may be nil, and only prefix. That is --distributed-prefix-cache=false,
|
||||
// where there is no index to drop from, and it stays a true no-op: nothing is
|
||||
// registered rather than a hook registered onto nothing. It is the CONCRETE
|
||||
// *prefixcache.Sync for that decision to be safe, since a nil provider inside
|
||||
// an interface would compare non-nil here and dereference on the first
|
||||
// departure.
|
||||
func registerDepartureEvictions(departures *nodes.DepartureNotifier, prefix *prefixcache.Sync, router *nodes.SmartRouter, gallery nodeProgressDropper, control *nodes.ControlClient, stager nodes.FileStager) error {
|
||||
if departures == nil {
|
||||
return fmt.Errorf("wiring departure evictions: no departure notifier, so a departed node would keep every per-node cache entry it has for the life of the process")
|
||||
}
|
||||
if router == nil {
|
||||
return fmt.Errorf("wiring departure evictions: no router, so a departed node would keep its probe-freshness entries and its staging operations")
|
||||
}
|
||||
if gallery == nil {
|
||||
return fmt.Errorf("wiring departure evictions: no gallery service, so a departed node would stay in every open operation's per-node breakdown")
|
||||
}
|
||||
if control == nil {
|
||||
return fmt.Errorf("wiring departure evictions: no control client, so a departed node would keep its cached HTTP client and that client's idle streams on a tunnel that is gone")
|
||||
}
|
||||
if stager == nil {
|
||||
return fmt.Errorf("wiring departure evictions: no file stager, so a departed node would keep the cached HTTP client its transfers ran on")
|
||||
}
|
||||
// S1. Inside a nil check and not inside the prefix-cache-enabled block, so
|
||||
// that "the disabled deployment registers nothing" is a fact a spec can
|
||||
// hold rather than a property of where a line was written.
|
||||
if prefix != nil {
|
||||
departures.OnDeparture(departurePrefixCache, func(node nodes.DepartedNode) {
|
||||
prefix.DropNode(node.ID)
|
||||
})
|
||||
}
|
||||
// S2 and S3 are two registrations and not one, because they are two rules:
|
||||
// a probe entry is keyed by node ID and a staging op by node NAME, and a
|
||||
// single hook doing both would hide which of them was lost.
|
||||
departures.OnDeparture(departureProbeCache, func(node nodes.DepartedNode) {
|
||||
router.InvalidateNodeProbes(node.ID)
|
||||
})
|
||||
departures.OnDeparture(departureStagingTracker, func(node nodes.DepartedNode) {
|
||||
router.StagingTracker().DropNode(node.Name)
|
||||
})
|
||||
// S4.
|
||||
departures.OnDeparture(departureGalleryNodes, func(node nodes.DepartedNode) {
|
||||
gallery.DropNodeProgress(node.ID)
|
||||
})
|
||||
// S5 and S6, the two per-node http.Client caches. Two registrations again,
|
||||
// because they are two caches with two owners: the control client's entry
|
||||
// is built on the first verb issued to a node and the stager's on the first
|
||||
// file staged to it, so a node can be in either without being in the other,
|
||||
// and one hook doing both would say only that some client was kept.
|
||||
//
|
||||
// Both are keyed by node ID and both are DROPPED rather than emptied. A
|
||||
// worker that comes back builds a fresh client on its next verb, over
|
||||
// whatever tunnel it has by then; keeping the old one would keep a
|
||||
// transport whose idle streams belong to a session that has ended.
|
||||
departures.OnDeparture(departureControlClients, func(node nodes.DepartedNode) {
|
||||
control.ForgetNode(node.ID)
|
||||
})
|
||||
departures.OnDeparture(departureStagerClients, func(node nodes.DepartedNode) {
|
||||
stager.ForgetNode(node.ID)
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// requireBroadcastCarrier refuses to hand back a distributed deployment whose
|
||||
// broadcast carrier is missing.
|
||||
//
|
||||
// The carrier reaches the deployment over two lines: the newBroadcastBus call
|
||||
// in initDistributed, and the Bus field in the twenty-three field literal
|
||||
// above. Deleting either one compiles and leaves every suite in this repository
|
||||
// green, and the two failures are different. Without the construction, nothing
|
||||
// can ever be published between replicas. Without the assignment the carrier is
|
||||
// opened and connected but Shutdown cannot see it, so every restart leaves a
|
||||
// pinned PostgreSQL session and its goroutines behind until the server runs out
|
||||
// of connections, and the operator sees the failure land on whatever connects
|
||||
// next rather than on LocalAI.
|
||||
//
|
||||
// Neither line can be reddened by a spec today: initDistributed opens NATS
|
||||
// before it reaches any of this, so it cannot be called from a unit test, and a
|
||||
// pointer field left out of a struct literal is not a compile error. What this
|
||||
// converts both omissions into is a deployment that refuses to start and names
|
||||
// what is missing, which is as far as they can be pinned until initDistributed
|
||||
// is testable. The guard itself is spec'd.
|
||||
func requireBroadcastCarrier(ds *DistributedServices) error {
|
||||
if ds == nil || ds.Bus == nil {
|
||||
return fmt.Errorf("distributed mode was initialized without a broadcast carrier: nothing could be published between replicas, and the PostgreSQL session it pins could not be closed on shutdown")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Broadcast is the ONE place a wiring site gets the deployment's fan-out
|
||||
// carrier, and it exists so that "this family travels on the broadcast carrier
|
||||
// and not on NATS" is decided once instead of at every adopter.
|
||||
//
|
||||
// It was five field reads before this: the fine-tune service, the quantization
|
||||
// service, the agent-task setter (twice, on two startup paths), the per-user
|
||||
// services manager and the Open Responses store. Every one of them takes a
|
||||
// messaging.Broadcaster, and the struct used to carry a second field that
|
||||
// satisfied it, so a site that reached for the wrong one compiled, started,
|
||||
// published and was delivered onto a carrier almost nothing read. That second
|
||||
// field is gone with the last family that needed a bus. Collapsing the choice
|
||||
// to one function is what keeps it a fact a spec can pin.
|
||||
//
|
||||
// The return is the interface and not *pgbus.Bus on purpose: handing a nil
|
||||
// *pgbus.Bus to an adopter would produce a non-nil interface wrapping a nil
|
||||
// pointer, and every adopter reads a nil carrier as "standalone, do not
|
||||
// broadcast". A typed nil would instead panic on the first Set. initDistributed
|
||||
// already refuses to return a deployment with no carrier (see
|
||||
// requireBroadcastCarrier), so the nil branch here is belt and braces for a
|
||||
// zero-valued struct in a test.
|
||||
func (ds *DistributedServices) Broadcast() messaging.Broadcaster {
|
||||
if ds == nil || ds.Bus == nil {
|
||||
return nil
|
||||
}
|
||||
return ds.Bus
|
||||
}
|
||||
|
||||
// newBroadcastBus opens the deployment's fan-out carrier on the auth database.
|
||||
//
|
||||
// The DSN is cfg.Auth.DatabaseURL and it may never be anything else. A second
|
||||
// source, a flag of its own or a value read from the environment, would let the
|
||||
// pinned LISTEN connection and the connection pool address two different
|
||||
// databases; that carrier publishes successfully, delivers nothing, on every
|
||||
// replica, and reports no error anywhere. isPostgresURL above has already
|
||||
// refused a value this carrier could not use.
|
||||
//
|
||||
// It is a function rather than four lines inside initDistributed so that the
|
||||
// equality can be pinned by a spec. initDistributed opens NATS before it
|
||||
// reaches this point and so cannot be called from a unit test, which would
|
||||
// leave the assignment as one line in a long function that compiles perfectly
|
||||
// well when it names the wrong field.
|
||||
func newBroadcastBus(ctx context.Context, cfg *config.ApplicationConfig, authDB *gorm.DB) (*pgbus.Bus, error) {
|
||||
if err := pgbus.Migrate(ctx, authDB); err != nil {
|
||||
return nil, fmt.Errorf("migrating the broadcast carrier: %w", err)
|
||||
}
|
||||
bus, err := pgbus.New(ctx, pgbus.Config{DSN: cfg.Auth.DatabaseURL, DB: authDB})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("opening the broadcast carrier: %w", err)
|
||||
}
|
||||
return bus, nil
|
||||
}
|
||||
|
||||
// unadvertisedNagInterval is how often a replica that could not advertise
|
||||
// itself says so again.
|
||||
//
|
||||
// Five minutes is chosen against the log it lands in, not against the urgency:
|
||||
// the condition never clears on its own, so this line is either read once and
|
||||
// acted on or it is noise for the life of the process, and a noisy line gets
|
||||
// filtered rather than fixed. It is still frequent enough that the state is
|
||||
// visible in any window of logs an operator pulls while investigating the
|
||||
// symptom it causes.
|
||||
const unadvertisedNagInterval = 5 * time.Minute
|
||||
|
||||
// nagUnadvertisedReplica repeats, for as long as the process runs, that this
|
||||
// replica is invisible to its peers, and names what that is currently costing.
|
||||
//
|
||||
// It exists because the deferral it accompanies changed cost between phases and
|
||||
// nothing about the deployment says so. Before workers held tunnels, a replica
|
||||
// with no advertised address was merely unreachable BY peers and could still
|
||||
// dial every worker directly, so a startup warning was proportionate. Now a
|
||||
// worker's tunnel lands on one replica and every other replica reaches it by
|
||||
// relaying to the owner, and the owner is resolved by joining the connection
|
||||
// row against a LIVE INSTANCES ROW - which this replica does not have. So every
|
||||
// worker that lands here is answered as unroutable everywhere else: on N
|
||||
// replicas behind round robin, (N-1)/N of that worker's traffic fails, while
|
||||
// this replica serves it perfectly and reports nothing.
|
||||
//
|
||||
// held is passed as a function rather than the registry so this can be driven
|
||||
// without one, and alarm is passed rather than logged inline so a spec can
|
||||
// observe the alarms instead of scraping a log.
|
||||
func nagUnadvertisedReplica(ctx context.Context, held func() []string, every time.Duration, alarm func([]string)) {
|
||||
ticker := time.NewTicker(every)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
alarm(held())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// logUnroutableWorkers says what the state costs RIGHT NOW.
|
||||
//
|
||||
// The two cases are kept apart because they call for different urgency and an
|
||||
// operator can tell them apart at a glance. With no worker held this is a
|
||||
// misconfiguration that has not been paid for yet; with workers held, every one
|
||||
// of them is named, because "which worker is broken" is the question the
|
||||
// symptom sends an operator to ask and the answer is that none of them is.
|
||||
func logUnroutableWorkers(held []string) {
|
||||
if len(held) == 0 {
|
||||
xlog.Warn("This replica is still not registered in the cluster: no advertised address. No worker holds a tunnel here yet; the first that does will be unroutable from every other replica",
|
||||
"knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR")
|
||||
return
|
||||
}
|
||||
xlog.Error("This replica is not registered in the cluster and holds worker tunnels: those workers are unroutable from every OTHER replica, and requests for their models fail there with no route. The workers are healthy; this replica is invisible",
|
||||
"workers", held, "worker_count", len(held), "knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR")
|
||||
}
|
||||
|
||||
// advertisedPeerAddr is the host:port peers dial to reach this replica.
|
||||
//
|
||||
// The operator's value wins outright. Otherwise it is derived from the port
|
||||
// this process serves on and the local address that routes to PostgreSQL, which
|
||||
// is only a peer-reachable answer when the database is on another host;
|
||||
// DiscoverAdvertisedAddr refuses rather than guessing when it is not.
|
||||
func advertisedPeerAddr(cfg *config.ApplicationConfig) (string, error) {
|
||||
if configured := cfg.Distributed.AdvertiseAddr; configured != "" {
|
||||
// A configured address skips discovery, so it also skips every check
|
||||
// discovery makes. Unusable is refused; merely questionable (a
|
||||
// loopback address, correct on one host and wrong on three) is said
|
||||
// once and honoured, because refusing it would refuse single-host
|
||||
// deployments that use it correctly.
|
||||
reason, err := cluster.CheckAdvertisedAddr(configured)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if reason != "" {
|
||||
xlog.Warn("Configured peer address is not one another host can dial",
|
||||
"address", configured, "reason", reason, "knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR")
|
||||
}
|
||||
return configured, nil
|
||||
}
|
||||
if cfg.APIAddress == "" {
|
||||
return "", fmt.Errorf("no API address to derive a peer port from")
|
||||
}
|
||||
_, port, err := net.SplitHostPort(cfg.APIAddress)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("reading the peer port out of API address %q: %w", cfg.APIAddress, err)
|
||||
}
|
||||
portNumber, err := strconv.Atoi(port)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("API address %q has a non-numeric port: %w", cfg.APIAddress, err)
|
||||
}
|
||||
return cluster.DiscoverAdvertisedAddr(cfg.Auth.DatabaseURL, portNumber)
|
||||
}
|
||||
|
||||
func isPostgresURL(url string) bool {
|
||||
|
||||
@@ -0,0 +1,294 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"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/pgbus"
|
||||
"github.com/mudler/LocalAI/core/services/syncstate"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// The guard on the one setting that decides whether any broadcast in the
|
||||
// deployment is ever delivered.
|
||||
//
|
||||
// The carrier holds a pinned LISTEN connection opened from a DSN, and publishes
|
||||
// travel on a pooled handle opened from another. When those two name different
|
||||
// databases every publish succeeds, every subscribe succeeds, and nothing
|
||||
// arrives, on every replica, with no error anywhere. There is exactly one
|
||||
// legitimate DSN, and these specs are what say so in a way that fails when it
|
||||
// stops being true.
|
||||
var _ = Describe("opening the deployment's broadcast carrier", func() {
|
||||
It("listens on the same database URL the auth pool was built from", func() {
|
||||
db, dsn := testutil.SetupTestDBWithDSN()
|
||||
cfg := &config.ApplicationConfig{}
|
||||
cfg.Auth.DatabaseURL = dsn
|
||||
|
||||
bus, err := newBroadcastBus(context.Background(), cfg, db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(bus.Close)
|
||||
|
||||
// Equality with the field, not "is a PostgreSQL URL": the failure being
|
||||
// excluded is two databases, and any DSN passes a shape check.
|
||||
Expect(bus.DSN()).To(Equal(cfg.Auth.DatabaseURL))
|
||||
})
|
||||
|
||||
It("migrates the spill table, so an oversized broadcast has somewhere to go", func() {
|
||||
db, dsn := testutil.SetupTestDBWithDSN()
|
||||
cfg := &config.ApplicationConfig{}
|
||||
cfg.Auth.DatabaseURL = dsn
|
||||
|
||||
bus, err := newBroadcastBus(context.Background(), cfg, db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(bus.Close)
|
||||
|
||||
Expect(db.Migrator().HasTable(&pgbus.BusMessage{})).To(BeTrue())
|
||||
})
|
||||
|
||||
It("refuses to open a carrier whose DSN is not the pool's database", func() {
|
||||
db, _ := testutil.SetupTestDBWithDSN()
|
||||
_, otherDSN := testutil.SetupTestDBWithDSN()
|
||||
cfg := &config.ApplicationConfig{}
|
||||
cfg.Auth.DatabaseURL = otherDSN
|
||||
|
||||
_, err := newBroadcastBus(context.Background(), cfg, db)
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
// The partial pin on two wiring lines that cannot be reddened by a spec: the
|
||||
// newBroadcastBus call, and `Bus: bus` in the returned literal. Neither is a
|
||||
// compile error when deleted and initDistributed cannot be unit tested while it
|
||||
// opens NATS first, so what is available is a boot refusal, and this is what
|
||||
// keeps that refusal honest.
|
||||
var _ = Describe("refusing a deployment with no broadcast carrier", func() {
|
||||
It("accepts services that carry one", func() {
|
||||
db, dsn := testutil.SetupTestDBWithDSN()
|
||||
cfg := &config.ApplicationConfig{}
|
||||
cfg.Auth.DatabaseURL = dsn
|
||||
bus, err := newBroadcastBus(context.Background(), cfg, db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(bus.Close)
|
||||
|
||||
Expect(requireBroadcastCarrier(&DistributedServices{Bus: bus})).To(Succeed())
|
||||
})
|
||||
|
||||
It("refuses services whose carrier was never assigned, and says what it costs", func() {
|
||||
err := requireBroadcastCarrier(&DistributedServices{})
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("published between replicas"))
|
||||
Expect(err.Error()).To(ContainSubstring("shutdown"))
|
||||
})
|
||||
|
||||
It("refuses a nil deployment rather than dereferencing it", func() {
|
||||
Expect(requireBroadcastCarrier(nil)).ToNot(Succeed())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("shutting the distributed services down", func() {
|
||||
It("closes the broadcast carrier", func() {
|
||||
// A pinned PostgreSQL session and the goroutine parked on it, per
|
||||
// replica restart. Nothing else in this process ever closes it, so the
|
||||
// line in the shutdown closure is the whole lifecycle.
|
||||
db, dsn := testutil.SetupTestDBWithDSN()
|
||||
cfg := &config.ApplicationConfig{}
|
||||
cfg.Auth.DatabaseURL = dsn
|
||||
bus, err := newBroadcastBus(context.Background(), cfg, db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(bus.IsConnected()).To(BeTrue())
|
||||
|
||||
(&DistributedServices{Bus: bus}).Shutdown()
|
||||
|
||||
Expect(bus.IsConnected()).To(BeFalse())
|
||||
})
|
||||
})
|
||||
|
||||
// The one place the four state.*.delta families are told which carrier they
|
||||
// travel on.
|
||||
//
|
||||
// It was five field reads before this: the fine-tune service, the quantization
|
||||
// service, the agent-task setter on two startup paths, the per-user services
|
||||
// manager and the Open Responses store. Every one of them takes a
|
||||
// messaging.Broadcaster, which the broker client on the same struct satisfied
|
||||
// too, so a site left holding that field compiled, started, published and was
|
||||
// delivered onto a carrier only agent workers read, and nothing failed until
|
||||
// the broker went away. That field and its type are now gone; collapsing the
|
||||
// choice into one function is what keeps it a fact these specs can hold, rather
|
||||
// than a property that lasted only as long as there was one carrier.
|
||||
var _ = Describe("handing the broadcast carrier to its adopters", func() {
|
||||
It("returns the carrier the deployment opened", func() {
|
||||
db, dsn := testutil.SetupTestDBWithDSN()
|
||||
cfg := &config.ApplicationConfig{}
|
||||
cfg.Auth.DatabaseURL = dsn
|
||||
bus, err := newBroadcastBus(context.Background(), cfg, db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(bus.Close)
|
||||
|
||||
// Identity and not "is a Broadcaster". There is no second carrier on
|
||||
// this struct any more: the family that needed one, agent.<name>.cancel,
|
||||
// rides the agent worker's own tunnel now. The identity assertion stays
|
||||
// because what it pins is that adopters get THIS bus rather than
|
||||
// anything else that satisfies the interface.
|
||||
ds := &DistributedServices{Bus: bus}
|
||||
|
||||
Expect(ds.Broadcast()).To(BeIdenticalTo(messaging.Broadcaster(bus)))
|
||||
})
|
||||
|
||||
It("returns an interface that reads as absent, not a typed nil, when there is no carrier", func() {
|
||||
// Every adopter branches on `bus == nil` to mean standalone. A nil
|
||||
// *pgbus.Bus placed in an interface is NOT nil, so that branch would be
|
||||
// skipped and the first Set would panic on a request rather than at
|
||||
// boot.
|
||||
//
|
||||
// Compared with == and not with BeNil(). Gomega's BeNil reports a nil
|
||||
// POINTER inside an interface as nil, so it passes on exactly the value
|
||||
// this spec exists to reject; the first draft of this spec did, and the
|
||||
// mutation that removed the guard stayed green.
|
||||
var ds *DistributedServices
|
||||
Expect(ds.Broadcast() == nil).To(BeTrue(), "a nil deployment must yield an interface that is itself nil")
|
||||
Expect((&DistributedServices{}).Broadcast() == nil).To(BeTrue(),
|
||||
"a deployment with no carrier must yield an interface that is itself nil, not one wrapping a nil *pgbus.Bus")
|
||||
})
|
||||
|
||||
It("gives an adopter a carrier-less map rather than one that panics on the first write", func() {
|
||||
// The consequence, driven through the component every adopter builds.
|
||||
// A typed nil satisfies `!= nil`, so Start subscribes on it and Set
|
||||
// publishes on it, and both dereference a nil *pgbus.Bus on a request
|
||||
// path rather than at boot.
|
||||
m := syncstate.New(syncstate.Config[string, string]{
|
||||
Name: "test.jobs",
|
||||
Key: func(v string) string { return v },
|
||||
Bus: (&DistributedServices{}).Broadcast(),
|
||||
})
|
||||
Expect(m.Start(context.Background())).To(Succeed())
|
||||
DeferCleanup(func() { Expect(m.Close()).To(Succeed()) })
|
||||
|
||||
Expect(func() { Expect(m.Set(context.Background(), "v")).To(Succeed()) }).ToNot(Panic())
|
||||
})
|
||||
})
|
||||
|
||||
// The registration of every per-node cache a node's departure evicts.
|
||||
//
|
||||
// Four subscribers, registered in one function, on a notifier the health
|
||||
// monitor is then handed. Every one of those is a line that compiles, starts
|
||||
// and serves when it is missing: a deployment whose departed nodes keep their
|
||||
// probe entries, staging rows, prefix affinity and per-node operation progress
|
||||
// does not fail, log or slow down, it just answers with state for a node that
|
||||
// left, for the life of the process.
|
||||
//
|
||||
// Asserted by NAME and not by count. A count says a cache was forgotten; only
|
||||
// the names say which, and "which" is the entire content of the failure.
|
||||
var _ = Describe("wiring the per-node caches a departure evicts", func() {
|
||||
// bootDistributed brings a real distributed deployment up against a fresh
|
||||
// database, which is what makes these assertions about production wiring
|
||||
// rather than about a notifier a spec assembled itself.
|
||||
bootDistributed := func(arm ...func(*config.ApplicationConfig)) *DistributedServices {
|
||||
GinkgoHelper()
|
||||
db, dsn := testutil.SetupTestDBWithDSN()
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
DeferCleanup(cancel)
|
||||
cfg := &config.ApplicationConfig{DataPath: GinkgoT().TempDir(), Context: ctx}
|
||||
cfg.Auth.Enabled = true
|
||||
cfg.Auth.DatabaseURL = dsn
|
||||
cfg.Distributed.Enabled = true
|
||||
for _, a := range arm {
|
||||
a(cfg)
|
||||
}
|
||||
ds, err := initDistributed(cfg, db, nil, galleryop.NewGalleryService(cfg, nil))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(ds.Shutdown)
|
||||
return ds
|
||||
}
|
||||
|
||||
It("registers every one of them, on the notifier the health monitor fires", func() {
|
||||
// Reached through the health monitor and not through a local variable,
|
||||
// because registering the four on a DIFFERENT notifier than the one the
|
||||
// monitor was built with evicts nothing while every count still reads
|
||||
// four.
|
||||
ds := bootDistributed()
|
||||
|
||||
Expect(ds.Health.Departures().SubscriberNames()).To(ConsistOf(
|
||||
departurePrefixCache,
|
||||
departureProbeCache,
|
||||
departureStagingTracker,
|
||||
departureGalleryNodes,
|
||||
departureControlClients,
|
||||
departureStagerClients,
|
||||
))
|
||||
})
|
||||
|
||||
It("registers no prefix-cache eviction when prefix-cache routing is disabled", func() {
|
||||
// --distributed-prefix-cache=false stays a TRUE no-op: there is no
|
||||
// index to drop from, so nothing is registered rather than a hook
|
||||
// registered onto nothing. The other three are unaffected, which is the
|
||||
// half that makes this a statement about S1 and not about the feature
|
||||
// flag switching the whole mechanism off.
|
||||
ds := bootDistributed(func(cfg *config.ApplicationConfig) {
|
||||
cfg.Distributed.PrefixCacheDisabled = true
|
||||
})
|
||||
|
||||
Expect(ds.Health.Departures().SubscriberNames()).To(ConsistOf(
|
||||
departureProbeCache,
|
||||
departureStagingTracker,
|
||||
departureGalleryNodes,
|
||||
departureControlClients,
|
||||
departureStagerClients,
|
||||
))
|
||||
})
|
||||
|
||||
It("refuses a deployment with no router, naming what its departed nodes would keep", func() {
|
||||
err := registerDepartureEvictions(nodes.NewDepartureNotifier(), nil, nil, galleryop.NewGalleryService(&config.ApplicationConfig{}, nil), specControlClient(), specFileStager())
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("probe-freshness"))
|
||||
})
|
||||
|
||||
It("refuses a deployment with no gallery service", func() {
|
||||
err := registerDepartureEvictions(nodes.NewDepartureNotifier(), nil, nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{}), nil, specControlClient(), specFileStager())
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("per-node breakdown"))
|
||||
})
|
||||
|
||||
It("refuses a deployment with no control client, naming the streams a departed node would keep", func() {
|
||||
err := registerDepartureEvictions(nodes.NewDepartureNotifier(), nil, nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{}), galleryop.NewGalleryService(&config.ApplicationConfig{}, nil), nil, specFileStager())
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("idle streams"))
|
||||
})
|
||||
|
||||
It("refuses a deployment with no file stager", func() {
|
||||
err := registerDepartureEvictions(nodes.NewDepartureNotifier(), nil, nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{}), galleryop.NewGalleryService(&config.ApplicationConfig{}, nil), specControlClient(), nil)
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("cached HTTP client"))
|
||||
})
|
||||
|
||||
It("refuses a deployment with no notifier at all", func() {
|
||||
err := registerDepartureEvictions(nil, nil, nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{}), galleryop.NewGalleryService(&config.ApplicationConfig{}, nil), specControlClient(), specFileStager())
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
// specControlClient and specFileStager build the two per-node client caches a
|
||||
// registration needs. Neither is dialled: what the refusal specs assert is that
|
||||
// the wiring refuses a MISSING one, and what the registration specs assert is
|
||||
// that a subscriber was registered for it.
|
||||
func specControlClient() *nodes.ControlClient {
|
||||
return nodes.NewControlClient(nil, "")
|
||||
}
|
||||
|
||||
func specFileStager() nodes.FileStager {
|
||||
return nodes.NewHTTPFileStager(func(string) (string, error) { return "", nil }, "", nil)
|
||||
}
|
||||
@@ -0,0 +1,92 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/agents"
|
||||
"github.com/mudler/LocalAI/core/services/jobs"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/pgbus"
|
||||
"github.com/mudler/xlog"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// newFanoutBridges builds the three surfaces whose cross-replica traffic is job
|
||||
// and agent fan-out, on the deployment's ONE broadcast carrier.
|
||||
//
|
||||
// One function rather than two constructor calls in the start-up path, and the
|
||||
// parameter type is the reason. jobs.NewDispatcher and agents.NewEventBridge
|
||||
// both take a messaging.Broadcaster, which they must: neither may know which
|
||||
// carrier a deployment runs, and their specs publish through a double. But that
|
||||
// also means ANY type satisfying the interface can be wired to them, so while a
|
||||
// second carrier existed, pointing either of them at it COMPILED, passed every
|
||||
// unit spec in both packages, and presented only as an SSE stream that stayed
|
||||
// empty while the work it was watching ran to completion on the other side of a
|
||||
// carrier nobody subscribed to. The second carrier went with the message
|
||||
// broker. Naming *pgbus.Bus here is what would make that a build failure rather
|
||||
// than a silent one again, and it is why this exists as a function instead of
|
||||
// as two lines and a comment asking the reader to be careful.
|
||||
//
|
||||
// The observable persister is STARTED here for the reason startJobDispatchLoop
|
||||
// starts its loop: a bridge that is built and never subscribed captures nothing
|
||||
// a worker publishes, and the only symptom is observables that quietly stop
|
||||
// being written on a deployment that has workers.
|
||||
//
|
||||
// The re-broadcaster is built HERE and handed to the dispatch loop rather than
|
||||
// built there from a carrier of its own, and that is the whole reason this
|
||||
// function returns three things. It is the surface a reviewer skips: it is what
|
||||
// turns an agent worker's progress line into a broadcast, so pointing it at a
|
||||
// carrier nobody subscribes to leaves every unit spec in every package passing.
|
||||
// The re-broadcaster publishes, the publish succeeds, Handle returns true, and
|
||||
// the only symptom in the deployment is an SSE stream with no progress in it.
|
||||
// Written as one expression shared with the dispatcher and the bridge, that
|
||||
// mis-wiring stops being a line a spec has to guess at: there is no second
|
||||
// carrier in scope to point it at.
|
||||
// canceller is NOT a carrier and must never become one. Every family this
|
||||
// function wires has both of its ends on a frontend replica and so travels on
|
||||
// the broadcast carrier, with one exception: the process that holds a
|
||||
// worker-run agent's cancel function is the agent WORKER, which has no database
|
||||
// and cannot join the PostgreSQL carrier at all. A cancel published on bus
|
||||
// would reach no worker, and every cancel of a worker-run agent would be lost
|
||||
// while reporting success. So it does not travel on bus, or on any bus: it is a
|
||||
// control RPC on the tunnel the worker already holds, issued by the agent
|
||||
// control client this takes.
|
||||
func newFanoutBridges(bus *pgbus.Bus, canceller agents.AgentWorkerCanceller,
|
||||
jobStore *jobs.JobStore, agentStore *agents.AgentStore,
|
||||
db *gorm.DB, instanceID string) (*jobs.Dispatcher, *agents.EventBridge, *nodes.Rebroadcaster, error) {
|
||||
// A nil check on the CONCRETE pointer, before it is widened. Once it is a
|
||||
// messaging.Broadcaster a nil *pgbus.Bus is a non-nil interface holding a
|
||||
// nil pointer, so every guard downstream reads it as a carrier that is
|
||||
// present and every publish through it panics on a request instead.
|
||||
if bus == nil {
|
||||
return nil, nil, nil, fmt.Errorf("the job and agent fan-out bridges were built with no broadcast carrier: every job's progress and every agent's events would reach no SSE stream in the deployment")
|
||||
}
|
||||
|
||||
if canceller == nil {
|
||||
return nil, nil, nil, fmt.Errorf("the agent event bridge was built with no way to cancel a worker-run agent: every cancel of one would be sent nowhere")
|
||||
}
|
||||
|
||||
dispatcher := jobs.NewDispatcher(jobStore, bus, db, instanceID)
|
||||
bridge := agents.NewEventBridge(bus, agentStore, instanceID, canceller)
|
||||
|
||||
// Warned rather than refused, and deliberately: the persister needs a store
|
||||
// and a deployment without one still serves live SSE correctly. What it
|
||||
// loses is the durable copy of a worker's observables, which is degraded
|
||||
// rather than broken.
|
||||
if err := bridge.StartObservablePersister(); err != nil {
|
||||
xlog.Warn("Failed to start observable persister", "error", err)
|
||||
} else {
|
||||
xlog.Info("Observable persister started")
|
||||
}
|
||||
|
||||
return dispatcher, bridge, nodes.NewRebroadcaster(bus), nil
|
||||
}
|
||||
|
||||
// The real canceller, asserted where both packages are already imported.
|
||||
//
|
||||
// agents may not import nodes, so the port there is satisfied structurally and
|
||||
// a signature drift between the two would otherwise surface as a nil field in
|
||||
// this file's argument list, which compiles. Here it is a build failure.
|
||||
var _ agents.AgentWorkerCanceller = (*nodes.AgentControlClient)(nil)
|
||||
@@ -0,0 +1,216 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"runtime"
|
||||
"strings"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/agents"
|
||||
"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/pgbus"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// The three fan-out surfaces, asserted from the OTHER replica's carrier.
|
||||
//
|
||||
// Everything here publishes on a second Bus and asserts on the effect the
|
||||
// surface built on the first one had. One bus talking to itself would pass with
|
||||
// the surfaces wired to any carrier at all, which is exactly the wiring defect
|
||||
// these exist to catch: a dispatcher on one carrier and a publisher on another
|
||||
// leaves every unit spec in both packages green and only an SSE stream empty.
|
||||
var _ = Describe("wiring the job and agent fan-out bridges", func() {
|
||||
var (
|
||||
ctx context.Context
|
||||
db *gorm.DB
|
||||
busA, busB *pgbus.Bus
|
||||
jobStore *jobs.JobStore
|
||||
agentStore *agents.AgentStore
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
ctx = context.Background()
|
||||
|
||||
var dsn string
|
||||
db, dsn = testutil.SetupTestDBWithDSN()
|
||||
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
|
||||
|
||||
newBus := func() *pgbus.Bus {
|
||||
b, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(b.Close)
|
||||
return b
|
||||
}
|
||||
busA, busB = newBus(), newBus()
|
||||
|
||||
var err error
|
||||
jobStore, err = jobs.NewJobStore(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
agentStore, err = agents.NewAgentStore(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
})
|
||||
|
||||
// The canceller is a SEPARATE argument, and its absence is refused
|
||||
// separately. A bridge built without one has nowhere to send the cancel of
|
||||
// an agent running on a worker, and the failure would present as
|
||||
// CancelExecution reporting success on a cancel that reached nobody.
|
||||
It("refuses to build with no way to cancel a worker-run agent", func() {
|
||||
_, _, _, err := newFanoutBridges(busA, nil, jobStore, agentStore, db, "replica-1")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("cancel a worker-run agent"))
|
||||
})
|
||||
|
||||
It("refuses to build with no carrier", func() {
|
||||
_, _, _, err := newFanoutBridges(nil, &stubCanceller{}, jobStore, agentStore, db, "replica-1")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("broadcast carrier"))
|
||||
})
|
||||
|
||||
// S1. The dispatcher persists a terminal result broadcast by a peer, which
|
||||
// it can only do if its wildcard subscription is on the carrier the peer
|
||||
// published to.
|
||||
It("subscribes the job dispatcher to results a peer replica broadcasts", func() {
|
||||
dispatcher, _, _, err := newFanoutBridges(busA, &stubCanceller{}, jobStore, agentStore, db, "replica-1")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(dispatcher.Start(ctx)).To(Succeed())
|
||||
DeferCleanup(dispatcher.Stop)
|
||||
|
||||
job := &jobs.JobRecord{TaskID: "t1", UserID: "u1", Status: "running", TriggeredBy: "manual"}
|
||||
Expect(jobStore.CreateJob(job)).To(Succeed())
|
||||
|
||||
Expect(busB.Publish(messaging.SubjectJobResult(job.ID), jobs.JobResultEvent{
|
||||
JobID: job.ID, Status: "completed", Result: "the answer",
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() string {
|
||||
stored, err := jobStore.GetJob(job.ID)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return stored.Status
|
||||
}, "20s").Should(Equal("completed"))
|
||||
})
|
||||
|
||||
// S1b. A cancel does NOT travel on a carrier, asserted by where it goes and
|
||||
// by where it does not.
|
||||
//
|
||||
// The nil refusal above only says a canceller was passed. What it cannot
|
||||
// say is that the bridge uses it instead of publishing onto the broadcast
|
||||
// carrier, which is the edit anyone finishing this migration would reach
|
||||
// for: it compiles, it publishes successfully onto PostgreSQL, and the
|
||||
// agent worker that has to act on the cancel is not and cannot be there.
|
||||
// Every unit suite stays green and every cancel of a worker-run agent is
|
||||
// lost while CancelExecution returns nil.
|
||||
//
|
||||
// So this asserts the cancel reaches the CANCELLER and, in the same spec,
|
||||
// that nothing is published on a peer replica's broadcast carrier. The
|
||||
// negative half is the load-bearing one: the positive half alone passes for
|
||||
// a bridge that does both.
|
||||
It("sends an agent cancel to the agent workers and publishes nothing on the broadcast carrier", func() {
|
||||
canceller := &stubCanceller{}
|
||||
_, bridge, _, err := newFanoutBridges(busA, canceller, jobStore, agentStore, db, "replica-1")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
onBroadcast := make(chan []byte, 4)
|
||||
_, err = busB.Subscribe(messaging.SubjectAgentCancelWildcard, func(data []byte) { onBroadcast <- data })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(bridge.CancelExecution(ctx, "a1", "u1", "msg-1")).To(Succeed())
|
||||
|
||||
Expect(canceller.requests).To(ConsistOf(messaging.AgentCancelRequest{
|
||||
AgentName: "a1", UserID: "u1", MessageID: "msg-1",
|
||||
}), "the cancel did not reach the agent workers, so it reached nobody and was reported as sent")
|
||||
Consistently(onBroadcast, "2s").ShouldNot(Receive(),
|
||||
"the agent cancel was published on the broadcast carrier, where no agent worker is or can be subscribed")
|
||||
})
|
||||
|
||||
// S2. The observable persister writes what a peer broadcast, which it can
|
||||
// only do if it was started AND is on the same carrier AND its filter has
|
||||
// the right number of tokens.
|
||||
It("subscribes the agent observable persister to events a peer replica broadcasts", func() {
|
||||
_, bridge, _, err := newFanoutBridges(busA, &stubCanceller{}, jobStore, agentStore, db, "replica-1")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(bridge).ToNot(BeNil())
|
||||
|
||||
Expect(busB.Publish(messaging.SubjectAgentEvents("a1", "u1"), agents.AgentEvent{
|
||||
AgentName: "a1",
|
||||
UserID: "u1",
|
||||
EventType: "observable_update",
|
||||
EventSubType: "tool_result",
|
||||
SourceInstance: "replica-2",
|
||||
MessageID: "obs-1",
|
||||
Metadata: `{"tool":"grep"}`,
|
||||
})).To(Succeed())
|
||||
|
||||
Eventually(func() int {
|
||||
records, err := agentStore.GetObservables(agents.AgentKey("u1", "a1"), 10)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return len(records)
|
||||
}, "20s").Should(Equal(1))
|
||||
})
|
||||
|
||||
// S3, and it is the one this whole arrangement is for. The re-broadcaster
|
||||
// is what turns an agent worker's progress line into a broadcast, and it is
|
||||
// the surface a reviewer skips, because pointing it at a carrier nobody
|
||||
// reads leaves every unit spec in every package green: it publishes, the
|
||||
// publish succeeds, and Handle returns true.
|
||||
//
|
||||
// So these assert the RECEIPT on a peer's carrier and never Handle's return
|
||||
// value, which is true for a publish that went nowhere.
|
||||
DescribeTable("re-broadcasts a worker's line onto the carrier a peer replica reads",
|
||||
func(subject string, payload string) {
|
||||
_, _, rebroadcast, err := newFanoutBridges(busA, &stubCanceller{}, jobStore, agentStore, db, "replica-1")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
delivered := make(chan []byte, 4)
|
||||
_, err = busB.Subscribe(subject, func(data []byte) { delivered <- data })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
rebroadcast.Handle(nodes.NodeTypeAgent, strings.ReplaceAll(subject, "*", "j1"), json.RawMessage(payload))
|
||||
|
||||
Eventually(delivered, "20s").Should(Receive(MatchJSON(payload)))
|
||||
},
|
||||
Entry("a job's progress", messaging.SubjectJobProgressWildcard, `{"job_id":"j1","status":"running"}`),
|
||||
Entry("a job's result", messaging.SubjectJobResultWildcard, `{"job_id":"j1","status":"completed"}`),
|
||||
)
|
||||
|
||||
It("re-broadcasts an agent's events onto the carrier a peer replica reads", func() {
|
||||
_, _, rebroadcast, err := newFanoutBridges(busA, &stubCanceller{}, jobStore, agentStore, db, "replica-1")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
delivered := make(chan []byte, 4)
|
||||
_, err = busB.Subscribe(messaging.SubjectAgentEventsWildcard, func(data []byte) { delivered <- data })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
rebroadcast.Handle(nodes.NodeTypeAgent, messaging.SubjectAgentEvents("a1", "u1"),
|
||||
json.RawMessage(`{"event_type":"json_message"}`))
|
||||
|
||||
Eventually(delivered, "20s").Should(Receive(MatchJSON(`{"event_type":"json_message"}`)))
|
||||
})
|
||||
})
|
||||
|
||||
// stubCanceller stands in for the frontend's agent control client, which
|
||||
// reaches workers over their tunnels and is driven over a real one in
|
||||
// core/services/nodes. What these specs need from it is only that the bridge
|
||||
// asks it at all.
|
||||
type stubCanceller struct {
|
||||
requests []messaging.AgentCancelRequest
|
||||
}
|
||||
|
||||
func (s *stubCanceller) CancelAgentRun(_ context.Context, req messaging.AgentCancelRequest) error {
|
||||
s.requests = append(s.requests, req)
|
||||
return nil
|
||||
}
|
||||
@@ -271,7 +271,11 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
// the model configs are loaded, so it is declared out here.
|
||||
var revisionStore modeladmin.RevisionStore
|
||||
|
||||
distSvc, err := initDistributed(options, application.authDB, application.ModelConfigLoader())
|
||||
// The gallery service is handed in rather than set afterwards: it owns one
|
||||
// of the per-node caches a node departure evicts, and every one of those is
|
||||
// registered inside initDistributed. It exists by now because start() built
|
||||
// it above.
|
||||
distSvc, err := initDistributed(options, application.authDB, application.ModelConfigLoader(), application.galleryService)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("distributed mode initialization failed: %w", err)
|
||||
}
|
||||
@@ -283,9 +287,15 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
// Wire ModelRouter so grpcModel() delegates to SmartRouter in distributed mode
|
||||
application.modelLoader.SetModelRouter(distSvc.ModelAdapter.AsModelRouter())
|
||||
// Wire DistributedModelStore so shutdown/list/watchdog can find remote models
|
||||
// The client factory is not optional here. Without it the store builds
|
||||
// remote models with no client, and pkg/model.Model.GRPC then dials the
|
||||
// worker's raw address with gRPC's own dialer, which is the direct dial
|
||||
// the tunnel replaces; ShutdownModel's Free and the backend monitor's
|
||||
// Status both reach it.
|
||||
distStore := nodes.NewDistributedModelStore(
|
||||
model.NewInMemoryModelStore(),
|
||||
distSvc.Registry,
|
||||
distSvc.BackendClients,
|
||||
)
|
||||
application.modelLoader.SetModelStore(distStore)
|
||||
// Drop the local stub when a model's last replica leaves the registry.
|
||||
@@ -322,15 +332,16 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
application.agentJobService.SetDistributedBackends(distSvc.Dispatcher)
|
||||
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)
|
||||
// dispatcher + DB read-through), on the deployment's broadcast carrier.
|
||||
application.agentJobService.SetTaskSyncBus(distSvc.Broadcast())
|
||||
}
|
||||
// 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
|
||||
// Wire the broadcast carrier and gallery store into GalleryService for
|
||||
// cross-instance progress/cancel. The carrier is wired below, next to
|
||||
// the subscriptions it feeds, so the two cannot name different carriers.
|
||||
if application.galleryService != nil {
|
||||
application.galleryService.SetNATSClient(distSvc.Nats)
|
||||
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 {
|
||||
@@ -390,7 +401,12 @@ func New(opts ...config.AppOption) (*Application, error) {
|
||||
xlog.Warn("Failed to apply peer model config change", "error", err)
|
||||
}
|
||||
}
|
||||
if err := application.galleryService.SubscribeBroadcasts(); err != nil {
|
||||
// S2. One call sets the carrier and opens the wildcard
|
||||
// subscriptions, and it names no carrier at all, so no carrier that
|
||||
// happens to hang off distSvc can be passed here by accident. See
|
||||
// cache_fanout_wiring.go for why that shape is kept now that the
|
||||
// broker's client is no longer one of them.
|
||||
if err := distSvc.wireGallery(application.galleryService); err != nil {
|
||||
xlog.Warn("Gallery service subscribe failed", "error", err)
|
||||
}
|
||||
// Wire distributed model/backend managers so delete propagates to workers
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package application
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// The runtime symptom for a deferral whose cost changed between phases.
|
||||
//
|
||||
// Not refusing to start without an advertised address stays deferred on
|
||||
// purpose: refusing would take out every single-host deployment. What is not
|
||||
// deferred is telling the operator, repeatedly, that this replica is invisible
|
||||
// and which workers that is costing - because the symptom it produces (a worker
|
||||
// that 5xxs from most of the fleet) reads as a worker problem, and a single
|
||||
// startup line has scrolled away long before anyone goes looking.
|
||||
var _ = Describe("the alarm for a replica with no advertised address", func() {
|
||||
It("keeps firing for as long as the state lasts, and names the workers it costs", func() {
|
||||
// Repetition is the property. A one-shot alarm is the startup line
|
||||
// again, which is what was already there and was not enough.
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
DeferCleanup(cancel)
|
||||
|
||||
alarms := make(chan []string, 8)
|
||||
go nagUnadvertisedReplica(ctx, func() []string { return []string{"w1", "w2"} },
|
||||
time.Millisecond, func(held []string) { alarms <- held })
|
||||
|
||||
// Two, not one: the second is what a one-shot implementation fails.
|
||||
var first, second []string
|
||||
Eventually(alarms, "10s").Should(Receive(&first))
|
||||
Eventually(alarms, "10s").Should(Receive(&second))
|
||||
Expect(first).To(ConsistOf("w1", "w2"),
|
||||
"the workers this is costing are the answer to the question the symptom provokes")
|
||||
Expect(second).To(ConsistOf("w1", "w2"))
|
||||
})
|
||||
|
||||
It("reads the held set on every tick rather than the one it started with", func() {
|
||||
// A replica accumulates tunnels while it runs, so an alarm bound to the
|
||||
// set at startup would name an empty list forever on exactly the
|
||||
// deployment where the cost is real.
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
DeferCleanup(cancel)
|
||||
|
||||
workers := make(chan []string, 32)
|
||||
for range 32 {
|
||||
workers <- []string{"w-late"}
|
||||
}
|
||||
alarms := make(chan []string, 8)
|
||||
go nagUnadvertisedReplica(ctx, func() []string { return <-workers },
|
||||
time.Millisecond, func(held []string) { alarms <- held })
|
||||
|
||||
var got []string
|
||||
Eventually(alarms, "10s").Should(Receive(&got))
|
||||
Expect(got).To(ConsistOf("w-late"))
|
||||
})
|
||||
|
||||
It("stops when the process context ends", func() {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
stopped := make(chan struct{})
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
nagUnadvertisedReplica(ctx, func() []string { return nil }, time.Hour, func([]string) {})
|
||||
close(stopped)
|
||||
}()
|
||||
|
||||
cancel()
|
||||
Eventually(stopped, "10s").Should(BeClosed())
|
||||
})
|
||||
})
|
||||
+312
-221
@@ -16,27 +16,46 @@ import (
|
||||
"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/jobs"
|
||||
mcpRemote "github.com/mudler/LocalAI/core/services/mcp"
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/pkg/sanitize"
|
||||
"github.com/mudler/cogito"
|
||||
"github.com/mudler/cogito/clients"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// 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 and serves agent execution and MCP CI runs as
|
||||
// STREAMING CONTROL VERBS on the tunnel it holds. The worker is a pure
|
||||
// executor: it receives the full agent config and skills in the request body,
|
||||
// so it does not need direct database access.
|
||||
//
|
||||
// It joins no queue group, and there is none left to join: a queue group only
|
||||
// ever selected one consumer out of a set, the frontend makes that selection
|
||||
// itself, and the work it hands over is a row it claimed on the job store.
|
||||
//
|
||||
// It also holds one tunnel to the frontend, so the frontend can reach its
|
||||
// control verbs by RPC without the worker opening an inbound port. No verb the
|
||||
// frontend addresses to THIS worker travels on the bus any more.
|
||||
//
|
||||
// It dials NO message bus. The last family that needed one was
|
||||
// agent.<name>.cancel, which ran the other way, from a frontend replica to
|
||||
// whichever worker held the execution; it is now a control verb on this
|
||||
// worker's own tunnel (workerctl.PathAgentCancel), so a cancel reaches the
|
||||
// worker running the agent without either side touching a broker.
|
||||
//
|
||||
// Usage:
|
||||
//
|
||||
// localai agent-worker --nats-url nats://... --register-to http://localai:8080
|
||||
// localai agent-worker --register-to http://localai:8080
|
||||
type AgentWorkerCMD struct {
|
||||
// NATS (required)
|
||||
NatsURL string `env:"LOCALAI_NATS_URL" required:"" help:"NATS server URL" group:"distributed"`
|
||||
// NatsURL is accepted and ignored, exactly as the backend worker's is (see
|
||||
// core/services/worker/config.go). An agent worker connects to no message
|
||||
// bus: every verb a frontend addresses to it arrives on the tunnel it
|
||||
// dials, and a cancel now arrives the same way. It stays here, without
|
||||
// required, so an existing command line or unit file that still carries
|
||||
// --nats-url starts rather than failing to parse.
|
||||
NatsURL string `env:"LOCALAI_NATS_URL" help:"Ignored. An agent worker connects to no message bus; the frontend reaches it over its outbound tunnel. Accepted so an existing worker command line still starts." group:"distributed" hidden:""`
|
||||
|
||||
// Registration (required)
|
||||
RegisterTo string `env:"LOCALAI_REGISTER_TO" required:"" help:"Frontend URL for registration" group:"registration"`
|
||||
@@ -48,34 +67,58 @@ type AgentWorkerCMD struct {
|
||||
APIURL string `env:"LOCALAI_API_URL" help:"LocalAI API URL for inference (auto-derived from RegisterTo if not set)" group:"api"`
|
||||
APIToken string `env:"LOCALAI_API_TOKEN" help:"API token for LocalAI inference (auto-provisioned during registration if not set)" group:"api"`
|
||||
|
||||
// NATS subjects
|
||||
Subject string `env:"LOCALAI_AGENT_SUBJECT" default:"agent.execute" help:"NATS subject for agent execution" group:"distributed"`
|
||||
Queue string `env:"LOCALAI_AGENT_QUEUE" default:"agent-workers" help:"NATS queue group name" group:"distributed"`
|
||||
|
||||
NatsJWT string `env:"LOCALAI_NATS_JWT" help:"NATS user JWT override (defaults to nats_jwt from registration)" group:"distributed"`
|
||||
NatsUserSeed string `env:"LOCALAI_NATS_USER_SEED" help:"NATS user seed override (defaults to nats_user_seed from registration)" group:"distributed"`
|
||||
NatsServiceJWT string `env:"LOCALAI_NATS_SERVICE_JWT" help:"Fallback NATS service JWT when registration does not mint agent JWT" group:"distributed"`
|
||||
NatsServiceSeed string `env:"LOCALAI_NATS_SERVICE_SEED" help:"Fallback NATS service seed paired with LOCALAI_NATS_SERVICE_JWT" group:"distributed"`
|
||||
NatsRequireAuth bool `env:"LOCALAI_NATS_REQUIRE_AUTH" default:"false" help:"Require NATS JWT+seed to connect" group:"distributed"`
|
||||
// The broker credential and TLS flags, accepted and ignored, hidden, on the
|
||||
// same terms as NatsURL above. There is no connection left to present a
|
||||
// credential on.
|
||||
NatsJWT string `env:"LOCALAI_NATS_JWT" help:"Ignored. An agent worker opens no bus connection to present a credential on." group:"distributed" hidden:""`
|
||||
NatsUserSeed string `env:"LOCALAI_NATS_USER_SEED" help:"Ignored. Paired with LOCALAI_NATS_JWT, which is itself ignored." group:"distributed" hidden:""`
|
||||
NatsServiceJWT string `env:"LOCALAI_NATS_SERVICE_JWT" help:"Ignored. An agent worker opens no bus connection to present a credential on." group:"distributed" hidden:""`
|
||||
NatsServiceSeed string `env:"LOCALAI_NATS_SERVICE_SEED" help:"Ignored. Paired with LOCALAI_NATS_SERVICE_JWT, which is itself ignored." group:"distributed" hidden:""`
|
||||
NatsRequireAuth bool `env:"LOCALAI_NATS_REQUIRE_AUTH" default:"false" help:"Ignored. Use --distributed-require-auth to make this worker wait through admin approval." group:"distributed" hidden:""`
|
||||
// DistributedRequireAuth is the umbrella switch; for the agent worker (which
|
||||
// has no file-transfer server) it implies NATS auth is required.
|
||||
DistributedRequireAuth bool `env:"LOCALAI_DISTRIBUTED_REQUIRE_AUTH" default:"false" help:"Umbrella switch implying --nats-require-auth (agent workers have no file-transfer server)" group:"distributed"`
|
||||
NatsTLSCA string `env:"LOCALAI_NATS_TLS_CA" type:"existingfile" help:"PEM file for NATS server CA (private PKI)" group:"distributed"`
|
||||
NatsTLSCert string `env:"LOCALAI_NATS_TLS_CERT" type:"existingfile" help:"Client certificate for NATS mTLS" group:"distributed"`
|
||||
NatsTLSKey string `env:"LOCALAI_NATS_TLS_KEY" type:"existingfile" help:"Client private key for NATS mTLS" group:"distributed"`
|
||||
// has no file-transfer server) it makes registration WAIT THROUGH ADMIN
|
||||
// APPROVAL rather than starting against a pending node.
|
||||
//
|
||||
// It used to imply --nats-require-auth as well, and the wait was a side
|
||||
// effect of that: the worker was waiting for a broker credential to be
|
||||
// minted. There is no credential and no broker, so the wait is now what the
|
||||
// switch is FOR, and it is described that way rather than by what it used
|
||||
// to imply.
|
||||
DistributedRequireAuth bool `env:"LOCALAI_DISTRIBUTED_REQUIRE_AUTH" default:"false" help:"Wait through admin approval at registration instead of starting against a node an admin has not approved" group:"distributed"`
|
||||
// type:"existingfile" is deliberately NOT kept: validating a path this
|
||||
// process never opens would fail a worker at startup over a certificate for
|
||||
// a broker the operator has already shut down.
|
||||
NatsTLSCA string `env:"LOCALAI_NATS_TLS_CA" help:"Ignored. No bus connection is opened, so no server certificate is verified." group:"distributed" hidden:""`
|
||||
NatsTLSCert string `env:"LOCALAI_NATS_TLS_CERT" help:"Ignored. No bus connection is opened, so no client certificate is presented." group:"distributed" hidden:""`
|
||||
NatsTLSKey string `env:"LOCALAI_NATS_TLS_KEY" help:"Ignored. Paired with LOCALAI_NATS_TLS_CERT, which is itself ignored." group:"distributed" hidden:""`
|
||||
|
||||
// Timeouts
|
||||
MCPCIJobTimeout string `env:"LOCALAI_MCP_CI_JOB_TIMEOUT" default:"10m" help:"Timeout for MCP CI job execution" group:"distributed"`
|
||||
}
|
||||
|
||||
// natsAuthRequired reports whether NATS JWT credentials must be present — the
|
||||
// granular flag or the umbrella (LOCALAI_DISTRIBUTED_REQUIRE_AUTH).
|
||||
func (cmd *AgentWorkerCMD) natsAuthRequired() bool {
|
||||
return cmd.NatsRequireAuth || cmd.DistributedRequireAuth
|
||||
// waitThroughApproval reports whether registration should block until an admin
|
||||
// approves this node, instead of returning a pending response and starting.
|
||||
//
|
||||
// A method rather than the field read it wraps, because the answer CHANGED and
|
||||
// the change is the one thing in this command an operator can be surprised by.
|
||||
// It used to be --distributed-require-auth OR --nats-require-auth, narrowed
|
||||
// further by whether the operator had supplied a broker JWT by hand. Every term
|
||||
// but the first was about a credential that no longer exists, so the gate is
|
||||
// now the first term alone.
|
||||
//
|
||||
// An operator who set ONLY --nats-require-auth therefore loses the wait and
|
||||
// gets the historical default: register, start, and let the tunnel dialer be
|
||||
// refused with 403 until an admin approves. That is a visible change, it is
|
||||
// documented in docs/content/features/distributed-mode.md, and it is a seam so
|
||||
// that it is also pinned: inlined at the call site it sat inside a Run that
|
||||
// dials a frontend, where no spec could reach it and swapping the two flags
|
||||
// back would have stayed green.
|
||||
func (cmd *AgentWorkerCMD) waitThroughApproval() bool {
|
||||
return cmd.DistributedRequireAuth
|
||||
}
|
||||
|
||||
func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
|
||||
xlog.Info("Starting agent worker", "nats", sanitize.URL(cmd.NatsURL), "register_to", cmd.RegisterTo)
|
||||
xlog.Info("Starting agent worker", "register_to", cmd.RegisterTo)
|
||||
|
||||
// Resolve API URL
|
||||
apiURL := cmp.Or(cmd.APIURL, strings.TrimRight(cmd.RegisterTo, "/"))
|
||||
@@ -104,14 +147,17 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
|
||||
shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
|
||||
defer shutdownCancel()
|
||||
|
||||
// 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(
|
||||
// Register, and obtain this node's identity and its tunnel credential.
|
||||
//
|
||||
// The manager holds one thing now: the tunnel token, which every
|
||||
// re-registration rotates. Its gate decides whether registration WAITS
|
||||
// THROUGH ADMIN APPROVAL instead of returning a pending response; which
|
||||
// flag decides that, and what changed about it, is on waitThroughApproval.
|
||||
credMgr := workerregistry.NewCredentialManager(
|
||||
func(ctx context.Context) (*workerregistry.RegisterResponse, error) {
|
||||
return regClient.RegisterFull(ctx, registrationBody)
|
||||
},
|
||||
cmd.natsAuthRequired() && cmd.NatsJWT == "" && cmd.NatsServiceJWT == "",
|
||||
cmd.waitThroughApproval(),
|
||||
)
|
||||
res, err := credMgr.Acquire(shutdownCtx)
|
||||
if err != nil {
|
||||
@@ -134,201 +180,173 @@ func (cmd *AgentWorkerCMD) Run(ctx *cliContext.Context) error {
|
||||
|
||||
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()
|
||||
// The executor and the event bridge the control plane serves, built BEFORE
|
||||
// the tunnel because a verb mounted with a nil handler answers a 404, which
|
||||
// a frontend reads as a worker too old to serve it.
|
||||
//
|
||||
// No ConfigProvider and no SkillStore: config and skills arrive in the
|
||||
// request body, exactly as they arrived in the job payload before, because
|
||||
// an agent worker still has no database.
|
||||
eventBridge := agents.NewWorkerEventBridge("agent-worker-" + nodeID)
|
||||
executor := agents.NewWorkerExecutor(eventBridge, nil, apiURL, cmd.APIToken)
|
||||
|
||||
// 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()
|
||||
}
|
||||
|
||||
// Create and start the NATS dispatcher.
|
||||
// No ConfigProvider or SkillStore needed — config and skills arrive in the job payload.
|
||||
dispatcher := agents.NewNATSDispatcher(
|
||||
natsClient,
|
||||
eventBridge,
|
||||
nil, // no ConfigProvider: config comes in the enriched NATS payload
|
||||
apiURL, cmd.APIToken,
|
||||
cmd.Subject, cmd.Queue,
|
||||
0, // no concurrency limit (CLI worker)
|
||||
)
|
||||
|
||||
if err := dispatcher.Start(shutdownCtx); err != nil {
|
||||
return fmt.Errorf("starting dispatcher: %w", err)
|
||||
}
|
||||
|
||||
// Subscribe to MCP tool execution requests (load-balanced across workers).
|
||||
// The frontend routes model-level MCP tool calls here via NATS request-reply.
|
||||
if _, err := natsClient.QueueSubscribeReply(messaging.SubjectMCPToolExecute, messaging.QueueAgentWorkers, func(data []byte, reply func([]byte)) {
|
||||
handleMCPToolRequest(data, reply)
|
||||
}); err != nil {
|
||||
return fmt.Errorf("subscribing to %s: %w", messaging.SubjectMCPToolExecute, err)
|
||||
}
|
||||
|
||||
// Subscribe to MCP discovery requests (load-balanced across workers).
|
||||
if _, err := natsClient.QueueSubscribeReply(messaging.SubjectMCPDiscovery, messaging.QueueAgentWorkers, func(data []byte, reply func([]byte)) {
|
||||
handleMCPDiscoveryRequest(data, reply)
|
||||
}); err != nil {
|
||||
return fmt.Errorf("subscribing to %s: %w", messaging.SubjectMCPDiscovery, 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 := natsClient.QueueSubscribe(messaging.SubjectMCPCIJobsNew, messaging.QueueWorkers, func(data []byte) {
|
||||
handleMCPCIJob(shutdownCtx, data, apiURL, cmd.APIToken, natsClient, mcpCIJobTimeout)
|
||||
}); err != nil {
|
||||
return fmt.Errorf("subscribing to %s: %w", messaging.SubjectMCPCIJobsNew, err)
|
||||
// The tunnel, and the loopback control plane behind it.
|
||||
//
|
||||
// It is now the ONLY way anything the frontend addresses to THIS worker
|
||||
// arrives: MCP tool execution, MCP discovery, backend.stop, agent execution
|
||||
// and MCP CI runs. Every one of their subjects is gone. The two queue
|
||||
// groups went last, because a queue group was only ever a way of SELECTING
|
||||
// a worker: the frontend makes that selection itself
|
||||
// (nodes.AgentSelector) and hands over a claim it took off the job store.
|
||||
//
|
||||
// The worker opens no inbound port for any of it: it dials out and the
|
||||
// control plane rides the tunnel it holds.
|
||||
//
|
||||
// The credential is read through credMgr rather than captured from res,
|
||||
// because every registration the manager performs ROTATES it: the frontend
|
||||
// stores only the hash of the newest one, so a captured value would lock
|
||||
// this worker out of its own tunnel after any re-registration.
|
||||
//
|
||||
// It is started AFTER registration, which is what supplies both the node
|
||||
// identity the dial names and the credential it presents.
|
||||
agentCtl, err := agentworker.Start(shutdownCtx, agentworker.Options{
|
||||
FrontendURL: cmd.RegisterTo,
|
||||
NodeID: nodeID,
|
||||
TunnelToken: credMgr.TunnelToken,
|
||||
ControlToken: cmd.RegistrationToken,
|
||||
Handlers: agentWorkerControlHandlers(executor, apiURL, cmd.APIToken, mcpCIJobTimeout),
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("starting the agent worker control plane: %w", err)
|
||||
}
|
||||
|
||||
// Subscribe to 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 — we listen for the NATS stop event instead.
|
||||
if _, err := natsClient.Subscribe(messaging.SubjectNodeBackendStop(nodeID), func(data []byte) {
|
||||
var req struct {
|
||||
Backend string `json:"backend"`
|
||||
defer func() {
|
||||
if err := agentCtl.Close(); err != nil {
|
||||
xlog.Warn("Closing the agent worker tunnel failed", "error", err)
|
||||
}
|
||||
if json.Unmarshal(data, &req) == nil && req.Backend != "" {
|
||||
mcpTools.CloseMCPSessions(req.Backend)
|
||||
}
|
||||
}); err != nil {
|
||||
return fmt.Errorf("subscribing to %s: %w", messaging.SubjectNodeBackendStop(nodeID), err)
|
||||
}
|
||||
}()
|
||||
|
||||
xlog.Info("Agent worker ready, waiting for jobs", "subject", cmd.Subject, "queue", cmd.Queue)
|
||||
xlog.Info("Agent worker ready, serving agent execution and MCP CI runs on its tunnel", "node", nodeID)
|
||||
|
||||
// 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.
|
||||
// Wait for an OS signal. There is no internal fatal condition left to wait
|
||||
// on: the one that existed was a broker credential this worker could no
|
||||
// longer renew, and there is no credential and no broker.
|
||||
sigCh := make(chan os.Signal, 1)
|
||||
signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
|
||||
var runErr error
|
||||
select {
|
||||
case <-sigCh:
|
||||
case <-shutdownCtx.Done():
|
||||
runErr = fmt.Errorf("agent worker shutting down: NATS credentials unavailable")
|
||||
xlog.Error("Internal shutdown requested", "error", runErr)
|
||||
}
|
||||
<-sigCh
|
||||
|
||||
xlog.Info("Shutting down agent worker")
|
||||
shutdownCancel() // stop heartbeat loop immediately
|
||||
dispatcher.Stop()
|
||||
mcpTools.CloseAllMCPSessions()
|
||||
regClient.GracefulDeregister(nodeID)
|
||||
return runErr
|
||||
return nil
|
||||
}
|
||||
|
||||
// handleMCPToolRequest handles a NATS request-reply for MCP tool execution.
|
||||
// The worker creates/caches MCP sessions from the serialized config and executes the tool.
|
||||
func handleMCPToolRequest(data []byte, reply func([]byte)) {
|
||||
// The MCP verbs, written ONCE and served on two carriers.
|
||||
//
|
||||
// The bus subscription and the tunnel's control route both call the same
|
||||
// serve* function and both send the same bytes, so a worker reached either way
|
||||
// answers identically. Two implementations of one verb is the shape that lets a
|
||||
// deployment behave differently depending on which carrier a frontend happened
|
||||
// to pick, and there is no version of this migration in which that is
|
||||
// acceptable: for the whole of it, both carriers are live at once.
|
||||
//
|
||||
// The distinction the return type carries: an MCP tool that RAN and failed is
|
||||
// this worker's own answer and travels as bytes with an error field set, on a
|
||||
// 200. A returned error is this worker failing to serve the verb at all, which
|
||||
// becomes a non-2xx over the tunnel and nothing the frontend may act on.
|
||||
|
||||
// dropMCPSessionsForBackend closes the MCP sessions this worker cached for a
|
||||
// backend that is going away.
|
||||
//
|
||||
// It is the agent worker's whole implementation of backend.stop, and it is
|
||||
// deliberately nothing like the backend worker's, which kills the process and
|
||||
// recycles its port. An agent worker runs no backend processes; what it holds
|
||||
// are sessions that were created against one.
|
||||
//
|
||||
// A backend nobody named is a no-op rather than an error. The event carries the
|
||||
// name, and a request without one asks this worker to forget nothing in
|
||||
// particular; failing it would put a malformed publish into the bucket the
|
||||
// frontend reads as a worker that could not be reached.
|
||||
func dropMCPSessionsForBackend(_ context.Context, req messaging.BackendStopRequest) error {
|
||||
if req.Backend == "" {
|
||||
return nil
|
||||
}
|
||||
mcpTools.CloseMCPSessions(req.Backend)
|
||||
return nil
|
||||
}
|
||||
|
||||
// serveMCPToolRequest answers an MCP tool execution request.
|
||||
func serveMCPToolRequest(ctx context.Context, raw json.RawMessage) (json.RawMessage, error) {
|
||||
return encodeMCPReply(runMCPTool(ctx, raw))
|
||||
}
|
||||
|
||||
// runMCPTool creates or reuses the named MCP sessions from the request's config
|
||||
// and executes the named tool against them.
|
||||
//
|
||||
// Every failure inside it is an answer rather than an error, because every one
|
||||
// of them is something this worker LEARNED by trying: a config it could not
|
||||
// build sessions from, a discovery that failed, a tool that returned an error.
|
||||
func runMCPTool(ctx context.Context, raw json.RawMessage) mcpRemote.MCPToolResponse {
|
||||
var req mcpRemote.MCPToolRequest
|
||||
if err := json.Unmarshal(data, &req); err != nil {
|
||||
sendMCPToolReply(reply, "", fmt.Sprintf("unmarshal error: %v", err))
|
||||
return
|
||||
if err := json.Unmarshal(raw, &req); err != nil {
|
||||
return mcpRemote.MCPToolResponse{Error: fmt.Sprintf("unmarshal error: %v", err)}
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), config.DefaultMCPToolTimeout)
|
||||
// Bounded here rather than by the caller, so the bus path and the tunnel
|
||||
// path give a stuck MCP server the same budget.
|
||||
ctx, cancel := context.WithTimeout(ctx, config.DefaultMCPToolTimeout)
|
||||
defer cancel()
|
||||
|
||||
// Create/cache named MCP sessions from the provided config
|
||||
namedSessions, err := mcpTools.NamedSessionsFromMCPConfig(req.ModelName, req.RemoteServers, req.StdioServers, nil)
|
||||
if err != nil {
|
||||
sendMCPToolReply(reply, "", fmt.Sprintf("session error: %v", err))
|
||||
return
|
||||
return mcpRemote.MCPToolResponse{Error: fmt.Sprintf("session error: %v", err)}
|
||||
}
|
||||
|
||||
// Discover tools to find the right session
|
||||
tools, err := mcpTools.DiscoverMCPTools(ctx, namedSessions)
|
||||
if err != nil {
|
||||
sendMCPToolReply(reply, "", fmt.Sprintf("discovery error: %v", err))
|
||||
return
|
||||
return mcpRemote.MCPToolResponse{Error: fmt.Sprintf("discovery error: %v", err)}
|
||||
}
|
||||
|
||||
// Execute the tool
|
||||
argsJSON, _ := json.Marshal(req.Arguments)
|
||||
result, err := mcpTools.ExecuteMCPToolCall(ctx, tools, req.ToolName, string(argsJSON))
|
||||
if err != nil {
|
||||
sendMCPToolReply(reply, "", err.Error())
|
||||
return
|
||||
return mcpRemote.MCPToolResponse{Error: err.Error()}
|
||||
}
|
||||
|
||||
sendMCPToolReply(reply, result, "")
|
||||
return mcpRemote.MCPToolResponse{Result: result}
|
||||
}
|
||||
|
||||
func sendMCPToolReply(reply func([]byte), result, errMsg string) {
|
||||
resp := mcpRemote.MCPToolResponse{Result: result, Error: errMsg}
|
||||
data, _ := json.Marshal(resp)
|
||||
reply(data)
|
||||
// serveMCPDiscoveryRequest answers an MCP tool/prompt/resource discovery
|
||||
// request.
|
||||
func serveMCPDiscoveryRequest(ctx context.Context, raw json.RawMessage) (json.RawMessage, error) {
|
||||
return encodeMCPReply(runMCPDiscovery(ctx, raw))
|
||||
}
|
||||
|
||||
// handleMCPDiscoveryRequest handles a NATS request-reply for MCP tool/prompt/resource discovery.
|
||||
func handleMCPDiscoveryRequest(data []byte, reply func([]byte)) {
|
||||
// runMCPDiscovery lists the servers this worker can reach for a model, with
|
||||
// their tools, prompts and resources.
|
||||
func runMCPDiscovery(ctx context.Context, raw json.RawMessage) mcpRemote.MCPDiscoveryResponse {
|
||||
var req mcpRemote.MCPDiscoveryRequest
|
||||
if err := json.Unmarshal(data, &req); err != nil {
|
||||
sendMCPDiscoveryReply(reply, nil, nil, fmt.Sprintf("unmarshal error: %v", err))
|
||||
return
|
||||
if err := json.Unmarshal(raw, &req); err != nil {
|
||||
return mcpRemote.MCPDiscoveryResponse{Error: fmt.Sprintf("unmarshal error: %v", err)}
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), config.DefaultMCPDiscoveryTimeout)
|
||||
ctx, cancel := context.WithTimeout(ctx, config.DefaultMCPDiscoveryTimeout)
|
||||
defer cancel()
|
||||
|
||||
// Create/cache named MCP sessions
|
||||
namedSessions, err := mcpTools.NamedSessionsFromMCPConfig(req.ModelName, req.RemoteServers, req.StdioServers, nil)
|
||||
if err != nil {
|
||||
sendMCPDiscoveryReply(reply, nil, nil, fmt.Sprintf("session error: %v", err))
|
||||
return
|
||||
return mcpRemote.MCPDiscoveryResponse{Error: fmt.Sprintf("session error: %v", err)}
|
||||
}
|
||||
|
||||
// List servers with their tools/prompts/resources
|
||||
serverInfos, err := mcpTools.ListMCPServers(ctx, namedSessions)
|
||||
if err != nil {
|
||||
sendMCPDiscoveryReply(reply, nil, nil, fmt.Sprintf("list error: %v", err))
|
||||
return
|
||||
return mcpRemote.MCPDiscoveryResponse{Error: fmt.Sprintf("list error: %v", err)}
|
||||
}
|
||||
|
||||
// Also get tool function schemas for the frontend
|
||||
@@ -342,68 +360,80 @@ func handleMCPDiscoveryRequest(data []byte, reply func([]byte)) {
|
||||
})
|
||||
}
|
||||
|
||||
// Convert server infos
|
||||
var servers []mcpRemote.MCPServerInfo
|
||||
for _, s := range serverInfos {
|
||||
for _, srv := range serverInfos {
|
||||
servers = append(servers, mcpRemote.MCPServerInfo{
|
||||
Name: s.Name,
|
||||
Type: s.Type,
|
||||
Tools: s.Tools,
|
||||
Prompts: s.Prompts,
|
||||
Resources: s.Resources,
|
||||
Error: s.Error,
|
||||
Name: srv.Name,
|
||||
Type: srv.Type,
|
||||
Tools: srv.Tools,
|
||||
Prompts: srv.Prompts,
|
||||
Resources: srv.Resources,
|
||||
Error: srv.Error,
|
||||
})
|
||||
}
|
||||
|
||||
sendMCPDiscoveryReply(reply, servers, toolDefs, "")
|
||||
return mcpRemote.MCPDiscoveryResponse{Servers: servers, Tools: toolDefs}
|
||||
}
|
||||
|
||||
func sendMCPDiscoveryReply(reply func([]byte), servers []mcpRemote.MCPServerInfo, tools []mcpRemote.MCPToolDef, errMsg string) {
|
||||
resp := mcpRemote.MCPDiscoveryResponse{Servers: servers, Tools: tools, Error: errMsg}
|
||||
data, _ := json.Marshal(resp)
|
||||
reply(data)
|
||||
// encodeMCPReply turns a verb's answer into the bytes both carriers send.
|
||||
//
|
||||
// A marshalling failure is the one thing here that is NOT an answer: this
|
||||
// worker has said nothing about the request, so it is returned as an error and
|
||||
// becomes a non-2xx over the tunnel rather than an empty 200.
|
||||
func encodeMCPReply(resp any) (json.RawMessage, error) {
|
||||
out, err := json.Marshal(resp)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encoding the reply: %w", err)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// handleMCPCIJob processes an MCP CI job on the agent worker.
|
||||
// The agent worker can create MCP sessions (has docker) and call the LocalAI API for inference.
|
||||
func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken string, natsClient messaging.MessagingClient, jobTimeout time.Duration) {
|
||||
//
|
||||
// Everything it publishes now goes onto pub, which is the response body of the
|
||||
// control verb rather than the bus. The subjects are unchanged, and they are
|
||||
// still what the claiming replica checks against this worker's allow list
|
||||
// before re-broadcasting, so an SSE stream open on any replica still sees the
|
||||
// same events on the same subjects.
|
||||
//
|
||||
// It returns the terminal answer rather than only publishing it. That is the
|
||||
// structural half of the fix: the claiming replica persists this before it
|
||||
// releases the claim, so a job that finished on a worker cannot be left
|
||||
// `running` because a result message went to a subject nobody was reading.
|
||||
func handleMCPCIJob(ctx context.Context, data []byte, apiURL, apiToken string, pub messaging.Publisher, jobTimeout time.Duration) jobs.ClaimReply {
|
||||
var evt jobs.JobEvent
|
||||
if err := json.Unmarshal(data, &evt); err != nil {
|
||||
xlog.Error("Failed to unmarshal job event", "error", err)
|
||||
return
|
||||
return jobs.ClaimReply{Status: "failed", Error: "unreadable job event"}
|
||||
}
|
||||
|
||||
job := evt.Job
|
||||
task := evt.Task
|
||||
if job == nil || task == nil {
|
||||
xlog.Error("MCP CI job missing enriched data", "jobID", evt.JobID)
|
||||
publishJobResult(natsClient, evt.JobID, "failed", "", "job or task data missing from NATS event")
|
||||
return
|
||||
return mcpCIAnswer(pub, evt.JobID, "failed", "", "job or task data missing from the job event")
|
||||
}
|
||||
|
||||
modelCfg := evt.ModelConfig
|
||||
if modelCfg == nil {
|
||||
publishJobResult(natsClient, evt.JobID, "failed", "", "model config missing from job event")
|
||||
return
|
||||
return mcpCIAnswer(pub, evt.JobID, "failed", "", "model config missing from job event")
|
||||
}
|
||||
|
||||
xlog.Info("Processing MCP CI job", "jobID", evt.JobID, "taskID", evt.TaskID, "model", task.Model)
|
||||
|
||||
// Publish running status
|
||||
natsClient.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
dropTrace(pub.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
JobID: evt.JobID, Status: "running", Message: "Job started on agent worker",
|
||||
})
|
||||
}), evt.JobID)
|
||||
|
||||
// Parse MCP config
|
||||
if modelCfg.MCP.Servers == "" && modelCfg.MCP.Stdio == "" {
|
||||
publishJobResult(natsClient, evt.JobID, "failed", "", "no MCP servers configured for model")
|
||||
return
|
||||
return mcpCIAnswer(pub, evt.JobID, "failed", "", "no MCP servers configured for model")
|
||||
}
|
||||
|
||||
remote, stdio, err := modelCfg.MCP.MCPConfigFromYAML()
|
||||
if err != nil {
|
||||
publishJobResult(natsClient, evt.JobID, "failed", "", fmt.Sprintf("failed to parse MCP config: %v", err))
|
||||
return
|
||||
return mcpCIAnswer(pub, evt.JobID, "failed", "", fmt.Sprintf("failed to parse MCP config: %v", err))
|
||||
}
|
||||
|
||||
// Create MCP sessions locally (agent worker has docker)
|
||||
@@ -413,8 +443,7 @@ func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken s
|
||||
if err != nil {
|
||||
errMsg = fmt.Sprintf("failed to create MCP sessions: %v", err)
|
||||
}
|
||||
publishJobResult(natsClient, evt.JobID, "failed", "", errMsg)
|
||||
return
|
||||
return mcpCIAnswer(pub, evt.JobID, "failed", "", errMsg)
|
||||
}
|
||||
|
||||
// Build prompt from template
|
||||
@@ -442,11 +471,11 @@ func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken s
|
||||
llm := clients.NewLocalAILLM(task.Model, apiToken, apiURL)
|
||||
|
||||
// Build cogito options
|
||||
ctx, cancel := context.WithTimeout(shutdownCtx, jobTimeout)
|
||||
ctx, cancel := context.WithTimeout(ctx, jobTimeout)
|
||||
defer cancel()
|
||||
|
||||
// Update job status to running in DB
|
||||
publishJobStatus(natsClient, evt.JobID, "running", "")
|
||||
publishJobStatus(pub, evt.JobID, "running", "")
|
||||
|
||||
// Buffer stream tokens and flush as complete blocks
|
||||
var reasoningBuf, contentBuf strings.Builder
|
||||
@@ -454,15 +483,15 @@ func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken s
|
||||
|
||||
flushStreamBuf := func() {
|
||||
if reasoningBuf.Len() > 0 {
|
||||
natsClient.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
dropTrace(pub.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
JobID: evt.JobID, TraceType: "reasoning", TraceContent: reasoningBuf.String(),
|
||||
})
|
||||
}), evt.JobID)
|
||||
reasoningBuf.Reset()
|
||||
}
|
||||
if contentBuf.Len() > 0 {
|
||||
natsClient.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
dropTrace(pub.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
JobID: evt.JobID, TraceType: "content", TraceContent: contentBuf.String(),
|
||||
})
|
||||
}), evt.JobID)
|
||||
contentBuf.Reset()
|
||||
}
|
||||
}
|
||||
@@ -473,15 +502,15 @@ func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken s
|
||||
cogito.WithMCPs(sessions...),
|
||||
cogito.WithStatusCallback(func(status string) {
|
||||
flushStreamBuf()
|
||||
natsClient.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
dropTrace(pub.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
JobID: evt.JobID, TraceType: "status", TraceContent: status,
|
||||
})
|
||||
}), evt.JobID)
|
||||
}),
|
||||
cogito.WithToolCallResultCallback(func(t cogito.ToolStatus) {
|
||||
flushStreamBuf()
|
||||
natsClient.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
dropTrace(pub.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
JobID: evt.JobID, TraceType: "tool_result", TraceContent: fmt.Sprintf("%s: %s", t.Name, t.Result),
|
||||
})
|
||||
}), evt.JobID)
|
||||
}),
|
||||
cogito.WithStreamCallback(func(ev cogito.StreamEvent) {
|
||||
// Flush if stream type changed (e.g., reasoning → content)
|
||||
@@ -495,9 +524,9 @@ func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken s
|
||||
case cogito.StreamEventContent:
|
||||
contentBuf.WriteString(ev.Content)
|
||||
case cogito.StreamEventToolCall:
|
||||
natsClient.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
dropTrace(pub.Publish(messaging.SubjectJobProgress(evt.JobID), jobs.ProgressEvent{
|
||||
JobID: evt.JobID, TraceType: "tool_call", TraceContent: fmt.Sprintf("%s(%s)", ev.ToolName, ev.ToolArgs),
|
||||
})
|
||||
}), evt.JobID)
|
||||
}
|
||||
}),
|
||||
)
|
||||
@@ -510,22 +539,84 @@ func handleMCPCIJob(shutdownCtx context.Context, data []byte, apiURL, apiToken s
|
||||
flushStreamBuf() // flush any remaining buffered tokens
|
||||
|
||||
if err != nil {
|
||||
publishJobResult(natsClient, evt.JobID, "failed", "", fmt.Sprintf("cogito execution failed: %v", err))
|
||||
return
|
||||
return mcpCIAnswer(pub, evt.JobID, "failed", "", fmt.Sprintf("cogito execution failed: %v", err))
|
||||
}
|
||||
|
||||
result := ""
|
||||
if msg := f.LastMessage(); msg != nil {
|
||||
result = msg.Content
|
||||
}
|
||||
publishJobResult(natsClient, evt.JobID, "completed", result, "")
|
||||
xlog.Info("MCP CI job completed", "jobID", evt.JobID, "resultLen", len(result))
|
||||
return mcpCIAnswer(pub, evt.JobID, "completed", result, "")
|
||||
}
|
||||
|
||||
func publishJobStatus(nc messaging.MessagingClient, jobID, status, message string) {
|
||||
jobs.PublishJobProgress(nc, jobID, status, message)
|
||||
// dropTrace logs a progress line that could not be written, and never returns
|
||||
// it. A progress line is a NOTIFICATION about a run; a failure to write one says
|
||||
// nothing about the run, and returning it would make the claiming replica read
|
||||
// a finished job as a verb this worker could not serve.
|
||||
func dropTrace(err error, jobID string) {
|
||||
if err != nil {
|
||||
xlog.Debug("An MCP CI progress line could not be written", "jobID", jobID, "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
func publishJobResult(nc messaging.MessagingClient, jobID, status, result, errMsg string) {
|
||||
jobs.PublishJobResult(nc, jobID, status, result, errMsg)
|
||||
// mcpCIAnswer is the ONE place an MCP CI run's terminal state is stated.
|
||||
//
|
||||
// It says it TWICE and on purpose, on two carriers with different jobs. The
|
||||
// publish drives the SSE streams, on the same jobs.<id>.result subject it always
|
||||
// used, re-broadcast by the claiming replica after its allow-list check. The
|
||||
// returned reply is the answer to the control verb, which the claiming replica
|
||||
// persists BEFORE it releases the claim; that is what makes a finished job
|
||||
// impossible to leave `running`, where the publish alone could reach nobody.
|
||||
func mcpCIAnswer(pub messaging.Publisher, jobID, status, result, errMsg string) jobs.ClaimReply {
|
||||
jobs.PublishJobResult(pub, jobID, status, result, errMsg)
|
||||
return jobs.ClaimReply{JobID: jobID, Status: status, Result: result, Error: errMsg}
|
||||
}
|
||||
|
||||
func publishJobStatus(pub messaging.Publisher, jobID, status, message string) {
|
||||
jobs.PublishJobProgress(pub, jobID, status, message)
|
||||
}
|
||||
|
||||
// agentWorkerControlHandlers is every verb this worker serves on the tunnel it
|
||||
// holds to the frontend.
|
||||
//
|
||||
// It is a function rather than a literal inside the start-up path so that a
|
||||
// spec can stand the same set up and post to it. Each of these is the ONLY
|
||||
// carrier for its verb: the queue subjects the two MCP verbs arrived on and the
|
||||
// node subject backend.stop arrived on are all gone, so a field silently
|
||||
// dropped here is a 404 at runtime, which the frontend reads as a worker too
|
||||
// old to serve the verb.
|
||||
func agentWorkerControlHandlers(executor *agents.WorkerExecutor, apiURL, apiToken string, mcpCITimeout time.Duration) agentworker.Config {
|
||||
return agentworker.Config{
|
||||
MCPTool: serveMCPToolRequest,
|
||||
MCPDiscovery: serveMCPDiscoveryRequest,
|
||||
// The verb that removed this process's last reason to dial a bus. It
|
||||
// reaches the SAME cancel registry the executor registers a run on,
|
||||
// because it is the same bridge: a cancel arriving on the tunnel has to
|
||||
// find an execution that is publishing onto a control stream.
|
||||
AgentCancel: executor.Cancel,
|
||||
// Drops the MCP sessions cached for a backend that went away, on the
|
||||
// path a backend worker serves by killing the process instead.
|
||||
BackendStop: dropMCPSessionsForBackend,
|
||||
// The two verbs that replace the queue groups. Both STREAM: their
|
||||
// progress, their agent events and their terminal answer all travel on
|
||||
// the response body the claiming replica is already reading, which is
|
||||
// what lets that replica persist the terminal line before it releases
|
||||
// the claim.
|
||||
AgentExecute: executor.Execute,
|
||||
MCPCIRun: serveMCPCIRun(apiURL, apiToken, mcpCITimeout),
|
||||
}
|
||||
}
|
||||
|
||||
// serveMCPCIRun answers workerctl.PathMCPCIRun.
|
||||
//
|
||||
// The handler's ctx is the REQUEST's, not this process's shutdown context, and
|
||||
// that is the point: when the claiming replica goes away the response body dies
|
||||
// with it, the run stops, and the claim is reaped for another replica to take.
|
||||
// Bound to this worker's shutdown context instead, the run would keep going
|
||||
// with nobody reading it.
|
||||
func serveMCPCIRun(apiURL, apiToken string, jobTimeout time.Duration) agentworker.StreamHandler {
|
||||
return func(ctx context.Context, raw json.RawMessage, pub messaging.Publisher) (json.RawMessage, error) {
|
||||
return json.Marshal(handleMCPCIJob(ctx, raw, apiURL, apiToken, pub, jobTimeout))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/alecthomas/kong"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/agents"
|
||||
mcpRemote "github.com/mudler/LocalAI/core/services/mcp"
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/workerctl"
|
||||
)
|
||||
|
||||
// The agent worker answers its MCP verbs on ONE carrier now: the control route
|
||||
// on the tunnel it holds. The queue-group subjects these used to arrive on are
|
||||
// gone, because a queue group was only ever a way of SELECTING a worker, and
|
||||
// the frontend now makes that selection itself (nodes.AgentSelector).
|
||||
//
|
||||
// What these specs pin is the split between an ANSWER and a FAILURE TO SERVE. A
|
||||
// tool that ran and failed is the worker's own verdict and travels inside the
|
||||
// reply, on a 200; a verb this worker could not serve at all is a Go error,
|
||||
// which becomes a non-2xx, and which nothing may read as evidence about
|
||||
// anything.
|
||||
var _ = Describe("The agent worker's MCP verbs", func() {
|
||||
It("answers a tool request it could not decode, rather than failing to serve it", func() {
|
||||
// The decode happened on this worker and its outcome is something the
|
||||
// worker LEARNED, so it belongs in the reply. Returned as an error it
|
||||
// would become "no route to that worker" at the frontend.
|
||||
raw, err := serveMCPToolRequest(context.Background(), json.RawMessage(`{"tool_name":`))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
var resp mcpRemote.MCPToolResponse
|
||||
Expect(json.Unmarshal(raw, &resp)).To(Succeed())
|
||||
Expect(resp.Error).To(ContainSubstring("unmarshal error"))
|
||||
Expect(resp.Result).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("answers a discovery request it could not decode the same way", func() {
|
||||
raw, err := serveMCPDiscoveryRequest(context.Background(), json.RawMessage(`not json`))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
var resp mcpRemote.MCPDiscoveryResponse
|
||||
Expect(json.Unmarshal(raw, &resp)).To(Succeed())
|
||||
Expect(resp.Error).To(ContainSubstring("unmarshal error"))
|
||||
})
|
||||
|
||||
})
|
||||
|
||||
var _ = Describe("The agent worker's backend stop", func() {
|
||||
// One carrier now: the tunnel's control route. The node subject this used
|
||||
// to arrive on is gone, so this implementation is reached one way only.
|
||||
|
||||
It("treats a stop that names no backend as a no-op rather than a failure", func() {
|
||||
// A malformed request must not become a non-2xx the frontend reads as
|
||||
// a worker it could not reach.
|
||||
Expect(dropMCPSessionsForBackend(context.Background(), messaging.BackendStopRequest{})).To(Succeed())
|
||||
})
|
||||
|
||||
It("succeeds for a backend it holds no sessions for", func() {
|
||||
// The ordinary case on a worker that never touched that backend. An
|
||||
// error here would be reported as this worker failing to serve the
|
||||
// verb, on every stop of every backend it does not know about.
|
||||
Expect(dropMCPSessionsForBackend(context.Background(),
|
||||
messaging.BackendStopRequest{Backend: "a-backend-this-worker-never-saw"})).To(Succeed())
|
||||
})
|
||||
})
|
||||
|
||||
// The wiring, pinned through the transport rather than by reading the struct.
|
||||
//
|
||||
// Every verb below has exactly one carrier: the control route. Dropping a field
|
||||
// from agentWorkerControlHandlers mounts nothing for that path, the catch-all
|
||||
// answers 404, and the frontend reads that 404 as a worker too old to serve the
|
||||
// verb rather than as a wiring mistake. Nothing else in this repo would notice.
|
||||
var _ = Describe("The agent worker's control-plane wiring", func() {
|
||||
var base string
|
||||
var bridge *agents.EventBridge
|
||||
|
||||
BeforeEach(func() {
|
||||
mux := http.NewServeMux()
|
||||
// Built exactly as Run builds it, from an executor and the MCP CI
|
||||
// timeout, so a field this function forgets to set is a 404 here.
|
||||
bridge = agents.NewWorkerEventBridge("agent-worker-spec")
|
||||
executor := agents.NewWorkerExecutor(bridge, nil, "http://127.0.0.1:1", "token")
|
||||
agentWorkerControlHandlers(executor, "http://127.0.0.1:1", "token", time.Second).Register(mux)
|
||||
srv := httptest.NewServer(mux)
|
||||
DeferCleanup(srv.Close)
|
||||
base = srv.URL
|
||||
})
|
||||
|
||||
DescribeTable("mounts the verb an agent worker is the only server of",
|
||||
func(path string) {
|
||||
resp, err := http.Post(base+path, "application/json", strings.NewReader(`{}`)) //nolint:gosec,noctx // httptest server, no redirects to follow
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = resp.Body.Close() })
|
||||
Expect(resp.StatusCode).ToNot(Equal(http.StatusNotFound),
|
||||
"%s is not mounted; agentWorkerControlHandlers does not wire its handler", path)
|
||||
},
|
||||
Entry("backend stop", workerctl.PathBackendStop),
|
||||
Entry("mcp tool execute", workerctl.PathMCPToolExecute),
|
||||
Entry("mcp discovery", workerctl.PathMCPDiscovery),
|
||||
// The two verbs that replaced the queue groups. An unwired one answers
|
||||
// a 404, which is EXACTLY what an older worker answers, so nothing else
|
||||
// in the tree can tell the two apart and only this spec can.
|
||||
Entry("agent execute", workerctl.PathAgentExecute),
|
||||
Entry("mcp ci run", workerctl.PathMCPCIRun),
|
||||
// The verb that removed this process's last reason to dial a bus.
|
||||
// Unwired it answers a 404, which the frontend reads as a worker too
|
||||
// old to serve it, and every cancel of an agent this worker is running
|
||||
// is then reported as one that could not be delivered - for ever.
|
||||
Entry("agent cancel", workerctl.PathAgentCancel),
|
||||
)
|
||||
|
||||
// The cancel verb end to end through the mux, because the thing that must
|
||||
// be true is that the path reaches the SAME cancel registry the executor
|
||||
// registers a run on. Two bridges would compile, mount, answer 200 and
|
||||
// cancel nothing.
|
||||
It("cancels a run registered on the executor's own bridge, and says so", func() {
|
||||
cancelled := make(chan struct{})
|
||||
bridge.RegisterCancel("msg-1", func() { close(cancelled) })
|
||||
|
||||
post := func(messageID string) messaging.AgentCancelReply {
|
||||
GinkgoHelper()
|
||||
body, err := json.Marshal(messaging.AgentCancelRequest{AgentName: "a1", UserID: "u1", MessageID: messageID})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
resp, err := http.Post(base+workerctl.PathAgentCancel, "application/json", bytes.NewReader(body)) //nolint:gosec,noctx // httptest server, no redirects to follow
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = resp.Body.Close() })
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
var reply messaging.AgentCancelReply
|
||||
Expect(json.NewDecoder(resp.Body).Decode(&reply)).To(Succeed())
|
||||
return reply
|
||||
}
|
||||
|
||||
Expect(post("msg-nobody-is-running").Cancelled).To(BeFalse(),
|
||||
"a worker answers only for itself, and false is that answer rather than a failure")
|
||||
Expect(post("msg-1").Cancelled).To(BeTrue())
|
||||
Eventually(cancelled, "20s").Should(BeClosed())
|
||||
})
|
||||
|
||||
It("fails to serve a cancel it cannot read, rather than answering that it found nothing", func() {
|
||||
// The two are different facts. A 200 with cancelled false would be read
|
||||
// as this worker's own answer about the run; a body it could not decode
|
||||
// is not an answer about anything.
|
||||
resp, err := http.Post(base+workerctl.PathAgentCancel, "application/json", strings.NewReader(`{"message_id":`)) //nolint:gosec,noctx // httptest server, no redirects to follow
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = resp.Body.Close() })
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusInternalServerError))
|
||||
})
|
||||
|
||||
DescribeTable("answers a dispatched verb as a stream, so progress and the terminal line share one body",
|
||||
func(path string) {
|
||||
resp, err := http.Post(base+path, "application/json", strings.NewReader(`{}`)) //nolint:gosec,noctx // httptest server, no redirects to follow
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = resp.Body.Close() })
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusOK))
|
||||
Expect(resp.Header.Get("Content-Type")).To(Equal(workerctl.ContentTypeStream))
|
||||
|
||||
// Exactly one reply line, and it is the last thing on the body.
|
||||
// That is what the claiming replica stops reading on.
|
||||
var envs []workerctl.Envelope
|
||||
sc := bufio.NewScanner(resp.Body)
|
||||
for sc.Scan() {
|
||||
if strings.TrimSpace(sc.Text()) == "" {
|
||||
continue
|
||||
}
|
||||
var env workerctl.Envelope
|
||||
Expect(json.Unmarshal(sc.Bytes(), &env)).To(Succeed())
|
||||
envs = append(envs, env)
|
||||
}
|
||||
Expect(sc.Err()).ToNot(HaveOccurred())
|
||||
Expect(envs).ToNot(BeEmpty())
|
||||
Expect(envs[len(envs)-1].Reply).ToNot(BeEmpty(), "the reply line must be last")
|
||||
for _, env := range envs[:len(envs)-1] {
|
||||
Expect(env.Reply).To(BeEmpty(), "only the last line may be a reply")
|
||||
}
|
||||
},
|
||||
Entry("agent execute", workerctl.PathAgentExecute),
|
||||
Entry("mcp ci run", workerctl.PathMCPCIRun),
|
||||
)
|
||||
})
|
||||
|
||||
// The agent worker dials NO message bus, pinned so that a flag reappearing as
|
||||
// required is a decision rather than an accident.
|
||||
//
|
||||
// Every verb a frontend addresses to this worker arrives on the tunnel it
|
||||
// dials, and that now includes the cancel: agent.<name>.cancel was the last
|
||||
// family in the other direction, from a frontend replica to whichever worker
|
||||
// held the execution, and it could not move to the broadcast carrier because
|
||||
// that carrier rides PostgreSQL and this process has no database. It is a
|
||||
// control verb on the worker's own tunnel instead.
|
||||
//
|
||||
// --nats-url is still ACCEPTED, and ignored, so an existing command line or
|
||||
// unit file starts unchanged.
|
||||
var _ = Describe("The agent worker's bus requirement", func() {
|
||||
parse := func(args ...string) error {
|
||||
// kong resolves env: tags from the process environment, so a
|
||||
// LOCALAI_NATS_URL inherited from a developer's shell would make the
|
||||
// first spec below pass for the wrong reason.
|
||||
if prior, had := os.LookupEnv("LOCALAI_NATS_URL"); had {
|
||||
Expect(os.Unsetenv("LOCALAI_NATS_URL")).To(Succeed())
|
||||
DeferCleanup(func() { _ = os.Setenv("LOCALAI_NATS_URL", prior) })
|
||||
}
|
||||
var cli struct {
|
||||
AgentWorker AgentWorkerCMD `cmd:""`
|
||||
}
|
||||
parser, err := kong.New(&cli)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
_, err = parser.Parse(append([]string{"agent-worker"}, args...))
|
||||
return err
|
||||
}
|
||||
|
||||
It("starts with no bus named at all", func() {
|
||||
Expect(parse("--register-to", "http://frontend:8080")).To(Succeed(),
|
||||
"an agent worker connects to no message bus and must not demand the URL of one")
|
||||
})
|
||||
|
||||
It("still accepts a command line that names one", func() {
|
||||
// Ignored, not rejected. An operator upgrading a fleet must not have to
|
||||
// edit every unit file in the same change.
|
||||
Expect(parse("--register-to", "http://frontend:8080", "--nats-url", "nats://bus:4222")).To(Succeed())
|
||||
})
|
||||
|
||||
It("still accepts the broker credentials that came with it", func() {
|
||||
// The whole set, because an operator's unit file carries the whole set:
|
||||
// a command line that parses --nats-url and then dies on --nats-jwt has
|
||||
// bought the fleet nothing.
|
||||
Expect(parse("--register-to", "http://frontend:8080",
|
||||
"--nats-url", "nats://bus:4222",
|
||||
"--nats-jwt", "eyJ0",
|
||||
"--nats-user-seed", "SUUSER",
|
||||
"--nats-service-jwt", "eyJ0",
|
||||
"--nats-service-seed", "SUSERVICE",
|
||||
"--nats-require-auth")).To(Succeed())
|
||||
})
|
||||
|
||||
It("does not stat the TLS material it no longer presents", func() {
|
||||
// These paths were validated as existing files while they were dialled
|
||||
// with. Keeping that on an ignored flag would fail a worker at startup
|
||||
// over a certificate for a broker the operator has already deleted,
|
||||
// which is the exact upgrade the acceptance exists to survive.
|
||||
missing := filepath.Join(GinkgoT().TempDir(), "a-broker-ca-that-was-deleted.pem")
|
||||
Expect(parse("--register-to", "http://frontend:8080",
|
||||
"--nats-tlsca", missing,
|
||||
"--nats-tls-cert", missing,
|
||||
"--nats-tls-key", missing)).To(Succeed())
|
||||
})
|
||||
|
||||
It("keeps every accepted bus flag hidden from --help", func() {
|
||||
var cli struct {
|
||||
AgentWorker AgentWorkerCMD `cmd:""`
|
||||
}
|
||||
parser, err := kong.New(&cli)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
var visible []string
|
||||
for _, node := range parser.Model.Children {
|
||||
for _, flag := range node.Flags {
|
||||
if strings.HasPrefix(flag.Name, "nats-") && !flag.Hidden {
|
||||
visible = append(visible, flag.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
Expect(visible).To(BeEmpty(),
|
||||
"%v are still offered in --help while doing nothing", visible)
|
||||
})
|
||||
})
|
||||
|
||||
// Which flag makes an agent worker wait through admin approval.
|
||||
//
|
||||
// The docs promise this specifically, and it is the one behavioural change in
|
||||
// the broker removal that an operator can be surprised by, so it is asserted
|
||||
// both ways round. The positive half alone would stay green if the gate were
|
||||
// widened back to OR --nats-require-auth; the negative half is what says the
|
||||
// change actually happened.
|
||||
var _ = Describe("The agent worker's approval gate", func() {
|
||||
It("waits when --distributed-require-auth is set", func() {
|
||||
cmd := &AgentWorkerCMD{DistributedRequireAuth: true}
|
||||
Expect(cmd.waitThroughApproval()).To(BeTrue())
|
||||
})
|
||||
|
||||
It("does not wait for the broker flag that used to imply it", func() {
|
||||
// --nats-require-auth named a bus this worker does not dial. An
|
||||
// operator who set only that one now gets the historical default:
|
||||
// register, start, and be refused at every tunnel dial with 403 until
|
||||
// an admin approves. Documented in the migration section of
|
||||
// docs/content/features/distributed-mode.md.
|
||||
cmd := &AgentWorkerCMD{NatsRequireAuth: true}
|
||||
Expect(cmd.waitThroughApproval()).To(BeFalse(),
|
||||
"an ignored flag is gating a real behaviour again")
|
||||
})
|
||||
|
||||
It("does not wait when neither is set", func() {
|
||||
Expect((&AgentWorkerCMD{}).waitThroughApproval()).To(BeFalse())
|
||||
})
|
||||
})
|
||||
+1
-1
@@ -18,7 +18,7 @@ var CLI struct {
|
||||
Transcript TranscriptCMD `cmd:"" help:"Convert audio to text"`
|
||||
P2PWorker worker.Worker `cmd:"" name:"p2p-worker" help:"Run workers to distribute workload via p2p (llama.cpp-only)"`
|
||||
Worker WorkerCMD `cmd:"" help:"Start a worker for distributed mode (generic, backend-agnostic)"`
|
||||
AgentWorker AgentWorkerCMD `cmd:"" name:"agent-worker" help:"Start an agent worker for distributed mode (executes agent chats via NATS)"`
|
||||
AgentWorker AgentWorkerCMD `cmd:"" name:"agent-worker" help:"Start an agent worker for distributed mode (executes agent chats over the tunnel it dials out to a frontend)"`
|
||||
Util UtilCMD `cmd:"" help:"Utility commands"`
|
||||
Agent AgentCMD `cmd:"" help:"Run agents standalone without the full LocalAI server"`
|
||||
MCPServer MCPServerCMD `cmd:"" name:"mcp-server" help:"Run the LocalAI admin tool surface as a stdio MCP server (controls a remote LocalAI instance over HTTP)"`
|
||||
|
||||
+38
-46
@@ -162,9 +162,10 @@ 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; no message bus)" group:"distributed"`
|
||||
InstanceID string `env:"LOCALAI_INSTANCE_ID" help:"Unique instance ID for distributed mode (auto-generated UUID if empty)" group:"distributed"`
|
||||
NatsURL string `env:"LOCALAI_NATS_URL" help:"NATS server URL (e.g., nats://localhost:4222)" group:"distributed"`
|
||||
NatsURL string `env:"LOCALAI_NATS_URL" help:"Ignored. No component of a distributed deployment connects to a message bus; state and fan-out ride PostgreSQL and workers are reached over their own tunnels. Accepted so an existing command line still starts." group:"distributed" hidden:""`
|
||||
DistributedAdvertiseAddr string `env:"LOCALAI_DISTRIBUTED_ADVERTISE_ADDR" help:"host:port other frontend replicas dial to reach this one (peer link). Empty = derived from the local address that routes to PostgreSQL, which only works when the database is on another host." 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"`
|
||||
@@ -172,29 +173,41 @@ type RunCMD struct {
|
||||
StorageSecretKey string `env:"LOCALAI_STORAGE_SECRET_KEY" help:"S3 secret access key" group:"distributed"`
|
||||
RegistrationToken string `env:"LOCALAI_REGISTRATION_TOKEN" help:"Token that backend nodes must provide to register (empty = no auth required)" group:"distributed"`
|
||||
RegistrationRequireAuth bool `env:"LOCALAI_REGISTRATION_REQUIRE_AUTH" default:"false" help:"Fail startup when distributed mode is enabled but LOCALAI_REGISTRATION_TOKEN is empty (node endpoints and worker file-transfer server would otherwise be unauthenticated)" group:"distributed"`
|
||||
DistributedRequireAuth bool `env:"LOCALAI_DISTRIBUTED_REQUIRE_AUTH" default:"false" help:"Umbrella switch: require BOTH NATS JWT credentials and a registration token when distributed mode is enabled (implies --nats-require-auth and --registration-require-auth)" group:"distributed"`
|
||||
DistributedRequireAuth bool `env:"LOCALAI_DISTRIBUTED_REQUIRE_AUTH" default:"false" help:"Umbrella switch: require a registration token when distributed mode is enabled (implies --registration-require-auth)" group:"distributed"`
|
||||
AutoApproveNodes bool `env:"LOCALAI_AUTO_APPROVE_NODES" default:"false" help:"Auto-approve new worker nodes (skip admin approval)" group:"distributed"`
|
||||
DistributedSharedModels bool `env:"LOCALAI_DISTRIBUTED_SHARED_MODELS" default:"false" help:"Assert that every node mounts the SAME models directory at the SAME path (shared volume). When true, the router skips staging model files to workers and loads them directly from the shared path, avoiding re-downloads." group:"distributed"`
|
||||
DistributedPrefixCache bool `env:"LOCALAI_DISTRIBUTED_PREFIX_CACHE" default:"true" help:"Enable prefix-cache-aware routing in distributed mode (default true). When false, routing falls back to round-robin." group:"distributed"`
|
||||
DistributedDiskHeadroomCheck bool `env:"LOCALAI_DISTRIBUTED_DISK_HEADROOM_CHECK" default:"true" help:"Reject worker nodes that lack free space to store the model, at scheduling time rather than partway through staging (default true). Free space is measured on the filesystem backing each worker's models directory, and compared against the model's own size plus a small margin. When false, node selection ignores free disk (pre-#11054 behaviour); the check still runs and warns when it would have rejected every node. Can also be toggled at runtime via the distributed_disk_headroom_check setting." group:"distributed"`
|
||||
DistributedPrefixCacheTTL string `env:"LOCALAI_DISTRIBUTED_PREFIX_CACHE_TTL" help:"Idle-timeout for prefix-cache index entries; also drives the background eviction cadence (every TTL/2). Default 5m." group:"distributed"`
|
||||
BackendInstallTimeout string `env:"LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT" help:"NATS round-trip timeout for backend.install requests sent to worker nodes (default 15m). Increase for slow links pulling multi-GB images." group:"distributed"`
|
||||
BackendUpgradeTimeout string `env:"LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT" help:"NATS round-trip timeout for backend.upgrade requests (default 15m)." group:"distributed"`
|
||||
BackendInstallTimeout string `env:"LOCALAI_NATS_BACKEND_INSTALL_TIMEOUT" help:"Timeout for a backend.install request the frontend sends a worker over its tunnel (default 15m). Increase for slow links pulling multi-GB images." group:"distributed"`
|
||||
BackendUpgradeTimeout string `env:"LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT" help:"Timeout for a backend.upgrade request the frontend sends a worker over its tunnel (default 15m)." group:"distributed"`
|
||||
ModelLoadTimeout string `env:"LOCALAI_NATS_MODEL_LOAD_TIMEOUT" help:"Fixed gRPC deadline for the remote LoadModel call sent to a worker node once its backend is installed and model files are staged. Unset (the default), the deadline is derived from the checkpoint size instead: 5m plus 20s per GiB, capped at 6h, so multi-tens-of-GB diffusion/video checkpoints get the minutes they need without a fixed cliff. Set this only to pin a specific budget; the value is used verbatim, including when it is shorter than the derived one." group:"distributed"`
|
||||
ModelLoadWait string `env:"LOCALAI_MODEL_LOAD_WAIT" help:"How long an inference request waits for a model that is still cold-loading onto a worker before it is answered with 503, a Retry-After header and live staging progress (default 60s). The request is served the moment the model becomes ready, so a model already most of the way staged needs no client retry. Set to 0 to wait as long as the load takes — only safe when no ingress or load balancer with an idle timeout sits in front." group:"distributed"`
|
||||
StaleNodeThreshold string `env:"LOCALAI_STALE_NODE_THRESHOLD" help:"How long a worker node may go without a durable heartbeat before the health monitor marks it offline (default 5m). Because a beat that only carries a fresher timestamp is held back by --node-heartbeat-checkpoint, this must stay comfortably wider than that interval; raise both together. Dead-node detection through the per-model gRPC health check and through request-time failure is unaffected by this knob." group:"distributed"`
|
||||
NodeHeartbeatCheckpoint string `env:"LOCALAI_NODE_HEARTBEAT_CHECKPOINT" help:"Minimum gap between durable heartbeat writes for a worker node (default 60s). A beat that only carries a fresher timestamp is dropped until this interval elapses; every field is compared against the value last written, so a node's first beat, a changed total VRAM/total disk/GPU vendor, and a free VRAM/RAM/disk reading that has moved more than 256 MiB from the written value all still write immediately, and a node that is not active is never suppressed. Set below the worker heartbeat interval to write on every beat." group:"distributed"`
|
||||
NatsAccountSeed string `env:"LOCALAI_NATS_ACCOUNT_SEED" help:"NATS account signing seed (SU...) used to mint per-node worker JWTs at registration" group:"distributed"`
|
||||
NatsServiceJWT string `env:"LOCALAI_NATS_SERVICE_JWT" help:"NATS user JWT for the frontend (and agent workers) to publish control-plane messages" group:"distributed"`
|
||||
NatsServiceSeed string `env:"LOCALAI_NATS_SERVICE_SEED" help:"NATS user signing seed (SU...) paired with LOCALAI_NATS_SERVICE_JWT" group:"distributed"`
|
||||
NatsWorkerJWTTTL string `env:"LOCALAI_NATS_WORKER_JWT_TTL" help:"Lifetime of minted per-node NATS JWTs (e.g. 24h, default 24h)" group:"distributed"`
|
||||
NatsRequireAuth bool `env:"LOCALAI_NATS_REQUIRE_AUTH" default:"false" help:"Require NATS JWT credentials (service JWT + account seed) when distributed mode is enabled" group:"distributed"`
|
||||
NatsTLSCA string `env:"LOCALAI_NATS_TLS_CA" type:"existingfile" help:"PEM file for NATS server CA (private PKI); use with tls:// in --nats-url" group:"distributed"`
|
||||
NatsTLSCert string `env:"LOCALAI_NATS_TLS_CERT" type:"existingfile" help:"Client certificate for NATS mTLS" group:"distributed"`
|
||||
NatsTLSKey string `env:"LOCALAI_NATS_TLS_KEY" type:"existingfile" help:"Client private key for NATS mTLS" group:"distributed"`
|
||||
ExposeNodeHeader bool `env:"LOCALAI_EXPOSE_NODE_HEADER" default:"false" help:"Set the X-LocalAI-Node response header on inference responses (OpenAI chat/completions/embeddings, Anthropic /v1/messages, Ollama /api/chat,/api/generate,/api/embed) with the ID of the worker that served the request. Disabled by default: the node ID reveals internal topology and should not be exposed on a public endpoint. Best-effort: under heavy concurrency the header may reflect a recent routing decision rather than this exact request's." group:"distributed"`
|
||||
ModelScheduling string `env:"LOCALAI_MODEL_SCHEDULING" help:"Declarative per-model scheduling config applied at startup (inline JSON list of {model_name,node_selector,min_replicas,max_replicas,replicas:\"all\"}). Authoritative: overwrites matching models on every boot. Distributed mode only." group:"distributed"`
|
||||
ModelSchedulingConfig string `env:"LOCALAI_MODEL_SCHEDULING_CONFIG" help:"Path to a YAML file with the same per-model scheduling list as LOCALAI_MODEL_SCHEDULING. Distributed mode only." group:"distributed"`
|
||||
WorkerReconnectGrace string `env:"LOCALAI_WORKER_RECONNECT_GRACE" help:"How long a worker whose tunnel was lost is treated as reconnecting rather than gone (default 90s, clear of two of the worker's own ceiling backoffs plus the dial between them). Only after this window may the scheduler stop placing work on that worker and clean up its rows, so a value below the worker's backoff condemns workers that are re-homing normally; raise it to make a rolling frontend restart safer, lower it to reap a genuinely dead worker sooner. Measured on the database clock, so every replica agrees." group:"distributed"`
|
||||
// The broker credential and TLS flags, accepted and ignored, hidden, on the
|
||||
// same terms as --nats-url above and for the same reason: kong fails on an
|
||||
// unknown flag, so deleting them turns every existing unit file, compose
|
||||
// file and Helm values file into a startup parse error at upgrade, in
|
||||
// exchange for nothing. There is no bus connection to present a credential
|
||||
// on and no minting left to do.
|
||||
//
|
||||
// type:"existingfile" is deliberately NOT kept on the TLS paths. Validating
|
||||
// a path this process never opens would fail a deployment at startup over a
|
||||
// certificate for a broker the operator has already shut down, which is
|
||||
// exactly the upgrade the acceptance exists to survive.
|
||||
NatsAccountSeed string `env:"LOCALAI_NATS_ACCOUNT_SEED" help:"Ignored. The frontend mints no per-node broker credential; nodes are authenticated by their registration and tunnel tokens." group:"distributed" hidden:""`
|
||||
NatsServiceJWT string `env:"LOCALAI_NATS_SERVICE_JWT" help:"Ignored. The frontend opens no bus connection to present a credential on." group:"distributed" hidden:""`
|
||||
NatsServiceSeed string `env:"LOCALAI_NATS_SERVICE_SEED" help:"Ignored. Paired with LOCALAI_NATS_SERVICE_JWT, which is itself ignored." group:"distributed" hidden:""`
|
||||
NatsWorkerJWTTTL string `env:"LOCALAI_NATS_WORKER_JWT_TTL" help:"Ignored. No per-node broker credential is minted, so none has a lifetime." group:"distributed" hidden:""`
|
||||
NatsRequireAuth bool `env:"LOCALAI_NATS_REQUIRE_AUTH" default:"false" help:"Ignored on the frontend. Use --registration-require-auth, or the --distributed-require-auth umbrella, to fail closed on a missing registration token." group:"distributed" hidden:""`
|
||||
NatsTLSCA string `env:"LOCALAI_NATS_TLS_CA" help:"Ignored. No bus connection is opened, so no server certificate is verified." group:"distributed" hidden:""`
|
||||
NatsTLSCert string `env:"LOCALAI_NATS_TLS_CERT" help:"Ignored. No bus connection is opened, so no client certificate is presented." group:"distributed" hidden:""`
|
||||
NatsTLSKey string `env:"LOCALAI_NATS_TLS_KEY" help:"Ignored. Paired with LOCALAI_NATS_TLS_CERT, which is itself ignored." group:"distributed" hidden:""`
|
||||
ExposeNodeHeader bool `env:"LOCALAI_EXPOSE_NODE_HEADER" default:"false" help:"Set the X-LocalAI-Node response header on inference responses (OpenAI chat/completions/embeddings, Anthropic /v1/messages, Ollama /api/chat,/api/generate,/api/embed) with the ID of the worker that served the request. Disabled by default: the node ID reveals internal topology and should not be exposed on a public endpoint. Best-effort: under heavy concurrency the header may reflect a recent routing decision rather than this exact request's." group:"distributed"`
|
||||
ModelScheduling string `env:"LOCALAI_MODEL_SCHEDULING" help:"Declarative per-model scheduling config applied at startup (inline JSON list of {model_name,node_selector,min_replicas,max_replicas,replicas:\"all\"}). Authoritative: overwrites matching models on every boot. Distributed mode only." group:"distributed"`
|
||||
ModelSchedulingConfig string `env:"LOCALAI_MODEL_SCHEDULING_CONFIG" help:"Path to a YAML file with the same per-model scheduling list as LOCALAI_MODEL_SCHEDULING. Distributed mode only." group:"distributed"`
|
||||
|
||||
Version bool
|
||||
|
||||
@@ -353,8 +366,8 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
|
||||
if r.InstanceID != "" {
|
||||
opts = append(opts, config.WithDistributedInstanceID(r.InstanceID))
|
||||
}
|
||||
if r.NatsURL != "" {
|
||||
opts = append(opts, config.WithNatsURL(r.NatsURL))
|
||||
if r.DistributedAdvertiseAddr != "" {
|
||||
opts = append(opts, config.WithDistributedAdvertiseAddr(r.DistributedAdvertiseAddr))
|
||||
}
|
||||
if r.StorageURL != "" {
|
||||
opts = append(opts, config.WithStorageURL(r.StorageURL))
|
||||
@@ -399,6 +412,13 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
|
||||
}
|
||||
opts = append(opts, config.WithModelLoadWait(d))
|
||||
}
|
||||
if r.WorkerReconnectGrace != "" {
|
||||
d, err := parseDistributedDuration("LOCALAI_WORKER_RECONNECT_GRACE", r.WorkerReconnectGrace)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
opts = append(opts, config.WithWorkerReconnectGrace(d))
|
||||
}
|
||||
if r.StaleNodeThreshold != "" {
|
||||
d, err := parseDistributedDuration("LOCALAI_STALE_NODE_THRESHOLD", r.StaleNodeThreshold)
|
||||
if err != nil {
|
||||
@@ -425,34 +445,6 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
|
||||
if r.DistributedSharedModels {
|
||||
opts = append(opts, config.EnableDistributedSharedModels)
|
||||
}
|
||||
if r.NatsAccountSeed != "" {
|
||||
opts = append(opts, config.WithNatsAccountSeed(r.NatsAccountSeed))
|
||||
}
|
||||
if r.NatsServiceJWT != "" {
|
||||
opts = append(opts, config.WithNatsServiceJWT(r.NatsServiceJWT))
|
||||
}
|
||||
if r.NatsServiceSeed != "" {
|
||||
opts = append(opts, config.WithNatsServiceSeed(r.NatsServiceSeed))
|
||||
}
|
||||
if r.NatsWorkerJWTTTL != "" {
|
||||
d, err := time.ParseDuration(r.NatsWorkerJWTTTL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("invalid LOCALAI_NATS_WORKER_JWT_TTL %q: %w", r.NatsWorkerJWTTTL, err)
|
||||
}
|
||||
opts = append(opts, config.WithNatsWorkerJWTTTL(d))
|
||||
}
|
||||
if r.NatsRequireAuth {
|
||||
opts = append(opts, config.EnableNatsRequireAuth)
|
||||
}
|
||||
if r.NatsTLSCA != "" {
|
||||
opts = append(opts, config.WithNatsTLSCA(r.NatsTLSCA))
|
||||
}
|
||||
if r.NatsTLSCert != "" {
|
||||
opts = append(opts, config.WithNatsTLSCert(r.NatsTLSCert))
|
||||
}
|
||||
if r.NatsTLSKey != "" {
|
||||
opts = append(opts, config.WithNatsTLSKey(r.NatsTLSKey))
|
||||
}
|
||||
if r.AutoApproveNodes {
|
||||
opts = append(opts, config.EnableAutoApproveNodes)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
package cli
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"github.com/alecthomas/kong"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// The frontend's broker flags are ACCEPTED and IGNORED, and the two halves are
|
||||
// separate Its on purpose.
|
||||
//
|
||||
// Accepted, because kong fails on an unknown flag: deleting --nats-url would
|
||||
// turn every existing unit file, compose file and Helm chart into a startup
|
||||
// parse error on the day of the upgrade, in exchange for nothing, since the
|
||||
// flag has nothing left to do. Ignored, because a flag that parses and is then
|
||||
// quietly honoured is the failure this spec exists to exclude, and a single
|
||||
// "it starts" assertion cannot tell the two apart.
|
||||
// runFlagVars supplies the kong variables cmd/local-ai/main.go supplies, so a
|
||||
// RunCMD can be parsed here at all: its path defaults interpolate ${basepath}.
|
||||
func runFlagVars() kong.Vars {
|
||||
return kong.Vars{
|
||||
"basepath": GinkgoT().TempDir(),
|
||||
"generatedcontentpath": DefaultGeneratedContentPath(),
|
||||
"uploadpath": DefaultUploadPath(),
|
||||
"galleries": config.DefaultGalleriesJSON,
|
||||
"backends": config.DefaultBackendGalleriesJSON,
|
||||
"version": "test",
|
||||
}
|
||||
}
|
||||
|
||||
var _ = Describe("The frontend's broker flags", func() {
|
||||
busFlags := []string{
|
||||
"--nats-url", "nats://bus:4222",
|
||||
"--nats-account-seed", "SUACCOUNT",
|
||||
"--nats-service-jwt", "eyJ0",
|
||||
"--nats-service-seed", "SUSERVICE",
|
||||
"--nats-worker-jwtttl", "24h",
|
||||
"--nats-require-auth",
|
||||
}
|
||||
|
||||
parse := func(args ...string) (*RunCMD, error) {
|
||||
// kong resolves env: tags from the process environment, so a
|
||||
// LOCALAI_NATS_URL inherited from a developer's shell would let the
|
||||
// first spec pass for the wrong reason.
|
||||
for _, name := range []string{"LOCALAI_NATS_URL", "LOCALAI_NATS_ACCOUNT_SEED", "LOCALAI_NATS_REQUIRE_AUTH"} {
|
||||
if prior, had := os.LookupEnv(name); had {
|
||||
Expect(os.Unsetenv(name)).To(Succeed())
|
||||
DeferCleanup(func() { _ = os.Setenv(name, prior) })
|
||||
}
|
||||
}
|
||||
var cli struct {
|
||||
Run RunCMD `cmd:""`
|
||||
}
|
||||
parser, err := kong.New(&cli, runFlagVars())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
_, err = parser.Parse(append([]string{"run"}, args...))
|
||||
return &cli.Run, err
|
||||
}
|
||||
|
||||
It("starts with no bus named at all", func() {
|
||||
_, err := parse("--distributed")
|
||||
Expect(err).To(Succeed(),
|
||||
"a distributed frontend dials no message bus and must not demand the URL of one")
|
||||
})
|
||||
|
||||
It("still accepts a command line that names one", func() {
|
||||
_, err := parse(append([]string{"--distributed"}, busFlags...)...)
|
||||
Expect(err).To(Succeed())
|
||||
})
|
||||
|
||||
It("does not stat the TLS material it no longer presents", func() {
|
||||
// The paths were validated as existing files while they were dialled
|
||||
// with. Keeping that validation on an ignored flag would fail a
|
||||
// deployment at startup over a certificate for a broker that is gone,
|
||||
// which is the exact upgrade the acceptance is meant to survive.
|
||||
missing := filepath.Join(GinkgoT().TempDir(), "a-broker-ca-that-was-deleted.pem")
|
||||
_, err := parse("--distributed",
|
||||
"--nats-tlsca", missing,
|
||||
"--nats-tls-cert", missing,
|
||||
"--nats-tls-key", missing)
|
||||
Expect(err).To(Succeed())
|
||||
})
|
||||
|
||||
It("keeps every accepted bus flag hidden from --help", func() {
|
||||
// Accepted for the upgrade, not offered to a new operator: a flag that
|
||||
// does nothing must not appear in the list of things to configure.
|
||||
var cli struct {
|
||||
Run RunCMD `cmd:""`
|
||||
}
|
||||
parser, err := kong.New(&cli, runFlagVars())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
var visible []string
|
||||
for _, node := range parser.Model.Children {
|
||||
for _, flag := range node.Flags {
|
||||
if len(flag.Name) >= 5 && flag.Name[:5] == "nats-" && !flag.Hidden {
|
||||
visible = append(visible, flag.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
Expect(visible).To(BeEmpty(),
|
||||
"%v are still offered in --help while doing nothing", visible)
|
||||
})
|
||||
})
|
||||
|
||||
// The serve-backend worker's half of the same promise, which nothing pinned.
|
||||
//
|
||||
// `local-ai worker` kept ONE broker flag and dropped the rest a release earlier,
|
||||
// when it stopped connecting to a broker at all. That asymmetry is documented
|
||||
// in docs/content/reference/cli-reference.md, and a documented promise with no
|
||||
// spec is how the wrong half gets deleted: --nats-url is the one an operator's
|
||||
// worker unit file actually carries, and it is the one whose removal would turn
|
||||
// an upgrade into a parse error on every worker in the fleet at once.
|
||||
//
|
||||
// The negative It is here for the same reason the frontend's two halves are
|
||||
// separate. Without it, "accepted" could be satisfied by quietly re-adding the
|
||||
// credential flags, and the docs would be describing a surface nobody checked.
|
||||
var _ = Describe("The serve-backend worker's broker flags", func() {
|
||||
parse := func(args ...string) error {
|
||||
for _, name := range []string{"LOCALAI_NATS_URL", "LOCALAI_ADDRESS"} {
|
||||
if prior, had := os.LookupEnv(name); had {
|
||||
Expect(os.Unsetenv(name)).To(Succeed())
|
||||
DeferCleanup(func() { _ = os.Setenv(name, prior) })
|
||||
}
|
||||
}
|
||||
var cli struct {
|
||||
Worker WorkerCMD `cmd:""`
|
||||
}
|
||||
parser, err := kong.New(&cli, runFlagVars())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
_, err = parser.Parse(append([]string{"worker"}, args...))
|
||||
return err
|
||||
}
|
||||
|
||||
It("still accepts the bus URL an existing worker unit file carries", func() {
|
||||
Expect(parse("--register-to", "http://frontend:8080", "--nats-url", "nats://bus:4222")).To(Succeed(),
|
||||
"a serve-backend worker dials no message bus, and an operator whose unit file still names one must still get a worker that starts")
|
||||
})
|
||||
|
||||
It("keeps it hidden from --help", func() {
|
||||
var cli struct {
|
||||
Worker WorkerCMD `cmd:""`
|
||||
}
|
||||
parser, err := kong.New(&cli, runFlagVars())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
var visible []string
|
||||
for _, node := range parser.Model.Children {
|
||||
for _, flag := range node.Flags {
|
||||
if strings.HasPrefix(flag.Name, "nats-") && !flag.Hidden {
|
||||
visible = append(visible, flag.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
Expect(visible).To(BeEmpty(),
|
||||
"%v are still offered in --help while doing nothing", visible)
|
||||
})
|
||||
|
||||
It("took no broker CREDENTIAL flag back", func() {
|
||||
// The credential and TLS flags left this command in phase 3 and must
|
||||
// stay gone: re-adding one would put a broker credential back on the
|
||||
// surface of a process that opens no broker connection, and the docs
|
||||
// say per-command which flags survive.
|
||||
Expect(parse("--register-to", "http://frontend:8080", "--nats-jwt", "eyJ0")).To(HaveOccurred())
|
||||
Expect(parse("--register-to", "http://frontend:8080", "--nats-service-jwt", "eyJ0")).To(HaveOccurred())
|
||||
Expect(parse("--register-to", "http://frontend:8080", "--nats-tlsca", "/dev/null")).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
@@ -96,7 +96,7 @@ func (r *VLLMDistributed) Run(ctx *cliContext.Context) error {
|
||||
FrontendURL: r.RegisterTo,
|
||||
RegistrationToken: r.RegistrationToken,
|
||||
}
|
||||
nodeID, _, _, _, regErr := regClient.RegisterWithRetry(context.Background(), r.registrationBody(), 10)
|
||||
nodeID, _, regErr := regClient.RegisterWithRetry(context.Background(), r.registrationBody(), 10)
|
||||
if regErr != nil {
|
||||
return fmt.Errorf("registering with frontend: %w", regErr)
|
||||
}
|
||||
|
||||
@@ -8,7 +8,9 @@ import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync"
|
||||
@@ -58,17 +60,27 @@ func (c *RegistrationClient) setAuth(req *http.Request) {
|
||||
|
||||
// RegisterResponse is the JSON body returned by /api/node/register.
|
||||
type RegisterResponse struct {
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status,omitempty"` // "pending" until an admin approves the node
|
||||
APIToken string `json:"api_token,omitempty"`
|
||||
NatsJWT string `json:"nats_jwt,omitempty"`
|
||||
NatsUserSeed string `json:"nats_user_seed,omitempty"`
|
||||
ID string `json:"id"`
|
||||
Status string `json:"status,omitempty"` // "pending" until an admin approves the node
|
||||
APIToken string `json:"api_token,omitempty"`
|
||||
// TunnelToken is this node's own credential for GET /api/cluster/connect.
|
||||
// The frontend mints a fresh one on every registration and keeps only its
|
||||
// hash, so this is the ONLY time the plaintext exists anywhere but in this
|
||||
// worker's memory: a worker that discards it cannot get it back without
|
||||
// registering again.
|
||||
TunnelToken string `json:"tunnel_token,omitempty"`
|
||||
// There are no nats_jwt / nats_user_seed fields. A frontend that predates
|
||||
// this release still sends them and this decodes fine: encoding/json
|
||||
// ignores a key with no field, so an old frontend talking to a new worker
|
||||
// is a no-op rather than a decode failure.
|
||||
}
|
||||
|
||||
// RegisterFull sends a single registration request and returns the full
|
||||
// response (node ID, approval status, and optional API token / NATS creds).
|
||||
// response (node ID, approval status, and optional API and tunnel tokens).
|
||||
// Re-registration is idempotent: the frontend preserves the node row and mints
|
||||
// a fresh NATS JWT each call, so this doubles as the credential-refresh call.
|
||||
// a fresh TUNNEL token each call, so this doubles as the rotation call. It is
|
||||
// the only credential a registration mints; the per-node broker JWT it used to
|
||||
// carry went with the bus.
|
||||
func (c *RegistrationClient) RegisterFull(ctx context.Context, body map[string]any) (*RegisterResponse, error) {
|
||||
jsonBody, _ := json.Marshal(body)
|
||||
url := c.baseURL() + "/api/node/register"
|
||||
@@ -87,7 +99,7 @@ func (c *RegistrationClient) RegisterFull(ctx context.Context, body map[string]a
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("registration failed with status %d", resp.StatusCode)
|
||||
return nil, registrationStatusError(resp)
|
||||
}
|
||||
|
||||
var result RegisterResponse
|
||||
@@ -97,38 +109,105 @@ 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)
|
||||
// ErrRegistrationRejected marks a registration the frontend REFUSED, as opposed
|
||||
// to one it could not answer.
|
||||
//
|
||||
// Retrying a refusal cannot change it: the request is wrong, or this worker is
|
||||
// not allowed to make it. The one that matters in practice is a worker of this
|
||||
// release registering against a frontend that predates it, which answers
|
||||
// "address is required for backend workers" with 400, because a worker no
|
||||
// longer has an address to send. Without this the retry ladder spends four
|
||||
// minutes on a verdict the frontend reached instantly, and the operator watches
|
||||
// it before being told anything.
|
||||
//
|
||||
// 408 and 429 are deliberately NOT rejections. Both are the frontend asking for
|
||||
// the same request again later, which is exactly what a retry does.
|
||||
var ErrRegistrationRejected = errors.New("the frontend refused this registration")
|
||||
|
||||
// maxRegistrationErrorBody bounds how much of a refusal's body is quoted back.
|
||||
// Enough for a message, not enough for an HTML error page to bury the log line
|
||||
// it is meant to explain.
|
||||
const maxRegistrationErrorBody = 512
|
||||
|
||||
// registrationStatusError turns a non-2xx response into an error that says WHY.
|
||||
//
|
||||
// The body is the point. The frontend explains its refusals there
|
||||
// ("address is required for backend workers", "invalid registration token"),
|
||||
// and discarding it left an operator with a bare status code: the one line that
|
||||
// would tell them which of several possible mistakes they made was read off the
|
||||
// socket and thrown away.
|
||||
func registrationStatusError(resp *http.Response) error {
|
||||
detail, err := io.ReadAll(io.LimitReader(resp.Body, maxRegistrationErrorBody))
|
||||
if err != nil {
|
||||
return "", "", "", "", err
|
||||
xlog.Debug("Could not read the frontend's registration error body", "status", resp.StatusCode, "error", err)
|
||||
}
|
||||
return res.ID, res.APIToken, res.NatsJWT, res.NatsUserSeed, nil
|
||||
msg := strings.Join(strings.Fields(string(detail)), " ")
|
||||
base := fmt.Sprintf("registration failed with status %d", resp.StatusCode)
|
||||
if msg != "" {
|
||||
base = fmt.Sprintf("%s: %s", base, msg)
|
||||
}
|
||||
if isRegistrationRejection(resp.StatusCode) {
|
||||
return fmt.Errorf("%s: %w", base, ErrRegistrationRejected)
|
||||
}
|
||||
return errors.New(base)
|
||||
}
|
||||
|
||||
// isRegistrationRejection reports whether a status is a verdict rather than a
|
||||
// condition that may pass.
|
||||
func isRegistrationRejection(status int) bool {
|
||||
if status == http.StatusRequestTimeout || status == http.StatusTooManyRequests {
|
||||
return false
|
||||
}
|
||||
return status >= 400 && status < 500
|
||||
}
|
||||
|
||||
// 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) {
|
||||
//
|
||||
// It drops every field of the response it does not name, the tunnel credential
|
||||
// among them. Callers that need one use RegisterFullWithRetry.
|
||||
//
|
||||
// The two broker-credential returns it used to carry are gone with the bus, and
|
||||
// so is the Register one-shot that existed only to carry them: it had no caller
|
||||
// left once nothing dialled a broker.
|
||||
func (c *RegistrationClient) RegisterWithRetry(ctx context.Context, body map[string]any, maxRetries int) (nodeID, apiToken string, err error) {
|
||||
res, err := c.RegisterFullWithRetry(ctx, body, maxRetries)
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
return res.ID, res.APIToken, nil
|
||||
}
|
||||
|
||||
// RegisterFullWithRetry retries registration with exponential backoff and
|
||||
// returns the whole response.
|
||||
func (c *RegistrationClient) RegisterFullWithRetry(ctx context.Context, body map[string]any, maxRetries int) (*RegisterResponse, error) {
|
||||
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 errors.Is(err, ErrRegistrationRejected) {
|
||||
// A verdict, not an outage. Reported on the first attempt so the
|
||||
// reason the frontend gave is the first thing in the log rather
|
||||
// than the last, after the ladder.
|
||||
return nil, err
|
||||
}
|
||||
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.
|
||||
|
||||
@@ -0,0 +1,177 @@
|
||||
package workerregistry
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"reflect"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// The case these specs exist for: a worker of this release registering against
|
||||
// a frontend that predates it. A worker no longer sends an address, the old
|
||||
// frontend requires one, and it answers 400 with the reason in the body. Two
|
||||
// things used to go wrong there at once. The reason was discarded, so the
|
||||
// operator saw only "status 400" and had to guess which of several mistakes
|
||||
// they had made; and the retry ladder spent four minutes on a verdict the
|
||||
// frontend reached instantly.
|
||||
var _ = Describe("Registration client refusals", func() {
|
||||
var (
|
||||
attempts atomic.Int32
|
||||
status atomic.Int32
|
||||
body atomic.Value // string
|
||||
server *httptest.Server
|
||||
client *RegistrationClient
|
||||
// seen carries one token per request the handler served, so a spec can
|
||||
// wait for the Nth attempt instead of sleeping for however long the
|
||||
// ladder's backoff happens to be.
|
||||
seen chan struct{}
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
attempts.Store(0)
|
||||
status.Store(int32(http.StatusBadRequest))
|
||||
body.Store(`{"error":{"code":400,"message":"address is required for backend workers"}}`)
|
||||
seen = make(chan struct{}, 64)
|
||||
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
attempts.Add(1)
|
||||
select {
|
||||
case seen <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
w.WriteHeader(int(status.Load()))
|
||||
_, _ = w.Write([]byte(body.Load().(string)))
|
||||
}))
|
||||
client = &RegistrationClient{FrontendURL: server.URL, HTTPTimeout: 2 * time.Second}
|
||||
})
|
||||
|
||||
AfterEach(func() { server.Close() })
|
||||
|
||||
It("quotes what the frontend said", func() {
|
||||
_, err := client.RegisterFull(context.Background(), map[string]any{"name": "w1"})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("status 400"))
|
||||
Expect(err.Error()).To(ContainSubstring("address is required for backend workers"))
|
||||
})
|
||||
|
||||
It("marks a 4xx as a refusal", func() {
|
||||
_, err := client.RegisterFull(context.Background(), map[string]any{"name": "w1"})
|
||||
Expect(err).To(MatchError(ErrRegistrationRejected))
|
||||
})
|
||||
|
||||
It("does not mark a 5xx as a refusal", func() {
|
||||
// A frontend that is restarting or wedged has not judged anything, and
|
||||
// retrying it is the whole reason the ladder exists.
|
||||
status.Store(int32(http.StatusBadGateway))
|
||||
body.Store("bad gateway")
|
||||
_, err := client.RegisterFull(context.Background(), map[string]any{"name": "w1"})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, ErrRegistrationRejected)).To(BeFalse())
|
||||
})
|
||||
|
||||
DescribeTable("treats a status that asks for the same request again as retryable",
|
||||
func(code int) {
|
||||
status.Store(int32(code))
|
||||
body.Store("later")
|
||||
_, err := client.RegisterFull(context.Background(), map[string]any{"name": "w1"})
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(errors.Is(err, ErrRegistrationRejected)).To(BeFalse())
|
||||
},
|
||||
Entry("408 Request Timeout", http.StatusRequestTimeout),
|
||||
Entry("429 Too Many Requests", http.StatusTooManyRequests),
|
||||
)
|
||||
|
||||
It("stops the retry ladder on the first refusal", func() {
|
||||
// Ten attempts on a 400 is roughly four minutes of backoff before the
|
||||
// operator is told anything, and the answer is the same one the
|
||||
// frontend gave immediately.
|
||||
_, err := client.RegisterFullWithRetry(context.Background(), map[string]any{"name": "w1"}, 10)
|
||||
Expect(err).To(MatchError(ErrRegistrationRejected))
|
||||
Expect(err.Error()).To(ContainSubstring("address is required for backend workers"))
|
||||
Expect(attempts.Load()).To(Equal(int32(1)))
|
||||
})
|
||||
|
||||
It("still retries something that is not a refusal", func() {
|
||||
// The control. Without it, a change that returned on EVERY error would
|
||||
// pass the spec above and silently delete the retry behaviour a worker
|
||||
// booting alongside its frontend depends on.
|
||||
status.Store(int32(http.StatusServiceUnavailable))
|
||||
body.Store("starting up")
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := client.RegisterFullWithRetry(ctx, map[string]any{"name": "w1"}, 10)
|
||||
done <- err
|
||||
}()
|
||||
|
||||
// Two tokens is the whole assertion: the ladder came back for a second
|
||||
// attempt on a status that is not a verdict. Waiting on the handler
|
||||
// rather than on a duration makes it exact instead of tolerant.
|
||||
Eventually(seen).Should(Receive())
|
||||
Eventually(seen, "10s").Should(Receive())
|
||||
cancel()
|
||||
|
||||
var ladderErr error
|
||||
Eventually(done).Should(Receive(&ladderErr))
|
||||
Expect(ladderErr).To(HaveOccurred())
|
||||
Expect(errors.Is(ladderErr, ErrRegistrationRejected)).To(BeFalse())
|
||||
Expect(attempts.Load()).To(BeNumerically(">=", 2))
|
||||
})
|
||||
|
||||
It("stops the credential manager's acquire loop on a refusal", func() {
|
||||
// The default worker path goes through Acquire, not the ladder above,
|
||||
// and its bound is 100 attempts rather than 10. A refusal there is the
|
||||
// same verdict and has to end the same way.
|
||||
mgr := NewCredentialManager(func(ctx context.Context) (*RegisterResponse, error) {
|
||||
return client.RegisterFull(ctx, map[string]any{"name": "w1"})
|
||||
}, true)
|
||||
_, err := mgr.Acquire(context.Background())
|
||||
Expect(err).To(MatchError(ErrRegistrationRejected))
|
||||
Expect(attempts.Load()).To(Equal(int32(1)))
|
||||
})
|
||||
})
|
||||
|
||||
// The other direction of the same upgrade: a worker of this release registering
|
||||
// against a frontend that still mints a per-node broker credential.
|
||||
//
|
||||
// The fields are gone from RegisterResponse, so the only question is what
|
||||
// happens to the keys still on the wire. encoding/json ignores a key with no
|
||||
// field, and that is asserted rather than assumed, because a decoder switched
|
||||
// to DisallowUnknownFields would turn every registration against an older
|
||||
// frontend into a hard failure with no other symptom.
|
||||
var _ = Describe("Registering against a frontend that still mints broker credentials", func() {
|
||||
It("decodes the response and drops the keys it no longer has fields for", func() {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusCreated)
|
||||
_, _ = w.Write([]byte(`{"id":"node-7","status":"healthy","api_token":"key-7","tunnel_token":"tunnel-7","nats_jwt":"eyJ0","nats_user_seed":"SUUSER"}`))
|
||||
}))
|
||||
DeferCleanup(server.Close)
|
||||
|
||||
client := &RegistrationClient{FrontendURL: server.URL, HTTPTimeout: 2 * time.Second}
|
||||
res, err := client.RegisterFull(context.Background(), map[string]any{"name": "w1"})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// The fields it DOES have, so the assertion below is about the two
|
||||
// unknown keys and not about a decode that produced nothing.
|
||||
Expect(res.ID).To(Equal("node-7"))
|
||||
Expect(res.APIToken).To(Equal("key-7"))
|
||||
Expect(res.TunnelToken).To(Equal("tunnel-7"))
|
||||
|
||||
// And nowhere for a broker credential to land: asserted on the struct's
|
||||
// own type, because a value assertion would need a field to read and
|
||||
// would stop compiling exactly when the field came back.
|
||||
t := reflect.TypeOf(*res)
|
||||
for _, gone := range []string{"NatsJWT", "NatsUserSeed"} {
|
||||
_, found := t.FieldByName(gone)
|
||||
Expect(found).To(BeFalse(),
|
||||
"RegisterResponse.%s is back: a worker that stores a broker credential is a worker something expects to dial a broker", gone)
|
||||
}
|
||||
})
|
||||
})
|
||||
@@ -2,11 +2,11 @@ package workerregistry
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/natsauth"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
@@ -15,128 +15,117 @@ import (
|
||||
// package (and its gorm/DB dependencies).
|
||||
const statusPending = "pending"
|
||||
|
||||
// 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
|
||||
// so an unauthorized/unapprovable worker exits and surfaces the problem (via a
|
||||
// non-zero exit and the resulting restart) rather than waiting forever.
|
||||
// defaultMaxAttempts bounds how many times Acquire registers before giving up.
|
||||
// It is high enough to ride out a slow admin approval or a transient frontend
|
||||
// outage, but finite so an unauthorized/unapprovable worker exits and surfaces
|
||||
// the problem (via a non-zero exit and the resulting restart) rather than
|
||||
// waiting forever.
|
||||
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 a node's own credentials at startup, waiting
|
||||
// through admin approval when that is required, and holds the tunnel token the
|
||||
// most recent registration minted. 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 {
|
||||
register RegisterFunc
|
||||
requireCreds bool // block until credentials are present (frontend minting in use)
|
||||
// Renamed from NATSCredentialManager and stripped rather than deleted. The JWT
|
||||
// half went with the message bus: nothing mints a broker credential and nothing
|
||||
// opens a connection to present one on. The tunnel token did not go with it,
|
||||
// and it is the reason a manager is still worth having: it is ROTATED by a
|
||||
// registration rather than expiring on a clock, and the frontend keeps only its
|
||||
// hash, so the dialer has to read the current value at dial time instead of
|
||||
// being handed one at startup.
|
||||
type CredentialManager struct {
|
||||
register RegisterFunc
|
||||
// requireApproval blocks Acquire until the node is out of pending.
|
||||
//
|
||||
// Narrower than the requireCreds it replaces: there is no credential left
|
||||
// to wait for being MINTED, only an admin decision to wait through. A
|
||||
// worker that proceeds while pending registers and heartbeats fine, and is
|
||||
// then refused at every tunnel dial, so an operator who wants the wait
|
||||
// rather than the refusal loop asks for it here.
|
||||
requireApproval bool
|
||||
|
||||
// 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)
|
||||
refreshLead float64 // refresh once this fraction of the JWT lifetime has elapsed
|
||||
refreshRetry time.Duration
|
||||
expiryOf func(jwt string) (time.Time, bool)
|
||||
maxAttempts int // bound on Acquire attempts (<=0 = unlimited)
|
||||
|
||||
mu sync.RWMutex
|
||||
jwt string
|
||||
seed string
|
||||
nodeID string
|
||||
// tunnelToken is the node's own tunnel credential from the most recent
|
||||
// registration. It is kept here because every re-registration this manager
|
||||
// performs ROTATES it, so the tunnel client has to read the current value
|
||||
// at dial time rather than be handed one at startup.
|
||||
tunnelToken string
|
||||
}
|
||||
|
||||
// NewNATSCredentialManager 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{
|
||||
register: register,
|
||||
requireCreds: requireCreds,
|
||||
initialBackoff: 2 * time.Second,
|
||||
maxBackoff: 30 * time.Second,
|
||||
maxAttempts: defaultMaxAttempts,
|
||||
refreshLead: 0.75,
|
||||
refreshRetry: 30 * time.Second,
|
||||
expiryOf: jwtExpiry,
|
||||
// NewCredentialManager builds a manager over register. When requireApproval is
|
||||
// true, Acquire blocks through admin approval instead of returning a pending
|
||||
// response.
|
||||
func NewCredentialManager(register RegisterFunc, requireApproval bool) *CredentialManager {
|
||||
return &CredentialManager{
|
||||
register: register,
|
||||
requireApproval: requireApproval,
|
||||
initialBackoff: 2 * time.Second,
|
||||
maxBackoff: 30 * time.Second,
|
||||
maxAttempts: defaultMaxAttempts,
|
||||
}
|
||||
}
|
||||
|
||||
// jwtExpiry decodes the expiry of a minted user JWT. ok is false when the token
|
||||
// is empty/undecodable or carries no expiry (e.g. a non-expiring service JWT).
|
||||
func jwtExpiry(token string) (time.Time, bool) {
|
||||
if token == "" {
|
||||
return time.Time{}, false
|
||||
}
|
||||
uc, err := natsauth.DecodeUserClaims(token)
|
||||
if err != nil || uc.Expires == 0 {
|
||||
return time.Time{}, false
|
||||
}
|
||||
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
|
||||
// A response that carries no tunnel token (a frontend that predates them,
|
||||
// or one whose minting failed) must not wipe a working credential this
|
||||
// worker already holds. Overwriting with "" would lock the tunnel out until
|
||||
// the next registration that did carry one, which is the opposite of what
|
||||
// an empty field means.
|
||||
if res.TunnelToken != "" {
|
||||
m.tunnelToken = res.TunnelToken
|
||||
}
|
||||
}
|
||||
|
||||
// Current returns the latest NATS credentials (both empty until acquired).
|
||||
func (m *NATSCredentialManager) Current() (jwt, seed string) {
|
||||
// TunnelToken returns the node's current tunnel credential, empty until one has
|
||||
// been issued. It is the callback the tunnel client reads on every dial.
|
||||
func (m *CredentialManager) TunnelToken() string {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
return m.jwt, m.seed
|
||||
return m.tunnelToken
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// Provider returns a callback compatible with messaging.WithUserJWTProvider,
|
||||
// supplying the current credentials on each (re)connect.
|
||||
func (m *NATSCredentialManager) Provider() func() (string, string) {
|
||||
return m.Current
|
||||
}
|
||||
|
||||
// HasCredentials reports whether complete NATS credentials have been obtained.
|
||||
func (m *NATSCredentialManager) HasCredentials() bool {
|
||||
jwt, seed := m.Current()
|
||||
return jwt != "" && seed != ""
|
||||
}
|
||||
|
||||
// Acquire registers and, when requireCreds is set, keeps re-registering with
|
||||
// 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) {
|
||||
// Acquire registers and, when requireApproval is set, keeps re-registering with
|
||||
// exponential backoff until the node is approved (status != pending). Without
|
||||
// requireApproval it returns the first successful response.
|
||||
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++ {
|
||||
res, err := m.register(ctx)
|
||||
switch {
|
||||
case errors.Is(err, ErrRegistrationRejected):
|
||||
// The frontend refused rather than failed. Waiting through the full
|
||||
// attempt ladder would delay the operator's only explanation by the
|
||||
// length of the ladder and change nothing about the answer.
|
||||
return nil, err
|
||||
case err != nil:
|
||||
lastReason = err
|
||||
xlog.Warn("Registration failed, retrying", "attempt", attempt, "next_retry", backoff, "error", err)
|
||||
case !m.requireCreds:
|
||||
case !m.requireApproval:
|
||||
m.store(res)
|
||||
return res, nil
|
||||
case res.Status == statusPending:
|
||||
lastReason = fmt.Errorf("node %s still pending admin approval", res.ID)
|
||||
xlog.Info("Node pending admin approval; waiting", "node", res.ID, "attempt", attempt, "next_retry", backoff)
|
||||
case res.NatsJWT == "" || res.NatsUserSeed == "":
|
||||
lastReason = fmt.Errorf("node %s approved but NATS credentials not minted", res.ID)
|
||||
xlog.Info("Node approved but NATS credentials not yet minted; waiting", "node", res.ID, "attempt", attempt, "next_retry", backoff)
|
||||
default:
|
||||
m.store(res)
|
||||
return res, nil
|
||||
@@ -148,53 +137,5 @@ func (m *NATSCredentialManager) Acquire(ctx context.Context) (*RegisterResponse,
|
||||
}
|
||||
backoff = min(backoff*2, m.maxBackoff)
|
||||
}
|
||||
return nil, fmt.Errorf("giving up acquiring NATS credentials after %d attempts: %w", m.maxAttempts, lastReason)
|
||||
}
|
||||
|
||||
// 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
|
||||
// when the current credential has no expiry (nothing to refresh), and a non-nil
|
||||
// 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 {
|
||||
failures := 0
|
||||
for {
|
||||
jwt, _ := m.Current()
|
||||
exp, ok := m.expiryOf(jwt)
|
||||
if !ok {
|
||||
xlog.Debug("NATS credential has no expiry; refresh loop exiting")
|
||||
return nil
|
||||
}
|
||||
wait := max(time.Duration(float64(time.Until(exp))*m.refreshLead), 0)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil
|
||||
case <-time.After(wait):
|
||||
}
|
||||
|
||||
res, err := m.register(ctx)
|
||||
if err == nil && res.NatsJWT != "" && res.NatsUserSeed != "" {
|
||||
m.store(res)
|
||||
failures = 0
|
||||
xlog.Info("Refreshed NATS credentials", "node", res.ID)
|
||||
continue
|
||||
}
|
||||
failures++
|
||||
if err != nil {
|
||||
xlog.Warn("NATS credential refresh failed; will retry", "attempt", failures, "error", err)
|
||||
} else {
|
||||
xlog.Warn("NATS credential refresh returned no credentials; will retry", "attempt", failures)
|
||||
}
|
||||
if m.maxAttempts > 0 && failures >= m.maxAttempts {
|
||||
return fmt.Errorf("NATS credential refresh failed %d times in a row", failures)
|
||||
}
|
||||
// Back off before retrying so a persistent failure near expiry does not spin.
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil
|
||||
case <-time.After(m.refreshRetry):
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("giving up registering after %d attempts: %w", m.maxAttempts, lastReason)
|
||||
}
|
||||
@@ -6,9 +6,6 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/natsauth"
|
||||
"github.com/nats-io/nkeys"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
@@ -50,20 +47,20 @@ func (f *fakeRegister) count() int {
|
||||
return f.calls
|
||||
}
|
||||
|
||||
var _ = Describe("NATSCredentialManager", func() {
|
||||
approved := func(jwt, seed string) *RegisterResponse {
|
||||
return &RegisterResponse{ID: "node-1", Status: "healthy", NatsJWT: jwt, NatsUserSeed: seed}
|
||||
var _ = Describe("CredentialManager", func() {
|
||||
approved := func(tunnelToken string) *RegisterResponse {
|
||||
return &RegisterResponse{ID: "node-1", Status: "healthy", TunnelToken: tunnelToken}
|
||||
}
|
||||
pending := &RegisterResponse{ID: "node-1", Status: "pending"}
|
||||
|
||||
Describe("Acquire (#4 — wait through admin approval)", func() {
|
||||
It("keeps re-registering until the node is approved and credentials are minted", func() {
|
||||
Describe("Acquire (wait through admin approval)", func() {
|
||||
It("keeps re-registering until the node is approved", func() {
|
||||
f := &fakeRegister{steps: []step{
|
||||
{res: pending}, // not approved yet
|
||||
{res: approved("", "")}, // approved but JWT not minted yet
|
||||
{res: approved("jwt-1", "seed-1")}, // finally minted
|
||||
{res: pending}, // not approved yet
|
||||
{res: pending}, // still not approved
|
||||
{res: approved("tunnel-1")}, // approved, and handed its tunnel credential
|
||||
}}
|
||||
m := NewNATSCredentialManager(f.fn(), true /* requireCreds */)
|
||||
m := NewCredentialManager(f.fn(), true /* requireApproval */)
|
||||
m.initialBackoff = time.Millisecond
|
||||
m.maxBackoff = time.Millisecond
|
||||
|
||||
@@ -71,28 +68,24 @@ var _ = Describe("NATSCredentialManager", func() {
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(res.ID).To(Equal("node-1"))
|
||||
Expect(f.count()).To(Equal(3))
|
||||
|
||||
jwt, seed := m.Current()
|
||||
Expect(jwt).To(Equal("jwt-1"))
|
||||
Expect(seed).To(Equal("seed-1"))
|
||||
Expect(m.HasCredentials()).To(BeTrue())
|
||||
Expect(m.TunnelToken()).To(Equal("tunnel-1"))
|
||||
Expect(m.NodeID()).To(Equal("node-1"))
|
||||
})
|
||||
|
||||
It("returns immediately on the first success when credentials are not required (anonymous NATS)", func() {
|
||||
It("returns immediately on the first success when approval is not required", func() {
|
||||
f := &fakeRegister{steps: []step{{res: pending}}}
|
||||
m := NewNATSCredentialManager(f.fn(), false /* requireCreds */)
|
||||
m := NewCredentialManager(f.fn(), false /* requireApproval */)
|
||||
|
||||
res, err := m.Acquire(context.Background())
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(res.Status).To(Equal("pending"))
|
||||
Expect(f.count()).To(Equal(1))
|
||||
Expect(m.HasCredentials()).To(BeFalse())
|
||||
Expect(m.TunnelToken()).To(BeEmpty())
|
||||
})
|
||||
|
||||
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 +96,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,83 +109,48 @@ var _ = Describe("NATSCredentialManager", func() {
|
||||
})
|
||||
})
|
||||
|
||||
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.refreshLead = 0.5
|
||||
m.refreshRetry = time.Millisecond
|
||||
// jwt-1 expires soon; jwt-2 is long-lived so the loop then idles.
|
||||
m.expiryOf = func(jwt string) (time.Time, bool) {
|
||||
switch jwt {
|
||||
case "jwt-1":
|
||||
return time.Now().Add(40 * time.Millisecond), true
|
||||
case "jwt-2":
|
||||
return time.Now().Add(time.Hour), true
|
||||
default:
|
||||
return time.Time{}, false
|
||||
}
|
||||
}
|
||||
m.store(approved("jwt-1", "seed-1"))
|
||||
// The behaviour that survived the rename, and which nothing pinned while
|
||||
// the manager was about a JWT.
|
||||
//
|
||||
// The frontend mints a FRESH tunnel token on every registration and keeps
|
||||
// only its hash, so the previous one stops working the moment a new one is
|
||||
// issued. A manager that handed out the first value it saw would lock the
|
||||
// worker out of its own tunnel after any re-registration, and the symptom
|
||||
// would be a 401 on a dial rather than anything at registration time.
|
||||
Describe("TunnelToken (rotation)", func() {
|
||||
It("serves the token from the most recent registration, not the first", func() {
|
||||
m := NewCredentialManager(nil, false)
|
||||
m.store(approved("tunnel-1"))
|
||||
Expect(m.TunnelToken()).To(Equal("tunnel-1"))
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
go func() { _ = m.RefreshLoop(ctx) }()
|
||||
|
||||
Eventually(func() string {
|
||||
jwt, _ := m.Current()
|
||||
return jwt
|
||||
}, "2s", "10ms").Should(Equal("jwt-2"))
|
||||
m.store(approved("tunnel-2"))
|
||||
Expect(m.TunnelToken()).To(Equal("tunnel-2"),
|
||||
"the frontend keeps only the newest token's hash, so serving the first one locks this worker out of its own tunnel")
|
||||
})
|
||||
|
||||
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.refreshLead = 0.5
|
||||
m.refreshRetry = time.Millisecond
|
||||
m.maxAttempts = 3
|
||||
m.expiryOf = func(string) (time.Time, bool) { return time.Now().Add(time.Millisecond), true }
|
||||
m.store(approved("jwt-1", "seed-1"))
|
||||
|
||||
errCh := make(chan error, 1)
|
||||
go func() { errCh <- m.RefreshLoop(context.Background()) }()
|
||||
Eventually(errCh, "2s").Should(Receive(MatchError(ContainSubstring("3 times in a row"))))
|
||||
})
|
||||
|
||||
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.expiryOf = func(string) (time.Time, bool) { return time.Time{}, false }
|
||||
m.store(approved("static", "seed"))
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() { _ = m.RefreshLoop(context.Background()); close(done) }()
|
||||
Eventually(done, "1s").Should(BeClosed())
|
||||
Expect(f.count()).To(Equal(0)) // never tried to re-register
|
||||
})
|
||||
})
|
||||
|
||||
Describe("jwtExpiry default", func() {
|
||||
It("decodes the expiry of a real minted worker JWT", func() {
|
||||
akp, err := nkeys.CreateAccount()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
seed, err := akp.Seed()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
cfg := natsauth.Config{AccountSeed: string(seed), WorkerJWTTTL: time.Hour}
|
||||
token, _, err := cfg.MintWorkerJWT("node-1", "backend")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
exp, ok := jwtExpiry(token)
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(exp).To(BeTemporally("~", time.Now().Add(time.Hour), 2*time.Minute))
|
||||
})
|
||||
|
||||
It("reports no expiry for an empty or undecodable token", func() {
|
||||
_, ok := jwtExpiry("")
|
||||
Expect(ok).To(BeFalse())
|
||||
_, ok = jwtExpiry("not-a-jwt")
|
||||
Expect(ok).To(BeFalse())
|
||||
It("keeps a working token when a registration carries none", func() {
|
||||
// A frontend that predates tunnel tokens, or one whose minting
|
||||
// failed, sends the field empty. Empty means "nothing new", not
|
||||
// "revoked": overwriting would lock the tunnel out until the next
|
||||
// registration that did carry one.
|
||||
m := NewCredentialManager(nil, false)
|
||||
m.store(approved("tunnel-1"))
|
||||
m.store(approved(""))
|
||||
Expect(m.TunnelToken()).To(Equal("tunnel-1"))
|
||||
})
|
||||
})
|
||||
})
|
||||
|
||||
// Deleted with the bus: the "RefreshLoop" Describe and the "jwtExpiry default"
|
||||
// Describe.
|
||||
//
|
||||
// RefreshLoop pinned that a worker re-registers before its minted broker JWT
|
||||
// expires and serves the new credential to the next connection. Every noun in
|
||||
// that sentence is gone: there is no JWT, no expiry to read, and no connection
|
||||
// to serve it to. It is retired rather than moved. The one credential a
|
||||
// registration still mints, the tunnel token, is rotated by a REGISTRATION
|
||||
// rather than by a clock, and the property that matters about it is that the
|
||||
// dialer reads the current value, which the rotation Describe above pins.
|
||||
//
|
||||
// jwtExpiry pinned that a real minted worker JWT's expiry decoded, which was
|
||||
// the only thing in this package that needed a broker library at all.
|
||||
+102
-136
@@ -5,30 +5,35 @@ import (
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/pkg/natsauth"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// DistributedConfig holds configuration for horizontal scaling mode.
|
||||
// When Enabled is true, PostgreSQL and NATS are required.
|
||||
// When Enabled is true, PostgreSQL is required. Nothing else is: fan-out rides
|
||||
// PostgreSQL and every worker is reached over the tunnel it dials out.
|
||||
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
|
||||
Enabled bool // --distributed / LOCALAI_DISTRIBUTED
|
||||
InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty)
|
||||
// AdvertiseAddr is the host:port OTHER REPLICAS dial to reach this one,
|
||||
// which is not the address this process binds: a replica behind a service
|
||||
// or a NAT binds one and is reached at another. Empty means "work it out",
|
||||
// by asking the kernel which local address routes to PostgreSQL; that
|
||||
// answer is only usable when the database is remote, so a deployment with
|
||||
// a local or sidecar database has to set this.
|
||||
AdvertiseAddr string // LOCALAI_DISTRIBUTED_ADVERTISE_ADDR
|
||||
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
|
||||
// node-register endpoints and the worker file-transfer server cannot run
|
||||
// unauthenticated. Mirrors NatsRequireAuth for the NATS bus.
|
||||
// unauthenticated.
|
||||
RegistrationRequireAuth bool // LOCALAI_REGISTRATION_REQUIRE_AUTH
|
||||
// RequireAuth is the umbrella switch (LOCALAI_DISTRIBUTED_REQUIRE_AUTH) for
|
||||
// distributed-mode auth: when true it implies BOTH NatsRequireAuth and
|
||||
// RegistrationRequireAuth, so a single knob locks down the bus and the
|
||||
// registration/file-transfer layer together. The granular flags remain
|
||||
// available to enforce just one layer.
|
||||
// distributed-mode auth. It implies RegistrationRequireAuth, and that is
|
||||
// now ALL it implies: it used to imply a NatsRequireAuth as well, and there
|
||||
// is no message bus left for that half to lock down. The granular flag
|
||||
// remains available.
|
||||
RequireAuth bool // LOCALAI_DISTRIBUTED_REQUIRE_AUTH
|
||||
AutoApproveNodes bool // --auto-approve-nodes / LOCALAI_AUTO_APPROVE_NODES (skip admin approval for new workers)
|
||||
// SharedModels asserts that every node (frontend and workers) mounts the
|
||||
@@ -40,16 +45,6 @@ type DistributedConfig struct {
|
||||
// false preserves the historical per-node staging behavior.
|
||||
SharedModels bool // --distributed-shared-models / LOCALAI_DISTRIBUTED_SHARED_MODELS
|
||||
|
||||
// NATS JWT auth (optional; see pkg/natsauth and docs/features/distributed-mode.md)
|
||||
NatsAccountSeed string // LOCALAI_NATS_ACCOUNT_SEED — account signing seed to mint per-node worker JWTs
|
||||
NatsServiceJWT string // LOCALAI_NATS_SERVICE_JWT — user JWT for frontends / agent workers
|
||||
NatsServiceSeed string // LOCALAI_NATS_SERVICE_SEED — signing seed paired with service JWT
|
||||
NatsWorkerJWTTTL time.Duration // LOCALAI_NATS_WORKER_JWT_TTL — minted worker JWT lifetime (default 24h)
|
||||
NatsRequireAuth bool // LOCALAI_NATS_REQUIRE_AUTH — fail startup if NATS credentials are missing
|
||||
NatsTLSCA string // LOCALAI_NATS_TLS_CA — PEM file for private CA (server verify)
|
||||
NatsTLSCert string // LOCALAI_NATS_TLS_CERT — client cert for NATS mTLS
|
||||
NatsTLSKey string // LOCALAI_NATS_TLS_KEY — client key paired with NatsTLSCert
|
||||
|
||||
// S3 configuration (used when StorageURL is set)
|
||||
StorageBucket string // --storage-bucket / LOCALAI_STORAGE_BUCKET
|
||||
StorageRegion string // --storage-region / LOCALAI_STORAGE_REGION
|
||||
@@ -75,8 +70,19 @@ type DistributedConfig struct {
|
||||
|
||||
MCPCIJobTimeout time.Duration // MCP CI job execution timeout (default 10m)
|
||||
|
||||
BackendInstallTimeout time.Duration // NATS round-trip timeout for backend.install (default 15m)
|
||||
BackendUpgradeTimeout time.Duration // NATS round-trip timeout for backend.upgrade (default 15m)
|
||||
// WorkerReconnectGrace is how long a worker whose tunnel was lost is
|
||||
// treated as reconnecting rather than gone. It is the ONLY thing that
|
||||
// separates a worker re-homing between frontend replicas from one that has
|
||||
// left, and absence is what makes the scheduler stop placing work and reap
|
||||
// the worker's rows, so a grace shorter than the worker's own reconnect
|
||||
// backoff condemns workers that are behaving exactly as designed.
|
||||
//
|
||||
// Zero means unset (DefaultWorkerReconnectGrace applies). Measured on the
|
||||
// database clock, so every replica agrees on when the window ends.
|
||||
WorkerReconnectGrace time.Duration // LOCALAI_WORKER_RECONNECT_GRACE
|
||||
|
||||
BackendInstallTimeout time.Duration // control round-trip timeout for backend.install (default 15m)
|
||||
BackendUpgradeTimeout time.Duration // control round-trip timeout for backend.upgrade (default 15m)
|
||||
// ModelLoadTimeout is the gRPC deadline for the remote LoadModel call the
|
||||
// router issues once a worker has the backend installed and the model files
|
||||
// staged. It therefore covers only the backend's own checkpoint load and
|
||||
@@ -139,9 +145,11 @@ 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 message-bus URL is required, and none is dialled. The last family
|
||||
// that needed one was agent.<agent>.cancel, which now rides the agent
|
||||
// worker's own tunnel as a control verb; a distributed deployment needs
|
||||
// PostgreSQL and the frontends' own HTTP listener. The flag is still
|
||||
// accepted so an existing command line starts unchanged.
|
||||
// S3 credentials must be paired
|
||||
if (c.StorageAccessKey != "" && c.StorageSecretKey == "") ||
|
||||
(c.StorageAccessKey == "" && c.StorageSecretKey != "") {
|
||||
@@ -157,13 +165,6 @@ func (c DistributedConfig) Validate() error {
|
||||
}
|
||||
xlog.Warn("distributed mode running without registration token — node endpoints and the worker file-transfer server are unprotected; set LOCALAI_REGISTRATION_TOKEN, or LOCALAI_DISTRIBUTED_REQUIRE_AUTH=true to fail closed")
|
||||
}
|
||||
if err := c.NatsAuthConfig().Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := c.NatsTLSFiles().Validate(); err != nil {
|
||||
return err
|
||||
}
|
||||
c.NatsAuthConfig().WarnIfInsecure(true)
|
||||
// Check for negative durations
|
||||
for name, d := range map[string]time.Duration{
|
||||
FlagMCPToolTimeout: c.MCPToolTimeout,
|
||||
@@ -172,11 +173,12 @@ func (c DistributedConfig) Validate() error {
|
||||
FlagDrainTimeout: c.DrainTimeout,
|
||||
FlagHealthCheckInterval: c.HealthCheckInterval,
|
||||
FlagStaleNodeThreshold: c.StaleNodeThreshold,
|
||||
FlagNodeHeartbeatCheckpoint: c.NodeHeartbeatCheckpoint,
|
||||
FlagMCPCIJobTimeout: c.MCPCIJobTimeout,
|
||||
FlagBackendInstallTimeout: c.BackendInstallTimeout,
|
||||
FlagBackendUpgradeTimeout: c.BackendUpgradeTimeout,
|
||||
FlagModelLoadTimeout: c.ModelLoadTimeout,
|
||||
FlagWorkerReconnectGrace: c.WorkerReconnectGrace,
|
||||
FlagNodeHeartbeatCheckpoint: c.NodeHeartbeatCheckpoint,
|
||||
} {
|
||||
if d < 0 {
|
||||
return fmt.Errorf("%s must not be negative", name)
|
||||
@@ -197,9 +199,11 @@ func WithDistributedInstanceID(id string) AppOption {
|
||||
}
|
||||
}
|
||||
|
||||
func WithNatsURL(url string) AppOption {
|
||||
// WithDistributedAdvertiseAddr pins the host:port peers dial to reach this
|
||||
// replica, overriding the route-based discovery.
|
||||
func WithDistributedAdvertiseAddr(addr string) AppOption {
|
||||
return func(o *ApplicationConfig) {
|
||||
o.Distributed.NatsURL = url
|
||||
o.Distributed.AdvertiseAddr = addr
|
||||
}
|
||||
}
|
||||
|
||||
@@ -209,42 +213,14 @@ func WithRegistrationToken(token string) AppOption {
|
||||
}
|
||||
}
|
||||
|
||||
func WithNatsAccountSeed(seed string) AppOption {
|
||||
return func(o *ApplicationConfig) {
|
||||
o.Distributed.NatsAccountSeed = seed
|
||||
}
|
||||
}
|
||||
|
||||
func WithNatsServiceJWT(jwt string) AppOption {
|
||||
return func(o *ApplicationConfig) {
|
||||
o.Distributed.NatsServiceJWT = jwt
|
||||
}
|
||||
}
|
||||
|
||||
func WithNatsServiceSeed(seed string) AppOption {
|
||||
return func(o *ApplicationConfig) {
|
||||
o.Distributed.NatsServiceSeed = seed
|
||||
}
|
||||
}
|
||||
|
||||
func WithNatsWorkerJWTTTL(d time.Duration) AppOption {
|
||||
return func(o *ApplicationConfig) {
|
||||
o.Distributed.NatsWorkerJWTTTL = d
|
||||
}
|
||||
}
|
||||
|
||||
var EnableNatsRequireAuth = func(o *ApplicationConfig) {
|
||||
o.Distributed.NatsRequireAuth = true
|
||||
}
|
||||
|
||||
// EnableRegistrationRequireAuth makes an empty registration token a hard error
|
||||
// in distributed mode (see DistributedConfig.RegistrationRequireAuth).
|
||||
var EnableRegistrationRequireAuth = func(o *ApplicationConfig) {
|
||||
o.Distributed.RegistrationRequireAuth = true
|
||||
}
|
||||
|
||||
// EnableDistributedRequireAuth is the umbrella switch implying both
|
||||
// NatsRequireAuth and RegistrationRequireAuth (see DistributedConfig.RequireAuth).
|
||||
// EnableDistributedRequireAuth is the umbrella switch implying
|
||||
// RegistrationRequireAuth (see DistributedConfig.RequireAuth).
|
||||
var EnableDistributedRequireAuth = func(o *ApplicationConfig) {
|
||||
o.Distributed.RequireAuth = true
|
||||
}
|
||||
@@ -255,30 +231,6 @@ func (c DistributedConfig) RegistrationAuthRequired() bool {
|
||||
return c.RegistrationRequireAuth || c.RequireAuth
|
||||
}
|
||||
|
||||
// NatsAuthRequired reports whether NATS JWT credentials must be present — the
|
||||
// granular flag or the umbrella.
|
||||
func (c DistributedConfig) NatsAuthRequired() bool {
|
||||
return c.NatsRequireAuth || c.RequireAuth
|
||||
}
|
||||
|
||||
func WithNatsTLSCA(path string) AppOption {
|
||||
return func(o *ApplicationConfig) {
|
||||
o.Distributed.NatsTLSCA = path
|
||||
}
|
||||
}
|
||||
|
||||
func WithNatsTLSCert(path string) AppOption {
|
||||
return func(o *ApplicationConfig) {
|
||||
o.Distributed.NatsTLSCert = path
|
||||
}
|
||||
}
|
||||
|
||||
func WithNatsTLSKey(path string) AppOption {
|
||||
return func(o *ApplicationConfig) {
|
||||
o.Distributed.NatsTLSKey = path
|
||||
}
|
||||
}
|
||||
|
||||
func WithStorageURL(url string) AppOption {
|
||||
return func(o *ApplicationConfig) {
|
||||
o.Distributed.StorageURL = url
|
||||
@@ -309,6 +261,14 @@ func WithStorageSecretKey(key string) AppOption {
|
||||
}
|
||||
}
|
||||
|
||||
// WithWorkerReconnectGrace sets how long a lost worker tunnel is read as
|
||||
// reconnecting rather than gone (see DistributedConfig.WorkerReconnectGrace).
|
||||
func WithWorkerReconnectGrace(d time.Duration) AppOption {
|
||||
return func(o *ApplicationConfig) {
|
||||
o.Distributed.WorkerReconnectGrace = d
|
||||
}
|
||||
}
|
||||
|
||||
func WithBackendInstallTimeout(d time.Duration) AppOption {
|
||||
return func(o *ApplicationConfig) {
|
||||
o.Distributed.BackendInstallTimeout = d
|
||||
@@ -414,18 +374,21 @@ func WithModelSchedulingConfigPath(path string) AppOption {
|
||||
// them as constants prevents the string from drifting from the actual
|
||||
// flag a future rename would produce.
|
||||
const (
|
||||
FlagMCPToolTimeout = "mcp-tool-timeout"
|
||||
FlagMCPDiscoveryTimeout = "mcp-discovery-timeout"
|
||||
FlagWorkerWaitTimeout = "worker-wait-timeout"
|
||||
FlagDrainTimeout = "drain-timeout"
|
||||
FlagHealthCheckInterval = "health-check-interval"
|
||||
FlagStaleNodeThreshold = "stale-node-threshold"
|
||||
FlagNodeHeartbeatCheckpoint = "node-heartbeat-checkpoint"
|
||||
FlagMCPCIJobTimeout = "mcp-ci-job-timeout"
|
||||
FlagBackendInstallTimeout = "backend-install-timeout"
|
||||
FlagBackendUpgradeTimeout = "backend-upgrade-timeout"
|
||||
FlagModelLoadTimeout = "model-load-timeout"
|
||||
FlagMCPToolTimeout = "mcp-tool-timeout"
|
||||
FlagMCPDiscoveryTimeout = "mcp-discovery-timeout"
|
||||
FlagWorkerWaitTimeout = "worker-wait-timeout"
|
||||
FlagDrainTimeout = "drain-timeout"
|
||||
FlagHealthCheckInterval = "health-check-interval"
|
||||
FlagStaleNodeThreshold = "stale-node-threshold"
|
||||
FlagMCPCIJobTimeout = "mcp-ci-job-timeout"
|
||||
FlagBackendInstallTimeout = "backend-install-timeout"
|
||||
FlagBackendUpgradeTimeout = "backend-upgrade-timeout"
|
||||
FlagModelLoadTimeout = "model-load-timeout"
|
||||
// FlagWorkerReconnectGrace names the reconnect-grace knob. Validate quotes
|
||||
// it when the operator hands it a negative duration.
|
||||
FlagWorkerReconnectGrace = "worker-reconnect-grace"
|
||||
FlagModelLoadWait = "model-load-wait"
|
||||
FlagNodeHeartbeatCheckpoint = "node-heartbeat-checkpoint"
|
||||
// FlagDiskHeadroomCheck names the disk-headroom toggle. It is quoted in
|
||||
// the warning the check emits while disabled, so the operator reading a
|
||||
// log line knows exactly which knob produced it.
|
||||
@@ -456,6 +419,30 @@ const (
|
||||
// LocalAI (with progress the client can act on) rather than from a proxy
|
||||
// dropping the connection.
|
||||
DefaultModelLoadWait = 60 * time.Second
|
||||
// DefaultWorkerReconnectGrace covers a worker that misses one reconnect at
|
||||
// the ceiling and lands on the next, with margin. The worker's own numbers
|
||||
// (core/services/worker/tunnel.go) are a 30s backoff ceiling
|
||||
// (tunnelBackoffMax) and a 10s dial budget (tunnelHandshakeTimeout), so two
|
||||
// ceiling waits with a hung dial between them puts the worker back at 70s,
|
||||
// not 60s: two waits alone is the boundary, not a bound.
|
||||
//
|
||||
// The ceiling is reachable precisely when it matters. The backoff resets
|
||||
// only after a session that lasted tunnelHealthyAfter (30s), which a
|
||||
// replica accepting a dial and then dying denies, so a worker crossing a
|
||||
// rolling frontend restart climbs to the ceiling rather than sitting near
|
||||
// the 500ms floor.
|
||||
//
|
||||
// 90s therefore has margin where 60s sat on the edge. The asymmetry is
|
||||
// deliberate: too short and a worker that is reconnecting exactly as
|
||||
// designed is reported GONE, which licenses a reap and costs a model
|
||||
// reload; too long and a worker that really has died is reaped later. The
|
||||
// second is cheaper, so the default errs long.
|
||||
//
|
||||
// Raising it further makes a rolling frontend restart safer still; lowering
|
||||
// it reaps a dead worker sooner. There IS a value at which a live worker is
|
||||
// reported as gone: any grace shorter than that worker's actual reconnect.
|
||||
// That is why this is a duration and not a boolean.
|
||||
DefaultWorkerReconnectGrace = 90 * time.Second
|
||||
)
|
||||
|
||||
// ModelLoadWaitUnbounded records LOCALAI_MODEL_LOAD_WAIT=0 — "wait as long as
|
||||
@@ -466,42 +453,21 @@ const ModelLoadWaitUnbounded = -1 * time.Second
|
||||
// DefaultMaxUploadSize is the default maximum upload body size (50 GB).
|
||||
const DefaultMaxUploadSize int64 = 50 << 30
|
||||
|
||||
// NatsTLSFiles returns NATS TLS/mTLS PEM paths for the messaging client.
|
||||
func (c DistributedConfig) NatsTLSFiles() messaging.TLSFiles {
|
||||
return messaging.TLSFiles{
|
||||
CA: c.NatsTLSCA,
|
||||
Cert: c.NatsTLSCert,
|
||||
Key: c.NatsTLSKey,
|
||||
}
|
||||
}
|
||||
|
||||
// NatsMessagingOptions builds messaging client options (JWT + TLS) for distributed components.
|
||||
// Pass explicit userJWT/userSeed when set (e.g. worker overrides); empty uses service JWT from config.
|
||||
func (c DistributedConfig) NatsMessagingOptions(userJWT, userSeed string) []messaging.Option {
|
||||
var opts []messaging.Option
|
||||
jwt, seed := userJWT, userSeed
|
||||
if jwt == "" && seed == "" {
|
||||
auth := c.NatsAuthConfig()
|
||||
jwt, seed = auth.ServiceUserJWT, auth.ServiceUserSeed
|
||||
}
|
||||
if jwt != "" && seed != "" {
|
||||
opts = append(opts, messaging.WithUserJWT(jwt, seed))
|
||||
}
|
||||
if tls := c.NatsTLSFiles(); tls.Enabled() {
|
||||
opts = append(opts, messaging.WithTLS(tls))
|
||||
}
|
||||
return opts
|
||||
}
|
||||
|
||||
// NatsAuthConfig builds pkg/natsauth settings from distributed configuration.
|
||||
func (c DistributedConfig) NatsAuthConfig() natsauth.Config {
|
||||
return natsauth.Config{
|
||||
AccountSeed: c.NatsAccountSeed,
|
||||
ServiceUserJWT: c.NatsServiceJWT,
|
||||
ServiceUserSeed: c.NatsServiceSeed,
|
||||
WorkerJWTTTL: c.NatsWorkerJWTTTL,
|
||||
RequireAuth: c.NatsAuthRequired(),
|
||||
// ReconnectGraceOrDefault returns the configured worker reconnect grace or the
|
||||
// default.
|
||||
//
|
||||
// A non-positive value falls back rather than being taken verbatim, which is
|
||||
// the opposite of what the timeout knobs above do, and deliberately. A
|
||||
// negative grace makes every departure older than the window the instant it is
|
||||
// stamped, so a worker two seconds into a normal reconnect reports as GONE, and
|
||||
// gone is the one answer a caller may reap and evict on. Validate rejects a
|
||||
// negative duration at startup; this is the second line, for a config built in
|
||||
// code that never went through it.
|
||||
func (c DistributedConfig) ReconnectGraceOrDefault() time.Duration {
|
||||
if c.WorkerReconnectGrace <= 0 {
|
||||
return DefaultWorkerReconnectGrace
|
||||
}
|
||||
return c.WorkerReconnectGrace
|
||||
}
|
||||
|
||||
// BackendInstallTimeoutOrDefault returns the configured timeout or the default.
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package config_test
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
@@ -95,7 +97,6 @@ var _ = Describe("DistributedConfig.Validate negative-duration errors", func() {
|
||||
It("rejects a negative BackendInstallTimeout with the flag name in the error", func() {
|
||||
c := config.DistributedConfig{
|
||||
Enabled: true,
|
||||
NatsURL: "nats://localhost:4222",
|
||||
BackendInstallTimeout: -1 * time.Second,
|
||||
}
|
||||
err := c.Validate()
|
||||
@@ -107,7 +108,6 @@ var _ = Describe("DistributedConfig.Validate negative-duration errors", func() {
|
||||
It("rejects a negative BackendUpgradeTimeout with the flag name in the error", func() {
|
||||
c := config.DistributedConfig{
|
||||
Enabled: true,
|
||||
NatsURL: "nats://localhost:4222",
|
||||
BackendUpgradeTimeout: -1 * time.Second,
|
||||
}
|
||||
err := c.Validate()
|
||||
@@ -118,7 +118,6 @@ var _ = Describe("DistributedConfig.Validate negative-duration errors", func() {
|
||||
It("rejects a negative ModelLoadTimeout with the flag name in the error", func() {
|
||||
c := config.DistributedConfig{
|
||||
Enabled: true,
|
||||
NatsURL: "nats://localhost:4222",
|
||||
ModelLoadTimeout: -1 * time.Second,
|
||||
}
|
||||
err := c.Validate()
|
||||
@@ -130,7 +129,6 @@ var _ = Describe("DistributedConfig.Validate negative-duration errors", func() {
|
||||
It("accepts all-zero durations as valid (defaults apply)", func() {
|
||||
c := config.DistributedConfig{
|
||||
Enabled: true,
|
||||
NatsURL: "nats://localhost:4222",
|
||||
}
|
||||
Expect(c.Validate()).To(Succeed())
|
||||
})
|
||||
@@ -140,7 +138,6 @@ var _ = Describe("DistributedConfig.Validate registration auth", func() {
|
||||
It("rejects an empty registration token when RequireAuth is set", func() {
|
||||
c := config.DistributedConfig{
|
||||
Enabled: true,
|
||||
NatsURL: "nats://localhost:4222",
|
||||
RegistrationRequireAuth: true,
|
||||
}
|
||||
err := c.Validate()
|
||||
@@ -152,7 +149,6 @@ var _ = Describe("DistributedConfig.Validate registration auth", func() {
|
||||
It("accepts a set registration token when RequireAuth is set", func() {
|
||||
c := config.DistributedConfig{
|
||||
Enabled: true,
|
||||
NatsURL: "nats://localhost:4222",
|
||||
RegistrationToken: "s3cret",
|
||||
RegistrationRequireAuth: true,
|
||||
}
|
||||
@@ -162,7 +158,6 @@ var _ = Describe("DistributedConfig.Validate registration auth", func() {
|
||||
It("warns but succeeds with an empty token when RequireAuth is unset", func() {
|
||||
c := config.DistributedConfig{
|
||||
Enabled: true,
|
||||
NatsURL: "nats://localhost:4222",
|
||||
}
|
||||
Expect(c.Validate()).To(Succeed())
|
||||
})
|
||||
@@ -170,12 +165,7 @@ var _ = Describe("DistributedConfig.Validate registration auth", func() {
|
||||
It("rejects an empty token when the umbrella RequireAuth is set", func() {
|
||||
c := config.DistributedConfig{
|
||||
Enabled: true,
|
||||
NatsURL: "nats://localhost:4222",
|
||||
RequireAuth: true,
|
||||
// Provide NATS creds so only the registration-token gap remains.
|
||||
NatsServiceJWT: "jwt",
|
||||
NatsServiceSeed: "seed",
|
||||
NatsAccountSeed: "acct",
|
||||
}
|
||||
err := c.Validate()
|
||||
Expect(err).To(HaveOccurred())
|
||||
@@ -183,18 +173,94 @@ var _ = Describe("DistributedConfig.Validate registration auth", func() {
|
||||
Expect(err.Error()).To(ContainSubstring("LOCALAI_REGISTRATION_TOKEN"))
|
||||
})
|
||||
|
||||
It("the umbrella implies NATS auth is required", func() {
|
||||
// The umbrella used to imply two things, and now implies one.
|
||||
//
|
||||
// The It that stood here pinned "LOCALAI_DISTRIBUTED_REQUIRE_AUTH makes a
|
||||
// missing broker service JWT fatal". That is retired, not moved: there is
|
||||
// no broker connection to demand a credential for, so a startup that failed
|
||||
// on a missing one would fail on the absence of something nothing uses. The
|
||||
// half of the umbrella that survives is the registration layer, and the It
|
||||
// above it pins exactly that: an empty registration token under the
|
||||
// umbrella is still fatal, and still says which knob to set.
|
||||
It("implies the registration layer and nothing else", func() {
|
||||
c := config.DistributedConfig{
|
||||
Enabled: true,
|
||||
NatsURL: "nats://localhost:4222",
|
||||
RegistrationToken: "tok", // registration layer satisfied
|
||||
RequireAuth: true, // umbrella → NATS creds now required
|
||||
RegistrationToken: "tok",
|
||||
RequireAuth: true,
|
||||
}
|
||||
Expect(c.NatsAuthRequired()).To(BeTrue())
|
||||
Expect(c.RegistrationAuthRequired()).To(BeTrue())
|
||||
// Missing NATS service JWT/seed must now be fatal.
|
||||
err := c.Validate()
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("LOCALAI_NATS_REQUIRE_AUTH"))
|
||||
// And with the registration layer satisfied there is nothing else left
|
||||
// for the umbrella to demand, so startup succeeds.
|
||||
Expect(c.Validate()).To(Succeed())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("DistributedConfig worker reconnect grace", func() {
|
||||
It("defaults clear of two ceiling backoffs plus the dial between them", func() {
|
||||
// The worker's own numbers (core/services/worker/tunnel.go): a 30s
|
||||
// backoff ceiling and a 10s dial budget, so two ceiling waits with a
|
||||
// hung dial between them puts the worker back at 70s. 60s would sit
|
||||
// under that and condemn a worker reconnecting exactly as designed;
|
||||
// 90s clears it with margin.
|
||||
Expect(config.DistributedConfig{}.ReconnectGraceOrDefault()).To(Equal(90 * time.Second))
|
||||
Expect(config.DefaultWorkerReconnectGrace).To(BeNumerically(">", 70*time.Second),
|
||||
"the default must clear two ceiling backoffs plus one handshake timeout")
|
||||
})
|
||||
|
||||
It("takes a configured worker reconnect grace verbatim", func() {
|
||||
cfg := config.DistributedConfig{WorkerReconnectGrace: 5 * time.Minute}
|
||||
Expect(cfg.ReconnectGraceOrDefault()).To(Equal(5 * time.Minute))
|
||||
})
|
||||
|
||||
It("falls back to the default rather than condemning every worker on a negative value", func() {
|
||||
// A negative grace makes every departure older than the window the
|
||||
// instant it is stamped, which reports a worker that has been gone for
|
||||
// two seconds as GONE, and gone is the one value a caller may reap on.
|
||||
cfg := config.DistributedConfig{WorkerReconnectGrace: -1 * time.Second}
|
||||
Expect(cfg.ReconnectGraceOrDefault()).To(Equal(config.DefaultWorkerReconnectGrace))
|
||||
})
|
||||
|
||||
It("refuses to start on a negative grace rather than reaping on it", func() {
|
||||
// The flag is in Validate's negative-duration table, and this is what
|
||||
// says so. A negative grace makes every departure older than the window
|
||||
// the instant it is stamped, so the deployment would answer GONE for
|
||||
// every worker that has ever lost a tunnel, and gone is the one answer
|
||||
// a caller may reap and evict on.
|
||||
c := config.DistributedConfig{
|
||||
Enabled: true,
|
||||
RegistrationToken: "tok",
|
||||
WorkerReconnectGrace: -1 * time.Second,
|
||||
}
|
||||
err := c.Validate()
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring(config.FlagWorkerReconnectGrace))
|
||||
})
|
||||
|
||||
It("is settable through the application option", func() {
|
||||
o := &config.ApplicationConfig{}
|
||||
config.WithWorkerReconnectGrace(90 * time.Second)(o)
|
||||
Expect(o.Distributed.ReconnectGraceOrDefault()).To(Equal(90 * time.Second))
|
||||
})
|
||||
})
|
||||
|
||||
// The frontend's distributed configuration has nowhere to put a broker URL, and
|
||||
// that absence is what makes the accepted-and-ignored CLI flags ignored.
|
||||
//
|
||||
// A help string saying "ignored" is a promise; a missing field is the mechanism.
|
||||
// Asserted by reflection rather than by reading a value, because a value
|
||||
// assertion needs a field to read and would therefore stop compiling exactly
|
||||
// when the property it guards is restored, which is the failure mode this
|
||||
// replaces: a spec that vanishes with the regression it was meant to catch.
|
||||
var _ = Describe("the distributed configuration's broker surface", func() {
|
||||
It("carries no NATS credential, TLS or URL field", func() {
|
||||
t := reflect.TypeOf(config.DistributedConfig{})
|
||||
var carried []string
|
||||
for i := range t.NumField() {
|
||||
if strings.HasPrefix(t.Field(i).Name, "Nats") {
|
||||
carried = append(carried, t.Field(i).Name)
|
||||
}
|
||||
}
|
||||
Expect(carried).To(BeEmpty(),
|
||||
"DistributedConfig grew %v back: a value the frontend can store is a value something can dial, and no component of a distributed deployment dials a message bus", carried)
|
||||
})
|
||||
})
|
||||
+96
-22
@@ -1,12 +1,14 @@
|
||||
package http
|
||||
|
||||
import (
|
||||
"context"
|
||||
"embed"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"math"
|
||||
"mime"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
@@ -28,6 +30,7 @@ import (
|
||||
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
clustersvc "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"
|
||||
@@ -499,17 +502,19 @@ func API(application *application.Application) (*echo.Echo, error) {
|
||||
var opcache *galleryop.OpCache
|
||||
if !application.ApplicationConfig().DisableWebUI {
|
||||
opcache = galleryop.NewOpCache(application.GalleryService())
|
||||
// In distributed mode, wire the NATS client + gallery store so this
|
||||
// replica's OpCache stays in sync with peers — without this the
|
||||
// In distributed mode, wire the broadcast carrier + gallery store so
|
||||
// this replica's OpCache stays in sync with peers. Without this the
|
||||
// /api/operations endpoint returns whatever this single replica
|
||||
// happened to admit, and a load-balanced UI poll alternates between
|
||||
// "operation visible" and "operation gone" between replicas.
|
||||
//
|
||||
// S1. The carrier choice lives in core/application with the other three
|
||||
// caches, and this call names no carrier at all, so nothing hanging off
|
||||
// the same struct can be handed over here by accident. See
|
||||
// core/application/cache_fanout_wiring.go for why that shape is kept
|
||||
// now that the broker's client is no longer one of those things.
|
||||
if d := application.Distributed(); d != nil {
|
||||
opcache.SetMessagingClient(d.Nats)
|
||||
if d.DistStores != nil && d.DistStores.Gallery != nil {
|
||||
opcache.SetGalleryStore(d.DistStores.Gallery)
|
||||
}
|
||||
if err := opcache.Start(application.ApplicationConfig().Context); err != nil {
|
||||
if err := d.WireOpCache(application.ApplicationConfig().Context, opcache); err != nil {
|
||||
xlog.Warn("OpCache distributed subscribe failed; running standalone", "error", err)
|
||||
}
|
||||
}
|
||||
@@ -520,13 +525,15 @@ func API(application *application.Application) (*echo.Echo, error) {
|
||||
routes.RegisterAgentPoolRoutes(e, application, agentsMw, skillsMw, collectionsMw)
|
||||
// Fine-tuning routes
|
||||
fineTuningMw := auth.RequireFeature(application.AuthDB(), auth.FeatureFineTuning)
|
||||
// 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
|
||||
// In distributed mode pass the deployment's broadcast carrier + PostgreSQL
|
||||
// store so fine-tune jobs stay consistent across replicas (the SyncedMap
|
||||
// broadcasts mutations and hydrates from the DB); standalone passes nil for
|
||||
// both. The carrier comes from Broadcast() and never from a field read here:
|
||||
// see the comment on that method for why the choice is made in one place.
|
||||
var ftBus messaging.Broadcaster
|
||||
var ftStore *distributed.FineTuneStore
|
||||
if d := application.Distributed(); d != nil {
|
||||
ftNats = d.Nats
|
||||
ftBus = d.Broadcast()
|
||||
if d.DistStores != nil && d.DistStores.FineTune != nil {
|
||||
ftStore = d.DistStores.FineTune
|
||||
}
|
||||
@@ -535,20 +542,21 @@ func API(application *application.Application) (*echo.Echo, error) {
|
||||
application.ApplicationConfig(),
|
||||
application.ModelLoader(),
|
||||
application.ModelConfigLoader(),
|
||||
ftNats,
|
||||
ftBus,
|
||||
ftStore,
|
||||
)
|
||||
routes.RegisterFineTuningRoutes(e, ftService, application.ApplicationConfig(), application, fineTuningMw)
|
||||
|
||||
// Quantization routes
|
||||
quantizationMw := auth.RequireFeature(application.AuthDB(), auth.FeatureQuantization)
|
||||
// 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
|
||||
// In distributed mode pass the deployment's broadcast carrier + PostgreSQL
|
||||
// store so quantization jobs stay consistent across replicas (the SyncedMap
|
||||
// broadcasts mutations and hydrates from the DB); standalone passes nil for
|
||||
// both. Same rule and same single source as the fine-tune wiring above.
|
||||
var quantBus messaging.Broadcaster
|
||||
var quantStore *distributed.QuantStore
|
||||
if d := application.Distributed(); d != nil {
|
||||
quantNats = d.Nats
|
||||
quantBus = d.Broadcast()
|
||||
if d.DistStores != nil && d.DistStores.Quant != nil {
|
||||
quantStore = d.DistStores.Quant
|
||||
}
|
||||
@@ -557,7 +565,7 @@ func API(application *application.Application) (*echo.Echo, error) {
|
||||
application.ApplicationConfig(),
|
||||
application.ModelLoader(),
|
||||
application.ModelConfigLoader(),
|
||||
quantNats,
|
||||
quantBus,
|
||||
quantStore,
|
||||
)
|
||||
routes.RegisterQuantizationRoutes(e, qService, application.ApplicationConfig(), application, quantizationMw)
|
||||
@@ -566,15 +574,81 @@ func API(application *application.Application) (*echo.Echo, error) {
|
||||
distCfg := application.ApplicationConfig().Distributed
|
||||
var registry *nodes.NodeRegistry
|
||||
var remoteUnloader nodes.NodeCommandSender
|
||||
// How the admin log-proxy routes reach a worker's own HTTP server. Left nil
|
||||
// outside distributed mode, where there are no workers and no tunnels; the
|
||||
// routes then refuse rather than dialling an address directly.
|
||||
var workerHTTPDialFor nodes.WorkerNetDialerFor
|
||||
if d := application.Distributed(); d != nil {
|
||||
registry = d.Registry
|
||||
if d.Router != nil {
|
||||
remoteUnloader = d.Router.Unloader()
|
||||
}
|
||||
if d.WorkerDialer != nil {
|
||||
workerHTTPDialFor = func(nodeID string) func(ctx context.Context, network, addr string) (net.Conn, error) {
|
||||
return d.WorkerDialer.DialerFor(nodeID, clustersvc.StreamTagHTTP)
|
||||
}
|
||||
}
|
||||
}
|
||||
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)
|
||||
routes.RegisterNodeSelfServiceRoutes(e, registry, distCfg.RegistrationToken, distCfg.AutoApproveNodes, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret)
|
||||
routes.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, workerHTTPDialFor)
|
||||
|
||||
// Replica-to-replica peer link. Registered only in distributed mode: in
|
||||
// single-node mode there are no peers, and the route authenticates with the
|
||||
// registration token, so publishing it unconditionally would put a
|
||||
// multiplexer on every single-binary install.
|
||||
if d := application.Distributed(); d != nil && d.PeerSessions != nil {
|
||||
if distCfg.RegistrationToken == "" {
|
||||
// The handler fails closed on an empty token, which is right and
|
||||
// invisible: without this line an operator sees only 401s on a
|
||||
// route they never configured, and nothing connecting them to the
|
||||
// token they did not set.
|
||||
xlog.Warn("Replica peer link will refuse every dial: no registration token is configured",
|
||||
"route", clustersvc.PeerPath, "knob", "LOCALAI_REGISTRATION_TOKEN")
|
||||
}
|
||||
// d.Cluster is what the handler resolves a dialling replica's id
|
||||
// against, so the route can check WHICH replica is on the far end and
|
||||
// not merely that it holds the deployment's shared token.
|
||||
routes.RegisterClusterRoutes(e, distCfg.RegistrationToken, d.Cluster, d.PeerSessions.Accept)
|
||||
}
|
||||
|
||||
// The worker tunnel, registered unconditionally. Both arguments are nil
|
||||
// outside distributed mode and the handler refuses every dial then, which
|
||||
// is what makes registering it always safe; what it buys is the
|
||||
// route-coverage test walking the route in a plain single-binary
|
||||
// application, and that test is what holds the rule that an unauthenticated
|
||||
// dial is refused BEFORE the WebSocket upgrade.
|
||||
var tunnels *clustersvc.TunnelRegistry
|
||||
if d := application.Distributed(); d != nil {
|
||||
tunnels = d.Tunnels
|
||||
if distCfg.RegistrationToken == "" {
|
||||
// A different warning from the peer link's, for the same missing
|
||||
// knob, because what breaks is different. Tunnels themselves work
|
||||
// without a registration token: each node is minted its own tunnel
|
||||
// credential at registration whether or not one is configured. What
|
||||
// is missing is the gate in FRONT of that. With no registration
|
||||
// token, RegisterNodeEndpoint validates nothing, so anyone who can
|
||||
// reach this frontend can register a node and be issued a tunnel
|
||||
// credential for it.
|
||||
//
|
||||
// How far that gets them depends on the OTHER knob. With
|
||||
// auto-approve on, the node is healthy at once and the credential
|
||||
// works immediately. With it off, the node is pending, and the
|
||||
// tunnel route refuses a pending node on every dial, so the
|
||||
// credential is inert until an admin approves it and approval is
|
||||
// the real gate. Worth stating precisely, because the same commit
|
||||
// argues exactly this distinction three files away to justify
|
||||
// minting for pending nodes at all.
|
||||
//
|
||||
// This warning replaced one that said the opposite, that tunnels
|
||||
// would refuse every dial without this token. That was true while
|
||||
// the tunnel authenticated against the registration token's own
|
||||
// hash, and stopped being true when nodes got credentials of their
|
||||
// own.
|
||||
xlog.Warn("Node registration is unauthenticated, so any caller that can reach this frontend can register a worker and be issued a tunnel credential",
|
||||
"route", clustersvc.ConnectPath, "knob", "LOCALAI_REGISTRATION_TOKEN")
|
||||
}
|
||||
}
|
||||
routes.RegisterWorkerTunnelRoute(e, registry, tunnels)
|
||||
|
||||
// Distributed SSE routes (job progress + agent events via NATS)
|
||||
if d := application.Distributed(); d != nil {
|
||||
|
||||
@@ -12,7 +12,6 @@ import (
|
||||
"github.com/mudler/LocalAI/core/http/auth"
|
||||
"github.com/mudler/LocalAI/core/http/routes"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/pkg/natsauth"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
@@ -114,7 +113,6 @@ func newNodeSelfServiceTestApp(db *gorm.DB, appConfig *config.ApplicationConfig,
|
||||
false,
|
||||
nil,
|
||||
"",
|
||||
natsauth.Config{},
|
||||
)
|
||||
return e
|
||||
}
|
||||
|
||||
@@ -74,8 +74,26 @@ func isPublicRoute(method, path string) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// ClusterPathPrefix is the machine-to-machine cluster namespace. It carries two
|
||||
// different trust relationships, on two different credentials: the
|
||||
// replica-to-replica peer link, which checks the shared cluster token, and the
|
||||
// worker-to-frontend tunnel, which checks the dialing node's own stored token
|
||||
// hash. What they have in common is the only thing this prefix asserts, that
|
||||
// each handler checks its own Authorization header, so the check below lets them
|
||||
// through the global session middleware rather than rejecting a caller that has
|
||||
// no session and no user.
|
||||
//
|
||||
// The cluster routes do NOT derive their paths from this constant: they are
|
||||
// registered from core/services/cluster's own literal, because that package
|
||||
// must not import core/http/auth. Nothing in the compiler holds the two
|
||||
// together, so a spec does instead, driving a peer request through this
|
||||
// middleware in core/http/endpoints/cluster/peer_test.go. Moving either string
|
||||
// without the other turns that spec red, which is the whole reason it exists.
|
||||
const ClusterPathPrefix = "/api/cluster/"
|
||||
|
||||
// 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/") ||
|
||||
strings.HasPrefix(path, ClusterPathPrefix)
|
||||
}
|
||||
@@ -29,7 +29,7 @@ import (
|
||||
// @Param request body schema.AnthropicRequest true "query params"
|
||||
// @Success 200 {object} schema.AnthropicResponse "Response"
|
||||
// @Router /v1/messages [post]
|
||||
func MessagesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, natsClient mcpTools.MCPNATSClient) echo.HandlerFunc {
|
||||
func MessagesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, agentControl mcpTools.AgentControl) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
id := uuid.New().String()
|
||||
|
||||
@@ -70,7 +70,7 @@ func MessagesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evalu
|
||||
if (len(mcpServers) > 0 || mcpPromptName != "" || len(mcpResourceURIs) > 0) && (cfg.MCP.Servers != "" || cfg.MCP.Stdio != "") {
|
||||
remote, stdio, mcpErr := cfg.MCP.MCPConfigFromYAML()
|
||||
if mcpErr == nil {
|
||||
mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), natsClient, cfg.Name, remote, stdio, mcpServers)
|
||||
mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), agentControl, cfg.Name, remote, stdio, mcpServers)
|
||||
|
||||
// Prompt and resource injection (pre-processing step — resolves locally regardless of distributed mode)
|
||||
namedSessions, sessErr := mcpTools.NamedSessionsFromMCPConfig(cfg.Name, remote, stdio, mcpServers)
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
package cluster_test
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
func TestClusterEndpoints(t *testing.T) {
|
||||
RegisterFailHandler(Fail)
|
||||
RunSpecs(t, "Cluster Endpoints Suite")
|
||||
}
|
||||
@@ -0,0 +1,263 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package cluster
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"crypto/subtle"
|
||||
"encoding/hex"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/libp2p/go-yamux/v5"
|
||||
clustersvc "github.com/mudler/LocalAI/core/services/cluster"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/xlog"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// ConnectHandler serves the door a worker knocks on: it authenticates the dial
|
||||
// against the node's OWN stored token, upgrades it to a WebSocket, wraps that as
|
||||
// a yamux server session and attaches it to the tunnel registry.
|
||||
//
|
||||
// The worker dials out and never listens, which is the whole point of the
|
||||
// tunnel: a worker behind NAT, in another cluster, or on a laptop needs no
|
||||
// inbound port. It is therefore the yamux CLIENT and this side the SERVER, so
|
||||
// this side owns the even stream IDs and is the side that opens streams.
|
||||
// Nothing here accepts streams: in this design the frontend asks and the worker
|
||||
// answers, so a worker that opened a stream into this session would park on the
|
||||
// accept backlog rather than be served.
|
||||
//
|
||||
// The route is registered in every deployment, including single-binary ones, so
|
||||
// that the route-coverage test under build tag `auth` walks it. A nil registry
|
||||
// or a nil tunnel registry therefore has to be a real answer rather than a
|
||||
// panic; see the 503 below.
|
||||
//
|
||||
// It is deliberately absent from auth.RouteFeatureRegistry. That registry gates
|
||||
// a route on the FEATURES OF AN AUTHENTICATED USER, resolved from auth.GetUser,
|
||||
// and there is no user here: the dialer is a worker process holding a machine
|
||||
// credential.
|
||||
//
|
||||
// The global auth middleware does RUN on this path; what it does not do is
|
||||
// reject. It attempts session, bearer and legacy-key authentication first, so it
|
||||
// may even have set auth_user from a worker token that happens to match an API
|
||||
// key, and then core/http/auth/middleware.go:90 lets the request through
|
||||
// because usesAlternativeAuthentication reports the path as one whose
|
||||
// credentials its own route checks. Nothing here reads what it set.
|
||||
func ConnectHandler(registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegistry) echo.HandlerFunc {
|
||||
// gorilla's default CheckOrigin restricts a browser to same-origin and lets
|
||||
// a header-less client (which every worker is) through, so the zero value
|
||||
// is what this link wants. The same choice PeerHandler makes.
|
||||
upgrader := websocket.Upgrader{}
|
||||
|
||||
return func(c echo.Context) error {
|
||||
// Everything below happens BEFORE the upgrade, and the order inside it
|
||||
// is load-bearing. The credential is read first because a dial with no
|
||||
// Authorization header at all is the anonymous case, and the
|
||||
// route-coverage test issues exactly that, with no query string: it
|
||||
// must see 401 rather than a 400 about a missing node id.
|
||||
token, ok := bearerToken(c.Request())
|
||||
if !ok {
|
||||
return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized")
|
||||
}
|
||||
|
||||
// Not 401. A frontend with no cluster cannot authenticate anybody, and
|
||||
// answering "unauthorized" would send the operator hunting a token
|
||||
// problem that does not exist. It is checked after the header so that
|
||||
// an anonymous dial still gets the 401 the coverage test requires.
|
||||
//
|
||||
// Only the registry half is covered by a spec. The two are read from one
|
||||
// application.Distributed() in core/http/app.go and initDistributed
|
||||
// returns an error rather than a partial struct, so a non-nil registry
|
||||
// beside a nil tunnel registry is unreachable and no spec constructs it;
|
||||
// the second half is defence against a future wiring that splits them,
|
||||
// where the cost would be a nil dereference in Attach after the
|
||||
// connection is already hijacked.
|
||||
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):
|
||||
// A node this frontend has never seen. Reported as 401 rather than
|
||||
// 404 so a caller cannot enumerate node IDs by status code.
|
||||
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. This
|
||||
// is the phase's standing rule in its HTTP form: telling a worker
|
||||
// its credentials are wrong when the database merely could not be
|
||||
// read sends it re-registering instead of retrying, and a worker
|
||||
// that re-registers has thrown away the identity its tunnel and its
|
||||
// loaded models are 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")
|
||||
}
|
||||
|
||||
// Split from the mismatch below because they are different operator
|
||||
// problems with different fixes. An empty stored hash means this node
|
||||
// last registered against a LocalAI that predates per-node tunnel
|
||||
// credentials, so it holds no secret this route can check and must
|
||||
// register again; a mismatch means the worker is presenting the wrong
|
||||
// one, usually a stale credential from before a rotation. One log line
|
||||
// for both leaves an operator reading "wrong token" while a whole fleet
|
||||
// of not-yet-restarted workers fails identically.
|
||||
if node.TunnelTokenHash == "" {
|
||||
// Debug, not Warn. Every worker in that state fails this way on
|
||||
// every reconnect, so warning per dial buries the log.
|
||||
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")
|
||||
}
|
||||
|
||||
// Authenticated but not authorised, so 403 rather than 401: the fix is an
|
||||
// admin approving the node, not a different credential, and answering
|
||||
// 401 would send an operator looking at tokens.
|
||||
//
|
||||
// Only StatusPending is refused. The rest of /api/node/ self-service
|
||||
// gates on nothing at all, but the one place that hands a node
|
||||
// something DURABLE refuses a pending one too: the agent worker's API
|
||||
// key (provisionAgentWorkerKey, guarded at its call site in
|
||||
// core/http/endpoints/localai/nodes.go). Cited by NAME, not by line:
|
||||
// an earlier version of this comment cited line numbers into a file the
|
||||
// same commit was editing, and both were stale before it landed. The
|
||||
// per-node broker credential this used to name alongside it is gone
|
||||
// with the bus, and so is the function that minted it.
|
||||
//
|
||||
// A tunnel is that kind of grant, not a heartbeat: it is
|
||||
// a standing pipe into the worker recorded in node_connections and
|
||||
// relayed to by every other replica. Draining and unhealthy nodes keep
|
||||
// their tunnels on purpose; draining means finish what you have, and a
|
||||
// node marked unhealthy for missed heartbeats needs the pipe to recover
|
||||
// through.
|
||||
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")
|
||||
}
|
||||
|
||||
ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil)
|
||||
if err != nil {
|
||||
// Upgrade has already written its own failure to the client.
|
||||
xlog.Debug("worker tunnel upgrade failed", "node", nodeID, "error", err)
|
||||
return nil
|
||||
}
|
||||
|
||||
sess, err := yamux.Server(clustersvc.WebsocketConn(ws), nil, nil)
|
||||
if err != nil {
|
||||
xlog.Error("Worker tunnel session setup failed", "node", nodeID, "error", err)
|
||||
_ = ws.Close()
|
||||
return nil
|
||||
}
|
||||
|
||||
// The same guard PeerHandler carries, for the same reason and one more.
|
||||
// net/http recovers a panic from this goroutine but does not close the
|
||||
// hijacked connection, and middleware.Recover does not either, so a
|
||||
// panic below would leave the worker holding a live session this replica
|
||||
// has no entry for and will never detach. The extra reason here is that
|
||||
// Attach does database work: a panic inside it, with the session left
|
||||
// open, is a tunnel nothing can reach and nothing will clean up.
|
||||
//
|
||||
// It re-panics rather than swallowing. Whatever it caught is a bug, and
|
||||
// the recovery middleware above is what should report it.
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
_ = sess.Close()
|
||||
panic(r)
|
||||
}
|
||||
}()
|
||||
|
||||
// From here the connection is hijacked, so no status can reach the
|
||||
// worker any more: a failure is a closed socket, which is what its
|
||||
// reconnect loop reads.
|
||||
epoch, err := tunnels.Attach(c.Request().Context(), nodeID, sess)
|
||||
if err != nil {
|
||||
xlog.Error("Attaching a worker tunnel failed", "node", nodeID, "error", err)
|
||||
_ = sess.Close()
|
||||
return nil
|
||||
}
|
||||
|
||||
xlog.Info("Worker tunnel established", "node", nodeID, "remote", ws.RemoteAddr().String())
|
||||
// The session outlives this handler, so something other than the
|
||||
// request goroutine has to notice it die. yamux closes shutdownCh from
|
||||
// its receive loop the moment the underlying conn fails
|
||||
// (go-yamux/v5@v5.1.0/session.go:691-695 calling close at
|
||||
// session.go:297-311), and its default config keepalives every 30s
|
||||
// (mux.go:73-74), so a worker that vanishes without a FIN is noticed
|
||||
// too rather than held forever.
|
||||
go func() {
|
||||
<-sess.CloseChan()
|
||||
// The token Attach returned, never a fresh or zero one. Detach
|
||||
// matches it by EQUALITY: it identifies THIS attachment, so a
|
||||
// worker that has already re-dialled onto this replica is not
|
||||
// evicted by its predecessor's teardown.
|
||||
tunnels.Detach(nodeID, epoch)
|
||||
xlog.Debug("worker tunnel closed", "node", nodeID)
|
||||
}()
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// bearerToken returns the token from an Authorization: Bearer header, and
|
||||
// whether one was present at all.
|
||||
//
|
||||
// The presence of a credential and its correctness are separate answers on
|
||||
// purpose: "no credential" is what decides the pre-upgrade 401, and it has to be
|
||||
// decidable before anything about the node is known.
|
||||
func bearerToken(r *http.Request) (string, bool) {
|
||||
// RFC 7235 makes the scheme case-insensitive; the token after it is not.
|
||||
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 presented token against the hash stored on the
|
||||
// node's own row, in constant time.
|
||||
//
|
||||
// Against the NODE's OWN tunnel credential, not the deployment's registration
|
||||
// token. A tunnel is a durable, multiplexed pipe into a worker, and a
|
||||
// credential that authorizes every worker at once would mean one leak lets an
|
||||
// attacker impersonate any worker whose ID it can read and take over that
|
||||
// worker's traffic by claiming its tunnel. The secret compared here is minted
|
||||
// per node at registration (attachTunnelToken in
|
||||
// core/http/endpoints/localai/nodes.go), returned to that worker once, and
|
||||
// stored only as this hash, so knowing the registration token no longer gets
|
||||
// anyone a tunnel.
|
||||
//
|
||||
// Note which column: BackendNode.TunnelTokenHash, not TokenHash. TokenHash is
|
||||
// still the hash of whatever token the worker registered WITH, which on most
|
||||
// deployments is the shared registration token, and comparing against it is
|
||||
// exactly the weakness this replaced.
|
||||
//
|
||||
// The empty-hash guard is defensive rather than deciding: a stored hash is
|
||||
// hex-encoded SHA-256, so 64 bytes or nothing, and ConstantTimeCompare already
|
||||
// returns 0 on a length mismatch (crypto/internal/fips140/subtle/constant_time.go:17-20
|
||||
// returns 0 outright when the lengths differ). It is kept because a reader should not have to
|
||||
// derive "a node with no credential authorizes nobody" from a length rule, and
|
||||
// because the caller logs that case separately.
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,391 @@
|
||||
package cluster_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/http/auth"
|
||||
"github.com/mudler/LocalAI/core/http/routes"
|
||||
clustersvc "github.com/mudler/LocalAI/core/services/cluster"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/libp2p/go-yamux/v5"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// workerToken is the tunnel credential one worker holds. Registration mints it
|
||||
// per node and keeps only its hash, which is the whole point of the check the
|
||||
// specs below pin: a second worker's credential, and the deployment-wide
|
||||
// registration token, are both wrong for this node.
|
||||
const workerToken = "worker-1-secret"
|
||||
|
||||
// registrationToken stands in for the shared secret every worker in a
|
||||
// deployment registers with. It is stored on the row too, in a DIFFERENT
|
||||
// column, and a spec below pins that presenting it does not open a tunnel.
|
||||
const registrationToken = "deployment-registration-token"
|
||||
|
||||
func tokenHash(token string) string {
|
||||
sum := sha256.Sum256([]byte(token))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
// bearer builds the header a worker dials with.
|
||||
func bearer(token string) http.Header {
|
||||
h := http.Header{}
|
||||
h.Set("Authorization", "Bearer "+token)
|
||||
return h
|
||||
}
|
||||
|
||||
// wsConnectURL is the worker tunnel route on a test server, named as nodeID.
|
||||
func wsConnectURL(s *httptest.Server, nodeID string) string {
|
||||
return "ws" + strings.TrimPrefix(s.URL, "http") + clustersvc.ConnectPath +
|
||||
"?id=" + url.QueryEscape(nodeID)
|
||||
}
|
||||
|
||||
var _ = Describe("Worker tunnel handler", func() {
|
||||
var (
|
||||
srv *httptest.Server
|
||||
db *gorm.DB
|
||||
reg *clustersvc.Registry
|
||||
tun *clustersvc.TunnelRegistry
|
||||
nodeID string
|
||||
ctx context.Context
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
ctx = context.Background()
|
||||
db = testutil.SetupTestDB()
|
||||
|
||||
nodeReg, err := nodes.NewNodeRegistry(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
node := &nodes.BackendNode{
|
||||
Name: "worker-1",
|
||||
Address: "10.0.0.9:50051",
|
||||
// Both hashes are set, and they differ. That is what a real
|
||||
// registration produces: TokenHash is the shared token the worker
|
||||
// registered WITH, TunnelTokenHash is the secret minted FOR it.
|
||||
TokenHash: tokenHash(registrationToken),
|
||||
TunnelTokenHash: tokenHash(workerToken),
|
||||
}
|
||||
Expect(nodeReg.Register(ctx, node, true)).To(Succeed())
|
||||
nodeID = node.ID
|
||||
Expect(nodeID).ToNot(BeEmpty())
|
||||
|
||||
reg = clustersvc.NewRegistry(db)
|
||||
Expect(reg.Register(ctx, "me", "10.0.0.1:8080", "v1", "")).To(Succeed())
|
||||
tun = clustersvc.NewTunnelRegistry(reg, "me")
|
||||
|
||||
e := echo.New()
|
||||
routes.RegisterWorkerTunnelRoute(e, nodeReg, tun)
|
||||
srv = httptest.NewServer(e)
|
||||
DeferCleanup(srv.Close)
|
||||
})
|
||||
|
||||
It("refuses an anonymous dial before upgrading", func() {
|
||||
// A plain GET, not a WebSocket dial: this is exactly what the
|
||||
// route-coverage test issues, and a handler that upgrades first answers
|
||||
// it with gorilla's own 400 handshake failure instead of a 401.
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL+clustersvc.ConnectPath+"?id="+nodeID, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized))
|
||||
Expect(resp.Header.Get("Upgrade")).To(BeEmpty(),
|
||||
"the handler upgraded an unauthenticated dial")
|
||||
Expect(tun.Held()).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("refuses a dial that carries no credentials at all", func() {
|
||||
_, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), nil)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized))
|
||||
Expect(tun.Held()).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("refuses the deployment's registration token, which this node also stores", func() {
|
||||
// The registration token is the credential every worker in the
|
||||
// deployment holds, and it IS on this node's row, in TokenHash.
|
||||
// Accepting it would mean one leaked shared secret impersonates any
|
||||
// worker whose ID an attacker can read; the node's own tunnel
|
||||
// credential is the only thing this route accepts.
|
||||
_, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(registrationToken))
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized))
|
||||
Expect(tun.Held()).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("refuses a dial that names a node it has never seen", func() {
|
||||
_, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, "no-such-node"), bearer(workerToken))
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized))
|
||||
Expect(tun.Held()).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("refuses an authenticated dial that names no node", func() {
|
||||
_, resp, err := websocket.DefaultDialer.Dial(
|
||||
"ws"+strings.TrimPrefix(srv.URL, "http")+clustersvc.ConnectPath, bearer(workerToken))
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusBadRequest))
|
||||
})
|
||||
|
||||
It("attaches an authenticated worker and carries bytes to it", func() {
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = conn.Close() })
|
||||
|
||||
// The worker is the side that dials, so its half of the mux is the
|
||||
// yamux CLIENT and the frontend's is the server.
|
||||
workerSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = workerSess.Close() })
|
||||
|
||||
Eventually(tun.Held, "10s").Should(ConsistOf(nodeID))
|
||||
|
||||
owner, _, err := reg.OwnerRow(ctx, nodeID)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(owner).To(Equal("me"),
|
||||
"the tunnel was stored without the claim that tells other replicas where it is")
|
||||
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
stream, aerr := workerSess.AcceptStream()
|
||||
if aerr != nil {
|
||||
return
|
||||
}
|
||||
defer func() { _ = stream.Close() }()
|
||||
buf := make([]byte, 4)
|
||||
if _, rerr := stream.Read(buf); rerr != nil {
|
||||
return
|
||||
}
|
||||
_, _ = stream.Write(buf)
|
||||
}()
|
||||
|
||||
stream, err := tun.Open(ctx, nodeID)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = stream.Close() })
|
||||
_, err = stream.Write([]byte("ping"))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
echoed := make([]byte, 4)
|
||||
_, err = stream.Read(echoed)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(string(echoed)).To(Equal("ping"))
|
||||
})
|
||||
|
||||
It("detaches the tunnel and drops its claim when the worker goes away", func() {
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Eventually(tun.Held, "10s").Should(ConsistOf(nodeID))
|
||||
|
||||
Expect(conn.Close()).To(Succeed())
|
||||
|
||||
Eventually(tun.Held, "10s").Should(BeEmpty(),
|
||||
"a dead tunnel is still held here, so every dialer routed to this replica gets a socket that carries nothing")
|
||||
Eventually(func() error {
|
||||
_, _, err := reg.OwnerRow(ctx, nodeID)
|
||||
return err
|
||||
}, "10s").Should(MatchError(clustersvc.ErrNoConnection),
|
||||
"the claim outlived the socket, so this replica keeps being named the owner of a worker it no longer holds")
|
||||
})
|
||||
|
||||
It("refuses a node with no tunnel credential, without falling back to its registration token", func() {
|
||||
// A node registered by a LocalAI predating per-node tunnel credentials
|
||||
// produces exactly this row: a registration-token hash in token_hash
|
||||
// and nothing in tunnel_token_hash. It cannot be back-filled, because
|
||||
// the plaintext only ever existed in the response that minted it, so
|
||||
// such a node must register again.
|
||||
//
|
||||
// The dial presents the REGISTRATION token, which is still on the row.
|
||||
// A handler that fell back to token_hash when the tunnel hash is empty
|
||||
// would let it in, which is the weakness this whole change removed.
|
||||
Expect(db.Exec(`UPDATE backend_nodes SET tunnel_token_hash = '' WHERE id = ?`, nodeID).Error).To(Succeed())
|
||||
|
||||
_, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(registrationToken))
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized))
|
||||
Expect(tun.Held()).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("refuses a node that is still awaiting admin approval", func() {
|
||||
// Approval is what gates a node's participation, and a tunnel is a
|
||||
// standing pipe recorded in node_connections, not a heartbeat. 403 and
|
||||
// not 401: the credential is right, the authorisation is missing, and
|
||||
// the fix is an admin rather than a different token.
|
||||
Expect(db.Exec(`UPDATE backend_nodes SET status = ? WHERE id = ?`, nodes.StatusPending, nodeID).Error).To(Succeed())
|
||||
|
||||
_, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken))
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusForbidden))
|
||||
Expect(tun.Held()).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("still admits a draining node, which has work to finish", func() {
|
||||
// Only pending is refused. Draining means "start nothing new", not
|
||||
// "lose the pipe your in-flight requests travel on", and a node marked
|
||||
// unhealthy for missed heartbeats needs the tunnel to recover through.
|
||||
Expect(db.Exec(`UPDATE backend_nodes SET status = ? WHERE id = ?`, nodes.StatusDraining, nodeID).Error).To(Succeed())
|
||||
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = conn.Close() })
|
||||
Eventually(tun.Held, "10s").Should(ConsistOf(nodeID))
|
||||
})
|
||||
|
||||
It("reports a lookup failure as a failure, not as a refusal", func() {
|
||||
// ErrNotOwner, 401 and 404 are all ANSWERS. A database that cannot be
|
||||
// read is none of them: telling a worker its credentials are wrong when
|
||||
// the frontend simply could not look them up sends it re-registering
|
||||
// instead of retrying.
|
||||
Expect(db.Exec(`DROP TABLE backend_nodes CASCADE`).Error).To(Succeed())
|
||||
|
||||
_, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken))
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusInternalServerError))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("Worker tunnel handler when the attach panics", func() {
|
||||
// net/http recovers a panic from the request goroutine but does NOT close a
|
||||
// hijacked connection, so without the handler's own recover the worker keeps
|
||||
// a live session this replica has no entry for and will never detach: its
|
||||
// opens fill yamux's 256-deep backlog and then hang with no error. The panic
|
||||
// is injected through a real path rather than a fake one, a registry built
|
||||
// over no database at all, which is what Attach's first database call
|
||||
// dereferences.
|
||||
var (
|
||||
srv *httptest.Server
|
||||
db *gorm.DB
|
||||
nodeID string
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
ctx := context.Background()
|
||||
db = testutil.SetupTestDB()
|
||||
nodeReg, err := nodes.NewNodeRegistry(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
node := &nodes.BackendNode{Name: "worker-1", Address: "10.0.0.9:50051", TunnelTokenHash: tokenHash(workerToken)}
|
||||
Expect(nodeReg.Register(ctx, node, true)).To(Succeed())
|
||||
nodeID = node.ID
|
||||
|
||||
e := echo.New()
|
||||
routes.RegisterWorkerTunnelRoute(e, nodeReg, clustersvc.NewTunnelRegistry(nil, "me"))
|
||||
srv = httptest.NewServer(e)
|
||||
DeferCleanup(func() {
|
||||
// A hijacked connection the handler never closed would park Close
|
||||
// forever, turning the assertion below into a suite hang.
|
||||
srv.CloseClientConnections()
|
||||
srv.Close()
|
||||
})
|
||||
})
|
||||
|
||||
It("closes the worker's session instead of stranding it", func() {
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, nodeID), bearer(workerToken))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = conn.Close() })
|
||||
|
||||
workerSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = workerSess.Close() })
|
||||
|
||||
// Asserting on OpenStream would hang rather than fail: yamux only
|
||||
// acknowledges a stream once the peer accepts it, and the leak this
|
||||
// pins is precisely that nobody ever will.
|
||||
Eventually(workerSess.IsClosed, "10s").Should(BeTrue())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("Worker tunnel handler without distributed mode", func() {
|
||||
// The route is registered in every deployment so that the route-coverage
|
||||
// test sees it, which is what pins the reject-before-upgrade rule. With no
|
||||
// node registry there is nothing to authenticate against, so it must refuse
|
||||
// every dial rather than publish an unauthenticated multiplexer.
|
||||
var srv *httptest.Server
|
||||
|
||||
BeforeEach(func() {
|
||||
e := echo.New()
|
||||
routes.RegisterWorkerTunnelRoute(e, nil, nil)
|
||||
srv = httptest.NewServer(e)
|
||||
DeferCleanup(srv.Close)
|
||||
})
|
||||
|
||||
It("refuses an anonymous dial with 401", func() {
|
||||
req, err := http.NewRequestWithContext(GinkgoT().Context(), http.MethodGet, srv.URL+clustersvc.ConnectPath, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized))
|
||||
})
|
||||
|
||||
It("tells a credentialed worker the frontend has no cluster, rather than rejecting it", func() {
|
||||
// A single-binary frontend cannot authenticate anybody, and saying
|
||||
// "unauthorized" would send an operator hunting a token problem that
|
||||
// does not exist.
|
||||
_, resp, err := websocket.DefaultDialer.Dial(wsConnectURL(srv, "w1"), bearer(workerToken))
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable))
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("Worker tunnel auth coverage", func() {
|
||||
// The same argument the peer link's coverage specs make: the tunnel route
|
||||
// authenticates a worker against its own stored token, not a session, so it
|
||||
// only works while its path sits under the prefix the global auth
|
||||
// middleware exempts.
|
||||
var srv *httptest.Server
|
||||
|
||||
BeforeEach(func() {
|
||||
e := echo.New()
|
||||
// A nil DB with one legacy API key is the cheapest configuration that
|
||||
// turns the middleware ON without a database.
|
||||
e.Use(auth.Middleware(nil, &config.ApplicationConfig{ApiKeys: []string{"an-api-key"}}))
|
||||
routes.RegisterWorkerTunnelRoute(e, nil, nil)
|
||||
e.GET("/api/nodes", func(c echo.Context) error { return c.NoContent(http.StatusOK) })
|
||||
srv = httptest.NewServer(e)
|
||||
DeferCleanup(srv.Close)
|
||||
})
|
||||
|
||||
It("refuses an uncredentialed request to a route outside the cluster prefix", func() {
|
||||
resp, err := http.Get(srv.URL + "/api/nodes")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized),
|
||||
"the global auth middleware is not actually guarding this server, so the assertion below would prove nothing")
|
||||
})
|
||||
|
||||
It("lets a worker dial reach the handler, which is the only thing that can authenticate it", func() {
|
||||
// The worker's token is not one of the API keys the middleware knows,
|
||||
// so a 503 from the handler's own no-cluster check can only mean the
|
||||
// request was let through by the middleware.
|
||||
req, err := http.NewRequestWithContext(GinkgoT().Context(), http.MethodGet, srv.URL+clustersvc.ConnectPath+"?id=w1", nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
req.Header.Set("Authorization", "Bearer "+workerToken)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusServiceUnavailable),
|
||||
"a worker dial must reach the handler; 401 here means the tunnel route left the auth-exempt prefix %q", auth.ClusterPathPrefix)
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,209 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
// Package cluster serves the replica-to-replica link that a LocalAI frontend
|
||||
// uses to reach a worker tunnel it does not own. A peer dials
|
||||
// GET /api/cluster/peer, the connection becomes one multiplexed yamux session,
|
||||
// and the relay opens a stream on it per request.
|
||||
package cluster
|
||||
|
||||
import (
|
||||
"crypto/subtle"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/libp2p/go-yamux/v5"
|
||||
clustersvc "github.com/mudler/LocalAI/core/services/cluster"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// PeerHandler upgrades an authenticated peer dial to a WebSocket, wraps it as
|
||||
// a yamux server session and hands it to onSession.
|
||||
//
|
||||
// TWO credentials are checked, and both must pass. token is the deployment's
|
||||
// shared cluster token and says the dialler belongs here at all; instances
|
||||
// resolves the replica id in ?id= to its row and checks the dialler's OWN peer
|
||||
// credential against the hash that row publishes. Neither replaces the other:
|
||||
// the shared check is unchanged, and identity is added in front of the
|
||||
// multiplexer it guards.
|
||||
//
|
||||
// onSession runs on the request goroutine, so it must return promptly; the
|
||||
// session outlives the handler because the upgrade hijacks the connection, and
|
||||
// closing it is the caller's job.
|
||||
func PeerHandler(token string, instances *clustersvc.Registry, onSession func(peerID string, sess *yamux.Session)) echo.HandlerFunc {
|
||||
// gorilla's default CheckOrigin already restricts a browser to same-origin
|
||||
// and lets a header-less client (which every peer is) through, so the
|
||||
// zero value is what this link wants.
|
||||
upgrader := websocket.Upgrader{}
|
||||
|
||||
return func(c echo.Context) error {
|
||||
// Reject before upgrading. Upgrading and then closing would give the
|
||||
// dialer a WebSocket error in place of an HTTP status, and both the
|
||||
// route-coverage test and a peer's own retry logic read the status.
|
||||
if !authorizedPeer(c.Request(), token) {
|
||||
return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized")
|
||||
}
|
||||
|
||||
// Not 401. A frontend with no cluster registry cannot resolve any
|
||||
// replica's identity, and answering "unauthorized" would send an operator
|
||||
// hunting a token problem that does not exist. Checked after the shared
|
||||
// token so an anonymous dial still gets the 401 first, which is the same
|
||||
// ordering ConnectHandler takes next door and for the same reason.
|
||||
//
|
||||
// This route is registered only in distributed mode, where the registry is
|
||||
// never nil, so the branch is defence against a future wiring that
|
||||
// registers it more widely. Failing closed is the point: a handler that
|
||||
// treated a nil registry as "nothing to check" would publish exactly the
|
||||
// unauthenticated multiplexer this argument was added to prevent.
|
||||
if instances == nil {
|
||||
return echo.NewHTTPError(http.StatusServiceUnavailable, "distributed mode not enabled")
|
||||
}
|
||||
|
||||
// Read before the credential so a dial that names nobody is still a 400.
|
||||
// The id is no longer taken on trust: everything below turns it from a
|
||||
// label into a claim, by resolving it to a row and requiring the dialler to
|
||||
// present that row's secret.
|
||||
peerID := c.QueryParam("id")
|
||||
if peerID == "" {
|
||||
return echo.NewHTTPError(http.StatusBadRequest, "missing peer id")
|
||||
}
|
||||
|
||||
// The identity half, read before any database work so a dial carrying no
|
||||
// credential costs no query.
|
||||
//
|
||||
// A missing header is refused, not waved through, and the choice is the
|
||||
// whole migration story of this route. Accepting a credential-less dial
|
||||
// "for compatibility" would leave the hole exactly as open as it was,
|
||||
// because an attacker simply omits the header too; there is no version of
|
||||
// a downgrade here that is safe, only versions that are quiet. So the
|
||||
// failure is loud instead: this is the one line that tells an operator
|
||||
// mid-rollout why an old replica cannot reach a new one, and it names the
|
||||
// upgrade rather than the network.
|
||||
presented := c.Request().Header.Get(clustersvc.PeerIdentityHeader)
|
||||
if presented == "" {
|
||||
xlog.Warn("Refusing a peer link: the dialling replica presented no peer credential. It is running a release from before per-replica peer identity, or it never registered a credential of its own. Upgrade it; this replica will not accept an unproven peer id",
|
||||
"peer", peerID, "header", clustersvc.PeerIdentityHeader)
|
||||
return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized")
|
||||
}
|
||||
|
||||
inst, err := instances.Get(c.Request().Context(), peerID)
|
||||
switch {
|
||||
case errors.Is(err, clustersvc.ErrInstanceNotFound):
|
||||
// A replica id this deployment has no row for. Reported as 401 rather
|
||||
// than 404 so a caller cannot enumerate replica ids by status code,
|
||||
// which is the same choice the worker tunnel makes for node ids.
|
||||
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, and this
|
||||
// is that rule in its HTTP form. Answering 401 on an unreadable
|
||||
// database would tell a healthy replica its credentials are wrong, and
|
||||
// nothing above the transport may conclude anything about a WORKER
|
||||
// from it either: a peer that cannot be linked to is not a worker that
|
||||
// has gone away.
|
||||
xlog.Error("Looking up a replica for its peer dial failed", "peer", peerID, "error", err)
|
||||
return echo.NewHTTPError(http.StatusInternalServerError, "peer lookup failed")
|
||||
}
|
||||
|
||||
// Split from the mismatch below because they are different operator
|
||||
// problems with different fixes, exactly as the worker tunnel splits them.
|
||||
// An empty stored hash means that replica last registered against a
|
||||
// LocalAI that predates peer credentials, so there is no secret to check
|
||||
// and it must register again, which a restart does. A mismatch means the
|
||||
// dialler is presenting the wrong one.
|
||||
//
|
||||
// Empty is refused. It is not "no restriction": an empty credential that
|
||||
// matched would mean every replica registered by an older frontend could
|
||||
// be impersonated by anyone holding the shared token, which is the whole
|
||||
// exposure being closed.
|
||||
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. That replica has not registered since per-replica peer identity was introduced",
|
||||
"peer", peerID)
|
||||
return echo.NewHTTPError(http.StatusUnauthorized, "unauthorized")
|
||||
}
|
||||
if !clustersvc.PeerTokenMatches(presented, inst.PeerTokenHash) {
|
||||
// This is the impostor case, and it is a WARN rather than a debug
|
||||
// line: holding the shared token and declaring somebody else's id is
|
||||
// precisely the attack this check exists for, and it should not be
|
||||
// invisible at default log level.
|
||||
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 already written its own failure to the client.
|
||||
xlog.Debug("cluster peer link upgrade failed", "peer", peerID, "error", err)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Server side of the mux: the dialing peer is the client, so it owns
|
||||
// the odd stream IDs and this side the even ones.
|
||||
//
|
||||
// The SAME configuration the dialler uses, and that is load bearing
|
||||
// rather than symmetry for its own sake. A yamux receive window is
|
||||
// advertised by the receiving side, so a nil here left this end on the
|
||||
// 256 KiB default while the dialler ran at 4 MiB, and the direction
|
||||
// governed by this end is the one that carries a relayed model artifact
|
||||
// INTO the replica that owns the worker's tunnel. That direction was
|
||||
// measured at roughly half the throughput of the same transfer without
|
||||
// a relay in it.
|
||||
//
|
||||
// It also puts the same ceiling on unread data at this end that
|
||||
// PeerLinkConfig already documents for the dialling end, so a replica
|
||||
// is now sized against that figure per link in BOTH directions. That
|
||||
// is the cost of the window being useful at all: a window is a bound
|
||||
// on data received and not yet read, so a receiver that will not
|
||||
// buffer cannot advertise one.
|
||||
sess, err := yamux.Server(clustersvc.WebsocketConn(ws), clustersvc.PeerLinkConfig(), nil)
|
||||
if err != nil {
|
||||
xlog.Error("cluster peer link session setup failed", "peer", peerID, "error", err)
|
||||
_ = ws.Close()
|
||||
return nil
|
||||
}
|
||||
|
||||
if onSession == nil {
|
||||
// Nothing will ever read from this session, so do not leave the
|
||||
// peer believing it has a live link.
|
||||
_ = sess.Close()
|
||||
return nil
|
||||
}
|
||||
|
||||
xlog.Debug("cluster peer link established", "peer", peerID, "remote", ws.RemoteAddr().String())
|
||||
// net/http recovers a panic from this goroutine but does not close a
|
||||
// hijacked connection afterwards, so a panicking callback would leave
|
||||
// the peer holding a link nobody accepts streams on: its opens would
|
||||
// fill the 256-deep backlog and then hang without an error.
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
_ = sess.Close()
|
||||
panic(r)
|
||||
}
|
||||
}()
|
||||
onSession(peerID, sess)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// authorizedPeer compares the request's bearer token with the cluster token in
|
||||
// constant time, matching the check the worker file-transfer server makes.
|
||||
//
|
||||
// Unlike that one, an empty configured token authorizes nobody: this route is
|
||||
// registered in every deployment, so failing open would publish an
|
||||
// unauthenticated mux to any caller that can reach the port.
|
||||
func authorizedPeer(r *http.Request, expected string) bool {
|
||||
if expected == "" {
|
||||
return false
|
||||
}
|
||||
// RFC 7235 makes the scheme case-insensitive; the token after it is not.
|
||||
const prefix = "Bearer "
|
||||
header := r.Header.Get("Authorization")
|
||||
if len(header) < len(prefix) || !strings.EqualFold(header[:len(prefix)], prefix) {
|
||||
return false
|
||||
}
|
||||
return subtle.ConstantTimeCompare([]byte(header[len(prefix):]), []byte(expected)) == 1
|
||||
}
|
||||
@@ -0,0 +1,407 @@
|
||||
package cluster_test
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/http/auth"
|
||||
"github.com/mudler/LocalAI/core/http/routes"
|
||||
clustersvc "github.com/mudler/LocalAI/core/services/cluster"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/libp2p/go-yamux/v5"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
)
|
||||
|
||||
// wsPeerURL is the peer route on a test server, named as peer-1.
|
||||
func wsPeerURL(s *httptest.Server) string {
|
||||
return "ws" + strings.TrimPrefix(s.URL, "http") + clustersvc.PeerPath + "?id=peer-1"
|
||||
}
|
||||
|
||||
// peerHeaders builds what a legitimate replica dials with: the deployment's
|
||||
// shared cluster token AND its own peer credential. Both are always sent,
|
||||
// because the handler checks both and a spec that sent one would be pinning
|
||||
// half the door.
|
||||
func peerHeaders(token string, cred clustersvc.PeerCredential) http.Header {
|
||||
h := http.Header{}
|
||||
h.Set("Authorization", "Bearer "+token)
|
||||
h.Set(clustersvc.PeerIdentityHeader, cred.Token())
|
||||
return h
|
||||
}
|
||||
|
||||
var _ = Describe("Peer link handler", func() {
|
||||
var (
|
||||
srv *httptest.Server
|
||||
sessions chan *yamux.Session
|
||||
reg *clustersvc.Registry
|
||||
cred clustersvc.PeerCredential
|
||||
ctx context.Context
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
ctx = context.Background()
|
||||
sessions = make(chan *yamux.Session, 1)
|
||||
|
||||
db := testutil.SetupTestDB()
|
||||
Expect(clustersvc.Migrate(ctx, db)).To(Succeed())
|
||||
reg = clustersvc.NewRegistry(db)
|
||||
// peer-1 is the replica every dial below claims to be, and this is the
|
||||
// credential that claim is checked against.
|
||||
cred = clustersvc.NewPeerCredential()
|
||||
Expect(reg.Register(ctx, "peer-1", "10.0.0.1:8080", "v1", cred.Hash())).To(Succeed())
|
||||
|
||||
e := echo.New()
|
||||
routes.RegisterClusterRoutes(e, "peer-token", reg, func(_ string, s *yamux.Session) {
|
||||
sessions <- s
|
||||
})
|
||||
srv = httptest.NewServer(e)
|
||||
DeferCleanup(srv.Close)
|
||||
})
|
||||
|
||||
It("rejects a connection with no token", func() {
|
||||
_, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), nil)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized))
|
||||
})
|
||||
|
||||
It("rejects a connection with the wrong token", func() {
|
||||
h := http.Header{}
|
||||
h.Set("Authorization", "Bearer wrong")
|
||||
_, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized))
|
||||
})
|
||||
|
||||
It("accepts an authenticated peer and yields a usable yamux session", func() {
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), peerHeaders("peer-token", cred))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = conn.Close() })
|
||||
|
||||
var serverSess *yamux.Session
|
||||
Eventually(sessions, "5s").Should(Receive(&serverSess))
|
||||
Expect(serverSess).ToNot(BeNil())
|
||||
|
||||
// The client wraps its side as a yamux CLIENT and opens a stream; the
|
||||
// server must accept it. This proves the WebSocket was adapted into a
|
||||
// stream-oriented conn correctly, which is the part most likely to be
|
||||
// subtly wrong.
|
||||
clientSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = clientSess.Close() })
|
||||
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
st, e := clientSess.OpenStream(GinkgoT().Context())
|
||||
if e == nil {
|
||||
_, _ = st.Write([]byte("hello"))
|
||||
}
|
||||
}()
|
||||
|
||||
accepted := make(chan []byte, 1)
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
st, e := serverSess.AcceptStream()
|
||||
if e != nil {
|
||||
return
|
||||
}
|
||||
buf := make([]byte, 5)
|
||||
if _, e := st.Read(buf); e == nil {
|
||||
accepted <- buf
|
||||
}
|
||||
}()
|
||||
Eventually(accepted, "10s").Should(Receive(Equal([]byte("hello"))))
|
||||
})
|
||||
|
||||
It("reports the peer id it proved", func() {
|
||||
ids := make(chan string, 1)
|
||||
e := echo.New()
|
||||
routes.RegisterClusterRoutes(e, "peer-token", reg, func(id string, _ *yamux.Session) { ids <- id })
|
||||
s2 := httptest.NewServer(e)
|
||||
DeferCleanup(s2.Close)
|
||||
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), peerHeaders("peer-token", cred))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = conn.Close() })
|
||||
|
||||
Eventually(ids, "5s").Should(Receive(Equal("peer-1")))
|
||||
})
|
||||
It("rejects every dial when no cluster token is configured", func() {
|
||||
// The route is registered in every deployment, so an empty configured
|
||||
// token must authorize nobody. Failing open the way the worker
|
||||
// file-transfer server's checkBearerToken does would publish an
|
||||
// unauthenticated yamux multiplexer to anyone who can reach the port.
|
||||
e := echo.New()
|
||||
accepted := make(chan *yamux.Session, 1)
|
||||
routes.RegisterClusterRoutes(e, "", reg, func(_ string, sess *yamux.Session) { accepted <- sess })
|
||||
s2 := httptest.NewServer(e)
|
||||
DeferCleanup(s2.Close)
|
||||
|
||||
for _, header := range []http.Header{nil, {"Authorization": []string{"Bearer "}}, {"Authorization": []string{"Bearer anything"}}, peerHeaders("", cred), peerHeaders("anything", cred)} {
|
||||
_, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), header)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized))
|
||||
}
|
||||
Expect(accepted).ToNot(Receive())
|
||||
})
|
||||
|
||||
It("accepts the bearer scheme in any case", func() {
|
||||
// RFC 7235 makes the scheme case-insensitive. The token after it is not.
|
||||
h := peerHeaders("peer-token", cred)
|
||||
h.Set("Authorization", "bearer peer-token")
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = conn.Close() })
|
||||
Eventually(sessions, "5s").Should(Receive())
|
||||
})
|
||||
|
||||
It("closes the session when the callback panics", func() {
|
||||
// net/http recovers the panic but leaves the hijacked socket open, so
|
||||
// without the handler's own recover the peer would keep a link nobody
|
||||
// ever accepts streams on.
|
||||
e := echo.New()
|
||||
routes.RegisterClusterRoutes(e, "peer-token", reg, func(_ string, _ *yamux.Session) {
|
||||
panic("callback exploded")
|
||||
})
|
||||
s2 := httptest.NewServer(e)
|
||||
DeferCleanup(func() {
|
||||
// A hijacked connection the handler never closed would park
|
||||
// httptest's Close forever, turning the assertion below into a
|
||||
// suite hang. Forcing the conns shut keeps the failure legible.
|
||||
s2.CloseClientConnections()
|
||||
s2.Close()
|
||||
})
|
||||
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), peerHeaders("peer-token", cred))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = conn.Close() })
|
||||
|
||||
clientSess, err := yamux.Client(clustersvc.WebsocketConn(conn), nil, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = clientSess.Close() })
|
||||
|
||||
// Asserting on OpenStream would hang rather than fail: yamux only
|
||||
// acknowledges a stream once the peer accepts it, and the leak this
|
||||
// pins is precisely that nobody ever will. The session's own liveness
|
||||
// is the observable that answers in both directions.
|
||||
Eventually(clientSess.IsClosed, "10s").Should(BeTrue())
|
||||
})
|
||||
|
||||
It("rejects an authenticated dial that names no peer", func() {
|
||||
// The session is keyed by peer id, so a nameless link could never be
|
||||
// looked up again; refusing it is cheaper than leaking it.
|
||||
_, resp, err := websocket.DefaultDialer.Dial(
|
||||
"ws"+strings.TrimPrefix(srv.URL, "http")+"/api/cluster/peer", peerHeaders("peer-token", cred))
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusBadRequest))
|
||||
Expect(sessions).ToNot(Receive())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("Peer link auth coverage", func() {
|
||||
// These specs put the REAL global auth middleware in front of the REAL
|
||||
// registrar and prove a peer dial reaches the handler anyway. The peer link
|
||||
// authenticates with the cluster token, not a session, so it only works
|
||||
// while its path sits under the prefix auth exempts; moving either one
|
||||
// alone 401s every peer dial, and the two live in packages that must not
|
||||
// import each other.
|
||||
//
|
||||
// The predicate that grants the exemption is unexported, so this asserts on
|
||||
// its effect rather than on it: what a caller can observe is whether the
|
||||
// request reaches the handler.
|
||||
var (
|
||||
srv *httptest.Server
|
||||
sessions chan *yamux.Session
|
||||
cred clustersvc.PeerCredential
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
ctx := context.Background()
|
||||
sessions = make(chan *yamux.Session, 1)
|
||||
db := testutil.SetupTestDB()
|
||||
Expect(clustersvc.Migrate(ctx, db)).To(Succeed())
|
||||
reg := clustersvc.NewRegistry(db)
|
||||
cred = clustersvc.NewPeerCredential()
|
||||
Expect(reg.Register(ctx, "peer-1", "10.0.0.1:8080", "v1", cred.Hash())).To(Succeed())
|
||||
|
||||
e := echo.New()
|
||||
// A nil DB with one legacy API key is the cheapest configuration that
|
||||
// turns the middleware ON without a database. With neither, Middleware
|
||||
// short-circuits to next() and every assertion below would pass against
|
||||
// a server that has no auth at all.
|
||||
e.Use(auth.Middleware(nil, &config.ApplicationConfig{ApiKeys: []string{"an-api-key"}}))
|
||||
routes.RegisterClusterRoutes(e, "peer-token", reg, func(_ string, s *yamux.Session) { sessions <- s })
|
||||
// A route outside the cluster prefix, registered on the same server, is
|
||||
// the control: it proves the middleware in front of both is live.
|
||||
e.GET("/api/nodes", func(c echo.Context) error { return c.NoContent(http.StatusOK) })
|
||||
srv = httptest.NewServer(e)
|
||||
DeferCleanup(srv.Close)
|
||||
})
|
||||
|
||||
It("refuses an uncredentialed request to a route outside the cluster prefix", func() {
|
||||
resp, err := http.Get(srv.URL + "/api/nodes")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusUnauthorized),
|
||||
"the global auth middleware is not actually guarding this server, so the peer-route assertions below would prove nothing")
|
||||
})
|
||||
|
||||
It("lets a peer dial reach the handler, which is the only thing that can authenticate it", func() {
|
||||
// The cluster token is not one of the API keys the middleware knows, so
|
||||
// a 400 from the handler's own missing-id check can only mean the
|
||||
// request was let through unauthenticated by the middleware.
|
||||
req, err := http.NewRequestWithContext(GinkgoT().Context(), http.MethodGet, srv.URL+clustersvc.PeerPath, nil)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
req.Header.Set("Authorization", "Bearer peer-token")
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
defer func() { _ = resp.Body.Close() }()
|
||||
Expect(resp.StatusCode).To(Equal(http.StatusBadRequest),
|
||||
"a peer dial must reach the handler; 401 here means the peer route left the auth-exempt prefix %q", auth.ClusterPathPrefix)
|
||||
})
|
||||
|
||||
It("completes a full peer handshake through the guarded server", func() {
|
||||
// The status-code assertion above cannot see the upgrade, and the
|
||||
// upgrade is what a peer actually does.
|
||||
conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), peerHeaders("peer-token", cred))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(func() { _ = conn.Close() })
|
||||
Eventually(sessions, "5s").Should(Receive())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("Peer link identity", func() {
|
||||
// The status codes a peer's own retry loop reads, and the order they are
|
||||
// decided in. Every refusal here happens BEFORE the WebSocket upgrade: a
|
||||
// handler that upgraded first would answer with a WebSocket error instead
|
||||
// of a status, and neither an operator nor a dialler could tell an
|
||||
// authorization failure from a transport one.
|
||||
var (
|
||||
srv *httptest.Server
|
||||
sessions chan *yamux.Session
|
||||
reg *clustersvc.Registry
|
||||
cred clustersvc.PeerCredential
|
||||
ctx context.Context
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
ctx = context.Background()
|
||||
sessions = make(chan *yamux.Session, 1)
|
||||
db := testutil.SetupTestDB()
|
||||
Expect(clustersvc.Migrate(ctx, db)).To(Succeed())
|
||||
reg = clustersvc.NewRegistry(db)
|
||||
cred = clustersvc.NewPeerCredential()
|
||||
Expect(reg.Register(ctx, "peer-1", "10.0.0.1:8080", "v1", cred.Hash())).To(Succeed())
|
||||
|
||||
e := echo.New()
|
||||
routes.RegisterClusterRoutes(e, "peer-token", reg, func(_ string, s *yamux.Session) { sessions <- s })
|
||||
srv = httptest.NewServer(e)
|
||||
DeferCleanup(srv.Close)
|
||||
})
|
||||
|
||||
// dialPeer returns the handshake status a dial was answered with.
|
||||
dialPeer := func(target *httptest.Server, h http.Header) int {
|
||||
GinkgoHelper()
|
||||
conn, resp, err := websocket.DefaultDialer.Dial(wsPeerURL(target), h)
|
||||
if err == nil {
|
||||
DeferCleanup(func() { _ = conn.Close() })
|
||||
Expect(resp).ToNot(BeNil())
|
||||
return resp.StatusCode
|
||||
}
|
||||
Expect(resp).ToNot(BeNil())
|
||||
Expect(resp.Header.Get("Upgrade")).To(BeEmpty(),
|
||||
"the handler upgraded a dial it then refused; a peer reads the status, not a WebSocket error")
|
||||
return resp.StatusCode
|
||||
}
|
||||
|
||||
It("refuses a dial that carries the shared token and no peer credential", func() {
|
||||
// This is the mixed-version case, and the answer is deliberately a
|
||||
// refusal rather than a fallback. A replica running a release from
|
||||
// before peer credentials sends no such header, and so does an attacker
|
||||
// holding the shared token: there is no way to accept the first without
|
||||
// accepting the second, so both are refused and the handler logs which
|
||||
// replica it refused and why.
|
||||
h := http.Header{}
|
||||
h.Set("Authorization", "Bearer peer-token")
|
||||
Expect(dialPeer(srv, h)).To(Equal(http.StatusUnauthorized))
|
||||
Expect(sessions).ToNot(Receive())
|
||||
})
|
||||
|
||||
It("refuses a dial that presents an empty peer credential", func() {
|
||||
h := peerHeaders("peer-token", clustersvc.PeerCredential{})
|
||||
Expect(dialPeer(srv, h)).To(Equal(http.StatusUnauthorized))
|
||||
Expect(sessions).ToNot(Receive())
|
||||
})
|
||||
|
||||
It("refuses a dial that declares a real replica with the wrong credential", func() {
|
||||
Expect(dialPeer(srv, peerHeaders("peer-token", clustersvc.NewPeerCredential()))).
|
||||
To(Equal(http.StatusUnauthorized))
|
||||
Expect(sessions).ToNot(Receive())
|
||||
})
|
||||
|
||||
It("refuses a dial whose declared replica publishes no credential", func() {
|
||||
Expect(reg.Register(ctx, "peer-1", "10.0.0.1:8080", "v1", "")).To(Succeed())
|
||||
Expect(dialPeer(srv, peerHeaders("peer-token", cred))).To(Equal(http.StatusUnauthorized))
|
||||
Expect(sessions).ToNot(Receive())
|
||||
})
|
||||
|
||||
It("accepts a dial that proves the replica it declares", func() {
|
||||
Expect(dialPeer(srv, peerHeaders("peer-token", cred))).To(Equal(http.StatusSwitchingProtocols))
|
||||
Eventually(sessions, "5s").Should(Receive())
|
||||
})
|
||||
|
||||
It("answers 503, not 401, when there is no cluster registry to check against", func() {
|
||||
// A frontend with nothing to resolve replica ids against cannot
|
||||
// authenticate anybody. Answering "unauthorized" would send an operator
|
||||
// hunting a token problem that does not exist, and accepting would
|
||||
// publish an unauthenticated multiplexer, so it does neither.
|
||||
e := echo.New()
|
||||
accepted := make(chan *yamux.Session, 1)
|
||||
routes.RegisterClusterRoutes(e, "peer-token", nil, func(_ string, s *yamux.Session) { accepted <- s })
|
||||
s2 := httptest.NewServer(e)
|
||||
DeferCleanup(s2.Close)
|
||||
|
||||
Expect(dialPeer(s2, peerHeaders("peer-token", cred))).To(Equal(http.StatusServiceUnavailable))
|
||||
Expect(accepted).ToNot(Receive())
|
||||
})
|
||||
|
||||
It("answers 500, not 401, when the replica lookup itself fails", func() {
|
||||
// A query that FAILED is neither a rejection nor an absence. Telling a
|
||||
// healthy replica its credentials are wrong because the database was
|
||||
// unreadable sends it re-registering; nothing above may read it as a
|
||||
// worker having gone away either.
|
||||
db := testutil.SetupTestDB()
|
||||
Expect(clustersvc.Migrate(ctx, db)).To(Succeed())
|
||||
broken := clustersvc.NewRegistry(db)
|
||||
Expect(broken.Register(ctx, "peer-1", "10.0.0.1:8080", "v1", cred.Hash())).To(Succeed())
|
||||
sqlDB, err := db.DB()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(sqlDB.Close()).To(Succeed())
|
||||
|
||||
e := echo.New()
|
||||
accepted := make(chan *yamux.Session, 1)
|
||||
routes.RegisterClusterRoutes(e, "peer-token", broken, func(_ string, s *yamux.Session) { accepted <- s })
|
||||
s2 := httptest.NewServer(e)
|
||||
DeferCleanup(s2.Close)
|
||||
|
||||
Expect(dialPeer(s2, peerHeaders("peer-token", cred))).To(Equal(http.StatusInternalServerError))
|
||||
Expect(accepted).ToNot(Receive())
|
||||
})
|
||||
|
||||
It("refuses an anonymous dial before it looks anything up", func() {
|
||||
// Ordering: the shared token is still the first gate, unchanged, so a
|
||||
// dial with no Authorization at all is a 401 and not a 400 about an id
|
||||
// or a 401 about a credential. The route-coverage test issues exactly
|
||||
// this shape.
|
||||
Expect(dialPeer(srv, nil)).To(Equal(http.StatusUnauthorized))
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,106 @@
|
||||
package localai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/http/auth"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/agentpool"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
"github.com/mudler/LocalAI/pkg/system"
|
||||
)
|
||||
|
||||
// DELETE /api/agent/tasks/:id takes the task id straight off the URL, so the
|
||||
// only thing standing between one tenant and another tenant's task is how the
|
||||
// service behind the route is scoped.
|
||||
//
|
||||
// The assertion is on the SURVIVING ROW rather than on the status code. Before
|
||||
// the task subject carried its tenant, the handler answered 200 and destroyed
|
||||
// the row; it now answers 404 because the task is not in this tenant's map. A
|
||||
// status assertion would have seen nothing wrong in the first case, and reading
|
||||
// the row keeps the spec honest whichever status a later change picks.
|
||||
var _ = Describe("DELETE /api/agent/tasks/:id across tenants", func() {
|
||||
var (
|
||||
e *echo.Echo
|
||||
app *application.Application
|
||||
other *agentpool.AgentJobService
|
||||
cancel context.CancelFunc
|
||||
taskID string
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
tmp := GinkgoT().TempDir()
|
||||
st, err := system.GetSystemState(
|
||||
system.WithModelPath(filepath.Join(tmp, "models")),
|
||||
system.WithBackendPath(filepath.Join(tmp, "backends")),
|
||||
)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
var ctx context.Context
|
||||
ctx, cancel = context.WithCancel(context.Background())
|
||||
app, err = application.New(config.WithContext(ctx), config.WithSystemState(st))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// One carrier, two tenants. The service the route reaches belongs to
|
||||
// u1; `other` is u2's, living on another replica as far as the bus is
|
||||
// concerned.
|
||||
bus := testutil.NewFakeBus()
|
||||
|
||||
mine := app.AgentJobService()
|
||||
Expect(mine).ToNot(BeNil())
|
||||
mine.SetUserID("u1")
|
||||
mine.SetTaskSyncBus(bus)
|
||||
Expect(mine.LoadTasksFromFile()).To(Succeed())
|
||||
|
||||
otherDir := GinkgoT().TempDir()
|
||||
otherCfg := config.NewApplicationConfig(
|
||||
config.WithDynamicConfigDir(otherDir),
|
||||
config.WithContext(ctx),
|
||||
)
|
||||
otherCfg.SystemState = st
|
||||
other = agentpool.NewAgentJobServiceWithPaths(otherCfg, nil, nil, nil,
|
||||
filepath.Join(otherDir, "tasks.json"), filepath.Join(otherDir, "jobs.json"))
|
||||
other.SetUserID("u2")
|
||||
other.SetTaskSyncBus(bus)
|
||||
Expect(other.LoadTasksFromFile()).To(Succeed())
|
||||
|
||||
taskID, err = other.CreateTask(schema.Task{Name: "u2 only", Model: "m", Prompt: "p"})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
e = echo.New()
|
||||
// Stand in for the auth middleware the real route runs behind: it is
|
||||
// what puts the caller into the echo context, and getJobService reads
|
||||
// the tenant from there.
|
||||
e.Use(func(next echo.HandlerFunc) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
c.Set("auth_user", &auth.User{ID: "u1", Role: auth.RoleUser})
|
||||
return next(c)
|
||||
}
|
||||
})
|
||||
e.DELETE("/api/agent/tasks/:id", DeleteTaskEndpoint(app))
|
||||
})
|
||||
|
||||
AfterEach(func() {
|
||||
cancel()
|
||||
})
|
||||
|
||||
It("leaves the owning tenant's task readable", func() {
|
||||
req := httptest.NewRequest(http.MethodDelete, "/api/agent/tasks/"+taskID, nil)
|
||||
rec := httptest.NewRecorder()
|
||||
e.ServeHTTP(rec, req)
|
||||
|
||||
got, err := other.GetTask(taskID)
|
||||
Expect(err).ToNot(HaveOccurred(), "tenant u2's task must survive a delete issued by tenant u1")
|
||||
Expect(got.Name).To(Equal("u2 only"))
|
||||
Expect(other.ListTasks()).To(HaveLen(1))
|
||||
})
|
||||
})
|
||||
@@ -121,6 +121,17 @@ func BackendLogsWebSocketEndpoint(ml *model.ModelLoader) echo.HandlerFunc {
|
||||
|
||||
conn := &backendLogsConn{Conn: ws}
|
||||
|
||||
// KNOWN RACE: the snapshot is sent before the subscription is registered, so
|
||||
// a line appended in that window is never streamed. A viewer attaching while
|
||||
// a model loads (when a backend is at its noisiest) can silently miss lines;
|
||||
// they stay in the buffer, so a reload shows them. Fixing it needs an atomic
|
||||
// snapshot-plus-subscribe held under the buffer's own lock (buf.mu in
|
||||
// pkg/model/backend_log_store.go), because that is the lock AppendLine takes
|
||||
// while it enqueues and fans out to subscribers. The store-level s.mu guards
|
||||
// only the buffers map and excludes nothing an appender does, so taking it
|
||||
// leaves this race exactly where it is. Reordering these two calls is not a
|
||||
// fix either: it would duplicate instead of drop.
|
||||
|
||||
// Send existing lines as initial batch
|
||||
existingLines := ml.BackendLogs().GetLines(modelID)
|
||||
initialMsg := map[string]any{
|
||||
|
||||
@@ -188,7 +188,7 @@ backend: llama-cpp
|
||||
lifecycle := &endpointLifecycleRecorder{}
|
||||
galleryService := galleryop.NewGalleryService(appConfig, nil)
|
||||
client := &endpointRecordingClient{}
|
||||
galleryService.SetNATSClient(client)
|
||||
galleryService.SetBroadcaster(client)
|
||||
endpointApp := echo.New()
|
||||
endpointApp.PATCH("/api/models/config-json/:name", PatchConfigEndpoint(configLoader, galleryService, appConfig, lifecycle))
|
||||
|
||||
|
||||
@@ -10,7 +10,6 @@ import (
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
@@ -46,20 +45,6 @@ func (c *endpointRecordingClient) Publish(subject string, data any) error {
|
||||
func (*endpointRecordingClient) Subscribe(string, func([]byte)) (messaging.Subscription, error) {
|
||||
return &endpointSubscription{}, nil
|
||||
}
|
||||
func (*endpointRecordingClient) QueueSubscribe(string, string, func([]byte)) (messaging.Subscription, error) {
|
||||
return &endpointSubscription{}, nil
|
||||
}
|
||||
func (*endpointRecordingClient) QueueSubscribeReply(string, string, func([]byte, func([]byte))) (messaging.Subscription, error) {
|
||||
return &endpointSubscription{}, nil
|
||||
}
|
||||
func (*endpointRecordingClient) SubscribeReply(string, func([]byte, func([]byte))) (messaging.Subscription, error) {
|
||||
return &endpointSubscription{}, nil
|
||||
}
|
||||
func (*endpointRecordingClient) Request(string, []byte, time.Duration) ([]byte, error) {
|
||||
return nil, nil
|
||||
}
|
||||
func (*endpointRecordingClient) IsConnected() bool { return true }
|
||||
func (*endpointRecordingClient) Close() {}
|
||||
|
||||
// testRenderer is a simple renderer for tests that returns JSON
|
||||
type testRenderer struct{}
|
||||
@@ -126,7 +111,7 @@ var _ = Describe("Edit Model test", func() {
|
||||
Expect(loader.LoadModelConfigsFromPath(tempDir)).To(Succeed())
|
||||
galleryService := galleryop.NewGalleryService(applicationConfig, nil)
|
||||
client := &endpointRecordingClient{}
|
||||
galleryService.SetNATSClient(client)
|
||||
galleryService.SetBroadcaster(client)
|
||||
|
||||
app := echo.New()
|
||||
app.POST("/models/edit/:name", EditModelEndpoint(loader, galleryService, applicationConfig))
|
||||
@@ -152,7 +137,7 @@ var _ = Describe("Edit Model test", func() {
|
||||
Expect(loader.LoadModelConfigsFromPath(tempDir)).To(Succeed())
|
||||
galleryService := galleryop.NewGalleryService(applicationConfig, nil)
|
||||
client := &endpointRecordingClient{}
|
||||
galleryService.SetNATSClient(client)
|
||||
galleryService.SetBroadcaster(client)
|
||||
|
||||
app := echo.New()
|
||||
app.POST("/models/edit/:name", EditModelEndpoint(loader, galleryService, applicationConfig))
|
||||
@@ -276,7 +261,7 @@ var _ = Describe("Edit Model test", func() {
|
||||
Expect(peerLoader.LoadModelConfigsFromPath(tempDir)).To(Succeed())
|
||||
galleryService := galleryop.NewGalleryService(applicationConfig, nil)
|
||||
client := &endpointRecordingClient{}
|
||||
galleryService.SetNATSClient(client)
|
||||
galleryService.SetBroadcaster(client)
|
||||
lifecycle := &endpointLifecycleRecorder{pendingCleanup: 2}
|
||||
app := echo.New()
|
||||
app.POST("/models/edit/:name", EditModelEndpoint(loader, galleryService, applicationConfig, lifecycle))
|
||||
|
||||
@@ -57,7 +57,7 @@ type MCPErrorEvent struct {
|
||||
// @Param request body schema.OpenAIRequest true "query params"
|
||||
// @Success 200 {object} schema.OpenAIResponse "Response"
|
||||
// @Router /v1/mcp/chat/completions [post]
|
||||
func MCPEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, natsClient mcpTools.MCPNATSClient, compressor middleware.ChatCompressor) echo.HandlerFunc {
|
||||
func MCPEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, agentControl mcpTools.AgentControl, compressor middleware.ChatCompressor) echo.HandlerFunc {
|
||||
// The legacy /v1/mcp/chat/completions endpoint never opts into the
|
||||
// in-process LocalAI Assistant tool surface — pass nil holder so the
|
||||
// assistant branch in chat.go is unreachable from this code path.
|
||||
@@ -65,7 +65,7 @@ func MCPEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
|
||||
// the per-model PII config and is kept for backward compatibility.
|
||||
// The request-side middleware on the main chat route handles
|
||||
// filtering for the standard /v1/chat/completions path.
|
||||
chatHandler := openai.ChatEndpoint(cl, ml, evaluator, appConfig, natsClient, nil, compressor)
|
||||
chatHandler := openai.ChatEndpoint(cl, ml, evaluator, appConfig, agentControl, nil, compressor)
|
||||
|
||||
return func(c echo.Context) error {
|
||||
input, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_LOCALAI_REQUEST).(*schema.OpenAIRequest)
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package localai
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
)
|
||||
|
||||
// mcpLocalSessionsOnly answers the request itself when this frontend cannot
|
||||
// serve it, and reports whether it did.
|
||||
//
|
||||
// MCP prompts and resources are served ONLY from MCP sessions this process
|
||||
// holds. In distributed mode the frontend holds none: creating them is what an
|
||||
// agent worker is for, because a stdio server usually means running docker, and
|
||||
// nothing in this programme carries prompts or resources to one. That is a hole
|
||||
// that PREDATES the removal of the message bus and is not closed by it; tools
|
||||
// and discovery had a carrier and these two never did.
|
||||
//
|
||||
// Until this task the endpoints did not say so. A model with MCP servers
|
||||
// configured answered 200 with an empty list, which reads as "this model has no
|
||||
// prompts" and is a different statement from "this deployment cannot tell you".
|
||||
// A client cannot act on the first and can act on the second, and an operator
|
||||
// reading an empty list has no reason to look further. 501 with the reason in
|
||||
// the body is the smallest honest answer, and it is deliberately NOT a 404 or
|
||||
// an empty 200: the resource may well exist, on a worker this frontend has no
|
||||
// verb to ask.
|
||||
//
|
||||
// One function for four endpoints, so the four cannot drift into four different
|
||||
// stories about the same limitation. Each endpoint still has to CALL it, and
|
||||
// each call is pinned by its own spec.
|
||||
func mcpLocalSessionsOnly(c echo.Context, appConfig *config.ApplicationConfig, surface string) (bool, error) {
|
||||
if appConfig == nil || !appConfig.Distributed.Enabled {
|
||||
return false, nil
|
||||
}
|
||||
return true, c.JSON(http.StatusNotImplemented, map[string]string{
|
||||
"error": "MCP " + surface + " are not available in distributed mode: they are served only from MCP sessions held by this frontend, and in distributed mode the sessions live on an agent worker, which serves no " + surface + " verb",
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package localai
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
)
|
||||
|
||||
// The MCP prompts and resources endpoints work only against sessions THIS
|
||||
// process holds, and in distributed mode it holds none. That is a hole older
|
||||
// than the removal of the message bus: tools and discovery had a carrier to an
|
||||
// agent worker and these two never did, and nothing in this programme gives
|
||||
// them one.
|
||||
//
|
||||
// What these specs pin is that the hole is now HONEST. The old answer was 200
|
||||
// with an empty list, which reads as "this model has no prompts": a client
|
||||
// cannot act on it, and an operator reading it has no reason to look further.
|
||||
//
|
||||
// Each endpoint is asserted separately on purpose. The refusal is one function
|
||||
// but four CALLS, and a call site that lost its call would still leave the
|
||||
// other three green.
|
||||
var _ = Describe("MCP prompts and resources in distributed mode", func() {
|
||||
distributed := &config.ApplicationConfig{}
|
||||
distributed.Distributed.Enabled = true
|
||||
|
||||
standalone := &config.ApplicationConfig{}
|
||||
|
||||
// loader with no models at all. Every endpoint below looks a model up, and
|
||||
// the refusal has to come FIRST: a 404 for an unknown model would be a
|
||||
// different, and wrong, story about the same request.
|
||||
emptyLoader := config.NewModelConfigLoader("")
|
||||
|
||||
call := func(h echo.HandlerFunc, method, path, body string, params map[string]string) *httptest.ResponseRecorder {
|
||||
GinkgoHelper()
|
||||
e := echo.New()
|
||||
req := httptest.NewRequest(method, path, strings.NewReader(body))
|
||||
req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
names := make([]string, 0, len(params))
|
||||
values := make([]string, 0, len(params))
|
||||
for k, v := range params {
|
||||
names = append(names, k)
|
||||
values = append(values, v)
|
||||
}
|
||||
c.SetParamNames(names...)
|
||||
c.SetParamValues(values...)
|
||||
Expect(h(c)).To(Succeed())
|
||||
return rec
|
||||
}
|
||||
|
||||
// assertRefused checks the STATUS and the BODY. A status-only assertion
|
||||
// cannot tell a 501 with an empty body from the empty 200 this replaces,
|
||||
// and the body is the whole reason the answer is worth giving.
|
||||
assertRefused := func(rec *httptest.ResponseRecorder, surface string) {
|
||||
GinkgoHelper()
|
||||
Expect(rec.Code).To(Equal(http.StatusNotImplemented))
|
||||
var body map[string]string
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &body)).To(Succeed())
|
||||
Expect(body["error"]).To(ContainSubstring(surface))
|
||||
Expect(body["error"]).To(ContainSubstring("distributed mode"))
|
||||
Expect(body["error"]).To(ContainSubstring("agent worker"))
|
||||
}
|
||||
|
||||
It("refuses to list prompts rather than answering with an empty list", func() {
|
||||
assertRefused(call(MCPPromptsEndpoint(emptyLoader, distributed),
|
||||
http.MethodGet, "/v1/mcp/prompts/m", "", map[string]string{"model": "m"}), "prompts")
|
||||
})
|
||||
|
||||
It("refuses to expand a prompt", func() {
|
||||
assertRefused(call(MCPGetPromptEndpoint(emptyLoader, distributed),
|
||||
http.MethodPost, "/v1/mcp/prompts/m/p", `{"arguments":{}}`,
|
||||
map[string]string{"model": "m", "prompt": "p"}), "prompts")
|
||||
})
|
||||
|
||||
It("refuses to list resources rather than answering with an empty list", func() {
|
||||
assertRefused(call(MCPResourcesEndpoint(emptyLoader, distributed),
|
||||
http.MethodGet, "/v1/mcp/resources/m", "", map[string]string{"model": "m"}), "resources")
|
||||
})
|
||||
|
||||
It("refuses to read a resource", func() {
|
||||
assertRefused(call(MCPReadResourceEndpoint(emptyLoader, distributed),
|
||||
http.MethodPost, "/v1/mcp/resources/m/read", `{"uri":"file:///x"}`,
|
||||
map[string]string{"model": "m"}), "resources")
|
||||
})
|
||||
|
||||
It("does not refuse in a standalone deployment, where the sessions are here", func() {
|
||||
// The negative control. A guard that fired unconditionally would pass
|
||||
// every assertion above and break MCP prompts for every single-binary
|
||||
// deployment, which is where they actually work.
|
||||
e := echo.New()
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/mcp/prompts/m", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
c.SetParamNames("model")
|
||||
c.SetParamValues("m")
|
||||
|
||||
err := MCPPromptsEndpoint(emptyLoader, standalone)(c)
|
||||
// The model does not exist, so this fails on the lookup. What matters
|
||||
// is that it got PAST the guard: a 501 here would mean the guard fired
|
||||
// with distributed mode off.
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(rec.Code).ToNot(Equal(http.StatusNotImplemented))
|
||||
})
|
||||
})
|
||||
@@ -17,6 +17,13 @@ func MCPPromptsEndpoint(cl *config.ModelConfigLoader, appConfig *config.Applicat
|
||||
return echo.ErrBadRequest
|
||||
}
|
||||
|
||||
// Before the model lookup and before the empty-config shortcut: both
|
||||
// answer 200 with an empty list, which is the very answer this refuses
|
||||
// to give in a deployment that cannot look.
|
||||
if refused, err := mcpLocalSessionsOnly(c, appConfig, "prompts"); refused {
|
||||
return err
|
||||
}
|
||||
|
||||
cfg, exists := cl.GetModelConfig(modelName)
|
||||
if !exists {
|
||||
return fmt.Errorf("model %q not found", modelName)
|
||||
@@ -86,6 +93,10 @@ func MCPGetPromptEndpoint(cl *config.ModelConfigLoader, appConfig *config.Applic
|
||||
return echo.ErrBadRequest
|
||||
}
|
||||
|
||||
if refused, err := mcpLocalSessionsOnly(c, appConfig, "prompts"); refused {
|
||||
return err
|
||||
}
|
||||
|
||||
cfg, exists := cl.GetModelConfig(modelName)
|
||||
if !exists {
|
||||
return fmt.Errorf("model %q not found", modelName)
|
||||
|
||||
@@ -17,6 +17,12 @@ func MCPResourcesEndpoint(cl *config.ModelConfigLoader, appConfig *config.Applic
|
||||
return echo.ErrBadRequest
|
||||
}
|
||||
|
||||
// See the prompts endpoint: this has to come before the empty-config
|
||||
// shortcut, which answers with the empty list this refuses to invent.
|
||||
if refused, err := mcpLocalSessionsOnly(c, appConfig, "resources"); refused {
|
||||
return err
|
||||
}
|
||||
|
||||
cfg, exists := cl.GetModelConfig(modelName)
|
||||
if !exists {
|
||||
return fmt.Errorf("model %q not found", modelName)
|
||||
@@ -73,6 +79,10 @@ func MCPReadResourceEndpoint(cl *config.ModelConfigLoader, appConfig *config.App
|
||||
return echo.ErrBadRequest
|
||||
}
|
||||
|
||||
if refused, err := mcpLocalSessionsOnly(c, appConfig, "resources"); refused {
|
||||
return err
|
||||
}
|
||||
|
||||
cfg, exists := cl.GetModelConfig(modelName)
|
||||
if !exists {
|
||||
return fmt.Errorf("model %q not found", modelName)
|
||||
|
||||
@@ -13,7 +13,7 @@ import (
|
||||
|
||||
// MCPServersEndpoint returns the list of MCP servers and their tools for a given model.
|
||||
// GET /v1/mcp/servers/:model
|
||||
func MCPServersEndpoint(cl *config.ModelConfigLoader, appConfig *config.ApplicationConfig, natsClient mcpTools.MCPNATSClient) echo.HandlerFunc {
|
||||
func MCPServersEndpoint(cl *config.ModelConfigLoader, appConfig *config.ApplicationConfig, agentControl mcpTools.AgentControl) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
modelName := c.Param("model")
|
||||
if modelName == "" {
|
||||
@@ -45,10 +45,10 @@ func MCPServersEndpoint(cl *config.ModelConfigLoader, appConfig *config.Applicat
|
||||
})
|
||||
}
|
||||
|
||||
// In distributed mode, route discovery through NATS to an agent worker
|
||||
// In distributed mode, ask an agent worker over the tunnel it holds
|
||||
// that can actually connect to the MCP servers.
|
||||
if natsClient != nil {
|
||||
resp, err := mcpTools.DiscoverMCPToolsRemote(c.Request().Context(), natsClient, cfg.Name, remote, stdio)
|
||||
if agentControl != nil {
|
||||
resp, err := mcpTools.DiscoverMCPToolsRemote(c.Request().Context(), agentControl, cfg.Name, remote, stdio)
|
||||
if err != nil {
|
||||
return c.JSON(http.StatusOK, map[string]any{
|
||||
"model": modelName,
|
||||
@@ -80,7 +80,7 @@ func MCPServersEndpoint(cl *config.ModelConfigLoader, appConfig *config.Applicat
|
||||
|
||||
// MCPServersEndpointFromMiddleware is a version that uses the middleware-resolved model config.
|
||||
// This allows it to use the same middleware chain as other endpoints.
|
||||
func MCPServersEndpointFromMiddleware(natsClient mcpTools.MCPNATSClient) echo.HandlerFunc {
|
||||
func MCPServersEndpointFromMiddleware(agentControl mcpTools.AgentControl) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
cfg, ok := c.Get(middleware.CONTEXT_LOCALS_KEY_MODEL_CONFIG).(*config.ModelConfig)
|
||||
if !ok || cfg == nil {
|
||||
@@ -103,9 +103,9 @@ func MCPServersEndpointFromMiddleware(natsClient mcpTools.MCPNATSClient) echo.Ha
|
||||
})
|
||||
}
|
||||
|
||||
// In distributed mode, route discovery through NATS to an agent worker.
|
||||
if natsClient != nil {
|
||||
resp, err := mcpTools.DiscoverMCPToolsRemote(c.Request().Context(), natsClient, cfg.Name, remote, stdio)
|
||||
// In distributed mode, ask an agent worker over the tunnel it holds.
|
||||
if agentControl != nil {
|
||||
resp, err := mcpTools.DiscoverMCPToolsRemote(c.Request().Context(), agentControl, cfg.Name, remote, stdio)
|
||||
if err != nil {
|
||||
return c.JSON(http.StatusOK, map[string]any{
|
||||
"model": cfg.Name,
|
||||
|
||||
@@ -2,6 +2,7 @@ package localai
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"crypto/subtle"
|
||||
"encoding/hex"
|
||||
@@ -9,6 +10,7 @@ import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sync"
|
||||
@@ -29,7 +31,6 @@ import (
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/core/services/nodes/prefixcache"
|
||||
"github.com/mudler/LocalAI/pkg/httpclient"
|
||||
"github.com/mudler/LocalAI/pkg/natsauth"
|
||||
"github.com/mudler/LocalAI/pkg/vrambudget"
|
||||
)
|
||||
|
||||
@@ -75,10 +76,14 @@ func GetNodeEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc {
|
||||
|
||||
// RegisterNodeRequest is the request body for registering a new worker node.
|
||||
type RegisterNodeRequest struct {
|
||||
Name string `json:"name"`
|
||||
NodeType string `json:"node_type,omitempty"` // "backend" (default) or "agent"
|
||||
Address string `json:"address"`
|
||||
HTTPAddress string `json:"http_address,omitempty"`
|
||||
Name string `json:"name"`
|
||||
NodeType string `json:"node_type,omitempty"` // "backend" (default) or "agent"
|
||||
// No address and no http_address. A worker has no inbound endpoint to
|
||||
// register: it holds one outbound tunnel to a frontend replica and every
|
||||
// protocol the frontend speaks to it travels on that. An older worker still
|
||||
// sends both keys and they are ignored, which is the intended outcome:
|
||||
// storing them would put a dialable-looking endpoint back in the API for
|
||||
// something nothing dials.
|
||||
Token string `json:"token,omitempty"`
|
||||
TotalVRAM uint64 `json:"total_vram,omitempty"`
|
||||
AvailableVRAM uint64 `json:"available_vram,omitempty"`
|
||||
@@ -111,7 +116,7 @@ type RegisterNodeRequest struct {
|
||||
// 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) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
var req RegisterNodeRequest
|
||||
if err := c.Bind(&req); err != nil {
|
||||
@@ -140,22 +145,15 @@ 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
|
||||
// A backend worker no longer has to state an address; the tunnel it
|
||||
// dials is what makes it reachable, and requiring one here would refuse
|
||||
// exactly the workers this design is for.
|
||||
if req.Name == "" {
|
||||
return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "name is required"))
|
||||
}
|
||||
if nodeType == nodes.NodeTypeBackend && req.Address == "" {
|
||||
return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "address is required for backend workers"))
|
||||
}
|
||||
if len(req.Name) > 255 {
|
||||
return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "name exceeds 255 characters"))
|
||||
}
|
||||
if len(req.Address) > 512 {
|
||||
return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "address exceeds 512 characters"))
|
||||
}
|
||||
if len(req.HTTPAddress) > 512 {
|
||||
return c.JSON(http.StatusBadRequest, nodeError(http.StatusBadRequest, "http_address exceeds 512 characters"))
|
||||
}
|
||||
|
||||
// Hash the token for storage (if provided)
|
||||
var tokenHash string
|
||||
@@ -175,8 +173,6 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au
|
||||
node := &nodes.BackendNode{
|
||||
Name: req.Name,
|
||||
NodeType: nodeType,
|
||||
Address: req.Address,
|
||||
HTTPAddress: req.HTTPAddress,
|
||||
TokenHash: tokenHash,
|
||||
TotalVRAM: req.TotalVRAM,
|
||||
AvailableVRAM: req.AvailableVRAM,
|
||||
@@ -244,7 +240,7 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au
|
||||
}
|
||||
}
|
||||
|
||||
attachNatsJWT(response, node, natsCfg)
|
||||
attachTunnelToken(ctx, response, registry, node)
|
||||
|
||||
return c.JSON(http.StatusCreated, response)
|
||||
}
|
||||
@@ -252,7 +248,7 @@ 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) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
ctx := c.Request().Context()
|
||||
id := c.Param("id")
|
||||
@@ -282,24 +278,101 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr
|
||||
}
|
||||
}
|
||||
|
||||
attachNatsJWT(response, node, natsCfg)
|
||||
|
||||
return c.JSON(http.StatusOK, response)
|
||||
}
|
||||
}
|
||||
|
||||
// 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 {
|
||||
// attachTunnelToken mints this node a fresh tunnel credential, stores only its
|
||||
// hash, and puts the plaintext in the registration response.
|
||||
//
|
||||
// It is minted for EVERY node that registers, pending ones included, which is a
|
||||
// deliberate divergence from the other per-node credential in this file: the
|
||||
// agent worker's API key (provisionAgentWorkerKey) is withheld from a node
|
||||
// awaiting approval. It is a bearer grant that WORKS the moment it is issued,
|
||||
// so issuing one to an unapproved node would route around the admin. A tunnel
|
||||
// credential is not:
|
||||
// core/http/endpoints/cluster/connect.go re-reads the node's status on every
|
||||
// dial and refuses a pending node with 403, so the credential is inert until an
|
||||
// admin approves and stays inert if approval is revoked. Withholding it would
|
||||
// instead strand every worker that registers exactly once (the static-NATS
|
||||
// path in core/services/worker/worker.go does), because approval alone does not
|
||||
// prompt a re-registration and nothing else can hand it the secret.
|
||||
//
|
||||
// BACKEND and AGENT nodes get one; anything else has its credential CLEARED.
|
||||
//
|
||||
// Both kinds of worker now dial out and hold one tunnel, and neither opens an
|
||||
// inbound port. A backend worker carries its gRPC backends and its file staging
|
||||
// on it; an agent worker carries only its own HTTP server, which is where its
|
||||
// MCP control verbs live. The gate was once backend-only, and correctly so:
|
||||
// nothing dialled into an agent worker, so a credential for it would have been
|
||||
// surface with no feature behind it. That stopped being true when the frontend
|
||||
// gained a reason to reach an agent worker by RPC.
|
||||
//
|
||||
// The gate lives HERE and not in ConnectHandler, which never looks at NodeType.
|
||||
// It does not need to, PROVIDED an ineligible node ends up with no credential
|
||||
// rather than merely being handed no new one, because the handler's empty-hash
|
||||
// branch is what does the refusing. So this CLEARS the column instead of
|
||||
// returning early, and the difference is not theoretical: Register upserts by
|
||||
// NAME, so a node re-registering under a different type keeps its ID, and
|
||||
// Register's struct Updates zero-skips TunnelTokenHash while writing the new
|
||||
// node_type. An early return left a live credential on a row whose type had
|
||||
// changed. Clearing is what makes "enforcement is structural" true, and it is
|
||||
// what the eligibility list widening rather than disappearing preserves: the
|
||||
// invariant is that a node's stored hash always matches the credential its
|
||||
// CURRENT type is entitled to, which for an ineligible type is none.
|
||||
//
|
||||
// It clears unconditionally rather than only when something is there, so the
|
||||
// invariant holds without depending on what the row happened to contain. The
|
||||
// cost is one UPDATE per registration of an ineligible node.
|
||||
//
|
||||
// No node type reaches the clearing branch through RegisterNodeEndpoint today,
|
||||
// which rejects any node_type that is neither backend nor agent. That is a
|
||||
// reason to keep the branch rather than to drop it: a row's node_type is also
|
||||
// written by older builds and will be written by future types, and the branch
|
||||
// is what makes adding one a decision about eligibility rather than a silent
|
||||
// grant.
|
||||
//
|
||||
// A failure to mint or to store is logged and the response goes out without the
|
||||
// token. Registration is what gets a worker into the cluster at all, and
|
||||
// failing it over a credential the worker does not need until it tunnels would
|
||||
// turn a tunnel problem into a node that cannot join. The worker sees no
|
||||
// tunnel_token, reports that it has no credential, and retries at its next
|
||||
// registration.
|
||||
func attachTunnelToken(ctx context.Context, response map[string]any, registry *nodes.NodeRegistry, node *nodes.BackendNode) {
|
||||
if node == nil {
|
||||
return
|
||||
}
|
||||
jwt, seed, err := natsCfg.MintWorkerJWT(node.ID, node.NodeType)
|
||||
if err != nil {
|
||||
xlog.Warn("Failed to mint NATS JWT for node", "node", node.Name, "id", node.ID, "error", err)
|
||||
if !tunnelEligible(node.NodeType) {
|
||||
// Cleared, not skipped. SetTunnelTokenHash writes the single column
|
||||
// directly rather than through a struct update, so unlike Register it
|
||||
// can write an empty value; see its doc.
|
||||
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
|
||||
}
|
||||
response["nats_jwt"] = jwt
|
||||
response["nats_user_seed"] = seed
|
||||
// crypto/rand.Text: at least 128 bits of randomness, no error to handle and
|
||||
// no length constant to get wrong.
|
||||
plaintext := rand.Text()
|
||||
sum := sha256.Sum256([]byte(plaintext))
|
||||
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"] = plaintext
|
||||
}
|
||||
|
||||
// tunnelEligible reports whether a node of this type holds a tunnel credential.
|
||||
//
|
||||
// One predicate rather than a condition written into attachTunnelToken, because
|
||||
// the mint branch and the clear branch are the two halves of ONE rule and must
|
||||
// not be able to disagree: written as two conditions, widening the mint without
|
||||
// widening the clear leaves a node type that is granted a credential and never
|
||||
// stripped of one, and widening the clear without the mint strips a node type
|
||||
// on every registration it makes. Neither has a symptom until a tunnel dial.
|
||||
func tunnelEligible(nodeType string) bool {
|
||||
return nodeType == nodes.NodeTypeBackend || nodeType == nodes.NodeTypeAgent
|
||||
}
|
||||
|
||||
// provisionAgentWorkerKey creates a dedicated user and API key for an agent worker node.
|
||||
@@ -634,12 +707,12 @@ func DeleteBackendOnNodeEndpoint(unloader nodes.NodeCommandSender) echo.HandlerF
|
||||
func ListBackendsOnNodeEndpoint(unloader nodes.NodeCommandSender, registry *nodes.NodeRegistry) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
nodeID := c.Param("id")
|
||||
// Agent-type workers don't run backends and never subscribe to the
|
||||
// nodes.<id>.backend.list NATS subject, so the request would hang
|
||||
// until timeout with "no responders". Their backend list is simply
|
||||
// empty. Mirror the aggregate-list guard in managers_distributed.go
|
||||
// (skip nodes whose NodeType is set and not "backend") so the
|
||||
// single-node and cluster-wide views stay consistent.
|
||||
// Agent-type workers don't run backends and mount no backend.list
|
||||
// route on the tunnel they hold, so asking one can only 404. Their
|
||||
// backend list is simply empty. Mirror the aggregate-list guard in
|
||||
// managers_distributed.go (skip nodes whose NodeType is set and not
|
||||
// "backend") so the single-node and cluster-wide views stay
|
||||
// consistent.
|
||||
if node, err := registry.Get(c.Request().Context(), nodeID); err == nil {
|
||||
if node.NodeType != "" && node.NodeType != nodes.NodeTypeBackend {
|
||||
return c.JSON(http.StatusOK, []messaging.NodeBackendInfo{})
|
||||
@@ -717,7 +790,7 @@ func DeleteModelOnNodeEndpoint(unloader nodes.NodeCommandSender, registry *nodes
|
||||
|
||||
// NodeBackendLogsListEndpoint proxies a request to a worker node's /v1/backend-logs
|
||||
// endpoint to list model IDs that have backend logs.
|
||||
func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken string) echo.HandlerFunc {
|
||||
func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken string, dialFor nodes.WorkerNetDialerFor) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
ctx := c.Request().Context()
|
||||
nodeID := c.Param("id")
|
||||
@@ -726,11 +799,11 @@ func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken
|
||||
return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found"))
|
||||
}
|
||||
|
||||
if node.HTTPAddress == "" {
|
||||
return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, "node has no HTTP address"))
|
||||
}
|
||||
|
||||
resp, err := proxyHTTPToWorker(node.HTTPAddress, "/v1/backend-logs", registrationToken)
|
||||
// No HTTPAddress guard: a tunnel-only worker reports none, and the
|
||||
// http stream tag ignores the target anyway. WorkerHTTPHost fills the
|
||||
// URL's host with something that identifies the node and resolves
|
||||
// nowhere; the tunnel decides where the bytes go.
|
||||
resp, err := proxyHTTPToWorker(ctx, dialFor, nodeID, nodes.WorkerHTTPHost(nodeID, node.HTTPAddress), "/v1/backend-logs", registrationToken)
|
||||
if err != nil {
|
||||
return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err)))
|
||||
}
|
||||
@@ -745,7 +818,7 @@ func NodeBackendLogsListEndpoint(registry *nodes.NodeRegistry, registrationToken
|
||||
|
||||
// NodeBackendLogsLinesEndpoint proxies a request to a worker node's
|
||||
// /v1/backend-logs/{modelId} endpoint to get buffered log lines.
|
||||
func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToken string) echo.HandlerFunc {
|
||||
func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToken string, dialFor nodes.WorkerNetDialerFor) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
ctx := c.Request().Context()
|
||||
nodeID := c.Param("id")
|
||||
@@ -756,12 +829,8 @@ func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToke
|
||||
return c.JSON(http.StatusNotFound, nodeError(http.StatusNotFound, "node not found"))
|
||||
}
|
||||
|
||||
if node.HTTPAddress == "" {
|
||||
return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, "node has no HTTP address"))
|
||||
}
|
||||
|
||||
path := "/v1/backend-logs/" + url.PathEscape(modelID)
|
||||
resp, err := proxyHTTPToWorker(node.HTTPAddress, path, registrationToken)
|
||||
resp, err := proxyHTTPToWorker(ctx, dialFor, nodeID, nodes.WorkerHTTPHost(nodeID, node.HTTPAddress), path, registrationToken)
|
||||
if err != nil {
|
||||
return c.JSON(http.StatusBadGateway, nodeError(http.StatusBadGateway, fmt.Sprintf("failed to reach worker: %v", err)))
|
||||
}
|
||||
@@ -776,7 +845,7 @@ func NodeBackendLogsLinesEndpoint(registry *nodes.NodeRegistry, registrationToke
|
||||
|
||||
// NodeBackendLogsWSEndpoint proxies a WebSocket connection to a worker node's
|
||||
// /v1/backend-logs/{modelId}/ws endpoint for real-time log streaming.
|
||||
func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken string) echo.HandlerFunc {
|
||||
func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken string, dialFor nodes.WorkerNetDialerFor) echo.HandlerFunc {
|
||||
browserUpgrader := websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool {
|
||||
origin := r.Header.Get("Origin")
|
||||
@@ -808,15 +877,41 @@ func NodeBackendLogsWSEndpoint(registry *nodes.NodeRegistry, registrationToken s
|
||||
return err
|
||||
}
|
||||
|
||||
// Dial the worker WebSocket
|
||||
workerURL := fmt.Sprintf("ws://%s/v1/backend-logs/%s/ws", node.HTTPAddress, url.PathEscape(modelID))
|
||||
// Dial the worker WebSocket over that worker's tunnel. The URL still
|
||||
// names the worker's registered address, for the Host header; the
|
||||
// NetDialContext below is what decides where the connection goes. A
|
||||
// missing dialer is a failure, not a direct dial: see
|
||||
// nodes.ErrNoWorkerDialer.
|
||||
workerURL := fmt.Sprintf("ws://%s/v1/backend-logs/%s/ws", nodes.WorkerHTTPHost(nodeID, node.HTTPAddress), url.PathEscape(modelID))
|
||||
workerHeaders := http.Header{}
|
||||
if registrationToken != "" {
|
||||
workerHeaders.Set("Authorization", "Bearer "+registrationToken)
|
||||
}
|
||||
|
||||
workerDialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second}
|
||||
workerWS, _, err := workerDialer.Dial(workerURL, workerHeaders)
|
||||
var workerDial func(ctx context.Context, network, addr string) (net.Conn, error)
|
||||
if dialFor != nil {
|
||||
workerDial = dialFor(nodeID)
|
||||
}
|
||||
if workerDial == nil {
|
||||
// A JSON body cannot be written here: the response writer was
|
||||
// hijacked by the upgrade above, so the status line is long gone
|
||||
// and the write lands nowhere. The browser has to be told the same
|
||||
// way every other failure past the upgrade tells it, with a close
|
||||
// frame, and the socket has to be closed or it leaks for the life
|
||||
// of the process.
|
||||
// Best-effort: the browser may already have gone, and there is
|
||||
// nothing left to report the failure to either way. The CLOSE is
|
||||
// what matters and it is unconditional.
|
||||
_ = browserWS.WriteMessage(websocket.CloseMessage,
|
||||
websocket.FormatCloseMessage(websocket.CloseInternalServerErr, "no route to worker"))
|
||||
_ = browserWS.Close()
|
||||
xlog.Error("Cannot stream backend logs: no way to reach the worker",
|
||||
"node", nodeID, "error", nodes.ErrNoWorkerDialer)
|
||||
return nil
|
||||
}
|
||||
|
||||
workerDialer := websocket.Dialer{HandshakeTimeout: 10 * time.Second, NetDialContext: workerDial}
|
||||
workerWS, _, err := workerDialer.DialContext(ctx, workerURL, workerHeaders)
|
||||
if err != nil {
|
||||
browserWS.WriteMessage(websocket.CloseMessage,
|
||||
websocket.FormatCloseMessage(websocket.CloseInternalServerErr, "failed to connect to worker"))
|
||||
@@ -1273,10 +1368,25 @@ func DeleteSchedulingEndpoint(registry *nodes.NodeRegistry) echo.HandlerFunc {
|
||||
}
|
||||
}
|
||||
|
||||
// proxyHTTPToWorker makes a GET request to a worker's HTTP server with bearer token auth.
|
||||
func proxyHTTPToWorker(httpAddress, path, token string) (*http.Response, error) {
|
||||
// proxyHTTPToWorker makes a GET request to a worker's HTTP server with bearer
|
||||
// token auth, over that worker's tunnel.
|
||||
//
|
||||
// The URL still names the worker's registered HTTP address, because that is
|
||||
// what the Host header and every error message should say; what it no longer
|
||||
// decides is where the bytes go. dialFor supplies the transport, and a nil one
|
||||
// is an error rather than a fall back to connecting to httpAddress: a worker
|
||||
// behind NAT has no address to connect to, and a direct dial is the bypass this
|
||||
// whole change removes.
|
||||
func proxyHTTPToWorker(ctx context.Context, dialFor nodes.WorkerNetDialerFor, nodeID, httpAddress, path, token string) (*http.Response, error) {
|
||||
if dialFor == nil {
|
||||
return nil, fmt.Errorf("reaching node %s: %w", nodeID, nodes.ErrNoWorkerDialer)
|
||||
}
|
||||
dial := dialFor(nodeID)
|
||||
if dial == nil {
|
||||
return nil, fmt.Errorf("reaching node %s: %w", nodeID, nodes.ErrNoWorkerDialer)
|
||||
}
|
||||
reqURL := fmt.Sprintf("http://%s%s", httpAddress, path)
|
||||
req, err := http.NewRequest("GET", reqURL, nil)
|
||||
req, err := http.NewRequestWithContext(ctx, "GET", reqURL, nil)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -1284,6 +1394,6 @@ func proxyHTTPToWorker(httpAddress, path, token string) (*http.Response, error)
|
||||
req.Header.Set("Authorization", "Bearer "+token)
|
||||
}
|
||||
|
||||
client := httpclient.NewWithTimeout(15 * time.Second)
|
||||
client := httpclient.NewWithTimeout(15*time.Second, httpclient.WithTransport(&http.Transport{DialContext: dial}))
|
||||
return client.Do(req)
|
||||
}
|
||||
@@ -42,8 +42,6 @@ func (s *stubNodeCommandSender) StopBackend(_, _ string) error { return nil }
|
||||
|
||||
func (s *stubNodeCommandSender) UnloadModelOnNode(_, _ string) error { return nil }
|
||||
|
||||
func (s *stubNodeCommandSender) PingNode(_ string) error { return nil }
|
||||
|
||||
var _ = Describe("ListBackendsOnNodeEndpoint", func() {
|
||||
var registry *nodes.NodeRegistry
|
||||
|
||||
|
||||
@@ -4,22 +4,28 @@ import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"crypto/subtle"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/labstack/echo/v4"
|
||||
"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"
|
||||
)
|
||||
|
||||
// hashOf is how the node registry stores a secret: hex-encoded SHA-256.
|
||||
func hashOf(secret string) string {
|
||||
sum := sha256.Sum256([]byte(secret))
|
||||
return hex.EncodeToString(sum[:])
|
||||
}
|
||||
|
||||
var _ = DescribeTable("token validation",
|
||||
func(expectedToken, providedToken string, wantMatch bool) {
|
||||
if expectedToken == "" {
|
||||
@@ -66,7 +72,7 @@ var _ = Describe("Node HTTP handlers", func() {
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsauth.Config{})
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "")
|
||||
Expect(handler(c)).To(Succeed())
|
||||
Expect(rec.Code).To(Equal(http.StatusCreated))
|
||||
|
||||
@@ -77,27 +83,252 @@ var _ = Describe("Node HTTP handlers", func() {
|
||||
Expect(resp["status"]).To(Equal(nodes.StatusHealthy))
|
||||
})
|
||||
|
||||
It("returns nats_jwt when account seed is configured", func() {
|
||||
akp, err := nkeys.CreateAccount()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
seed, err := akp.Seed()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
// register posts one registration and returns the decoded response.
|
||||
register := func(body string, expectedToken string, autoApprove bool) map[string]any {
|
||||
e := echo.New()
|
||||
body := `{"name":"worker-nats","address":"10.0.0.2:50051"}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
|
||||
req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
|
||||
natsCfg := natsauth.Config{AccountSeed: string(seed)}
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsCfg)
|
||||
handler := RegisterNodeEndpoint(registry, expectedToken, autoApprove, nil, "")
|
||||
ExpectWithOffset(1, handler(c)).To(Succeed())
|
||||
ExpectWithOffset(1, rec.Code).To(Equal(http.StatusCreated))
|
||||
|
||||
var resp map[string]any
|
||||
ExpectWithOffset(1, json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed())
|
||||
return resp
|
||||
}
|
||||
|
||||
It("mints a per-node tunnel credential and stores only its hash", func() {
|
||||
resp := register(`{"name":"worker-tunnel","address":"10.0.0.3:50051","token":"shared-registration-token"}`,
|
||||
"shared-registration-token", true)
|
||||
|
||||
plaintext, _ := resp["tunnel_token"].(string)
|
||||
Expect(plaintext).ToNot(BeEmpty())
|
||||
// Not the registration token. That is the whole point: a leaked
|
||||
// registration token plus a known node ID used to open a tunnel,
|
||||
// because the tunnel authenticated against the hash of exactly the
|
||||
// value every worker in the deployment holds.
|
||||
Expect(plaintext).ToNot(Equal("shared-registration-token"))
|
||||
|
||||
node, err := registry.Get(context.Background(), resp["id"].(string))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
// Stored as a hash, never as the secret.
|
||||
Expect(node.TunnelTokenHash).To(Equal(hashOf(plaintext)))
|
||||
Expect(node.TunnelTokenHash).ToNot(Equal(plaintext))
|
||||
// And it is a DIFFERENT column from the registration token's hash,
|
||||
// which is what the tunnel used to compare against.
|
||||
Expect(node.TunnelTokenHash).ToNot(Equal(node.TokenHash))
|
||||
Expect(node.TokenHash).To(Equal(hashOf("shared-registration-token")))
|
||||
|
||||
// The security property, which none of the above actually pins: a
|
||||
// second node registering with the SAME shared token gets a
|
||||
// DIFFERENT credential. Everything above is satisfied by a secret
|
||||
// derived deterministically from the registration token, which
|
||||
// would isolate nothing; a mutation that did exactly that passed
|
||||
// every assertion before this one.
|
||||
other := register(`{"name":"worker-tunnel-2","address":"10.0.0.3:50052","token":"shared-registration-token"}`,
|
||||
"shared-registration-token", true)
|
||||
Expect(other["tunnel_token"]).ToNot(Equal(plaintext))
|
||||
})
|
||||
|
||||
It("rotates the tunnel credential on every re-registration", func() {
|
||||
body := `{"name":"worker-rotate","address":"10.0.0.4:50051"}`
|
||||
first := register(body, "", true)
|
||||
second := register(body, "", true)
|
||||
|
||||
Expect(second["id"]).To(Equal(first["id"]), "re-registration must keep the node identity")
|
||||
firstToken := first["tunnel_token"].(string)
|
||||
secondToken := second["tunnel_token"].(string)
|
||||
// Only the hash is stored, so a re-registering worker cannot be told
|
||||
// the secret it already holds; the alternative to rotating would be
|
||||
// storing the plaintext.
|
||||
Expect(secondToken).ToNot(Equal(firstToken))
|
||||
|
||||
node, err := registry.Get(context.Background(), first["id"].(string))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(node.TunnelTokenHash).To(Equal(hashOf(secondToken)))
|
||||
Expect(node.TunnelTokenHash).ToNot(Equal(hashOf(firstToken)))
|
||||
})
|
||||
|
||||
It("issues a tunnel credential to a node still awaiting approval", func() {
|
||||
// Deliberately unlike the agent API key and the NATS JWT, which are
|
||||
// both withheld from a pending node. Those work the moment they are
|
||||
// issued; this one does not, because the tunnel endpoint re-reads
|
||||
// the node's status on every dial and refuses a pending node. A
|
||||
// worker that registers exactly once would otherwise never receive
|
||||
// one, since approval alone prompts no re-registration.
|
||||
first := register(`{"name":"worker-pending","address":"10.0.0.5:50051","token":"shared"}`, "shared", false)
|
||||
Expect(first["status"]).To(Equal(nodes.StatusPending))
|
||||
plaintext, _ := first["tunnel_token"].(string)
|
||||
Expect(plaintext).ToNot(BeEmpty())
|
||||
|
||||
// Non-empty alone does not pin per-node-ness, and a review's
|
||||
// variant of the "derived from the shared token" mutation stayed
|
||||
// green on exactly that gap. A pending node's credential has to be
|
||||
// as unpredictable and as per-node as an approved one's, since it
|
||||
// becomes live the moment an admin approves.
|
||||
Expect(plaintext).ToNot(Equal("shared"))
|
||||
node, err := registry.Get(context.Background(), first["id"].(string))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(node.TunnelTokenHash).To(Equal(hashOf(plaintext)))
|
||||
Expect(node.TunnelTokenHash).ToNot(Equal(node.TokenHash))
|
||||
|
||||
second := register(`{"name":"worker-pending-2","address":"10.0.0.5:50052","token":"shared"}`, "shared", false)
|
||||
Expect(second["status"]).To(Equal(nodes.StatusPending))
|
||||
Expect(second["tunnel_token"]).ToNot(Equal(plaintext))
|
||||
})
|
||||
|
||||
It("issues a tunnel credential to an agent node", func() {
|
||||
// An agent worker holds a tunnel too. It runs no backends and
|
||||
// stages no files, so what its tunnel carries is only its own HTTP
|
||||
// server, but the frontend reaches its control verbs over it and
|
||||
// therefore has to be able to dial it at all.
|
||||
//
|
||||
// The gate that used to refuse this was correct while nothing
|
||||
// dialled into an agent worker. Reopening it is deliberate, and the
|
||||
// spec below is what keeps its other half honest.
|
||||
resp := register(`{"name":"agent-1","node_type":"agent"}`, "", true)
|
||||
Expect(resp["node_type"]).To(Equal(nodes.NodeTypeAgent))
|
||||
|
||||
plaintext, _ := resp["tunnel_token"].(string)
|
||||
Expect(plaintext).ToNot(BeEmpty())
|
||||
|
||||
node, err := registry.Get(context.Background(), resp["id"].(string))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
// Stored as a hash, never as the secret, exactly as a backend
|
||||
// node's is; ConnectHandler compares against this column and does
|
||||
// not look at node_type at all.
|
||||
Expect(node.TunnelTokenHash).To(Equal(hashOf(plaintext)))
|
||||
Expect(node.TunnelTokenHash).ToNot(Equal(plaintext))
|
||||
|
||||
// Per-node, not derived from anything shared. A second agent
|
||||
// registering gets a different credential.
|
||||
other := register(`{"name":"agent-2","node_type":"agent"}`, "", true)
|
||||
Expect(other["tunnel_token"]).ToNot(Equal(plaintext))
|
||||
})
|
||||
|
||||
It("rewrites the hash, rather than leaving a stale one, when a node changes type", func() {
|
||||
// Register upserts BY NAME, so a node can change node_type in
|
||||
// place, and Register's struct Updates zero-skips TunnelTokenHash
|
||||
// while writing the new node_type. The invariant is that the stored
|
||||
// hash always matches the node's CURRENT credential: a node whose
|
||||
// type changed must not be left holding the one it was handed
|
||||
// under its old type, because that is a secret the worker still
|
||||
// knows and nothing would ever retire.
|
||||
//
|
||||
// This is the same shape as the Register-upserts-by-name hazard
|
||||
// already carried forward: a name is not an identity.
|
||||
backend := register(`{"name":"shifty","address":"10.0.0.7:50051"}`, "", true)
|
||||
// Read with a comma-ok rather than a bare assertion: a build that
|
||||
// issues no credential must fail this spec on the assertion below,
|
||||
// naming what it is missing, rather than panic on a nil interface.
|
||||
backendToken, _ := backend["tunnel_token"].(string)
|
||||
Expect(backendToken).ToNot(BeEmpty())
|
||||
|
||||
agent := register(`{"name":"shifty","node_type":"agent"}`, "", true)
|
||||
Expect(agent["id"]).To(Equal(backend["id"]), "re-registration must keep the node identity")
|
||||
Expect(agent["node_type"]).To(Equal(nodes.NodeTypeAgent))
|
||||
agentToken, _ := agent["tunnel_token"].(string)
|
||||
Expect(agentToken).ToNot(BeEmpty(),
|
||||
"the node changed type and was handed no credential, so its stored hash is whatever its previous type left behind")
|
||||
|
||||
node, err := registry.Get(context.Background(), backend["id"].(string))
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(node.TunnelTokenHash).To(Equal(hashOf(agentToken)))
|
||||
Expect(node.TunnelTokenHash).ToNot(Equal(hashOf(backendToken)),
|
||||
"the node kept the credential it earned under its previous type")
|
||||
})
|
||||
|
||||
It("clears the credential of a node whose type is entitled to none", func() {
|
||||
// The other half of the gate, and the half that makes enforcement
|
||||
// STRUCTURAL: ConnectHandler never looks at node_type, so what
|
||||
// refuses an ineligible node is its empty hash. Skipping the mint
|
||||
// would leave a live credential on the row.
|
||||
//
|
||||
// Driven through the registry rather than through
|
||||
// RegisterNodeEndpoint, which rejects any node_type that is neither
|
||||
// backend nor agent. That validation is exactly why the branch
|
||||
// cannot be reached from the wire today, and exactly why the branch
|
||||
// has to stay: a row's node_type is also written by other builds,
|
||||
// and adding a type must be a decision about eligibility rather
|
||||
// than a silent grant.
|
||||
node := &nodes.BackendNode{Name: "shifty-future", NodeType: nodes.NodeTypeBackend}
|
||||
Expect(registry.Register(context.Background(), node, true)).To(Succeed())
|
||||
response := map[string]any{}
|
||||
attachTunnelToken(context.Background(), response, registry, node)
|
||||
Expect(response).To(HaveKey("tunnel_token"))
|
||||
stored, err := registry.Get(context.Background(), node.ID)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(stored.TunnelTokenHash).ToNot(BeEmpty())
|
||||
|
||||
node.NodeType = "some-future-worker-kind"
|
||||
response = map[string]any{}
|
||||
attachTunnelToken(context.Background(), response, registry, node)
|
||||
|
||||
Expect(response).ToNot(HaveKey("tunnel_token"))
|
||||
stored, err = registry.Get(context.Background(), node.ID)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(stored.TunnelTokenHash).To(BeEmpty(),
|
||||
"an ineligible node kept a usable tunnel credential, so the mint-site gate is not structural")
|
||||
})
|
||||
|
||||
// The register path, and the approve path is asserted separately below:
|
||||
// the helper that attached a per-node broker credential had two call
|
||||
// sites, and one restored call site is exactly the regression a single
|
||||
// spec would miss.
|
||||
//
|
||||
// Asserted on the DECODED MAP with ToNot(HaveKey(...)), not on a struct.
|
||||
// A struct assertion would pass vacuously: the response is assembled as a
|
||||
// map[string]any, so the only way a key can be present is if something put
|
||||
// it there, and the only way to see that is to look at the keys.
|
||||
It("mints no broker credential into a registration response", func() {
|
||||
e := echo.New()
|
||||
body := `{"name":"worker-no-bus","address":"10.0.0.2:50051"}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
|
||||
req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "")
|
||||
Expect(handler(c)).To(Succeed())
|
||||
Expect(rec.Code).To(Equal(http.StatusCreated))
|
||||
|
||||
var resp map[string]any
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed())
|
||||
Expect(resp["nats_jwt"]).ToNot(BeEmpty())
|
||||
// The keys it DOES carry, so the two absences below are facts about
|
||||
// the response and not about a decode that produced an empty map.
|
||||
Expect(resp).To(HaveKey("id"))
|
||||
Expect(resp).To(HaveKey("tunnel_token"))
|
||||
Expect(resp).ToNot(HaveKey("nats_jwt"))
|
||||
Expect(resp).ToNot(HaveKey("nats_user_seed"))
|
||||
})
|
||||
|
||||
It("mints no broker credential into an approval response either", func() {
|
||||
node := &nodes.BackendNode{
|
||||
ID: uuid.New().String(),
|
||||
Name: "worker-approve-no-bus",
|
||||
Address: "10.0.0.3:50051",
|
||||
NodeType: nodes.NodeTypeBackend,
|
||||
}
|
||||
Expect(registry.Register(context.Background(), node, false)).To(Succeed())
|
||||
|
||||
e := echo.New()
|
||||
req := httptest.NewRequest(http.MethodPost, "/", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
c.SetParamNames("id")
|
||||
c.SetParamValues(node.ID)
|
||||
|
||||
Expect(ApproveNodeEndpoint(registry, nil, "")(c)).To(Succeed())
|
||||
Expect(rec.Code).To(Equal(http.StatusOK))
|
||||
|
||||
var resp map[string]any
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed())
|
||||
Expect(resp).To(HaveKey("id"))
|
||||
Expect(resp).ToNot(HaveKey("nats_jwt"))
|
||||
Expect(resp).ToNot(HaveKey("nats_user_seed"))
|
||||
})
|
||||
|
||||
It("returns 400 when name is missing", func() {
|
||||
@@ -108,7 +339,7 @@ var _ = Describe("Node HTTP handlers", func() {
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsauth.Config{})
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "")
|
||||
Expect(handler(c)).To(Succeed())
|
||||
Expect(rec.Code).To(Equal(http.StatusBadRequest))
|
||||
|
||||
@@ -128,7 +359,7 @@ var _ = Describe("Node HTTP handlers", func() {
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsauth.Config{})
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "")
|
||||
Expect(handler(c)).To(Succeed())
|
||||
Expect(rec.Code).To(Equal(http.StatusBadRequest))
|
||||
|
||||
@@ -139,7 +370,11 @@ var _ = Describe("Node HTTP handlers", func() {
|
||||
Expect(errObj["message"]).To(ContainSubstring("exceeds 255 characters"))
|
||||
})
|
||||
|
||||
It("returns 400 when address is missing for backend node type", func() {
|
||||
It("registers a backend worker that states no address", func() {
|
||||
// This used to be a 400. It is the shape every worker now
|
||||
// registers with: it has no inbound endpoint, it holds one outbound
|
||||
// tunnel, and refusing it here would refuse exactly the workers the
|
||||
// tunnel exists for.
|
||||
e := echo.New()
|
||||
body := `{"name":"worker-no-addr"}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
|
||||
@@ -147,15 +382,37 @@ var _ = Describe("Node HTTP handlers", func() {
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsauth.Config{})
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "")
|
||||
Expect(handler(c)).To(Succeed())
|
||||
Expect(rec.Code).To(Equal(http.StatusBadRequest))
|
||||
Expect(rec.Code).To(Equal(http.StatusCreated))
|
||||
|
||||
var resp map[string]any
|
||||
Expect(json.Unmarshal(rec.Body.Bytes(), &resp)).To(Succeed())
|
||||
errObj, ok := resp["error"].(map[string]any)
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(errObj["message"]).To(ContainSubstring("address is required"))
|
||||
stored, err := registry.GetByName(context.Background(), "worker-no-addr")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(stored.NodeType).To(Equal(nodes.NodeTypeBackend))
|
||||
Expect(stored.Address).To(BeEmpty())
|
||||
Expect(stored.HTTPAddress).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("stores no address even when a worker still sends one", func() {
|
||||
// An older worker keeps sending both keys. Storing them would put a
|
||||
// dialable-looking endpoint back into the API and the Nodes page for
|
||||
// something nothing dials, and would leave a reader of either one
|
||||
// unsure which workers are reached how.
|
||||
e := echo.New()
|
||||
body := `{"name":"worker-legacy-addr","address":"10.0.0.9:50051","http_address":"10.0.0.9:50050"}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body))
|
||||
req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "")
|
||||
Expect(handler(c)).To(Succeed())
|
||||
Expect(rec.Code).To(Equal(http.StatusCreated))
|
||||
|
||||
stored, err := registry.GetByName(context.Background(), "worker-legacy-addr")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(stored.Address).To(BeEmpty())
|
||||
Expect(stored.HTTPAddress).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("returns 400 when node_type is invalid", func() {
|
||||
@@ -166,7 +423,7 @@ var _ = Describe("Node HTTP handlers", func() {
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsauth.Config{})
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "")
|
||||
Expect(handler(c)).To(Succeed())
|
||||
Expect(rec.Code).To(Equal(http.StatusBadRequest))
|
||||
|
||||
@@ -185,7 +442,7 @@ var _ = Describe("Node HTTP handlers", func() {
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
|
||||
handler := RegisterNodeEndpoint(registry, "correct-token", true, nil, "", natsauth.Config{})
|
||||
handler := RegisterNodeEndpoint(registry, "correct-token", true, nil, "")
|
||||
Expect(handler(c)).To(Succeed())
|
||||
Expect(rec.Code).To(Equal(http.StatusUnauthorized))
|
||||
})
|
||||
@@ -198,7 +455,7 @@ var _ = Describe("Node HTTP handlers", func() {
|
||||
rec := httptest.NewRecorder()
|
||||
c := e.NewContext(req, rec)
|
||||
|
||||
handler := RegisterNodeEndpoint(registry, "", false, nil, "", natsauth.Config{})
|
||||
handler := RegisterNodeEndpoint(registry, "", false, nil, "")
|
||||
Expect(handler(c)).To(Succeed())
|
||||
Expect(rec.Code).To(Equal(http.StatusCreated))
|
||||
|
||||
@@ -221,7 +478,7 @@ var _ = Describe("Node HTTP handlers", func() {
|
||||
req1 := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body1))
|
||||
req1.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
|
||||
rec1 := httptest.NewRecorder()
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsauth.Config{})
|
||||
handler := RegisterNodeEndpoint(registry, "", true, nil, "")
|
||||
Expect(handler(e.NewContext(req1, rec1))).To(Succeed())
|
||||
Expect(rec1.Code).To(Equal(http.StatusCreated))
|
||||
|
||||
|
||||
@@ -11,7 +11,8 @@ import (
|
||||
)
|
||||
|
||||
// ToolExecutor abstracts MCP tool discovery and execution.
|
||||
// Implementations handle local (in-process sessions) vs distributed (NATS) modes.
|
||||
// Implementations handle local (in-process sessions) vs distributed (an agent
|
||||
// worker reached over its tunnel) modes.
|
||||
type ToolExecutor interface {
|
||||
// DiscoverTools returns the tool function schemas available from MCP servers.
|
||||
DiscoverTools(ctx context.Context) ([]functions.Function, error)
|
||||
@@ -58,28 +59,29 @@ func (e *LocalToolExecutor) HasTools() bool {
|
||||
return len(e.tools) > 0
|
||||
}
|
||||
|
||||
// DistributedToolExecutor routes tool operations through NATS to agent workers.
|
||||
// DistributedToolExecutor routes tool operations to agent workers over the
|
||||
// tunnels they hold.
|
||||
type DistributedToolExecutor struct {
|
||||
natsClient MCPNATSClient
|
||||
modelName string
|
||||
remote config.MCPGenericConfig[config.MCPRemoteServers]
|
||||
stdio config.MCPGenericConfig[config.MCPSTDIOServers]
|
||||
toolDefs []mcpRemote.MCPToolDef
|
||||
agent AgentControl
|
||||
modelName string
|
||||
remote config.MCPGenericConfig[config.MCPRemoteServers]
|
||||
stdio config.MCPGenericConfig[config.MCPSTDIOServers]
|
||||
toolDefs []mcpRemote.MCPToolDef
|
||||
}
|
||||
|
||||
// NewDistributedToolExecutor creates a ToolExecutor that routes through NATS.
|
||||
// It discovers tools immediately via a NATS request-reply to an agent worker.
|
||||
func NewDistributedToolExecutor(ctx context.Context, natsClient MCPNATSClient, modelName string,
|
||||
// NewDistributedToolExecutor creates a ToolExecutor that routes to agent
|
||||
// workers. It discovers tools immediately with a control RPC to one of them.
|
||||
func NewDistributedToolExecutor(ctx context.Context, agent AgentControl, modelName string,
|
||||
remote config.MCPGenericConfig[config.MCPRemoteServers],
|
||||
stdio config.MCPGenericConfig[config.MCPSTDIOServers],
|
||||
) *DistributedToolExecutor {
|
||||
e := &DistributedToolExecutor{
|
||||
natsClient: natsClient,
|
||||
modelName: modelName,
|
||||
remote: remote,
|
||||
stdio: stdio,
|
||||
agent: agent,
|
||||
modelName: modelName,
|
||||
remote: remote,
|
||||
stdio: stdio,
|
||||
}
|
||||
resp, err := DiscoverMCPToolsRemote(ctx, natsClient, modelName, remote, stdio)
|
||||
resp, err := DiscoverMCPToolsRemote(ctx, agent, modelName, remote, stdio)
|
||||
if err != nil {
|
||||
xlog.Error("Failed to discover MCP tools (distributed)", "error", err)
|
||||
} else if resp != nil {
|
||||
@@ -103,7 +105,7 @@ func (e *DistributedToolExecutor) IsTool(name string) bool {
|
||||
}
|
||||
|
||||
func (e *DistributedToolExecutor) ExecuteTool(ctx context.Context, toolName, arguments string) (string, error) {
|
||||
return ExecuteMCPToolCallRemote(ctx, e.natsClient, e.modelName, e.remote, e.stdio, toolName, arguments)
|
||||
return ExecuteMCPToolCallRemote(ctx, e.agent, e.modelName, e.remote, e.stdio, toolName, arguments)
|
||||
}
|
||||
|
||||
func (e *DistributedToolExecutor) HasTools() bool {
|
||||
@@ -111,15 +113,16 @@ func (e *DistributedToolExecutor) HasTools() bool {
|
||||
}
|
||||
|
||||
// NewToolExecutor creates the appropriate ToolExecutor based on the current mode.
|
||||
// When natsClient is non-nil, returns a DistributedToolExecutor that routes through NATS.
|
||||
// When natsClient is nil, creates local sessions and returns a LocalToolExecutor.
|
||||
func NewToolExecutor(ctx context.Context, natsClient MCPNATSClient, modelName string,
|
||||
// When agent is non-nil, returns a DistributedToolExecutor that reaches an agent
|
||||
// worker over its tunnel. When agent is nil, creates local sessions and returns
|
||||
// a LocalToolExecutor.
|
||||
func NewToolExecutor(ctx context.Context, agent AgentControl, modelName string,
|
||||
remote config.MCPGenericConfig[config.MCPRemoteServers],
|
||||
stdio config.MCPGenericConfig[config.MCPSTDIOServers],
|
||||
enabledServers []string,
|
||||
) ToolExecutor {
|
||||
if natsClient != nil {
|
||||
return NewDistributedToolExecutor(ctx, natsClient, modelName, remote, stdio)
|
||||
if agent != nil {
|
||||
return NewDistributedToolExecutor(ctx, agent, modelName, remote, stdio)
|
||||
}
|
||||
sessions, err := NamedSessionsFromMCPConfig(modelName, remote, stdio, enabledServers)
|
||||
if err != nil || len(sessions) == 0 {
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package mcp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
mcpRemote "github.com/mudler/LocalAI/core/services/mcp"
|
||||
)
|
||||
|
||||
// recordingAgent records the context each verb was given, so a spec can read
|
||||
// the budget the caller applied.
|
||||
//
|
||||
// The budget is invisible any other way. It travels as a context deadline, not
|
||||
// as anything on the wire, so the only place it can be observed is inside the
|
||||
// call; a spec that waited for it to expire would be a spec that sleeps for
|
||||
// minutes, and one that read a constant would pass with the constant unused.
|
||||
type recordingAgent struct {
|
||||
toolDeadline time.Time
|
||||
toolHasDL bool
|
||||
discoDeadline time.Time
|
||||
discoHasDL bool
|
||||
|
||||
toolReply *mcpRemote.MCPToolResponse
|
||||
discoReply *mcpRemote.MCPDiscoveryResponse
|
||||
err error
|
||||
}
|
||||
|
||||
func (a *recordingAgent) ExecuteMCPTool(ctx context.Context, _ mcpRemote.MCPToolRequest) (*mcpRemote.MCPToolResponse, error) {
|
||||
a.toolDeadline, a.toolHasDL = ctx.Deadline()
|
||||
if a.err != nil {
|
||||
return nil, a.err
|
||||
}
|
||||
return a.toolReply, nil
|
||||
}
|
||||
|
||||
func (a *recordingAgent) DiscoverMCPTools(ctx context.Context, _ mcpRemote.MCPDiscoveryRequest) (*mcpRemote.MCPDiscoveryResponse, error) {
|
||||
a.discoDeadline, a.discoHasDL = ctx.Deadline()
|
||||
if a.err != nil {
|
||||
return nil, a.err
|
||||
}
|
||||
return a.discoReply, nil
|
||||
}
|
||||
|
||||
var (
|
||||
noRemote = config.MCPGenericConfig[config.MCPRemoteServers]{}
|
||||
noStdio = config.MCPGenericConfig[config.MCPSTDIOServers]{}
|
||||
)
|
||||
|
||||
var _ = Describe("Routing MCP to an agent worker", func() {
|
||||
ctx := context.Background()
|
||||
|
||||
Describe("with no agent control client wired", func() {
|
||||
// A frontend in distributed mode with nothing to reach an agent worker
|
||||
// through must say so. The alternative shape, which this replaces, was
|
||||
// a nil-pointer dereference inside a chat request.
|
||||
It("names the missing client rather than dialling nothing", func() {
|
||||
_, err := ExecuteMCPToolCallRemote(ctx, nil, "m", noRemote, noStdio, "weather", "{}")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("agent control client"))
|
||||
})
|
||||
|
||||
It("names it for discovery too", func() {
|
||||
_, err := DiscoverMCPToolsRemote(ctx, nil, "m", noRemote, noStdio)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("agent control client"))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("the budget it gives the call", func() {
|
||||
// Two verbs, two constants, two assertions. The rule is written at both
|
||||
// sites, so a single spec would leave whichever site it did not cover
|
||||
// free to lose its deadline: a tool call whose worker went quiet would
|
||||
// then hold the caller until the tunnel's keepalive noticed.
|
||||
//
|
||||
// Asserted as a WINDOW around the deadline rather than an equality,
|
||||
// because the deadline is stamped from a clock this spec does not hold.
|
||||
// The window is far tighter than the difference between the two
|
||||
// constants, so a call given the wrong one still fails here.
|
||||
const slack = 5 * time.Second
|
||||
|
||||
It("bounds a tool call by the documented tool timeout", func() {
|
||||
agent := &recordingAgent{toolReply: &mcpRemote.MCPToolResponse{Result: "ok"}}
|
||||
start := time.Now()
|
||||
|
||||
out, err := ExecuteMCPToolCallRemote(ctx, agent, "m", noRemote, noStdio, "weather", `{"city":"London"}`)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(Equal("ok"))
|
||||
|
||||
Expect(agent.toolHasDL).To(BeTrue(), "the tool call was given no deadline at all")
|
||||
Expect(agent.toolDeadline).To(BeTemporally("~", start.Add(config.DefaultMCPToolTimeout), slack))
|
||||
})
|
||||
|
||||
It("bounds discovery by the documented discovery timeout", func() {
|
||||
agent := &recordingAgent{discoReply: &mcpRemote.MCPDiscoveryResponse{}}
|
||||
start := time.Now()
|
||||
|
||||
_, err := DiscoverMCPToolsRemote(ctx, agent, "m", noRemote, noStdio)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(agent.discoHasDL).To(BeTrue(), "discovery was given no deadline at all")
|
||||
Expect(agent.discoDeadline).To(BeTemporally("~", start.Add(config.DefaultMCPDiscoveryTimeout), slack))
|
||||
})
|
||||
|
||||
It("does not extend a caller's shorter deadline", func() {
|
||||
// The caller's own budget wins. context.WithTimeout keeps the
|
||||
// earlier of the two, and a hand-rolled deadline would not.
|
||||
short, cancel := context.WithTimeout(ctx, time.Second)
|
||||
defer cancel()
|
||||
agent := &recordingAgent{toolReply: &mcpRemote.MCPToolResponse{Result: "ok"}}
|
||||
|
||||
_, err := ExecuteMCPToolCallRemote(short, agent, "m", noRemote, noStdio, "weather", "{}")
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(agent.toolDeadline).To(BeTemporally("<", time.Now().Add(config.DefaultMCPToolTimeout)))
|
||||
})
|
||||
})
|
||||
|
||||
It("passes the agent's failure through with its identity intact", func() {
|
||||
// The classification is the control client's, and re-deciding it here
|
||||
// would be the same rule in two places. What this pins is that
|
||||
// wrapping it for a human does not hide it from errors.Is.
|
||||
boom := errors.New("the fleet said no")
|
||||
agent := &recordingAgent{err: boom}
|
||||
|
||||
_, err := ExecuteMCPToolCallRemote(ctx, agent, "m", noRemote, noStdio, "weather", "{}")
|
||||
Expect(err).To(MatchError(boom))
|
||||
_, err = DiscoverMCPToolsRemote(ctx, agent, "m", noRemote, noStdio)
|
||||
Expect(err).To(MatchError(boom))
|
||||
})
|
||||
|
||||
It("refuses tool arguments that are not JSON before reaching for a worker", func() {
|
||||
agent := &recordingAgent{toolReply: &mcpRemote.MCPToolResponse{Result: "ok"}}
|
||||
|
||||
_, err := ExecuteMCPToolCallRemote(ctx, agent, "m", noRemote, noStdio, "weather", "{not json")
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("invalid tool arguments JSON"))
|
||||
Expect(agent.toolHasDL).To(BeFalse(), "no worker should have been asked")
|
||||
})
|
||||
})
|
||||
@@ -16,7 +16,6 @@ import (
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
mcpRemote "github.com/mudler/LocalAI/core/services/mcp"
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
|
||||
"github.com/mudler/LocalAI/pkg/functions"
|
||||
"github.com/mudler/LocalAI/pkg/httpclient"
|
||||
@@ -100,9 +99,23 @@ var (
|
||||
client = mcp.NewClient(&mcp.Implementation{Name: "LocalAI", Version: "v1.0.0"}, nil)
|
||||
)
|
||||
|
||||
// MCPNATSClient is the interface for NATS request-reply operations needed by MCP routing.
|
||||
type MCPNATSClient interface {
|
||||
Request(subject string, data []byte, timeout time.Duration) ([]byte, error)
|
||||
// AgentControl is the frontend's port onto the MCP verbs an agent worker serves
|
||||
// over the tunnel it holds. *nodes.AgentControlClient is the only production
|
||||
// implementation, and it is what decides WHICH agent worker answers.
|
||||
//
|
||||
// The contract a caller here relies on, and the reason these functions do no
|
||||
// classification of their own: an implementation returns the worker's own Error
|
||||
// field AS A GO ERROR, so a nil error means the verb succeeded. Re-reading the
|
||||
// Error field at every call site would be the same rule written twice, and the
|
||||
// copy that gets forgotten is the one that reports a failed tool call as an
|
||||
// empty success.
|
||||
//
|
||||
// It replaces an interface over NATS request-reply. What that carried was a
|
||||
// subject and a queue group, which between them chose a worker; the choosing is
|
||||
// now a query and the carrying is an ordinary control RPC.
|
||||
type AgentControl interface {
|
||||
ExecuteMCPTool(ctx context.Context, req mcpRemote.MCPToolRequest) (*mcpRemote.MCPToolResponse, error)
|
||||
DiscoverMCPTools(ctx context.Context, req mcpRemote.MCPDiscoveryRequest) (*mcpRemote.MCPDiscoveryResponse, error)
|
||||
}
|
||||
|
||||
// MetadataKeyLocalAIAssistant is the request-metadata key the chat handler
|
||||
@@ -510,18 +523,28 @@ func ExecuteMCPToolCall(ctx context.Context, tools []MCPToolInfo, toolName strin
|
||||
return string(combined), nil
|
||||
}
|
||||
|
||||
// ExecuteMCPToolCallRemote routes an MCP tool execution request to an agent worker via NATS.
|
||||
// Used in distributed mode when the frontend doesn't hold MCP sessions locally.
|
||||
// ExecuteMCPToolCallRemote runs one MCP tool on an agent worker.
|
||||
//
|
||||
// Used in distributed mode, where the frontend holds no MCP sessions of its
|
||||
// own: an agent worker is what can create them (stdio servers under docker),
|
||||
// so the frontend serialises the model's MCP configuration and asks one.
|
||||
//
|
||||
// The budget is applied HERE, as a context deadline, and that is the whole of
|
||||
// the change in where it lives. On the bus it was the request-reply timeout,
|
||||
// which was the only thing bounding a worker that never answered; the control
|
||||
// RPC carries no deadline of its own (see nodes.ControlClient.clientFor), so
|
||||
// without this a tool call whose worker went quiet holds the caller until the
|
||||
// tunnel's own keepalive notices, which is far longer than any caller expects.
|
||||
func ExecuteMCPToolCallRemote(
|
||||
ctx context.Context,
|
||||
natsClient MCPNATSClient,
|
||||
agent AgentControl,
|
||||
modelName string,
|
||||
remote config.MCPGenericConfig[config.MCPRemoteServers],
|
||||
stdio config.MCPGenericConfig[config.MCPSTDIOServers],
|
||||
toolName, arguments string,
|
||||
) (string, error) {
|
||||
if natsClient == nil {
|
||||
return "", fmt.Errorf("NATS client not configured for distributed MCP")
|
||||
if agent == nil {
|
||||
return "", fmt.Errorf("no agent control client is configured for distributed MCP: this frontend cannot reach an agent worker to run tool %q", toolName)
|
||||
}
|
||||
|
||||
var args map[string]any
|
||||
@@ -531,63 +554,50 @@ func ExecuteMCPToolCallRemote(
|
||||
}
|
||||
}
|
||||
|
||||
req := mcpRemote.MCPToolRequest{
|
||||
ctx, cancel := context.WithTimeout(ctx, config.DefaultMCPToolTimeout)
|
||||
defer cancel()
|
||||
|
||||
resp, err := agent.ExecuteMCPTool(ctx, mcpRemote.MCPToolRequest{
|
||||
ModelName: modelName,
|
||||
ToolName: toolName,
|
||||
Arguments: args,
|
||||
RemoteServers: remote,
|
||||
StdioServers: stdio,
|
||||
}
|
||||
reqData, _ := json.Marshal(req)
|
||||
|
||||
replyData, err := natsClient.Request(messaging.SubjectMCPToolExecute, reqData, config.DefaultMCPToolTimeout)
|
||||
})
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("NATS MCP tool request failed: %w", err)
|
||||
}
|
||||
|
||||
var resp mcpRemote.MCPToolResponse
|
||||
if err := json.Unmarshal(replyData, &resp); err != nil {
|
||||
return "", fmt.Errorf("unmarshal MCP reply: %w", err)
|
||||
}
|
||||
if resp.Error != "" {
|
||||
return "", fmt.Errorf("remote MCP tool error: %s", resp.Error)
|
||||
return "", fmt.Errorf("the MCP tool call could not be run on an agent worker: %w", err)
|
||||
}
|
||||
return resp.Result, nil
|
||||
}
|
||||
|
||||
// DiscoverMCPToolsRemote routes an MCP discovery request to an agent worker via NATS.
|
||||
// Returns server info and tool function schemas from the remote worker.
|
||||
// DiscoverMCPToolsRemote asks an agent worker which MCP servers and tool
|
||||
// schemas a model's configuration reaches.
|
||||
func DiscoverMCPToolsRemote(
|
||||
ctx context.Context,
|
||||
natsClient MCPNATSClient,
|
||||
agent AgentControl,
|
||||
modelName string,
|
||||
remote config.MCPGenericConfig[config.MCPRemoteServers],
|
||||
stdio config.MCPGenericConfig[config.MCPSTDIOServers],
|
||||
) (*mcpRemote.MCPDiscoveryResponse, error) {
|
||||
if natsClient == nil {
|
||||
return nil, fmt.Errorf("NATS client not configured for distributed MCP")
|
||||
if agent == nil {
|
||||
return nil, fmt.Errorf("no agent control client is configured for distributed MCP: this frontend cannot reach an agent worker to discover the tools of model %q", modelName)
|
||||
}
|
||||
|
||||
req := mcpRemote.MCPDiscoveryRequest{
|
||||
// Its own budget, and its own constant. Discovery opens every configured
|
||||
// MCP server, which a tool call on an already-open session does not, so the
|
||||
// two are not the same wait and never were.
|
||||
ctx, cancel := context.WithTimeout(ctx, config.DefaultMCPDiscoveryTimeout)
|
||||
defer cancel()
|
||||
|
||||
resp, err := agent.DiscoverMCPTools(ctx, mcpRemote.MCPDiscoveryRequest{
|
||||
ModelName: modelName,
|
||||
RemoteServers: remote,
|
||||
StdioServers: stdio,
|
||||
}
|
||||
reqData, _ := json.Marshal(req)
|
||||
|
||||
replyData, err := natsClient.Request(messaging.SubjectMCPDiscovery, reqData, config.DefaultMCPDiscoveryTimeout)
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("NATS MCP discovery request failed: %w", err)
|
||||
return nil, fmt.Errorf("MCP discovery could not be run on an agent worker: %w", err)
|
||||
}
|
||||
|
||||
var resp mcpRemote.MCPDiscoveryResponse
|
||||
if err := json.Unmarshal(replyData, &resp); err != nil {
|
||||
return nil, fmt.Errorf("unmarshal MCP discovery reply: %w", err)
|
||||
}
|
||||
if resp.Error != "" {
|
||||
return nil, fmt.Errorf("remote MCP discovery error: %s", resp.Error)
|
||||
}
|
||||
return &resp, nil
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// ListMCPServers returns server info with tool, prompt, and resource names for each session.
|
||||
|
||||
@@ -159,7 +159,7 @@ func applyAutoparserOverride(
|
||||
// @Param request body schema.OpenAIRequest true "query params"
|
||||
// @Success 200 {object} schema.OpenAIResponse "Response"
|
||||
// @Router /v1/chat/completions [post]
|
||||
func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, startupOptions *config.ApplicationConfig, natsClient mcpTools.MCPNATSClient, assistantHolder *mcpTools.LocalAIAssistantHolder, compressor middleware.ChatCompressor) echo.HandlerFunc {
|
||||
func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, startupOptions *config.ApplicationConfig, agentControl mcpTools.AgentControl, assistantHolder *mcpTools.LocalAIAssistantHolder, compressor middleware.ChatCompressor) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
var textContentToReturn string
|
||||
id := uuid.New().String()
|
||||
@@ -260,7 +260,7 @@ func ChatEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator
|
||||
if (len(mcpServers) > 0 || mcpPromptName != "" || len(mcpResourceURIs) > 0) && (config.MCP.Servers != "" || config.MCP.Stdio != "") {
|
||||
remote, stdio, mcpErr := config.MCP.MCPConfigFromYAML()
|
||||
if mcpErr == nil {
|
||||
mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), natsClient, config.Name, remote, stdio, mcpServers)
|
||||
mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), agentControl, config.Name, remote, stdio, mcpServers)
|
||||
|
||||
// Prompt and resource injection (pre-processing step — resolves locally regardless of distributed mode)
|
||||
namedSessions, sessErr := mcpTools.NamedSessionsFromMCPConfig(config.Name, remote, stdio, mcpServers)
|
||||
|
||||
@@ -45,7 +45,7 @@ var classifierTestHistory = schema.Messages{
|
||||
|
||||
func classifierResultEvents(t *fakeTransport) []types.ClassifierResultEvent {
|
||||
var out []types.ClassifierResultEvent
|
||||
for _, e := range t.events {
|
||||
for _, e := range t.recordedEvents() {
|
||||
if ev, ok := e.(types.ClassifierResultEvent); ok {
|
||||
out = append(out, ev)
|
||||
}
|
||||
@@ -57,7 +57,7 @@ func classifierResultEvents(t *fakeTransport) []types.ClassifierResultEvent {
|
||||
// item — what a classifier response actually "spoke".
|
||||
func replyTexts(t *fakeTransport) []string {
|
||||
var out []string
|
||||
for _, e := range t.events {
|
||||
for _, e := range t.recordedEvents() {
|
||||
if ev, ok := e.(types.ResponseOutputTextDoneEvent); ok {
|
||||
out = append(out, ev.Text)
|
||||
}
|
||||
@@ -277,7 +277,7 @@ var _ = Describe("classifierRespond", func() {
|
||||
Expect(t.countEvents(types.ServerEventTypeResponseOutputTextDone)).To(Equal(1))
|
||||
Expect(t.countEvents(types.ServerEventTypeResponseFunctionCallArgumentsDone)).To(Equal(1))
|
||||
var fcArgs string
|
||||
for _, e := range t.events {
|
||||
for _, e := range t.recordedEvents() {
|
||||
if done, ok := e.(types.ResponseFunctionCallArgumentsDoneEvent); ok {
|
||||
fcArgs = done.Arguments
|
||||
}
|
||||
@@ -656,7 +656,7 @@ var _ = Describe("classifierRespond slot filling", func() {
|
||||
Expect(results[0].Arguments).To(MatchJSON(`{"direction":"up","distance":3,"units":"meters"}`))
|
||||
|
||||
var fcArgs string
|
||||
for _, e := range t.events {
|
||||
for _, e := range t.recordedEvents() {
|
||||
if done, ok := e.(types.ResponseFunctionCallArgumentsDoneEvent); ok {
|
||||
fcArgs = done.Arguments
|
||||
}
|
||||
@@ -695,7 +695,7 @@ var _ = Describe("classifierRespond slot filling", func() {
|
||||
|
||||
Expect(handled).To(BeTrue())
|
||||
var fcArgs string
|
||||
for _, e := range t.events {
|
||||
for _, e := range t.recordedEvents() {
|
||||
if done, ok := e.(types.ResponseFunctionCallArgumentsDoneEvent); ok {
|
||||
fcArgs = done.Arguments
|
||||
}
|
||||
|
||||
@@ -16,9 +16,19 @@ import (
|
||||
// fakeTransport records the server events and audio sent to a realtime client
|
||||
// so streaming behaviour can be asserted without a real WebSocket/WebRTC peer.
|
||||
// It is not a *WebRTCTransport, so handler code takes the WebSocket path.
|
||||
//
|
||||
// Every field is behind the mutex, and the events are read only through
|
||||
// recordedEvents. A real transport is written to by the response and turn
|
||||
// coordinators' goroutines while the spec goroutine reads what has arrived so
|
||||
// far, so a double that appended without a lock could not be driven the way
|
||||
// production drives it. Both fields are named with a `Log` suffix so a raw read
|
||||
// from another spec file does not compile; audioLog has no reader yet, and the
|
||||
// accessor for it is left to whichever spec first needs one, because an unread
|
||||
// one does not build.
|
||||
type fakeTransport struct {
|
||||
events []types.ServerEvent
|
||||
audio []fakeAudioChunk
|
||||
mu sync.Mutex
|
||||
eventLog []types.ServerEvent
|
||||
audioLog []fakeAudioChunk
|
||||
}
|
||||
|
||||
type fakeAudioChunk struct {
|
||||
@@ -27,23 +37,35 @@ type fakeAudioChunk struct {
|
||||
}
|
||||
|
||||
func (f *fakeTransport) SendEvent(e types.ServerEvent) error {
|
||||
f.events = append(f.events, e)
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.eventLog = append(f.eventLog, e)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeTransport) ReadEvent() ([]byte, error) { return nil, nil }
|
||||
|
||||
func (f *fakeTransport) SendAudio(_ context.Context, pcm []byte, sampleRate int) error {
|
||||
f.audio = append(f.audio, fakeAudioChunk{pcm: pcm, sampleRate: sampleRate})
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.audioLog = append(f.audioLog, fakeAudioChunk{pcm: pcm, sampleRate: sampleRate})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *fakeTransport) Close() error { return nil }
|
||||
|
||||
// recordedEvents returns a snapshot of the events sent so far. A COPY, because
|
||||
// the caller ranges over it while the coordinators may still be sending.
|
||||
func (f *fakeTransport) recordedEvents() []types.ServerEvent {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return append([]types.ServerEvent(nil), f.eventLog...)
|
||||
}
|
||||
|
||||
// countEvents returns how many recorded events have the given type.
|
||||
func (f *fakeTransport) countEvents(et types.ServerEventType) int {
|
||||
n := 0
|
||||
for _, e := range f.events {
|
||||
for _, e := range f.recordedEvents() {
|
||||
if e.ServerEventType() == et {
|
||||
n++
|
||||
}
|
||||
@@ -55,7 +77,7 @@ func (f *fakeTransport) countEvents(et types.ServerEventType) int {
|
||||
// delta event — i.e. the text streamed to the client as it is generated.
|
||||
func (f *fakeTransport) transcriptDeltaText() string {
|
||||
var b strings.Builder
|
||||
for _, e := range f.events {
|
||||
for _, e := range f.recordedEvents() {
|
||||
if d, ok := e.(types.ResponseOutputAudioTranscriptDeltaEvent); ok {
|
||||
b.WriteString(d.Delta)
|
||||
}
|
||||
|
||||
@@ -269,7 +269,7 @@ var _ = Describe("liveTurnState", func() {
|
||||
lts.drainEvents(1.0)
|
||||
|
||||
var got []types.ConversationItemInputAudioTranscriptionDeltaEvent
|
||||
for _, e := range ftr.events {
|
||||
for _, e := range ftr.recordedEvents() {
|
||||
if d, ok := e.(types.ConversationItemInputAudioTranscriptionDeltaEvent); ok {
|
||||
got = append(got, d)
|
||||
}
|
||||
@@ -335,7 +335,7 @@ var _ = Describe("commitUtteranceWithTranscript", func() {
|
||||
Expect(tr.countEvents(types.ServerEventTypeConversationItemInputAudioTranscriptionCompleted)).To(Equal(1))
|
||||
|
||||
var completed types.ConversationItemInputAudioTranscriptionCompletedEvent
|
||||
for _, e := range tr.events {
|
||||
for _, e := range tr.recordedEvents() {
|
||||
if c, ok := e.(types.ConversationItemInputAudioTranscriptionCompletedEvent); ok {
|
||||
completed = c
|
||||
}
|
||||
@@ -394,7 +394,7 @@ var _ = Describe("emitPrecomputedTranscription", func() {
|
||||
|
||||
Expect(tr.countEvents(types.ServerEventTypeConversationItemInputAudioTranscriptionDelta)).To(Equal(2), "empty deltas skipped")
|
||||
Expect(tr.countEvents(types.ServerEventTypeConversationItemInputAudioTranscriptionCompleted)).To(Equal(1))
|
||||
for _, e := range tr.events {
|
||||
for _, e := range tr.recordedEvents() {
|
||||
switch ev := e.(type) {
|
||||
case types.ConversationItemInputAudioTranscriptionDeltaEvent:
|
||||
Expect(ev.ItemID).To(Equal("item42"))
|
||||
|
||||
@@ -38,7 +38,7 @@ var _ = Describe("emitSoundDetection", func() {
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(t.countEvents(types.ServerEventTypeConversationItemSoundDetection)).To(Equal(1))
|
||||
|
||||
ev, ok := t.events[0].(types.ConversationItemSoundDetectionEvent)
|
||||
ev, ok := t.recordedEvents()[0].(types.ConversationItemSoundDetectionEvent)
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(ev.ItemID).To(Equal("item1"))
|
||||
Expect(ev.ContentIndex).To(Equal(0))
|
||||
@@ -62,7 +62,7 @@ var _ = Describe("emitSoundDetection", func() {
|
||||
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(t.countEvents(types.ServerEventTypeConversationItemSoundDetection)).To(Equal(1))
|
||||
ev, ok := t.events[0].(types.ConversationItemSoundDetectionEvent)
|
||||
ev, ok := t.recordedEvents()[0].(types.ConversationItemSoundDetectionEvent)
|
||||
Expect(ok).To(BeTrue())
|
||||
Expect(ev.Detections).To(BeEmpty())
|
||||
})
|
||||
|
||||
@@ -250,8 +250,8 @@ var _ = Describe("triggerResponse", func() {
|
||||
// The single terminal carries the produced output item and the usage —
|
||||
// both empty in the legacy code.
|
||||
var done *types.ResponseDoneEvent
|
||||
for i := range t.events {
|
||||
if d, ok := t.events[i].(types.ResponseDoneEvent); ok {
|
||||
for _, e := range t.recordedEvents() {
|
||||
if d, ok := e.(types.ResponseDoneEvent); ok {
|
||||
done = &d
|
||||
}
|
||||
}
|
||||
@@ -287,8 +287,8 @@ var _ = Describe("triggerResponse", func() {
|
||||
|
||||
var created *types.ResponseCreatedEvent
|
||||
var done *types.ResponseDoneEvent
|
||||
for i := range t.events {
|
||||
switch e := t.events[i].(type) {
|
||||
for _, sent := range t.recordedEvents() {
|
||||
switch e := sent.(type) {
|
||||
case types.ResponseCreatedEvent:
|
||||
created = &e
|
||||
case types.ResponseDoneEvent:
|
||||
@@ -317,8 +317,8 @@ var _ = Describe("triggerResponse", func() {
|
||||
|
||||
triggerResponse(context.Background(), session, &Conversation{}, t, nil)
|
||||
|
||||
for i := range t.events {
|
||||
if d, ok := t.events[i].(types.ResponseDoneEvent); ok {
|
||||
for _, e := range t.recordedEvents() {
|
||||
if d, ok := e.(types.ResponseDoneEvent); ok {
|
||||
Expect(d.Response.Metadata).To(BeEmpty())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -67,7 +67,7 @@ func itSession(gate *voiceGate) (*Session, *fakeModel) {
|
||||
// hasSpeakerNotAuthorized reports whether a speaker_not_authorized error event
|
||||
// was emitted to the client.
|
||||
func hasSpeakerNotAuthorized(tr *fakeTransport) bool {
|
||||
for _, e := range tr.events {
|
||||
for _, e := range tr.recordedEvents() {
|
||||
if ev, ok := e.(types.ErrorEvent); ok && ev.Error.Code == "speaker_not_authorized" {
|
||||
return true
|
||||
}
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
package openresponses
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/distributed"
|
||||
"github.com/mudler/LocalAI/core/services/syncstate"
|
||||
)
|
||||
|
||||
// responseMetadataStoreAdapter bridges the durable store onto the generic
|
||||
// syncstate.Store the responses.metadata map consumes. It mirrors
|
||||
// finetune.fineTuneStoreAdapter and agentpool.taskStoreAdapter, which is the
|
||||
// shape every other adopter already uses.
|
||||
//
|
||||
// It lives in this package rather than beside the store because syncedResponse
|
||||
// is unexported: the projection of a response that a peer may act on is defined
|
||||
// here and nowhere else, and the durable row carries it as opaque JSON so the
|
||||
// two cannot drift.
|
||||
type responseMetadataStoreAdapter struct {
|
||||
store *distributed.ResponseMetadataStore
|
||||
}
|
||||
|
||||
// compile-time assertion that the adapter satisfies the component's Store.
|
||||
var _ syncstate.Store[string, *syncedResponse] = (*responseMetadataStoreAdapter)(nil)
|
||||
|
||||
// List re-hydrates the map from the durable rows that are still live.
|
||||
//
|
||||
// A decode failure is returned rather than skipped. Skipping would drop exactly
|
||||
// one response with nothing failing anywhere, which is the invisible-404 this
|
||||
// whole store exists to prevent; returning the error leaves the map holding
|
||||
// whatever it already had, because syncstate.hydrate replaces nothing when the
|
||||
// source errors.
|
||||
func (a *responseMetadataStoreAdapter) List(ctx context.Context) ([]*syncedResponse, error) {
|
||||
records, err := a.store.ListUnexpired(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]*syncedResponse, 0, len(records))
|
||||
for i := range records {
|
||||
v := &syncedResponse{}
|
||||
if err := json.Unmarshal(records[i].PayloadJSON, v); err != nil {
|
||||
return nil, fmt.Errorf("decoding replicated response metadata %q: %w", records[i].ID, err)
|
||||
}
|
||||
out = append(out, v)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Upsert writes the whole projection as JSON, plus the columns a hydrate and a
|
||||
// purge filter on.
|
||||
//
|
||||
// ExpiresAt is lifted out of the payload into its own column because
|
||||
// ListUnexpired and PurgeExpired compare against it on the database clock.
|
||||
//
|
||||
// It is null whenever the deployment runs the default Open Responses TTL of 0,
|
||||
// which is the ordinary case and not an error: the store then falls back to
|
||||
// distributed.DefaultResponseMetadataRetention, so the row is still swept and
|
||||
// still drops out of a hydrate. What the column buys is the other direction. A
|
||||
// deployment that DOES configure a TTL gets that TTL honoured here, rather than
|
||||
// having its responses outlive the map they mirror or die before it.
|
||||
func (a *responseMetadataStoreAdapter) Upsert(ctx context.Context, v *syncedResponse) error {
|
||||
if v == nil {
|
||||
return fmt.Errorf("replicating response metadata: nil value")
|
||||
}
|
||||
payload, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return fmt.Errorf("encoding replicated response metadata %q: %w", v.ID, err)
|
||||
}
|
||||
return a.store.Upsert(ctx, &distributed.ResponseMetadataRecord{
|
||||
ID: v.ID,
|
||||
OwnerReplica: v.OwnerReplica,
|
||||
Owner: v.Owner,
|
||||
PayloadJSON: payload,
|
||||
ExpiresAt: v.ExpiresAt,
|
||||
})
|
||||
}
|
||||
|
||||
func (a *responseMetadataStoreAdapter) Delete(ctx context.Context, k string) error {
|
||||
return a.store.Delete(ctx, k)
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
package openresponses
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/distributed"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
var _ = Describe("responseMetadataStoreAdapter", func() {
|
||||
var (
|
||||
db *gorm.DB
|
||||
adapter *responseMetadataStoreAdapter
|
||||
ctx context.Context
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
ctx = context.Background()
|
||||
db = testutil.SetupTestDB()
|
||||
store, err := distributed.NewResponseMetadataStore(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
adapter = &responseMetadataStoreAdapter{store: store}
|
||||
})
|
||||
|
||||
It("round-trips every field of the replicated projection", func() {
|
||||
storedAt := time.Now().UTC().Truncate(time.Second)
|
||||
expiresAt := storedAt.Add(time.Hour)
|
||||
completedAt := storedAt.Unix()
|
||||
|
||||
in := &syncedResponse{
|
||||
ID: "resp_roundtrip",
|
||||
OwnerReplica: "replica-a",
|
||||
Owner: "user-1",
|
||||
Request: &schema.OpenResponsesRequest{
|
||||
Model: "test-model",
|
||||
Input: "Hello",
|
||||
},
|
||||
Response: &schema.ORResponseResource{
|
||||
ID: "resp_roundtrip",
|
||||
Object: "response",
|
||||
CreatedAt: storedAt.Unix(),
|
||||
CompletedAt: &completedAt,
|
||||
Status: schema.ORStatusCompleted,
|
||||
Model: "test-model",
|
||||
Output: []schema.ORItemField{
|
||||
{Type: "message", ID: "msg_roundtrip", Role: "assistant"},
|
||||
},
|
||||
},
|
||||
StoredAt: storedAt,
|
||||
ExpiresAt: &expiresAt,
|
||||
StreamEnabled: true,
|
||||
IsBackground: true,
|
||||
}
|
||||
Expect(adapter.Upsert(ctx, in)).To(Succeed())
|
||||
|
||||
out, err := adapter.List(ctx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(HaveLen(1))
|
||||
got := out[0]
|
||||
|
||||
// Asserted field by field rather than with one Equal on the struct: a
|
||||
// single Equal against a value whose Request is nil passes for an
|
||||
// adapter that dropped the field, and previous_response_id chaining is
|
||||
// exactly what that field feeds.
|
||||
Expect(got.ID).To(Equal(in.ID))
|
||||
Expect(got.OwnerReplica).To(Equal(in.OwnerReplica))
|
||||
Expect(got.Owner).To(Equal(in.Owner))
|
||||
Expect(got.Request).ToNot(BeNil())
|
||||
Expect(got.Request.Model).To(Equal("test-model"))
|
||||
Expect(got.Request.Input).To(Equal("Hello"))
|
||||
Expect(got.Response).ToNot(BeNil())
|
||||
Expect(got.Response.ID).To(Equal("resp_roundtrip"))
|
||||
Expect(got.Response.Status).To(Equal(schema.ORStatusCompleted))
|
||||
Expect(got.Response.Model).To(Equal("test-model"))
|
||||
Expect(got.Response.CompletedAt).ToNot(BeNil())
|
||||
Expect(*got.Response.CompletedAt).To(Equal(completedAt))
|
||||
Expect(got.Response.Output).To(HaveLen(1))
|
||||
Expect(got.Response.Output[0].ID).To(Equal("msg_roundtrip"))
|
||||
Expect(got.StoredAt.UTC()).To(BeTemporally("==", storedAt))
|
||||
Expect(got.ExpiresAt).ToNot(BeNil())
|
||||
Expect(got.ExpiresAt.UTC()).To(BeTemporally("==", expiresAt))
|
||||
Expect(got.StreamEnabled).To(BeTrue())
|
||||
Expect(got.IsBackground).To(BeTrue())
|
||||
})
|
||||
|
||||
It("lifts ExpiresAt into the row's own column", func() {
|
||||
expiresAt := time.Now().UTC().Add(time.Hour).Truncate(time.Second)
|
||||
Expect(adapter.Upsert(ctx, &syncedResponse{ID: "resp_expiry", ExpiresAt: &expiresAt})).To(Succeed())
|
||||
|
||||
// Read the column directly. The round-trip above is served entirely from
|
||||
// the JSON payload and stays green for an adapter that left this column
|
||||
// null, and a null here is what makes ListUnexpired and PurgeExpired
|
||||
// treat the response as immortal.
|
||||
var row distributed.ResponseMetadataRecord
|
||||
Expect(db.First(&row, "id = ?", "resp_expiry").Error).To(Succeed())
|
||||
Expect(row.ExpiresAt).ToNot(BeNil())
|
||||
Expect(row.ExpiresAt.UTC()).To(BeTemporally("==", expiresAt))
|
||||
})
|
||||
|
||||
It("copies the owner columns out of the payload so an operator can filter on them", func() {
|
||||
Expect(adapter.Upsert(ctx, &syncedResponse{
|
||||
ID: "resp_owner", OwnerReplica: "replica-b", Owner: "user-9",
|
||||
})).To(Succeed())
|
||||
|
||||
var row distributed.ResponseMetadataRecord
|
||||
Expect(db.First(&row, "id = ?", "resp_owner").Error).To(Succeed())
|
||||
Expect(row.OwnerReplica).To(Equal("replica-b"))
|
||||
Expect(row.Owner).To(Equal("user-9"))
|
||||
})
|
||||
|
||||
It("does not return a row whose expiry has passed", func() {
|
||||
past := time.Now().Add(-time.Hour)
|
||||
Expect(adapter.Upsert(ctx, &syncedResponse{ID: "resp_dead", ExpiresAt: &past})).To(Succeed())
|
||||
Expect(adapter.Upsert(ctx, &syncedResponse{ID: "resp_live"})).To(Succeed())
|
||||
|
||||
out, err := adapter.List(ctx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(HaveLen(1))
|
||||
Expect(out[0].ID).To(Equal("resp_live"))
|
||||
})
|
||||
|
||||
It("removes a deleted response from the durable rows", func() {
|
||||
Expect(adapter.Upsert(ctx, &syncedResponse{ID: "resp_gone"})).To(Succeed())
|
||||
Expect(adapter.Delete(ctx, "resp_gone")).To(Succeed())
|
||||
|
||||
out, err := adapter.List(ctx)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(out).To(BeEmpty())
|
||||
})
|
||||
|
||||
It("reports an unreachable database as an error and never as an empty list", func() {
|
||||
Expect(adapter.Upsert(ctx, &syncedResponse{ID: "resp_before_outage"})).To(Succeed())
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(sqlDB.Close()).To(Succeed())
|
||||
|
||||
// "No such response" and "the database could not be reached" are
|
||||
// different facts. An empty slice here would make a re-hydrate replace
|
||||
// the whole map with nothing, so a transient outage would 404 every
|
||||
// response on this replica.
|
||||
out, err := adapter.List(ctx)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(out).To(BeNil())
|
||||
})
|
||||
})
|
||||
@@ -31,7 +31,7 @@ import (
|
||||
// @Param request body schema.OpenResponsesRequest true "Request body"
|
||||
// @Success 200 {object} schema.ORResponseResource "Response"
|
||||
// @Router /v1/responses [post]
|
||||
func ResponsesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, natsClient mcpTools.MCPNATSClient) echo.HandlerFunc {
|
||||
func ResponsesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, evaluator *templates.Evaluator, appConfig *config.ApplicationConfig, agentControl mcpTools.AgentControl) echo.HandlerFunc {
|
||||
return func(c echo.Context) error {
|
||||
createdAt := time.Now().Unix()
|
||||
responseID := fmt.Sprintf("resp_%s", uuid.New().String())
|
||||
@@ -124,7 +124,7 @@ func ResponsesEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, eval
|
||||
if !hasMCPRequest {
|
||||
enabledServers = nil // backward compat: auto-activate all servers
|
||||
}
|
||||
mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), natsClient, cfg.Name, remote, stdio, enabledServers)
|
||||
mcpExecutor = mcpTools.NewToolExecutor(c.Request().Context(), agentControl, cfg.Name, remote, stdio, enabledServers)
|
||||
|
||||
// Prompt and resource injection (pre-processing step — resolves locally regardless of distributed mode)
|
||||
if hasMCPRequest {
|
||||
|
||||
@@ -56,11 +56,19 @@ type ResponseStore struct {
|
||||
// (see sync.go), which is how a standalone deployment keeps exactly the
|
||||
// previous process-local behaviour. Guarded by mu.
|
||||
synced *syncstate.SyncedMap[string, *syncedResponse]
|
||||
nats messaging.MessagingClient
|
||||
bus messaging.Broadcaster
|
||||
cancelSub messaging.Subscription
|
||||
replicaID string
|
||||
lifeCtx context.Context
|
||||
lifeCancel context.CancelFunc
|
||||
|
||||
// purgeWG tracks the durable-metadata purge sweep so Close cannot return
|
||||
// while it is still issuing statements. purgeTicks, when non-nil, replaces
|
||||
// that sweep's wall-clock ticker: only a spec sets it, so the sweep can be
|
||||
// driven one tick at a time instead of waited out or slept through.
|
||||
// Guarded by mu, and read once when the sweep starts.
|
||||
purgeWG sync.WaitGroup
|
||||
purgeTicks <-chan time.Time
|
||||
}
|
||||
|
||||
// StreamedEvent represents a buffered SSE event for streaming resume
|
||||
|
||||
@@ -7,13 +7,15 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/distributed"
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/syncstate"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
// syncStateName is the syncstate namespace for replicated response metadata.
|
||||
// It becomes the NATS subject "state.responses.metadata.delta".
|
||||
// It becomes the broadcast subject "state.responses-metadata.delta" (the '.' in
|
||||
// the name is sanitized to '-', so the subject keeps three tokens).
|
||||
const syncStateName = "responses.metadata"
|
||||
|
||||
// ErrResponseNotLocal is returned by the stream-resume accessors when the
|
||||
@@ -75,13 +77,29 @@ type responseCancelEvent struct {
|
||||
// - a wildcard subscription on the response-cancel subject, so a cancel that
|
||||
// lands on the wrong replica still reaches the context.CancelFunc.
|
||||
//
|
||||
// The SyncedMap has no durable Store: responses are ephemeral, TTL-bounded
|
||||
// state that today does not survive a process restart either, so peers converge
|
||||
// through deltas alone. A replica that joins later does not learn about
|
||||
// responses created before it started; that is the same visibility a client had
|
||||
// before this change and strictly better than the 404 it got from every peer.
|
||||
func (s *ResponseStore) EnableDistributed(ctx context.Context, nats messaging.MessagingClient, replicaID string) error {
|
||||
if nats == nil {
|
||||
// The SyncedMap is backed by a durable Store, and that is what makes a gap in
|
||||
// the carrier survivable. Both carriers deliver at most once to CONNECTED
|
||||
// subscribers and neither replays: a replica whose subscription was down while a
|
||||
// response was created never receives that delta, and with deltas as the only
|
||||
// convergence path it would answer 404 for that response forever while its peers
|
||||
// answered 200. The Store is what the reconnect re-hydrate and the periodic
|
||||
// reconcile actually read, so the gap closes rather than becoming permanent.
|
||||
//
|
||||
// What is NOT durable is unchanged and deliberate: the resume buffer and the
|
||||
// CancelFunc never leave the owning replica (see syncedResponse), so the write
|
||||
// through is one row per response state change, not one per generated token.
|
||||
//
|
||||
// The store is required rather than optional. This is called only from the
|
||||
// distributed branch of route registration, so a nil store there is a wiring bug
|
||||
// and not a deployment shape; tolerating it would silently restore the
|
||||
// deltas-only map this parameter exists to replace, and a non-nil interface
|
||||
// wrapping a nil pointer would instead surface as a panic on a request.
|
||||
func (s *ResponseStore) EnableDistributed(ctx context.Context, bus messaging.Broadcaster,
|
||||
replicaID string, store *distributed.ResponseMetadataStore) error {
|
||||
if store == nil {
|
||||
return errors.New("enabling cross-replica Open Responses: the store parameter is nil, so a reconnecting replica would have nothing to re-hydrate from")
|
||||
}
|
||||
if bus == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -91,9 +109,10 @@ func (s *ResponseStore) EnableDistributed(ctx context.Context, nats messaging.Me
|
||||
lifeCtx, lifeCancel := context.WithCancel(context.Background()) //#nosec G118 -- cancelled in Close()
|
||||
|
||||
synced := syncstate.New(syncstate.Config[string, *syncedResponse]{
|
||||
Name: syncStateName,
|
||||
Key: func(v *syncedResponse) string { return v.ID },
|
||||
Nats: nats,
|
||||
Name: syncStateName,
|
||||
Key: func(v *syncedResponse) string { return v.ID },
|
||||
Bus: bus,
|
||||
Store: &responseMetadataStoreAdapter{store: store},
|
||||
})
|
||||
if err := synced.Start(ctx); err != nil {
|
||||
lifeCancel()
|
||||
@@ -104,12 +123,12 @@ func (s *ResponseStore) EnableDistributed(ctx context.Context, nats messaging.Me
|
||||
// the store, so everything it reads has to be in place first.
|
||||
s.mu.Lock()
|
||||
s.replicaID = replicaID
|
||||
s.nats = nats
|
||||
s.bus = bus
|
||||
s.lifeCtx, s.lifeCancel = lifeCtx, lifeCancel
|
||||
s.synced = synced
|
||||
s.mu.Unlock()
|
||||
|
||||
sub, err := messaging.SubscribeJSON(nats, messaging.SubjectResponseCancelWildcard, s.applyRemoteCancel)
|
||||
sub, err := messaging.SubscribeJSON(bus, messaging.SubjectResponseCancelWildcard, s.applyRemoteCancel)
|
||||
if err != nil {
|
||||
if cerr := s.Close(); cerr != nil {
|
||||
xlog.Warn("failed to tear down response metadata sync after subscribe error", "error", cerr)
|
||||
@@ -121,10 +140,70 @@ func (s *ResponseStore) EnableDistributed(ctx context.Context, nats messaging.Me
|
||||
s.cancelSub = sub
|
||||
s.mu.Unlock()
|
||||
|
||||
s.startMetadataPurge(lifeCtx, store)
|
||||
|
||||
xlog.Info("Open Responses store replicating across replicas", "replica_id", replicaID)
|
||||
return nil
|
||||
}
|
||||
|
||||
// DefaultResponseMetadataPurgeInterval is how often a replica sweeps rows whose
|
||||
// TTL has passed out of the durable metadata table.
|
||||
//
|
||||
// Every replica runs the sweep; the DELETE is idempotent and rows are selected
|
||||
// on the database clock, so two replicas sweeping at once cost one extra
|
||||
// statement and never disagree about which rows are dead.
|
||||
const DefaultResponseMetadataPurgeInterval = 5 * time.Minute
|
||||
|
||||
// startMetadataPurge launches the sweep goroutine on the store's own lifetime
|
||||
// context, so Close stops it.
|
||||
//
|
||||
// purgeTicks is the seam that makes the sweep testable without a clock: a spec
|
||||
// installs its own channel and sends one tick, which is deterministic where a
|
||||
// five-minute ticker and a sleep are not.
|
||||
func (s *ResponseStore) startMetadataPurge(lifeCtx context.Context, store *distributed.ResponseMetadataStore) {
|
||||
s.mu.Lock()
|
||||
ticks := s.purgeTicks
|
||||
s.mu.Unlock()
|
||||
|
||||
var stop func()
|
||||
if ticks == nil {
|
||||
t := time.NewTicker(DefaultResponseMetadataPurgeInterval)
|
||||
ticks, stop = t.C, t.Stop
|
||||
}
|
||||
|
||||
s.purgeWG.Add(1)
|
||||
go s.purgeExpiredMetadata(lifeCtx, store, ticks, stop)
|
||||
}
|
||||
|
||||
// purgeExpiredMetadata sweeps expired rows until the store is closed.
|
||||
//
|
||||
// A failed sweep is logged and retried on the next tick rather than ending the
|
||||
// loop: an unreachable database is not evidence that there is nothing to purge,
|
||||
// and a loop that exited on the first transient error would leave the table
|
||||
// growing for the rest of the process's life.
|
||||
func (s *ResponseStore) purgeExpiredMetadata(ctx context.Context, store *distributed.ResponseMetadataStore,
|
||||
ticks <-chan time.Time, stop func()) {
|
||||
defer s.purgeWG.Done()
|
||||
if stop != nil {
|
||||
defer stop()
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticks:
|
||||
n, err := store.PurgeExpired(ctx)
|
||||
if err != nil {
|
||||
xlog.Warn("failed to purge expired Open Responses metadata", "error", err)
|
||||
continue
|
||||
}
|
||||
if n > 0 {
|
||||
xlog.Debug("Purged expired Open Responses metadata", "rows", n)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Close tears down the distributed wiring. It is idempotent so a test (or a
|
||||
// double shutdown) can call it more than once, and is a no-op for a standalone
|
||||
// store.
|
||||
@@ -141,6 +220,11 @@ func (s *ResponseStore) Close() error {
|
||||
if cancel != nil {
|
||||
cancel()
|
||||
}
|
||||
// Wait for the purge sweep before returning: a Close that raced its own
|
||||
// goroutine would let a sweep issue a statement against a database the
|
||||
// caller has already torn down, which a spec sees as a flake and a
|
||||
// deployment sees as an error line during shutdown.
|
||||
s.purgeWG.Wait()
|
||||
if sub != nil {
|
||||
if err := sub.Unsubscribe(); err != nil {
|
||||
return err
|
||||
@@ -162,14 +246,14 @@ func (s *ResponseStore) syncMap() *syncstate.SyncedMap[string, *syncedResponse]
|
||||
// distributed returns the replication handles as a consistent snapshot. Every
|
||||
// path that broadcasts reads them through here so a concurrent Close cannot be
|
||||
// observed half-applied. A nil map means standalone mode.
|
||||
func (s *ResponseStore) distributed() (*syncstate.SyncedMap[string, *syncedResponse], context.Context, messaging.MessagingClient, string) {
|
||||
func (s *ResponseStore) distributed() (*syncstate.SyncedMap[string, *syncedResponse], context.Context, messaging.Broadcaster, string) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
ctx := s.lifeCtx
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
return s.synced, ctx, s.nats, s.replicaID
|
||||
return s.synced, ctx, s.bus, s.replicaID
|
||||
}
|
||||
|
||||
// replicaIdentity returns this process's replica ID (empty in standalone mode).
|
||||
@@ -290,9 +374,9 @@ func (s *ResponseStore) delegateCancel(v *syncedResponse) (*schema.ORResponseRes
|
||||
return v.Response, nil
|
||||
}
|
||||
|
||||
m, ctx, nats, replicaID := s.distributed()
|
||||
if nats != nil {
|
||||
if err := nats.Publish(messaging.SubjectResponseCancel(v.ID),
|
||||
m, ctx, bus, replicaID := s.distributed()
|
||||
if bus != nil {
|
||||
if err := bus.Publish(messaging.SubjectResponseCancel(v.ID),
|
||||
responseCancelEvent{ResponseID: v.ID, Origin: replicaID}); err != nil {
|
||||
xlog.Warn("failed to broadcast Open Responses cancel", "response_id", v.ID, "error", err)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package openresponses
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/distributed"
|
||||
"github.com/mudler/LocalAI/core/services/pgbus"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// The same two-replica topology, on the carrier a deployment actually runs.
|
||||
//
|
||||
// The specs beside this one share ONE in-memory double, so a replica hears its
|
||||
// peer through a function call. Here each replica holds its own LISTEN
|
||||
// connection, which is the only arrangement in which a cancel can be published
|
||||
// on a channel nobody listened to, arrive after the handler that would have
|
||||
// applied it was closed, or be dropped for a subscriber that fell behind. A
|
||||
// double cannot fail any of those ways.
|
||||
var _ = Describe("ResponseStore cross-replica on the broadcast carrier", func() {
|
||||
var (
|
||||
ctx context.Context
|
||||
db *gorm.DB
|
||||
store *distributed.ResponseMetadataStore
|
||||
replicaA *ResponseStore
|
||||
replicaB *ResponseStore
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
ctx = context.Background()
|
||||
|
||||
var dsn string
|
||||
db, dsn = testutil.SetupTestDBWithDSN()
|
||||
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
|
||||
|
||||
newBus := func() *pgbus.Bus {
|
||||
b, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(b.Close)
|
||||
return b
|
||||
}
|
||||
|
||||
var err error
|
||||
store, err = distributed.NewResponseMetadataStore(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
replicaA = NewResponseStore(0)
|
||||
replicaB = NewResponseStore(0)
|
||||
Expect(replicaA.EnableDistributed(ctx, newBus(), "replica-a", store)).To(Succeed())
|
||||
Expect(replicaB.EnableDistributed(ctx, newBus(), "replica-b", store)).To(Succeed())
|
||||
})
|
||||
|
||||
AfterEach(func() {
|
||||
Expect(replicaA.Close()).To(Succeed())
|
||||
Expect(replicaB.Close()).To(Succeed())
|
||||
})
|
||||
|
||||
It("reaches the CancelFunc held by the owning replica over two LISTEN connections", func() {
|
||||
const id = "resp_cancel_pg"
|
||||
cancelled := make(chan struct{})
|
||||
replicaA.StoreBackground(id, &schema.OpenResponsesRequest{Model: "test-model"},
|
||||
&schema.ORResponseResource{
|
||||
ID: id, Object: "response", CreatedAt: time.Now().Unix(),
|
||||
Status: schema.ORStatusInProgress, Model: "test-model",
|
||||
}, func() { close(cancelled) }, false)
|
||||
|
||||
// Waited for, not assumed. On the real carrier the metadata delta is
|
||||
// asynchronous, so a cancel issued before it lands answers "not found"
|
||||
// even though the durable row exists: the SyncedMap reads its own
|
||||
// memory and goes to the table only on re-hydrate. The in-memory double
|
||||
// the sibling specs share delivers synchronously and hides that
|
||||
// entirely, which is why this spec exists.
|
||||
Eventually(func() error {
|
||||
_, err := replicaB.Get(id)
|
||||
return err
|
||||
}, "20s").Should(Succeed())
|
||||
|
||||
// The cancel lands on the replica that does NOT hold the CancelFunc.
|
||||
resp, err := replicaB.Cancel(id)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(resp.Status).To(Equal(schema.ORStatusCancelled))
|
||||
|
||||
Eventually(cancelled, "20s").Should(BeClosed())
|
||||
})
|
||||
|
||||
It("does not report a cancel that reached nobody as a cancel that was refused", func() {
|
||||
// The owner is gone, so nothing applies the broadcast. This carrier is
|
||||
// at-most-once with no replay, and there is no reply to wait for, so the
|
||||
// caller must still get a prompt terminal answer rather than an error
|
||||
// that reads as the generation having declined to stop.
|
||||
const id = "resp_dead_owner_pg"
|
||||
replicaA.StoreBackground(id, &schema.OpenResponsesRequest{Model: "test-model"},
|
||||
&schema.ORResponseResource{
|
||||
ID: id, Object: "response", CreatedAt: time.Now().Unix(),
|
||||
Status: schema.ORStatusInProgress, Model: "test-model",
|
||||
}, func() {}, false)
|
||||
Eventually(func() error {
|
||||
_, err := replicaB.Get(id)
|
||||
return err
|
||||
}, "20s").Should(Succeed())
|
||||
Expect(replicaA.Close()).To(Succeed())
|
||||
|
||||
resp, err := replicaB.Cancel(id)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(resp.Status).To(Equal(schema.ORStatusCancelled))
|
||||
})
|
||||
})
|
||||
@@ -6,9 +6,12 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/distributed"
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// These specs model the two-replica topology from issue #10993: two independent
|
||||
@@ -18,6 +21,8 @@ import (
|
||||
var _ = Describe("ResponseStore cross-replica", func() {
|
||||
var (
|
||||
bus *testutil.FakeBus
|
||||
db *gorm.DB
|
||||
store *distributed.ResponseMetadataStore
|
||||
replicaA *ResponseStore
|
||||
replicaB *ResponseStore
|
||||
ctx context.Context
|
||||
@@ -27,11 +32,18 @@ var _ = Describe("ResponseStore cross-replica", func() {
|
||||
ctx = context.Background()
|
||||
bus = testutil.NewFakeBus()
|
||||
|
||||
// One database, two replicas, exactly as a deployment has it: the
|
||||
// durable rows are shared and the in-memory maps are not.
|
||||
db = testutil.SetupTestDB()
|
||||
var err error
|
||||
store, err = distributed.NewResponseMetadataStore(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
replicaA = NewResponseStore(0)
|
||||
replicaB = NewResponseStore(0)
|
||||
|
||||
Expect(replicaA.EnableDistributed(ctx, bus, "replica-a")).To(Succeed())
|
||||
Expect(replicaB.EnableDistributed(ctx, bus, "replica-b")).To(Succeed())
|
||||
Expect(replicaA.EnableDistributed(ctx, bus, "replica-a", store)).To(Succeed())
|
||||
Expect(replicaB.EnableDistributed(ctx, bus, "replica-b", store)).To(Succeed())
|
||||
})
|
||||
|
||||
AfterEach(func() {
|
||||
@@ -190,6 +202,143 @@ var _ = Describe("ResponseStore cross-replica", func() {
|
||||
})
|
||||
})
|
||||
|
||||
Describe("re-hydrating after a gap in the carrier", func() {
|
||||
// Both carriers deliver at most once to CONNECTED subscribers and
|
||||
// neither replays. This is the only path by which a replica learns about
|
||||
// a response whose delta it never received, and it is what Task 11's
|
||||
// move of this family onto pgbus depends on.
|
||||
It("restores from the durable rows what a missed delta had removed from memory", func() {
|
||||
const id = "resp_rehydrate"
|
||||
replicaA.Store(id, &schema.OpenResponsesRequest{Model: "test-model"}, newResponse(id, schema.ORStatusCompleted))
|
||||
Expect(replicaB.Get(id)).ToNot(BeNil())
|
||||
|
||||
// Drop it from both maps' memory the way a peer delta would, which
|
||||
// leaves the durable row untouched: the apply path is memory-only.
|
||||
Expect(bus.Publish(messaging.SubjectSyncStateDelta(syncStateName),
|
||||
map[string]any{"op": "delete", "key": id})).To(Succeed())
|
||||
_, err := replicaB.Get(id)
|
||||
Expect(err).To(HaveOccurred())
|
||||
|
||||
bus.TriggerReconnect()
|
||||
|
||||
stored, err := replicaB.Get(id)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(stored).ToNot(BeNil())
|
||||
Expect(stored.Response.Status).To(Equal(schema.ORStatusCompleted))
|
||||
Expect(stored.Request).ToNot(BeNil())
|
||||
})
|
||||
|
||||
It("keeps what it already had when the durable source cannot be read", func() {
|
||||
const id = "resp_outage"
|
||||
replicaA.Store(id, &schema.OpenResponsesRequest{Model: "test-model"}, newResponse(id, schema.ORStatusCompleted))
|
||||
Expect(replicaB.Get(id)).ToNot(BeNil())
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(sqlDB.Close()).To(Succeed())
|
||||
|
||||
// An unreachable database is not the same fact as an empty table. A
|
||||
// re-hydrate that could not read must change nothing, or a transient
|
||||
// outage would 404 every response this replica knows about.
|
||||
bus.TriggerReconnect()
|
||||
|
||||
stored, err := replicaB.Get(id)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(stored).ToNot(BeNil())
|
||||
})
|
||||
})
|
||||
|
||||
Describe("purging expired durable metadata", func() {
|
||||
It("sweeps rows whose TTL has passed and stops when the store is closed", func() {
|
||||
ticks := make(chan time.Time)
|
||||
replica := NewResponseStore(0)
|
||||
replica.purgeTicks = ticks
|
||||
Expect(replica.EnableDistributed(ctx, testutil.NewFakeBus(), "replica-purge", store)).To(Succeed())
|
||||
|
||||
past := time.Now().Add(-time.Hour)
|
||||
Expect(store.Upsert(ctx, &distributed.ResponseMetadataRecord{
|
||||
ID: "resp_expired_row", PayloadJSON: []byte(`{"id":"resp_expired_row"}`), ExpiresAt: &past,
|
||||
})).To(Succeed())
|
||||
Expect(store.Upsert(ctx, &distributed.ResponseMetadataRecord{
|
||||
ID: "resp_live_row", PayloadJSON: []byte(`{"id":"resp_live_row"}`),
|
||||
})).To(Succeed())
|
||||
|
||||
// A non-blocking send that retries: if EnableDistributed never
|
||||
// started the sweep there is no receiver, and this fails by name
|
||||
// instead of hanging the suite on an unbuffered send.
|
||||
Eventually(ticks).Should(BeSent(time.Now()), "the purge sweep must be running to receive a tick")
|
||||
|
||||
Eventually(func() ([]string, error) {
|
||||
recs, err := store.ListUnexpired(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]string, 0, len(recs))
|
||||
for i := range recs {
|
||||
out = append(out, recs[i].ID)
|
||||
}
|
||||
return out, nil
|
||||
}).Should(ConsistOf("resp_live_row"))
|
||||
|
||||
var count int64
|
||||
Eventually(func() (int64, error) {
|
||||
err := db.Model(&distributed.ResponseMetadataRecord{}).
|
||||
Where("id = ?", "resp_expired_row").Count(&count).Error
|
||||
return count, err
|
||||
}).Should(BeZero())
|
||||
|
||||
// Close waits for the sweep, so a Close that returns is proof the
|
||||
// goroutine is gone. If it never exited this would hang rather than
|
||||
// pass, which is why nothing here polls for a flag.
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
defer close(done)
|
||||
Expect(replica.Close()).To(Succeed())
|
||||
}()
|
||||
Eventually(done).Should(BeClosed())
|
||||
})
|
||||
})
|
||||
|
||||
Describe("wiring", func() {
|
||||
It("publishes response metadata on the carrier it was handed, on the responses family's subject", func() {
|
||||
// S5 in the wiring table. EnableDistributed takes a
|
||||
// messaging.Broadcaster, and route registration hands it the
|
||||
// deployment's broadcast carrier; the subject is asserted by name
|
||||
// because every state.* family shares one LISTEN channel and the
|
||||
// subject is the only thing that separates them.
|
||||
const id = "resp_subject"
|
||||
replicaA.Store(id, &schema.OpenResponsesRequest{Model: "test-model"}, newResponse(id, schema.ORStatusCompleted))
|
||||
|
||||
Expect(bus.PublishCount(messaging.SubjectSyncStateDelta(syncStateName))).To(BeNumerically(">=", 1))
|
||||
Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("finetune.jobs"))).To(Equal(0))
|
||||
Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("agent.tasks"))).To(Equal(0))
|
||||
})
|
||||
|
||||
It("subscribes for delegated cancels on that same carrier", func() {
|
||||
// The second leg EnableDistributed wires. It rides the same carrier
|
||||
// as the metadata map, and a cancel that lands on the wrong replica
|
||||
// reaches nothing without it.
|
||||
solo := NewResponseStore(0)
|
||||
own := testutil.NewFakeBus()
|
||||
Expect(solo.EnableDistributed(ctx, own, "replica-solo", store)).To(Succeed())
|
||||
DeferCleanup(func() { Expect(solo.Close()).To(Succeed()) })
|
||||
|
||||
Expect(own.Subscribers()).To(BeNumerically(">=", 2),
|
||||
"one subscription for the metadata map and one for the cancel wildcard")
|
||||
})
|
||||
|
||||
It("refuses to enable replication without a durable store", func() {
|
||||
// A nil store here is a wiring bug, not a deployment shape: this is
|
||||
// reached only from the distributed branch of route registration.
|
||||
// Tolerating it would silently restore a deltas-only map, whose gap
|
||||
// is permanent.
|
||||
err := NewResponseStore(0).EnableDistributed(ctx, bus, "replica-c", nil)
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("store"))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("standalone mode", func() {
|
||||
It("keeps a store with no messaging client purely local", func() {
|
||||
standalone := NewResponseStore(0)
|
||||
|
||||
@@ -5,7 +5,7 @@ import { test, expect } from './coverage-fixtures.js'
|
||||
// - a delete affordance is present and goes through ConfirmDialog
|
||||
//
|
||||
// We mock the distributed-mode API so the tests can run against the
|
||||
// standalone ui-test-server without spinning up workers/NATS.
|
||||
// standalone ui-test-server without spinning up workers.
|
||||
|
||||
const NODE_ID = 'test-node-1'
|
||||
const NODE_NAME = 'worker-test'
|
||||
|
||||
@@ -62,3 +62,67 @@ test.describe('Nodes roster panels', () => {
|
||||
await expect(page.getByText('alpha')).toHaveCount(0)
|
||||
})
|
||||
})
|
||||
|
||||
test.describe('Nodes join command', () => {
|
||||
// The panel emits BOTH the backend and the agent join command from one
|
||||
// component. Neither worker kind dials a message bus any more: each holds one
|
||||
// outward tunnel to --register-to and takes every verb on it. A join command
|
||||
// carrying --nats-url would tell an operator to stand up, secure and pay for a
|
||||
// broker that nothing in the deployment connects to, which is the one way this
|
||||
// migration can still cost money after the code stopped using it.
|
||||
//
|
||||
// Asserted on the RENDERED command text rather than on the component's
|
||||
// variables, because the variables are what the fix deletes: a spec reading
|
||||
// them would stop compiling instead of failing, and a compile error is not
|
||||
// evidence about what an operator is shown.
|
||||
test('emits no bus flag for either worker kind', async ({ page }) => {
|
||||
await mockCluster(page, [])
|
||||
await page.goto('/app/nodes')
|
||||
|
||||
await page.getByRole('radio', { name: /^Backend$/ }).click()
|
||||
const backendCli = page.locator('.p2p-cmd pre').first()
|
||||
await expect(backendCli).toContainText('local-ai worker', { timeout: 15_000 })
|
||||
await expect(backendCli).not.toContainText('--nats-url')
|
||||
const backendDocker = page.locator('.p2p-cmd pre').nth(1)
|
||||
await expect(backendDocker).toContainText('LOCALAI_REGISTER_TO')
|
||||
await expect(backendDocker).not.toContainText('LOCALAI_NATS_URL')
|
||||
|
||||
// The agent tab is the one that regressed: it was the last surface still
|
||||
// emitting the flag, and it kept emitting it for two tasks after the agent
|
||||
// worker stopped dialling.
|
||||
await page.getByRole('radio', { name: /^Agent$/ }).click()
|
||||
const agentCli = page.locator('.p2p-cmd pre').first()
|
||||
await expect(agentCli).toContainText('local-ai agent-worker', { timeout: 15_000 })
|
||||
await expect(agentCli).toContainText('--register-to',
|
||||
)
|
||||
await expect(agentCli).not.toContainText('--nats-url')
|
||||
const agentDocker = page.locator('.p2p-cmd pre').nth(1)
|
||||
await expect(agentDocker).toContainText('LOCALAI_REGISTER_TO')
|
||||
await expect(agentDocker).not.toContainText('LOCALAI_NATS_URL')
|
||||
})
|
||||
|
||||
test('does not advertise flags the CLI does not have', async ({ page }) => {
|
||||
// The "How to Enable Distributed Mode" card renders ONLY on the disabled
|
||||
// state, which the page enters when /api/nodes answers 503. Mocking a
|
||||
// healthy cluster here would assert absence against a card that was never
|
||||
// on the page.
|
||||
await page.route('**/api/nodes', r => r.fulfill({ status: 503, contentType: 'application/json', body: '{}' }))
|
||||
await page.route('**/api/nodes/models', r => r.fulfill({ status: 503, contentType: 'application/json', body: '{}' }))
|
||||
await page.route('**/api/nodes/scheduling', r => r.fulfill({ status: 503, contentType: 'application/json', body: '{}' }))
|
||||
await page.goto('/app/nodes')
|
||||
|
||||
const card = page.locator('.p2p-enable')
|
||||
await expect(card).toBeVisible({ timeout: 15_000 })
|
||||
// --distributed-nats and --distributed-db were never real flags; a copied
|
||||
// command carrying them fails at kong before LocalAI does anything.
|
||||
await expect(card).not.toContainText('--distributed-nats')
|
||||
await expect(card).not.toContainText('--distributed-db')
|
||||
// Neither step tells an operator to point anything at a bus. The FRONTEND
|
||||
// command is asserted first and by itself: it is the one that used to carry
|
||||
// --nats-url as a required flag, so an operator following this card would
|
||||
// have stood a broker up before starting LocalAI at all.
|
||||
await expect(card.locator('.p2p-cmd pre').nth(0)).toContainText('--auth-database-url')
|
||||
await expect(card.locator('.p2p-cmd pre').nth(0)).not.toContainText('--nats-url')
|
||||
await expect(card.locator('.p2p-cmd pre').nth(1)).not.toContainText('--nats-url')
|
||||
})
|
||||
})
|
||||
@@ -19,7 +19,11 @@ export default function NodePanel({ node, models = [], onApprove, onDrain, onRes
|
||||
<div className="node-panel__id">
|
||||
<StatusPill status={node.status} />
|
||||
<span className="node-panel__name">{node.name}</span>
|
||||
<span className="cell-mono cell-muted">{node.address}</span>
|
||||
{/* A worker has no address to show: it holds an outbound tunnel and
|
||||
binds nothing routable. Its id is what identifies it in routing
|
||||
logs, so that is what an operator needs here. Pre-tunnel nodes
|
||||
may still carry an address until they re-register. */}
|
||||
<span className="cell-mono cell-muted">{node.address || node.id}</span>
|
||||
</div>
|
||||
<div className="node-panel__actions" onClick={(e) => e.stopPropagation()}>
|
||||
{node.status === 'pending' && (
|
||||
|
||||
@@ -78,7 +78,7 @@ export default function NodeDetail() {
|
||||
<PageHeader
|
||||
eyebrow={<a onClick={() => navigate('/app/nodes')} className="link-plain"><i className="fas fa-arrow-left icon-before" aria-hidden="true" />Cluster</a>}
|
||||
title={<><StatusPill status={node.status} /> {node.name}</>}
|
||||
supporting={node.address}
|
||||
supporting={node.address || node.id}
|
||||
actions={
|
||||
<>
|
||||
{node.status === 'draining'
|
||||
|
||||
@@ -67,14 +67,14 @@ function WorkerHintCard({ addToast, activeTab, hasWorkers }) {
|
||||
<div>
|
||||
<p className="form-label">CLI</p>
|
||||
<CommandBlock
|
||||
command={`local-ai ${workerCmd} \\\n --register-to "${frontendUrl}" \\\n --nats-url "nats://nats:4222" \\\n --registration-token "$LOCALAI_REGISTRATION_TOKEN"`}
|
||||
command={`local-ai ${workerCmd} \\\n --register-to "${frontendUrl}" \\\n --registration-token "$LOCALAI_REGISTRATION_TOKEN"`}
|
||||
addToast={addToast}
|
||||
/>
|
||||
</div>
|
||||
<div>
|
||||
<p className="form-label">Docker</p>
|
||||
<CommandBlock
|
||||
command={`docker run --net host ${flagsStr}\\\n -e LOCALAI_REGISTER_TO="${frontendUrl}" \\\n -e LOCALAI_NATS_URL="nats://nats:4222" \\\n -e LOCALAI_REGISTRATION_TOKEN="$TOKEN" \\\n ${dockerImage(option, dev)} ${workerCmd}`}
|
||||
command={`docker run --net host ${flagsStr}\\\n -e LOCALAI_REGISTER_TO="${frontendUrl}" \\\n -e LOCALAI_REGISTRATION_TOKEN="$TOKEN" \\\n ${dockerImage(option, dev)} ${workerCmd}`}
|
||||
addToast={addToast}
|
||||
/>
|
||||
</div>
|
||||
@@ -240,7 +240,7 @@ export default function Nodes() {
|
||||
<div className="flex-1">
|
||||
<p className="fw-medium mb-xs">Start LocalAI with distributed mode</p>
|
||||
<CommandBlock
|
||||
command={`local-ai run --distributed \\\n --distributed-db "postgres://user:pass@host/db" \\\n --distributed-nats "nats://host:4222"`}
|
||||
command={`local-ai run --distributed \\\n --auth-database-url "postgres://user:pass@host/db"`}
|
||||
addToast={addToast}
|
||||
/>
|
||||
</div>
|
||||
@@ -250,7 +250,7 @@ export default function Nodes() {
|
||||
<div className="flex-1">
|
||||
<p className="fw-medium mb-xs">Register backend nodes</p>
|
||||
<CommandBlock
|
||||
command={`local-ai worker \\\n --register-to "http://localai-host:8080" \\\n --nats-url "nats://nats:4222" \\\n --node-name "gpu-node-1"`}
|
||||
command={`local-ai worker \\\n --register-to "http://localai-host:8080" \\\n --node-name "gpu-node-1"`}
|
||||
addToast={addToast}
|
||||
/>
|
||||
</div>
|
||||
|
||||
+2
-1
@@ -16,7 +16,8 @@ export function percentColor(pct) {
|
||||
// milliseconds, regardless of its encoding. The agent SSE bridge emits the
|
||||
// json_message timestamp in three different shapes depending on deploy mode:
|
||||
// an RFC3339 string (standalone agent pool), Unix milliseconds (local
|
||||
// dispatcher), or Unix nanoseconds (older NATS path). A numeric value is
|
||||
// dispatcher), or Unix nanoseconds (releases before the tunnel migration,
|
||||
// which a rolling upgrade still has in flight). A numeric value is
|
||||
// classified by magnitude (s / ms / us / ns) so any of them yields a sane
|
||||
// epoch. Falls back to Date.now() for null/empty/unparseable input.
|
||||
export function normalizeTimestampMs(ts) {
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package routes
|
||||
|
||||
import (
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
|
||||
)
|
||||
|
||||
// mcpAgentControl returns the port the MCP endpoints reach an agent worker
|
||||
// through, or a nil interface in a deployment that has no agent workers to
|
||||
// reach.
|
||||
//
|
||||
// One function for four route files, and it exists because of the trap rather
|
||||
// than to save three lines. Every one of those endpoints decides between its
|
||||
// local and its distributed path by asking whether this value is nil, and a nil
|
||||
// *nodes.AgentControlClient assigned straight into the interface is NOT nil:
|
||||
// the interface carries a type, so the check passes, the distributed path is
|
||||
// taken, and every MCP request in a standalone deployment fails instead of
|
||||
// using the in-process sessions it has. Returning the untyped nil explicitly is
|
||||
// the only spelling that keeps that check meaning what its call sites think it
|
||||
// means.
|
||||
func mcpAgentControl(app *application.Application) mcpTools.AgentControl {
|
||||
return mcpAgentControlOf(app.Distributed())
|
||||
}
|
||||
|
||||
// mcpAgentControlOf is the half of the rule above that a spec can reach.
|
||||
//
|
||||
// It is separate because an *application.Application carries its distributed
|
||||
// services in an unexported field with no way in from outside the package, so
|
||||
// the typed-nil conversion could not otherwise be asserted at all, and it is
|
||||
// precisely the conversion that is easy to get wrong and silent when it is.
|
||||
func mcpAgentControlOf(d *application.DistributedServices) mcpTools.AgentControl {
|
||||
if d == nil || d.AgentControl == nil {
|
||||
return nil
|
||||
}
|
||||
return d.AgentControl
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package routes
|
||||
|
||||
import (
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
)
|
||||
|
||||
// Every MCP endpoint decides between its LOCAL sessions and an agent worker by
|
||||
// asking whether this value is nil. That check means what its call sites think
|
||||
// it means only if the nil that reaches it is an untyped one.
|
||||
var _ = Describe("the MCP endpoints' agent control port", func() {
|
||||
It("is nil in a deployment with no distributed services", func() {
|
||||
Expect(mcpAgentControlOf(nil)).To(BeNil())
|
||||
})
|
||||
|
||||
It("is nil, and not a non-nil interface holding a nil pointer, when nothing built the client", func() {
|
||||
// The trap, stated as an assertion. A nil *nodes.AgentControlClient
|
||||
// assigned straight into the interface compares NON-nil, so every MCP
|
||||
// request would take the distributed path and fail, in a deployment
|
||||
// that has perfectly good in-process sessions.
|
||||
Expect(mcpAgentControlOf(&application.DistributedServices{})).To(BeNil())
|
||||
})
|
||||
|
||||
It("carries the client through when there is one", func() {
|
||||
// The negative control: a function that returned nil unconditionally
|
||||
// would pass both assertions above and disable distributed MCP
|
||||
// entirely, with the local path silently taken instead.
|
||||
client := nodes.NewAgentControlClient(nil, nil)
|
||||
Expect(mcpAgentControlOf(&application.DistributedServices{AgentControl: client})).To(BeIdenticalTo(client))
|
||||
})
|
||||
})
|
||||
@@ -10,7 +10,6 @@ import (
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/anthropic"
|
||||
mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
|
||||
"github.com/mudler/LocalAI/core/http/middleware"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/routing/pii"
|
||||
@@ -25,17 +24,14 @@ func RegisterAnthropicRoutes(app *echo.Echo,
|
||||
application *application.Application,
|
||||
) {
|
||||
// Anthropic Messages API endpoint
|
||||
var natsClient mcpTools.MCPNATSClient
|
||||
if d := application.Distributed(); d != nil {
|
||||
natsClient = d.Nats
|
||||
}
|
||||
agentControl := mcpAgentControl(application)
|
||||
|
||||
messagesHandler := anthropic.MessagesEndpoint(
|
||||
application.ModelConfigLoader(),
|
||||
application.ModelLoader(),
|
||||
application.TemplatesEvaluator(),
|
||||
application.ApplicationConfig(),
|
||||
natsClient,
|
||||
agentControl,
|
||||
)
|
||||
|
||||
messagesMiddleware := []echo.MiddlewareFunc{
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
clusterep "github.com/mudler/LocalAI/core/http/endpoints/cluster"
|
||||
clustersvc "github.com/mudler/LocalAI/core/services/cluster"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/libp2p/go-yamux/v5"
|
||||
)
|
||||
|
||||
// RegisterClusterRoutes registers the replica-to-replica peer link. onPeer
|
||||
// receives every authenticated session; see clusterep.PeerHandler for what it
|
||||
// is expected to do with it.
|
||||
//
|
||||
// instances is what turns the ?id= on that route from a self-declared label
|
||||
// into a claim: the handler resolves the id to a replica row and checks the
|
||||
// dialler's own peer credential against the hash it publishes. It is a required
|
||||
// argument rather than an option, so a deployment cannot register this route
|
||||
// with nothing to verify against; a nil one makes the handler answer 503.
|
||||
//
|
||||
// The path is core/services/cluster's own constant, so the handler and the
|
||||
// dialler cannot be registered and dialled at different paths. That the path
|
||||
// also falls under auth.ClusterPathPrefix, and so bypasses the global session
|
||||
// middleware, is asserted by driving a request through that middleware in
|
||||
// core/http/endpoints/cluster/peer_test.go.
|
||||
//
|
||||
// The route carries no auth middleware: it authenticates itself against the
|
||||
// cluster token, because a peer replica has no session and no user.
|
||||
func RegisterClusterRoutes(e *echo.Echo, token string, instances *clustersvc.Registry, onPeer func(string, *yamux.Session)) {
|
||||
e.GET(clustersvc.PeerPath, clusterep.PeerHandler(token, instances, onPeer))
|
||||
}
|
||||
|
||||
// RegisterWorkerTunnelRoute registers the endpoint a worker dials to open its
|
||||
// tunnel. registry authenticates the dial against the node's own stored token;
|
||||
// tunnels is what the resulting session is attached to.
|
||||
//
|
||||
// Unlike the peer link this is registered in EVERY deployment, single-binary
|
||||
// ones included, and both arguments may be nil there. Two reasons. The handler
|
||||
// fails closed without a registry, since a token can only be checked against a
|
||||
// node row and there are none; and being registered unconditionally is what
|
||||
// puts the route in front of the route-coverage test under build tag `auth`,
|
||||
// which is the thing that holds the reject-before-upgrade rule in place. A
|
||||
// route registered only in distributed mode is invisible to that test.
|
||||
//
|
||||
// Like the peer link, it carries no auth middleware and derives its path from
|
||||
// core/services/cluster's own constant, so the handler and the worker's dialler
|
||||
// cannot end up on different paths.
|
||||
func RegisterWorkerTunnelRoute(e *echo.Echo, registry *nodes.NodeRegistry, tunnels *clustersvc.TunnelRegistry) {
|
||||
e.GET(clustersvc.ConnectPath, clusterep.ConnectHandler(registry, tunnels))
|
||||
}
|
||||
@@ -6,7 +6,6 @@ import (
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/localai"
|
||||
mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
|
||||
"github.com/mudler/LocalAI/core/http/middleware"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
compressionservice "github.com/mudler/LocalAI/core/services/compression"
|
||||
@@ -455,11 +454,8 @@ func RegisterLocalAIRoutes(router *echo.Echo,
|
||||
compressionservice.CounterFunc(tokens.CountMessages),
|
||||
compressionservice.NewInferenceSummarizer(cl, ml, appConfig),
|
||||
)
|
||||
var mcpNATS mcpTools.MCPNATSClient
|
||||
if d := app.Distributed(); d != nil {
|
||||
mcpNATS = d.Nats
|
||||
}
|
||||
mcpStreamHandler := localai.MCPEndpoint(cl, ml, evaluator, appConfig, mcpNATS, chatCompressor)
|
||||
agentControl := mcpAgentControl(app)
|
||||
mcpStreamHandler := localai.MCPEndpoint(cl, ml, evaluator, appConfig, agentControl, chatCompressor)
|
||||
mcpStreamMiddleware := []echo.MiddlewareFunc{
|
||||
requestExtractor.BuildFilteredFirstAvailableDefaultModel(config.BuildUsecaseFilterFn(config.FLAG_CHAT)),
|
||||
requestExtractor.SetModelAndConfig(func() schema.LocalAIRequest { return new(schema.OpenAIRequest) }),
|
||||
@@ -478,7 +474,7 @@ func RegisterLocalAIRoutes(router *echo.Echo,
|
||||
router.POST("/mcp/chat/completions", mcpStreamHandler, mcpStreamMiddleware...)
|
||||
|
||||
// MCP server listing endpoint
|
||||
router.GET("/v1/mcp/servers/:model", localai.MCPServersEndpoint(cl, appConfig, mcpNATS), mcpMw)
|
||||
router.GET("/v1/mcp/servers/:model", localai.MCPServersEndpoint(cl, appConfig, agentControl), mcpMw)
|
||||
|
||||
// MCP prompts endpoints
|
||||
router.GET("/v1/mcp/prompts/:model", localai.MCPPromptsEndpoint(cl, appConfig), mcpMw)
|
||||
|
||||
@@ -10,7 +10,6 @@ import (
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/localai"
|
||||
"github.com/mudler/LocalAI/core/services/galleryop"
|
||||
"github.com/mudler/LocalAI/core/services/nodes"
|
||||
"github.com/mudler/LocalAI/pkg/natsauth"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
@@ -36,7 +35,7 @@ func nodeReadyMiddleware(registry *nodes.NodeRegistry) echo.MiddlewareFunc {
|
||||
// token but do not verify per-node identity. A compromised worker can heartbeat/drain/
|
||||
// deregister other nodes. Future: issue per-node JWT at registration, validate node
|
||||
// identity on subsequent requests (compare :id param with token subject).
|
||||
func RegisterNodeSelfServiceRoutes(e *echo.Echo, registry *nodes.NodeRegistry, registrationToken string, autoApprove bool, authDB *gorm.DB, hmacSecret string, natsCfg natsauth.Config) {
|
||||
func RegisterNodeSelfServiceRoutes(e *echo.Echo, registry *nodes.NodeRegistry, registrationToken string, autoApprove bool, authDB *gorm.DB, hmacSecret string) {
|
||||
if registry == nil {
|
||||
return
|
||||
}
|
||||
@@ -45,7 +44,7 @@ func RegisterNodeSelfServiceRoutes(e *echo.Echo, registry *nodes.NodeRegistry, r
|
||||
tokenAuthMw := nodeTokenAuth(registrationToken)
|
||||
|
||||
node := e.Group("/api/node", readyMw, tokenAuthMw)
|
||||
node.POST("/register", localai.RegisterNodeEndpoint(registry, registrationToken, autoApprove, authDB, hmacSecret, natsCfg))
|
||||
node.POST("/register", localai.RegisterNodeEndpoint(registry, registrationToken, autoApprove, authDB, hmacSecret))
|
||||
node.POST("/:id/heartbeat", localai.HeartbeatEndpoint(registry))
|
||||
node.POST("/:id/drain", localai.DrainNodeEndpoint(registry))
|
||||
node.POST("/:id/resume", localai.ResumeNodeEndpoint(registry))
|
||||
@@ -61,7 +60,13 @@ func RegisterNodeSelfServiceRoutes(e *echo.Echo, registry *nodes.NodeRegistry, r
|
||||
// backend install path (POST /:id/backends/install). That handler enqueues a
|
||||
// ManagementOp on the gallery channel rather than blocking on a NATS reply, so
|
||||
// the browser gets HTTP 202 + jobID immediately instead of waiting up to 3 minutes.
|
||||
func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloader nodes.NodeCommandSender, galleryService *galleryop.GalleryService, opcache *galleryop.OpCache, appConfig *config.ApplicationConfig, adminMw echo.MiddlewareFunc, authDB *gorm.DB, hmacSecret string, registrationToken string, natsCfg natsauth.Config) {
|
||||
//
|
||||
// workerDialFor is how the log-proxy routes reach a worker's own HTTP server:
|
||||
// over the tunnel that worker holds, never by connecting to the address it
|
||||
// registered. It is nil outside distributed mode, and those two routes then
|
||||
// answer 502 rather than dialling, because a worker with no tunnel has nothing
|
||||
// for them to proxy to.
|
||||
func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloader nodes.NodeCommandSender, galleryService *galleryop.GalleryService, opcache *galleryop.OpCache, appConfig *config.ApplicationConfig, adminMw echo.MiddlewareFunc, authDB *gorm.DB, hmacSecret string, registrationToken string, workerDialFor nodes.WorkerNetDialerFor) {
|
||||
if registry == nil {
|
||||
return
|
||||
}
|
||||
@@ -85,7 +90,7 @@ func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloade
|
||||
admin.DELETE("/:id", localai.DeregisterNodeEndpoint(registry))
|
||||
admin.POST("/:id/drain", localai.DrainNodeEndpoint(registry))
|
||||
admin.POST("/:id/resume", localai.ResumeNodeEndpoint(registry))
|
||||
admin.POST("/:id/approve", localai.ApproveNodeEndpoint(registry, authDB, hmacSecret, natsCfg))
|
||||
admin.POST("/:id/approve", localai.ApproveNodeEndpoint(registry, authDB, hmacSecret))
|
||||
|
||||
// Backend management on workers
|
||||
admin.GET("/:id/backends", localai.ListBackendsOnNodeEndpoint(unloader, registry))
|
||||
@@ -101,8 +106,8 @@ func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloade
|
||||
admin.POST("/:id/models/delete", localai.DeleteModelOnNodeEndpoint(unloader, registry))
|
||||
|
||||
// Backend log streaming (proxied from worker HTTP server)
|
||||
admin.GET("/:id/backend-logs", localai.NodeBackendLogsListEndpoint(registry, registrationToken))
|
||||
admin.GET("/:id/backend-logs/:modelId", localai.NodeBackendLogsLinesEndpoint(registry, registrationToken))
|
||||
admin.GET("/:id/backend-logs", localai.NodeBackendLogsListEndpoint(registry, registrationToken, workerDialFor))
|
||||
admin.GET("/:id/backend-logs/:modelId", localai.NodeBackendLogsLinesEndpoint(registry, registrationToken, workerDialFor))
|
||||
|
||||
// Label management
|
||||
admin.GET("/:id/labels", localai.GetNodeLabelsEndpoint(registry))
|
||||
@@ -123,7 +128,7 @@ func RegisterNodeAdminRoutes(e *echo.Echo, registry *nodes.NodeRegistry, unloade
|
||||
admin.DELETE("/:id/vram-budget", localai.ResetVRAMBudgetEndpoint(registry))
|
||||
|
||||
// WebSocket proxy for real-time log streaming from workers
|
||||
e.GET("/ws/nodes/:id/backend-logs/:modelId", localai.NodeBackendLogsWSEndpoint(registry, registrationToken), readyMw, adminMw)
|
||||
e.GET("/ws/nodes/:id/backend-logs/:modelId", localai.NodeBackendLogsWSEndpoint(registry, registrationToken, workerDialFor), readyMw, adminMw)
|
||||
}
|
||||
|
||||
// nodeTokenAuth validates the registration token for node self-service endpoints.
|
||||
|
||||
@@ -5,7 +5,6 @@ import (
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/localai"
|
||||
mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/openai"
|
||||
"github.com/mudler/LocalAI/core/http/middleware"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
@@ -38,18 +37,15 @@ func RegisterOpenAIRoutes(app *echo.Echo,
|
||||
app.POST("/v1/realtime/transcription_session", openai.RealtimeTranscriptionSession(application), traceMiddleware)
|
||||
app.POST("/v1/realtime/calls", openai.RealtimeCalls(application), traceMiddleware)
|
||||
|
||||
// NATS client for distributed MCP tool routing (nil when not in distributed mode)
|
||||
var natsClient mcpTools.MCPNATSClient
|
||||
if d := application.Distributed(); d != nil {
|
||||
natsClient = d.Nats
|
||||
}
|
||||
// How the MCP endpoints reach an agent worker; nil outside distributed mode.
|
||||
agentControl := mcpAgentControl(application)
|
||||
|
||||
// chat
|
||||
chatCompressor := compressionservice.New(
|
||||
compressionservice.CounterFunc(tokens.CountMessages),
|
||||
compressionservice.NewInferenceSummarizer(application.ModelConfigLoader(), application.ModelLoader(), application.ApplicationConfig()),
|
||||
)
|
||||
chatHandler := openai.ChatEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig(), natsClient, application.LocalAIAssistant(), chatCompressor)
|
||||
chatHandler := openai.ChatEndpoint(application.ModelConfigLoader(), application.ModelLoader(), application.TemplatesEvaluator(), application.ApplicationConfig(), agentControl, application.LocalAIAssistant(), chatCompressor)
|
||||
chatMiddleware := []echo.MiddlewareFunc{
|
||||
nodeHeaderMiddleware,
|
||||
usageMiddleware,
|
||||
|
||||
@@ -1,14 +1,16 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/labstack/echo/v4"
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/config"
|
||||
localai "github.com/mudler/LocalAI/core/http/endpoints/localai"
|
||||
mcpTools "github.com/mudler/LocalAI/core/http/endpoints/mcp"
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/openresponses"
|
||||
"github.com/mudler/LocalAI/core/http/middleware"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/distributed"
|
||||
"github.com/mudler/xlog"
|
||||
)
|
||||
|
||||
@@ -16,18 +18,11 @@ func RegisterOpenResponsesRoutes(app *echo.Echo,
|
||||
re *middleware.RequestExtractor,
|
||||
application *application.Application) {
|
||||
|
||||
// NATS client for distributed MCP tool routing (nil when not in distributed mode)
|
||||
var natsClient mcpTools.MCPNATSClient
|
||||
// How the MCP endpoints reach an agent worker; nil outside distributed mode.
|
||||
agentControl := mcpAgentControl(application)
|
||||
if d := application.Distributed(); d != nil {
|
||||
natsClient = d.Nats
|
||||
|
||||
// Replicate response metadata across frontend replicas and subscribe to
|
||||
// delegated cancels. Without this a GET, a previous_response_id lookup or
|
||||
// a cancel that the load balancer sends to a replica other than the
|
||||
// creator 404s, and the cancel never reaches the CancelFunc (#10993).
|
||||
// Standalone deployments skip this entirely and stay process-local.
|
||||
if err := openresponses.GetGlobalStore().EnableDistributed(
|
||||
application.ApplicationConfig().Context, d.Nats, application.InstanceID()); err != nil {
|
||||
if err := enableDistributedResponses(application.ApplicationConfig().Context, d,
|
||||
openresponses.GetGlobalStore(), application.InstanceID()); err != nil {
|
||||
xlog.Error("Failed to enable cross-replica Open Responses store", "error", err)
|
||||
}
|
||||
}
|
||||
@@ -38,7 +33,7 @@ func RegisterOpenResponsesRoutes(app *echo.Echo,
|
||||
application.ModelLoader(),
|
||||
application.TemplatesEvaluator(),
|
||||
application.ApplicationConfig(),
|
||||
natsClient,
|
||||
agentControl,
|
||||
)
|
||||
|
||||
responsesMiddleware := []echo.MiddlewareFunc{
|
||||
@@ -85,3 +80,32 @@ func setOpenResponsesRequestContext(re *middleware.RequestExtractor) echo.Middle
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// enableDistributedResponses replicates response metadata across frontend
|
||||
// replicas and subscribes to delegated cancels. Without it a GET, a
|
||||
// previous_response_id lookup or a cancel that the load balancer sends to a
|
||||
// replica other than the creator 404s, and the cancel never reaches the
|
||||
// CancelFunc (#10993). Standalone deployments never reach here and stay
|
||||
// process-local.
|
||||
//
|
||||
// The durable store is what a replica re-hydrates from after its subscription
|
||||
// missed a delta; without it the same response_id answers 404 here and 200 on
|
||||
// the peer that created it, forever.
|
||||
//
|
||||
// A named function rather than a block inside route registration, and that is
|
||||
// the point of it. EnableDistributed takes a messaging.Broadcaster, as it must:
|
||||
// its own specs publish through a double. So handing it any carrier other than
|
||||
// the deployment's COMPILES and reddens nothing anywhere, and the only symptom
|
||||
// is a cancel that answers 404 on every replica but one. The broker client that
|
||||
// used to be the second carrier in scope is gone; what pins the choice is the
|
||||
// spec beside this file, which drives it from the OTHER carrier. Registering
|
||||
// routes needs a whole Application and therefore has no spec; this needs a
|
||||
// DistributedServices and a store, and therefore has one.
|
||||
func enableDistributedResponses(ctx context.Context, d *application.DistributedServices,
|
||||
store *openresponses.ResponseStore, replicaID string) error {
|
||||
var responseStore *distributed.ResponseMetadataStore
|
||||
if d.DistStores != nil {
|
||||
responseStore = d.DistStores.Responses
|
||||
}
|
||||
return store.EnableDistributed(ctx, d.Broadcast(), replicaID, responseStore)
|
||||
}
|
||||
@@ -0,0 +1,98 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package routes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"runtime"
|
||||
"time"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/application"
|
||||
"github.com/mudler/LocalAI/core/http/endpoints/openresponses"
|
||||
"github.com/mudler/LocalAI/core/schema"
|
||||
"github.com/mudler/LocalAI/core/services/distributed"
|
||||
"github.com/mudler/LocalAI/core/services/messaging"
|
||||
"github.com/mudler/LocalAI/core/services/pgbus"
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// Which carrier the Open Responses store is enabled on.
|
||||
//
|
||||
// EnableDistributed takes a messaging.Broadcaster, which it must: its own specs
|
||||
// publish through a double, and it cannot be made to name a concrete carrier
|
||||
// without dragging that dependency through the whole endpoint package. The
|
||||
// consequence is that handing it any carrier other than the deployment's
|
||||
// compiles and reddens nothing, and the only symptom is a cancel that answers
|
||||
// 404 on every replica but the creator. So it is pinned here, by watching what
|
||||
// actually arrives on the carrier: busB below IS the other carrier, which is
|
||||
// why this spec keeps its force now that the broker's client is gone.
|
||||
var _ = Describe("wiring the Open Responses store to a carrier", func() {
|
||||
var (
|
||||
ctx context.Context
|
||||
busA, busB *pgbus.Bus
|
||||
store *distributed.ResponseMetadataStore
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
if runtime.GOOS == "darwin" {
|
||||
Skip("testcontainers requires Docker, not available on macOS CI")
|
||||
}
|
||||
ctx = context.Background()
|
||||
|
||||
db, dsn := testutil.SetupTestDBWithDSN()
|
||||
Expect(pgbus.Migrate(ctx, db)).To(Succeed())
|
||||
newBus := func() *pgbus.Bus {
|
||||
b, err := pgbus.New(ctx, pgbus.Config{DSN: dsn, DB: db})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(b.Close)
|
||||
return b
|
||||
}
|
||||
busA, busB = newBus(), newBus()
|
||||
|
||||
var err error
|
||||
store, err = distributed.NewResponseMetadataStore(db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
})
|
||||
|
||||
It("enables it on the deployment's broadcast carrier and not on anything else it holds", func() {
|
||||
// A DistributedServices holding BOTH, exactly as a running deployment
|
||||
// does. That is what makes this an assertion about which one was
|
||||
// chosen rather than about there being one at all.
|
||||
d := &application.DistributedServices{
|
||||
Bus: busA,
|
||||
DistStores: &distributed.Stores{Responses: store},
|
||||
}
|
||||
|
||||
responses := openresponses.NewResponseStore(0)
|
||||
Expect(enableDistributedResponses(ctx, d, responses, "replica-a")).To(Succeed())
|
||||
DeferCleanup(func() { _ = responses.Close() })
|
||||
|
||||
// A peer's carrier sees the metadata this replica mirrors, which it can
|
||||
// only do if the store was enabled on the carrier and not on the other
|
||||
// thing DistributedServices is holding.
|
||||
mirrored := make(chan []byte, 8)
|
||||
_, err := busB.Subscribe(messaging.SubjectSyncStateDelta("responses.metadata"), func(data []byte) {
|
||||
mirrored <- data
|
||||
})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
const id = "resp_wiring"
|
||||
responses.StoreBackground(id, &schema.OpenResponsesRequest{Model: "test-model"},
|
||||
&schema.ORResponseResource{
|
||||
ID: id, Object: "response", CreatedAt: time.Now().Unix(),
|
||||
Status: schema.ORStatusInProgress, Model: "test-model",
|
||||
}, func() {}, false)
|
||||
|
||||
Eventually(mirrored, "20s").Should(Receive(ContainSubstring(id)))
|
||||
})
|
||||
|
||||
It("refuses to enable without the durable store a reconnecting replica re-hydrates from", func() {
|
||||
d := &application.DistributedServices{Bus: busA}
|
||||
|
||||
responses := openresponses.NewResponseStore(0)
|
||||
Expect(enableDistributedResponses(ctx, d, responses, "replica-a")).ToNot(Succeed())
|
||||
})
|
||||
})
|
||||
@@ -2,7 +2,9 @@ package advisorylock
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"runtime"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
@@ -14,6 +16,51 @@ import (
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
// alterThisDatabase applies a server-side setting to the database this handle is
|
||||
// actually connected to, and proves it landed.
|
||||
//
|
||||
// The name is read back from the connection rather than written as a literal.
|
||||
// The test helper hands each spec its own database on a shared server, so a
|
||||
// hard-coded name ALTERs a database this handle never touches: the statement
|
||||
// succeeds, the override does nothing, and the two specs below go green having
|
||||
// exercised none of the condition they exist for. They regress a model-load
|
||||
// advisory-lock wedge that has already shipped to production once, so a green
|
||||
// spec that proves nothing is the worst outcome available here.
|
||||
//
|
||||
// The read-back is the guard. Idle connections are dropped first so the next one
|
||||
// is opened fresh and inherits the new database-level default; SHOW then reports
|
||||
// what a waiter's own connection would inherit. If that ever stops matching, the
|
||||
// spec fails here rather than passing for the wrong reason.
|
||||
func alterThisDatabase(db *gorm.DB, setting, value string) {
|
||||
GinkgoHelper()
|
||||
|
||||
var name string
|
||||
Expect(db.Raw("SELECT current_database()").Scan(&name).Error).ToNot(HaveOccurred())
|
||||
Expect(name).ToNot(BeEmpty())
|
||||
|
||||
Expect(db.Exec(fmt.Sprintf("ALTER DATABASE %q SET %s = %s", name, setting, quoteLiteral(value))).Error).
|
||||
ToNot(HaveOccurred())
|
||||
|
||||
sqlDB, err := db.DB()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
// database/sql retains no idle connections at 0, closing the ones it is
|
||||
// already holding, so every connection after this point is opened fresh and
|
||||
// inherits the new database-level default.
|
||||
sqlDB.SetMaxIdleConns(0)
|
||||
|
||||
var applied string
|
||||
Expect(db.Raw("SHOW " + setting).Scan(&applied).Error).ToNot(HaveOccurred())
|
||||
Expect(applied).To(Equal(value),
|
||||
"the %s override did not reach the database this spec is holding (%s), so the spec below would pass without ever reproducing the condition it regresses",
|
||||
setting, name)
|
||||
}
|
||||
|
||||
// quoteLiteral wraps a settings value as a SQL string literal. The values here
|
||||
// are spec constants, so this only has to be correct, not hostile-input-proof.
|
||||
func quoteLiteral(v string) string {
|
||||
return "'" + strings.ReplaceAll(v, "'", "''") + "'"
|
||||
}
|
||||
|
||||
var _ = Describe("AdvisoryLock", func() {
|
||||
Context("PostgreSQL advisory locks", func() {
|
||||
var db *gorm.DB
|
||||
@@ -166,12 +213,7 @@ var _ = Describe("AdvisoryLock", func() {
|
||||
// blocked on pg_advisory_lock() is aborted by the server after this
|
||||
// window and surfaces SQLSTATE 55P03 ("canceling statement due to
|
||||
// lock timeout") to the caller instead of waiting for its turn.
|
||||
Expect(db.Exec("ALTER DATABASE testdb SET lock_timeout = '300ms'").Error).ToNot(HaveOccurred())
|
||||
sqlDB, err := db.DB()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
// Drop pooled connections so subsequent ones reconnect and inherit
|
||||
// the new database-level lock_timeout default.
|
||||
sqlDB.SetMaxIdleConns(0)
|
||||
alterThisDatabase(db, "lock_timeout", "300ms")
|
||||
|
||||
holding := make(chan struct{})
|
||||
released := make(chan struct{})
|
||||
@@ -214,12 +256,7 @@ var _ = Describe("AdvisoryLock", func() {
|
||||
// statement_timeout=60s; a cold model load holds the lock far longer,
|
||||
// so every concurrent caller died with SQLSTATE 57014 ("canceling
|
||||
// statement due to statement timeout") rather than waiting its turn.
|
||||
Expect(db.Exec("ALTER DATABASE testdb SET statement_timeout = '300ms'").Error).ToNot(HaveOccurred())
|
||||
sqlDB, err := db.DB()
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
// Drop pooled connections so subsequent ones reconnect and inherit
|
||||
// the new database-level statement_timeout default.
|
||||
sqlDB.SetMaxIdleConns(0)
|
||||
alterThisDatabase(db, "statement_timeout", "300ms")
|
||||
|
||||
holding := make(chan struct{})
|
||||
released := make(chan struct{})
|
||||
|
||||
@@ -46,15 +46,17 @@ type AgentJobService struct {
|
||||
evaluator *templates.Evaluator
|
||||
|
||||
// tasks is the cross-replica task store: an in-memory map kept consistent
|
||||
// across replicas via NATS, with read-through to the configured persister
|
||||
// across replicas over the deployment's fan-out carrier, with read-through
|
||||
// to the configured persister
|
||||
// (file in standalone, PostgreSQL in distributed). Unlike jobs - which already
|
||||
// converge via the dispatcher + DB read-through - tasks previously read
|
||||
// in-memory only, so ListTasks went stale on non-originating replicas.
|
||||
tasks *syncstate.SyncedMap[string, schema.Task]
|
||||
// taskNats is the distributed NATS client backing the tasks SyncedMap. It is
|
||||
// not available at construction time, so it is injected via SetTaskSyncNATS
|
||||
// during distributed wiring; nil keeps tasks in-memory-only (standalone).
|
||||
taskNats messaging.MessagingClient
|
||||
// taskBus is the deployment's broadcast carrier backing the tasks SyncedMap.
|
||||
// It is not available at construction time, so it is injected via
|
||||
// SetTaskSyncBus during distributed wiring; nil keeps tasks in-memory-only
|
||||
// (standalone).
|
||||
taskBus messaging.Broadcaster
|
||||
|
||||
// Storage (in-memory primary, persister for secondary persistence)
|
||||
jobs *xsync.SyncedMap[string, schema.Job]
|
||||
@@ -96,9 +98,18 @@ func (s *AgentJobService) SetDistributedBackends(dispatcher DistributedDispatche
|
||||
s.dispatcher = dispatcher
|
||||
}
|
||||
|
||||
// SetUserID sets the user ID for per-user scoping of DB queries.
|
||||
// SetUserID sets the user ID for per-user scoping of DB queries, and rebuilds
|
||||
// the tasks map so the id lands in its subject.
|
||||
//
|
||||
// The rebuild is what makes the two setters order-independent. Without it the
|
||||
// map keeps whichever tenant it was built with, so wiring that happened to call
|
||||
// SetTaskSyncBus first would publish this user's tasks on the CLUSTER-WIDE
|
||||
// subject and every other tenant would apply them. Like SetTaskSyncBus, this
|
||||
// is only ever called before Start / hydrate, while the map is still empty, so
|
||||
// rebuilding loses no state.
|
||||
func (s *AgentJobService) SetUserID(id string) {
|
||||
s.userID = id
|
||||
s.buildTasksMap()
|
||||
}
|
||||
|
||||
// SetDistributedJobStore sets the database-backed job store for persisting tasks/jobs.
|
||||
@@ -108,28 +119,39 @@ func (s *AgentJobService) SetDistributedJobStore(store *jobs.JobStore) {
|
||||
s.persister = &dbJobPersister{store: store}
|
||||
}
|
||||
|
||||
// SetTaskSyncNATS wires the distributed NATS client used to keep agent *tasks*
|
||||
// consistent across replicas (jobs already converge via the dispatcher + DB
|
||||
// read-through, so they are left untouched). The client is not available when the
|
||||
// service is constructed, so it is injected here during distributed wiring and the
|
||||
// tasks SyncedMap is rebuilt to pick it up. It is always called before Start /
|
||||
// SetTaskSyncBus wires the deployment's broadcast carrier used to keep agent
|
||||
// *tasks* consistent across replicas (jobs already converge via the dispatcher +
|
||||
// DB read-through, so they are left untouched). The carrier is not available when
|
||||
// the service is constructed, so it is injected here during distributed wiring and
|
||||
// the tasks SyncedMap is rebuilt to pick it up. It is always called before Start /
|
||||
// hydrate, while the map is still empty, so rebuilding loses no state. Passing nil
|
||||
// (standalone) keeps the map in-memory-only with no broadcast.
|
||||
func (s *AgentJobService) SetTaskSyncNATS(nats messaging.MessagingClient) {
|
||||
s.taskNats = nats
|
||||
//
|
||||
// The parameter is messaging.Broadcaster, so in distributed mode this family
|
||||
// travels on PostgreSQL LISTEN/NOTIFY. The name says Bus and not NATS because
|
||||
// the two are no longer the same thing and a stale name here would be the only
|
||||
// documentation a wiring site reads.
|
||||
func (s *AgentJobService) SetTaskSyncBus(bus messaging.Broadcaster) {
|
||||
s.taskBus = bus
|
||||
s.buildTasksMap()
|
||||
}
|
||||
|
||||
// buildTasksMap (re)constructs the cross-replica tasks SyncedMap from the current
|
||||
// taskNats. The Store adapter reads s.persister/s.userID live, so a persister swap
|
||||
// (SetDistributedJobStore) needs no rebuild; only the NATS client, fixed at
|
||||
// New-time, forces one - hence SetTaskSyncNATS calls this.
|
||||
// taskBus. The Store adapter reads s.persister/s.userID live, so a persister swap
|
||||
// (SetDistributedJobStore) needs no rebuild; only the carrier, fixed at map-build
|
||||
// time, forces one - hence SetTaskSyncBus calls this.
|
||||
func (s *AgentJobService) buildTasksMap() {
|
||||
s.tasks = syncstate.New(syncstate.Config[string, schema.Task]{
|
||||
Name: "agent.tasks",
|
||||
Key: func(t schema.Task) string { return t.ID },
|
||||
Nats: s.taskNats,
|
||||
Bus: s.taskBus,
|
||||
Store: &taskStoreAdapter{svc: s},
|
||||
// There is one AgentJobService per user, so this map is per-tenant and
|
||||
// its deltas must not reach another tenant's copy. The empty userID is
|
||||
// the cluster-wide administrative service, which hydrates from every
|
||||
// tenant's rows and therefore also applies every tenant's deltas.
|
||||
PerTenant: true,
|
||||
Tenant: s.userID,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -261,7 +283,7 @@ func NewAgentJobServiceWithPaths(
|
||||
cronEntries: xsync.NewSyncedMap[string, cron.EntryID](),
|
||||
retentionDays: retentionDays,
|
||||
}
|
||||
// Build the cross-replica tasks map standalone (nil NATS); SetTaskSyncNATS
|
||||
// Build the cross-replica tasks map standalone (nil carrier); SetTaskSyncBus
|
||||
// rebuilds it with the distributed client once that is available, before Start.
|
||||
s.buildTasksMap()
|
||||
return s
|
||||
|
||||
Loaded 100 of 416 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user