diff --git a/.agents/building-and-testing.md b/.agents/building-and-testing.md index 021d555ec993..eee5123ac82c 100644 --- a/.agents/building-and-testing.md +++ b/.agents/building-and-testing.md @@ -45,3 +45,23 @@ 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, NATS), split by a Ginkgo label: + +- `make test-e2e-distributed` runs `Distributed && !VLLMMultinode && !Cluster` over `./tests/e2e/distributed` recursively. Services are wired directly into the test binary. ~240 specs, ~75s. +- `make test-e2e-cluster` runs `Cluster` and spawns real `local-ai` child processes through the `tests/e2e/distributed/cluster` helper package. 6 specs, about 8m30s measured over three consecutive runs (509.1s / 509.8s / 512.3s, so 8m29s to 8m32s). + +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. + +- **Containers are suite-scoped, not spec-scoped.** `SetupInfra` used to start a PostgreSQL (~10s) and a NATS (~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`.) Containers now start once in `BeforeSuite` and each spec gets 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:** the NATS bus is now *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 pair, keeping NATS subjects isolated per process. A single shared NATS 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, no PostgreSQL and no NATS, 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). +- **`--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. +- **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. diff --git a/.agents/ci-caching.md b/.agents/ci-caching.md index 6742049e68ff..8dc243747384 100644 --- a/.agents/ci-caching.md +++ b/.agents/ci-caching.md @@ -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 diff --git a/.github/workflows/tests-e2e-distributed.yml b/.github/workflows/tests-e2e-distributed.yml new file mode 100644 index 000000000000..8f17e72f169a --- /dev/null +++ b/.github/workflows/tests-e2e-distributed.yml @@ -0,0 +1,204 @@ +--- +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. These two are the only images + # the suite needs once the testcontainers reaper is disabled below. + run: | + docker pull postgres:16-alpine + docker pull nats:2-alpine + - name: Distributed E2E + # TESTCONTAINERS_RYUK_DISABLED keeps the pre-pull above meaningful. The + # reaper exists to clean up leaked containers on a long-lived host, but + # this runner is ephemeral and every container dies with the VM. Leaving + # it enabled would pull 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 and NATS 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 'stub\n' > core/http/react-ui/dist/index.html + - name: Build local-ai + # Not `make build`: that target pulls in the React UI build. The specs + # exec this binary directly via LOCALAI_E2E_BINARY. + run: | + PATH="$PATH:$HOME/go/bin" go build -o local-ai ./cmd/local-ai + - 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 + docker pull nats:2-alpine + - name: Cluster E2E + env: + LOCALAI_E2E_BINARY: ${{ github.workspace }}/local-ai + # 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 diff --git a/.golangci.yml b/.golangci.yml index d25d1ccb4789..7a68e69a75f2 100644 --- a/.golangci.yml +++ b/.golangci.yml @@ -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] diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index d87db37eae63..f48c7c4a337c 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -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, PostgreSQL and NATS) has two suites. Both bring up their PostgreSQL and NATS with testcontainers, so Docker has to be available: + +```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 (around 240 specs in roughly 75 seconds). It starts one PostgreSQL and one NATS 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 8m30s (measured 509.1s / 509.8s / 512.3s over three consecutive runs): three of its six specs wait out real staleness and health-check windows. 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. diff --git a/Makefile b/Makefile index ebedb2c98248..e45930585049 100644 --- a/Makefile +++ b/Makefile @@ -340,12 +340,59 @@ 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 +# 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 + NATS via testcontainers). # Includes NatsJWT specs (JWT-enabled NATS). Requires Docker. # 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". 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 + $(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 + +# Cluster e2e: runs local-ai as real child processes (frontend replicas + +# workers) against PostgreSQL and NATS, and kills them to assert failover. +# Needs a built ./local-ai (or LOCALAI_E2E_BINARY) plus the mock backend. +# +# The argument-validation specs in tests/e2e/distributed/cluster deliberately +# stay in test-e2e-distributed above: they need no binary, no PostgreSQL and no +# NATS, 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: 20 specs, measured at 800 to 830 seconds of Ginkgo time (13 to 14 +# minutes wall including the compile) on a fast developer box. It was 591 to 612 +# seconds before the phase 3 control-plane specs; those five added roughly 200 +# seconds, most of it in the two that wait out real windows rather than poll for +# a state change (cluster.InstanceLiveness is 30s, and a departed worker cannot +# be demoted before its reconnect grace). +# +# --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. +test-e2e-cluster: protogen-go build-mock-backend + @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 diff --git a/core/application/absence_wiring.go b/core/application/absence_wiring.go new file mode 100644 index 000000000000..b5f3df99532e --- /dev/null +++ b/core/application/absence_wiring.go @@ -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 +} diff --git a/core/application/absence_wiring_test.go b/core/application/absence_wiring_test.go new file mode 100644 index 000000000000..5ca33122c2f6 --- /dev/null +++ b/core/application/absence_wiring_test.go @@ -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) + 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) + + 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) + router, health := present() + + Expect(requireAbsenceWiring(blindRouter, health)).ToNot(Succeed()) + Expect(requireAbsenceWiring(router, blindHealth)).ToNot(Succeed()) + }) +}) diff --git a/core/application/distributed.go b/core/application/distributed.go index b7dc0bf91351..c814af22488c 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -5,6 +5,8 @@ import ( "encoding/json" "fmt" "io" + "net" + "strconv" "strings" "sync" "time" @@ -12,12 +14,14 @@ 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/nodes" "github.com/mudler/LocalAI/core/services/nodes/prefixcache" "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" @@ -43,6 +47,35 @@ type DistributedServices struct { Unloader *nodes.RemoteUnloaderAdapter ModelCleanup *nodes.ModelCleanupService + // 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 + shutdownOnce sync.Once } @@ -53,6 +86,22 @@ 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() } @@ -162,6 +211,100 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("Node registry initialized") + // Replica membership. NewNodeRegistry has just migrated the tables this + // reads, so it has to come after it. + clusterRegistry := cluster.NewRegistry(authDB) + 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. + xlog.Error("This replica is not registered in the cluster: no advertised address. Peers cannot reach it, and any worker whose tunnel lands here will be unroutable from every other replica", + "error", err, "knob", "LOCALAI_DISTRIBUTED_ADVERTISE_ADDR") + } else { + membership = cluster.NewMembership(clusterRegistry, cfg.Distributed.InstanceID, advertised, internal.PrintableVersion()) + // 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. + peers := cluster.NewPeerPool(cfg.Distributed.InstanceID, cfg.Distributed.RegistrationToken, 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 @@ -197,11 +340,19 @@ 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. healthMon := nodes.NewHealthMonitor(registry, authDB, cfg.Distributed.HealthCheckIntervalOrDefault(), cfg.Distributed.StaleNodeThresholdOrDefault(), routerAuthToken, !cfg.Distributed.DisablePerModelHealthCheck, + clusterRegistry, + cfg.Distributed.ReconnectGraceOrDefault(), + backendClients, ) // Initialize job store @@ -246,28 +397,42 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade } xlog.Info("File manager initialized", "cacheDir", cacheDir) + // 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) + // 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(), ) @@ -357,12 +522,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, @@ -393,7 +565,7 @@ 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 @@ -421,6 +593,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, @@ -429,30 +602,138 @@ 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, + 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, + Cluster: clusterRegistry, + Membership: membership, + PeerSessions: peerSessions, + Peers: peers, + Tunnels: tunnels, + WorkerDialer: workerDialer, + BackendClients: backendClients, }, 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 { return strings.HasPrefix(url, "postgres://") || strings.HasPrefix(url, "postgresql://") } diff --git a/core/application/startup.go b/core/application/startup.go index abc2f4a17571..c790a1403617 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -283,9 +283,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. diff --git a/core/application/unadvertised_replica_test.go b/core/application/unadvertised_replica_test.go new file mode 100644 index 000000000000..ed9eeb4d6399 --- /dev/null +++ b/core/application/unadvertised_replica_test.go @@ -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()) + }) +}) diff --git a/core/cli/run.go b/core/cli/run.go index 6b9b3e4dc0b9..7af51b84612b 100644 --- a/core/cli/run.go +++ b/core/cli/run.go @@ -165,6 +165,7 @@ type RunCMD struct { Distributed bool `env:"LOCALAI_DISTRIBUTED" default:"false" help:"Enable distributed mode (requires PostgreSQL + NATS)" 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"` + 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"` @@ -182,6 +183,7 @@ type RunCMD struct { BackendUpgradeTimeout string `env:"LOCALAI_NATS_BACKEND_UPGRADE_TIMEOUT" help:"NATS round-trip timeout for backend.upgrade requests (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"` + 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"` 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"` @@ -351,6 +353,9 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error { if r.InstanceID != "" { opts = append(opts, config.WithDistributedInstanceID(r.InstanceID)) } + if r.DistributedAdvertiseAddr != "" { + opts = append(opts, config.WithDistributedAdvertiseAddr(r.DistributedAdvertiseAddr)) + } if r.NatsURL != "" { opts = append(opts, config.WithNatsURL(r.NatsURL)) } @@ -397,6 +402,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.RegistrationToken != "" { opts = append(opts, config.WithRegistrationToken(r.RegistrationToken)) } diff --git a/core/cli/workerregistry/client.go b/core/cli/workerregistry/client.go index cf46455c95c0..fb00fb3f1d86 100644 --- a/core/cli/workerregistry/client.go +++ b/core/cli/workerregistry/client.go @@ -8,7 +8,9 @@ import ( "bytes" "context" "encoding/json" + "errors" "fmt" + "io" "net/http" "strings" "sync" @@ -58,9 +60,15 @@ 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"` + 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"` NatsJWT string `json:"nats_jwt,omitempty"` NatsUserSeed string `json:"nats_user_seed,omitempty"` } @@ -87,7 +95,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,6 +105,58 @@ func (c *RegistrationClient) RegisterFull(ctx context.Context, body map[string]a return &result, nil } +// 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 { + xlog.Debug("Could not read the frontend's registration error body", "status", resp.StatusCode, "error", err) + } + 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 +} + // 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) { @@ -108,27 +168,48 @@ func (c *RegistrationClient) Register(ctx context.Context, body map[string]any) } // RegisterWithRetry retries registration with exponential backoff. +// +// It drops every field of the response it does not name, the tunnel credential +// among them. Callers that need one use RegisterFullWithRetry. func (c *RegistrationClient) RegisterWithRetry(ctx context.Context, body map[string]any, maxRetries int) (nodeID, apiToken, natsJWT, natsSeed string, err error) { + res, err := c.RegisterFullWithRetry(ctx, body, maxRetries) + if err != nil { + return "", "", "", "", err + } + return res.ID, res.APIToken, res.NatsJWT, res.NatsUserSeed, nil +} + +// 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. diff --git a/core/cli/workerregistry/client_test.go b/core/cli/workerregistry/client_test.go new file mode 100644 index 000000000000..5870d2524b5c --- /dev/null +++ b/core/cli/workerregistry/client_test.go @@ -0,0 +1,138 @@ +package workerregistry + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "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 := NewNATSCredentialManager(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))) + }) +}) diff --git a/core/cli/workerregistry/credentials.go b/core/cli/workerregistry/credentials.go index 24dd6f3c8ed7..b023b9916918 100644 --- a/core/cli/workerregistry/credentials.go +++ b/core/cli/workerregistry/credentials.go @@ -2,6 +2,7 @@ package workerregistry import ( "context" + "errors" "fmt" "sync" "time" @@ -50,6 +51,11 @@ type NATSCredentialManager struct { 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 @@ -87,6 +93,22 @@ func (m *NATSCredentialManager) store(res *RegisterResponse) { if res.NatsJWT != "" && res.NatsUserSeed != "" { m.jwt, m.seed = res.NatsJWT, res.NatsUserSeed } + // Guarded the same way the NATS pair is: 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 + } +} + +// 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 *NATSCredentialManager) TunnelToken() string { + m.mu.RLock() + defer m.mu.RUnlock() + return m.tunnelToken } // Current returns the latest NATS credentials (both empty until acquired). @@ -125,6 +147,11 @@ func (m *NATSCredentialManager) Acquire(ctx context.Context) (*RegisterResponse, 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) diff --git a/core/config/distributed_config.go b/core/config/distributed_config.go index 5a48a84e9b44..af2b375bf8a6 100644 --- a/core/config/distributed_config.go +++ b/core/config/distributed_config.go @@ -13,8 +13,15 @@ import ( // DistributedConfig holds configuration for horizontal scaling mode. // When Enabled is true, PostgreSQL and NATS are required. type DistributedConfig struct { - Enabled bool // --distributed / LOCALAI_DISTRIBUTED - InstanceID string // --instance-id / LOCALAI_INSTANCE_ID (auto-generated UUID if empty) + 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 NatsURL string // --nats-url / LOCALAI_NATS_URL StorageURL string // --storage-url / LOCALAI_STORAGE_URL (S3 endpoint) RegistrationToken string // --registration-token / LOCALAI_REGISTRATION_TOKEN (required token for node registration) @@ -74,6 +81,17 @@ type DistributedConfig struct { MCPCIJobTimeout time.Duration // MCP CI job execution timeout (default 10m) + // 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 // NATS round-trip timeout for backend.install (default 15m) BackendUpgradeTimeout time.Duration // NATS round-trip timeout for backend.upgrade (default 15m) // ModelLoadTimeout is the gRPC deadline for the remote LoadModel call the @@ -175,6 +193,7 @@ func (c DistributedConfig) Validate() error { FlagBackendInstallTimeout: c.BackendInstallTimeout, FlagBackendUpgradeTimeout: c.BackendUpgradeTimeout, FlagModelLoadTimeout: c.ModelLoadTimeout, + FlagWorkerReconnectGrace: c.WorkerReconnectGrace, } { if d < 0 { return fmt.Errorf("%s must not be negative", name) @@ -195,6 +214,14 @@ func WithDistributedInstanceID(id 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.AdvertiseAddr = addr + } +} + func WithNatsURL(url string) AppOption { return func(o *ApplicationConfig) { o.Distributed.NatsURL = url @@ -307,6 +334,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 @@ -401,7 +436,10 @@ const ( FlagBackendInstallTimeout = "backend-install-timeout" FlagBackendUpgradeTimeout = "backend-upgrade-timeout" FlagModelLoadTimeout = "model-load-timeout" - FlagModelLoadWait = "model-load-wait" + // 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" // 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. @@ -426,6 +464,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 @@ -474,6 +536,23 @@ func (c DistributedConfig) NatsAuthConfig() natsauth.Config { } } +// 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. func (c DistributedConfig) BackendInstallTimeoutOrDefault() time.Duration { return cmp.Or(c.BackendInstallTimeout, DefaultBackendInstallTimeout) diff --git a/core/config/distributed_config_test.go b/core/config/distributed_config_test.go index ec7fbe8dc7e9..deab290fa461 100644 --- a/core/config/distributed_config_test.go +++ b/core/config/distributed_config_test.go @@ -176,3 +176,52 @@ var _ = Describe("DistributedConfig.Validate registration auth", func() { Expect(err.Error()).To(ContainSubstring("LOCALAI_NATS_REQUIRE_AUTH")) }) }) + +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, + NatsURL: "nats://localhost:4222", + 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)) + }) +}) diff --git a/core/http/app.go b/core/http/app.go index 2e1453ac0e38..ec2b864ee8b7 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -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" @@ -566,15 +569,79 @@ 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.RegisterNodeAdminRoutes(e, registry, remoteUnloader, application.GalleryService(), opcache, application.ApplicationConfig(), adminMiddleware, application.AuthDB(), application.ApplicationConfig().Auth.APIKeyHMACSecret, application.ApplicationConfig().Distributed.RegistrationToken, natsCfg, 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") + } + routes.RegisterClusterRoutes(e, distCfg.RegistrationToken, 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 { diff --git a/core/http/auth/public_routes.go b/core/http/auth/public_routes.go index 658205a78f8f..2c1e1dc3cdd1 100644 --- a/core/http/auth/public_routes.go +++ b/core/http/auth/public_routes.go @@ -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) } diff --git a/core/http/endpoints/cluster/cluster_suite_test.go b/core/http/endpoints/cluster/cluster_suite_test.go new file mode 100644 index 000000000000..00f45e37a672 --- /dev/null +++ b/core/http/endpoints/cluster/cluster_suite_test.go @@ -0,0 +1,13 @@ +package cluster_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestClusterEndpoints(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Cluster Endpoints Suite") +} diff --git a/core/http/endpoints/cluster/connect.go b/core/http/endpoints/cluster/connect.go new file mode 100644 index 000000000000..8d175704669c --- /dev/null +++ b/core/http/endpoints/cluster/connect.go @@ -0,0 +1,262 @@ +// 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 two places that hand a node something + // DURABLE both refuse a pending one: the agent worker's API key + // (provisionAgentWorkerKey, guarded at its call site in + // core/http/endpoints/localai/nodes.go) and its NATS credential + // (attachNatsJWT in the same file). Cited by NAME, not by line: the + // previous version of this comment cited line numbers into a file this + // same commit was editing, and both were stale before it landed. + // + // 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 +} diff --git a/core/http/endpoints/cluster/connect_test.go b/core/http/endpoints/cluster/connect_test.go new file mode 100644 index 000000000000..15206d50444c --- /dev/null +++ b/core/http/endpoints/cluster/connect_test.go @@ -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) + }) +}) diff --git a/core/http/endpoints/cluster/peer.go b/core/http/endpoints/cluster/peer.go new file mode 100644 index 000000000000..9dcaea8b38de --- /dev/null +++ b/core/http/endpoints/cluster/peer.go @@ -0,0 +1,146 @@ +// 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" + "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. +// +// 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, 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") + } + + // SELF-DECLARED, and knowingly so. Unlike the worker route next door, + // which resolves ?id= to a node row and checks that node's OWN minted + // credential, this route has only the shared cluster token to check, + // so the id is a label and not a claim anything verifies. + // + // What that costs, exactly, for anything already holding the shared + // token (every worker holds it, and it is the same token that + // authenticates registration): it can relay to every worker tunnel this + // replica owns, reaching every backend gRPC process and every worker's + // file-transfer server; by declaring a legitimate replica's id it can + // make SessionStore.Accept evict that replica's inbound link, at will; + // and it can aim the per-session receive window, which PeerLinkConfig + // sizes at roughly 31 GiB of unread data per session, at one replica's + // memory. The first two are not new capabilities in KIND - before + // workers stopped listening, a holder of that token could already dial + // any worker's advertised ports directly - but the token is now the + // only thing between an attacker and the whole fleet's tunnels, and the + // third is a figure written down as sizing guidance that is also a + // budget on a route this open. + // + // It is deferred rather than patched, because the cheap patch does not + // work: checking ?id= against the instances table stops an invented id + // and stops nothing else, since the attack declares a REAL replica's + // id, and it would buy a false sense of a closed hole. Closing it takes + // a credential per replica, minted where a replica joins the instances + // table and presented here, which is a design with its own migration + // and its own specs. Tracked as the phase-3 item named at + // nodes.BackendNode.TunnelTokenHash. + peerID := c.QueryParam("id") + if peerID == "" { + return echo.NewHTTPError(http.StatusBadRequest, "missing peer id") + } + + 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 +} diff --git a/core/http/endpoints/cluster/peer_test.go b/core/http/endpoints/cluster/peer_test.go new file mode 100644 index 000000000000..dc8cba3da9e6 --- /dev/null +++ b/core/http/endpoints/cluster/peer_test.go @@ -0,0 +1,256 @@ +package cluster_test + +import ( + "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/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" +} + +var _ = Describe("Peer link handler", func() { + var ( + srv *httptest.Server + sessions chan *yamux.Session + ) + + BeforeEach(func() { + sessions = make(chan *yamux.Session, 1) + e := echo.New() + routes.RegisterClusterRoutes(e, "peer-token", 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() { + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(srv), h) + 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 was given", func() { + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + ids := make(chan string, 1) + e := echo.New() + routes.RegisterClusterRoutes(e, "peer-token", func(id string, _ *yamux.Session) { ids <- id }) + s2 := httptest.NewServer(e) + DeferCleanup(s2.Close) + + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), h) + 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, "", 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"}}} { + _, 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 := http.Header{} + 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", 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() + }) + + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + conn, _, err := websocket.DefaultDialer.Dial(wsPeerURL(s2), h) + 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. + h := http.Header{} + h.Set("Authorization", "Bearer peer-token") + _, resp, err := websocket.DefaultDialer.Dial( + "ws"+strings.TrimPrefix(srv.URL, "http")+"/api/cluster/peer", h) + 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 + ) + + BeforeEach(func() { + sessions = make(chan *yamux.Session, 1) + 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", 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. + h := http.Header{} + 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()) + }) +}) diff --git a/core/http/endpoints/localai/backend_logs.go b/core/http/endpoints/localai/backend_logs.go index 6072b8483708..f2112502d253 100644 --- a/core/http/endpoints/localai/backend_logs.go +++ b/core/http/endpoints/localai/backend_logs.go @@ -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{ diff --git a/core/http/endpoints/localai/nodes.go b/core/http/endpoints/localai/nodes.go index bbae523b1025..0571e7d7f8df 100644 --- a/core/http/endpoints/localai/nodes.go +++ b/core/http/endpoints/localai/nodes.go @@ -2,6 +2,7 @@ package localai import ( "context" + "crypto/rand" "crypto/sha256" "crypto/subtle" "encoding/hex" @@ -9,6 +10,7 @@ import ( "errors" "fmt" "io" + "net" "net/http" "net/url" "sync" @@ -75,10 +77,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"` @@ -140,22 +146,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 +174,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,6 +241,7 @@ func RegisterNodeEndpoint(registry *nodes.NodeRegistry, expectedToken string, au } } + attachTunnelToken(ctx, response, registry, node) attachNatsJWT(response, node, natsCfg) return c.JSON(http.StatusCreated, response) @@ -288,6 +286,78 @@ func ApproveNodeEndpoint(registry *nodes.NodeRegistry, authDB *gorm.DB, hmacSecr } } +// 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 two other per-node credentials in this file: +// the agent worker's API key (provisionAgentWorkerKey) and its NATS JWT +// (attachNatsJWT) are both withheld from a node awaiting approval. Those two +// are bearer grants that WORK the moment they are 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. +// +// Only BACKEND nodes get one, and that is a decision rather than an oversight. +// An agent worker serves no gRPC backends and no file staging; nothing dials +// into it at all, so a tunnel replaces nothing for it and there is no client on +// the agent side that would ever open one. Minting anyway would hand out a +// working credential for a pipe nobody drives, which is surface without a +// feature, and it would contradict every comment in this change that says +// "backend workers, the ones that tunnel". +// +// 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 backend node re-registering as an agent 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 that had become an +// agent. Clearing is what makes "enforcement is structural" true. +// +// 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 agent registration. +// +// The day agent workers want a tunnel, relaxing the eligibility condition is +// the whole change, and it has to be a deliberate one. +// +// 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 + } + if node.NodeType != nodes.NodeTypeBackend { + // 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 that is not a backend worker", + "node", node.Name, "type", node.NodeType, "error", err) + } + return + } + // 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 +} + // 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 { @@ -717,7 +787,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 +796,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 +815,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 +826,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 +842,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 +874,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 +1365,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 +1391,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) } diff --git a/core/http/endpoints/localai/nodes_backends_list_test.go b/core/http/endpoints/localai/nodes_backends_list_test.go index 636ab58b818e..c625e8e9510f 100644 --- a/core/http/endpoints/localai/nodes_backends_list_test.go +++ b/core/http/endpoints/localai/nodes_backends_list_test.go @@ -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 diff --git a/core/http/endpoints/localai/nodes_test.go b/core/http/endpoints/localai/nodes_test.go index 19e6a6b07eea..2255116aba8c 100644 --- a/core/http/endpoints/localai/nodes_test.go +++ b/core/http/endpoints/localai/nodes_test.go @@ -4,6 +4,7 @@ import ( "context" "crypto/sha256" "crypto/subtle" + "encoding/hex" "encoding/json" "net/http" "net/http/httptest" @@ -20,6 +21,12 @@ import ( . "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 == "" { @@ -77,6 +84,149 @@ var _ = Describe("Node HTTP handlers", func() { Expect(resp["status"]).To(Equal(nodes.StatusHealthy)) }) + // register posts one registration and returns the decoded response. + register := func(body string, expectedToken string, autoApprove bool) map[string]any { + e := echo.New() + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set(echo.HeaderContentType, echo.MIMEApplicationJSON) + rec := httptest.NewRecorder() + c := e.NewContext(req, rec) + + handler := RegisterNodeEndpoint(registry, expectedToken, autoApprove, nil, "", natsauth.Config{}) + 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("does not issue a tunnel credential to an agent node", func() { + // An agent worker serves no gRPC backends and no file staging; + // nothing dials into it, so a tunnel replaces nothing for it and no + // client on its side would open one. Minting anyway would be + // credential surface with no feature behind it. + // + // Enforcement is structural rather than a second check: with no + // credential minted, the node's hash stays empty and the tunnel + // route refuses it like any other node without one. + resp := register(`{"name":"agent-1","node_type":"agent"}`, "", true) + Expect(resp["node_type"]).To(Equal(nodes.NodeTypeAgent)) + Expect(resp).ToNot(HaveKey("tunnel_token")) + + node, err := registry.Get(context.Background(), resp["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + Expect(node.TunnelTokenHash).To(BeEmpty()) + }) + + It("clears a tunnel credential when a node stops being a backend node", func() { + // Register upserts BY NAME, so a node can change node_type in place. + // Skipping the mint on the way through leaves the credential the + // node earned as a backend sitting on a row that is now an agent: + // Register's struct Updates zero-skips the column while writing the + // new node_type, so nothing else clears it. ConnectHandler never + // looks at node_type, so that stale hash is a usable tunnel + // credential for a node type that is not supposed to hold one. + // + // 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) + Expect(backend["tunnel_token"]).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)) + Expect(agent).ToNot(HaveKey("tunnel_token")) + + node, err := registry.Get(context.Background(), backend["id"].(string)) + Expect(err).ToNot(HaveOccurred()) + // The claim the gate makes is that an ineligible node HAS no + // credential, not merely that it was not handed a new one. Only + // then is the empty-hash refusal in ConnectHandler the enforcement. + Expect(node.TunnelTokenHash).To(BeEmpty(), + "the node kept the credential it earned as a backend, so the mint-site gate is not structural") + }) + It("returns nats_jwt when account seed is configured", func() { akp, err := nkeys.CreateAccount() Expect(err).ToNot(HaveOccurred()) @@ -139,7 +289,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)) @@ -149,13 +303,35 @@ var _ = Describe("Node HTTP handlers", func() { handler := RegisterNodeEndpoint(registry, "", true, nil, "", natsauth.Config{}) 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, "", natsauth.Config{}) + 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() { diff --git a/core/http/react-ui/e2e/nodes-roster.spec.js b/core/http/react-ui/e2e/nodes-roster.spec.js index c6396d5b3468..0f229d43321d 100644 --- a/core/http/react-ui/e2e/nodes-roster.spec.js +++ b/core/http/react-ui/e2e/nodes-roster.spec.js @@ -62,3 +62,48 @@ 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, so the bus flag has to differ per tab rather than be deleted. + // Backend workers connect to no NATS server; agent workers still do. + test('omits the NATS flag for a backend worker and keeps it for an agent worker', 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).not.toContainText('LOCALAI_NATS_URL') + + 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('--nats-url') + const agentDocker = page.locator('.p2p-cmd pre').nth(1) + await expect(agentDocker).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') + // And the worker step no longer tells an operator to point a backend + // worker at a bus it does not dial. + await expect(card.locator('.p2p-cmd pre').nth(1)).not.toContainText('--nats-url') + }) +}) diff --git a/core/http/react-ui/src/components/nodes/NodePanel.jsx b/core/http/react-ui/src/components/nodes/NodePanel.jsx index 623db00936bf..b73714ab9972 100644 --- a/core/http/react-ui/src/components/nodes/NodePanel.jsx +++ b/core/http/react-ui/src/components/nodes/NodePanel.jsx @@ -19,7 +19,11 @@ export default function NodePanel({ node, models = [], onApprove, onDrain, onRes
{node.name} - {node.address} + {/* 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. */} + {node.address || node.id}
e.stopPropagation()}> {node.status === 'pending' && ( diff --git a/core/http/react-ui/src/pages/NodeDetail.jsx b/core/http/react-ui/src/pages/NodeDetail.jsx index bff7db526ad2..52d3fea1d5f2 100644 --- a/core/http/react-ui/src/pages/NodeDetail.jsx +++ b/core/http/react-ui/src/pages/NodeDetail.jsx @@ -78,7 +78,7 @@ export default function NodeDetail() { navigate('/app/nodes')} className="link-plain">