[verified] feat: complete Ralph task workstreams
This commit is contained in:
21
.vscode/launch.json
vendored
21
.vscode/launch.json
vendored
@@ -7,7 +7,26 @@
|
|||||||
"request": "launch",
|
"request": "launch",
|
||||||
"python": "${workspaceFolder}/.venv-rocm/bin/python",
|
"python": "${workspaceFolder}/.venv-rocm/bin/python",
|
||||||
"module": "meshnet_tracker.cli",
|
"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",
|
"console": "integratedTerminal",
|
||||||
"justMyCode": false
|
"justMyCode": false
|
||||||
},
|
},
|
||||||
|
|||||||
130
QUICKSTART.md
130
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
|
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.
|
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 |
|
| Role | `--model` / alias | HF repo | Notes |
|
||||||
|------|-------------------|---------|-------|
|
|------|-------------------|---------|-------|
|
||||||
@@ -288,8 +296,8 @@ HF_HOME=/path/to/models .venv-rocm/bin/meshnet-node start \
|
|||||||
--quantization bfloat16
|
--quantization bfloat16
|
||||||
```
|
```
|
||||||
|
|
||||||
For the Qwen3.6 alpha model on Linux ROCm, install the optional FLA ROCm fast
|
A model whose recipe uses an optional accelerator needs that package in the *same*
|
||||||
path in the same env:
|
env. For example, `qwen3.6-35b-a3b` on Linux ROCm can use FLA:
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
.venv-rocm/bin/pip install 'flash-linear-attention[rocm]'
|
.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
|
--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 <model>` before starting the node; startup
|
||||||
|
will refuse to register anyway if the shard cannot execute.
|
||||||
|
|
||||||
### Linux ROCm: Triton JIT compiler prerequisite
|
### Linux ROCm: Triton JIT compiler prerequisite
|
||||||
|
|
||||||
Some model/runtime paths invoke Triton at the first real forward. Triton builds a local HIP
|
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.
|
no HuggingFace gating. Best for first-time setup.
|
||||||
|
|
||||||
**Alpha model:** `qwen3.6-35b-a3b` — 40 layers, ~72 GB BF16 download, MoE with hybrid
|
**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
|
linear attention. Tracker accepts the alias or full repo id (`unsloth/Qwen3.6-35B-A3B`).
|
||||||
GPU use `flash-linear-attention[cuda]`. 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`).
|
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
|
fetches only `config.json`. Override with `--shard-start` / `--shard-end` for partial
|
||||||
shards or multi-node splits.
|
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 <model> --quantization bfloat16
|
||||||
|
```
|
||||||
|
|
||||||
|
```
|
||||||
|
meshnet-node doctor
|
||||||
|
Model: <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 <id>` 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
|
### Core command
|
||||||
|
|
||||||
Replace `<tracker-url>` and adjust the prefix for your shell (see table above).
|
Replace `<tracker-url>` 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
|
.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
|
```powershell
|
||||||
$env:HF_HOME = "D:\DEV\models"
|
$env:HF_HOME = "D:\DEV\models"
|
||||||
pip install triton-windows
|
pip install triton-windows
|
||||||
pip install -U flash-linear-attention
|
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
|
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).
|
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
|
```bash
|
||||||
HF_HOME=/path/to/models .venv/bin/meshnet-node start --tracker <tracker-url> --model qwen3.6-35b-a3b --quantization bfloat16
|
|
||||||
# Install once on that machine: pip install flash-linear-attention[cuda]
|
# Install once on that machine: pip install flash-linear-attention[cuda]
|
||||||
|
HF_HOME=/path/to/models .venv/bin/meshnet-node start --tracker <tracker-url> --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
|
```bash
|
||||||
HF_HOME=/path/to/models .venv-rocm/bin/meshnet-node start --tracker <tracker-url> --model qwen3.6-35b-a3b --quantization bfloat16
|
|
||||||
# Install once on that machine: .venv-rocm/bin/pip install 'flash-linear-attention[rocm]'
|
# 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 <tracker-url> --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
|
After the first node registers a model, later nodes can join with only the tracker
|
||||||
URL (shard auto-assigned):
|
URL (shard auto-assigned):
|
||||||
|
|
||||||
@@ -786,3 +890,9 @@ failure the node logs a warning and falls back to direct HTTP before erroring.
|
|||||||
```bash
|
```bash
|
||||||
.venv/bin/python -m pytest -q
|
.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.
|
||||||
|
|||||||
@@ -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.
|
- 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.
|
- 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
|
## Consequences
|
||||||
|
|
||||||
- First startup has a bounded validation cost before registration, but failures occur before traffic rather than under a paid request.
|
- First startup has a bounded validation cost before registration, but failures occur before traffic rather than under a paid request.
|
||||||
|
|||||||
105
docs/dev/certified-hardware-lanes.md
Normal file
105
docs/dev/certified-hardware-lanes.md
Normal file
@@ -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.
|
||||||
93
docs/dev/dashboard-test-runner.md
Normal file
93
docs/dev/dashboard-test-runner.md
Normal file
@@ -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.<hostname>` 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:<name>` 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 |
|
||||||
@@ -38,6 +38,12 @@ even without installing `packages/node`.
|
|||||||
.venv/bin/python -m pytest
|
.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
|
## Optional-dependency tests
|
||||||
|
|
||||||
Some tests import heavyweight or optional third-party packages and guard
|
Some tests import heavyweight or optional third-party packages and guard
|
||||||
|
|||||||
130
packages/gateway/meshnet_gateway/prefill_backpressure.py
Normal file
130
packages/gateway/meshnet_gateway/prefill_backpressure.py
Normal file
@@ -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
|
||||||
@@ -14,6 +14,8 @@ import urllib.request
|
|||||||
import uuid
|
import uuid
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
|
from .prefill_backpressure import BoundedPrefillSender, PrefillTransferLimits
|
||||||
|
|
||||||
_STUB_HIDDEN_DIM = 64
|
_STUB_HIDDEN_DIM = 64
|
||||||
_STUB_DTYPE = "bfloat16"
|
_STUB_DTYPE = "bfloat16"
|
||||||
_WIRE_VERSION = "2"
|
_WIRE_VERSION = "2"
|
||||||
@@ -653,15 +655,16 @@ def _last_message_content(messages: object) -> str:
|
|||||||
|
|
||||||
def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) -> list[_BinaryActivation]:
|
def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) -> list[_BinaryActivation]:
|
||||||
session = str(uuid.uuid4())
|
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))
|
total_tokens = max(1, _prompt_token_count(prompt))
|
||||||
chunk_total = max(1, (total_tokens + chunk_token_count - 1) // chunk_token_count)
|
chunk_total = max(1, (total_tokens + chunk_token_count - 1) // chunk_token_count)
|
||||||
responses: list[_BinaryActivation] = []
|
|
||||||
|
|
||||||
|
def chunks():
|
||||||
for chunk_index in range(chunk_total):
|
for chunk_index in range(chunk_total):
|
||||||
remaining_tokens = total_tokens - (chunk_index * chunk_token_count)
|
remaining_tokens = total_tokens - (chunk_index * chunk_token_count)
|
||||||
seq_len = min(chunk_token_count, remaining_tokens)
|
seq_len = min(chunk_token_count, remaining_tokens)
|
||||||
activation = _BinaryActivation(
|
yield _BinaryActivation(
|
||||||
body=_make_stub_binary_activation([1, seq_len, _STUB_HIDDEN_DIM], _STUB_DTYPE),
|
body=_make_stub_binary_activation([1, seq_len, _STUB_HIDDEN_DIM], _STUB_DTYPE),
|
||||||
shape=[1, seq_len, _STUB_HIDDEN_DIM],
|
shape=[1, seq_len, _STUB_HIDDEN_DIM],
|
||||||
dtype=_STUB_DTYPE,
|
dtype=_STUB_DTYPE,
|
||||||
@@ -671,6 +674,8 @@ def _run_binary_pipeline(route: list[str], prompt: str, timeout: float = 5.0) ->
|
|||||||
encoding=_preferred_binary_encoding(),
|
encoding=_preferred_binary_encoding(),
|
||||||
headers={},
|
headers={},
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def forward(activation: _BinaryActivation) -> _BinaryActivation:
|
||||||
for hop_index, node_url in enumerate(route):
|
for hop_index, node_url in enumerate(route):
|
||||||
activation = _post_binary_forward(
|
activation = _post_binary_forward(
|
||||||
f"{node_url}/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,
|
hop_index=hop_index,
|
||||||
timeout=timeout,
|
timeout=timeout,
|
||||||
)
|
)
|
||||||
responses.append(activation)
|
return activation
|
||||||
return responses
|
|
||||||
|
# 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:
|
def _chunk_token_count() -> int:
|
||||||
raw_value = os.environ.get("MESHNET_CHUNK_TOKENS", "128")
|
return PrefillTransferLimits.from_env().chunk_tokens
|
||||||
try:
|
|
||||||
value = int(raw_value)
|
|
||||||
except ValueError:
|
|
||||||
return 128
|
|
||||||
return value if value > 0 else 128
|
|
||||||
|
|
||||||
|
|
||||||
def _prompt_token_count(prompt: str) -> int:
|
def _prompt_token_count(prompt: str) -> int:
|
||||||
@@ -737,6 +740,18 @@ def _post_binary_forward(
|
|||||||
|
|
||||||
encoding = response_headers.get("x-meshnet-encoding")
|
encoding = response_headers.get("x-meshnet-encoding")
|
||||||
raw_body = _decompress_body(response_body, encoding)
|
raw_body = _decompress_body(response_body, encoding)
|
||||||
|
# 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"])
|
shape = _parse_shape(response_headers["x-meshnet-shape"])
|
||||||
dtype = response_headers["x-meshnet-dtype"]
|
dtype = response_headers["x-meshnet-dtype"]
|
||||||
session = response_headers["x-meshnet-session"]
|
session = response_headers["x-meshnet-session"]
|
||||||
|
|||||||
142
packages/node/meshnet_node/activation_compression.py
Normal file
142
packages/node/meshnet_node/activation_compression.py
Normal file
@@ -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_<ROUTE>_<PHASE>_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"}
|
||||||
225
packages/node/meshnet_node/admission.py
Normal file
225
packages/node/meshnet_node/admission.py
Normal file
@@ -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
|
||||||
@@ -237,6 +237,85 @@ def _cmd_config(args) -> int:
|
|||||||
return 0
|
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:
|
def _cmd_start(args) -> int:
|
||||||
"""Legacy `start` subcommand — preserves backward compatibility with existing tests."""
|
"""Legacy `start` subcommand — preserves backward compatibility with existing tests."""
|
||||||
from .config import DEFAULTS
|
from .config import DEFAULTS
|
||||||
@@ -322,6 +401,7 @@ def main() -> None:
|
|||||||
" models List supported models\n"
|
" models List supported models\n"
|
||||||
" models --browse Browse HuggingFace Hub\n"
|
" models --browse Browse HuggingFace Hub\n"
|
||||||
" config Show current config\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
|
# config subcommand
|
||||||
subparsers.add_parser("config", help="Show current saved config")
|
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 subcommand (legacy / backward-compat)
|
||||||
start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)")
|
start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)")
|
||||||
start_cmd.add_argument("--tracker")
|
start_cmd.add_argument("--tracker")
|
||||||
@@ -406,6 +520,8 @@ def main() -> None:
|
|||||||
sys.exit(_cmd_models(args))
|
sys.exit(_cmd_models(args))
|
||||||
elif args.command == "config":
|
elif args.command == "config":
|
||||||
sys.exit(_cmd_config(args))
|
sys.exit(_cmd_config(args))
|
||||||
|
elif args.command == "doctor":
|
||||||
|
sys.exit(_cmd_doctor(args))
|
||||||
elif args.command == "start":
|
elif args.command == "start":
|
||||||
sys.exit(_cmd_start(args))
|
sys.exit(_cmd_start(args))
|
||||||
else:
|
else:
|
||||||
|
|||||||
633
packages/node/meshnet_node/doctor.py
Normal file
633
packages/node/meshnet_node/doctor.py
Normal file
@@ -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 <repo-or-name>, 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)
|
||||||
@@ -9,16 +9,26 @@ import json
|
|||||||
import os
|
import os
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
|
import warnings
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, Literal
|
from typing import Any, Literal, Mapping
|
||||||
|
|
||||||
Quantization = Literal["auto", "bfloat16", "int8", "nf4"]
|
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):
|
class ModelBackendError(RuntimeError):
|
||||||
"""Base class for real model backend startup and execution failures."""
|
"""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):
|
class MissingModelDependencyError(ModelBackendError):
|
||||||
"""Raised when optional model dependencies are not installed."""
|
"""Raised when optional model dependencies are not installed."""
|
||||||
|
|
||||||
@@ -61,6 +71,14 @@ def _torch_cuda_is_executable(torch_module: Any) -> bool:
|
|||||||
|
|
||||||
@dataclass(frozen=True)
|
@dataclass(frozen=True)
|
||||||
class TensorPayload:
|
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
|
body: bytes
|
||||||
shape: list[int]
|
shape: list[int]
|
||||||
attention_mask_header: str | None
|
attention_mask_header: str | None
|
||||||
@@ -213,6 +231,7 @@ class TorchModelShard:
|
|||||||
quantization: Quantization = "auto",
|
quantization: Quantization = "auto",
|
||||||
cache_dir: Path | None = None,
|
cache_dir: Path | None = None,
|
||||||
force_cpu: bool = False,
|
force_cpu: bool = False,
|
||||||
|
recipe_params: Mapping[str, Any] | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
|
if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
|
||||||
raise ValueError("shard_start must be <= shard_end and non-negative")
|
raise ValueError("shard_start must be <= shard_end and non-negative")
|
||||||
@@ -220,6 +239,8 @@ class TorchModelShard:
|
|||||||
self.shard_start = shard_start
|
self.shard_start = shard_start
|
||||||
self.shard_end = shard_end
|
self.shard_end = shard_end
|
||||||
self.quantization = quantization
|
self.quantization = quantization
|
||||||
|
self.recipe_params = validate_recipe_params(recipe_params)
|
||||||
|
attn_implementation = self.recipe_params.get("attn_implementation")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
import torch
|
import torch
|
||||||
@@ -260,6 +281,7 @@ class TorchModelShard:
|
|||||||
shard_end,
|
shard_end,
|
||||||
dtype,
|
dtype,
|
||||||
self.device,
|
self.device,
|
||||||
|
attn_implementation=attn_implementation,
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
load_kwargs = {
|
load_kwargs = {
|
||||||
@@ -270,6 +292,8 @@ class TorchModelShard:
|
|||||||
}
|
}
|
||||||
if quant_config is not None:
|
if quant_config is not None:
|
||||||
load_kwargs["quantization_config"] = quant_config
|
load_kwargs["quantization_config"] = quant_config
|
||||||
|
if attn_implementation is not None:
|
||||||
|
load_kwargs["attn_implementation"] = attn_implementation
|
||||||
self.model = AutoModelForCausalLM.from_pretrained(
|
self.model = AutoModelForCausalLM.from_pretrained(
|
||||||
load_source,
|
load_source,
|
||||||
**load_kwargs,
|
**load_kwargs,
|
||||||
@@ -313,6 +337,8 @@ class TorchModelShard:
|
|||||||
# consume CPU tensors ("Pointer argument cannot be accessed from Triton"),
|
# consume CPU tensors ("Pointer argument cannot be accessed from Triton"),
|
||||||
# so CPU shards intentionally stay on the stateless prefill path.
|
# so CPU shards intentionally stay on the stateless prefill path.
|
||||||
self.supports_kv_cache = self.device.type != "cpu"
|
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(
|
self.kv_sessions = SessionCacheStore(
|
||||||
max_sessions=int(os.environ.get("MESHNET_KV_MAX_SESSIONS", "8")),
|
max_sessions=int(os.environ.get("MESHNET_KV_MAX_SESSIONS", "8")),
|
||||||
ttl_seconds=float(os.environ.get("MESHNET_KV_TTL_SECONDS", "600")),
|
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(
|
def load_torch_shard(
|
||||||
model_id: str,
|
model_id: str,
|
||||||
shard_start: int,
|
shard_start: int,
|
||||||
@@ -695,9 +734,16 @@ def load_torch_shard(
|
|||||||
quantization: Quantization = "auto",
|
quantization: Quantization = "auto",
|
||||||
cache_dir: Path | None = None,
|
cache_dir: Path | None = None,
|
||||||
force_cpu: bool = False,
|
force_cpu: bool = False,
|
||||||
|
recipe_params: Mapping[str, Any] | None = None,
|
||||||
) -> TorchModelShard:
|
) -> TorchModelShard:
|
||||||
return 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,
|
init_empty_weights_fn: Any | None = None,
|
||||||
set_tensor_fn: Any | None = None,
|
set_tensor_fn: Any | None = None,
|
||||||
safe_open_fn: Any | None = None,
|
safe_open_fn: Any | None = None,
|
||||||
|
attn_implementation: str | None = None,
|
||||||
) -> Any:
|
) -> Any:
|
||||||
from .model_catalog import layers_from_config
|
from .model_catalog import layers_from_config
|
||||||
from .safetensors_selection import (
|
from .safetensors_selection import (
|
||||||
@@ -763,6 +810,10 @@ def _load_partial_model_from_snapshot(
|
|||||||
|
|
||||||
snapshot_dir = Path(load_source)
|
snapshot_dir = Path(load_source)
|
||||||
cfg = auto_config.from_pretrained(str(snapshot_dir))
|
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)
|
total_layers = layers_from_config(cfg)
|
||||||
if total_layers is None:
|
if total_layers is None:
|
||||||
raise PartialModelLoadUnsupported(
|
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:
|
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)
|
return tensor.reshape(shape)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -3,6 +3,7 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import base64
|
import base64
|
||||||
|
import http.client
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
@@ -10,8 +11,6 @@ import re
|
|||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
import urllib.parse
|
import urllib.parse
|
||||||
import urllib.error
|
|
||||||
import urllib.request
|
|
||||||
from concurrent.futures import ThreadPoolExecutor
|
from concurrent.futures import ThreadPoolExecutor
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
|
|
||||||
@@ -44,8 +43,14 @@ BINARY_FRAME_MAGIC = b"MRF1"
|
|||||||
|
|
||||||
|
|
||||||
def encode_binary_frame(header: dict, body: bytes) -> bytes:
|
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()
|
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]:
|
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")
|
raise ValueError("not a meshnet binary relay frame")
|
||||||
header_len = int.from_bytes(frame[4:8], "big")
|
header_len = int.from_bytes(frame[4:8], "big")
|
||||||
header = json.loads(frame[8:8 + header_len].decode())
|
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)
|
@dataclass(frozen=True)
|
||||||
@@ -82,6 +89,62 @@ def _max_concurrency_from_env() -> int:
|
|||||||
return max(1, value)
|
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:
|
class RelayHttpBridge:
|
||||||
"""Connect outbound to a relay and proxy relay HTTP requests to localhost.
|
"""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_log_lock = threading.Lock()
|
||||||
self._decode_steps: dict[str, int] = {}
|
self._decode_steps: dict[str, int] = {}
|
||||||
self._ws = None
|
self._ws = None
|
||||||
|
self._loopback_clients = _LoopbackHttpClientPool(self.local_base_url)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def relay_addr(self) -> str:
|
def relay_addr(self) -> str:
|
||||||
@@ -141,6 +205,7 @@ class RelayHttpBridge:
|
|||||||
self._thread.join(timeout=3.0)
|
self._thread.join(timeout=3.0)
|
||||||
if self._executor is not None:
|
if self._executor is not None:
|
||||||
self._executor.shutdown(wait=False)
|
self._executor.shutdown(wait=False)
|
||||||
|
self._loopback_clients.close()
|
||||||
|
|
||||||
def _run(self) -> None:
|
def _run(self) -> None:
|
||||||
import websockets.sync.client as wsc # type: ignore[import]
|
import websockets.sync.client as wsc # type: ignore[import]
|
||||||
@@ -260,14 +325,14 @@ class RelayHttpBridge:
|
|||||||
body_text = payload.get("body") or ""
|
body_text = payload.get("body") or ""
|
||||||
data = body_text.encode() if isinstance(body_text, str) else bytes(body_text)
|
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:
|
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)
|
resp_headers = dict(resp.headers)
|
||||||
content_type = resp.headers.get("Content-Type", "")
|
content_type = resp.headers.get("Content-Type", "")
|
||||||
if "text/event-stream" in 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
|
return
|
||||||
resp_bytes = resp.read()
|
resp_bytes = resp.read()
|
||||||
# Forward all X-Meshnet-* headers so the caller can reconstruct the activation.
|
# Forward all X-Meshnet-* headers so the caller can reconstruct the activation.
|
||||||
@@ -289,14 +354,18 @@ class RelayHttpBridge:
|
|||||||
else:
|
else:
|
||||||
result["body"] = resp_bytes.decode(errors="replace")
|
result["body"] = resp_bytes.decode(errors="replace")
|
||||||
self._send_response_frame(result)
|
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({
|
self._send_response_frame({
|
||||||
"request_id": request_id,
|
"request_id": request_id,
|
||||||
"status": exc.code,
|
"status": 503,
|
||||||
"headers": {"Content-Type": exc.headers.get("Content-Type", "application/json")},
|
"headers": {"Content-Type": "application/json"},
|
||||||
"body": exc.read().decode(errors="replace"),
|
"body": json.dumps({"error": f"relay bridge local request failed: {exc}"}),
|
||||||
})
|
})
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
|
self._loopback_clients.discard()
|
||||||
self._send_response_frame({
|
self._send_response_frame({
|
||||||
"request_id": request_id,
|
"request_id": request_id,
|
||||||
"status": 503,
|
"status": 503,
|
||||||
@@ -304,7 +373,7 @@ class RelayHttpBridge:
|
|||||||
"body": json.dumps({"error": f"relay bridge local request failed: {exc}"}),
|
"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.
|
"""Forward an SSE response as chunk frames, one per complete SSE event.
|
||||||
|
|
||||||
Frame order: header frame (status + headers), chunk frames, done frame.
|
Frame order: header frame (status + headers), chunk frames, done frame.
|
||||||
@@ -319,7 +388,7 @@ class RelayHttpBridge:
|
|||||||
"done": False,
|
"done": False,
|
||||||
})
|
})
|
||||||
if not sent:
|
if not sent:
|
||||||
return
|
return False
|
||||||
event_lines: list[str] = []
|
event_lines: list[str] = []
|
||||||
for raw_line in resp:
|
for raw_line in resp:
|
||||||
line = raw_line.decode(errors="replace")
|
line = raw_line.decode(errors="replace")
|
||||||
@@ -333,7 +402,7 @@ class RelayHttpBridge:
|
|||||||
"chunk": "".join(event_lines),
|
"chunk": "".join(event_lines),
|
||||||
"done": False,
|
"done": False,
|
||||||
}):
|
}):
|
||||||
return
|
return False
|
||||||
event_lines = []
|
event_lines = []
|
||||||
if event_lines:
|
if event_lines:
|
||||||
if not self._send_response_frame({
|
if not self._send_response_frame({
|
||||||
@@ -342,8 +411,8 @@ class RelayHttpBridge:
|
|||||||
"chunk": "".join(event_lines),
|
"chunk": "".join(event_lines),
|
||||||
"done": False,
|
"done": False,
|
||||||
}):
|
}):
|
||||||
return
|
return False
|
||||||
self._send_response_frame({
|
return self._send_response_frame({
|
||||||
"request_id": request_id,
|
"request_id": request_id,
|
||||||
"stream": True,
|
"stream": True,
|
||||||
"done": True,
|
"done": True,
|
||||||
|
|||||||
385
packages/node/meshnet_node/route_session_benchmark.py
Normal file
385
packages/node/meshnet_node/route_session_benchmark.py
Normal file
@@ -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())
|
||||||
155
packages/node/meshnet_node/seam_telemetry.py
Normal file
155
packages/node/meshnet_node/seam_telemetry.py
Normal file
@@ -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
|
||||||
@@ -8,6 +8,7 @@ import urllib.parse
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from .downloader import compute_shard_checksum, write_shard_archive
|
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):
|
# Binary activation wire format (contract for all shard nodes):
|
||||||
# POST /forward with raw tensor bytes in the body and tensor/session/chunk
|
# POST /forward with raw tensor bytes in the body and tensor/session/chunk
|
||||||
@@ -21,6 +22,7 @@ _DTYPE_SIZES = {
|
|||||||
"bfloat16": 2,
|
"bfloat16": 2,
|
||||||
"float32": 4,
|
"float32": 4,
|
||||||
}
|
}
|
||||||
|
_COMPRESSION_POLICIES = CompressionPolicies()
|
||||||
|
|
||||||
|
|
||||||
def _make_stub_binary_activation(shape: list[int], dtype: str) -> bytes:
|
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:
|
def _decompress_body(body: bytes, encoding: str | None) -> bytes:
|
||||||
if not encoding:
|
return decompress_activation(body, encoding).body
|
||||||
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
|
|
||||||
|
|
||||||
|
|
||||||
def _compress_body(body: bytes, encoding: str | None) -> bytes:
|
def _compress_body(body: bytes, encoding: str | None) -> bytes:
|
||||||
@@ -307,7 +300,14 @@ class _StubHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
server.received_activations = True
|
server.received_activations = True
|
||||||
|
|
||||||
raw_payload = _make_stub_binary_activation(shape, dtype)
|
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_response(200)
|
||||||
self.send_header("Content-Type", "application/octet-stream")
|
self.send_header("Content-Type", "application/octet-stream")
|
||||||
self.send_header("Content-Length", str(len(payload)))
|
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-Session", session)
|
||||||
self.send_header("X-Meshnet-Chunk-Index", chunk_index)
|
self.send_header("X-Meshnet-Chunk-Index", chunk_index)
|
||||||
self.send_header("X-Meshnet-Chunk-Total", chunk_total)
|
self.send_header("X-Meshnet-Chunk-Total", chunk_total)
|
||||||
if encoding:
|
if compression.encoding:
|
||||||
self.send_header("X-Meshnet-Encoding", encoding)
|
self.send_header("X-Meshnet-Encoding", compression.encoding)
|
||||||
if server.is_last_shard:
|
if server.is_last_shard:
|
||||||
self.send_header("X-Meshnet-Stub-Response-Prefix", server.response_prefix)
|
self.send_header("X-Meshnet-Stub-Response-Prefix", server.response_prefix)
|
||||||
self.end_headers()
|
self.end_headers()
|
||||||
|
|||||||
@@ -14,9 +14,19 @@ import urllib.request
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
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 .downloader import compute_shard_checksum, download_shard
|
||||||
from .hardware import detect_hardware, benchmark_throughput_checked, with_forced_cpu
|
from .hardware import detect_hardware, benchmark_throughput_checked, with_forced_cpu
|
||||||
from .model_catalog import model_metadata_for
|
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 .relay_bridge import RelayHttpBridge, peer_id_from_wallet
|
||||||
from .server import StubNodeServer
|
from .server import StubNodeServer
|
||||||
from .torch_server import TorchNodeServer
|
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}"
|
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(
|
def run_startup(
|
||||||
tracker_url: str,
|
tracker_url: str,
|
||||||
port: int = 0,
|
port: int = 0,
|
||||||
@@ -668,6 +740,8 @@ def run_startup(
|
|||||||
torch_interop_threads: int | None = None,
|
torch_interop_threads: int | None = None,
|
||||||
node_name: str | None = None,
|
node_name: str | None = None,
|
||||||
force_cpu: bool = False,
|
force_cpu: bool = False,
|
||||||
|
recipe_id: str | None = None,
|
||||||
|
capability_validator: CapabilityValidator | None = None,
|
||||||
) -> StubNodeServer | TorchNodeServer:
|
) -> StubNodeServer | TorchNodeServer:
|
||||||
"""Execute the full startup sequence and return a running node server.
|
"""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
|
2. Load or generate Solana wallet keypair
|
||||||
3. Query tracker for optimal shard assignment
|
3. Query tracker for optimal shard assignment
|
||||||
4. Download (or stub) the assigned shard from peers, then HuggingFace
|
4. Download (or stub) the assigned shard from peers, then HuggingFace
|
||||||
5. Start local HTTP server
|
5. Prove the loaded shard runs — a failure here exits before step 6
|
||||||
6. Register with tracker
|
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.
|
Prints a compact status summary on completion.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
tracker_url = tracker_url.rstrip("/")
|
tracker_url = tracker_url.rstrip("/")
|
||||||
|
manifest, recipe = _resolve_recipe(recipe_id)
|
||||||
relay_url = _discover_relay_url(tracker_url)
|
relay_url = _discover_relay_url(tracker_url)
|
||||||
display_fields = _registration_display_fields(node_name)
|
display_fields = _registration_display_fields(node_name)
|
||||||
if max_loaded_shards < 1:
|
if max_loaded_shards < 1:
|
||||||
@@ -874,6 +953,20 @@ def run_startup(
|
|||||||
debug=debug,
|
debug=debug,
|
||||||
max_loaded_shards=max_loaded_shards,
|
max_loaded_shards=max_loaded_shards,
|
||||||
force_cpu=force_cpu,
|
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()
|
_node_start_time = time.monotonic()
|
||||||
actual_port = node.start()
|
actual_port = node.start()
|
||||||
@@ -910,6 +1003,11 @@ def run_startup(
|
|||||||
"tracker_mode": (shard_start == 0),
|
"tracker_mode": (shard_start == 0),
|
||||||
"managed_assignment": not user_pinned_shard,
|
"managed_assignment": not user_pinned_shard,
|
||||||
"model_metadata": model_metadata_for(model_id, total_layers, cache_dir=cache_dir),
|
"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_models": (
|
||||||
_downloaded_model_inventory(
|
_downloaded_model_inventory(
|
||||||
model_id.split("/")[-1],
|
model_id.split("/")[-1],
|
||||||
@@ -1030,6 +1128,20 @@ def run_startup(
|
|||||||
debug=debug,
|
debug=debug,
|
||||||
max_loaded_shards=max_loaded_shards,
|
max_loaded_shards=max_loaded_shards,
|
||||||
force_cpu=force_cpu,
|
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()
|
_node_start_time = time.monotonic()
|
||||||
actual_port = node.start()
|
actual_port = node.start()
|
||||||
@@ -1062,6 +1174,11 @@ def run_startup(
|
|||||||
"tracker_mode": (assigned_shard_start == 0),
|
"tracker_mode": (assigned_shard_start == 0),
|
||||||
"managed_assignment": True,
|
"managed_assignment": True,
|
||||||
"model_metadata": model_metadata_for(assigned_hf_repo, assigned_num_layers, cache_dir=cache_dir),
|
"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_models": (
|
||||||
_downloaded_model_inventory(
|
_downloaded_model_inventory(
|
||||||
assigned_hf_repo.split("/")[-1],
|
assigned_hf_repo.split("/")[-1],
|
||||||
@@ -1212,6 +1329,20 @@ def run_startup(
|
|||||||
debug=debug,
|
debug=debug,
|
||||||
max_loaded_shards=max_loaded_shards,
|
max_loaded_shards=max_loaded_shards,
|
||||||
force_cpu=force_cpu,
|
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()
|
actual_port = node.start()
|
||||||
total_layers = getattr(getattr(node, "backend", None), "total_layers", None) or assigned_total_layers
|
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),
|
"tracker_mode": (shard_start == 0),
|
||||||
"managed_assignment": not user_pinned_shard,
|
"managed_assignment": not user_pinned_shard,
|
||||||
"model_metadata": model_metadata_for(hf_repo, total_layers, cache_dir=shard_path),
|
"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,
|
**registration_capabilities,
|
||||||
**relay_fields,
|
**relay_fields,
|
||||||
**display_fields,
|
**display_fields,
|
||||||
@@ -1282,6 +1418,19 @@ def run_startup(
|
|||||||
model=assigned_model,
|
model=assigned_model,
|
||||||
shard_path=shard_path,
|
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()
|
actual_port = node.start()
|
||||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||||
endpoint = f"http://{public_host}:{actual_port}"
|
endpoint = f"http://{public_host}:{actual_port}"
|
||||||
@@ -1304,6 +1453,11 @@ def run_startup(
|
|||||||
"shard_start": shard_start,
|
"shard_start": shard_start,
|
||||||
"shard_end": shard_end,
|
"shard_end": shard_end,
|
||||||
"shard_checksum": shard_checksum,
|
"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,
|
"downloaded_models": downloaded_models,
|
||||||
"hardware_profile": hw,
|
"hardware_profile": hw,
|
||||||
"wallet_address": address,
|
"wallet_address": address,
|
||||||
|
|||||||
70
packages/node/meshnet_node/testing.py
Normal file
70
packages/node/meshnet_node/testing.py
Normal file
@@ -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
|
||||||
@@ -3,17 +3,17 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import base64
|
import base64
|
||||||
|
import http.client
|
||||||
import http.server
|
import http.server
|
||||||
import json
|
import json
|
||||||
import sys
|
import sys
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
import urllib.error
|
|
||||||
import urllib.parse
|
import urllib.parse
|
||||||
import urllib.request
|
import urllib.request
|
||||||
import uuid
|
import uuid
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
from typing import Any, Mapping
|
||||||
|
|
||||||
from .model_backend import (
|
from .model_backend import (
|
||||||
InsufficientVRAMError,
|
InsufficientVRAMError,
|
||||||
@@ -22,16 +22,32 @@ from .model_backend import (
|
|||||||
Quantization,
|
Quantization,
|
||||||
TailTokenResult,
|
TailTokenResult,
|
||||||
TorchModelShard,
|
TorchModelShard,
|
||||||
|
_tensor_from_bfloat16_bytes,
|
||||||
validate_quantization,
|
validate_quantization,
|
||||||
)
|
)
|
||||||
|
from .seam_telemetry import GenerationTelemetry
|
||||||
|
from .activation_compression import (
|
||||||
|
CompressionPolicies,
|
||||||
|
CompressionPolicy,
|
||||||
|
compress_activation,
|
||||||
|
decompress_activation,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class _PipelineCacheMiss(Exception):
|
class _PipelineCacheMiss(Exception):
|
||||||
"""A downstream hop reported 409 cache_miss — head must re-prefill."""
|
"""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 (
|
from .server import (
|
||||||
_WIRE_VERSION,
|
_WIRE_VERSION,
|
||||||
_compress_body,
|
|
||||||
_decompress_body,
|
|
||||||
_parse_shape,
|
_parse_shape,
|
||||||
_validate_activation_body,
|
_validate_activation_body,
|
||||||
)
|
)
|
||||||
@@ -107,6 +123,7 @@ class _RelayHopClient:
|
|||||||
self.relay_addr = relay_addr
|
self.relay_addr = relay_addr
|
||||||
self.timeout = timeout
|
self.timeout = timeout
|
||||||
self._ws = None
|
self._ws = None
|
||||||
|
self._lock = threading.RLock()
|
||||||
|
|
||||||
def request(
|
def request(
|
||||||
self,
|
self,
|
||||||
@@ -118,6 +135,7 @@ class _RelayHopClient:
|
|||||||
|
|
||||||
from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size
|
from .relay_bridge import decode_binary_frame, encode_binary_frame, ws_max_size
|
||||||
|
|
||||||
|
with self._lock:
|
||||||
if self._ws is None:
|
if self._ws is None:
|
||||||
self._ws = wsc.connect(
|
self._ws = wsc.connect(
|
||||||
self.relay_addr,
|
self.relay_addr,
|
||||||
@@ -125,28 +143,48 @@ class _RelayHopClient:
|
|||||||
max_size=ws_max_size(),
|
max_size=ws_max_size(),
|
||||||
compression=None,
|
compression=None,
|
||||||
)
|
)
|
||||||
request_id = f"{time.time_ns():x}"
|
request_id = headers.get("X-Meshnet-Activation-Id") or uuid.uuid4().hex
|
||||||
frame = encode_binary_frame({
|
frame = encode_binary_frame({
|
||||||
"request_id": request_id,
|
"request_id": request_id,
|
||||||
"method": "POST",
|
"method": "POST",
|
||||||
"path": path,
|
"path": path,
|
||||||
"headers": headers,
|
"headers": headers,
|
||||||
}, body)
|
}, 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)
|
self._ws.send(frame)
|
||||||
raw = self._ws.recv(timeout=self.timeout)
|
raw = self._ws.recv(timeout=self.timeout)
|
||||||
if isinstance(raw, (bytes, bytearray)):
|
if isinstance(raw, (bytes, bytearray)):
|
||||||
resp_header, resp_body = decode_binary_frame(bytes(raw))
|
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))
|
status = int(resp_header.get("status", 503))
|
||||||
resp_headers = {k.lower(): v for k, v in (resp_header.get("headers") or {}).items()}
|
resp_headers = {
|
||||||
return status, resp_headers, resp_body
|
k.lower(): v for k, v in (resp_header.get("headers") or {}).items()
|
||||||
|
}
|
||||||
|
else:
|
||||||
resp = json.loads(raw)
|
resp = json.loads(raw)
|
||||||
|
response_id = str(resp.get("request_id") or "")
|
||||||
status = int(resp.get("status", 503))
|
status = int(resp.get("status", 503))
|
||||||
resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()}
|
resp_headers = {
|
||||||
|
k.lower(): v for k, v in (resp.get("headers") or {}).items()
|
||||||
|
}
|
||||||
body_b64 = resp.get("body_base64")
|
body_b64 = resp.get("body_base64")
|
||||||
resp_body = base64.b64decode(body_b64) if body_b64 else (resp.get("body") or "").encode()
|
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
|
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:
|
def close(self) -> None:
|
||||||
|
with self._lock:
|
||||||
if self._ws is not None:
|
if self._ws is not None:
|
||||||
try:
|
try:
|
||||||
self._ws.close()
|
self._ws.close()
|
||||||
@@ -156,6 +194,59 @@ class _RelayHopClient:
|
|||||||
self._ws = None
|
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._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(
|
def _relay_hop(
|
||||||
relay_addr: str,
|
relay_addr: str,
|
||||||
path: str,
|
path: str,
|
||||||
@@ -178,18 +269,15 @@ def _relay_hop(
|
|||||||
client.close()
|
client.close()
|
||||||
|
|
||||||
|
|
||||||
# Below this, zstd overhead outweighs the win (per-token decode bodies are ~KBs).
|
_COMPRESSION_POLICIES = CompressionPolicies()
|
||||||
_COMPRESS_MIN_BYTES = 64 * 1024
|
|
||||||
|
|
||||||
|
|
||||||
def _maybe_compress_activation(body: bytes) -> tuple[bytes, str | None]:
|
def _maybe_compress_activation(
|
||||||
"""zstd-compress large activation bodies; returns (wire_body, encoding)."""
|
body: bytes, policy: CompressionPolicy | None = None,
|
||||||
if len(body) < _COMPRESS_MIN_BYTES:
|
) -> tuple[bytes, str | None]:
|
||||||
return body, None
|
"""Compatibility wrapper for callers that only need wire body and encoding."""
|
||||||
try:
|
result = compress_activation(body, policy or _COMPRESSION_POLICIES.for_condition("lan", "prefill"))
|
||||||
return _compress_body(body, "zstd"), "zstd"
|
return result.body, result.encoding
|
||||||
except Exception:
|
|
||||||
return body, None
|
|
||||||
|
|
||||||
|
|
||||||
def _is_cache_miss_body(body: bytes) -> bool:
|
def _is_cache_miss_body(body: bytes) -> bool:
|
||||||
@@ -285,6 +373,7 @@ class _TorchHTTPServer(http.server.HTTPServer):
|
|||||||
"elapsed_seconds": round(elapsed, 1),
|
"elapsed_seconds": round(elapsed, 1),
|
||||||
"tokens_per_sec": round(tokens / elapsed, 2) if tokens > 0 else 0.0,
|
"tokens_per_sec": round(tokens / elapsed, 2) if tokens > 0 else 0.0,
|
||||||
"routing_complete": bool(rec.get("routing_complete")),
|
"routing_complete": bool(rec.get("routing_complete")),
|
||||||
|
"telemetry": rec["telemetry"].snapshot(now=now) if rec.get("telemetry") else None,
|
||||||
})
|
})
|
||||||
return out
|
return out
|
||||||
|
|
||||||
@@ -306,6 +395,10 @@ class _TorchHTTPServer(http.server.HTTPServer):
|
|||||||
|
|
||||||
|
|
||||||
class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
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
|
def log_message(self, fmt, *args): # noqa: suppress request logs in tests
|
||||||
pass
|
pass
|
||||||
|
|
||||||
@@ -317,6 +410,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
)
|
)
|
||||||
|
|
||||||
def _request_log_suffix(self) -> str:
|
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")
|
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 ""
|
return f" request_id={req_id}" if req_id else ""
|
||||||
|
|
||||||
@@ -334,6 +430,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
"started": time.monotonic(),
|
"started": time.monotonic(),
|
||||||
"tokens": 0,
|
"tokens": 0,
|
||||||
"routing_complete": False,
|
"routing_complete": False,
|
||||||
|
"telemetry": None,
|
||||||
}
|
}
|
||||||
|
|
||||||
def _track_request_progress(
|
def _track_request_progress(
|
||||||
@@ -366,6 +463,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self._handle_chat_completions()
|
self._handle_chat_completions()
|
||||||
else:
|
else:
|
||||||
self.send_response(404)
|
self.send_response(404)
|
||||||
|
self.send_header("Content-Length", "0")
|
||||||
self.end_headers()
|
self.end_headers()
|
||||||
|
|
||||||
def _handle_infer(self) -> None:
|
def _handle_infer(self) -> None:
|
||||||
@@ -427,7 +525,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
encoding = self.headers.get("X-Meshnet-Encoding")
|
encoding = self.headers.get("X-Meshnet-Encoding")
|
||||||
length = int(self.headers.get("Content-Length", 0))
|
length = int(self.headers.get("Content-Length", 0))
|
||||||
body = self.rfile.read(length)
|
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)
|
_validate_activation_body(raw_body, shape, dtype)
|
||||||
if dtype != "bfloat16":
|
if dtype != "bfloat16":
|
||||||
raise ValueError("real model backend requires bfloat16 activation input")
|
raise ValueError("real model backend requires bfloat16 activation input")
|
||||||
@@ -438,6 +536,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
except (KeyError, ValueError, TypeError):
|
except (KeyError, ValueError, TypeError):
|
||||||
self.send_response(400)
|
self.send_response(400)
|
||||||
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
|
self.send_header("X-Meshnet-Wire", _WIRE_VERSION)
|
||||||
|
self.send_header("Content-Length", "0")
|
||||||
self.end_headers()
|
self.end_headers()
|
||||||
return
|
return
|
||||||
|
|
||||||
@@ -511,7 +610,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self._send_json(200, {"text": result})
|
self._send_json(200, {"text": result})
|
||||||
return
|
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_response(200)
|
||||||
self.send_header("Content-Type", "application/octet-stream")
|
self.send_header("Content-Type", "application/octet-stream")
|
||||||
self.send_header("Content-Length", str(len(response_body)))
|
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-Session", session)
|
||||||
self.send_header("X-Meshnet-Chunk-Index", chunk_index)
|
self.send_header("X-Meshnet-Chunk-Index", chunk_index)
|
||||||
self.send_header("X-Meshnet-Chunk-Total", chunk_total)
|
self.send_header("X-Meshnet-Chunk-Total", chunk_total)
|
||||||
if encoding:
|
if response_compression.encoding:
|
||||||
self.send_header("X-Meshnet-Encoding", encoding)
|
self.send_header("X-Meshnet-Encoding", response_compression.encoding)
|
||||||
if result.attention_mask_header:
|
if result.attention_mask_header:
|
||||||
self.send_header("X-Meshnet-Attn-Mask", result.attention_mask_header)
|
self.send_header("X-Meshnet-Attn-Mask", result.attention_mask_header)
|
||||||
if result.position_ids_header:
|
if result.position_ids_header:
|
||||||
@@ -700,6 +804,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
current_text = prompt_text
|
current_text = prompt_text
|
||||||
|
|
||||||
session_id = str(uuid.uuid4())
|
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))
|
use_kv = bool(getattr(backend, "supports_kv_cache", False))
|
||||||
# EOS detection by id must work on the stateless path too: the tail
|
# EOS detection by id must work on the stateless path too: the tail
|
||||||
# returns token_id regardless of caching, and EOS usually decodes to
|
# 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
|
last_token_id: int | None = None
|
||||||
failure_reason: str | None = None
|
failure_reason: str | None = None
|
||||||
relay_clients: dict[str, _RelayHopClient] = {}
|
relay_clients: dict[str, _RelayHopClient] = {}
|
||||||
|
direct_clients: dict[str, _DirectHopClient] = {}
|
||||||
|
|
||||||
def _prefill_step() -> tuple[str, int | None]:
|
def _prefill_step() -> tuple[str, int | None]:
|
||||||
"""Full-sequence prefill: initial step and cache-miss recovery."""
|
"""Full-sequence prefill: initial step and cache-miss recovery."""
|
||||||
@@ -734,6 +844,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
payload, remaining_route, backend=backend,
|
payload, remaining_route, backend=backend,
|
||||||
session=session_id, cache_mode="prefill" if use_kv else None,
|
session=session_id, cache_mode="prefill" if use_kv else None,
|
||||||
relay_clients=relay_clients,
|
relay_clients=relay_clients,
|
||||||
|
direct_clients=direct_clients if use_kv else None,
|
||||||
|
telemetry=telemetry,
|
||||||
)
|
)
|
||||||
|
|
||||||
for step in range(max_tokens):
|
for step in range(max_tokens):
|
||||||
@@ -745,6 +857,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
payload, remaining_route, backend=backend,
|
payload, remaining_route, backend=backend,
|
||||||
session=session_id, cache_mode="decode",
|
session=session_id, cache_mode="decode",
|
||||||
relay_clients=relay_clients,
|
relay_clients=relay_clients,
|
||||||
|
direct_clients=direct_clients,
|
||||||
|
telemetry=telemetry,
|
||||||
)
|
)
|
||||||
except (KVCacheMiss, _PipelineCacheMiss) as miss:
|
except (KVCacheMiss, _PipelineCacheMiss) as miss:
|
||||||
# Evicted/restarted node or head lost its own session:
|
# Evicted/restarted node or head lost its own session:
|
||||||
@@ -758,11 +872,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
else:
|
else:
|
||||||
token_str, token_id = _prefill_step()
|
token_str, token_id = _prefill_step()
|
||||||
except _PipelineCacheMiss as exc:
|
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}"
|
failure_reason = f"cache miss on prefill: {exc}"
|
||||||
break
|
break
|
||||||
except Exception as exc:
|
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}"
|
failure_reason = f"distributed encode error: {exc}"
|
||||||
break
|
break
|
||||||
# Stop on error responses or EOS.
|
# Stop on error responses or EOS.
|
||||||
@@ -789,14 +903,32 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
tokens=len(generated),
|
tokens=len(generated),
|
||||||
routing_complete=True,
|
routing_complete=True,
|
||||||
)
|
)
|
||||||
|
telemetry.note_tokens(len(generated))
|
||||||
now = time.monotonic()
|
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:
|
if step == 0 or now - last_gen_log >= _GENERATION_LOG_INTERVAL:
|
||||||
elapsed = now - gen_started
|
elapsed = now - gen_started
|
||||||
token_count = len(generated)
|
token_count = len(generated)
|
||||||
tps = token_count / max(elapsed, 1e-6)
|
tps = token_count / max(elapsed, 1e-6)
|
||||||
_write_progress_line(
|
_write_progress_line(
|
||||||
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}",
|
f"tokens={token_count} elapsed_s={elapsed:.1f} tps={tps:.2f}",
|
||||||
)
|
)
|
||||||
last_gen_log = now
|
last_gen_log = now
|
||||||
@@ -808,6 +940,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
pass
|
pass
|
||||||
for relay_client in relay_clients.values():
|
for relay_client in relay_clients.values():
|
||||||
relay_client.close()
|
relay_client.close()
|
||||||
|
for direct_client in direct_clients.values():
|
||||||
|
direct_client.close()
|
||||||
|
|
||||||
if generated:
|
if generated:
|
||||||
elapsed = time.monotonic() - gen_started
|
elapsed = time.monotonic() - gen_started
|
||||||
@@ -815,10 +949,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
tps = token_count / max(elapsed, 1e-6)
|
tps = token_count / max(elapsed, 1e-6)
|
||||||
_write_progress_line(
|
_write_progress_line(
|
||||||
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}",
|
f"elapsed_s={elapsed:.1f} tps={tps:.2f}",
|
||||||
final=True,
|
final=True,
|
||||||
)
|
)
|
||||||
|
telemetry.close()
|
||||||
|
|
||||||
result_text = "".join(generated)
|
result_text = "".join(generated)
|
||||||
# A failure before the first token is an upstream error, not an empty
|
# 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,
|
session: str | None = None,
|
||||||
cache_mode: str | None = None,
|
cache_mode: str | None = None,
|
||||||
relay_clients: dict[str, _RelayHopClient] | None = None,
|
relay_clients: dict[str, _RelayHopClient] | None = None,
|
||||||
|
direct_clients: dict[str, _DirectHopClient] | None = None,
|
||||||
|
telemetry: GenerationTelemetry | None = None,
|
||||||
) -> tuple[str, int | None]:
|
) -> tuple[str, int | None]:
|
||||||
"""Forward an activation through the downstream route.
|
"""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.
|
# Full single-node (head+tail) is handled before entering this method.
|
||||||
if active_backend.is_tail:
|
if active_backend.is_tail:
|
||||||
try:
|
try:
|
||||||
tensor = active_backend.torch.frombuffer(
|
tensor = _tensor_from_bfloat16_bytes(
|
||||||
bytearray(payload.body), # type: ignore[union-attr]
|
payload.body, payload.shape, active_backend.torch, # type: ignore[union-attr]
|
||||||
dtype=active_backend.torch.bfloat16,
|
).to(active_backend.device)
|
||||||
).reshape(payload.shape).to(active_backend.device) # type: ignore[union-attr]
|
|
||||||
if hasattr(active_backend, "decode_tail_token"):
|
if hasattr(active_backend, "decode_tail_token"):
|
||||||
tail = active_backend.decode_tail_token(tensor)
|
tail = active_backend.decode_tail_token(tensor)
|
||||||
return tail.text, tail.token_id
|
return tail.text, tail.token_id
|
||||||
@@ -968,7 +1104,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
+ (f" relay={relay_addr}" if relay_addr else ""),
|
+ (f" relay={relay_addr}" if relay_addr else ""),
|
||||||
flush=True,
|
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] = {
|
headers: dict[str, str] = {
|
||||||
"Content-Type": "application/octet-stream",
|
"Content-Type": "application/octet-stream",
|
||||||
"X-Meshnet-Wire": _WIRE_VERSION,
|
"X-Meshnet-Wire": _WIRE_VERSION,
|
||||||
@@ -979,9 +1120,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
"X-Meshnet-Chunk-Total": "1",
|
"X-Meshnet-Chunk-Total": "1",
|
||||||
"X-Meshnet-Hop-Index": str(hop_index),
|
"X-Meshnet-Hop-Index": str(hop_index),
|
||||||
"X-Meshnet-Start-Layer": str(start_layer),
|
"X-Meshnet-Start-Layer": str(start_layer),
|
||||||
|
"X-Meshnet-Activation-Id": uuid.uuid4().hex,
|
||||||
|
"X-Meshnet-Compression-Route": route_condition,
|
||||||
}
|
}
|
||||||
if wire_encoding:
|
if wire_encoding:
|
||||||
headers["X-Meshnet-Encoding"] = 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:
|
if cache_mode:
|
||||||
headers["X-Meshnet-Cache"] = cache_mode
|
headers["X-Meshnet-Cache"] = cache_mode
|
||||||
past_len = getattr(payload, "past_len", None)
|
past_len = getattr(payload, "past_len", None)
|
||||||
@@ -992,6 +1141,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
if current_pos:
|
if current_pos:
|
||||||
headers["X-Meshnet-Position-Ids"] = current_pos
|
headers["X-Meshnet-Position-Ids"] = current_pos
|
||||||
if relay_addr:
|
if relay_addr:
|
||||||
|
connection_reused = bool(
|
||||||
|
relay_clients is not None and relay_addr in relay_clients
|
||||||
|
)
|
||||||
|
seam_started = time.monotonic()
|
||||||
try:
|
try:
|
||||||
if relay_clients is None:
|
if relay_clients is None:
|
||||||
status, resp_headers, resp_body = _relay_hop(
|
status, resp_headers, resp_body = _relay_hop(
|
||||||
@@ -1004,44 +1157,100 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
status, resp_headers, resp_body = relay_client.request(
|
status, resp_headers, resp_body = relay_client.request(
|
||||||
"/forward", wire_body, headers,
|
"/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):
|
if status == 409 and _is_cache_miss_body(resp_body):
|
||||||
raise _PipelineCacheMiss(node_url)
|
raise _PipelineCacheMiss(node_url)
|
||||||
if status >= 400:
|
if status >= 400:
|
||||||
detail = _response_error_snippet(resp_body)
|
detail = _response_error_snippet(resp_body)
|
||||||
print(
|
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,
|
flush=True,
|
||||||
)
|
)
|
||||||
return f"pipeline error at {node_url} via relay: status {status}: {detail}", None
|
return f"pipeline error at {node_url} via relay: status {status}: {detail}", None
|
||||||
except _PipelineCacheMiss:
|
except _PipelineCacheMiss:
|
||||||
raise
|
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:
|
except Exception as exc:
|
||||||
print(
|
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}",
|
f"falling back to direct {node_url}",
|
||||||
flush=True,
|
flush=True,
|
||||||
)
|
)
|
||||||
relay_addr = None # fall through to direct
|
relay_addr = None # fall through to direct
|
||||||
if not relay_addr:
|
if not relay_addr:
|
||||||
req = urllib.request.Request(
|
connection_reused = bool(
|
||||||
f"{node_url}/forward",
|
direct_clients is not None and node_url in direct_clients
|
||||||
data=wire_body,
|
|
||||||
headers=headers,
|
|
||||||
method="POST",
|
|
||||||
)
|
)
|
||||||
|
seam_started = time.monotonic()
|
||||||
try:
|
try:
|
||||||
with urllib.request.urlopen(req, timeout=120.0) as r:
|
if direct_clients is None:
|
||||||
resp_body = r.read()
|
direct_client = _DirectHopClient(node_url, timeout=120.0)
|
||||||
resp_headers = {k.lower(): v for k, v in r.headers.items()}
|
try:
|
||||||
except urllib.error.HTTPError as exc:
|
status, resp_headers, resp_body = direct_client.request(
|
||||||
body = exc.read()
|
"/forward", wire_body, headers,
|
||||||
if exc.code == 409 and _is_cache_miss_body(body):
|
)
|
||||||
raise _PipelineCacheMiss(node_url) from exc
|
finally:
|
||||||
detail = _response_error_snippet(body)
|
direct_client.close()
|
||||||
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}: {detail}", flush=True)
|
else:
|
||||||
return f"pipeline error at {node_url}: {exc}: {detail}", None
|
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:
|
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
|
return f"pipeline error at {node_url}: {exc}", None
|
||||||
content_type = resp_headers.get("content-type", "")
|
content_type = resp_headers.get("content-type", "")
|
||||||
if "application/json" in 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))
|
shape_header = resp_headers.get("x-meshnet-shape", ",".join(str(d) for d in current_shape))
|
||||||
current_shape = _parse_shape(shape_header)
|
current_shape = _parse_shape(shape_header)
|
||||||
try:
|
try:
|
||||||
current_body = _decompress_body(
|
decompression = decompress_activation(
|
||||||
resp_body, resp_headers.get("x-meshnet-encoding")
|
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:
|
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
|
return f"pipeline error at {node_url}: {exc}", None
|
||||||
current_attn = resp_headers.get("x-meshnet-attn-mask")
|
current_attn = resp_headers.get("x-meshnet-attn-mask")
|
||||||
current_pos = resp_headers.get("x-meshnet-position-ids")
|
current_pos = resp_headers.get("x-meshnet-position-ids")
|
||||||
@@ -1084,14 +1303,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self.send_response(200)
|
self.send_response(200)
|
||||||
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
|
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
|
||||||
self.send_header("Cache-Control", "no-cache")
|
self.send_header("Cache-Control", "no-cache")
|
||||||
|
self.send_header("Transfer-Encoding", "chunked")
|
||||||
self.end_headers()
|
self.end_headers()
|
||||||
|
|
||||||
def _emit(data: str) -> None:
|
def _emit(data: str) -> bool:
|
||||||
try:
|
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()
|
self.wfile.flush()
|
||||||
|
return True
|
||||||
except (BrokenPipeError, ConnectionResetError):
|
except (BrokenPipeError, ConnectionResetError):
|
||||||
pass
|
return False
|
||||||
|
|
||||||
_emit(json.dumps({
|
_emit(json.dumps({
|
||||||
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
||||||
@@ -1105,7 +1327,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
# instead of showing an empty completion.
|
# instead of showing an empty completion.
|
||||||
_emit(json.dumps({"error": {"message": error, "type": "upstream_error"}}))
|
_emit(json.dumps({"error": {"message": error, "type": "upstream_error"}}))
|
||||||
try:
|
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()
|
self.wfile.flush()
|
||||||
except (BrokenPipeError, ConnectionResetError):
|
except (BrokenPipeError, ConnectionResetError):
|
||||||
pass
|
pass
|
||||||
@@ -1117,7 +1339,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||||
}))
|
}))
|
||||||
try:
|
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()
|
self.wfile.flush()
|
||||||
except (BrokenPipeError, ConnectionResetError):
|
except (BrokenPipeError, ConnectionResetError):
|
||||||
pass
|
pass
|
||||||
@@ -1159,14 +1381,17 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self.send_response(200)
|
self.send_response(200)
|
||||||
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
|
self.send_header("Content-Type", "text/event-stream; charset=utf-8")
|
||||||
self.send_header("Cache-Control", "no-cache")
|
self.send_header("Cache-Control", "no-cache")
|
||||||
|
self.send_header("Transfer-Encoding", "chunked")
|
||||||
self.end_headers()
|
self.end_headers()
|
||||||
|
|
||||||
def _emit(data: str) -> None:
|
def _emit(data: str) -> bool:
|
||||||
try:
|
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()
|
self.wfile.flush()
|
||||||
except BrokenPipeError:
|
return True
|
||||||
pass
|
except (BrokenPipeError, ConnectionResetError):
|
||||||
|
return False
|
||||||
|
|
||||||
_emit(json.dumps({
|
_emit(json.dumps({
|
||||||
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
"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"}],
|
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||||
}))
|
}))
|
||||||
try:
|
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()
|
self.wfile.flush()
|
||||||
except BrokenPipeError:
|
except (BrokenPipeError, ConnectionResetError):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
@@ -1255,6 +1480,7 @@ class TorchNodeServer:
|
|||||||
debug: bool = False,
|
debug: bool = False,
|
||||||
max_loaded_shards: int = 1,
|
max_loaded_shards: int = 1,
|
||||||
force_cpu: bool = False,
|
force_cpu: bool = False,
|
||||||
|
recipe_params: Mapping[str, Any] | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self._host = host
|
self._host = host
|
||||||
self._requested_port = port
|
self._requested_port = port
|
||||||
@@ -1266,6 +1492,7 @@ class TorchNodeServer:
|
|||||||
quantization,
|
quantization,
|
||||||
cache_dir,
|
cache_dir,
|
||||||
force_cpu=force_cpu,
|
force_cpu=force_cpu,
|
||||||
|
recipe_params=recipe_params,
|
||||||
)
|
)
|
||||||
self._backends: dict[str, TorchModelShard] = {self._backend.model_id: self._backend}
|
self._backends: dict[str, TorchModelShard] = {self._backend.model_id: self._backend}
|
||||||
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
|
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
|
||||||
@@ -1415,13 +1642,20 @@ def _load_backend(
|
|||||||
quantization: str,
|
quantization: str,
|
||||||
cache_dir: Path | None = None,
|
cache_dir: Path | None = None,
|
||||||
force_cpu: bool = False,
|
force_cpu: bool = False,
|
||||||
|
recipe_params: Mapping[str, Any] | None = None,
|
||||||
) -> TorchModelShard:
|
) -> TorchModelShard:
|
||||||
from .model_backend import load_torch_shard
|
from .model_backend import load_torch_shard
|
||||||
|
|
||||||
quant = validate_quantization(quantization)
|
quant = validate_quantization(quantization)
|
||||||
try:
|
try:
|
||||||
return load_torch_shard(
|
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:
|
except MissingModelDependencyError:
|
||||||
raise
|
raise
|
||||||
|
|||||||
@@ -46,11 +46,12 @@ def ws_max_size() -> int | None:
|
|||||||
# inflation, no JSON re-encode of megabytes per hop. Text JSON frames remain the
|
# inflation, no JSON re-encode of megabytes per hop. Text JSON frames remain the
|
||||||
# control plane (gossip, peer-register, streamed SSE responses).
|
# control plane (gossip, peer-register, streamed SSE responses).
|
||||||
BINARY_FRAME_MAGIC = b"MRF1"
|
BINARY_FRAME_MAGIC = b"MRF1"
|
||||||
|
_RPC_PEER_DISCONNECTED = object()
|
||||||
|
|
||||||
|
|
||||||
def encode_binary_frame(header: dict, body: bytes) -> bytes:
|
def encode_binary_frame(header: dict, body: bytes) -> bytes:
|
||||||
header_bytes = json.dumps(header, separators=(",", ":")).encode()
|
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]:
|
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")
|
raise ValueError("not a meshnet binary relay frame")
|
||||||
header_len = int.from_bytes(frame[4:8], "big")
|
header_len = int.from_bytes(frame[4:8], "big")
|
||||||
header = json.loads(frame[8:8 + header_len].decode())
|
header = json.loads(frame[8:8 + header_len].decode())
|
||||||
return header, bytes(frame[8 + header_len:])
|
return header, frame[8 + header_len:]
|
||||||
|
|
||||||
|
|
||||||
class RelayServer:
|
class RelayServer:
|
||||||
@@ -79,12 +80,16 @@ class RelayServer:
|
|||||||
ssl_cert: Path | None = None,
|
ssl_cert: Path | None = None,
|
||||||
ssl_key: Path | None = None,
|
ssl_key: Path | None = None,
|
||||||
max_peers: int = 500,
|
max_peers: int = 500,
|
||||||
|
rpc_timeout: float = 310.0,
|
||||||
|
rpc_idle_timeout: float = 120.0,
|
||||||
):
|
):
|
||||||
self.host = host
|
self.host = host
|
||||||
self.port = port
|
self.port = port
|
||||||
self.ssl_cert = ssl_cert
|
self.ssl_cert = ssl_cert
|
||||||
self.ssl_key = ssl_key
|
self.ssl_key = ssl_key
|
||||||
self.max_peers = max_peers
|
self.max_peers = max_peers
|
||||||
|
self.rpc_timeout = rpc_timeout
|
||||||
|
self.rpc_idle_timeout = rpc_idle_timeout
|
||||||
|
|
||||||
self._registry = PeerRegistry()
|
self._registry = PeerRegistry()
|
||||||
self._loop: asyncio.AbstractEventLoop | None = None
|
self._loop: asyncio.AbstractEventLoop | None = None
|
||||||
@@ -94,9 +99,10 @@ class RelayServer:
|
|||||||
self._ready = threading.Event()
|
self._ready = threading.Event()
|
||||||
self._running = False
|
self._running = False
|
||||||
self._stop_event: asyncio.Event | None = None
|
self._stop_event: asyncio.Event | None = None
|
||||||
# request_id → queue of relay-http-response frames (US-036: a streamed
|
# relay request id → (target peer, requester request id, response queue).
|
||||||
# response is a sequence of frames; a frame without "stream" is terminal).
|
# The relay-generated id prevents two Route Sessions using the same
|
||||||
self._pending_rpc: dict[str, asyncio.Queue] = {}
|
# legacy request_id from overwriting each other's pending response.
|
||||||
|
self._pending_rpc: dict[str, tuple[str, str, asyncio.Queue]] = {}
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def registry(self) -> PeerRegistry:
|
def registry(self) -> PeerRegistry:
|
||||||
@@ -118,6 +124,12 @@ class RelayServer:
|
|||||||
if self._thread:
|
if self._thread:
|
||||||
self._thread.join(timeout=3.0)
|
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:
|
def _run(self) -> None:
|
||||||
asyncio.set_event_loop(self._loop)
|
asyncio.set_event_loop(self._loop)
|
||||||
self._loop.run_until_complete(self._serve())
|
self._loop.run_until_complete(self._serve())
|
||||||
@@ -192,9 +204,9 @@ class RelayServer:
|
|||||||
header, _ = decode_binary_frame(bytes(raw))
|
header, _ = decode_binary_frame(bytes(raw))
|
||||||
except (ValueError, json.JSONDecodeError):
|
except (ValueError, json.JSONDecodeError):
|
||||||
continue
|
continue
|
||||||
queue = self._pending_rpc.get(header.get("request_id"))
|
pending = self._pending_rpc.get(header.get("request_id"))
|
||||||
if queue is not None:
|
if pending is not None:
|
||||||
queue.put_nowait(bytes(raw))
|
pending[2].put_nowait(bytes(raw))
|
||||||
continue
|
continue
|
||||||
try:
|
try:
|
||||||
envelope = json.loads(raw)
|
envelope = json.loads(raw)
|
||||||
@@ -226,9 +238,9 @@ class RelayServer:
|
|||||||
if topic == "relay-http-response":
|
if topic == "relay-http-response":
|
||||||
payload = envelope.get("payload", {})
|
payload = envelope.get("payload", {})
|
||||||
request_id = payload.get("request_id")
|
request_id = payload.get("request_id")
|
||||||
queue = self._pending_rpc.get(request_id)
|
pending = self._pending_rpc.get(request_id)
|
||||||
if queue is not None:
|
if pending is not None:
|
||||||
queue.put_nowait(payload)
|
pending[2].put_nowait(payload)
|
||||||
continue
|
continue
|
||||||
|
|
||||||
# Fan out to all other registered peers
|
# Fan out to all other registered peers
|
||||||
@@ -242,6 +254,9 @@ class RelayServer:
|
|||||||
finally:
|
finally:
|
||||||
if peer_id:
|
if peer_id:
|
||||||
self._registry.unregister(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)
|
log.debug("Peer unregistered: %s", peer_id)
|
||||||
|
|
||||||
async def _handle_circuit_relay(self, ws_requester, target_peer_id: str) -> None:
|
async def _handle_circuit_relay(self, ws_requester, target_peer_id: str) -> None:
|
||||||
@@ -290,18 +305,19 @@ class RelayServer:
|
|||||||
except Exception:
|
except Exception:
|
||||||
return
|
return
|
||||||
|
|
||||||
request_id: str | None = None
|
requester_request_id: str | None = None
|
||||||
|
relay_request_id = uuid.uuid4().hex
|
||||||
try:
|
try:
|
||||||
if isinstance(raw, (bytes, bytearray)):
|
if isinstance(raw, (bytes, bytearray)):
|
||||||
header, body = decode_binary_frame(bytes(raw))
|
header, body = decode_binary_frame(bytes(raw))
|
||||||
request_id = str(header.get("request_id") or uuid.uuid4())
|
requester_request_id = str(header.get("request_id") or uuid.uuid4())
|
||||||
header["request_id"] = request_id
|
header["request_id"] = relay_request_id
|
||||||
header["target_peer"] = target_peer_id
|
header["target_peer"] = target_peer_id
|
||||||
outbound: str | bytes = encode_binary_frame(header, body)
|
outbound: str | bytes = encode_binary_frame(header, body)
|
||||||
else:
|
else:
|
||||||
payload = json.loads(raw)
|
payload = json.loads(raw)
|
||||||
request_id = str(payload.get("request_id") or uuid.uuid4())
|
requester_request_id = str(payload.get("request_id") or uuid.uuid4())
|
||||||
payload["request_id"] = request_id
|
payload["request_id"] = relay_request_id
|
||||||
payload["target_peer"] = target_peer_id
|
payload["target_peer"] = target_peer_id
|
||||||
outbound = json.dumps({
|
outbound = json.dumps({
|
||||||
"topic": "relay-http-request",
|
"topic": "relay-http-request",
|
||||||
@@ -314,16 +330,18 @@ class RelayServer:
|
|||||||
return
|
return
|
||||||
|
|
||||||
queue: asyncio.Queue = asyncio.Queue()
|
queue: asyncio.Queue = asyncio.Queue()
|
||||||
self._pending_rpc[request_id] = queue
|
self._pending_rpc[relay_request_id] = (
|
||||||
overall_timeout = 310.0
|
target_peer_id, requester_request_id, queue,
|
||||||
idle_timeout = 120.0
|
)
|
||||||
|
overall_timeout = self.rpc_timeout
|
||||||
|
idle_timeout = self.rpc_idle_timeout
|
||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
deadline = loop.time() + overall_timeout
|
deadline = loop.time() + overall_timeout
|
||||||
target = self._registry.get(target_peer_id)
|
target = self._registry.get(target_peer_id)
|
||||||
try:
|
try:
|
||||||
if target is None:
|
if target is None:
|
||||||
await ws_requester.send(json.dumps({
|
await ws_requester.send(json.dumps({
|
||||||
"request_id": request_id,
|
"request_id": requester_request_id,
|
||||||
"status": 503,
|
"status": 503,
|
||||||
"headers": {"Content-Type": "application/json"},
|
"headers": {"Content-Type": "application/json"},
|
||||||
"body": json.dumps({"error": f"peer {target_peer_id!r} disconnected"}),
|
"body": json.dumps({"error": f"peer {target_peer_id!r} disconnected"}),
|
||||||
@@ -339,21 +357,34 @@ class RelayServer:
|
|||||||
frame = await asyncio.wait_for(
|
frame = await asyncio.wait_for(
|
||||||
queue.get(), timeout=min(idle_timeout, remaining)
|
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)):
|
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
|
break
|
||||||
await ws_requester.send(json.dumps(frame))
|
response = dict(frame)
|
||||||
if not frame.get("stream") or frame.get("done"):
|
response["request_id"] = requester_request_id
|
||||||
|
await ws_requester.send(json.dumps(response))
|
||||||
|
if not response.get("stream") or response.get("done"):
|
||||||
break
|
break
|
||||||
except asyncio.TimeoutError:
|
except asyncio.TimeoutError:
|
||||||
await ws_requester.send(json.dumps({
|
await ws_requester.send(json.dumps({
|
||||||
"request_id": request_id,
|
"request_id": requester_request_id,
|
||||||
"status": 504,
|
"status": 504,
|
||||||
"headers": {"Content-Type": "application/json"},
|
"headers": {"Content-Type": "application/json"},
|
||||||
"body": json.dumps({"error": "relay rpc timed out"}),
|
"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:
|
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:
|
async def _broadcast(raw: str | bytes, peers: list) -> None:
|
||||||
|
|||||||
415
packages/tracker/meshnet_tracker/capability.py
Normal file
415
packages/tracker/meshnet_tracker/capability.py
Normal file
@@ -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)
|
||||||
@@ -9,6 +9,7 @@ from pathlib import Path
|
|||||||
|
|
||||||
from .accounts import DEFAULT_ACCOUNTS_DB_PATH
|
from .accounts import DEFAULT_ACCOUNTS_DB_PATH
|
||||||
from .billing import DEFAULT_BILLING_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 .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH
|
||||||
from .logging_setup import (
|
from .logging_setup import (
|
||||||
DEFAULT_LOG_BACKUP_COUNT,
|
DEFAULT_LOG_BACKUP_COUNT,
|
||||||
@@ -105,6 +106,18 @@ def main() -> None:
|
|||||||
metavar="PATH",
|
metavar="PATH",
|
||||||
help="SQLite database path for persistent model usage statistics",
|
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(
|
common.add_argument(
|
||||||
"--relay-url",
|
"--relay-url",
|
||||||
default=None,
|
default=None,
|
||||||
@@ -416,6 +429,7 @@ def main() -> None:
|
|||||||
cluster_peers=cluster_peers or None,
|
cluster_peers=cluster_peers or None,
|
||||||
cluster_self_url=args.self_url,
|
cluster_self_url=args.self_url,
|
||||||
stats_db=getattr(args, "stats_db", None),
|
stats_db=getattr(args, "stats_db", None),
|
||||||
|
capability_policy=getattr(args, "capability_policy", None),
|
||||||
relay_url=relay_url,
|
relay_url=relay_url,
|
||||||
embedded_relay=args.embedded_relay,
|
embedded_relay=args.embedded_relay,
|
||||||
embedded_relay_host=args.relay_host,
|
embedded_relay_host=args.relay_host,
|
||||||
|
|||||||
@@ -6,10 +6,14 @@ HTTP API contract:
|
|||||||
"endpoint": "http://host:port", "shard_start": int, "shard_end": int,
|
"endpoint": "http://host:port", "shard_start": int, "shard_end": int,
|
||||||
"model": str optional, "shard_checksum": str optional,
|
"model": str optional, "shard_checksum": str optional,
|
||||||
"hardware_profile": object, "wallet_address": 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}
|
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
|
- POST /v1/nodes/{node_id}/heartbeat
|
||||||
Response 200: {}
|
Response 200: {}
|
||||||
Response 404: {"error": "node not found"}
|
Response 404: {"error": "node not found"}
|
||||||
@@ -48,6 +52,20 @@ from typing import Any
|
|||||||
|
|
||||||
from .accounts import DEFAULT_ACCOUNTS_DB_PATH, AccountStore
|
from .accounts import DEFAULT_ACCOUNTS_DB_PATH, AccountStore
|
||||||
from .auth import is_validator_token, sign_hive_request, verify_hive_request
|
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 .wallet_proof import binding_message, verify_wallet_signature
|
||||||
from .billing import DEFAULT_BILLING_DB_PATH, BillingLedger
|
from .billing import DEFAULT_BILLING_DB_PATH, BillingLedger
|
||||||
from .calibration import DEFAULT_CALIBRATION_DB_PATH, ToplocCalibrationStore
|
from .calibration import DEFAULT_CALIBRATION_DB_PATH, ToplocCalibrationStore
|
||||||
@@ -587,6 +605,8 @@ class _NodeEntry:
|
|||||||
"heartbeats_expected", "heartbeats_received",
|
"heartbeats_expected", "heartbeats_received",
|
||||||
# dynamic reassignment queued by the tracker
|
# dynamic reassignment queued by the tracker
|
||||||
"pending_new_assignment",
|
"pending_new_assignment",
|
||||||
|
# the tracker's verdict on the capability proof this node presented
|
||||||
|
"capability",
|
||||||
)
|
)
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
@@ -616,6 +636,7 @@ class _NodeEntry:
|
|||||||
cert_fingerprint: str | None = None,
|
cert_fingerprint: str | None = None,
|
||||||
peer_id: str | None = None,
|
peer_id: str | None = None,
|
||||||
friendly_name: str | None = None,
|
friendly_name: str | None = None,
|
||||||
|
capability: "CapabilityState | None" = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
self.node_id = node_id
|
self.node_id = node_id
|
||||||
self.endpoint = endpoint
|
self.endpoint = endpoint
|
||||||
@@ -643,6 +664,9 @@ class _NodeEntry:
|
|||||||
self.cert_fingerprint = cert_fingerprint
|
self.cert_fingerprint = cert_fingerprint
|
||||||
self.peer_id = peer_id
|
self.peer_id = peer_id
|
||||||
self.friendly_name = friendly_name
|
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.pending_directives: list[dict] = []
|
||||||
self.last_heartbeat: float = time.monotonic()
|
self.last_heartbeat: float = time.monotonic()
|
||||||
self.total_requests: int = 0
|
self.total_requests: int = 0
|
||||||
@@ -731,12 +755,88 @@ def _reputation_multiplier(node: "_NodeEntry", contracts: Any | None) -> float:
|
|||||||
return 0.5 + 0.5 * reputation
|
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(
|
def _select_route(
|
||||||
nodes: list[_NodeEntry],
|
nodes: list[_NodeEntry],
|
||||||
required_start: int,
|
required_start: int,
|
||||||
required_end: int,
|
required_end: int,
|
||||||
model: str | None = None,
|
model: str | None = None,
|
||||||
contracts: Any | None = None,
|
contracts: Any | None = None,
|
||||||
|
policy: str | None = None,
|
||||||
) -> tuple[list[_NodeEntry], str]:
|
) -> tuple[list[_NodeEntry], str]:
|
||||||
"""Greedy interval-cover biased toward fast, lightly-loaded, reputable nodes.
|
"""Greedy interval-cover biased toward fast, lightly-loaded, reputable nodes.
|
||||||
|
|
||||||
@@ -747,7 +847,10 @@ def _select_route(
|
|||||||
Tiebreak: higher shard_end (fewer hops).
|
Tiebreak: higher shard_end (fewer hops).
|
||||||
"""
|
"""
|
||||||
candidates = sorted(
|
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]
|
key=lambda n: (n.shard_start, -n.shard_end), # type: ignore[operator]
|
||||||
)
|
)
|
||||||
route: list[_NodeEntry] = []
|
route: list[_NodeEntry] = []
|
||||||
@@ -783,6 +886,7 @@ def _enumerate_routes(
|
|||||||
model: str | None = None,
|
model: str | None = None,
|
||||||
contracts: Any | None = None,
|
contracts: Any | None = None,
|
||||||
max_candidates: int = 8,
|
max_candidates: int = 8,
|
||||||
|
policy: str | None = None,
|
||||||
) -> list["RouteCandidate"]:
|
) -> list["RouteCandidate"]:
|
||||||
"""Enumerate viable route candidates for bandit selection (ADR-0021).
|
"""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
|
with the longest-advancing hops. The route's prior throughput estimate is
|
||||||
its bottleneck hop's queue-adjusted effective throughput — used only until
|
its bottleneck hop's queue-adjusted effective throughput — used only until
|
||||||
observed route samples exist.
|
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 = [
|
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
|
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
|
# 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,
|
route_stats: "RouteStatsStore | None" = None,
|
||||||
relay_status: dict | None = None,
|
relay_status: dict | None = None,
|
||||||
test_runner: "TestRunManager | None" = None,
|
test_runner: "TestRunManager | None" = None,
|
||||||
|
capability_policy: str | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
super().__init__(addr, handler)
|
super().__init__(addr, handler)
|
||||||
self.registry = registry
|
self.registry = registry
|
||||||
|
self.capability_policy = normalize_policy(
|
||||||
|
capability_policy if capability_policy is not None else policy_from_env()
|
||||||
|
)
|
||||||
self.lock = lock
|
self.lock = lock
|
||||||
self.heartbeat_timeout = heartbeat_timeout
|
self.heartbeat_timeout = heartbeat_timeout
|
||||||
self.model_presets = model_presets
|
self.model_presets = model_presets
|
||||||
@@ -3177,9 +3289,19 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
def model_supply_for(node: _NodeEntry) -> dict:
|
def model_supply_for(node: _NodeEntry) -> dict:
|
||||||
return _model_health_summary(server, node.model, node.hf_repo)
|
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, {
|
self._send_json(200, {
|
||||||
"relay_url": server.relay_url,
|
"relay_url": server.relay_url,
|
||||||
"relay": dict(server.relay_status),
|
"relay": dict(server.relay_status),
|
||||||
|
"capability_policy": server.capability_policy,
|
||||||
"pool": _pool_summary(nodes),
|
"pool": _pool_summary(nodes),
|
||||||
"memory_pool": memory_pool,
|
"memory_pool": memory_pool,
|
||||||
"recommended_models": [
|
"recommended_models": [
|
||||||
@@ -3213,6 +3335,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
"capacity": capacity_for(node),
|
"capacity": capacity_for(node),
|
||||||
"model_supply": model_supply_for(node),
|
"model_supply": model_supply_for(node),
|
||||||
"throughput": throughput_for(node),
|
"throughput": throughput_for(node),
|
||||||
|
"capability": capability_for(node),
|
||||||
"stats": _node_health(node, server.heartbeat_timeout),
|
"stats": _node_health(node, server.heartbeat_timeout),
|
||||||
}
|
}
|
||||||
for node in nodes
|
for node in nodes
|
||||||
@@ -3365,6 +3488,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
if n.tracker_mode
|
if n.tracker_mode
|
||||||
and _node_matches_model(n, model)
|
and _node_matches_model(n, model)
|
||||||
and _quantization_satisfies(n.quantization, requested_quantization)
|
and _quantization_satisfies(n.quantization, requested_quantization)
|
||||||
|
and _capability_routable(n, server.capability_policy)
|
||||||
]
|
]
|
||||||
|
|
||||||
if not candidates:
|
if not candidates:
|
||||||
@@ -3375,6 +3499,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
if n.shard_start == 0
|
if n.shard_start == 0
|
||||||
and _node_matches_model(n, model)
|
and _node_matches_model(n, model)
|
||||||
and _quantization_satisfies(n.quantization, requested_quantization)
|
and _quantization_satisfies(n.quantization, requested_quantization)
|
||||||
|
and _capability_routable(n, server.capability_policy)
|
||||||
]
|
]
|
||||||
|
|
||||||
if not candidates:
|
if not candidates:
|
||||||
@@ -3480,6 +3605,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
model=route_model,
|
model=route_model,
|
||||||
contracts=server.contracts,
|
contracts=server.contracts,
|
||||||
max_candidates=server.route_stats.config.max_candidate_routes,
|
max_candidates=server.route_stats.config.max_candidate_routes,
|
||||||
|
policy=server.capability_policy,
|
||||||
)
|
)
|
||||||
picked, routing_decision = choose_route(
|
picked, routing_decision = choose_route(
|
||||||
route_candidates, server.route_stats, route_model, rng=server.route_rng,
|
route_candidates, server.route_stats, route_model, rng=server.route_rng,
|
||||||
@@ -3489,7 +3615,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
else:
|
else:
|
||||||
# No head-anchored candidate — legacy greedy cover as fallback
|
# No head-anchored candidate — legacy greedy cover as fallback
|
||||||
# (also produces the layer-gap error message).
|
# (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"}
|
routing_decision = {"mode": "greedy-fallback"}
|
||||||
if route_error:
|
if route_error:
|
||||||
_tracker_log(
|
_tracker_log(
|
||||||
@@ -4344,6 +4475,25 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
self._send_json(400, {"error": str(exc)})
|
self._send_json(400, {"error": str(exc)})
|
||||||
return
|
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(
|
node_id = _node_id_for_registration(
|
||||||
endpoint,
|
endpoint,
|
||||||
model,
|
model,
|
||||||
@@ -4378,6 +4528,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
cert_fingerprint=cert_fingerprint,
|
cert_fingerprint=cert_fingerprint,
|
||||||
peer_id=peer_id,
|
peer_id=peer_id,
|
||||||
friendly_name=friendly_name,
|
friendly_name=friendly_name,
|
||||||
|
capability=capability,
|
||||||
)
|
)
|
||||||
with server.lock:
|
with server.lock:
|
||||||
self._purge_expired_nodes()
|
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),
|
old_model_health=_model_health_summary(server, old.model, old.hf_repo),
|
||||||
model_health=model_health,
|
model_health=model_health,
|
||||||
)
|
)
|
||||||
|
routable = _capability_routable(entry, server.capability_policy)
|
||||||
_tracker_log(
|
_tracker_log(
|
||||||
server,
|
server,
|
||||||
"info",
|
"info",
|
||||||
@@ -4444,6 +4596,18 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
shard=f"{entry.shard_start}-{entry.shard_end}",
|
shard=f"{entry.shard_start}-{entry.shard_end}",
|
||||||
tracker_mode=entry.tracker_mode,
|
tracker_mode=entry.tracker_mode,
|
||||||
model_health=model_health,
|
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"
|
shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded"
|
||||||
@@ -4456,7 +4620,17 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
flush=True,
|
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:
|
if assignment_directive is not None:
|
||||||
payload["directive"] = assignment_directive
|
payload["directive"] = assignment_directive
|
||||||
self._send_json(200, payload)
|
self._send_json(200, payload)
|
||||||
@@ -5989,6 +6163,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
model=route_model,
|
model=route_model,
|
||||||
contracts=server.contracts,
|
contracts=server.contracts,
|
||||||
max_candidates=server.route_stats.config.max_candidate_routes,
|
max_candidates=server.route_stats.config.max_candidate_routes,
|
||||||
|
policy=server.capability_policy,
|
||||||
)
|
)
|
||||||
if candidates:
|
if candidates:
|
||||||
# Prefer a distributed multi-hop route when available. Greedy
|
# Prefer a distributed multi-hop route when available. Greedy
|
||||||
@@ -5998,7 +6173,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
error = ""
|
error = ""
|
||||||
else:
|
else:
|
||||||
route, error = _select_route(
|
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:
|
if error:
|
||||||
_tracker_log(
|
_tracker_log(
|
||||||
@@ -6084,6 +6261,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
model=stats_key,
|
model=stats_key,
|
||||||
contracts=server.contracts,
|
contracts=server.contracts,
|
||||||
max_candidates=cfg.max_candidate_routes,
|
max_candidates=cfg.max_candidate_routes,
|
||||||
|
policy=server.capability_policy,
|
||||||
)
|
)
|
||||||
out[stats_key] = {
|
out[stats_key] = {
|
||||||
"epoch": server.route_stats.epoch(stats_key),
|
"epoch": server.route_stats.epoch(stats_key),
|
||||||
@@ -6177,7 +6355,11 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
|||||||
routes = []
|
routes = []
|
||||||
remaining = list(candidates)
|
remaining = list(candidates)
|
||||||
for _ in range(redundancy):
|
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:
|
if error:
|
||||||
self._send_json(503, {"error": error})
|
self._send_json(503, {"error": error})
|
||||||
return
|
return
|
||||||
@@ -6262,7 +6444,11 @@ class TrackerServer:
|
|||||||
routing_config: RoutingConfig | None = None,
|
routing_config: RoutingConfig | None = None,
|
||||||
enable_test_runner: bool = False,
|
enable_test_runner: bool = False,
|
||||||
test_runner: TestRunManager | None = None,
|
test_runner: TestRunManager | None = None,
|
||||||
|
capability_policy: str | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
|
self._capability_policy = normalize_policy(
|
||||||
|
capability_policy if capability_policy is not None else policy_from_env()
|
||||||
|
)
|
||||||
self._host = host
|
self._host = host
|
||||||
self._requested_port = port
|
self._requested_port = port
|
||||||
self._heartbeat_timeout = heartbeat_timeout
|
self._heartbeat_timeout = heartbeat_timeout
|
||||||
@@ -6475,6 +6661,7 @@ class TrackerServer:
|
|||||||
route_stats=self._route_stats,
|
route_stats=self._route_stats,
|
||||||
relay_status=http_relay_status,
|
relay_status=http_relay_status,
|
||||||
test_runner=self._test_runner,
|
test_runner=self._test_runner,
|
||||||
|
capability_policy=self._capability_policy,
|
||||||
)
|
)
|
||||||
self.port = self._server.server_address[1]
|
self.port = self._server.server_address[1]
|
||||||
|
|
||||||
@@ -6708,6 +6895,15 @@ class TrackerServer:
|
|||||||
else None
|
else None
|
||||||
),
|
),
|
||||||
friendly_name=_normalize_friendly_name(payload.get("friendly_name")),
|
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:
|
with self._lock:
|
||||||
self._registry[node_id] = entry
|
self._registry[node_id] = entry
|
||||||
|
|||||||
63
tests/test_activation_compression.py
Normal file
63
tests/test_activation_compression.py
Normal file
@@ -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
|
||||||
@@ -15,6 +15,7 @@ from meshnet_tracker.routing_stats import (
|
|||||||
route_signature,
|
route_signature,
|
||||||
route_table,
|
route_table,
|
||||||
)
|
)
|
||||||
|
from meshnet_tracker.capability import absent_state
|
||||||
from meshnet_tracker.server import TrackerServer, _enumerate_routes
|
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,
|
proxy_inflight=0,
|
||||||
wallet_address=None,
|
wallet_address=None,
|
||||||
relay_addr=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_start": 0,
|
||||||
"shard_end": shard_end,
|
"shard_end": shard_end,
|
||||||
"tracker_mode": True,
|
"tracker_mode": True,
|
||||||
|
"quantization": "bfloat16",
|
||||||
"benchmark_tokens_per_sec": bench,
|
"benchmark_tokens_per_sec": bench,
|
||||||
"hardware_profile": {},
|
"hardware_profile": {},
|
||||||
"score": 1.0},
|
"score": 1.0},
|
||||||
|
|||||||
@@ -454,6 +454,137 @@ def test_relay_rpc_reuses_connection_for_sequential_requests(monkeypatch):
|
|||||||
assert connection_attempts == 1
|
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():
|
def test_binary_relay_frame_codecs_interoperate():
|
||||||
"Node and relay ship the same binary frame format as separate copies.\n\nTags: gossip, network, relay"
|
"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")
|
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():
|
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"
|
"Pipeline hops zstd-compress large activations; tiny decode bodies pass raw.\n\nTags: gossip, network, relay"
|
||||||
|
|
||||||
|
|||||||
134
tests/test_http_keepalive.py
Normal file
134
tests/test_http_keepalive.py
Normal file
@@ -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)
|
||||||
@@ -9,6 +9,10 @@ import types
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from unittest.mock import MagicMock, patch
|
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
|
# model_catalog tests
|
||||||
@@ -336,6 +340,7 @@ def test_startup_auto_detects_shard_range(monkeypatch, tmp_path):
|
|||||||
# shard_start and shard_end intentionally omitted
|
# shard_start and shard_end intentionally omitted
|
||||||
quantization="bfloat16",
|
quantization="bfloat16",
|
||||||
host="127.0.0.1",
|
host="127.0.0.1",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
assert calls == ["Qwen/Qwen2.5-0.5B-Instruct"]
|
assert calls == ["Qwen/Qwen2.5-0.5B-Instruct"]
|
||||||
assert isinstance(node, FakeNode)
|
assert isinstance(node, FakeNode)
|
||||||
|
|||||||
396
tests/test_node_admission.py
Normal file
396
tests/test_node_admission.py
Normal file
@@ -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
|
||||||
644
tests/test_node_doctor.py
Normal file
644
tests/test_node_doctor.py
Normal file
@@ -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="<test manifest>",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
@@ -26,6 +26,10 @@ from meshnet_node.startup import (
|
|||||||
_tracker_http_error_message,
|
_tracker_http_error_message,
|
||||||
run_startup,
|
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_node.wallet import _b58encode, load_or_create_wallet
|
||||||
from meshnet_contracts import LocalSolanaContracts
|
from meshnet_contracts import LocalSolanaContracts
|
||||||
from meshnet_tracker.server import TrackerServer
|
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",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
torch_threads=8,
|
torch_threads=8,
|
||||||
torch_interop_threads=1,
|
torch_interop_threads=1,
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
node.stop()
|
node.stop()
|
||||||
|
|
||||||
@@ -394,6 +399,7 @@ def test_real_model_startup_passes_download_dir_and_kimi_metadata(monkeypatch, t
|
|||||||
shard_end=60,
|
shard_end=60,
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
cache_dir=cache_dir,
|
cache_dir=cache_dir,
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
node.stop()
|
node.stop()
|
||||||
|
|
||||||
@@ -448,6 +454,7 @@ def test_cuda_benchmark_failure_is_registered_for_inventory_only_gpu(monkeypatch
|
|||||||
shard_start=0,
|
shard_start=0,
|
||||||
shard_end=23,
|
shard_end=23,
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
node.stop()
|
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",
|
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||||||
advertise_host="172.29.104.23",
|
advertise_host="172.29.104.23",
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
pass
|
pass
|
||||||
@@ -1216,6 +1224,7 @@ def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, ca
|
|||||||
vram_mb_override=6144,
|
vram_mb_override=6144,
|
||||||
max_loaded_shards=2,
|
max_loaded_shards=2,
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
|
|
||||||
assert node.backend.total_layers == 24
|
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_start=0,
|
||||||
shard_end=23,
|
shard_end=23,
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
pass
|
pass
|
||||||
@@ -1359,6 +1369,7 @@ def test_public_tracker_model_node_registers_relay_metadata_from_tracker_url_onl
|
|||||||
tracker_url=tracker_url,
|
tracker_url=tracker_url,
|
||||||
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
network_map = _get_json(f"{tracker_url}/v1/network/map")
|
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",
|
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||||||
advertise_host="172.29.104.23",
|
advertise_host="172.29.104.23",
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
network_map = _get_json(f"{tracker_url}/v1/network/map")
|
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,
|
tracker_url=tracker_url,
|
||||||
advertise_host="203.0.113.21",
|
advertise_host="203.0.113.21",
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
route_resp = _get_json(
|
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,
|
tracker_url=tracker_url,
|
||||||
advertise_host="203.0.113.32",
|
advertise_host="203.0.113.32",
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||||
@@ -1637,6 +1651,7 @@ def test_full_startup_sequence(tmp_path):
|
|||||||
model="stub-model",
|
model="stub-model",
|
||||||
wallet_path=wallet_path,
|
wallet_path=wallet_path,
|
||||||
cache_dir=cache_dir,
|
cache_dir=cache_dir,
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
# Wallet was created on disk
|
# Wallet was created on disk
|
||||||
@@ -1699,6 +1714,7 @@ def test_preset_model_startup_starts_heartbeat(tmp_path, monkeypatch):
|
|||||||
model="stub-model",
|
model="stub-model",
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
cache_dir=tmp_path / "shards",
|
cache_dir=tmp_path / "shards",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
assert len(heartbeat_calls) == 1
|
assert len(heartbeat_calls) == 1
|
||||||
@@ -1739,6 +1755,7 @@ def test_preset_model_startup_honors_pinned_shard_range(tmp_path, monkeypatch):
|
|||||||
shard_end=5,
|
shard_end=5,
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
cache_dir=tmp_path / "shards",
|
cache_dir=tmp_path / "shards",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
assert len(heartbeat_calls) == 1
|
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,
|
shard_end=39,
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
cache_dir=tmp_path / "shards",
|
cache_dir=tmp_path / "shards",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
finally:
|
finally:
|
||||||
tracker.stop()
|
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",
|
tracker_url="http://127.0.0.1:8080",
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
tracker_source_disabled=True,
|
tracker_source_disabled=True,
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
assert torch_calls[0]["shard_start"] == 0
|
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",
|
model="tiny-llama",
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
cache_dir=tmp_path / "node-shards",
|
cache_dir=tmp_path / "node-shards",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
assert len(torch_calls) == 1
|
assert len(torch_calls) == 1
|
||||||
@@ -1972,6 +1992,7 @@ def test_torch_startup_retries_registration_when_tracker_unreachable(
|
|||||||
tracker_url=tracker_url,
|
tracker_url=tracker_url,
|
||||||
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
assert register_calls["count"] == 1
|
assert register_calls["count"] == 1
|
||||||
@@ -2052,6 +2073,7 @@ def test_real_model_startup_registers_downloaded_inventory_without_checksum(
|
|||||||
model="tiny-llama",
|
model="tiny-llama",
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
cache_dir=tmp_path / "node-shards",
|
cache_dir=tmp_path / "node-shards",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
assert len(hf_calls) == 1
|
assert len(hf_calls) == 1
|
||||||
@@ -2342,6 +2364,7 @@ def test_startup_cpu_fallback(tmp_path, monkeypatch):
|
|||||||
model="stub-model",
|
model="stub-model",
|
||||||
wallet_path=tmp_path / "wallet.json",
|
wallet_path=tmp_path / "wallet.json",
|
||||||
cache_dir=tmp_path / "shards",
|
cache_dir=tmp_path / "shards",
|
||||||
|
capability_validator=assume_capability,
|
||||||
)
|
)
|
||||||
try:
|
try:
|
||||||
# Node is running even on CPU
|
# Node is running even on CPU
|
||||||
|
|||||||
93
tests/test_prefill_backpressure.py
Normal file
93
tests/test_prefill_backpressure.py
Normal file
@@ -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
|
||||||
@@ -24,6 +24,7 @@ from meshnet_node.model_backend import (
|
|||||||
_should_partial_materialize_shard,
|
_should_partial_materialize_shard,
|
||||||
_decoder_attention_mask,
|
_decoder_attention_mask,
|
||||||
_int_tensor_header,
|
_int_tensor_header,
|
||||||
|
_tensor_from_bfloat16_bytes,
|
||||||
_torch_cuda_is_executable,
|
_torch_cuda_is_executable,
|
||||||
build_quantization_config,
|
build_quantization_config,
|
||||||
validate_quantization,
|
validate_quantization,
|
||||||
@@ -538,6 +539,20 @@ def test_int_tensor_header_serializes_torch_tensors():
|
|||||||
assert header.startswith("1,3:")
|
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():
|
def test_decoder_attention_mask_is_causal_float_mask():
|
||||||
"Decoder attention mask is causal float mask\n\nTags: model, node, real-inference"
|
"Decoder attention mask is causal float mask\n\nTags: model, node, real-inference"
|
||||||
torch = pytest.importorskip("torch")
|
torch = pytest.importorskip("torch")
|
||||||
|
|||||||
126
tests/test_route_session_benchmark.py
Normal file
126
tests/test_route_session_benchmark.py
Normal file
@@ -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
|
||||||
78
tests/test_seam_telemetry.py
Normal file
78
tests/test_seam_telemetry.py
Normal file
@@ -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,
|
||||||
|
)
|
||||||
525
tests/test_tracker_capability_admission.py
Normal file
525
tests/test_tracker_capability_admission.py
Normal file
@@ -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()
|
||||||
Reference in New Issue
Block a user