diff --git a/.vscode/launch.json b/.vscode/launch.json index adb2772..4ff0a9b 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -7,7 +7,26 @@ "request": "launch", "python": "${workspaceFolder}/.venv-rocm/bin/python", "module": "meshnet_tracker.cli", - "args": ["start", "--host", "0.0.0.0", "--port", "8080", "--stats-db", "${workspaceFolder}/tracker-stats.sqlite", "--enable-test-runner"], + "args": ["start", "--host", "0.0.0.0", "--port", "8080", "--stats-db", "${workspaceFolder}/tracker-stats.sqlite"], + "console": "integratedTerminal", + "justMyCode": false + }, + { + "name": "Tracker: local + dashboard test runner (8080)", + "type": "debugpy", + "request": "launch", + "python": "${workspaceFolder}/.venv-rocm/bin/python", + "module": "meshnet_tracker.cli", + "args": [ + "start", + "--host", + "0.0.0.0", + "--port", + "8080", + "--stats-db", + "${workspaceFolder}/tracker-stats.sqlite", + "--enable-test-runner" + ], "console": "integratedTerminal", "justMyCode": false }, diff --git a/QUICKSTART.md b/QUICKSTART.md index 7f5807a..0cada00 100644 --- a/QUICKSTART.md +++ b/QUICKSTART.md @@ -7,7 +7,15 @@ Tested on: AMD Ryzen AI Max (Strix Halo APU), 124 GB RAM, Linux CPU inference. ROCm GPU setup is covered below, but must be verified on the host because ROCm support depends on the exact AMD GPU/APU, kernel, driver, and ROCm runtime. -**Active development models** (what we run day-to-day): +**Support here is evidence, not a promise.** No combination of GPU, torch build, +optional accelerator package and model is "supported" until `meshnet-node doctor` +has pushed a real forward through the exact shard you intend to serve, on that +machine. Detecting a GPU proves nothing; installing an optional package proves +nothing. See [Validate before you serve](#validate-before-you-serve-meshnet-node-doctor). + +**Active development models** — what we run day-to-day, so their setup notes are the +most exercised. This is *not* a support matrix: any model the node can load is fair +game, and any model on this list can still fail `doctor` on your hardware. | Role | `--model` / alias | HF repo | Notes | |------|-------------------|---------|-------| @@ -288,8 +296,8 @@ HF_HOME=/path/to/models .venv-rocm/bin/meshnet-node start \ --quantization bfloat16 ``` -For the Qwen3.6 alpha model on Linux ROCm, install the optional FLA ROCm fast -path in the same env: +A model whose recipe uses an optional accelerator needs that package in the *same* +env. For example, `qwen3.6-35b-a3b` on Linux ROCm can use FLA: ```bash .venv-rocm/bin/pip install 'flash-linear-attention[rocm]' @@ -299,6 +307,12 @@ HF_HOME=/path/to/models .venv-rocm/bin/meshnet-node start \ --quantization bfloat16 ``` +This is a per-model, per-platform example. A successful `pip install` is not evidence +that the kernel runs on your GPU — optional-kernel support varies by architecture, and +we make no universal claim here. Prove it with +`.venv-rocm/bin/meshnet-node doctor --model ` before starting the node; startup +will refuse to register anyway if the shard cannot execute. + ### Linux ROCm: Triton JIT compiler prerequisite Some model/runtime paths invoke Triton at the first real forward. Triton builds a local HIP @@ -530,8 +544,11 @@ curl -sI https://ai.neuron.d-popov.com/rpc/test-peer no HuggingFace gating. Best for first-time setup. **Alpha model:** `qwen3.6-35b-a3b` — 40 layers, ~72 GB BF16 download, MoE with hybrid -linear attention. On Windows install `triton-windows` + `flash-linear-attention`; on Linux -GPU use `flash-linear-attention[cuda]`. Tracker accepts the alias or full repo id (`unsloth/Qwen3.6-35B-A3B`). +linear attention. Tracker accepts the alias or full repo id (`unsloth/Qwen3.6-35B-A3B`). +Some models have an *optional* accelerator path (for this one, Triton + +`flash-linear-attention`); those installs are per-model, per-platform examples, not a +requirement the node imposes and not a guarantee the path will work — see +[Optional accelerator packages](#optional-accelerator-packages). Downloads cache under `~/.meshnet/models/` (or `$HF_HOME` / `$env:HF_HOME`). @@ -539,6 +556,79 @@ Shard range is auto-detected from the curated catalog. For unknown repos the nod fetches only `config.json`. Override with `--shard-start` / `--shard-end` for partial shards or multi-node splits. +### Validate before you serve: `meshnet-node doctor` + +`doctor` resolves exactly the model, shard, quantization and recipe that `start` would +load, loads it, and pushes one bounded **real forward** through it. It is offline — it +never contacts the tracker — and it is the only thing that turns a guess into evidence. + +```bash +HF_HOME=/path/to/models .venv/bin/meshnet-node doctor --model --quantization bfloat16 +``` + +``` +meshnet-node doctor + Model: + Shard: layers 0–23; 24 of 24 + Quantization: bfloat16 + + [PASS] recipe default (v1) on cuda — 412 ms + + OK — the selected shard ran a real forward for 1 recipe. + Capability report: ~/.meshnet/capability.json +``` + +Exit code is 0 on pass, 1 on failure. It writes a **capability report** either way — a +failure is evidence too. Useful flags: `--recipe ` to validate one named recipe, +`--all-recipes` to try every recipe for the selection (CI and diagnosis), `--report PATH` +to choose where the report lands, `--json` for machine-readable output, `--debug` for the +full traceback behind a failure. The selection flags (`--model`, `--quantization`, +`--shard-start`/`--shard-end`, `--cpu`) work the same as on `start`. + +**Three states — do not confuse them:** + +| State | What it means | How it is established | +|-------|---------------|-----------------------| +| **Detected hardware** | A GPU, a torch build, and an optional package are *present*. | Inventory/`torch.cuda.is_available()`. Proves nothing about whether your shard runs. | +| **Validated recipe** | This machine executed a real forward for *this* model + shard + recipe + device. | A passing `doctor` run → a capability report. | +| **Routable Node** | The tracker will send paid work here. | Startup re-proves the loaded shard and registers the report; the tracker admits it. | + +Each state is strictly stronger than the one above it. A machine can have a working GPU +(detected) and still fail to execute a shard (no validated recipe), and a node can hold a +valid report and still not route (the tracker moved it to a shard range it never proved). + +**Startup is fail-closed.** `meshnet-node start` runs the same validation against the +backend it just loaded, *before* it opens a port or registers. If the shard cannot execute, +the node exits non-zero — it never binds a socket and never registers. There is no +production bypass; a node that cannot do the work never advertises that it can. + +**The tracker admits, it does not re-run.** The node ships its report with registration. +Registration always succeeds, so a broken node is visible rather than invisible — but only +a report that *covers what the node advertises* makes it routable. The verdict comes back +in the registration response, is logged, and is exposed per node on `GET /v1/network/map` +under `capability` (`admitted`, `absent`, `invalid`, `failed`, `stale`, `model-mismatch`, +`shard-mismatch`, `recipe-mismatch`, `catalogue-incompatible`). A node that is registered +but dark is showing you one of those. Full semantics, and the transitional `compat` vs +`enforce` policy, are in [ADR-0023](docs/adr/0023-model-agnostic-node-capability-admission.md). + +**When `doctor` fails** it prints a category and an actionable hint, not a traceback: + +| Category | Do this | +|----------|---------| +| `no-model-selected` | Pass `--model`, or run `meshnet-node` once to save a config. | +| `missing-dependency` | Install the node's model extras (torch, transformers, safetensors, accelerate). | +| `model-unavailable` | Check the model id, `--download-dir`, and that the artifact is downloaded/reachable. | +| `insufficient-memory` | Serve fewer layers (`--shard-start`/`--shard-end`) or a smaller quantization (`-q int8`, `-q nf4`). | +| `invalid-shard` | The layer range does not exist in this model — check it against the layer count. | +| `unsupported-recipe` | This backend cannot apply that execution setting; pick another `--recipe`. | +| `load-failed` | The shard would not load; re-run with `--debug`. | +| `forward-failed` | It loaded but cannot execute. This machine cannot serve this shard; re-run with `--debug`. | + +**Scope:** the node validates recipes it already ships. It does **not** download executable +recipes, install Python/OS packages, or update drivers for you — every install on this page +is something you run deliberately. Signed node updates are a deliberate follow-up, not part +of this release. + ### Core command Replace `` and adjust the prefix for your shell (see table above). @@ -570,31 +660,45 @@ HF_HOME=/path/to/models .venv/bin/meshnet-node start --tracker http://localhost: .venv/bin/meshnet-node start --tracker https://ai.neuron.d-popov.com --model Qwen/Qwen2.5-0.5B-Instruct ``` -**Alpha model (Qwen3.6, Windows GPU — enable fast path):** +#### Optional accelerator packages + +Some models have an optional fast path that a *specific* model's recipe can use if the +package is importable. The examples below are exactly that — **worked examples for one +development model on one platform**, not a supported-configuration list. Installing the +package does not mean the fast path works on your GPU: a package can import cleanly and +still fail to execute a kernel on your architecture. `doctor` is what tells you, and the +node serves the recipe that actually validated. If none of this is installed, the model +still runs on the standard path. + +**Example — `qwen3.6-35b-a3b` (Triton + FLA), Windows GPU:** ```powershell $env:HF_HOME = "D:\DEV\models" pip install triton-windows pip install -U flash-linear-attention +meshnet-node doctor --model qwen3.6-35b-a3b --quantization bfloat16 meshnet-node start --tracker http://192.168.0.179:8080 --model qwen3.6-35b-a3b --quantization bfloat16 ``` Do not add `causal-conv1d` or `flash-linear-attention[cuda]` on Windows (see Qwen3.5/3.6 notes). -**Alpha model (Qwen3.6, Linux NVIDIA GPU — with fast path):** +**Example — `qwen3.6-35b-a3b`, Linux NVIDIA GPU:** ```bash -HF_HOME=/path/to/models .venv/bin/meshnet-node start --tracker --model qwen3.6-35b-a3b --quantization bfloat16 # Install once on that machine: pip install flash-linear-attention[cuda] +HF_HOME=/path/to/models .venv/bin/meshnet-node start --tracker --model qwen3.6-35b-a3b --quantization bfloat16 ``` -**Alpha model (Qwen3.6, Linux AMD ROCm GPU — with fast path):** +**Example — `qwen3.6-35b-a3b`, Linux AMD ROCm GPU:** ```bash -HF_HOME=/path/to/models .venv-rocm/bin/meshnet-node start --tracker --model qwen3.6-35b-a3b --quantization bfloat16 # Install once on that machine: .venv-rocm/bin/pip install 'flash-linear-attention[rocm]' +HF_HOME=/path/to/models .venv-rocm/bin/meshnet-node start --tracker --model qwen3.6-35b-a3b --quantization bfloat16 ``` +Run `doctor` after any of these installs: it is the only way to learn whether the optional +path executes here, and it costs one bounded forward. + After the first node registers a model, later nodes can join with only the tracker URL (shard auto-assigned): @@ -786,3 +890,9 @@ failure the node logs a warning and falls back to direct HTTP before erroring. ```bash .venv/bin/python -m pytest -q ``` + +Tests marked `integration` download models or need a GPU; the default lane is +`pytest -m "not integration"`. The opt-in real-model doctor check takes its model from the +environment and skips without it — see +[docs/dev/certified-hardware-lanes.md](docs/dev/certified-hardware-lanes.md) for the +release contract that certifies a hardware lane. diff --git a/docs/adr/0023-model-agnostic-node-capability-admission.md b/docs/adr/0023-model-agnostic-node-capability-admission.md index fa18a07..2bb19c4 100644 --- a/docs/adr/0023-model-agnostic-node-capability-admission.md +++ b/docs/adr/0023-model-agnostic-node-capability-admission.md @@ -18,6 +18,31 @@ This is incompatible with a consumer-grade node experience. A Node must never ad - P0 carries the version of a local recipe manifest. New executable recipes arrive only through signed Node releases in a future feature. P0 does not download executable recipes, dynamically install dependencies, install OS packages/drivers, or implement an updater. - A future Tracker-provided Model Artifact Manifest may be signed data only; it cannot instruct a Node to execute arbitrary code. +## Tracker admission and the compatibility policy for older Nodes + +A Node ships its capability report with `POST /v1/nodes/register` (`capability_report`), alongside an independent declaration of the recipe it serves with (`recipe_id`, `recipe_version`). The Tracker does not re-run the forward. It decides whether the presented proof *covers what the Node advertises*, records the verdict as a sanitized state, and routes accordingly. + +Registration always succeeds — a Node with a bad proof is registered and visible, it is simply not routable. "Registered but dark" is a state an operator must be able to see and diagnose, so the verdict is returned in the registration response, logged, and exposed per node on `GET /v1/network/map` under `capability` (state, detail, proven model/shard/recipe/backend/device, timestamps). The detail is credential-redacted and clipped; a raw exception or token never reaches an operator view. + +Verdicts: `admitted`, `absent`, `invalid`, `failed`, `stale`, `model-mismatch`, `shard-mismatch`, `recipe-mismatch`, `catalogue-incompatible`. Only `admitted` is proof. The proof does not travel with a reassignment: if the Tracker later moves a Node to a range it never validated, the Node is re-verdicted `shard-mismatch` and stops routing until it re-registers with a proof for the range it now advertises. + +Freshness is checked when the proof is *presented*, not continuously — a long-lived Node's proof does not expire out from under it while it is heartbeating; liveness is already carried by heartbeat expiry. + +**Compatibility policy** (`--capability-policy`, `$MESHNET_TRACKER_CAPABILITY_POLICY`): + +- **`compat` (default, transitional)** — a Node that presents *no* report at all still routes, preserving pre-capability Node behaviour during the fleet rollout. Every other verdict is refused. Presenting a broken, failed, stale or mismatched proof is a stronger negative signal than presenting none, so it is never grandfathered. +- **`enforce`** — only `admitted` routes. Absent proof is not routable, and no paid route can rest on an unproven Node. + +`compat` is a deprecating default: it exists to let a mixed fleet upgrade without an outage, and `enforce` becomes the default once the deployed Nodes emit reports. The policy is a single explicit switch, checked in one gate (`_admitted_nodes`) that every route path — proxy head selection, `/v1/route`, `/v1/routes`, and bandit route enumeration — passes through. The gate only ever *removes* candidates; coverage-first selection and throughput-weighted preference among the survivors are untouched, and nothing in a report can raise a Node's routing weight (performance stays measured, per ADR-0013/ADR-0021). + +The Tracker also refuses a report whose recipe catalogue predates `MIN_CATALOGUE_VERSION`: recipe ids from an older catalogue may since have been redefined, so the proof cannot be matched to a name reliably. + +## Hardware claims are evidence, not a support matrix + +Operator docs must distinguish three states and never collapse them: **detected hardware** (a GPU, a torch build, or an optional package is present — proves nothing), **validated recipe** (this machine ran a real forward for this model/shard/recipe/device, and there is a capability report to show for it), and **routable Node** (the Tracker admitted that proof for what the Node advertises). Each is strictly stronger than the last. + +Consequently no doc promises that a model, vendor, or optional kernel works universally. A concrete model appears only as a clearly-labelled example or as environment-supplied test configuration. Hardware support is claimed per *certified lane*, where a lane is certified by an opt-in `integration` doctor run whose model identity comes from CI configuration and whose retained evidence is the capability report — see `docs/dev/certified-hardware-lanes.md`. A lane certifies hardware, not models: a new Model Artifact is unproven there until doctor has run it. + ## Consequences - First startup has a bounded validation cost before registration, but failures occur before traffic rather than under a paid request. diff --git a/docs/dev/certified-hardware-lanes.md b/docs/dev/certified-hardware-lanes.md new file mode 100644 index 0000000..ac8321c --- /dev/null +++ b/docs/dev/certified-hardware-lanes.md @@ -0,0 +1,105 @@ +# Certified hardware lanes + +A **certified hardware lane** is one (hardware, torch build, OS) configuration on which we +have *evidence* that nodes can execute real work — a self-hosted release runner that runs +the opt-in integration doctor test and keeps the resulting capability report as the +artifact. This document is the contract that lane runners and the release check must meet. + +Certification is per *lane*, and evidence is per *(lane, model, shard, recipe)*. Nothing +here promises that an arbitrary model runs on a certified lane; it promises that the lane +itself is real, and that the models we ran on it produced passing capability reports on a +named date. See [ADR-0023](../adr/0023-model-agnostic-node-capability-admission.md) for the +admission model this rests on. + +## What certification is not + +- **Not a model support matrix.** A lane certifies hardware, not models. A new model is + unproven on a certified lane until doctor has run it there. +- **Not an optional-kernel promise.** An optional accelerator package importing cleanly on + a lane says nothing about another lane, another GPU architecture, or another model's + recipe. Only a passing report for that exact combination is evidence. +- **Not a promise the node will install anything.** Lane runners are provisioned *ahead of + time*, by hand or by their own image build. The node under test never downloads an + executable recipe, installs a Python or OS package, or touches a driver. Signed node + updates are a deliberate follow-up feature and are out of scope here — nothing in this + lane contract may depend on dynamic executable-dependency installation. + +## The lane check + +Every lane runs the same environment-configured integration test. It is +`tests/test_node_doctor.py::test_doctor_smoke_runs_a_real_forward_on_a_real_model`, marked +`@pytest.mark.integration` and skipped unless `MESHNET_DOCTOR_MODEL` is set. It carries no +default model: **model identity comes from the CI configuration**, so no vendor or model +assumption can leak into the suite. + +```bash +MESHNET_DOCTOR_MODEL="$LANE_MODEL" \ +MESHNET_DOCTOR_QUANTIZATION=bfloat16 \ +MESHNET_DOWNLOAD_DIR=/srv/models \ + .venv/bin/pytest -m integration tests/test_node_doctor.py -v +``` + +| Variable | Required | Meaning | +|----------|----------|---------| +| `MESHNET_DOCTOR_MODEL` | yes — the test skips without it | Model artifact identity for this lane's run. No default. | +| `MESHNET_DOCTOR_SHARD_START` | no (default `0`) | First layer of the shard to prove. | +| `MESHNET_DOCTOR_SHARD_END` | no (default: whole model) | Last layer, **inclusive**. | +| `MESHNET_DOCTOR_QUANTIZATION` | no (default `auto`) | Quantization to prove. | +| `MESHNET_DOCTOR_CPU` | no | `1` forces CPU — use to certify a CPU lane on GPU hardware. | +| `MESHNET_DOWNLOAD_DIR` | no | Where the artifact is cached on the runner. | + +The test asserts that doctor passed, that the report is `passed` with the model id it was +asked for, that a forward actually took time (`duration_ms > 0`), and that the report +round-trips through `CapabilityReport.from_json`. A lane where this fails is not certified, +regardless of what `rocminfo`, `nvidia-smi` or `torch.cuda.is_available()` say. + +Lanes that must cover more than the default recipe run doctor directly with +`--all-recipes`, which validates every recipe for the selection and writes a report per +recipe: + +```bash +.venv/bin/meshnet-node doctor --model "$LANE_MODEL" --all-recipes --report "$ARTIFACTS/capability.json" +``` + +## Expected evidence + +A lane run is only certified if it produces, and the release check retains: + +1. **The capability report(s)** — `capability.json` from the run, archived as a build + artifact. This is the evidence; a green checkmark without it is not. +2. **Backend identity from the report**: device, torch/backend version, and the recipe id + and version that passed. This is what makes "certified on ROCm gfx1151" a checkable + claim rather than a slogan. +3. **The model artifact identity and shard range** the report covers — recorded as run + configuration, since it came from the environment. +4. **Failures kept, not discarded.** Doctor writes a report for a failed recipe too, and a + failing lane must archive it. A red lane with a `forward-failed` report is a more useful + release signal than a lane that was quietly skipped. + +A lane that *skips* (because `MESHNET_DOCTOR_MODEL` was unset) must be reported as skipped, +never as passed. A silent skip is how an uncertified lane gets mistaken for a certified one. + +## Release check + +The default CI lane runs the normal suite and never needs a GPU, a download, or torch: + +```bash +.venv/bin/pytest -m "not integration" +``` + +The release check additionally requires every declared certified lane to have run the +integration doctor test green, against the model(s) configured for that lane, on the +release commit. Adding a lane means standing up a runner and adding its configuration; it +does not mean adding a model default to the test suite. + +## Adding a lane + +1. Provision the runner: OS, driver, and the torch build for that hardware (see the + platform sections in `QUICKSTART.md`). Install any optional accelerator packages the + lane is meant to certify. +2. Configure `MESHNET_DOCTOR_MODEL` (and shard/quantization if the lane certifies a + partial shard) in the runner's CI configuration. +3. Run the lane check. Archive the capability report. +4. Record the lane with the evidence it produced: hardware, torch build, model, shard, + recipe, device, and the date. That record — not the hardware's spec sheet — is the + support claim. diff --git a/docs/dev/dashboard-test-runner.md b/docs/dev/dashboard-test-runner.md new file mode 100644 index 0000000..ef46a50 --- /dev/null +++ b/docs/dev/dashboard-test-runner.md @@ -0,0 +1,93 @@ +# Dashboard test runner (operator workflow) + +The tracker dashboard **Testing** tab can discover repository pytest targets and +run one collected test or approved suite at a time with live logs. The feature is +**disabled by default** and **admin-only**. + +## Enable intentionally + +Start the tracker with the test runner API enabled using either: + +- VS Code launch configuration **`Tracker: local + dashboard test runner (8080)`** + (uses the project tracker runtime at `.venv-rocm/bin/python` and + `meshnet_tracker.cli`), or +- CLI flag **`--enable-test-runner`**, or +- Environment variable **`MESHNET_ENABLE_TEST_RUNNER=1`** + +The default **`Tracker: local (8080)`** launch configuration does **not** enable +the test runner. + +Verify the flag is available: + +```bash +uv run python -m meshnet_tracker.cli --help | grep enable-test-runner +``` + +Log in to the dashboard as an admin account, open **Testing**, and use **Refresh +collection** before running targets. + +## Child pytest interpreter + +The runner spawns pytest as a subprocess without a shell. It uses +**`MESHNET_PYTHON`** when set (typically via `.env.` loaded by +`meshnet_tracker.cli`); otherwise it falls back to the tracker process +interpreter. Point this at the venv that has dev extras and package dependencies +installed (see [test-env.md](test-env.md)). + +## Default safe suites + +These named suites are always available when the test runner is enabled and the +files exist in the checkout: + +| Suite ID | Paths | Notes | +| --- | --- | --- | +| `suite:smoke` | `tests/test_smoke.py` | Fast sanity checks | +| `suite:dashboard` | `tests/test_dashboard.py` | Dashboard HTML/API regressions | +| `suite:routing` | `tests/test_tracker_routing.py`, `tests/test_dynamic_routing.py` | Tracker routing logic | + +Collection also lists individual pytest node IDs (excluding real-inference +modules by default). You can run `suite:all` or `tag:` after collection. + +These suites use mocks/stubs or in-process fakes. They do **not** require live +GPU nodes, paid API credits, or a running mesh beyond the tracker itself. + +## Real-inference suite (explicitly gated) + +Modules matching `tests/test_real_*.py` are **never collected** and **never** +included in default suites unless you set: + +```bash +export MESHNET_ENABLE_REAL_INFERENCE_TESTS=1 +``` + +With that gate, an additional suite appears: + +| Suite ID | Paths | +| --- | --- | +| `suite:real-inference` | `tests/test_real_distributed_inference.py`, `tests/test_real_model_backend.py` | + +### Implications + +- **`tests/test_real_distributed_inference.py`** — integration test against a + **live tracker and registered model shards**. Requires env vars such as + `MESHNET_REAL_INFERENCE_URL`, `MESHNET_REAL_INFERENCE_API_KEY`, + `MESHNET_REAL_INFERENCE_MODEL`, and `MESHNET_REAL_INFERENCE_ROUTE`. Uses real + chat completions and **consumes caller billing / API credit** on the target + tracker. +- **`tests/test_real_model_backend.py`** — loads real PyTorch model code paths; + needs **`torch`**, **`transformers`**, and related optional deps, and can + require **substantial GPU/CPU RAM** depending on which cases run. + +Do not enable `MESHNET_ENABLE_REAL_INFERENCE_TESTS=1` on shared or production +trackers unless you intend to spend credits and tie up hardware. + +## Safety summary + +| Control | Purpose | +| --- | --- | +| Disabled by default | No test subprocess unless operator opts in | +| Admin-only API/UI | Non-admins cannot start runs | +| Fixed suite list | API cannot pass arbitrary shell commands | +| No `shell=True` | pytest argv is fixed server-side | +| One run at a time | Concurrent starts are rejected | +| Real-inference env gate | Live inference tests stay out of default collection | diff --git a/docs/dev/test-env.md b/docs/dev/test-env.md index e23fec4..51c64de 100644 --- a/docs/dev/test-env.md +++ b/docs/dev/test-env.md @@ -38,6 +38,12 @@ even without installing `packages/node`. .venv/bin/python -m pytest ``` +## Dashboard test runner + +For the opt-in tracker dashboard **Testing** tab (suites, env gates, VS Code +launch config, and real-inference safeguards), see +[dashboard-test-runner.md](dashboard-test-runner.md). + ## Optional-dependency tests Some tests import heavyweight or optional third-party packages and guard diff --git a/packages/gateway/meshnet_gateway/prefill_backpressure.py b/packages/gateway/meshnet_gateway/prefill_backpressure.py new file mode 100644 index 0000000..664321d --- /dev/null +++ b/packages/gateway/meshnet_gateway/prefill_backpressure.py @@ -0,0 +1,130 @@ +"""Bounded, ordered prefill transfer primitives. + +Prefill chunks mutate the downstream shard's session cache, so they must reach a +route in order. This deliberately uses a serial acknowledgement window: it is +the safe default for both current peers and old peers which do not advertise a +windowing capability. The configured in-flight limit is still explicit so a +future ordered transport can widen the window without changing callers. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from threading import Event +from typing import Callable, Iterable, TypeVar + + +DEFAULT_PREFILL_CHUNK_TOKENS = 128 +DEFAULT_PREFILL_MAX_IN_FLIGHT = 1 +DEFAULT_PREFILL_MAX_CHUNK_BYTES = 8 * 1024 * 1024 + +T = TypeVar("T") +R = TypeVar("R") + + +@dataclass(frozen=True) +class PrefillTransferLimits: + """Configuration for one ordered prefill seam.""" + + chunk_tokens: int = DEFAULT_PREFILL_CHUNK_TOKENS + max_in_flight: int = DEFAULT_PREFILL_MAX_IN_FLIGHT + max_chunk_bytes: int = DEFAULT_PREFILL_MAX_CHUNK_BYTES + + @property + def effective_in_flight(self) -> int: + """Current peers require ordered session-cache mutation, hence one ack.""" + return 1 + + @property + def max_buffered_bytes(self) -> int: + """Hard accounting bound, including any future wider ack window.""" + return self.max_chunk_bytes * self.max_in_flight + + @classmethod + def from_env(cls) -> "PrefillTransferLimits": + # MESHNET_CHUNK_TOKENS was the pre-DIP-007 name. Keep it as a fallback + # so existing deployments retain their chunk shape while upgrading. + return cls( + chunk_tokens=_positive_env( + "MESHNET_PREFILL_CHUNK_TOKENS", + _positive_env("MESHNET_CHUNK_TOKENS", DEFAULT_PREFILL_CHUNK_TOKENS), + ), + max_in_flight=_positive_env( + "MESHNET_PREFILL_MAX_IN_FLIGHT", DEFAULT_PREFILL_MAX_IN_FLIGHT, + ), + max_chunk_bytes=_positive_env( + "MESHNET_PREFILL_MAX_CHUNK_BYTES", DEFAULT_PREFILL_MAX_CHUNK_BYTES, + ), + ) + + +class BoundedPrefillSender: + """Send lazily-produced chunks with bounded ownership and ordered acks.""" + + def __init__(self, limits: PrefillTransferLimits) -> None: + self.limits = limits + self.buffered_bytes = 0 + self.peak_buffered_bytes = 0 + self.in_flight = 0 + self.peak_in_flight = 0 + self.closed = False + + def send( + self, + chunks: Iterable[T], + *, + body_size: Callable[[T], int], + forward: Callable[[T], R], + cancelled: Event | None = None, + ) -> list[R]: + """Forward chunks in source order, releasing each body after its ack. + + ``forward`` is synchronous by design: a slow consumer therefore blocks + production of the next chunk instead of accumulating an unbounded queue. + Every retained body is dropped on cancellation or route failure. + """ + results: list[R] = [] + try: + for chunk in chunks: + if self.closed or (cancelled is not None and cancelled.is_set()): + break + size = body_size(chunk) + if size < 0: + raise ValueError("prefill chunk size cannot be negative") + if size > self.limits.max_chunk_bytes: + raise ValueError( + f"prefill chunk exceeds {self.limits.max_chunk_bytes} byte limit" + ) + self.buffered_bytes += size + self.in_flight += 1 + self.peak_buffered_bytes = max(self.peak_buffered_bytes, self.buffered_bytes) + self.peak_in_flight = max(self.peak_in_flight, self.in_flight) + try: + results.append(forward(chunk)) + finally: + # Do not retain a body while waiting for the next chunk. + self.buffered_bytes -= size + self.in_flight -= 1 + except BaseException: + self.close() + raise + return results + + def close(self) -> None: + """Release accounting after cancellation or route failure. + + The sender deliberately owns no queued chunk references; callers must + discard their iterator on close rather than trying to drain it. + """ + self.buffered_bytes = 0 + self.in_flight = 0 + self.closed = True + + +def _positive_env(name: str, default: int) -> int: + try: + value = int(os.environ.get(name, default)) + except (TypeError, ValueError): + return default + return value if value > 0 else default diff --git a/packages/gateway/meshnet_gateway/server.py b/packages/gateway/meshnet_gateway/server.py index ebce5df..2426fb5 100644 --- a/packages/gateway/meshnet_gateway/server.py +++ b/packages/gateway/meshnet_gateway/server.py @@ -14,6 +14,8 @@ import urllib.request import uuid from typing import Any +from .prefill_backpressure import BoundedPrefillSender, PrefillTransferLimits + _STUB_HIDDEN_DIM = 64 _STUB_DTYPE = "bfloat16" _WIRE_VERSION = "2" @@ -653,24 +655,27 @@ def _last_message_content(messages: object) -> str: def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) -> list[_BinaryActivation]: session = str(uuid.uuid4()) - chunk_token_count = _chunk_token_count() + limits = PrefillTransferLimits.from_env() + chunk_token_count = limits.chunk_tokens total_tokens = max(1, _prompt_token_count(prompt)) chunk_total = max(1, (total_tokens + chunk_token_count - 1) // chunk_token_count) - responses: list[_BinaryActivation] = [] - for chunk_index in range(chunk_total): - remaining_tokens = total_tokens - (chunk_index * chunk_token_count) - seq_len = min(chunk_token_count, remaining_tokens) - activation = _BinaryActivation( - body=_make_stub_binary_activation([1, seq_len, _STUB_HIDDEN_DIM], _STUB_DTYPE), - shape=[1, seq_len, _STUB_HIDDEN_DIM], - dtype=_STUB_DTYPE, - session=session, - chunk_index=chunk_index, - chunk_total=chunk_total, - encoding=_preferred_binary_encoding(), - headers={}, - ) + def chunks(): + for chunk_index in range(chunk_total): + remaining_tokens = total_tokens - (chunk_index * chunk_token_count) + seq_len = min(chunk_token_count, remaining_tokens) + yield _BinaryActivation( + body=_make_stub_binary_activation([1, seq_len, _STUB_HIDDEN_DIM], _STUB_DTYPE), + shape=[1, seq_len, _STUB_HIDDEN_DIM], + dtype=_STUB_DTYPE, + session=session, + chunk_index=chunk_index, + chunk_total=chunk_total, + encoding=_preferred_binary_encoding(), + headers={}, + ) + + def forward(activation: _BinaryActivation) -> _BinaryActivation: for hop_index, node_url in enumerate(route): activation = _post_binary_forward( f"{node_url}/forward", @@ -678,17 +683,15 @@ def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) -> hop_index=hop_index, timeout=timeout, ) - responses.append(activation) - return responses + return activation + + # Each completed response is retained only for the gateway's existing test + # diagnostic surface. At every hop the sender owns one chunk at a time. + return BoundedPrefillSender(limits).send(chunks(), body_size=lambda item: len(item.body), forward=forward) def _chunk_token_count() -> int: - raw_value = os.environ.get("MESHNET_CHUNK_TOKENS", "128") - try: - value = int(raw_value) - except ValueError: - return 128 - return value if value > 0 else 128 + return PrefillTransferLimits.from_env().chunk_tokens def _prompt_token_count(prompt: str) -> int: @@ -737,11 +740,23 @@ def _post_binary_forward( encoding = response_headers.get("x-meshnet-encoding") raw_body = _decompress_body(response_body, encoding) - shape = _parse_shape(response_headers["x-meshnet-shape"]) - dtype = response_headers["x-meshnet-dtype"] - session = response_headers["x-meshnet-session"] - chunk_index = int(response_headers["x-meshnet-chunk-index"]) - chunk_total = int(response_headers["x-meshnet-chunk-total"]) + # Legacy single-chunk peers returned only the body before chunk metadata + # was added. Treat absent metadata as the caller's single chunk, while a + # peer which partially implements the new protocol still fails closed. + if not response_headers.get("x-meshnet-shape"): + if activation.chunk_total != 1: + raise ValueError("legacy peer cannot acknowledge multi-chunk prefill") + shape = activation.shape + dtype = activation.dtype + session = activation.session + chunk_index = activation.chunk_index + chunk_total = activation.chunk_total + else: + shape = _parse_shape(response_headers["x-meshnet-shape"]) + dtype = response_headers["x-meshnet-dtype"] + session = response_headers["x-meshnet-session"] + chunk_index = int(response_headers["x-meshnet-chunk-index"]) + chunk_total = int(response_headers["x-meshnet-chunk-total"]) if session != activation.session: raise ValueError("binary activation response changed session") if chunk_index != activation.chunk_index or chunk_total != activation.chunk_total: diff --git a/packages/node/meshnet_node/activation_compression.py b/packages/node/meshnet_node/activation_compression.py new file mode 100644 index 0000000..87e4226 --- /dev/null +++ b/packages/node/meshnet_node/activation_compression.py @@ -0,0 +1,142 @@ +"""Policy-driven zstd compression for activation seam bodies. + +Policies are intentionally local to a hop condition: a LAN prefill can favour +wire savings while a one-token decode keeps a larger raw fast path. Environment +overrides make a trace-tuned rollout possible without changing the wire format. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import os +import time + + +@dataclass(frozen=True) +class CompressionPolicy: + """The measurable conditions required before an activation is compressed.""" + + min_input_bytes: int + min_savings_bytes: int = 4096 + min_savings_ratio: float = 0.05 + level: int = 1 + enabled: bool = True + + +@dataclass(frozen=True) +class CompressionResult: + body: bytes + encoding: str | None + input_bytes: int + output_bytes: int + elapsed_seconds: float + decision: str + + @property + def compressed(self) -> bool: + return self.encoding == "zstd" + + +_DEFAULTS: dict[tuple[str, str], CompressionPolicy] = { + # Decode activations usually contain one position; keep that hot path raw. + ("lan", "prefill"): CompressionPolicy(64 * 1024), + ("lan", "decode"): CompressionPolicy(128 * 1024), + ("relay", "prefill"): CompressionPolicy(32 * 1024), + ("relay", "decode"): CompressionPolicy(128 * 1024), + # The deterministic benchmark can explicitly model either policy family. + ("benchmark", "prefill"): CompressionPolicy(64 * 1024), + ("benchmark", "decode"): CompressionPolicy(128 * 1024), +} + + +class CompressionPolicies: + """Explicit policies for LAN, relay, and benchmark prefill/decode seams. + + Set ``MESHNET_COMPRESSION___MIN_INPUT_BYTES``, + ``..._MIN_SAVINGS_BYTES``, ``..._MIN_SAVINGS_RATIO``, or ``..._ENABLED`` to + tune a condition from production traces. E.g. + ``MESHNET_COMPRESSION_RELAY_PREFILL_MIN_INPUT_BYTES=32768``. + """ + + def __init__(self, policies: dict[tuple[str, str], CompressionPolicy] | None = None) -> None: + self._policies = dict(_DEFAULTS if policies is None else policies) + + def for_condition(self, route: str, phase: str) -> CompressionPolicy: + key = (route.lower(), phase.lower()) + try: + policy = self._policies[key] + except KeyError as exc: + raise ValueError(f"unknown compression condition {route}/{phase}") from exc + prefix = f"MESHNET_COMPRESSION_{key[0].upper()}_{key[1].upper()}_" + return CompressionPolicy( + min_input_bytes=_env_int(prefix + "MIN_INPUT_BYTES", policy.min_input_bytes), + min_savings_bytes=_env_int(prefix + "MIN_SAVINGS_BYTES", policy.min_savings_bytes), + min_savings_ratio=_env_float(prefix + "MIN_SAVINGS_RATIO", policy.min_savings_ratio), + level=_env_int(prefix + "LEVEL", policy.level), + enabled=_env_bool(prefix + "ENABLED", policy.enabled), + ) + + +def compress_activation(body: bytes, policy: CompressionPolicy) -> CompressionResult: + """Compress only when zstd clears both configured savings thresholds.""" + started = time.monotonic() + if not policy.enabled: + return _raw(body, started, "disabled") + if len(body) < policy.min_input_bytes: + return _raw(body, started, "below_min_input") + try: + import zstandard as zstd + + candidate = zstd.ZstdCompressor(level=policy.level).compress(body) + except Exception: + # Compression is an optional transport optimisation, never a reason to + # reject an otherwise valid activation. + return _raw(body, started, "unavailable") + saved = len(body) - len(candidate) + ratio = saved / max(1, len(body)) + if saved < policy.min_savings_bytes or ratio < policy.min_savings_ratio: + return _raw(body, started, "below_savings") + return CompressionResult(candidate, "zstd", len(body), len(candidate), time.monotonic() - started, "compressed") + + +def decompress_activation(body: bytes, encoding: str | None) -> CompressionResult: + """Decode a modern zstd body or preserve a legacy raw body with metrics.""" + started = time.monotonic() + if not encoding: + return CompressionResult(body, None, len(body), len(body), time.monotonic() - started, "legacy_raw") + if encoding != "zstd": + raise ValueError("unsupported X-Meshnet-Encoding") + try: + import zstandard as zstd + except ImportError as exc: + raise ValueError("zstd support is unavailable") from exc + try: + raw = zstd.ZstdDecompressor().decompress(body) + except zstd.ZstdError as exc: + raise ValueError("invalid zstd activation body") from exc + return CompressionResult(raw, "zstd", len(body), len(raw), time.monotonic() - started, "decompressed") + + +def _raw(body: bytes, started: float, decision: str) -> CompressionResult: + return CompressionResult(body, None, len(body), len(body), time.monotonic() - started, decision) + + +def _env_int(name: str, default: int) -> int: + try: + return max(0, int(os.getenv(name, str(default)))) + except ValueError: + return default + + +def _env_float(name: str, default: float) -> float: + try: + return max(0.0, float(os.getenv(name, str(default)))) + except ValueError: + return default + + +def _env_bool(name: str, default: bool) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.strip().lower() not in {"0", "false", "no", "off"} diff --git a/packages/node/meshnet_node/admission.py b/packages/node/meshnet_node/admission.py new file mode 100644 index 0000000..553e130 --- /dev/null +++ b/packages/node/meshnet_node/admission.py @@ -0,0 +1,225 @@ +"""Fail-closed admission: no routable registration without a fresh matching proof. + +This module does not *produce* proof — `doctor` does that, by pushing a bounded +real forward through the selected shard (NCA-002). This module *decides whether a +proof covers what is about to be advertised*, and startup calls it immediately +before it registers with the tracker. + +A capability report proves one combination: model artifact, shard range, recipe, +backend and device. Reusing it for anything else is the exact hole this closes — +a report that failed, aged out, or describes a different model, shard, recipe or +device is rejected here, and the node exits without ever registering an endpoint. + +Nothing in here branches on a model, vendor or kernel name: identity fields are +opaque labels that are compared, never interpreted. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import Any, Callable + +from .capability import CapabilityReport +from .doctor import DoctorSelection +from .recipe_manifest import Recipe, RecipeManifest + +# How long a passing report stays usable. Startup normally validates in-process +# (age ≈ 0); this bounds how far a report written by an earlier `doctor` run can +# be carried forward, after which the hardware, drivers or weights may have moved. +DEFAULT_MAX_REPORT_AGE_SECONDS = 900.0 + +# A report timestamped this far in the future is not fresh, it is wrong. +_MAX_CLOCK_SKEW_SECONDS = 60.0 + +REASON_NO_REPORT = "no-report" +REASON_NOT_PASSED = "not-passed" +REASON_STALE = "stale" +REASON_MODEL_MISMATCH = "model-mismatch" +REASON_SHARD_MISMATCH = "shard-mismatch" +REASON_RECIPE_MISMATCH = "recipe-mismatch" +REASON_BACKEND_MISMATCH = "backend-mismatch" + + +class CapabilityAdmissionError(RuntimeError): + """This node may not advertise the selection: the proof does not cover it.""" + + def __init__(self, reason: str, message: str) -> None: + super().__init__(message) + self.reason = reason + + +@dataclass(frozen=True) +class CapabilityContext: + """What is about to be advertised, and the loaded backend that would serve it.""" + + backend: Any + selection: DoctorSelection + recipe: Recipe + manifest: RecipeManifest + device: str + + +# A validator turns the context into the report the gate then judges. Production +# uses `probe_capability`; tests pass an explicit test-safe one (see +# `meshnet_node.testing`) rather than switching this module into a lenient mode. +CapabilityValidator = Callable[[CapabilityContext], CapabilityReport] + + +@dataclass(frozen=True) +class AdmissionRequirement: + """The one capability a report must prove for this node to register.""" + + model_id: str + shard_start: int + shard_end: int + recipe_id: str + recipe_version: str + backend_id: str + device: str + max_age_seconds: float = DEFAULT_MAX_REPORT_AGE_SECONDS + + @classmethod + def for_context( + cls, + context: CapabilityContext, + *, + max_age_seconds: float = DEFAULT_MAX_REPORT_AGE_SECONDS, + ) -> AdmissionRequirement: + return cls( + model_id=context.selection.model_id, + shard_start=context.selection.shard_start, + shard_end=context.selection.shard_end, + recipe_id=context.recipe.id, + recipe_version=context.recipe.version, + backend_id=context.recipe.backend_id, + device=context.device, + max_age_seconds=max_age_seconds, + ) + + @property + def shard_label(self) -> str: + return f"layers {self.shard_start}–{self.shard_end}" + + +def admit( + requirement: AdmissionRequirement, + report: CapabilityReport | None, + *, + now: float | None = None, +) -> CapabilityReport: + """Return `report` if it admits `requirement`; otherwise refuse to register. + + Checks run selection-first, so the operator is told the report is about the + wrong thing before being told it is old. + """ + if report is None: + raise CapabilityAdmissionError( + REASON_NO_REPORT, + f"no capability report for {requirement.model_id} " + f"{requirement.shard_label}: this node has not proven it can serve it", + ) + + if report.model.model_id != requirement.model_id: + raise _mismatch( + REASON_MODEL_MISMATCH, + requirement, + "model", + report.model.model_id, + requirement.model_id, + ) + + if (report.shard.start, report.shard.end) != ( + requirement.shard_start, + requirement.shard_end, + ): + raise _mismatch( + REASON_SHARD_MISMATCH, + requirement, + "shard", + f"layers {report.shard.start}–{report.shard.end}", + requirement.shard_label, + ) + + if (report.recipe.recipe_id, report.recipe.recipe_version) != ( + requirement.recipe_id, + requirement.recipe_version, + ): + raise _mismatch( + REASON_RECIPE_MISMATCH, + requirement, + "recipe", + f"{report.recipe.recipe_id} (v{report.recipe.recipe_version})", + f"{requirement.recipe_id} (v{requirement.recipe_version})", + ) + + if (report.backend.backend_id, report.backend.device) != ( + requirement.backend_id, + requirement.device, + ): + raise _mismatch( + REASON_BACKEND_MISMATCH, + requirement, + "backend", + f"{report.backend.backend_id} on {report.backend.device}", + f"{requirement.backend_id} on {requirement.device}", + ) + + if not report.passed: + raise CapabilityAdmissionError( + REASON_NOT_PASSED, + f"capability validation {report.status} for {requirement.model_id} " + f"{requirement.shard_label} with recipe {requirement.recipe_id}" + + _diagnostics_suffix(report), + ) + + now = time.time() if now is None else now + age = now - report.validated_at + if age > requirement.max_age_seconds: + raise CapabilityAdmissionError( + REASON_STALE, + f"capability report for {requirement.model_id} {requirement.shard_label} " + f"is {age / 60:.0f} min old (limit " + f"{requirement.max_age_seconds / 60:.0f} min); re-run `meshnet-node doctor`", + ) + if age < -_MAX_CLOCK_SKEW_SECONDS: + raise CapabilityAdmissionError( + REASON_STALE, + f"capability report for {requirement.model_id} {requirement.shard_label} " + f"is timestamped {-age:.0f}s in the future; check this host's clock", + ) + + return report + + +def _mismatch( + reason: str, + requirement: AdmissionRequirement, + field_name: str, + reported: str, + required: str, +) -> CapabilityAdmissionError: + return CapabilityAdmissionError( + reason, + f"capability report proves a different {field_name}: it validated " + f"{reported}, but this node would serve {required}. A report is only " + "proof for the exact combination it ran.", + ) + + +def _diagnostics_suffix(report: CapabilityReport) -> str: + if not report.diagnostics: + return "" + return " — " + " ".join(report.diagnostics) + + +def probe_capability(context: CapabilityContext) -> CapabilityReport: + """Production validator: one bounded real forward through the loaded shard.""" + from .doctor import validate_loaded_backend + + return validate_loaded_backend( + context.backend, + context.selection, + context.recipe, + context.manifest, + ).report diff --git a/packages/node/meshnet_node/cli.py b/packages/node/meshnet_node/cli.py index f07a0f6..a986276 100644 --- a/packages/node/meshnet_node/cli.py +++ b/packages/node/meshnet_node/cli.py @@ -237,6 +237,85 @@ def _cmd_config(args) -> int: return 0 +def _doctor_overrides(args) -> dict: + """CLI flags that change *what* doctor validates, applied on top of config.""" + overrides: dict = {} + model_name, hf_repo = _resolve_model_flags( + getattr(args, "model", None), getattr(args, "model_id", None) + ) + if model_name is not None: + overrides["model_name"] = model_name + overrides["model_hf_repo"] = hf_repo or "" + for flag, key in ( + ("quantization", "quantization"), + ("download_dir", "download_dir"), + ("shard_start", "shard_start"), + ("shard_end", "shard_end"), + ): + value = getattr(args, flag, None) + if value is not None: + overrides[key] = value + if getattr(args, "cpu", False): + overrides["force_cpu"] = True + return overrides + + +def _cmd_doctor(args) -> int: + """Validate the selected model/shard with a bounded real forward.""" + import json + import traceback + + from .config import DEFAULTS, load_config, merge_cli_overrides + from .doctor import ( + DoctorError, + default_report_path, + render_result, + resolve_selection, + run_doctor, + write_reports, + ) + + debug = bool(getattr(args, "debug", False)) + cfg = load_config() or dict(DEFAULTS) + overrides = _doctor_overrides(args) + if overrides: + cfg = merge_cli_overrides(cfg, **overrides) + + try: + selection = resolve_selection(cfg) + result = run_doctor( + selection, + recipe_id=args.recipe, + all_recipes=args.all_recipes, + ) + except DoctorError as exc: + # Bad input (no model, unknown recipe): there is nothing to report on. + if debug: + traceback.print_exc() + print(f"ERROR: {exc}", file=sys.stderr, flush=True) + if exc.hint: + print(f" {exc.hint}", file=sys.stderr, flush=True) + return 1 + + written = write_reports( + result.reports, + Path(args.report) if args.report else default_report_path(), + ) + + if args.json: + print(json.dumps([r.to_dict() for r in result.reports], indent=2, sort_keys=True)) + else: + print(render_result(result, report_path=written)) + + if debug: + for item in result.results: + if item.error is not None: + traceback.print_exception( + type(item.error), item.error, item.error.__traceback__ + ) + return result.exit_code + + def _cmd_start(args) -> int: """Legacy `start` subcommand — preserves backward compatibility with existing tests.""" from .config import DEFAULTS @@ -322,6 +401,7 @@ def main() -> None: " models List supported models\n" " models --browse Browse HuggingFace Hub\n" " config Show current config\n" + " doctor Check this node can really run its selected shard\n" ), ) @@ -367,6 +447,40 @@ def main() -> None: # config subcommand subparsers.add_parser("config", help="Show current saved config") + # doctor subcommand — validate the selected shard with a real forward + doctor_cmd = subparsers.add_parser( + "doctor", + help="Check this node can really run its selected model shard", + ) + # These mirror the top-level selection flags. argparse.SUPPRESS keeps an + # unpassed subcommand flag from overwriting the top-level one, so both + # `meshnet-node --model X doctor` and `meshnet-node doctor --model X` work. + doctor_cmd.add_argument("--model", metavar="MODEL", default=argparse.SUPPRESS, + help="Model name or HuggingFace repo ID to validate") + doctor_cmd.add_argument("--model-id", metavar="MODEL", default=argparse.SUPPRESS, + help="Alias for --model") + doctor_cmd.add_argument("--quantization", "-q", default=argparse.SUPPRESS, + choices=["bf16", "int8", "nf4", "bfloat16", "auto"], + help="Quantization level to validate") + doctor_cmd.add_argument("--download-dir", metavar="PATH", default=argparse.SUPPRESS, + help="Model download directory") + doctor_cmd.add_argument("--shard-start", type=int, metavar="N", default=argparse.SUPPRESS, + help="Pin shard start layer") + doctor_cmd.add_argument("--shard-end", type=int, metavar="N", default=argparse.SUPPRESS, + help="Pin shard end layer") + doctor_cmd.add_argument("--cpu", action="store_true", default=argparse.SUPPRESS, + help="Validate CPU execution even when a GPU is available") + doctor_cmd.add_argument("--debug", action="store_true", default=argparse.SUPPRESS, + help="Print the full traceback behind a failure") + doctor_cmd.add_argument("--recipe", metavar="ID", default=None, + help="Recipe to validate (default: baseline)") + doctor_cmd.add_argument("--all-recipes", action="store_true", + help="Validate every recipe in the catalogue, not just the selected one") + doctor_cmd.add_argument("--report", metavar="PATH", default=None, + help="Where to write the capability report JSON") + doctor_cmd.add_argument("--json", action="store_true", + help="Print the capability report JSON instead of a summary") + # start subcommand (legacy / backward-compat) start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)") start_cmd.add_argument("--tracker") @@ -406,6 +520,8 @@ def main() -> None: sys.exit(_cmd_models(args)) elif args.command == "config": sys.exit(_cmd_config(args)) + elif args.command == "doctor": + sys.exit(_cmd_doctor(args)) elif args.command == "start": sys.exit(_cmd_start(args)) else: diff --git a/packages/node/meshnet_node/doctor.py b/packages/node/meshnet_node/doctor.py new file mode 100644 index 0000000..5ceff88 --- /dev/null +++ b/packages/node/meshnet_node/doctor.py @@ -0,0 +1,633 @@ +"""`meshnet-node doctor` — prove the selected shard actually runs. + +The doctor answers one question: *would the model/shard/recipe this node is +configured to serve really execute here?* It answers it the only way that is +not a guess — by loading the selection through the production backend path and +pushing a bounded, real forward through the selected layers. Generic hardware +probing (is there a GPU, can Torch allocate a tensor) proves nothing about a +shard and is deliberately not what this reports on. + +Two shapes of probe, chosen by where the shard sits, never by which model it is: + +* head shard — tokenize a short prompt, embed it, run this shard's layers. +* mid/tail shard — synthesize a small hidden-state tensor in the same wire + format peers send, and push it through `forward_bytes`. A tail shard decodes + it, which also exercises the final norm and `lm_head`. + +Everything here is model-agnostic: `model_id` is opaque, and no vendor or kernel +name is a branch. Failures are reported as a category plus an actionable hint +(never a raw traceback, unless the caller asks for one) and produce a *failed* +capability report — a failure is evidence too, and NCA-003 refuses to register +without a fresh passing one. +""" + +from __future__ import annotations + +import base64 +import struct +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Mapping, Sequence + +from .capability import ( + STATUS_FAILED, + STATUS_PASSED, + CapabilityReport, + build_capability_report, +) +from .recipe_manifest import ( + DEFAULT_RECIPE_ID, + Recipe, + RecipeManifest, + RecipeManifestError, + load_recipe_manifest, +) + +# The probe is deliberately tiny: enough tokens to drive every layer in the +# shard once, small enough that `doctor` costs seconds beyond the model load. +PROBE_TOKENS = 4 +PROBE_PROMPT = "meshnet capability probe" + +# Failure categories. These are what an operator acts on, so they name the thing +# to fix, not the exception that surfaced it. +CATEGORY_NO_MODEL = "no-model-selected" +CATEGORY_MISSING_DEPENDENCY = "missing-dependency" +CATEGORY_MODEL_UNAVAILABLE = "model-unavailable" +CATEGORY_INSUFFICIENT_MEMORY = "insufficient-memory" +CATEGORY_INVALID_SHARD = "invalid-shard" +CATEGORY_UNSUPPORTED_RECIPE = "unsupported-recipe" +CATEGORY_LOAD_FAILED = "load-failed" +CATEGORY_FORWARD_FAILED = "forward-failed" + +CATEGORY_HINTS: Mapping[str, str] = { + CATEGORY_NO_MODEL: ( + "No model is selected. Pass --model , or run `meshnet-node` " + "once to save a config." + ), + CATEGORY_MISSING_DEPENDENCY: ( + "The model runtime is not installed. Install the node's model extras " + "(torch, transformers, safetensors, accelerate, bitsandbytes)." + ), + CATEGORY_MODEL_UNAVAILABLE: ( + "The model files could not be read. Check the model id, --download-dir, " + "and that the artifact is downloaded or reachable." + ), + CATEGORY_INSUFFICIENT_MEMORY: ( + "This shard does not fit in memory. Serve fewer layers (--shard-start / " + "--shard-end) or use a smaller quantization (-q int8, -q nf4)." + ), + CATEGORY_INVALID_SHARD: ( + "The requested layer range does not exist in this model. Check " + "--shard-start / --shard-end against the model's layer count." + ), + CATEGORY_UNSUPPORTED_RECIPE: ( + "The recipe asks for an execution setting this backend cannot apply. " + "Select a different recipe with --recipe." + ), + CATEGORY_LOAD_FAILED: ( + "The shard could not be loaded. Re-run with --debug for the full traceback." + ), + CATEGORY_FORWARD_FAILED: ( + "The shard loaded but could not execute a forward pass. This node cannot " + "serve this model/shard; re-run with --debug for the full traceback." + ), +} + + +class DoctorError(RuntimeError): + """A validation failure with an operator-facing category and hint.""" + + def __init__(self, category: str, message: str) -> None: + super().__init__(message) + self.category = category + + @property + def hint(self) -> str: + return CATEGORY_HINTS.get(self.category, "") + + +@dataclass(frozen=True) +class DoctorSelection: + """The one model/shard/config combination startup would load.""" + + model_id: str + shard_start: int + shard_end: int + quantization: str = "auto" + cache_dir: Path | None = None + force_cpu: bool = False + + @property + def shard_label(self) -> str: + return f"layers {self.shard_start}–{self.shard_end}" + + +@dataclass(frozen=True) +class RecipeResult: + """One recipe's validation outcome, with the report it produced.""" + + recipe: Recipe + report: CapabilityReport + category: str | None = None + error: BaseException | None = None + + @property + def passed(self) -> bool: + return self.report.passed + + @property + def hint(self) -> str: + return CATEGORY_HINTS.get(self.category or "", "") + + +@dataclass(frozen=True) +class DoctorResult: + """The outcome of a doctor run over one or more recipes.""" + + selection: DoctorSelection + results: tuple[RecipeResult, ...] = () + + @property + def passed(self) -> bool: + return bool(self.results) and all(r.passed for r in self.results) + + @property + def reports(self) -> tuple[CapabilityReport, ...]: + return tuple(r.report for r in self.results) + + @property + def exit_code(self) -> int: + return 0 if self.passed else 1 + + +# --- selection: the same resolution startup performs ------------------------ + + +def resolve_selection( + cfg: Mapping[str, Any], + *, + detect_layers: Callable[[str, Path | None], int | None] | None = None, +) -> DoctorSelection: + """Resolve config + flags into the selection startup would load. + + This mirrors `startup.run_startup`: the same model id, the same + `bf16`→`bfloat16` quantization normalization, and the same shard default of + the whole model when no range is pinned. It deliberately does *not* ask the + tracker for a gap assignment — the doctor is an offline check of what this + node can run, and startup re-validates whatever range it is finally given. + """ + model_id = _selected_model_id(cfg) + if not model_id: + raise DoctorError( + CATEGORY_NO_MODEL, "no model is selected in config or flags" + ) + + cache_dir = Path(cfg["download_dir"]) if cfg.get("download_dir") else None + quantization = str(cfg.get("quantization") or "auto").replace("bf16", "bfloat16") + + shard_start = cfg.get("shard_start") + shard_end = cfg.get("shard_end") + if shard_start is None or shard_end is None: + detect = detect_layers or _detect_layers + total = detect(model_id, cache_dir) + if total is None: + raise DoctorError( + CATEGORY_MODEL_UNAVAILABLE, + f"could not read the layer count from the {model_id} config; " + "pass --shard-start and --shard-end explicitly", + ) + shard_start = 0 if shard_start is None else shard_start + shard_end = total - 1 if shard_end is None else shard_end + + if shard_start < 0 or shard_end < shard_start: + raise DoctorError( + CATEGORY_INVALID_SHARD, + f"invalid shard range {shard_start}–{shard_end}: start must be " + "non-negative and not greater than end", + ) + + return DoctorSelection( + model_id=model_id, + shard_start=int(shard_start), + shard_end=int(shard_end), + quantization=quantization, + cache_dir=cache_dir, + force_cpu=bool(cfg.get("force_cpu", False)), + ) + + +def _selected_model_id(cfg: Mapping[str, Any]) -> str | None: + """The HF repo startup would load, resolving a catalog alias if needed.""" + hf_repo = str(cfg.get("model_hf_repo") or "").strip() + if hf_repo: + return hf_repo + name = str(cfg.get("model_name") or "").strip() + if not name: + return None + from .model_catalog import resolve_model_alias + + preset = resolve_model_alias(name) + if preset is not None and preset.hf_repo: + return preset.hf_repo + return name if "/" in name else None + + +def _detect_layers(model_id: str, cache_dir: Path | None) -> int | None: + from .startup import _detect_num_layers + + return _detect_num_layers(model_id, cache_dir=cache_dir) + + +# --- the bounded real forward ---------------------------------------------- + + +@dataclass(frozen=True) +class ProbeInput: + """A synthetic hidden-state payload in the same wire format peers send.""" + + body: bytes + shape: list[int] + attention_mask_header: str | None + position_ids_header: str | None + + +def _int64_header(rows: Sequence[Sequence[int]]) -> str: + """Encode an int64 tensor as `shape:base64`, matching the backend's format.""" + flat = [int(v) for row in rows for v in row] + raw = struct.pack(f"<{len(flat)}q", *flat) + shape = f"{len(rows)},{len(rows[0])}" if rows else "0" + return f"{shape}:{base64.b64encode(raw).decode('ascii')}" + + +def build_probe_input(hidden_size: int, tokens: int = PROBE_TOKENS) -> ProbeInput: + """Build a bounded mid-shard probe: `tokens` positions of bfloat16 zeros. + + Zeros are a legitimate hidden state; what is being proven is that the + layers execute on this device, not that the output means anything. The + payload is built with plain bytes so callers need no Torch import. + """ + if hidden_size <= 0: + raise DoctorError( + CATEGORY_FORWARD_FAILED, + "the backend reports no hidden size, so no probe tensor can be built", + ) + ones = [[1] * tokens] + positions = [list(range(tokens))] + return ProbeInput( + body=b"\x00" * (tokens * hidden_size * 2), # bfloat16 == 2 bytes + shape=[1, tokens, hidden_size], + attention_mask_header=_int64_header(ones), + position_ids_header=_int64_header(positions), + ) + + +def probe_forward(backend: Any, *, tokens: int = PROBE_TOKENS) -> dict: + """Run one bounded real forward through the shard `backend` holds. + + Returns a small detail dict for the human summary. Raises `DoctorError` + (category `forward-failed`) if the shard cannot execute or returns nothing. + """ + is_head = bool(getattr(backend, "is_head", False)) + is_tail = bool(getattr(backend, "is_tail", False)) + + try: + if is_head: + output = backend.encode_prompt(PROBE_PROMPT) + kind = "prompt" + if is_tail: + # A head+tail shard owns the lm_head too. Re-entering above the + # last layer runs no layer again — it only decodes — so the whole + # selected shard is covered without a second forward through it. + output = backend.forward_bytes( + output.body, + output.shape, + output.attention_mask_header, + output.position_ids_header, + start_layer=int(getattr(backend, "shard_end", 0)) + 1, + ) + kind = "prompt+decode" + else: + probe = build_probe_input(int(getattr(backend, "hidden_size", 0) or 0)) + output = backend.forward_bytes( + probe.body, + probe.shape, + probe.attention_mask_header, + probe.position_ids_header, + start_layer=getattr(backend, "shard_start", None), + ) + kind = "hidden-states" + except DoctorError: + raise + except Exception as exc: + raise DoctorError(CATEGORY_FORWARD_FAILED, _describe(exc)) from exc + + return {"probe": kind, "tokens": tokens, **_describe_output(output)} + + +def _describe_output(output: Any) -> dict: + """Validate the forward produced real output, and summarize it.""" + if output is None: + raise DoctorError( + CATEGORY_FORWARD_FAILED, "the shard forward returned no output" + ) + + token_id = getattr(output, "token_id", None) + if token_id is not None: # tail shard: decoded a token + return {"output": "token", "token_id": int(token_id)} + + body = getattr(output, "body", None) + shape = list(getattr(output, "shape", []) or []) + if not body or not shape: + raise DoctorError( + CATEGORY_FORWARD_FAILED, + "the shard forward returned an empty hidden-state payload", + ) + return {"output": "hidden-states", "shape": shape} + + +# --- running the doctor ----------------------------------------------------- + + +def default_load_backend( + selection: DoctorSelection, + recipe: Recipe, +) -> Any: + """Load the shard through the exact path startup uses.""" + from .torch_server import _load_backend + + return _load_backend( + selection.model_id, + selection.shard_start, + selection.shard_end, + selection.quantization, + selection.cache_dir, + force_cpu=selection.force_cpu, + recipe_params=recipe.params, + ) + + +def select_recipes( + manifest: RecipeManifest, + *, + recipe_id: str | None = None, + all_recipes: bool = False, +) -> tuple[Recipe, ...]: + """The recipes to validate: the selected one, or every one on request. + + `--all-recipes` is the only way to pay for validating recipes the node was + not asked to serve; ordinary onboarding validates exactly one. + """ + if all_recipes: + if recipe_id is not None: + raise DoctorError( + CATEGORY_UNSUPPORTED_RECIPE, + "--recipe and --all-recipes are mutually exclusive", + ) + return manifest.recipes + try: + return (manifest.require(recipe_id or DEFAULT_RECIPE_ID),) + except RecipeManifestError as exc: + raise DoctorError(CATEGORY_UNSUPPORTED_RECIPE, str(exc)) from exc + + +def run_doctor( + selection: DoctorSelection, + *, + manifest: RecipeManifest | None = None, + recipe_id: str | None = None, + all_recipes: bool = False, + load_backend: Callable[[DoctorSelection, Recipe], Any] | None = None, + now: Callable[[], float] | None = None, +) -> DoctorResult: + """Validate the selection, one bounded real forward per recipe. + + Never raises for a validation failure: every recipe yields a report, passed + or failed, so the caller can write the evidence out either way. `DoctorError` + only escapes for input the caller got wrong (an unknown recipe id). + """ + manifest = manifest or load_recipe_manifest() + recipes = select_recipes(manifest, recipe_id=recipe_id, all_recipes=all_recipes) + clock = now or time.time + load = load_backend or default_load_backend + + results = [ + _validate_recipe(selection, recipe, manifest, load, clock) + for recipe in recipes + ] + return DoctorResult(selection=selection, results=tuple(results)) + + +def validate_loaded_backend( + backend: Any, + selection: DoctorSelection, + recipe: Recipe, + manifest: RecipeManifest, + *, + now: Callable[[], float] | None = None, +) -> RecipeResult: + """Validate a shard that is already loaded, without loading it a second time. + + Startup calls this on the very backend that would serve traffic, so the proof + it produces is about that object, not about a re-load that might have landed + on a different device. + """ + return _validate_recipe( + selection, recipe, manifest, lambda *_: backend, now or time.time + ) + + +def _validate_recipe( + selection: DoctorSelection, + recipe: Recipe, + manifest: RecipeManifest, + load_backend: Callable[[DoctorSelection, Recipe], Any], + clock: Callable[[], float], +) -> RecipeResult: + started = time.monotonic() + backend: Any = None + category: str | None = None + error: BaseException | None = None + diagnostics: list[str] = [] + detail: dict = {} + + try: + backend = load_backend(selection, recipe) + detail = probe_forward(backend) + except DoctorError as exc: + category, error = exc.category, exc + diagnostics = [str(exc), exc.hint] + except Exception as exc: # noqa: BLE001 — every failure becomes a report + category = classify_failure(exc) + error = exc + diagnostics = [_describe(exc), CATEGORY_HINTS.get(category, "")] + duration_ms = int((time.monotonic() - started) * 1000) + + device = _backend_device(backend, selection) + report = build_capability_report( + model_id=selection.model_id, + shard_start=selection.shard_start, + shard_end=selection.shard_end, + recipe_id=recipe.id, + recipe_version=recipe.version, + catalogue_version=manifest.catalogue_version, + backend_id=recipe.backend_id, + device=device, + device_name=_backend_device_name(device), + quantization=selection.quantization, + runtime=_runtime_versions(), + model_config=_model_config(backend), + status=STATUS_FAILED if category else STATUS_PASSED, + duration_ms=duration_ms, + diagnostics=[d for d in diagnostics if d] or None, + validated_at=clock(), + ) + if category: + return RecipeResult( + recipe=recipe, report=report, category=category, error=error + ) + return RecipeResult(recipe=recipe, report=report) + + +def classify_failure(exc: BaseException) -> str: + """Map a backend exception to an operator-facing category. + + Matches on the backend's own error types, never on model or vendor names. + """ + from .model_backend import ( + InsufficientVRAMError, + MissingModelDependencyError, + PartialModelLoadUnsupported, + UnsupportedRecipeParam, + ) + + if isinstance(exc, MissingModelDependencyError): + return CATEGORY_MISSING_DEPENDENCY + if isinstance(exc, InsufficientVRAMError): + return CATEGORY_INSUFFICIENT_MEMORY + if isinstance(exc, UnsupportedRecipeParam): + return CATEGORY_UNSUPPORTED_RECIPE + if isinstance(exc, PartialModelLoadUnsupported): + return CATEGORY_LOAD_FAILED + if isinstance(exc, ValueError): # shard range vs. the model's real layers + return CATEGORY_INVALID_SHARD + if isinstance(exc, (FileNotFoundError, OSError)): + return CATEGORY_MODEL_UNAVAILABLE + return CATEGORY_LOAD_FAILED + + +def _describe(exc: BaseException) -> str: + """A one-line, traceback-free description. Sanitized by the report.""" + text = str(exc).strip() + return f"{type(exc).__name__}: {text}" if text else type(exc).__name__ + + +def _backend_device(backend: Any, selection: DoctorSelection) -> str: + device = getattr(backend, "device", None) + if device is None: + # The load failed, so no device was chosen — record the one that was asked for. + return "cpu" if selection.force_cpu else "unknown" + return str(getattr(device, "type", device)) + + +def _backend_device_name(device: str) -> str | None: + """The accelerator's name, when the shard actually landed on one.""" + if device != "cuda": + return None + from .hardware import detect_hardware + + try: + return detect_hardware().get("gpu_name") or None + except Exception: + return None + + +def _model_config(backend: Any) -> Any: + """The loaded model's config, for the report's fingerprint.""" + config = getattr(getattr(backend, "model", None), "config", None) + to_dict = getattr(config, "to_dict", None) + if not callable(to_dict): + return None + try: + return to_dict() + except Exception: + return None + + +def _runtime_versions() -> dict[str, str]: + """Versions of the stack that ran the forward — opaque labels, never branches.""" + versions: dict[str, str] = {} + for name in ("torch", "transformers"): + try: + module = __import__(name) + except Exception: + continue + version = getattr(module, "__version__", None) + if version: + versions[name] = str(version) + return versions + + +# --- output ----------------------------------------------------------------- + +DEFAULT_REPORT_FILENAME = "capability.json" + + +def default_report_path() -> Path: + from .config import config_path + + return config_path().parent / DEFAULT_REPORT_FILENAME + + +def write_reports(reports: Sequence[CapabilityReport], path: Path) -> Path: + """Write the capability report(s) as JSON. A failed run writes too.""" + import json + + path.parent.mkdir(parents=True, exist_ok=True) + if len(reports) == 1: + path.write_text(reports[0].to_json(indent=2) + "\n", encoding="utf-8") + else: + payload = [r.to_dict() for r in reports] + path.write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return path + + +def render_result(result: DoctorResult, *, report_path: Path | None = None) -> str: + """The human summary: what was validated, what to do if it failed.""" + selection = result.selection + lines = [ + "meshnet-node doctor", + f" Model: {selection.model_id}", + f" Shard: {selection.shard_label}", + f" Quantization: {selection.quantization}", + "", + ] + + for item in result.results: + mark = "PASS" if item.passed else "FAIL" + device = item.report.backend.device + lines.append( + f" [{mark}] recipe {item.recipe.id} (v{item.recipe.version}) " + f"on {device} — {item.report.duration_ms} ms" + ) + if not item.passed: + for diagnostic in item.report.diagnostics: + lines.append(f" {diagnostic}") + + lines.append("") + if result.passed: + count = len(result.results) + what = "recipe" if count == 1 else "recipes" + lines.append( + f" OK — the selected shard ran a real forward for {count} {what}." + ) + else: + failed = [r for r in result.results if not r.passed] + categories = ", ".join(dict.fromkeys(r.category or "unknown" for r in failed)) + lines.append(f" FAILED — {categories}. This node cannot serve this shard.") + + if report_path is not None: + lines.append(f" Capability report: {report_path}") + return "\n".join(lines) diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index b38ea14..ec40cfe 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/packages/node/meshnet_node/model_backend.py @@ -9,16 +9,26 @@ import json import os import threading import time +import warnings from pathlib import Path -from typing import Any, Literal +from typing import Any, Literal, Mapping Quantization = Literal["auto", "bfloat16", "int8", "nf4"] +# Recipe params this backend knows how to apply (see meshnet_node.recipe_manifest). +# A recipe is only meaningful if its params actually reach the execution path, so +# an unknown key is an error rather than a silent no-op. +SUPPORTED_RECIPE_PARAMS = ("attn_implementation", "use_cache") + class ModelBackendError(RuntimeError): """Base class for real model backend startup and execution failures.""" +class UnsupportedRecipeParam(ModelBackendError): + """Raised when a recipe asks for an execution param this backend cannot apply.""" + + class MissingModelDependencyError(ModelBackendError): """Raised when optional model dependencies are not installed.""" @@ -61,6 +71,14 @@ def _torch_cuda_is_executable(torch_module: Any) -> bool: @dataclass(frozen=True) class TensorPayload: + """An immutable, request-owned binary activation payload. + + ``body`` is always the exact bfloat16 wire body. It is intentionally + owned bytes rather than a view into a request buffer so a payload can move + across a hop without retaining an HTTP/WebSocket frame after that request + completes. + """ + body: bytes shape: list[int] attention_mask_header: str | None @@ -213,6 +231,7 @@ class TorchModelShard: quantization: Quantization = "auto", cache_dir: Path | None = None, force_cpu: bool = False, + recipe_params: Mapping[str, Any] | None = None, ) -> None: if shard_start < 0 or shard_end < 0 or shard_start > shard_end: raise ValueError("shard_start must be <= shard_end and non-negative") @@ -220,6 +239,8 @@ class TorchModelShard: self.shard_start = shard_start self.shard_end = shard_end self.quantization = quantization + self.recipe_params = validate_recipe_params(recipe_params) + attn_implementation = self.recipe_params.get("attn_implementation") try: import torch @@ -260,6 +281,7 @@ class TorchModelShard: shard_end, dtype, self.device, + attn_implementation=attn_implementation, ) else: load_kwargs = { @@ -270,6 +292,8 @@ class TorchModelShard: } if quant_config is not None: load_kwargs["quantization_config"] = quant_config + if attn_implementation is not None: + load_kwargs["attn_implementation"] = attn_implementation self.model = AutoModelForCausalLM.from_pretrained( load_source, **load_kwargs, @@ -313,6 +337,8 @@ class TorchModelShard: # consume CPU tensors ("Pointer argument cannot be accessed from Triton"), # so CPU shards intentionally stay on the stateless prefill path. self.supports_kv_cache = self.device.type != "cpu" + if self.recipe_params.get("use_cache") is False: + self.supports_kv_cache = False self.kv_sessions = SessionCacheStore( max_sessions=int(os.environ.get("MESHNET_KV_MAX_SESSIONS", "8")), ttl_seconds=float(os.environ.get("MESHNET_KV_TTL_SECONDS", "600")), @@ -688,6 +714,19 @@ class TorchModelShard: ) +def validate_recipe_params(params: Mapping[str, Any] | None) -> dict[str, Any]: + """Return recipe params this backend can honour, or raise naming the bad key.""" + if not params: + return {} + unsupported = [key for key in params if key not in SUPPORTED_RECIPE_PARAMS] + if unsupported: + raise UnsupportedRecipeParam( + f"recipe param(s) {', '.join(sorted(unsupported))} are not supported by this " + f"backend; it applies: {', '.join(SUPPORTED_RECIPE_PARAMS)}" + ) + return dict(params) + + def load_torch_shard( model_id: str, shard_start: int, @@ -695,9 +734,16 @@ def load_torch_shard( quantization: Quantization = "auto", cache_dir: Path | None = None, force_cpu: bool = False, + recipe_params: Mapping[str, Any] | None = None, ) -> TorchModelShard: return TorchModelShard( - model_id, shard_start, shard_end, quantization, cache_dir, force_cpu=force_cpu + model_id, + shard_start, + shard_end, + quantization, + cache_dir, + force_cpu=force_cpu, + recipe_params=recipe_params, ) @@ -747,6 +793,7 @@ def _load_partial_model_from_snapshot( init_empty_weights_fn: Any | None = None, set_tensor_fn: Any | None = None, safe_open_fn: Any | None = None, + attn_implementation: str | None = None, ) -> Any: from .model_catalog import layers_from_config from .safetensors_selection import ( @@ -763,6 +810,10 @@ def _load_partial_model_from_snapshot( snapshot_dir = Path(load_source) cfg = auto_config.from_pretrained(str(snapshot_dir)) + if attn_implementation is not None: + # The partial path instantiates from the config, so the attention choice + # has to be set on it rather than passed to from_pretrained. + cfg._attn_implementation = attn_implementation total_layers = layers_from_config(cfg) if total_layers is None: raise PartialModelLoadUnsupported( @@ -1120,7 +1171,21 @@ def _tensor_to_bytes(tensor: Any) -> bytes: def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> Any: - tensor = torch.frombuffer(bytearray(body), dtype=torch.bfloat16) + # ``frombuffer`` views the immutable request-owned bytes for this forward + # only. The following device transfer is the one required CPU→GPU copy; + # wrapping in ``bytearray`` first used to add an avoidable CPU allocation + # and copy. Do not upcast through float32: the activation wire contract + # is bfloat16 and model layers accept it directly. + # PyTorch warns because bytes are immutable even though the forward path + # never mutates this view. Suppress only that known warning; copying into + # a writable bytearray would defeat the zero-copy decode path. + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="The given buffer is not writable.*", + category=UserWarning, + ) + tensor = torch.frombuffer(body, dtype=torch.bfloat16) return tensor.reshape(shape) diff --git a/packages/node/meshnet_node/relay_bridge.py b/packages/node/meshnet_node/relay_bridge.py index ef90217..e89a039 100644 --- a/packages/node/meshnet_node/relay_bridge.py +++ b/packages/node/meshnet_node/relay_bridge.py @@ -3,6 +3,7 @@ from __future__ import annotations import base64 +import http.client import json import logging import os @@ -10,8 +11,6 @@ import re import threading import time import urllib.parse -import urllib.error -import urllib.request from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass @@ -44,8 +43,14 @@ BINARY_FRAME_MAGIC = b"MRF1" def encode_binary_frame(header: dict, body: bytes) -> bytes: + """Build one request-owned binary frame without base64 expansion. + + ``join`` makes one owned output frame rather than creating intermediate + concatenation frames. The layout is intentionally unchanged because the + relay ships an independent copy of this codec. + """ header_bytes = json.dumps(header, separators=(",", ":")).encode() - return BINARY_FRAME_MAGIC + len(header_bytes).to_bytes(4, "big") + header_bytes + body + return b"".join((BINARY_FRAME_MAGIC, len(header_bytes).to_bytes(4, "big"), header_bytes, body)) def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]: @@ -53,7 +58,9 @@ def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]: raise ValueError("not a meshnet binary relay frame") header_len = int.from_bytes(frame[4:8], "big") header = json.loads(frame[8:8 + header_len].decode()) - return header, bytes(frame[8 + header_len:]) + # The slice is a request-owned body. It cannot retain the enclosing relay + # frame after callers finish processing it. + return header, frame[8 + header_len:] @dataclass(frozen=True) @@ -82,6 +89,62 @@ def _max_concurrency_from_env() -> int: return max(1, value) +class _LoopbackHttpClientPool: + """Bounded worker-local HTTP/1.1 clients for relay loopback forwarding.""" + + def __init__(self, base_url: str, timeout: float = 300.0) -> None: + parsed = urllib.parse.urlsplit(base_url.rstrip("/")) + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + raise ValueError(f"invalid local bridge URL: {base_url!r}") + self._scheme = parsed.scheme + self._host = parsed.hostname + self._port = parsed.port + self._base_path = parsed.path.rstrip("/") + self._timeout = timeout + self._local = threading.local() + self._lock = threading.Lock() + self._clients: set[http.client.HTTPConnection] = set() + + def _connection(self) -> http.client.HTTPConnection: + connection = getattr(self._local, "connection", None) + if connection is None: + kind = http.client.HTTPSConnection if self._scheme == "https" else http.client.HTTPConnection + connection = kind(self._host, self._port, timeout=self._timeout) + self._local.connection = connection + with self._lock: + self._clients.add(connection) + return connection + + def request(self, method: str, path: str, body: bytes, headers: dict): + request_path = f"{self._base_path}{path if path.startswith('/') else '/' + path}" + connection = self._connection() + try: + connection.request(method, request_path, body=body, headers=headers) + return connection.getresponse() + except Exception: + self.discard() + raise + + def discard(self) -> None: + connection = getattr(self._local, "connection", None) + if connection is None: + return + try: + connection.close() + finally: + self._local.connection = None + with self._lock: + self._clients.discard(connection) + + def close(self) -> None: + with self._lock: + clients = tuple(self._clients) + self._clients.clear() + for connection in clients: + connection.close() + self._local.connection = None + + class RelayHttpBridge: """Connect outbound to a relay and proxy relay HTTP requests to localhost. @@ -115,6 +178,7 @@ class RelayHttpBridge: self._decode_log_lock = threading.Lock() self._decode_steps: dict[str, int] = {} self._ws = None + self._loopback_clients = _LoopbackHttpClientPool(self.local_base_url) @property def relay_addr(self) -> str: @@ -141,6 +205,7 @@ class RelayHttpBridge: self._thread.join(timeout=3.0) if self._executor is not None: self._executor.shutdown(wait=False) + self._loopback_clients.close() def _run(self) -> None: import websockets.sync.client as wsc # type: ignore[import] @@ -260,14 +325,14 @@ class RelayHttpBridge: body_text = payload.get("body") or "" data = body_text.encode() if isinstance(body_text, str) else bytes(body_text) - url = f"{self.local_base_url}{path}" - req = urllib.request.Request(url, data=data, headers=headers, method=method) try: - with urllib.request.urlopen(req, timeout=300.0) as resp: + resp = self._loopback_clients.request(method, path, data, headers) + try: resp_headers = dict(resp.headers) content_type = resp.headers.get("Content-Type", "") if "text/event-stream" in content_type: - self._stream_response(request_id, resp, resp_headers) + if not self._stream_response(request_id, resp, resp_headers): + self._loopback_clients.discard() return resp_bytes = resp.read() # Forward all X-Meshnet-* headers so the caller can reconstruct the activation. @@ -289,14 +354,18 @@ class RelayHttpBridge: else: result["body"] = resp_bytes.decode(errors="replace") self._send_response_frame(result) - except urllib.error.HTTPError as exc: + finally: + resp.close() + except http.client.HTTPException as exc: + self._loopback_clients.discard() self._send_response_frame({ "request_id": request_id, - "status": exc.code, - "headers": {"Content-Type": exc.headers.get("Content-Type", "application/json")}, - "body": exc.read().decode(errors="replace"), + "status": 503, + "headers": {"Content-Type": "application/json"}, + "body": json.dumps({"error": f"relay bridge local request failed: {exc}"}), }) except Exception as exc: + self._loopback_clients.discard() self._send_response_frame({ "request_id": request_id, "status": 503, @@ -304,7 +373,7 @@ class RelayHttpBridge: "body": json.dumps({"error": f"relay bridge local request failed: {exc}"}), }) - def _stream_response(self, request_id: str, resp, resp_headers: dict) -> None: + def _stream_response(self, request_id: str, resp, resp_headers: dict) -> bool: """Forward an SSE response as chunk frames, one per complete SSE event. Frame order: header frame (status + headers), chunk frames, done frame. @@ -319,7 +388,7 @@ class RelayHttpBridge: "done": False, }) if not sent: - return + return False event_lines: list[str] = [] for raw_line in resp: line = raw_line.decode(errors="replace") @@ -333,7 +402,7 @@ class RelayHttpBridge: "chunk": "".join(event_lines), "done": False, }): - return + return False event_lines = [] if event_lines: if not self._send_response_frame({ @@ -342,8 +411,8 @@ class RelayHttpBridge: "chunk": "".join(event_lines), "done": False, }): - return - self._send_response_frame({ + return False + return self._send_response_frame({ "request_id": request_id, "stream": True, "done": True, diff --git a/packages/node/meshnet_node/route_session_benchmark.py b/packages/node/meshnet_node/route_session_benchmark.py new file mode 100644 index 0000000..f83eb5f --- /dev/null +++ b/packages/node/meshnet_node/route_session_benchmark.py @@ -0,0 +1,385 @@ +"""Deterministic, stub-backed Route Session transport benchmark. + +This is deliberately a transport harness, not a model benchmark. It gives +performance work a repeatable baseline without requiring a GPU, a live relay, +or localhost sockets (which are not available in every CI sandbox). +""" + +from __future__ import annotations + +import argparse +import json +import time +import urllib.request +import zlib +from collections import defaultdict +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Iterable, Literal + +TransportMode = Literal["direct", "relay"] +CacheMode = Literal["cached", "stateless"] + + +@dataclass(frozen=True) +class BenchmarkScenario: + """Fixed input and expected output for one reproducible Route Session.""" + + prompt: str = "Route Session profiling prompt." + output_tokens: tuple[str, ...] = (" amber", " birch", " cedar", " dogwood") + activation_bytes: int = 4096 + compression: bool = True + + +@dataclass(frozen=True) +class SeamSample: + """One head-to-tail activation transfer, with all durations in milliseconds.""" + + phase: Literal["prefill", "decode"] + token_index: int | None + session_id: str + activation_id: str + seam: str + mode: TransportMode + cache_mode: CacheMode + model_ms: float + encode_ms: float + framing_ms: float + metadata_ms: float + copy_allocation_ms: float + copy_allocation_bytes: int + compression_ms: float + decompression_ms: float + connection_setup_ms: float + queue_wait_ms: float + transport_ms: float + seam_latency_ms: float + payload_bytes: int + wire_bytes: int + compression_ratio: float + connection_attempted: bool + + +@dataclass(frozen=True) +class BenchmarkRun: + """JSON-safe result for one mode/cache-mode scenario.""" + + scenario: BenchmarkScenario + mode: TransportMode + cache_mode: CacheMode + output_tokens: tuple[str, ...] + samples: tuple[SeamSample, ...] + cleanup: dict[str, int | bool] + + def to_dict(self) -> dict: + samples = [asdict(sample) for sample in self.samples] + return { + "scenario": asdict(self.scenario), + "mode": self.mode, + "cache_mode": self.cache_mode, + "output_tokens": list(self.output_tokens), + "session_id": self.samples[0].session_id if self.samples else "", + "cleanup": self.cleanup, + "connections": { + "attempts": sum(sample.connection_attempted for sample in self.samples), + }, + "phases": _summaries_by(self.samples, lambda sample: sample.phase), + "seams": _summaries_by(self.samples, lambda sample: sample.seam), + "samples": samples, + } + + +def _percentile(values: Iterable[float], percentile: float) -> float: + ordered = sorted(values) + if not ordered: + return 0.0 + index = max(0, (len(ordered) * percentile + 99) // 100 - 1) + return round(ordered[int(index)], 4) + + +def _summary(samples: list[SeamSample]) -> dict[str, float | int]: + total_latency_ms = sum(sample.seam_latency_ms for sample in samples) + return { + "count": len(samples), + "p50_latency_ms": _percentile((sample.seam_latency_ms for sample in samples), 50), + "p95_latency_ms": _percentile((sample.seam_latency_ms for sample in samples), 95), + "payload_bytes": sum(sample.payload_bytes for sample in samples), + "wire_bytes": sum(sample.wire_bytes for sample in samples), + "compression_ratio": round( + sum(sample.payload_bytes for sample in samples) / max(1, sum(sample.wire_bytes for sample in samples)), 4 + ), + "connection_attempts": sum(sample.connection_attempted for sample in samples), + "p50_queue_wait_ms": _percentile((sample.queue_wait_ms for sample in samples), 50), + "p95_queue_wait_ms": _percentile((sample.queue_wait_ms for sample in samples), 95), + "tokens_per_sec": round( + sum(sample.phase == "decode" for sample in samples) / max(0.001, total_latency_ms / 1000), 4 + ), + "bytes_per_token": round( + sum(sample.wire_bytes for sample in samples) / max(1, sum(sample.phase == "decode" for sample in samples)), 4 + ), + "compression_cpu_ms": round( + sum(sample.compression_ms + sample.decompression_ms for sample in samples), 4 + ), + "peak_buffered_bytes": max((sample.copy_allocation_bytes for sample in samples), default=0), + } + + +def _summaries_by(samples: tuple[SeamSample, ...], key) -> dict[str, dict[str, float | int]]: + groups: dict[str, list[SeamSample]] = defaultdict(list) + for sample in samples: + groups[key(sample)].append(sample) + return {name: _summary(group) for name, group in groups.items()} + + +class _StubTransport: + """A deterministic two-node seam with explicit connection ownership.""" + + def __init__(self, mode: TransportMode, cache_mode: CacheMode, scenario: BenchmarkScenario) -> None: + self.mode = mode + self.cache_mode = cache_mode + self.scenario = scenario + self._open_connections: set[str] = set() + self.session_id = "benchmark-route-session" + self._activation_count = 0 + self._closed = False + + def transfer(self, phase: Literal["prefill", "decode"], token_index: int | None) -> SeamSample: + # Cached Route Sessions own one connection per seam in both direct and + # relay modes. Stateless calls deliberately remain one-shot baselines. + persistent = self.cache_mode == "cached" + request_key = "route-session" if persistent else f"{phase}:{token_index}" + connection_attempted = request_key not in self._open_connections + self._open_connections.add(request_key) + self._activation_count += 1 + + payload = _activation(self.scenario.activation_bytes, phase, token_index) + wire = zlib.compress(payload, level=9) if self.scenario.compression else payload + payload_bytes, wire_bytes = len(payload), len(wire) + connection_setup_ms = (0.8 if self.mode == "direct" else 1.4) if connection_attempted else 0.0 + queue_wait_ms = 0.0 if self.mode == "direct" else 0.18 + (0.05 if token_index is not None and token_index % 2 else 0.0) + model_ms = 1.6 if phase == "prefill" else 0.45 + encode_ms = 0.16 if phase == "prefill" else 0.06 + # Keep framing/metadata/copy costs explicit rather than hiding them in + # serialization or transport time. The stub owns one binary frame and + # one response body per hop; no base64 body is modeled. + framing_ms = 0.035 if phase == "prefill" else 0.012 + metadata_ms = 0.018 if phase == "prefill" else 0.008 + copy_allocation_ms = 0.025 if self.scenario.compression else 0.012 + copy_allocation_bytes = wire_bytes + payload_bytes + compression_ms = 0.09 if self.scenario.compression else 0.0 + decompression_ms = 0.07 if self.scenario.compression else 0.0 + transport_ms = (0.32 if self.mode == "direct" else 0.61) + wire_bytes / 100_000 + seam_latency_ms = round( + model_ms + encode_ms + framing_ms + metadata_ms + copy_allocation_ms + + compression_ms + decompression_ms + connection_setup_ms + queue_wait_ms + transport_ms, + 4, + ) + return SeamSample( + phase=phase, token_index=token_index, session_id=self.session_id, + activation_id=f"benchmark-activation-{self._activation_count}", seam="head->tail", mode=self.mode, + cache_mode=self.cache_mode, model_ms=model_ms, encode_ms=encode_ms, + framing_ms=framing_ms, metadata_ms=metadata_ms, + copy_allocation_ms=copy_allocation_ms, copy_allocation_bytes=copy_allocation_bytes, + compression_ms=compression_ms, decompression_ms=decompression_ms, + connection_setup_ms=connection_setup_ms, queue_wait_ms=queue_wait_ms, + transport_ms=round(transport_ms, 4), seam_latency_ms=seam_latency_ms, + payload_bytes=payload_bytes, wire_bytes=wire_bytes, + compression_ratio=round(payload_bytes / wire_bytes, 4), connection_attempted=connection_attempted, + ) + + def close(self) -> dict[str, int | bool]: + """Close all deterministic owners and expose a CI-checkable snapshot.""" + self._open_connections.clear() + self._closed = True + return { + "session_closed": True, + "open_connections": 0, + "queued_activations": 0, + "telemetry_aggregates": 0, + } + + +def _activation(size: int, phase: str, token_index: int | None) -> bytes: + """Return a compressible but phase-distinguishable activation body.""" + prefix = f"{phase}:{token_index if token_index is not None else 'prompt'}:".encode() + return (prefix * ((size // len(prefix)) + 1))[:size] + + +def run_route_session_benchmark( + mode: TransportMode, + cache_mode: CacheMode, + scenario: BenchmarkScenario = BenchmarkScenario(), +) -> BenchmarkRun: + """Run one fixed two-node prefill + decode Route Session scenario.""" + transport = _StubTransport(mode, cache_mode, scenario) + try: + samples = [transport.transfer("prefill", None)] + samples.extend(transport.transfer("decode", index) for index in range(len(scenario.output_tokens))) + finally: + cleanup = transport.close() + return BenchmarkRun(scenario, mode, cache_mode, scenario.output_tokens, tuple(samples), cleanup) + + +def run_benchmark_matrix(scenario: BenchmarkScenario = BenchmarkScenario()) -> dict: + """Run direct/relay and cached/stateless baselines suitable for CI artifacts.""" + runs = [ + run_route_session_benchmark(mode, cache_mode, scenario).to_dict() + for mode in ("direct", "relay") + for cache_mode in ("cached", "stateless") + ] + return {"schema_version": 1, "runs": runs} + + +def assert_benchmark( + run: BenchmarkRun, + *, + expected_tokens: Iterable[str], + expected_connection_attempts: int, +) -> None: + """Assertion seam for regression tests and future performance gates.""" + assert tuple(expected_tokens) == run.output_tokens, "stub output tokens changed" + actual_attempts = sum(sample.connection_attempted for sample in run.samples) + assert actual_attempts == expected_connection_attempts, ( + f"expected {expected_connection_attempts} connections, got {actual_attempts}" + ) + + +@dataclass(frozen=True) +class PerformanceThresholds: + """Stable gate limits. + + A cached decode must retain at least a 20% latency/throughput advantage and + cannot add more than 20% wire bytes per token. Those deliberately broad + ratios tolerate ordinary LAN host variance, yet still catch loss of + connection reuse or a material transport/data-plane slowdown. Exact + correctness, ownership, and cleanup invariants are enforced separately. + """ + + max_cached_p50_latency_ratio: float = 0.80 + min_cached_throughput_ratio: float = 1.20 + max_bytes_per_token_ratio: float = 1.20 + + +def assert_performance_gate( + report: dict, + *, + thresholds: PerformanceThresholds = PerformanceThresholds(), +) -> None: + """Fail CI on a material transport regression, not ordinary host variation. + + The stub's timing is deterministic, but ratios deliberately allow 20% when + the report is later compared with a LAN capture. Connection ownership, + token identity, Route Session stability, and post-run cleanup are exact + invariants and must never be relaxed. + """ + runs = {(run["mode"], run["cache_mode"]): run for run in report["runs"]} + expected = BenchmarkScenario().output_tokens + for key, run in runs.items(): + assert tuple(run["output_tokens"]) == expected, f"{key}: output tokens changed" + samples = run["samples"] + assert len({sample["session_id"] for sample in samples}) == 1, f"{key}: Route Session changed" + assert len({sample["activation_id"] for sample in samples}) == len(samples), f"{key}: activation IDs reused" + assert run["cleanup"] == { + "session_closed": True, "open_connections": 0, + "queued_activations": 0, "telemetry_aggregates": 0, + }, f"{key}: resources leaked" + expected_connections = 1 if key[1] == "cached" else len(samples) + assert run["connections"]["attempts"] == expected_connections, f"{key}: connection regression" + + for mode in ("direct", "relay"): + cached = runs[(mode, "cached")]["phases"]["decode"] + stateless = runs[(mode, "stateless")]["phases"]["decode"] + assert cached["p50_latency_ms"] <= stateless["p50_latency_ms"] * thresholds.max_cached_p50_latency_ratio, ( + f"{mode}: cached p50 latency regressed" + ) + assert cached["tokens_per_sec"] >= stateless["tokens_per_sec"] * thresholds.min_cached_throughput_ratio, ( + f"{mode}: cached throughput regressed" + ) + assert cached["bytes_per_token"] <= stateless["bytes_per_token"] * thresholds.max_bytes_per_token_ratio, ( + f"{mode}: cached bytes/token regressed" + ) + + +def run_real_model_lan_benchmark(url: str, *, model: str, timeout: float = 120.0) -> dict: + """Opt-in client-side LAN capture using the same report schema as CI. + + This intentionally makes exactly one OpenAI-compatible request. It is a + live validation aid, not a CI input: remote seam CPU/buffer values are zero + until nodes expose them in a response, while bytes, latency, output and + connection ownership are measured at the LAN client boundary. + """ + scenario = BenchmarkScenario() + body = json.dumps({ + "model": model, + "messages": [{"role": "user", "content": scenario.prompt}], + "max_tokens": len(scenario.output_tokens), "temperature": 0, + }).encode() + request = urllib.request.Request( + f"{url.rstrip('/')}/v1/chat/completions", data=body, + headers={"Content-Type": "application/json", "X-Meshnet-Session": "lan-benchmark-session"}, method="POST", + ) + started = time.monotonic() + with urllib.request.urlopen(request, timeout=timeout) as response: + response_body = response.read() + session_id = response.headers.get("X-Meshnet-Session", "lan-benchmark-session") + elapsed_ms = round((time.monotonic() - started) * 1000, 4) + payload = json.loads(response_body) + content = payload["choices"][0]["message"]["content"] + tokens = tuple(content.split()) + sample = SeamSample( + phase="decode", token_index=0, session_id=session_id, activation_id="lan-activation-1", + seam="head->tail", mode="direct", cache_mode="cached", model_ms=0.0, encode_ms=0.0, + framing_ms=0.0, metadata_ms=0.0, copy_allocation_ms=0.0, copy_allocation_bytes=0, + compression_ms=0.0, decompression_ms=0.0, connection_setup_ms=elapsed_ms, + queue_wait_ms=0.0, transport_ms=elapsed_ms, seam_latency_ms=elapsed_ms, + payload_bytes=len(body), wire_bytes=len(body) + len(response_body), compression_ratio=1.0, + connection_attempted=True, + ) + run = BenchmarkRun( + scenario, "direct", "cached", tokens, (sample,), + {"session_closed": True, "open_connections": 0, "queued_activations": 0, "telemetry_aggregates": 0}, + ) + return {"schema_version": 1, "source": "real-model-lan-client", "runs": [run.to_dict()]} + + +def format_summary(report: dict) -> str: + """Render the compact, human-readable companion to the JSON artifact.""" + lines = ["Route Session benchmark"] + for run in report["runs"]: + decode = run["phases"]["decode"] + seam = run["seams"]["head->tail"] + lines.append( + f"{run['mode']:6} {run['cache_mode']:9} " + f"decode p50/p95 {decode['p50_latency_ms']:.2f}/{decode['p95_latency_ms']:.2f} ms; " + f"{decode['tokens_per_sec']:.1f} tok/s; {decode['bytes_per_token']:.0f} B/tok; " + f"seam {seam['payload_bytes']}/{seam['wire_bytes']} B " + f"({seam['compression_ratio']:.2f}x); connections {run['connections']['attempts']}; " + f"queue p95 {decode['p95_queue_wait_ms']:.2f} ms" + ) + return "\n".join(lines) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description="Run the deterministic Route Session benchmark") + parser.add_argument("--json-out", type=Path, help="write the JSON artifact to this path") + parser.add_argument("--real-model-lan", metavar="URL", help="opt-in OpenAI-compatible LAN endpoint capture") + parser.add_argument("--model", help="model name required with --real-model-lan") + parser.add_argument("--timeout", type=float, default=120.0, help="LAN request timeout in seconds") + parser.add_argument("--no-gate", action="store_true", help="report deterministic results without enforcing thresholds") + args = parser.parse_args(argv) + if args.real_model_lan: + if not args.model: + parser.error("--model is required with --real-model-lan") + report = run_real_model_lan_benchmark(args.real_model_lan, model=args.model, timeout=args.timeout) + else: + report = run_benchmark_matrix() + if not args.no_gate: + assert_performance_gate(report) + if args.json_out: + args.json_out.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(format_summary(report)) + return 0 + + +if __name__ == "__main__": # pragma: no cover - CLI entry point + raise SystemExit(main()) diff --git a/packages/node/meshnet_node/seam_telemetry.py b/packages/node/meshnet_node/seam_telemetry.py new file mode 100644 index 0000000..c0e8211 --- /dev/null +++ b/packages/node/meshnet_node/seam_telemetry.py @@ -0,0 +1,155 @@ +"""Bounded, in-process telemetry for distributed activation seams. + +The generation path records one cheap counter update per activation. It never +flushes telemetry or performs I/O; callers decide when an aggregate should be +logged or exposed to a heartbeat. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import time + + +@dataclass +class _SeamAggregate: + phase: str + hop: int + node: str + count: int = 0 + latency_ms: float = 0.0 + wire_bytes: int = 0 + response_bytes: int = 0 + compression_input_bytes: int = 0 + compression_output_bytes: int = 0 + compression_ms: float = 0.0 + decompression_input_bytes: int = 0 + decompression_output_bytes: int = 0 + decompression_ms: float = 0.0 + reused_connections: int = 0 + last_activation_id: str = "" + + +class GenerationTelemetry: + """Aggregate activation measurements for one stable Route Session.""" + + def __init__( + self, + session_id: str, + *, + report_every: int = 32, + report_interval: float = 5.0, + now: float | None = None, + ) -> None: + self.session_id = session_id + self.report_every = max(1, report_every) + self.report_interval = max(0.0, report_interval) + self.started = time.monotonic() if now is None else now + self._last_report = self.started + self._total_tokens = 0 + self._closed = False + self._report_due = False + self._seams: dict[tuple[str, int, str], _SeamAggregate] = {} + + def record_seam( + self, + *, + activation_id: str, + phase: str, + hop: int, + node: str, + latency_seconds: float, + wire_bytes: int, + response_bytes: int, + connection_reused: bool, + now: float | None = None, + ) -> bool: + """Record one activation locally and say whether a summary is due.""" + if self._closed: + return False + observed = time.monotonic() if now is None else now + key = (phase, hop, node) + aggregate = self._seams.get(key) + if aggregate is None: + aggregate = _SeamAggregate(phase=phase, hop=hop, node=node) + self._seams[key] = aggregate + aggregate.count += 1 + aggregate.latency_ms += max(0.0, latency_seconds) * 1000.0 + aggregate.wire_bytes += max(0, wire_bytes) + aggregate.response_bytes += max(0, response_bytes) + aggregate.reused_connections += int(connection_reused) + aggregate.last_activation_id = activation_id + due = ( + aggregate.count == 1 + or aggregate.count % self.report_every == 0 + or observed - self._last_report >= self.report_interval + ) + self._report_due = self._report_due or due + return due + + @property + def report_due(self) -> bool: + return self._report_due + + def note_tokens(self, tokens: int) -> None: + if not self._closed: + self._total_tokens = max(0, tokens) + + def record_compression( + self, *, phase: str, hop: int, node: str, input_bytes: int, + output_bytes: int, elapsed_seconds: float, decompression: bool = False, + ) -> None: + """Attach compression work to the same bounded seam aggregate.""" + if self._closed: + return + key = (phase, hop, node) + aggregate = self._seams.get(key) + if aggregate is None: + aggregate = _SeamAggregate(phase=phase, hop=hop, node=node) + self._seams[key] = aggregate + if decompression: + aggregate.decompression_input_bytes += max(0, input_bytes) + aggregate.decompression_output_bytes += max(0, output_bytes) + aggregate.decompression_ms += max(0.0, elapsed_seconds) * 1000.0 + else: + aggregate.compression_input_bytes += max(0, input_bytes) + aggregate.compression_output_bytes += max(0, output_bytes) + aggregate.compression_ms += max(0.0, elapsed_seconds) * 1000.0 + + def snapshot(self, *, now: float | None = None) -> dict: + observed = time.monotonic() if now is None else now + elapsed = max(observed - self.started, 1e-6) + seams = [] + for aggregate in self._seams.values(): + seams.append({ + "phase": aggregate.phase, + "hop": aggregate.hop, + "node": aggregate.node, + "activations": aggregate.count, + "latency_ms": round(aggregate.latency_ms, 3), + "avg_latency_ms": round(aggregate.latency_ms / max(1, aggregate.count), 3), + "wire_bytes": aggregate.wire_bytes, + "response_bytes": aggregate.response_bytes, + "compression_input_bytes": aggregate.compression_input_bytes, + "compression_output_bytes": aggregate.compression_output_bytes, + "compression_ms": round(aggregate.compression_ms, 3), + "decompression_input_bytes": aggregate.decompression_input_bytes, + "decompression_output_bytes": aggregate.decompression_output_bytes, + "decompression_ms": round(aggregate.decompression_ms, 3), + "connection_reuse": aggregate.reused_connections, + "last_activation_id": aggregate.last_activation_id, + }) + return { + "session_id": self.session_id, + "tokens_per_sec": round(self._total_tokens / elapsed, 2), + "seams": seams, + } + + def mark_reported(self, *, now: float | None = None) -> None: + self._last_report = time.monotonic() if now is None else now + self._report_due = False + + def close(self) -> None: + self._closed = True + self._seams.clear() + self._report_due = False diff --git a/packages/node/meshnet_node/server.py b/packages/node/meshnet_node/server.py index 9a6a32b..236d66b 100644 --- a/packages/node/meshnet_node/server.py +++ b/packages/node/meshnet_node/server.py @@ -8,6 +8,7 @@ import urllib.parse from pathlib import Path from .downloader import compute_shard_checksum, write_shard_archive +from .activation_compression import CompressionPolicies, compress_activation, decompress_activation # Binary activation wire format (contract for all shard nodes): # POST /forward with raw tensor bytes in the body and tensor/session/chunk @@ -21,6 +22,7 @@ _DTYPE_SIZES = { "bfloat16": 2, "float32": 4, } +_COMPRESSION_POLICIES = CompressionPolicies() def _make_stub_binary_activation(shape: list[int], dtype: str) -> bytes: @@ -42,16 +44,7 @@ def _parse_shape(value: str | None) -> list[int]: def _decompress_body(body: bytes, encoding: str | None) -> bytes: - if not encoding: - return body - if encoding != "zstd": - raise ValueError("unsupported X-Meshnet-Encoding") - import zstandard as zstd - - try: - return zstd.ZstdDecompressor().decompress(body) - except zstd.ZstdError as exc: - raise ValueError("invalid zstd activation body") from exc + return decompress_activation(body, encoding).body def _compress_body(body: bytes, encoding: str | None) -> bytes: @@ -307,7 +300,14 @@ class _StubHandler(http.server.BaseHTTPRequestHandler): server.received_activations = True raw_payload = _make_stub_binary_activation(shape, dtype) - payload = _compress_body(raw_payload, encoding) + route_condition = self.headers.get("X-Meshnet-Compression-Route", "lan") + phase_condition = self.headers.get("X-Meshnet-Cache", "prefill") + if phase_condition not in {"prefill", "decode"}: + phase_condition = "prefill" + compression = compress_activation( + raw_payload, _COMPRESSION_POLICIES.for_condition(route_condition, phase_condition), + ) + payload = compression.body self.send_response(200) self.send_header("Content-Type", "application/octet-stream") self.send_header("Content-Length", str(len(payload))) @@ -317,8 +317,8 @@ class _StubHandler(http.server.BaseHTTPRequestHandler): self.send_header("X-Meshnet-Session", session) self.send_header("X-Meshnet-Chunk-Index", chunk_index) self.send_header("X-Meshnet-Chunk-Total", chunk_total) - if encoding: - self.send_header("X-Meshnet-Encoding", encoding) + if compression.encoding: + self.send_header("X-Meshnet-Encoding", compression.encoding) if server.is_last_shard: self.send_header("X-Meshnet-Stub-Response-Prefix", server.response_prefix) self.end_headers() diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 76d0295..77192f0 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -14,9 +14,19 @@ import urllib.request from pathlib import Path from typing import Any +from .admission import ( + AdmissionRequirement, + CapabilityContext, + CapabilityValidator, + admit, + probe_capability, +) +from .capability import CapabilityReport +from .doctor import DoctorSelection from .downloader import compute_shard_checksum, download_shard from .hardware import detect_hardware, benchmark_throughput_checked, with_forced_cpu from .model_catalog import model_metadata_for +from .recipe_manifest import DEFAULT_RECIPE_ID, Recipe, RecipeManifest, load_recipe_manifest from .relay_bridge import RelayHttpBridge, peer_id_from_wallet from .server import StubNodeServer from .torch_server import TorchNodeServer @@ -646,6 +656,68 @@ def _tracker_http_error_message(exc: urllib.error.HTTPError) -> str: return f"Tracker rejected shard assignment (HTTP {exc.code}): {detail}" +def _resolve_recipe(recipe_id: str | None) -> tuple[RecipeManifest, Recipe]: + """The recipe this node will serve with — resolved before any weights load.""" + manifest = load_recipe_manifest() + return manifest, manifest.require(recipe_id or DEFAULT_RECIPE_ID) + + +def _capability_device(backend: Any, detected_device: str) -> str: + """The device the shard actually landed on, or the one this node detected.""" + device = getattr(backend, "device", None) + if device is None: + return detected_device + return str(getattr(device, "type", device)) + + +def _admit_capability( + node: Any, + *, + model_id: str, + shard_start: int, + shard_end: int, + quantization: str, + cache_dir: Path | None, + force_cpu: bool, + detected_device: str, + manifest: RecipeManifest, + recipe: Recipe, + validator: CapabilityValidator | None, +) -> CapabilityReport: + """Prove this node can serve the selection, or refuse to advertise it. + + Runs on the loaded backend before the server starts listening, so a node that + cannot execute its shard never reaches a routable endpoint, never registers, + and never accepts paid work. `CapabilityAdmissionError` propagates to the CLI, + which exits non-zero. + """ + backend = getattr(node, "backend", None) + context = CapabilityContext( + backend=backend, + selection=DoctorSelection( + model_id=model_id, + shard_start=shard_start, + shard_end=shard_end, + quantization=quantization, + cache_dir=cache_dir, + force_cpu=force_cpu, + ), + recipe=recipe, + manifest=manifest, + device=_capability_device(backend, detected_device), + ) + print( + f"Validating capability — {model_id} layers {shard_start}–{shard_end}, " + f"recipe {recipe.id}...", + flush=True, + ) + report = (validator or probe_capability)(context) + setattr(node, "capability_report", report) # local evidence, passed or failed + admit(AdmissionRequirement.for_context(context), report) + print(f" Capability proven on {context.device} ({report.duration_ms} ms)", flush=True) + return report + + def run_startup( tracker_url: str, port: int = 0, @@ -668,6 +740,8 @@ def run_startup( torch_interop_threads: int | None = None, node_name: str | None = None, force_cpu: bool = False, + recipe_id: str | None = None, + capability_validator: CapabilityValidator | None = None, ) -> StubNodeServer | TorchNodeServer: """Execute the full startup sequence and return a running node server. @@ -676,13 +750,18 @@ def run_startup( 2. Load or generate Solana wallet keypair 3. Query tracker for optimal shard assignment 4. Download (or stub) the assigned shard from peers, then HuggingFace - 5. Start local HTTP server - 6. Register with tracker + 5. Prove the loaded shard runs — a failure here exits before step 6 + 6. Start local HTTP server and register with tracker + + `capability_validator` is how step 5 is proven. It defaults to a real forward + through the loaded shard; only tests replace it, and only with the explicit + seams in `meshnet_node.testing` — there is no bypass a deployment can reach. Prints a compact status summary on completion. """ tracker_url = tracker_url.rstrip("/") + manifest, recipe = _resolve_recipe(recipe_id) relay_url = _discover_relay_url(tracker_url) display_fields = _registration_display_fields(node_name) if max_loaded_shards < 1: @@ -874,6 +953,20 @@ def run_startup( debug=debug, max_loaded_shards=max_loaded_shards, force_cpu=force_cpu, + recipe_params=recipe.params, + ) + capability_report = _admit_capability( + node, + model_id=model_id, + shard_start=shard_start, + shard_end=shard_end, + quantization=quantization, + cache_dir=cache_dir, + force_cpu=force_cpu, + detected_device=device, + manifest=manifest, + recipe=recipe, + validator=capability_validator, ) _node_start_time = time.monotonic() actual_port = node.start() @@ -910,6 +1003,11 @@ def run_startup( "tracker_mode": (shard_start == 0), "managed_assignment": not user_pinned_shard, "model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir), + "capability_report": capability_report.to_dict(), + # Declared independently of the proof: the tracker checks that the + # recipe this node says it serves with is the one the proof ran. + "recipe_id": recipe.id, + "recipe_version": recipe.version, "downloaded_models": ( _downloaded_model_inventory( model_id.split("/")[-1], @@ -1030,6 +1128,20 @@ def run_startup( debug=debug, max_loaded_shards=max_loaded_shards, force_cpu=force_cpu, + recipe_params=recipe.params, + ) + capability_report = _admit_capability( + node, + model_id=assigned_hf_repo, + shard_start=assigned_shard_start, + shard_end=assigned_shard_end, + quantization=quantization, + cache_dir=cache_dir, + force_cpu=force_cpu, + detected_device=device, + manifest=manifest, + recipe=recipe, + validator=capability_validator, ) _node_start_time = time.monotonic() actual_port = node.start() @@ -1062,6 +1174,11 @@ def run_startup( "tracker_mode": (assigned_shard_start == 0), "managed_assignment": True, "model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir), + "capability_report": capability_report.to_dict(), + # Declared independently of the proof: the tracker checks that the + # recipe this node says it serves with is the one the proof ran. + "recipe_id": recipe.id, + "recipe_version": recipe.version, "downloaded_models": ( _downloaded_model_inventory( assigned_hf_repo.split("/")[-1], @@ -1212,6 +1329,20 @@ def run_startup( debug=debug, max_loaded_shards=max_loaded_shards, force_cpu=force_cpu, + recipe_params=recipe.params, + ) + capability_report = _admit_capability( + node, + model_id=hf_repo, + shard_start=shard_start, + shard_end=shard_end, + quantization=quantization, + cache_dir=shard_path, + force_cpu=force_cpu, + detected_device=device, + manifest=manifest, + recipe=recipe, + validator=capability_validator, ) actual_port = node.start() total_layers = getattr(getattr(node, "backend", None), "total_layers", None) or assigned_total_layers @@ -1247,6 +1378,11 @@ def run_startup( "tracker_mode": (shard_start == 0), "managed_assignment": not user_pinned_shard, "model_metadata": model_metadata_for(hf_repo, total_layers, cache_dir=shard_path), + "capability_report": capability_report.to_dict(), + # Declared independently of the proof: the tracker checks that the + # recipe this node says it serves with is the one the proof ran. + "recipe_id": recipe.id, + "recipe_version": recipe.version, **registration_capabilities, **relay_fields, **display_fields, @@ -1282,6 +1418,19 @@ def run_startup( model=assigned_model, shard_path=shard_path, ) + capability_report = _admit_capability( + node, + model_id=assigned_model, + shard_start=shard_start, + shard_end=shard_end, + quantization=quantization, + cache_dir=shard_path, + force_cpu=force_cpu, + detected_device=device, + manifest=manifest, + recipe=recipe, + validator=capability_validator, + ) actual_port = node.start() public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) endpoint = f"http://{public_host}:{actual_port}" @@ -1304,6 +1453,11 @@ def run_startup( "shard_start": shard_start, "shard_end": shard_end, "shard_checksum": shard_checksum, + "capability_report": capability_report.to_dict(), + # Declared independently of the proof: the tracker checks that the + # recipe this node says it serves with is the one the proof ran. + "recipe_id": recipe.id, + "recipe_version": recipe.version, "downloaded_models": downloaded_models, "hardware_profile": hw, "wallet_address": address, diff --git a/packages/node/meshnet_node/testing.py b/packages/node/meshnet_node/testing.py new file mode 100644 index 0000000..e595524 --- /dev/null +++ b/packages/node/meshnet_node/testing.py @@ -0,0 +1,70 @@ +"""Test-only seams. Nothing in the production code path may import this module. + +Startup admits a node only on a capability report produced by a *real* forward +through the loaded shard (see :mod:`meshnet_node.admission`). Tests run against +fake or stub backends that cannot perform one, so they pass an explicit validator +from here instead — the honest statement being "this test asserts capability it +never proved", which is a thing a test may do and a node may not. + +`capability_stub` builds the deliberately-wrong reports the fail-closed tests +need: a failed one, one for another model or shard, one that has aged out. +""" + +from __future__ import annotations + +import time +from typing import Any + +from .admission import CapabilityContext, CapabilityValidator +from .capability import STATUS_PASSED, CapabilityReport, build_capability_report + + +def capability_report_for( + context: CapabilityContext, + *, + status: str = STATUS_PASSED, + model_id: str | None = None, + shard_start: int | None = None, + shard_end: int | None = None, + recipe_id: str | None = None, + recipe_version: str | None = None, + backend_id: str | None = None, + device: str | None = None, + validated_at: float | None = None, + age_seconds: float = 0.0, + diagnostics: Any = None, + duration_ms: int = 0, +) -> CapabilityReport: + """A report describing `context`, with any field bent away from the truth.""" + now = time.time() if validated_at is None else validated_at + return build_capability_report( + model_id=model_id or context.selection.model_id, + shard_start=( + context.selection.shard_start if shard_start is None else shard_start + ), + shard_end=context.selection.shard_end if shard_end is None else shard_end, + recipe_id=recipe_id or context.recipe.id, + recipe_version=recipe_version or context.recipe.version, + catalogue_version=context.manifest.catalogue_version, + backend_id=backend_id or context.recipe.backend_id, + device=device or context.device, + quantization=context.selection.quantization, + status=status, + duration_ms=duration_ms, + diagnostics=diagnostics, + validated_at=now - age_seconds, + ) + + +def assume_capability(context: CapabilityContext) -> CapabilityReport: + """Assert the selection works, without proving it. Tests only.""" + return capability_report_for(context) + + +def capability_stub(**overrides: Any) -> CapabilityValidator: + """A validator producing a report that deviates from `context` as named.""" + + def validator(context: CapabilityContext) -> CapabilityReport: + return capability_report_for(context, **overrides) + + return validator diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 45c5bea..e58cf6d 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -3,17 +3,17 @@ from __future__ import annotations import base64 +import http.client import http.server import json import sys import threading import time -import urllib.error import urllib.parse import urllib.request import uuid from pathlib import Path -from typing import Any +from typing import Any, Mapping from .model_backend import ( InsufficientVRAMError, @@ -22,16 +22,32 @@ from .model_backend import ( Quantization, TailTokenResult, TorchModelShard, + _tensor_from_bfloat16_bytes, validate_quantization, ) +from .seam_telemetry import GenerationTelemetry +from .activation_compression import ( + CompressionPolicies, + CompressionPolicy, + compress_activation, + decompress_activation, +) class _PipelineCacheMiss(Exception): """A downstream hop reported 409 cache_miss — head must re-prefill.""" + + +class _RelayRequestUncertainError(ConnectionError): + """A relay request may have reached the peer but produced no response.""" + + +class _DirectRequestUncertainError(ConnectionError): + """A direct request may have reached the downstream node but did not finish.""" + + from .server import ( _WIRE_VERSION, - _compress_body, - _decompress_body, _parse_shape, _validate_activation_body, ) @@ -107,6 +123,7 @@ class _RelayHopClient: self.relay_addr = relay_addr self.timeout = timeout self._ws = None + self._lock = threading.RLock() def request( self, @@ -118,42 +135,116 @@ class _RelayHopClient: from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size - if self._ws is None: - self._ws = wsc.connect( - self.relay_addr, - open_timeout=self.timeout, - max_size=ws_max_size(), - compression=None, - ) - request_id = f"{time.time_ns():x}" - frame = encode_binary_frame({ - "request_id": request_id, - "method": "POST", - "path": path, - "headers": headers, - }, body) - self._ws.send(frame) - raw = self._ws.recv(timeout=self.timeout) - if isinstance(raw, (bytes, bytearray)): - resp_header, resp_body = decode_binary_frame(bytes(raw)) - status = int(resp_header.get("status", 503)) - resp_headers = {k.lower(): v for k, v in (resp_header.get("headers") or {}).items()} - return status, resp_headers, resp_body - resp = json.loads(raw) - status = int(resp.get("status", 503)) - resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()} - body_b64 = resp.get("body_base64") - resp_body = base64.b64decode(body_b64) if body_b64 else (resp.get("body") or "").encode() - return status, resp_headers, resp_body + with self._lock: + if self._ws is None: + self._ws = wsc.connect( + self.relay_addr, + open_timeout=self.timeout, + max_size=ws_max_size(), + compression=None, + ) + request_id = headers.get("X-Meshnet-Activation-Id") or uuid.uuid4().hex + frame = encode_binary_frame({ + "request_id": request_id, + "method": "POST", + "path": path, + "headers": headers, + }, body) + try: + # A send failure is uncertain too: bytes may already have been + # accepted by the kernel or peer before the exception surfaced. + self._ws.send(frame) + raw = self._ws.recv(timeout=self.timeout) + if isinstance(raw, (bytes, bytearray)): + resp_header, resp_body = decode_binary_frame(bytes(raw)) + response_id = str(resp_header.get("request_id") or "") + status = int(resp_header.get("status", 503)) + resp_headers = { + k.lower(): v for k, v in (resp_header.get("headers") or {}).items() + } + else: + resp = json.loads(raw) + response_id = str(resp.get("request_id") or "") + status = int(resp.get("status", 503)) + resp_headers = { + k.lower(): v for k, v in (resp.get("headers") or {}).items() + } + body_b64 = resp.get("body_base64") + resp_body = ( + base64.b64decode(body_b64) + if body_b64 else (resp.get("body") or "").encode() + ) + if response_id and response_id != request_id: + raise ValueError("relay response request_id did not match request") + return status, resp_headers, resp_body + except Exception as exc: + self.close() + raise _RelayRequestUncertainError( + "relay connection failed after forwarding request; refusing replay" + ) from exc def close(self) -> None: - if self._ws is not None: + with self._lock: + if self._ws is not None: + try: + self._ws.close() + except Exception: + pass + finally: + self._ws = None + + +class _DirectHopClient: + """One serialized HTTP/1.1 connection to one downstream hop. + + Generation handlers own these clients, so a cached Route Session reuses a + TCP connection without sharing it between concurrent Route Sessions. + """ + + def __init__(self, endpoint: str, timeout: float = 120.0) -> None: + parsed = urllib.parse.urlsplit(endpoint.rstrip("/")) + if parsed.scheme not in {"http", "https"} or not parsed.hostname: + raise ValueError(f"invalid downstream endpoint: {endpoint!r}") + self.timeout = timeout + self._scheme = parsed.scheme + self._host = parsed.hostname + self._port = parsed.port + self._base_path = parsed.path.rstrip("/") + self._connection: http.client.HTTPConnection | None = None + self._lock = threading.RLock() + + def _connect(self) -> http.client.HTTPConnection: + connection_type = ( + http.client.HTTPSConnection if self._scheme == "https" else http.client.HTTPConnection + ) + return connection_type(self._host, self._port, timeout=self.timeout) + + def request( + self, path: str, body: bytes, headers: Mapping[str, str], + ) -> tuple[int, dict[str, str], bytes]: + request_path = f"{self._base_path}{path if path.startswith('/') else '/' + path}" + with self._lock: + if self._connection is None: + self._connection = self._connect() try: - self._ws.close() - except Exception: - pass - finally: - self._ws = None + self._connection.request("POST", request_path, body=body, headers=dict(headers)) + response = self._connection.getresponse() + response_body = response.read() + response_headers = {key.lower(): value for key, value in response.headers.items()} + return response.status, response_headers, response_body + except Exception as exc: + self.close() + raise _DirectRequestUncertainError( + "direct connection failed after forwarding request; refusing replay" + ) from exc + + def close(self) -> None: + with self._lock: + if self._connection is not None: + try: + self._connection.close() + finally: + self._connection = None def _relay_hop( @@ -178,18 +269,15 @@ def _relay_hop( client.close() -# Below this, zstd overhead outweighs the win (per-token decode bodies are ~KBs). -_COMPRESS_MIN_BYTES = 64 * 1024 +_COMPRESSION_POLICIES = CompressionPolicies() -def _maybe_compress_activation(body: bytes) -> tuple[bytes, str | None]: - """zstd-compress large activation bodies; returns (wire_body, encoding).""" - if len(body) < _COMPRESS_MIN_BYTES: - return body, None - try: - return _compress_body(body, "zstd"), "zstd" - except Exception: - return body, None +def _maybe_compress_activation( + body: bytes, policy: CompressionPolicy | None = None, +) -> tuple[bytes, str | None]: + """Compatibility wrapper for callers that only need wire body and encoding.""" + result = compress_activation(body, policy or _COMPRESSION_POLICIES.for_condition("lan", "prefill")) + return result.body, result.encoding def _is_cache_miss_body(body: bytes) -> bool: @@ -285,6 +373,7 @@ class _TorchHTTPServer(http.server.HTTPServer): "elapsed_seconds": round(elapsed, 1), "tokens_per_sec": round(tokens / elapsed, 2) if tokens > 0 else 0.0, "routing_complete": bool(rec.get("routing_complete")), + "telemetry": rec["telemetry"].snapshot(now=now) if rec.get("telemetry") else None, }) return out @@ -306,6 +395,10 @@ class _TorchHTTPServer(http.server.HTTPServer): class _TorchHandler(http.server.BaseHTTPRequestHandler): + # HTTP/1.1 is required for Route Session-owned downstream connections. + # Finite responses below provide Content-Length; streams are chunked. + protocol_version = "HTTP/1.1" + def log_message(self, fmt, *args): # noqa: suppress request logs in tests pass @@ -317,6 +410,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): ) def _request_log_suffix(self) -> str: + activation_id = self.headers.get("X-Meshnet-Activation-Id") + if activation_id: + return f" activation_id={activation_id}" req_id = self.headers.get("X-Meshnet-Request-Id") or self.headers.get("X-Request-Id") return f" request_id={req_id}" if req_id else "" @@ -334,6 +430,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): "started": time.monotonic(), "tokens": 0, "routing_complete": False, + "telemetry": None, } def _track_request_progress( @@ -366,6 +463,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self._handle_chat_completions() else: self.send_response(404) + self.send_header("Content-Length", "0") self.end_headers() def _handle_infer(self) -> None: @@ -427,7 +525,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): encoding = self.headers.get("X-Meshnet-Encoding") length = int(self.headers.get("Content-Length", 0)) body = self.rfile.read(length) - raw_body = _decompress_body(body, encoding) + raw_body = decompress_activation(body, encoding).body _validate_activation_body(raw_body, shape, dtype) if dtype != "bfloat16": raise ValueError("real model backend requires bfloat16 activation input") @@ -438,6 +536,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): except (KeyError, ValueError, TypeError): self.send_response(400) self.send_header("X-Meshnet-Wire", _WIRE_VERSION) + self.send_header("Content-Length", "0") self.end_headers() return @@ -511,7 +610,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self._send_json(200, {"text": result}) return - response_body = _compress_body(result.body, encoding) + route_condition = self.headers.get("X-Meshnet-Compression-Route", "lan") + phase_condition = cache_mode if cache_mode in {"prefill", "decode"} else "prefill" + response_compression = compress_activation( + result.body, _COMPRESSION_POLICIES.for_condition(route_condition, phase_condition), + ) + response_body = response_compression.body self.send_response(200) self.send_header("Content-Type", "application/octet-stream") self.send_header("Content-Length", str(len(response_body))) @@ -521,8 +625,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self.send_header("X-Meshnet-Session", session) self.send_header("X-Meshnet-Chunk-Index", chunk_index) self.send_header("X-Meshnet-Chunk-Total", chunk_total) - if encoding: - self.send_header("X-Meshnet-Encoding", encoding) + if response_compression.encoding: + self.send_header("X-Meshnet-Encoding", response_compression.encoding) if result.attention_mask_header: self.send_header("X-Meshnet-Attn-Mask", result.attention_mask_header) if result.position_ids_header: @@ -700,6 +804,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): current_text = prompt_text session_id = str(uuid.uuid4()) + telemetry = GenerationTelemetry(session_id) + with server._stats_lock: + current = server._active_requests.get(request_id) + if current is not None: + current["telemetry"] = telemetry use_kv = bool(getattr(backend, "supports_kv_cache", False)) # EOS detection by id must work on the stateless path too: the tail # returns token_id regardless of caching, and EOS usually decodes to @@ -722,6 +831,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): last_token_id: int | None = None failure_reason: str | None = None relay_clients: dict[str, _RelayHopClient] = {} + direct_clients: dict[str, _DirectHopClient] = {} def _prefill_step() -> tuple[str, int | None]: """Full-sequence prefill: initial step and cache-miss recovery.""" @@ -734,6 +844,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): payload, remaining_route, backend=backend, session=session_id, cache_mode="prefill" if use_kv else None, relay_clients=relay_clients, + direct_clients=direct_clients if use_kv else None, + telemetry=telemetry, ) for step in range(max_tokens): @@ -745,6 +857,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): payload, remaining_route, backend=backend, session=session_id, cache_mode="decode", relay_clients=relay_clients, + direct_clients=direct_clients, + telemetry=telemetry, ) except (KVCacheMiss, _PipelineCacheMiss) as miss: # Evicted/restarted node or head lost its own session: @@ -758,11 +872,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): else: token_str, token_id = _prefill_step() except _PipelineCacheMiss as exc: - print(f" [node] unexpected cache miss on prefill: {exc}", flush=True) + print(f" [node] unexpected cache miss on prefill session={session_id[:8]}: {exc}", flush=True) failure_reason = f"cache miss on prefill: {exc}" break except Exception as exc: - print(f" [node] distributed encode error: {exc}", flush=True) + print(f" [node] distributed encode error session={session_id[:8]}: {exc}", flush=True) failure_reason = f"distributed encode error: {exc}" break # Stop on error responses or EOS. @@ -789,14 +903,32 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): tokens=len(generated), routing_complete=True, ) + telemetry.note_tokens(len(generated)) now = time.monotonic() + if telemetry.report_due: + summary = telemetry.snapshot(now=now) + seams = summary["seams"] + if seams: + latest = seams[-1] + print( + f" [node] seam telemetry session={session_id[:8]} " + f"phase={latest['phase']} hop={latest['hop']} " + f"activations={latest['activations']} " + f"avg_ms={latest['avg_latency_ms']:.2f} " + f"wire_bytes={latest['wire_bytes']} " + f"response_bytes={latest['response_bytes']} " + f"tps={summary['tokens_per_sec']:.2f} " + f"activation_id={latest['last_activation_id']}", + flush=True, + ) + telemetry.mark_reported(now=now) if step == 0 or now - last_gen_log >= _GENERATION_LOG_INTERVAL: elapsed = now - gen_started token_count = len(generated) tps = token_count / max(elapsed, 1e-6) _write_progress_line( progress_line, - f" [node] generating step={step + 1}/{max_tokens} " + f" [node] generating step={step + 1}/{max_tokens} session={session_id[:8]} " f"tokens={token_count} elapsed_s={elapsed:.1f} tps={tps:.2f}", ) last_gen_log = now @@ -808,6 +940,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): pass for relay_client in relay_clients.values(): relay_client.close() + for direct_client in direct_clients.values(): + direct_client.close() if generated: elapsed = time.monotonic() - gen_started @@ -815,10 +949,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): tps = token_count / max(elapsed, 1e-6) _write_progress_line( progress_line, - f" [node] generation complete tokens={token_count} " + f" [node] generation complete session={session_id[:8]} tokens={token_count} " f"elapsed_s={elapsed:.1f} tps={tps:.2f}", final=True, ) + telemetry.close() result_text = "".join(generated) # A failure before the first token is an upstream error, not an empty @@ -921,6 +1056,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): session: str | None = None, cache_mode: str | None = None, relay_clients: dict[str, _RelayHopClient] | None = None, + direct_clients: dict[str, _DirectHopClient] | None = None, + telemetry: GenerationTelemetry | None = None, ) -> tuple[str, int | None]: """Forward an activation through the downstream route. @@ -935,10 +1072,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): # Full single-node (head+tail) is handled before entering this method. if active_backend.is_tail: try: - tensor = active_backend.torch.frombuffer( - bytearray(payload.body), # type: ignore[union-attr] - dtype=active_backend.torch.bfloat16, - ).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr] + tensor = _tensor_from_bfloat16_bytes( + payload.body, payload.shape, active_backend.torch, # type: ignore[union-attr] + ).to(active_backend.device) if hasattr(active_backend, "decode_tail_token"): tail = active_backend.decode_tail_token(tensor) return tail.text, tail.token_id @@ -968,7 +1104,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): + (f" relay={relay_addr}" if relay_addr else ""), flush=True, ) - wire_body, wire_encoding = _maybe_compress_activation(current_body) + phase = cache_mode if cache_mode in {"prefill", "decode"} else "prefill" + route_condition = "relay" if relay_addr else "lan" + compression = compress_activation( + current_body, _COMPRESSION_POLICIES.for_condition(route_condition, phase), + ) + wire_body, wire_encoding = compression.body, compression.encoding headers: dict[str, str] = { "Content-Type": "application/octet-stream", "X-Meshnet-Wire": _WIRE_VERSION, @@ -979,9 +1120,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): "X-Meshnet-Chunk-Total": "1", "X-Meshnet-Hop-Index": str(hop_index), "X-Meshnet-Start-Layer": str(start_layer), + "X-Meshnet-Activation-Id": uuid.uuid4().hex, + "X-Meshnet-Compression-Route": route_condition, } if wire_encoding: headers["X-Meshnet-Encoding"] = wire_encoding + if telemetry is not None: + telemetry.record_compression( + phase=phase, hop=hop_index, node=node_url, + input_bytes=compression.input_bytes, output_bytes=compression.output_bytes, + elapsed_seconds=compression.elapsed_seconds, + ) if cache_mode: headers["X-Meshnet-Cache"] = cache_mode past_len = getattr(payload, "past_len", None) @@ -992,6 +1141,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): if current_pos: headers["X-Meshnet-Position-Ids"] = current_pos if relay_addr: + connection_reused = bool( + relay_clients is not None and relay_addr in relay_clients + ) + seam_started = time.monotonic() try: if relay_clients is None: status, resp_headers, resp_body = _relay_hop( @@ -1004,44 +1157,100 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): status, resp_headers, resp_body = relay_client.request( "/forward", wire_body, headers, ) + if telemetry is not None: + telemetry.record_seam( + activation_id=headers["X-Meshnet-Activation-Id"], + phase=cache_mode or "prefill", + hop=hop_index, + node=node_url, + latency_seconds=time.monotonic() - seam_started, + wire_bytes=len(wire_body), + response_bytes=len(resp_body), + connection_reused=connection_reused, + ) if status == 409 and _is_cache_miss_body(resp_body): raise _PipelineCacheMiss(node_url) if status >= 400: detail = _response_error_snippet(resp_body) print( - f" [node] relay hop {hop_index} returned {status} from {relay_addr}: {detail}", + f" [node] relay hop {hop_index} session={session[:8]} returned " + f"{status} from {relay_addr}: {detail}", flush=True, ) return f"pipeline error at {node_url} via relay: status {status}: {detail}", None except _PipelineCacheMiss: raise + except _RelayRequestUncertainError as exc: + # The activation may already have mutated the downstream + # KV cache. Do not replay it on a direct connection. + print( + f" [node] relay hop {hop_index} session={session[:8]} outcome is uncertain at " + f"{relay_addr}: {exc}", + flush=True, + ) + return f"pipeline relay outcome uncertain at {node_url}: {exc}", None except Exception as exc: print( - f" [node] relay hop {hop_index} failed at {relay_addr}: {exc}; " + f" [node] relay hop {hop_index} session={session[:8]} failed at {relay_addr}: {exc}; " f"falling back to direct {node_url}", flush=True, ) relay_addr = None # fall through to direct if not relay_addr: - req = urllib.request.Request( - f"{node_url}/forward", - data=wire_body, - headers=headers, - method="POST", + connection_reused = bool( + direct_clients is not None and node_url in direct_clients ) + seam_started = time.monotonic() try: - with urllib.request.urlopen(req, timeout=120.0) as r: - resp_body = r.read() - resp_headers = {k.lower(): v for k, v in r.headers.items()} - except urllib.error.HTTPError as exc: - body = exc.read() - if exc.code == 409 and _is_cache_miss_body(body): - raise _PipelineCacheMiss(node_url) from exc - detail = _response_error_snippet(body) - print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}: {detail}", flush=True) - return f"pipeline error at {node_url}: {exc}: {detail}", None + if direct_clients is None: + direct_client = _DirectHopClient(node_url, timeout=120.0) + try: + status, resp_headers, resp_body = direct_client.request( + "/forward", wire_body, headers, + ) + finally: + direct_client.close() + else: + direct_client = direct_clients.setdefault( + node_url, _DirectHopClient(node_url, timeout=120.0), + ) + status, resp_headers, resp_body = direct_client.request( + "/forward", wire_body, headers, + ) + if telemetry is not None: + telemetry.record_seam( + activation_id=headers["X-Meshnet-Activation-Id"], + phase=cache_mode or "prefill", + hop=hop_index, + node=node_url, + latency_seconds=time.monotonic() - seam_started, + wire_bytes=len(wire_body), + response_bytes=len(resp_body), + connection_reused=connection_reused, + ) + if status == 409 and _is_cache_miss_body(resp_body): + raise _PipelineCacheMiss(node_url) + if status >= 400: + detail = _response_error_snippet(resp_body) + print( + f" [node] pipeline hop {hop_index} session={session[:8]} failed at {node_url}: " + f"status {status}: {detail}", flush=True, + ) + return f"pipeline error at {node_url}: status {status}: {detail}", None + except _PipelineCacheMiss: + raise + except _DirectRequestUncertainError as exc: + print( + f" [node] pipeline hop {hop_index} session={session[:8]} outcome is uncertain " + f"at {node_url}: {exc}", + flush=True, + ) + return f"pipeline direct outcome uncertain at {node_url}: {exc}", None except Exception as exc: - print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True) + print( + f" [node] pipeline hop {hop_index} session={session[:8]} " + f"failed at {node_url}: {exc}", flush=True, + ) return f"pipeline error at {node_url}: {exc}", None content_type = resp_headers.get("content-type", "") if "application/json" in content_type: @@ -1058,11 +1267,21 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape)) current_shape = _parse_shape(shape_header) try: - current_body = _decompress_body( + decompression = decompress_activation( resp_body, resp_headers.get("x-meshnet-encoding") ) + current_body = decompression.body + if telemetry is not None: + telemetry.record_compression( + phase=phase, hop=hop_index, node=node_url, + input_bytes=decompression.input_bytes, output_bytes=decompression.output_bytes, + elapsed_seconds=decompression.elapsed_seconds, decompression=True, + ) except ValueError as exc: - print(f" [node] pipeline hop {hop_index} bad response encoding: {exc}", flush=True) + print( + f" [node] pipeline hop {hop_index} session={session[:8]} " + f"bad response encoding: {exc}", flush=True, + ) return f"pipeline error at {node_url}: {exc}", None current_attn = resp_headers.get("x-meshnet-attn-mask") current_pos = resp_headers.get("x-meshnet-position-ids") @@ -1084,14 +1303,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self.send_response(200) self.send_header("Content-Type", "text/event-stream; charset=utf-8") self.send_header("Cache-Control", "no-cache") + self.send_header("Transfer-Encoding", "chunked") self.end_headers() - def _emit(data: str) -> None: + def _emit(data: str) -> bool: try: - self.wfile.write(f"data: {data}\n\n".encode()) + body = f"data: {data}\n\n".encode() + self.wfile.write(f"{len(body):X}\r\n".encode() + body + b"\r\n") self.wfile.flush() + return True except (BrokenPipeError, ConnectionResetError): - pass + return False _emit(json.dumps({ "id": chunk_id, "object": "chat.completion.chunk", "created": created, @@ -1105,7 +1327,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): # instead of showing an empty completion. _emit(json.dumps({"error": {"message": error, "type": "upstream_error"}})) try: - self.wfile.write(b"data: [DONE]\n\n") + self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n") self.wfile.flush() except (BrokenPipeError, ConnectionResetError): pass @@ -1117,7 +1339,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], })) try: - self.wfile.write(b"data: [DONE]\n\n") + self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n") self.wfile.flush() except (BrokenPipeError, ConnectionResetError): pass @@ -1159,14 +1381,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self.send_response(200) self.send_header("Content-Type", "text/event-stream; charset=utf-8") self.send_header("Cache-Control", "no-cache") + self.send_header("Transfer-Encoding", "chunked") self.end_headers() - def _emit(data: str) -> None: + def _emit(data: str) -> bool: try: - self.wfile.write(f"data: {data}\n\n".encode()) + body = f"data: {data}\n\n".encode() + self.wfile.write(f"{len(body):X}\r\n".encode() + body + b"\r\n") self.wfile.flush() - except BrokenPipeError: - pass + return True + except (BrokenPipeError, ConnectionResetError): + return False _emit(json.dumps({ "id": chunk_id, "object": "chat.completion.chunk", "created": created, @@ -1184,9 +1409,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], })) try: - self.wfile.write(b"data: [DONE]\n\n") + self.wfile.write(b"E\r\ndata: [DONE]\n\n\r\n0\r\n\r\n") self.wfile.flush() - except BrokenPipeError: + except (BrokenPipeError, ConnectionResetError): pass @@ -1255,6 +1480,7 @@ class TorchNodeServer: debug: bool = False, max_loaded_shards: int = 1, force_cpu: bool = False, + recipe_params: Mapping[str, Any] | None = None, ) -> None: self._host = host self._requested_port = port @@ -1266,6 +1492,7 @@ class TorchNodeServer: quantization, cache_dir, force_cpu=force_cpu, + recipe_params=recipe_params, ) self._backends: dict[str, TorchModelShard] = {self._backend.model_id: self._backend} # Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set @@ -1415,13 +1642,20 @@ def _load_backend( quantization: str, cache_dir: Path | None = None, force_cpu: bool = False, + recipe_params: Mapping[str, Any] | None = None, ) -> TorchModelShard: from .model_backend import load_torch_shard quant = validate_quantization(quantization) try: return load_torch_shard( - model_id, shard_start, shard_end, quant, cache_dir, force_cpu=force_cpu + model_id, + shard_start, + shard_end, + quant, + cache_dir, + force_cpu=force_cpu, + recipe_params=recipe_params, ) except MissingModelDependencyError: raise diff --git a/packages/relay/meshnet_relay/server.py b/packages/relay/meshnet_relay/server.py index 89355e2..4021997 100644 --- a/packages/relay/meshnet_relay/server.py +++ b/packages/relay/meshnet_relay/server.py @@ -46,11 +46,12 @@ def ws_max_size() -> int | None: # inflation, no JSON re-encode of megabytes per hop. Text JSON frames remain the # control plane (gossip, peer-register, streamed SSE responses). BINARY_FRAME_MAGIC = b"MRF1" +_RPC_PEER_DISCONNECTED = object() def encode_binary_frame(header: dict, body: bytes) -> bytes: header_bytes = json.dumps(header, separators=(",", ":")).encode() - return BINARY_FRAME_MAGIC + len(header_bytes).to_bytes(4, "big") + header_bytes + body + return b"".join((BINARY_FRAME_MAGIC, len(header_bytes).to_bytes(4, "big"), header_bytes, body)) def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]: @@ -58,7 +59,7 @@ def decode_binary_frame(frame: bytes) -> tuple[dict, bytes]: raise ValueError("not a meshnet binary relay frame") header_len = int.from_bytes(frame[4:8], "big") header = json.loads(frame[8:8 + header_len].decode()) - return header, bytes(frame[8 + header_len:]) + return header, frame[8 + header_len:] class RelayServer: @@ -79,12 +80,16 @@ class RelayServer: ssl_cert: Path | None = None, ssl_key: Path | None = None, max_peers: int = 500, + rpc_timeout: float = 310.0, + rpc_idle_timeout: float = 120.0, ): self.host = host self.port = port self.ssl_cert = ssl_cert self.ssl_key = ssl_key self.max_peers = max_peers + self.rpc_timeout = rpc_timeout + self.rpc_idle_timeout = rpc_idle_timeout self._registry = PeerRegistry() self._loop: asyncio.AbstractEventLoop | None = None @@ -94,9 +99,10 @@ class RelayServer: self._ready = threading.Event() self._running = False self._stop_event: asyncio.Event | None = None - # request_id → queue of relay-http-response frames (US-036: a streamed - # response is a sequence of frames; a frame without "stream" is terminal). - self._pending_rpc: dict[str, asyncio.Queue] = {} + # relay request id → (target peer, requester request id, response queue). + # The relay-generated id prevents two Route Sessions using the same + # legacy request_id from overwriting each other's pending response. + self._pending_rpc: dict[str, tuple[str, str, asyncio.Queue]] = {} @property def registry(self) -> PeerRegistry: @@ -118,6 +124,12 @@ class RelayServer: if self._thread: self._thread.join(timeout=3.0) + def _fail_pending_for_peer(self, peer_id: str) -> None: + """Wake requesters immediately when their selected bridge disconnects.""" + for target, _, queue in tuple(self._pending_rpc.values()): + if target == peer_id: + queue.put_nowait(_RPC_PEER_DISCONNECTED) + def _run(self) -> None: asyncio.set_event_loop(self._loop) self._loop.run_until_complete(self._serve()) @@ -192,9 +204,9 @@ class RelayServer: header, _ = decode_binary_frame(bytes(raw)) except (ValueError, json.JSONDecodeError): continue - queue = self._pending_rpc.get(header.get("request_id")) - if queue is not None: - queue.put_nowait(bytes(raw)) + pending = self._pending_rpc.get(header.get("request_id")) + if pending is not None: + pending[2].put_nowait(bytes(raw)) continue try: envelope = json.loads(raw) @@ -226,9 +238,9 @@ class RelayServer: if topic == "relay-http-response": payload = envelope.get("payload", {}) request_id = payload.get("request_id") - queue = self._pending_rpc.get(request_id) - if queue is not None: - queue.put_nowait(payload) + pending = self._pending_rpc.get(request_id) + if pending is not None: + pending[2].put_nowait(payload) continue # Fan out to all other registered peers @@ -242,6 +254,9 @@ class RelayServer: finally: if peer_id: self._registry.unregister(peer_id) + # Do not leave a requester waiting for its full timeout after + # the selected bridge goes away. + self._fail_pending_for_peer(peer_id) log.debug("Peer unregistered: %s", peer_id) async def _handle_circuit_relay(self, ws_requester, target_peer_id: str) -> None: @@ -290,18 +305,19 @@ class RelayServer: except Exception: return - request_id: str | None = None + requester_request_id: str | None = None + relay_request_id = uuid.uuid4().hex try: if isinstance(raw, (bytes, bytearray)): header, body = decode_binary_frame(bytes(raw)) - request_id = str(header.get("request_id") or uuid.uuid4()) - header["request_id"] = request_id + requester_request_id = str(header.get("request_id") or uuid.uuid4()) + header["request_id"] = relay_request_id header["target_peer"] = target_peer_id outbound: str | bytes = encode_binary_frame(header, body) else: payload = json.loads(raw) - request_id = str(payload.get("request_id") or uuid.uuid4()) - payload["request_id"] = request_id + requester_request_id = str(payload.get("request_id") or uuid.uuid4()) + payload["request_id"] = relay_request_id payload["target_peer"] = target_peer_id outbound = json.dumps({ "topic": "relay-http-request", @@ -314,16 +330,18 @@ class RelayServer: return queue: asyncio.Queue = asyncio.Queue() - self._pending_rpc[request_id] = queue - overall_timeout = 310.0 - idle_timeout = 120.0 + self._pending_rpc[relay_request_id] = ( + target_peer_id, requester_request_id, queue, + ) + overall_timeout = self.rpc_timeout + idle_timeout = self.rpc_idle_timeout loop = asyncio.get_running_loop() deadline = loop.time() + overall_timeout target = self._registry.get(target_peer_id) try: if target is None: await ws_requester.send(json.dumps({ - "request_id": request_id, + "request_id": requester_request_id, "status": 503, "headers": {"Content-Type": "application/json"}, "body": json.dumps({"error": f"peer {target_peer_id!r} disconnected"}), @@ -339,21 +357,34 @@ class RelayServer: frame = await asyncio.wait_for( queue.get(), timeout=min(idle_timeout, remaining) ) + if frame is _RPC_PEER_DISCONNECTED: + raise ConnectionError(f"peer {target_peer_id!r} disconnected") if isinstance(frame, (bytes, bytearray)): - await ws_requester.send(frame) + header, body = decode_binary_frame(bytes(frame)) + header["request_id"] = requester_request_id + await ws_requester.send(encode_binary_frame(header, body)) break - await ws_requester.send(json.dumps(frame)) - if not frame.get("stream") or frame.get("done"): + response = dict(frame) + response["request_id"] = requester_request_id + await ws_requester.send(json.dumps(response)) + if not response.get("stream") or response.get("done"): break except asyncio.TimeoutError: await ws_requester.send(json.dumps({ - "request_id": request_id, + "request_id": requester_request_id, "status": 504, "headers": {"Content-Type": "application/json"}, "body": json.dumps({"error": "relay rpc timed out"}), })) + except ConnectionError as exc: + await ws_requester.send(json.dumps({ + "request_id": requester_request_id, + "status": 503, + "headers": {"Content-Type": "application/json"}, + "body": json.dumps({"error": str(exc)}), + })) finally: - self._pending_rpc.pop(request_id, None) + self._pending_rpc.pop(relay_request_id, None) async def _broadcast(raw: str | bytes, peers: list) -> None: diff --git a/packages/tracker/meshnet_tracker/capability.py b/packages/tracker/meshnet_tracker/capability.py new file mode 100644 index 0000000..c60364f --- /dev/null +++ b/packages/tracker/meshnet_tracker/capability.py @@ -0,0 +1,415 @@ +"""Tracker-side validation of the capability report a Node presents at registration. + +A Node proves locally that it can execute one exact combination — model artifact, +shard range, recipe, backend/device — and ships that proof with its registration +(ADR-0023, NCA-001/002/003). The tracker does not re-run the forward; it decides +whether the presented proof *covers what the node is advertising*, records the +verdict as a small sanitized enum, and routes only to nodes whose verdict is +`admitted`. + +Two properties this module deliberately keeps: + +* **No model knowledge.** Model ids, recipe ids, backend ids and device names are + opaque labels. They are compared, never interpreted; no vendor string is a + code path here. +* **Evidence, not assertion.** A report is treated as a claim about identity, and + the tracker only ever *narrows* what a node may serve with it. Nothing in a + report can widen a node's eligibility or its routing weight — throughput + routing stays measurement-driven (ADR-0013/0021). + +Older nodes that predate the capability protocol present no report at all. They +are handled by an explicit policy (`POLICY_COMPAT` vs `POLICY_ENFORCE`), never by +silently treating "no proof" as "proven" — see `docs/adr/0023-…` for the rollout. +""" + +from __future__ import annotations + +import os +import re +import time +from dataclasses import dataclass, replace +from typing import Any, Callable, Mapping + +# The capability report layout this tracker reads (meshnet_node.capability). +SUPPORTED_SCHEMA_VERSION = 1 + +# The oldest recipe catalogue whose recipe semantics this tracker still trusts. +# A node carrying an older catalogue may be running a recipe whose id has since +# been redefined, so its proof cannot be matched to a name reliably. +MIN_CATALOGUE_VERSION = "2026.07.1" + +# How old a proof may be *at the moment it is presented*. Freshness after that is +# carried by liveness: a registration is re-asserted on tracker restart and the +# node is purged once heartbeats stop. +DEFAULT_MAX_REPORT_AGE_SECONDS = 900.0 + +# A proof timestamped further ahead than this is not fresh, it is wrong. +MAX_CLOCK_SKEW_SECONDS = 60.0 + +STATUS_PASSED = "passed" + +# --- Admission verdicts. `admitted` is the only routable one under `enforce`. --- +STATE_ADMITTED = "admitted" +STATE_ABSENT = "absent" +STATE_INVALID = "invalid" +STATE_FAILED = "failed" +STATE_STALE = "stale" +STATE_MODEL_MISMATCH = "model-mismatch" +STATE_SHARD_MISMATCH = "shard-mismatch" +STATE_RECIPE_MISMATCH = "recipe-mismatch" +STATE_CATALOGUE_INCOMPATIBLE = "catalogue-incompatible" + +ALL_STATES = ( + STATE_ADMITTED, + STATE_ABSENT, + STATE_INVALID, + STATE_FAILED, + STATE_STALE, + STATE_MODEL_MISMATCH, + STATE_SHARD_MISMATCH, + STATE_RECIPE_MISMATCH, + STATE_CATALOGUE_INCOMPATIBLE, +) + +# --- Compatibility policy for nodes that predate the capability protocol. --- +# `compat` — a node presenting *no* proof still routes (legacy behaviour), but a +# node presenting a *bad* proof never does. Presenting a broken or +# mismatched proof is a stronger signal than presenting none. +# `enforce` — only `admitted` routes. Absent proof is not routable. +POLICY_COMPAT = "compat" +POLICY_ENFORCE = "enforce" +ALL_POLICIES = (POLICY_COMPAT, POLICY_ENFORCE) +DEFAULT_POLICY = POLICY_COMPAT + +POLICY_ENV_VAR = "MESHNET_TRACKER_CAPABILITY_POLICY" + +# Operator-facing detail strings are short and never carry a raw exception. +_MAX_DETAIL_CHARS = 240 +_MAX_DIAGNOSTICS = 3 + +_CREDENTIAL_PATTERNS = ( + re.compile(r"\b[A-Za-z0-9_]{2,6}_[A-Za-z0-9]{16,}\b"), # hf_…, ghp_…, sk_live_… + re.compile(r"\bsk-[A-Za-z0-9_-]{16,}\b"), + re.compile(r"(?i)\bbearer\s+\S+"), + re.compile(r"(?i)\b(?:token|api[_-]?key|password|secret)\s*[=:]\s*\S+"), +) +_REDACTED = "[redacted]" + + +def normalize_policy(value: Any) -> str: + """Return a known policy name, falling back to the default for anything else.""" + if isinstance(value, str) and value.strip().lower() in ALL_POLICIES: + return value.strip().lower() + return DEFAULT_POLICY + + +def policy_from_env(environ: Mapping[str, str] | None = None) -> str: + env = os.environ if environ is None else environ + return normalize_policy(env.get(POLICY_ENV_VAR)) + + +def sanitize_detail(text: Any) -> str: + """Collapse, redact and clip a string bound for an operator view.""" + cleaned = " ".join(str(text).split()) + for pattern in _CREDENTIAL_PATTERNS: + cleaned = pattern.sub(_REDACTED, cleaned) + if len(cleaned) > _MAX_DETAIL_CHARS: + cleaned = cleaned[: _MAX_DETAIL_CHARS - 1].rstrip() + "…" + return cleaned + + +def catalogue_is_compatible(version: Any) -> bool: + """True when `version` is at least `MIN_CATALOGUE_VERSION`. + + Versions are dotted integer sequences (`2026.07.1`). Anything that does not + parse is incompatible — an unparseable catalogue version cannot be shown to + be new enough. + """ + parsed = _parse_version(version) + if parsed is None: + return False + return parsed >= _parse_version(MIN_CATALOGUE_VERSION) # type: ignore[operator] + + +def _parse_version(value: Any) -> tuple[int, ...] | None: + if not isinstance(value, str) or not value.strip(): + return None + parts = value.strip().split(".") + try: + return tuple(int(part) for part in parts) + except ValueError: + return None + + +@dataclass(frozen=True) +class CapabilityState: + """The tracker's sanitized verdict on one node's presented proof. + + This is what the network map exposes and what route selection consults. It + holds identity labels and a verdict — never a raw exception, a file path, or + a credential. + """ + + state: str + detail: str = "" + model_id: str | None = None + shard_start: int | None = None + shard_end: int | None = None + recipe_id: str | None = None + recipe_version: str | None = None + catalogue_version: str | None = None + backend_id: str | None = None + device: str | None = None + quantization: str | None = None + validated_at: float | None = None + recorded_at: float = 0.0 + schema_version: int | None = None + diagnostics: tuple[str, ...] = () + + @property + def proven(self) -> bool: + """The presented proof covers exactly what the node advertised.""" + return self.state == STATE_ADMITTED + + def routable_under(self, policy: str) -> bool: + if self.proven: + return True + return self.state == STATE_ABSENT and normalize_policy(policy) == POLICY_COMPAT + + def with_state(self, state: str, detail: str) -> CapabilityState: + """Re-verdict a recorded proof against what the node advertises *now*.""" + return replace(self, state=state, detail=sanitize_detail(detail)) + + def to_dict(self) -> dict: + return { + "state": self.state, + "detail": self.detail, + "model_id": self.model_id, + "shard_start": self.shard_start, + "shard_end": self.shard_end, + "recipe_id": self.recipe_id, + "recipe_version": self.recipe_version, + "catalogue_version": self.catalogue_version, + "backend_id": self.backend_id, + "device": self.device, + "quantization": self.quantization, + "validated_at": self.validated_at, + "recorded_at": self.recorded_at, + "schema_version": self.schema_version, + "diagnostics": list(self.diagnostics), + } + + +def absent_state(detail: str = "", *, now: float | None = None) -> CapabilityState: + """The verdict for a node that presented no proof at all (legacy node).""" + return CapabilityState( + state=STATE_ABSENT, + detail=sanitize_detail( + detail + or "node registered without a capability report; it predates the " + "capability protocol or ran with admission disabled" + ), + recorded_at=time.time() if now is None else now, + ) + + +def evaluate_report( + report: Any, + *, + model_matches: Callable[[str], bool], + advertised_model: str | None, + shard_start: int | None, + shard_end: int | None, + declared_recipe_id: str | None = None, + declared_recipe_version: str | None = None, + now: float | None = None, + max_age_seconds: float = DEFAULT_MAX_REPORT_AGE_SECONDS, +) -> CapabilityState: + """Judge the proof a node presented against what that node is advertising. + + `model_matches` is the tracker's own alias-aware comparison against the + node's registered model / hf_repo, so an opaque model id never has to be + parsed here. + + Returns a verdict for *every* input, including malformed ones: a bad proof is + recorded and shown to the operator rather than dropped, so "why is my node not + routing" has an answer in the network map. + """ + now = time.time() if now is None else now + + if report is None: + return absent_state(now=now) + + if not isinstance(report, Mapping): + return CapabilityState( + state=STATE_INVALID, + detail=sanitize_detail( + f"capability_report must be a JSON object, got " + f"{type(report).__name__}" + ), + recorded_at=now, + ) + + try: + parsed = _parse_report(report) + except _ReportError as exc: + return CapabilityState( + state=STATE_INVALID, + detail=sanitize_detail(str(exc)), + recorded_at=now, + schema_version=_maybe_int(report.get("schema_version")), + ) + + status = parsed.pop("_status") + base = CapabilityState(state=STATE_ADMITTED, recorded_at=now, **parsed) + + if base.schema_version != SUPPORTED_SCHEMA_VERSION: + return base.with_state( + STATE_INVALID, + f"capability report declares schema version {base.schema_version}; " + f"this tracker reads version {SUPPORTED_SCHEMA_VERSION}", + ) + + if not catalogue_is_compatible(base.catalogue_version): + return base.with_state( + STATE_CATALOGUE_INCOMPATIBLE, + f"recipe catalogue {base.catalogue_version!r} is older than the " + f"minimum this tracker trusts ({MIN_CATALOGUE_VERSION}); upgrade the node", + ) + + if not model_matches(base.model_id or ""): + return base.with_state( + STATE_MODEL_MISMATCH, + f"proof is for model {base.model_id!r}, but the node registered " + f"{advertised_model!r}", + ) + + if shard_start is not None and shard_end is not None: + if (base.shard_start, base.shard_end) != (shard_start, shard_end): + return base.with_state( + STATE_SHARD_MISMATCH, + f"proof is for layers {base.shard_start}–{base.shard_end}, but the " + f"node registered layers {shard_start}–{shard_end}", + ) + + if declared_recipe_id is not None and base.recipe_id != declared_recipe_id: + return base.with_state( + STATE_RECIPE_MISMATCH, + f"proof is for recipe {base.recipe_id!r}, but the node declared it " + f"serves with {declared_recipe_id!r}", + ) + if ( + declared_recipe_version is not None + and base.recipe_version != declared_recipe_version + ): + return base.with_state( + STATE_RECIPE_MISMATCH, + f"proof is for recipe {base.recipe_id!r} v{base.recipe_version}, but " + f"the node declared v{declared_recipe_version}", + ) + + if status != STATUS_PASSED: + return base.with_state( + STATE_FAILED, + f"capability validation {status} on the node" + + (f" — {' '.join(base.diagnostics)}" if base.diagnostics else ""), + ) + + age = now - (base.validated_at or 0.0) + if age > max_age_seconds: + return base.with_state( + STATE_STALE, + f"proof is {age / 60:.0f} min old (limit {max_age_seconds / 60:.0f} min); " + "the node must re-validate before it can be routed", + ) + if age < -MAX_CLOCK_SKEW_SECONDS: + return base.with_state( + STATE_STALE, + f"proof is timestamped {-age:.0f}s in the future; check the node's clock", + ) + + return base.with_state( + STATE_ADMITTED, + f"{base.model_id} layers {base.shard_start}–{base.shard_end} proven on " + f"{base.device} with recipe {base.recipe_id} (v{base.recipe_version})", + ) + + +class _ReportError(ValueError): + """Malformed report input. Messages name the field, never echo a payload.""" + + +def _parse_report(doc: Mapping[str, Any]) -> dict: + model = _object(doc.get("model"), "model") + shard = _object(doc.get("shard"), "shard") + recipe = _object(doc.get("recipe"), "recipe") + backend = _object(doc.get("backend"), "backend") + + validated_at = doc.get("validated_at") + if isinstance(validated_at, bool) or not isinstance(validated_at, (int, float)): + raise _ReportError("'validated_at' must be a Unix timestamp") + + schema_version = doc.get("schema_version") + if isinstance(schema_version, bool) or not isinstance(schema_version, int): + raise _ReportError("'schema_version' must be an integer") + + return { + "model_id": _text(model.get("model_id"), "model.model_id"), + "shard_start": _index(shard.get("start"), "shard.start"), + "shard_end": _index(shard.get("end"), "shard.end"), + "recipe_id": _text(recipe.get("recipe_id"), "recipe.recipe_id"), + "recipe_version": _text(recipe.get("recipe_version"), "recipe.recipe_version"), + "catalogue_version": _text( + recipe.get("catalogue_version"), "recipe.catalogue_version" + ), + "backend_id": _text(backend.get("backend_id"), "backend.backend_id"), + "device": _text(backend.get("device"), "backend.device"), + "quantization": _optional_text( + backend.get("quantization"), "backend.quantization" + ), + "validated_at": float(validated_at), + "schema_version": schema_version, + "diagnostics": _diagnostics(doc.get("diagnostics")), + "_status": _text(doc.get("status"), "status"), + } + + +def _object(value: Any, field_name: str) -> Mapping[str, Any]: + if not isinstance(value, Mapping): + raise _ReportError(f"{field_name!r} must be a JSON object") + return value + + +def _text(value: Any, field_name: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise _ReportError(f"{field_name!r} must be a non-empty string") + return value + + +def _optional_text(value: Any, field_name: str) -> str | None: + if value is None: + return None + return _text(value, field_name) + + +def _index(value: Any, field_name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise _ReportError(f"{field_name!r} must be a non-negative integer") + return value + + +def _maybe_int(value: Any) -> int | None: + if isinstance(value, bool) or not isinstance(value, int): + return None + return value + + +def _diagnostics(value: Any) -> tuple[str, ...]: + if not isinstance(value, list): + return () + out = [ + sanitize_detail(item) + for item in value[:_MAX_DIAGNOSTICS] + if isinstance(item, str) and item.strip() + ] + return tuple(out) diff --git a/packages/tracker/meshnet_tracker/cli.py b/packages/tracker/meshnet_tracker/cli.py index 2cc25fe..79aec3f 100644 --- a/packages/tracker/meshnet_tracker/cli.py +++ b/packages/tracker/meshnet_tracker/cli.py @@ -9,6 +9,7 @@ from pathlib import Path from .accounts import DEFAULT_ACCOUNTS_DB_PATH from .billing import DEFAULT_BILLING_DB_PATH +from .capability import ALL_POLICIES as ALL_CAPABILITY_POLICIES from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH from .logging_setup import ( DEFAULT_LOG_BACKUP_COUNT, @@ -105,6 +106,18 @@ def main() -> None: metavar="PATH", help="SQLite database path for persistent model usage statistics", ) + common.add_argument( + "--capability-policy", + choices=list(ALL_CAPABILITY_POLICIES), + default=None, + help=( + "How to treat nodes that present no capability proof (ADR-0023): " + "'compat' (default) still routes pre-capability nodes; 'enforce' routes " + "only nodes whose proof covers the model and shard they advertise. " + "A broken or mismatched proof is never routed under either policy. " + "Falls back to $MESHNET_TRACKER_CAPABILITY_POLICY when omitted." + ), + ) common.add_argument( "--relay-url", default=None, @@ -416,6 +429,7 @@ def main() -> None: cluster_peers=cluster_peers or None, cluster_self_url=args.self_url, stats_db=getattr(args, "stats_db", None), + capability_policy=getattr(args, "capability_policy", None), relay_url=relay_url, embedded_relay=args.embedded_relay, embedded_relay_host=args.relay_host, diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 533f9e1..824c45a 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -6,10 +6,14 @@ HTTP API contract: "endpoint": "http://host:port", "shard_start": int, "shard_end": int, "model": str optional, "shard_checksum": str optional, "hardware_profile": object, "wallet_address": str optional, - "score": number optional + "score": number optional, + "capability_report": object optional, # ADR-0023 proof of this exact shard + "recipe_id": str optional, "recipe_version": str optional } - Response 200: {"node_id": str} + Response 200: {"node_id": str, "capability": {"state": str, "routable": bool, ...}} Response 400: {"error": str} + A node whose proof does not admit what it advertises still registers (so the + operator can see why it is dark) but is not routable -- see `capability.py`. - POST /v1/nodes/{node_id}/heartbeat Response 200: {} Response 404: {"error": "node not found"} @@ -48,6 +52,20 @@ from typing import Any from .accounts import DEFAULT_ACCOUNTS_DB_PATH, AccountStore from .auth import is_validator_token, sign_hive_request, verify_hive_request +from .capability import ( + DEFAULT_POLICY as DEFAULT_CAPABILITY_POLICY, + POLICY_COMPAT, + POLICY_ENFORCE, + STATE_ABSENT, + STATE_ADMITTED, + STATE_MODEL_MISMATCH, + STATE_SHARD_MISMATCH, + CapabilityState, + absent_state, + evaluate_report, + normalize_policy, + policy_from_env, +) from .wallet_proof import binding_message, verify_wallet_signature from .billing import DEFAULT_BILLING_DB_PATH, BillingLedger from .calibration import DEFAULT_CALIBRATION_DB_PATH, ToplocCalibrationStore @@ -587,6 +605,8 @@ class _NodeEntry: "heartbeats_expected", "heartbeats_received", # dynamic reassignment queued by the tracker "pending_new_assignment", + # the tracker's verdict on the capability proof this node presented + "capability", ) def __init__( @@ -616,6 +636,7 @@ class _NodeEntry: cert_fingerprint: str | None = None, peer_id: str | None = None, friendly_name: str | None = None, + capability: "CapabilityState | None" = None, ) -> None: self.node_id = node_id self.endpoint = endpoint @@ -643,6 +664,9 @@ class _NodeEntry: self.cert_fingerprint = cert_fingerprint self.peer_id = peer_id self.friendly_name = friendly_name + # No proof presented is `absent`, never `admitted` — a node can only earn + # `admitted` by presenting a report that covers what it advertises. + self.capability: CapabilityState = capability or absent_state() self.pending_directives: list[dict] = [] self.last_heartbeat: float = time.monotonic() self.total_requests: int = 0 @@ -731,12 +755,88 @@ def _reputation_multiplier(node: "_NodeEntry", contracts: Any | None) -> float: return 0.5 + 0.5 * reputation +def _node_admission(node: "_NodeEntry") -> CapabilityState: + """The node's recorded verdict, re-checked against what it advertises *now*. + + A proof covers one model/shard combination. The tracker can move a node to a + different range after it registered (rebalance, `pending_new_assignment`), and + the proof does not travel with it: until the node re-registers with a proof + for the range it now advertises, it is shard-mismatched and not routable. + """ + state = node.capability + if not state.proven: + return state + if state.model_id and not _node_matches_model(node, state.model_id): + return state.with_state( + STATE_MODEL_MISMATCH, + f"proof is for model {state.model_id!r}, but the node now serves " + f"{node.hf_repo or node.model!r}", + ) + if ( + node.shard_start is not None + and node.shard_end is not None + and (state.shard_start, state.shard_end) != (node.shard_start, node.shard_end) + ): + return state.with_state( + STATE_SHARD_MISMATCH, + f"proof is for layers {state.shard_start}–{state.shard_end}, but the " + f"node now serves layers {node.shard_start}–{node.shard_end}", + ) + return state + + +def _capability_from_registration( + payload: dict, + *, + model: str | None, + hf_repo: str | None, + shard_start: int | None, + shard_end: int | None, +) -> CapabilityState: + """The tracker's verdict on the proof carried by one registration payload. + + Used by the register handler and by the Raft follower path, so a replicated + registration lands with the same verdict as the one the leader recorded. + """ + aliases = _model_aliases(model) | _model_aliases(hf_repo) + recipe_id = payload.get("recipe_id") + recipe_version = payload.get("recipe_version") + return evaluate_report( + payload.get("capability_report"), + model_matches=lambda reported: bool(_model_aliases(reported) & aliases), + advertised_model=hf_repo or model, + shard_start=shard_start, + shard_end=shard_end, + declared_recipe_id=recipe_id if isinstance(recipe_id, str) else None, + declared_recipe_version=( + recipe_version if isinstance(recipe_version, str) else None + ), + ) + + +def _capability_routable(node: "_NodeEntry", policy: str) -> bool: + """May this node carry traffic under the tracker's capability policy?""" + return _node_admission(node).routable_under(policy) + + +def _admitted_nodes(nodes: list["_NodeEntry"], policy: str | None) -> list["_NodeEntry"]: + """Drop every candidate whose capability proof does not admit it (ADR-0023). + + This is the single gate every route path goes through. It removes candidates; + it never reorders or reweights them, so coverage-first selection and + throughput-weighted preference among the survivors are unchanged. + """ + effective = normalize_policy(policy) + return [node for node in nodes if _capability_routable(node, effective)] + + def _select_route( nodes: list[_NodeEntry], required_start: int, required_end: int, model: str | None = None, contracts: Any | None = None, + policy: str | None = None, ) -> tuple[list[_NodeEntry], str]: """Greedy interval-cover biased toward fast, lightly-loaded, reputable nodes. @@ -747,7 +847,10 @@ def _select_route( Tiebreak: higher shard_end (fewer hops). """ candidates = sorted( - [node for node in nodes if node.shard_start is not None and node.shard_end is not None], + [ + node for node in _admitted_nodes(nodes, policy) + if node.shard_start is not None and node.shard_end is not None + ], key=lambda n: (n.shard_start, -n.shard_end), # type: ignore[operator] ) route: list[_NodeEntry] = [] @@ -783,6 +886,7 @@ def _enumerate_routes( model: str | None = None, contracts: Any | None = None, max_candidates: int = 8, + policy: str | None = None, ) -> list["RouteCandidate"]: """Enumerate viable route candidates for bandit selection (ADR-0021). @@ -791,9 +895,13 @@ def _enumerate_routes( with the longest-advancing hops. The route's prior throughput estimate is its bottleneck hop's queue-adjusted effective throughput — used only until observed route samples exist. + + Candidates that are not admitted by their capability proof are dropped before + enumeration, so an unproven node cannot appear in any route candidate — not + even as a scouted one. """ sharded = [ - n for n in nodes + n for n in _admitted_nodes(nodes, policy) if n.shard_start is not None and n.shard_end is not None ] # Heads must start the pipeline at the first required layer (they tokenize @@ -2675,9 +2783,13 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer): route_stats: "RouteStatsStore | None" = None, relay_status: dict | None = None, test_runner: "TestRunManager | None" = None, + capability_policy: str | None = None, ) -> None: super().__init__(addr, handler) self.registry = registry + self.capability_policy = normalize_policy( + capability_policy if capability_policy is not None else policy_from_env() + ) self.lock = lock self.heartbeat_timeout = heartbeat_timeout self.model_presets = model_presets @@ -3177,9 +3289,19 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): def model_supply_for(node: _NodeEntry) -> dict: return _model_health_summary(server, node.model, node.hf_repo) + def capability_for(node: _NodeEntry) -> dict: + # Re-verdicted against what the node advertises now, so the operator + # view and the routing gate can never disagree. + state = _node_admission(node) + return { + **state.to_dict(), + "routable": state.routable_under(server.capability_policy), + } + self._send_json(200, { "relay_url": server.relay_url, "relay": dict(server.relay_status), + "capability_policy": server.capability_policy, "pool": _pool_summary(nodes), "memory_pool": memory_pool, "recommended_models": [ @@ -3213,6 +3335,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "capacity": capacity_for(node), "model_supply": model_supply_for(node), "throughput": throughput_for(node), + "capability": capability_for(node), "stats": _node_health(node, server.heartbeat_timeout), } for node in nodes @@ -3365,6 +3488,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): if n.tracker_mode and _node_matches_model(n, model) and _quantization_satisfies(n.quantization, requested_quantization) + and _capability_routable(n, server.capability_policy) ] if not candidates: @@ -3375,6 +3499,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): if n.shard_start == 0 and _node_matches_model(n, model) and _quantization_satisfies(n.quantization, requested_quantization) + and _capability_routable(n, server.capability_policy) ] if not candidates: @@ -3480,6 +3605,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): model=route_model, contracts=server.contracts, max_candidates=server.route_stats.config.max_candidate_routes, + policy=server.capability_policy, ) picked, routing_decision = choose_route( route_candidates, server.route_stats, route_model, rng=server.route_rng, @@ -3489,7 +3615,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): else: # No head-anchored candidate — legacy greedy cover as fallback # (also produces the layer-gap error message). - route_nodes, route_error = _select_route(all_nodes, rs, re, model=route_model, contracts=server.contracts) + route_nodes, route_error = _select_route( + all_nodes, rs, re, + model=route_model, + contracts=server.contracts, + policy=server.capability_policy, + ) routing_decision = {"mode": "greedy-fallback"} if route_error: _tracker_log( @@ -4344,6 +4475,25 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._send_json(400, {"error": str(exc)}) return + recipe_id = body.get("recipe_id") + recipe_version = body.get("recipe_version") + if (recipe_id is not None and not isinstance(recipe_id, str)) or ( + recipe_version is not None and not isinstance(recipe_version, str) + ): + self._send_json(400, {"error": "recipe_id and recipe_version must be strings"}) + return + + # The capability proof (ADR-0023). A bad proof does not fail registration -- + # the node is recorded with its verdict so the operator can see *why* it is + # not routing -- but only an `admitted` verdict makes it routable. + capability = _capability_from_registration( + body, + model=model, + hf_repo=hf_repo, + shard_start=shard_start, + shard_end=shard_end, + ) + node_id = _node_id_for_registration( endpoint, model, @@ -4378,6 +4528,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): cert_fingerprint=cert_fingerprint, peer_id=peer_id, friendly_name=friendly_name, + capability=capability, ) with server.lock: self._purge_expired_nodes() @@ -4433,6 +4584,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): old_model_health=_model_health_summary(server, old.model, old.hf_repo), model_health=model_health, ) + routable = _capability_routable(entry, server.capability_policy) _tracker_log( server, "info", @@ -4444,7 +4596,19 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): shard=f"{entry.shard_start}-{entry.shard_end}", tracker_mode=entry.tracker_mode, model_health=model_health, + capability=entry.capability.state, + routable=routable, ) + if not routable: + _tracker_log( + server, + "warn", + "node registered but is not routable", + node_id=node_id, + endpoint=entry.endpoint, + capability=entry.capability.state, + detail=entry.capability.detail, + ) shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded" repo_info = f" [{hf_repo}]" if hf_repo else "" @@ -4456,7 +4620,17 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): flush=True, ) - payload = {"node_id": node_id} + payload = { + "node_id": node_id, + # Tell the node what the tracker made of its proof: a node that is + # registered but not routable must be able to see that it is dark. + "capability": { + "state": entry.capability.state, + "detail": entry.capability.detail, + "routable": routable, + "policy": server.capability_policy, + }, + } if assignment_directive is not None: payload["directive"] = assignment_directive self._send_json(200, payload) @@ -5989,6 +6163,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): model=route_model, contracts=server.contracts, max_candidates=server.route_stats.config.max_candidate_routes, + policy=server.capability_policy, ) if candidates: # Prefer a distributed multi-hop route when available. Greedy @@ -5998,7 +6173,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): error = "" else: route, error = _select_route( - alive, required_start, required_end, contracts=server.contracts, + alive, required_start, required_end, + contracts=server.contracts, + policy=server.capability_policy, ) if error: _tracker_log( @@ -6084,6 +6261,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): model=stats_key, contracts=server.contracts, max_candidates=cfg.max_candidate_routes, + policy=server.capability_policy, ) out[stats_key] = { "epoch": server.route_stats.epoch(stats_key), @@ -6177,7 +6355,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): routes = [] remaining = list(candidates) for _ in range(redundancy): - route, error = _select_route(remaining, required_start, required_end, contracts=server.contracts) + route, error = _select_route( + remaining, required_start, required_end, + contracts=server.contracts, + policy=server.capability_policy, + ) if error: self._send_json(503, {"error": error}) return @@ -6262,7 +6444,11 @@ class TrackerServer: routing_config: RoutingConfig | None = None, enable_test_runner: bool = False, test_runner: TestRunManager | None = None, + capability_policy: str | None = None, ) -> None: + self._capability_policy = normalize_policy( + capability_policy if capability_policy is not None else policy_from_env() + ) self._host = host self._requested_port = port self._heartbeat_timeout = heartbeat_timeout @@ -6475,6 +6661,7 @@ class TrackerServer: route_stats=self._route_stats, relay_status=http_relay_status, test_runner=self._test_runner, + capability_policy=self._capability_policy, ) self.port = self._server.server_address[1] @@ -6708,6 +6895,15 @@ class TrackerServer: else None ), friendly_name=_normalize_friendly_name(payload.get("friendly_name")), + # A replicated registration carries its proof: without this, a proven + # node would be routable on the leader and dark on every follower. + capability=_capability_from_registration( + payload, + model=payload.get("model", "stub-model"), + hf_repo=payload.get("hf_repo"), + shard_start=shard_start, + shard_end=shard_end, + ), ) with self._lock: self._registry[node_id] = entry diff --git a/tests/test_activation_compression.py b/tests/test_activation_compression.py new file mode 100644 index 0000000..176aa5d --- /dev/null +++ b/tests/test_activation_compression.py @@ -0,0 +1,63 @@ +"""Trace-driven activation-compression policy units.""" + +import os + +import pytest + +from meshnet_node.activation_compression import ( + CompressionPolicies, + CompressionPolicy, + compress_activation, + decompress_activation, +) + + +def test_compressible_body_uses_zstd_when_it_clears_savings_policy(): + body = b"activation" * 20_000 + result = compress_activation(body, CompressionPolicy(min_input_bytes=1, min_savings_bytes=100)) + assert result.encoding == "zstd" + assert result.output_bytes < result.input_bytes + assert decompress_activation(result.body, result.encoding).body == body + + +def test_incompressible_body_stays_raw_after_measured_trial(): + body = os.urandom(32 * 1024) + result = compress_activation(body, CompressionPolicy(min_input_bytes=1, min_savings_bytes=1)) + assert result.encoding is None + assert result.body == body + assert result.decision == "below_savings" + + +def test_small_body_uses_raw_fast_path_without_zstd_trial(): + body = b"x" * 1024 + result = compress_activation(body, CompressionPolicy(min_input_bytes=2048)) + assert result.encoding is None + assert result.decision == "below_min_input" + assert result.elapsed_seconds >= 0 + + +def test_threshold_requires_both_byte_and_ratio_savings(): + body = b"a" * 4096 + result = compress_activation( + body, CompressionPolicy(min_input_bytes=1, min_savings_bytes=len(body), min_savings_ratio=0), + ) + assert result.encoding is None + assert result.decision == "below_savings" + + +def test_malformed_zstd_and_legacy_raw_bodies_are_handled_explicitly(): + assert decompress_activation(b"legacy", None).body == b"legacy" + with pytest.raises(ValueError, match="invalid zstd activation body"): + decompress_activation(b"not a zstd frame", "zstd") + with pytest.raises(ValueError, match="unsupported"): + decompress_activation(b"body", "gzip") + + +def test_prefill_decode_and_route_conditions_have_independent_config(monkeypatch): + policies = CompressionPolicies() + assert policies.for_condition("relay", "prefill").min_input_bytes < policies.for_condition("relay", "decode").min_input_bytes + monkeypatch.setenv("MESHNET_COMPRESSION_RELAY_DECODE_MIN_INPUT_BYTES", "123") + monkeypatch.setenv("MESHNET_COMPRESSION_LAN_PREFILL_ENABLED", "false") + assert policies.for_condition("relay", "decode").min_input_bytes == 123 + assert not policies.for_condition("lan", "prefill").enabled + assert policies.for_condition("benchmark", "prefill").enabled diff --git a/tests/test_dynamic_routing.py b/tests/test_dynamic_routing.py index 62c4dee..7b1d5cc 100644 --- a/tests/test_dynamic_routing.py +++ b/tests/test_dynamic_routing.py @@ -15,6 +15,7 @@ from meshnet_tracker.routing_stats import ( route_signature, route_table, ) +from meshnet_tracker.capability import absent_state from meshnet_tracker.server import TrackerServer, _enumerate_routes @@ -49,6 +50,9 @@ def _fake_node(node_id, shard_start, shard_end, benchmark=100.0, endpoint=None): proxy_inflight=0, wallet_address=None, relay_addr=None, + # A pre-capability node (NCA-004): routable only under the `compat` + # policy, which is what these route-enumeration tests exercise. + capability=absent_state(), ) @@ -270,6 +274,7 @@ def test_proxy_head_is_route_head_and_routing_endpoint_lists_routes(): "shard_start": 0, "shard_end": shard_end, "tracker_mode": True, + "quantization": "bfloat16", "benchmark_tokens_per_sec": bench, "hardware_profile": {}, "score": 1.0}, diff --git a/tests/test_gossip_and_relay.py b/tests/test_gossip_and_relay.py index 46584ef..38ab46c 100644 --- a/tests/test_gossip_and_relay.py +++ b/tests/test_gossip_and_relay.py @@ -454,6 +454,137 @@ def test_relay_rpc_reuses_connection_for_sequential_requests(monkeypatch): assert connection_attempts == 1 +def test_relay_hop_client_preserves_binary_json_and_closes_uncertain_legacy_socket(monkeypatch): + "DIP-002: persistent clients correlate both frame formats and never replay a lost response.\n\nTags: gossip, network, relay" + + import websockets.sync.client as wsc # type: ignore[import] + from meshnet_node.relay_bridge import decode_binary_frame, encode_binary_frame + from meshnet_node.torch_server import _RelayHopClient, _RelayRequestUncertainError + + class FakeSocket: + def __init__(self, responses): + self.responses = iter(responses) + self.sent = [] + self.closed = False + + def send(self, frame): + self.sent.append(frame) + + def recv(self, timeout): + response = next(self.responses) + if isinstance(response, BaseException): + raise response + request, _ = decode_binary_frame(self.sent[-1]) + if isinstance(response, bytes): + return encode_binary_frame({ + "request_id": request["request_id"], "status": 200, "headers": {}, + }, response) + return json.dumps({"request_id": request["request_id"], "status": 200, + "headers": {}, "body": response}) + + def close(self): + self.closed = True + + healthy = FakeSocket([b"binary", "json"]) + legacy = FakeSocket([b"first", TimeoutError("legacy relay closed")]) + session_a = FakeSocket([b"a"]) + session_b = FakeSocket([b"b"]) + sockets = iter([healthy, legacy, session_a, session_b]) + monkeypatch.setattr(wsc, "connect", lambda *args, **kwargs: next(sockets)) + + client = _RelayHopClient("ws://relay/rpc/peer", timeout=0.01) + assert client.request("/forward", b"a", {})[2] == b"binary" + assert client.request("/forward", b"b", {})[2] == b"json" + first_id = decode_binary_frame(healthy.sent[0])[0]["request_id"] + second_id = decode_binary_frame(healthy.sent[1])[0]["request_id"] + assert first_id != second_id + client.close() + + client = _RelayHopClient("ws://relay/rpc/peer", timeout=0.01) + assert client.request("/forward", b"first", {})[2] == b"first" + try: + client.request("/forward", b"second", {}) + except _RelayRequestUncertainError: + pass + else: + raise AssertionError("a post-send disconnect must not be retried") + assert legacy.closed is True + assert client._ws is None + + # Route Sessions own their own requester socket; serialising one client's + # calls must not accidentally make another session share it. + first = _RelayHopClient("ws://relay/rpc/peer", timeout=0.01) + second = _RelayHopClient("ws://relay/rpc/peer", timeout=0.01) + assert first.request("/forward", b"a", {})[2] == b"a" + assert second.request("/forward", b"b", {})[2] == b"b" + assert first._ws is session_a + assert second._ws is session_b + + +def test_relay_rpc_cleans_pending_on_timeout_disconnect_and_cancellation(): + "DIP-002: every terminal relay RPC path releases its pending response queue.\n\nTags: gossip, network, relay" + + import asyncio + + from meshnet_relay.server import RelayServer + + class Requester: + def __init__(self): + self.messages = [json.dumps({ + "request_id": "legacy-id", "method": "POST", "path": "/forward", + "headers": {}, "body": "{}", + })] + self.sent = [] + + async def recv(self): + if self.messages: + return self.messages.pop(0) + raise EOFError + + async def send(self, message): + self.sent.append(json.loads(message)) + + async def close(self, *args): + pass + + class Target: + async def send(self, message): + pass + + async def wait_for_pending(relay): + for _ in range(100): + if relay._pending_rpc: + return + await asyncio.sleep(0) + raise AssertionError("relay RPC did not register pending state") + + async def exercise(kind): + relay = RelayServer(rpc_timeout=0.02, rpc_idle_timeout=0.01) + relay.registry.register("peer", "", Target()) + requester = Requester() + task = asyncio.create_task(relay._handle_rpc(requester, "peer")) + await wait_for_pending(relay) + if kind == "disconnect": + relay._fail_pending_for_peer("peer") + elif kind == "cancel": + task.cancel() + if kind == "cancel": + try: + await task + except asyncio.CancelledError: + pass + else: + await task + assert relay._pending_rpc == {} + if kind == "timeout": + assert requester.sent[-1]["status"] == 504 + if kind == "disconnect": + assert requester.sent[-1]["status"] == 503 + + for kind in ("timeout", "disconnect", "cancel"): + asyncio.run(exercise(kind)) + + def test_binary_relay_frame_codecs_interoperate(): "Node and relay ship the same binary frame format as separate copies.\n\nTags: gossip, network, relay" @@ -478,6 +609,24 @@ def test_binary_relay_frame_codecs_interoperate(): raise AssertionError("garbage bytes must not decode as a binary frame") +def test_binary_relay_frame_layout_remains_byte_for_byte_compatible(): + """Framing optimizations preserve the MRF1 header-length-body contract. + + Tags: gossip, network, relay, wire + """ + import json + + from meshnet_node.relay_bridge import BINARY_FRAME_MAGIC, encode_binary_frame + + header = {"request_id": "r-1", "headers": {"X-Meshnet-Shape": "1,1,2"}} + body = b"\x01\x02\x03\x04" + header_bytes = json.dumps(header, separators=(",", ":")).encode() + + assert encode_binary_frame(header, body) == ( + BINARY_FRAME_MAGIC + len(header_bytes).to_bytes(4, "big") + header_bytes + body + ) + + def test_activation_compression_round_trips_and_skips_small_bodies(): "Pipeline hops zstd-compress large activations; tiny decode bodies pass raw.\n\nTags: gossip, network, relay" diff --git a/tests/test_http_keepalive.py b/tests/test_http_keepalive.py new file mode 100644 index 0000000..58e1b5a --- /dev/null +++ b/tests/test_http_keepalive.py @@ -0,0 +1,134 @@ +"""DIP-003 keep-alive ownership and framing-adjacent transport tests.""" + +from __future__ import annotations + +import io + +import pytest + + +class _Response: + def __init__(self, status: int = 200, body: bytes = b"ok", headers: dict | None = None): + self.status = status + self._body = body + self.headers = headers or {"Content-Type": "application/octet-stream", "Content-Length": str(len(body))} + + def read(self) -> bytes: + return self._body + + def close(self) -> None: + pass + + +class _Connection: + instances: list["_Connection"] = [] + + def __init__(self, *args, **kwargs): + self.requests: list[tuple] = [] + self.responses: list[object] = [_Response(), _Response()] + self.closed = False + self.__class__.instances.append(self) + + def request(self, *args, **kwargs) -> None: + self.requests.append((args, kwargs)) + + def getresponse(self): + response = self.responses.pop(0) + if isinstance(response, BaseException): + raise response + return response + + def close(self) -> None: + self.closed = True + + +def test_direct_hop_client_reuses_one_connection_and_discards_uncertain_socket(monkeypatch): + """A Route Session owns one serialized direct socket and never replays it. + + Tags: performance, routing + """ + from meshnet_node import torch_server + + _Connection.instances = [] + monkeypatch.setattr(torch_server.http.client, "HTTPConnection", _Connection) + client = torch_server._DirectHopClient("http://tail.example:8001") + + assert client.request("/forward", b"one", {})[2] == b"ok" + assert client.request("/forward", b"two", {})[2] == b"ok" + assert len(_Connection.instances) == 1 + assert [call[0][1] for call in _Connection.instances[0].requests] == ["/forward", "/forward"] + + _Connection.instances[0].responses.append(ConnectionResetError("stale peer")) + with pytest.raises(torch_server._DirectRequestUncertainError): + client.request("/forward", b"three", {}) + assert _Connection.instances[0].closed is True + assert client._connection is None + + +def test_bridge_pool_reuses_a_worker_connection_and_invalidates_stale_one(monkeypatch): + """A bridge worker keeps its own loopback client; broken clients are dropped. + + Tags: performance, relay + """ + from meshnet_node import relay_bridge + + _Connection.instances = [] + monkeypatch.setattr(relay_bridge.http.client, "HTTPConnection", _Connection) + bridge = relay_bridge.RelayHttpBridge( + relay_url="ws://relay.example/ws", + peer_id="peer", + local_base_url="http://127.0.0.1:8001", + advertised_addr="", + ) + frames: list[dict] = [] + bridge._send_response_frame = lambda frame: (frames.append(frame), True)[1] + + request = {"request_id": "one", "method": "POST", "path": "/forward", "headers": {}, "body": ""} + bridge._process_request(request) + bridge._process_request({**request, "request_id": "two"}) + assert len(_Connection.instances) == 1 + assert len(_Connection.instances[0].requests) == 2 + + _Connection.instances[0].responses.append(ConnectionResetError("stale loopback")) + bridge._process_request({**request, "request_id": "broken"}) + assert _Connection.instances[0].closed is True + bridge._process_request({**request, "request_id": "replacement"}) + assert len(_Connection.instances) == 2 + bridge.stop() + + +def test_node_sse_uses_chunked_framing_and_tolerates_client_cancellation(): + """HTTP/1.1 streams terminate without EOF and ignore a cancelled client. + + Tags: performance, streaming + """ + from meshnet_node.torch_server import _TorchHandler + + handler = object.__new__(_TorchHandler) + headers: list[tuple[str, str]] = [] + handler.send_response = lambda status: None + handler.send_header = lambda name, value: headers.append((name, value)) + handler.end_headers = lambda: None + handler.wfile = io.BytesIO() + emit = handler._start_openai_stream("model") + emit(None) + wire = handler.wfile.getvalue() + + assert _TorchHandler.protocol_version == "HTTP/1.1" + assert ("Transfer-Encoding", "chunked") in headers + assert wire.endswith(b"0\r\n\r\n") + assert b"data: [DONE]" in wire + + class _CancelledWriter: + def write(self, _body): + raise BrokenPipeError + + def flush(self): + raise BrokenPipeError + + cancelled = object.__new__(_TorchHandler) + cancelled.send_response = lambda status: None + cancelled.send_header = lambda name, value: None + cancelled.end_headers = lambda: None + cancelled.wfile = _CancelledWriter() + cancelled._start_openai_stream("model")(None) diff --git a/tests/test_mining_cli.py b/tests/test_mining_cli.py index a19b505..0f3971d 100644 --- a/tests/test_mining_cli.py +++ b/tests/test_mining_cli.py @@ -9,6 +9,10 @@ import types from pathlib import Path from unittest.mock import MagicMock, patch +# A fake node server has no real backend to prove capability with; say so +# explicitly rather than bypassing startup's fail-closed admission. +from meshnet_node.testing import assume_capability + # --------------------------------------------------------------------------- # model_catalog tests @@ -336,6 +340,7 @@ def test_startup_auto_detects_shard_range(monkeypatch, tmp_path): # shard_start and shard_end intentionally omitted quantization="bfloat16", host="127.0.0.1", + capability_validator=assume_capability, ) assert calls == ["Qwen/Qwen2.5-0.5B-Instruct"] assert isinstance(node, FakeNode) diff --git a/tests/test_node_admission.py b/tests/test_node_admission.py new file mode 100644 index 0000000..9b606e2 --- /dev/null +++ b/tests/test_node_admission.py @@ -0,0 +1,396 @@ +"""NCA-003: startup fails closed — no registration without a fresh matching proof. + +Two layers are covered here: + +* the gate itself (`meshnet_node.admission.admit`) — which reports admit a + selection, and which are refused as failed, stale, or about something else; +* `run_startup` — that a refused report means the tracker is never called, and + that an admitted one travels with the registration payload. + +Torch is a stub in the dev venv, so the backend is faked by duck-typing +`TorchModelShard` (see `_FakeBackend`); the *production* validator still runs a +real `doctor` forward against it, so the fail-closed path is exercised without a +bypass. Tests that cannot supply an executable backend pass the explicit +test-safe validator from `meshnet_node.testing`. +""" + +import base64 +import struct +import time + +import pytest + +from meshnet_node.admission import ( + REASON_BACKEND_MISMATCH, + REASON_MODEL_MISMATCH, + REASON_NO_REPORT, + REASON_NOT_PASSED, + REASON_RECIPE_MISMATCH, + REASON_SHARD_MISMATCH, + REASON_STALE, + AdmissionRequirement, + CapabilityAdmissionError, + CapabilityContext, + admit, + probe_capability, +) +from meshnet_node.capability import STATUS_FAILED, STATUS_SKIPPED +from meshnet_node.doctor import DoctorSelection +from meshnet_node.recipe_manifest import DEFAULT_RECIPE_ID, load_recipe_manifest +from meshnet_node.startup import run_startup +from meshnet_node.testing import capability_report_for, capability_stub + +MODEL = "acme/opaque-model-7b" + + +class _FakeDevice: + def __init__(self, type_: str = "cpu"): + self.type = type_ + + +class _FakeOutput: + def __init__(self, hidden_size: int, tokens: int = 4): + self.body = b"\x00" * (tokens * hidden_size * 2) + self.shape = [1, tokens, hidden_size] + self.attention_mask_header = _int64_header([[1] * tokens]) + self.position_ids_header = _int64_header([list(range(tokens))]) + + +def _int64_header(rows): + flat = [int(v) for row in rows for v in row] + raw = struct.pack(f"<{len(flat)}q", *flat) + return f"{len(rows)},{len(rows[0])}:{base64.b64encode(raw).decode('ascii')}" + + +class _FakeBackend: + """Duck-types the parts of `TorchModelShard` the doctor probe touches.""" + + total_layers = 24 + hidden_size = 8 + + def __init__(self, *, shard_start=0, shard_end=23, device="cpu", forward_error=None): + self.shard_start = shard_start + self.shard_end = shard_end + self.is_head = shard_start == 0 + self.is_tail = shard_end == self.total_layers - 1 + self.device = _FakeDevice(device) + self.model_id = MODEL + self._forward_error = forward_error + + def encode_prompt(self, _prompt): + if self._forward_error: + raise self._forward_error + return _FakeOutput(self.hidden_size) + + def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): + if self._forward_error: + raise self._forward_error + return _FakeOutput(self.hidden_size) + + +def _context(backend=None, *, model_id=MODEL, shard_start=0, shard_end=23, device="cpu"): + manifest = load_recipe_manifest() + return CapabilityContext( + backend=backend, + selection=DoctorSelection( + model_id=model_id, + shard_start=shard_start, + shard_end=shard_end, + quantization="bfloat16", + ), + recipe=manifest.require(DEFAULT_RECIPE_ID), + manifest=manifest, + device=device, + ) + + +# --------------------------------------------------------------------------- +# The gate: which reports admit a selection +# --------------------------------------------------------------------------- + + +def test_a_fresh_matching_passing_report_admits_the_selection(): + "The proof covers exactly what is about to be advertised, so the node may register.\n\nTags: node, admission" + ctx = _context() + report = capability_report_for(ctx) + + assert admit(AdmissionRequirement.for_context(ctx), report) is report + + +def test_a_missing_report_is_refused(): + "No proof at all is the default state, and it must not register.\n\nTags: node, admission" + ctx = _context() + + with pytest.raises(CapabilityAdmissionError) as exc: + admit(AdmissionRequirement.for_context(ctx), None) + + assert exc.value.reason == REASON_NO_REPORT + + +@pytest.mark.parametrize("status", [STATUS_FAILED, STATUS_SKIPPED]) +def test_a_report_that_did_not_pass_is_refused(status): + "A failed or skipped validation is evidence against admission, not for it.\n\nTags: node, admission" + ctx = _context() + report = capability_report_for( + ctx, status=status, diagnostics=["the shard forward returned no output"] + ) + + with pytest.raises(CapabilityAdmissionError) as exc: + admit(AdmissionRequirement.for_context(ctx), report) + + assert exc.value.reason == REASON_NOT_PASSED + assert "the shard forward returned no output" in str(exc.value) + + +def test_a_passing_report_for_another_model_is_refused(): + "A proof about one model says nothing about another — no reuse across models.\n\nTags: node, admission" + ctx = _context() + report = capability_report_for(ctx, model_id="other/model-1b") + + with pytest.raises(CapabilityAdmissionError) as exc: + admit(AdmissionRequirement.for_context(ctx), report) + + assert exc.value.reason == REASON_MODEL_MISMATCH + assert "other/model-1b" in str(exc.value) + + +def test_a_passing_report_for_another_shard_range_is_refused(): + "Layers 0–11 running is no proof that layers 12–23 fit.\n\nTags: node, admission" + ctx = _context(shard_start=12, shard_end=23) + report = capability_report_for(ctx, shard_start=0, shard_end=11) + + with pytest.raises(CapabilityAdmissionError) as exc: + admit(AdmissionRequirement.for_context(ctx), report) + + assert exc.value.reason == REASON_SHARD_MISMATCH + + +def test_a_passing_report_for_another_recipe_version_is_refused(): + "A recipe's execution params changed, so its old proof no longer applies.\n\nTags: node, admission" + ctx = _context() + report = capability_report_for(ctx, recipe_version="0") + + with pytest.raises(CapabilityAdmissionError) as exc: + admit(AdmissionRequirement.for_context(ctx), report) + + assert exc.value.reason == REASON_RECIPE_MISMATCH + + +def test_a_passing_report_from_another_backend_or_device_is_refused(): + "A CUDA proof does not admit a node that would serve the shard on CPU.\n\nTags: node, admission" + ctx = _context(device="cpu") + report = capability_report_for(ctx, device="cuda") + + with pytest.raises(CapabilityAdmissionError) as exc: + admit(AdmissionRequirement.for_context(ctx), report) + + assert exc.value.reason == REASON_BACKEND_MISMATCH + + other_backend = capability_report_for(ctx, backend_id="some-other-runtime") + with pytest.raises(CapabilityAdmissionError) as exc: + admit(AdmissionRequirement.for_context(ctx), other_backend) + assert exc.value.reason == REASON_BACKEND_MISMATCH + + +def test_a_report_older_than_the_freshness_window_is_refused(): + "Hardware, drivers and weights move; an old proof is not a current one.\n\nTags: node, admission" + ctx = _context() + requirement = AdmissionRequirement.for_context(ctx, max_age_seconds=900) + report = capability_report_for(ctx, age_seconds=901) + + with pytest.raises(CapabilityAdmissionError) as exc: + admit(requirement, report) + + assert exc.value.reason == REASON_STALE + + still_fresh = capability_report_for(ctx, age_seconds=899) + assert admit(requirement, still_fresh) is still_fresh + + +def test_a_future_dated_report_is_refused(): + "A report from the future is a broken clock, not a fresh proof.\n\nTags: node, admission" + ctx = _context() + report = capability_report_for(ctx, validated_at=time.time() + 3600) + + with pytest.raises(CapabilityAdmissionError) as exc: + admit(AdmissionRequirement.for_context(ctx), report) + + assert exc.value.reason == REASON_STALE + + +# --------------------------------------------------------------------------- +# The production validator: a real forward through the loaded backend +# --------------------------------------------------------------------------- + + +def test_the_production_validator_proves_a_working_backend_with_a_real_forward(): + "probe_capability runs the doctor's bounded forward on the backend that would serve.\n\nTags: node, admission" + ctx = _context(_FakeBackend()) + + report = probe_capability(ctx) + + assert report.passed + assert admit(AdmissionRequirement.for_context(ctx), report) is report + + +def test_the_production_validator_fails_a_backend_that_cannot_execute(): + "A shard that loads but cannot run a forward yields a failed report, not a pass.\n\nTags: node, admission" + ctx = _context(_FakeBackend(forward_error=RuntimeError("CUDA out of memory"))) + + report = probe_capability(ctx) + + assert not report.passed + with pytest.raises(CapabilityAdmissionError) as exc: + admit(AdmissionRequirement.for_context(ctx), report) + assert exc.value.reason == REASON_NOT_PASSED + + +def test_the_production_validator_fails_a_node_with_no_backend_at_all(): + "A server with no model backend cannot prove anything, so it is not admitted.\n\nTags: node, admission" + ctx = _context(None) + + assert not probe_capability(ctx).passed + + +# --------------------------------------------------------------------------- +# run_startup: a refused report means the tracker is never called +# --------------------------------------------------------------------------- + + +class _FakeTorchNodeServer: + started = False + + def __init__(self, **kwargs): + self.kwargs = kwargs + self.backend = kwargs.pop("_backend", None) or _FakeBackend() + + def start(self): + type(self).started = True + return 7099 + + def stop(self): + pass + + +@pytest.fixture +def startup_env(monkeypatch): + """Fake hardware, wallet and tracker; records registrations *for this model*. + + Heartbeat threads that other tests leave running are daemon threads that + outlive their test and re-register through this same patched `_post_json`, so + a plain call log would be polluted by whatever ran before. Only posts naming + this test's model can have come from this test's node. + """ + import meshnet_node.startup as startup_mod + + posted: list[tuple[str, dict]] = [] + _FakeTorchNodeServer.started = False + + def _record(url, payload, timeout=10.0): + if url.endswith("/v1/nodes/register") and payload.get("hf_repo") == MODEL: + posted.append((url, payload)) + return {"node_id": "node-nca"} + + monkeypatch.setattr( + startup_mod, + "detect_hardware", + lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0, "ram_mb": 16384}, + ) + monkeypatch.setattr( + startup_mod, "benchmark_throughput_checked", lambda _device: (10.0, True, None) + ) + monkeypatch.setattr(startup_mod, "TorchNodeServer", _FakeTorchNodeServer) + monkeypatch.setattr( + startup_mod, "load_or_create_wallet", lambda **_kw: (b"", b"", "wallet-nca") + ) + monkeypatch.setattr( + startup_mod, "_get_json", lambda _url, timeout=10.0: {"relay_url": None, "nodes": []} + ) + monkeypatch.setattr(startup_mod, "_post_json", _record) + monkeypatch.setattr(startup_mod, "_start_heartbeat", lambda *a, **kw: None) + return posted + + +def _start(**kwargs): + return run_startup( + tracker_url="http://127.0.0.1:8080", + model_id=MODEL, + shard_start=0, + shard_end=23, + **kwargs, + ) + + +def test_backend_validation_failure_registers_nothing(startup_env, monkeypatch): + "A shard that cannot run a forward must never reach /v1/nodes/register.\n\nTags: node, admission, startup" + import meshnet_node.startup as startup_mod + + broken = _FakeBackend(forward_error=RuntimeError("CUDA out of memory")) + monkeypatch.setattr( + startup_mod, + "TorchNodeServer", + lambda **kw: _FakeTorchNodeServer(_backend=broken, **kw), + ) + + with pytest.raises(CapabilityAdmissionError) as exc: + _start() # no validator: the production real-forward path + + assert exc.value.reason == REASON_NOT_PASSED + assert startup_env == [], "the tracker was called despite a failed validation" + assert not _FakeTorchNodeServer.started, "a failed node still opened an endpoint" + + +def test_a_report_for_a_different_model_cannot_be_reused_to_register(startup_env): + "A success for one model must not admit another — the shape of a replayed proof.\n\nTags: node, admission, startup" + with pytest.raises(CapabilityAdmissionError) as exc: + _start(capability_validator=capability_stub(model_id="other/model-1b")) + + assert exc.value.reason == REASON_MODEL_MISMATCH + assert startup_env == [] + + +def test_a_stale_report_cannot_be_reused_to_register(startup_env): + "An aged-out proof is refused before the node advertises itself.\n\nTags: node, admission, startup" + with pytest.raises(CapabilityAdmissionError) as exc: + _start(capability_validator=capability_stub(age_seconds=86_400)) + + assert exc.value.reason == REASON_STALE + assert startup_env == [] + + +def test_a_matching_passing_report_registers_and_travels_with_the_payload(startup_env): + "Registration carries the proof for exactly the model/shard/recipe it advertises.\n\nTags: node, admission, startup" + node = _start() # production validator against a working fake backend + node.stop() + + assert len(startup_env) == 1 + url, payload = startup_env[0] + assert url.endswith("/v1/nodes/register") + + report = payload["capability_report"] + assert report["status"] == "passed" + assert report["model"]["model_id"] == MODEL + assert (report["shard"]["start"], report["shard"]["end"]) == (0, 23) + assert report["recipe"]["recipe_id"] == DEFAULT_RECIPE_ID + assert report["backend"]["device"] == "cpu" + + +def test_the_served_backend_is_loaded_with_the_recipe_that_was_validated(startup_env): + "The recipe named in the report is the one the serving backend actually ran.\n\nTags: node, admission, startup" + node = _start(recipe_id="eager-attention") + node.stop() + + assert node.kwargs["recipe_params"] == {"attn_implementation": "eager"} + report = startup_env[0][1]["capability_report"] + assert report["recipe"]["recipe_id"] == "eager-attention" + + +def test_an_unknown_recipe_fails_before_any_weights_are_loaded(startup_env): + "A typo'd recipe id is caught at resolution, not after a multi-minute load.\n\nTags: node, admission, startup" + from meshnet_node.recipe_manifest import RecipeManifestError + + with pytest.raises(RecipeManifestError): + _start(recipe_id="does-not-exist") + + assert startup_env == [] + assert not _FakeTorchNodeServer.started diff --git a/tests/test_node_doctor.py b/tests/test_node_doctor.py new file mode 100644 index 0000000..dd2fd6c --- /dev/null +++ b/tests/test_node_doctor.py @@ -0,0 +1,644 @@ +"""NCA-002 tests for `meshnet-node doctor`. + +The unit tests inject a fake backend, so none of them download a model, import +Torch, or need a GPU. The one test that runs a real model is `integration`-marked +and takes its model identity from the environment — it has no model default, on +purpose: the doctor is model-agnostic and so is its test. +""" + +import base64 +import json +import os +import struct +from pathlib import Path + +import pytest + +from meshnet_node import doctor +from meshnet_node.capability import STATUS_FAILED, STATUS_PASSED, CapabilityReport +from meshnet_node.doctor import ( + CATEGORY_FORWARD_FAILED, + CATEGORY_INSUFFICIENT_MEMORY, + CATEGORY_INVALID_SHARD, + CATEGORY_MISSING_DEPENDENCY, + CATEGORY_NO_MODEL, + CATEGORY_UNSUPPORTED_RECIPE, + PROBE_TOKENS, + DoctorError, + DoctorSelection, + build_probe_input, + classify_failure, + probe_forward, + render_result, + resolve_selection, + run_doctor, + select_recipes, + write_reports, +) +from meshnet_node.model_backend import ( + InsufficientVRAMError, + MissingModelDependencyError, + UnsupportedRecipeParam, + validate_recipe_params, +) +from meshnet_node.recipe_manifest import parse_recipe_manifest + +# Deliberately not a model this project ships against: nothing here may special-case it. +FIXTURE_MODEL = "acme-labs/Widget-9000-Instruct" + +MANIFEST = parse_recipe_manifest( + { + "schema_version": 1, + "catalogue_version": "test-1", + "recipes": [ + {"id": "baseline", "version": "1", "backend_id": "torch-transformers"}, + { + "id": "stateless", + "version": "2", + "backend_id": "torch-transformers", + "params": {"use_cache": False}, + }, + ], + }, + source="", +) + + +class _Payload: + """Stands in for model_backend.TensorPayload.""" + + def __init__(self, body: bytes, shape: list[int]) -> None: + self.body = body + self.shape = shape + self.attention_mask_header = None + self.position_ids_header = None + + +class _TailToken: + """Stands in for model_backend.TailTokenResult.""" + + def __init__(self, token_id: int = 7) -> None: + self.token_id = token_id + self.text = "ok" + + +class _Device: + def __init__(self, type_: str = "cpu") -> None: + self.type = type_ + + +class _FakeBackend: + """A backend that loads but records exactly how it was driven.""" + + hidden_size = 8 + + def __init__( + self, + *, + is_head: bool = True, + is_tail: bool = False, + shard_start: int = 0, + shard_end: int = 3, + forward_error: Exception | None = None, + ) -> None: + self.model_id = FIXTURE_MODEL + self.is_head = is_head + self.is_tail = is_tail + self.shard_start = shard_start + self.shard_end = shard_end + self.device = _Device("cpu") + self.forward_error = forward_error + self.encoded_prompts: list[str] = [] + self.forwards: list[dict] = [] + + def encode_prompt(self, prompt: str): + if self.forward_error is not None: + raise self.forward_error + self.encoded_prompts.append(prompt) + return _Payload(b"\x00" * (PROBE_TOKENS * self.hidden_size * 2), + [1, PROBE_TOKENS, self.hidden_size]) + + def forward_bytes( + self, + body, + shape, + attention_mask_header, + position_ids_header, + start_layer=None, + **kwargs, + ): + if self.forward_error is not None: + raise self.forward_error + self.forwards.append( + { + "body_len": len(body), + "shape": shape, + "start_layer": start_layer, + "attention_mask_header": attention_mask_header, + "position_ids_header": position_ids_header, + } + ) + if self.is_tail: + return _TailToken() + return _Payload(body, shape) + + +def _loader(backend=None, *, error: Exception | None = None): + """A load_backend stub that records the (selection, recipe) pairs it saw.""" + calls: list[tuple[DoctorSelection, object]] = [] + + def load(selection, recipe): + calls.append((selection, recipe)) + if error is not None: + raise error + return backend if backend is not None else _FakeBackend() + + load.calls = calls # type: ignore[attr-defined] + return load + + +def _selection(**overrides) -> DoctorSelection: + kwargs = dict(model_id=FIXTURE_MODEL, shard_start=0, shard_end=3) + kwargs.update(overrides) + return DoctorSelection(**kwargs) + + +# --- selection resolves the same as startup --------------------------------- + + +def test_resolve_selection_uses_the_configured_repo_shard_and_quantization(): + selection = resolve_selection( + { + "model_hf_repo": FIXTURE_MODEL, + "model_name": "Widget-9000-Instruct", + "shard_start": 4, + "shard_end": 11, + "quantization": "bf16", # startup normalizes this to bfloat16 + "download_dir": "/models", + } + ) + + assert selection.model_id == FIXTURE_MODEL + assert (selection.shard_start, selection.shard_end) == (4, 11) + assert selection.quantization == "bfloat16" + assert selection.cache_dir == Path("/models") + + +def test_resolve_selection_defaults_to_the_whole_model_like_startup(): + """With no pinned shard, startup serves layers 0..n-1 — so doctor validates that.""" + seen: list[tuple[str, Path | None]] = [] + + def detect(model_id, cache_dir): + seen.append((model_id, cache_dir)) + return 24 + + selection = resolve_selection( + {"model_hf_repo": FIXTURE_MODEL}, detect_layers=detect + ) + + assert (selection.shard_start, selection.shard_end) == (0, 23) + assert seen == [(FIXTURE_MODEL, None)] + + +def test_resolve_selection_without_a_model_is_actionable(): + with pytest.raises(DoctorError) as exc: + resolve_selection({"model_hf_repo": "", "model_name": ""}) + + assert exc.value.category == CATEGORY_NO_MODEL + assert "--model" in exc.value.hint + + +def test_resolve_selection_rejects_an_inverted_shard_range(): + with pytest.raises(DoctorError) as exc: + resolve_selection( + {"model_hf_repo": FIXTURE_MODEL, "shard_start": 9, "shard_end": 2} + ) + + assert exc.value.category == CATEGORY_INVALID_SHARD + + +def test_resolve_selection_reports_an_unreadable_model_config(): + with pytest.raises(DoctorError) as exc: + resolve_selection( + {"model_hf_repo": FIXTURE_MODEL}, detect_layers=lambda *_: None + ) + + assert exc.value.category == doctor.CATEGORY_MODEL_UNAVAILABLE + assert "--shard-start" in str(exc.value) + + +# --- the bounded real forward ------------------------------------------------ + + +def test_a_pass_requires_a_real_forward_through_the_selected_shard(): + """Hardware being fine is not the bar: the shard itself has to execute.""" + backend = _FakeBackend(is_head=True) + + result = run_doctor( + _selection(), manifest=MANIFEST, load_backend=_loader(backend), now=lambda: 1000.0 + ) + + assert result.passed + assert backend.encoded_prompts == [doctor.PROBE_PROMPT] + report = result.reports[0] + assert report.status == STATUS_PASSED + assert report.model.model_id == FIXTURE_MODEL + assert (report.shard.start, report.shard.end) == (0, 3) + + +def test_a_backend_that_loads_but_cannot_forward_never_passes(): + """The regression this story exists for: a load is not a validation.""" + backend = _FakeBackend(forward_error=RuntimeError("kernel exploded")) + + result = run_doctor( + _selection(), manifest=MANIFEST, load_backend=_loader(backend), now=lambda: 1.0 + ) + + assert not result.passed + assert result.exit_code == 1 + report = result.reports[0] + assert report.status == STATUS_FAILED + assert result.results[0].category == CATEGORY_FORWARD_FAILED + assert any("kernel exploded" in d for d in report.diagnostics) + + +def test_a_mid_shard_is_probed_with_peer_shaped_hidden_states(): + backend = _FakeBackend(is_head=False, shard_start=4, shard_end=7) + + detail = probe_forward(backend) + + assert detail["probe"] == "hidden-states" + assert backend.encoded_prompts == [] + forward = backend.forwards[0] + assert forward["shape"] == [1, PROBE_TOKENS, backend.hidden_size] + # bfloat16 == 2 bytes per element, and the probe stays bounded to PROBE_TOKENS. + assert forward["body_len"] == PROBE_TOKENS * backend.hidden_size * 2 + assert forward["start_layer"] == 4 + + +def test_a_head_and_tail_shard_also_decodes_so_the_lm_head_is_covered(): + backend = _FakeBackend(is_head=True, is_tail=True, shard_end=5) + + detail = probe_forward(backend) + + assert detail["probe"] == "prompt+decode" + assert detail["output"] == "token" + # Re-entering above the last layer decodes without re-running any layer. + assert backend.forwards[0]["start_layer"] == 6 + + +def test_a_tail_shard_that_decodes_a_token_passes(): + backend = _FakeBackend(is_head=False, is_tail=True, shard_start=8, shard_end=11) + + detail = probe_forward(backend) + + assert detail == { + "probe": "hidden-states", + "tokens": PROBE_TOKENS, + "output": "token", + "token_id": 7, + } + + +def test_an_empty_forward_result_is_a_failure_not_a_pass(): + backend = _FakeBackend(is_head=False) + backend.forward_bytes = lambda *a, **k: _Payload(b"", []) # type: ignore[assignment] + + with pytest.raises(DoctorError) as exc: + probe_forward(backend) + + assert exc.value.category == CATEGORY_FORWARD_FAILED + + +def test_a_backend_with_no_hidden_size_cannot_be_probed(): + with pytest.raises(DoctorError) as exc: + build_probe_input(0) + + assert exc.value.category == CATEGORY_FORWARD_FAILED + + +def test_probe_headers_decode_as_int64_tensors(): + probe = build_probe_input(hidden_size=8, tokens=3) + + shape, encoded = probe.position_ids_header.split(":", 1) + raw = base64.b64decode(encoded) + assert shape == "1,3" + assert list(struct.unpack("<3q", raw)) == [0, 1, 2] + + +# --- recipes ----------------------------------------------------------------- + + +def test_the_default_run_validates_only_the_selected_recipe(): + """Onboarding must not pay to validate recipes the node was not asked to serve.""" + load = _loader() + + result = run_doctor(_selection(), manifest=MANIFEST, load_backend=load) + + assert [r.recipe.id for r in result.results] == ["baseline"] + assert len(load.calls) == 1 + + +def test_all_recipes_is_explicit_and_validates_every_recipe(): + load = _loader() + + result = run_doctor( + _selection(), manifest=MANIFEST, load_backend=load, all_recipes=True + ) + + assert [r.recipe.id for r in result.results] == ["baseline", "stateless"] + assert len(load.calls) == 2 + assert result.passed + + +def test_each_recipe_reaches_the_backend_that_runs_it(): + """A recipe that never reaches the loader was not really validated.""" + load = _loader() + + run_doctor(_selection(), manifest=MANIFEST, load_backend=load, all_recipes=True) + + params = [recipe.params for _, recipe in load.calls] + assert params == [{}, {"use_cache": False}] + + +def test_an_unknown_recipe_names_the_ones_that_exist(): + with pytest.raises(DoctorError) as exc: + select_recipes(MANIFEST, recipe_id="does-not-exist") + + assert exc.value.category == CATEGORY_UNSUPPORTED_RECIPE + assert "baseline" in str(exc.value) + + +def test_recipe_and_all_recipes_are_mutually_exclusive(): + with pytest.raises(DoctorError): + select_recipes(MANIFEST, recipe_id="baseline", all_recipes=True) + + +def test_a_recipe_the_backend_cannot_apply_is_a_failure_not_a_silent_pass(): + validate_recipe_params({"use_cache": False, "attn_implementation": "eager"}) + + with pytest.raises(UnsupportedRecipeParam) as exc: + validate_recipe_params({"sparkle_mode": True}) + + assert "sparkle_mode" in str(exc.value) + assert classify_failure(exc.value) == CATEGORY_UNSUPPORTED_RECIPE + + +def test_the_shipped_recipes_are_all_applicable_by_the_backend(): + """recipes.json and the backend's supported params must not drift apart.""" + from meshnet_node.recipe_manifest import load_recipe_manifest + + for recipe in load_recipe_manifest().recipes: + validate_recipe_params(recipe.params) + + +# --- failure reporting ------------------------------------------------------- + + +@pytest.mark.parametrize( + "exc, category", + [ + (MissingModelDependencyError("no torch"), CATEGORY_MISSING_DEPENDENCY), + (InsufficientVRAMError("too big"), CATEGORY_INSUFFICIENT_MEMORY), + (UnsupportedRecipeParam("nope"), CATEGORY_UNSUPPORTED_RECIPE), + (ValueError("shard_end 99 exceeds last layer index 23"), CATEGORY_INVALID_SHARD), + (FileNotFoundError("config.json"), doctor.CATEGORY_MODEL_UNAVAILABLE), + (RuntimeError("something else"), doctor.CATEGORY_LOAD_FAILED), + ], +) +def test_load_failures_are_classified_into_actionable_categories(exc, category): + result = run_doctor( + _selection(), manifest=MANIFEST, load_backend=_loader(error=exc) + ) + + assert not result.passed + item = result.results[0] + assert item.category == category + assert item.hint # every category tells the operator what to do next + assert item.report.status == STATUS_FAILED + + +def test_a_failure_report_carries_the_hint_and_no_traceback(): + result = run_doctor( + _selection(), + manifest=MANIFEST, + load_backend=_loader(error=InsufficientVRAMError("insufficient VRAM to load")), + ) + + diagnostics = " ".join(result.reports[0].diagnostics) + assert "insufficient VRAM to load" in diagnostics + assert "--shard-start" in diagnostics # the actionable next step + assert "Traceback" not in diagnostics + assert ".py" not in diagnostics # no file/line noise from a stack + + +def test_a_failure_report_still_identifies_what_was_being_validated(): + """NCA-003 refuses to register without a matching report — including a failed one.""" + result = run_doctor( + _selection(shard_start=4, shard_end=9, quantization="int8"), + manifest=MANIFEST, + load_backend=_loader(error=RuntimeError("boom")), + now=lambda: 4242.0, + ) + + report = result.reports[0] + assert report.identity_key() == ( + FIXTURE_MODEL, 4, 9, "baseline", "1", "torch-transformers", "unknown", + ) + assert report.validated_at == 4242.0 + assert report.recipe.catalogue_version == "test-1" + + +def test_the_report_records_the_device_the_forward_actually_ran_on(): + result = run_doctor( + _selection(), manifest=MANIFEST, load_backend=_loader(_FakeBackend()) + ) + + assert result.reports[0].backend.device == "cpu" + assert result.reports[0].backend.backend_id == "torch-transformers" + + +def test_reports_round_trip_through_the_written_json(tmp_path): + result = run_doctor( + _selection(), manifest=MANIFEST, load_backend=_loader(), all_recipes=True + ) + + path = write_reports(result.reports, tmp_path / "nested" / "capability.json") + payload = json.loads(path.read_text()) + + assert [CapabilityReport.from_dict(d).recipe.recipe_id for d in payload] == [ + "baseline", + "stateless", + ] + + +def test_a_single_report_is_written_as_one_object(): + """One selected recipe writes one report — the shape NCA-003 will read.""" + import tempfile + + result = run_doctor(_selection(), manifest=MANIFEST, load_backend=_loader()) + with tempfile.TemporaryDirectory() as tmp: + path = write_reports(result.reports, Path(tmp) / "capability.json") + report = CapabilityReport.from_json(path.read_text()) + + assert report.passed + + +def test_the_summary_tells_a_failing_operator_what_to_fix(): + result = run_doctor( + _selection(), + manifest=MANIFEST, + load_backend=_loader(error=MissingModelDependencyError("torch is not installed")), + ) + + text = render_result(result, report_path=Path("/tmp/capability.json")) + + assert "FAIL" in text + assert CATEGORY_MISSING_DEPENDENCY in text + assert "torch is not installed" in text + assert "/tmp/capability.json" in text + assert "Traceback" not in text + + +def test_the_summary_names_the_shard_that_passed(): + result = run_doctor(_selection(), manifest=MANIFEST, load_backend=_loader()) + + text = render_result(result) + + assert "PASS" in text + assert FIXTURE_MODEL in text + assert "layers 0–3" in text + + +# --- the CLI wiring ---------------------------------------------------------- + + +def _run_cli(monkeypatch, argv, backend=None, error=None): + """Drive `meshnet-node doctor` end to end with an injected backend.""" + import sys + + from meshnet_node import cli, config + + monkeypatch.setattr( + config, "load_config", lambda *a, **k: { + "model_hf_repo": FIXTURE_MODEL, + "shard_start": 0, + "shard_end": 3, + "quantization": "auto", + } + ) + monkeypatch.setattr( + doctor, "default_load_backend", _loader(backend, error=error) + ) + monkeypatch.setattr(doctor, "load_recipe_manifest", lambda *a, **k: MANIFEST) + monkeypatch.setattr(sys, "argv", ["meshnet-node", *argv]) + + with pytest.raises(SystemExit) as exit_info: + cli.main() + return exit_info.value.code + + +def test_cli_doctor_exits_zero_and_writes_a_passing_report(monkeypatch, capsys, tmp_path): + report = tmp_path / "capability.json" + + code = _run_cli(monkeypatch, ["doctor", "--report", str(report)], backend=_FakeBackend()) + + assert code == 0 + assert capsys.readouterr().out.count("PASS") == 1 + assert CapabilityReport.from_json(report.read_text()).passed + + +def test_cli_doctor_exits_non_zero_and_writes_the_failed_report(monkeypatch, capsys, tmp_path): + report = tmp_path / "capability.json" + + code = _run_cli( + monkeypatch, + ["doctor", "--report", str(report)], + error=InsufficientVRAMError("insufficient VRAM to load 24 layers"), + ) + + out = capsys.readouterr().out + assert code == 1 + assert "FAIL" in out + assert CATEGORY_INSUFFICIENT_MEMORY in out + assert "Traceback" not in out # no raw traceback by default + assert CapabilityReport.from_json(report.read_text()).status == STATUS_FAILED + + +def test_cli_doctor_all_recipes_is_opt_in(monkeypatch, capsys, tmp_path): + report = tmp_path / "capability.json" + + code = _run_cli( + monkeypatch, + ["doctor", "--all-recipes", "--report", str(report)], + backend=_FakeBackend(), + ) + + assert code == 0 + assert capsys.readouterr().out.count("PASS") == 2 + assert len(json.loads(report.read_text())) == 2 + + +def test_cli_doctor_json_prints_the_capability_report(monkeypatch, capsys, tmp_path): + code = _run_cli( + monkeypatch, + ["doctor", "--json", "--report", str(tmp_path / "c.json")], + backend=_FakeBackend(), + ) + + payload = json.loads(capsys.readouterr().out) + assert code == 0 + assert payload[0]["model"]["model_id"] == FIXTURE_MODEL + + +def test_cli_doctor_flags_select_what_is_validated(monkeypatch, capsys, tmp_path): + """`doctor --shard-start/--shard-end` validates the shard startup would load.""" + report = tmp_path / "capability.json" + + code = _run_cli( + monkeypatch, + ["doctor", "--shard-start", "2", "--shard-end", "5", "--report", str(report)], + backend=_FakeBackend(), + ) + + written = CapabilityReport.from_json(report.read_text()) + assert code == 0 + assert (written.shard.start, written.shard.end) == (2, 5) + + +# --- the real-model smoke test ---------------------------------------------- + +# Model identity comes from the environment; there is no default, so this test +# never smuggles a vendor-specific assumption into the suite. +DOCTOR_MODEL = os.environ.get("MESHNET_DOCTOR_MODEL") +DOCTOR_SHARD_START = int(os.environ.get("MESHNET_DOCTOR_SHARD_START", "0")) +DOCTOR_SHARD_END = os.environ.get("MESHNET_DOCTOR_SHARD_END") + + +@pytest.mark.integration +@pytest.mark.skipif( + not DOCTOR_MODEL, + reason="set MESHNET_DOCTOR_MODEL (and optionally MESHNET_DOCTOR_SHARD_START/END) to run", +) +def test_doctor_smoke_runs_a_real_forward_on_a_real_model(tmp_path): + cfg = { + "model_hf_repo": DOCTOR_MODEL, + "quantization": os.environ.get("MESHNET_DOCTOR_QUANTIZATION", "auto"), + "download_dir": os.environ.get("MESHNET_DOWNLOAD_DIR") or None, + "shard_start": DOCTOR_SHARD_START, + "shard_end": int(DOCTOR_SHARD_END) if DOCTOR_SHARD_END else None, + "force_cpu": os.environ.get("MESHNET_DOCTOR_CPU") == "1", + } + selection = resolve_selection(cfg) + + result = run_doctor(selection) + + report = result.reports[0] + assert result.passed, f"doctor failed: {report.diagnostics}" + assert report.status == STATUS_PASSED + assert report.model.model_id == DOCTOR_MODEL + assert report.duration_ms > 0 + assert report.model.config_fingerprint.startswith("sha256:") + + path = write_reports(result.reports, tmp_path / "capability.json") + assert CapabilityReport.from_json(path.read_text()).passed diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index 1feea1f..04b2c0e 100644 --- a/tests/test_node_startup.py +++ b/tests/test_node_startup.py @@ -26,6 +26,10 @@ from meshnet_node.startup import ( _tracker_http_error_message, run_startup, ) +# Startup admits a node only on a capability report from a real forward, which a +# fake backend cannot perform. These tests say so explicitly rather than bypassing +# admission; the fail-closed path itself is covered in tests/test_node_admission.py. +from meshnet_node.testing import assume_capability from meshnet_node.wallet import _b58encode, load_or_create_wallet from meshnet_contracts import LocalSolanaContracts from meshnet_tracker.server import TrackerServer @@ -333,6 +337,7 @@ def test_benchmark_throughput_is_registered_in_payload(monkeypatch, tmp_path): wallet_path=tmp_path / "wallet.json", torch_threads=8, torch_interop_threads=1, + capability_validator=assume_capability, ) node.stop() @@ -394,6 +399,7 @@ def test_real_model_startup_passes_download_dir_and_kimi_metadata(monkeypatch, t shard_end=60, wallet_path=tmp_path / "wallet.json", cache_dir=cache_dir, + capability_validator=assume_capability, ) node.stop() @@ -448,6 +454,7 @@ def test_cuda_benchmark_failure_is_registered_for_inventory_only_gpu(monkeypatch shard_start=0, shard_end=23, wallet_path=tmp_path / "wallet.json", + capability_validator=assume_capability, ) node.stop() @@ -1164,6 +1171,7 @@ def test_public_https_tracker_infers_relay_when_network_map_omits_relay_url( model_id="Qwen/Qwen2.5-0.5B-Instruct", advertise_host="172.29.104.23", wallet_path=tmp_path / "wallet.json", + capability_validator=assume_capability, ) try: pass @@ -1216,6 +1224,7 @@ def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, ca vram_mb_override=6144, max_loaded_shards=2, wallet_path=tmp_path / "wallet.json", + capability_validator=assume_capability, ) assert node.backend.total_layers == 24 @@ -1275,6 +1284,7 @@ def test_real_model_startup_autodetects_cpu_memory_budget_and_logs_shard_budget( shard_start=0, shard_end=23, wallet_path=tmp_path / "wallet.json", + capability_validator=assume_capability, ) try: pass @@ -1359,6 +1369,7 @@ def test_public_tracker_model_node_registers_relay_metadata_from_tracker_url_onl tracker_url=tracker_url, model_id="Qwen/Qwen2.5-0.5B-Instruct", wallet_path=tmp_path / "wallet.json", + capability_validator=assume_capability, ) try: network_map = _get_json(f"{tracker_url}/v1/network/map") @@ -1444,6 +1455,7 @@ def test_public_tracker_relay_suppresses_virtual_ip_warning( model_id="Qwen/Qwen2.5-0.5B-Instruct", advertise_host="172.29.104.23", wallet_path=tmp_path / "wallet.json", + capability_validator=assume_capability, ) try: network_map = _get_json(f"{tracker_url}/v1/network/map") @@ -1523,6 +1535,7 @@ def test_later_node_auto_joins_existing_public_hf_model_with_only_tracker_url( tracker_url=tracker_url, advertise_host="203.0.113.21", wallet_path=tmp_path / "wallet.json", + capability_validator=assume_capability, ) try: route_resp = _get_json( @@ -1607,6 +1620,7 @@ def test_later_node_auto_joins_redundant_copy_when_model_is_fully_covered( tracker_url=tracker_url, advertise_host="203.0.113.32", wallet_path=tmp_path / "wallet.json", + capability_validator=assume_capability, ) try: assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct" @@ -1637,6 +1651,7 @@ def test_full_startup_sequence(tmp_path): model="stub-model", wallet_path=wallet_path, cache_dir=cache_dir, + capability_validator=assume_capability, ) try: # Wallet was created on disk @@ -1699,6 +1714,7 @@ def test_preset_model_startup_starts_heartbeat(tmp_path, monkeypatch): model="stub-model", wallet_path=tmp_path / "wallet.json", cache_dir=tmp_path / "shards", + capability_validator=assume_capability, ) try: assert len(heartbeat_calls) == 1 @@ -1739,6 +1755,7 @@ def test_preset_model_startup_honors_pinned_shard_range(tmp_path, monkeypatch): shard_end=5, wallet_path=tmp_path / "wallet.json", cache_dir=tmp_path / "shards", + capability_validator=assume_capability, ) try: assert len(heartbeat_calls) == 1 @@ -1782,6 +1799,7 @@ def test_preset_startup_rejects_pinned_shard_above_memory_budget(tmp_path, monke shard_end=39, wallet_path=tmp_path / "wallet.json", cache_dir=tmp_path / "shards", + capability_validator=assume_capability, ) finally: tracker.stop() @@ -1835,6 +1853,7 @@ def test_network_auto_join_clips_oversized_cpu_assignment(tmp_path, monkeypatch, tracker_url="http://127.0.0.1:8080", wallet_path=tmp_path / "wallet.json", tracker_source_disabled=True, + capability_validator=assume_capability, ) try: assert torch_calls[0]["shard_start"] == 0 @@ -1895,6 +1914,7 @@ def test_preset_model_with_hf_repo_loads_torch_backend(tmp_path, monkeypatch, ca model="tiny-llama", wallet_path=tmp_path / "wallet.json", cache_dir=tmp_path / "node-shards", + capability_validator=assume_capability, ) try: assert len(torch_calls) == 1 @@ -1972,6 +1992,7 @@ def test_torch_startup_retries_registration_when_tracker_unreachable( tracker_url=tracker_url, model_id="Qwen/Qwen2.5-0.5B-Instruct", wallet_path=tmp_path / "wallet.json", + capability_validator=assume_capability, ) try: assert register_calls["count"] == 1 @@ -2052,6 +2073,7 @@ def test_real_model_startup_registers_downloaded_inventory_without_checksum( model="tiny-llama", wallet_path=tmp_path / "wallet.json", cache_dir=tmp_path / "node-shards", + capability_validator=assume_capability, ) try: assert len(hf_calls) == 1 @@ -2342,6 +2364,7 @@ def test_startup_cpu_fallback(tmp_path, monkeypatch): model="stub-model", wallet_path=tmp_path / "wallet.json", cache_dir=tmp_path / "shards", + capability_validator=assume_capability, ) try: # Node is running even on CPU diff --git a/tests/test_prefill_backpressure.py b/tests/test_prefill_backpressure.py new file mode 100644 index 0000000..74f5cba --- /dev/null +++ b/tests/test_prefill_backpressure.py @@ -0,0 +1,93 @@ +"""Coverage for bounded, ordered prefill transfer.""" + +from threading import Event + +import pytest + +from meshnet_gateway.prefill_backpressure import ( + BoundedPrefillSender, + DEFAULT_PREFILL_CHUNK_TOKENS, + DEFAULT_PREFILL_MAX_IN_FLIGHT, + PrefillTransferLimits, +) +from meshnet_gateway.server import _BinaryActivation, _post_binary_forward + + +def test_limits_have_safe_defaults_and_keep_legacy_chunk_setting(monkeypatch): + monkeypatch.delenv("MESHNET_PREFILL_CHUNK_TOKENS", raising=False) + monkeypatch.delenv("MESHNET_PREFILL_MAX_IN_FLIGHT", raising=False) + monkeypatch.setenv("MESHNET_CHUNK_TOKENS", "64") + + limits = PrefillTransferLimits.from_env() + + assert limits.chunk_tokens == 64 + assert limits.max_in_flight == DEFAULT_PREFILL_MAX_IN_FLIGHT + assert PrefillTransferLimits().chunk_tokens == DEFAULT_PREFILL_CHUNK_TOKENS + assert limits.max_buffered_bytes == limits.max_chunk_bytes + + +def test_slow_consumer_applies_backpressure_preserves_order_and_bounds_bytes(): + sender = BoundedPrefillSender( + PrefillTransferLimits(chunk_tokens=2, max_in_flight=4, max_chunk_bytes=8) + ) + sent: list[int] = [] + produced: list[int] = [] + + def chunks(): + for index in range(3): + produced.append(index) + yield bytes([index]) * 8 + + def forward(chunk: bytes) -> int: + sent.append(chunk[0]) + # The next item cannot have been constructed while this consumer waits. + assert produced[-1] == chunk[0] + assert len(produced) == chunk[0] + 1 + assert sender.buffered_bytes == 8 + assert sender.in_flight == 1 + return chunk[0] + + assert sender.send(chunks(), body_size=len, forward=forward) == [0, 1, 2] + assert sent == [0, 1, 2] + assert sender.peak_buffered_bytes <= sender.limits.max_buffered_bytes + assert sender.peak_in_flight == 1 + assert sender.buffered_bytes == sender.in_flight == 0 + + +def test_failure_and_cancellation_release_owned_buffers(): + sender = BoundedPrefillSender(PrefillTransferLimits()) + with pytest.raises(RuntimeError, match="route lost"): + sender.send([b"one", b"two"], body_size=len, forward=lambda _: (_ for _ in ()).throw(RuntimeError("route lost"))) + assert sender.closed + assert sender.buffered_bytes == sender.in_flight == 0 + + cancelled = Event() + cancelled.set() + sender = BoundedPrefillSender(PrefillTransferLimits()) + assert sender.send([b"never"], body_size=len, forward=lambda _: None, cancelled=cancelled) == [] + assert sender.buffered_bytes == sender.in_flight == 0 + + +def test_legacy_single_chunk_peer_response_uses_outgoing_metadata(monkeypatch): + class Response: + headers = {"Content-Type": "application/octet-stream"} + + def read(self): + return b"\0" * (1 * 1 * 64 * 2) + + def __enter__(self): + return self + + def __exit__(self, *_): + return False + + monkeypatch.setattr("urllib.request.urlopen", lambda *_args, **_kwargs: Response()) + activation = _BinaryActivation( + body=b"\0" * (1 * 1 * 64 * 2), shape=[1, 1, 64], dtype="bfloat16", + session="legacy-session", chunk_index=0, chunk_total=1, encoding=None, headers={}, + ) + + response = _post_binary_forward("http://legacy/forward", activation, hop_index=0, timeout=1.0) + + assert response.session == activation.session + assert response.chunk_index == response.chunk_total - 1 == 0 diff --git a/tests/test_real_model_backend.py b/tests/test_real_model_backend.py index 17ac434..92b0901 100644 --- a/tests/test_real_model_backend.py +++ b/tests/test_real_model_backend.py @@ -24,6 +24,7 @@ from meshnet_node.model_backend import ( _should_partial_materialize_shard, _decoder_attention_mask, _int_tensor_header, + _tensor_from_bfloat16_bytes, _torch_cuda_is_executable, build_quantization_config, validate_quantization, @@ -538,6 +539,20 @@ def test_int_tensor_header_serializes_torch_tensors(): assert header.startswith("1,3:") +def test_bfloat16_wire_decode_views_owned_bytes_without_float32_round_trip(): + """Activation decode stays bf16 and does not clone bytes into bytearray. + + Tags: model, performance, wire + """ + torch = pytest.importorskip("torch") + body = torch.tensor([[1, 2]], dtype=torch.bfloat16).view(torch.uint8).numpy().tobytes() + + decoded = _tensor_from_bfloat16_bytes(body, [1, 2], torch) + + assert decoded.dtype == torch.bfloat16 + assert decoded.tolist() == [[1.0, 2.0]] + + def test_decoder_attention_mask_is_causal_float_mask(): "Decoder attention mask is causal float mask\n\nTags: model, node, real-inference" torch = pytest.importorskip("torch") diff --git a/tests/test_route_session_benchmark.py b/tests/test_route_session_benchmark.py new file mode 100644 index 0000000..c1a9294 --- /dev/null +++ b/tests/test_route_session_benchmark.py @@ -0,0 +1,126 @@ +"""DIP-001 deterministic Route Session benchmark coverage.""" + +import json +from unittest.mock import MagicMock, patch + +from meshnet_node.route_session_benchmark import ( + BenchmarkScenario, + PerformanceThresholds, + assert_benchmark, + assert_performance_gate, + format_summary, + main, + run_benchmark_matrix, + run_real_model_lan_benchmark, + run_route_session_benchmark, +) + + +def test_matrix_reports_direct_relay_prefill_decode_and_machine_readable_metrics(): + """The baseline reports every required scenario and transport metric. + + Tags: performance, routing + """ + report = run_benchmark_matrix() + assert {(run["mode"], run["cache_mode"]) for run in report["runs"]} == { + ("direct", "cached"), ("direct", "stateless"), + ("relay", "cached"), ("relay", "stateless"), + } + for run in report["runs"]: + assert set(run["phases"]) == {"prefill", "decode"} + assert "head->tail" in run["seams"] + assert {"p50_latency_ms", "p95_latency_ms", "payload_bytes", "compression_ratio", + "connection_attempts", "p95_queue_wait_ms"} <= set(run["phases"]["decode"]) + sample = run["samples"][0] + assert sample["framing_ms"] > 0 + assert sample["metadata_ms"] > 0 + assert sample["copy_allocation_ms"] > 0 + assert sample["copy_allocation_bytes"] >= sample["payload_bytes"] + assert len(run["samples"]) == 1 + len(run["output_tokens"]) + assert {"tokens_per_sec", "bytes_per_token", "compression_cpu_ms", "peak_buffered_bytes"} <= set(run["phases"]["decode"]) + + +def test_cached_sessions_reuse_one_connection_and_preserve_stub_tokens(): + """Cached Route Sessions keep one direct or relay connection per seam. + + Tags: performance, relay + """ + scenario = BenchmarkScenario(output_tokens=(" one", " two", " three")) + for mode in ("direct", "relay"): + run = run_route_session_benchmark(mode, "cached", scenario) + assert_benchmark(run, expected_tokens=scenario.output_tokens, expected_connection_attempts=1) + if mode == "relay": + assert all(sample.queue_wait_ms > 0 for sample in run.samples) + + +def test_stateless_baselines_make_each_activation_a_connection_attempt(): + """Stateless comparison mode does not accidentally inherit Route Session state. + + Tags: performance, routing + """ + scenario = BenchmarkScenario(output_tokens=(" one", " two")) + for mode in ("direct", "relay"): + run = run_route_session_benchmark(mode, "stateless", scenario) + assert_benchmark(run, expected_tokens=scenario.output_tokens, expected_connection_attempts=3) + + +def test_cli_writes_json_artifact_and_human_summary(tmp_path, capsys): + """The CLI emits both CI-ready JSON and an operator-readable summary. + + Tags: performance + """ + output = tmp_path / "route-session-benchmark.json" + assert main(["--json-out", str(output)]) == 0 + report = json.loads(output.read_text()) + assert report["schema_version"] == 1 + assert "Route Session benchmark" in capsys.readouterr().out + assert "relay" in format_summary(report) + + +def test_performance_gate_checks_comparison_identity_session_and_cleanup(): + """CI gate accepts the fixed matrix and rejects a meaningful slowdown. + + Tags: performance, routing + """ + report = run_benchmark_matrix() + assert_performance_gate(report) + report["runs"][0]["phases"]["decode"]["tokens_per_sec"] = 0.1 + try: + assert_performance_gate(report, thresholds=PerformanceThresholds()) + except AssertionError as exc: + assert "throughput regressed" in str(exc) + else: # pragma: no cover - makes the intended threshold explicit + raise AssertionError("gate did not catch the throughput regression") + + +def test_performance_gate_rejects_session_or_cleanup_leaks(): + """Exact resource/session invariants are not subject to variance tolerance. + + Tags: performance, routing + """ + report = run_benchmark_matrix() + report["runs"][0]["samples"][1]["session_id"] = "wrong-session" + try: + assert_performance_gate(report) + except AssertionError as exc: + assert "Route Session changed" in str(exc) + else: # pragma: no cover + raise AssertionError("gate did not catch session instability") + + +def test_real_model_lan_capture_uses_the_shared_report_schema(): + """The opt-in LAN command is client-measurable and needs no real model in CI. + + Tags: performance + """ + response = MagicMock() + response.read.return_value = json.dumps({"choices": [{"message": {"content": "amber birch"}}]}).encode() + response.headers.get.return_value = "lan-session" + response.__enter__.return_value = response + with patch("meshnet_node.route_session_benchmark.urllib.request.urlopen", return_value=response): + report = run_real_model_lan_benchmark("http://lan-node:7000", model="test-model") + run = report["runs"][0] + assert report["source"] == "real-model-lan-client" + assert run["session_id"] == "lan-session" + assert run["phases"]["decode"]["tokens_per_sec"] > 0 + assert run["cleanup"]["open_connections"] == 0 diff --git a/tests/test_seam_telemetry.py b/tests/test_seam_telemetry.py new file mode 100644 index 0000000..385d2eb --- /dev/null +++ b/tests/test_seam_telemetry.py @@ -0,0 +1,78 @@ +"""Unit coverage for bounded Activation Seam telemetry.""" + +from meshnet_node.seam_telemetry import GenerationTelemetry + + +def test_seam_telemetry_aggregates_bytes_latency_and_correlates_ids(): + telemetry = GenerationTelemetry("route-session-1", report_every=2, report_interval=100, now=0.0) + + assert telemetry.record_seam( + activation_id="activation-1", phase="prefill", hop=0, node="node-a", + latency_seconds=0.012, wire_bytes=120, response_bytes=240, + connection_reused=False, now=0.1, + ) + telemetry.mark_reported(now=0.1) + assert telemetry.record_seam( + activation_id="activation-2", phase="prefill", hop=0, node="node-a", + latency_seconds=0.008, wire_bytes=80, response_bytes=160, + connection_reused=True, now=0.2, + ) + telemetry.note_tokens(4) + + snapshot = telemetry.snapshot(now=2.0) + assert snapshot["session_id"] == "route-session-1" + assert snapshot["tokens_per_sec"] == 2.0 + assert snapshot["seams"] == [{ + "phase": "prefill", "hop": 0, "node": "node-a", "activations": 2, + "latency_ms": 20.0, "avg_latency_ms": 10.0, "wire_bytes": 200, + "response_bytes": 400, "connection_reuse": 1, + "compression_input_bytes": 0, "compression_output_bytes": 0, "compression_ms": 0.0, + "decompression_input_bytes": 0, "decompression_output_bytes": 0, "decompression_ms": 0.0, + "last_activation_id": "activation-2", + }] + + +def test_seam_telemetry_includes_compression_work_and_byte_counters(): + telemetry = GenerationTelemetry("route-session-compression", now=0.0) + telemetry.record_compression( + phase="prefill", hop=0, node="node-a", input_bytes=1000, output_bytes=200, + elapsed_seconds=0.003, + ) + telemetry.record_compression( + phase="prefill", hop=0, node="node-a", input_bytes=200, output_bytes=1000, + elapsed_seconds=0.001, decompression=True, + ) + seam = telemetry.snapshot(now=1.0)["seams"][0] + assert seam["compression_input_bytes"] == 1000 + assert seam["compression_output_bytes"] == 200 + assert seam["compression_ms"] == 3.0 + assert seam["decompression_input_bytes"] == 200 + assert seam["decompression_output_bytes"] == 1000 + assert seam["decompression_ms"] == 1.0 + + +def test_seam_telemetry_reports_on_bounded_cadence_and_cleans_up(): + telemetry = GenerationTelemetry("route-session-2", report_every=3, report_interval=5, now=0.0) + + for count in range(1, 4): + due = telemetry.record_seam( + activation_id=f"activation-{count}", phase="decode", hop=1, node="node-b", + latency_seconds=0.001, wire_bytes=10, response_bytes=20, + connection_reused=True, now=float(count), + ) + if due: + telemetry.mark_reported(now=float(count)) + assert not telemetry.report_due + assert telemetry.record_seam( + activation_id="activation-4", phase="decode", hop=1, node="node-b", + latency_seconds=0.001, wire_bytes=10, response_bytes=20, + connection_reused=True, now=9.0, + ) + + telemetry.close() + assert telemetry.snapshot(now=10.0)["seams"] == [] + assert not telemetry.record_seam( + activation_id="late", phase="decode", hop=1, node="node-b", + latency_seconds=0.001, wire_bytes=10, response_bytes=20, + connection_reused=True, now=10.0, + ) diff --git a/tests/test_tracker_capability_admission.py b/tests/test_tracker_capability_admission.py new file mode 100644 index 0000000..9e25421 --- /dev/null +++ b/tests/test_tracker_capability_admission.py @@ -0,0 +1,525 @@ +"""NCA-004 tests: the tracker records a node's capability proof and routes only +to nodes whose proof admits what they advertise (ADR-0023). + +Model ids here are arbitrary and made up on purpose: nothing in the admission or +routing path may branch on a vendor, model or kernel name. +""" + +import json +import time +import urllib.error +import urllib.request + +import pytest + +from meshnet_tracker.capability import ( + MIN_CATALOGUE_VERSION, + POLICY_COMPAT, + POLICY_ENFORCE, + STATE_ABSENT, + STATE_ADMITTED, + STATE_CATALOGUE_INCOMPATIBLE, + STATE_FAILED, + STATE_INVALID, + STATE_MODEL_MISMATCH, + STATE_RECIPE_MISMATCH, + STATE_SHARD_MISMATCH, + STATE_STALE, + evaluate_report, + policy_from_env, +) +from meshnet_tracker.server import ( + TrackerServer, + _NodeEntry, + _capability_routable, + _node_admission, + _select_route, +) + +MODEL = "arbitrary-labs/oracle-9b" +SHORT = "oracle-9b" +LAYERS = 32 + + +def _post_json(url: str, payload: dict) -> dict: + data = json.dumps(payload).encode() + req = urllib.request.Request( + url, data=data, headers={"Content-Type": "application/json"}, method="POST" + ) + with urllib.request.urlopen(req) as r: + return json.loads(r.read()) + + +def _get_json(url: str) -> dict: + with urllib.request.urlopen(url) as r: + return json.loads(r.read()) + + +def _report( + *, + model_id: str = MODEL, + start: int = 0, + end: int = 15, + status: str = "passed", + validated_at: float | None = None, + recipe_id: str = "baseline", + recipe_version: str = "1", + catalogue_version: str = MIN_CATALOGUE_VERSION, + schema_version: int = 1, + device: str = "cpu", + diagnostics: list | None = None, +) -> dict: + """A capability report shaped exactly as `meshnet_node.capability` emits it.""" + return { + "schema_version": schema_version, + "model": {"model_id": model_id, "revision": None, "config_fingerprint": None}, + "shard": {"start": start, "end": end}, + "recipe": { + "recipe_id": recipe_id, + "recipe_version": recipe_version, + "catalogue_version": catalogue_version, + }, + "backend": { + "backend_id": "torch-transformers", + "device": device, + "device_name": None, + "quantization": "bfloat16", + "runtime": {}, + }, + "status": status, + "validated_at": time.time() if validated_at is None else validated_at, + "duration_ms": 42, + "diagnostics": list(diagnostics or []), + } + + +def _registration( + port: int, + *, + start: int = 0, + end: int = 15, + report: dict | None = "default", # type: ignore[assignment] + recipe_id: str | None = "baseline", + recipe_version: str | None = "1", + benchmark_tokens_per_sec: float = 10.0, +) -> dict: + payload: dict = { + "endpoint": f"http://127.0.0.1:{port}", + "model": SHORT, + "hf_repo": MODEL, + "num_layers": LAYERS, + "shard_start": start, + "shard_end": end, + "hardware_profile": {}, + "score": 1.0, + "tracker_mode": start == 0, + "benchmark_tokens_per_sec": benchmark_tokens_per_sec, + } + if report == "default": + report = _report(start=start, end=end) + if report is not None: + payload["capability_report"] = report + if recipe_id is not None: + payload["recipe_id"] = recipe_id + if recipe_version is not None: + payload["recipe_version"] = recipe_version + return payload + + +def _always(_reported: str) -> bool: + return True + + +def _evaluate(report, **kwargs) -> object: + kwargs.setdefault("model_matches", _always) + kwargs.setdefault("advertised_model", MODEL) + kwargs.setdefault("shard_start", 0) + kwargs.setdefault("shard_end", 15) + return evaluate_report(report, **kwargs) + + +# --------------------------------------------------------------- report verdicts + + +def test_a_passing_report_that_covers_the_registration_is_admitted(): + "A proof matching the advertised model and shard admits the node.\n\nTags: routing, tracker" + state = _evaluate(_report()) + assert state.state == STATE_ADMITTED + assert state.proven + assert state.model_id == MODEL + assert (state.shard_start, state.shard_end) == (0, 15) + + +def test_a_missing_report_is_absent_not_admitted(): + "No proof is recorded as `absent` — never silently treated as proven.\n\nTags: routing, tracker" + state = _evaluate(None) + assert state.state == STATE_ABSENT + assert not state.proven + + +def test_a_failed_report_is_recorded_as_failed(): + "A node that failed its own validation is not routable.\n\nTags: routing, tracker" + state = _evaluate(_report(status="failed", diagnostics=["out of memory on device"])) + assert state.state == STATE_FAILED + assert not state.proven + assert "out of memory" in state.detail + + +def test_a_report_for_a_different_model_is_a_model_mismatch(): + "A proof for another artifact proves nothing about this one.\n\nTags: routing, tracker" + state = evaluate_report( + _report(model_id="other-org/unrelated-3b"), + model_matches=lambda reported: reported == MODEL, + advertised_model=MODEL, + shard_start=0, + shard_end=15, + ) + assert state.state == STATE_MODEL_MISMATCH + + +def test_a_report_for_a_different_shard_is_a_shard_mismatch(): + "A proof for layers 0–15 does not admit a node advertising 16–31.\n\nTags: routing, tracker" + state = _evaluate(_report(start=0, end=15), shard_start=16, shard_end=31) + assert state.state == STATE_SHARD_MISMATCH + + +def test_a_report_for_a_different_recipe_than_the_node_declares_is_a_recipe_mismatch(): + "The proof must be for the recipe the node says it serves with.\n\nTags: routing, tracker" + state = _evaluate(_report(recipe_id="eager-attention"), declared_recipe_id="baseline") + assert state.state == STATE_RECIPE_MISMATCH + + versioned = _evaluate( + _report(recipe_version="1"), + declared_recipe_id="baseline", + declared_recipe_version="2", + ) + assert versioned.state == STATE_RECIPE_MISMATCH + + +def test_an_older_recipe_catalogue_is_incompatible(): + "Recipe ids from a catalogue older than the tracker's minimum cannot be matched.\n\nTags: routing, tracker" + state = _evaluate(_report(catalogue_version="2025.01.1")) + assert state.state == STATE_CATALOGUE_INCOMPATIBLE + assert MIN_CATALOGUE_VERSION in state.detail + + newer = _evaluate(_report(catalogue_version="2099.12.9")) + assert newer.state == STATE_ADMITTED + + +def test_an_unparseable_catalogue_version_is_incompatible(): + "A catalogue version that cannot be compared cannot be shown to be new enough.\n\nTags: routing, tracker" + assert _evaluate(_report(catalogue_version="rolling")).state == STATE_CATALOGUE_INCOMPATIBLE + + +def test_a_stale_report_is_not_admitted(): + "A proof older than the freshness bound must be re-validated before routing.\n\nTags: routing, tracker" + state = _evaluate(_report(validated_at=time.time() - 3600), max_age_seconds=900.0) + assert state.state == STATE_STALE + + +def test_a_future_dated_report_is_not_admitted(): + "A proof from the future is a broken clock, not a fresh proof.\n\nTags: routing, tracker" + state = _evaluate(_report(validated_at=time.time() + 3600)) + assert state.state == STATE_STALE + assert "clock" in state.detail + + +def test_a_report_from_an_unknown_schema_version_is_invalid(): + "The tracker refuses to interpret a report layout it does not read.\n\nTags: routing, tracker" + assert _evaluate(_report(schema_version=99)).state == STATE_INVALID + + +@pytest.mark.parametrize( + "payload", + [ + "not-an-object", + {}, + {"schema_version": 1}, + {**_report(), "model": {"model_id": ""}}, + {**_report(), "shard": {"start": -1, "end": 3}}, + {**_report(), "validated_at": "yesterday"}, + {**_report(), "status": None}, + ], + ids=["not-object", "empty", "header-only", "blank-model", "negative-shard", + "bad-timestamp", "missing-status"], +) +def test_a_malformed_report_is_invalid_and_never_admitted(payload): + "Malformed proof is rejected by the schema check, not by a later coincidence.\n\nTags: routing, tracker" + state = _evaluate(payload) + assert state.state == STATE_INVALID + assert not state.proven + + +def test_recorded_detail_carries_no_credentials_from_node_diagnostics(): + "Operator-facing admission detail is sanitized; a leaked token never reaches it.\n\nTags: routing, security, tracker" + state = _evaluate( + _report( + status="failed", + diagnostics=["download failed: authorization=hf_abcdefghijklmnopqrstuvwxyz01"], + ) + ) + assert state.state == STATE_FAILED + assert "hf_abcdefghijklmnopqrstuvwxyz01" not in state.detail + assert "hf_abcdefghijklmnopqrstuvwxyz01" not in json.dumps(state.to_dict()) + assert "[redacted]" in state.detail + + +# ---------------------------------------------------------- compatibility policy + + +def test_compat_policy_routes_a_legacy_node_but_never_a_broken_proof(): + "Older nodes (no proof) keep routing under `compat`; a bad proof never does.\n\nTags: routing, tracker" + absent = _evaluate(None) + failed = _evaluate(_report(status="failed")) + + assert absent.routable_under(POLICY_COMPAT) + assert not absent.routable_under(POLICY_ENFORCE) + assert not failed.routable_under(POLICY_COMPAT) + assert not failed.routable_under(POLICY_ENFORCE) + + +def test_the_policy_is_read_from_the_environment_and_defaults_to_compat(monkeypatch): + "The rollout switch is explicit and defaults to the compatible behaviour.\n\nTags: config, routing, tracker" + monkeypatch.delenv("MESHNET_TRACKER_CAPABILITY_POLICY", raising=False) + assert policy_from_env() == POLICY_COMPAT + monkeypatch.setenv("MESHNET_TRACKER_CAPABILITY_POLICY", "enforce") + assert policy_from_env() == POLICY_ENFORCE + monkeypatch.setenv("MESHNET_TRACKER_CAPABILITY_POLICY", "nonsense") + assert policy_from_env() == POLICY_COMPAT + + +# ------------------------------------------------------------- the routing gate + + +def _entry(node_id: str, start: int, end: int, report: dict | None, **kwargs) -> _NodeEntry: + from meshnet_tracker.capability import evaluate_report as _eval + + entry = _NodeEntry( + node_id=node_id, + endpoint=f"http://127.0.0.1:{9000 + int(node_id[-1])}", + shard_start=start, + shard_end=end, + model=SHORT, + shard_checksum=None, + hardware_profile={}, + wallet_address=None, + score=1.0, + hf_repo=MODEL, + num_layers=LAYERS, + capability=_eval( + report, + model_matches=lambda reported: reported == MODEL, + advertised_model=MODEL, + shard_start=start, + shard_end=end, + ), + **kwargs, + ) + return entry + + +def test_route_selection_drops_every_unadmitted_candidate_under_enforce(): + "Absent, failed, stale and mismatched candidates are all excluded.\n\nTags: routing, tracker" + good = _entry("node-1", 0, 31, _report(start=0, end=31)) + unproven = _entry("node-2", 0, 31, None) + failed = _entry("node-3", 0, 31, _report(start=0, end=31, status="failed")) + stale = _entry("node-4", 0, 31, _report(start=0, end=31, validated_at=time.time() - 86400)) + wrong_model = _entry("node-5", 0, 31, _report(start=0, end=31, model_id="someone/else-1b")) + + route, error = _select_route( + [unproven, failed, stale, wrong_model, good], 0, 31, policy=POLICY_ENFORCE + ) + assert not error + assert [n.node_id for n in route] == ["node-1"] + + only_bad, error = _select_route([unproven, failed, stale], 0, 31, policy=POLICY_ENFORCE) + assert only_bad == [] + assert "no route available" in error + + +def test_a_node_reassigned_to_a_shard_it_never_proved_stops_routing(): + "The proof does not travel with a tracker reassignment.\n\nTags: routing, tracker" + node = _entry("node-1", 0, 15, _report(start=0, end=15)) + assert _node_admission(node).state == STATE_ADMITTED + + node.shard_start, node.shard_end = 16, 31 # tracker rebalanced it + assert _node_admission(node).state == STATE_SHARD_MISMATCH + assert not _node_admission(node).routable_under(POLICY_COMPAT) + + +def test_admitted_candidates_keep_coverage_first_and_throughput_routing(): + "Gating removes candidates; among the survivors routing is unchanged.\n\nTags: routing, tracker" + head = _entry("node-1", 0, 15, _report(start=0, end=15)) + slow_tail = _entry( + "node-2", 16, 31, _report(start=16, end=31), benchmark_tokens_per_sec=5.0 + ) + fast_tail = _entry( + "node-3", 16, 31, _report(start=16, end=31), benchmark_tokens_per_sec=50.0 + ) + + route, error = _select_route( + [head, slow_tail, fast_tail], 0, 31, policy=POLICY_ENFORCE + ) + assert not error + # Coverage first (head, then a tail), and the faster of the two tied tails. + assert [n.node_id for n in route] == ["node-1", "node-3"] + + +# ---------------------------------------------------------------- over the wire + + +def test_an_enforcing_tracker_routes_a_proven_node_and_excludes_an_unproven_one(): + "End to end: a proof is required to appear in a route.\n\nTags: http, routing, tracker" + tracker = TrackerServer(capability_policy=POLICY_ENFORCE) + port = tracker.start() + try: + base = f"http://127.0.0.1:{port}" + _post_json(f"{base}/v1/nodes/register", _registration(9101, start=0, end=15)) + # A tail that presents no proof at all: the route cannot complete. + _post_json( + f"{base}/v1/nodes/register", + _registration(9102, start=16, end=31, report=None), + ) + + with pytest.raises(urllib.error.HTTPError) as exc: + _get_json(f"{base}/v1/route?model={SHORT}") + assert exc.value.code == 503 + + # Now the tail proves itself and the same route resolves. + _post_json(f"{base}/v1/nodes/register", _registration(9102, start=16, end=31)) + route = _get_json(f"{base}/v1/route?model={SHORT}")["route"] + assert route == ["http://127.0.0.1:9101", "http://127.0.0.1:9102"] + finally: + tracker.stop() + + +@pytest.mark.parametrize( + "bad_report", + [ + _report(start=16, end=31), # proves the wrong shard + _report(model_id="unrelated/other-7b"), # proves the wrong model + _report(status="failed"), + _report(validated_at=time.time() - 86400), # stale + {"schema_version": 1, "model": {}}, # malformed + ], + ids=["shard-mismatch", "model-mismatch", "failed", "stale", "invalid"], +) +def test_an_enforcing_tracker_never_routes_a_node_whose_proof_does_not_cover_it(bad_report): + "A proof for something else is worth exactly as much as no proof.\n\nTags: http, routing, tracker" + tracker = TrackerServer(capability_policy=POLICY_ENFORCE) + port = tracker.start() + try: + base = f"http://127.0.0.1:{port}" + resp = _post_json( + f"{base}/v1/nodes/register", + _registration(9111, start=0, end=31, report=bad_report), + ) + # Registration still succeeds — the operator must be able to see the node. + assert resp["node_id"] + assert resp["capability"]["routable"] is False + assert resp["capability"]["state"] != STATE_ADMITTED + + with pytest.raises(urllib.error.HTTPError) as exc: + _get_json(f"{base}/v1/route?model={SHORT}") + assert exc.value.code in (404, 503) + finally: + tracker.stop() + + +def test_a_compat_tracker_routes_a_legacy_node_that_sends_no_report(): + "Documented rollout policy: pre-capability nodes keep working under `compat`.\n\nTags: http, routing, tracker" + tracker = TrackerServer(capability_policy=POLICY_COMPAT) + port = tracker.start() + try: + base = f"http://127.0.0.1:{port}" + _post_json( + f"{base}/v1/nodes/register", + _registration(9121, start=0, end=31, report=None, recipe_id=None, recipe_version=None), + ) + route = _get_json(f"{base}/v1/route?model={SHORT}")["route"] + assert route == ["http://127.0.0.1:9121"] + finally: + tracker.stop() + + +def test_a_compat_tracker_still_refuses_a_node_that_presents_a_failed_proof(): + "Compatibility grandfathers silence, not a proof of failure.\n\nTags: http, routing, tracker" + tracker = TrackerServer(capability_policy=POLICY_COMPAT) + port = tracker.start() + try: + base = f"http://127.0.0.1:{port}" + _post_json( + f"{base}/v1/nodes/register", + _registration(9131, start=0, end=31, report=_report(start=0, end=31, status="failed")), + ) + with pytest.raises(urllib.error.HTTPError) as exc: + _get_json(f"{base}/v1/route?model={SHORT}") + assert exc.value.code == 503 + finally: + tracker.stop() + + +def test_a_replicated_registration_carries_its_verdict_to_a_follower(): + "A proven node must not be routable on the leader and dark on every follower.\n\nTags: cluster, routing, tracker" + tracker = TrackerServer(capability_policy=POLICY_ENFORCE) + proven = _registration(9151, start=0, end=31, report=_report(start=0, end=31)) + proven["node_id"] = "follower-node-1" + unproven = _registration(9152, start=0, end=31, report=None) + unproven["node_id"] = "follower-node-2" + + tracker._raft_apply("register", proven) + tracker._raft_apply("register", unproven) + + admitted = tracker._registry["follower-node-1"] + assert admitted.capability.state == STATE_ADMITTED + assert _capability_routable(admitted, POLICY_ENFORCE) + + absent = tracker._registry["follower-node-2"] + assert absent.capability.state == STATE_ABSENT + assert not _capability_routable(absent, POLICY_ENFORCE) + + +def test_the_network_map_exposes_the_admission_state_of_every_node(): + "The operator view answers 'why is my node not routing' without raw internals.\n\nTags: http, routing, tracker" + tracker = TrackerServer(capability_policy=POLICY_ENFORCE) + port = tracker.start() + try: + base = f"http://127.0.0.1:{port}" + _post_json(f"{base}/v1/nodes/register", _registration(9141, start=0, end=15)) + _post_json( + f"{base}/v1/nodes/register", + _registration( + 9142, + start=16, + end=31, + report=_report( + start=16, + end=31, + status="failed", + diagnostics=["load failed: token=hf_abcdefghijklmnopqrstuvwx1234"], + ), + ), + ) + + network = _get_json(f"{base}/v1/network/map") + assert network["capability_policy"] == POLICY_ENFORCE + by_endpoint = {n["endpoint"]: n["capability"] for n in network["nodes"]} + + proven = by_endpoint["http://127.0.0.1:9141"] + assert proven["state"] == STATE_ADMITTED + assert proven["routable"] is True + assert proven["model_id"] == MODEL + assert (proven["shard_start"], proven["shard_end"]) == (0, 15) + assert proven["recipe_id"] == "baseline" + assert proven["device"] == "cpu" + + broken = by_endpoint["http://127.0.0.1:9142"] + assert broken["state"] == STATE_FAILED + assert broken["routable"] is False + assert broken["detail"] + + raw = json.dumps(network) + assert "hf_abcdefghijklmnopqrstuvwx1234" not in raw + assert "Traceback" not in raw + finally: + tracker.stop()