diff --git a/.claude/memory/alpha-hardening-navigation.md b/.claude/memory/alpha-hardening-navigation.md index 057c33c..4ce5bee 100644 --- a/.claude/memory/alpha-hardening-navigation.md +++ b/.claude/memory/alpha-hardening-navigation.md @@ -30,3 +30,7 @@ Both are already migrated into `.scratch/alpha-hardening/prd.json` (AH-021 updat **Why:** three audits agreed the alpha blockers are unauthenticated gossip (anyone can inject billing events), the free-credit faucet, and ephemeral bans. **How to apply:** work test-first per issue acceptance criteria; use `.venv`; `cryptography` belongs in node deps (wallet.py imports it — causes many of the 24 "failures" in a fresh env). See [[project-status]] and [[autonomous-work-style]]. + +## Routing telemetry resume (2026-07-07) + +`.scratch/alpha-hardening/issues/24-routing-telemetry-resume.md` / AH-024 captures the interrupted Claude handoff. Learned routing is already committed at `518c259`; the dirty tree contains live-progress/current-request heartbeat/dashboard telemetry. First known blocker: `packages/tracker/meshnet_tracker/server.py:1490` uses `threading.Lock | None`, which crashes import because `threading.Lock` is a factory function at runtime. Fix that before running the targeted telemetry tests. Keep `.claude/settings.local.json` uncommitted unless explicitly approved. diff --git a/.claude/memory/project-status.md b/.claude/memory/project-status.md index 8da1b6f..4334a9b 100644 --- a/.claude/memory/project-status.md +++ b/.claude/memory/project-status.md @@ -42,7 +42,8 @@ Historical handoff note: `/mnt/c/Users/popov/Downloads/neuron-tai-alpha-handoff- - Verification: downloader/startup targeted subset passes (`pytest tests/test_node_startup.py -k "download_shard or same_shard"`). Full `tests/test_node_startup.py` has 46 passed and 4 unrelated Windows chmod/path separator failures. - Live Windows confirmation: `meshnet-node start --tracker http://192.168.0.179:8080 --model Qwen3.6-35B-A3B` reuses `F:\_STORAGE\models\qwen3.6-35b-a3b`, prints `Cached at`, registers, and reaches ready as node `5gMLrmyB-26b1f8a4204a`. - Follow-up fix: preset-model startup now starts the heartbeat thread after registration; without this, the node appeared briefly on the dashboard and was purged on first inference/route after heartbeat expiry. Tracker dashboard now has a "Console output" panel backed by `/v1/console` for node register/expiry, routing failures, and proxy events. -- Qwen3.6-35B-A3B reserve-based split is expected: an 79 GB CPU node may be assigned layers 0-36, and a second node fills 37-39. Do not "fix" this by bypassing the 20% assignment reserve unless the shard-planning policy changes. +- Qwen3.6-35B-A3B CPU runtime cap (2026-07-08): the old reserve-based split could assign an 79 GB CPU node layers 0-36, but real partial loading can exceed that budget and die without a Python traceback. Node startup now clips oversized CPU auto-assignments before loading, and tracker CPU assignment uses a stricter runtime headroom factor; do not revert this to the old 20% reserve-only policy. - Route hardening: tracker chat proxy and `/v1/route` diagnostics now use alias-aware preset node matching for split Qwen3.6 routes; dashboard derives grouped inference history from proxy route/complete console events and shows observed TPS after completion. - Live proxy hardening: model lookup trims outer whitespace before alias matching (`qwen3.6-35b-a3b ` resolves), and tracker route logs/dashboard queue depth combine heartbeat queue with tracker-local proxy in-flight counts so Postman-style bursts no longer show every selected route as queue `0`. - Split-shard streaming hardening: Qwen3.6-style distributed generation now emits SSE chunks token-by-token from the head node instead of buffering all generated text until completion. Tracker direct/relay stream proxy logs `proxy progress` with live tokens/TPS, dashboard Inference history shows currently processing requests with live TPS/tokens/queue, and relay stream completion no longer references an undefined `session_id`. +- Native Windows Qwen3.6-MoE import fix: `flash-linear-attention` imports `triton`; without `triton-windows`, startup fails with misleading `Could not import module 'Qwen3_5MoeForCausalLM'`. Installed `triton-windows` in `C:\Users\popov\miniforge3` and added it as a Windows-only node dependency. diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..0a081b4 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,28 @@ +# Normalize line endings across Windows/Linux checkouts. +# All text files are stored as LF in the repo and checked out as LF +# on every OS. Git auto-detects text vs binary. +* text=auto eol=lf + +# Explicitly binary — never touch these bytes. +*.png binary +*.jpg binary +*.jpeg binary +*.gif binary +*.ico binary +*.pdf binary +*.zip binary +*.gz binary +*.tar binary +*.wasm binary +*.sqlite binary +*.sqlite3 binary +*.safetensors binary +*.gguf binary + +# Scripts that must stay LF even if someone forces CRLF locally. +*.sh text eol=lf +*.py text eol=lf + +# Windows batch files genuinely need CRLF. +*.bat text eol=crlf +*.cmd text eol=crlf diff --git a/.scratch/alpha-hardening/README.md b/.scratch/alpha-hardening/README.md index 65df940..2d2651e 100644 --- a/.scratch/alpha-hardening/README.md +++ b/.scratch/alpha-hardening/README.md @@ -9,6 +9,8 @@ Pre-release alpha audit + grilling (2026-07-04). Bucket 1 trust-boundary blocker Locked scope: one settlement tracker, open node join, devnet mock-USDT, reputation carries forward → fraud must be bounded. See [ADR-0016](../../docs/adr/0016-alpha-scope-and-known-limitations.md). +**Resume task (2026-07-07):** [24 - Routing telemetry resume](./issues/24-routing-telemetry-resume.md) is `ready-for-agent`. Learned-routing commit `518c259` is already present; dirty tree contains current-request heartbeat/dashboard telemetry and a known import-time annotation crash in `server.py:1490`. + ## Artifacts | Path | Status | diff --git a/.scratch/alpha-hardening/issues/24-routing-telemetry-resume.md b/.scratch/alpha-hardening/issues/24-routing-telemetry-resume.md new file mode 100644 index 0000000..5246152 --- /dev/null +++ b/.scratch/alpha-hardening/issues/24-routing-telemetry-resume.md @@ -0,0 +1,92 @@ +Status: ready-for-agent + +Scoped 2026-07-07 from an interrupted Claude session. This is a resume/cleanup task for routing and live-progress work that is partly committed and partly left dirty in the working tree. + +# 24 - Finish learned-routing telemetry and live-progress cleanup + +## Current state + +The main dynamic routing feature is already committed at `518c259` (`routing improvements - dynamic (wip)`): + +- `packages/tracker/meshnet_tracker/routing_stats.py` - decayed-EWMA route stats store, epsilon-greedy route selection, diagnostics. +- `packages/tracker/meshnet_tracker/server.py` - route enumeration per head, bandit selection in the chat proxy, epoch bumps on node join/leave, `/v1/routing`, route sample recording with 8-token hygiene. +- `packages/tracker/meshnet_tracker/cli.py` - `--route-explore-share`, `--route-weight-alpha`, `--route-stats-half-life` and env vars. +- `packages/tracker/meshnet_tracker/dashboard.html` - "Routing (learned)" panel. +- `docs/adr/0021-dynamic-statistical-routing.md` - design record. +- `tests/test_dynamic_routing.py` - includes the exact GPU(0-21)+CPU(0-39) topology, hybrid downstream `start_layer=22`, 0.6/0.4 traffic split for a 1.5 TPS ratio, and scout-rate behavior. + +The current working tree still has uncommitted follow-up work: + +- `packages/node/meshnet_node/torch_server.py` - tracks in-flight chat requests, exposes `TorchNodeServer.current_requests`, prints generation progress with TPS. +- `packages/node/meshnet_node/startup.py` - sends `current_requests` in heartbeat payloads and increases heartbeat cadence while busy. +- `packages/tracker/meshnet_tracker/server.py` - accepts heartbeat `current_requests`, includes them in `/v1/network/map`, and logs `proxy connecting` before upstream connection. +- `packages/tracker/meshnet_tracker/dashboard.html` - enriches the call wall from heartbeat `current_requests` so active requests remain visible even before terminal proxy events. +- `tests/test_real_model_backend.py` and `tests/test_tracker_routing.py` - targeted coverage for current-request snapshots, heartbeat sanitization/storage, and TPS progress logging. +- `QUICKSTART.md` - documents optional linear-attention fast-path packages for Qwen3.5/3.6 GPU nodes. + +There is also an untracked local file, `.claude/settings.local.json`, which should not be included unless the owner explicitly wants local Claude settings committed. + +## Known blocker found during resume + +Targeted pytest currently fails during import before reaching the new tests: + +```text +TypeError: unsupported operand type(s) for |: 'builtin_function_or_method' and 'NoneType' +``` + +Immediate cause: `packages/tracker/meshnet_tracker/server.py:1490` annotates `ws_lock: threading.Lock | None = None`. `threading.Lock` is a factory function at runtime, not a type, so `| None` evaluates eagerly and crashes. This exists on `HEAD` too, not just in the dirty telemetry changes. + +Fix options: + +- Add `from __future__ import annotations` at the top of `server.py`, then run enough tests to catch any annotation side effects. +- Or change that annotation to a safe runtime type such as `Any | None` / remove the union annotation. Keep the change minimal. + +## What to do next + +1. Fix the import-time `threading.Lock | None` crash. +2. Re-run the targeted tests: + + ```bash + .\.venv\Scripts\python.exe -m pytest tests/test_tracker_routing.py::test_tracker_heartbeat_stores_current_requests tests/test_tracker_routing.py::test_normalize_current_requests_sanitizes_payload tests/test_real_model_backend.py::test_current_requests_snapshot_while_generating tests/test_real_model_backend.py::test_distributed_generating_log_includes_tps -q + ``` + +3. Run the relevant routing regression tests: + + ```bash + .\.venv\Scripts\python.exe -m pytest tests/test_dynamic_routing.py tests/test_tracker_routing.py -q + ``` + +4. If practical, run the non-integration suite: + + ```bash + .\.venv\Scripts\python.exe -m pytest tests/ -q -m "not integration" + ``` + +5. Confirm or document the pre-existing failure from the interrupted session: `test_proxy_chat_splits_payout_by_tracker_assigned_route_span` reportedly failed on `HEAD` too and was unrelated. +6. Commit the intentional work in two commits if it remains naturally split: + - learned routing is already committed in `518c259`; leave it alone unless fixing regressions there. + - commit the live-progress/current-request telemetry cleanup separately after tests pass. + +## Acceptance criteria + +- [ ] Importing `meshnet_tracker.server` no longer crashes on the lock annotation. +- [ ] Current-request heartbeat payloads are sanitized and surfaced in `/v1/network/map`. +- [ ] Node-side in-flight chat snapshots report request id, model, token count, elapsed seconds, tokens/sec, and routing completion. +- [ ] Dashboard call wall can show active requests from heartbeat data, not only tracker console terminal events. +- [ ] Targeted telemetry tests pass. +- [ ] Dynamic routing tests still pass, including GPU(0-21)+CPU(0-39) hybrid-route enumeration and traffic split behavior. +- [ ] Full or non-integration suite result is recorded; unrelated pre-existing failures are named explicitly. +- [ ] `.claude/settings.local.json` remains uncommitted unless intentionally approved. + +## ADR links + +- [ADR-0020](../../docs/adr/0020-chat-streaming-live-progress-and-mixed-topology-routing.md) +- [ADR-0021](../../docs/adr/0021-dynamic-statistical-routing.md) + +## Blocked by + +None. The import-time annotation crash is the first fix. + +## Blocks + +Clean handoff/commit of the interrupted live routing progress work. diff --git a/.scratch/alpha-hardening/prd.json b/.scratch/alpha-hardening/prd.json index f53aedf..af339cc 100644 --- a/.scratch/alpha-hardening/prd.json +++ b/.scratch/alpha-hardening/prd.json @@ -483,9 +483,29 @@ "notes": "Source issue: .scratch/alpha-hardening/issues/23-dynamic-hf-pricing.md. High priority, ship-soon for launch, NOT an alpha-release blocker (unlike AH-021).", "dependsOn": [], "completionNotes": "Completed by agent" + }, + { + "id": "AH-024", + "title": "24 - Finish learned-routing telemetry and live-progress cleanup", + "description": "Status: ready-for-agent\n\nScoped 2026-07-07 from an interrupted Claude session. The learned-routing feature is already committed at 518c259 (`routing improvements - dynamic (wip)`): routing_stats.py, tracker route enumeration and bandit selection, CLI routing flags, `/v1/routing`, dashboard Routing (learned), ADR-0021, and tests/test_dynamic_routing.py including the GPU(0-21)+CPU(0-39) hybrid topology. The dirty working tree contains follow-up live-progress/current-request telemetry in torch_server.py, startup.py, tracker server/dashboard, tests, and QUICKSTART. Known blocker found during resume: importing meshnet_tracker.server currently crashes at `server.py:1490` because `ws_lock: threading.Lock | None = None` evaluates `threading.Lock` as a factory function, not a type. Fix that first, then verify and commit the telemetry cleanup separately from the already-committed dynamic-routing work. Leave `.claude/settings.local.json` uncommitted unless explicitly approved.\n\nSource issue has exact file list, commands, and the reported pre-existing unrelated failure (`test_proxy_chat_splits_payout_by_tracker_assigned_route_span`).", + "acceptanceCriteria": [ + "Importing `meshnet_tracker.server` no longer crashes on the lock annotation", + "Current-request heartbeat payloads are sanitized and surfaced in `/v1/network/map`", + "Node-side in-flight chat snapshots report request id, model, token count, elapsed seconds, tokens/sec, and routing completion", + "Dashboard call wall can show active requests from heartbeat data, not only tracker console terminal events", + "Targeted telemetry tests pass", + "Dynamic routing tests still pass, including GPU(0-21)+CPU(0-39) hybrid-route enumeration and traffic split behavior", + "Full or non-integration suite result is recorded; unrelated pre-existing failures are named explicitly", + "`.claude/settings.local.json` remains uncommitted unless intentionally approved" + ], + "priority": 24, + "passes": true, + "notes": "Source issue: .scratch/alpha-hardening/issues/24-routing-telemetry-resume.md. Resume task for interrupted 2026-07-07 Claude session; first known fix is server.py:1490 annotation crash.", + "dependsOn": [], + "completionNotes": "" } ], "metadata": { - "updatedAt": "2026-07-06T06:01:25.474Z" + "updatedAt": "2026-07-07T21:30:00.000Z" } -} \ No newline at end of file +} diff --git a/CONTEXT.md b/CONTEXT.md index 9adab50..a3b9ee7 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -1,23 +1,23 @@ -# Distributed Inference Network - -A volunteer GPU network where nodes independently load model shards, a tracker routes inference through optimal node chains, and contributors earn tokens for serving compute. - -## Language - -### Nodes & compute - -**Node**: -A volunteer machine that runs the node client, holds one or more shards on disk, and serves inference requests for those shards. -_Avoid_: worker, peer, miner, server - -**Shard**: -A contiguous range of transformer layers from a model that a node loads and serves. Shards are the unit of storage, assignment, and reward. -_Avoid_: partition, slice, chunk, segment - -**Shard Swarm**: -The P2P group of nodes that collectively seed and download a specific shard. One swarm exists per shard. -_Avoid_: torrent, cluster, pool - +# Distributed Inference Network + +A volunteer GPU network where nodes independently load model shards, a tracker routes inference through optimal node chains, and contributors earn tokens for serving compute. + +## Language + +### Nodes & compute + +**Node**: +A volunteer machine that runs the node client, holds one or more shards on disk, and serves inference requests for those shards. +_Avoid_: worker, peer, miner, server + +**Shard**: +A contiguous range of transformer layers from a model that a node loads and serves. Shards are the unit of storage, assignment, and reward. +_Avoid_: partition, slice, chunk, segment + +**Shard Swarm**: +The P2P group of nodes that collectively seed and download a specific shard. One swarm exists per shard. +_Avoid_: torrent, cluster, pool + **Inference Route**: An ordered sequence of nodes whose shards together cover all layers of a model. The tracker selects the optimal route per request. _Avoid_: pipeline, chain, path @@ -55,115 +55,115 @@ Realtime progress information for an active Route Session, including phase, gene _Avoid_: logs, debug output ### Tracker - -**Tracker**: -The coordinator service that maintains the node registry, scores nodes by throughput/latency, and assigns inference routes. Runs as a centralized service with a P2P gossip fallback. -_Avoid_: coordinator, scheduler, director - + +**Tracker**: +The coordinator service that maintains the node registry, scores nodes by throughput/latency, and assigns inference routes. Runs as a centralized service with a P2P gossip fallback. +_Avoid_: coordinator, scheduler, director + **Tracker Node**: A node that serves at least the first-layer shard (`layers[0..k]`) for a model and acts as the inference entry point for that model. Tracker nodes own the tokenizer and `embed_tokens`, receive client requests directly, select the onward route from the coverage map, and stream results and progress when possible. Any node advertising a new model to the network becomes its tracker node. _Avoid_: primary node, master node, gateway node - -**Coverage Map**: -The tracker's per-model mapping of layer ranges to node counts: `[(start_layer, end_layer, node_count), ...]`. A layer range with `node_count=0` is a coverage gap — the model is unroutable until the gap is filled. Coverage-first bin-packing fills all gaps before adding redundancy. -_Avoid_: shard map, assignment table, coverage report - -**Rebalance Directive**: -A `LOAD_SHARD` or `DROP_SHARD` instruction the tracker issues to a node when the coverage map changes (node joins, node leaves, or load-balance reoptimization). Delivered as part of the node's heartbeat response. -_Avoid_: rebalance command, shard instruction, migration order - -**Node Score**: -A throughput/latency rating the tracker maintains per node, used for route selection. Updated continuously from inference telemetry. -_Avoid_: reputation, rating, rank - -### Payments & fraud - -**Stake**: -Collateral a node stands to lose for fraud. In the current design the node's Pending Balance serves as stake — no upfront deposit is required. An optional USDT/TAI deposit may return later for routing priority. -_Avoid_: deposit, bond, escrow - -**Treasury**: -The single project-owned Solana wallet that custodially holds client deposits, pays node payouts, and accumulates the Protocol Cut. Its keypair is loaded only on settlement-capable trackers. -_Avoid_: escrow, vault, hot wallet - -**Pending Balance**: -A node's accrued, not-yet-paid USDT earnings on the tracker ledger. Doubles as the node's fraud collateral: it is forfeited in full when a validator catches a divergent output. -_Avoid_: unpaid rewards, accrual, balance due - -**Settlement Period**: -The dynamic interval driving on-chain payouts: a node is paid when its Pending Balance exceeds the Payout Threshold or the period elapses, whichever comes first. Short in development (seconds), long in production (daily), configurable to grow with volume. -_Avoid_: epoch, payout cycle, billing cycle - -**Payout Threshold**: -The minimum Pending Balance that triggers an immediate payout before the Settlement Period elapses. Includes a dust floor so payouts are never smaller than they are worth. -_Avoid_: minimum payout, dust limit - -**Protocol Cut**: -The 10% of inference fees retained by the project for infrastructure; the remaining 90% goes to the nodes that served the request. Accumulates in the Treasury as the future TAI liquidity reserve. -_Avoid_: spread, commission, house fee - -**Deposit Watcher**: -The tracker component that observes the Treasury's on-chain USDT deposits and credits the sending client's API-key ledger balance. -_Avoid_: payment listener, chain scanner - -**Mock USDT**: -The self-created 6-decimal SPL mint that stands in for USDT on devnet, where real USDT does not exist. The mint address is configuration, so mainnet cutover is a config change. -_Avoid_: test token, fake USDT, devnet dollar - -**Tax**: -The share of caller payments distributed to compute nodes as rewards. Taxes are weighted by completed work and historical node speed so faster, larger nodes earn proportionally more. -_Avoid_: fee, toll, commission - -**Caller Credit**: -Free starting balance granted to a new caller/API key so they can try the network before topping up. -_Avoid_: signup bonus, faucet, airdrop - -**Free Compute Job**: -Work a compute node performs without earning immediate rewards, usually during probation or bootstrap phases. -_Avoid_: unpaid labor, warmup request - -**Slash**: -The penalty for a proven fraud incident: the node's entire Pending Balance is forfeited to the Treasury and a Strike is recorded. -_Avoid_: penalize, burn, fine, forfeit - -**Strike**: -A fraud incident recorded on-chain against a node. Enough strikes result in a ban. -_Avoid_: infraction, violation, flag - -**Ban**: -Permanent exclusion of a wallet from the network after exceeding the strike threshold. Recorded on-chain. -_Avoid_: blacklist, block, suspension - -**Probationary Period**: -The first N jobs a new wallet must complete without earning, to raise the cost of re-entering after a ban. -_Avoid_: trial period, warmup, grace period - -**Token**: -TAI, our native Solana SPL token. Deferred (ADR-0015): nodes are currently paid directly in USDT; TAI returns as the reward/upside layer once volume exists, funded by the accumulated Protocol Cut. Clients never need to hold it. -_Avoid_: coin, reward token, native token - -**Contract Boundary**: -The Python interface in `packages/contracts` that represents registry, payment, and settlement behavior. During the prototype it is implemented by deterministic local wrappers; later the same boundary is backed by real Solana programs. -_Avoid_: mock contract, fake chain, temporary hack - -**Validator**: -A trusted node (or the tracker itself) that re-runs a sample of inference requests to detect fraud. -_Avoid_: auditor, checker, referee - -**Validation Event**: -A completed inference record that contains enough information for a validator to decide whether to sample and re-run the request: session id, model preset, messages, inference route, node wallets, and observed output. -_Avoid_: audit log, trace, receipt - -**Slash Proof**: -The record submitted by a validator when a sampled re-run diverges from the observed output beyond tolerance. In the prototype this is deterministic local contract state; later it maps to an on-chain proof transaction. -_Avoid_: accusation, report, claim - -### Client-facing - -**Client**: -Any application or user that sends inference requests to the gateway. Prepays USDT into the Treasury; each request is metered against the resulting ledger balance at a per-1K-tokens price set per model. -_Avoid_: user, caller, consumer - -**Model Preset**: -A named, versioned model available on the network (e.g. `llama-3-70b`). The tracker knows which nodes hold which shards for each preset. -_Avoid_: model, checkpoint, version + +**Coverage Map**: +The tracker's per-model mapping of layer ranges to node counts: `[(start_layer, end_layer, node_count), ...]`. A layer range with `node_count=0` is a coverage gap — the model is unroutable until the gap is filled. Coverage-first bin-packing fills all gaps before adding redundancy. +_Avoid_: shard map, assignment table, coverage report + +**Rebalance Directive**: +A `LOAD_SHARD` or `DROP_SHARD` instruction the tracker issues to a node when the coverage map changes (node joins, node leaves, or load-balance reoptimization). Delivered as part of the node's heartbeat response. +_Avoid_: rebalance command, shard instruction, migration order + +**Node Score**: +A throughput/latency rating the tracker maintains per node, used for route selection. Updated continuously from inference telemetry. +_Avoid_: reputation, rating, rank + +### Payments & fraud + +**Stake**: +Collateral a node stands to lose for fraud. In the current design the node's Pending Balance serves as stake — no upfront deposit is required. An optional USDT/TAI deposit may return later for routing priority. +_Avoid_: deposit, bond, escrow + +**Treasury**: +The single project-owned Solana wallet that custodially holds client deposits, pays node payouts, and accumulates the Protocol Cut. Its keypair is loaded only on settlement-capable trackers. +_Avoid_: escrow, vault, hot wallet + +**Pending Balance**: +A node's accrued, not-yet-paid USDT earnings on the tracker ledger. Doubles as the node's fraud collateral: it is forfeited in full when a validator catches a divergent output. +_Avoid_: unpaid rewards, accrual, balance due + +**Settlement Period**: +The dynamic interval driving on-chain payouts: a node is paid when its Pending Balance exceeds the Payout Threshold or the period elapses, whichever comes first. Short in development (seconds), long in production (daily), configurable to grow with volume. +_Avoid_: epoch, payout cycle, billing cycle + +**Payout Threshold**: +The minimum Pending Balance that triggers an immediate payout before the Settlement Period elapses. Includes a dust floor so payouts are never smaller than they are worth. +_Avoid_: minimum payout, dust limit + +**Protocol Cut**: +The 10% of inference fees retained by the project for infrastructure; the remaining 90% goes to the nodes that served the request. Accumulates in the Treasury as the future TAI liquidity reserve. +_Avoid_: spread, commission, house fee + +**Deposit Watcher**: +The tracker component that observes the Treasury's on-chain USDT deposits and credits the sending client's API-key ledger balance. +_Avoid_: payment listener, chain scanner + +**Mock USDT**: +The self-created 6-decimal SPL mint that stands in for USDT on devnet, where real USDT does not exist. The mint address is configuration, so mainnet cutover is a config change. +_Avoid_: test token, fake USDT, devnet dollar + +**Tax**: +The share of caller payments distributed to compute nodes as rewards. Taxes are weighted by completed work and historical node speed so faster, larger nodes earn proportionally more. +_Avoid_: fee, toll, commission + +**Caller Credit**: +Free starting balance granted to a new caller/API key so they can try the network before topping up. +_Avoid_: signup bonus, faucet, airdrop + +**Free Compute Job**: +Work a compute node performs without earning immediate rewards, usually during probation or bootstrap phases. +_Avoid_: unpaid labor, warmup request + +**Slash**: +The penalty for a proven fraud incident: the node's entire Pending Balance is forfeited to the Treasury and a Strike is recorded. +_Avoid_: penalize, burn, fine, forfeit + +**Strike**: +A fraud incident recorded on-chain against a node. Enough strikes result in a ban. +_Avoid_: infraction, violation, flag + +**Ban**: +Permanent exclusion of a wallet from the network after exceeding the strike threshold. Recorded on-chain. +_Avoid_: blacklist, block, suspension + +**Probationary Period**: +The first N jobs a new wallet must complete without earning, to raise the cost of re-entering after a ban. +_Avoid_: trial period, warmup, grace period + +**Token**: +TAI, our native Solana SPL token. Deferred (ADR-0015): nodes are currently paid directly in USDT; TAI returns as the reward/upside layer once volume exists, funded by the accumulated Protocol Cut. Clients never need to hold it. +_Avoid_: coin, reward token, native token + +**Contract Boundary**: +The Python interface in `packages/contracts` that represents registry, payment, and settlement behavior. During the prototype it is implemented by deterministic local wrappers; later the same boundary is backed by real Solana programs. +_Avoid_: mock contract, fake chain, temporary hack + +**Validator**: +A trusted node (or the tracker itself) that re-runs a sample of inference requests to detect fraud. +_Avoid_: auditor, checker, referee + +**Validation Event**: +A completed inference record that contains enough information for a validator to decide whether to sample and re-run the request: session id, model preset, messages, inference route, node wallets, and observed output. +_Avoid_: audit log, trace, receipt + +**Slash Proof**: +The record submitted by a validator when a sampled re-run diverges from the observed output beyond tolerance. In the prototype this is deterministic local contract state; later it maps to an on-chain proof transaction. +_Avoid_: accusation, report, claim + +### Client-facing + +**Client**: +Any application or user that sends inference requests to the gateway. Prepays USDT into the Treasury; each request is metered against the resulting ledger balance at a per-1K-tokens price set per model. +_Avoid_: user, caller, consumer + +**Model Preset**: +A named, versioned model available on the network (e.g. `llama-3-70b`). The tracker knows which nodes hold which shards for each preset. +_Avoid_: model, checkpoint, version diff --git a/QUICKSTART.md b/QUICKSTART.md index 6f4b452..eabb88f 100644 --- a/QUICKSTART.md +++ b/QUICKSTART.md @@ -1,783 +1,552 @@ -# Quickstart — Running a node and testing inference - -This guide gets you from zero to a live inference request in three terminals. -Tested on: AMD Ryzen AI Max (Strix Halo APU), 124 GB RAM, Linux, CPU inference. - ---- - -## Prerequisites - -```bash -# Clone and enter repo -cd /run/media/popov/d/DEV/repos/d-popov.com/AI - -# Create the virtualenv if it does not exist yet -python3 -m venv .venv - -# Keep packaging tools current enough for editable installs -.venv/bin/python -m pip install --upgrade pip setuptools wheel - -# Install Python packages (editable — picks up code changes immediately) -.venv/bin/pip install -e packages/tracker -e packages/node -e packages/p2p -e packages/gateway -e packages/relay - -# CPU-only PyTorch (skip if you have CUDA/ROCm already) -.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu - -# HuggingFace model libraries -.venv/bin/pip install "transformers>=5.12" accelerate -``` - -> **NVIDIA GPU (CUDA):** replace the torch line with `pip install torch` (default index). -> **AMD GPU (ROCm):** `pip install torch --index-url https://download.pytorch.org/whl/rocm6.2` - -### Version and library notes for Qwen3.5/3.6-MoE models - -- **transformers ≥ 5.12 is required** for Qwen3.5/3.6-MoE (e.g. `Qwen3.6-35B-A3B`). - Older versions fail at load time with - `'Qwen3_5MoeConfig' object has no attribute 'vocab_size'`. Check with - `python -c "import transformers; print(transformers.__version__)"` and upgrade - with `pip install -U transformers` in the environment that runs `meshnet-node` - (conda/miniforge users: upgrade inside that env, not a layered `.venv`). -- The startup warning - `The fast path is not available because one of the required library is not installed` - is **harmless** — transformers falls back to a pure-torch implementation of the - linear-attention layers. The fast-path packages (`flash-linear-attention`, - `causal-conv1d`) are CUDA-only kernels: install them for GPU speed if you want, - skip them entirely on CPU nodes. -- `pip install nvidia-ml-py` silences the pynvml deprecation warning on NVIDIA hosts. - -## Bootstrap a tracker on a new machine - -Use this when provisioning a fresh LAN/public tracker host. The tracker itself is -lightweight; install the relay too if nodes will connect from NAT, WSL2, mobile, -or other networks where inbound node ports are not reachable. - -```bash -# 1. Get the repo onto the tracker host -git clone https://git.d-popov.com/popov/neuron-tai.git AI -cd AI - -# 2. Create an isolated Python environment -python3 -m venv .venv -.venv/bin/python -m pip install --upgrade pip setuptools wheel - -# 3. Install only the services needed by the tracker host -.venv/bin/pip install -e packages/tracker -e packages/relay -e packages/gateway -``` - -For a private LAN tracker, start only the tracker and open the selected TCP port -on the host firewall if other machines will join: - -```bash -.venv/bin/meshnet-tracker start --host 0.0.0.0 --port 8080 -# --starting-credit 1 --devnet-topup 10 -``` - -Verify from the tracker host: - -```bash -curl -s http://localhost:8080/v1/network/map | python3 -m json.tool -``` - -Verify from another LAN machine, replacing the IP with the tracker host's LAN IP: - -```bash -curl -s http://192.168.0.179:8080/v1/network/map | python3 -m json.tool -``` - -For a public tracker with relay support, run both services. The relay listens on -`8765`; the tracker below listens on `8081` and advertises the public WebSocket -URL that nodes should use for outbound relay connections: - -```bash -# Terminal 1 — relay -.venv/bin/meshnet-relay --host 0.0.0.0 --port 8765 - -# Terminal 2 — tracker -.venv/bin/meshnet-tracker start \ - --host 0.0.0.0 \ - --port 8081 \ - --relay-url wss://ai.neuron.d-popov.com/ws -``` - -If this host sits behind Nginx Proxy Manager, point `/` and `/v1/*` at tracker -port `8081`, and point `/ws` plus `/rpc` at relay port `8765` as shown in the -public tracker section below. After the proxy is configured, verify the public -bootstrap endpoint: - -```bash -curl -s https://ai.neuron.d-popov.com/v1/network/map | python3 -m json.tool -``` - -Nodes can then join with either the LAN tracker URL or the public URL: - -```bash -.venv/bin/meshnet-node start --tracker http://192.168.0.179:8080 --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 -.venv/bin/meshnet-node start --tracker https://ai.d-popov.com --model Qwen/Qwen2.5-0.5B-Instruct -``` - -### Windows / WSL2 - -Run the Linux commands from WSL, not Git Bash. From the repo opened in Git Bash: - -```bash -wsl -cd /mnt/d/DEV/workspace/REPOS/git.d-popov.com/neuron-tai -python3 -m venv .venv -.venv/bin/python -m pip install --upgrade pip setuptools wheel -.venv/bin/pip install -e packages/tracker -e packages/node -e packages/p2p -e packages/gateway -e packages/relay -.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu -.venv/bin/pip install "transformers>=5.12" accelerate -.venv/bin/meshnet-node --help -``` - -If `.venv/bin/meshnet-node` is missing, the editable install step did not finish -successfully. Re-run the `.venv/bin/pip install -e ...` command above inside WSL. - -WSL2 sits behind Windows NAT and is **not directly reachable** from other LAN machines. -Direct cross-host hops time out. The relay path (see below) solves this: the WSL2 node -opens an outbound WebSocket to the relay server and all traffic flows through that tunnel. -No firewall rules, no `--advertise-host` needed — just point at the public tracker URL. - -### Native Windows PowerShell node (not WSL) - -Use this when the tracker is on another machine and you want Windows to host a -reachable node on the LAN. - -#### Option A — existing conda/miniforge environment with CUDA torch (recommended if you already have it) - -First, make sure the conda base environment is active so that `python` and `pip` both -resolve to the same miniforge installation: - -```powershell -conda activate base -deactivate # drop any .venv that may be layered on top; safe no-op if none active -``` - -Install project packages into the active conda/miniforge env: - -```powershell -cd D:\DEV\workspace\REPOS\git.d-popov.com\neuron-tai - -pip install -e packages\tracker -e packages\node -e packages\p2p -e packages\gateway -e packages\relay -pip install "transformers>=5.12" accelerate safetensors # torch is already present -``` - -> Conda/miniforge envs often carry an older `transformers` pinned by other tools -> (aider, etc.). Qwen3.5/3.6-MoE models need **transformers ≥ 5.12** — verify with -> `python -c "import transformers; print(transformers.__version__)"`. The pip -> resolver may print dependency-conflict warnings for those other tools; they don't -> affect `meshnet-node`. - -Verify torch is importable and CUDA is live **before** starting the node: - -```powershell -python -c "import torch; print(torch.__version__, torch.cuda.is_available())" -# Expected: 2.x.x+cuXXX True -``` - -If you get `ModuleNotFoundError: No module named 'torch'` even though `pip install torch` -says "already satisfied", the `torch/` package directory is missing while the metadata -stub remains (can happen after a conda-managed install). Force-reinstall all three -PyTorch packages together so their versions stay in sync: - -```powershell -pip install --force-reinstall torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 -``` - -> **Important:** always reinstall `torch`, `torchvision`, and `torchaudio` as a group. -> Upgrading only `torch` leaves `torchvision` on an incompatible version, which causes -> `RuntimeError: operator torchvision::nms does not exist` and makes transformers fail -> to import any model class (the error surfaces as `Could not import module 'Qwen2ForCausalLM'`). - -Then re-run the verify step above. - -If that prints `True` but `meshnet-node` still can't find torch, the venv entry point -is shadowing the conda one. Check which binary wins: - -```powershell -(Get-Command meshnet-node).Source -# Should show: C:\Users\\miniforge3\Scripts\meshnet-node.exe -# If it shows .venv\Scripts\meshnet-node.exe, use the full path below instead -``` - -To start a node: - -```powershell -$env:HF_HOME = "D:\DEV\models" -meshnet-node start --tracker https://ai.neuron.d-popov.com --model Qwen/Qwen2.5-0.5B-Instruct -``` - -If the wrong entry point is shadowing, invoke via the full conda path: - -```powershell -C:\Users\popov\miniforge3\Scripts\meshnet-node.exe start ` - --tracker https://ai.neuron.d-popov.com ` - --model Qwen/Qwen2.5-0.5B-Instruct -``` - -#### Option B — isolated virtualenv (fresh machine, no existing torch) - -1. Install prerequisites on Windows: - - Python 3.11 or 3.12 from - - Git for Windows from - -2. Open **PowerShell** in the cloned repo and install the node packages: - -```powershell -# Example repo path; adjust to wherever you cloned it -cd D:\DEV\workspace\REPOS\git.d-popov.com\neuron-tai - -python -m venv .venv -.\.venv\Scripts\python.exe -m pip install --upgrade pip setuptools wheel -.\.venv\Scripts\pip.exe install -e packages\tracker -e packages\node -e packages\p2p -e packages\gateway -e packages\relay - -# CPU-only PyTorch. For NVIDIA CUDA, use `pip install torch` instead. -.\.venv\Scripts\pip.exe install torch --index-url https://download.pytorch.org/whl/cpu -.\.venv\Scripts\pip.exe install "transformers>=5.12" accelerate - -.\.venv\Scripts\meshnet-node.exe --help -``` - -For `start`-specific flags, run: - -```powershell -.\.venv\Scripts\meshnet-node.exe start --help -``` - -3. Find the Windows LAN IP address: - -```powershell -ipconfig -``` - -Use the IPv4 address on the active Ethernet/Wi-Fi adapter, for example -`192.168.0.42`. Avoid WSL/Docker/Hyper-V adapter addresses like `172.16.x.x`, -`172.17.x.x`, or other virtual adapter IPs. - -4. Allow inbound traffic for the node port in Windows Firewall. Run PowerShell as -Administrator once: - -```powershell -New-NetFirewallRule ` - -DisplayName "Meshnet node 8005" ` - -Direction Inbound ` - -Action Allow ` - -Protocol TCP ` - -LocalPort 8005 -``` - -5. Start the Windows node from normal PowerShell. Replace the tracker and -advertised host values with your actual LAN addresses: - -```powershell -$env:HF_HOME = "D:\DEV\models" - -.\.venv\Scripts\meshnet-node.exe start ` - --tracker http://192.168.0.179:8081 ` - --model Qwen/Qwen2.5-0.5B-Instruct ` - --shard-start 12 --shard-end 23 ` - --quantization bfloat16 ` - --host 0.0.0.0 ` - --advertise-host 192.168.0.42 ` - --port 8005 -``` - -One-line variants (direct LAN — node must be reachable by IP from other machines): - -```powershell -.\.venv\Scripts\meshnet-node.exe start --tracker http://192.168.0.179:8081 --model Qwen/Qwen2.5-0.5B-Instruct --advertise-host 192.168.0.20 -``` - -Via public hostname with relay (works from behind NAT, WSL2, 5G — no `--advertise-host` needed): - -```powershell -.\.venv\Scripts\meshnet-node.exe start --tracker https://ai.neuron.d-popov.com --model Qwen/Qwen2.5-0.5B-Instruct -``` - -WSL (same relay path — no `--advertise-host`): - -```bash -.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 http://192.168.0.179:8081 --model Qwen/Qwen2.5-0.5B-Instruct -``` - -`--host 0.0.0.0` binds the node to all Windows interfaces. `--advertise-host` -is what the tracker gives to other nodes for direct connections; omit it when using -the relay path since all traffic flows through the relay tunnel instead. - -If you want verbose per-hop pipeline logs while debugging a split model, add -`--debug`. Leave it off for normal runs; otherwise every generated token logs -lines like: - -```text - [node] pipeline hop 0: http://127.0.0.1:8005 start_layer=22 - [node] pipeline hop 0 returned text=' token' - [node] pipeline hop 1: wss://ai.neuron.d-popov.com/rpc/abc123 relay start_layer=12 -``` - -6. From the tracker machine, verify Windows is reachable: - -```bash -curl http://192.168.0.42:8005/v1/health -``` - -If that endpoint returns 404 or 501, that is okay: it still proves the TCP -connection reached the node process. If it times out or connection-refuses, check -the Windows Firewall rule, `--host 0.0.0.0`, the selected LAN IP, and that the -node is still running. - ---- - -## Public tracker + relay (internet / NAT nodes) - -This setup lets nodes connect from anywhere — behind home NAT, 5G, WSL2, or -on a different continent — without opening firewall ports. - -### Architecture - -``` -Client → HTTPS → ai.neuron.d-popov.com (nginx) - ├─ /v1/* → meshnet-tracker :8081 - ├─ /ws → meshnet-relay :8765 (node persistent outbound WS) - └─ /rpc/* → meshnet-relay :8765 (caller opens WS per hop) -``` - -### Nginx Proxy Manager (Docker) - -Use **one** proxy host for the domain. Do not create a second host for the same -domain to reach another port — path routing is done with **Custom locations** on -that same host. - -**1. Details tab** (default `/` → tracker) - -| Field | Value | -|-------|--------| -| Domain Names | `ai.neuron.d-popov.com` | -| Scheme | `http` | -| Forward Hostname / IP | LAN IP of the tracker machine (e.g. `192.168.0.179`) | -| Forward Port | `8081` | -| Websockets Support | ON | - -This serves `/v1/network/map`, `/v1/chat/completions`, and the rest of the tracker API. - -**2. Custom locations tab** (sub-paths → relay) - -The Custom locations form has **no separate Websockets toggle** — only location, -scheme, forward host, optional sub-folder path, and port. Add **two** locations -(both pointing at the relay process on port `8765`). Leave **“Add a path for -sub-folder forwarding”** empty so the full URI reaches the relay -(`/ws`, `/rpc/`). - -Location A — persistent node connections: - -| Field | Value | -|-------|--------| -| Define location | `/ws` | -| Scheme | `http` | -| Forward Hostname / IP | `192.168.0.179` | -| Forward Port | `8765` | -| Sub-folder path | *(leave empty)* | - -Location B — per-hop RPC: - -| Field | Value | -|-------|--------| -| Define location | `/rpc` | -| Scheme | `http` | -| Forward Hostname / IP | `192.168.0.179` | -| Forward Port | `8765` | -| Sub-folder path | *(leave empty)* | - -Nginx matches the longer prefixes first: `/ws` and `/rpc/…` go to relay; everything -else stays on `8081`. - -**3. SSL tab** - -Use your existing Let’s Encrypt certificate (unchanged). - -**4. Advanced tab** (only if WebSocket upgrade fails on `/ws` or `/rpc`) - -Custom locations do not expose a Websockets checkbox. If nodes show -`Relay configured but not connected yet` while `/v1/network/map` works, add this -snippet on the **proxy host** Advanced tab: - -```nginx -proxy_http_version 1.1; -proxy_set_header Upgrade $http_upgrade; -proxy_set_header Connection $http_connection; -proxy_read_timeout 3600s; -proxy_send_timeout 3600s; -``` - -**5. Verify routing** - -```bash -# Tracker (8081 via default location) -curl -s https://ai.neuron.d-popov.com/v1/network/map | python3 -m json.tool - -# Relay paths should not 502/404 from the tracker — check response headers/status -curl -sI https://ai.neuron.d-popov.com/ws -curl -sI https://ai.neuron.d-popov.com/rpc/test-peer -``` - -After NPM is correct, start relay and tracker on the LAN machine: - -```bash -# Terminal 1 — relay -.venv/bin/meshnet-relay --host 0.0.0.0 --port 8765 - -# Terminal 2 — tracker (advertises relay to nodes) -.venv/bin/meshnet-tracker start \ - --host 0.0.0.0 \ - --port 8081 \ - --relay-url wss://ai.neuron.d-popov.com/ws -``` - -Nodes using `https://ai.neuron.d-popov.com` should then log: - -```text -Relay advertised by tracker — using outbound tunnel wss://ai.neuron.d-popov.com/ws - Relay connected — wss://ai.neuron.d-popov.com/rpc/ -``` - -The `--relay-url` flag embeds the relay address in `/v1/network/map`. Every node -queries that endpoint on startup and auto-connects if a relay URL is present. - -### Start a node (any machine, any network) - -No `--advertise-host`, firewall rule, port forwarding, relay URL, or peer URL is -needed on the node. The public tracker is the only bootstrap URL the user types. -The node queries the tracker for `/v1/network/map`, discovers the relay URL, and -opens a persistent outbound WebSocket. If the relay connection drops, the node -keeps retrying it. - -```bash -.venv/bin/meshnet-node start \ - --tracker https://ai.neuron.d-popov.com \ - --model Qwen/Qwen2.5-0.5B-Instruct -``` - -No authentication is required in the prototype. The first public node for a model -must still choose that model with `--model` or a saved wizard config. After at -least one HF model node is registered, later nodes can join the public network -with only the tracker URL; the tracker assigns an uncovered shard if one exists: - -```bash -.venv/bin/meshnet-node start --tracker https://ai.neuron.d-popov.com -``` - -Use the public tracker to verify registration and routing: - -```bash -curl -s "https://ai.neuron.d-popov.com/v1/network/map" | python3 -m json.tool -curl -s "https://ai.neuron.d-popov.com/v1/route?model=qwen2.5-0.5b" | python3 -m json.tool -``` - -Expected startup output (relay path): - -``` - Auto-detected 24 layers → shard 0–23 - Relay connected — wss://ai.neuron.d-popov.com/rpc/abc1def2ef3f4567 -================================ -meshnet-node ready - Wallet:
- Model ID: Qwen/Qwen2.5-0.5B-Instruct - Shard: layers 0–23; 24 of 24 - Quantization: bfloat16 - Endpoint: http://172.29.104.23:7001 - Node ID: - Hardware: CPU -================================ -``` - -The `Endpoint` shown is the local IP (unreachable from outside). Other nodes reach -this one via `wss://ai.neuron.d-popov.com/rpc/` instead. - -### How relay hops work - -When node A needs to forward activations to node B (behind NAT): - -1. Tracker injects `X-Meshnet-Route` with `relay_addr` for each behind-NAT hop. -2. Node A opens a WebSocket to `wss://relay/rpc/{peer_id_B}`. -3. Relay forwards the `relay-http-request` envelope to Node B's persistent connection. -4. Node B processes `/forward` locally, returns `relay-http-response`. -5. Relay sends the response back to Node A over the same WebSocket. -6. Node A closes the WebSocket and continues the pipeline. - -Binary activation tensors (bfloat16) are Base64-encoded through the relay JSON -protocol and decoded on both sides — no precision loss. - -If the relay hop fails (relay down, peer disconnected), the node logs a warning and -falls back to a direct HTTP attempt before returning an error. - -### Test from WSL2 using the public tracker - -In WSL2 (which gets a `172.x.x.x` virtual IP — unreachable from other machines): - -```bash -# WSL2 Terminal 1 — head node (layers 0–11, handles chat requests) -.venv/bin/meshnet-node start \ - --tracker https://ai.neuron.d-popov.com \ - --model Qwen/Qwen2.5-0.5B-Instruct \ - --shard-start 0 --shard-end 11 - -# WSL2 Terminal 2 — tail node (layers 12–23) -.venv/bin/meshnet-node start \ - --tracker https://ai.neuron.d-popov.com \ - --model Qwen/Qwen2.5-0.5B-Instruct \ - --shard-start 12 --shard-end 23 -``` - -Both nodes connect to the relay automatically. When a chat request arrives at Node A, -it forwards activations to Node B via `wss://ai.neuron.d-popov.com/rpc/{peer_id_B}`. - -Send inference through the tracker (which picks the head node and injects the route): - -```bash -curl -s https://ai.neuron.d-popov.com/v1/chat/completions \ - -H "Content-Type: application/json" \ - -H "Authorization: Bearer sk-mesh-" \ - -d '{ - "model": "Qwen/Qwen2.5-0.5B-Instruct", - "messages": [{"role": "user", "content": "What is 7 times 8?"}], - "stream": false - }' | python3 -m json.tool -``` - -Or send directly to Node A's local port (within WSL): - -```bash -curl -s http://localhost:7001/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{"model": "Qwen/Qwen2.5-0.5B-Instruct", "messages": [{"role": "user", "content": "Hi"}]}' -``` - -## Accounts, API keys, and credit (billing-enabled trackers) - -Public trackers run with billing on: `/v1/chat/completions` requires a real -API key from a registered account. Unknown bearer strings get `401`; a key -with no balance gets `402 insufficient balance`. - -**Dashboard flow (easiest):** open `https:///dashboard`, register with -an email + password, then click **+ new key**. The key (`sk-mesh-…`) shows its -balance next to it. If the tracker was started with `--starting-credit`, your -first key arrives pre-funded (Caller Credit, once per account). If it was -started with `--devnet-topup`, every key row has a **+$N (devnet)** button to -refill during testing. - -**Curl flow:** - -```bash -# 1. Register (once) -curl -s https:///v1/auth/register \ - -H "Content-Type: application/json" \ - -d '{"email": "you@example.com", "password": "hunter22-or-better"}' -# → {"session_token": "...", ...} - -# 2. Create an API key (session token from step 1) -curl -s https:///v1/account/keys -X POST \ - -H "Authorization: Bearer " -# → {"api_key": "sk-mesh-...", "caller_credit_granted": true} - -# 3. Check balance / usage -curl -s https:///v1/account -H "Authorization: Bearer " - -# 4. (devnet trackers only) top up a key -curl -s https:///v1/account/topup -X POST \ - -H "Authorization: Bearer " \ - -H "Content-Type: application/json" \ - -d '{"api_key": "sk-mesh-..."}' -``` - -Operator side: both features default to 1 USDT (`--starting-credit` / -`--devnet-topup`). Set both to 0 on mainnet deployments — real deposits flow -through the on-chain USDT treasury watcher instead. - ---- - -## Step 1 — Start the tracker (Terminal 1) - -```bash -cd /run/media/popov/d/DEV/repos/d-popov.com/AI -.venv/bin/meshnet-tracker start --port 8080 -``` - -Expected output: -``` -Tracker listening on 0.0.0.0:8080 -``` - -Keep this terminal open. - ---- - -## Step 2 — Start a node (Terminal 2) - -### Recommended model: Qwen2.5-0.5B-Instruct - -- 0.5B parameters, ~1 GB in BF16 -- No HuggingFace account or license required -- Downloads once to `~/.meshnet/models/`, cached for future runs -- 24 transformer layers (auto-detected — no need to specify) - -```bash -cd /run/media/popov/d/DEV/repos/d-popov.com/AI -HF_HOME=/run/media/popov/d/DEV/models \ -.venv/bin/meshnet-node start \ - --model Qwen/Qwen2.5-0.5B-Instruct \ - --quantization bfloat16 \ - --tracker http://localhost:8080 \ - --port 8001 -``` - -Shard range is **auto-detected** from the curated catalog (no network call for known -models). For unknown repos, the node fetches only `config.json` (~1 KB) to read -`num_hidden_layers`. You can still pass `--shard-start` / `--shard-end` explicitly -to run a partial shard on one machine. - -Expected output (after model loads): -``` - Auto-detected 24 layers → shard 0–23 -================================ -meshnet-node ready - Wallet:
- Model ID: Qwen/Qwen2.5-0.5B-Instruct - Shard: layers 0–23 - Quantization: bfloat16 - Endpoint: http://:8001 - Hardware: CPU -================================ -``` - -### Other model options (all CPU-friendly) - -| Model | HF repo | Layers | BF16 size | Notes | -|-------|---------|--------|-----------|-------| -| Qwen2.5-0.5B | `Qwen/Qwen2.5-0.5B-Instruct` | 24 | ~1 GB | Fastest, no gating | -| Qwen2.5-1.5B | `Qwen/Qwen2.5-1.5B-Instruct` | 28 | ~3 GB | Better quality | -| Phi-3-mini | `microsoft/Phi-3-mini-4k-instruct` | 32 | ~7.5 GB | Best CPU quality | -| Llama-3.2-1B | `meta-llama/Llama-3.2-1B-Instruct` | 16 | ~2 GB | Requires HF login | -| Llama-3.2-3B | `meta-llama/Llama-3.2-3B-Instruct` | 28 | ~6 GB | Requires HF login | - -For gated models (Llama), run `huggingface-cli login` first. - ---- - -## Step 3 — Send an inference request (Terminal 3) - -If you started the node with `--port 8001`, send the request directly to that -head node: - -```bash Qwen2.5-0.5B-Instruct -curl -s http://localhost:8001/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "qwen2.5-0.5b", - "messages": [{"role": "user", "content": "What is 7 times 8? Answer in one word."}], - "stream": false - }' | python3 -m json.tool -``` - -If you did not pass `--port`, `meshnet-node start` uses the first free port at -or above `7000`. Use the `Endpoint:` printed by the node instead of `8001`. - -To test tracker routing/proxying, send the same OpenAI-compatible request to the -tracker, using either the full HuggingFace repo or the quick alias: - -```bash -curl -s http://localhost:8080/v1/chat/completions \ - -H "Content-Type: application/json" \ - -d '{ - "model": "qwen2.5-0.5b", - "messages": [{"role": "user", "content": "What is 7 times 8? Answer in one word."}], - "stream": false - }' | python3 -m json.tool -``` - -Or use the test script: - -```bash -.venv/bin/python scripts/test_lan_inference.py \ - --tracker http://localhost:8080 \ - --gateway http://localhost:8001 -``` - ---- - -## Two-node split (same machine, two terminals) - -Split Qwen2.5-0.5B's 24 layers across two node processes to test the sharded pipeline: - -**Node A — layers 0–11 (tracker mode, serves chat completions):** -```bash -HF_HOME=/run/media/popov/d/DEV/models \ -.venv/bin/meshnet-node start \ - --model Qwen/Qwen2.5-0.5B-Instruct \ - --shard-start 0 --shard-end 11 \ - --quantization bfloat16 \ - --tracker http://localhost:8080 \ - --port 8001 -``` - -**Node B — layers 12–23:** -```bash -HF_HOME=/run/media/popov/d/DEV/models \ -.venv/bin/meshnet-node start \ - --model Qwen/Qwen2.5-0.5B-Instruct \ - --shard-start 12 --shard-end 23 \ - --quantization bfloat16 \ - --tracker http://localhost:8080 \ - --port 8002 -``` - -Send the request to Node A — it tokenizes, runs layers 0–11, passes binary -activations to Node B, and streams the final response back. - ---- - -## Two-machine LAN test (Linux + Windows/WSL2) - -See `docs/TWO_MACHINE_TEST.md` (created by US-018). - -For WSL2 nodes, registration only proves the node can reach the tracker -outbound. Tracker-routed inference also requires the tracker to reach the node's -advertised endpoint inbound. Either run the node in native Windows PowerShell, -configure Windows port forwarding into WSL for the node port, or start the -tracker with a relay URL so the node registers a `relay_addr`. - ---- - -## Browse available models - -```bash -# Show curated list with VRAM requirements -.venv/bin/meshnet-node models - -# Browse HuggingFace Hub top-20 text-generation models -.venv/bin/meshnet-node models --browse -``` - ---- - -## Start with the interactive wizard - -```bash -# First run: wizard detects GPU, shows model list, saves config -.venv/bin/meshnet-node - -# Subsequent runs: starts directly from saved config -.venv/bin/meshnet-node - -# Re-run wizard even with saved config -.venv/bin/meshnet-node --reset-config -``` - ---- - -## Run all tests - -```bash -.venv/bin/python -m pytest -q -``` +# Quickstart — Running a node and testing inference + +Get from zero to a live inference request in **three terminals**: install once, start +the tracker, start a node, send a request. + +Tested on: AMD Ryzen AI Max (Strix Halo APU), 124 GB RAM, Linux, CPU inference. + +**Active development models** (what we run day-to-day): + +| Role | `--model` / alias | HF repo | Notes | +|------|-------------------|---------|-------| +| Smoke tests, small splits | `Qwen/Qwen2.5-0.5B-Instruct` | same | 24 layers, ~1 GB BF16, no gating — default for new setups | +| Alpha / production target | `qwen3.6-35b-a3b` | `unsloth/Qwen3.6-35B-A3B` | 40 layers, ~72 GB BF16, hybrid linear-attention MoE; aliases include `Qwen3.6-35B-A3B`, `Qwen/Qwen3.6-35B-A3B` | + +--- + +## At a glance + +| Step | What | Terminals | +|------|------|-----------| +| **0** | Install Python packages | once per machine | +| **1** | Start tracker (and relay if needed) | 1–2 | +| **2** | Start node(s) | 1+ | +| **3** | Send inference request | 1 | + +**Pick your connectivity mode** — this determines which flags you need on the node: + +| Mode | When to use | Tracker URL | Node extras | +|------|-------------|-------------|-------------| +| **Local dev** | Everything on one machine | `http://localhost:8080` | none | +| **Direct LAN** | Node has a real LAN IP other machines can reach | `http://:8080` | `--host 0.0.0.0 --advertise-host ` + firewall | +| **Relay / public** | WSL2, NAT, 5G, or any unreachable inbound port | `https://ai.neuron.d-popov.com` (or your public URL) | none — relay handles routing | + +> **WSL2:** not reachable from other LAN machines by default. Use the **relay / public** +> tracker URL, or run the node in native Windows PowerShell with direct LAN mode. + +**Command prefix by shell** (used in examples below): + +| Shell | Prefix | Model cache env | +|-------|--------|-----------------| +| Linux / WSL | `.venv/bin/` | `HF_HOME=/path/to/models` | +| Windows PowerShell | `.\.venv\Scripts\` | `$env:HF_HOME = "D:\DEV\models"` | + +--- + +## 0. Install prerequisites (once per machine) + +Editable installs point wrappers at this source tree — code edits apply without +reinstalling. + +### Node machine — full install + +
+Linux / WSL + +```bash +cd /path/to/neuron-tai +python3 -m venv .venv +source .venv/bin/activate +.venv/bin/python -m pip install --upgrade pip setuptools wheel +.venv/bin/python -m pip install -e packages/tracker -e packages/node -e packages/p2p -e packages/gateway -e packages/relay +.venv/bin/python -m pip install "transformers>=5.12" accelerate +.venv/bin/meshnet-node --help +``` + +
+ +
+Windows PowerShell (.venv) + +Requires Python 3.11+ and Git for Windows. + +```powershell +cd D:\DEV\workspace\REPOS\git.d-popov.com\neuron-tai +python -m venv .venv +.\.venv\Scripts\Activate.ps1 +.\.venv\Scripts\python.exe -m pip install --upgrade pip setuptools wheel +.\.venv\Scripts\python.exe -m pip install -e .\packages\tracker -e .\packages\node -e .\packages\p2p -e .\packages\gateway -e .\packages\relay +.\.venv\Scripts\python.exe -m pip install "transformers>=5.12" accelerate +.\.venv\Scripts\meshnet-node.exe --help +``` + +
+ +
+Windows — conda/miniforge with CUDA (skip if using .venv above) + +```powershell +conda activate base +deactivate # drop any layered .venv; safe no-op if none active +cd D:\DEV\workspace\REPOS\git.d-popov.com\neuron-tai +pip install -e packages\tracker -e packages\node -e packages\p2p -e packages\gateway -e packages\relay +pip install "transformers>=5.12" accelerate safetensors +python -c "import torch; print(torch.__version__, torch.cuda.is_available())" +# Expected: 2.x.x+cuXXX True +``` + +If `torch` import fails despite pip saying "already satisfied", force-reinstall all +three together (never upgrade `torch` alone — breaks `torchvision`): + +```powershell +pip install --force-reinstall torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 +``` + +If `.venv\Scripts\meshnet-node.exe` shadows the conda binary, use the full path: +`C:\Users\\miniforge3\Scripts\meshnet-node.exe`. + +
+ +> Run Linux/WSL commands from **WSL**, not Git Bash. From Git Bash: `wsl`, then `cd` +> to the repo under `/mnt/d/...`. + +### Tracker host — lightweight install + +Tracker + relay only; skip node packages unless this machine also runs nodes. + +```bash +git clone https://git.d-popov.com/popov/neuron-tai.git AI && cd AI +python3 -m venv .venv +.venv/bin/python -m pip install --upgrade pip setuptools wheel +.venv/bin/pip install -e packages/tracker -e packages/relay -e packages/gateway +``` + +### PyTorch variant + +Install **one** torch line into the same env as `meshnet-node`: + +| Hardware | Install | +|----------|---------| +| NVIDIA CUDA | `pip install torch` (default index) | +| CPU only | `pip install torch --index-url https://download.pytorch.org/whl/cpu` | +| AMD ROCm | `pip install torch --index-url https://download.pytorch.org/whl/rocm6.2` | + +On Windows `.venv`, prefix with `.\.venv\Scripts\pip.exe`. Conda users with CUDA +torch already installed can skip this step. + +### Qwen3.5/3.6-MoE notes + +Applies to **`qwen3.6-35b-a3b`** and other hybrid linear-attention models. **`Qwen2.5-0.5B`** +does not need any of this — it is a standard transformer with no FLA fast path. + +- **transformers ≥ 5.12 required** — older versions fail with + `'Qwen3_5MoeConfig' object has no attribute 'vocab_size'`. Check: + `python -c "import transformers; print(transformers.__version__)"`. +- **GPU fast path (optional)** — without it inference still works; startup prints + `The fast path is not available…` and linear-attention layers use a slower PyTorch + fallback. Install **only for your platform**: + + | Platform | Install | Notes | + |----------|---------|-------| + | **Native Windows + NVIDIA** | `pip install triton-windows` then `pip install flash-linear-attention` | **Fast path works.** FLA [officially supports `triton-windows`](https://github.com/fla-org/flash-linear-attention/pull/757) (tested Win11, PyTorch 2.10, triton-windows 3.6). Do **not** use the `[cuda]` extra on Windows — pip looks for Linux `triton` and fails. Do **not** install `causal-conv1d` — FLA ≥0.3.2 ships Triton conv1d; the separate package is Linux-only and breaks on Windows (`bare_metal_version` / nvcc errors). | + | **Linux + NVIDIA CUDA** | `pip install flash-linear-attention[cuda]` | `causal-conv1d` optional (same FLA built-in conv1d note). Needs CUDA toolkit (`nvcc`) matching torch, or a prebuilt wheel. | + | **Linux + AMD ROCm** | `pip install flash-linear-attention[rocm]` | Same optional `causal-conv1d` note. | + + **Windows verify** (after install): + + ```powershell + python -c "import triton; import fla; print('triton', triton.__version__, 'fla ok')" + ``` + + `triton-windows` is also pulled by `meshnet-node` on Windows. Without it, Qwen3.6-MoE + startup fails with misleading `Could not import module 'Qwen3_5MoeForCausalLM'`. + +
+Windows fast path — what failed and what actually works + +The command that failed — `pip install flash-linear-attention[cuda] causal-conv1d` — mixes +two different things: + +1. **`flash-linear-attention[cuda]` on Windows** — wrong extra. `[cuda]` pulls PyPI + `triton>=3.3`, which does not exist for Windows (`No matching distribution found`). + Use plain `pip install flash-linear-attention` **after** `triton-windows` is already + installed; FLA detects `triton-windows` and uses it. +2. **`causal-conv1d`** — separate Dao-AILab CUDA extension, **not required** for FLA or + Qwen3.6 when FLA is installed. No official Windows wheels. Source builds need `nvcc` + whose major version matches torch's CUDA (e.g. torch `+cu118` needs CUDA 11.8 toolkit, + not 12.5). Community wheels exist for narrow Python/torch combos + ([PR #46](https://github.com/Dao-AILab/causal-conv1d/pull/46)) but we skip them. + +**Working Windows stack** (confirmed on this repo's dev machine: Python 3.12, torch +2.7.1+cu118, triton-windows 3.7.1, flash-linear-attention 0.5.0): + +```powershell +pip install triton-windows +pip install -U flash-linear-attention +python -c "import triton; import fla; print('ok')" +``` + +If the fast-path warning persists after that, upgrade FLA to ≥0.5.1 (includes the +`triton-windows` detection from PR #757) and restart the node. + +
+ +- `pip install nvidia-ml-py` silences the pynvml deprecation warning on NVIDIA hosts. + +--- + +## 1. Start the tracker + +### LAN tracker (private network, direct node reachability) + +**Terminal 1:** + +```bash +.venv/bin/meshnet-tracker start --host 0.0.0.0 --port 8080 +# Optional devnet billing: --starting-credit 1 --devnet-topup 10 +``` + +Expected: `Tracker listening on 0.0.0.0:8080`. Open the port on the host firewall +if other machines will join. + +**Verify:** + +```bash +curl -s http://localhost:8080/v1/network/map | python3 -m json.tool +curl -s http://192.168.0.179:8080/v1/network/map | python3 -m json.tool # from another LAN machine +``` + +### Public tracker + relay (NAT / WSL2 / internet nodes) + +Nodes behind NAT cannot receive inbound connections. Run **both** services on the +tracker host; the tracker advertises the relay URL in `/v1/network/map`. + +**Terminal 1 — relay:** + +```bash +.venv/bin/meshnet-relay --host 0.0.0.0 --port 8765 +``` + +**Terminal 2 — tracker:** + +```bash +.venv/bin/meshnet-tracker start --host 0.0.0.0 --port 8081 --relay-url wss://ai.neuron.d-popov.com/ws +``` + +**Verify:** + +```bash +curl -s https://ai.neuron.d-popov.com/v1/network/map | python3 -m json.tool +``` + +Nodes should log `Relay connected — wss://…/rpc/` on startup. + +
+Nginx Proxy Manager setup (public hostname) + +Architecture: + +``` +Client → HTTPS → ai.neuron.d-popov.com (nginx) + ├─ /v1/* → meshnet-tracker :8081 + ├─ /ws → meshnet-relay :8765 (node persistent outbound WS) + └─ /rpc/* → meshnet-relay :8765 (caller opens WS per hop) +``` + +Use **one** proxy host. Route sub-paths via **Custom locations** — do not create +a second host for the same domain. + +**Details tab** (default `/` → tracker): + +| Field | Value | +|-------|--------| +| Domain Names | `ai.neuron.d-popov.com` | +| Scheme | `http` | +| Forward Hostname / IP | LAN IP of tracker machine (e.g. `192.168.0.179`) | +| Forward Port | `8081` | +| Websockets Support | ON | + +**Custom locations** (both → relay port `8765`, sub-folder path empty): + +| Location | Forward to | +|----------|--------------| +| `/ws` | `192.168.0.179:8765` | +| `/rpc` | `192.168.0.179:8765` | + +**Advanced tab** (only if WebSocket upgrade fails): + +```nginx +proxy_http_version 1.1; +proxy_set_header Upgrade $http_upgrade; +proxy_set_header Connection $http_connection; +proxy_read_timeout 3600s; +proxy_send_timeout 3600s; +``` + +**Verify routing:** + +```bash +curl -s https://ai.neuron.d-popov.com/v1/network/map | python3 -m json.tool +curl -sI https://ai.neuron.d-popov.com/ws +curl -sI https://ai.neuron.d-popov.com/rpc/test-peer +``` + +
+ +--- + +## 2. Start a node + +**Starter model:** `Qwen/Qwen2.5-0.5B-Instruct` — 0.5B params, ~1 GB BF16, 24 layers, +no HuggingFace gating. Best for first-time setup. + +**Alpha model:** `qwen3.6-35b-a3b` — 40 layers, ~72 GB BF16 download, MoE with hybrid +linear attention. On Windows install `triton-windows` + `flash-linear-attention`; on Linux +GPU use `flash-linear-attention[cuda]`. Tracker accepts the alias or full repo id (`unsloth/Qwen3.6-35B-A3B`). + +Downloads cache under `~/.meshnet/models/` (or `$HF_HOME` / `$env:HF_HOME`). + +Shard range is auto-detected from the curated catalog. For unknown repos the node +fetches only `config.json`. Override with `--shard-start` / `--shard-end` for partial +shards or multi-node splits. + +### Core command + +Replace `` and adjust the prefix for your shell (see table above). + +**Linux / WSL:** + +```bash +HF_HOME=/path/to/models .venv/bin/meshnet-node start --tracker --model Qwen/Qwen2.5-0.5B-Instruct --quantization bfloat16 +``` + +**Windows PowerShell:** + +```powershell +$env:HF_HOME = "D:\DEV\models" +.\.venv\Scripts\meshnet-node.exe start --tracker --model Qwen/Qwen2.5-0.5B-Instruct --quantization bfloat16 +``` + +### Ready-to-run examples + +**Local dev (same machine as tracker):** + +```bash +HF_HOME=/path/to/models .venv/bin/meshnet-node start --tracker http://localhost:8080 --model Qwen/Qwen2.5-0.5B-Instruct --quantization bfloat16 --port 8001 +``` + +**Public / relay (works from WSL2, NAT, 5G — no extra flags):** + +```bash +.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):** + +```powershell +$env:HF_HOME = "D:\DEV\models" +pip install triton-windows +pip install -U flash-linear-attention +meshnet-node start --tracker http://192.168.0.179:8080 --model qwen3.6-35b-a3b --quantization bfloat16 +``` + +Do not add `causal-conv1d` or `flash-linear-attention[cuda]` on Windows (see Qwen3.5/3.6 notes). + +**Alpha model (Qwen3.6, Linux GPU — with fast path):** + +```bash +HF_HOME=/path/to/models .venv/bin/meshnet-node start --tracker --model qwen3.6-35b-a3b --quantization bfloat16 +# Install once on that machine: pip install flash-linear-attention[cuda] +``` + +After the first node registers a model, later nodes can join with only the tracker +URL (shard auto-assigned): + +```bash +.venv/bin/meshnet-node start --tracker https://ai.neuron.d-popov.com +``` + +**Direct LAN (Windows node reachable by IP):** + +```powershell +$env:HF_HOME = "D:\DEV\models" +.\.venv\Scripts\meshnet-node.exe start --tracker http://192.168.0.179:8081 --model Qwen/Qwen2.5-0.5B-Instruct --shard-start 12 --shard-end 23 --quantization bfloat16 --host 0.0.0.0 --advertise-host 192.168.0.42 --port 8005 +``` + +
+Windows direct LAN — firewall and IP checklist + +1. Find LAN IP: `ipconfig` — use active Ethernet/Wi-Fi IPv4 (e.g. `192.168.0.42`). + Avoid WSL/Docker/Hyper-V addresses (`172.x.x.x`). +2. Allow inbound port (Administrator PowerShell, once): + +```powershell +New-NetFirewallRule -DisplayName "Meshnet node 8005" -Direction Inbound -Action Allow -Protocol TCP -LocalPort 8005 +``` + +3. Verify from tracker machine: + +```bash +curl http://192.168.0.42:8005/v1/health +``` + +404/501 is fine — it proves TCP reached the node. Timeout = check firewall, +`--host 0.0.0.0`, and `--advertise-host`. + +
+ +
+Two-node split (same or different machines) + +**Node A — layers 0–11 (head, serves chat):** + +```bash +HF_HOME=/path/to/models .venv/bin/meshnet-node start --tracker --model Qwen/Qwen2.5-0.5B-Instruct --shard-start 0 --shard-end 11 --quantization bfloat16 --port 8001 +``` + +**Node B — layers 12–23:** + +```bash +HF_HOME=/path/to/models .venv/bin/meshnet-node start --tracker --model Qwen/Qwen2.5-0.5B-Instruct --shard-start 12 --shard-end 23 --quantization bfloat16 --port 8002 +``` + +Send inference to Node A. For cross-machine LAN tests see `docs/TWO_MACHINE_TEST.md`. + +
+ +### Useful flags + +| Flag | Purpose | +|------|---------| +| `--port 8001` | Fixed listen port (default: first free ≥ 7000) | +| `--host 0.0.0.0` | Bind all interfaces (needed for direct LAN) | +| `--advertise-host ` | LAN IP the tracker tells other nodes (direct LAN only) | +| `--shard-start N --shard-end M` | Partial layer range | +| `--debug` | Verbose per-hop pipeline logs (noisy; off by default) | + +`--host 0.0.0.0` binds locally; `--advertise-host` is what peers use for direct +hops. Omit both when using the relay path. + +### Expected output + +``` + Auto-detected 24 layers → shard 0–23 + Relay connected — wss://ai.neuron.d-popov.com/rpc/abc1def2ef3f4567 # relay mode only +================================ +meshnet-node ready + Wallet:
+ Model ID: Qwen/Qwen2.5-0.5B-Instruct + Shard: layers 0–23; 24 of 24 + Quantization: bfloat16 + Endpoint: http://:8001 + Node ID: + Hardware: CPU +================================ +``` + +The `Endpoint` is the local address. In relay mode, peers reach this node via +`wss:///rpc/` instead. + +### Other CPU-friendly models + +| Model | `--model` / alias | Layers | BF16 size | Notes | +|-------|-------------------|--------|-----------|-------| +| **Qwen2.5-0.5B** (dev default) | `Qwen/Qwen2.5-0.5B-Instruct` | 24 | ~1 GB | Fastest, no gating | +| **Qwen3.6-35B-A3B** (alpha) | `qwen3.6-35b-a3b` | 40 | ~72 GB | MoE; needs transformers ≥5.12; see Qwen3.5/3.6 notes | +| Qwen2.5-1.5B | `Qwen/Qwen2.5-1.5B-Instruct` | 28 | ~3 GB | Better quality | +| Phi-3-mini | `microsoft/Phi-3-mini-4k-instruct` | 32 | ~7.5 GB | Best CPU quality | +| Llama-3.2-1B | `meta-llama/Llama-3.2-1B-Instruct` | 16 | ~2 GB | Requires HF login | +| Llama-3.2-3B | `meta-llama/Llama-3.2-3B-Instruct` | 28 | ~6 GB | Requires HF login | + +For gated models (Llama): `huggingface-cli login` first. + +Browse more: `.venv/bin/meshnet-node models` or `.venv/bin/meshnet-node models --browse`. + +--- + +## 3. Send an inference request + +**Terminal 3** — direct to the head node (replace port if you omitted `--port`): + +```bash +curl -s http://localhost:8001/v1/chat/completions -H "Content-Type: application/json" -d '{"model": "qwen2.5-0.5b", "messages": [{"role": "user", "content": "What is 7 times 8? Answer in one word."}], "stream": false}' | python3 -m json.tool +``` + +**Via tracker** (tests routing / proxying): + +```bash +curl -s http://localhost:8080/v1/chat/completions -H "Content-Type: application/json" -d '{"model": "qwen2.5-0.5b", "messages": [{"role": "user", "content": "What is 7 times 8? Answer in one word."}], "stream": false}' | python3 -m json.tool +``` + +**Public tracker:** + +```bash +curl -s https://ai.neuron.d-popov.com/v1/chat/completions -H "Content-Type: application/json" -H "Authorization: Bearer sk-mesh-" -d '{"model": "Qwen/Qwen2.5-0.5B-Instruct", "messages": [{"role": "user", "content": "What is 7 times 8?"}], "stream": false}' | python3 -m json.tool +``` + +**Test script:** + +```bash +.venv/bin/python scripts/test_lan_inference.py --tracker http://localhost:8080 --gateway http://localhost:8001 +``` + +**Verify registration on public tracker:** + +```bash +curl -s "https://ai.neuron.d-popov.com/v1/network/map" | python3 -m json.tool +curl -s "https://ai.neuron.d-popov.com/v1/route?model=qwen2.5-0.5b" | python3 -m json.tool +``` + +
+Accounts, API keys, and credit (billing-enabled trackers) + +Public trackers require a real API key for `/v1/chat/completions`. Unknown bearer → +`401`; zero balance → `402 insufficient balance`. + +**Dashboard:** open `https:///dashboard`, register, click **+ new key**. +With `--starting-credit`, the first key is pre-funded. With `--devnet-topup`, use +**+$N (devnet)** to refill during testing. + +**Curl flow:** + +```bash +curl -s https:///v1/auth/register -H "Content-Type: application/json" -d '{"email": "you@example.com", "password": "hunter22-or-better"}' +curl -s https:///v1/account/keys -X POST -H "Authorization: Bearer " +curl -s https:///v1/account -H "Authorization: Bearer " +curl -s https:///v1/account/topup -X POST -H "Authorization: Bearer " -H "Content-Type: application/json" -d '{"api_key": "sk-mesh-..."}' +``` + +Operator defaults: `--starting-credit` and `--devnet-topup` both default to 1 USDT. +Set both to 0 on mainnet. + +
+ +--- + +## Reference + +### How relay hops work + +When node A forwards activations to node B (behind NAT): + +1. Tracker injects `X-Meshnet-Route` with `relay_addr` for behind-NAT hops. +2. Node A opens WebSocket to `wss://relay/rpc/{peer_id_B}`. +3. Relay forwards `relay-http-request` to Node B's persistent connection. +4. Node B processes `/forward`, returns `relay-http-response`. +5. Relay sends response back to Node A; Node A continues the pipeline. + +Activations (bfloat16) are Base64-encoded in JSON — no precision loss. On relay +failure the node logs a warning and falls back to direct HTTP before erroring. + +### Interactive wizard + +```bash +.venv/bin/meshnet-node # first run: wizard; later: saved config +.venv/bin/meshnet-node --reset-config +``` + +### Run all tests + +```bash +.venv/bin/python -m pytest -q +``` diff --git a/docs/INSTALL_WINDOWS.md b/docs/INSTALL_WINDOWS.md index 6775845..b7730dc 100644 --- a/docs/INSTALL_WINDOWS.md +++ b/docs/INSTALL_WINDOWS.md @@ -1,220 +1,230 @@ -# Installing meshnet-node on Windows 11 with WSL2 - -This guide covers setting up a meshnet-node on a Windows 11 machine using WSL2 with CUDA passthrough so it can join an existing inference network over LAN. - -## Prerequisites - -- Windows 11 with WSL2 support (most systems with Windows 10 version 2004+ qualify) -- NVIDIA GPU with CUDA support (driver ≥ 525.x recommended for WSL2 CUDA) -- At least 8 GB RAM + enough VRAM for the model shard you intend to serve -- The Linux machine (other node) is reachable on your LAN - ---- - -## Step 1 — Enable WSL2 and install Ubuntu - -Open **PowerShell as Administrator** and run: - -```powershell -wsl --install -d Ubuntu-24.04 -``` - -This installs WSL2 with Ubuntu 24.04. Reboot when prompted. - -After reboot, Ubuntu starts and asks you to create a UNIX username/password. Choose anything convenient. - -Verify WSL version: - -```powershell -wsl -l -v -``` - -Output should show `VERSION 2`. - ---- - -## Step 2 — Install NVIDIA GPU driver on Windows (NOT inside WSL) - -WSL2 CUDA passthrough works through the Windows host driver. **Do not install CUDA inside WSL2.** - -1. Download the latest Game Ready or Studio driver for your GPU from https://www.nvidia.com/drivers -2. Install on Windows normally (standard installer). -3. Inside WSL2 (Ubuntu terminal), verify: - -```bash -nvidia-smi -``` - -Expected output: your GPU name, driver version, CUDA version. If this command fails, the Windows driver is too old — update it. - -> **Note:** The `cuda-toolkit` package inside WSL2 is optional and only needed if you compile CUDA kernels. For inference with `torch`, the Windows host driver is sufficient. - ---- - -## Step 3 — Install Python 3.11+ inside WSL2 - -Ubuntu 24.04 ships Python 3.12. Confirm: - -```bash -python3 --version -``` - -If it shows 3.10 or older: - -```bash -sudo add-apt-repository ppa:deadsnakes/ppa -sudo apt update -sudo apt install python3.12 python3.12-venv python3.12-dev -``` - -Install pip: - -```bash -curl -sS https://bootstrap.pypa.io/get-pip.py | python3 -``` - ---- - -## Step 4 — Clone the repository - -Inside WSL2: - -```bash -# Store the repo in the Linux filesystem (faster I/O than /mnt/c) -cd ~ -git clone https://github.com/YOUR_ORG/d-popov.com.git -cd d-popov.com/AI -``` - ---- - -## Step 5 — Create a virtualenv and install meshnet-node - -```bash -python3 -m venv .venv -source .venv/bin/activate - -# Install node + PyTorch (CUDA build) -pip install torch --index-url https://download.pytorch.org/whl/cu124 -pip install -e "packages/node[torch]" -``` - -Verify the install: - -```bash -meshnet-node --help -python -c "import transformers; print(transformers.__version__)" -``` - -`transformers` must be **≥ 5.12** for Qwen3.5/3.6-MoE models (older versions fail -with `'Qwen3_5MoeConfig' object has no attribute 'vocab_size'`). If you install -into an existing conda/miniforge env instead of a fresh venv, run -`pip install -U transformers` there. The startup warning about -`flash-linear-attention` / `causal-conv1d` ("fast path is not available") is -harmless on CPU — those are optional CUDA-only kernels. - ---- - -## Step 6 — Pre-download the model shard - -Download the model before starting the node so the startup process doesn't time out on the tracker side: - -```bash -python3 - <<'EOF' -from transformers import AutoConfig -AutoConfig.from_pretrained("microsoft/Phi-3-medium-128k-instruct") -EOF -``` - -For the full model weights (needed at runtime), `transformers` downloads them automatically on first `meshnet-node` start. If you want to pre-fetch: - -```bash -python3 -c " -from transformers import AutoModelForCausalLM -AutoModelForCausalLM.from_pretrained('microsoft/Phi-3-medium-128k-instruct', device_map='cpu') -" -``` - -This can take 10–30 minutes on first run. - ---- - -## Step 7 — Expose the node port to your LAN - -WSL2 runs behind a NAT with a virtual IP (typically `172.x.x.x`). Your LAN sees the Windows host IP. You need to forward the node port. - -**Option A — Windows port proxy (recommended for simple setups):** - -In **PowerShell as Administrator**: - -```powershell -# Get the current WSL2 IP (changes on each WSL restart) -$wslIp = (wsl hostname -I).Trim() - -# Forward Windows host port 8001 → WSL2 port 8001 -netsh interface portproxy add v4tov4 ` - listenport=8001 listenaddress=0.0.0.0 ` - connectport=8001 connectaddress=$wslIp - -# Allow inbound on Windows Firewall -New-NetFirewallRule -DisplayName "meshnet-node" ` - -Direction Inbound -Protocol TCP -LocalPort 8001 -Action Allow -``` - -Verify: from the Linux machine, `curl http://WINDOWS_LAN_IP:8001/v1/health` should return a response once the node is running. - -**Redo this after every WSL2 restart** — the WSL2 IP changes. - -**Option B — P2P relay (US-017, no port forwarding needed):** - -Start a relay node on the Linux machine. The WSL2 node connects outbound through the relay. No firewall rules needed. See `docs/TWO_MACHINE_TEST.md` for details. - ---- - -## Step 8 — Start the node - -Replace `192.168.1.10` with the actual LAN IP of the Linux machine running the tracker. -Replace shard range with the complementary range to what the Linux node is serving. - -```bash -source .venv/bin/activate - -meshnet-node \ - --model microsoft/Phi-3-medium-128k-instruct \ - --quantization bf16 \ - --shard-start 20 --shard-end 39 \ - --tracker http://192.168.1.10:8080 \ - --port 8001 \ - --host 0.0.0.0 \ - --advertise-host WINDOWS_LAN_IP -``` - -The `--advertise-host` flag tells the tracker what IP the Linux machine should use to reach this node. Use your Windows machine's LAN IP (e.g. `192.168.1.20`), **not** the WSL2 internal IP. - -Expected startup output: - -``` -Detecting hardware... - GPU: NVIDIA GeForce RTX 3080 (10240 MB VRAM) -Loading wallet... - Wallet: 5K7r... -Loading real PyTorch model shard... - Auto-detected 40 layers → shard 20–39 -================================ -meshnet-node ready - Model ID: microsoft/Phi-3-medium-128k-instruct - Shard: layers 20–39; 20 of 40 - Endpoint: http://192.168.1.20:8001 - Hardware: CUDA -================================ -``` - ---- - -## Known issues - -- **WSL2 IP changes on restart.** Always re-run the `netsh` port-proxy command after restarting WSL2 or Windows. -- **CUDA not visible in WSL2.** If `nvidia-smi` fails inside WSL2, update the Windows host GPU driver to ≥ 525.x. Installing CUDA inside WSL2 will not fix it. -- **Model download is slow.** HuggingFace downloads happen over HTTPS. Pre-fetch the model before a timed test (see Step 6). -- **Port 8001 already in use.** Change `--port` to another value and update the firewall/portproxy rules accordingly. -- **`bf16` not supported on older GPUs.** Use `--quantization int8` on Turing (RTX 20xx) cards or earlier if bfloat16 ops fail. +# Installing meshnet-node on Windows 11 with WSL2 + +This guide covers setting up a meshnet-node on a Windows 11 machine using WSL2 with CUDA passthrough so it can join an existing inference network over LAN. + +## Prerequisites + +- Windows 11 with WSL2 support (most systems with Windows 10 version 2004+ qualify) +- NVIDIA GPU with CUDA support (driver ≥ 525.x recommended for WSL2 CUDA) +- At least 8 GB RAM + enough VRAM for the model shard you intend to serve +- The Linux machine (other node) is reachable on your LAN + +--- + +## Step 1 — Enable WSL2 and install Ubuntu + +Open **PowerShell as Administrator** and run: + +```powershell +wsl --install -d Ubuntu-24.04 +``` + +This installs WSL2 with Ubuntu 24.04. Reboot when prompted. + +After reboot, Ubuntu starts and asks you to create a UNIX username/password. Choose anything convenient. + +Verify WSL version: + +```powershell +wsl -l -v +``` + +Output should show `VERSION 2`. + +--- + +## Step 2 — Install NVIDIA GPU driver on Windows (NOT inside WSL) + +WSL2 CUDA passthrough works through the Windows host driver. **Do not install CUDA inside WSL2.** + +1. Download the latest Game Ready or Studio driver for your GPU from https://www.nvidia.com/drivers +2. Install on Windows normally (standard installer). +3. Inside WSL2 (Ubuntu terminal), verify: + +```bash +nvidia-smi +``` + +Expected output: your GPU name, driver version, CUDA version. If this command fails, the Windows driver is too old — update it. + +> **Note:** The `cuda-toolkit` package inside WSL2 is optional and only needed if you compile CUDA kernels. For inference with `torch`, the Windows host driver is sufficient. + +--- + +## Step 3 — Install Python 3.11+ inside WSL2 + +Ubuntu 24.04 ships Python 3.12. Confirm: + +```bash +python3 --version +``` + +If it shows 3.10 or older: + +```bash +sudo add-apt-repository ppa:deadsnakes/ppa +sudo apt update +sudo apt install python3.12 python3.12-venv python3.12-dev +``` + +Install pip: + +```bash +curl -sS https://bootstrap.pypa.io/get-pip.py | python3 +``` + +--- + +## Step 4 — Clone the repository + +Inside WSL2: + +```bash +# Store the repo in the Linux filesystem (faster I/O than /mnt/c) +cd ~ +git clone https://github.com/YOUR_ORG/d-popov.com.git +cd d-popov.com/AI +``` + +--- + +## Step 5 — Create a virtualenv and install meshnet-node + +```bash +python3 -m venv .venv +source .venv/bin/activate + +# Install node + PyTorch (CUDA build) +pip install torch --index-url https://download.pytorch.org/whl/cu124 +pip install -e "packages/node[torch]" +``` + +Verify the install: + +```bash +meshnet-node --help +python -c "import transformers; print(transformers.__version__)" +``` + +`transformers` must be **≥ 5.12** for Qwen3.5/3.6-MoE models (older versions fail +with `'Qwen3_5MoeConfig' object has no attribute 'vocab_size'`). If you install +into an existing conda/miniforge env instead of a fresh venv, run +`pip install -U transformers` there. The startup warning about +`flash-linear-attention` / `causal-conv1d` ("fast path is not available") is +harmless on CPU — those are optional CUDA-only kernels. + +If you run the node from native Windows instead of WSL2, install the Triton shim +in the same environment: + +```powershell +python -m pip install triton-windows +``` + +Without it, Qwen3.5/3.6-MoE startup can fail with the misleading message +`Could not import module 'Qwen3_5MoeForCausalLM'`. + +--- + +## Step 6 — Pre-download the model shard + +Download the model before starting the node so the startup process doesn't time out on the tracker side: + +```bash +python3 - <<'EOF' +from transformers import AutoConfig +AutoConfig.from_pretrained("microsoft/Phi-3-medium-128k-instruct") +EOF +``` + +For the full model weights (needed at runtime), `transformers` downloads them automatically on first `meshnet-node` start. If you want to pre-fetch: + +```bash +python3 -c " +from transformers import AutoModelForCausalLM +AutoModelForCausalLM.from_pretrained('microsoft/Phi-3-medium-128k-instruct', device_map='cpu') +" +``` + +This can take 10–30 minutes on first run. + +--- + +## Step 7 — Expose the node port to your LAN + +WSL2 runs behind a NAT with a virtual IP (typically `172.x.x.x`). Your LAN sees the Windows host IP. You need to forward the node port. + +**Option A — Windows port proxy (recommended for simple setups):** + +In **PowerShell as Administrator**: + +```powershell +# Get the current WSL2 IP (changes on each WSL restart) +$wslIp = (wsl hostname -I).Trim() + +# Forward Windows host port 8001 → WSL2 port 8001 +netsh interface portproxy add v4tov4 ` + listenport=8001 listenaddress=0.0.0.0 ` + connectport=8001 connectaddress=$wslIp + +# Allow inbound on Windows Firewall +New-NetFirewallRule -DisplayName "meshnet-node" ` + -Direction Inbound -Protocol TCP -LocalPort 8001 -Action Allow +``` + +Verify: from the Linux machine, `curl http://WINDOWS_LAN_IP:8001/v1/health` should return a response once the node is running. + +**Redo this after every WSL2 restart** — the WSL2 IP changes. + +**Option B — P2P relay (US-017, no port forwarding needed):** + +Start a relay node on the Linux machine. The WSL2 node connects outbound through the relay. No firewall rules needed. See `docs/TWO_MACHINE_TEST.md` for details. + +--- + +## Step 8 — Start the node + +Replace `192.168.1.10` with the actual LAN IP of the Linux machine running the tracker. +Replace shard range with the complementary range to what the Linux node is serving. + +```bash +source .venv/bin/activate + +meshnet-node \ + --model microsoft/Phi-3-medium-128k-instruct \ + --quantization bf16 \ + --shard-start 20 --shard-end 39 \ + --tracker http://192.168.1.10:8080 \ + --port 8001 \ + --host 0.0.0.0 \ + --advertise-host WINDOWS_LAN_IP +``` + +The `--advertise-host` flag tells the tracker what IP the Linux machine should use to reach this node. Use your Windows machine's LAN IP (e.g. `192.168.1.20`), **not** the WSL2 internal IP. + +Expected startup output: + +``` +Detecting hardware... + GPU: NVIDIA GeForce RTX 3080 (10240 MB VRAM) +Loading wallet... + Wallet: 5K7r... +Loading real PyTorch model shard... + Auto-detected 40 layers → shard 20–39 +================================ +meshnet-node ready + Model ID: microsoft/Phi-3-medium-128k-instruct + Shard: layers 20–39; 20 of 40 + Endpoint: http://192.168.1.20:8001 + Hardware: CUDA +================================ +``` + +--- + +## Known issues + +- **WSL2 IP changes on restart.** Always re-run the `netsh` port-proxy command after restarting WSL2 or Windows. +- **CUDA not visible in WSL2.** If `nvidia-smi` fails inside WSL2, update the Windows host GPU driver to ≥ 525.x. Installing CUDA inside WSL2 will not fix it. +- **Model download is slow.** HuggingFace downloads happen over HTTPS. Pre-fetch the model before a timed test (see Step 6). +- **Port 8001 already in use.** Change `--port` to another value and update the firewall/portproxy rules accordingly. +- **`bf16` not supported on older GPUs.** Use `--quantization int8` on Turing (RTX 20xx) cards or earlier if bfloat16 ops fail. diff --git a/docs/adr/0020-chat-streaming-live-progress-and-mixed-topology-routing.md b/docs/adr/0020-chat-streaming-live-progress-and-mixed-topology-routing.md new file mode 100644 index 0000000..62fe4da --- /dev/null +++ b/docs/adr/0020-chat-streaming-live-progress-and-mixed-topology-routing.md @@ -0,0 +1,127 @@ +# ADR-0020: Dashboard chat streaming, live request progress, and the mixed-topology routing flaw + +## Status: Accepted (chat/streaming/styles implemented); routing flaw documented, fix pending + +## Context + +Live alpha testing (2026-07-07) with `Qwen3.6-35B-A3B` split across two LAN nodes surfaced +three UX gaps and one routing correctness flaw: + +1. **No visibility while a request is processing.** The Call wall showed + "no in-flight requests" during a 52-second generation. Cause: the dashboard chat sent + `stream: false`, and the tracker only emits `proxy progress` console events (the Call + wall's live-status source, `_tracker_log_proxy_progress`, `server.py` ~2199) for + **streamed** requests. Non-streamed proxying produces only + `route selected → connected → complete`, and short requests complete inside the + dashboard's 4-second poll window. +2. **Chat did not stream.** The nodes support SSE token-by-token generation + (`generate_text_streaming`, hardened earlier for split shards), and the tracker proxy + passes `text/event-stream` through (`server.py` ~3256), but the chat panel blocked on + full JSON and showed nothing until completion. +3. **Chat panel styles drifted.** The "new chat layout" redesign left hardcoded one-off + colors (`#1f4788`, `#2563b8`, `#10151d`, `#1a1012`, `#5c2020`, `#ffb4b4`) mixed with + the CSS custom-property palette. + +## Decisions + +### 1. Chat streams by default (SSE) + +`dashboard.html` `sendChat()` now sends `stream: true` and consumes the SSE body with a +`ReadableStream` reader: + +- Assistant tokens render incrementally into the last bubble (direct DOM update, full + re-render only at boundaries), with a blinking `▍` cursor while streaming. +- Chat status shows live progress: `generating… N tokens · X tok/s`. +- The send button becomes a stop button (`■`) during generation, backed by an + `AbortController`; a stopped generation keeps the partial text. +- Non-SSE responses (JSON fallback, errors) are still handled; `data: {"error": ...}` + stream events surface as error bubbles. +- `streaming` flags are stripped when loading persisted sessions so an interrupted + generation never leaves a stuck cursor. + +### 2. Live in-flight visibility rides on streaming + +No tracker change was needed: because chat now streams, the tracker emits `proxy progress` +events (throttled to stdout, updated in place in the console ring via +`update_console_key`), and the existing Call wall state machine +(`buildCallWallStates`) renders processing rows with live tokens/TPS/queue. + +**Known limitation (accepted):** non-streamed API requests still show no progress between +`proxy connected` and `proxy complete` — there is nothing to report until the node +returns. Callers wanting live visibility should use `stream: true`. + +### 3. Chat style tokens + +All chat colors route through `:root` custom properties (`--hover-bg`, `--chat-user-bg` +`#1f6feb`, `--chat-user-border`, `--chat-error-bg/border/fg`). No hardcoded hex values +remain in chat rules, so future palette changes are single-line edits. + +## Documented flaw: mixed-topology routing (partial GPU head + full CPU node) + +### Observed (2026-07-07, tracker 192.168.0.179:8080) + +Two nodes registered for `qwen3.6-35b-a3b`: + +| node | hardware | shard | benchmark | +|---|---|---|---| +| `5gMLrmyB-ec3afe6f1a03` (192.168.0.20) | RTX 4060, CUDA | 0–21 (partial, fast) | 11,164 | +| `7j77FsPY-55249b0583e5` (192.168.0.179) | CPU | 0–39 (full, slow) | 425 | + +When the tracker selected the GPU node as head, it injected: + +``` +downstream=[{"endpoint": "http://192.168.0.179:7000", "start_layer": 0}] +``` + +`start_layer: 0` — not 22. The downstream full node re-ran **all 40 layers from layer 0 +on hidden states that had already passed through the head's layers 0–21**, producing +garbage logits. Evidence from the logs: + +- GPU-headed requests: `generation complete tokens=1` and billed `out=0`/`out=1`/`out=3` + — near-instant EOS from corrupt activations. +- The same prompt routed directly to the CPU full node: 209 tokens over 52 s (healthy). +- Observed TPS for GPU-headed requests was meaningless (2.5–19.0 "tok/s" on 0–3 token + outputs), and those samples now pollute the rolling per-`(node, model)` throughput + stats used for routing preference. +- Clients were **billed** for these broken 1-token responses. + +### Root cause + +The route planner treats the full-coverage node as a standalone complete route +(`route=7j77FsPY…[0-39]`) but still injects it as the head's downstream with the +downstream node's own `shard_start` (0) instead of `head.shard_end + 1` (22). A partial +head + full-model downstream is a topology the planner never had to handle before — +prior split tests used disjoint shards (0–11 + 12–23) where `shard_start` happened to +equal the correct continuation layer. + +### Required fix (not yet implemented) + +1. **Correct continuation layer:** when hop N ends at layer `e`, hop N+1 must execute + from `start_layer = e + 1` regardless of the downstream node's own `shard_start` + (the `X-Meshnet-Start-Layer` overlapping-shard mechanism from ADR-0012 exists for + exactly this; the planner must set it for full-model downstream nodes too). +2. **Route preference sanity:** with a healthy single-node full route available, prefer + it over a multi-hop route unless the pipeline is estimated faster; a fast head that + forces a slow full-model tail wins nothing (every token still crosses the CPU node). +3. **Stat hygiene:** exclude or flag throughput samples from responses with ≤ a few + output tokens, so broken routes don't skew routing preference. +4. **Billing guard (consider):** suspiciously short completions from multi-hop routes + during this window were billed; a minimum-viability check (or refund path) may be + warranted once audits land. + +### Verification for the fix + +Reproduce with a partial GPU head (0–21) + full CPU node (0–39): a chat request routed +through the GPU head must produce output equivalent to the direct CPU route, with +`downstream start_layer=22` visible in `proxy route selected`, and multi-token streamed +output on the Call wall. + +## Verification of this ADR's implemented changes + +- `pytest tests/test_dashboard.py` — 5 passed (stale "Chat / inference" panel assertion + updated to the tabbed layout). +- Embedded dashboard JS parses (`new Function(script)` under Node 22). +- Live check: open `/dashboard` → Chat, send a prompt to `qwen3.6-35b-a3b` — tokens + must appear incrementally with live tok/s in the status line, the Call wall must show + the request as `processing` with live TPS, and the send button must stop generation + mid-stream keeping partial text. diff --git a/docs/adr/0021-dynamic-statistical-routing.md b/docs/adr/0021-dynamic-statistical-routing.md new file mode 100644 index 0000000..7d94553 --- /dev/null +++ b/docs/adr/0021-dynamic-statistical-routing.md @@ -0,0 +1,119 @@ +# ADR-0021: Dynamic statistical routing (bandit-style route selection) + +## Status: Accepted, implemented + +## Context + +ADR-0020 documented the mixed-topology flaw: with a fast GPU node serving layers 0–21 and +a slow CPU node serving 0–39 of `Qwen3.6-35B-A3B`, the tracker picked the GPU node as +proxy head *independently* of route planning, injecting a downstream hop with the wrong +`start_layer` (0 instead of 22) and corrupting generation. + +Beyond the bug, the deeper issue is that the tracker **cannot know a priori** which route +is faster. Is one CPU node running all 40 layers faster than a GPU running 0–21 plus a +CPU hop for 22–39? Benchmarks don't answer that — network hops, MoE expert loading, and +queue dynamics only show up in real end-to-end requests. The router must *measure*. + +## Decision + +Route selection is a **multi-armed bandit** over enumerated candidate routes, implemented +in `packages/tracker/meshnet_tracker/routing_stats.py` and wired into the chat proxy in +`server.py`. + +### Arms: route signatures + +A route's identity is `model_key | node_id[shard] -> node_id[shard] -> …`. Node ids embed +wallet + shard, so a node re-registering with a different shard produces a new arm +automatically. The proxy target is **always the route's own head** (`route_nodes[0]`), +and each hop's `start_layer` is `previous_hop.shard_end + 1` — this fixes ADR-0020's flaw +structurally: head choice and route planning can no longer disagree. + +### Candidate enumeration (`_enumerate_routes`) + +One candidate per distinct head (a node whose `shard_start` equals the model's first +layer — it must tokenize/embed), greedily completed with longest-advancing hops. Each +candidate carries a `prior_tps`: its bottleneck hop's queue-adjusted effective throughput +× reputation. Capped at 8 candidates ranked by prior. + +### Statistics: decayed EWMA + topology epochs + +Per (model, signature), `RouteStatsStore` keeps an EWMA of observed end-to-end tokens/sec +with **time-decayed sample mass** (half-life default 600 s). Two staleness mechanisms +handle the morphing network: + +- **Continuous**: sample mass decays; a route unproven for a while (mass < 0.5) drops out + of the exploit pool and gets re-scouted. +- **Abrupt**: any node join/leave/shard-change bumps the model's *topology epoch*. Stats + from an older epoch keep their EWMA as a display prior but are demoted to the scout + pool ("stale") until re-measured under the new topology. + +Sample hygiene: completions below `min_sample_tokens` (default 8) are rejected — the +1-token garbage responses from the ADR-0020 bug would otherwise poison arms with +meaningless tps values. Routes with no samples for 24 h are pruned. + +### Selection policy (`choose_route`) + +1. **Scout** (probability `explore_share`, default 0.3): if any candidate is unproven / + stale / decayed, route the request there — least-measured first, tiebreak on prior. + These are the user's "discovery/scout routes". With *no* proven arms at all, selection + is deterministic best-prior (matches the old benchmark-based behavior, keeps cold + start sane and tests deterministic). +2. **Exploit** (otherwise): weighted random among proven arms with + `P(route) ∝ tps^alpha`, `alpha` default 1.0 — a 1.5×-faster route gets 1.5× the + traffic. `alpha` is a config knob: >1 shifts toward winner-takes-most as the network + matures, without redesign. (Proportional split is not throughput-optimal in queueing + terms, but it keeps every arm warm with fresh samples; tune alpha up when traffic + justifies it.) + +Pinned routes (`"route": [...]` in the request body) bypass the bandit but still record +samples. + +### Configuration + +| CLI flag | env var | default | +|---|---|---| +| `--route-explore-share` | `MESHNET_ROUTE_EXPLORE_SHARE` | 0.3 | +| `--route-weight-alpha` | `MESHNET_ROUTE_WEIGHT_ALPHA` | 1.0 | +| `--route-stats-half-life` | `MESHNET_ROUTE_STATS_HALF_LIFE` | 600 | +| — | `MESHNET_ROUTE_MIN_SAMPLE_TOKENS` | 8 | + +High explore share now (development, few requests); drop toward 0.05–0.1 once real +traffic provides passive coverage. + +### Visibility + +- **`GET /v1/routing`** (optionally `?model=`): per model — topology epoch and the full + candidate table: hops, learned tps, **coefficient** (tps ÷ best proven route's tps), + **expected traffic share**, sample count, decayed weight, status + (proven / unsampled / stale / decayed). +- **Dashboard → Overview → "Routing (learned)"**: renders that table live (4 s poll), + with the active config in the header line. +- **Console/`proxy route selected`** events now include the routing decision + (`{"mode": "scout"|"exploit"|"pinned"|"greedy-fallback", "signature": …}`), so the Call + wall history shows which arm served each request. + +## Storage considerations + +Stats are **in-memory per tracker** for alpha: they are cheap to relearn (a few requests +per route), and gossiping them would import ADR-0019's consistency questions for data +that is intentionally ephemeral. If multi-tracker route learning is needed later, ship +route samples over the existing stats gossip and merge EWMAs by decayed weight — the +store's (value, mass, timestamp) representation merges cleanly. + +## Consequences + +- The GPU(0–21)+CPU(0–39) topology now works: both routes get measured, the coefficient + is visible on the dashboard, and traffic shifts to whichever is actually faster. +- Routing is no longer deterministic once samples exist. Tests needing determinism seed + `server.route_rng` or rely on the cold-start deterministic path. +- The billing-relevant fix: heads are always part of the planned route, so per-hop + `start_layer` and work-unit spans are consistent. + +## Verification + +`tests/test_dynamic_routing.py` (11 tests): EWMA/decay/epoch semantics, near-empty sample +rejection, traffic split ≈ tps ratio at alpha=1 (0.6/0.4 over 4000 seeded draws), scout +rate ≈ explore share, mixed-topology enumeration (both routes, hybrid prior = bottleneck), +head-is-route-head regression with `start_layer=22` on the hybrid route, and `/v1/routing` +table shape. Live: start both nodes, run several chats, open the dashboard "Routing +(learned)" panel and watch coefficients converge. diff --git a/docs/issues/30-manual-route-and-hop-benchmark.md b/docs/issues/30-manual-route-and-hop-benchmark.md index 99209fe..c664fd9 100644 --- a/docs/issues/30-manual-route-and-hop-benchmark.md +++ b/docs/issues/30-manual-route-and-hop-benchmark.md @@ -1,48 +1,48 @@ -# US-020 — Manual route selection + hop-penalty benchmarking - -## Context - -The tracker auto-selects inference routes based on synthetic benchmark scores. To measure -the real cost of adding hops (latency per node boundary), we need: - -1. A way to pin a request to a specific route so we control the variable. -2. A benchmark endpoint that runs the same prompt through 1-node, 2-node, and 3-node - routes and records per-hop latency. - -Results are stored to disk. Routing algorithm is **not** changed in this story — this is -data collection only. The data will inform a future routing optimisation story. - -## Design decisions (grilled 2026-07-01) - -| Decision | Choice | -|---|---| -| Route spec | Optional `route` field in JSON request body (list of node IDs) | -| Trigger | Explicit only — `POST /v1/benchmark/hop-penalty` endpoint | -| Auth | Header-presence stub (`Authorization` must be non-empty); real auth in future story | -| Routing integration | Store data only; routing algorithm unchanged | -| Persistence | Append to `benchmark_results.json` in tracker working dir; in-memory queryable | - -## Acceptance criteria - -- `POST /v1/chat/completions` accepts optional `"route": ["", ...]` in the - request body. If present, the tracker uses those nodes in order instead of auto-selecting. - If absent, existing routing is unchanged (no breaking change for unaware clients). -- Missing or invalid node IDs in `route` return HTTP 400 with a descriptive error. -- `POST /v1/benchmark/hop-penalty` is auth-gated: requests without a non-empty - `Authorization` header return HTTP 401. Body: `{"model": "...", "prompt": "...", - "max_new_tokens": 64}`. -- Benchmark fans out to up to three routes: 1-node (single node covering all layers), - 2-node (two consecutive shard nodes), 3-node (three nodes) — using whatever is - currently registered. Routes with insufficient coverage are skipped, not errored. -- Response includes per-route breakdown: `total_ms`, `per_hop_ms: [...]`, - `tokens_generated`, `route: [node_id, ...]`. -- Results are appended to `/benchmark_results.json` (created if - absent) as a JSON array. Each entry includes timestamp, model, prompt hash, and the - per-route breakdown. -- `GET /v1/benchmark/results` returns the stored results array. Also auth-gated. -- Clients that never send `route` or call `/v1/benchmark/*` are completely unaffected. -- Integration test: send the same prompt via a pinned 1-node route and a pinned 2-node - route; assert 2-node result has 2 entries in `per_hop_ms`; assert both records appear - in `benchmark_results.json`. -- `python -m pytest` passes from repo root. -- Commit only this story's changes. +# US-020 — Manual route selection + hop-penalty benchmarking + +## Context + +The tracker auto-selects inference routes based on synthetic benchmark scores. To measure +the real cost of adding hops (latency per node boundary), we need: + +1. A way to pin a request to a specific route so we control the variable. +2. A benchmark endpoint that runs the same prompt through 1-node, 2-node, and 3-node + routes and records per-hop latency. + +Results are stored to disk. Routing algorithm is **not** changed in this story — this is +data collection only. The data will inform a future routing optimisation story. + +## Design decisions (grilled 2026-07-01) + +| Decision | Choice | +|---|---| +| Route spec | Optional `route` field in JSON request body (list of node IDs) | +| Trigger | Explicit only — `POST /v1/benchmark/hop-penalty` endpoint | +| Auth | Header-presence stub (`Authorization` must be non-empty); real auth in future story | +| Routing integration | Store data only; routing algorithm unchanged | +| Persistence | Append to `benchmark_results.json` in tracker working dir; in-memory queryable | + +## Acceptance criteria + +- `POST /v1/chat/completions` accepts optional `"route": ["", ...]` in the + request body. If present, the tracker uses those nodes in order instead of auto-selecting. + If absent, existing routing is unchanged (no breaking change for unaware clients). +- Missing or invalid node IDs in `route` return HTTP 400 with a descriptive error. +- `POST /v1/benchmark/hop-penalty` is auth-gated: requests without a non-empty + `Authorization` header return HTTP 401. Body: `{"model": "...", "prompt": "...", + "max_new_tokens": 64}`. +- Benchmark fans out to up to three routes: 1-node (single node covering all layers), + 2-node (two consecutive shard nodes), 3-node (three nodes) — using whatever is + currently registered. Routes with insufficient coverage are skipped, not errored. +- Response includes per-route breakdown: `total_ms`, `per_hop_ms: [...]`, + `tokens_generated`, `route: [node_id, ...]`. +- Results are appended to `/benchmark_results.json` (created if + absent) as a JSON array. Each entry includes timestamp, model, prompt hash, and the + per-route breakdown. +- `GET /v1/benchmark/results` returns the stored results array. Also auth-gated. +- Clients that never send `route` or call `/v1/benchmark/*` are completely unaffected. +- Integration test: send the same prompt via a pinned 1-node route and a pinned 2-node + route; assert 2-node result has 2 entries in `per_hop_ms`; assert both records appear + in `benchmark_results.json`. +- `python -m pytest` passes from repo root. +- Commit only this story's changes. diff --git a/packages/node/meshnet_node/model_backend.py b/packages/node/meshnet_node/model_backend.py index 0fdadba..49d18a2 100644 --- a/packages/node/meshnet_node/model_backend.py +++ b/packages/node/meshnet_node/model_backend.py @@ -1,827 +1,827 @@ -"""HuggingFace/PyTorch shard backend for real node inference.""" - -from __future__ import annotations - -import base64 -from dataclasses import dataclass -import json -from pathlib import Path -from typing import Any, Literal - -Quantization = Literal["auto", "bfloat16", "int8", "nf4"] - - -class ModelBackendError(RuntimeError): - """Base class for real model backend startup and execution failures.""" - - -class MissingModelDependencyError(ModelBackendError): - """Raised when optional model dependencies are not installed.""" - - -class InsufficientVRAMError(ModelBackendError): - """Raised when a requested shard cannot fit in available CUDA memory.""" - - -class PartialModelLoadUnsupported(ModelBackendError): - """Raised when a shard cannot be materialized from a local snapshot subset.""" - - -@dataclass(frozen=True) -class TensorPayload: - body: bytes - shape: list[int] - attention_mask_header: str | None - position_ids_header: str | None - - -def validate_quantization(value: str) -> Quantization: - if value not in {"auto", "bfloat16", "int8", "nf4"}: - raise ValueError("quantization must be one of: auto, bfloat16, int8, nf4") - return value # type: ignore[return-value] - - -def build_quantization_config(quantization: Quantization) -> Any | None: - """Return a transformers BitsAndBytesConfig for quantized weights.""" - if quantization in {"auto", "bfloat16"}: - return None - try: - import torch - from transformers import BitsAndBytesConfig - except ModuleNotFoundError as exc: - raise MissingModelDependencyError( - "transformers and torch are required for int8/nf4 quantization" - ) from exc - - if quantization == "int8": - return BitsAndBytesConfig(load_in_8bit=True) - return BitsAndBytesConfig( - load_in_4bit=True, - bnb_4bit_quant_type="nf4", - bnb_4bit_compute_dtype=torch.bfloat16, - ) - - -class TorchModelShard: - """Executable subset of a HuggingFace causal language model.""" - - def __init__( - self, - model_id: str, - shard_start: int, - shard_end: int, - quantization: Quantization = "auto", - cache_dir: Path | None = None, - ) -> None: - if shard_start < 0 or shard_end < 0 or shard_start > shard_end: - raise ValueError("shard_start must be <= shard_end and non-negative") - self.model_id = model_id - self.shard_start = shard_start - self.shard_end = shard_end - self.quantization = quantization - - try: - import torch - from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer - except ModuleNotFoundError as exc: - raise MissingModelDependencyError( - "real model backend requires torch, transformers, safetensors, accelerate, and bitsandbytes" - ) from exc - - self.torch = torch - self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - load_source = str(cache_dir) if cache_dir is not None and (cache_dir / "config.json").exists() else model_id - quant_config, dtype, uses_quantized_weights = _model_load_plan( - AutoConfig, - load_source, - quantization, - torch, - None if load_source != model_id else cache_dir, - ) - try: - total_layers_hint = _total_layers_for_local_snapshot(AutoConfig, load_source) - if _should_partial_materialize_shard( - load_source, - shard_start, - shard_end, - total_layers_hint=total_layers_hint, - uses_quantized_weights=uses_quantized_weights, - ): - self.model = _load_partial_model_from_snapshot( - AutoConfig, - AutoModelForCausalLM, - torch, - load_source, - shard_start, - shard_end, - dtype, - self.device, - ) - else: - load_kwargs = { - "device_map": "auto" if uses_quantized_weights else None, - "dtype": dtype, - "low_cpu_mem_usage": True, - "cache_dir": str(cache_dir) if cache_dir is not None and load_source == model_id else None, - } - if quant_config is not None: - load_kwargs["quantization_config"] = quant_config - self.model = AutoModelForCausalLM.from_pretrained( - load_source, - **load_kwargs, - ) - if not uses_quantized_weights: - self.model.to(self.device) - except Exception as exc: - if _looks_like_oom(exc): - memory_kind = "VRAM" if self.device.type == "cuda" else "RAM" - raise InsufficientVRAMError( - f"insufficient {memory_kind} to load {model_id} layers {shard_start}:{shard_end} " - f"with {quantization} quantization; choose a smaller shard or lower quantization" - ) from exc - raise - - self.model.eval() - self.tokenizer = AutoTokenizer.from_pretrained( - load_source, - cache_dir=str(cache_dir) if cache_dir is not None and load_source == model_id else None, - ) - self.layers = _model_layers(self.model) - self.total_layers = len(self.layers) - # shard_end is INCLUSIVE (last layer index, 0-based), matching the CLI convention. - if shard_end >= self.total_layers: - raise ValueError( - f"shard_end {shard_end} exceeds last layer index {self.total_layers - 1}" - ) - self.is_head = shard_start == 0 - self.is_tail = shard_end >= self.total_layers - 1 - self.hidden_size = int( - getattr(self.model.config, "hidden_size", 0) - or getattr(self.model.config, "n_embd", 0) - ) - self._embed_tokens = _embed_tokens(self.model) if self.is_head else None - self._position_embeddings = _position_embeddings(self.model) - self._norm = _final_norm(self.model) if self.is_tail else None - self._lm_head = getattr(self.model, "lm_head", None) if self.is_tail else None - - def encode_prompt(self, prompt: str) -> TensorPayload: - if not self.is_head or self._embed_tokens is None: - raise ModelBackendError("text prompts can only be accepted by the head shard") - encoded = self.tokenizer(prompt, return_tensors="pt") - input_ids = encoded["input_ids"].to(self.device) - attention_mask = encoded.get("attention_mask") - if attention_mask is None: - attention_mask = self.torch.ones_like(input_ids) - attention_mask = attention_mask.to(self.device) - position_ids = _position_ids(attention_mask, self.torch) - hidden_states = self._embed_tokens(input_ids) - if self._position_embeddings is not None: - hidden_states = hidden_states + self._position_embeddings(position_ids) - hidden_states = self._run_layers(hidden_states, attention_mask, position_ids) - return self._payload(hidden_states, attention_mask, position_ids) - - def forward_bytes( - self, - body: bytes, - shape: list[int], - attention_mask_header: str | None, - position_ids_header: str | None, - start_layer: int | None = None, - ) -> TensorPayload | str: - hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to( - self.device - ) - attention_mask = _tensor_from_int64_header( - attention_mask_header, self.torch, self.device - ) - position_ids = _tensor_from_int64_header( - position_ids_header, self.torch, self.device - ) - hidden_states = self._run_layers( - hidden_states, attention_mask, position_ids, start_layer=start_layer - ) - if self.is_tail: - return self.decode_tail(hidden_states) - return self._payload(hidden_states, attention_mask, position_ids) - - def decode_tail(self, hidden_states: Any) -> str: - if self._norm is not None: - hidden_states = self._norm(hidden_states) - if self._lm_head is None: - raise ModelBackendError("tail shard has no lm_head") - logits = self._lm_head(hidden_states) - token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item()) - return self.tokenizer.decode([token_id], skip_special_tokens=True) - - def generate_text( - self, - messages: list[dict], - max_new_tokens: int = 5120, - temperature: float = 1.0, - top_p: float = 1.0, - ) -> str: - """Autoregressive generation using HF generate() — single-node (head+tail) mode.""" - if not self.is_head or not self.is_tail: - raise ModelBackendError("local generation requires a full head+tail shard") - encoded = self._encode_messages(messages) - input_ids = encoded["input_ids"].to(self.device) - attention_mask = encoded.get("attention_mask") - if attention_mask is not None: - attention_mask = attention_mask.to(self.device) - pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None) - do_sample = temperature != 1.0 or top_p != 1.0 - with self.torch.inference_mode(): - generated = self.model.generate( - input_ids=input_ids, - attention_mask=attention_mask, - max_new_tokens=max(1, int(max_new_tokens)), - do_sample=do_sample, - temperature=temperature if do_sample else None, - top_p=top_p if do_sample else None, - pad_token_id=pad_token_id, - ) - new_tokens = generated[0, input_ids.shape[-1]:] - return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip() - - def generate_text_streaming( - self, - messages: list[dict], - max_new_tokens: int = 5000, - temperature: float = 1.0, - top_p: float = 1.0, - ): - """Yield decoded token strings one at a time using HF TextIteratorStreamer.""" - if not self.is_head or not self.is_tail: - raise ModelBackendError("streaming generation requires a full head+tail shard") - import threading - try: - from transformers import TextIteratorStreamer # type: ignore[import] - except ImportError: - yield self.generate_text(messages, max_new_tokens, temperature, top_p) - return - - encoded = self._encode_messages(messages) - input_ids = encoded["input_ids"].to(self.device) - attention_mask = encoded.get("attention_mask") - if attention_mask is not None: - attention_mask = attention_mask.to(self.device) - pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None) - do_sample = temperature != 1.0 or top_p != 1.0 - - streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True) - gen_kwargs = dict( - input_ids=input_ids, - attention_mask=attention_mask, - max_new_tokens=max(1, int(max_new_tokens)), - do_sample=do_sample, - temperature=temperature if do_sample else None, - top_p=top_p if do_sample else None, - pad_token_id=pad_token_id, - streamer=streamer, - ) - t = threading.Thread(target=self.model.generate, kwargs=gen_kwargs, daemon=True) - t.start() - for token_text in streamer: - yield token_text - t.join() - - def count_prompt_tokens(self, messages: list[dict]) -> int: - """Return tokenizer-backed prompt token count for OpenAI usage metadata.""" - encoded = self._encode_messages(messages) - input_ids = encoded["input_ids"] - return int(input_ids.shape[-1]) - - def count_text_tokens(self, text: str) -> int: - """Return tokenizer-backed completion token count for OpenAI usage metadata.""" - try: - encoded = self.tokenizer( - text, - return_tensors="pt", - add_special_tokens=False, - ) - except TypeError: - encoded = self.tokenizer(text, return_tensors="pt") - return int(encoded["input_ids"].shape[-1]) - - def _encode_messages(self, messages: list[dict]) -> dict: - """Format messages with chat template (if available) and tokenize.""" - if hasattr(self.tokenizer, "apply_chat_template"): - try: - prompt_str = self.tokenizer.apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=False, - ) - return dict(self.tokenizer(prompt_str, return_tensors="pt")) - except Exception: - pass - prompt = " ".join( - str(m.get("content", "")) - for m in messages - if isinstance(m, dict) and m.get("role") == "user" - ) - return dict(self.tokenizer(prompt, return_tensors="pt")) - - def _run_layers( - self, - hidden_states: Any, - attention_mask: Any, - position_ids: Any, - start_layer: int | None = None, - ) -> Any: - # start_layer overrides shard_start for overlapping-shard routing - # (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent - # indexing outside the loaded weights. - effective_start = ( - max(self.shard_start, start_layer) - if start_layer is not None - else self.shard_start - ) - position_embeddings = _rotary_position_embeddings( - self.model, - hidden_states, - position_ids, - ) - layer_attention_mask = _decoder_attention_mask( - attention_mask, - hidden_states, - self.torch, - ) - with self.torch.inference_mode(): - for layer in self.layers[effective_start:self.shard_end + 1]: - hidden_states = _call_layer( - layer, - hidden_states, - layer_attention_mask, - position_ids, - position_embeddings, - ) - return hidden_states.to(self.torch.bfloat16) - - def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload: - hidden_states = hidden_states.to(self.torch.bfloat16).contiguous() - return TensorPayload( - body=_tensor_to_bytes(hidden_states), - shape=list(hidden_states.shape), - attention_mask_header=_int_tensor_header(attention_mask) - if attention_mask is not None - else None, - position_ids_header=_int_tensor_header(position_ids) - if position_ids is not None - else None, - ) - - -def load_torch_shard( - model_id: str, - shard_start: int, - shard_end: int, - quantization: Quantization = "auto", - cache_dir: Path | None = None, -) -> TorchModelShard: - return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir) - - -def _total_layers_for_local_snapshot(auto_config: Any, load_source: str) -> int | None: - snapshot_dir = Path(load_source) - if not (snapshot_dir / "config.json").exists(): - return None - from .model_catalog import layers_from_config - - try: - cfg = auto_config.from_pretrained(str(snapshot_dir)) - except Exception: - return None - return layers_from_config(cfg) - - -def _should_partial_materialize_shard( - load_source: str, - shard_start: int, - shard_end: int, - *, - total_layers_hint: int | None, - uses_quantized_weights: bool, -) -> bool: - if uses_quantized_weights: - return False - snapshot_dir = Path(load_source) - if not snapshot_dir.exists() or not (snapshot_dir / "config.json").exists(): - return False - if not (snapshot_dir / "model.safetensors.index.json").exists(): - return False - if total_layers_hint is None: - return False - return True - - -def _load_partial_model_from_snapshot( - auto_config: Any, - auto_model_for_causal_lm: Any, - torch: Any, - load_source: str, - shard_start: int, - shard_end: int, - dtype: Any, - device: Any, - *, - init_empty_weights_fn: Any | None = None, - set_tensor_fn: Any | None = None, - safe_open_fn: Any | None = None, -) -> Any: - from .model_catalog import layers_from_config - from .safetensors_selection import ( - INDEX_FILENAME, - select_tensor_names_for_layers_from_index, - ) - - if init_empty_weights_fn is None: - from accelerate import init_empty_weights as init_empty_weights_fn - if set_tensor_fn is None: - from accelerate.utils import set_module_tensor_to_device as set_tensor_fn - if safe_open_fn is None: - from safetensors import safe_open as safe_open_fn - - snapshot_dir = Path(load_source) - cfg = auto_config.from_pretrained(str(snapshot_dir)) - total_layers = layers_from_config(cfg) - if total_layers is None: - raise PartialModelLoadUnsupported( - f"could not determine num_hidden_layers for local snapshot {snapshot_dir}" - ) - if shard_end >= total_layers: - raise ValueError( - f"shard_end {shard_end} exceeds last layer index {total_layers - 1}" - ) - - index_path = snapshot_dir / INDEX_FILENAME - try: - index = json.loads(index_path.read_text(encoding="utf-8")) - except FileNotFoundError as exc: - raise PartialModelLoadUnsupported( - f"missing SafeTensors index for partial load: {index_path}" - ) from exc - weight_map = index.get("weight_map") - if not isinstance(weight_map, dict): - raise PartialModelLoadUnsupported(f"{INDEX_FILENAME} must contain a weight_map object") - - tensor_names = select_tensor_names_for_layers_from_index( - weight_map, - shard_start, - shard_end, - total_layers=total_layers, - ) - if not tensor_names: - raise PartialModelLoadUnsupported( - f"no checkpoint tensors matched layers {shard_start}-{shard_end} in {snapshot_dir}" - ) - - with init_empty_weights_fn(): - model = auto_model_for_causal_lm.from_config(_causal_lm_config(cfg), torch_dtype=dtype) - tie_weights = getattr(model, "tie_weights", None) - if callable(tie_weights): - tie_weights() - - # Multimodal/MTP checkpoints (e.g. Qwen3.5/3.6-MoE) carry vision and - # multi-token-prediction tensors the text-only CausalLM never builds; - # transformers' from_pretrained drops them via _keys_to_ignore_on_load_unexpected, - # so the manual loader must skip them too. - expected_keys = _model_state_dict_keys(model) - tensors_by_file: dict[str, list[str]] = {} - skipped: list[str] = [] - for tensor_name in sorted(tensor_names): - rel_file = weight_map.get(tensor_name) - if not isinstance(rel_file, str): - continue - if ( - expected_keys is not None - and _checkpoint_tensor_name_for_model(model, tensor_name) not in expected_keys - ): - skipped.append(tensor_name) - continue - tensors_by_file.setdefault(rel_file, []).append(tensor_name) - if skipped: - preview = ", ".join(skipped[:3]) - print( - f" Skipping {len(skipped)} checkpoint tensors absent from the causal LM " - f"(e.g. {preview})", - flush=True, - ) - if not tensors_by_file: - raise PartialModelLoadUnsupported( - f"no checkpoint tensors for layers {shard_start}-{shard_end} match the " - f"causal LM built from {snapshot_dir}" - ) - - for rel_file, names in tensors_by_file.items(): - checkpoint_file = snapshot_dir / rel_file - if not checkpoint_file.exists(): - raise PartialModelLoadUnsupported( - f"checkpoint file advertised in {INDEX_FILENAME} is missing: {checkpoint_file}" - ) - with safe_open_fn(str(checkpoint_file), framework="pt", device="cpu") as handle: - for tensor_name in names: - set_tensor_fn( - model, - _checkpoint_tensor_name_for_model(model, tensor_name), - device, - value=handle.get_tensor(tensor_name), - dtype=dtype, - ) - - for module in _active_modules_for_shard(model, shard_start, shard_end): - if hasattr(module, "to"): - module.to(device) - return model - - -def _model_load_plan( - auto_config: Any, - model_id: str, - quantization: Quantization, - torch: Any, - cache_dir: Path | None = None, -) -> tuple[Any | None, Any, bool]: - """Return (explicit quant config, dtype, uses quantized weights).""" - if quantization != "auto": - quant_config = build_quantization_config(quantization) - return quant_config, torch.bfloat16, quant_config is not None - - cfg = auto_config.from_pretrained( - model_id, - cache_dir=str(cache_dir) if cache_dir is not None else None, - ) - if _native_quantization_config(cfg) is not None: - return None, _native_torch_dtype(cfg, torch), True - return None, _native_torch_dtype(cfg, torch), False - - -def _config_candidates(cfg: Any) -> list[Any]: - candidates = [cfg] - get_text_config = getattr(cfg, "get_text_config", None) - if callable(get_text_config): - try: - candidates.append(get_text_config()) - except Exception: - pass - text_config = getattr(cfg, "text_config", None) - if text_config is not None: - candidates.append(text_config) - return candidates - - -def _native_quantization_config(cfg: Any) -> Any | None: - for candidate in _config_candidates(cfg): - quant_config = getattr(candidate, "quantization_config", None) - if quant_config: - return quant_config - return None - - -def _native_torch_dtype(cfg: Any, torch: Any) -> Any: - for candidate in _config_candidates(cfg): - for attr in ("dtype", "torch_dtype"): - dtype = getattr(candidate, attr, None) - if dtype is None: - continue - if isinstance(dtype, str): - dtype_name = dtype.removeprefix("torch.") - dtype_value = getattr(torch, dtype_name, None) - if dtype_value is not None: - return dtype_value - else: - return dtype - return torch.bfloat16 - - -def _causal_lm_config(cfg: Any) -> Any: - """Use the text decoder config for composite VLM/MoE presets.""" - get_text_config = getattr(cfg, "get_text_config", None) - if callable(get_text_config): - try: - return get_text_config() - except Exception: - pass - text_config = getattr(cfg, "text_config", None) - if text_config is not None: - return text_config - return cfg - - -def _model_state_dict_keys(model: Any) -> set[str] | None: - """Expected parameter/buffer names, or None when the model can't report them.""" - state_dict = getattr(model, "state_dict", None) - if not callable(state_dict): - return None - try: - return set(state_dict().keys()) - except Exception: - return None - - -def _checkpoint_tensor_name_for_model(model: Any, tensor_name: str) -> str: - """Map multimodal checkpoint keys onto text-only CausalLM modules when needed.""" - inner = getattr(model, "model", None) - if inner is not None and hasattr(inner, "language_model"): - return tensor_name - if ".language_model." in tensor_name: - return tensor_name.replace(".language_model.", ".") - return tensor_name - - -def _transformer_backbone(model: Any) -> Any: - if hasattr(model, "model"): - inner = model.model - language_model = getattr(inner, "language_model", None) - if language_model is not None: - return language_model - return inner - if hasattr(model, "transformer"): - return model.transformer - raise ModelBackendError( - "unsupported HuggingFace model architecture: no transformer backbone found" - ) - - -def _model_layers(model: Any) -> Any: - backbone = _transformer_backbone(model) - for attr in ("layers", "h", "blocks"): - layers = getattr(backbone, attr, None) - if layers is not None: - return layers - raise ModelBackendError( - "unsupported HuggingFace model architecture: no transformer layers found" - ) - - -def _embed_tokens(model: Any) -> Any: - backbone = _transformer_backbone(model) - for attr in ("embed_tokens", "wte"): - embed = getattr(backbone, attr, None) - if embed is not None: - return embed - raise ModelBackendError( - "unsupported HuggingFace model architecture: no token embeddings found" - ) - - -def _position_embeddings(model: Any) -> Any | None: - backbone = _transformer_backbone(model) - return getattr(backbone, "wpe", None) - - -def _rotary_embedding_module(model: Any) -> Any | None: - backbone = _transformer_backbone(model) - return getattr(backbone, "rotary_emb", None) - - -def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> list[Any]: - active: list[Any] = [] - - def add(module: Any | None) -> None: - if module is None: - return - if any(existing is module for existing in active): - return - active.append(module) - - if shard_start == 0: - add(_embed_tokens(model)) - add(_position_embeddings(model)) - add(_rotary_embedding_module(model)) - for layer in _model_layers(model)[shard_start:shard_end + 1]: - add(layer) - total_layers = len(_model_layers(model)) - if shard_end >= total_layers - 1: - add(_final_norm(model)) - add(getattr(model, "lm_head", None)) - return active - - -def _final_norm(model: Any) -> Any | None: - backbone = _transformer_backbone(model) - for attr in ("norm", "ln_f", "final_layer_norm"): - norm = getattr(backbone, attr, None) - if norm is not None: - return norm - return None - - -def _position_ids(attention_mask: Any, torch: Any) -> Any: - position_ids = attention_mask.long().cumsum(-1) - 1 - return position_ids.masked_fill(attention_mask == 0, 0).to(torch.long) - - -def _decoder_attention_mask(attention_mask: Any, hidden_states: Any, torch: Any) -> Any: - """Build a causal additive mask for decoder layers called outside model.forward.""" - if attention_mask is None: - return None - if len(getattr(attention_mask, "shape", ())) != 2: - return attention_mask - batch_size, seq_len = attention_mask.shape - if seq_len <= 1: - return None if bool(attention_mask.all()) else attention_mask.to(hidden_states.dtype) - - min_value = torch.finfo(hidden_states.dtype).min - causal = torch.full( - (seq_len, seq_len), - min_value, - dtype=hidden_states.dtype, - device=hidden_states.device, - ) - causal = torch.triu(causal, diagonal=1) - causal = causal[None, None, :, :].expand(batch_size, 1, seq_len, seq_len).clone() - - padding = attention_mask.to(device=hidden_states.device) - if not bool(padding.all()): - causal = causal.masked_fill(padding[:, None, None, :] == 0, min_value) - return causal - - -def _rotary_position_embeddings(model: Any, hidden_states: Any, position_ids: Any) -> Any | None: - """Return model-level rotary embeddings required by newer HF decoder layers.""" - if position_ids is None: - return None - rotary = _rotary_embedding_module(model) - if rotary is None: - return None - return rotary(hidden_states, position_ids) - - -def _call_layer( - layer: Any, - hidden_states: Any, - attention_mask: Any, - position_ids: Any, - position_embeddings: Any | None = None, -) -> Any: - attempts = ( - { - "attention_mask": attention_mask, - "position_ids": position_ids, - "position_embeddings": position_embeddings, - "use_cache": False, - }, - { - "attention_mask": attention_mask, - "position_ids": position_ids, - "use_cache": False, - }, - {"attention_mask": attention_mask, "use_cache": False}, - {"use_cache": False}, - {}, - ) - last_exc: Exception | None = None - for kwargs in attempts: - filtered = {key: value for key, value in kwargs.items() if value is not None} - try: - output = layer(hidden_states, **filtered) - return output[0] if isinstance(output, tuple) else output - except TypeError as exc: - last_exc = exc - if last_exc is not None: - raise last_exc - return layer(hidden_states)[0] - - -def _tensor_to_bytes(tensor: Any) -> bytes: - import torch - - return tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() - - -def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> Any: - tensor = torch.frombuffer(bytearray(body), dtype=torch.bfloat16) - return tensor.reshape(shape) - - -def _int_tensor_header(tensor: Any) -> str: - data = tensor.detach().cpu().long().contiguous() - raw = data.numpy().tobytes() - shape = ",".join(str(dim) for dim in data.shape) - encoded = base64.b64encode(raw).decode("ascii") - return f"{shape}:{encoded}" - - -def _tensor_from_int64_header(value: str | None, torch: Any, device: Any) -> Any | None: - if not value: - return None - shape_text, encoded = value.split(":", 1) - shape = [int(part) for part in shape_text.split(",") if part] - raw = base64.b64decode(encoded.encode("ascii")) - return torch.frombuffer(bytearray(raw), dtype=torch.int64).reshape(shape).to(device) - - -def _looks_like_oom(exc: BaseException) -> bool: - current: BaseException | None = exc - while current is not None: - text = str(current).lower() - if ( - "out of memory" in text - or "cuda error: out of memory" in text - or "paging file is too small" in text - or "os error 1455" in text - ): - return True - current = current.__cause__ or current.__context__ - return False +"""HuggingFace/PyTorch shard backend for real node inference.""" + +from __future__ import annotations + +import base64 +from dataclasses import dataclass +import json +from pathlib import Path +from typing import Any, Literal + +Quantization = Literal["auto", "bfloat16", "int8", "nf4"] + + +class ModelBackendError(RuntimeError): + """Base class for real model backend startup and execution failures.""" + + +class MissingModelDependencyError(ModelBackendError): + """Raised when optional model dependencies are not installed.""" + + +class InsufficientVRAMError(ModelBackendError): + """Raised when a requested shard cannot fit in available CUDA memory.""" + + +class PartialModelLoadUnsupported(ModelBackendError): + """Raised when a shard cannot be materialized from a local snapshot subset.""" + + +@dataclass(frozen=True) +class TensorPayload: + body: bytes + shape: list[int] + attention_mask_header: str | None + position_ids_header: str | None + + +def validate_quantization(value: str) -> Quantization: + if value not in {"auto", "bfloat16", "int8", "nf4"}: + raise ValueError("quantization must be one of: auto, bfloat16, int8, nf4") + return value # type: ignore[return-value] + + +def build_quantization_config(quantization: Quantization) -> Any | None: + """Return a transformers BitsAndBytesConfig for quantized weights.""" + if quantization in {"auto", "bfloat16"}: + return None + try: + import torch + from transformers import BitsAndBytesConfig + except ModuleNotFoundError as exc: + raise MissingModelDependencyError( + "transformers and torch are required for int8/nf4 quantization" + ) from exc + + if quantization == "int8": + return BitsAndBytesConfig(load_in_8bit=True) + return BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_compute_dtype=torch.bfloat16, + ) + + +class TorchModelShard: + """Executable subset of a HuggingFace causal language model.""" + + def __init__( + self, + model_id: str, + shard_start: int, + shard_end: int, + quantization: Quantization = "auto", + cache_dir: Path | None = None, + ) -> None: + if shard_start < 0 or shard_end < 0 or shard_start > shard_end: + raise ValueError("shard_start must be <= shard_end and non-negative") + self.model_id = model_id + self.shard_start = shard_start + self.shard_end = shard_end + self.quantization = quantization + + try: + import torch + from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer + except ModuleNotFoundError as exc: + raise MissingModelDependencyError( + "real model backend requires torch, transformers, safetensors, accelerate, and bitsandbytes" + ) from exc + + self.torch = torch + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + load_source = str(cache_dir) if cache_dir is not None and (cache_dir / "config.json").exists() else model_id + quant_config, dtype, uses_quantized_weights = _model_load_plan( + AutoConfig, + load_source, + quantization, + torch, + None if load_source != model_id else cache_dir, + ) + try: + total_layers_hint = _total_layers_for_local_snapshot(AutoConfig, load_source) + if _should_partial_materialize_shard( + load_source, + shard_start, + shard_end, + total_layers_hint=total_layers_hint, + uses_quantized_weights=uses_quantized_weights, + ): + self.model = _load_partial_model_from_snapshot( + AutoConfig, + AutoModelForCausalLM, + torch, + load_source, + shard_start, + shard_end, + dtype, + self.device, + ) + else: + load_kwargs = { + "device_map": "auto" if uses_quantized_weights else None, + "dtype": dtype, + "low_cpu_mem_usage": True, + "cache_dir": str(cache_dir) if cache_dir is not None and load_source == model_id else None, + } + if quant_config is not None: + load_kwargs["quantization_config"] = quant_config + self.model = AutoModelForCausalLM.from_pretrained( + load_source, + **load_kwargs, + ) + if not uses_quantized_weights: + self.model.to(self.device) + except Exception as exc: + if _looks_like_oom(exc): + memory_kind = "VRAM" if self.device.type == "cuda" else "RAM" + raise InsufficientVRAMError( + f"insufficient {memory_kind} to load {model_id} layers {shard_start}:{shard_end} " + f"with {quantization} quantization; choose a smaller shard or lower quantization" + ) from exc + raise + + self.model.eval() + self.tokenizer = AutoTokenizer.from_pretrained( + load_source, + cache_dir=str(cache_dir) if cache_dir is not None and load_source == model_id else None, + ) + self.layers = _model_layers(self.model) + self.total_layers = len(self.layers) + # shard_end is INCLUSIVE (last layer index, 0-based), matching the CLI convention. + if shard_end >= self.total_layers: + raise ValueError( + f"shard_end {shard_end} exceeds last layer index {self.total_layers - 1}" + ) + self.is_head = shard_start == 0 + self.is_tail = shard_end >= self.total_layers - 1 + self.hidden_size = int( + getattr(self.model.config, "hidden_size", 0) + or getattr(self.model.config, "n_embd", 0) + ) + self._embed_tokens = _embed_tokens(self.model) if self.is_head else None + self._position_embeddings = _position_embeddings(self.model) + self._norm = _final_norm(self.model) if self.is_tail else None + self._lm_head = getattr(self.model, "lm_head", None) if self.is_tail else None + + def encode_prompt(self, prompt: str) -> TensorPayload: + if not self.is_head or self._embed_tokens is None: + raise ModelBackendError("text prompts can only be accepted by the head shard") + encoded = self.tokenizer(prompt, return_tensors="pt") + input_ids = encoded["input_ids"].to(self.device) + attention_mask = encoded.get("attention_mask") + if attention_mask is None: + attention_mask = self.torch.ones_like(input_ids) + attention_mask = attention_mask.to(self.device) + position_ids = _position_ids(attention_mask, self.torch) + hidden_states = self._embed_tokens(input_ids) + if self._position_embeddings is not None: + hidden_states = hidden_states + self._position_embeddings(position_ids) + hidden_states = self._run_layers(hidden_states, attention_mask, position_ids) + return self._payload(hidden_states, attention_mask, position_ids) + + def forward_bytes( + self, + body: bytes, + shape: list[int], + attention_mask_header: str | None, + position_ids_header: str | None, + start_layer: int | None = None, + ) -> TensorPayload | str: + hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to( + self.device + ) + attention_mask = _tensor_from_int64_header( + attention_mask_header, self.torch, self.device + ) + position_ids = _tensor_from_int64_header( + position_ids_header, self.torch, self.device + ) + hidden_states = self._run_layers( + hidden_states, attention_mask, position_ids, start_layer=start_layer + ) + if self.is_tail: + return self.decode_tail(hidden_states) + return self._payload(hidden_states, attention_mask, position_ids) + + def decode_tail(self, hidden_states: Any) -> str: + if self._norm is not None: + hidden_states = self._norm(hidden_states) + if self._lm_head is None: + raise ModelBackendError("tail shard has no lm_head") + logits = self._lm_head(hidden_states) + token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item()) + return self.tokenizer.decode([token_id], skip_special_tokens=True) + + def generate_text( + self, + messages: list[dict], + max_new_tokens: int = 5120, + temperature: float = 1.0, + top_p: float = 1.0, + ) -> str: + """Autoregressive generation using HF generate() — single-node (head+tail) mode.""" + if not self.is_head or not self.is_tail: + raise ModelBackendError("local generation requires a full head+tail shard") + encoded = self._encode_messages(messages) + input_ids = encoded["input_ids"].to(self.device) + attention_mask = encoded.get("attention_mask") + if attention_mask is not None: + attention_mask = attention_mask.to(self.device) + pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None) + do_sample = temperature != 1.0 or top_p != 1.0 + with self.torch.inference_mode(): + generated = self.model.generate( + input_ids=input_ids, + attention_mask=attention_mask, + max_new_tokens=max(1, int(max_new_tokens)), + do_sample=do_sample, + temperature=temperature if do_sample else None, + top_p=top_p if do_sample else None, + pad_token_id=pad_token_id, + ) + new_tokens = generated[0, input_ids.shape[-1]:] + return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip() + + def generate_text_streaming( + self, + messages: list[dict], + max_new_tokens: int = 5000, + temperature: float = 1.0, + top_p: float = 1.0, + ): + """Yield decoded token strings one at a time using HF TextIteratorStreamer.""" + if not self.is_head or not self.is_tail: + raise ModelBackendError("streaming generation requires a full head+tail shard") + import threading + try: + from transformers import TextIteratorStreamer # type: ignore[import] + except ImportError: + yield self.generate_text(messages, max_new_tokens, temperature, top_p) + return + + encoded = self._encode_messages(messages) + input_ids = encoded["input_ids"].to(self.device) + attention_mask = encoded.get("attention_mask") + if attention_mask is not None: + attention_mask = attention_mask.to(self.device) + pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None) + do_sample = temperature != 1.0 or top_p != 1.0 + + streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True) + gen_kwargs = dict( + input_ids=input_ids, + attention_mask=attention_mask, + max_new_tokens=max(1, int(max_new_tokens)), + do_sample=do_sample, + temperature=temperature if do_sample else None, + top_p=top_p if do_sample else None, + pad_token_id=pad_token_id, + streamer=streamer, + ) + t = threading.Thread(target=self.model.generate, kwargs=gen_kwargs, daemon=True) + t.start() + for token_text in streamer: + yield token_text + t.join() + + def count_prompt_tokens(self, messages: list[dict]) -> int: + """Return tokenizer-backed prompt token count for OpenAI usage metadata.""" + encoded = self._encode_messages(messages) + input_ids = encoded["input_ids"] + return int(input_ids.shape[-1]) + + def count_text_tokens(self, text: str) -> int: + """Return tokenizer-backed completion token count for OpenAI usage metadata.""" + try: + encoded = self.tokenizer( + text, + return_tensors="pt", + add_special_tokens=False, + ) + except TypeError: + encoded = self.tokenizer(text, return_tensors="pt") + return int(encoded["input_ids"].shape[-1]) + + def _encode_messages(self, messages: list[dict]) -> dict: + """Format messages with chat template (if available) and tokenize.""" + if hasattr(self.tokenizer, "apply_chat_template"): + try: + prompt_str = self.tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=False, + ) + return dict(self.tokenizer(prompt_str, return_tensors="pt")) + except Exception: + pass + prompt = " ".join( + str(m.get("content", "")) + for m in messages + if isinstance(m, dict) and m.get("role") == "user" + ) + return dict(self.tokenizer(prompt, return_tensors="pt")) + + def _run_layers( + self, + hidden_states: Any, + attention_mask: Any, + position_ids: Any, + start_layer: int | None = None, + ) -> Any: + # start_layer overrides shard_start for overlapping-shard routing + # (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent + # indexing outside the loaded weights. + effective_start = ( + max(self.shard_start, start_layer) + if start_layer is not None + else self.shard_start + ) + position_embeddings = _rotary_position_embeddings( + self.model, + hidden_states, + position_ids, + ) + layer_attention_mask = _decoder_attention_mask( + attention_mask, + hidden_states, + self.torch, + ) + with self.torch.inference_mode(): + for layer in self.layers[effective_start:self.shard_end + 1]: + hidden_states = _call_layer( + layer, + hidden_states, + layer_attention_mask, + position_ids, + position_embeddings, + ) + return hidden_states.to(self.torch.bfloat16) + + def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload: + hidden_states = hidden_states.to(self.torch.bfloat16).contiguous() + return TensorPayload( + body=_tensor_to_bytes(hidden_states), + shape=list(hidden_states.shape), + attention_mask_header=_int_tensor_header(attention_mask) + if attention_mask is not None + else None, + position_ids_header=_int_tensor_header(position_ids) + if position_ids is not None + else None, + ) + + +def load_torch_shard( + model_id: str, + shard_start: int, + shard_end: int, + quantization: Quantization = "auto", + cache_dir: Path | None = None, +) -> TorchModelShard: + return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir) + + +def _total_layers_for_local_snapshot(auto_config: Any, load_source: str) -> int | None: + snapshot_dir = Path(load_source) + if not (snapshot_dir / "config.json").exists(): + return None + from .model_catalog import layers_from_config + + try: + cfg = auto_config.from_pretrained(str(snapshot_dir)) + except Exception: + return None + return layers_from_config(cfg) + + +def _should_partial_materialize_shard( + load_source: str, + shard_start: int, + shard_end: int, + *, + total_layers_hint: int | None, + uses_quantized_weights: bool, +) -> bool: + if uses_quantized_weights: + return False + snapshot_dir = Path(load_source) + if not snapshot_dir.exists() or not (snapshot_dir / "config.json").exists(): + return False + if not (snapshot_dir / "model.safetensors.index.json").exists(): + return False + if total_layers_hint is None: + return False + return True + + +def _load_partial_model_from_snapshot( + auto_config: Any, + auto_model_for_causal_lm: Any, + torch: Any, + load_source: str, + shard_start: int, + shard_end: int, + dtype: Any, + device: Any, + *, + init_empty_weights_fn: Any | None = None, + set_tensor_fn: Any | None = None, + safe_open_fn: Any | None = None, +) -> Any: + from .model_catalog import layers_from_config + from .safetensors_selection import ( + INDEX_FILENAME, + select_tensor_names_for_layers_from_index, + ) + + if init_empty_weights_fn is None: + from accelerate import init_empty_weights as init_empty_weights_fn + if set_tensor_fn is None: + from accelerate.utils import set_module_tensor_to_device as set_tensor_fn + if safe_open_fn is None: + from safetensors import safe_open as safe_open_fn + + snapshot_dir = Path(load_source) + cfg = auto_config.from_pretrained(str(snapshot_dir)) + total_layers = layers_from_config(cfg) + if total_layers is None: + raise PartialModelLoadUnsupported( + f"could not determine num_hidden_layers for local snapshot {snapshot_dir}" + ) + if shard_end >= total_layers: + raise ValueError( + f"shard_end {shard_end} exceeds last layer index {total_layers - 1}" + ) + + index_path = snapshot_dir / INDEX_FILENAME + try: + index = json.loads(index_path.read_text(encoding="utf-8")) + except FileNotFoundError as exc: + raise PartialModelLoadUnsupported( + f"missing SafeTensors index for partial load: {index_path}" + ) from exc + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + raise PartialModelLoadUnsupported(f"{INDEX_FILENAME} must contain a weight_map object") + + tensor_names = select_tensor_names_for_layers_from_index( + weight_map, + shard_start, + shard_end, + total_layers=total_layers, + ) + if not tensor_names: + raise PartialModelLoadUnsupported( + f"no checkpoint tensors matched layers {shard_start}-{shard_end} in {snapshot_dir}" + ) + + with init_empty_weights_fn(): + model = auto_model_for_causal_lm.from_config(_causal_lm_config(cfg), torch_dtype=dtype) + tie_weights = getattr(model, "tie_weights", None) + if callable(tie_weights): + tie_weights() + + # Multimodal/MTP checkpoints (e.g. Qwen3.5/3.6-MoE) carry vision and + # multi-token-prediction tensors the text-only CausalLM never builds; + # transformers' from_pretrained drops them via _keys_to_ignore_on_load_unexpected, + # so the manual loader must skip them too. + expected_keys = _model_state_dict_keys(model) + tensors_by_file: dict[str, list[str]] = {} + skipped: list[str] = [] + for tensor_name in sorted(tensor_names): + rel_file = weight_map.get(tensor_name) + if not isinstance(rel_file, str): + continue + if ( + expected_keys is not None + and _checkpoint_tensor_name_for_model(model, tensor_name) not in expected_keys + ): + skipped.append(tensor_name) + continue + tensors_by_file.setdefault(rel_file, []).append(tensor_name) + if skipped: + preview = ", ".join(skipped[:3]) + print( + f" Skipping {len(skipped)} checkpoint tensors absent from the causal LM " + f"(e.g. {preview})", + flush=True, + ) + if not tensors_by_file: + raise PartialModelLoadUnsupported( + f"no checkpoint tensors for layers {shard_start}-{shard_end} match the " + f"causal LM built from {snapshot_dir}" + ) + + for rel_file, names in tensors_by_file.items(): + checkpoint_file = snapshot_dir / rel_file + if not checkpoint_file.exists(): + raise PartialModelLoadUnsupported( + f"checkpoint file advertised in {INDEX_FILENAME} is missing: {checkpoint_file}" + ) + with safe_open_fn(str(checkpoint_file), framework="pt", device="cpu") as handle: + for tensor_name in names: + set_tensor_fn( + model, + _checkpoint_tensor_name_for_model(model, tensor_name), + device, + value=handle.get_tensor(tensor_name), + dtype=dtype, + ) + + for module in _active_modules_for_shard(model, shard_start, shard_end): + if hasattr(module, "to"): + module.to(device) + return model + + +def _model_load_plan( + auto_config: Any, + model_id: str, + quantization: Quantization, + torch: Any, + cache_dir: Path | None = None, +) -> tuple[Any | None, Any, bool]: + """Return (explicit quant config, dtype, uses quantized weights).""" + if quantization != "auto": + quant_config = build_quantization_config(quantization) + return quant_config, torch.bfloat16, quant_config is not None + + cfg = auto_config.from_pretrained( + model_id, + cache_dir=str(cache_dir) if cache_dir is not None else None, + ) + if _native_quantization_config(cfg) is not None: + return None, _native_torch_dtype(cfg, torch), True + return None, _native_torch_dtype(cfg, torch), False + + +def _config_candidates(cfg: Any) -> list[Any]: + candidates = [cfg] + get_text_config = getattr(cfg, "get_text_config", None) + if callable(get_text_config): + try: + candidates.append(get_text_config()) + except Exception: + pass + text_config = getattr(cfg, "text_config", None) + if text_config is not None: + candidates.append(text_config) + return candidates + + +def _native_quantization_config(cfg: Any) -> Any | None: + for candidate in _config_candidates(cfg): + quant_config = getattr(candidate, "quantization_config", None) + if quant_config: + return quant_config + return None + + +def _native_torch_dtype(cfg: Any, torch: Any) -> Any: + for candidate in _config_candidates(cfg): + for attr in ("dtype", "torch_dtype"): + dtype = getattr(candidate, attr, None) + if dtype is None: + continue + if isinstance(dtype, str): + dtype_name = dtype.removeprefix("torch.") + dtype_value = getattr(torch, dtype_name, None) + if dtype_value is not None: + return dtype_value + else: + return dtype + return torch.bfloat16 + + +def _causal_lm_config(cfg: Any) -> Any: + """Use the text decoder config for composite VLM/MoE presets.""" + get_text_config = getattr(cfg, "get_text_config", None) + if callable(get_text_config): + try: + return get_text_config() + except Exception: + pass + text_config = getattr(cfg, "text_config", None) + if text_config is not None: + return text_config + return cfg + + +def _model_state_dict_keys(model: Any) -> set[str] | None: + """Expected parameter/buffer names, or None when the model can't report them.""" + state_dict = getattr(model, "state_dict", None) + if not callable(state_dict): + return None + try: + return set(state_dict().keys()) + except Exception: + return None + + +def _checkpoint_tensor_name_for_model(model: Any, tensor_name: str) -> str: + """Map multimodal checkpoint keys onto text-only CausalLM modules when needed.""" + inner = getattr(model, "model", None) + if inner is not None and hasattr(inner, "language_model"): + return tensor_name + if ".language_model." in tensor_name: + return tensor_name.replace(".language_model.", ".") + return tensor_name + + +def _transformer_backbone(model: Any) -> Any: + if hasattr(model, "model"): + inner = model.model + language_model = getattr(inner, "language_model", None) + if language_model is not None: + return language_model + return inner + if hasattr(model, "transformer"): + return model.transformer + raise ModelBackendError( + "unsupported HuggingFace model architecture: no transformer backbone found" + ) + + +def _model_layers(model: Any) -> Any: + backbone = _transformer_backbone(model) + for attr in ("layers", "h", "blocks"): + layers = getattr(backbone, attr, None) + if layers is not None: + return layers + raise ModelBackendError( + "unsupported HuggingFace model architecture: no transformer layers found" + ) + + +def _embed_tokens(model: Any) -> Any: + backbone = _transformer_backbone(model) + for attr in ("embed_tokens", "wte"): + embed = getattr(backbone, attr, None) + if embed is not None: + return embed + raise ModelBackendError( + "unsupported HuggingFace model architecture: no token embeddings found" + ) + + +def _position_embeddings(model: Any) -> Any | None: + backbone = _transformer_backbone(model) + return getattr(backbone, "wpe", None) + + +def _rotary_embedding_module(model: Any) -> Any | None: + backbone = _transformer_backbone(model) + return getattr(backbone, "rotary_emb", None) + + +def _active_modules_for_shard(model: Any, shard_start: int, shard_end: int) -> list[Any]: + active: list[Any] = [] + + def add(module: Any | None) -> None: + if module is None: + return + if any(existing is module for existing in active): + return + active.append(module) + + if shard_start == 0: + add(_embed_tokens(model)) + add(_position_embeddings(model)) + add(_rotary_embedding_module(model)) + for layer in _model_layers(model)[shard_start:shard_end + 1]: + add(layer) + total_layers = len(_model_layers(model)) + if shard_end >= total_layers - 1: + add(_final_norm(model)) + add(getattr(model, "lm_head", None)) + return active + + +def _final_norm(model: Any) -> Any | None: + backbone = _transformer_backbone(model) + for attr in ("norm", "ln_f", "final_layer_norm"): + norm = getattr(backbone, attr, None) + if norm is not None: + return norm + return None + + +def _position_ids(attention_mask: Any, torch: Any) -> Any: + position_ids = attention_mask.long().cumsum(-1) - 1 + return position_ids.masked_fill(attention_mask == 0, 0).to(torch.long) + + +def _decoder_attention_mask(attention_mask: Any, hidden_states: Any, torch: Any) -> Any: + """Build a causal additive mask for decoder layers called outside model.forward.""" + if attention_mask is None: + return None + if len(getattr(attention_mask, "shape", ())) != 2: + return attention_mask + batch_size, seq_len = attention_mask.shape + if seq_len <= 1: + return None if bool(attention_mask.all()) else attention_mask.to(hidden_states.dtype) + + min_value = torch.finfo(hidden_states.dtype).min + causal = torch.full( + (seq_len, seq_len), + min_value, + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + causal = torch.triu(causal, diagonal=1) + causal = causal[None, None, :, :].expand(batch_size, 1, seq_len, seq_len).clone() + + padding = attention_mask.to(device=hidden_states.device) + if not bool(padding.all()): + causal = causal.masked_fill(padding[:, None, None, :] == 0, min_value) + return causal + + +def _rotary_position_embeddings(model: Any, hidden_states: Any, position_ids: Any) -> Any | None: + """Return model-level rotary embeddings required by newer HF decoder layers.""" + if position_ids is None: + return None + rotary = _rotary_embedding_module(model) + if rotary is None: + return None + return rotary(hidden_states, position_ids) + + +def _call_layer( + layer: Any, + hidden_states: Any, + attention_mask: Any, + position_ids: Any, + position_embeddings: Any | None = None, +) -> Any: + attempts = ( + { + "attention_mask": attention_mask, + "position_ids": position_ids, + "position_embeddings": position_embeddings, + "use_cache": False, + }, + { + "attention_mask": attention_mask, + "position_ids": position_ids, + "use_cache": False, + }, + {"attention_mask": attention_mask, "use_cache": False}, + {"use_cache": False}, + {}, + ) + last_exc: Exception | None = None + for kwargs in attempts: + filtered = {key: value for key, value in kwargs.items() if value is not None} + try: + output = layer(hidden_states, **filtered) + return output[0] if isinstance(output, tuple) else output + except TypeError as exc: + last_exc = exc + if last_exc is not None: + raise last_exc + return layer(hidden_states)[0] + + +def _tensor_to_bytes(tensor: Any) -> bytes: + import torch + + return tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + + +def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> Any: + tensor = torch.frombuffer(bytearray(body), dtype=torch.bfloat16) + return tensor.reshape(shape) + + +def _int_tensor_header(tensor: Any) -> str: + data = tensor.detach().cpu().long().contiguous() + raw = data.numpy().tobytes() + shape = ",".join(str(dim) for dim in data.shape) + encoded = base64.b64encode(raw).decode("ascii") + return f"{shape}:{encoded}" + + +def _tensor_from_int64_header(value: str | None, torch: Any, device: Any) -> Any | None: + if not value: + return None + shape_text, encoded = value.split(":", 1) + shape = [int(part) for part in shape_text.split(",") if part] + raw = base64.b64decode(encoded.encode("ascii")) + return torch.frombuffer(bytearray(raw), dtype=torch.int64).reshape(shape).to(device) + + +def _looks_like_oom(exc: BaseException) -> bool: + current: BaseException | None = exc + while current is not None: + text = str(current).lower() + if ( + "out of memory" in text + or "cuda error: out of memory" in text + or "paging file is too small" in text + or "os error 1455" in text + ): + return True + current = current.__cause__ or current.__context__ + return False diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index 54fa139..c119474 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -197,12 +197,43 @@ def _max_assignable_layers( memory_mb: int, total_layers: int | None, bytes_per_layer: int | None = None, + *, + safety_fraction: float = 0.8, ) -> int: if total_layers is None or total_layers <= 0 or memory_mb <= 0: return 0 budget_bytes = memory_mb * 1024 * 1024 layer_bytes = bytes_per_layer or _DEFAULT_BYTES_PER_LAYER - return min(total_layers, int((budget_bytes * 0.8) // layer_bytes)) + return min(total_layers, int((budget_bytes * safety_fraction) // layer_bytes)) + + +def _runtime_shard_safety_fraction(device: str) -> float: + """CPU partial loads need room for model skeletons, tokenizer, and allocator peaks.""" + return 0.55 if device != "cuda" else 0.8 + + +def _cap_auto_assigned_shard( + shard_start: int, + shard_end: int, + total_layers: int | None, + memory_mb: int, + bytes_per_layer: int | None, + device: str, +) -> tuple[int, bool, int]: + if bytes_per_layer is None or total_layers is None: + return shard_end, False, 0 + max_layers = _max_assignable_layers( + memory_mb, + total_layers, + bytes_per_layer=bytes_per_layer, + safety_fraction=_runtime_shard_safety_fraction(device), + ) + if max_layers <= 0: + return shard_end, False, max_layers + assigned_layers = shard_end - shard_start + 1 + if assigned_layers <= max_layers: + return shard_end, False, max_layers + return min(total_layers - 1, shard_start + max_layers - 1), True, max_layers def _format_shard_label( @@ -226,13 +257,19 @@ def _shard_budget_line( total_layers: int | None, quantization: str, bytes_per_layer: int | None = None, + safety_fraction: float = 0.8, ) -> str: memory_gb = memory_mb / 1024 gb_str = f"{memory_gb:.1f} GB" budget_quantization = "bfloat16" if quantization == "auto" else quantization if total_layers is None or total_layers <= 0: return f"Memory budget: {gb_str} {memory_source}; shard budget: unknown model layer count" - max_layers = _max_assignable_layers(memory_mb, total_layers, bytes_per_layer=bytes_per_layer) + max_layers = _max_assignable_layers( + memory_mb, + total_layers, + bytes_per_layer=bytes_per_layer, + safety_fraction=safety_fraction, + ) # Remaining capacity after one full model load (rough estimate) shard_bytes = max_layers * (bytes_per_layer or _DEFAULT_BYTES_PER_LAYER) remaining_gb = (memory_mb * 1024 * 1024 - shard_bytes) / (1024 ** 3) @@ -349,25 +386,38 @@ def _attach_relay_bridge(node: StubNodeServer | TorchNodeServer, bridge: RelayHt _PENDING_NODE_ID = "pending" +_HEARTBEAT_INTERVAL_IDLE = 20.0 +_HEARTBEAT_INTERVAL_BUSY = 3.0 + + def _start_heartbeat( tracker_url: str, node_id: str, register_payload: dict, - interval: float = 20.0, + interval: float = _HEARTBEAT_INTERVAL_IDLE, node_ref: Any | None = None, start_time: float | None = None, ) -> threading.Thread: """Daemon thread: sends heartbeats and re-registers automatically after tracker restarts. Heartbeat body carries cumulative stats (total_requests, failed_requests, - queue_depth, uptime_seconds, status). Stats are buffered locally during - outage and flushed on next successful heartbeat. + queue_depth, current_requests, uptime_seconds, status). Stats are buffered + locally during outage and flushed on next successful heartbeat. Heartbeat response may include new_assignment: {model, shard_start, shard_end} which is logged for now (hot-reload implemented in US-026). """ _start_time = start_time or time.monotonic() + def _current_requests_snapshot() -> list[dict]: + if node_ref is None: + return [] + getter = getattr(node_ref, "current_requests", None) + if getter is None: + return [] + current = getter() if callable(getter) else getter + return list(current) if isinstance(current, list) else [] + def _get_stats() -> dict: uptime = time.monotonic() - _start_time stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"} @@ -379,8 +429,16 @@ def _start_heartbeat( ) stats["failed_requests"] = getattr(node_ref, "failed_requests", 0) stats["queue_depth"] = getattr(node_ref, "queue_depth", 0) + current_requests = _current_requests_snapshot() + if current_requests: + stats["current_requests"] = current_requests return stats + def _sleep_interval() -> float: + if _current_requests_snapshot() or (node_ref is not None and getattr(node_ref, "queue_depth", 0) > 0): + return _HEARTBEAT_INTERVAL_BUSY + return interval + def _reregister() -> bool: nonlocal node_id try: @@ -442,7 +500,7 @@ def _start_heartbeat( outage_streak = 1 if node_id == _PENDING_NODE_ID else 0 while True: - time.sleep(interval) + time.sleep(_sleep_interval()) if outage_streak > 0: # Tracker was down — attempt re-registration first (it may have restarted @@ -708,6 +766,25 @@ def run_startup( if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"): shard_start = net_asgn["shard_start"] shard_end = net_asgn["shard_end"] + asgn_total_layers = int(net_asgn.get("num_layers") or detected) + asgn_bytes_per_layer = _assignment_bytes_per_layer(net_asgn, quantization) + capped_shard_end, was_capped, max_runtime_layers = _cap_auto_assigned_shard( + shard_start, + shard_end, + asgn_total_layers, + memory_budget_mb, + asgn_bytes_per_layer, + device, + ) + if was_capped: + original_end = shard_end + shard_end = capped_shard_end + print( + f" WARNING: tracker assigned layers {shard_start}-{original_end}, " + f"but CPU-safe runtime budget fits {max_runtime_layers}/{asgn_total_layers} layers; " + f"loading layers {shard_start}-{shard_end} instead.", + flush=True, + ) full_sources = ( [] if tracker_source_disabled else _full_model_sources(net_asgn.get("model_sources", [])) @@ -723,7 +800,7 @@ def run_startup( ) print( f" Tracker found uncovered shard: " - f"layers {shard_start}–{shard_end} (of {detected})", + f"layers {shard_start}-{shard_end} (of {detected})", flush=True, ) except Exception: @@ -752,6 +829,7 @@ def run_startup( shard_label = _format_shard_label(shard_start, shard_end, total_layers) public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) endpoint = f"http://{public_host}:{actual_port}" + node.set_advertised_endpoint(endpoint) local_base_url = f"http://127.0.0.1:{actual_port}" relay_bridge, relay_fields = _start_relay_bridge_if_available( tracker_url, @@ -838,18 +916,36 @@ def run_startup( assigned_shard_start: int = net_assignment["shard_start"] assigned_shard_end: int = net_assignment["shard_end"] assigned_num_layers: int = net_assignment["num_layers"] + assigned_bytes_per_layer = _assignment_bytes_per_layer(net_assignment, quantization) + capped_shard_end, was_capped, max_runtime_layers = _cap_auto_assigned_shard( + assigned_shard_start, + assigned_shard_end, + assigned_num_layers, + memory_budget_mb, + assigned_bytes_per_layer, + device, + ) + if was_capped: + original_end = assigned_shard_end + assigned_shard_end = capped_shard_end + print( + f" WARNING: tracker assigned layers {assigned_shard_start}-{original_end}, " + f"but CPU-safe runtime budget fits {max_runtime_layers}/{assigned_num_layers} layers; " + f"loading layers {assigned_shard_start}-{assigned_shard_end} instead.", + flush=True, + ) assigned_model_sources: list[dict] = net_assignment.get("model_sources", []) if _gap_found: print( f" Assigned gap: {assigned_hf_repo} " - f"layers {assigned_shard_start}–{assigned_shard_end} " + f"layers {assigned_shard_start}-{assigned_shard_end} " f"(of {assigned_num_layers})", flush=True, ) else: print( f" Assigned redundant copy: {assigned_hf_repo} " - f"layers {assigned_shard_start}–{assigned_shard_end} " + f"layers {assigned_shard_start}-{assigned_shard_end} " f"(of {assigned_num_layers})", flush=True, ) @@ -882,6 +978,7 @@ def run_startup( actual_port = node.start() public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) endpoint = f"http://{public_host}:{actual_port}" + node.set_advertised_endpoint(endpoint) local_base_url = f"http://127.0.0.1:{actual_port}" relay_bridge, relay_fields = _start_relay_bridge_if_available( tracker_url, @@ -935,7 +1032,7 @@ def run_startup( f" Wallet: {address}\n" f" Model ID: {assigned_hf_repo}\n" f" Shard: {shard_label}\n" - f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization)}\n" + f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_num_layers, quantization, bytes_per_layer=assigned_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n" f" Quantization: {quantization}\n" f" Endpoint: {endpoint}\n" f" Node ID: {tracker_node_id or 'unregistered'}\n" @@ -1001,6 +1098,7 @@ def run_startup( memory_budget_mb, assigned_total_layers, assignment_bytes_per_layer, + safety_fraction=_runtime_shard_safety_fraction(device), ) if pinned_layers > max_layers: raise ValueError( @@ -1058,6 +1156,7 @@ def run_startup( shard_label = f"{shard_label} (pinned)" public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) endpoint = f"http://{public_host}:{actual_port}" + node.set_advertised_endpoint(endpoint) local_base_url = f"http://127.0.0.1:{actual_port}" relay_bridge, relay_fields = _start_relay_bridge_if_available( tracker_url, @@ -1094,7 +1193,7 @@ def run_startup( f" Wallet: {address}\n" f" Model ID: {hf_repo}\n" f" Shard: {shard_label}\n" - f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n" + f" {_shard_budget_line(memory_budget_mb, memory_budget_source, total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n" f" Quantization: {quantization}\n" f" Endpoint: {endpoint}\n" f" Node ID: {tracker_node_id or 'unregistered'}\n" @@ -1171,7 +1270,7 @@ def run_startup( f"meshnet-node ready\n" f" Wallet: {address}\n" f" Shard: {shard_label}\n" - f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer)}\n" + f" {_shard_budget_line(memory_budget_mb, memory_budget_source, assigned_total_layers, quantization, bytes_per_layer=assignment_bytes_per_layer, safety_fraction=_runtime_shard_safety_fraction(device))}\n" f" Endpoint: {endpoint}\n" f" Node ID: {node_id}\n" f" Hardware: {hw_str}\n" diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 3e9f5ba..3bc8605 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -31,6 +31,69 @@ from .server import ( ) +def _endpoint_key(url: str) -> str: + """Normalize http(s) endpoints for host:port comparison.""" + parsed = urllib.parse.urlparse(url.rstrip("/")) + host = (parsed.hostname or "").lower() + if not host: + return url.rstrip("/").lower() + port = parsed.port + if port is None: + port = 443 if parsed.scheme == "https" else 80 + return f"{host}:{port}" + + +def _own_endpoint_key(server: _TorchHTTPServer) -> str: + advertised = getattr(server, "advertised_endpoint", None) + if advertised: + return _endpoint_key(advertised) + host, port = server.server_address + return _endpoint_key(f"http://{host}:{port}") + + +def _clamp_downstream_hops( + hops: list[dict], + backend: TorchModelShard | None, +) -> list[dict]: + """Ensure downstream start_layer continues after this shard's layers.""" + if not hops or backend is None: + return hops + shard_end = getattr(backend, "shard_end", None) + if shard_end is None: + return hops + min_start = int(shard_end) + 1 + clamped: list[dict] = [] + for hop in hops: + adjusted = dict(hop) + if int(adjusted.get("start_layer", 0)) < min_start: + adjusted["start_layer"] = min_start + clamped.append(adjusted) + return clamped + + +def _format_downstream_route(hops: list[dict]) -> str: + return ", ".join( + f"{h['endpoint']}@{h.get('start_layer', 0)}" for h in hops + ) + + +def _write_progress_line(state: list[bool], message: str, *, final: bool = False) -> None: + """Rewrite one in-place progress line (\\r) or finish with a newline.""" + if final: + if state[0]: + sys.stdout.write("\r" + message + "\n") + state[0] = False + else: + print(message, flush=True) + return + if state[0]: + sys.stdout.write("\r" + message) + else: + sys.stdout.write(message) + state[0] = True + sys.stdout.flush() + + def _relay_hop( relay_addr: str, path: str, @@ -87,10 +150,31 @@ class _TorchHTTPServer(http.server.HTTPServer): self.route_timeout = route_timeout self.debug = debug self.max_loaded_shards = max(1, max_loaded_shards) + self.advertised_endpoint: str | None = None self.total_requests: int = 0 self.failed_requests: int = 0 self.queue_depth: int = 0 self._stats_lock = threading.Lock() + self._active_requests: dict[str, dict[str, Any]] = {} + + def snapshot_current_requests(self) -> list[dict[str, Any]]: + """In-flight request snapshots for tracker heartbeats.""" + now = time.monotonic() + with self._stats_lock: + out: list[dict[str, Any]] = [] + for rec in self._active_requests.values(): + elapsed = max(now - float(rec["started"]), 1e-6) + tokens = int(rec.get("tokens") or 0) + out.append({ + "request_id": str(rec["request_id"]), + "model": str(rec.get("model") or ""), + "kind": str(rec.get("kind") or "chat"), + "tokens": tokens, + "elapsed_seconds": round(elapsed, 1), + "tokens_per_sec": round(tokens / elapsed, 2) if tokens > 0 else 0.0, + "routing_complete": bool(rec.get("routing_complete")), + }) + return out def resolve_backend(self, model_name: str | None) -> TorchModelShard | None: if not model_name: @@ -113,10 +197,53 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): def log_message(self, fmt, *args): # noqa: suppress request logs in tests pass + def _request_id(self) -> str: + return ( + self.headers.get("X-Meshnet-Request-Id") + or self.headers.get("X-Request-Id") + or f"local-{time.time_ns():x}" + ) + def _request_log_suffix(self) -> str: 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 "" + def _track_request_begin( + self, + server: "_TorchHTTPServer", + request_id: str, + model: str, + ) -> None: + with server._stats_lock: + server._active_requests[request_id] = { + "request_id": request_id, + "model": model, + "kind": "chat", + "started": time.monotonic(), + "tokens": 0, + "routing_complete": False, + } + + def _track_request_progress( + self, + server: "_TorchHTTPServer", + request_id: str, + *, + tokens: int, + routing_complete: bool = False, + ) -> None: + with server._stats_lock: + rec = server._active_requests.get(request_id) + if rec is None: + return + rec["tokens"] = tokens + if routing_complete: + rec["routing_complete"] = True + + def _track_request_end(self, server: "_TorchHTTPServer", request_id: str) -> None: + with server._stats_lock: + server._active_requests.pop(request_id, None) + def do_POST(self): server: _TorchHTTPServer = self.server # type: ignore[assignment] if self.path == "/forward": @@ -294,12 +421,14 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): def _handle_chat_completions(self) -> None: server: _TorchHTTPServer = self.server # type: ignore[assignment] + request_id = self._request_id() with server._stats_lock: server.total_requests += 1 server.queue_depth += 1 try: - self._do_chat_completions(server) + self._do_chat_completions(server, request_id) finally: + self._track_request_end(server, request_id) with server._stats_lock: server.queue_depth -= 1 @@ -308,7 +437,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): with server._stats_lock: server.failed_requests += 1 - def _do_chat_completions(self, server: "_TorchHTTPServer") -> None: + def _do_chat_completions(self, server: "_TorchHTTPServer", request_id: str) -> None: body = self._read_json_body() if body is None: return @@ -325,6 +454,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): temperature = float(body.get("temperature") or 1.0) top_p = float(body.get("top_p") or 1.0) + self._track_request_begin(server, request_id, model_name) print( f" [node] processing chat model={model_name!r} stream={stream} " f"max_tokens={max_tokens}{self._request_log_suffix()}", @@ -335,6 +465,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): # Avoids the single-token-per-forward-pass limitation of the distributed path. if backend.is_head and backend.is_tail: gen_started = time.monotonic() + progress_line = [False] try: if stream: token_count = 0 @@ -346,13 +477,19 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): ): if token_text: token_count += 1 + self._track_request_progress( + server, request_id, tokens=token_count, routing_complete=True, + ) yield token_text self._stream_openai_response(_counting_stream(), model_name) - print( + elapsed = time.monotonic() - gen_started + tps = token_count / max(elapsed, 1e-6) + _write_progress_line( + progress_line, f" [node] chat complete (stream) tokens={token_count} " - f"elapsed_s={time.monotonic() - gen_started:.1f}{self._request_log_suffix()}", - flush=True, + f"elapsed_s={elapsed:.1f} tps={tps:.2f}{self._request_log_suffix()}", + final=True, ) else: text = backend.generate_text(messages, max_tokens, temperature, top_p) @@ -414,10 +551,12 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): stream_emit = None if stream: stream_emit = self._start_openai_stream(model_name) + self._track_request_progress(server, request_id, tokens=0, routing_complete=True) _GENERATION_LOG_INTERVAL = 5.0 gen_started = time.monotonic() last_gen_log = gen_started + progress_line = [False] for step in range(max_tokens): try: @@ -437,20 +576,33 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): if stream_emit is not None: stream_emit(token_str) current_text = current_text + token_str + self._track_request_progress( + server, + request_id, + tokens=len(generated), + routing_complete=True, + ) now = time.monotonic() if step == 0 or now - last_gen_log >= _GENERATION_LOG_INTERVAL: - print( + elapsed = now - gen_started + token_count = len(generated) + tps = token_count / max(elapsed, 1e-6) + _write_progress_line( + progress_line, f" [node] generating step={step + 1}/{max_tokens} " - f"tokens={len(generated)} elapsed_s={now - gen_started:.1f}", - flush=True, + f"tokens={token_count} elapsed_s={elapsed:.1f} tps={tps:.2f}", ) last_gen_log = now if generated: - print( - f" [node] generation complete tokens={len(generated)} " - f"elapsed_s={time.monotonic() - gen_started:.1f}", - flush=True, + elapsed = time.monotonic() - gen_started + token_count = len(generated) + tps = token_count / max(elapsed, 1e-6) + _write_progress_line( + progress_line, + f" [node] generation complete tokens={token_count} " + f"elapsed_s={elapsed:.1f} tps={tps:.2f}", + final=True, ) result_text = "".join(generated) @@ -467,6 +619,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): start_layer tells each downstream node which layer to begin from, enabling correct execution when shard ranges overlap. """ + server: _TorchHTTPServer = self.server # type: ignore[assignment] + active_backend = backend or server.backend + # Fast path: tracker pre-resolved the downstream route and injected it as a header. injected = self.headers.get("X-Meshnet-Route") if injected: @@ -485,14 +640,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): hops.append(hop) elif isinstance(item, str): hops.append({"endpoint": item, "start_layer": 0}) - print(f" [node] using injected downstream route: {[h['endpoint'] for h in hops]}", flush=True) + hops = _clamp_downstream_hops(hops, active_backend) + print( + f" [node] using injected downstream route: {_format_downstream_route(hops)}", + flush=True, + ) return hops except (json.JSONDecodeError, TypeError, KeyError): pass # Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject). - server: _TorchHTTPServer = self.server # type: ignore[assignment] - active_backend = backend or server.backend if server.tracker_url is None: return [] route_model = getattr(active_backend, "model_id", None) or model @@ -500,23 +657,31 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}" with urllib.request.urlopen(url, timeout=server.route_timeout) as r: route_resp = json.loads(r.read()) - own_port = server.server_address[1] + own_key = _own_endpoint_key(server) nodes_info = route_resp.get("nodes", []) - hops = [] - covered_up_to: int | None = None + hops: list[dict] = [] + passed_self = False for node_info in nodes_info: ep = node_info.get("endpoint", "") - if ep.rstrip("/").endswith(f":{own_port}"): - covered_up_to = node_info.get("shard_end") + if not ep: continue - if covered_up_to is None: - covered_up_to = (node_info.get("shard_start") or 1) - 1 - hop = {"endpoint": ep, "start_layer": covered_up_to + 1} + if _endpoint_key(ep) == own_key: + passed_self = True + continue + if not passed_self: + continue + hop = { + "endpoint": ep, + "start_layer": int(node_info.get("start_layer", 0)), + } if node_info.get("relay_addr"): hop["relay_addr"] = str(node_info["relay_addr"]) hops.append(hop) - covered_up_to = node_info.get("shard_end", covered_up_to) - print(f" [node] tracker downstream route: {[h['endpoint'] for h in hops]}", flush=True) + hops = _clamp_downstream_hops(hops, active_backend) + print( + f" [node] tracker downstream route: {_format_downstream_route(hops)}", + flush=True, + ) return hops except Exception as exc: print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True) @@ -849,6 +1014,12 @@ class TorchNodeServer: def queue_depth(self) -> int: return self._server.queue_depth if self._server is not None else 0 + @property + def current_requests(self) -> list[dict[str, Any]]: + if self._server is None: + return [] + return self._server.snapshot_current_requests() + @property def loaded_model_ids(self) -> list[str]: return list(self._backends.keys()) @@ -928,6 +1099,11 @@ class TorchNodeServer: self._thread.start() return self.port + def set_advertised_endpoint(self, endpoint: str) -> None: + """Set the LAN-facing endpoint used for route self-detection.""" + if self._server is not None: + self._server.advertised_endpoint = endpoint + def stop(self) -> None: if self._server is None: return diff --git a/packages/node/pyproject.toml b/packages/node/pyproject.toml index 454aaf2..3ba2d56 100644 --- a/packages/node/pyproject.toml +++ b/packages/node/pyproject.toml @@ -1,33 +1,34 @@ -[build-system] -requires = ["setuptools>=64"] -build-backend = "setuptools.build_meta" - -[project] -name = "meshnet-node" -version = "0.1.0" -description = "Distributed Inference Network node client" -requires-python = ">=3.10" - -dependencies = [ - "cryptography>=41", - "huggingface-hub>=0.20", - "accelerate>=0.28", - "bitsandbytes>=0.43", - "rich>=13", - "safetensors>=0.4", - "torch>=2.1", - "transformers>=5.12", - "websockets>=13", - "zstandard>=0.22", - "kernels>=0.11.1,<0.16", -] - -[project.scripts] -meshnet-node = "meshnet_node.cli:main" - -[tool.setuptools.packages.find] -where = ["."] -include = ["meshnet_node*"] - -[tool.setuptools.package-data] -meshnet_node = ["*.json"] +[build-system] +requires = ["setuptools>=64"] +build-backend = "setuptools.build_meta" + +[project] +name = "meshnet-node" +version = "0.1.0" +description = "Distributed Inference Network node client" +requires-python = ">=3.10" + +dependencies = [ + "cryptography>=41", + "huggingface-hub>=0.20", + "accelerate>=0.28", + "bitsandbytes>=0.43", + "rich>=13", + "safetensors>=0.4", + "torch>=2.1", + "transformers>=5.12", + "triton-windows>=3.7; platform_system == 'Windows'", + "websockets>=13", + "zstandard>=0.22", + "kernels>=0.11.1,<0.16", +] + +[project.scripts] +meshnet-node = "meshnet_node.cli:main" + +[tool.setuptools.packages.find] +where = ["."] +include = ["meshnet_node*"] + +[tool.setuptools.package-data] +meshnet_node = ["*.json"] diff --git a/packages/tracker/meshnet_tracker/cli.py b/packages/tracker/meshnet_tracker/cli.py index 642ccd5..0b40b55 100644 --- a/packages/tracker/meshnet_tracker/cli.py +++ b/packages/tracker/meshnet_tracker/cli.py @@ -1,13 +1,13 @@ -"""meshnet-tracker CLI entry point.""" - -import argparse -import os -import sys -import time -from pathlib import Path - -from .accounts import DEFAULT_ACCOUNTS_DB_PATH -from .billing import DEFAULT_BILLING_DB_PATH +"""meshnet-tracker CLI entry point.""" + +import argparse +import os +import sys +import time +from pathlib import Path + +from .accounts import DEFAULT_ACCOUNTS_DB_PATH +from .billing import DEFAULT_BILLING_DB_PATH from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH from .logging_setup import ( DEFAULT_LOG_BACKUP_COUNT, @@ -15,258 +15,299 @@ from .logging_setup import ( DEFAULT_LOG_MAX_BYTES, configure_tracker_file_logging, ) +from .routing_stats import RoutingConfig from .server import ( - DEFAULT_CALLER_CREDIT_USDT, - DEFAULT_DEVNET_TOPUP_USDT, - TrackerServer, - derive_relay_url_from_public_tracker_url, -) - -DEFAULT_REGISTRY_DB_PATH = "meshnet_registry.sqlite3" - - -def _load_env_file(path: Path) -> None: - """Load simple KEY=VALUE pairs from an env file without overriding env vars.""" - if not path.exists(): - return - try: - lines = path.read_text().splitlines() - except OSError: - return - for line in lines: - text = line.strip() - if not text or text.startswith("#"): - continue - if text.startswith("export "): - text = text[len("export "):].strip() - if "=" not in text: - continue - key, value = text.split("=", 1) - key = key.strip() - if not key or key in os.environ: - continue - value = value.strip() - if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}: - value = value[1:-1] - os.environ[key] = value - - -def _load_env_defaults() -> None: - """Load local and user-level tracker env defaults before parsing arguments.""" - _load_env_file(Path.cwd() / ".env") - _load_env_file(Path.home() / ".config" / "meshnet" / "secrets.env") - - -def main() -> None: - _load_env_defaults() - common = argparse.ArgumentParser(add_help=False) - common.add_argument("--host", default="0.0.0.0", help="Host interface to listen on") - common.add_argument("--port", type=int, default=8080, help="Port to listen on") - common.add_argument( - "--heartbeat-timeout", - type=float, - default=30.0, - help="Seconds before a node is removed from the registry after missed heartbeat", - ) - common.add_argument( - "--cluster-peers", - default="", - help="Comma-separated URLs of peer tracker nodes (enables Raft cluster mode)", - ) - common.add_argument( - "--self-url", - default=None, - help="This tracker's own URL as seen by peers (auto-derived from --host/--port if omitted)", - ) - common.add_argument( - "--stats-db", - default=None, - metavar="PATH", - help="SQLite database path for persistent model usage statistics", - ) - common.add_argument( - "--relay-url", - default=None, - help="Public ws(s):// relay URL advertised to nodes, for example wss://ai.neuron.d-popov.com/ws", - ) - common.add_argument( - "--billing-db", - default=DEFAULT_BILLING_DB_PATH, - metavar="PATH", - help=( - "SQLite database path for the USDT billing ledger " - f"(default: {DEFAULT_BILLING_DB_PATH}; ADR-0015)" - ), - ) - common.add_argument( - "--no-billing", - action="store_true", - help="Disable the USDT billing ledger", - ) - common.add_argument( - "--max-charge-per-request", - type=float, - default=None, - help=( - "Reject chat completion requests whose prompt plus requested completion " - "token bound would cost more than this many USDT" - ), - ) - common.add_argument( - "--starting-credit", - type=float, - default=DEFAULT_CALLER_CREDIT_USDT, - metavar="USDT", - help=( - "One-time Caller Credit granted when an account creates its first " - f"API key (default: {DEFAULT_CALLER_CREDIT_USDT}; set 0 to require " - "deposits before inference)" - ), - ) - common.add_argument( - "--devnet-topup", - type=float, - default=DEFAULT_DEVNET_TOPUP_USDT, - metavar="USDT", - help=( - "Dashboard devnet top-up faucet: each click credits this many USDT " - f"to one of the account's keys (default: {DEFAULT_DEVNET_TOPUP_USDT}; " - "MUST be 0 on mainnet deployments)" - ), - ) - common.add_argument( - "--registry-db", - default=DEFAULT_REGISTRY_DB_PATH, - metavar="PATH", - help=( - "SQLite database path for persisted strike/ban/reputation registry " - f"state (default: {DEFAULT_REGISTRY_DB_PATH})" - ), - ) - common.add_argument( - "--no-registry-contracts", - action="store_true", - help="Disable the local contract registry used for strike/ban/reputation enforcement", - ) - common.add_argument( - "--accounts-db", - default=DEFAULT_ACCOUNTS_DB_PATH, - metavar="PATH", - help=( - "SQLite database path for dashboard user accounts " - f"(default: {DEFAULT_ACCOUNTS_DB_PATH})" - ), - ) - common.add_argument( - "--no-accounts", - action="store_true", - help="Disable dashboard user accounts (registration/login)", - ) - common.add_argument( - "--solana-rpc-url", - default=None, - help="Solana RPC URL (e.g. https://api.devnet.solana.com); enables the on-chain treasury", - ) - common.add_argument( - "--usdt-mint", - default=None, - help="SPL mint address of (mock) USDT — see scripts/devnet_setup.py", - ) - common.add_argument( - "--treasury-keypair", - default=None, - metavar="PATH", - help="Treasury keypair JSON path (only on settlement-capable trackers)", - ) - common.add_argument( - "--settle-period", - type=float, - default=86400.0, - help="Max seconds between payouts to a node (dev: 60, prod: 86400)", - ) - common.add_argument( - "--payout-threshold", - type=float, - default=5.0, - help="Pending USDT that triggers an immediate payout (dev: 0)", - ) - common.add_argument( - "--payout-dust-floor", - type=float, - default=0.01, - help="Never pay out less than this many USDT", - ) - common.add_argument( - "--validator-service-token", - default=None, - help=( - "Service token the validator uses on POST /v1/billing/forfeit " - "(default: MESHNET_VALIDATOR_SERVICE_TOKEN env; ADR-0017)" - ), - ) - common.add_argument( - "--hive-secret", - default=None, - help=( - "Shared secret authenticating gossip between tracker peers " - "(default: MESHNET_HIVE_SECRET env; required for multi-tracker replication)" - ), - ) - common.add_argument( - "--toploc-calibration-db", - default=None, - metavar="PATH", - help=( - "SQLite path for the AH-021 honest-noise TOPLOC calibration corpus " - "(enables POST /v1/calibration/toploc/run + GET /v1/calibration/toploc/results)" - ), - ) - common.add_argument( - "--toploc-reference-node-url", - default=None, - help="Reference node the calibration job teacher-forces claimed tokens against (see validator README)", - ) - common.add_argument( - "--toploc-calibration-gate-min-hardware-profiles", - type=int, - default=1, - help=( - "Distinct (GPU model, dtype) profiles the corpus must cover before " - "gate_status.ready is true (alpha exception: fleet size is acceptable)" - ), - ) - common.add_argument( - "--enable-hf-pricing", - action="store_true", - help=( - "Enable the daily dynamic pricing refresh (issue 23): for presets with a " - "curated hf_aliases list, sets the client price to 80%% of the cheapest " - "matching HuggingFace inference-marketplace rate. Presets without " - "hf_aliases are unaffected and keep their static price." - ), - ) - common.add_argument( - "--hf-pricing-log-db", - default=None, - metavar="PATH", - help=( - "SQLite database path for the dynamic pricing change log " - f"(default when --enable-hf-pricing is set: {DEFAULT_HF_PRICING_LOG_DB_PATH}; " - "enables GET /v1/pricing/hf/history)" - ), - ) - common.add_argument( - "--hf-pricing-refresh-interval", - type=float, - default=86400.0, - help="Seconds between dynamic pricing refresh passes (default: daily)", - ) + DEFAULT_CALLER_CREDIT_USDT, + DEFAULT_DEVNET_TOPUP_USDT, + TrackerServer, + derive_relay_url_from_public_tracker_url, +) + +DEFAULT_REGISTRY_DB_PATH = "meshnet_registry.sqlite3" + + +def _load_env_file(path: Path) -> None: + """Load simple KEY=VALUE pairs from an env file without overriding env vars.""" + if not path.exists(): + return + try: + lines = path.read_text().splitlines() + except OSError: + return + for line in lines: + text = line.strip() + if not text or text.startswith("#"): + continue + if text.startswith("export "): + text = text[len("export "):].strip() + if "=" not in text: + continue + key, value = text.split("=", 1) + key = key.strip() + if not key or key in os.environ: + continue + value = value.strip() + if len(value) >= 2 and value[0] == value[-1] and value[0] in {"'", '"'}: + value = value[1:-1] + os.environ[key] = value + + +def _load_env_defaults() -> None: + """Load local and user-level tracker env defaults before parsing arguments.""" + _load_env_file(Path.cwd() / ".env") + _load_env_file(Path.home() / ".config" / "meshnet" / "secrets.env") + + +def _routing_config_from_args(args: argparse.Namespace) -> RoutingConfig | None: + """Build a RoutingConfig from CLI flags; None keeps env-var/server defaults.""" + overrides = { + "explore_share": args.route_explore_share, + "weight_alpha": args.route_weight_alpha, + "stats_half_life_seconds": args.route_stats_half_life, + } + set_values = {key: value for key, value in overrides.items() if value is not None} + if not set_values: + return None + return RoutingConfig(**set_values) + + +def main() -> None: + _load_env_defaults() + common = argparse.ArgumentParser(add_help=False) + common.add_argument("--host", default="0.0.0.0", help="Host interface to listen on") + common.add_argument("--port", type=int, default=8080, help="Port to listen on") + common.add_argument( + "--heartbeat-timeout", + type=float, + default=30.0, + help="Seconds before a node is removed from the registry after missed heartbeat", + ) + common.add_argument( + "--cluster-peers", + default="", + help="Comma-separated URLs of peer tracker nodes (enables Raft cluster mode)", + ) + common.add_argument( + "--self-url", + default=None, + help="This tracker's own URL as seen by peers (auto-derived from --host/--port if omitted)", + ) + common.add_argument( + "--stats-db", + default=None, + metavar="PATH", + help="SQLite database path for persistent model usage statistics", + ) + common.add_argument( + "--relay-url", + default=None, + help="Public ws(s):// relay URL advertised to nodes, for example wss://ai.neuron.d-popov.com/ws", + ) + common.add_argument( + "--billing-db", + default=DEFAULT_BILLING_DB_PATH, + metavar="PATH", + help=( + "SQLite database path for the USDT billing ledger " + f"(default: {DEFAULT_BILLING_DB_PATH}; ADR-0015)" + ), + ) + common.add_argument( + "--no-billing", + action="store_true", + help="Disable the USDT billing ledger", + ) + common.add_argument( + "--max-charge-per-request", + type=float, + default=None, + help=( + "Reject chat completion requests whose prompt plus requested completion " + "token bound would cost more than this many USDT" + ), + ) + common.add_argument( + "--starting-credit", + type=float, + default=DEFAULT_CALLER_CREDIT_USDT, + metavar="USDT", + help=( + "One-time Caller Credit granted when an account creates its first " + f"API key (default: {DEFAULT_CALLER_CREDIT_USDT}; set 0 to require " + "deposits before inference)" + ), + ) + common.add_argument( + "--devnet-topup", + type=float, + default=DEFAULT_DEVNET_TOPUP_USDT, + metavar="USDT", + help=( + "Dashboard devnet top-up faucet: each click credits this many USDT " + f"to one of the account's keys (default: {DEFAULT_DEVNET_TOPUP_USDT}; " + "MUST be 0 on mainnet deployments)" + ), + ) + common.add_argument( + "--registry-db", + default=DEFAULT_REGISTRY_DB_PATH, + metavar="PATH", + help=( + "SQLite database path for persisted strike/ban/reputation registry " + f"state (default: {DEFAULT_REGISTRY_DB_PATH})" + ), + ) + common.add_argument( + "--no-registry-contracts", + action="store_true", + help="Disable the local contract registry used for strike/ban/reputation enforcement", + ) + common.add_argument( + "--accounts-db", + default=DEFAULT_ACCOUNTS_DB_PATH, + metavar="PATH", + help=( + "SQLite database path for dashboard user accounts " + f"(default: {DEFAULT_ACCOUNTS_DB_PATH})" + ), + ) + common.add_argument( + "--no-accounts", + action="store_true", + help="Disable dashboard user accounts (registration/login)", + ) + common.add_argument( + "--solana-rpc-url", + default=None, + help="Solana RPC URL (e.g. https://api.devnet.solana.com); enables the on-chain treasury", + ) + common.add_argument( + "--usdt-mint", + default=None, + help="SPL mint address of (mock) USDT — see scripts/devnet_setup.py", + ) + common.add_argument( + "--treasury-keypair", + default=None, + metavar="PATH", + help="Treasury keypair JSON path (only on settlement-capable trackers)", + ) + common.add_argument( + "--settle-period", + type=float, + default=86400.0, + help="Max seconds between payouts to a node (dev: 60, prod: 86400)", + ) + common.add_argument( + "--payout-threshold", + type=float, + default=5.0, + help="Pending USDT that triggers an immediate payout (dev: 0)", + ) + common.add_argument( + "--payout-dust-floor", + type=float, + default=0.01, + help="Never pay out less than this many USDT", + ) + common.add_argument( + "--validator-service-token", + default=None, + help=( + "Service token the validator uses on POST /v1/billing/forfeit " + "(default: MESHNET_VALIDATOR_SERVICE_TOKEN env; ADR-0017)" + ), + ) + common.add_argument( + "--hive-secret", + default=None, + help=( + "Shared secret authenticating gossip between tracker peers " + "(default: MESHNET_HIVE_SECRET env; required for multi-tracker replication)" + ), + ) + common.add_argument( + "--toploc-calibration-db", + default=None, + metavar="PATH", + help=( + "SQLite path for the AH-021 honest-noise TOPLOC calibration corpus " + "(enables POST /v1/calibration/toploc/run + GET /v1/calibration/toploc/results)" + ), + ) + common.add_argument( + "--toploc-reference-node-url", + default=None, + help="Reference node the calibration job teacher-forces claimed tokens against (see validator README)", + ) + common.add_argument( + "--toploc-calibration-gate-min-hardware-profiles", + type=int, + default=1, + help=( + "Distinct (GPU model, dtype) profiles the corpus must cover before " + "gate_status.ready is true (alpha exception: fleet size is acceptable)" + ), + ) + common.add_argument( + "--enable-hf-pricing", + action="store_true", + help=( + "Enable the daily dynamic pricing refresh (issue 23): for presets with a " + "curated hf_aliases list, sets the client price to 80%% of the cheapest " + "matching HuggingFace inference-marketplace rate. Presets without " + "hf_aliases are unaffected and keep their static price." + ), + ) + common.add_argument( + "--hf-pricing-log-db", + default=None, + metavar="PATH", + help=( + "SQLite database path for the dynamic pricing change log " + f"(default when --enable-hf-pricing is set: {DEFAULT_HF_PRICING_LOG_DB_PATH}; " + "enables GET /v1/pricing/hf/history)" + ), + ) + common.add_argument( + "--hf-pricing-refresh-interval", + type=float, + default=86400.0, + help="Seconds between dynamic pricing refresh passes (default: daily)", + ) common.add_argument( "--models-dir", default=None, metavar="PATH", help="Local HuggingFace snapshot root advertised as tracker model-file source (default: MESHNET_MODELS_DIR)", ) + common.add_argument( + "--route-explore-share", + type=float, + default=None, + metavar="FRACTION", + help=( + "Fraction of requests routed down unproven/stale routes to gather " + "throughput statistics (ADR-0021; default 0.3, lower once traffic grows)" + ), + ) + common.add_argument( + "--route-weight-alpha", + type=float, + default=None, + metavar="ALPHA", + help=( + "Traffic weight exponent among proven routes: share ∝ tps^alpha " + "(default 1.0 — a 1.5x-faster route gets 1.5x the traffic)" + ), + ) + common.add_argument( + "--route-stats-half-life", + type=float, + default=None, + metavar="SECONDS", + help="Half-life for decaying route throughput observations (default 600)", + ) common.add_argument( "--log-dir", default=DEFAULT_LOG_DIR, @@ -295,18 +336,18 @@ def main() -> None: action="store_true", help="Disable rotating tracker log files and only write to the terminal", ) - - parser = argparse.ArgumentParser( - prog="meshnet-tracker", - description="Distributed Inference Network node registry and route selection", - parents=[common], - ) - subparsers = parser.add_subparsers(dest="command") - - subparsers.add_parser("start", help="Start the tracker server", parents=[common]) - - args = parser.parse_args() - + + parser = argparse.ArgumentParser( + prog="meshnet-tracker", + description="Distributed Inference Network node registry and route selection", + parents=[common], + ) + subparsers = parser.add_subparsers(dest="command") + + subparsers.add_parser("start", help="Start the tracker server", parents=[common]) + + args = parser.parse_args() + if args.command in {None, "start"}: if not args.no_file_logs: log_dir = configure_tracker_file_logging( @@ -316,62 +357,63 @@ def main() -> None: ) print(f"meshnet-tracker logs: {log_dir}", flush=True) cluster_peers = [u.strip() for u in args.cluster_peers.split(",") if u.strip()] - relay_url = args.relay_url or derive_relay_url_from_public_tracker_url(args.self_url) - treasury = None - if args.solana_rpc_url and args.usdt_mint and args.treasury_keypair: - from meshnet_contracts.solana_adapter import SolanaCustodialTreasury - - treasury = SolanaCustodialTreasury( - args.solana_rpc_url, args.usdt_mint, args.treasury_keypair, - ) - contracts = None - if not args.no_registry_contracts: - from meshnet_contracts import LocalSolanaContracts # type: ignore[import-not-found] - - contracts = LocalSolanaContracts(registry_db=args.registry_db) - server = TrackerServer( - host=args.host, - port=args.port, - heartbeat_timeout=args.heartbeat_timeout, - cluster_peers=cluster_peers or None, - cluster_self_url=args.self_url, - stats_db=getattr(args, "stats_db", None), - relay_url=relay_url, - enable_billing=not args.no_billing, - billing_db=None if args.no_billing else args.billing_db, - max_charge_per_request=args.max_charge_per_request, - starting_credit=args.starting_credit, - devnet_topup_amount=args.devnet_topup, - contracts=contracts, - accounts_db=None if args.no_accounts else args.accounts_db, - treasury=treasury, - settle_period=args.settle_period, - payout_threshold=args.payout_threshold, - payout_dust_floor=args.payout_dust_floor, - validator_service_token=args.validator_service_token, - hive_secret=args.hive_secret, - toploc_calibration_db=args.toploc_calibration_db, - toploc_reference_node_url=args.toploc_reference_node_url, - toploc_calibration_gate_min_hardware_profiles=args.toploc_calibration_gate_min_hardware_profiles, - enable_hf_pricing=args.enable_hf_pricing, - hf_pricing_log_db=( - args.hf_pricing_log_db - or (DEFAULT_HF_PRICING_LOG_DB_PATH if args.enable_hf_pricing else None) - ), - hf_pricing_refresh_interval=args.hf_pricing_refresh_interval, - models_dir=args.models_dir, - ) - port = server.start() - print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True) - try: - while True: - time.sleep(1) - except KeyboardInterrupt: - server.stop() - sys.exit(0) - else: - parser.print_help() - - -if __name__ == "__main__": - main() + relay_url = args.relay_url or derive_relay_url_from_public_tracker_url(args.self_url) + treasury = None + if args.solana_rpc_url and args.usdt_mint and args.treasury_keypair: + from meshnet_contracts.solana_adapter import SolanaCustodialTreasury + + treasury = SolanaCustodialTreasury( + args.solana_rpc_url, args.usdt_mint, args.treasury_keypair, + ) + contracts = None + if not args.no_registry_contracts: + from meshnet_contracts import LocalSolanaContracts # type: ignore[import-not-found] + + contracts = LocalSolanaContracts(registry_db=args.registry_db) + server = TrackerServer( + host=args.host, + port=args.port, + heartbeat_timeout=args.heartbeat_timeout, + cluster_peers=cluster_peers or None, + cluster_self_url=args.self_url, + stats_db=getattr(args, "stats_db", None), + relay_url=relay_url, + enable_billing=not args.no_billing, + billing_db=None if args.no_billing else args.billing_db, + max_charge_per_request=args.max_charge_per_request, + starting_credit=args.starting_credit, + devnet_topup_amount=args.devnet_topup, + contracts=contracts, + accounts_db=None if args.no_accounts else args.accounts_db, + treasury=treasury, + settle_period=args.settle_period, + payout_threshold=args.payout_threshold, + payout_dust_floor=args.payout_dust_floor, + validator_service_token=args.validator_service_token, + hive_secret=args.hive_secret, + toploc_calibration_db=args.toploc_calibration_db, + toploc_reference_node_url=args.toploc_reference_node_url, + toploc_calibration_gate_min_hardware_profiles=args.toploc_calibration_gate_min_hardware_profiles, + enable_hf_pricing=args.enable_hf_pricing, + hf_pricing_log_db=( + args.hf_pricing_log_db + or (DEFAULT_HF_PRICING_LOG_DB_PATH if args.enable_hf_pricing else None) + ), + hf_pricing_refresh_interval=args.hf_pricing_refresh_interval, + models_dir=args.models_dir, + routing_config=_routing_config_from_args(args), + ) + port = server.start() + print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True) + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + server.stop() + sys.exit(0) + else: + parser.print_help() + + +if __name__ == "__main__": + main() diff --git a/packages/tracker/meshnet_tracker/dashboard.html b/packages/tracker/meshnet_tracker/dashboard.html index 2c15801..dc055e8 100644 --- a/packages/tracker/meshnet_tracker/dashboard.html +++ b/packages/tracker/meshnet_tracker/dashboard.html @@ -1,1427 +1,1627 @@ - - - - - -meshnet tracker - - - -
-

meshnet tracker

- - -
- -
-

Account

loading…
-

Tracker hive

loading…
-

Nodes & coverage

loading…
-

Model usage (RPM)

loading…
-

Call wall

loading...
-
-

Chat / inference

-
- -
-
- -
select a model to start
-
-
-
Send a message to start this conversation.
-
-
-
- - -
-
-
-
-
-

Usage summary

login required
-

Node throughput

login required
-

Request history

login required
-

Node pending payouts

admin login required
-

Settlement history

admin login required
-

All accounts (admin)

-

Strikes / bans / forfeitures

admin login required
-

Client balances

admin login required
-

Console output

admin login required
-
- - - + + + + + +meshnet tracker + + + +
+

meshnet tracker

+ + +
+ +
+

Account

loading…
+

Tracker hive

loading…
+

Nodes & coverage

loading…
+

Model usage (RPM)

loading…
+

Routing (learned)

loading…
+

Call wall

loading...
+
+

Chat / inference

+
+ +
+
+ +
select a model to start
+
+
+
Send a message to start this conversation.
+
+
+
+ + +
+
+
+
+
+

Usage summary

login required
+

Node throughput

login required
+

Request history

login required
+

Node pending payouts

admin login required
+

Settlement history

admin login required
+

All accounts (admin)

+

Strikes / bans / forfeitures

admin login required
+

Client balances

admin login required
+

Console output

admin login required
+
+ + + diff --git a/packages/tracker/meshnet_tracker/routing_stats.py b/packages/tracker/meshnet_tracker/routing_stats.py new file mode 100644 index 0000000..5fca3b8 --- /dev/null +++ b/packages/tracker/meshnet_tracker/routing_stats.py @@ -0,0 +1,257 @@ +"""Learned route statistics for dynamic bandit-style route selection (ADR-0021). + +The tracker treats each viable route (ordered chain of node shards covering a +model) as a bandit arm. Observed end-to-end tokens/sec per route is kept as a +time-decayed EWMA. Selection splits traffic between: + +- **exploit**: weighted-random among *proven* routes, weight ∝ tps ** alpha + (alpha=1.0 → a 1.5x-faster route gets 1.5x the traffic); +- **scout**: with probability `explore_share`, the least-measured unproven or + stale route is chosen so the tracker keeps learning as the network morphs. + +Staleness has two mechanisms: +- continuous: sample mass decays with `stats_half_life_seconds`, so old + observations fade; +- abrupt: every node join/leave bumps the model's *topology epoch*; stats from + an older epoch keep their EWMA as a prior but drop back into the scout pool + until re-measured. + +Route signatures embed node ids and shard ranges, so a node re-registering +with a different shard produces a new arm automatically. +""" + +from __future__ import annotations + +import random +import threading +import time +from dataclasses import dataclass, field +from typing import Any, Iterable + + +@dataclass(frozen=True) +class RoutingConfig: + explore_share: float = 0.3 + weight_alpha: float = 1.0 + stats_half_life_seconds: float = 600.0 + min_sample_tokens: int = 8 + # One fresh sample has mass 1.0 and decays from there; 0.5 keeps a single + # observation "proven" for one half-life before demoting it to the scout pool. + min_proven_weight: float = 0.5 + max_candidate_routes: int = 8 + prune_after_seconds: float = 86400.0 + + +@dataclass +class RouteStat: + ewma_tps: float = 0.0 + weight: float = 0.0 # decayed effective sample mass + last_sample_ts: float = 0.0 + epoch: int = 0 + samples: int = 0 # lifetime raw sample count (display only) + + def decayed_weight(self, now: float, half_life: float) -> float: + if self.weight <= 0.0: + return 0.0 + age = max(0.0, now - self.last_sample_ts) + return self.weight * 0.5 ** (age / half_life) + + +@dataclass +class RouteCandidate: + nodes: list[Any] + signature: str + prior_tps: float = 0.0 + + +def route_signature(model_key: str, nodes: Iterable[Any]) -> str: + hops = "->".join( + f"{getattr(n, 'node_id', '?')}[{getattr(n, 'shard_start', '?')}-{getattr(n, 'shard_end', '?')}]" + for n in nodes + ) + return f"{model_key}|{hops}" + + +class RouteStatsStore: + """Thread-safe per-route decayed throughput statistics.""" + + def __init__(self, config: RoutingConfig | None = None) -> None: + self.config = config or RoutingConfig() + self._lock = threading.Lock() + self._stats: dict[str, RouteStat] = {} + self._epochs: dict[str, int] = {} + + def epoch(self, model_key: str) -> int: + with self._lock: + return self._epochs.get(model_key, 0) + + def bump_epoch(self, model_keys: Iterable[str | None]) -> None: + """Mark the topology changed for the given model keys (node join/leave).""" + with self._lock: + for key in model_keys: + if key: + self._epochs[key] = self._epochs.get(key, 0) + 1 + + def record_sample( + self, + model_key: str, + signature: str, + tokens: int, + elapsed_seconds: float, + now: float | None = None, + ) -> bool: + """Fold one completed request into the route's EWMA. + + Returns False (and records nothing) for samples below + `min_sample_tokens` — near-empty completions come from broken routes + and would poison the arm with meaningless throughput values. + """ + cfg = self.config + if tokens < cfg.min_sample_tokens or elapsed_seconds <= 0.0: + return False + tps = tokens / elapsed_seconds + ts = time.time() if now is None else now + with self._lock: + stat = self._stats.get(signature) + if stat is None: + stat = RouteStat() + self._stats[signature] = stat + carried = stat.decayed_weight(ts, cfg.stats_half_life_seconds) + total = carried + 1.0 + stat.ewma_tps = (stat.ewma_tps * carried + tps) / total + stat.weight = total + stat.last_sample_ts = ts + stat.epoch = self._epochs.get(model_key, 0) + stat.samples += 1 + return True + + def snapshot(self, signature: str, model_key: str, now: float | None = None) -> dict: + """Point-in-time view of one route's learned state.""" + ts = time.time() if now is None else now + cfg = self.config + with self._lock: + stat = self._stats.get(signature) + current_epoch = self._epochs.get(model_key, 0) + if stat is None: + return {"tps": None, "weight": 0.0, "samples": 0, "status": "unsampled"} + weight = stat.decayed_weight(ts, cfg.stats_half_life_seconds) + if stat.epoch != current_epoch: + status = "stale" + elif weight < cfg.min_proven_weight: + status = "decayed" if stat.samples else "unsampled" + else: + status = "proven" + return { + "tps": round(stat.ewma_tps, 4) if stat.samples else None, + "weight": round(weight, 4), + "samples": stat.samples, + "status": status, + } + + def prune(self, now: float | None = None) -> int: + """Drop routes with no samples for `prune_after_seconds`.""" + ts = time.time() if now is None else now + cutoff = ts - self.config.prune_after_seconds + with self._lock: + dead = [sig for sig, stat in self._stats.items() if stat.last_sample_ts < cutoff] + for sig in dead: + del self._stats[sig] + return len(dead) + + +def choose_route( + candidates: list[RouteCandidate], + store: RouteStatsStore, + model_key: str, + rng: random.Random | None = None, + now: float | None = None, +) -> tuple[RouteCandidate | None, dict]: + """Pick a route: ε-scout among unproven arms, else weighted ∝ tps**alpha. + + Returns (candidate, decision) where decision explains the pick for logs + and diagnostics: {"mode": "scout"|"exploit"|"prior", ...}. + """ + if not candidates: + return None, {"mode": "none"} + rng = rng or random + cfg = store.config + proven: list[tuple[RouteCandidate, float]] = [] + scouts: list[tuple[RouteCandidate, float]] = [] + for cand in candidates: + snap = store.snapshot(cand.signature, model_key, now=now) + if snap["status"] == "proven": + proven.append((cand, max(float(snap["tps"] or 0.0), 1e-6))) + else: + scouts.append((cand, float(snap["weight"]))) + if scouts and (not proven or rng.random() < cfg.explore_share): + # Least-measured first so new/stale arms accumulate samples fastest; + # tiebreak on prior estimate so plausible routes get scouted first. + scouts.sort(key=lambda item: (item[1], -item[0].prior_tps)) + pick = scouts[0][0] + return pick, {"mode": "scout", "signature": pick.signature} + if proven: + weights = [tps ** cfg.weight_alpha for _, tps in proven] + pick = rng.choices([cand for cand, _ in proven], weights=weights, k=1)[0] + return pick, { + "mode": "exploit", + "signature": pick.signature, + "candidates": len(proven), + } + # No stats anywhere yet — fall back to the prior (benchmark-derived) estimate. + weights = [max(cand.prior_tps, 1e-6) ** cfg.weight_alpha for cand in candidates] + pick = rng.choices(candidates, weights=weights, k=1)[0] + return pick, {"mode": "prior", "signature": pick.signature} + + +def route_table( + candidates: list[RouteCandidate], + store: RouteStatsStore, + model_key: str, + now: float | None = None, +) -> list[dict]: + """Diagnostics rows: learned tps, coefficient vs best, expected traffic share.""" + cfg = store.config + rows = [] + for cand in candidates: + snap = store.snapshot(cand.signature, model_key, now=now) + rows.append({"candidate": cand, **snap}) + proven = [r for r in rows if r["status"] == "proven"] + scouts = [r for r in rows if r["status"] != "proven"] + best_tps = max((float(r["tps"]) for r in proven), default=0.0) + exploit_budget = 1.0 - (cfg.explore_share if scouts and proven else 0.0) + if not proven: + exploit_budget = 0.0 + weight_sum = sum(float(r["tps"]) ** cfg.weight_alpha for r in proven) or 1.0 + out = [] + for r in rows: + cand: RouteCandidate = r["candidate"] + if r["status"] == "proven": + share = exploit_budget * (float(r["tps"]) ** cfg.weight_alpha) / weight_sum + coefficient = round(float(r["tps"]) / best_tps, 3) if best_tps else None + else: + share = ( + (cfg.explore_share if proven else 1.0) / len(scouts) + if scouts + else 0.0 + ) + coefficient = None + out.append({ + "signature": cand.signature, + "hops": [ + { + "node_id": getattr(n, "node_id", "?"), + "shard": f"{getattr(n, 'shard_start', '?')}-{getattr(n, 'shard_end', '?')}", + "endpoint": getattr(n, "endpoint", "?"), + } + for n in cand.nodes + ], + "tps": r["tps"], + "coefficient": coefficient, + "expected_share": round(share, 4), + "samples": r["samples"], + "weight": r["weight"], + "status": r["status"], + "prior_tps": round(cand.prior_tps, 4), + }) + out.sort(key=lambda r: (-(r["tps"] or 0.0), -r["prior_tps"])) + return out diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 1e53cb2..97cfb7c 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -1,5959 +1,6207 @@ -"""Tracker HTTP server — node registry and route selection for the inference mesh. - -HTTP API contract: -- POST /v1/nodes/register - Request JSON: { - "endpoint": "http://host:port", "shard_start": int, "shard_end": int, - "model": str optional, "shard_checksum": str optional, - "hardware_profile": object, "wallet_address": str optional, - "score": number optional - } - Response 200: {"node_id": str} - Response 400: {"error": str} -- POST /v1/nodes/{node_id}/heartbeat - Response 200: {} - Response 404: {"error": "node not found"} -- GET /v1/route?model= - Response 200: {"route": ["http://node-a", "http://node-b"]} - Response 400/404/503: {"error": str} -- GET /v1/routes?model=&redundancy= - Response 200: {"routes": [{"route": [...], "nodes": [...]}]} - Response 400/404/503: {"error": str} -""" - -import http.cookies -import http.server -import hashlib -import itertools -import json -import os -import select -import socketserver -import sqlite3 -import tarfile -import threading -import time -import urllib.parse -import urllib.request -import uuid -from collections import deque -from dataclasses import dataclass, field -from importlib.resources import files -from pathlib import Path -from typing import Any - -from .accounts import DEFAULT_ACCOUNTS_DB_PATH, AccountStore -from .auth import is_validator_token, sign_hive_request, verify_hive_request -from .wallet_proof import binding_message, verify_wallet_signature -from .billing import DEFAULT_BILLING_DB_PATH, BillingLedger -from .calibration import DEFAULT_CALIBRATION_DB_PATH, ToplocCalibrationStore -from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH, HfPricingLog, refresh_preset_price -from .gossip import NodeGossip -from .logging_setup import tracker_logger -from .model_files import files_for_layer_range, snapshot_dir_for_repo -from .raft import RaftNode - - -_CONSOLE_LIMIT = 300 -_PROXY_PROGRESS_LOG_INTERVAL = 5.0 -_SESSION_COOKIE_NAME = "meshnet_session" - - -def _preset_price_keys(name: str, preset: dict) -> set[str]: - """All model strings a client may bill under for one preset. - - ``BillingLedger.price_for`` is keyed by the raw ``model`` string in the - request, so the preset price must be registered under the preset name, - its ``hf_repo``, and every alias — otherwise ``unsloth/Qwen…`` style - requests silently fall back to the default rate. - """ - keys = {name} - hf_repo = preset.get("hf_repo") - if isinstance(hf_repo, str) and hf_repo: - keys.add(hf_repo) - for alias in preset.get("aliases") or []: - if isinstance(alias, str) and alias: - keys.add(alias) - return keys - - -def derive_relay_url_from_public_tracker_url(url: str | None) -> str | None: - """Return wss://host/ws when url is a public HTTPS tracker origin.""" - if not url: - return None - parsed = urllib.parse.urlparse(url) - if parsed.scheme != "https": - return None - host = parsed.hostname - if not host or host in ("127.0.0.1", "localhost"): - return None - return f"wss://{parsed.netloc}/ws" - - -def _load_model_presets_from_data() -> dict[str, dict]: - """Load recommended model presets from package JSON data.""" - try: - raw = files("meshnet_tracker").joinpath("model_presets.json").read_text() - data = json.loads(raw) - except Exception: - return {} - models = data.get("models", {}) - if not isinstance(models, dict): - return {} - return { - str(name): preset - for name, preset in models.items() - if isinstance(preset, dict) - } - - -DEFAULT_MODEL_PRESETS: dict[str, dict] = { - "stub-model": { - "layers_start": 0, - "layers_end": 31, - "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, - }, - "openai-community/gpt2": { - "layers_start": 0, - "layers_end": 11, - "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024}, - }, - **_load_model_presets_from_data(), -} - -def _clone_model_presets(presets: dict[str, dict]) -> dict[str, dict]: - """Shallow-copy each preset dict so a runtime mutation (e.g. issue 23's - dynamic pricing refresh writing hf_last_price_per_1k/hf_last_updated) - never leaks into the shared module-level DEFAULT_MODEL_PRESETS and from - there into other TrackerServer instances in the same process.""" - return {name: dict(preset) for name, preset in presets.items()} - - -DEFAULT_VRAM_BYTES = 8 * 1024 * 1024 * 1024 -DEFAULT_RAM_BYTES = 16 * 1024 * 1024 * 1024 - - -def _snapshot_regular_files(snapshot_dir: Path) -> list[str]: - return sorted( - path.relative_to(snapshot_dir).as_posix() - for path in snapshot_dir.rglob("*") - if path.is_file() - ) -DEFAULT_QUANTIZATIONS = ["bfloat16"] -DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0 -# US-039/US-040 — single source of truth for the credit defaults (referenced by -# TrackerServer, _TrackerHTTPServer, and the CLI). Alpha runs devnet-friendly; -# flip both to 0 before any deployment holding a mainnet treasury. -DEFAULT_CALLER_CREDIT_USDT = 1.0 -DEFAULT_DEVNET_TOPUP_USDT = 1.0 - - -def _model_aliases(model: str | None) -> set[str]: - """Return stable lookup aliases for a model repo or display name.""" - if not model: - return set() - normalized = model.strip() - if not normalized: - return set() - aliases = {normalized} - short = normalized.rsplit("/", 1)[-1] - aliases.add(short) - lowered = short.lower() - aliases.add(lowered) - if lowered.endswith("-instruct"): - aliases.add(lowered.removesuffix("-instruct")) - return aliases - - -def _preset_aliases(name: str, preset: dict | None) -> set[str]: - aliases = _model_aliases(name) - if not preset: - return aliases - hf_repo = preset.get("hf_repo") - if isinstance(hf_repo, str): - aliases |= _model_aliases(hf_repo) - for alias in preset.get("aliases", []) or []: - if isinstance(alias, str): - aliases |= _model_aliases(alias) - return aliases - - -def _resolve_model_preset(model_presets: dict, model: str) -> tuple[str, dict] | tuple[None, None]: - requested = _model_aliases(model) - for name, preset in model_presets.items(): - if requested & _preset_aliases(name, preset): - return name, preset - return None, None - - -def _node_matches_model(node: "_NodeEntry", model: str) -> bool: - requested = _model_aliases(model) - if not requested: - return False - return bool(requested & (_model_aliases(node.model) | _model_aliases(node.hf_repo))) - - -def _node_matches_preset(node: "_NodeEntry", name: str, preset: dict) -> bool: - requested = _preset_aliases(name, preset) - return bool(requested & (_model_aliases(node.model) | _model_aliases(node.hf_repo))) - - -class _RollingCounter: - """Circular-bucket request counter. - - Tracks events in `num_buckets` slots, each covering `bucket_seconds` of wall time. - Old buckets are silently discarded when their epoch expires. - """ - - def __init__(self, num_buckets: int, bucket_seconds: int) -> None: - self._num = num_buckets - self._bsec = bucket_seconds - self._counts: list[int] = [0] * num_buckets - self._epochs: list[int] = [-1] * num_buckets # which epoch each slot holds - - def _epoch(self, now: float) -> int: - return int(now) // self._bsec - - def record(self, now: float | None = None) -> None: - t = now if now is not None else time.time() - ep = self._epoch(t) - idx = ep % self._num - if self._epochs[idx] != ep: - self._counts[idx] = 0 - self._epochs[idx] = ep - self._counts[idx] += 1 - - def rpm(self, now: float | None = None) -> float: - t = now if now is not None else time.time() - cutoff = self._epoch(t) - self._num # epochs <= this are stale - total = sum(self._counts[i] for i in range(self._num) if self._epochs[i] > cutoff) - window_minutes = (self._num * self._bsec) / 60.0 - return total / window_minutes if window_minutes > 0 else 0.0 - - def buckets(self) -> list[tuple[int, int]]: - return [(self._epochs[i], self._counts[i]) for i in range(self._num)] - - def restore_buckets(self, data: list[tuple[int, int]]) -> None: - for i, (ep, cnt) in enumerate(data): - if i < self._num: - self._epochs[i] = ep - self._counts[i] = cnt - - -class _RollingThroughput: - """Circular buckets for observed tokens/sec. - - Each bucket stores total output tokens and total elapsed seconds. TPS for a - window is the ratio across non-stale buckets, which avoids over-weighting - small fast requests. - """ - - def __init__(self, num_buckets: int, bucket_seconds: int) -> None: - self._num = num_buckets - self._bsec = bucket_seconds - self._tokens: list[int] = [0] * num_buckets - self._seconds: list[float] = [0.0] * num_buckets - self._samples: list[int] = [0] * num_buckets - self._epochs: list[int] = [-1] * num_buckets - - def _epoch(self, now: float) -> int: - return int(now) // self._bsec - - def record(self, total_tokens: int, elapsed_seconds: float, now: float | None = None) -> None: - if total_tokens <= 0 or elapsed_seconds <= 0: - return - t = now if now is not None else time.time() - ep = self._epoch(t) - idx = ep % self._num - if self._epochs[idx] != ep: - self._tokens[idx] = 0 - self._seconds[idx] = 0.0 - self._samples[idx] = 0 - self._epochs[idx] = ep - self._tokens[idx] += int(total_tokens) - self._seconds[idx] += float(elapsed_seconds) - self._samples[idx] += 1 - - def stats(self, now: float | None = None) -> dict: - t = now if now is not None else time.time() - cutoff = self._epoch(t) - self._num - tokens = 0 - seconds = 0.0 - samples = 0 - for i in range(self._num): - if self._epochs[i] > cutoff: - tokens += self._tokens[i] - seconds += self._seconds[i] - samples += self._samples[i] - return { - "tokens_per_sec": round(tokens / seconds, 4) if seconds > 0 else None, - "tokens": tokens, - "seconds": round(seconds, 6), - "sample_count": samples, - } - - def buckets(self) -> list[tuple[int, int, float, int]]: - return [ - (self._epochs[i], self._tokens[i], self._seconds[i], self._samples[i]) - for i in range(self._num) - ] - - def restore_buckets(self, data: list[tuple[int, int, float, int]]) -> None: - for i, (ep, tokens, seconds, samples) in enumerate(data): - if i < self._num: - self._epochs[i] = ep - self._tokens[i] = tokens - self._seconds[i] = seconds - self._samples[i] = samples - - -class _ModelStats: - """Three rolling windows for one model: last hour, last day, last month.""" - - def __init__(self) -> None: - self.per_minute = _RollingCounter(60, 60) # 60 × 1-min buckets = 1 hour - self.per_hour = _RollingCounter(24, 3600) # 24 × 1-hr buckets = 1 day - self.per_day = _RollingCounter(30, 86400) # 30 × 1-day buckets = ~1 month - - def record(self, now: float | None = None) -> None: - t = now if now is not None else time.time() - self.per_minute.record(t) - self.per_hour.record(t) - self.per_day.record(t) - - def rpms(self, now: float | None = None) -> dict: - t = now if now is not None else time.time() - return { - "rpm_last_hour": round(self.per_minute.rpm(t), 4), - "rpm_last_day": round(self.per_hour.rpm(t), 4), - "rpm_last_month": round(self.per_day.rpm(t), 4), - } - - -class _NodeModelThroughput: - """Observed throughput windows for one node/model pair.""" - - def __init__(self) -> None: - self.per_minute = _RollingThroughput(60, 60) - self.per_hour = _RollingThroughput(24, 3600) - self.per_day = _RollingThroughput(30, 86400) - - def record(self, total_tokens: int, elapsed_seconds: float, now: float | None = None) -> None: - t = now if now is not None else time.time() - self.per_minute.record(total_tokens, elapsed_seconds, t) - self.per_hour.record(total_tokens, elapsed_seconds, t) - self.per_day.record(total_tokens, elapsed_seconds, t) - - def stats(self, now: float | None = None) -> dict: - hour = self.per_minute.stats(now) - day = self.per_hour.stats(now) - month = self.per_day.stats(now) - return { - "tokens_per_sec_last_hour": hour["tokens_per_sec"], - "tokens_per_sec_last_day": day["tokens_per_sec"], - "tokens_per_sec_last_month": month["tokens_per_sec"], - "sample_count_last_hour": hour["sample_count"], - "tokens_last_hour": hour["tokens"], - "seconds_last_hour": hour["seconds"], - } - - -class _StatsCollector: - """Thread-safe model request stats with SQLite persistence and peer slice merging.""" - - SAVE_INTERVAL = 60.0 # seconds between DB saves - - def __init__(self, db_path: str | None = None) -> None: - self._lock = threading.Lock() - self._local: dict[str, _ModelStats] = {} - self._node_model: dict[tuple[str, str], _NodeModelThroughput] = {} - self._peer_rpms: dict[str, dict[str, dict]] = {} # tracker_url -> model -> rpms - self._db_path = db_path - if db_path: - self._init_db() - self._load_from_db() - - # ---------- public API ---------- - - def record_request(self, model: str, now: float | None = None) -> None: - t = now if now is not None else time.time() - with self._lock: - if model not in self._local: - self._local[model] = _ModelStats() - self._local[model].record(t) - - def record_node_throughput( - self, - node_id: str, - model: str, - total_tokens: int, - elapsed_seconds: float, - now: float | None = None, - ) -> None: - if not node_id or not model or total_tokens <= 0 or elapsed_seconds <= 0: - return - t = now if now is not None else time.time() - with self._lock: - key = (node_id, model) - if key not in self._node_model: - self._node_model[key] = _NodeModelThroughput() - self._node_model[key].record(total_tokens, elapsed_seconds, t) - - def get_local_rpms(self, now: float | None = None) -> dict[str, dict]: - t = now if now is not None else time.time() - with self._lock: - return {m: s.rpms(t) for m, s in self._local.items()} - - def get_node_model_stats(self, node_id: str, model: str, now: float | None = None) -> dict: - t = now if now is not None else time.time() - empty = { - "tokens_per_sec_last_hour": None, - "tokens_per_sec_last_day": None, - "tokens_per_sec_last_month": None, - "sample_count_last_hour": 0, - "tokens_last_hour": 0, - "seconds_last_hour": 0.0, - } - with self._lock: - stat = self._node_model.get((node_id, model)) - return stat.stats(t) if stat is not None else dict(empty) - - def get_node_throughput_stats(self, now: float | None = None) -> dict[str, dict]: - t = now if now is not None else time.time() - with self._lock: - result: dict[str, dict] = {} - for (node_id, model), stat in self._node_model.items(): - result.setdefault(node_id, {"models": {}})["models"][model] = stat.stats(t) - return result - - def merge_peer_rpms(self, tracker_url: str, rpms: dict[str, dict]) -> None: - with self._lock: - self._peer_rpms[tracker_url] = dict(rpms) - - def get_combined_stats(self, now: float | None = None) -> dict: - t = now if now is not None else time.time() - with self._lock: - combined: dict[str, dict] = {} - for model, ms in self._local.items(): - combined[model] = ms.rpms(t) - for _url, peer in self._peer_rpms.items(): - for model, rpms in peer.items(): - if model not in combined: - combined[model] = {"rpm_last_hour": 0.0, "rpm_last_day": 0.0, "rpm_last_month": 0.0} - for key in ("rpm_last_hour", "rpm_last_day", "rpm_last_month"): - combined[model][key] = round(combined[model].get(key, 0.0) + rpms.get(key, 0.0), 4) - return combined - - # ---------- persistence ---------- - - def _init_db(self) -> None: - con = sqlite3.connect(self._db_path) # type: ignore[arg-type] - con.execute( - "CREATE TABLE IF NOT EXISTS model_rpm_buckets " - "(model TEXT, window TEXT, bucket_idx INTEGER, bucket_epoch INTEGER, count INTEGER, " - "PRIMARY KEY (model, window, bucket_idx))" - ) - con.execute( - "CREATE TABLE IF NOT EXISTS node_model_tps_buckets " - "(node_id TEXT, model TEXT, window TEXT, bucket_idx INTEGER, bucket_epoch INTEGER, " - "tokens INTEGER, seconds REAL, samples INTEGER, " - "PRIMARY KEY (node_id, model, window, bucket_idx))" - ) - con.commit() - con.close() - - def save_to_db(self) -> None: - if not self._db_path: - return - with self._lock: - rows = [] - for model, ms in self._local.items(): - for window_name, counter in ( - ("hour", ms.per_minute), - ("day", ms.per_hour), - ("month", ms.per_day), - ): - for idx, (ep, cnt) in enumerate(counter.buckets()): - if ep >= 0: - rows.append((model, window_name, idx, ep, cnt)) - tps_rows = [] - for (node_id, model), stat in self._node_model.items(): - for window_name, counter in ( - ("hour", stat.per_minute), - ("day", stat.per_hour), - ("month", stat.per_day), - ): - for idx, (ep, tokens, seconds, samples) in enumerate(counter.buckets()): - if ep >= 0: - tps_rows.append((node_id, model, window_name, idx, ep, tokens, seconds, samples)) - con = sqlite3.connect(self._db_path) # type: ignore[arg-type] - con.executemany( - "INSERT OR REPLACE INTO model_rpm_buckets VALUES (?,?,?,?,?)", rows - ) - con.executemany( - "INSERT OR REPLACE INTO node_model_tps_buckets VALUES (?,?,?,?,?,?,?,?)", - tps_rows, - ) - con.commit() - con.close() - - def _load_from_db(self) -> None: - con = sqlite3.connect(self._db_path) # type: ignore[arg-type] - rows = con.execute("SELECT model, window, bucket_idx, bucket_epoch, count FROM model_rpm_buckets").fetchall() - tps_rows = con.execute( - "SELECT node_id, model, window, bucket_idx, bucket_epoch, tokens, seconds, samples " - "FROM node_model_tps_buckets" - ).fetchall() - con.close() - grouped: dict[str, dict[str, list[tuple[int, int]]]] = {} - for model, window, idx, ep, cnt in rows: - grouped.setdefault(model, {}).setdefault(window, []).append((idx, ep, cnt)) - for model, windows in grouped.items(): - ms = _ModelStats() - for window_name, entries in windows.items(): - counter = {"hour": ms.per_minute, "day": ms.per_hour, "month": ms.per_day}.get(window_name) - if counter is None: - continue - data = [(0, -1)] * counter._num - for idx, ep, cnt in entries: - if 0 <= idx < counter._num: - data[idx] = (ep, cnt) - counter.restore_buckets(data) - self._local[model] = ms - tps_grouped: dict[tuple[str, str], dict[str, list[tuple[int, int, int, float, int]]]] = {} - for node_id, model, window, idx, ep, tokens, seconds, samples in tps_rows: - tps_grouped.setdefault((node_id, model), {}).setdefault(window, []).append( - (idx, ep, tokens, seconds, samples) - ) - for key, windows in tps_grouped.items(): - stat = _NodeModelThroughput() - for window_name, entries in windows.items(): - counter = {"hour": stat.per_minute, "day": stat.per_hour, "month": stat.per_day}.get(window_name) - if counter is None: - continue - data = [(0, 0, 0.0, 0)] * counter._num - for idx, ep, tokens, seconds, samples in entries: - if 0 <= idx < counter._num: - data[idx] = (ep, int(tokens), float(seconds), int(samples)) - counter.restore_buckets(data) - self._node_model[key] = stat - - -class _NodeEntry: - __slots__ = ( - "node_id", "endpoint", "shard_start", "shard_end", - "model", "hf_repo", "num_layers", "model_metadata", "shard_checksum", "downloaded_models", "hardware_profile", "wallet_address", - "score", "vram_bytes", "ram_bytes", "quantizations", "max_loaded_shards", - "benchmark_tokens_per_sec", "quantization", "managed_assignment", - "model_tokens_per_sec", - "pending_directives", "last_heartbeat", "tracker_mode", - "relay_addr", "cert_fingerprint", "peer_id", - # heartbeat stats (reported by node, cumulative) - "total_requests", "failed_requests", "queue_depth", "proxy_inflight", "uptime_seconds", - "current_requests", - "status", # "ready" | "loading" - "heartbeats_expected", "heartbeats_received", - # dynamic reassignment queued by the tracker - "pending_new_assignment", - ) - - def __init__( - self, - node_id: str, - endpoint: str, - shard_start: int | None, - shard_end: int | None, - model: str | None, - shard_checksum: str | None, - hardware_profile: dict, - wallet_address: str | None, - score: float, - vram_bytes: int = DEFAULT_VRAM_BYTES, - ram_bytes: int = DEFAULT_RAM_BYTES, - quantizations: list[str] | None = None, - max_loaded_shards: int = 1, - benchmark_tokens_per_sec: float = DEFAULT_BENCHMARK_TOKENS_PER_SEC, - quantization: str | None = None, - managed_assignment: bool = False, - tracker_mode: bool = False, - hf_repo: str | None = None, - num_layers: int | None = None, - model_metadata: dict | None = None, - downloaded_models: list[dict] | None = None, - relay_addr: str | None = None, - cert_fingerprint: str | None = None, - peer_id: str | None = None, - ) -> None: - self.node_id = node_id - self.endpoint = endpoint - self.shard_start = shard_start - self.shard_end = shard_end - self.model = model - self.shard_checksum = shard_checksum - self.hardware_profile = hardware_profile - self.wallet_address = wallet_address - self.score = score - self.vram_bytes = vram_bytes - self.ram_bytes = ram_bytes - self.quantizations = quantizations or list(DEFAULT_QUANTIZATIONS) - self.max_loaded_shards = max_loaded_shards - self.benchmark_tokens_per_sec = benchmark_tokens_per_sec - self.quantization = quantization - self.managed_assignment = managed_assignment - self.model_tokens_per_sec: dict[str, float] = {} - self.tracker_mode = tracker_mode - self.hf_repo = hf_repo - self.num_layers = num_layers - self.model_metadata = dict(model_metadata or {}) - self.downloaded_models = [dict(item) for item in (downloaded_models or []) if isinstance(item, dict)] - self.relay_addr = relay_addr - self.cert_fingerprint = cert_fingerprint - self.peer_id = peer_id - self.pending_directives: list[dict] = [] - self.last_heartbeat: float = time.monotonic() - self.total_requests: int = 0 - self.failed_requests: int = 0 - self.queue_depth: int = 0 - self.proxy_inflight: int = 0 - self.current_requests: list[dict] = [] - self.uptime_seconds: float = 0.0 - self.status: str = "ready" - self.heartbeats_expected: int = 0 - self.heartbeats_received: int = 0 - self.pending_new_assignment: dict | None = None - - -def _effective_throughput(node: "_NodeEntry", model: str | None = None) -> float: - """Effective tokens/s accounting for current queue depth.""" - observed = None - if model: - observed = node.model_tokens_per_sec.get(model) - if observed is None: - for alias in _model_aliases(model): - observed = node.model_tokens_per_sec.get(alias) - if observed is not None: - break - base = observed if observed is not None and observed > 0 else node.benchmark_tokens_per_sec - return base / (_effective_queue_depth(node) + 1) - - -def _effective_queue_depth(node: "_NodeEntry") -> int: - """Best current load estimate: heartbeat queue or tracker-routed in-flight.""" - return max(node.queue_depth, node.proxy_inflight) - - -_CURRENT_REQUEST_FIELDS = frozenset({ - "request_id", "model", "kind", "tokens", "tokens_per_sec", - "elapsed_seconds", "routing_complete", -}) - - -def _normalize_current_requests(items: object, *, limit: int = 32) -> list[dict]: - """Sanitize node-reported in-flight request snapshots from heartbeats.""" - if not isinstance(items, list): - return [] - out: list[dict] = [] - for item in items[:limit]: - if not isinstance(item, dict): - continue - request_id = item.get("request_id") - if not request_id: - continue - rec: dict = {"request_id": str(request_id)} - for key in _CURRENT_REQUEST_FIELDS: - if key == "request_id" or key not in item: - continue - value = item[key] - if key in {"tokens"}: - rec[key] = int(value) - elif key in {"tokens_per_sec", "elapsed_seconds"}: - rec[key] = float(value) - elif key == "routing_complete": - rec[key] = bool(value) - else: - rec[key] = str(value) - out.append(rec) - return out - - -def _record_proxy_inflight( - server: "_TrackerHTTPServer", - nodes: list["_NodeEntry"], - delta: int, -) -> None: - with server.lock: - for node in nodes: - node.proxy_inflight = max(0, node.proxy_inflight + delta) - - -def _reputation_multiplier(node: "_NodeEntry", contracts: Any | None) -> float: - """ADR-0018 §6: veteran, good-standing nodes route more -- earnings scale - with tenure/reputation. Maps reputation [0, 1] to a [0.5, 1.0] routing - weight so a damaged-but-not-banned wallet loses traffic gradually rather - than being cut off outright (banning is handled separately).""" - if contracts is None or not node.wallet_address: - return 1.0 - reputation = contracts.registry.get_wallet(node.wallet_address).reputation - return 0.5 + 0.5 * reputation - - -def _select_route( - nodes: list[_NodeEntry], - required_start: int, - required_end: int, - model: str | None = None, - contracts: Any | None = None, -) -> tuple[list[_NodeEntry], str]: - """Greedy interval-cover biased toward fast, lightly-loaded, reputable nodes. - - Among nodes that equally advance coverage, prefer the one with higher - effective throughput: observed per-model tokens/sec / (queue_depth + 1), - falling back to startup benchmark_tokens_per_sec until observations exist, - weighted by the node's reputation multiplier (see `_reputation_multiplier`). - Tiebreak: higher shard_end (fewer hops). - """ - candidates = sorted( - [node for node in nodes if node.shard_start is not None and node.shard_end is not None], - key=lambda n: (n.shard_start, -n.shard_end), # type: ignore[operator] - ) - route: list[_NodeEntry] = [] - covered_up_to = required_start - 1 - - def _routing_score(node: "_NodeEntry") -> float: - return _effective_throughput(node, model) * _reputation_multiplier(node, contracts) - - while covered_up_to < required_end: - best: _NodeEntry | None = None - for node in candidates: - if node.shard_start <= covered_up_to + 1 and node.shard_end > covered_up_to: - if best is None: - best = node - elif node.shard_end > best.shard_end: - best = node - elif node.shard_end == best.shard_end and _routing_score(node) > _routing_score(best): - best = node - if best is None: - missing = covered_up_to + 1 - return [], f"no route available: no registered node covers layer {missing}" - route.append(best) - covered_up_to = best.shard_end - candidates = [n for n in candidates if n is not best] - - return route, "" - - -def _node_route_summary(nodes: list["_NodeEntry"]) -> list[dict]: - return [ - { - "node_id": node.node_id, - "model": node.model, - "hf_repo": node.hf_repo, - "endpoint": node.endpoint, - "shard": f"{node.shard_start}-{node.shard_end}", - "num_layers": node.num_layers, - "queue_depth": _effective_queue_depth(node), - "heartbeat_queue_depth": node.queue_depth, - "proxy_inflight": node.proxy_inflight, - "current_requests": list(node.current_requests), - } - for node in nodes - ] - - -def _coverage_percentage( - nodes: list[_NodeEntry], - required_start: int, - required_end: int, -) -> float: - required_layers = required_end - required_start + 1 - if required_layers <= 0: - return 0.0 - - intervals = sorted( - ( - (max(required_start, node.shard_start), min(required_end, node.shard_end)) - for node in nodes - if node.shard_start is not None - and node.shard_end is not None - if node.shard_end >= required_start and node.shard_start <= required_end - ), - key=lambda interval: interval[0], - ) - covered = 0 - covered_until = required_start - 1 - for start, end in intervals: - if end <= covered_until: - continue - next_start = max(start, covered_until + 1) - if next_start > end: - continue - covered += end - next_start + 1 - covered_until = end - return round((covered / required_layers) * 100, 2) - - -def _served_model_copies( - nodes: list[_NodeEntry], - required_start: int, - required_end: int, -) -> float: - required_layers = required_end - required_start + 1 - if required_layers <= 0: - return 0.0 - - layer_counts = [] - for layer in range(required_start, required_end + 1): - count = 0 - for node in nodes: - if node.shard_start is None or node.shard_end is None: - continue - if node.shard_start <= layer <= node.shard_end: - count += 1 - layer_counts.append(count) - - if not layer_counts: - return 0.0 - - complete_copies = min(layer_counts) - residual_layers = sum(1 for count in layer_counts if count > complete_copies) - return round(complete_copies + (residual_layers / required_layers), 2) - - -def _model_health_summary( - server: "_TrackerHTTPServer", - model: str | None, - hf_repo: str | None = None, -) -> dict: - lookup = hf_repo or model - if not lookup: - return {} - - resolved_name, preset = _resolve_model_preset(server.model_presets, lookup) - if preset is not None: - required_start, required_end = _preset_layer_bounds(preset) - model_nodes = [ - node for node in server.registry.values() - if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] - ] - model_id = resolved_name or lookup - else: - model_nodes = [ - node for node in server.registry.values() - if _node_matches_model(node, lookup) - and node.shard_start is not None - and node.shard_end is not None - and node.num_layers is not None - ] - if not model_nodes: - return { - "model": lookup, - "served_model_copies": 0.0, - "coverage_percentage": 0.0, - "node_count": 0, - } - required_start = 0 - required_end = max(node.num_layers for node in model_nodes if node.num_layers is not None) - 1 - model_id = hf_repo or model or lookup - - return { - "model": model_id, - "required_start": required_start, - "required_end": required_end, - "served_model_copies": _served_model_copies(model_nodes, required_start, required_end), - "coverage_percentage": _coverage_percentage(model_nodes, required_start, required_end), - "node_count": len(model_nodes), - } - - -def _preset_layer_bounds(preset: dict) -> tuple[int, int]: - start = int(preset.get("layers_start", 0)) - if "layers_end" in preset: - return start, int(preset["layers_end"]) - return start, start + int(preset["total_layers"]) - 1 - - -def _preset_bytes_per_layer(preset: dict) -> dict[str, int]: - raw = preset.get("bytes_per_layer", preset.get("bytes_per_layer_at_quant", {})) - if isinstance(raw, dict) and raw: - return {str(quant): int(value) for quant, value in raw.items()} - return {"bfloat16": 30 * 1024 * 1024} - - -def _node_quantization(node: _NodeEntry, preset: dict) -> str: - bytes_per_layer = _preset_bytes_per_layer(preset) - if node.quantization in bytes_per_layer: - return node.quantization - for quantization in node.quantizations: - if quantization in bytes_per_layer: - return quantization - return next(iter(bytes_per_layer)) - - -def _node_memory_budget_bytes(node: _NodeEntry) -> tuple[int, str]: - """Return the memory pool used for shard-capacity planning.""" - if node.vram_bytes > 0: - return node.vram_bytes, "vram" - if node.ram_bytes > 0: - return node.ram_bytes, "ram" - return DEFAULT_RAM_BYTES, "ram-default" - - -def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int: - bytes_per_layer = _preset_bytes_per_layer(preset) - quantization = _node_quantization(node, preset) - layer_bytes = bytes_per_layer[quantization] - if layer_bytes <= 0: - return 0 - memory_budget_bytes, _ = _node_memory_budget_bytes(node) - return int((memory_budget_bytes * 0.8) // layer_bytes) - - -def _node_capacity_summary(node: _NodeEntry, preset: dict | None = None) -> dict: - """Operator-facing capacity fields for inspection endpoints.""" - memory_budget_bytes, memory_budget_source = _node_memory_budget_bytes(node) - summary = { - "vram_bytes": node.vram_bytes, - "ram_bytes": node.ram_bytes, - "memory_budget_bytes": memory_budget_bytes, - "memory_budget_source": memory_budget_source, - "max_loaded_shards": node.max_loaded_shards, - "quantizations": list(node.quantizations), - "quantization": node.quantization, - "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec, - "effective_throughput": round(_effective_throughput(node), 4), - } - if preset is not None: - summary["max_assignable_layers"] = _node_layer_capacity(node, preset) - return summary - - -def _node_memory_budget_for_preset(node: _NodeEntry, preset: dict | None = None) -> int: - budget, _source = _node_memory_budget_bytes(node) - if preset is None: - return int(budget * 0.8) - return _node_layer_capacity(node, preset) * max(1, next(iter(_preset_bytes_per_layer(preset).values()))) - - -def _pool_summary(nodes: list[_NodeEntry], preset: dict | None = None) -> dict: - total_vram = sum(max(0, node.vram_bytes) for node in nodes) - total_ram = sum(max(0, node.ram_bytes) for node in nodes) - total_budget = sum(_node_memory_budget_bytes(node)[0] for node in nodes) - effective_budget = sum(_node_memory_budget_for_preset(node, preset) for node in nodes) - return { - "node_count": len(nodes), - "total_vram_bytes": total_vram, - "total_ram_bytes": total_ram, - "total_memory_budget_bytes": total_budget, - "effective_assignable_memory_bytes": effective_budget, - "total_benchmark_tokens_per_sec": round(sum(node.benchmark_tokens_per_sec for node in nodes), 4), - "total_effective_throughput": round(sum(_effective_throughput(node) for node in nodes), 4), - } - - -def _repo_demand_rpm( - server: "_TrackerHTTPServer", - repo: str, - local_rpms: dict[str, dict], -) -> float: - """Look up recent request rate by hf_repo or preset alias.""" - if repo in local_rpms: - return float(local_rpms[repo].get("rpm_last_hour", 0.0) or 0.0) - resolved_name, _preset = _resolve_model_preset(server.model_presets, repo) - if resolved_name and resolved_name in local_rpms: - return float(local_rpms[resolved_name].get("rpm_last_hour", 0.0) or 0.0) - return 0.0 - - -def _preset_for_node(server: "_TrackerHTTPServer", node: _NodeEntry) -> dict | None: - if node.hf_repo: - _resolved, preset = _resolve_model_preset(server.model_presets, node.hf_repo) - if preset is not None: - return preset - if node.model: - preset = server.model_presets.get(node.model) - if preset is not None: - return preset - if node.hf_repo and node.num_layers: - return _hf_rebalance_preset([node]) - return None - - -def _assignment_memory_bytes(node: _NodeEntry, preset: dict | None) -> int: - if preset is None or node.shard_start is None or node.shard_end is None: - return 0 - layers = node.shard_end - node.shard_start + 1 - if layers <= 0: - return 0 - quantization = _node_quantization(node, preset) - layer_bytes = _preset_bytes_per_layer(preset).get(quantization) - if not layer_bytes: - layer_bytes = next(iter(_preset_bytes_per_layer(preset).values())) - return layers * int(layer_bytes) - - -def _endpoint_memory_pool( - server: "_TrackerHTTPServer", - nodes: list[_NodeEntry], -) -> dict: - """Per-host RAM/VRAM budget, usage, spare capacity, and loaded assignments.""" - if not nodes: - return {} - host = nodes[0] - budget_bytes, budget_source = _node_memory_budget_bytes(host) - reserve_bytes = int(budget_bytes * 0.8) - loaded: list[dict] = [] - used_bytes = 0 - for node in nodes: - preset = _preset_for_node(server, node) - assignment_bytes = _assignment_memory_bytes(node, preset) - used_bytes += assignment_bytes - loaded.append({ - "node_id": node.node_id, - "model": node.model, - "hf_repo": node.hf_repo, - "shard_start": node.shard_start, - "shard_end": node.shard_end, - "memory_bytes": assignment_bytes, - "status": node.status, - }) - max_slots = max(node.max_loaded_shards for node in nodes) - loaded_slots = len([node for node in nodes if node.shard_start is not None]) - spare_bytes = max(0, reserve_bytes - used_bytes) - return { - "endpoint": host.endpoint, - "memory_budget_bytes": budget_bytes, - "memory_budget_source": budget_source, - "memory_reserve_bytes": reserve_bytes, - "memory_used_bytes": used_bytes, - "memory_spare_bytes": spare_bytes, - "loaded_slots": loaded_slots, - "max_loaded_shards": max_slots, - "spare_slots": max(0, max_slots - loaded_slots), - "loaded": loaded, - } - - -def _memory_pool_map(server: "_TrackerHTTPServer") -> dict: - """Aggregate and per-endpoint view of assignable RAM across the hive.""" - from collections import defaultdict - - groups: dict[str, list[_NodeEntry]] = defaultdict(list) - for node in server.registry.values(): - groups[node.endpoint.rstrip("/")].append(node) - - hosts = [_endpoint_memory_pool(server, group) for group in groups.values()] - hosts.sort(key=lambda item: (-item.get("memory_spare_bytes", 0), item.get("endpoint", ""))) - total_budget = sum(item["memory_budget_bytes"] for item in hosts) - total_used = sum(item["memory_used_bytes"] for item in hosts) - total_spare = sum(item["memory_spare_bytes"] for item in hosts) - total_slots = sum(item["max_loaded_shards"] for item in hosts) - loaded_slots = sum(item["loaded_slots"] for item in hosts) - return { - "total_memory_budget_bytes": total_budget, - "total_memory_used_bytes": total_used, - "total_memory_spare_bytes": total_spare, - "total_loaded_slots": loaded_slots, - "total_max_loaded_shards": total_slots, - "total_spare_slots": max(0, total_slots - loaded_slots), - "hosts": hosts, - } - - -def _add_shard_directive( - node: _NodeEntry, - model: str, - start: int, - end: int, - quantization: str, - *, - model_sources: list[dict] | None = None, -) -> dict: - directive = { - "action": "ADD_SHARD", - "model": model, - "start_layer": start, - "end_layer": end, - "shard_start": start, - "shard_end": end, - "quantization": quantization, - } - if model_sources: - directive["model_sources"] = model_sources - return directive - - -def _model_demand_and_supply( - server: "_TrackerHTTPServer", - model_key: str, - preset: dict, -) -> dict: - resolved_name, _ = _resolve_model_preset(server.model_presets, model_key) - lookup = preset.get("hf_repo") or resolved_name or model_key - required_start, required_end = _preset_layer_bounds(preset) - model_nodes = [ - node for node in server.registry.values() - if _node_matches_preset(node, resolved_name or model_key, preset) # type: ignore[arg-type] - ] - health = _model_health_summary(server, model_key, preset.get("hf_repo")) - local_rpms = server.stats.get_local_rpms() if server.stats is not None else {} - demand_rpm = _repo_demand_rpm(server, str(lookup), local_rpms) - return { - "model": resolved_name or model_key, - "hf_repo": preset.get("hf_repo"), - "preset": preset, - "demand_rpm": demand_rpm, - "served_model_copies": float(health.get("served_model_copies") or 0.0), - "coverage_percentage": float(health.get("coverage_percentage") or 0.0), - "required_start": required_start, - "required_end": required_end, - "model_nodes": model_nodes, - } - - -def _scale_demanded_models_locked(server: "_TrackerHTTPServer") -> None: - """Use spare host RAM/slots to add copies of high-demand models.""" - pool = _memory_pool_map(server) - candidates: list[tuple[float, str, dict]] = [] - for model_key, preset in server.model_presets.items(): - if not isinstance(preset, dict) or not preset.get("hf_repo"): - continue - info = _model_demand_and_supply(server, model_key, preset) - demand = info["demand_rpm"] - copies = info["served_model_copies"] - coverage = info["coverage_percentage"] - if coverage < 100.0: - continue - target_copies = max(1.0, demand / 5.0) - if copies >= target_copies: - continue - candidates.append((demand, model_key, info)) - - if not candidates: - return - - candidates.sort(key=lambda item: (-item[0], item[2]["served_model_copies"])) - for _demand, model_key, info in candidates: - preset = info["preset"] - hf_repo = str(info["hf_repo"]) - required_start = info["required_start"] - required_end = info["required_end"] - total_layers = required_end - required_start + 1 - layer_bytes = next(iter(_preset_bytes_per_layer(preset).values())) - full_copy_bytes = total_layers * layer_bytes - - for host in pool["hosts"]: - if host["spare_slots"] <= 0 or host["memory_spare_bytes"] < full_copy_bytes: - continue - host_nodes = [ - server.registry[node["node_id"]] - for node in host["loaded"] - if node["node_id"] in server.registry - ] - if not host_nodes: - continue - if any((n.hf_repo or n.model) == hf_repo for n in host_nodes): - continue - anchor = max(host_nodes, key=lambda n: n.benchmark_tokens_per_sec) - if anchor.status != "ready" or anchor.pending_new_assignment is not None: - continue - capacity = min(_node_layer_capacity(anchor, preset), total_layers) - if capacity <= 0: - continue - quantization = _node_quantization(anchor, preset) - shard_end = min(required_end, required_start + capacity - 1) - assignment = _add_shard_directive( - anchor, - hf_repo, - required_start, - shard_end, - quantization, - ) - anchor.pending_new_assignment = assignment - anchor.pending_directives.append(assignment) - _tracker_log( - server, - "info", - "scale demanded model copy", - endpoint=anchor.endpoint, - node_id=anchor.node_id, - model=model_key, - hf_repo=hf_repo, - shard=f"{required_start}-{shard_end}", - demand_rpm=round(info["demand_rpm"], 4), - served_model_copies=info["served_model_copies"], - target_copies=round(max(1.0, info["demand_rpm"] / 5.0), 2), - host_spare_bytes=host["memory_spare_bytes"], - ) - break - - -def _deployment_summary(nodes: list[_NodeEntry], preset: dict | None) -> dict: - if preset is None: - return {"recommended": False} - pool = _pool_summary(nodes, preset) - required = int(preset.get("required_model_bytes", 0) or 0) - deployable = required > 0 and pool["effective_assignable_memory_bytes"] >= required - missing = max(0, required - pool["effective_assignable_memory_bytes"]) if required > 0 else 0 - return { - "recommended": bool(preset.get("recommended", False)), - "status": preset.get("deployment_status", "available"), - "required_model_bytes": required or None, - "download_size_bytes": preset.get("download_size_bytes"), - "native_quantization": preset.get("native_quantization"), - "pool": pool, - "deployable": deployable, - "missing_effective_memory_bytes": missing, - } - - -def _max_layers_for_memory(memory_mb: int, total_layers: int, preset: dict | None = None) -> int: - if total_layers <= 0: - return 0 - if memory_mb <= 0: - return max(1, total_layers // 2) - memory_bytes = memory_mb * 1024 * 1024 - bytes_per_layer = next(iter(_preset_bytes_per_layer(preset).values())) if preset is not None else 30 * 1024 * 1024 - return min( - total_layers, - max(1, int((memory_bytes * 0.8) // bytes_per_layer)), - ) - - -def _model_metadata_from_nodes(nodes: list[_NodeEntry]) -> dict: - metadata: dict = {} - for node in nodes: - if node.model_metadata: - metadata.update(node.model_metadata) - if "num_layers" not in metadata: - layers = [node.num_layers for node in nodes if node.num_layers is not None] - if layers: - metadata["num_layers"] = max(layers) - return metadata - - -def _coverage_map( - nodes: list[_NodeEntry], - required_start: int, - required_end: int, -) -> list[dict]: - layer_counts = [] - for layer in range(required_start, required_end + 1): - count = 0 - for node in nodes: - if node.shard_start is None or node.shard_end is None: - continue - if node.shard_start <= layer <= node.shard_end: - count += 1 - layer_counts.append((layer, count)) - - coverage: list[dict] = [] - for layer, count in layer_counts: - if coverage and coverage[-1]["node_count"] == count and coverage[-1]["end_layer"] == layer - 1: - coverage[-1]["end_layer"] = layer - else: - coverage.append({"start_layer": layer, "end_layer": layer, "node_count": count}) - return coverage - - -def _node_health(node: "_NodeEntry", heartbeat_timeout: float) -> dict: - """Per-node health detail for the availability map.""" - age = time.monotonic() - node.last_heartbeat - alive = age <= heartbeat_timeout - hb_expected = max(1, round(node.uptime_seconds / 20.0)) # assume ~20s interval - hb_rate = round(min(1.0, node.heartbeats_received / hb_expected), 4) if node.heartbeats_received else 0.0 - total = node.total_requests - inf_rate = round((total - node.failed_requests) / total, 4) if total > 0 else 1.0 - return { - "node_id": node.node_id, - "endpoint": node.endpoint, - "alive": alive, - "last_seen_seconds_ago": round(age, 1), - "status": node.status, - "queue_depth": _effective_queue_depth(node), - "heartbeat_queue_depth": node.queue_depth, - "proxy_inflight": node.proxy_inflight, - "current_requests": list(node.current_requests), - "total_requests": node.total_requests, - "heartbeat_success_rate": hb_rate, - "inference_success_rate": inf_rate, - "capacity": _node_capacity_summary(node), - } - - -def _coverage_map_detailed( - nodes: list["_NodeEntry"], - required_start: int, - required_end: int, - heartbeat_timeout: float, -) -> list[dict]: - """Like _coverage_map but with per-node identity and health in each band. - - Includes all nodes (alive and stale) so operators can see both coverage - holes and which dead nodes used to fill them. - """ - now = time.monotonic() - - def covers(node: "_NodeEntry", layer: int) -> bool: - return ( - node.shard_start is not None - and node.shard_end is not None - and node.shard_start <= layer <= node.shard_end - ) - - # Build per-layer list of nodes (with alive flag) - layer_nodes: list[list[tuple["_NodeEntry", bool]]] = [] - for layer in range(required_start, required_end + 1): - ns = [ - (n, (now - n.last_heartbeat) <= heartbeat_timeout) - for n in nodes - if covers(n, layer) - ] - layer_nodes.append(ns) - - # Merge consecutive layers with identical node-set (same node_ids, same alive) - coverage: list[dict] = [] - for i, layer_idx in enumerate(range(required_start, required_end + 1)): - ns = layer_nodes[i] - node_ids = [n.node_id for n, _ in ns] - alive_flags = [a for _, a in ns] - # Compare with last band - if ( - coverage - and coverage[-1]["end_layer"] == layer_idx - 1 - and [nd["node_id"] for nd in coverage[-1]["nodes"]] == node_ids - and [nd["alive"] for nd in coverage[-1]["nodes"]] == alive_flags - ): - coverage[-1]["end_layer"] = layer_idx - else: - coverage.append({ - "start_layer": layer_idx, - "end_layer": layer_idx, - "node_count": len(ns), - "nodes": [_node_health(n, heartbeat_timeout) for n, _ in ns], - }) - return coverage - - -def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]: - return [ - (segment["start_layer"], segment["end_layer"]) - for segment in coverage - if segment["node_count"] == 0 - ] - - -def _unassigned_managed_nodes(nodes: list["_NodeEntry"]) -> list["_NodeEntry"]: - return [ - node for node in nodes - if node.managed_assignment - and (node.shard_start is None or node.shard_end is None) - ] - - -def _emit_shard_change_directives( - node: "_NodeEntry", - model: str, - previous_range: tuple[int | None, int | None, str | None], - preset: dict, -) -> None: - """Queue DROP/LOAD directives when a managed node's shard assignment changes.""" - previous_start, previous_end, previous_quantization = previous_range - current_range = (node.shard_start, node.shard_end, node.quantization) - if node.shard_start is None or node.shard_end is None or current_range == previous_range: - return - if previous_start is not None and previous_end is not None: - node.pending_directives.append( - _drop_directive( - node, - model, - previous_start, - previous_end, - previous_quantization or _node_quantization(node, preset), - ) - ) - node.pending_directives.append( - _load_directive( - node, - model, - node.shard_start, - node.shard_end, - node.quantization or _node_quantization(node, preset), - ) - ) - - -def _assign_redundant_managed_nodes( - managed_nodes: list["_NodeEntry"], - model: str, - preset: dict, - required_start: int, - required_end: int, -) -> None: - """Give newly joined managed nodes their own copy without reshuffling incumbents.""" - total_layers = required_end - required_start + 1 - for node in sorted( - managed_nodes, - key=lambda entry: ( - -entry.benchmark_tokens_per_sec, - -_node_layer_capacity(entry, preset), - entry.node_id, - ), - ): - if node.shard_start is not None and node.shard_end is not None: - continue - capacity = min(_node_layer_capacity(node, preset), total_layers) - if capacity <= 0: - continue - previous_range = (node.shard_start, node.shard_end, node.quantization) - quantization = _node_quantization(node, preset) - node.quantization = quantization - node.shard_start = required_start - node.shard_end = min(required_end, required_start + capacity - 1) - _emit_shard_change_directives(node, model, previous_range, preset) - - -def _relay_http_request_frames( - relay_addr: str, - path: str, - body: bytes, - headers: dict[str, str], - timeout: float = 310.0, - idle_timeout: float = 120.0, - *, - cancel_event: threading.Event | None = None, - ws_holder: list[Any] | None = None, - ws_lock: threading.Lock | None = None, -): - """Send an HTTP-shaped request through a relay RPC WebSocket, yielding - response frames until a terminal one (US-036). - - A frame with ``stream: true`` is part of a chunked SSE response ending with - ``done: true``; a frame without ``stream`` is a complete single response. - Yields nothing when the relay is unreachable; stops silently on idle or - overall timeout (the caller bills whatever was observed). - """ - try: - import websockets.sync.client as wsc # type: ignore[import] - except Exception: - return - request_id = str(uuid.uuid4()) - deadline = time.monotonic() + timeout - try: - with wsc.connect(relay_addr, open_timeout=10, close_timeout=5) as ws: - if ws_holder is not None: - if ws_lock is not None: - with ws_lock: - ws_holder.clear() - ws_holder.append(ws) - else: - ws_holder.clear() - ws_holder.append(ws) - ws.send(json.dumps({ - "request_id": request_id, - "method": "POST", - "path": path, - "headers": headers, - "body": body.decode(errors="replace"), - })) - while True: - if cancel_event is not None and cancel_event.is_set(): - return - remaining = deadline - time.monotonic() - if remaining <= 0: - return - raw = ws.recv(timeout=min(idle_timeout, remaining)) - frame = json.loads(raw) - if frame.get("request_id") not in {None, request_id}: - continue - yield frame - if not frame.get("stream") or frame.get("done"): - return - except Exception: - return - - -def _relay_http_request( - relay_addr: str, - path: str, - body: bytes, - headers: dict[str, str], - timeout: float = 310.0, -) -> dict | None: - """Send an HTTP-shaped request through a relay and buffer the response. - - Streamed frame sequences are collapsed into one response dict whose body is - the concatenated SSE text — used by non-chat callers and kept for - backward compatibility; the chat proxy streams frames directly. - """ - frames = _relay_http_request_frames(relay_addr, path, body, headers, timeout=timeout) - first = next(frames, None) - if first is None: - return None - if not first.get("stream"): - return first - chunks = [first.get("chunk") or ""] - for frame in frames: - chunks.append(frame.get("chunk") or "") - return { - "request_id": first.get("request_id"), - "status": first.get("status", 200), - "headers": first.get("headers") or {}, - "body": "".join(chunks), - } - - -def _usage_split(payload: dict) -> dict | None: - """Parse a usage block into {"prompt", "completion", "total"} (ints or None).""" - usage = payload.get("usage") - if not isinstance(usage, dict): - return None - - def _num(key: str) -> int | None: - value = usage.get(key) - return int(value) if isinstance(value, (int, float)) else None - - return { - "prompt": _num("prompt_tokens"), - "completion": _num("completion_tokens"), - "total": _usage_total_tokens(payload), - } - - -def _stream_line_tokens(line: bytes) -> tuple[int, dict | None]: - """Token accounting for one SSE line: (observed output delta, usage split or None).""" - if not line.startswith(b"data:"): - return 0, None - payload = line[5:].strip() - if not payload or payload == b"[DONE]": - return 0, None - try: - chunk_payload = json.loads(payload) - except json.JSONDecodeError: - return 1, None - return _observed_stream_tokens(chunk_payload), _usage_split(chunk_payload) - - -def _stream_billable_split( - observed_output: int, usage: dict | None, request_body: dict -) -> tuple[int, int]: - """(input_tokens, output_tokens) for a streamed response (US-045). - - Output: observed deltas, capped by reported completion count when the - stream carried a usage chunk. Input: reported prompt count, else the - prompt estimate from the request body. - """ - prompt = (usage or {}).get("prompt") - completion = (usage or {}).get("completion") - total = (usage or {}).get("total") - if prompt is None: - prompt = _estimate_prompt_tokens(request_body) or 0 - if completion is None and total is not None: - completion = max(0, total - prompt) - return max(0, prompt), _billable_stream_tokens(observed_output, completion) - - -def _billable_non_stream_split(payload: dict, request_body: dict) -> tuple[int, int]: - """(input_tokens, output_tokens) for a buffered response (US-045). - - Prefers the response usage block; falls back to content estimates. - Completion stays capped by the request's max-tokens bound, as before. - """ - usage = _usage_split(payload) - prompt_estimate = _estimate_prompt_tokens(request_body) or 0 - prompt = (usage or {}).get("prompt") - completion = (usage or {}).get("completion") - if prompt is None: - prompt = prompt_estimate - if completion is None: - total = (usage or {}).get("total") - if total is not None: - completion = max(0, total - prompt) - else: - completion = _observed_non_stream_completion_tokens(payload) - limit = _requested_completion_token_limit(request_body) - if limit is not None and completion > limit: - completion = min(completion, limit) - prompt = max(prompt, prompt_estimate) - return max(0, prompt), max(0, completion) - - -def _find_pinned_route( - nodes: list[_NodeEntry], - required_start: int, - required_end: int, - hop_count: int, -) -> list[_NodeEntry] | None: - """First combination of exactly ``hop_count`` distinct nodes covering the - layer range, where every node extends coverage (US-030 benchmark routes).""" - for combo in itertools.permutations(nodes, hop_count): - covered = required_start - 1 - valid = True - for candidate in combo: - if candidate.shard_start is None or candidate.shard_end is None: - valid = False - break - if candidate.shard_start > covered + 1 or candidate.shard_end <= covered: - valid = False - break - covered = candidate.shard_end - if valid and covered >= required_end: - return list(combo) - return None - - -def _nodes_and_bounds_for_model( - server: "_TrackerHTTPServer", - model: str, -) -> tuple[list[_NodeEntry], int, int] | None: - resolved_name, preset = _resolve_model_preset(server.model_presets, model) - if preset is not None: - required_start, required_end = _preset_layer_bounds(preset) - return [ - node for node in server.registry.values() - if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] - ], required_start, required_end - - nodes = [ - node for node in server.registry.values() - if _node_matches_model(node, model) - and node.shard_start is not None - and node.shard_end is not None - and node.num_layers is not None - ] - if not nodes: - return None - return nodes, 0, max(node.num_layers for node in nodes) - 1 - - -def _fetch_toploc_commitment( - node: _NodeEntry, - *, - session_id: str, - model: str, - messages: list[dict], -) -> dict | None: - """Fetch a node's own on-demand TOPLOC boundary commitment (ADR-0018 §3), - same protocol as `ValidatorProcess._fetch_hop_commitment`.""" - endpoint = node.endpoint - if not isinstance(endpoint, str) or not endpoint: - return None - try: - req = urllib.request.Request( - f"{endpoint.rstrip('/')}/v1/audit/toploc/commitment", - data=json.dumps({ - "session_id": session_id, - "model": model, - "messages": messages, - "shard_start": node.shard_start, - "shard_end": node.shard_end, - }).encode(), - headers={"Content-Type": "application/json"}, - method="POST", - ) - with urllib.request.urlopen(req, timeout=5.0) as resp: - response = json.loads(resp.read()) - except (OSError, ValueError, json.JSONDecodeError): - return None - proof = response.get("toploc_proof") or response.get("activation_proof") - token_ids = response.get("claimed_token_ids") or response.get("output_token_ids") - if not isinstance(proof, dict): - return None - if not isinstance(token_ids, list) or not all(isinstance(t, int) for t in token_ids): - return None - return {"toploc_proof": proof, "claimed_token_ids": token_ids} - - -def _fetch_toploc_reference_activations( - reference_node_url: str, - *, - model: str, - messages: list[dict], - claimed_token_ids: list[int], - claim: Any, -) -> list | None: - """Teacher-force the claimed tokens on the reference node (same contract - as `ValidatorProcess._run_teacher_forced_prefill` / validator README's - "TOPLOC audit contract").""" - try: - req = urllib.request.Request( - f"{reference_node_url.rstrip('/')}/v1/audit/toploc", - data=json.dumps({ - "model": model, - "messages": messages, - "claimed_token_ids": claimed_token_ids, - "dtype": claim.dtype, - "quantization": claim.quantization, - "decode_batching_size": claim.decode_batching_size, - "topk": claim.topk, - "skip_prefill": claim.skip_prefill, - }).encode(), - headers={"Content-Type": "application/json"}, - method="POST", - ) - with urllib.request.urlopen(req, timeout=300.0) as resp: - response = json.loads(resp.read()) - except (OSError, ValueError, json.JSONDecodeError): - return None - activations = response.get("activations") - if not isinstance(activations, list): - return None - return activations - - -def _load_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict: - return { - "action": "LOAD_SHARD", - "model": model, - "start_layer": start, - "end_layer": end, - "shard_start": start, - "shard_end": end, - "quantization": quantization, - } - - -def _drop_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict: - return { - "action": "DROP_SHARD", - "model": model, - "start_layer": start, - "end_layer": end, - "shard_start": start, - "shard_end": end, - "quantization": quantization, - } - - -def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]: - now = time.monotonic() - expired_ids = [ - node_id for node_id, entry in server.registry.items() - if (now - entry.last_heartbeat) > server.heartbeat_timeout - ] - expired_entries: list[tuple[str, _NodeEntry]] = [] - for node_id in expired_ids: - entry = server.registry.pop(node_id) - expired_entries.append((node_id, entry)) - if expired_ids: - _rebalance_all_locked(server) - for node_id, entry in expired_entries: - _tracker_log( - server, - "warn", - "node expired", - node_id=node_id, - endpoint=entry.endpoint, - model=entry.model, - hf_repo=entry.hf_repo, - shard=f"{entry.shard_start}-{entry.shard_end}", - heartbeat_timeout_seconds=server.heartbeat_timeout, - model_health=_model_health_summary(server, entry.model, entry.hf_repo), - ) - return expired_ids - - -def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None: - resolved_name, preset = _resolve_model_preset(server.model_presets, model) - if preset is None: - return - required_start, required_end = _preset_layer_bounds(preset) - total_layers = required_end - required_start + 1 - model_nodes = [ - node for node in server.registry.values() - if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] - ] - managed_nodes = [node for node in model_nodes if node.managed_assignment] - if not managed_nodes: - return - - coverage = _coverage_map(model_nodes, required_start, required_end) - gaps = _coverage_gaps(coverage) - unassigned = _unassigned_managed_nodes(managed_nodes) - if not gaps and not unassigned: - return - if not gaps and unassigned: - _assign_redundant_managed_nodes( - unassigned, model, preset, required_start, required_end, - ) - return - - previous_ranges = { - node.node_id: (node.shard_start, node.shard_end, node.quantization) - for node in managed_nodes - } - for node in managed_nodes: - node.shard_start = None - node.shard_end = None - - managed_nodes.sort( - key=lambda node: ( - -node.benchmark_tokens_per_sec, - -_node_layer_capacity(node, preset), - node.node_id, - ) - ) - base_nodes = [node for node in model_nodes if not node.managed_assignment] - base_coverage = _coverage_map(base_nodes, required_start, required_end) - gaps = _coverage_gaps(base_coverage) - - eligible_nodes = [ - node for node in managed_nodes - if _node_layer_capacity(node, preset) > 0 - ] - node_index = 0 - for gap_start, gap_end in gaps: - cursor = gap_start - while cursor <= gap_end and node_index < len(eligible_nodes): - node = eligible_nodes[node_index] - remaining_layers = gap_end - cursor + 1 - remaining_nodes_after = len(eligible_nodes) - node_index - 1 - capacity = min( - _node_layer_capacity(node, preset), - total_layers, - max(1, remaining_layers - remaining_nodes_after), - ) - if capacity <= 0: - node_index += 1 - continue - quantization = _node_quantization(node, preset) - node.quantization = quantization - node.shard_start = cursor - node.shard_end = min(gap_end, cursor + capacity - 1) - cursor = node.shard_end + 1 - node_index += 1 - - for node in managed_nodes: - _emit_shard_change_directives( - node, - model, - previous_ranges[node.node_id], - preset, - ) - - -def _hf_rebalance_preset(nodes: list[_NodeEntry]) -> dict: - total_layers = max(node.num_layers or 0 for node in nodes) - return { - "layers_start": 0, - "layers_end": total_layers - 1, - "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024}, - } - - -def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> None: - model_nodes = [ - node for node in server.registry.values() - if node.hf_repo == hf_repo - and node.num_layers is not None - ] - managed_nodes = [node for node in model_nodes if node.managed_assignment] - if not model_nodes or not managed_nodes: - return - - preset = _hf_rebalance_preset(model_nodes) - required_start, required_end = _preset_layer_bounds(preset) - total_layers = required_end - required_start + 1 - if total_layers <= 0: - return - - coverage = _coverage_map(model_nodes, required_start, required_end) - gaps = _coverage_gaps(coverage) - unassigned = _unassigned_managed_nodes(managed_nodes) - if not gaps and not unassigned: - return - if not gaps and unassigned: - _assign_redundant_managed_nodes( - unassigned, hf_repo, preset, required_start, required_end, - ) - return - - previous_ranges = { - node.node_id: (node.shard_start, node.shard_end, node.quantization) - for node in managed_nodes - } - for node in managed_nodes: - node.shard_start = None - node.shard_end = None - - managed_nodes.sort( - key=lambda node: ( - -node.benchmark_tokens_per_sec, - -_node_layer_capacity(node, preset), - node.node_id, - ) - ) - base_nodes = [node for node in model_nodes if not node.managed_assignment] - base_coverage = _coverage_map(base_nodes, required_start, required_end) - gaps = _coverage_gaps(base_coverage) - - eligible_nodes = [ - node for node in managed_nodes - if _node_layer_capacity(node, preset) > 0 - ] - node_index = 0 - for gap_start, gap_end in gaps: - cursor = gap_start - while cursor <= gap_end and node_index < len(eligible_nodes): - node = eligible_nodes[node_index] - remaining_layers = gap_end - cursor + 1 - remaining_nodes_after = len(eligible_nodes) - node_index - 1 - capacity = min( - _node_layer_capacity(node, preset), - total_layers, - max(1, remaining_layers - remaining_nodes_after), - ) - if capacity <= 0: - node_index += 1 - continue - quantization = _node_quantization(node, preset) - node.quantization = quantization - node.shard_start = cursor - node.shard_end = min(gap_end, cursor + capacity - 1) - cursor = node.shard_end + 1 - node_index += 1 - - for node in managed_nodes: - _emit_shard_change_directives( - node, - hf_repo, - previous_ranges[node.node_id], - preset, - ) - - -def _rebalance_all_locked(server: "_TrackerHTTPServer") -> None: - for model in list(server.model_presets): - _rebalance_model_locked(server, model) - for hf_repo in sorted({node.hf_repo for node in server.registry.values() if node.hf_repo}): - _rebalance_hf_model_locked(server, hf_repo) - - -def _api_key_from_headers(headers) -> str | None: - auth = headers.get("Authorization") - if not auth: - return None - if auth.lower().startswith("bearer "): - return auth.split(" ", 1)[1].strip() or None - return auth.strip() or None - - -def _session_token_from_headers(headers) -> str | None: - token = _api_key_from_headers(headers) - if token: - return token - cookie_header = headers.get("Cookie") - if not cookie_header: - return None - cookie = http.cookies.SimpleCookie() - try: - cookie.load(cookie_header) - except http.cookies.CookieError: - return None - morsel = cookie.get(_SESSION_COOKIE_NAME) - if morsel is None: - return None - return morsel.value.strip() or None - - -def _session_cookie_header(token: str | None) -> str: - cookie = http.cookies.SimpleCookie() - cookie[_SESSION_COOKIE_NAME] = token or "" - morsel = cookie[_SESSION_COOKIE_NAME] - morsel["path"] = "/" - morsel["httponly"] = True - morsel["samesite"] = "Lax" - if token: - morsel["max-age"] = str(int(7 * 86400)) - else: - morsel["max-age"] = "0" - return morsel.OutputString() - - -def _usage_total_tokens(payload: dict) -> int | None: - usage = payload.get("usage") - if not isinstance(usage, dict): - return None - total = usage.get("total_tokens") - if isinstance(total, (int, float)): - return int(total) - prompt = usage.get("prompt_tokens") - completion = usage.get("completion_tokens") - if isinstance(prompt, (int, float)) or isinstance(completion, (int, float)): - return int(prompt or 0) + int(completion or 0) - return None - - -def _estimate_text_tokens(value: Any) -> int | None: - if isinstance(value, str): - text = value.strip() - if not text: - return 0 - return len(text.split()) - if isinstance(value, list): - total = 0 - found = False - for item in value: - if isinstance(item, str): - estimated = _estimate_text_tokens(item) - elif isinstance(item, dict): - estimated = _estimate_text_tokens(item.get("text")) - else: - estimated = None - if estimated is not None: - total += estimated - found = True - return total if found else None - return None - - -def _estimate_prompt_tokens(body: dict) -> int | None: - messages = body.get("messages") - if isinstance(messages, list): - total = 0 - found = False - for message in messages: - if not isinstance(message, dict): - continue - estimated = _estimate_text_tokens(message.get("content")) - if estimated is not None: - total += estimated - found = True - return total if found else None - prompt = body.get("prompt") - return _estimate_text_tokens(prompt) - - -def _requested_completion_token_limit(body: dict) -> int | None: - for field in ("max_completion_tokens", "max_tokens"): - value = body.get(field) - if isinstance(value, bool): - return None - if isinstance(value, (int, float)): - return max(0, int(value)) - return None - - -def _request_total_token_upper_bound(body: dict) -> int | None: - completion_limit = _requested_completion_token_limit(body) - if completion_limit is None: - return None - prompt_estimate = _estimate_prompt_tokens(body) - return completion_limit + (prompt_estimate or 0) - - -def _billable_non_stream_tokens(payload: dict, request_body: dict) -> int: - reported = _usage_total_tokens(payload) - upper_bound = _request_total_token_upper_bound(request_body) - if reported is None: - completion_estimate = _observed_non_stream_completion_tokens(payload) - if completion_estimate > 0: - billable = completion_estimate + (_estimate_prompt_tokens(request_body) or 0) - elif upper_bound is not None: - billable = upper_bound - else: - return 0 - return min(billable, upper_bound) if upper_bound is not None else billable - billable = max(0, reported) - if upper_bound is not None: - billable = min(billable, upper_bound) - return billable - - -def _observed_non_stream_completion_tokens(payload: dict) -> int: - choices = payload.get("choices") - if not isinstance(choices, list): - return 0 - observed = 0 - for choice in choices: - if not isinstance(choice, dict): - continue - message = choice.get("message") - estimated: int | None = None - if isinstance(message, dict): - estimated = _estimate_text_tokens(message.get("content")) - if estimated is None: - estimated = _estimate_text_tokens(choice.get("text")) - if estimated is not None: - observed += estimated - elif choice: - observed += 1 - return observed - - -def _observed_output_from_non_stream_payload(payload: dict) -> str: - choices = payload.get("choices") - if not isinstance(choices, list) or not choices: - return "" - first = choices[0] - if not isinstance(first, dict): - return "" - message = first.get("message") - if isinstance(message, dict) and isinstance(message.get("content"), str): - return message["content"] - text = first.get("text") - return text if isinstance(text, str) else "" - - -def _observed_stream_tokens(payload: dict) -> int: - choices = payload.get("choices") - if not isinstance(choices, list): - return 0 - observed = 0 - for choice in choices: - if not isinstance(choice, dict): - continue - delta = choice.get("delta") - if isinstance(delta, dict): - estimated = _estimate_text_tokens(delta.get("content")) - if estimated: - observed += estimated - continue - if choice: - observed += 1 - return observed - - -def _billable_stream_tokens(observed_tokens: int, reported_tokens: int | None) -> int: - observed = max(0, observed_tokens) - if reported_tokens is None: - return observed - return min(observed, max(0, reported_tokens)) - - -def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None: - if contracts is None or not wallet_address: - return None - if contracts.registry.get_wallet(wallet_address).banned: - return "wallet is banned" - return None - - -def _tracker_log( - server: "_TrackerHTTPServer", - level: str, - message: str, - *, - stdout: bool = True, - update_console_key: str | None = None, - **fields: Any, -) -> None: - log_level = { - "debug": 10, - "info": 20, - "warn": 30, - "warning": 30, - "error": 40, - }.get(level.lower(), 20) - event = { - "ts": time.time(), - "level": level, - "message": message, - "fields": { - key: value - for key, value in fields.items() - if value is not None - }, - } - with server.console_lock: - if update_console_key is not None: - updated = False - for existing in reversed(server.console_events): - if ( - existing.get("message") == message - and existing.get("fields", {}).get("request_id") == update_console_key - ): - existing["ts"] = event["ts"] - existing["fields"] = event["fields"] - updated = True - break - if not updated: - server.console_events.append(event) - else: - server.console_events.append(event) - extras = " ".join(f"{key}={value}" for key, value in event["fields"].items()) - suffix = f" {extras}" if extras else "" - tracker_logger().log(log_level, f"{message}{suffix}") - if stdout: - print(f"[tracker] {level}: {message}{suffix}", flush=True) - - -@dataclass -class _ActiveProxyContext: - request_id: str - cancel_event: threading.Event = field(default_factory=threading.Event) - upstream: Any | None = None - upstream_lock: threading.Lock = field(default_factory=threading.Lock) - relay_ws: Any | None = None - relay_ws_lock: threading.Lock = field(default_factory=threading.Lock) - - -def _register_active_proxy(server: "_TrackerHTTPServer", request_id: str) -> _ActiveProxyContext: - ctx = _ActiveProxyContext(request_id=request_id) - with server.active_proxies_lock: - server.active_proxies[request_id] = ctx - return ctx - - -def _unregister_active_proxy(server: "_TrackerHTTPServer", request_id: str) -> None: - with server.active_proxies_lock: - server.active_proxies.pop(request_id, None) - - -def _request_proxy_cancel(server: "_TrackerHTTPServer", request_id: str) -> bool: - with server.active_proxies_lock: - ctx = server.active_proxies.get(request_id) - if ctx is None: - return False - ctx.cancel_event.set() - - def _close_resources() -> None: - with ctx.upstream_lock: - upstream = ctx.upstream - if upstream is not None: - try: - upstream.close() - except Exception: - pass - with ctx.relay_ws_lock: - relay_ws = ctx.relay_ws - if relay_ws is not None: - try: - relay_ws.close() - except Exception: - pass - - threading.Thread(target=_close_resources, daemon=True).start() - return True - - -def _upstream_socket(upstream: Any) -> Any | None: - fp = getattr(upstream, "fp", None) - raw = getattr(fp, "raw", None) if fp is not None else None - return getattr(raw, "_sock", None) if raw is not None else None - - -def _set_upstream_read_timeout(upstream: Any, timeout: float | None) -> None: - sock = _upstream_socket(upstream) - if sock is not None: - sock.settimeout(timeout) - - -def _clear_proxy_progress_log_state(server: "_TrackerHTTPServer", request_id: str) -> None: - state = getattr(server, "_proxy_progress_log_state", None) - if state is not None: - state.pop(request_id, None) - - -def _tracker_log_proxy_progress( - server: "_TrackerHTTPServer", - *, - request_id: str, - model: str, - route_model: str, - tokens: int, - started: float, - route_nodes: list["_NodeEntry"], - stream: bool = True, - relay: bool = False, -) -> None: - elapsed = time.monotonic() - started - effective_elapsed = max(elapsed, 1e-6) - now = time.monotonic() - state = getattr(server, "_proxy_progress_log_state", None) - if state is None: - state = {} - server._proxy_progress_log_state = state - last_stdout = state.get(request_id) - stdout = last_stdout is None or (now - last_stdout) >= _PROXY_PROGRESS_LOG_INTERVAL - if stdout: - state[request_id] = now - _tracker_log( - server, - "info", - "proxy progress", - stdout=stdout, - update_console_key=request_id, - request_id=request_id, - model=model, - route_model=route_model, - stream=stream, - relay=relay or None, - tokens=tokens, - elapsed_seconds=round(elapsed, 4), - tokens_per_sec=round(tokens / effective_elapsed, 4) if tokens > 0 else 0.0, - route=_node_route_summary(route_nodes), - ) - - -def _node_id_for_registration( - endpoint: str, - model: str, - wallet_address: str | None, - shard_start: int | None, - shard_end: int | None, - hf_repo: str | None, -) -> str: - wallet_prefix = wallet_address[:8] if wallet_address else "anon" - stable_key = "|".join([ - wallet_address or "", - endpoint.rstrip("/"), - model, - hf_repo or "", - "" if shard_start is None else str(shard_start), - "" if shard_end is None else str(shard_end), - ]) - digest = hashlib.sha256(stable_key.encode()).hexdigest()[:12] - return f"{wallet_prefix}-{digest}" - - -class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer): - daemon_threads = True - - def __init__( - self, - addr: tuple, - handler, - registry: dict, - lock: threading.Lock, - heartbeat_timeout: float, - model_presets: dict, - contracts: Any | None, - minimum_stake: int, - relay_url: str | None = None, - raft: "RaftNode | None" = None, - gossip: "NodeGossip | None" = None, - stats: "_StatsCollector | None" = None, - billing: "BillingLedger | None" = None, - accounts: "AccountStore | None" = None, - benchmark_results_path: str | None = None, - validator_service_token: str | None = None, - hive_secret: str | None = None, - max_charge_per_request: float | None = None, - starting_credit: float = DEFAULT_CALLER_CREDIT_USDT, - devnet_topup_amount: float = DEFAULT_DEVNET_TOPUP_USDT, - toploc_calibration: "ToplocCalibrationStore | None" = None, - toploc_reference_node_url: str | None = None, - toploc_calibration_gate_min_hardware_profiles: int = 1, - toploc_backend: Any | None = None, - hf_pricing_log: "HfPricingLog | None" = None, - models_dir: Path | None = None, - ) -> None: - super().__init__(addr, handler) - self.registry = registry - self.lock = lock - self.heartbeat_timeout = heartbeat_timeout - self.model_presets = model_presets - self.contracts = contracts - self.minimum_stake = minimum_stake - self.relay_url = relay_url.rstrip("/") if relay_url else None - self.raft = raft - self.gossip = gossip - self.stats: _StatsCollector | None = stats - self.billing: BillingLedger | None = billing - self.accounts: AccountStore | None = accounts - self.benchmark_results_path = benchmark_results_path or os.path.join( - os.getcwd(), "benchmark_results.json" - ) - self.benchmark_lock = threading.Lock() - self.validator_service_token = validator_service_token - self.hive_secret = hive_secret - self.max_charge_per_request = max_charge_per_request - self.starting_credit = starting_credit - self.devnet_topup_amount = devnet_topup_amount - self.toploc_calibration: ToplocCalibrationStore | None = toploc_calibration - self.toploc_reference_node_url = ( - toploc_reference_node_url.rstrip("/") if toploc_reference_node_url else None - ) - self.toploc_calibration_gate_min_hardware_profiles = toploc_calibration_gate_min_hardware_profiles - self.toploc_backend = toploc_backend - self.hf_pricing_log: HfPricingLog | None = hf_pricing_log - self.models_dir = models_dir - self.console_events = deque(maxlen=_CONSOLE_LIMIT) - self.console_lock = threading.Lock() - self.active_proxies: dict[str, _ActiveProxyContext] = {} - self.active_proxies_lock = threading.Lock() - - -class _TrackerHandler(http.server.BaseHTTPRequestHandler): - def log_message(self, fmt, *args): # noqa: suppress request logs in tests - pass - - def _send_json(self, status: int, data: dict, headers: dict[str, str] | None = None) -> None: - body = json.dumps(data).encode() - self.send_response(status) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - for name, value in (headers or {}).items(): - self.send_header(name, value) - self.end_headers() - try: - self.wfile.write(body) - except BrokenPipeError: - pass - - # ---- unified auth boundary (ADR-0017) ---- - - def _resolve_identity(self) -> tuple[str | None, dict | None]: - """Resolve the caller to (role, account). - - Roles: "validator" (service token), "admin"/"user" (session token). - Client API keys resolve to no privileged role — they authorize - inference and wallet binding only, never operator endpoints. - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - token = _session_token_from_headers(self.headers) - if not token: - return None, None - if is_validator_token(token, server.validator_service_token): - return "validator", None - if server.accounts is not None: - account = server.accounts.session_account(token) - if account is not None: - return account.get("role", "user"), account - return None, None - - def _require_role(self, *allowed: str) -> bool: - """Gate a privileged handler; sends 401/403 and returns False on failure. - - 401 when no credential was presented; 403 when a credential was - presented but does not resolve to an allowed role — this covers - client API keys and garbage bearer strings on operator endpoints. - """ - role, _account = self._resolve_identity() - if role in allowed: - return True - if _api_key_from_headers(self.headers) is None: - self._send_json(401, {"error": "authentication required (admin session or service token)"}) - else: - self._send_json(403, {"error": "this endpoint requires an admin session or service token"}) - return False - - def _read_hive_authenticated_body(self) -> dict | None: - """Read + verify a hive gossip body (HMAC per ADR-0017 §3). - - Fails closed: without a configured --hive-secret no gossip is - accepted. Sends the error response itself and returns None on failure. - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - length = int(self.headers.get("Content-Length", 0)) - raw = self.rfile.read(length) if length else b"{}" - if not verify_hive_request(server.hive_secret, self.headers, raw): - self._send_json(401, {"error": "valid hive signature required"}) - return None - try: - body = json.loads(raw or b"{}") - except json.JSONDecodeError: - self._send_json(400, {"error": "invalid JSON body"}) - return None - if not isinstance(body, dict): - self._send_json(400, {"error": "JSON body must be an object"}) - return None - return body - - def _read_json_body(self) -> dict | None: - length = int(self.headers.get("Content-Length", 0)) - try: - body = json.loads(self.rfile.read(length) or b"{}") - except json.JSONDecodeError: - self._send_json(400, {"error": "invalid JSON body"}) - return None - if not isinstance(body, dict): - self._send_json(400, {"error": "JSON body must be an object"}) - return None - return body - - def _purge_expired_nodes(self) -> None: - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - _purge_expired_nodes_locked(server) - - def do_POST(self): - if self.path == "/v1/chat/completions": - self._handle_proxy_chat() - return - if self.path == "/v1/nodes/register": - self._handle_register() - return - if self.path == "/v1/raft/vote": - self._handle_raft_vote() - return - if self.path == "/v1/raft/append": - self._handle_raft_append() - return - if self.path == "/v1/gossip": - self._handle_gossip() - return - if self.path == "/v1/stats/gossip": - self._handle_stats_gossip() - return - if self.path == "/v1/billing/gossip": - self._handle_billing_gossip() - return - if self.path == "/v1/billing/forfeit": - self._handle_billing_forfeit() - return - if self.path == "/v1/auth/register": - self._handle_auth_register() - return - if self.path == "/v1/auth/login": - self._handle_auth_login() - return - if self.path == "/v1/auth/logout": - self._handle_auth_logout() - return - if self.path == "/v1/account/keys": - self._handle_account_key_create() - return - if self.path == "/v1/account/keys/revoke": - self._handle_account_key_revoke() - return - if self.path == "/v1/account/topup": - self._handle_account_topup() - return - if self.path == "/v1/accounts/gossip": - self._handle_accounts_gossip() - return - if self.path == "/v1/registry/gossip": - self._handle_registry_gossip() - return - if self.path == "/v1/benchmark/hop-penalty": - self._handle_benchmark_hop_penalty() - return - if self.path == "/v1/calibration/toploc/run": - self._handle_toploc_calibration_run() - return - if self.path == "/v1/wallet/register": - self._handle_wallet_register() - return - parts = self.path.split("/") - # /v1/nodes//heartbeat -> ['', 'v1', 'nodes', '', 'heartbeat'] - if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat": - self._handle_heartbeat(parts[3]) - return - # /v1/proxy/requests//cancel - if ( - len(parts) == 6 - and parts[1] == "v1" - and parts[2] == "proxy" - and parts[3] == "requests" - and parts[5] == "cancel" - and parts[4] - ): - self._handle_proxy_request_cancel(urllib.parse.unquote(parts[4])) - return - self.send_response(404) - self.end_headers() - - def do_GET(self): - parsed = urllib.parse.urlparse(self.path) - if parsed.path == "/v1/route": - self._handle_route(parsed) - elif parsed.path == "/v1/routes": - self._handle_routes(parsed) - elif parsed.path == "/v1/nodes/assign": - self._handle_assign(parsed) - elif parsed.path == "/v1/network/assign": - self._handle_network_assign(parsed) - elif parsed.path == "/v1/network/map": - self._handle_network_map() - elif parsed.path == "/v1/models": - self._handle_models() - elif parsed.path == "/v1/model-files/download": - self._handle_model_files_download(parsed) - elif parsed.path.startswith("/v1/coverage/"): - model = urllib.parse.unquote(parsed.path.removeprefix("/v1/coverage/")) - self._handle_coverage(model) - elif parsed.path.startswith("/v1/tracker-nodes/"): - model = urllib.parse.unquote(parsed.path.removeprefix("/v1/tracker-nodes/")) - self._handle_tracker_nodes(model) - elif parsed.path.startswith("/v1/head-workers/"): - model = urllib.parse.unquote(parsed.path.removeprefix("/v1/head-workers/")) - self._handle_tracker_nodes(model) - elif parsed.path == "/v1/raft/status": - self._handle_raft_status() - elif parsed.path == "/v1/stats": - self._handle_stats() - elif parsed.path == "/v1/console": - self._handle_console() - elif parsed.path == "/v1/billing/summary": - self._handle_billing_summary() - elif parsed.path == "/v1/billing/settlements": - self._handle_billing_settlements() - elif parsed.path == "/v1/account": - self._handle_account_me() - elif parsed.path == "/v1/admin/accounts": - self._handle_admin_accounts() - elif parsed.path == "/v1/benchmark/results": - self._handle_benchmark_results() - elif parsed.path == "/v1/calibration/toploc/results": - self._handle_toploc_calibration_results() - elif parsed.path == "/v1/pricing/hf/history": - self._handle_hf_pricing_history(parsed) - elif parsed.path == "/v1/registry/wallets": - self._handle_registry_wallets() - elif parsed.path in ("/dashboard", "/dashboard/"): - self._handle_dashboard() - elif parsed.path == "/v1/health": - self._send_json(200, {"status": "ok"}) - else: - self.send_response(404) - self.end_headers() - - def _handle_models(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - created = int(time.time()) - with server.lock: - self._purge_expired_nodes() - alive = list(server.registry.values()) - if server.contracts is not None: - alive = [ - node for node in alive - if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned - ] - data = [] - seen_ids: set[str] = set() - for name, preset in server.model_presets.items(): - model_nodes = [node for node in alive if _node_matches_preset(node, name, preset)] - if not model_nodes and not preset.get("recommended"): - continue - required_start, required_end = _preset_layer_bounds(preset) - coverage = _coverage_percentage( - model_nodes, - required_start, - required_end, - ) - aliases = [name] - hf_repo = preset.get("hf_repo") - if hf_repo and hf_repo not in aliases: - aliases.append(hf_repo) - for alias in preset.get("aliases", []) or []: - if isinstance(alias, str) and alias not in aliases: - aliases.append(alias) - data.append({ - "id": name, - "object": "model", - "created": created, - "owned_by": "meshnet", - "name": name, - "hf_repo": hf_repo, - "aliases": aliases, - "metadata": dict(preset.get("metadata") or _model_metadata_from_nodes(model_nodes)), - "recommended": bool(preset.get("recommended", False)), - "deployment": _deployment_summary(alive, preset), - "shard_coverage_percentage": coverage, - }) - seen_ids.add(name) - if hf_repo: - seen_ids.add(hf_repo) - - hf_model_ids = sorted({ - node.hf_repo or node.model - for node in alive - if node.model is not None - and node.model not in server.model_presets - and node.shard_start is not None - and node.shard_end is not None - and node.num_layers is not None - }) - for model_id in hf_model_ids: - if model_id is None or model_id in seen_ids: - continue - model_nodes = [ - node for node in alive - if node.shard_start is not None - and node.shard_end is not None - and node.num_layers is not None - and (node.hf_repo == model_id or (node.hf_repo is None and node.model == model_id)) - ] - if not model_nodes: - continue - short_names = sorted({node.model for node in model_nodes if node.model}) - aliases = [model_id, *[name for name in short_names if name != model_id]] - required_start = 0 - required_end = max(node.num_layers for node in model_nodes) - 1 - data.append({ - "id": model_id, - "object": "model", - "created": created, - "owned_by": "meshnet", - "name": short_names[0] if short_names else model_id, - "hf_repo": model_id if any(node.hf_repo == model_id for node in model_nodes) else None, - "aliases": aliases, - "metadata": _model_metadata_from_nodes(model_nodes), - "shard_coverage_percentage": _coverage_percentage( - model_nodes, - required_start, - required_end, - ), - }) - seen_ids.add(model_id) - self._send_json(200, {"object": "list", "data": data}) - - def _handle_coverage(self, model: str): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - # Do NOT purge before coverage — dead nodes are included with alive=false - # so operators can see what was covering each layer band before failure. - with server.lock: - resolved = _nodes_and_bounds_for_model(server, model) - if resolved is None: - self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) - return - all_nodes, required_start, required_end = resolved - if server.contracts is not None: - all_nodes = [ - node for node in all_nodes - if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned - ] - coverage = _coverage_map_detailed(all_nodes, required_start, required_end, server.heartbeat_timeout) - self._send_json(200, {"model": model, "coverage": coverage}) - - def _handle_tracker_nodes(self, model: str): - """Return head workers: worker nodes that can start inference for a model. - - The historical endpoint name is /v1/tracker-nodes, but these are not - tracker processes and they are the only machines that load model shards. - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - resolved_name, preset = _resolve_model_preset(server.model_presets, model) - if preset is None: - self._send_json(404, {"error": f"unknown model preset: {model!r}"}) - return - required_start, _ = _preset_layer_bounds(preset) - with server.lock: - self._purge_expired_nodes() - alive = [ - node for node in server.registry.values() - if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] - ] - if server.contracts is not None: - alive = [ - node for node in alive - if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned - ] - tracker_nodes = [ - node for node in alive - if node.shard_start is not None - and node.shard_start == required_start - and node.tracker_mode - ] - self._send_json(200, { - "model": resolved_name, - "head_workers": [ - { - "node_id": node.node_id, - "endpoint": node.endpoint, - "relay_addr": node.relay_addr, - "peer_id": node.peer_id, - } - for node in tracker_nodes - ], - "tracker_nodes": [ - { - "node_id": node.node_id, - "endpoint": node.endpoint, - "relay_addr": node.relay_addr, - "peer_id": node.peer_id, - "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec, - } - for node in tracker_nodes - ], - }) - - def _handle_network_map(self) -> None: - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - with server.lock: - self._purge_expired_nodes() - nodes = list(server.registry.values()) - memory_pool = _memory_pool_map(server) - - def capacity_for(node: _NodeEntry) -> dict: - preset = None - if node.model: - preset = server.model_presets.get(node.model) - if preset is None and node.hf_repo and node.num_layers: - preset = _hf_rebalance_preset([node]) - return _node_capacity_summary(node, preset) - - def throughput_for(node: _NodeEntry) -> dict: - if server.stats is None: - return {} - models = [m for m in (node.hf_repo, node.model) if m] - result = {} - for model in models: - result[model] = server.stats.get_node_model_stats(node.node_id, model) - return result - - def model_supply_for(node: _NodeEntry) -> dict: - return _model_health_summary(server, node.model, node.hf_repo) - - self._send_json(200, { - "relay_url": server.relay_url, - "pool": _pool_summary(nodes), - "memory_pool": memory_pool, - "recommended_models": [ - { - "id": name, - "hf_repo": preset.get("hf_repo"), - "aliases": list(preset.get("aliases", []) or []), - "metadata": dict(preset.get("metadata") or {}), - "deployment": _deployment_summary(nodes, preset), - } - for name, preset in server.model_presets.items() - if preset.get("recommended") - ], - "nodes": [ - { - "node_id": node.node_id, - "endpoint": node.endpoint, - "relay_addr": node.relay_addr, - "peer_id": node.peer_id, - "model": node.model, - "hf_repo": node.hf_repo, - "num_layers": node.num_layers, - "model_metadata": dict(node.model_metadata), - "downloaded_models": [dict(item) for item in node.downloaded_models], - "shard_start": node.shard_start, - "shard_end": node.shard_end, - "tracker_mode": node.tracker_mode, - "last_heartbeat": node.last_heartbeat, - "capacity": capacity_for(node), - "model_supply": model_supply_for(node), - "throughput": throughput_for(node), - "stats": _node_health(node, server.heartbeat_timeout), - } - for node in nodes - ], - }) - - # ---------------------------------------------------------------- OpenAI proxy - - def _handle_proxy_chat(self) -> None: - """Proxy POST /v1/chat/completions to a tracker-mode (first-shard) node. - - Picks a live tracker-mode node for the requested model using round-robin, - then forwards the request verbatim and relays the response (including - streaming SSE chunks) back to the caller. - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - - length = int(self.headers.get("Content-Length", 0)) - raw_body = self.rfile.read(length) if length else b"{}" - - try: - body = json.loads(raw_body) - except json.JSONDecodeError: - self._send_json(400, {"error": {"message": "invalid JSON", "type": "invalid_request_error", "code": "invalid_request"}}) - return - - model: str = body.get("model", "") - is_stream: bool = bool(body.get("stream", False)) - - if model and server.stats is not None: - server.stats.record_request(model) - - # Billing gate (ADR-0015): reject before any routing — no free work. - api_key = _api_key_from_headers(self.headers) - if server.billing is not None: - if api_key is None: - self._send_json(401, {"error": { - "message": "missing API key: send Authorization: Bearer ", - "type": "invalid_request_error", - "code": "missing_api_key", - }}) - return - if server.accounts is not None and server.accounts.is_key_revoked(api_key): - self._send_json(401, {"error": { - "message": "API key has been revoked", - "type": "invalid_request_error", - "code": "invalid_api_key", - }}) - return - # US-039: with accounts enabled, only real account keys may spend — - # arbitrary bearer strings must never become billable clients. - if server.accounts is not None and not server.accounts.is_active_key(api_key): - self._send_json(401, {"error": { - "message": "unknown API key: create one at /dashboard (register, then + new key)", - "type": "invalid_request_error", - "code": "invalid_api_key", - }}) - return - if not server.billing.has_funds(api_key): - self._send_json(402, {"error": { - "message": "insufficient balance: deposit USDT to continue", - "type": "insufficient_quota", - "code": "insufficient_balance", - }}) - return - if server.max_charge_per_request is not None: - token_limit = _requested_completion_token_limit(body) - if token_limit is None: - self._send_json(400, {"error": { - "message": ( - "max_charge_per_request is enabled; include max_tokens " - "or max_completion_tokens so this request can be capped" - ), - "type": "invalid_request_error", - "code": "missing_token_limit", - }}) - return - in_rate, out_rate = server.billing.prices_for(model) - prompt_estimate = _estimate_prompt_tokens(body) or 0 - estimated_charge = (in_rate * prompt_estimate + out_rate * token_limit) / 1000.0 - if estimated_charge > server.max_charge_per_request: - self._send_json(402, {"error": { - "message": ( - f"request exceeds max_charge_per_request " - f"({estimated_charge:.6f} USDT estimated > " - f"{server.max_charge_per_request:.6f} USDT cap)" - ), - "type": "insufficient_quota", - "code": "spend_cap_exceeded", - }}) - return - - # US-030: optional pinned route — "route": [node_id, ...] uses those - # nodes in order instead of auto-selection. Absent field: unchanged. - pinned_ids = body.get("route") - pinned_nodes: list[_NodeEntry] | None = None - if pinned_ids is not None: - if ( - not isinstance(pinned_ids, list) - or not pinned_ids - or not all(isinstance(nid, str) and nid for nid in pinned_ids) - ): - self._send_json(400, {"error": { - "message": "route must be a non-empty list of node id strings", - "type": "invalid_request_error", - "code": "invalid_route", - }}) - return - with server.lock: - self._purge_expired_nodes() - missing = [nid for nid in pinned_ids if nid not in server.registry] - if missing: - self._send_json(400, {"error": { - "message": f"unknown node ids in route: {missing}", - "type": "invalid_request_error", - "code": "unknown_route_nodes", - }}) - return - pinned_nodes = [server.registry[nid] for nid in pinned_ids] - - if pinned_nodes is not None: - node = pinned_nodes[0] - else: - # Find a live tracker-mode node for this model - with server.lock: - self._purge_expired_nodes() - candidates = [ - n for n in server.registry.values() - if n.tracker_mode and _node_matches_model(n, model) - ] - - if not candidates: - # Fall back: any node serving shard_start=0 for this model - with server.lock: - candidates = [ - n for n in server.registry.values() - if n.shard_start == 0 and _node_matches_model(n, model) - ] - - if not candidates: - with server.lock: - registered = [ - { - "node_id": n.node_id, - "model": n.model, - "hf_repo": n.hf_repo, - "shard": f"{n.shard_start}-{n.shard_end}", - "tracker_mode": n.tracker_mode, - } - for n in server.registry.values() - ] - _tracker_log( - server, - "warn", - "no nodes available for model", - model=model, - registered_nodes=registered, - ) - self._send_json(503, {"error": { - "message": f"no nodes available for model {model!r}", - "type": "service_unavailable", - "code": "model_not_available", - }}) - return - - node = max(candidates, key=lambda n: _effective_throughput(n, model)) - target_url = f"{node.endpoint}/v1/chat/completions" - request_id = str(body.get("id") or f"req-{time.time_ns():x}") - body["id"] = request_id - raw_body = json.dumps(body).encode() - - # Pre-resolve the downstream route so the first-shard node skips its own - # tracker query. We already hold the full registry picture — no need for - # a second round-trip. - route_model = node.hf_repo or node.model or model - with server.lock: - resolved_route_model, route_preset = _resolve_model_preset(server.model_presets, route_model) - if route_preset is not None: - route_model = resolved_route_model or route_model - preset = route_preset - rs, re = _preset_layer_bounds(preset) - all_nodes: list = [ - n for n in server.registry.values() - if _node_matches_preset(n, route_model, preset) - and n.shard_start is not None - and n.shard_end is not None - ] - else: - all_nodes = [ - n for n in server.registry.values() - if _node_matches_model(n, route_model) - and n.shard_start is not None and n.num_layers is not None - ] - rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1) - if pinned_nodes is not None: - route_nodes = pinned_nodes - else: - route_nodes, route_error = _select_route(all_nodes, rs, re, model=route_model, contracts=server.contracts) - if route_error: - _tracker_log( - server, - "warn", - "route unavailable", - model=model, - route_model=route_model, - error=route_error, - candidate_count=len(all_nodes), - candidates=_node_route_summary(all_nodes), - ) - self._send_json(503, {"error": { - "message": route_error, - "type": "service_unavailable", - "code": "route_not_available", - }}) - return - # Compute start_layer for each hop: each node begins where the previous ended + 1. - # This allows overlapping shard registrations without double-computation. - covered_up_to = rs - 1 - route_hops: list[dict] = [] - node_work: list[tuple[str | None, int]] = [] - for rn in route_nodes: - hop: dict = {"endpoint": rn.endpoint, "start_layer": covered_up_to + 1} - if rn.relay_addr: - hop["relay_addr"] = rn.relay_addr - route_hops.append(hop) - effective_end = rn.shard_end if rn.shard_end is not None else covered_up_to - node_work.append((rn.wallet_address, max(0, effective_end - covered_up_to))) - covered_up_to = effective_end - # Strip the first-shard node we're about to proxy to — it's already handling the request. - downstream_hops = [h for h in route_hops if h["endpoint"].rstrip("/") != node.endpoint.rstrip("/")] - downstream_urls = json.dumps(downstream_hops) - route_debug = " -> ".join( - f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]" - for n in route_nodes - ) - inflight_nodes = route_nodes or [node] - inflight_recorded = True - _record_proxy_inflight(server, inflight_nodes, 1) - - def finish_proxy_inflight() -> None: - nonlocal inflight_recorded - if inflight_recorded: - _record_proxy_inflight(server, inflight_nodes, -1) - inflight_recorded = False - _unregister_active_proxy(server, request_id) - - proxy_ctx = _register_active_proxy(server, request_id) - - _tracker_log( - server, - "info", - "proxy route selected", - request_id=request_id, - model=model, - head_node_id=node.node_id, - head_endpoint=node.endpoint, - downstream=downstream_urls, - route=route_debug or "", - nodes=_node_route_summary(route_nodes), - ) - - req = urllib.request.Request( - target_url, - data=raw_body, - headers={ - "Content-Type": "application/json", - "X-Meshnet-Route": downstream_urls, - "X-Meshnet-Request-Id": request_id, - }, - method="POST", - ) - # Copy Authorization header from client if present - auth = self.headers.get("Authorization") - if auth: - req.add_header("Authorization", auth) - - relay_headers = { - "Content-Type": "application/json", - "X-Meshnet-Route": downstream_urls, - "X-Meshnet-Request-Id": request_id, - **({"Authorization": auth} if auth else {}), - } - - if node.relay_addr: - _tracker_log( - server, - "info", - "proxy via relay", - request_id=request_id, - relay_addr=node.relay_addr, - direct_endpoint=target_url, - ) - started = time.monotonic() - relay_ws_holder: list[Any] = [] - frames = _relay_http_request_frames( - node.relay_addr, - path="/v1/chat/completions", - body=raw_body, - headers=relay_headers, - cancel_event=proxy_ctx.cancel_event, - ws_holder=relay_ws_holder, - ws_lock=proxy_ctx.relay_ws_lock, - ) - first = next(frames, None) - with proxy_ctx.relay_ws_lock: - proxy_ctx.relay_ws = relay_ws_holder[0] if relay_ws_holder else None - if proxy_ctx.cancel_event.is_set(): - if self._finalize_proxy_cancel( - proxy_ctx=proxy_ctx, - server=server, - request_id=request_id, - started=started, - model=model, - route_model=route_model, - route_nodes=route_nodes, - api_key=api_key, - node_work=node_work, - body=body, - finish_proxy_inflight=finish_proxy_inflight, - ): - return - if first is not None and first.get("stream"): - # Streamed response (US-036): forward SSE chunks as they arrive - # and run the same token accounting as the direct stream path. - self._stream_relayed_frames( - first, frames, started, - model, route_model, route_nodes, api_key, node_work, - request_body=body, - request_id=request_id, - proxy_ctx=proxy_ctx, - finish_proxy_inflight=finish_proxy_inflight, - ) - finish_proxy_inflight() - return - if first is not None: - elapsed = time.monotonic() - started - self._send_relayed_response(first) - if int(first.get("status", 503)) < 400: - body_text = first.get("body") or "" - try: - in_tokens, out_tokens = _billable_non_stream_split(json.loads(body_text), body) - except (json.JSONDecodeError, TypeError): - in_tokens, out_tokens = 0, 0 - tokens = in_tokens + out_tokens - self._record_observed_throughput(model, route_model, tokens, elapsed, route_nodes) - _clear_proxy_progress_log_state(server, request_id) - _tracker_log( - server, - "info", - "proxy complete", - request_id=request_id, - model=model, - route_model=route_model, - status=int(first.get("status", 503)), - tokens=tokens, - elapsed_seconds=round(elapsed, 4), - tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, - route=_node_route_summary(route_nodes), - ) - self._bill_completed( - api_key, model, tokens, node_work, - input_tokens=in_tokens, output_tokens=out_tokens, - ) - finish_proxy_inflight() - return - _tracker_log( - server, - "warn", - "relay proxy failed, trying direct", - request_id=request_id, - relay_addr=node.relay_addr, - direct_endpoint=target_url, - ) - - try: - started = time.monotonic() - upstream_result: list[Any] = [] - connect_errors: list[BaseException] = [] - - def _connect_upstream() -> None: - try: - upstream_result.append(urllib.request.urlopen(req, timeout=300.0)) - except BaseException as exc: - connect_errors.append(exc) - - connect_thread = threading.Thread(target=_connect_upstream, daemon=True) - connect_thread.start() - while connect_thread.is_alive(): - if proxy_ctx.cancel_event.is_set(): - connect_thread.join(timeout=310.0) - if upstream_result: - try: - upstream_result[0].close() - except Exception: - pass - if self._finalize_proxy_cancel( - proxy_ctx=proxy_ctx, - server=server, - request_id=request_id, - started=started, - model=model, - route_model=route_model, - route_nodes=route_nodes, - api_key=api_key, - node_work=node_work, - body=body, - finish_proxy_inflight=finish_proxy_inflight, - ): - return - connect_thread.join(0.2) - - if proxy_ctx.cancel_event.is_set(): - if upstream_result: - try: - upstream_result[0].close() - except Exception: - pass - if self._finalize_proxy_cancel( - proxy_ctx=proxy_ctx, - server=server, - request_id=request_id, - started=started, - model=model, - route_model=route_model, - route_nodes=route_nodes, - api_key=api_key, - node_work=node_work, - body=body, - finish_proxy_inflight=finish_proxy_inflight, - ): - return - - if connect_errors: - raise connect_errors[0] - - upstream = upstream_result[0] - with proxy_ctx.upstream_lock: - proxy_ctx.upstream = upstream - upstream_sock = _upstream_socket(upstream) - if upstream_sock is not None: - _set_upstream_read_timeout(upstream, None) - else: - _set_upstream_read_timeout(upstream, 0.5) - _tracker_log(server, "info", "proxy connected", request_id=request_id, target_url=target_url) - except urllib.error.HTTPError as exc: - # Relay error status + body from node - err_body = exc.read() - self.send_response(exc.code) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(err_body))) - self.end_headers() - try: - self.wfile.write(err_body) - except BrokenPipeError: - pass - _clear_proxy_progress_log_state(server, request_id) - finish_proxy_inflight() - return - except Exception as exc: - _clear_proxy_progress_log_state(server, request_id) - if node.relay_addr: - _tracker_log( - server, - "error", - "direct proxy failed after relay", - request_id=request_id, - target_url=target_url, - relay_addr=node.relay_addr, - error=repr(exc), - ) - else: - _tracker_log( - server, - "error", - "proxy failed", - request_id=request_id, - target_url=target_url, - error=repr(exc), - ) - self._send_json(503, {"error": { - "message": f"upstream node unreachable: {exc}", - "type": "service_unavailable", - "code": "upstream_error", - }}) - finish_proxy_inflight() - return - - with upstream: - content_type = upstream.headers.get("Content-Type", "application/json") - if is_stream or "text/event-stream" in content_type: - # Relay SSE stream chunk-by-chunk - self.send_response(200) - self.send_header("Content-Type", "text/event-stream; charset=utf-8") - self.send_header("Cache-Control", "no-cache") - self.send_header("X-Accel-Buffering", "no") - self.end_headers() - stream_usage: dict | None = None - observed_stream_tokens = 0 - client_gone = False - try: - while True: - if proxy_ctx.cancel_event.is_set(): - break - if upstream_sock is not None: - readable, _, _ = select.select([upstream_sock], [], [], 0.5) - if not readable: - continue - try: - line = upstream.readline() - except TimeoutError: - continue - if not line: - if proxy_ctx.cancel_event.is_set(): - break - break - if not client_gone: - try: - self.wfile.write(line) - self.wfile.flush() - except (BrokenPipeError, ConnectionResetError): - # Keep draining upstream so completed node work is still billed. - client_gone = True - observed, usage = _stream_line_tokens(line) - observed_stream_tokens += observed - if observed: - _tracker_log_proxy_progress( - server, - request_id=request_id, - model=model, - route_model=route_model, - tokens=observed_stream_tokens, - started=started, - route_nodes=route_nodes, - ) - if usage is not None: - stream_usage = usage - except (BrokenPipeError, ConnectionResetError): - client_gone = True - if self._finalize_proxy_cancel( - proxy_ctx=proxy_ctx, - server=server, - request_id=request_id, - started=started, - model=model, - route_model=route_model, - route_nodes=route_nodes, - api_key=api_key, - node_work=node_work, - body=body, - finish_proxy_inflight=finish_proxy_inflight, - observed_stream_tokens=observed_stream_tokens, - stream_usage=stream_usage, - ): - return - elapsed = time.monotonic() - started - # Bill even on client disconnect — the nodes did the work. - # Observed stream chunks are authoritative for the upper bound; - # upstream usage may only lower that count. - in_tokens, out_tokens = _stream_billable_split( - observed_stream_tokens, stream_usage, body - ) - self._record_observed_throughput( - model, route_model, in_tokens + out_tokens, elapsed, route_nodes - ) - tokens = in_tokens + out_tokens - _clear_proxy_progress_log_state(server, request_id) - _tracker_log( - server, - "info", - "proxy complete", - request_id=request_id, - model=model, - route_model=route_model, - status=200, - stream=True, - client_disconnected=client_gone, - tokens=tokens, - elapsed_seconds=round(elapsed, 4), - tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, - route=_node_route_summary(route_nodes), - ) - self._bill_completed( - api_key, model, tokens, node_work, - input_tokens=in_tokens, output_tokens=out_tokens, - ) - else: - # Non-streaming: buffer and relay - resp_body = upstream.read() - elapsed = time.monotonic() - started - observed_output = "" - try: - response_payload = json.loads(resp_body) - in_tokens, out_tokens = _billable_non_stream_split(response_payload, body) - observed_output = _observed_output_from_non_stream_payload(response_payload) - except json.JSONDecodeError: - in_tokens, out_tokens = 0, 0 - tokens = in_tokens + out_tokens - self._record_observed_throughput(model, route_model, tokens, elapsed, route_nodes) - _clear_proxy_progress_log_state(server, request_id) - _tracker_log( - server, - "info", - "proxy complete", - request_id=request_id, - model=model, - route_model=route_model, - target_url=target_url, - status=200, - bytes=len(resp_body), - tokens=tokens, - elapsed_seconds=round(elapsed, 4), - tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, - route=_node_route_summary(route_nodes), - ) - self._record_validation_event(request_id, model, body, observed_output, route_nodes) - self._bill_completed( - api_key, model, tokens, node_work, - input_tokens=in_tokens, output_tokens=out_tokens, - ) - self.send_response(200) - self.send_header("Content-Type", content_type) - self.send_header("Content-Length", str(len(resp_body))) - self.end_headers() - try: - self.wfile.write(resp_body) - except (BrokenPipeError, ConnectionResetError): - pass - finish_proxy_inflight() - - def _record_observed_throughput( - self, - requested_model: str, - route_model: str, - total_tokens: int, - elapsed_seconds: float, - route_nodes: list[_NodeEntry], - ) -> None: - """Record observed route TPS for participating nodes. - - The tracker sees end-to-end request duration, not per-hop timings, so - each hop gets the same route-level observation for now. Per-hop telemetry - can refine this later without changing the external stats shape. - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if server.stats is None or total_tokens <= 0: - return - elapsed_seconds = max(elapsed_seconds, 1e-6) - models = [m for m in (requested_model, route_model) if m] - if len(models) == 2 and models[0] == models[1]: - models = [models[0]] - for node in route_nodes: - for model in models: - server.stats.record_node_throughput( - node.node_id, - model, - total_tokens=total_tokens, - elapsed_seconds=elapsed_seconds, - ) - stats = server.stats.get_node_model_stats(node.node_id, model) - observed = stats.get("tokens_per_sec_last_hour") - if observed is not None: - node.model_tokens_per_sec[model] = float(observed) - - def _record_validation_event( - self, - session_id: str, - model: str, - request_body: dict, - observed_output: str, - route_nodes: list[_NodeEntry], - ) -> None: - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - validation = getattr(server.contracts, "validation", None) if server.contracts is not None else None - if validation is None: - return - messages = request_body.get("messages") - if not isinstance(messages, list): - messages = [] - event_nodes = [ - { - "node_id": node.node_id, - "endpoint": node.endpoint, - "wallet_address": node.wallet_address, - "shard_start": node.shard_start, - "shard_end": node.shard_end, - } - for node in route_nodes - ] - try: - validation.record_completed_inference( - session_id=session_id, - model=model, - messages=[dict(message) for message in messages if isinstance(message, dict)], - observed_output=observed_output, - route_nodes=event_nodes, - ) - except Exception as exc: - print(f"[tracker] validation event recording failed for {session_id}: {exc}", flush=True) - - def _bill_completed( - self, - api_key: str | None, - model: str, - total_tokens: int, - node_work: list[tuple[str | None, int]], - *, - input_tokens: int | None = None, - output_tokens: int | None = None, - ) -> None: - """Charge a completed request against the billing ledger (ADR-0015). - - With ``input_tokens``/``output_tokens`` the ledger bills each side at - its own rate (US-045); ``total_tokens`` alone falls back to the - blended rate. - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if server.billing is None or api_key is None: - return - # Probationary period (issue 08): a wallet's first N jobs earn nothing — - # its share stays in the protocol cut. Job counts live in the registry - # contract, so this only applies when the tracker runs with contracts. - if server.contracts is not None: - adjusted: list[tuple[str | None, int]] = [] - for wallet, work in node_work: - if wallet: - in_probation = server.contracts.registry.probationary_jobs_remaining(wallet) > 0 - server.contracts.registry.record_completed_job(wallet) - if in_probation: - wallet = None - adjusted.append((wallet, work)) - node_work = adjusted - try: - event = server.billing.charge_request( - api_key, model, total_tokens, node_work, - input_tokens=input_tokens, output_tokens=output_tokens, - ) - print( - f"[tracker] billed api_key=…{api_key[-6:]}: model={model!r} " - f"tokens={event['total_tokens']} " - f"(in={event.get('input_tokens', '?')} out={event.get('output_tokens', '?')}) " - f"cost={event['cost']:.6f} USDT shares={event['shares']}", - flush=True, - ) - except Exception as exc: - print(f"[tracker] billing failed for model={model!r}: {exc}", flush=True) - - def _stream_relayed_frames( - self, - first: dict, - frames, - started: float, - model: str, - route_model: str, - route_nodes: list, - api_key: str | None, - node_work: list, - request_body: dict, - request_id: str, - *, - proxy_ctx: _ActiveProxyContext | None = None, - finish_proxy_inflight: Any = None, - ) -> None: - """Forward a streamed relay response (US-036) to the client as SSE, - billing with the same accounting as the direct stream path.""" - headers = first.get("headers") if isinstance(first.get("headers"), dict) else {} - self.send_response(int(first.get("status", 200))) - self.send_header("Content-Type", headers.get("Content-Type", "text/event-stream; charset=utf-8")) - self.send_header("Cache-Control", "no-cache") - self.send_header("X-Accel-Buffering", "no") - self.end_headers() - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - stream_usage: dict | None = None - observed_stream_tokens = 0 - client_gone = False - for frame in itertools.chain([first], frames): - if proxy_ctx is not None and proxy_ctx.cancel_event.is_set(): - break - chunk = frame.get("chunk") or "" - if not chunk: - continue - data = chunk.encode() - if not client_gone: - try: - self.wfile.write(data) - self.wfile.flush() - except BrokenPipeError: - # Keep draining frames — the nodes did the work; bill it. - client_gone = True - for line in data.splitlines(): - observed, usage = _stream_line_tokens(line) - observed_stream_tokens += observed - if observed: - _tracker_log_proxy_progress( - server, - request_id=request_id, - model=model, - route_model=route_model, - tokens=observed_stream_tokens, - started=started, - route_nodes=route_nodes, - relay=True, - ) - if usage is not None: - stream_usage = usage - if ( - proxy_ctx is not None - and finish_proxy_inflight is not None - and self._finalize_proxy_cancel( - proxy_ctx=proxy_ctx, - server=server, - request_id=request_id, - started=started, - model=model, - route_model=route_model, - route_nodes=route_nodes, - api_key=api_key, - node_work=node_work, - body=request_body, - finish_proxy_inflight=finish_proxy_inflight, - observed_stream_tokens=observed_stream_tokens, - stream_usage=stream_usage, - ) - ): - return - elapsed = time.monotonic() - started - in_tokens, out_tokens = _stream_billable_split( - observed_stream_tokens, stream_usage, request_body - ) - self._record_observed_throughput( - model, route_model, in_tokens + out_tokens, elapsed, route_nodes - ) - tokens = in_tokens + out_tokens - _clear_proxy_progress_log_state(server, request_id) - _tracker_log( - server, - "info", - "proxy complete", - request_id=request_id, - model=model, - route_model=route_model, - status=200, - stream=True, - client_disconnected=client_gone, - relay=True, - tokens=tokens, - elapsed_seconds=round(elapsed, 4), - tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, - route=_node_route_summary(route_nodes), - ) - self._bill_completed( - api_key, model, tokens, node_work, - input_tokens=in_tokens, output_tokens=out_tokens, - ) - - def _send_relayed_response(self, response: dict) -> None: - status = int(response.get("status", 503)) - headers = response.get("headers") if isinstance(response.get("headers"), dict) else {} - body_text = response.get("body") or "" - body = body_text.encode() if isinstance(body_text, str) else bytes(body_text) - self.send_response(status) - self.send_header("Content-Type", headers.get("Content-Type", "application/json")) - self.send_header("Content-Length", str(len(body))) - self.end_headers() - try: - self.wfile.write(body) - except BrokenPipeError: - pass - - def _handle_register(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - body = self._read_json_body() - if body is None: - return - - # --- Raft cluster mode: forward to leader or propose via Raft --- - if server.raft is not None: - if not server.raft.is_leader: - leader = server.raft.leader() - if leader is None: - self._send_json(503, {"error": "no leader elected — retry in a moment"}) - return - # Proxy to leader - try: - data = json.dumps(body).encode() - req = urllib.request.Request( - f"{leader}/v1/nodes/register", - data=data, - headers={"Content-Type": "application/json"}, - method="POST", - ) - with urllib.request.urlopen(req, timeout=5.0) as r: - resp_body = r.read() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(resp_body))) - self.end_headers() - self.wfile.write(resp_body) - except Exception as exc: - self._send_json(503, {"error": f"leader proxy failed: {exc}"}) - return - # Leader path: fall through to normal registration, then replicate via Raft. - # We let the registration run first (to generate node_id), then asynchronously - # propose to the Raft log so followers can replicate the entry. - # Raft proposal happens after the response is sent (fire-and-forget replication). - - endpoint = body.get("endpoint") - if not isinstance(endpoint, str) or not endpoint: - self._send_json(400, {"error": "endpoint is required"}) - return - parsed_endpoint = urllib.parse.urlparse(endpoint) - if parsed_endpoint.scheme not in {"http", "https"} or not parsed_endpoint.netloc: - self._send_json(400, {"error": "endpoint must be an http(s) URL"}) - return - - shard_start: int | None - shard_end: int | None - explicit_shard = "shard_start" in body or "shard_end" in body - if explicit_shard: - try: - shard_start = int(body["shard_start"]) - shard_end = int(body["shard_end"]) - except (KeyError, TypeError, ValueError): - self._send_json(400, {"error": "shard_start and shard_end must be numeric"}) - return - if shard_start < 0 or shard_end < 0 or shard_start > shard_end: - self._send_json(400, {"error": "shard range must be non-negative and ordered"}) - return - else: - shard_start = None - shard_end = None - try: - score = float(body.get("score", 1.0)) - except (TypeError, ValueError): - self._send_json(400, {"error": "score must be numeric"}) - return - - hardware_profile = body.get("hardware_profile", {}) - if not isinstance(hardware_profile, dict): - self._send_json(400, {"error": "hardware_profile must be an object"}) - return - model = body.get("model") - if model is None: - model = "stub-model" - if not isinstance(model, str): - self._send_json(400, {"error": "model must be a string"}) - return - shard_checksum = body.get("shard_checksum") - if shard_checksum is not None and not isinstance(shard_checksum, str): - self._send_json(400, {"error": "shard_checksum must be a string"}) - return - try: - vram_bytes = int(body.get("vram_bytes", DEFAULT_VRAM_BYTES)) - ram_bytes = int(body.get("ram_bytes", DEFAULT_RAM_BYTES)) - max_loaded_shards = int(body.get("max_loaded_shards", 1)) - benchmark_tokens_per_sec = float( - body.get("benchmark_tokens_per_sec", DEFAULT_BENCHMARK_TOKENS_PER_SEC) - ) - except (TypeError, ValueError): - self._send_json(400, {"error": "vram_bytes, ram_bytes, max_loaded_shards, and benchmark_tokens_per_sec must be numeric"}) - return - if vram_bytes < 0 or ram_bytes < 0 or max_loaded_shards < 1 or benchmark_tokens_per_sec <= 0: - self._send_json(400, {"error": "capability values must be positive"}) - return - quantizations_body = body.get("quantizations", DEFAULT_QUANTIZATIONS) - if not ( - isinstance(quantizations_body, list) - and quantizations_body - and all(isinstance(item, str) and item for item in quantizations_body) - ): - self._send_json(400, {"error": "quantizations must be a non-empty string array"}) - return - quantizations = list(quantizations_body) - quantization = body.get("quantization") - if quantization is not None and not isinstance(quantization, str): - self._send_json(400, {"error": "quantization must be a string"}) - return - wallet_address = body.get("wallet_address") - if wallet_address is not None and not isinstance(wallet_address, str): - self._send_json(400, {"error": "wallet_address must be a string"}) - return - ban_error = _registration_ban_error(server.contracts, wallet_address) - if ban_error: - self._send_json(403, {"error": ban_error}) - return - - tracker_mode = bool(body.get("tracker_mode", False)) - managed_assignment = bool(body.get("managed_assignment", False)) - hf_repo = body.get("hf_repo") - if hf_repo is not None and not isinstance(hf_repo, str): - self._send_json(400, {"error": "hf_repo must be a string"}) - return - num_layers_body = body.get("num_layers") - num_layers: int | None = None - if num_layers_body is not None: - try: - num_layers = int(num_layers_body) - except (TypeError, ValueError): - self._send_json(400, {"error": "num_layers must be an integer"}) - return - model_metadata = body.get("model_metadata", {}) - if model_metadata is None: - model_metadata = {} - if not isinstance(model_metadata, dict): - self._send_json(400, {"error": "model_metadata must be an object"}) - return - downloaded_models = body.get("downloaded_models", []) - if downloaded_models is None: - downloaded_models = [] - if not ( - isinstance(downloaded_models, list) - and all(isinstance(item, dict) for item in downloaded_models) - ): - self._send_json(400, {"error": "downloaded_models must be an array of objects"}) - return - relay_addr = body.get("relay_addr") or None - cert_fingerprint = body.get("cert_fingerprint") or None - peer_id = body.get("peer_id") or None - - node_id = _node_id_for_registration( - endpoint, - model, - wallet_address, - shard_start, - shard_end, - hf_repo, - ) - entry = _NodeEntry( - node_id=node_id, - endpoint=endpoint.rstrip("/"), - shard_start=shard_start, - shard_end=shard_end, - model=model, - shard_checksum=shard_checksum, - hardware_profile=hardware_profile, - wallet_address=wallet_address, - score=score, - vram_bytes=vram_bytes, - ram_bytes=ram_bytes, - quantizations=quantizations, - max_loaded_shards=max_loaded_shards, - benchmark_tokens_per_sec=benchmark_tokens_per_sec, - quantization=quantization, - managed_assignment=managed_assignment or not explicit_shard, - tracker_mode=tracker_mode, - hf_repo=hf_repo, - num_layers=num_layers, - model_metadata=model_metadata, - downloaded_models=downloaded_models, - relay_addr=relay_addr, - cert_fingerprint=cert_fingerprint, - peer_id=peer_id, - ) - with server.lock: - self._purge_expired_nodes() - # Dedup: replace the same node id or the same endpoint+model assignment. - endpoint_key = entry.endpoint.rstrip("/") - model_key = entry.hf_repo or entry.model - stale_ids = [ - eid for eid, e in server.registry.items() - if eid == node_id - or ( - e.endpoint.rstrip("/") == endpoint_key - and (e.hf_repo or e.model) == model_key - ) - ] - stale_entries: list[tuple[str, _NodeEntry]] = [] - for eid in stale_ids: - old = server.registry.pop(eid) - stale_entries.append((eid, old)) - server.registry[node_id] = entry - if entry.managed_assignment and not explicit_shard: - if entry.hf_repo: - _rebalance_hf_model_locked(server, entry.hf_repo) - else: - _rebalance_model_locked(server, model) - assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None - if assignment_directive is not None: - entry.pending_directives.clear() - - model_health = _model_health_summary(server, entry.model, entry.hf_repo) - for eid, old in stale_entries: - _tracker_log( - server, - "info", - "node re-registered", - old_node_id=eid, - node_id=node_id, - endpoint=old.endpoint, - old_model=old.model, - old_hf_repo=old.hf_repo, - old_model_health=_model_health_summary(server, old.model, old.hf_repo), - model_health=model_health, - ) - _tracker_log( - server, - "info", - "node registered", - node_id=node_id, - endpoint=entry.endpoint, - model=entry.model, - hf_repo=entry.hf_repo, - shard=f"{entry.shard_start}-{entry.shard_end}", - tracker_mode=entry.tracker_mode, - model_health=model_health, - ) - - shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded" - repo_info = f" [{hf_repo}]" if hf_repo else "" - budget_bytes, budget_source = _node_memory_budget_bytes(entry) - budget_gb = budget_bytes / (1024 ** 3) - print( - f"[tracker] node registered: {node_id} {endpoint} {model}{repo_info} {shard_info} " - f"capacity={budget_gb:.1f}GB {budget_source} slots={max_loaded_shards}", - flush=True, - ) - - payload = {"node_id": node_id} - if assignment_directive is not None: - payload["directive"] = assignment_directive - self._send_json(200, payload) - - # Raft replication: leader proposes this registration to followers so their - # registries stay in sync. Fire-and-forget (async) — the client already - # got its node_id; replication happens in the background. - if server.raft is not None and server.raft.is_leader: - raft_payload = dict(body) - raft_payload["node_id"] = node_id # include the generated ID - threading.Thread( - target=server.raft.propose, - args=("register", raft_payload), - daemon=True, - ).start() - - def _handle_heartbeat(self, node_id: str): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - body: dict = {} - content_length = int(self.headers.get("Content-Length", 0)) - if content_length > 0: - try: - body = json.loads(self.rfile.read(content_length)) - except Exception: - pass - with server.lock: - self._purge_expired_nodes() - entry = server.registry.get(node_id) - if entry is None: - _tracker_log(server, "warn", "heartbeat for unknown node", node_id=node_id) - self._send_json(404, {"error": "node not found"}) - return - entry.last_heartbeat = time.monotonic() - entry.heartbeats_received += 1 - # P2P metadata - if body.get("relay_addr"): - entry.relay_addr = body["relay_addr"] - if body.get("cert_fingerprint"): - entry.cert_fingerprint = body["cert_fingerprint"] - if body.get("peer_id"): - entry.peer_id = body["peer_id"] - # Node stats (cumulative — node always sends totals, tracker replaces) - if "total_requests" in body: - entry.total_requests = int(body["total_requests"]) - if "failed_requests" in body: - entry.failed_requests = int(body["failed_requests"]) - if "queue_depth" in body: - entry.queue_depth = int(body["queue_depth"]) - if "current_requests" in body: - entry.current_requests = _normalize_current_requests(body["current_requests"]) - if "uptime_seconds" in body: - entry.uptime_seconds = float(body["uptime_seconds"]) - if "status" in body and body["status"] in ("ready", "loading"): - entry.status = body["status"] - directives = list(entry.pending_directives) - entry.pending_directives.clear() - new_assignment = entry.pending_new_assignment - if new_assignment is not None: - entry.pending_new_assignment = None - if server.gossip is not None: - server.gossip.record(node_id) - resp: dict = {} - if directives: - resp["directives"] = directives - if new_assignment is not None: - resp["new_assignment"] = new_assignment - self._send_json(200, resp) - - # ---------------------------------------------------------------- Raft handlers - - def _handle_raft_vote(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - body = self._read_json_body() - if body is None: - return - if server.raft is None: - self._send_json(503, {"error": "raft not enabled"}) - return - result = server.raft.handle_request_vote(body) - self._send_json(200, result) - - def _handle_raft_append(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - body = self._read_json_body() - if body is None: - return - if server.raft is None: - self._send_json(503, {"error": "raft not enabled"}) - return - result = server.raft.handle_append_entries(body) - self._send_json(200, result) - - def _handle_raft_status(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if server.raft is None: - self._send_json(200, {"role": "standalone", "term": 0, "leader": None}) - return - self._send_json(200, server.raft.status()) - - def _handle_gossip(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - body = self._read_json_body() - if body is None: - return - if server.gossip is not None and isinstance(body, dict): - server.gossip.merge({k: float(v) for k, v in body.items()}) - self._send_json(200, {}) - - def _handle_stats(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if server.stats is None: - self._send_json(200, {"models": {}, "nodes": {}}) - return - self._send_json(200, { - "models": server.stats.get_combined_stats(), - "nodes": server.stats.get_node_throughput_stats(), - }) - - def _handle_stats_gossip(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - body = self._read_hive_authenticated_body() - if body is None: - return - tracker_url = body.get("tracker_url", "") - rpms = body.get("stats", {}) - if server.stats is not None and tracker_url and isinstance(rpms, dict): - server.stats.merge_peer_rpms(tracker_url, rpms) - self._send_json(200, {}) - - def _handle_billing_summary(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if not self._require_role("admin"): - return - if server.billing is None: - self._send_json(404, {"error": "billing is not enabled on this tracker"}) - return - self._send_json(200, server.billing.snapshot()) - - def _handle_dashboard(self): - """Serve the read-only web dashboard (US-035). Any tracker in the - mesh — leader or follower — serves it from its replicated state.""" - try: - html = files("meshnet_tracker").joinpath("dashboard.html").read_text() - except (FileNotFoundError, OSError): - self._send_json(404, {"error": "dashboard asset missing"}) - return - body = html.encode() - self.send_response(200) - self.send_header("Content-Type", "text/html; charset=utf-8") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - try: - self.wfile.write(body) - except BrokenPipeError: - pass - - def _handle_console(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - with server.console_lock: - events = [dict(event) for event in server.console_events] - self._send_json(200, {"events": events}) - - def _handle_proxy_request_cancel(self, request_id: str) -> None: - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if server.accounts is not None and not self._require_role("admin"): - return - if not _request_proxy_cancel(server, request_id): - self._send_json(404, {"error": f"no active proxy for request {request_id!r}"}) - return - self._send_json(200, {"status": "canceled", "request_id": request_id}) - - def _finalize_proxy_cancel( - self, - *, - proxy_ctx: _ActiveProxyContext, - server: "_TrackerHTTPServer", - request_id: str, - started: float, - model: str, - route_model: str, - route_nodes: list, - api_key: str | None, - node_work: list, - body: dict, - finish_proxy_inflight: Any, - observed_stream_tokens: int = 0, - stream_usage: dict | None = None, - ) -> bool: - if not proxy_ctx.cancel_event.is_set(): - return False - elapsed = time.monotonic() - started - _clear_proxy_progress_log_state(server, request_id) - tokens = observed_stream_tokens - if observed_stream_tokens > 0: - in_tokens, out_tokens = _stream_billable_split( - observed_stream_tokens, stream_usage, body, - ) - tokens = in_tokens + out_tokens - self._record_observed_throughput( - model, route_model, tokens, elapsed, route_nodes, - ) - _tracker_log( - server, - "info", - "proxy canceled", - request_id=request_id, - model=model, - route_model=route_model, - tokens=tokens, - elapsed_seconds=round(elapsed, 4), - tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 and tokens > 0 else 0.0, - route=_node_route_summary(route_nodes), - ) - if observed_stream_tokens > 0: - in_tokens, out_tokens = _stream_billable_split( - observed_stream_tokens, stream_usage, body, - ) - self._bill_completed( - api_key, model, in_tokens + out_tokens, node_work, - input_tokens=in_tokens, output_tokens=out_tokens, - ) - finish_proxy_inflight() - return True - - def _handle_registry_wallets(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if not self._require_role("admin"): - return - if server.contracts is None: - self._send_json(200, {"wallets": {}}) - return - wallets = server.contracts.registry.list_wallets() - self._send_json(200, {"wallets": { - wallet: { - "stake_balance": state.stake_balance, - "strike_count": state.strike_count, - "banned": state.banned, - "completed_jobs": state.completed_job_count, - "reputation": state.reputation, - "last_audit_ts": state.last_audit_ts, - "probationary_jobs_remaining": ( - server.contracts.registry.probationary_jobs_remaining(wallet) - ), - } - for wallet, state in wallets.items() - }}) - - def _handle_billing_settlements(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if not self._require_role("admin"): - return - if server.billing is None: - self._send_json(404, {"error": "billing is not enabled on this tracker"}) - return - self._send_json(200, {"settlements": server.billing.settlement_history()}) - - def _handle_billing_gossip(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - body = self._read_hive_authenticated_body() - if body is None: - return - if server.billing is None: - self._send_json(200, {"applied": 0}) - return - events = body.get("events") - if not isinstance(events, list): - self._send_json(400, {"error": "events must be a list"}) - return - applied = server.billing.apply_events([e for e in events if isinstance(e, dict)]) - self._send_json(200, {"applied": applied}) - - def _handle_billing_forfeit(self): - """Privileged: forfeit a node's pending balance + record a strike (US-034). - - ADR-0017 §4: validator service token or admin session only. Client - API keys — valid or not — are rejected. - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if not self._require_role("validator", "admin"): - return - if server.billing is None: - self._send_json(404, {"error": "billing is not enabled on this tracker"}) - return - body = self._read_json_body() - if body is None: - return - wallet = body.get("wallet") - if not wallet or not isinstance(wallet, str): - self._send_json(400, {"error": "wallet is required"}) - return - reason = body.get("reason") or "fraud" - event = server.billing.forfeit_pending(wallet, reason=str(reason)) - strike_state: dict = {} - if server.contracts is not None: - server.contracts.registry.record_strike(wallet) - wallet_state = server.contracts.registry.get_wallet(wallet) - strike_state = { - "strike_count": wallet_state.strike_count, - "banned": wallet_state.banned, - } - print( - f"[tracker] forfeited pending balance of {wallet}: " - f"{event['amount']:.6f} USDT ({reason}) strikes={strike_state.get('strike_count', 'n/a')}", - flush=True, - ) - self._send_json(200, {"forfeited": event["amount"], **strike_state}) - - # ---- user accounts (registration, login, API keys) ---- - - def _session_account(self) -> dict | None: - """Resolve the session token in the Authorization header, or None.""" - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if server.accounts is None: - return None - return server.accounts.session_account(_session_token_from_headers(self.headers)) - - def _require_accounts(self) -> "AccountStore | None": - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if server.accounts is None: - self._send_json(404, {"error": "accounts are not enabled on this tracker"}) - return None - return server.accounts - - def _handle_auth_register(self): - accounts = self._require_accounts() - if accounts is None: - return - body = self._read_json_body() - if body is None: - return - try: - account = accounts.register( - email=body.get("email"), - wallet=body.get("wallet"), - password=str(body.get("password") or ""), - ) - except ValueError as exc: - self._send_json(400, {"error": str(exc)}) - return - token = accounts.create_session(account["account_id"]) - api_key = accounts.create_api_key(account["account_id"]) - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if server.billing is not None: - # US-039: registration creates the account's first key, so Caller - # Credit lands here; the account-derived event id keeps later - # grants (e.g. via /v1/account/keys) no-ops. - server.billing.grant_caller_credit( - api_key, account["account_id"], server.starting_credit - ) - server.billing.ensure_client(api_key) - print( - f"[tracker] account registered: {account.get('email') or account.get('wallet')} " - f"role={account['role']}", flush=True, - ) - self._send_json( - 200, - {"account": account, "session_token": token, "api_key": api_key}, - headers={"Set-Cookie": _session_cookie_header(token)}, - ) - - def _handle_auth_login(self): - accounts = self._require_accounts() - if accounts is None: - return - body = self._read_json_body() - if body is None: - return - identifier = body.get("identifier") or body.get("email") or body.get("wallet") or "" - account = accounts.verify_login(str(identifier), str(body.get("password") or "")) - if account is None: - self._send_json(401, {"error": "invalid credentials"}) - return - token = accounts.create_session(account["account_id"]) - self._send_json( - 200, - {"account": account, "session_token": token}, - headers={"Set-Cookie": _session_cookie_header(token)}, - ) - - def _handle_auth_logout(self): - accounts = self._require_accounts() - if accounts is None: - return - accounts.destroy_session(_session_token_from_headers(self.headers)) - self._send_json(200, {"ok": True}, headers={"Set-Cookie": _session_cookie_header(None)}) - - def _handle_account_me(self): - """Balance, usage, and API keys for the logged-in account.""" - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if self._require_accounts() is None: - return - account = self._session_account() - if account is None: - self._send_json(401, {"error": "login required"}) - return - keys = server.accounts.keys_for(account["account_id"]) # type: ignore[union-attr] - balances = {} - usage: dict = {"requests": 0, "total_tokens": 0, "total_cost": 0.0, "recent": [], "records": []} - if server.billing is not None: - balances = {key: server.billing.get_client_balance(key) for key in keys} - usage = server.billing.usage_for(keys) - self._send_json(200, { - "account": account, - "api_keys": keys, - "balances": balances, - "total_balance": sum(balances.values()), - "usage": usage, - "topup_amount": server.devnet_topup_amount if server.billing is not None else 0.0, - }) - - def _handle_account_key_create(self): - accounts = self._require_accounts() - if accounts is None: - return - account = self._session_account() - if account is None: - self._send_json(401, {"error": "login required"}) - return - api_key = accounts.create_api_key(account["account_id"]) - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - credited = False - if server.billing is not None: - # US-039: Caller Credit lands with the account's first key; the - # account-derived event id keeps it once-per-account forever. - credited = server.billing.grant_caller_credit( - api_key, account["account_id"], server.starting_credit - ) - server.billing.ensure_client(api_key) - self._send_json(200, {"api_key": api_key, "caller_credit_granted": credited}) - - def _handle_account_topup(self): - """Devnet faucet (US-040): credit the configured amount to one of the - logged-in account's keys. 404 unless the operator enabled it.""" - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - accounts = self._require_accounts() - if accounts is None: - return - if server.billing is None or server.devnet_topup_amount <= 0: - self._send_json(404, {"error": "top-up is not enabled on this tracker"}) - return - account = self._session_account() - if account is None: - self._send_json(401, {"error": "login required"}) - return - body = self._read_json_body() - if body is None: - return - api_key = body.get("api_key") - if not api_key or not isinstance(api_key, str): - self._send_json(400, {"error": "api_key is required"}) - return - if accounts.owner_of_key(api_key) != account["account_id"]: - self._send_json(403, {"error": "key not found on this account"}) - return - balance = server.billing.credit_client( - api_key, server.devnet_topup_amount, note="devnet-topup" - ) - self._send_json(200, { - "api_key": api_key, - "credited": server.devnet_topup_amount, - "balance": balance, - }) - - def _handle_account_key_revoke(self): - accounts = self._require_accounts() - if accounts is None: - return - account = self._session_account() - if account is None: - self._send_json(401, {"error": "login required"}) - return - body = self._read_json_body() - if body is None: - return - api_key = body.get("api_key") - if not api_key or not isinstance(api_key, str): - self._send_json(400, {"error": "api_key is required"}) - return - if not accounts.revoke_api_key(account["account_id"], api_key): - self._send_json(404, {"error": "key not found on this account"}) - return - self._send_json(200, {"revoked": api_key}) - - def _handle_admin_accounts(self): - """Admin-only: all accounts with their keys and balances.""" - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - accounts = self._require_accounts() - if accounts is None: - return - account = self._session_account() - if account is None: - self._send_json(401, {"error": "login required"}) - return - if account["role"] != "admin": - self._send_json(403, {"error": "admin role required"}) - return - listing = accounts.list_accounts() - if server.billing is not None: - for entry in listing: - entry["balances"] = { - key: server.billing.get_client_balance(key) - for key in entry.get("api_keys", []) - } - self._send_json(200, {"accounts": listing}) - - def _handle_accounts_gossip(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - body = self._read_hive_authenticated_body() - if body is None: - return - if server.accounts is None: - self._send_json(200, {"applied": 0}) - return - events = body.get("events") - if not isinstance(events, list): - self._send_json(400, {"error": "events must be a list"}) - return - applied = server.accounts.apply_events([e for e in events if isinstance(e, dict)]) - self._send_json(200, {"applied": applied}) - - def _handle_registry_gossip(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - body = self._read_hive_authenticated_body() - if body is None: - return - registry_log = getattr(server.contracts, "registry_log", None) - if registry_log is None: - self._send_json(200, {"applied": 0}) - return - events = body.get("events") - if not isinstance(events, list): - self._send_json(400, {"error": "events must be a list"}) - return - applied = registry_log.apply_events([e for e in events if isinstance(e, dict)]) - self._send_json(200, {"applied": applied}) - - def _handle_wallet_register(self): - """Bind a client wallet pubkey to an API key (US-032, C6). - - Deposits from that wallet into the treasury are then credited to the - API key's ledger balance by the deposit watcher. - - Requires a signature proving ownership of the wallet's private key — - the caller signs ``binding_message(api_key, wallet)`` with the - wallet's ed25519 key. An admin session may force-rebind a wallet - already bound to a different API key (documented override; no - signed-release flow is implemented). - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if server.billing is None: - self._send_json(404, {"error": "billing is not enabled on this tracker"}) - return - role, _account = self._resolve_identity() - body = self._read_json_body() - if body is None: - return - wallet = body.get("wallet") - if not wallet or not isinstance(wallet, str): - self._send_json(400, {"error": "wallet is required"}) - return - - if role == "admin": - api_key = body.get("api_key") - if not api_key or not isinstance(api_key, str): - self._send_json(400, {"error": "api_key is required for admin override"}) - return - event = server.billing.bind_wallet(api_key, wallet, admin_override=True) - else: - api_key = _api_key_from_headers(self.headers) - if not api_key: - self._send_json(401, {"error": "Authorization header required"}) - return - signature = body.get("signature") - if not signature or not isinstance(signature, str): - self._send_json(400, {"error": "signature is required to prove wallet ownership"}) - return - if not verify_wallet_signature(wallet, binding_message(api_key, wallet), signature): - self._send_json(401, {"error": "invalid wallet signature"}) - return - event = server.billing.bind_wallet(api_key, wallet) - - if event.get("rejected"): - self._send_json(409, {"error": "wallet already bound to a different API key"}) - return - print(f"[tracker] wallet bound: {wallet} -> api_key …{api_key[-6:]}", flush=True) - self._send_json(200, {"wallet": wallet, "bound": True}) - - def _handle_benchmark_hop_penalty(self): - """Privileged: run the same prompt through 1/2/3-node pinned routes (US-030). - - Data collection only — the routing algorithm is unchanged. Per-hop - latency is derived incrementally: the k-node route's final hop penalty - is its total minus the (k-1)-node route's total. - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if not self._require_role("admin", "validator"): - return - auth = self.headers.get("Authorization") - body = self._read_json_body() - if body is None: - return - model = body.get("model", "") - if not model: - self._send_json(400, {"error": "model is required"}) - return - prompt = body.get("prompt") or "Benchmark: say OK." - max_new_tokens = int(body.get("max_new_tokens", 64)) - - with server.lock: - self._purge_expired_nodes() - resolved = _nodes_and_bounds_for_model(server, model) - if resolved is None or not resolved[0]: - self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) - return - all_nodes, rs, re = resolved - - self_url = f"http://127.0.0.1:{self.server.server_address[1]}" - route_results: list[dict] = [] - prev_total_ms: float | None = None - prev_per_hop: list[float] = [] - for hop_count in (1, 2, 3): - combo = _find_pinned_route(all_nodes, rs, re, hop_count) - if combo is None: - continue # insufficient coverage for this hop count — skip - request_body = json.dumps({ - "model": model, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": max_new_tokens, - "route": [n.node_id for n in combo], - }).encode() - req = urllib.request.Request( - f"{self_url}/v1/chat/completions", - data=request_body, - headers={"Content-Type": "application/json", "Authorization": auth}, - method="POST", - ) - started = time.monotonic() - try: - with urllib.request.urlopen(req, timeout=300.0) as resp: - payload = json.loads(resp.read()) - except Exception as exc: - route_results.append({ - "route": [n.node_id for n in combo], - "error": str(exc), - }) - continue - total_ms = (time.monotonic() - started) * 1000.0 - if prev_total_ms is not None and len(prev_per_hop) == hop_count - 1: - per_hop_ms = prev_per_hop + [max(0.0, total_ms - prev_total_ms)] - else: - per_hop_ms = [total_ms / hop_count] * hop_count - prev_total_ms = total_ms - prev_per_hop = per_hop_ms - route_results.append({ - "route": [n.node_id for n in combo], - "total_ms": total_ms, - "per_hop_ms": per_hop_ms, - "tokens_generated": _usage_total_tokens(payload) or 0, - }) - - record = { - "timestamp": time.time(), - "model": model, - "prompt_hash": hashlib.sha256(prompt.encode()).hexdigest()[:16], - "routes": route_results, - } - with server.benchmark_lock: - existing: list = [] - if os.path.exists(server.benchmark_results_path): - try: - with open(server.benchmark_results_path, encoding="utf-8") as fh: - existing = json.load(fh) - except (json.JSONDecodeError, OSError): - existing = [] - if not isinstance(existing, list): - existing = [] - existing.append(record) - with open(server.benchmark_results_path, "w", encoding="utf-8") as fh: - json.dump(existing, fh, indent=2) - self._send_json(200, record) - - def _handle_benchmark_results(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if not self._require_role("admin", "validator"): - return - results: list = [] - with server.benchmark_lock: - if os.path.exists(server.benchmark_results_path): - try: - with open(server.benchmark_results_path, encoding="utf-8") as fh: - results = json.load(fh) - except (json.JSONDecodeError, OSError): - results = [] - self._send_json(200, {"results": results if isinstance(results, list) else []}) - - def _handle_toploc_calibration_run(self): - """Privileged: honest-noise TOPLOC calibration dispatch (issue 21). - - Fans the same fixed prompt through every currently registered node - that can solo-serve the full model (one pinned-route hop, mirroring - `_handle_benchmark_hop_penalty`'s dispatch pattern), then verifies - each node's own on-demand TOPLOC commitment against a teacher-forced - replay on the reference node — same audit contract the validator - uses (`packages/validator/README.md` "TOPLOC audit contract"). Each - node's raw divergence (not just pass/fail) is recorded into the - calibration corpus, keyed by wallet + GPU model + dtype, so - thresholds can eventually be derived instead of guessed. - - Nodes that only hold a partial shard (need a multi-hop route) are - skipped for this pass — solo dispatch isolates one node's hardware - noise without a route composition confound — and nodes that don't - answer the on-demand commitment fetch (endpoint down, or node-side - TOPLOC serving not yet wired) are skipped and reported, not treated - as a pass or a fail. - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if not self._require_role("admin", "validator"): - return - if server.toploc_calibration is None: - self._send_json(503, {"error": "toploc calibration store is not enabled on this tracker"}) - return - if not server.toploc_reference_node_url: - self._send_json(503, {"error": "toploc_reference_node_url is not configured on this tracker"}) - return - auth = self.headers.get("Authorization") - body = self._read_json_body() - if body is None: - return - model = body.get("model", "") - if not model: - self._send_json(400, {"error": "model is required"}) - return - prompt = body.get("prompt") or "Calibration: say OK." - max_new_tokens = int(body.get("max_new_tokens", 32)) - seed = body.get("seed", 0) - - with server.lock: - self._purge_expired_nodes() - resolved = _nodes_and_bounds_for_model(server, model) - if resolved is None or not resolved[0]: - self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) - return - all_nodes, rs, re = resolved - if server.contracts is not None: - all_nodes = [ - node for node in all_nodes - if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned - ] - solo_nodes = [ - node for node in all_nodes - if node.shard_start is not None and node.shard_end is not None - and node.shard_start <= rs and node.shard_end >= re - ] - - self_url = f"http://127.0.0.1:{self.server.server_address[1]}" - messages = [{"role": "user", "content": prompt}] - node_results: list[dict] = [] - for node in solo_nodes: - request_id = f"cal-{uuid.uuid4().hex}" - request_body = json.dumps({ - "id": request_id, - "model": model, - "messages": messages, - "max_tokens": max_new_tokens, - "temperature": 0, - "seed": seed, - "route": [node.node_id], - }).encode() - req = urllib.request.Request( - f"{self_url}/v1/chat/completions", - data=request_body, - headers={"Content-Type": "application/json", "Authorization": auth}, - method="POST", - ) - try: - with urllib.request.urlopen(req, timeout=300.0) as resp: - json.loads(resp.read()) - except Exception as exc: - node_results.append({"node_id": node.node_id, "wallet_address": node.wallet_address, "error": str(exc)}) - continue - - outcome = self._verify_node_toploc_calibration( - server, node, request_id=request_id, model=model, messages=messages, - ) - node_results.append(outcome) - - skipped_partial_shard = [ - node.node_id for node in all_nodes if node not in solo_nodes - ] - record = { - "timestamp": time.time(), - "model": model, - "prompt_hash": hashlib.sha256(prompt.encode()).hexdigest()[:16], - "nodes": node_results, - "skipped_partial_shard_node_ids": skipped_partial_shard, - "gate_status": server.toploc_calibration.gate_status( - min_hardware_profiles=self._toploc_calibration_gate_min_hardware_profiles(), - ), - } - self._send_json(200, record) - - def _toploc_calibration_gate_min_hardware_profiles(self) -> int: - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - return server.toploc_calibration_gate_min_hardware_profiles - - def _verify_node_toploc_calibration( - self, - server: "_TrackerHTTPServer", - node: "_NodeEntry", - *, - request_id: str, - model: str, - messages: list[dict], - ) -> dict: - """One node's calibration outcome: fetch its on-demand commitment, - teacher-force the claimed tokens on the reference node, verify, and - persist the raw divergence into the corpus.""" - from meshnet_validator.audit import ToplocProofClaim, verify_activation_proofs_detailed - - gpu_model = (node.hardware_profile or {}).get("gpu_name") or (node.hardware_profile or {}).get("device") or "unknown" - dtype = node.quantization or "unknown" - base_result = { - "node_id": node.node_id, - "wallet_address": node.wallet_address, - "gpu_model": gpu_model, - "dtype": dtype, - } - commitment = _fetch_toploc_commitment( - node, session_id=request_id, model=model, messages=messages, - ) - if commitment is None: - return {**base_result, "skipped": "no on-demand toploc commitment available"} - - try: - claim = ToplocProofClaim.from_mapping(commitment["toploc_proof"]) - except (KeyError, TypeError, ValueError): - return {**base_result, "skipped": "malformed toploc commitment"} - - reference_activations = _fetch_toploc_reference_activations( - server.toploc_reference_node_url, - model=model, - messages=messages, - claimed_token_ids=commitment["claimed_token_ids"], - claim=claim, - ) - if reference_activations is None: - return {**base_result, "skipped": "reference node teacher-forced replay failed"} - - result = verify_activation_proofs_detailed(reference_activations, claim, backend=server.toploc_backend) - if node.wallet_address: - server.toploc_calibration.record_run( - node_wallet=node.wallet_address, - gpu_model=gpu_model, - dtype=dtype, - model=model, - passed=result.passed, - exp_intersections=result.exp_intersections, - mant_err_mean=result.mant_err_mean, - mant_err_median=result.mant_err_median, - ) - return { - **base_result, - "passed": result.passed, - "exp_intersections": result.exp_intersections, - "mant_err_mean": result.mant_err_mean, - "mant_err_median": result.mant_err_median, - "chunk_count": result.chunk_count, - } - - def _handle_toploc_calibration_results(self): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if not self._require_role("admin", "validator"): - return - if server.toploc_calibration is None: - self._send_json(503, {"error": "toploc calibration store is not enabled on this tracker"}) - return - min_profiles = self._toploc_calibration_gate_min_hardware_profiles() - self._send_json(200, { - "runs": server.toploc_calibration.runs(), - "envelope": server.toploc_calibration.envelope(), - "gate_status": server.toploc_calibration.gate_status(min_hardware_profiles=min_profiles), - }) - - def _handle_hf_pricing_history(self, parsed: urllib.parse.ParseResult): - """Dispute-auditability log for the dynamic HF-benchmarked pricing (issue 23).""" - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if not self._require_role("admin", "validator"): - return - if server.hf_pricing_log is None: - self._send_json(503, {"error": "hf pricing log is not enabled on this tracker"}) - return - params = urllib.parse.parse_qs(parsed.query) - model = params.get("model", [None])[0] - self._send_json(200, {"changes": server.hf_pricing_log.history(model)}) - - def _handle_assign(self, parsed: urllib.parse.ParseResult): - """Return an optimal shard assignment for a node given its hardware profile. - - Query params: - model — model preset name (default: first preset) - device — "cuda" | "cpu" - vram_mb — integer VRAM in MB (0 for CPU) - ram_mb — integer system RAM in MB, used when vram_mb=0 - - The greedy strategy: find the first gap in current layer coverage - and assign it. If no gap exists, assign the full model range so the - node provides redundant coverage. - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - params = urllib.parse.parse_qs(parsed.query) - - model_list = params.get("model") - if not model_list: - model = next(iter(server.model_presets), None) - if model is None: - self._send_json(503, {"error": "no model presets configured"}) - return - else: - model = model_list[0] - - resolved_name, preset = _resolve_model_preset(server.model_presets, model) - if preset is None: - self._send_json(404, {"error": f"unknown model preset: {model!r}"}) - return - - required_start, required_end = _preset_layer_bounds(preset) - - with server.lock: - self._purge_expired_nodes() - alive = [ - node for node in server.registry.values() - if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] - ] - if server.contracts is not None: - alive = [ - node for node in alive - if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned - ] - - device = params.get("device", ["cpu"])[0] - try: - vram_mb = int(params.get("vram_mb", ["0"])[0]) - except ValueError: - vram_mb = 0 - try: - ram_mb = int(params.get("ram_mb", ["0"])[0]) - except ValueError: - ram_mb = 0 - max_layers = required_end - required_start + 1 - # VRAM only bounds the shard for CUDA nodes; a CPU node may still report - # a detected-but-unusable GPU, and must be sized by system RAM. - memory_mb = vram_mb if (device == "cuda" and vram_mb > 0) else ram_mb - if memory_mb > 0: - layer_bytes = _preset_bytes_per_layer(preset).get("bfloat16", 30 * 1024 * 1024) - max_layers = min(max_layers, max(1, int(((memory_mb * 1024 * 1024) * 0.8) // layer_bytes))) - elif device != "cuda" or vram_mb < 8192: - max_layers = min(max_layers, 16) - - # Collect covered intervals sorted by start layer. - covered = sorted( - [ - (n.shard_start, n.shard_end) - for n in alive - if n.shard_start is not None and n.shard_end is not None - ], - key=lambda t: t[0], - ) - - # Walk from required_start to find the first uncovered layer. - gap_start = required_start - for s, e in covered: - if s <= gap_start: - gap_start = max(gap_start, e + 1) - else: - break # gap found before this node - - if gap_start > required_end: - # Full coverage exists — assign the full range for redundancy. - shard_start = required_start - shard_end = min(required_end, shard_start + max_layers - 1) - else: - shard_start = gap_start - shard_end = min(required_end, shard_start + max_layers - 1) - - peers = [ - {"endpoint": node.endpoint, "checksum": node.shard_checksum} - for node in alive - if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] - and node.shard_start == shard_start - and node.shard_end == shard_end - and node.shard_checksum - ] - - self._send_json(200, { - "shard_start": shard_start, - "shard_end": shard_end, - "model": resolved_name, - "model_layers_end": required_end, - "peers": peers, - "bytes_per_layer": _preset_bytes_per_layer(preset), - "model_sources": self._model_sources( - resolved_name, - preset, - shard_start, - shard_end, - ), - **({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}), - }) - - def _model_sources(self, model: str, preset: dict, shard_start: int, shard_end: int) -> list[dict]: - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - hf_repo = preset.get("hf_repo") - if not server.models_dir or not isinstance(hf_repo, str) or not hf_repo: - return [] - snapshot_dir = snapshot_dir_for_repo(server.models_dir, hf_repo) - if snapshot_dir is None: - return [] - files = files_for_layer_range(snapshot_dir, shard_start, shard_end) - if not files: - return [] - host = self.headers.get("Host") or f"{self.server.server_address[0]}:{self.server.server_address[1]}" - base_url = f"http://{host}" - query = urllib.parse.urlencode({ - "model": model, - "shard_start": shard_start, - "shard_end": shard_end, - }) - full_query = urllib.parse.urlencode({"model": model, "full": 1}) - full_files = _snapshot_regular_files(snapshot_dir) - file_sizes: dict[str, int] = {} - for rel in set(files) | set(full_files): - try: - file_sizes[rel] = (snapshot_dir / rel).resolve().stat().st_size - except OSError: - continue - return [{ - "type": "tracker", - "endpoint": base_url, - "url": f"{base_url}/v1/model-files/download?{query}", - "full_url": f"{base_url}/v1/model-files/download?{full_query}", - "files": files, - "full_files": full_files, - "file_sizes": file_sizes, - }] - - def _handle_model_files_download(self, parsed: urllib.parse.ParseResult) -> None: - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - if server.models_dir is None: - self._send_json(404, {"error": "tracker model-file source is not enabled"}) - return - params = urllib.parse.parse_qs(parsed.query) - model = params.get("model", [""])[0] - resolved_name, preset = _resolve_model_preset(server.model_presets, model) - if preset is None: - self._send_json(404, {"error": f"unknown model preset: {model!r}"}) - return - hf_repo = preset.get("hf_repo") - if not isinstance(hf_repo, str) or not hf_repo: - self._send_json(404, {"error": f"model preset has no hf_repo: {resolved_name!r}"}) - return - snapshot_dir = snapshot_dir_for_repo(server.models_dir, hf_repo) - if snapshot_dir is None: - self._send_json(404, {"error": f"local snapshot not found for {hf_repo}"}) - return - full_download = params.get("full", ["0"])[0] in {"1", "true", "yes"} - if full_download: - rel_files = _snapshot_regular_files(snapshot_dir) - else: - try: - shard_start = int(params.get("shard_start", [""])[0]) - shard_end = int(params.get("shard_end", [""])[0]) - except ValueError: - self._send_json(400, {"error": "shard_start and shard_end must be integers"}) - return - rel_files = files_for_layer_range(snapshot_dir, shard_start, shard_end) - if not rel_files: - self._send_json(404, {"error": "no local files matched the assigned shard"}) - return - - # Single-file mode: ?file= streams one file with Content-Length so - # clients can verify completeness and retry per file instead of - # restarting one giant tar stream after a network hiccup. - single = params.get("file", [""])[0] - if single: - if single not in set(rel_files): - self._send_json(404, {"error": f"file not in requested shard: {single!r}"}) - return - file_path = snapshot_dir / single - try: - # resolve() dereferences HF snapshot symlinks into blobs/. - size = file_path.resolve().stat().st_size - except OSError: - self._send_json(404, {"error": f"file unreadable: {single!r}"}) - return - self.send_response(200) - self.send_header("Content-Type", "application/octet-stream") - self.send_header("Content-Length", str(size)) - self.send_header("X-Meshnet-Model-Source", "tracker") - self.end_headers() - try: - with file_path.open("rb") as f: - while True: - chunk = f.read(1024 * 1024) - if not chunk: - break - self.wfile.write(chunk) - except (BrokenPipeError, ConnectionResetError): - print( - f"model-file download aborted by client " - f"({self.client_address[0]}, model={resolved_name}, file={single})", - flush=True, - ) - return - - self.send_response(200) - self.send_header("Content-Type", "application/x-tar") - self.send_header("X-Meshnet-Model-Source", "tracker") - self.end_headers() - try: - # dereference: HF cache snapshots are symlinks into blobs/ — ship contents. - with tarfile.open(fileobj=self.wfile, mode="w|", dereference=True) as archive: - for rel in rel_files: - archive.add(snapshot_dir / rel, arcname=rel) - except (BrokenPipeError, ConnectionResetError): - print( - f"model-file download aborted by client " - f"({self.client_address[0]}, model={resolved_name})", - flush=True, - ) - - def _handle_network_assign(self, parsed: urllib.parse.ParseResult): - """Assign a new node to fill the biggest uncovered shard gap across HF-model nodes. - - Query params: - vram_mb — integer VRAM in MB (0 = CPU-only node) - ram_mb — integer system RAM in MB, used when vram_mb=0 - device — "cuda" | "cpu" - hf_repo — optional; if set, restrict search to this repo only - - Returns: - {hf_repo, shard_start, shard_end, num_layers, gap_found} - gap_found=true means a real uncovered gap was assigned; false means redundancy. - """ - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - params = urllib.parse.parse_qs(parsed.query) - try: - vram_mb = int(params.get("vram_mb", ["0"])[0]) - except ValueError: - vram_mb = 0 - try: - ram_mb = int(params.get("ram_mb", ["0"])[0]) - except ValueError: - ram_mb = 0 - device = params.get("device", ["cpu"])[0] - filter_repo = params.get("hf_repo", [None])[0] # optional repo filter - - with server.lock: - self._purge_expired_nodes() - all_nodes = list(server.registry.values()) - - # Collect only nodes that registered a real HF model (have hf_repo + shard bounds). - hf_nodes = [ - n for n in all_nodes - if n.hf_repo - and n.shard_start is not None - and n.shard_end is not None - and n.num_layers is not None - and (filter_repo is None or n.hf_repo == filter_repo) - ] - - if not hf_nodes: - # The caller is not registered yet — count its memory toward - # deployability, or the first node can never bootstrap the network. - caller = _NodeEntry( - node_id="assign-candidate", - endpoint="", - shard_start=None, - shard_end=None, - model=None, - shard_checksum=None, - hardware_profile={}, - wallet_address=None, - score=1.0, - vram_bytes=vram_mb * 1024 * 1024, - ram_bytes=ram_mb * 1024 * 1024, - ) - pool_with_caller = all_nodes + [caller] - resolved_name = None - preset = None - if filter_repo: - resolved_name, preset = _resolve_model_preset(server.model_presets, filter_repo) - else: - deployable = [ - (name, preset) - for name, preset in server.model_presets.items() - if preset.get("recommended") and _deployment_summary(pool_with_caller, preset)["deployable"] - ] - if deployable: - resolved_name, preset = deployable[0] - if preset is not None and preset.get("hf_repo"): - required_start, required_end = _preset_layer_bounds(preset) - total_l = required_end - required_start + 1 - memory_mb = vram_mb if (device == "cuda" and vram_mb > 0) else ram_mb - max_layers = _max_layers_for_memory(memory_mb, total_l, preset) - shard_start = required_start - shard_end = min(required_end, shard_start + max_layers - 1) - self._send_json(200, { - "hf_repo": preset["hf_repo"], - "model": resolved_name, - "shard_start": shard_start, - "shard_end": shard_end, - "num_layers": total_l, - "gap_found": True, - "price_per_token": 0.0, - "bytes_per_layer": _preset_bytes_per_layer(preset), - "deployment": _deployment_summary(pool_with_caller, preset), - "model_sources": self._model_sources( - str(resolved_name), - preset, - shard_start, - shard_end, - ), - }) - return - - msg = ( - f"no HF-model nodes registered for {filter_repo!r}" - if filter_repo - else "no HF-model nodes registered; cannot assign shards" - ) - self._send_json(503, {"error": msg}) - return - - # Group by hf_repo; pick the one with the largest total_layers and biggest gap. - from collections import defaultdict - repo_groups: dict = defaultdict(list) - repo_layers: dict = {} - for n in hf_nodes: - repo_groups[n.hf_repo].append(n) - # Use the largest num_layers seen for this repo. - if n.hf_repo not in repo_layers or n.num_layers > repo_layers[n.hf_repo]: - repo_layers[n.hf_repo] = n.num_layers - - # Smart scoring: demand_rpm × coverage_deficit - # coverage_deficit = uncovered_layers / total_layers (0 = fully covered) - # demand_rpm comes from the stats collector; defaults to 0.0 when no traffic yet. - demand_rpms: dict[str, float] = {} - if server.stats is not None: - local_rpms = server.stats.get_local_rpms() - for repo in repo_groups: - demand_rpms[repo] = _repo_demand_rpm(server, repo, local_rpms) - - best_repo = None - best_score = -1.0 - best_gap_size = -1 - best_gap_start = 0 - best_num_layers = 0 - - for repo, nodes in repo_groups.items(): - total = repo_layers[repo] - served_copies = _served_model_copies(nodes, 0, total - 1) - covered = sorted( - [(n.shard_start, n.shard_end) for n in nodes], - key=lambda t: t[0], - ) - # Walk from 0 to find first uncovered layer. - gap_start = 0 - for s, e in covered: - if s <= gap_start: - gap_start = max(gap_start, e + 1) - else: - break - gap_size = max(0, (total - 1) - gap_start + 1) - coverage_deficit = gap_size / max(total, 1) - demand = demand_rpms.get(repo, 0.0) - # Prefer demanded models already being served; coverage gaps still dominate. - score = (demand + 1.0) * (coverage_deficit + 0.01) + served_copies * 0.01 - if score > best_score or (score == best_score and gap_size > best_gap_size): - best_score = score - best_gap_size = gap_size - best_gap_start = gap_start - best_repo = repo - best_num_layers = total - - gap_found = best_gap_size > 0 - if not gap_found: - # All shards covered — add redundancy on the most demanded/served model. - best_repo = max( - repo_groups, - key=lambda repo: ( - demand_rpms.get(repo, 0.0), - _served_model_copies(repo_groups[repo], 0, repo_layers[repo] - 1), - len(repo_groups[repo]), - ), - ) - best_gap_start = 0 - best_num_layers = repo_layers[best_repo] - - # Capacity: use the same 80%-of-memory rule as registered node planning. - total_l = best_num_layers - memory_mb = vram_mb if (device == "cuda" and vram_mb > 0) else ram_mb - resolved_name, best_preset = _resolve_model_preset(server.model_presets, str(best_repo)) - if memory_mb > 0: - max_layers = _max_layers_for_memory(memory_mb, total_l, best_preset) - elif device == "cuda" and vram_mb >= 8192: - max_layers = total_l - else: - max_layers = _max_layers_for_memory(memory_mb, total_l, best_preset) - - shard_start = best_gap_start - shard_end = min(total_l - 1, shard_start + max_layers - 1) - - self._send_json(200, { - "hf_repo": best_repo, - "model": resolved_name, - "shard_start": shard_start, - "shard_end": shard_end, - "num_layers": total_l, - "gap_found": gap_found, - "price_per_token": 0.0, - "bytes_per_layer": _preset_bytes_per_layer(best_preset) if best_preset is not None else None, - "model_sources": self._model_sources( - str(resolved_name or best_repo), - best_preset or {"hf_repo": best_repo}, - shard_start, - shard_end, - ), - }) - - def _handle_route(self, parsed: urllib.parse.ParseResult): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - params = urllib.parse.parse_qs(parsed.query) - model_list = params.get("model") - if not model_list: - self._send_json(400, {"error": "missing 'model' query parameter"}) - return - - model = model_list[0] - resolved_name, preset = _resolve_model_preset(server.model_presets, model) - - with server.lock: - self._purge_expired_nodes() - if preset is not None: - # Preset-based routing (stub-model system). - alive = [ - node for node in server.registry.values() - if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] - ] - required_start, required_end = _preset_layer_bounds(preset) - else: - # HF model routing: match by hf_repo (full) or model short name. - alive = [ - node for node in server.registry.values() - if _node_matches_model(node, model) - and node.shard_start is not None - and node.shard_end is not None - and node.num_layers is not None - ] - if not alive: - self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) - return - required_start = 0 - required_end = max(n.num_layers for n in alive) - 1 # type: ignore[type-var] - - if server.contracts is not None: - alive = [ - node for node in alive - if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned - ] - - route, error = _select_route(alive, required_start, required_end, contracts=server.contracts) - if error: - _tracker_log( - server, - "warn", - "route unavailable", - model=model, - route_model=resolved_name or model, - error=error, - candidate_count=len(alive), - candidates=_node_route_summary(alive), - ) - self._send_json(503, {"error": error}) - return - - covered_up_to = required_start - 1 - route_with_start: list[tuple] = [] - for rn in route: - route_with_start.append((rn, covered_up_to + 1)) - covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to - - self._send_json(200, { - "route": [e.endpoint for e, _ in route_with_start], - "nodes": [ - { - "node_id": e.node_id, - "endpoint": e.endpoint, - "start_layer": start, - "relay_addr": e.relay_addr, - "peer_id": e.peer_id, - "wallet_address": e.wallet_address, - "shard_start": e.shard_start, - "shard_end": e.shard_end, - "model": e.model, - "hf_repo": e.hf_repo, - "num_layers": e.num_layers, - "model_metadata": dict(e.model_metadata), - "downloaded_models": [dict(item) for item in e.downloaded_models], - "shard_checksum": e.shard_checksum, - "score": e.score, - } - for e, start in route_with_start - ], - }) - - def _handle_routes(self, parsed: urllib.parse.ParseResult): - server: _TrackerHTTPServer = self.server # type: ignore[assignment] - params = urllib.parse.parse_qs(parsed.query) - model_list = params.get("model") - if not model_list: - self._send_json(400, {"error": "missing 'model' query parameter"}) - return - - try: - redundancy = int(params.get("redundancy", ["1"])[0]) - except ValueError: - self._send_json(400, {"error": "redundancy must be an integer"}) - return - if redundancy < 1: - self._send_json(400, {"error": "redundancy must be at least 1"}) - return - - model = model_list[0] - resolved_name, preset = _resolve_model_preset(server.model_presets, model) - if preset is None: - self._send_json(404, {"error": f"unknown model preset: {model!r}"}) - return - - required_start, required_end = _preset_layer_bounds(preset) - - with server.lock: - self._purge_expired_nodes() - candidates = [ - node for node in server.registry.values() - if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] - ] - if server.contracts is not None: - candidates = [ - node for node in candidates - if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned - ] - - routes = [] - remaining = list(candidates) - for _ in range(redundancy): - route, error = _select_route(remaining, required_start, required_end, contracts=server.contracts) - if error: - self._send_json(503, {"error": error}) - return - route_endpoints = {node.endpoint for node in route} - remaining = [node for node in remaining if node.endpoint not in route_endpoints] - routes.append({ - "route": [e.endpoint for e in route], - "nodes": [ - { - "node_id": e.node_id, - "endpoint": e.endpoint, - "relay_addr": e.relay_addr, - "peer_id": e.peer_id, - "wallet_address": e.wallet_address, - "shard_start": e.shard_start, - "shard_end": e.shard_end, - "model": e.model, - "downloaded_models": [dict(item) for item in e.downloaded_models], - "shard_checksum": e.shard_checksum, - "score": e.score, - } - for e in route - ], - }) - - self._send_json(200, {"routes": routes}) - - -class TrackerServer: - """HTTP tracker that manages node registration and resolves inference routes. - - Nodes register via POST /v1/nodes/register and keep themselves alive with - POST /v1/nodes//heartbeat. The gateway queries GET /v1/route?model= - to obtain an ordered list of node endpoints whose shards cover all layers. - """ - - def __init__( - self, - host: str = "127.0.0.1", - port: int = 0, - heartbeat_timeout: float = 90.0, - rebalance_interval: float = 30.0, - model_presets: dict | None = None, - contracts: Any | None = None, - minimum_stake: int = 0, - cluster_peers: list[str] | None = None, - cluster_self_url: str | None = None, - stats_db: str | None = None, - relay_url: str | None = None, - billing: BillingLedger | None = None, - enable_billing: bool = False, - billing_db: str | None = None, - accounts: AccountStore | None = None, - accounts_db: str | None = None, - benchmark_results_path: str | None = None, - treasury: Any | None = None, - deposit_poll_interval: float = 15.0, - settle_period: float = 86400.0, - payout_threshold: float = 5.0, - payout_dust_floor: float = 0.01, - settlement_check_interval: float | None = None, - validator_service_token: str | None = None, - hive_secret: str | None = None, - max_charge_per_request: float | None = None, - starting_credit: float = DEFAULT_CALLER_CREDIT_USDT, - devnet_topup_amount: float = DEFAULT_DEVNET_TOPUP_USDT, - toploc_calibration: ToplocCalibrationStore | None = None, - toploc_calibration_db: str | None = None, - toploc_reference_node_url: str | None = None, - toploc_calibration_gate_min_hardware_profiles: int = 1, - toploc_backend: Any | None = None, - enable_hf_pricing: bool = False, - hf_pricing_log: HfPricingLog | None = None, - hf_pricing_log_db: str | None = None, - hf_pricing_refresh_interval: float = 86400.0, - hf_pricing_fetch_html: Any | None = None, - models_dir: str | Path | None = None, - ) -> None: - self._host = host - self._requested_port = port - self._heartbeat_timeout = heartbeat_timeout - self._rebalance_interval = rebalance_interval - self._model_presets: dict = ( - model_presets if model_presets is not None else _clone_model_presets(DEFAULT_MODEL_PRESETS) - ) - self._contracts = contracts - self._minimum_stake = minimum_stake - self._cluster_peers: list[str] = list(cluster_peers) if cluster_peers else [] - self._cluster_self_url: str | None = cluster_self_url - self._relay_url = relay_url - self._registry: dict[str, _NodeEntry] = {} - self._lock = threading.Lock() - self._server: _TrackerHTTPServer | None = None - self._thread: threading.Thread | None = None - self._rebalance_stop = threading.Event() - self._rebalance_thread: threading.Thread | None = None - self._raft: RaftNode | None = None - self._gossip: NodeGossip | None = None - self._stats: _StatsCollector | None = _StatsCollector(db_path=stats_db) if stats_db else _StatsCollector() - self._stats_stop = threading.Event() - self._stats_thread: threading.Thread | None = None - if billing is None: - db_path = billing_db - if db_path is None and enable_billing: - db_path = DEFAULT_BILLING_DB_PATH - if db_path: - preset_prices: dict[str, tuple[float, float]] = {} - for name, preset in self._model_presets.items(): - if not isinstance(preset, dict): - continue - base = preset.get("price_per_1k_tokens") - input_rate = preset.get("input_price_per_1k_tokens", base) - output_rate = preset.get("output_price_per_1k_tokens", base) - if input_rate is None or output_rate is None: - continue - for key in _preset_price_keys(name, preset): - preset_prices[key] = (float(input_rate), float(output_rate)) - billing = BillingLedger(db_path=db_path, prices=preset_prices) - self._billing: BillingLedger | None = billing - self._billing_gossip_cursor = 0 - if accounts is None: - accounts_path = accounts_db - if accounts_path is None and enable_billing: - accounts_path = DEFAULT_ACCOUNTS_DB_PATH - if accounts_path: - accounts = AccountStore(db_path=accounts_path) - self._accounts: AccountStore | None = accounts - self._accounts_gossip_cursor = 0 - self._registry_gossip_cursor = 0 - self._benchmark_results_path = benchmark_results_path - self._treasury = treasury - self._deposit_poll_interval = deposit_poll_interval - self._deposit_stop = threading.Event() - self._deposit_thread: threading.Thread | None = None - self._settle_period = settle_period - self._payout_threshold = payout_threshold - self._payout_dust_floor = payout_dust_floor - self._settlement_check_interval = settlement_check_interval or max( - 1.0, min(settle_period / 4.0, 15.0) - ) - self._settlement_stop = threading.Event() - self._settlement_thread: threading.Thread | None = None - if max_charge_per_request is not None and max_charge_per_request <= 0.0: - raise ValueError("max_charge_per_request must be positive") - self._max_charge_per_request = max_charge_per_request - self._starting_credit = max(0.0, starting_credit) - self._devnet_topup_amount = max(0.0, devnet_topup_amount) - self._validator_service_token = ( - validator_service_token - if validator_service_token is not None - else os.environ.get("MESHNET_VALIDATOR_SERVICE_TOKEN") or None - ) - self._hive_secret = ( - hive_secret - if hive_secret is not None - else os.environ.get("MESHNET_HIVE_SECRET") or None - ) - if toploc_calibration is None and toploc_calibration_db: - toploc_calibration = ToplocCalibrationStore(db_path=toploc_calibration_db) - self._toploc_calibration: ToplocCalibrationStore | None = toploc_calibration - self._toploc_reference_node_url = toploc_reference_node_url - self._toploc_calibration_gate_min_hardware_profiles = toploc_calibration_gate_min_hardware_profiles - self._toploc_backend = toploc_backend - if hf_pricing_log is None and (enable_hf_pricing or hf_pricing_log_db): - hf_pricing_log = HfPricingLog(db_path=hf_pricing_log_db or DEFAULT_HF_PRICING_LOG_DB_PATH) - self._hf_pricing_log: HfPricingLog | None = hf_pricing_log - self._enable_hf_pricing = enable_hf_pricing - self._hf_pricing_refresh_interval = hf_pricing_refresh_interval - self._hf_pricing_fetch_html = hf_pricing_fetch_html - # MESHNET_DOWNLOAD_DIR is the node-side model store; on a box running - # both tracker and node it doubles as the tracker's snapshot source. - raw_models_dir = ( - models_dir - if models_dir is not None - else os.environ.get("MESHNET_MODELS_DIR") or os.environ.get("MESHNET_DOWNLOAD_DIR") - ) - self._models_dir = Path(raw_models_dir).expanduser() if raw_models_dir else None - self._hf_pricing_stop = threading.Event() - self._hf_pricing_thread: threading.Thread | None = None - self.port: int | None = None - - def start(self) -> int: - if self._server is not None: - raise RuntimeError("TrackerServer is already running") - - effective_relay_url = ( - self._relay_url - or derive_relay_url_from_public_tracker_url(self._cluster_self_url) - ) - - # Start HTTP server first so we know our port - self._server = _TrackerHTTPServer( - (self._host, self._requested_port), - _TrackerHandler, - self._registry, - self._lock, - self._heartbeat_timeout, - self._model_presets, - self._contracts, - self._minimum_stake, - relay_url=effective_relay_url, - stats=self._stats, - billing=self._billing, - accounts=self._accounts, - benchmark_results_path=self._benchmark_results_path, - validator_service_token=self._validator_service_token, - hive_secret=self._hive_secret, - max_charge_per_request=self._max_charge_per_request, - starting_credit=self._starting_credit, - devnet_topup_amount=self._devnet_topup_amount, - toploc_calibration=self._toploc_calibration, - toploc_reference_node_url=self._toploc_reference_node_url, - toploc_calibration_gate_min_hardware_profiles=self._toploc_calibration_gate_min_hardware_profiles, - toploc_backend=self._toploc_backend, - hf_pricing_log=self._hf_pricing_log, - models_dir=self._models_dir, - ) - self.port = self._server.server_address[1] - - # Start Raft + gossip if cluster peers are configured - if self._cluster_peers: - self_url = self._cluster_self_url or f"http://{self._host}:{self.port}" - self._raft = RaftNode( - self_url=self_url, - peers=self._cluster_peers, - apply_fn=self._raft_apply, - ) - self._gossip = NodeGossip(peers=self._cluster_peers) - self._server.raft = self._raft - self._server.gossip = self._gossip - self._raft.start() - self._gossip.start() - - self._rebalance_stop.clear() - self._stats_stop.clear() - self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) - self._thread.start() - self._rebalance_thread = threading.Thread(target=self._rebalance_loop, daemon=True) - self._rebalance_thread.start() - self._stats_thread = threading.Thread(target=self._stats_loop, daemon=True) - self._stats_thread.start() - if self._treasury is not None and self._billing is not None: - self._deposit_stop.clear() - self._deposit_thread = threading.Thread(target=self._deposit_loop, daemon=True) - self._deposit_thread.start() - self._settlement_stop.clear() - self._settlement_thread = threading.Thread(target=self._settlement_loop, daemon=True) - self._settlement_thread.start() - if self._enable_hf_pricing and self._billing is not None: - self._hf_pricing_stop.clear() - self._hf_pricing_thread = threading.Thread(target=self._hf_pricing_loop, daemon=True) - self._hf_pricing_thread.start() - return self.port - - def _settlement_loop(self) -> None: - """Leader-only on-chain settlement (US-033, ADR-0015). - - Pay a node when pending ≥ threshold OR its pending age ≥ max period, - with a dust floor. Pending is debited (payout events, replicated) - before the transaction is sent; unconfirmed batches are resent with - the same settlement id, so retries never double-pay. - """ - billing = self._billing - treasury = self._treasury - assert billing is not None and treasury is not None - while not self._settlement_stop.wait(self._settlement_check_interval): - if self._raft is not None and not self._raft.is_leader: - continue # followers replicate the ledger but never sign - # resend batches whose transaction never confirmed - for settlement_id, payouts in billing.unconfirmed_settlements().items(): - self._send_settlement(settlement_id, payouts) - banned: set[str] = set() - if self._server is not None and self._server.contracts is not None: - registry = self._server.contracts.registry - banned = { - wallet for wallet, _ in billing.payables( - threshold=0.0, max_period=0.0, dust_floor=0.0, - ) - if registry.get_wallet(wallet).banned - } - due = billing.payables( - threshold=self._payout_threshold, - max_period=self._settle_period, - dust_floor=self._payout_dust_floor, - exclude=banned, - ) - if not due: - continue - settlement_id = uuid.uuid4().hex - registry = ( - self._server.contracts.registry - if self._server is not None and self._server.contracts is not None - else None - ) - settled: list[tuple[str, float]] = [] - for wallet, amount in due: - # ADR-0015: a fraud forfeiture landing between the payables() - # snapshot above and this debit must never be paid out on top - # of — recheck the ban and let settle_node_payout clamp to - # whatever is still actually pending. - if registry is not None and registry.get_wallet(wallet).banned: - continue - event = billing.settle_node_payout(wallet, amount, reference=settlement_id) - if event["amount"] > 0: - settled.append((wallet, event["amount"])) - if settled: - self._send_settlement(settlement_id, settled) - - def _send_settlement(self, settlement_id: str, payouts: list) -> None: - billing = self._billing - treasury = self._treasury - assert billing is not None and treasury is not None - try: - signature = treasury.send_payouts([(w, a) for w, a in payouts]) - except Exception as exc: - print( - f"[tracker] settlement {settlement_id} failed (will retry): {exc}", - flush=True, - ) - return - billing.confirm_settlement(settlement_id, signature) - total = sum(a for _, a in payouts) - print( - f"[tracker] settled {settlement_id}: {len(payouts)} payout(s), " - f"{total:.6f} USDT total (tx {signature})", - flush=True, - ) - - def _deposit_loop(self) -> None: - """Poll the treasury token account and credit confirmed deposits (US-032).""" - billing = self._billing - treasury = self._treasury - assert billing is not None and treasury is not None - while not self._deposit_stop.wait(self._deposit_poll_interval): - try: - deposits = treasury.list_new_deposits( - lambda sig: billing.has_event(f"deposit-{sig}") - ) - except Exception as exc: - print(f"[tracker] deposit poll failed: {exc}", flush=True) - continue - for deposit in deposits: - api_key = billing.api_key_for_wallet(deposit.sender_wallet) - if api_key is None: - print( - f"[tracker] unattributed deposit {deposit.signature}: " - f"{deposit.amount_usdt} USDT from unbound wallet " - f"{deposit.sender_wallet} — register via /v1/wallet/register", - flush=True, - ) - continue - credited = billing.credit_deposit( - api_key, deposit.amount_usdt, deposit.signature - ) - if credited is not None: - print( - f"[tracker] deposit credited: {deposit.amount_usdt} USDT " - f"-> api_key …{api_key[-6:]} (tx {deposit.signature})", - flush=True, - ) - - def _hf_pricing_loop(self) -> None: - """Daily dynamic pricing refresh benchmarked against HF inference rates (issue 23). - - For every preset with a curated, human-verified ``hf_aliases`` list, - fetch current HF marketplace pricing and set the client price to 80% - of the cheapest matching alias. Presets with no (or an empty) - ``hf_aliases`` are left entirely alone — they keep the static - default price. Any single preset's fetch/parse failure is logged and - skipped; it never raises into this loop or the request path. - """ - billing = self._billing - assert billing is not None - while not self._hf_pricing_stop.wait(self._hf_pricing_refresh_interval): - for name, preset in list(self._model_presets.items()): - if not isinstance(preset, dict) or not preset.get("hf_aliases"): - continue - try: - result = refresh_preset_price( - model_name=name, - preset=preset, - current_price=billing.price_for(name), - fetch_html=self._hf_pricing_fetch_html, - ) - except Exception as exc: - print(f"[tracker] hf pricing refresh failed for {name!r}: {exc}", flush=True) - continue - if result is None: - continue - new_input = result.get("new_input_price_per_1k", result["new_price_per_1k"]) - new_output = result.get("new_output_price_per_1k", result["new_price_per_1k"]) - for key in _preset_price_keys(name, preset): - billing.set_prices(key, new_input, new_output) - preset["hf_last_price_per_1k"] = result["new_price_per_1k"] - preset["hf_last_updated"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) - if self._hf_pricing_log is not None: - self._hf_pricing_log.record_change( - model=name, - old_price_per_1k=result["old_price_per_1k"], - new_price_per_1k=result["new_price_per_1k"], - source_repo_id=result["source_repo_id"], - source_provider=result["source_provider"], - ) - print( - f"[tracker] hf pricing: {name} {result['old_price_per_1k']:.6f} -> " - f"{result['new_price_per_1k']:.6f} USDT/1k tokens " - f"(80% of {result['source_repo_id']}::{result['source_provider']})", - flush=True, - ) - - def _raft_apply(self, command: str, payload: dict) -> None: - """Called by RaftNode when a log entry is committed — replicate to local registry.""" - if command != "register": - return - # Re-apply the registration to our local registry (follower path). - # On the leader this is a no-op since it already registered locally. - node_id = payload.get("node_id") - if not node_id or node_id in self._registry: - return # already present (leader case) or malformed - endpoint = payload.get("endpoint", "") - try: - shard_start = int(payload["shard_start"]) if payload.get("shard_start") is not None else None - shard_end = int(payload["shard_end"]) if payload.get("shard_end") is not None else None - except (TypeError, ValueError): - return - entry = _NodeEntry( - node_id=node_id, - endpoint=endpoint.rstrip("/"), - shard_start=shard_start, - shard_end=shard_end, - model=payload.get("model", "stub-model"), - shard_checksum=payload.get("shard_checksum"), - hardware_profile=payload.get("hardware_profile", {}), - wallet_address=payload.get("wallet_address"), - score=float(payload.get("score", 1.0)), - tracker_mode=bool(payload.get("tracker_mode", False)), - hf_repo=payload.get("hf_repo"), - num_layers=int(payload["num_layers"]) if payload.get("num_layers") is not None else None, - model_metadata=payload.get("model_metadata") if isinstance(payload.get("model_metadata"), dict) else None, - downloaded_models=( - payload.get("downloaded_models") - if isinstance(payload.get("downloaded_models"), list) - else None - ), - ) - with self._lock: - self._registry[node_id] = entry - - def _rebalance_loop(self) -> None: - while not self._rebalance_stop.wait(self._rebalance_interval): - server = self._server - if server is None: - return - with self._lock: - _purge_expired_nodes_locked(server) - _rebalance_all_locked(server) - _scale_demanded_models_locked(server) - - def _push_to_peers(self, path: str, body: bytes) -> bool: - """POST a hive-signed payload to every cluster peer; True if all succeeded.""" - headers = {"Content-Type": "application/json"} - if self._hive_secret: - headers.update(sign_hive_request(self._hive_secret, body)) - delivered_all = True - for peer in self._cluster_peers: - try: - req = urllib.request.Request( - f"{peer}{path}", data=body, headers=headers, method="POST", - ) - with urllib.request.urlopen(req, timeout=2.0) as r: - r.read() - except Exception: - delivered_all = False - return delivered_all - - def _stats_loop(self) -> None: - """Periodically save stats/billing to DB and push local slices to cluster peers.""" - while not self._stats_stop.wait(_StatsCollector.SAVE_INTERVAL): - if self._stats is not None: - self._stats.save_to_db() - if self._billing is not None: - self._billing.save_to_db() - if self._accounts is not None: - self._accounts.save_to_db() - registry_log = getattr(self._contracts, "registry_log", None) - if registry_log is not None: - registry_log.save_to_db() - if self._cluster_peers and not self._hive_secret: - print( - "[tracker] WARNING: cluster peers configured without --hive-secret — " - "gossip pushes will be rejected (ADR-0017)", - flush=True, - ) - if self._stats is not None and self._cluster_peers: - self_url = self._cluster_self_url or f"http://{self._host}:{self.port}" - local_rpms = self._stats.get_local_rpms() - body = json.dumps({"tracker_url": self_url, "stats": local_rpms}).encode() - self._push_to_peers("/v1/stats/gossip", body) - if self._billing is not None and self._cluster_peers: - events, cursor = self._billing.events_since(self._billing_gossip_cursor) - if events: - body = json.dumps({"events": events}).encode() - # Only advance past events every peer has seen; unreachable - # peers get the full backlog on the next tick (idempotent - # via event-id dedupe on the receiving side). - if self._push_to_peers("/v1/billing/gossip", body): - self._billing_gossip_cursor = cursor - if self._accounts is not None and self._cluster_peers: - events, cursor = self._accounts.events_since(self._accounts_gossip_cursor) - if events: - body = json.dumps({"events": events}).encode() - if self._push_to_peers("/v1/accounts/gossip", body): - self._accounts_gossip_cursor = cursor - registry_log = getattr(self._contracts, "registry_log", None) - if registry_log is not None and self._cluster_peers: - events, cursor = registry_log.events_since(self._registry_gossip_cursor) - if events: - body = json.dumps({"events": events}).encode() - if self._push_to_peers("/v1/registry/gossip", body): - self._registry_gossip_cursor = cursor - - def stop(self) -> None: - if self._raft is not None: - self._raft.stop() - if self._gossip is not None: - self._gossip.stop() - if self._server is None: - return - self._rebalance_stop.set() - self._stats_stop.set() - self._deposit_stop.set() - self._settlement_stop.set() - self._hf_pricing_stop.set() - if self._stats is not None: - self._stats.save_to_db() - if self._billing is not None: - self._billing.save_to_db() - if self._accounts is not None: - self._accounts.save_to_db() - registry_log = getattr(self._contracts, "registry_log", None) - if registry_log is not None: - registry_log.save_to_db() - self._server.shutdown() - self._server.server_close() - if self._thread is not None: - self._thread.join(timeout=1) - if self._rebalance_thread is not None: - self._rebalance_thread.join(timeout=1) - if self._stats_thread is not None: - self._stats_thread.join(timeout=1) - if self._deposit_thread is not None: - self._deposit_thread.join(timeout=1) - self._deposit_thread = None - if self._settlement_thread is not None: - self._settlement_thread.join(timeout=1) - self._settlement_thread = None - if self._hf_pricing_thread is not None: - self._hf_pricing_thread.join(timeout=1) - self._hf_pricing_thread = None - self._server = None - self._thread = None - self._rebalance_thread = None - self._stats_thread = None - self._raft = None - self._gossip = None - self.port = None +"""Tracker HTTP server — node registry and route selection for the inference mesh. + +HTTP API contract: +- POST /v1/nodes/register + Request JSON: { + "endpoint": "http://host:port", "shard_start": int, "shard_end": int, + "model": str optional, "shard_checksum": str optional, + "hardware_profile": object, "wallet_address": str optional, + "score": number optional + } + Response 200: {"node_id": str} + Response 400: {"error": str} +- POST /v1/nodes/{node_id}/heartbeat + Response 200: {} + Response 404: {"error": "node not found"} +- GET /v1/route?model= + Response 200: {"route": ["http://node-a", "http://node-b"]} + Response 400/404/503: {"error": str} +- GET /v1/routes?model=&redundancy= + Response 200: {"routes": [{"route": [...], "nodes": [...]}]} + Response 400/404/503: {"error": str} +- GET /v1/routing?model= (ADR-0021 learned route table) + Response 200: {"config": {...}, "models": {model: {"epoch": int, "routes": [...]}}} +""" + +import http.cookies +import http.server +import hashlib +import itertools +import json +import os +import random +import select +import socketserver +import sqlite3 +import tarfile +import threading +import time +import urllib.parse +import urllib.request +import uuid +from collections import deque +from dataclasses import dataclass, field +from importlib.resources import files +from pathlib import Path +from typing import Any + +from .accounts import DEFAULT_ACCOUNTS_DB_PATH, AccountStore +from .auth import is_validator_token, sign_hive_request, verify_hive_request +from .wallet_proof import binding_message, verify_wallet_signature +from .billing import DEFAULT_BILLING_DB_PATH, BillingLedger +from .calibration import DEFAULT_CALIBRATION_DB_PATH, ToplocCalibrationStore +from .hf_pricing import DEFAULT_HF_PRICING_LOG_DB_PATH, HfPricingLog, refresh_preset_price +from .gossip import NodeGossip +from .logging_setup import tracker_logger +from .routing_stats import ( + RouteCandidate, + RouteStatsStore, + RoutingConfig, + choose_route, + route_signature, + route_table, +) +from .model_files import files_for_layer_range, snapshot_dir_for_repo +from .raft import RaftNode + + +_CONSOLE_LIMIT = 300 +_PROXY_PROGRESS_LOG_INTERVAL = 5.0 +_SESSION_COOKIE_NAME = "meshnet_session" + + +def _preset_price_keys(name: str, preset: dict) -> set[str]: + """All model strings a client may bill under for one preset. + + ``BillingLedger.price_for`` is keyed by the raw ``model`` string in the + request, so the preset price must be registered under the preset name, + its ``hf_repo``, and every alias — otherwise ``unsloth/Qwen…`` style + requests silently fall back to the default rate. + """ + keys = {name} + hf_repo = preset.get("hf_repo") + if isinstance(hf_repo, str) and hf_repo: + keys.add(hf_repo) + for alias in preset.get("aliases") or []: + if isinstance(alias, str) and alias: + keys.add(alias) + return keys + + +def derive_relay_url_from_public_tracker_url(url: str | None) -> str | None: + """Return wss://host/ws when url is a public HTTPS tracker origin.""" + if not url: + return None + parsed = urllib.parse.urlparse(url) + if parsed.scheme != "https": + return None + host = parsed.hostname + if not host or host in ("127.0.0.1", "localhost"): + return None + return f"wss://{parsed.netloc}/ws" + + +def _load_model_presets_from_data() -> dict[str, dict]: + """Load recommended model presets from package JSON data.""" + try: + raw = files("meshnet_tracker").joinpath("model_presets.json").read_text() + data = json.loads(raw) + except Exception: + return {} + models = data.get("models", {}) + if not isinstance(models, dict): + return {} + return { + str(name): preset + for name, preset in models.items() + if isinstance(preset, dict) + } + + +DEFAULT_MODEL_PRESETS: dict[str, dict] = { + "stub-model": { + "layers_start": 0, + "layers_end": 31, + "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, + }, + "openai-community/gpt2": { + "layers_start": 0, + "layers_end": 11, + "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024}, + }, + **_load_model_presets_from_data(), +} + +def _clone_model_presets(presets: dict[str, dict]) -> dict[str, dict]: + """Shallow-copy each preset dict so a runtime mutation (e.g. issue 23's + dynamic pricing refresh writing hf_last_price_per_1k/hf_last_updated) + never leaks into the shared module-level DEFAULT_MODEL_PRESETS and from + there into other TrackerServer instances in the same process.""" + return {name: dict(preset) for name, preset in presets.items()} + + +DEFAULT_VRAM_BYTES = 8 * 1024 * 1024 * 1024 +DEFAULT_RAM_BYTES = 16 * 1024 * 1024 * 1024 + + +def _snapshot_regular_files(snapshot_dir: Path) -> list[str]: + return sorted( + path.relative_to(snapshot_dir).as_posix() + for path in snapshot_dir.rglob("*") + if path.is_file() + ) +DEFAULT_QUANTIZATIONS = ["bfloat16"] +DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0 +# US-039/US-040 — single source of truth for the credit defaults (referenced by +# TrackerServer, _TrackerHTTPServer, and the CLI). Alpha runs devnet-friendly; +# flip both to 0 before any deployment holding a mainnet treasury. +DEFAULT_CALLER_CREDIT_USDT = 1.0 +DEFAULT_DEVNET_TOPUP_USDT = 1.0 + + +def _endpoint_key(url: str) -> str: + """Normalize http(s) endpoints for host:port comparison.""" + parsed = urllib.parse.urlparse(url.rstrip("/")) + host = (parsed.hostname or "").lower() + if not host: + return url.rstrip("/").lower() + port = parsed.port + if port is None: + port = 443 if parsed.scheme == "https" else 80 + return f"{host}:{port}" + + +def _model_aliases(model: str | None) -> set[str]: + """Return stable lookup aliases for a model repo or display name.""" + if not model: + return set() + normalized = model.strip() + if not normalized: + return set() + aliases = {normalized} + short = normalized.rsplit("/", 1)[-1] + aliases.add(short) + lowered = short.lower() + aliases.add(lowered) + if lowered.endswith("-instruct"): + aliases.add(lowered.removesuffix("-instruct")) + return aliases + + +def _preset_aliases(name: str, preset: dict | None) -> set[str]: + aliases = _model_aliases(name) + if not preset: + return aliases + hf_repo = preset.get("hf_repo") + if isinstance(hf_repo, str): + aliases |= _model_aliases(hf_repo) + for alias in preset.get("aliases", []) or []: + if isinstance(alias, str): + aliases |= _model_aliases(alias) + return aliases + + +def _resolve_model_preset(model_presets: dict, model: str) -> tuple[str, dict] | tuple[None, None]: + requested = _model_aliases(model) + for name, preset in model_presets.items(): + if requested & _preset_aliases(name, preset): + return name, preset + return None, None + + +def _node_matches_model(node: "_NodeEntry", model: str) -> bool: + requested = _model_aliases(model) + if not requested: + return False + return bool(requested & (_model_aliases(node.model) | _model_aliases(node.hf_repo))) + + +def _node_matches_preset(node: "_NodeEntry", name: str, preset: dict) -> bool: + requested = _preset_aliases(name, preset) + return bool(requested & (_model_aliases(node.model) | _model_aliases(node.hf_repo))) + + +class _RollingCounter: + """Circular-bucket request counter. + + Tracks events in `num_buckets` slots, each covering `bucket_seconds` of wall time. + Old buckets are silently discarded when their epoch expires. + """ + + def __init__(self, num_buckets: int, bucket_seconds: int) -> None: + self._num = num_buckets + self._bsec = bucket_seconds + self._counts: list[int] = [0] * num_buckets + self._epochs: list[int] = [-1] * num_buckets # which epoch each slot holds + + def _epoch(self, now: float) -> int: + return int(now) // self._bsec + + def record(self, now: float | None = None) -> None: + t = now if now is not None else time.time() + ep = self._epoch(t) + idx = ep % self._num + if self._epochs[idx] != ep: + self._counts[idx] = 0 + self._epochs[idx] = ep + self._counts[idx] += 1 + + def rpm(self, now: float | None = None) -> float: + t = now if now is not None else time.time() + cutoff = self._epoch(t) - self._num # epochs <= this are stale + total = sum(self._counts[i] for i in range(self._num) if self._epochs[i] > cutoff) + window_minutes = (self._num * self._bsec) / 60.0 + return total / window_minutes if window_minutes > 0 else 0.0 + + def buckets(self) -> list[tuple[int, int]]: + return [(self._epochs[i], self._counts[i]) for i in range(self._num)] + + def restore_buckets(self, data: list[tuple[int, int]]) -> None: + for i, (ep, cnt) in enumerate(data): + if i < self._num: + self._epochs[i] = ep + self._counts[i] = cnt + + +class _RollingThroughput: + """Circular buckets for observed tokens/sec. + + Each bucket stores total output tokens and total elapsed seconds. TPS for a + window is the ratio across non-stale buckets, which avoids over-weighting + small fast requests. + """ + + def __init__(self, num_buckets: int, bucket_seconds: int) -> None: + self._num = num_buckets + self._bsec = bucket_seconds + self._tokens: list[int] = [0] * num_buckets + self._seconds: list[float] = [0.0] * num_buckets + self._samples: list[int] = [0] * num_buckets + self._epochs: list[int] = [-1] * num_buckets + + def _epoch(self, now: float) -> int: + return int(now) // self._bsec + + def record(self, total_tokens: int, elapsed_seconds: float, now: float | None = None) -> None: + if total_tokens <= 0 or elapsed_seconds <= 0: + return + t = now if now is not None else time.time() + ep = self._epoch(t) + idx = ep % self._num + if self._epochs[idx] != ep: + self._tokens[idx] = 0 + self._seconds[idx] = 0.0 + self._samples[idx] = 0 + self._epochs[idx] = ep + self._tokens[idx] += int(total_tokens) + self._seconds[idx] += float(elapsed_seconds) + self._samples[idx] += 1 + + def stats(self, now: float | None = None) -> dict: + t = now if now is not None else time.time() + cutoff = self._epoch(t) - self._num + tokens = 0 + seconds = 0.0 + samples = 0 + for i in range(self._num): + if self._epochs[i] > cutoff: + tokens += self._tokens[i] + seconds += self._seconds[i] + samples += self._samples[i] + return { + "tokens_per_sec": round(tokens / seconds, 4) if seconds > 0 else None, + "tokens": tokens, + "seconds": round(seconds, 6), + "sample_count": samples, + } + + def buckets(self) -> list[tuple[int, int, float, int]]: + return [ + (self._epochs[i], self._tokens[i], self._seconds[i], self._samples[i]) + for i in range(self._num) + ] + + def restore_buckets(self, data: list[tuple[int, int, float, int]]) -> None: + for i, (ep, tokens, seconds, samples) in enumerate(data): + if i < self._num: + self._epochs[i] = ep + self._tokens[i] = tokens + self._seconds[i] = seconds + self._samples[i] = samples + + +class _ModelStats: + """Three rolling windows for one model: last hour, last day, last month.""" + + def __init__(self) -> None: + self.per_minute = _RollingCounter(60, 60) # 60 × 1-min buckets = 1 hour + self.per_hour = _RollingCounter(24, 3600) # 24 × 1-hr buckets = 1 day + self.per_day = _RollingCounter(30, 86400) # 30 × 1-day buckets = ~1 month + + def record(self, now: float | None = None) -> None: + t = now if now is not None else time.time() + self.per_minute.record(t) + self.per_hour.record(t) + self.per_day.record(t) + + def rpms(self, now: float | None = None) -> dict: + t = now if now is not None else time.time() + return { + "rpm_last_hour": round(self.per_minute.rpm(t), 4), + "rpm_last_day": round(self.per_hour.rpm(t), 4), + "rpm_last_month": round(self.per_day.rpm(t), 4), + } + + +class _NodeModelThroughput: + """Observed throughput windows for one node/model pair.""" + + def __init__(self) -> None: + self.per_minute = _RollingThroughput(60, 60) + self.per_hour = _RollingThroughput(24, 3600) + self.per_day = _RollingThroughput(30, 86400) + + def record(self, total_tokens: int, elapsed_seconds: float, now: float | None = None) -> None: + t = now if now is not None else time.time() + self.per_minute.record(total_tokens, elapsed_seconds, t) + self.per_hour.record(total_tokens, elapsed_seconds, t) + self.per_day.record(total_tokens, elapsed_seconds, t) + + def stats(self, now: float | None = None) -> dict: + hour = self.per_minute.stats(now) + day = self.per_hour.stats(now) + month = self.per_day.stats(now) + return { + "tokens_per_sec_last_hour": hour["tokens_per_sec"], + "tokens_per_sec_last_day": day["tokens_per_sec"], + "tokens_per_sec_last_month": month["tokens_per_sec"], + "sample_count_last_hour": hour["sample_count"], + "tokens_last_hour": hour["tokens"], + "seconds_last_hour": hour["seconds"], + } + + +class _StatsCollector: + """Thread-safe model request stats with SQLite persistence and peer slice merging.""" + + SAVE_INTERVAL = 60.0 # seconds between DB saves + + def __init__(self, db_path: str | None = None) -> None: + self._lock = threading.Lock() + self._local: dict[str, _ModelStats] = {} + self._node_model: dict[tuple[str, str], _NodeModelThroughput] = {} + self._peer_rpms: dict[str, dict[str, dict]] = {} # tracker_url -> model -> rpms + self._db_path = db_path + if db_path: + self._init_db() + self._load_from_db() + + # ---------- public API ---------- + + def record_request(self, model: str, now: float | None = None) -> None: + t = now if now is not None else time.time() + with self._lock: + if model not in self._local: + self._local[model] = _ModelStats() + self._local[model].record(t) + + def record_node_throughput( + self, + node_id: str, + model: str, + total_tokens: int, + elapsed_seconds: float, + now: float | None = None, + ) -> None: + if not node_id or not model or total_tokens <= 0 or elapsed_seconds <= 0: + return + t = now if now is not None else time.time() + with self._lock: + key = (node_id, model) + if key not in self._node_model: + self._node_model[key] = _NodeModelThroughput() + self._node_model[key].record(total_tokens, elapsed_seconds, t) + + def get_local_rpms(self, now: float | None = None) -> dict[str, dict]: + t = now if now is not None else time.time() + with self._lock: + return {m: s.rpms(t) for m, s in self._local.items()} + + def get_node_model_stats(self, node_id: str, model: str, now: float | None = None) -> dict: + t = now if now is not None else time.time() + empty = { + "tokens_per_sec_last_hour": None, + "tokens_per_sec_last_day": None, + "tokens_per_sec_last_month": None, + "sample_count_last_hour": 0, + "tokens_last_hour": 0, + "seconds_last_hour": 0.0, + } + with self._lock: + stat = self._node_model.get((node_id, model)) + return stat.stats(t) if stat is not None else dict(empty) + + def get_node_throughput_stats(self, now: float | None = None) -> dict[str, dict]: + t = now if now is not None else time.time() + with self._lock: + result: dict[str, dict] = {} + for (node_id, model), stat in self._node_model.items(): + result.setdefault(node_id, {"models": {}})["models"][model] = stat.stats(t) + return result + + def merge_peer_rpms(self, tracker_url: str, rpms: dict[str, dict]) -> None: + with self._lock: + self._peer_rpms[tracker_url] = dict(rpms) + + def get_combined_stats(self, now: float | None = None) -> dict: + t = now if now is not None else time.time() + with self._lock: + combined: dict[str, dict] = {} + for model, ms in self._local.items(): + combined[model] = ms.rpms(t) + for _url, peer in self._peer_rpms.items(): + for model, rpms in peer.items(): + if model not in combined: + combined[model] = {"rpm_last_hour": 0.0, "rpm_last_day": 0.0, "rpm_last_month": 0.0} + for key in ("rpm_last_hour", "rpm_last_day", "rpm_last_month"): + combined[model][key] = round(combined[model].get(key, 0.0) + rpms.get(key, 0.0), 4) + return combined + + # ---------- persistence ---------- + + def _init_db(self) -> None: + con = sqlite3.connect(self._db_path) # type: ignore[arg-type] + con.execute( + "CREATE TABLE IF NOT EXISTS model_rpm_buckets " + "(model TEXT, window TEXT, bucket_idx INTEGER, bucket_epoch INTEGER, count INTEGER, " + "PRIMARY KEY (model, window, bucket_idx))" + ) + con.execute( + "CREATE TABLE IF NOT EXISTS node_model_tps_buckets " + "(node_id TEXT, model TEXT, window TEXT, bucket_idx INTEGER, bucket_epoch INTEGER, " + "tokens INTEGER, seconds REAL, samples INTEGER, " + "PRIMARY KEY (node_id, model, window, bucket_idx))" + ) + con.commit() + con.close() + + def save_to_db(self) -> None: + if not self._db_path: + return + with self._lock: + rows = [] + for model, ms in self._local.items(): + for window_name, counter in ( + ("hour", ms.per_minute), + ("day", ms.per_hour), + ("month", ms.per_day), + ): + for idx, (ep, cnt) in enumerate(counter.buckets()): + if ep >= 0: + rows.append((model, window_name, idx, ep, cnt)) + tps_rows = [] + for (node_id, model), stat in self._node_model.items(): + for window_name, counter in ( + ("hour", stat.per_minute), + ("day", stat.per_hour), + ("month", stat.per_day), + ): + for idx, (ep, tokens, seconds, samples) in enumerate(counter.buckets()): + if ep >= 0: + tps_rows.append((node_id, model, window_name, idx, ep, tokens, seconds, samples)) + con = sqlite3.connect(self._db_path) # type: ignore[arg-type] + con.executemany( + "INSERT OR REPLACE INTO model_rpm_buckets VALUES (?,?,?,?,?)", rows + ) + con.executemany( + "INSERT OR REPLACE INTO node_model_tps_buckets VALUES (?,?,?,?,?,?,?,?)", + tps_rows, + ) + con.commit() + con.close() + + def _load_from_db(self) -> None: + con = sqlite3.connect(self._db_path) # type: ignore[arg-type] + rows = con.execute("SELECT model, window, bucket_idx, bucket_epoch, count FROM model_rpm_buckets").fetchall() + tps_rows = con.execute( + "SELECT node_id, model, window, bucket_idx, bucket_epoch, tokens, seconds, samples " + "FROM node_model_tps_buckets" + ).fetchall() + con.close() + grouped: dict[str, dict[str, list[tuple[int, int]]]] = {} + for model, window, idx, ep, cnt in rows: + grouped.setdefault(model, {}).setdefault(window, []).append((idx, ep, cnt)) + for model, windows in grouped.items(): + ms = _ModelStats() + for window_name, entries in windows.items(): + counter = {"hour": ms.per_minute, "day": ms.per_hour, "month": ms.per_day}.get(window_name) + if counter is None: + continue + data = [(0, -1)] * counter._num + for idx, ep, cnt in entries: + if 0 <= idx < counter._num: + data[idx] = (ep, cnt) + counter.restore_buckets(data) + self._local[model] = ms + tps_grouped: dict[tuple[str, str], dict[str, list[tuple[int, int, int, float, int]]]] = {} + for node_id, model, window, idx, ep, tokens, seconds, samples in tps_rows: + tps_grouped.setdefault((node_id, model), {}).setdefault(window, []).append( + (idx, ep, tokens, seconds, samples) + ) + for key, windows in tps_grouped.items(): + stat = _NodeModelThroughput() + for window_name, entries in windows.items(): + counter = {"hour": stat.per_minute, "day": stat.per_hour, "month": stat.per_day}.get(window_name) + if counter is None: + continue + data = [(0, 0, 0.0, 0)] * counter._num + for idx, ep, tokens, seconds, samples in entries: + if 0 <= idx < counter._num: + data[idx] = (ep, int(tokens), float(seconds), int(samples)) + counter.restore_buckets(data) + self._node_model[key] = stat + + +class _NodeEntry: + __slots__ = ( + "node_id", "endpoint", "shard_start", "shard_end", + "model", "hf_repo", "num_layers", "model_metadata", "shard_checksum", "downloaded_models", "hardware_profile", "wallet_address", + "score", "vram_bytes", "ram_bytes", "quantizations", "max_loaded_shards", + "benchmark_tokens_per_sec", "quantization", "managed_assignment", + "model_tokens_per_sec", + "pending_directives", "last_heartbeat", "tracker_mode", + "relay_addr", "cert_fingerprint", "peer_id", + # heartbeat stats (reported by node, cumulative) + "total_requests", "failed_requests", "queue_depth", "proxy_inflight", "uptime_seconds", + "current_requests", + "status", # "ready" | "loading" + "heartbeats_expected", "heartbeats_received", + # dynamic reassignment queued by the tracker + "pending_new_assignment", + ) + + def __init__( + self, + node_id: str, + endpoint: str, + shard_start: int | None, + shard_end: int | None, + model: str | None, + shard_checksum: str | None, + hardware_profile: dict, + wallet_address: str | None, + score: float, + vram_bytes: int = DEFAULT_VRAM_BYTES, + ram_bytes: int = DEFAULT_RAM_BYTES, + quantizations: list[str] | None = None, + max_loaded_shards: int = 1, + benchmark_tokens_per_sec: float = DEFAULT_BENCHMARK_TOKENS_PER_SEC, + quantization: str | None = None, + managed_assignment: bool = False, + tracker_mode: bool = False, + hf_repo: str | None = None, + num_layers: int | None = None, + model_metadata: dict | None = None, + downloaded_models: list[dict] | None = None, + relay_addr: str | None = None, + cert_fingerprint: str | None = None, + peer_id: str | None = None, + ) -> None: + self.node_id = node_id + self.endpoint = endpoint + self.shard_start = shard_start + self.shard_end = shard_end + self.model = model + self.shard_checksum = shard_checksum + self.hardware_profile = hardware_profile + self.wallet_address = wallet_address + self.score = score + self.vram_bytes = vram_bytes + self.ram_bytes = ram_bytes + self.quantizations = quantizations or list(DEFAULT_QUANTIZATIONS) + self.max_loaded_shards = max_loaded_shards + self.benchmark_tokens_per_sec = benchmark_tokens_per_sec + self.quantization = quantization + self.managed_assignment = managed_assignment + self.model_tokens_per_sec: dict[str, float] = {} + self.tracker_mode = tracker_mode + self.hf_repo = hf_repo + self.num_layers = num_layers + self.model_metadata = dict(model_metadata or {}) + self.downloaded_models = [dict(item) for item in (downloaded_models or []) if isinstance(item, dict)] + self.relay_addr = relay_addr + self.cert_fingerprint = cert_fingerprint + self.peer_id = peer_id + self.pending_directives: list[dict] = [] + self.last_heartbeat: float = time.monotonic() + self.total_requests: int = 0 + self.failed_requests: int = 0 + self.queue_depth: int = 0 + self.proxy_inflight: int = 0 + self.current_requests: list[dict] = [] + self.uptime_seconds: float = 0.0 + self.status: str = "ready" + self.heartbeats_expected: int = 0 + self.heartbeats_received: int = 0 + self.pending_new_assignment: dict | None = None + + +def _effective_throughput(node: "_NodeEntry", model: str | None = None) -> float: + """Effective tokens/s accounting for current queue depth.""" + observed = None + if model: + observed = node.model_tokens_per_sec.get(model) + if observed is None: + for alias in _model_aliases(model): + observed = node.model_tokens_per_sec.get(alias) + if observed is not None: + break + base = observed if observed is not None and observed > 0 else node.benchmark_tokens_per_sec + return base / (_effective_queue_depth(node) + 1) + + +def _effective_queue_depth(node: "_NodeEntry") -> int: + """Best current load estimate: heartbeat queue or tracker-routed in-flight.""" + return max(node.queue_depth, node.proxy_inflight) + + +_CURRENT_REQUEST_FIELDS = frozenset({ + "request_id", "model", "kind", "tokens", "tokens_per_sec", + "elapsed_seconds", "routing_complete", +}) + + +def _normalize_current_requests(items: object, *, limit: int = 32) -> list[dict]: + """Sanitize node-reported in-flight request snapshots from heartbeats.""" + if not isinstance(items, list): + return [] + out: list[dict] = [] + for item in items[:limit]: + if not isinstance(item, dict): + continue + request_id = item.get("request_id") + if not request_id: + continue + rec: dict = {"request_id": str(request_id)} + for key in _CURRENT_REQUEST_FIELDS: + if key == "request_id" or key not in item: + continue + value = item[key] + if key in {"tokens"}: + rec[key] = int(value) + elif key in {"tokens_per_sec", "elapsed_seconds"}: + rec[key] = float(value) + elif key == "routing_complete": + rec[key] = bool(value) + else: + rec[key] = str(value) + out.append(rec) + return out + + +def _record_proxy_inflight( + server: "_TrackerHTTPServer", + nodes: list["_NodeEntry"], + delta: int, +) -> None: + with server.lock: + for node in nodes: + node.proxy_inflight = max(0, node.proxy_inflight + delta) + + +def _reputation_multiplier(node: "_NodeEntry", contracts: Any | None) -> float: + """ADR-0018 §6: veteran, good-standing nodes route more -- earnings scale + with tenure/reputation. Maps reputation [0, 1] to a [0.5, 1.0] routing + weight so a damaged-but-not-banned wallet loses traffic gradually rather + than being cut off outright (banning is handled separately).""" + if contracts is None or not node.wallet_address: + return 1.0 + reputation = contracts.registry.get_wallet(node.wallet_address).reputation + return 0.5 + 0.5 * reputation + + +def _select_route( + nodes: list[_NodeEntry], + required_start: int, + required_end: int, + model: str | None = None, + contracts: Any | None = None, +) -> tuple[list[_NodeEntry], str]: + """Greedy interval-cover biased toward fast, lightly-loaded, reputable nodes. + + Among nodes that equally advance coverage, prefer the one with higher + effective throughput: observed per-model tokens/sec / (queue_depth + 1), + falling back to startup benchmark_tokens_per_sec until observations exist, + weighted by the node's reputation multiplier (see `_reputation_multiplier`). + Tiebreak: higher shard_end (fewer hops). + """ + candidates = sorted( + [node for node in nodes if node.shard_start is not None and node.shard_end is not None], + key=lambda n: (n.shard_start, -n.shard_end), # type: ignore[operator] + ) + route: list[_NodeEntry] = [] + covered_up_to = required_start - 1 + + def _routing_score(node: "_NodeEntry") -> float: + return _effective_throughput(node, model) * _reputation_multiplier(node, contracts) + + while covered_up_to < required_end: + best: _NodeEntry | None = None + for node in candidates: + if node.shard_start <= covered_up_to + 1 and node.shard_end > covered_up_to: + if best is None: + best = node + elif node.shard_end > best.shard_end: + best = node + elif node.shard_end == best.shard_end and _routing_score(node) > _routing_score(best): + best = node + if best is None: + missing = covered_up_to + 1 + return [], f"no route available: no registered node covers layer {missing}" + route.append(best) + covered_up_to = best.shard_end + candidates = [n for n in candidates if n is not best] + + return route, "" + + +def _enumerate_routes( + nodes: list["_NodeEntry"], + required_start: int, + required_end: int, + model: str | None = None, + contracts: Any | None = None, + max_candidates: int = 8, +) -> list["RouteCandidate"]: + """Enumerate viable route candidates for bandit selection (ADR-0021). + + One candidate per distinct head (a node that can embed the prompt, i.e. + covers `required_start` from layer 0 of the range), each greedily completed + with the longest-advancing hops. The route's prior throughput estimate is + its bottleneck hop's queue-adjusted effective throughput — used only until + observed route samples exist. + """ + sharded = [ + n for n in nodes + 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 + # and embed the prompt — same condition as tracker_mode registration). + heads = [n for n in sharded if n.shard_start == required_start] + candidates: dict[str, RouteCandidate] = {} + for head in heads: + route = [head] + covered_up_to = head.shard_end + pool = [n for n in sharded if n is not head] + while covered_up_to < required_end: + best = None + for n in pool: + if n.shard_start <= covered_up_to + 1 and n.shard_end > covered_up_to: + if best is None or n.shard_end > best.shard_end or ( + n.shard_end == best.shard_end + and _effective_throughput(n, model) * _reputation_multiplier(n, contracts) + > _effective_throughput(best, model) * _reputation_multiplier(best, contracts) + ): + best = n + if best is None: + route = [] + break + route.append(best) + covered_up_to = best.shard_end + pool = [n for n in pool if n is not best] + if not route: + continue + signature = route_signature(model or "?", route) + if signature in candidates: + continue + prior = min( + _effective_throughput(n, model) * _reputation_multiplier(n, contracts) + for n in route + ) + candidates[signature] = RouteCandidate(nodes=route, signature=signature, prior_tps=prior) + ranked = sorted(candidates.values(), key=lambda c: -c.prior_tps) + return ranked[:max_candidates] + + +def _node_route_summary(nodes: list["_NodeEntry"]) -> list[dict]: + return [ + { + "node_id": node.node_id, + "model": node.model, + "hf_repo": node.hf_repo, + "endpoint": node.endpoint, + "shard": f"{node.shard_start}-{node.shard_end}", + "num_layers": node.num_layers, + "queue_depth": _effective_queue_depth(node), + "heartbeat_queue_depth": node.queue_depth, + "proxy_inflight": node.proxy_inflight, + "current_requests": list(node.current_requests), + } + for node in nodes + ] + + +def _coverage_percentage( + nodes: list[_NodeEntry], + required_start: int, + required_end: int, +) -> float: + required_layers = required_end - required_start + 1 + if required_layers <= 0: + return 0.0 + + intervals = sorted( + ( + (max(required_start, node.shard_start), min(required_end, node.shard_end)) + for node in nodes + if node.shard_start is not None + and node.shard_end is not None + if node.shard_end >= required_start and node.shard_start <= required_end + ), + key=lambda interval: interval[0], + ) + covered = 0 + covered_until = required_start - 1 + for start, end in intervals: + if end <= covered_until: + continue + next_start = max(start, covered_until + 1) + if next_start > end: + continue + covered += end - next_start + 1 + covered_until = end + return round((covered / required_layers) * 100, 2) + + +def _served_model_copies( + nodes: list[_NodeEntry], + required_start: int, + required_end: int, +) -> float: + required_layers = required_end - required_start + 1 + if required_layers <= 0: + return 0.0 + + layer_counts = [] + for layer in range(required_start, required_end + 1): + count = 0 + for node in nodes: + if node.shard_start is None or node.shard_end is None: + continue + if node.shard_start <= layer <= node.shard_end: + count += 1 + layer_counts.append(count) + + if not layer_counts: + return 0.0 + + complete_copies = min(layer_counts) + residual_layers = sum(1 for count in layer_counts if count > complete_copies) + return round(complete_copies + (residual_layers / required_layers), 2) + + +def _model_health_summary( + server: "_TrackerHTTPServer", + model: str | None, + hf_repo: str | None = None, +) -> dict: + lookup = hf_repo or model + if not lookup: + return {} + + resolved_name, preset = _resolve_model_preset(server.model_presets, lookup) + if preset is not None: + required_start, required_end = _preset_layer_bounds(preset) + model_nodes = [ + node for node in server.registry.values() + if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] + ] + model_id = resolved_name or lookup + else: + model_nodes = [ + node for node in server.registry.values() + if _node_matches_model(node, lookup) + and node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + ] + if not model_nodes: + return { + "model": lookup, + "served_model_copies": 0.0, + "coverage_percentage": 0.0, + "node_count": 0, + } + required_start = 0 + required_end = max(node.num_layers for node in model_nodes if node.num_layers is not None) - 1 + model_id = hf_repo or model or lookup + + return { + "model": model_id, + "required_start": required_start, + "required_end": required_end, + "served_model_copies": _served_model_copies(model_nodes, required_start, required_end), + "coverage_percentage": _coverage_percentage(model_nodes, required_start, required_end), + "node_count": len(model_nodes), + } + + +def _preset_layer_bounds(preset: dict) -> tuple[int, int]: + start = int(preset.get("layers_start", 0)) + if "layers_end" in preset: + return start, int(preset["layers_end"]) + return start, start + int(preset["total_layers"]) - 1 + + +def _preset_bytes_per_layer(preset: dict) -> dict[str, int]: + raw = preset.get("bytes_per_layer", preset.get("bytes_per_layer_at_quant", {})) + if isinstance(raw, dict) and raw: + return {str(quant): int(value) for quant, value in raw.items()} + return {"bfloat16": 30 * 1024 * 1024} + + +def _node_quantization(node: _NodeEntry, preset: dict) -> str: + bytes_per_layer = _preset_bytes_per_layer(preset) + if node.quantization in bytes_per_layer: + return node.quantization + for quantization in node.quantizations: + if quantization in bytes_per_layer: + return quantization + return next(iter(bytes_per_layer)) + + +def _node_memory_budget_bytes(node: _NodeEntry) -> tuple[int, str]: + """Return the memory pool used for shard-capacity planning.""" + if node.vram_bytes > 0: + return node.vram_bytes, "vram" + if node.ram_bytes > 0: + return node.ram_bytes, "ram" + return DEFAULT_RAM_BYTES, "ram-default" + + +def _node_layer_capacity(node: _NodeEntry, preset: dict) -> int: + bytes_per_layer = _preset_bytes_per_layer(preset) + quantization = _node_quantization(node, preset) + layer_bytes = bytes_per_layer[quantization] + if layer_bytes <= 0: + return 0 + memory_budget_bytes, _ = _node_memory_budget_bytes(node) + return int((memory_budget_bytes * 0.8) // layer_bytes) + + +def _node_capacity_summary(node: _NodeEntry, preset: dict | None = None) -> dict: + """Operator-facing capacity fields for inspection endpoints.""" + memory_budget_bytes, memory_budget_source = _node_memory_budget_bytes(node) + summary = { + "vram_bytes": node.vram_bytes, + "ram_bytes": node.ram_bytes, + "memory_budget_bytes": memory_budget_bytes, + "memory_budget_source": memory_budget_source, + "max_loaded_shards": node.max_loaded_shards, + "quantizations": list(node.quantizations), + "quantization": node.quantization, + "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec, + "effective_throughput": round(_effective_throughput(node), 4), + } + if preset is not None: + summary["max_assignable_layers"] = _node_layer_capacity(node, preset) + return summary + + +def _node_memory_budget_for_preset(node: _NodeEntry, preset: dict | None = None) -> int: + budget, _source = _node_memory_budget_bytes(node) + if preset is None: + return int(budget * 0.8) + return _node_layer_capacity(node, preset) * max(1, next(iter(_preset_bytes_per_layer(preset).values()))) + + +def _pool_summary(nodes: list[_NodeEntry], preset: dict | None = None) -> dict: + total_vram = sum(max(0, node.vram_bytes) for node in nodes) + total_ram = sum(max(0, node.ram_bytes) for node in nodes) + total_budget = sum(_node_memory_budget_bytes(node)[0] for node in nodes) + effective_budget = sum(_node_memory_budget_for_preset(node, preset) for node in nodes) + return { + "node_count": len(nodes), + "total_vram_bytes": total_vram, + "total_ram_bytes": total_ram, + "total_memory_budget_bytes": total_budget, + "effective_assignable_memory_bytes": effective_budget, + "total_benchmark_tokens_per_sec": round(sum(node.benchmark_tokens_per_sec for node in nodes), 4), + "total_effective_throughput": round(sum(_effective_throughput(node) for node in nodes), 4), + } + + +def _repo_demand_rpm( + server: "_TrackerHTTPServer", + repo: str, + local_rpms: dict[str, dict], +) -> float: + """Look up recent request rate by hf_repo or preset alias.""" + if repo in local_rpms: + return float(local_rpms[repo].get("rpm_last_hour", 0.0) or 0.0) + resolved_name, _preset = _resolve_model_preset(server.model_presets, repo) + if resolved_name and resolved_name in local_rpms: + return float(local_rpms[resolved_name].get("rpm_last_hour", 0.0) or 0.0) + return 0.0 + + +def _preset_for_node(server: "_TrackerHTTPServer", node: _NodeEntry) -> dict | None: + if node.hf_repo: + _resolved, preset = _resolve_model_preset(server.model_presets, node.hf_repo) + if preset is not None: + return preset + if node.model: + preset = server.model_presets.get(node.model) + if preset is not None: + return preset + if node.hf_repo and node.num_layers: + return _hf_rebalance_preset([node]) + return None + + +def _assignment_memory_bytes(node: _NodeEntry, preset: dict | None) -> int: + if preset is None or node.shard_start is None or node.shard_end is None: + return 0 + layers = node.shard_end - node.shard_start + 1 + if layers <= 0: + return 0 + quantization = _node_quantization(node, preset) + layer_bytes = _preset_bytes_per_layer(preset).get(quantization) + if not layer_bytes: + layer_bytes = next(iter(_preset_bytes_per_layer(preset).values())) + return layers * int(layer_bytes) + + +def _endpoint_memory_pool( + server: "_TrackerHTTPServer", + nodes: list[_NodeEntry], +) -> dict: + """Per-host RAM/VRAM budget, usage, spare capacity, and loaded assignments.""" + if not nodes: + return {} + host = nodes[0] + budget_bytes, budget_source = _node_memory_budget_bytes(host) + reserve_bytes = int(budget_bytes * 0.8) + loaded: list[dict] = [] + used_bytes = 0 + for node in nodes: + preset = _preset_for_node(server, node) + assignment_bytes = _assignment_memory_bytes(node, preset) + used_bytes += assignment_bytes + loaded.append({ + "node_id": node.node_id, + "model": node.model, + "hf_repo": node.hf_repo, + "shard_start": node.shard_start, + "shard_end": node.shard_end, + "memory_bytes": assignment_bytes, + "status": node.status, + }) + max_slots = max(node.max_loaded_shards for node in nodes) + loaded_slots = len([node for node in nodes if node.shard_start is not None]) + spare_bytes = max(0, reserve_bytes - used_bytes) + return { + "endpoint": host.endpoint, + "memory_budget_bytes": budget_bytes, + "memory_budget_source": budget_source, + "memory_reserve_bytes": reserve_bytes, + "memory_used_bytes": used_bytes, + "memory_spare_bytes": spare_bytes, + "loaded_slots": loaded_slots, + "max_loaded_shards": max_slots, + "spare_slots": max(0, max_slots - loaded_slots), + "loaded": loaded, + } + + +def _memory_pool_map(server: "_TrackerHTTPServer") -> dict: + """Aggregate and per-endpoint view of assignable RAM across the hive.""" + from collections import defaultdict + + groups: dict[str, list[_NodeEntry]] = defaultdict(list) + for node in server.registry.values(): + groups[node.endpoint.rstrip("/")].append(node) + + hosts = [_endpoint_memory_pool(server, group) for group in groups.values()] + hosts.sort(key=lambda item: (-item.get("memory_spare_bytes", 0), item.get("endpoint", ""))) + total_budget = sum(item["memory_budget_bytes"] for item in hosts) + total_used = sum(item["memory_used_bytes"] for item in hosts) + total_spare = sum(item["memory_spare_bytes"] for item in hosts) + total_slots = sum(item["max_loaded_shards"] for item in hosts) + loaded_slots = sum(item["loaded_slots"] for item in hosts) + return { + "total_memory_budget_bytes": total_budget, + "total_memory_used_bytes": total_used, + "total_memory_spare_bytes": total_spare, + "total_loaded_slots": loaded_slots, + "total_max_loaded_shards": total_slots, + "total_spare_slots": max(0, total_slots - loaded_slots), + "hosts": hosts, + } + + +def _add_shard_directive( + node: _NodeEntry, + model: str, + start: int, + end: int, + quantization: str, + *, + model_sources: list[dict] | None = None, +) -> dict: + directive = { + "action": "ADD_SHARD", + "model": model, + "start_layer": start, + "end_layer": end, + "shard_start": start, + "shard_end": end, + "quantization": quantization, + } + if model_sources: + directive["model_sources"] = model_sources + return directive + + +def _model_demand_and_supply( + server: "_TrackerHTTPServer", + model_key: str, + preset: dict, +) -> dict: + resolved_name, _ = _resolve_model_preset(server.model_presets, model_key) + lookup = preset.get("hf_repo") or resolved_name or model_key + required_start, required_end = _preset_layer_bounds(preset) + model_nodes = [ + node for node in server.registry.values() + if _node_matches_preset(node, resolved_name or model_key, preset) # type: ignore[arg-type] + ] + health = _model_health_summary(server, model_key, preset.get("hf_repo")) + local_rpms = server.stats.get_local_rpms() if server.stats is not None else {} + demand_rpm = _repo_demand_rpm(server, str(lookup), local_rpms) + return { + "model": resolved_name or model_key, + "hf_repo": preset.get("hf_repo"), + "preset": preset, + "demand_rpm": demand_rpm, + "served_model_copies": float(health.get("served_model_copies") or 0.0), + "coverage_percentage": float(health.get("coverage_percentage") or 0.0), + "required_start": required_start, + "required_end": required_end, + "model_nodes": model_nodes, + } + + +def _scale_demanded_models_locked(server: "_TrackerHTTPServer") -> None: + """Use spare host RAM/slots to add copies of high-demand models.""" + pool = _memory_pool_map(server) + candidates: list[tuple[float, str, dict]] = [] + for model_key, preset in server.model_presets.items(): + if not isinstance(preset, dict) or not preset.get("hf_repo"): + continue + info = _model_demand_and_supply(server, model_key, preset) + demand = info["demand_rpm"] + copies = info["served_model_copies"] + coverage = info["coverage_percentage"] + if coverage < 100.0: + continue + target_copies = max(1.0, demand / 5.0) + if copies >= target_copies: + continue + candidates.append((demand, model_key, info)) + + if not candidates: + return + + candidates.sort(key=lambda item: (-item[0], item[2]["served_model_copies"])) + for _demand, model_key, info in candidates: + preset = info["preset"] + hf_repo = str(info["hf_repo"]) + required_start = info["required_start"] + required_end = info["required_end"] + total_layers = required_end - required_start + 1 + layer_bytes = next(iter(_preset_bytes_per_layer(preset).values())) + full_copy_bytes = total_layers * layer_bytes + + for host in pool["hosts"]: + if host["spare_slots"] <= 0 or host["memory_spare_bytes"] < full_copy_bytes: + continue + host_nodes = [ + server.registry[node["node_id"]] + for node in host["loaded"] + if node["node_id"] in server.registry + ] + if not host_nodes: + continue + if any((n.hf_repo or n.model) == hf_repo for n in host_nodes): + continue + anchor = max(host_nodes, key=lambda n: n.benchmark_tokens_per_sec) + if anchor.status != "ready" or anchor.pending_new_assignment is not None: + continue + capacity = min(_node_layer_capacity(anchor, preset), total_layers) + if capacity <= 0: + continue + quantization = _node_quantization(anchor, preset) + shard_end = min(required_end, required_start + capacity - 1) + assignment = _add_shard_directive( + anchor, + hf_repo, + required_start, + shard_end, + quantization, + ) + anchor.pending_new_assignment = assignment + anchor.pending_directives.append(assignment) + _tracker_log( + server, + "info", + "scale demanded model copy", + endpoint=anchor.endpoint, + node_id=anchor.node_id, + model=model_key, + hf_repo=hf_repo, + shard=f"{required_start}-{shard_end}", + demand_rpm=round(info["demand_rpm"], 4), + served_model_copies=info["served_model_copies"], + target_copies=round(max(1.0, info["demand_rpm"] / 5.0), 2), + host_spare_bytes=host["memory_spare_bytes"], + ) + break + + +def _deployment_summary(nodes: list[_NodeEntry], preset: dict | None) -> dict: + if preset is None: + return {"recommended": False} + pool = _pool_summary(nodes, preset) + required = int(preset.get("required_model_bytes", 0) or 0) + deployable = required > 0 and pool["effective_assignable_memory_bytes"] >= required + missing = max(0, required - pool["effective_assignable_memory_bytes"]) if required > 0 else 0 + return { + "recommended": bool(preset.get("recommended", False)), + "status": preset.get("deployment_status", "available"), + "required_model_bytes": required or None, + "download_size_bytes": preset.get("download_size_bytes"), + "native_quantization": preset.get("native_quantization"), + "pool": pool, + "deployable": deployable, + "missing_effective_memory_bytes": missing, + } + + +def _max_layers_for_memory( + memory_mb: int, + total_layers: int, + preset: dict | None = None, + *, + device: str | None = None, +) -> int: + if total_layers <= 0: + return 0 + if memory_mb <= 0: + return max(1, total_layers // 2) + memory_bytes = memory_mb * 1024 * 1024 + bytes_per_layer = next(iter(_preset_bytes_per_layer(preset).values())) if preset is not None else 30 * 1024 * 1024 + safety_fraction = 0.55 if device == "cpu" else 0.8 + return min( + total_layers, + max(1, int((memory_bytes * safety_fraction) // bytes_per_layer)), + ) + + +def _model_metadata_from_nodes(nodes: list[_NodeEntry]) -> dict: + metadata: dict = {} + for node in nodes: + if node.model_metadata: + metadata.update(node.model_metadata) + if "num_layers" not in metadata: + layers = [node.num_layers for node in nodes if node.num_layers is not None] + if layers: + metadata["num_layers"] = max(layers) + return metadata + + +def _coverage_map( + nodes: list[_NodeEntry], + required_start: int, + required_end: int, +) -> list[dict]: + layer_counts = [] + for layer in range(required_start, required_end + 1): + count = 0 + for node in nodes: + if node.shard_start is None or node.shard_end is None: + continue + if node.shard_start <= layer <= node.shard_end: + count += 1 + layer_counts.append((layer, count)) + + coverage: list[dict] = [] + for layer, count in layer_counts: + if coverage and coverage[-1]["node_count"] == count and coverage[-1]["end_layer"] == layer - 1: + coverage[-1]["end_layer"] = layer + else: + coverage.append({"start_layer": layer, "end_layer": layer, "node_count": count}) + return coverage + + +def _node_health(node: "_NodeEntry", heartbeat_timeout: float) -> dict: + """Per-node health detail for the availability map.""" + age = time.monotonic() - node.last_heartbeat + alive = age <= heartbeat_timeout + hb_expected = max(1, round(node.uptime_seconds / 20.0)) # assume ~20s interval + hb_rate = round(min(1.0, node.heartbeats_received / hb_expected), 4) if node.heartbeats_received else 0.0 + total = node.total_requests + inf_rate = round((total - node.failed_requests) / total, 4) if total > 0 else 1.0 + return { + "node_id": node.node_id, + "endpoint": node.endpoint, + "alive": alive, + "last_seen_seconds_ago": round(age, 1), + "status": node.status, + "queue_depth": _effective_queue_depth(node), + "heartbeat_queue_depth": node.queue_depth, + "proxy_inflight": node.proxy_inflight, + "current_requests": list(node.current_requests), + "total_requests": node.total_requests, + "heartbeat_success_rate": hb_rate, + "inference_success_rate": inf_rate, + "capacity": _node_capacity_summary(node), + } + + +def _coverage_map_detailed( + nodes: list["_NodeEntry"], + required_start: int, + required_end: int, + heartbeat_timeout: float, +) -> list[dict]: + """Like _coverage_map but with per-node identity and health in each band. + + Includes all nodes (alive and stale) so operators can see both coverage + holes and which dead nodes used to fill them. + """ + now = time.monotonic() + + def covers(node: "_NodeEntry", layer: int) -> bool: + return ( + node.shard_start is not None + and node.shard_end is not None + and node.shard_start <= layer <= node.shard_end + ) + + # Build per-layer list of nodes (with alive flag) + layer_nodes: list[list[tuple["_NodeEntry", bool]]] = [] + for layer in range(required_start, required_end + 1): + ns = [ + (n, (now - n.last_heartbeat) <= heartbeat_timeout) + for n in nodes + if covers(n, layer) + ] + layer_nodes.append(ns) + + # Merge consecutive layers with identical node-set (same node_ids, same alive) + coverage: list[dict] = [] + for i, layer_idx in enumerate(range(required_start, required_end + 1)): + ns = layer_nodes[i] + node_ids = [n.node_id for n, _ in ns] + alive_flags = [a for _, a in ns] + # Compare with last band + if ( + coverage + and coverage[-1]["end_layer"] == layer_idx - 1 + and [nd["node_id"] for nd in coverage[-1]["nodes"]] == node_ids + and [nd["alive"] for nd in coverage[-1]["nodes"]] == alive_flags + ): + coverage[-1]["end_layer"] = layer_idx + else: + coverage.append({ + "start_layer": layer_idx, + "end_layer": layer_idx, + "node_count": len(ns), + "nodes": [_node_health(n, heartbeat_timeout) for n, _ in ns], + }) + return coverage + + +def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]: + return [ + (segment["start_layer"], segment["end_layer"]) + for segment in coverage + if segment["node_count"] == 0 + ] + + +def _unassigned_managed_nodes(nodes: list["_NodeEntry"]) -> list["_NodeEntry"]: + return [ + node for node in nodes + if node.managed_assignment + and (node.shard_start is None or node.shard_end is None) + ] + + +def _emit_shard_change_directives( + node: "_NodeEntry", + model: str, + previous_range: tuple[int | None, int | None, str | None], + preset: dict, +) -> None: + """Queue DROP/LOAD directives when a managed node's shard assignment changes.""" + previous_start, previous_end, previous_quantization = previous_range + current_range = (node.shard_start, node.shard_end, node.quantization) + if node.shard_start is None or node.shard_end is None or current_range == previous_range: + return + if previous_start is not None and previous_end is not None: + node.pending_directives.append( + _drop_directive( + node, + model, + previous_start, + previous_end, + previous_quantization or _node_quantization(node, preset), + ) + ) + node.pending_directives.append( + _load_directive( + node, + model, + node.shard_start, + node.shard_end, + node.quantization or _node_quantization(node, preset), + ) + ) + + +def _assign_redundant_managed_nodes( + managed_nodes: list["_NodeEntry"], + model: str, + preset: dict, + required_start: int, + required_end: int, +) -> None: + """Give newly joined managed nodes their own copy without reshuffling incumbents.""" + total_layers = required_end - required_start + 1 + for node in sorted( + managed_nodes, + key=lambda entry: ( + -entry.benchmark_tokens_per_sec, + -_node_layer_capacity(entry, preset), + entry.node_id, + ), + ): + if node.shard_start is not None and node.shard_end is not None: + continue + capacity = min(_node_layer_capacity(node, preset), total_layers) + if capacity <= 0: + continue + previous_range = (node.shard_start, node.shard_end, node.quantization) + quantization = _node_quantization(node, preset) + node.quantization = quantization + node.shard_start = required_start + node.shard_end = min(required_end, required_start + capacity - 1) + _emit_shard_change_directives(node, model, previous_range, preset) + + +def _relay_http_request_frames( + relay_addr: str, + path: str, + body: bytes, + headers: dict[str, str], + timeout: float = 310.0, + idle_timeout: float = 120.0, + *, + cancel_event: threading.Event | None = None, + ws_holder: list[Any] | None = None, + # Quoted: threading.Lock is a factory function (not a class) before + # Python 3.13, so an unquoted `| None` union crashes at import time. + ws_lock: "threading.Lock | None" = None, +): + """Send an HTTP-shaped request through a relay RPC WebSocket, yielding + response frames until a terminal one (US-036). + + A frame with ``stream: true`` is part of a chunked SSE response ending with + ``done: true``; a frame without ``stream`` is a complete single response. + Yields nothing when the relay is unreachable; stops silently on idle or + overall timeout (the caller bills whatever was observed). + """ + try: + import websockets.sync.client as wsc # type: ignore[import] + except Exception: + return + request_id = str(uuid.uuid4()) + deadline = time.monotonic() + timeout + try: + with wsc.connect(relay_addr, open_timeout=10, close_timeout=5) as ws: + if ws_holder is not None: + if ws_lock is not None: + with ws_lock: + ws_holder.clear() + ws_holder.append(ws) + else: + ws_holder.clear() + ws_holder.append(ws) + ws.send(json.dumps({ + "request_id": request_id, + "method": "POST", + "path": path, + "headers": headers, + "body": body.decode(errors="replace"), + })) + while True: + if cancel_event is not None and cancel_event.is_set(): + return + remaining = deadline - time.monotonic() + if remaining <= 0: + return + raw = ws.recv(timeout=min(idle_timeout, remaining)) + frame = json.loads(raw) + if frame.get("request_id") not in {None, request_id}: + continue + yield frame + if not frame.get("stream") or frame.get("done"): + return + except Exception: + return + + +def _relay_http_request( + relay_addr: str, + path: str, + body: bytes, + headers: dict[str, str], + timeout: float = 310.0, +) -> dict | None: + """Send an HTTP-shaped request through a relay and buffer the response. + + Streamed frame sequences are collapsed into one response dict whose body is + the concatenated SSE text — used by non-chat callers and kept for + backward compatibility; the chat proxy streams frames directly. + """ + frames = _relay_http_request_frames(relay_addr, path, body, headers, timeout=timeout) + first = next(frames, None) + if first is None: + return None + if not first.get("stream"): + return first + chunks = [first.get("chunk") or ""] + for frame in frames: + chunks.append(frame.get("chunk") or "") + return { + "request_id": first.get("request_id"), + "status": first.get("status", 200), + "headers": first.get("headers") or {}, + "body": "".join(chunks), + } + + +def _usage_split(payload: dict) -> dict | None: + """Parse a usage block into {"prompt", "completion", "total"} (ints or None).""" + usage = payload.get("usage") + if not isinstance(usage, dict): + return None + + def _num(key: str) -> int | None: + value = usage.get(key) + return int(value) if isinstance(value, (int, float)) else None + + return { + "prompt": _num("prompt_tokens"), + "completion": _num("completion_tokens"), + "total": _usage_total_tokens(payload), + } + + +def _stream_line_tokens(line: bytes) -> tuple[int, dict | None]: + """Token accounting for one SSE line: (observed output delta, usage split or None).""" + if not line.startswith(b"data:"): + return 0, None + payload = line[5:].strip() + if not payload or payload == b"[DONE]": + return 0, None + try: + chunk_payload = json.loads(payload) + except json.JSONDecodeError: + return 1, None + return _observed_stream_tokens(chunk_payload), _usage_split(chunk_payload) + + +def _stream_billable_split( + observed_output: int, usage: dict | None, request_body: dict +) -> tuple[int, int]: + """(input_tokens, output_tokens) for a streamed response (US-045). + + Output: observed deltas, capped by reported completion count when the + stream carried a usage chunk. Input: reported prompt count, else the + prompt estimate from the request body. + """ + prompt = (usage or {}).get("prompt") + completion = (usage or {}).get("completion") + total = (usage or {}).get("total") + if prompt is None: + prompt = _estimate_prompt_tokens(request_body) or 0 + if completion is None and total is not None: + completion = max(0, total - prompt) + return max(0, prompt), _billable_stream_tokens(observed_output, completion) + + +def _billable_non_stream_split(payload: dict, request_body: dict) -> tuple[int, int]: + """(input_tokens, output_tokens) for a buffered response (US-045). + + Prefers the response usage block; falls back to content estimates. + Completion stays capped by the request's max-tokens bound, as before. + """ + usage = _usage_split(payload) + prompt_estimate = _estimate_prompt_tokens(request_body) or 0 + prompt = (usage or {}).get("prompt") + completion = (usage or {}).get("completion") + if prompt is None: + prompt = prompt_estimate + if completion is None: + total = (usage or {}).get("total") + if total is not None: + completion = max(0, total - prompt) + else: + completion = _observed_non_stream_completion_tokens(payload) + limit = _requested_completion_token_limit(request_body) + if limit is not None and completion > limit: + completion = min(completion, limit) + prompt = max(prompt, prompt_estimate) + return max(0, prompt), max(0, completion) + + +def _find_pinned_route( + nodes: list[_NodeEntry], + required_start: int, + required_end: int, + hop_count: int, +) -> list[_NodeEntry] | None: + """First combination of exactly ``hop_count`` distinct nodes covering the + layer range, where every node extends coverage (US-030 benchmark routes).""" + for combo in itertools.permutations(nodes, hop_count): + covered = required_start - 1 + valid = True + for candidate in combo: + if candidate.shard_start is None or candidate.shard_end is None: + valid = False + break + if candidate.shard_start > covered + 1 or candidate.shard_end <= covered: + valid = False + break + covered = candidate.shard_end + if valid and covered >= required_end: + return list(combo) + return None + + +def _nodes_and_bounds_for_model( + server: "_TrackerHTTPServer", + model: str, +) -> tuple[list[_NodeEntry], int, int] | None: + resolved_name, preset = _resolve_model_preset(server.model_presets, model) + if preset is not None: + required_start, required_end = _preset_layer_bounds(preset) + return [ + node for node in server.registry.values() + if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] + ], required_start, required_end + + nodes = [ + node for node in server.registry.values() + if _node_matches_model(node, model) + and node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + ] + if not nodes: + return None + return nodes, 0, max(node.num_layers for node in nodes) - 1 + + +def _fetch_toploc_commitment( + node: _NodeEntry, + *, + session_id: str, + model: str, + messages: list[dict], +) -> dict | None: + """Fetch a node's own on-demand TOPLOC boundary commitment (ADR-0018 §3), + same protocol as `ValidatorProcess._fetch_hop_commitment`.""" + endpoint = node.endpoint + if not isinstance(endpoint, str) or not endpoint: + return None + try: + req = urllib.request.Request( + f"{endpoint.rstrip('/')}/v1/audit/toploc/commitment", + data=json.dumps({ + "session_id": session_id, + "model": model, + "messages": messages, + "shard_start": node.shard_start, + "shard_end": node.shard_end, + }).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5.0) as resp: + response = json.loads(resp.read()) + except (OSError, ValueError, json.JSONDecodeError): + return None + proof = response.get("toploc_proof") or response.get("activation_proof") + token_ids = response.get("claimed_token_ids") or response.get("output_token_ids") + if not isinstance(proof, dict): + return None + if not isinstance(token_ids, list) or not all(isinstance(t, int) for t in token_ids): + return None + return {"toploc_proof": proof, "claimed_token_ids": token_ids} + + +def _fetch_toploc_reference_activations( + reference_node_url: str, + *, + model: str, + messages: list[dict], + claimed_token_ids: list[int], + claim: Any, +) -> list | None: + """Teacher-force the claimed tokens on the reference node (same contract + as `ValidatorProcess._run_teacher_forced_prefill` / validator README's + "TOPLOC audit contract").""" + try: + req = urllib.request.Request( + f"{reference_node_url.rstrip('/')}/v1/audit/toploc", + data=json.dumps({ + "model": model, + "messages": messages, + "claimed_token_ids": claimed_token_ids, + "dtype": claim.dtype, + "quantization": claim.quantization, + "decode_batching_size": claim.decode_batching_size, + "topk": claim.topk, + "skip_prefill": claim.skip_prefill, + }).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=300.0) as resp: + response = json.loads(resp.read()) + except (OSError, ValueError, json.JSONDecodeError): + return None + activations = response.get("activations") + if not isinstance(activations, list): + return None + return activations + + +def _load_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict: + return { + "action": "LOAD_SHARD", + "model": model, + "start_layer": start, + "end_layer": end, + "shard_start": start, + "shard_end": end, + "quantization": quantization, + } + + +def _drop_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict: + return { + "action": "DROP_SHARD", + "model": model, + "start_layer": start, + "end_layer": end, + "shard_start": start, + "shard_end": end, + "quantization": quantization, + } + + +def _route_stats_keys(server: "_TrackerHTTPServer", entry: "_NodeEntry") -> list[str]: + """All stats keys a node's routes may be recorded under (model, repo, resolved preset).""" + keys = {entry.model, entry.hf_repo} + resolved, _ = _resolve_model_preset(server.model_presets, entry.hf_repo or entry.model) + keys.add(resolved) + return [key for key in keys if key] + + +def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]: + now = time.monotonic() + expired_ids = [ + node_id for node_id, entry in server.registry.items() + if (now - entry.last_heartbeat) > server.heartbeat_timeout + ] + expired_entries: list[tuple[str, _NodeEntry]] = [] + for node_id in expired_ids: + entry = server.registry.pop(node_id) + expired_entries.append((node_id, entry)) + if expired_ids: + _rebalance_all_locked(server) + server.route_stats.bump_epoch( + key + for _, entry in expired_entries + for key in _route_stats_keys(server, entry) + ) + for node_id, entry in expired_entries: + _tracker_log( + server, + "warn", + "node expired", + node_id=node_id, + endpoint=entry.endpoint, + model=entry.model, + hf_repo=entry.hf_repo, + shard=f"{entry.shard_start}-{entry.shard_end}", + heartbeat_timeout_seconds=server.heartbeat_timeout, + model_health=_model_health_summary(server, entry.model, entry.hf_repo), + ) + return expired_ids + + +def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None: + resolved_name, preset = _resolve_model_preset(server.model_presets, model) + if preset is None: + return + required_start, required_end = _preset_layer_bounds(preset) + total_layers = required_end - required_start + 1 + model_nodes = [ + node for node in server.registry.values() + if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] + ] + managed_nodes = [node for node in model_nodes if node.managed_assignment] + if not managed_nodes: + return + + coverage = _coverage_map(model_nodes, required_start, required_end) + gaps = _coverage_gaps(coverage) + unassigned = _unassigned_managed_nodes(managed_nodes) + if not gaps and not unassigned: + return + if not gaps and unassigned: + _assign_redundant_managed_nodes( + unassigned, model, preset, required_start, required_end, + ) + return + + previous_ranges = { + node.node_id: (node.shard_start, node.shard_end, node.quantization) + for node in managed_nodes + } + for node in managed_nodes: + node.shard_start = None + node.shard_end = None + + managed_nodes.sort( + key=lambda node: ( + -node.benchmark_tokens_per_sec, + -_node_layer_capacity(node, preset), + node.node_id, + ) + ) + base_nodes = [node for node in model_nodes if not node.managed_assignment] + base_coverage = _coverage_map(base_nodes, required_start, required_end) + gaps = _coverage_gaps(base_coverage) + + eligible_nodes = [ + node for node in managed_nodes + if _node_layer_capacity(node, preset) > 0 + ] + node_index = 0 + for gap_start, gap_end in gaps: + cursor = gap_start + while cursor <= gap_end and node_index < len(eligible_nodes): + node = eligible_nodes[node_index] + remaining_layers = gap_end - cursor + 1 + remaining_nodes_after = len(eligible_nodes) - node_index - 1 + capacity = min( + _node_layer_capacity(node, preset), + total_layers, + max(1, remaining_layers - remaining_nodes_after), + ) + if capacity <= 0: + node_index += 1 + continue + quantization = _node_quantization(node, preset) + node.quantization = quantization + node.shard_start = cursor + node.shard_end = min(gap_end, cursor + capacity - 1) + cursor = node.shard_end + 1 + node_index += 1 + + for node in managed_nodes: + _emit_shard_change_directives( + node, + model, + previous_ranges[node.node_id], + preset, + ) + + +def _hf_rebalance_preset(nodes: list[_NodeEntry]) -> dict: + total_layers = max(node.num_layers or 0 for node in nodes) + return { + "layers_start": 0, + "layers_end": total_layers - 1, + "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024}, + } + + +def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> None: + model_nodes = [ + node for node in server.registry.values() + if node.hf_repo == hf_repo + and node.num_layers is not None + ] + managed_nodes = [node for node in model_nodes if node.managed_assignment] + if not model_nodes or not managed_nodes: + return + + preset = _hf_rebalance_preset(model_nodes) + required_start, required_end = _preset_layer_bounds(preset) + total_layers = required_end - required_start + 1 + if total_layers <= 0: + return + + coverage = _coverage_map(model_nodes, required_start, required_end) + gaps = _coverage_gaps(coverage) + unassigned = _unassigned_managed_nodes(managed_nodes) + if not gaps and not unassigned: + return + if not gaps and unassigned: + _assign_redundant_managed_nodes( + unassigned, hf_repo, preset, required_start, required_end, + ) + return + + previous_ranges = { + node.node_id: (node.shard_start, node.shard_end, node.quantization) + for node in managed_nodes + } + for node in managed_nodes: + node.shard_start = None + node.shard_end = None + + managed_nodes.sort( + key=lambda node: ( + -node.benchmark_tokens_per_sec, + -_node_layer_capacity(node, preset), + node.node_id, + ) + ) + base_nodes = [node for node in model_nodes if not node.managed_assignment] + base_coverage = _coverage_map(base_nodes, required_start, required_end) + gaps = _coverage_gaps(base_coverage) + + eligible_nodes = [ + node for node in managed_nodes + if _node_layer_capacity(node, preset) > 0 + ] + node_index = 0 + for gap_start, gap_end in gaps: + cursor = gap_start + while cursor <= gap_end and node_index < len(eligible_nodes): + node = eligible_nodes[node_index] + remaining_layers = gap_end - cursor + 1 + remaining_nodes_after = len(eligible_nodes) - node_index - 1 + capacity = min( + _node_layer_capacity(node, preset), + total_layers, + max(1, remaining_layers - remaining_nodes_after), + ) + if capacity <= 0: + node_index += 1 + continue + quantization = _node_quantization(node, preset) + node.quantization = quantization + node.shard_start = cursor + node.shard_end = min(gap_end, cursor + capacity - 1) + cursor = node.shard_end + 1 + node_index += 1 + + for node in managed_nodes: + _emit_shard_change_directives( + node, + hf_repo, + previous_ranges[node.node_id], + preset, + ) + + +def _rebalance_all_locked(server: "_TrackerHTTPServer") -> None: + for model in list(server.model_presets): + _rebalance_model_locked(server, model) + for hf_repo in sorted({node.hf_repo for node in server.registry.values() if node.hf_repo}): + _rebalance_hf_model_locked(server, hf_repo) + + +def _api_key_from_headers(headers) -> str | None: + auth = headers.get("Authorization") + if not auth: + return None + if auth.lower().startswith("bearer "): + return auth.split(" ", 1)[1].strip() or None + return auth.strip() or None + + +def _session_token_from_headers(headers) -> str | None: + token = _api_key_from_headers(headers) + if token: + return token + cookie_header = headers.get("Cookie") + if not cookie_header: + return None + cookie = http.cookies.SimpleCookie() + try: + cookie.load(cookie_header) + except http.cookies.CookieError: + return None + morsel = cookie.get(_SESSION_COOKIE_NAME) + if morsel is None: + return None + return morsel.value.strip() or None + + +def _session_cookie_header(token: str | None) -> str: + cookie = http.cookies.SimpleCookie() + cookie[_SESSION_COOKIE_NAME] = token or "" + morsel = cookie[_SESSION_COOKIE_NAME] + morsel["path"] = "/" + morsel["httponly"] = True + morsel["samesite"] = "Lax" + if token: + morsel["max-age"] = str(int(7 * 86400)) + else: + morsel["max-age"] = "0" + return morsel.OutputString() + + +def _usage_total_tokens(payload: dict) -> int | None: + usage = payload.get("usage") + if not isinstance(usage, dict): + return None + total = usage.get("total_tokens") + if isinstance(total, (int, float)): + return int(total) + prompt = usage.get("prompt_tokens") + completion = usage.get("completion_tokens") + if isinstance(prompt, (int, float)) or isinstance(completion, (int, float)): + return int(prompt or 0) + int(completion or 0) + return None + + +def _estimate_text_tokens(value: Any) -> int | None: + if isinstance(value, str): + text = value.strip() + if not text: + return 0 + return len(text.split()) + if isinstance(value, list): + total = 0 + found = False + for item in value: + if isinstance(item, str): + estimated = _estimate_text_tokens(item) + elif isinstance(item, dict): + estimated = _estimate_text_tokens(item.get("text")) + else: + estimated = None + if estimated is not None: + total += estimated + found = True + return total if found else None + return None + + +def _estimate_prompt_tokens(body: dict) -> int | None: + messages = body.get("messages") + if isinstance(messages, list): + total = 0 + found = False + for message in messages: + if not isinstance(message, dict): + continue + estimated = _estimate_text_tokens(message.get("content")) + if estimated is not None: + total += estimated + found = True + return total if found else None + prompt = body.get("prompt") + return _estimate_text_tokens(prompt) + + +def _requested_completion_token_limit(body: dict) -> int | None: + for field in ("max_completion_tokens", "max_tokens"): + value = body.get(field) + if isinstance(value, bool): + return None + if isinstance(value, (int, float)): + return max(0, int(value)) + return None + + +def _request_total_token_upper_bound(body: dict) -> int | None: + completion_limit = _requested_completion_token_limit(body) + if completion_limit is None: + return None + prompt_estimate = _estimate_prompt_tokens(body) + return completion_limit + (prompt_estimate or 0) + + +def _billable_non_stream_tokens(payload: dict, request_body: dict) -> int: + reported = _usage_total_tokens(payload) + upper_bound = _request_total_token_upper_bound(request_body) + if reported is None: + completion_estimate = _observed_non_stream_completion_tokens(payload) + if completion_estimate > 0: + billable = completion_estimate + (_estimate_prompt_tokens(request_body) or 0) + elif upper_bound is not None: + billable = upper_bound + else: + return 0 + return min(billable, upper_bound) if upper_bound is not None else billable + billable = max(0, reported) + if upper_bound is not None: + billable = min(billable, upper_bound) + return billable + + +def _observed_non_stream_completion_tokens(payload: dict) -> int: + choices = payload.get("choices") + if not isinstance(choices, list): + return 0 + observed = 0 + for choice in choices: + if not isinstance(choice, dict): + continue + message = choice.get("message") + estimated: int | None = None + if isinstance(message, dict): + estimated = _estimate_text_tokens(message.get("content")) + if estimated is None: + estimated = _estimate_text_tokens(choice.get("text")) + if estimated is not None: + observed += estimated + elif choice: + observed += 1 + return observed + + +def _observed_output_from_non_stream_payload(payload: dict) -> str: + choices = payload.get("choices") + if not isinstance(choices, list) or not choices: + return "" + first = choices[0] + if not isinstance(first, dict): + return "" + message = first.get("message") + if isinstance(message, dict) and isinstance(message.get("content"), str): + return message["content"] + text = first.get("text") + return text if isinstance(text, str) else "" + + +def _observed_stream_tokens(payload: dict) -> int: + choices = payload.get("choices") + if not isinstance(choices, list): + return 0 + observed = 0 + for choice in choices: + if not isinstance(choice, dict): + continue + delta = choice.get("delta") + if isinstance(delta, dict): + estimated = _estimate_text_tokens(delta.get("content")) + if estimated: + observed += estimated + continue + if choice: + observed += 1 + return observed + + +def _billable_stream_tokens(observed_tokens: int, reported_tokens: int | None) -> int: + observed = max(0, observed_tokens) + if reported_tokens is None: + return observed + return min(observed, max(0, reported_tokens)) + + +def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None: + if contracts is None or not wallet_address: + return None + if contracts.registry.get_wallet(wallet_address).banned: + return "wallet is banned" + return None + + +def _tracker_log( + server: "_TrackerHTTPServer", + level: str, + message: str, + *, + stdout: bool = True, + update_console_key: str | None = None, + **fields: Any, +) -> None: + log_level = { + "debug": 10, + "info": 20, + "warn": 30, + "warning": 30, + "error": 40, + }.get(level.lower(), 20) + event = { + "ts": time.time(), + "level": level, + "message": message, + "fields": { + key: value + for key, value in fields.items() + if value is not None + }, + } + with server.console_lock: + if update_console_key is not None: + updated = False + for existing in reversed(server.console_events): + if ( + existing.get("message") == message + and existing.get("fields", {}).get("request_id") == update_console_key + ): + existing["ts"] = event["ts"] + existing["fields"] = event["fields"] + updated = True + break + if not updated: + server.console_events.append(event) + else: + server.console_events.append(event) + extras = " ".join(f"{key}={value}" for key, value in event["fields"].items()) + suffix = f" {extras}" if extras else "" + tracker_logger().log(log_level, f"{message}{suffix}") + if stdout: + print(f"[tracker] {level}: {message}{suffix}", flush=True) + + +@dataclass +class _ActiveProxyContext: + request_id: str + cancel_event: threading.Event = field(default_factory=threading.Event) + upstream: Any | None = None + upstream_lock: threading.Lock = field(default_factory=threading.Lock) + relay_ws: Any | None = None + relay_ws_lock: threading.Lock = field(default_factory=threading.Lock) + + +def _register_active_proxy(server: "_TrackerHTTPServer", request_id: str) -> _ActiveProxyContext: + ctx = _ActiveProxyContext(request_id=request_id) + with server.active_proxies_lock: + server.active_proxies[request_id] = ctx + return ctx + + +def _unregister_active_proxy(server: "_TrackerHTTPServer", request_id: str) -> None: + with server.active_proxies_lock: + server.active_proxies.pop(request_id, None) + + +def _request_proxy_cancel(server: "_TrackerHTTPServer", request_id: str) -> bool: + with server.active_proxies_lock: + ctx = server.active_proxies.get(request_id) + if ctx is None: + return False + ctx.cancel_event.set() + + def _close_resources() -> None: + with ctx.upstream_lock: + upstream = ctx.upstream + if upstream is not None: + try: + upstream.close() + except Exception: + pass + with ctx.relay_ws_lock: + relay_ws = ctx.relay_ws + if relay_ws is not None: + try: + relay_ws.close() + except Exception: + pass + + threading.Thread(target=_close_resources, daemon=True).start() + return True + + +def _upstream_socket(upstream: Any) -> Any | None: + fp = getattr(upstream, "fp", None) + raw = getattr(fp, "raw", None) if fp is not None else None + return getattr(raw, "_sock", None) if raw is not None else None + + +def _set_upstream_read_timeout(upstream: Any, timeout: float | None) -> None: + sock = _upstream_socket(upstream) + if sock is not None: + sock.settimeout(timeout) + + +def _clear_proxy_progress_log_state(server: "_TrackerHTTPServer", request_id: str) -> None: + state = getattr(server, "_proxy_progress_log_state", None) + if state is not None: + state.pop(request_id, None) + + +def _tracker_log_proxy_progress( + server: "_TrackerHTTPServer", + *, + request_id: str, + model: str, + route_model: str, + tokens: int, + started: float, + route_nodes: list["_NodeEntry"], + stream: bool = True, + relay: bool = False, +) -> None: + elapsed = time.monotonic() - started + effective_elapsed = max(elapsed, 1e-6) + now = time.monotonic() + state = getattr(server, "_proxy_progress_log_state", None) + if state is None: + state = {} + server._proxy_progress_log_state = state + last_stdout = state.get(request_id) + stdout = last_stdout is None or (now - last_stdout) >= _PROXY_PROGRESS_LOG_INTERVAL + if stdout: + state[request_id] = now + _tracker_log( + server, + "info", + "proxy progress", + stdout=stdout, + update_console_key=request_id, + request_id=request_id, + model=model, + route_model=route_model, + stream=stream, + relay=relay or None, + tokens=tokens, + elapsed_seconds=round(elapsed, 4), + tokens_per_sec=round(tokens / effective_elapsed, 4) if tokens > 0 else 0.0, + route=_node_route_summary(route_nodes), + ) + + +def _node_id_for_registration( + endpoint: str, + model: str, + wallet_address: str | None, + shard_start: int | None, + shard_end: int | None, + hf_repo: str | None, +) -> str: + wallet_prefix = wallet_address[:8] if wallet_address else "anon" + stable_key = "|".join([ + wallet_address or "", + endpoint.rstrip("/"), + model, + hf_repo or "", + "" if shard_start is None else str(shard_start), + "" if shard_end is None else str(shard_end), + ]) + digest = hashlib.sha256(stable_key.encode()).hexdigest()[:12] + return f"{wallet_prefix}-{digest}" + + +class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer): + daemon_threads = True + + def __init__( + self, + addr: tuple, + handler, + registry: dict, + lock: threading.Lock, + heartbeat_timeout: float, + model_presets: dict, + contracts: Any | None, + minimum_stake: int, + relay_url: str | None = None, + raft: "RaftNode | None" = None, + gossip: "NodeGossip | None" = None, + stats: "_StatsCollector | None" = None, + billing: "BillingLedger | None" = None, + accounts: "AccountStore | None" = None, + benchmark_results_path: str | None = None, + validator_service_token: str | None = None, + hive_secret: str | None = None, + max_charge_per_request: float | None = None, + starting_credit: float = DEFAULT_CALLER_CREDIT_USDT, + devnet_topup_amount: float = DEFAULT_DEVNET_TOPUP_USDT, + toploc_calibration: "ToplocCalibrationStore | None" = None, + toploc_reference_node_url: str | None = None, + toploc_calibration_gate_min_hardware_profiles: int = 1, + toploc_backend: Any | None = None, + hf_pricing_log: "HfPricingLog | None" = None, + models_dir: Path | None = None, + route_stats: "RouteStatsStore | None" = None, + ) -> None: + super().__init__(addr, handler) + self.registry = registry + self.lock = lock + self.heartbeat_timeout = heartbeat_timeout + self.model_presets = model_presets + self.contracts = contracts + self.minimum_stake = minimum_stake + self.relay_url = relay_url.rstrip("/") if relay_url else None + self.raft = raft + self.gossip = gossip + self.stats: _StatsCollector | None = stats + self.billing: BillingLedger | None = billing + self.accounts: AccountStore | None = accounts + self.benchmark_results_path = benchmark_results_path or os.path.join( + os.getcwd(), "benchmark_results.json" + ) + self.benchmark_lock = threading.Lock() + self.validator_service_token = validator_service_token + self.hive_secret = hive_secret + self.max_charge_per_request = max_charge_per_request + self.starting_credit = starting_credit + self.devnet_topup_amount = devnet_topup_amount + self.toploc_calibration: ToplocCalibrationStore | None = toploc_calibration + self.toploc_reference_node_url = ( + toploc_reference_node_url.rstrip("/") if toploc_reference_node_url else None + ) + self.toploc_calibration_gate_min_hardware_profiles = toploc_calibration_gate_min_hardware_profiles + self.toploc_backend = toploc_backend + self.hf_pricing_log: HfPricingLog | None = hf_pricing_log + self.models_dir = models_dir + self.console_events = deque(maxlen=_CONSOLE_LIMIT) + self.console_lock = threading.Lock() + self.active_proxies: dict[str, _ActiveProxyContext] = {} + self.active_proxies_lock = threading.Lock() + self.route_stats: RouteStatsStore = route_stats or RouteStatsStore() + self.route_rng = random.Random() + + +class _TrackerHandler(http.server.BaseHTTPRequestHandler): + def log_message(self, fmt, *args): # noqa: suppress request logs in tests + pass + + def _send_json(self, status: int, data: dict, headers: dict[str, str] | None = None) -> None: + body = json.dumps(data).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + for name, value in (headers or {}).items(): + self.send_header(name, value) + self.end_headers() + try: + self.wfile.write(body) + except BrokenPipeError: + pass + + # ---- unified auth boundary (ADR-0017) ---- + + def _resolve_identity(self) -> tuple[str | None, dict | None]: + """Resolve the caller to (role, account). + + Roles: "validator" (service token), "admin"/"user" (session token). + Client API keys resolve to no privileged role — they authorize + inference and wallet binding only, never operator endpoints. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + token = _session_token_from_headers(self.headers) + if not token: + return None, None + if is_validator_token(token, server.validator_service_token): + return "validator", None + if server.accounts is not None: + account = server.accounts.session_account(token) + if account is not None: + return account.get("role", "user"), account + return None, None + + def _require_role(self, *allowed: str) -> bool: + """Gate a privileged handler; sends 401/403 and returns False on failure. + + 401 when no credential was presented; 403 when a credential was + presented but does not resolve to an allowed role — this covers + client API keys and garbage bearer strings on operator endpoints. + """ + role, _account = self._resolve_identity() + if role in allowed: + return True + if _api_key_from_headers(self.headers) is None: + self._send_json(401, {"error": "authentication required (admin session or service token)"}) + else: + self._send_json(403, {"error": "this endpoint requires an admin session or service token"}) + return False + + def _read_hive_authenticated_body(self) -> dict | None: + """Read + verify a hive gossip body (HMAC per ADR-0017 §3). + + Fails closed: without a configured --hive-secret no gossip is + accepted. Sends the error response itself and returns None on failure. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + length = int(self.headers.get("Content-Length", 0)) + raw = self.rfile.read(length) if length else b"{}" + if not verify_hive_request(server.hive_secret, self.headers, raw): + self._send_json(401, {"error": "valid hive signature required"}) + return None + try: + body = json.loads(raw or b"{}") + except json.JSONDecodeError: + self._send_json(400, {"error": "invalid JSON body"}) + return None + if not isinstance(body, dict): + self._send_json(400, {"error": "JSON body must be an object"}) + return None + return body + + def _read_json_body(self) -> dict | None: + length = int(self.headers.get("Content-Length", 0)) + try: + body = json.loads(self.rfile.read(length) or b"{}") + except json.JSONDecodeError: + self._send_json(400, {"error": "invalid JSON body"}) + return None + if not isinstance(body, dict): + self._send_json(400, {"error": "JSON body must be an object"}) + return None + return body + + def _purge_expired_nodes(self) -> None: + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + _purge_expired_nodes_locked(server) + + def do_POST(self): + if self.path == "/v1/chat/completions": + self._handle_proxy_chat() + return + if self.path == "/v1/nodes/register": + self._handle_register() + return + if self.path == "/v1/raft/vote": + self._handle_raft_vote() + return + if self.path == "/v1/raft/append": + self._handle_raft_append() + return + if self.path == "/v1/gossip": + self._handle_gossip() + return + if self.path == "/v1/stats/gossip": + self._handle_stats_gossip() + return + if self.path == "/v1/billing/gossip": + self._handle_billing_gossip() + return + if self.path == "/v1/billing/forfeit": + self._handle_billing_forfeit() + return + if self.path == "/v1/auth/register": + self._handle_auth_register() + return + if self.path == "/v1/auth/login": + self._handle_auth_login() + return + if self.path == "/v1/auth/logout": + self._handle_auth_logout() + return + if self.path == "/v1/account/keys": + self._handle_account_key_create() + return + if self.path == "/v1/account/keys/revoke": + self._handle_account_key_revoke() + return + if self.path == "/v1/account/topup": + self._handle_account_topup() + return + if self.path == "/v1/accounts/gossip": + self._handle_accounts_gossip() + return + if self.path == "/v1/registry/gossip": + self._handle_registry_gossip() + return + if self.path == "/v1/benchmark/hop-penalty": + self._handle_benchmark_hop_penalty() + return + if self.path == "/v1/calibration/toploc/run": + self._handle_toploc_calibration_run() + return + if self.path == "/v1/wallet/register": + self._handle_wallet_register() + return + parts = self.path.split("/") + # /v1/nodes//heartbeat -> ['', 'v1', 'nodes', '', 'heartbeat'] + if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat": + self._handle_heartbeat(parts[3]) + return + # /v1/proxy/requests//cancel + if ( + len(parts) == 6 + and parts[1] == "v1" + and parts[2] == "proxy" + and parts[3] == "requests" + and parts[5] == "cancel" + and parts[4] + ): + self._handle_proxy_request_cancel(urllib.parse.unquote(parts[4])) + return + self.send_response(404) + self.end_headers() + + def do_GET(self): + parsed = urllib.parse.urlparse(self.path) + if parsed.path == "/v1/route": + self._handle_route(parsed) + elif parsed.path == "/v1/routes": + self._handle_routes(parsed) + elif parsed.path == "/v1/routing": + self._handle_routing(parsed) + elif parsed.path == "/v1/nodes/assign": + self._handle_assign(parsed) + elif parsed.path == "/v1/network/assign": + self._handle_network_assign(parsed) + elif parsed.path == "/v1/network/map": + self._handle_network_map() + elif parsed.path == "/v1/models": + self._handle_models() + elif parsed.path == "/v1/model-files/download": + self._handle_model_files_download(parsed) + elif parsed.path.startswith("/v1/coverage/"): + model = urllib.parse.unquote(parsed.path.removeprefix("/v1/coverage/")) + self._handle_coverage(model) + elif parsed.path.startswith("/v1/tracker-nodes/"): + model = urllib.parse.unquote(parsed.path.removeprefix("/v1/tracker-nodes/")) + self._handle_tracker_nodes(model) + elif parsed.path.startswith("/v1/head-workers/"): + model = urllib.parse.unquote(parsed.path.removeprefix("/v1/head-workers/")) + self._handle_tracker_nodes(model) + elif parsed.path == "/v1/raft/status": + self._handle_raft_status() + elif parsed.path == "/v1/stats": + self._handle_stats() + elif parsed.path == "/v1/console": + self._handle_console() + elif parsed.path == "/v1/billing/summary": + self._handle_billing_summary() + elif parsed.path == "/v1/billing/settlements": + self._handle_billing_settlements() + elif parsed.path == "/v1/account": + self._handle_account_me() + elif parsed.path == "/v1/admin/accounts": + self._handle_admin_accounts() + elif parsed.path == "/v1/benchmark/results": + self._handle_benchmark_results() + elif parsed.path == "/v1/calibration/toploc/results": + self._handle_toploc_calibration_results() + elif parsed.path == "/v1/pricing/hf/history": + self._handle_hf_pricing_history(parsed) + elif parsed.path == "/v1/registry/wallets": + self._handle_registry_wallets() + elif parsed.path in ("/dashboard", "/dashboard/"): + self._handle_dashboard() + elif parsed.path == "/v1/health": + self._send_json(200, {"status": "ok"}) + else: + self.send_response(404) + self.end_headers() + + def _handle_models(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + created = int(time.time()) + with server.lock: + self._purge_expired_nodes() + alive = list(server.registry.values()) + if server.contracts is not None: + alive = [ + node for node in alive + if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned + ] + data = [] + seen_ids: set[str] = set() + for name, preset in server.model_presets.items(): + model_nodes = [node for node in alive if _node_matches_preset(node, name, preset)] + if not model_nodes and not preset.get("recommended"): + continue + required_start, required_end = _preset_layer_bounds(preset) + coverage = _coverage_percentage( + model_nodes, + required_start, + required_end, + ) + aliases = [name] + hf_repo = preset.get("hf_repo") + if hf_repo and hf_repo not in aliases: + aliases.append(hf_repo) + for alias in preset.get("aliases", []) or []: + if isinstance(alias, str) and alias not in aliases: + aliases.append(alias) + data.append({ + "id": name, + "object": "model", + "created": created, + "owned_by": "meshnet", + "name": name, + "hf_repo": hf_repo, + "aliases": aliases, + "metadata": dict(preset.get("metadata") or _model_metadata_from_nodes(model_nodes)), + "recommended": bool(preset.get("recommended", False)), + "deployment": _deployment_summary(alive, preset), + "shard_coverage_percentage": coverage, + }) + seen_ids.add(name) + if hf_repo: + seen_ids.add(hf_repo) + + hf_model_ids = sorted({ + node.hf_repo or node.model + for node in alive + if node.model is not None + and node.model not in server.model_presets + and node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + }) + for model_id in hf_model_ids: + if model_id is None or model_id in seen_ids: + continue + model_nodes = [ + node for node in alive + if node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + and (node.hf_repo == model_id or (node.hf_repo is None and node.model == model_id)) + ] + if not model_nodes: + continue + short_names = sorted({node.model for node in model_nodes if node.model}) + aliases = [model_id, *[name for name in short_names if name != model_id]] + required_start = 0 + required_end = max(node.num_layers for node in model_nodes) - 1 + data.append({ + "id": model_id, + "object": "model", + "created": created, + "owned_by": "meshnet", + "name": short_names[0] if short_names else model_id, + "hf_repo": model_id if any(node.hf_repo == model_id for node in model_nodes) else None, + "aliases": aliases, + "metadata": _model_metadata_from_nodes(model_nodes), + "shard_coverage_percentage": _coverage_percentage( + model_nodes, + required_start, + required_end, + ), + }) + seen_ids.add(model_id) + self._send_json(200, {"object": "list", "data": data}) + + def _handle_coverage(self, model: str): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + # Do NOT purge before coverage — dead nodes are included with alive=false + # so operators can see what was covering each layer band before failure. + with server.lock: + resolved = _nodes_and_bounds_for_model(server, model) + if resolved is None: + self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) + return + all_nodes, required_start, required_end = resolved + if server.contracts is not None: + all_nodes = [ + node for node in all_nodes + if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned + ] + coverage = _coverage_map_detailed(all_nodes, required_start, required_end, server.heartbeat_timeout) + self._send_json(200, {"model": model, "coverage": coverage}) + + def _handle_tracker_nodes(self, model: str): + """Return head workers: worker nodes that can start inference for a model. + + The historical endpoint name is /v1/tracker-nodes, but these are not + tracker processes and they are the only machines that load model shards. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + resolved_name, preset = _resolve_model_preset(server.model_presets, model) + if preset is None: + self._send_json(404, {"error": f"unknown model preset: {model!r}"}) + return + required_start, _ = _preset_layer_bounds(preset) + with server.lock: + self._purge_expired_nodes() + alive = [ + node for node in server.registry.values() + if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] + ] + if server.contracts is not None: + alive = [ + node for node in alive + if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned + ] + tracker_nodes = [ + node for node in alive + if node.shard_start is not None + and node.shard_start == required_start + and node.tracker_mode + ] + self._send_json(200, { + "model": resolved_name, + "head_workers": [ + { + "node_id": node.node_id, + "endpoint": node.endpoint, + "relay_addr": node.relay_addr, + "peer_id": node.peer_id, + } + for node in tracker_nodes + ], + "tracker_nodes": [ + { + "node_id": node.node_id, + "endpoint": node.endpoint, + "relay_addr": node.relay_addr, + "peer_id": node.peer_id, + "benchmark_tokens_per_sec": node.benchmark_tokens_per_sec, + } + for node in tracker_nodes + ], + }) + + def _handle_network_map(self) -> None: + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + with server.lock: + self._purge_expired_nodes() + nodes = list(server.registry.values()) + memory_pool = _memory_pool_map(server) + + def capacity_for(node: _NodeEntry) -> dict: + preset = None + if node.model: + preset = server.model_presets.get(node.model) + if preset is None and node.hf_repo and node.num_layers: + preset = _hf_rebalance_preset([node]) + return _node_capacity_summary(node, preset) + + def throughput_for(node: _NodeEntry) -> dict: + if server.stats is None: + return {} + models = [m for m in (node.hf_repo, node.model) if m] + result = {} + for model in models: + result[model] = server.stats.get_node_model_stats(node.node_id, model) + return result + + def model_supply_for(node: _NodeEntry) -> dict: + return _model_health_summary(server, node.model, node.hf_repo) + + self._send_json(200, { + "relay_url": server.relay_url, + "pool": _pool_summary(nodes), + "memory_pool": memory_pool, + "recommended_models": [ + { + "id": name, + "hf_repo": preset.get("hf_repo"), + "aliases": list(preset.get("aliases", []) or []), + "metadata": dict(preset.get("metadata") or {}), + "deployment": _deployment_summary(nodes, preset), + } + for name, preset in server.model_presets.items() + if preset.get("recommended") + ], + "nodes": [ + { + "node_id": node.node_id, + "endpoint": node.endpoint, + "relay_addr": node.relay_addr, + "peer_id": node.peer_id, + "model": node.model, + "hf_repo": node.hf_repo, + "num_layers": node.num_layers, + "model_metadata": dict(node.model_metadata), + "downloaded_models": [dict(item) for item in node.downloaded_models], + "shard_start": node.shard_start, + "shard_end": node.shard_end, + "tracker_mode": node.tracker_mode, + "last_heartbeat": node.last_heartbeat, + "capacity": capacity_for(node), + "model_supply": model_supply_for(node), + "throughput": throughput_for(node), + "stats": _node_health(node, server.heartbeat_timeout), + } + for node in nodes + ], + }) + + # ---------------------------------------------------------------- OpenAI proxy + + def _handle_proxy_chat(self) -> None: + """Proxy POST /v1/chat/completions to a tracker-mode (first-shard) node. + + Picks a live tracker-mode node for the requested model using round-robin, + then forwards the request verbatim and relays the response (including + streaming SSE chunks) back to the caller. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + + length = int(self.headers.get("Content-Length", 0)) + raw_body = self.rfile.read(length) if length else b"{}" + + try: + body = json.loads(raw_body) + except json.JSONDecodeError: + self._send_json(400, {"error": {"message": "invalid JSON", "type": "invalid_request_error", "code": "invalid_request"}}) + return + + model: str = body.get("model", "") + is_stream: bool = bool(body.get("stream", False)) + + if model and server.stats is not None: + server.stats.record_request(model) + + # Billing gate (ADR-0015): reject before any routing — no free work. + api_key = _api_key_from_headers(self.headers) + if server.billing is not None: + if api_key is None: + self._send_json(401, {"error": { + "message": "missing API key: send Authorization: Bearer ", + "type": "invalid_request_error", + "code": "missing_api_key", + }}) + return + if server.accounts is not None and server.accounts.is_key_revoked(api_key): + self._send_json(401, {"error": { + "message": "API key has been revoked", + "type": "invalid_request_error", + "code": "invalid_api_key", + }}) + return + # US-039: with accounts enabled, only real account keys may spend — + # arbitrary bearer strings must never become billable clients. + if server.accounts is not None and not server.accounts.is_active_key(api_key): + self._send_json(401, {"error": { + "message": "unknown API key: create one at /dashboard (register, then + new key)", + "type": "invalid_request_error", + "code": "invalid_api_key", + }}) + return + if not server.billing.has_funds(api_key): + self._send_json(402, {"error": { + "message": "insufficient balance: deposit USDT to continue", + "type": "insufficient_quota", + "code": "insufficient_balance", + }}) + return + if server.max_charge_per_request is not None: + token_limit = _requested_completion_token_limit(body) + if token_limit is None: + self._send_json(400, {"error": { + "message": ( + "max_charge_per_request is enabled; include max_tokens " + "or max_completion_tokens so this request can be capped" + ), + "type": "invalid_request_error", + "code": "missing_token_limit", + }}) + return + in_rate, out_rate = server.billing.prices_for(model) + prompt_estimate = _estimate_prompt_tokens(body) or 0 + estimated_charge = (in_rate * prompt_estimate + out_rate * token_limit) / 1000.0 + if estimated_charge > server.max_charge_per_request: + self._send_json(402, {"error": { + "message": ( + f"request exceeds max_charge_per_request " + f"({estimated_charge:.6f} USDT estimated > " + f"{server.max_charge_per_request:.6f} USDT cap)" + ), + "type": "insufficient_quota", + "code": "spend_cap_exceeded", + }}) + return + + # US-030: optional pinned route — "route": [node_id, ...] uses those + # nodes in order instead of auto-selection. Absent field: unchanged. + pinned_ids = body.get("route") + pinned_nodes: list[_NodeEntry] | None = None + if pinned_ids is not None: + if ( + not isinstance(pinned_ids, list) + or not pinned_ids + or not all(isinstance(nid, str) and nid for nid in pinned_ids) + ): + self._send_json(400, {"error": { + "message": "route must be a non-empty list of node id strings", + "type": "invalid_request_error", + "code": "invalid_route", + }}) + return + with server.lock: + self._purge_expired_nodes() + missing = [nid for nid in pinned_ids if nid not in server.registry] + if missing: + self._send_json(400, {"error": { + "message": f"unknown node ids in route: {missing}", + "type": "invalid_request_error", + "code": "unknown_route_nodes", + }}) + return + pinned_nodes = [server.registry[nid] for nid in pinned_ids] + + if pinned_nodes is not None: + node = pinned_nodes[0] + else: + # Find a live tracker-mode node for this model + with server.lock: + self._purge_expired_nodes() + candidates = [ + n for n in server.registry.values() + if n.tracker_mode and _node_matches_model(n, model) + ] + + if not candidates: + # Fall back: any node serving shard_start=0 for this model + with server.lock: + candidates = [ + n for n in server.registry.values() + if n.shard_start == 0 and _node_matches_model(n, model) + ] + + if not candidates: + with server.lock: + registered = [ + { + "node_id": n.node_id, + "model": n.model, + "hf_repo": n.hf_repo, + "shard": f"{n.shard_start}-{n.shard_end}", + "tracker_mode": n.tracker_mode, + } + for n in server.registry.values() + ] + _tracker_log( + server, + "warn", + "no nodes available for model", + model=model, + registered_nodes=registered, + ) + self._send_json(503, {"error": { + "message": f"no nodes available for model {model!r}", + "type": "service_unavailable", + "code": "model_not_available", + }}) + return + + node = max(candidates, key=lambda n: _effective_throughput(n, model)) + target_url = f"{node.endpoint}/v1/chat/completions" + request_id = str(body.get("id") or f"req-{time.time_ns():x}") + body["id"] = request_id + raw_body = json.dumps(body).encode() + + # Pre-resolve the downstream route so the first-shard node skips its own + # tracker query. We already hold the full registry picture — no need for + # a second round-trip. + route_model = node.hf_repo or node.model or model + with server.lock: + resolved_route_model, route_preset = _resolve_model_preset(server.model_presets, route_model) + if route_preset is not None: + route_model = resolved_route_model or route_model + preset = route_preset + rs, re = _preset_layer_bounds(preset) + all_nodes: list = [ + n for n in server.registry.values() + if _node_matches_preset(n, route_model, preset) + and n.shard_start is not None + and n.shard_end is not None + ] + else: + all_nodes = [ + n for n in server.registry.values() + if _node_matches_model(n, route_model) + and n.shard_start is not None and n.num_layers is not None + ] + rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1) + if pinned_nodes is not None: + route_nodes = pinned_nodes + routing_decision = {"mode": "pinned"} + else: + # ADR-0021: enumerate viable routes and pick one bandit-style — + # ε-scout among unproven routes, otherwise weighted ∝ observed tps^α. + route_candidates = _enumerate_routes( + all_nodes, rs, re, + model=route_model, + contracts=server.contracts, + max_candidates=server.route_stats.config.max_candidate_routes, + ) + picked, routing_decision = choose_route( + route_candidates, server.route_stats, route_model, rng=server.route_rng, + ) + if picked is not None: + route_nodes = picked.nodes + else: + # No head-anchored candidate — legacy greedy cover as fallback + # (also produces the layer-gap error message). + route_nodes, route_error = _select_route(all_nodes, rs, re, model=route_model, contracts=server.contracts) + routing_decision = {"mode": "greedy-fallback"} + if route_error: + _tracker_log( + server, + "warn", + "route unavailable", + model=model, + route_model=route_model, + error=route_error, + candidate_count=len(all_nodes), + candidates=_node_route_summary(all_nodes), + ) + self._send_json(503, {"error": { + "message": route_error, + "type": "service_unavailable", + "code": "route_not_available", + }}) + return + # The proxy target must be the route's own head: an independently + # chosen fastest node may not be part of the planned route, which + # previously injected downstream hops with wrong start layers + # (ADR-0020 mixed-topology flaw). + if route_nodes: + node = route_nodes[0] + target_url = f"{node.endpoint}/v1/chat/completions" + # Compute start_layer for each hop: each node begins where the previous ended + 1. + # This allows overlapping shard registrations without double-computation. + covered_up_to = rs - 1 + route_hops: list[dict] = [] + node_work: list[tuple[str | None, int]] = [] + for rn in route_nodes: + hop: dict = {"endpoint": rn.endpoint, "start_layer": covered_up_to + 1} + if rn.relay_addr: + hop["relay_addr"] = rn.relay_addr + route_hops.append(hop) + effective_end = rn.shard_end if rn.shard_end is not None else covered_up_to + node_work.append((rn.wallet_address, max(0, effective_end - covered_up_to))) + covered_up_to = effective_end + # Strip the first-shard node we're about to proxy to — it's already handling the request. + downstream_hops = [ + h for h in route_hops + if _endpoint_key(h["endpoint"]) != _endpoint_key(node.endpoint) + ] + downstream_urls = json.dumps(downstream_hops) + route_debug = " -> ".join( + f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]" + for n in route_nodes + ) + inflight_nodes = route_nodes or [node] + inflight_recorded = True + _record_proxy_inflight(server, inflight_nodes, 1) + + def finish_proxy_inflight() -> None: + nonlocal inflight_recorded + if inflight_recorded: + _record_proxy_inflight(server, inflight_nodes, -1) + inflight_recorded = False + _unregister_active_proxy(server, request_id) + + proxy_ctx = _register_active_proxy(server, request_id) + + _tracker_log( + server, + "info", + "proxy route selected", + request_id=request_id, + model=model, + head_node_id=node.node_id, + head_endpoint=node.endpoint, + downstream=downstream_urls, + route=route_debug or "", + routing=routing_decision, + nodes=_node_route_summary(route_nodes), + ) + + req = urllib.request.Request( + target_url, + data=raw_body, + headers={ + "Content-Type": "application/json", + "X-Meshnet-Route": downstream_urls, + "X-Meshnet-Request-Id": request_id, + }, + method="POST", + ) + # Copy Authorization header from client if present + auth = self.headers.get("Authorization") + if auth: + req.add_header("Authorization", auth) + + relay_headers = { + "Content-Type": "application/json", + "X-Meshnet-Route": downstream_urls, + "X-Meshnet-Request-Id": request_id, + **({"Authorization": auth} if auth else {}), + } + + if node.relay_addr: + _tracker_log( + server, + "info", + "proxy via relay", + request_id=request_id, + relay_addr=node.relay_addr, + direct_endpoint=target_url, + ) + started = time.monotonic() + relay_ws_holder: list[Any] = [] + frames = _relay_http_request_frames( + node.relay_addr, + path="/v1/chat/completions", + body=raw_body, + headers=relay_headers, + cancel_event=proxy_ctx.cancel_event, + ws_holder=relay_ws_holder, + ws_lock=proxy_ctx.relay_ws_lock, + ) + first = next(frames, None) + with proxy_ctx.relay_ws_lock: + proxy_ctx.relay_ws = relay_ws_holder[0] if relay_ws_holder else None + if proxy_ctx.cancel_event.is_set(): + if self._finalize_proxy_cancel( + proxy_ctx=proxy_ctx, + server=server, + request_id=request_id, + started=started, + model=model, + route_model=route_model, + route_nodes=route_nodes, + api_key=api_key, + node_work=node_work, + body=body, + finish_proxy_inflight=finish_proxy_inflight, + ): + return + if first is not None and first.get("stream"): + # Streamed response (US-036): forward SSE chunks as they arrive + # and run the same token accounting as the direct stream path. + self._stream_relayed_frames( + first, frames, started, + model, route_model, route_nodes, api_key, node_work, + request_body=body, + request_id=request_id, + proxy_ctx=proxy_ctx, + finish_proxy_inflight=finish_proxy_inflight, + ) + finish_proxy_inflight() + return + if first is not None: + elapsed = time.monotonic() - started + self._send_relayed_response(first) + if int(first.get("status", 503)) < 400: + body_text = first.get("body") or "" + try: + in_tokens, out_tokens = _billable_non_stream_split(json.loads(body_text), body) + except (json.JSONDecodeError, TypeError): + in_tokens, out_tokens = 0, 0 + tokens = in_tokens + out_tokens + self._record_observed_throughput(model, route_model, tokens, elapsed, route_nodes) + _clear_proxy_progress_log_state(server, request_id) + _tracker_log( + server, + "info", + "proxy complete", + request_id=request_id, + model=model, + route_model=route_model, + status=int(first.get("status", 503)), + tokens=tokens, + elapsed_seconds=round(elapsed, 4), + tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, + route=_node_route_summary(route_nodes), + ) + self._bill_completed( + api_key, model, tokens, node_work, + input_tokens=in_tokens, output_tokens=out_tokens, + ) + finish_proxy_inflight() + return + _tracker_log( + server, + "warn", + "relay proxy failed, trying direct", + request_id=request_id, + relay_addr=node.relay_addr, + direct_endpoint=target_url, + ) + + try: + started = time.monotonic() + _tracker_log( + server, + "info", + "proxy connecting", + request_id=request_id, + target_url=target_url, + stream=is_stream or None, + ) + upstream_result: list[Any] = [] + connect_errors: list[BaseException] = [] + + def _connect_upstream() -> None: + try: + upstream_result.append(urllib.request.urlopen(req, timeout=300.0)) + except BaseException as exc: + connect_errors.append(exc) + + connect_thread = threading.Thread(target=_connect_upstream, daemon=True) + connect_thread.start() + while connect_thread.is_alive(): + if proxy_ctx.cancel_event.is_set(): + connect_thread.join(timeout=310.0) + if upstream_result: + try: + upstream_result[0].close() + except Exception: + pass + if self._finalize_proxy_cancel( + proxy_ctx=proxy_ctx, + server=server, + request_id=request_id, + started=started, + model=model, + route_model=route_model, + route_nodes=route_nodes, + api_key=api_key, + node_work=node_work, + body=body, + finish_proxy_inflight=finish_proxy_inflight, + ): + return + connect_thread.join(0.2) + + if proxy_ctx.cancel_event.is_set(): + if upstream_result: + try: + upstream_result[0].close() + except Exception: + pass + if self._finalize_proxy_cancel( + proxy_ctx=proxy_ctx, + server=server, + request_id=request_id, + started=started, + model=model, + route_model=route_model, + route_nodes=route_nodes, + api_key=api_key, + node_work=node_work, + body=body, + finish_proxy_inflight=finish_proxy_inflight, + ): + return + + if connect_errors: + raise connect_errors[0] + + upstream = upstream_result[0] + with proxy_ctx.upstream_lock: + proxy_ctx.upstream = upstream + upstream_sock = _upstream_socket(upstream) + if upstream_sock is not None: + _set_upstream_read_timeout(upstream, None) + else: + _set_upstream_read_timeout(upstream, 0.5) + _tracker_log(server, "info", "proxy connected", request_id=request_id, target_url=target_url) + except urllib.error.HTTPError as exc: + # Relay error status + body from node + err_body = exc.read() + self.send_response(exc.code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(err_body))) + self.end_headers() + try: + self.wfile.write(err_body) + except BrokenPipeError: + pass + _clear_proxy_progress_log_state(server, request_id) + finish_proxy_inflight() + return + except Exception as exc: + _clear_proxy_progress_log_state(server, request_id) + if node.relay_addr: + _tracker_log( + server, + "error", + "direct proxy failed after relay", + request_id=request_id, + target_url=target_url, + relay_addr=node.relay_addr, + error=repr(exc), + ) + else: + _tracker_log( + server, + "error", + "proxy failed", + request_id=request_id, + target_url=target_url, + error=repr(exc), + ) + self._send_json(503, {"error": { + "message": f"upstream node unreachable: {exc}", + "type": "service_unavailable", + "code": "upstream_error", + }}) + finish_proxy_inflight() + return + + with upstream: + content_type = upstream.headers.get("Content-Type", "application/json") + if is_stream or "text/event-stream" in content_type: + # Relay SSE stream chunk-by-chunk + self.send_response(200) + self.send_header("Content-Type", "text/event-stream; charset=utf-8") + self.send_header("Cache-Control", "no-cache") + self.send_header("X-Accel-Buffering", "no") + self.end_headers() + stream_usage: dict | None = None + observed_stream_tokens = 0 + client_gone = False + try: + while True: + if proxy_ctx.cancel_event.is_set(): + break + if upstream_sock is not None: + readable, _, _ = select.select([upstream_sock], [], [], 0.5) + if not readable: + continue + try: + line = upstream.readline() + except TimeoutError: + continue + if not line: + if proxy_ctx.cancel_event.is_set(): + break + break + if not client_gone: + try: + self.wfile.write(line) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + # Keep draining upstream so completed node work is still billed. + client_gone = True + observed, usage = _stream_line_tokens(line) + observed_stream_tokens += observed + if observed: + _tracker_log_proxy_progress( + server, + request_id=request_id, + model=model, + route_model=route_model, + tokens=observed_stream_tokens, + started=started, + route_nodes=route_nodes, + ) + if usage is not None: + stream_usage = usage + except (BrokenPipeError, ConnectionResetError): + client_gone = True + if self._finalize_proxy_cancel( + proxy_ctx=proxy_ctx, + server=server, + request_id=request_id, + started=started, + model=model, + route_model=route_model, + route_nodes=route_nodes, + api_key=api_key, + node_work=node_work, + body=body, + finish_proxy_inflight=finish_proxy_inflight, + observed_stream_tokens=observed_stream_tokens, + stream_usage=stream_usage, + ): + return + elapsed = time.monotonic() - started + # Bill even on client disconnect — the nodes did the work. + # Observed stream chunks are authoritative for the upper bound; + # upstream usage may only lower that count. + in_tokens, out_tokens = _stream_billable_split( + observed_stream_tokens, stream_usage, body + ) + self._record_observed_throughput( + model, route_model, in_tokens + out_tokens, elapsed, route_nodes + ) + tokens = in_tokens + out_tokens + _clear_proxy_progress_log_state(server, request_id) + _tracker_log( + server, + "info", + "proxy complete", + request_id=request_id, + model=model, + route_model=route_model, + status=200, + stream=True, + client_disconnected=client_gone, + tokens=tokens, + elapsed_seconds=round(elapsed, 4), + tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, + route=_node_route_summary(route_nodes), + ) + self._bill_completed( + api_key, model, tokens, node_work, + input_tokens=in_tokens, output_tokens=out_tokens, + ) + else: + # Non-streaming: buffer and relay + resp_body = upstream.read() + elapsed = time.monotonic() - started + observed_output = "" + try: + response_payload = json.loads(resp_body) + in_tokens, out_tokens = _billable_non_stream_split(response_payload, body) + observed_output = _observed_output_from_non_stream_payload(response_payload) + except json.JSONDecodeError: + in_tokens, out_tokens = 0, 0 + tokens = in_tokens + out_tokens + self._record_observed_throughput(model, route_model, tokens, elapsed, route_nodes) + _clear_proxy_progress_log_state(server, request_id) + _tracker_log( + server, + "info", + "proxy complete", + request_id=request_id, + model=model, + route_model=route_model, + target_url=target_url, + status=200, + bytes=len(resp_body), + tokens=tokens, + elapsed_seconds=round(elapsed, 4), + tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, + route=_node_route_summary(route_nodes), + ) + self._record_validation_event(request_id, model, body, observed_output, route_nodes) + self._bill_completed( + api_key, model, tokens, node_work, + input_tokens=in_tokens, output_tokens=out_tokens, + ) + self.send_response(200) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(resp_body))) + self.end_headers() + try: + self.wfile.write(resp_body) + except (BrokenPipeError, ConnectionResetError): + pass + finish_proxy_inflight() + + def _record_observed_throughput( + self, + requested_model: str, + route_model: str, + total_tokens: int, + elapsed_seconds: float, + route_nodes: list[_NodeEntry], + ) -> None: + """Record observed route TPS for participating nodes. + + The tracker sees end-to-end request duration, not per-hop timings, so + each hop gets the same route-level observation for now. Per-hop telemetry + can refine this later without changing the external stats shape. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if route_model and route_nodes: + # Route-level bandit sample (ADR-0021); the store itself rejects + # near-empty completions that would poison the arm. + server.route_stats.record_sample( + route_model, + route_signature(route_model, route_nodes), + total_tokens, + elapsed_seconds, + ) + if server.stats is None or total_tokens <= 0: + return + elapsed_seconds = max(elapsed_seconds, 1e-6) + models = [m for m in (requested_model, route_model) if m] + if len(models) == 2 and models[0] == models[1]: + models = [models[0]] + for node in route_nodes: + for model in models: + server.stats.record_node_throughput( + node.node_id, + model, + total_tokens=total_tokens, + elapsed_seconds=elapsed_seconds, + ) + stats = server.stats.get_node_model_stats(node.node_id, model) + observed = stats.get("tokens_per_sec_last_hour") + if observed is not None: + node.model_tokens_per_sec[model] = float(observed) + + def _record_validation_event( + self, + session_id: str, + model: str, + request_body: dict, + observed_output: str, + route_nodes: list[_NodeEntry], + ) -> None: + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + validation = getattr(server.contracts, "validation", None) if server.contracts is not None else None + if validation is None: + return + messages = request_body.get("messages") + if not isinstance(messages, list): + messages = [] + event_nodes = [ + { + "node_id": node.node_id, + "endpoint": node.endpoint, + "wallet_address": node.wallet_address, + "shard_start": node.shard_start, + "shard_end": node.shard_end, + } + for node in route_nodes + ] + try: + validation.record_completed_inference( + session_id=session_id, + model=model, + messages=[dict(message) for message in messages if isinstance(message, dict)], + observed_output=observed_output, + route_nodes=event_nodes, + ) + except Exception as exc: + print(f"[tracker] validation event recording failed for {session_id}: {exc}", flush=True) + + def _bill_completed( + self, + api_key: str | None, + model: str, + total_tokens: int, + node_work: list[tuple[str | None, int]], + *, + input_tokens: int | None = None, + output_tokens: int | None = None, + ) -> None: + """Charge a completed request against the billing ledger (ADR-0015). + + With ``input_tokens``/``output_tokens`` the ledger bills each side at + its own rate (US-045); ``total_tokens`` alone falls back to the + blended rate. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.billing is None or api_key is None: + return + # Probationary period (issue 08): a wallet's first N jobs earn nothing — + # its share stays in the protocol cut. Job counts live in the registry + # contract, so this only applies when the tracker runs with contracts. + if server.contracts is not None: + adjusted: list[tuple[str | None, int]] = [] + for wallet, work in node_work: + if wallet: + in_probation = server.contracts.registry.probationary_jobs_remaining(wallet) > 0 + server.contracts.registry.record_completed_job(wallet) + if in_probation: + wallet = None + adjusted.append((wallet, work)) + node_work = adjusted + try: + event = server.billing.charge_request( + api_key, model, total_tokens, node_work, + input_tokens=input_tokens, output_tokens=output_tokens, + ) + print( + f"[tracker] billed api_key=…{api_key[-6:]}: model={model!r} " + f"tokens={event['total_tokens']} " + f"(in={event.get('input_tokens', '?')} out={event.get('output_tokens', '?')}) " + f"cost={event['cost']:.6f} USDT shares={event['shares']}", + flush=True, + ) + except Exception as exc: + print(f"[tracker] billing failed for model={model!r}: {exc}", flush=True) + + def _stream_relayed_frames( + self, + first: dict, + frames, + started: float, + model: str, + route_model: str, + route_nodes: list, + api_key: str | None, + node_work: list, + request_body: dict, + request_id: str, + *, + proxy_ctx: _ActiveProxyContext | None = None, + finish_proxy_inflight: Any = None, + ) -> None: + """Forward a streamed relay response (US-036) to the client as SSE, + billing with the same accounting as the direct stream path.""" + headers = first.get("headers") if isinstance(first.get("headers"), dict) else {} + self.send_response(int(first.get("status", 200))) + self.send_header("Content-Type", headers.get("Content-Type", "text/event-stream; charset=utf-8")) + self.send_header("Cache-Control", "no-cache") + self.send_header("X-Accel-Buffering", "no") + self.end_headers() + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + stream_usage: dict | None = None + observed_stream_tokens = 0 + client_gone = False + for frame in itertools.chain([first], frames): + if proxy_ctx is not None and proxy_ctx.cancel_event.is_set(): + break + chunk = frame.get("chunk") or "" + if not chunk: + continue + data = chunk.encode() + if not client_gone: + try: + self.wfile.write(data) + self.wfile.flush() + except BrokenPipeError: + # Keep draining frames — the nodes did the work; bill it. + client_gone = True + for line in data.splitlines(): + observed, usage = _stream_line_tokens(line) + observed_stream_tokens += observed + if observed: + _tracker_log_proxy_progress( + server, + request_id=request_id, + model=model, + route_model=route_model, + tokens=observed_stream_tokens, + started=started, + route_nodes=route_nodes, + relay=True, + ) + if usage is not None: + stream_usage = usage + if ( + proxy_ctx is not None + and finish_proxy_inflight is not None + and self._finalize_proxy_cancel( + proxy_ctx=proxy_ctx, + server=server, + request_id=request_id, + started=started, + model=model, + route_model=route_model, + route_nodes=route_nodes, + api_key=api_key, + node_work=node_work, + body=request_body, + finish_proxy_inflight=finish_proxy_inflight, + observed_stream_tokens=observed_stream_tokens, + stream_usage=stream_usage, + ) + ): + return + elapsed = time.monotonic() - started + in_tokens, out_tokens = _stream_billable_split( + observed_stream_tokens, stream_usage, request_body + ) + self._record_observed_throughput( + model, route_model, in_tokens + out_tokens, elapsed, route_nodes + ) + tokens = in_tokens + out_tokens + _clear_proxy_progress_log_state(server, request_id) + _tracker_log( + server, + "info", + "proxy complete", + request_id=request_id, + model=model, + route_model=route_model, + status=200, + stream=True, + client_disconnected=client_gone, + relay=True, + tokens=tokens, + elapsed_seconds=round(elapsed, 4), + tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 else 0.0, + route=_node_route_summary(route_nodes), + ) + self._bill_completed( + api_key, model, tokens, node_work, + input_tokens=in_tokens, output_tokens=out_tokens, + ) + + def _send_relayed_response(self, response: dict) -> None: + status = int(response.get("status", 503)) + headers = response.get("headers") if isinstance(response.get("headers"), dict) else {} + body_text = response.get("body") or "" + body = body_text.encode() if isinstance(body_text, str) else bytes(body_text) + self.send_response(status) + self.send_header("Content-Type", headers.get("Content-Type", "application/json")) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + try: + self.wfile.write(body) + except BrokenPipeError: + pass + + def _handle_register(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_json_body() + if body is None: + return + + # --- Raft cluster mode: forward to leader or propose via Raft --- + if server.raft is not None: + if not server.raft.is_leader: + leader = server.raft.leader() + if leader is None: + self._send_json(503, {"error": "no leader elected — retry in a moment"}) + return + # Proxy to leader + try: + data = json.dumps(body).encode() + req = urllib.request.Request( + f"{leader}/v1/nodes/register", + data=data, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5.0) as r: + resp_body = r.read() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(resp_body))) + self.end_headers() + self.wfile.write(resp_body) + except Exception as exc: + self._send_json(503, {"error": f"leader proxy failed: {exc}"}) + return + # Leader path: fall through to normal registration, then replicate via Raft. + # We let the registration run first (to generate node_id), then asynchronously + # propose to the Raft log so followers can replicate the entry. + # Raft proposal happens after the response is sent (fire-and-forget replication). + + endpoint = body.get("endpoint") + if not isinstance(endpoint, str) or not endpoint: + self._send_json(400, {"error": "endpoint is required"}) + return + parsed_endpoint = urllib.parse.urlparse(endpoint) + if parsed_endpoint.scheme not in {"http", "https"} or not parsed_endpoint.netloc: + self._send_json(400, {"error": "endpoint must be an http(s) URL"}) + return + + shard_start: int | None + shard_end: int | None + explicit_shard = "shard_start" in body or "shard_end" in body + if explicit_shard: + try: + shard_start = int(body["shard_start"]) + shard_end = int(body["shard_end"]) + except (KeyError, TypeError, ValueError): + self._send_json(400, {"error": "shard_start and shard_end must be numeric"}) + return + if shard_start < 0 or shard_end < 0 or shard_start > shard_end: + self._send_json(400, {"error": "shard range must be non-negative and ordered"}) + return + else: + shard_start = None + shard_end = None + try: + score = float(body.get("score", 1.0)) + except (TypeError, ValueError): + self._send_json(400, {"error": "score must be numeric"}) + return + + hardware_profile = body.get("hardware_profile", {}) + if not isinstance(hardware_profile, dict): + self._send_json(400, {"error": "hardware_profile must be an object"}) + return + model = body.get("model") + if model is None: + model = "stub-model" + if not isinstance(model, str): + self._send_json(400, {"error": "model must be a string"}) + return + shard_checksum = body.get("shard_checksum") + if shard_checksum is not None and not isinstance(shard_checksum, str): + self._send_json(400, {"error": "shard_checksum must be a string"}) + return + try: + vram_bytes = int(body.get("vram_bytes", DEFAULT_VRAM_BYTES)) + ram_bytes = int(body.get("ram_bytes", DEFAULT_RAM_BYTES)) + max_loaded_shards = int(body.get("max_loaded_shards", 1)) + benchmark_tokens_per_sec = float( + body.get("benchmark_tokens_per_sec", DEFAULT_BENCHMARK_TOKENS_PER_SEC) + ) + except (TypeError, ValueError): + self._send_json(400, {"error": "vram_bytes, ram_bytes, max_loaded_shards, and benchmark_tokens_per_sec must be numeric"}) + return + if vram_bytes < 0 or ram_bytes < 0 or max_loaded_shards < 1 or benchmark_tokens_per_sec <= 0: + self._send_json(400, {"error": "capability values must be positive"}) + return + quantizations_body = body.get("quantizations", DEFAULT_QUANTIZATIONS) + if not ( + isinstance(quantizations_body, list) + and quantizations_body + and all(isinstance(item, str) and item for item in quantizations_body) + ): + self._send_json(400, {"error": "quantizations must be a non-empty string array"}) + return + quantizations = list(quantizations_body) + quantization = body.get("quantization") + if quantization is not None and not isinstance(quantization, str): + self._send_json(400, {"error": "quantization must be a string"}) + return + wallet_address = body.get("wallet_address") + if wallet_address is not None and not isinstance(wallet_address, str): + self._send_json(400, {"error": "wallet_address must be a string"}) + return + ban_error = _registration_ban_error(server.contracts, wallet_address) + if ban_error: + self._send_json(403, {"error": ban_error}) + return + + tracker_mode = bool(body.get("tracker_mode", False)) + managed_assignment = bool(body.get("managed_assignment", False)) + hf_repo = body.get("hf_repo") + if hf_repo is not None and not isinstance(hf_repo, str): + self._send_json(400, {"error": "hf_repo must be a string"}) + return + num_layers_body = body.get("num_layers") + num_layers: int | None = None + if num_layers_body is not None: + try: + num_layers = int(num_layers_body) + except (TypeError, ValueError): + self._send_json(400, {"error": "num_layers must be an integer"}) + return + model_metadata = body.get("model_metadata", {}) + if model_metadata is None: + model_metadata = {} + if not isinstance(model_metadata, dict): + self._send_json(400, {"error": "model_metadata must be an object"}) + return + downloaded_models = body.get("downloaded_models", []) + if downloaded_models is None: + downloaded_models = [] + if not ( + isinstance(downloaded_models, list) + and all(isinstance(item, dict) for item in downloaded_models) + ): + self._send_json(400, {"error": "downloaded_models must be an array of objects"}) + return + relay_addr = body.get("relay_addr") or None + cert_fingerprint = body.get("cert_fingerprint") or None + peer_id = body.get("peer_id") or None + + node_id = _node_id_for_registration( + endpoint, + model, + wallet_address, + shard_start, + shard_end, + hf_repo, + ) + entry = _NodeEntry( + node_id=node_id, + endpoint=endpoint.rstrip("/"), + shard_start=shard_start, + shard_end=shard_end, + model=model, + shard_checksum=shard_checksum, + hardware_profile=hardware_profile, + wallet_address=wallet_address, + score=score, + vram_bytes=vram_bytes, + ram_bytes=ram_bytes, + quantizations=quantizations, + max_loaded_shards=max_loaded_shards, + benchmark_tokens_per_sec=benchmark_tokens_per_sec, + quantization=quantization, + managed_assignment=managed_assignment or not explicit_shard, + tracker_mode=tracker_mode, + hf_repo=hf_repo, + num_layers=num_layers, + model_metadata=model_metadata, + downloaded_models=downloaded_models, + relay_addr=relay_addr, + cert_fingerprint=cert_fingerprint, + peer_id=peer_id, + ) + with server.lock: + self._purge_expired_nodes() + # Dedup: replace the same node id or the same endpoint+model assignment. + endpoint_key = entry.endpoint.rstrip("/") + model_key = entry.hf_repo or entry.model + stale_ids = [ + eid for eid, e in server.registry.items() + if eid == node_id + or ( + e.endpoint.rstrip("/") == endpoint_key + and (e.hf_repo or e.model) == model_key + ) + ] + stale_entries: list[tuple[str, _NodeEntry]] = [] + for eid in stale_ids: + old = server.registry.pop(eid) + stale_entries.append((eid, old)) + is_topology_change = node_id not in stale_ids or any( + (old.shard_start, old.shard_end) != (entry.shard_start, entry.shard_end) + for eid, old in stale_entries + if eid == node_id + ) + server.registry[node_id] = entry + if is_topology_change: + server.route_stats.bump_epoch(_route_stats_keys(server, entry)) + if entry.managed_assignment and not explicit_shard: + if entry.hf_repo: + _rebalance_hf_model_locked(server, entry.hf_repo) + else: + _rebalance_model_locked(server, model) + assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None + if assignment_directive is not None: + entry.pending_directives.clear() + + model_health = _model_health_summary(server, entry.model, entry.hf_repo) + for eid, old in stale_entries: + _tracker_log( + server, + "info", + "node re-registered", + old_node_id=eid, + node_id=node_id, + endpoint=old.endpoint, + old_model=old.model, + old_hf_repo=old.hf_repo, + old_model_health=_model_health_summary(server, old.model, old.hf_repo), + model_health=model_health, + ) + _tracker_log( + server, + "info", + "node registered", + node_id=node_id, + endpoint=entry.endpoint, + model=entry.model, + hf_repo=entry.hf_repo, + shard=f"{entry.shard_start}-{entry.shard_end}", + tracker_mode=entry.tracker_mode, + model_health=model_health, + ) + + shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded" + repo_info = f" [{hf_repo}]" if hf_repo else "" + budget_bytes, budget_source = _node_memory_budget_bytes(entry) + budget_gb = budget_bytes / (1024 ** 3) + print( + f"[tracker] node registered: {node_id} {endpoint} {model}{repo_info} {shard_info} " + f"capacity={budget_gb:.1f}GB {budget_source} slots={max_loaded_shards}", + flush=True, + ) + + payload = {"node_id": node_id} + if assignment_directive is not None: + payload["directive"] = assignment_directive + self._send_json(200, payload) + + # Raft replication: leader proposes this registration to followers so their + # registries stay in sync. Fire-and-forget (async) — the client already + # got its node_id; replication happens in the background. + if server.raft is not None and server.raft.is_leader: + raft_payload = dict(body) + raft_payload["node_id"] = node_id # include the generated ID + threading.Thread( + target=server.raft.propose, + args=("register", raft_payload), + daemon=True, + ).start() + + def _handle_heartbeat(self, node_id: str): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body: dict = {} + content_length = int(self.headers.get("Content-Length", 0)) + if content_length > 0: + try: + body = json.loads(self.rfile.read(content_length)) + except Exception: + pass + with server.lock: + self._purge_expired_nodes() + entry = server.registry.get(node_id) + if entry is None: + _tracker_log(server, "warn", "heartbeat for unknown node", node_id=node_id) + self._send_json(404, {"error": "node not found"}) + return + entry.last_heartbeat = time.monotonic() + entry.heartbeats_received += 1 + # P2P metadata + if body.get("relay_addr"): + entry.relay_addr = body["relay_addr"] + if body.get("cert_fingerprint"): + entry.cert_fingerprint = body["cert_fingerprint"] + if body.get("peer_id"): + entry.peer_id = body["peer_id"] + # Node stats (cumulative — node always sends totals, tracker replaces) + if "total_requests" in body: + entry.total_requests = int(body["total_requests"]) + if "failed_requests" in body: + entry.failed_requests = int(body["failed_requests"]) + if "queue_depth" in body: + entry.queue_depth = int(body["queue_depth"]) + if "current_requests" in body: + entry.current_requests = _normalize_current_requests(body["current_requests"]) + if "uptime_seconds" in body: + entry.uptime_seconds = float(body["uptime_seconds"]) + if "status" in body and body["status"] in ("ready", "loading"): + entry.status = body["status"] + directives = list(entry.pending_directives) + entry.pending_directives.clear() + new_assignment = entry.pending_new_assignment + if new_assignment is not None: + entry.pending_new_assignment = None + if server.gossip is not None: + server.gossip.record(node_id) + resp: dict = {} + if directives: + resp["directives"] = directives + if new_assignment is not None: + resp["new_assignment"] = new_assignment + self._send_json(200, resp) + + # ---------------------------------------------------------------- Raft handlers + + def _handle_raft_vote(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_json_body() + if body is None: + return + if server.raft is None: + self._send_json(503, {"error": "raft not enabled"}) + return + result = server.raft.handle_request_vote(body) + self._send_json(200, result) + + def _handle_raft_append(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_json_body() + if body is None: + return + if server.raft is None: + self._send_json(503, {"error": "raft not enabled"}) + return + result = server.raft.handle_append_entries(body) + self._send_json(200, result) + + def _handle_raft_status(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.raft is None: + self._send_json(200, {"role": "standalone", "term": 0, "leader": None}) + return + self._send_json(200, server.raft.status()) + + def _handle_gossip(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_json_body() + if body is None: + return + if server.gossip is not None and isinstance(body, dict): + server.gossip.merge({k: float(v) for k, v in body.items()}) + self._send_json(200, {}) + + def _handle_stats(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.stats is None: + self._send_json(200, {"models": {}, "nodes": {}}) + return + self._send_json(200, { + "models": server.stats.get_combined_stats(), + "nodes": server.stats.get_node_throughput_stats(), + }) + + def _handle_stats_gossip(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_hive_authenticated_body() + if body is None: + return + tracker_url = body.get("tracker_url", "") + rpms = body.get("stats", {}) + if server.stats is not None and tracker_url and isinstance(rpms, dict): + server.stats.merge_peer_rpms(tracker_url, rpms) + self._send_json(200, {}) + + def _handle_billing_summary(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if not self._require_role("admin"): + return + if server.billing is None: + self._send_json(404, {"error": "billing is not enabled on this tracker"}) + return + self._send_json(200, server.billing.snapshot()) + + def _handle_dashboard(self): + """Serve the read-only web dashboard (US-035). Any tracker in the + mesh — leader or follower — serves it from its replicated state.""" + try: + html = files("meshnet_tracker").joinpath("dashboard.html").read_text() + except (FileNotFoundError, OSError): + self._send_json(404, {"error": "dashboard asset missing"}) + return + body = html.encode() + self.send_response(200) + self.send_header("Content-Type", "text/html; charset=utf-8") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + try: + self.wfile.write(body) + except BrokenPipeError: + pass + + def _handle_console(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + with server.console_lock: + events = [dict(event) for event in server.console_events] + self._send_json(200, {"events": events}) + + def _handle_proxy_request_cancel(self, request_id: str) -> None: + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.accounts is not None and not self._require_role("admin"): + return + if not _request_proxy_cancel(server, request_id): + self._send_json(404, {"error": f"no active proxy for request {request_id!r}"}) + return + self._send_json(200, {"status": "canceled", "request_id": request_id}) + + def _finalize_proxy_cancel( + self, + *, + proxy_ctx: _ActiveProxyContext, + server: "_TrackerHTTPServer", + request_id: str, + started: float, + model: str, + route_model: str, + route_nodes: list, + api_key: str | None, + node_work: list, + body: dict, + finish_proxy_inflight: Any, + observed_stream_tokens: int = 0, + stream_usage: dict | None = None, + ) -> bool: + if not proxy_ctx.cancel_event.is_set(): + return False + elapsed = time.monotonic() - started + _clear_proxy_progress_log_state(server, request_id) + tokens = observed_stream_tokens + if observed_stream_tokens > 0: + in_tokens, out_tokens = _stream_billable_split( + observed_stream_tokens, stream_usage, body, + ) + tokens = in_tokens + out_tokens + self._record_observed_throughput( + model, route_model, tokens, elapsed, route_nodes, + ) + _tracker_log( + server, + "info", + "proxy canceled", + request_id=request_id, + model=model, + route_model=route_model, + tokens=tokens, + elapsed_seconds=round(elapsed, 4), + tokens_per_sec=round(tokens / elapsed, 4) if elapsed > 0 and tokens > 0 else 0.0, + route=_node_route_summary(route_nodes), + ) + if observed_stream_tokens > 0: + in_tokens, out_tokens = _stream_billable_split( + observed_stream_tokens, stream_usage, body, + ) + self._bill_completed( + api_key, model, in_tokens + out_tokens, node_work, + input_tokens=in_tokens, output_tokens=out_tokens, + ) + finish_proxy_inflight() + return True + + def _handle_registry_wallets(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if not self._require_role("admin"): + return + if server.contracts is None: + self._send_json(200, {"wallets": {}}) + return + wallets = server.contracts.registry.list_wallets() + self._send_json(200, {"wallets": { + wallet: { + "stake_balance": state.stake_balance, + "strike_count": state.strike_count, + "banned": state.banned, + "completed_jobs": state.completed_job_count, + "reputation": state.reputation, + "last_audit_ts": state.last_audit_ts, + "probationary_jobs_remaining": ( + server.contracts.registry.probationary_jobs_remaining(wallet) + ), + } + for wallet, state in wallets.items() + }}) + + def _handle_billing_settlements(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if not self._require_role("admin"): + return + if server.billing is None: + self._send_json(404, {"error": "billing is not enabled on this tracker"}) + return + self._send_json(200, {"settlements": server.billing.settlement_history()}) + + def _handle_billing_gossip(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_hive_authenticated_body() + if body is None: + return + if server.billing is None: + self._send_json(200, {"applied": 0}) + return + events = body.get("events") + if not isinstance(events, list): + self._send_json(400, {"error": "events must be a list"}) + return + applied = server.billing.apply_events([e for e in events if isinstance(e, dict)]) + self._send_json(200, {"applied": applied}) + + def _handle_billing_forfeit(self): + """Privileged: forfeit a node's pending balance + record a strike (US-034). + + ADR-0017 §4: validator service token or admin session only. Client + API keys — valid or not — are rejected. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if not self._require_role("validator", "admin"): + return + if server.billing is None: + self._send_json(404, {"error": "billing is not enabled on this tracker"}) + return + body = self._read_json_body() + if body is None: + return + wallet = body.get("wallet") + if not wallet or not isinstance(wallet, str): + self._send_json(400, {"error": "wallet is required"}) + return + reason = body.get("reason") or "fraud" + event = server.billing.forfeit_pending(wallet, reason=str(reason)) + strike_state: dict = {} + if server.contracts is not None: + server.contracts.registry.record_strike(wallet) + wallet_state = server.contracts.registry.get_wallet(wallet) + strike_state = { + "strike_count": wallet_state.strike_count, + "banned": wallet_state.banned, + } + print( + f"[tracker] forfeited pending balance of {wallet}: " + f"{event['amount']:.6f} USDT ({reason}) strikes={strike_state.get('strike_count', 'n/a')}", + flush=True, + ) + self._send_json(200, {"forfeited": event["amount"], **strike_state}) + + # ---- user accounts (registration, login, API keys) ---- + + def _session_account(self) -> dict | None: + """Resolve the session token in the Authorization header, or None.""" + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.accounts is None: + return None + return server.accounts.session_account(_session_token_from_headers(self.headers)) + + def _require_accounts(self) -> "AccountStore | None": + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.accounts is None: + self._send_json(404, {"error": "accounts are not enabled on this tracker"}) + return None + return server.accounts + + def _handle_auth_register(self): + accounts = self._require_accounts() + if accounts is None: + return + body = self._read_json_body() + if body is None: + return + try: + account = accounts.register( + email=body.get("email"), + wallet=body.get("wallet"), + password=str(body.get("password") or ""), + ) + except ValueError as exc: + self._send_json(400, {"error": str(exc)}) + return + token = accounts.create_session(account["account_id"]) + api_key = accounts.create_api_key(account["account_id"]) + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.billing is not None: + # US-039: registration creates the account's first key, so Caller + # Credit lands here; the account-derived event id keeps later + # grants (e.g. via /v1/account/keys) no-ops. + server.billing.grant_caller_credit( + api_key, account["account_id"], server.starting_credit + ) + server.billing.ensure_client(api_key) + print( + f"[tracker] account registered: {account.get('email') or account.get('wallet')} " + f"role={account['role']}", flush=True, + ) + self._send_json( + 200, + {"account": account, "session_token": token, "api_key": api_key}, + headers={"Set-Cookie": _session_cookie_header(token)}, + ) + + def _handle_auth_login(self): + accounts = self._require_accounts() + if accounts is None: + return + body = self._read_json_body() + if body is None: + return + identifier = body.get("identifier") or body.get("email") or body.get("wallet") or "" + account = accounts.verify_login(str(identifier), str(body.get("password") or "")) + if account is None: + self._send_json(401, {"error": "invalid credentials"}) + return + token = accounts.create_session(account["account_id"]) + self._send_json( + 200, + {"account": account, "session_token": token}, + headers={"Set-Cookie": _session_cookie_header(token)}, + ) + + def _handle_auth_logout(self): + accounts = self._require_accounts() + if accounts is None: + return + accounts.destroy_session(_session_token_from_headers(self.headers)) + self._send_json(200, {"ok": True}, headers={"Set-Cookie": _session_cookie_header(None)}) + + def _handle_account_me(self): + """Balance, usage, and API keys for the logged-in account.""" + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if self._require_accounts() is None: + return + account = self._session_account() + if account is None: + self._send_json(401, {"error": "login required"}) + return + keys = server.accounts.keys_for(account["account_id"]) # type: ignore[union-attr] + balances = {} + usage: dict = {"requests": 0, "total_tokens": 0, "total_cost": 0.0, "recent": [], "records": []} + if server.billing is not None: + balances = {key: server.billing.get_client_balance(key) for key in keys} + usage = server.billing.usage_for(keys) + self._send_json(200, { + "account": account, + "api_keys": keys, + "balances": balances, + "total_balance": sum(balances.values()), + "usage": usage, + "topup_amount": server.devnet_topup_amount if server.billing is not None else 0.0, + }) + + def _handle_account_key_create(self): + accounts = self._require_accounts() + if accounts is None: + return + account = self._session_account() + if account is None: + self._send_json(401, {"error": "login required"}) + return + api_key = accounts.create_api_key(account["account_id"]) + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + credited = False + if server.billing is not None: + # US-039: Caller Credit lands with the account's first key; the + # account-derived event id keeps it once-per-account forever. + credited = server.billing.grant_caller_credit( + api_key, account["account_id"], server.starting_credit + ) + server.billing.ensure_client(api_key) + self._send_json(200, {"api_key": api_key, "caller_credit_granted": credited}) + + def _handle_account_topup(self): + """Devnet faucet (US-040): credit the configured amount to one of the + logged-in account's keys. 404 unless the operator enabled it.""" + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + accounts = self._require_accounts() + if accounts is None: + return + if server.billing is None or server.devnet_topup_amount <= 0: + self._send_json(404, {"error": "top-up is not enabled on this tracker"}) + return + account = self._session_account() + if account is None: + self._send_json(401, {"error": "login required"}) + return + body = self._read_json_body() + if body is None: + return + api_key = body.get("api_key") + if not api_key or not isinstance(api_key, str): + self._send_json(400, {"error": "api_key is required"}) + return + if accounts.owner_of_key(api_key) != account["account_id"]: + self._send_json(403, {"error": "key not found on this account"}) + return + balance = server.billing.credit_client( + api_key, server.devnet_topup_amount, note="devnet-topup" + ) + self._send_json(200, { + "api_key": api_key, + "credited": server.devnet_topup_amount, + "balance": balance, + }) + + def _handle_account_key_revoke(self): + accounts = self._require_accounts() + if accounts is None: + return + account = self._session_account() + if account is None: + self._send_json(401, {"error": "login required"}) + return + body = self._read_json_body() + if body is None: + return + api_key = body.get("api_key") + if not api_key or not isinstance(api_key, str): + self._send_json(400, {"error": "api_key is required"}) + return + if not accounts.revoke_api_key(account["account_id"], api_key): + self._send_json(404, {"error": "key not found on this account"}) + return + self._send_json(200, {"revoked": api_key}) + + def _handle_admin_accounts(self): + """Admin-only: all accounts with their keys and balances.""" + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + accounts = self._require_accounts() + if accounts is None: + return + account = self._session_account() + if account is None: + self._send_json(401, {"error": "login required"}) + return + if account["role"] != "admin": + self._send_json(403, {"error": "admin role required"}) + return + listing = accounts.list_accounts() + if server.billing is not None: + for entry in listing: + entry["balances"] = { + key: server.billing.get_client_balance(key) + for key in entry.get("api_keys", []) + } + self._send_json(200, {"accounts": listing}) + + def _handle_accounts_gossip(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_hive_authenticated_body() + if body is None: + return + if server.accounts is None: + self._send_json(200, {"applied": 0}) + return + events = body.get("events") + if not isinstance(events, list): + self._send_json(400, {"error": "events must be a list"}) + return + applied = server.accounts.apply_events([e for e in events if isinstance(e, dict)]) + self._send_json(200, {"applied": applied}) + + def _handle_registry_gossip(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_hive_authenticated_body() + if body is None: + return + registry_log = getattr(server.contracts, "registry_log", None) + if registry_log is None: + self._send_json(200, {"applied": 0}) + return + events = body.get("events") + if not isinstance(events, list): + self._send_json(400, {"error": "events must be a list"}) + return + applied = registry_log.apply_events([e for e in events if isinstance(e, dict)]) + self._send_json(200, {"applied": applied}) + + def _handle_wallet_register(self): + """Bind a client wallet pubkey to an API key (US-032, C6). + + Deposits from that wallet into the treasury are then credited to the + API key's ledger balance by the deposit watcher. + + Requires a signature proving ownership of the wallet's private key — + the caller signs ``binding_message(api_key, wallet)`` with the + wallet's ed25519 key. An admin session may force-rebind a wallet + already bound to a different API key (documented override; no + signed-release flow is implemented). + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.billing is None: + self._send_json(404, {"error": "billing is not enabled on this tracker"}) + return + role, _account = self._resolve_identity() + body = self._read_json_body() + if body is None: + return + wallet = body.get("wallet") + if not wallet or not isinstance(wallet, str): + self._send_json(400, {"error": "wallet is required"}) + return + + if role == "admin": + api_key = body.get("api_key") + if not api_key or not isinstance(api_key, str): + self._send_json(400, {"error": "api_key is required for admin override"}) + return + event = server.billing.bind_wallet(api_key, wallet, admin_override=True) + else: + api_key = _api_key_from_headers(self.headers) + if not api_key: + self._send_json(401, {"error": "Authorization header required"}) + return + signature = body.get("signature") + if not signature or not isinstance(signature, str): + self._send_json(400, {"error": "signature is required to prove wallet ownership"}) + return + if not verify_wallet_signature(wallet, binding_message(api_key, wallet), signature): + self._send_json(401, {"error": "invalid wallet signature"}) + return + event = server.billing.bind_wallet(api_key, wallet) + + if event.get("rejected"): + self._send_json(409, {"error": "wallet already bound to a different API key"}) + return + print(f"[tracker] wallet bound: {wallet} -> api_key …{api_key[-6:]}", flush=True) + self._send_json(200, {"wallet": wallet, "bound": True}) + + def _handle_benchmark_hop_penalty(self): + """Privileged: run the same prompt through 1/2/3-node pinned routes (US-030). + + Data collection only — the routing algorithm is unchanged. Per-hop + latency is derived incrementally: the k-node route's final hop penalty + is its total minus the (k-1)-node route's total. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if not self._require_role("admin", "validator"): + return + auth = self.headers.get("Authorization") + body = self._read_json_body() + if body is None: + return + model = body.get("model", "") + if not model: + self._send_json(400, {"error": "model is required"}) + return + prompt = body.get("prompt") or "Benchmark: say OK." + max_new_tokens = int(body.get("max_new_tokens", 64)) + + with server.lock: + self._purge_expired_nodes() + resolved = _nodes_and_bounds_for_model(server, model) + if resolved is None or not resolved[0]: + self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) + return + all_nodes, rs, re = resolved + + self_url = f"http://127.0.0.1:{self.server.server_address[1]}" + route_results: list[dict] = [] + prev_total_ms: float | None = None + prev_per_hop: list[float] = [] + for hop_count in (1, 2, 3): + combo = _find_pinned_route(all_nodes, rs, re, hop_count) + if combo is None: + continue # insufficient coverage for this hop count — skip + request_body = json.dumps({ + "model": model, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_new_tokens, + "route": [n.node_id for n in combo], + }).encode() + req = urllib.request.Request( + f"{self_url}/v1/chat/completions", + data=request_body, + headers={"Content-Type": "application/json", "Authorization": auth}, + method="POST", + ) + started = time.monotonic() + try: + with urllib.request.urlopen(req, timeout=300.0) as resp: + payload = json.loads(resp.read()) + except Exception as exc: + route_results.append({ + "route": [n.node_id for n in combo], + "error": str(exc), + }) + continue + total_ms = (time.monotonic() - started) * 1000.0 + if prev_total_ms is not None and len(prev_per_hop) == hop_count - 1: + per_hop_ms = prev_per_hop + [max(0.0, total_ms - prev_total_ms)] + else: + per_hop_ms = [total_ms / hop_count] * hop_count + prev_total_ms = total_ms + prev_per_hop = per_hop_ms + route_results.append({ + "route": [n.node_id for n in combo], + "total_ms": total_ms, + "per_hop_ms": per_hop_ms, + "tokens_generated": _usage_total_tokens(payload) or 0, + }) + + record = { + "timestamp": time.time(), + "model": model, + "prompt_hash": hashlib.sha256(prompt.encode()).hexdigest()[:16], + "routes": route_results, + } + with server.benchmark_lock: + existing: list = [] + if os.path.exists(server.benchmark_results_path): + try: + with open(server.benchmark_results_path, encoding="utf-8") as fh: + existing = json.load(fh) + except (json.JSONDecodeError, OSError): + existing = [] + if not isinstance(existing, list): + existing = [] + existing.append(record) + with open(server.benchmark_results_path, "w", encoding="utf-8") as fh: + json.dump(existing, fh, indent=2) + self._send_json(200, record) + + def _handle_benchmark_results(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if not self._require_role("admin", "validator"): + return + results: list = [] + with server.benchmark_lock: + if os.path.exists(server.benchmark_results_path): + try: + with open(server.benchmark_results_path, encoding="utf-8") as fh: + results = json.load(fh) + except (json.JSONDecodeError, OSError): + results = [] + self._send_json(200, {"results": results if isinstance(results, list) else []}) + + def _handle_toploc_calibration_run(self): + """Privileged: honest-noise TOPLOC calibration dispatch (issue 21). + + Fans the same fixed prompt through every currently registered node + that can solo-serve the full model (one pinned-route hop, mirroring + `_handle_benchmark_hop_penalty`'s dispatch pattern), then verifies + each node's own on-demand TOPLOC commitment against a teacher-forced + replay on the reference node — same audit contract the validator + uses (`packages/validator/README.md` "TOPLOC audit contract"). Each + node's raw divergence (not just pass/fail) is recorded into the + calibration corpus, keyed by wallet + GPU model + dtype, so + thresholds can eventually be derived instead of guessed. + + Nodes that only hold a partial shard (need a multi-hop route) are + skipped for this pass — solo dispatch isolates one node's hardware + noise without a route composition confound — and nodes that don't + answer the on-demand commitment fetch (endpoint down, or node-side + TOPLOC serving not yet wired) are skipped and reported, not treated + as a pass or a fail. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if not self._require_role("admin", "validator"): + return + if server.toploc_calibration is None: + self._send_json(503, {"error": "toploc calibration store is not enabled on this tracker"}) + return + if not server.toploc_reference_node_url: + self._send_json(503, {"error": "toploc_reference_node_url is not configured on this tracker"}) + return + auth = self.headers.get("Authorization") + body = self._read_json_body() + if body is None: + return + model = body.get("model", "") + if not model: + self._send_json(400, {"error": "model is required"}) + return + prompt = body.get("prompt") or "Calibration: say OK." + max_new_tokens = int(body.get("max_new_tokens", 32)) + seed = body.get("seed", 0) + + with server.lock: + self._purge_expired_nodes() + resolved = _nodes_and_bounds_for_model(server, model) + if resolved is None or not resolved[0]: + self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) + return + all_nodes, rs, re = resolved + if server.contracts is not None: + all_nodes = [ + node for node in all_nodes + if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned + ] + solo_nodes = [ + node for node in all_nodes + if node.shard_start is not None and node.shard_end is not None + and node.shard_start <= rs and node.shard_end >= re + ] + + self_url = f"http://127.0.0.1:{self.server.server_address[1]}" + messages = [{"role": "user", "content": prompt}] + node_results: list[dict] = [] + for node in solo_nodes: + request_id = f"cal-{uuid.uuid4().hex}" + request_body = json.dumps({ + "id": request_id, + "model": model, + "messages": messages, + "max_tokens": max_new_tokens, + "temperature": 0, + "seed": seed, + "route": [node.node_id], + }).encode() + req = urllib.request.Request( + f"{self_url}/v1/chat/completions", + data=request_body, + headers={"Content-Type": "application/json", "Authorization": auth}, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=300.0) as resp: + json.loads(resp.read()) + except Exception as exc: + node_results.append({"node_id": node.node_id, "wallet_address": node.wallet_address, "error": str(exc)}) + continue + + outcome = self._verify_node_toploc_calibration( + server, node, request_id=request_id, model=model, messages=messages, + ) + node_results.append(outcome) + + skipped_partial_shard = [ + node.node_id for node in all_nodes if node not in solo_nodes + ] + record = { + "timestamp": time.time(), + "model": model, + "prompt_hash": hashlib.sha256(prompt.encode()).hexdigest()[:16], + "nodes": node_results, + "skipped_partial_shard_node_ids": skipped_partial_shard, + "gate_status": server.toploc_calibration.gate_status( + min_hardware_profiles=self._toploc_calibration_gate_min_hardware_profiles(), + ), + } + self._send_json(200, record) + + def _toploc_calibration_gate_min_hardware_profiles(self) -> int: + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + return server.toploc_calibration_gate_min_hardware_profiles + + def _verify_node_toploc_calibration( + self, + server: "_TrackerHTTPServer", + node: "_NodeEntry", + *, + request_id: str, + model: str, + messages: list[dict], + ) -> dict: + """One node's calibration outcome: fetch its on-demand commitment, + teacher-force the claimed tokens on the reference node, verify, and + persist the raw divergence into the corpus.""" + from meshnet_validator.audit import ToplocProofClaim, verify_activation_proofs_detailed + + gpu_model = (node.hardware_profile or {}).get("gpu_name") or (node.hardware_profile or {}).get("device") or "unknown" + dtype = node.quantization or "unknown" + base_result = { + "node_id": node.node_id, + "wallet_address": node.wallet_address, + "gpu_model": gpu_model, + "dtype": dtype, + } + commitment = _fetch_toploc_commitment( + node, session_id=request_id, model=model, messages=messages, + ) + if commitment is None: + return {**base_result, "skipped": "no on-demand toploc commitment available"} + + try: + claim = ToplocProofClaim.from_mapping(commitment["toploc_proof"]) + except (KeyError, TypeError, ValueError): + return {**base_result, "skipped": "malformed toploc commitment"} + + reference_activations = _fetch_toploc_reference_activations( + server.toploc_reference_node_url, + model=model, + messages=messages, + claimed_token_ids=commitment["claimed_token_ids"], + claim=claim, + ) + if reference_activations is None: + return {**base_result, "skipped": "reference node teacher-forced replay failed"} + + result = verify_activation_proofs_detailed(reference_activations, claim, backend=server.toploc_backend) + if node.wallet_address: + server.toploc_calibration.record_run( + node_wallet=node.wallet_address, + gpu_model=gpu_model, + dtype=dtype, + model=model, + passed=result.passed, + exp_intersections=result.exp_intersections, + mant_err_mean=result.mant_err_mean, + mant_err_median=result.mant_err_median, + ) + return { + **base_result, + "passed": result.passed, + "exp_intersections": result.exp_intersections, + "mant_err_mean": result.mant_err_mean, + "mant_err_median": result.mant_err_median, + "chunk_count": result.chunk_count, + } + + def _handle_toploc_calibration_results(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if not self._require_role("admin", "validator"): + return + if server.toploc_calibration is None: + self._send_json(503, {"error": "toploc calibration store is not enabled on this tracker"}) + return + min_profiles = self._toploc_calibration_gate_min_hardware_profiles() + self._send_json(200, { + "runs": server.toploc_calibration.runs(), + "envelope": server.toploc_calibration.envelope(), + "gate_status": server.toploc_calibration.gate_status(min_hardware_profiles=min_profiles), + }) + + def _handle_hf_pricing_history(self, parsed: urllib.parse.ParseResult): + """Dispute-auditability log for the dynamic HF-benchmarked pricing (issue 23).""" + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if not self._require_role("admin", "validator"): + return + if server.hf_pricing_log is None: + self._send_json(503, {"error": "hf pricing log is not enabled on this tracker"}) + return + params = urllib.parse.parse_qs(parsed.query) + model = params.get("model", [None])[0] + self._send_json(200, {"changes": server.hf_pricing_log.history(model)}) + + def _handle_assign(self, parsed: urllib.parse.ParseResult): + """Return an optimal shard assignment for a node given its hardware profile. + + Query params: + model — model preset name (default: first preset) + device — "cuda" | "cpu" + vram_mb — integer VRAM in MB (0 for CPU) + ram_mb — integer system RAM in MB, used when vram_mb=0 + + The greedy strategy: find the first gap in current layer coverage + and assign it. If no gap exists, assign the full model range so the + node provides redundant coverage. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + params = urllib.parse.parse_qs(parsed.query) + + model_list = params.get("model") + if not model_list: + model = next(iter(server.model_presets), None) + if model is None: + self._send_json(503, {"error": "no model presets configured"}) + return + else: + model = model_list[0] + + resolved_name, preset = _resolve_model_preset(server.model_presets, model) + if preset is None: + self._send_json(404, {"error": f"unknown model preset: {model!r}"}) + return + + required_start, required_end = _preset_layer_bounds(preset) + + with server.lock: + self._purge_expired_nodes() + alive = [ + node for node in server.registry.values() + if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] + ] + if server.contracts is not None: + alive = [ + node for node in alive + if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned + ] + + device = params.get("device", ["cpu"])[0] + try: + vram_mb = int(params.get("vram_mb", ["0"])[0]) + except ValueError: + vram_mb = 0 + try: + ram_mb = int(params.get("ram_mb", ["0"])[0]) + except ValueError: + ram_mb = 0 + max_layers = required_end - required_start + 1 + # VRAM only bounds the shard for CUDA nodes; a CPU node may still report + # a detected-but-unusable GPU, and must be sized by system RAM. + memory_mb = vram_mb if (device == "cuda" and vram_mb > 0) else ram_mb + if memory_mb > 0: + layer_bytes = _preset_bytes_per_layer(preset).get("bfloat16", 30 * 1024 * 1024) + max_layers = min(max_layers, max(1, int(((memory_mb * 1024 * 1024) * 0.8) // layer_bytes))) + elif device != "cuda" or vram_mb < 8192: + max_layers = min(max_layers, 16) + + # Collect covered intervals sorted by start layer. + covered = sorted( + [ + (n.shard_start, n.shard_end) + for n in alive + if n.shard_start is not None and n.shard_end is not None + ], + key=lambda t: t[0], + ) + + # Walk from required_start to find the first uncovered layer. + gap_start = required_start + for s, e in covered: + if s <= gap_start: + gap_start = max(gap_start, e + 1) + else: + break # gap found before this node + + if gap_start > required_end: + # Full coverage exists — assign the full range for redundancy. + shard_start = required_start + shard_end = min(required_end, shard_start + max_layers - 1) + else: + shard_start = gap_start + shard_end = min(required_end, shard_start + max_layers - 1) + + peers = [ + {"endpoint": node.endpoint, "checksum": node.shard_checksum} + for node in alive + if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] + and node.shard_start == shard_start + and node.shard_end == shard_end + and node.shard_checksum + ] + + self._send_json(200, { + "shard_start": shard_start, + "shard_end": shard_end, + "model": resolved_name, + "model_layers_end": required_end, + "peers": peers, + "bytes_per_layer": _preset_bytes_per_layer(preset), + "model_sources": self._model_sources( + resolved_name, + preset, + shard_start, + shard_end, + ), + **({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}), + }) + + def _model_sources(self, model: str, preset: dict, shard_start: int, shard_end: int) -> list[dict]: + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + hf_repo = preset.get("hf_repo") + if not server.models_dir or not isinstance(hf_repo, str) or not hf_repo: + return [] + snapshot_dir = snapshot_dir_for_repo(server.models_dir, hf_repo) + if snapshot_dir is None: + return [] + files = files_for_layer_range(snapshot_dir, shard_start, shard_end) + if not files: + return [] + host = self.headers.get("Host") or f"{self.server.server_address[0]}:{self.server.server_address[1]}" + base_url = f"http://{host}" + query = urllib.parse.urlencode({ + "model": model, + "shard_start": shard_start, + "shard_end": shard_end, + }) + full_query = urllib.parse.urlencode({"model": model, "full": 1}) + full_files = _snapshot_regular_files(snapshot_dir) + file_sizes: dict[str, int] = {} + for rel in set(files) | set(full_files): + try: + file_sizes[rel] = (snapshot_dir / rel).resolve().stat().st_size + except OSError: + continue + return [{ + "type": "tracker", + "endpoint": base_url, + "url": f"{base_url}/v1/model-files/download?{query}", + "full_url": f"{base_url}/v1/model-files/download?{full_query}", + "files": files, + "full_files": full_files, + "file_sizes": file_sizes, + }] + + def _handle_model_files_download(self, parsed: urllib.parse.ParseResult) -> None: + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.models_dir is None: + self._send_json(404, {"error": "tracker model-file source is not enabled"}) + return + params = urllib.parse.parse_qs(parsed.query) + model = params.get("model", [""])[0] + resolved_name, preset = _resolve_model_preset(server.model_presets, model) + if preset is None: + self._send_json(404, {"error": f"unknown model preset: {model!r}"}) + return + hf_repo = preset.get("hf_repo") + if not isinstance(hf_repo, str) or not hf_repo: + self._send_json(404, {"error": f"model preset has no hf_repo: {resolved_name!r}"}) + return + snapshot_dir = snapshot_dir_for_repo(server.models_dir, hf_repo) + if snapshot_dir is None: + self._send_json(404, {"error": f"local snapshot not found for {hf_repo}"}) + return + full_download = params.get("full", ["0"])[0] in {"1", "true", "yes"} + if full_download: + rel_files = _snapshot_regular_files(snapshot_dir) + else: + try: + shard_start = int(params.get("shard_start", [""])[0]) + shard_end = int(params.get("shard_end", [""])[0]) + except ValueError: + self._send_json(400, {"error": "shard_start and shard_end must be integers"}) + return + rel_files = files_for_layer_range(snapshot_dir, shard_start, shard_end) + if not rel_files: + self._send_json(404, {"error": "no local files matched the assigned shard"}) + return + + # Single-file mode: ?file= streams one file with Content-Length so + # clients can verify completeness and retry per file instead of + # restarting one giant tar stream after a network hiccup. + single = params.get("file", [""])[0] + if single: + if single not in set(rel_files): + self._send_json(404, {"error": f"file not in requested shard: {single!r}"}) + return + file_path = snapshot_dir / single + try: + # resolve() dereferences HF snapshot symlinks into blobs/. + size = file_path.resolve().stat().st_size + except OSError: + self._send_json(404, {"error": f"file unreadable: {single!r}"}) + return + self.send_response(200) + self.send_header("Content-Type", "application/octet-stream") + self.send_header("Content-Length", str(size)) + self.send_header("X-Meshnet-Model-Source", "tracker") + self.end_headers() + try: + with file_path.open("rb") as f: + while True: + chunk = f.read(1024 * 1024) + if not chunk: + break + self.wfile.write(chunk) + except (BrokenPipeError, ConnectionResetError): + print( + f"model-file download aborted by client " + f"({self.client_address[0]}, model={resolved_name}, file={single})", + flush=True, + ) + return + + self.send_response(200) + self.send_header("Content-Type", "application/x-tar") + self.send_header("X-Meshnet-Model-Source", "tracker") + self.end_headers() + try: + # dereference: HF cache snapshots are symlinks into blobs/ — ship contents. + with tarfile.open(fileobj=self.wfile, mode="w|", dereference=True) as archive: + for rel in rel_files: + archive.add(snapshot_dir / rel, arcname=rel) + except (BrokenPipeError, ConnectionResetError): + print( + f"model-file download aborted by client " + f"({self.client_address[0]}, model={resolved_name})", + flush=True, + ) + + def _handle_network_assign(self, parsed: urllib.parse.ParseResult): + """Assign a new node to fill the biggest uncovered shard gap across HF-model nodes. + + Query params: + vram_mb — integer VRAM in MB (0 = CPU-only node) + ram_mb — integer system RAM in MB, used when vram_mb=0 + device — "cuda" | "cpu" + hf_repo — optional; if set, restrict search to this repo only + + Returns: + {hf_repo, shard_start, shard_end, num_layers, gap_found} + gap_found=true means a real uncovered gap was assigned; false means redundancy. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + params = urllib.parse.parse_qs(parsed.query) + try: + vram_mb = int(params.get("vram_mb", ["0"])[0]) + except ValueError: + vram_mb = 0 + try: + ram_mb = int(params.get("ram_mb", ["0"])[0]) + except ValueError: + ram_mb = 0 + device = params.get("device", ["cpu"])[0] + filter_repo = params.get("hf_repo", [None])[0] # optional repo filter + + with server.lock: + self._purge_expired_nodes() + all_nodes = list(server.registry.values()) + + # Collect only nodes that registered a real HF model (have hf_repo + shard bounds). + hf_nodes = [ + n for n in all_nodes + if n.hf_repo + and n.shard_start is not None + and n.shard_end is not None + and n.num_layers is not None + and (filter_repo is None or n.hf_repo == filter_repo) + ] + + if not hf_nodes: + # The caller is not registered yet — count its memory toward + # deployability, or the first node can never bootstrap the network. + caller = _NodeEntry( + node_id="assign-candidate", + endpoint="", + shard_start=None, + shard_end=None, + model=None, + shard_checksum=None, + hardware_profile={}, + wallet_address=None, + score=1.0, + vram_bytes=vram_mb * 1024 * 1024, + ram_bytes=ram_mb * 1024 * 1024, + ) + pool_with_caller = all_nodes + [caller] + resolved_name = None + preset = None + if filter_repo: + resolved_name, preset = _resolve_model_preset(server.model_presets, filter_repo) + else: + deployable = [ + (name, preset) + for name, preset in server.model_presets.items() + if preset.get("recommended") and _deployment_summary(pool_with_caller, preset)["deployable"] + ] + if deployable: + resolved_name, preset = deployable[0] + if preset is not None and preset.get("hf_repo"): + required_start, required_end = _preset_layer_bounds(preset) + total_l = required_end - required_start + 1 + memory_mb = vram_mb if (device == "cuda" and vram_mb > 0) else ram_mb + max_layers = _max_layers_for_memory(memory_mb, total_l, preset, device=device) + shard_start = required_start + shard_end = min(required_end, shard_start + max_layers - 1) + self._send_json(200, { + "hf_repo": preset["hf_repo"], + "model": resolved_name, + "shard_start": shard_start, + "shard_end": shard_end, + "num_layers": total_l, + "gap_found": True, + "price_per_token": 0.0, + "bytes_per_layer": _preset_bytes_per_layer(preset), + "deployment": _deployment_summary(pool_with_caller, preset), + "model_sources": self._model_sources( + str(resolved_name), + preset, + shard_start, + shard_end, + ), + }) + return + + msg = ( + f"no HF-model nodes registered for {filter_repo!r}" + if filter_repo + else "no HF-model nodes registered; cannot assign shards" + ) + self._send_json(503, {"error": msg}) + return + + # Group by hf_repo; pick the one with the largest total_layers and biggest gap. + from collections import defaultdict + repo_groups: dict = defaultdict(list) + repo_layers: dict = {} + for n in hf_nodes: + repo_groups[n.hf_repo].append(n) + # Use the largest num_layers seen for this repo. + if n.hf_repo not in repo_layers or n.num_layers > repo_layers[n.hf_repo]: + repo_layers[n.hf_repo] = n.num_layers + + # Smart scoring: demand_rpm × coverage_deficit + # coverage_deficit = uncovered_layers / total_layers (0 = fully covered) + # demand_rpm comes from the stats collector; defaults to 0.0 when no traffic yet. + demand_rpms: dict[str, float] = {} + if server.stats is not None: + local_rpms = server.stats.get_local_rpms() + for repo in repo_groups: + demand_rpms[repo] = _repo_demand_rpm(server, repo, local_rpms) + + best_repo = None + best_score = -1.0 + best_gap_size = -1 + best_gap_start = 0 + best_num_layers = 0 + + for repo, nodes in repo_groups.items(): + total = repo_layers[repo] + served_copies = _served_model_copies(nodes, 0, total - 1) + covered = sorted( + [(n.shard_start, n.shard_end) for n in nodes], + key=lambda t: t[0], + ) + # Walk from 0 to find first uncovered layer. + gap_start = 0 + for s, e in covered: + if s <= gap_start: + gap_start = max(gap_start, e + 1) + else: + break + gap_size = max(0, (total - 1) - gap_start + 1) + coverage_deficit = gap_size / max(total, 1) + demand = demand_rpms.get(repo, 0.0) + # Prefer demanded models already being served; coverage gaps still dominate. + score = (demand + 1.0) * (coverage_deficit + 0.01) + served_copies * 0.01 + if score > best_score or (score == best_score and gap_size > best_gap_size): + best_score = score + best_gap_size = gap_size + best_gap_start = gap_start + best_repo = repo + best_num_layers = total + + gap_found = best_gap_size > 0 + if not gap_found: + # All shards covered — add redundancy on the most demanded/served model. + best_repo = max( + repo_groups, + key=lambda repo: ( + demand_rpms.get(repo, 0.0), + _served_model_copies(repo_groups[repo], 0, repo_layers[repo] - 1), + len(repo_groups[repo]), + ), + ) + best_gap_start = 0 + best_num_layers = repo_layers[best_repo] + + # Capacity: use the same 80%-of-memory rule as registered node planning. + total_l = best_num_layers + memory_mb = vram_mb if (device == "cuda" and vram_mb > 0) else ram_mb + resolved_name, best_preset = _resolve_model_preset(server.model_presets, str(best_repo)) + if memory_mb > 0: + max_layers = _max_layers_for_memory(memory_mb, total_l, best_preset, device=device) + elif device == "cuda" and vram_mb >= 8192: + max_layers = total_l + else: + max_layers = _max_layers_for_memory(memory_mb, total_l, best_preset, device=device) + + shard_start = best_gap_start + shard_end = min(total_l - 1, shard_start + max_layers - 1) + + self._send_json(200, { + "hf_repo": best_repo, + "model": resolved_name, + "shard_start": shard_start, + "shard_end": shard_end, + "num_layers": total_l, + "gap_found": gap_found, + "price_per_token": 0.0, + "bytes_per_layer": _preset_bytes_per_layer(best_preset) if best_preset is not None else None, + "model_sources": self._model_sources( + str(resolved_name or best_repo), + best_preset or {"hf_repo": best_repo}, + shard_start, + shard_end, + ), + }) + + def _handle_route(self, parsed: urllib.parse.ParseResult): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + params = urllib.parse.parse_qs(parsed.query) + model_list = params.get("model") + if not model_list: + self._send_json(400, {"error": "missing 'model' query parameter"}) + return + + model = model_list[0] + resolved_name, preset = _resolve_model_preset(server.model_presets, model) + + with server.lock: + self._purge_expired_nodes() + if preset is not None: + # Preset-based routing (stub-model system). + alive = [ + node for node in server.registry.values() + if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] + ] + required_start, required_end = _preset_layer_bounds(preset) + else: + # HF model routing: match by hf_repo (full) or model short name. + alive = [ + node for node in server.registry.values() + if _node_matches_model(node, model) + and node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + ] + if not alive: + self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) + return + required_start = 0 + required_end = max(n.num_layers for n in alive) - 1 # type: ignore[type-var] + + if server.contracts is not None: + alive = [ + node for node in alive + if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned + ] + + route_model = resolved_name or model + candidates = _enumerate_routes( + alive, + required_start, + required_end, + model=route_model, + contracts=server.contracts, + max_candidates=server.route_stats.config.max_candidate_routes, + ) + if candidates: + # Prefer a distributed multi-hop route when available. Greedy + # _select_route alone would pick a single full-shard node and omit + # partial head shards that share the same port on another host. + route = max(candidates, key=lambda cand: len(cand.nodes)).nodes + error = "" + else: + route, error = _select_route( + alive, required_start, required_end, contracts=server.contracts, + ) + if error: + _tracker_log( + server, + "warn", + "route unavailable", + model=model, + route_model=resolved_name or model, + error=error, + candidate_count=len(alive), + candidates=_node_route_summary(alive), + ) + self._send_json(503, {"error": error}) + return + + covered_up_to = required_start - 1 + route_with_start: list[tuple] = [] + for rn in route: + route_with_start.append((rn, covered_up_to + 1)) + covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to + + self._send_json(200, { + "route": [e.endpoint for e, _ in route_with_start], + "nodes": [ + { + "node_id": e.node_id, + "endpoint": e.endpoint, + "start_layer": start, + "relay_addr": e.relay_addr, + "peer_id": e.peer_id, + "wallet_address": e.wallet_address, + "shard_start": e.shard_start, + "shard_end": e.shard_end, + "model": e.model, + "hf_repo": e.hf_repo, + "num_layers": e.num_layers, + "model_metadata": dict(e.model_metadata), + "downloaded_models": [dict(item) for item in e.downloaded_models], + "shard_checksum": e.shard_checksum, + "score": e.score, + } + for e, start in route_with_start + ], + }) + + def _handle_routing(self, parsed: urllib.parse.ParseResult): + """Learned route table (ADR-0021): per-model candidates with observed + tps, coefficient vs the best proven route, and expected traffic share.""" + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + params = urllib.parse.parse_qs(parsed.query) + only_model = (params.get("model") or [None])[0] + with server.lock: + self._purge_expired_nodes() + alive = list(server.registry.values()) + model_keys: dict[str, list] = {} + for node in alive: + if node.shard_start is None or node.shard_end is None: + continue + key = node.hf_repo or node.model + if not key: + continue + model_keys.setdefault(key, []).append(node) + cfg = server.route_stats.config + out: dict[str, dict] = {} + for key, nodes in model_keys.items(): + if only_model and key != only_model and not any( + key == alias for alias in _model_aliases(only_model) + ): + continue + resolved_name, preset = _resolve_model_preset(server.model_presets, key) + # Stats are recorded under the proxy's resolved route_model — + # use the same key here or lookups always miss. + stats_key = resolved_name or key + if preset is not None: + rs, re = _preset_layer_bounds(preset) + else: + layer_counts = [n.num_layers for n in nodes if n.num_layers is not None] + if not layer_counts: + continue + rs, re = 0, max(layer_counts) - 1 + candidates = _enumerate_routes( + nodes, rs, re, + model=stats_key, + contracts=server.contracts, + max_candidates=cfg.max_candidate_routes, + ) + out[stats_key] = { + "epoch": server.route_stats.epoch(stats_key), + "routes": route_table(candidates, server.route_stats, stats_key), + } + self._send_json(200, { + "config": { + "explore_share": cfg.explore_share, + "weight_alpha": cfg.weight_alpha, + "stats_half_life_seconds": cfg.stats_half_life_seconds, + "min_sample_tokens": cfg.min_sample_tokens, + }, + "models": out, + }) + + def _handle_routes(self, parsed: urllib.parse.ParseResult): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + params = urllib.parse.parse_qs(parsed.query) + model_list = params.get("model") + if not model_list: + self._send_json(400, {"error": "missing 'model' query parameter"}) + return + + try: + redundancy = int(params.get("redundancy", ["1"])[0]) + except ValueError: + self._send_json(400, {"error": "redundancy must be an integer"}) + return + if redundancy < 1: + self._send_json(400, {"error": "redundancy must be at least 1"}) + return + + model = model_list[0] + resolved_name, preset = _resolve_model_preset(server.model_presets, model) + if preset is None: + self._send_json(404, {"error": f"unknown model preset: {model!r}"}) + return + + required_start, required_end = _preset_layer_bounds(preset) + + with server.lock: + self._purge_expired_nodes() + candidates = [ + node for node in server.registry.values() + if _node_matches_preset(node, resolved_name, preset) # type: ignore[arg-type] + ] + if server.contracts is not None: + candidates = [ + node for node in candidates + if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned + ] + + routes = [] + remaining = list(candidates) + for _ in range(redundancy): + route, error = _select_route(remaining, required_start, required_end, contracts=server.contracts) + if error: + self._send_json(503, {"error": error}) + return + route_endpoints = {node.endpoint for node in route} + remaining = [node for node in remaining if node.endpoint not in route_endpoints] + routes.append({ + "route": [e.endpoint for e in route], + "nodes": [ + { + "node_id": e.node_id, + "endpoint": e.endpoint, + "relay_addr": e.relay_addr, + "peer_id": e.peer_id, + "wallet_address": e.wallet_address, + "shard_start": e.shard_start, + "shard_end": e.shard_end, + "model": e.model, + "downloaded_models": [dict(item) for item in e.downloaded_models], + "shard_checksum": e.shard_checksum, + "score": e.score, + } + for e in route + ], + }) + + self._send_json(200, {"routes": routes}) + + +class TrackerServer: + """HTTP tracker that manages node registration and resolves inference routes. + + Nodes register via POST /v1/nodes/register and keep themselves alive with + POST /v1/nodes//heartbeat. The gateway queries GET /v1/route?model= + to obtain an ordered list of node endpoints whose shards cover all layers. + """ + + def __init__( + self, + host: str = "127.0.0.1", + port: int = 0, + heartbeat_timeout: float = 90.0, + rebalance_interval: float = 30.0, + model_presets: dict | None = None, + contracts: Any | None = None, + minimum_stake: int = 0, + cluster_peers: list[str] | None = None, + cluster_self_url: str | None = None, + stats_db: str | None = None, + relay_url: str | None = None, + billing: BillingLedger | None = None, + enable_billing: bool = False, + billing_db: str | None = None, + accounts: AccountStore | None = None, + accounts_db: str | None = None, + benchmark_results_path: str | None = None, + treasury: Any | None = None, + deposit_poll_interval: float = 15.0, + settle_period: float = 86400.0, + payout_threshold: float = 5.0, + payout_dust_floor: float = 0.01, + settlement_check_interval: float | None = None, + validator_service_token: str | None = None, + hive_secret: str | None = None, + max_charge_per_request: float | None = None, + starting_credit: float = DEFAULT_CALLER_CREDIT_USDT, + devnet_topup_amount: float = DEFAULT_DEVNET_TOPUP_USDT, + toploc_calibration: ToplocCalibrationStore | None = None, + toploc_calibration_db: str | None = None, + toploc_reference_node_url: str | None = None, + toploc_calibration_gate_min_hardware_profiles: int = 1, + toploc_backend: Any | None = None, + enable_hf_pricing: bool = False, + hf_pricing_log: HfPricingLog | None = None, + hf_pricing_log_db: str | None = None, + hf_pricing_refresh_interval: float = 86400.0, + hf_pricing_fetch_html: Any | None = None, + models_dir: str | Path | None = None, + routing_config: RoutingConfig | None = None, + ) -> None: + self._host = host + self._requested_port = port + self._heartbeat_timeout = heartbeat_timeout + self._rebalance_interval = rebalance_interval + self._model_presets: dict = ( + model_presets if model_presets is not None else _clone_model_presets(DEFAULT_MODEL_PRESETS) + ) + self._contracts = contracts + self._minimum_stake = minimum_stake + self._cluster_peers: list[str] = list(cluster_peers) if cluster_peers else [] + self._cluster_self_url: str | None = cluster_self_url + self._relay_url = relay_url + self._registry: dict[str, _NodeEntry] = {} + self._lock = threading.Lock() + self._server: _TrackerHTTPServer | None = None + self._thread: threading.Thread | None = None + self._rebalance_stop = threading.Event() + self._rebalance_thread: threading.Thread | None = None + self._raft: RaftNode | None = None + self._gossip: NodeGossip | None = None + self._stats: _StatsCollector | None = _StatsCollector(db_path=stats_db) if stats_db else _StatsCollector() + self._stats_stop = threading.Event() + self._stats_thread: threading.Thread | None = None + if billing is None: + db_path = billing_db + if db_path is None and enable_billing: + db_path = DEFAULT_BILLING_DB_PATH + if db_path: + preset_prices: dict[str, tuple[float, float]] = {} + for name, preset in self._model_presets.items(): + if not isinstance(preset, dict): + continue + base = preset.get("price_per_1k_tokens") + input_rate = preset.get("input_price_per_1k_tokens", base) + output_rate = preset.get("output_price_per_1k_tokens", base) + if input_rate is None or output_rate is None: + continue + for key in _preset_price_keys(name, preset): + preset_prices[key] = (float(input_rate), float(output_rate)) + billing = BillingLedger(db_path=db_path, prices=preset_prices) + self._billing: BillingLedger | None = billing + self._billing_gossip_cursor = 0 + if accounts is None: + accounts_path = accounts_db + if accounts_path is None and enable_billing: + accounts_path = DEFAULT_ACCOUNTS_DB_PATH + if accounts_path: + accounts = AccountStore(db_path=accounts_path) + self._accounts: AccountStore | None = accounts + self._accounts_gossip_cursor = 0 + self._registry_gossip_cursor = 0 + self._benchmark_results_path = benchmark_results_path + self._treasury = treasury + self._deposit_poll_interval = deposit_poll_interval + self._deposit_stop = threading.Event() + self._deposit_thread: threading.Thread | None = None + self._settle_period = settle_period + self._payout_threshold = payout_threshold + self._payout_dust_floor = payout_dust_floor + self._settlement_check_interval = settlement_check_interval or max( + 1.0, min(settle_period / 4.0, 15.0) + ) + self._settlement_stop = threading.Event() + self._settlement_thread: threading.Thread | None = None + if max_charge_per_request is not None and max_charge_per_request <= 0.0: + raise ValueError("max_charge_per_request must be positive") + self._max_charge_per_request = max_charge_per_request + self._starting_credit = max(0.0, starting_credit) + self._devnet_topup_amount = max(0.0, devnet_topup_amount) + self._validator_service_token = ( + validator_service_token + if validator_service_token is not None + else os.environ.get("MESHNET_VALIDATOR_SERVICE_TOKEN") or None + ) + self._hive_secret = ( + hive_secret + if hive_secret is not None + else os.environ.get("MESHNET_HIVE_SECRET") or None + ) + if toploc_calibration is None and toploc_calibration_db: + toploc_calibration = ToplocCalibrationStore(db_path=toploc_calibration_db) + self._toploc_calibration: ToplocCalibrationStore | None = toploc_calibration + self._toploc_reference_node_url = toploc_reference_node_url + self._toploc_calibration_gate_min_hardware_profiles = toploc_calibration_gate_min_hardware_profiles + self._toploc_backend = toploc_backend + if hf_pricing_log is None and (enable_hf_pricing or hf_pricing_log_db): + hf_pricing_log = HfPricingLog(db_path=hf_pricing_log_db or DEFAULT_HF_PRICING_LOG_DB_PATH) + self._hf_pricing_log: HfPricingLog | None = hf_pricing_log + self._enable_hf_pricing = enable_hf_pricing + self._hf_pricing_refresh_interval = hf_pricing_refresh_interval + self._hf_pricing_fetch_html = hf_pricing_fetch_html + # MESHNET_DOWNLOAD_DIR is the node-side model store; on a box running + # both tracker and node it doubles as the tracker's snapshot source. + raw_models_dir = ( + models_dir + if models_dir is not None + else os.environ.get("MESHNET_MODELS_DIR") or os.environ.get("MESHNET_DOWNLOAD_DIR") + ) + self._models_dir = Path(raw_models_dir).expanduser() if raw_models_dir else None + self._hf_pricing_stop = threading.Event() + self._hf_pricing_thread: threading.Thread | None = None + if routing_config is None: + routing_config = RoutingConfig( + explore_share=float(os.environ.get("MESHNET_ROUTE_EXPLORE_SHARE", RoutingConfig.explore_share)), + weight_alpha=float(os.environ.get("MESHNET_ROUTE_WEIGHT_ALPHA", RoutingConfig.weight_alpha)), + stats_half_life_seconds=float( + os.environ.get("MESHNET_ROUTE_STATS_HALF_LIFE", RoutingConfig.stats_half_life_seconds) + ), + min_sample_tokens=int( + os.environ.get("MESHNET_ROUTE_MIN_SAMPLE_TOKENS", RoutingConfig.min_sample_tokens) + ), + ) + self._route_stats = RouteStatsStore(routing_config) + self.port: int | None = None + + def start(self) -> int: + if self._server is not None: + raise RuntimeError("TrackerServer is already running") + + effective_relay_url = ( + self._relay_url + or derive_relay_url_from_public_tracker_url(self._cluster_self_url) + ) + + # Start HTTP server first so we know our port + self._server = _TrackerHTTPServer( + (self._host, self._requested_port), + _TrackerHandler, + self._registry, + self._lock, + self._heartbeat_timeout, + self._model_presets, + self._contracts, + self._minimum_stake, + relay_url=effective_relay_url, + stats=self._stats, + billing=self._billing, + accounts=self._accounts, + benchmark_results_path=self._benchmark_results_path, + validator_service_token=self._validator_service_token, + hive_secret=self._hive_secret, + max_charge_per_request=self._max_charge_per_request, + starting_credit=self._starting_credit, + devnet_topup_amount=self._devnet_topup_amount, + toploc_calibration=self._toploc_calibration, + toploc_reference_node_url=self._toploc_reference_node_url, + toploc_calibration_gate_min_hardware_profiles=self._toploc_calibration_gate_min_hardware_profiles, + toploc_backend=self._toploc_backend, + hf_pricing_log=self._hf_pricing_log, + models_dir=self._models_dir, + route_stats=self._route_stats, + ) + self.port = self._server.server_address[1] + + # Start Raft + gossip if cluster peers are configured + if self._cluster_peers: + self_url = self._cluster_self_url or f"http://{self._host}:{self.port}" + self._raft = RaftNode( + self_url=self_url, + peers=self._cluster_peers, + apply_fn=self._raft_apply, + ) + self._gossip = NodeGossip(peers=self._cluster_peers) + self._server.raft = self._raft + self._server.gossip = self._gossip + self._raft.start() + self._gossip.start() + + self._rebalance_stop.clear() + self._stats_stop.clear() + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + self._thread.start() + self._rebalance_thread = threading.Thread(target=self._rebalance_loop, daemon=True) + self._rebalance_thread.start() + self._stats_thread = threading.Thread(target=self._stats_loop, daemon=True) + self._stats_thread.start() + if self._treasury is not None and self._billing is not None: + self._deposit_stop.clear() + self._deposit_thread = threading.Thread(target=self._deposit_loop, daemon=True) + self._deposit_thread.start() + self._settlement_stop.clear() + self._settlement_thread = threading.Thread(target=self._settlement_loop, daemon=True) + self._settlement_thread.start() + if self._enable_hf_pricing and self._billing is not None: + self._hf_pricing_stop.clear() + self._hf_pricing_thread = threading.Thread(target=self._hf_pricing_loop, daemon=True) + self._hf_pricing_thread.start() + return self.port + + def _settlement_loop(self) -> None: + """Leader-only on-chain settlement (US-033, ADR-0015). + + Pay a node when pending ≥ threshold OR its pending age ≥ max period, + with a dust floor. Pending is debited (payout events, replicated) + before the transaction is sent; unconfirmed batches are resent with + the same settlement id, so retries never double-pay. + """ + billing = self._billing + treasury = self._treasury + assert billing is not None and treasury is not None + while not self._settlement_stop.wait(self._settlement_check_interval): + if self._raft is not None and not self._raft.is_leader: + continue # followers replicate the ledger but never sign + # resend batches whose transaction never confirmed + for settlement_id, payouts in billing.unconfirmed_settlements().items(): + self._send_settlement(settlement_id, payouts) + banned: set[str] = set() + if self._server is not None and self._server.contracts is not None: + registry = self._server.contracts.registry + banned = { + wallet for wallet, _ in billing.payables( + threshold=0.0, max_period=0.0, dust_floor=0.0, + ) + if registry.get_wallet(wallet).banned + } + due = billing.payables( + threshold=self._payout_threshold, + max_period=self._settle_period, + dust_floor=self._payout_dust_floor, + exclude=banned, + ) + if not due: + continue + settlement_id = uuid.uuid4().hex + registry = ( + self._server.contracts.registry + if self._server is not None and self._server.contracts is not None + else None + ) + settled: list[tuple[str, float]] = [] + for wallet, amount in due: + # ADR-0015: a fraud forfeiture landing between the payables() + # snapshot above and this debit must never be paid out on top + # of — recheck the ban and let settle_node_payout clamp to + # whatever is still actually pending. + if registry is not None and registry.get_wallet(wallet).banned: + continue + event = billing.settle_node_payout(wallet, amount, reference=settlement_id) + if event["amount"] > 0: + settled.append((wallet, event["amount"])) + if settled: + self._send_settlement(settlement_id, settled) + + def _send_settlement(self, settlement_id: str, payouts: list) -> None: + billing = self._billing + treasury = self._treasury + assert billing is not None and treasury is not None + try: + signature = treasury.send_payouts([(w, a) for w, a in payouts]) + except Exception as exc: + print( + f"[tracker] settlement {settlement_id} failed (will retry): {exc}", + flush=True, + ) + return + billing.confirm_settlement(settlement_id, signature) + total = sum(a for _, a in payouts) + print( + f"[tracker] settled {settlement_id}: {len(payouts)} payout(s), " + f"{total:.6f} USDT total (tx {signature})", + flush=True, + ) + + def _deposit_loop(self) -> None: + """Poll the treasury token account and credit confirmed deposits (US-032).""" + billing = self._billing + treasury = self._treasury + assert billing is not None and treasury is not None + while not self._deposit_stop.wait(self._deposit_poll_interval): + try: + deposits = treasury.list_new_deposits( + lambda sig: billing.has_event(f"deposit-{sig}") + ) + except Exception as exc: + print(f"[tracker] deposit poll failed: {exc}", flush=True) + continue + for deposit in deposits: + api_key = billing.api_key_for_wallet(deposit.sender_wallet) + if api_key is None: + print( + f"[tracker] unattributed deposit {deposit.signature}: " + f"{deposit.amount_usdt} USDT from unbound wallet " + f"{deposit.sender_wallet} — register via /v1/wallet/register", + flush=True, + ) + continue + credited = billing.credit_deposit( + api_key, deposit.amount_usdt, deposit.signature + ) + if credited is not None: + print( + f"[tracker] deposit credited: {deposit.amount_usdt} USDT " + f"-> api_key …{api_key[-6:]} (tx {deposit.signature})", + flush=True, + ) + + def _hf_pricing_loop(self) -> None: + """Daily dynamic pricing refresh benchmarked against HF inference rates (issue 23). + + For every preset with a curated, human-verified ``hf_aliases`` list, + fetch current HF marketplace pricing and set the client price to 80% + of the cheapest matching alias. Presets with no (or an empty) + ``hf_aliases`` are left entirely alone — they keep the static + default price. Any single preset's fetch/parse failure is logged and + skipped; it never raises into this loop or the request path. + """ + billing = self._billing + assert billing is not None + while not self._hf_pricing_stop.wait(self._hf_pricing_refresh_interval): + for name, preset in list(self._model_presets.items()): + if not isinstance(preset, dict) or not preset.get("hf_aliases"): + continue + try: + result = refresh_preset_price( + model_name=name, + preset=preset, + current_price=billing.price_for(name), + fetch_html=self._hf_pricing_fetch_html, + ) + except Exception as exc: + print(f"[tracker] hf pricing refresh failed for {name!r}: {exc}", flush=True) + continue + if result is None: + continue + new_input = result.get("new_input_price_per_1k", result["new_price_per_1k"]) + new_output = result.get("new_output_price_per_1k", result["new_price_per_1k"]) + for key in _preset_price_keys(name, preset): + billing.set_prices(key, new_input, new_output) + preset["hf_last_price_per_1k"] = result["new_price_per_1k"] + preset["hf_last_updated"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) + if self._hf_pricing_log is not None: + self._hf_pricing_log.record_change( + model=name, + old_price_per_1k=result["old_price_per_1k"], + new_price_per_1k=result["new_price_per_1k"], + source_repo_id=result["source_repo_id"], + source_provider=result["source_provider"], + ) + print( + f"[tracker] hf pricing: {name} {result['old_price_per_1k']:.6f} -> " + f"{result['new_price_per_1k']:.6f} USDT/1k tokens " + f"(80% of {result['source_repo_id']}::{result['source_provider']})", + flush=True, + ) + + def _raft_apply(self, command: str, payload: dict) -> None: + """Called by RaftNode when a log entry is committed — replicate to local registry.""" + if command != "register": + return + # Re-apply the registration to our local registry (follower path). + # On the leader this is a no-op since it already registered locally. + node_id = payload.get("node_id") + if not node_id or node_id in self._registry: + return # already present (leader case) or malformed + endpoint = payload.get("endpoint", "") + try: + shard_start = int(payload["shard_start"]) if payload.get("shard_start") is not None else None + shard_end = int(payload["shard_end"]) if payload.get("shard_end") is not None else None + except (TypeError, ValueError): + return + entry = _NodeEntry( + node_id=node_id, + endpoint=endpoint.rstrip("/"), + shard_start=shard_start, + shard_end=shard_end, + model=payload.get("model", "stub-model"), + shard_checksum=payload.get("shard_checksum"), + hardware_profile=payload.get("hardware_profile", {}), + wallet_address=payload.get("wallet_address"), + score=float(payload.get("score", 1.0)), + tracker_mode=bool(payload.get("tracker_mode", False)), + hf_repo=payload.get("hf_repo"), + num_layers=int(payload["num_layers"]) if payload.get("num_layers") is not None else None, + model_metadata=payload.get("model_metadata") if isinstance(payload.get("model_metadata"), dict) else None, + downloaded_models=( + payload.get("downloaded_models") + if isinstance(payload.get("downloaded_models"), list) + else None + ), + ) + with self._lock: + self._registry[node_id] = entry + + def _rebalance_loop(self) -> None: + while not self._rebalance_stop.wait(self._rebalance_interval): + server = self._server + if server is None: + return + with self._lock: + _purge_expired_nodes_locked(server) + _rebalance_all_locked(server) + _scale_demanded_models_locked(server) + + def _push_to_peers(self, path: str, body: bytes) -> bool: + """POST a hive-signed payload to every cluster peer; True if all succeeded.""" + headers = {"Content-Type": "application/json"} + if self._hive_secret: + headers.update(sign_hive_request(self._hive_secret, body)) + delivered_all = True + for peer in self._cluster_peers: + try: + req = urllib.request.Request( + f"{peer}{path}", data=body, headers=headers, method="POST", + ) + with urllib.request.urlopen(req, timeout=2.0) as r: + r.read() + except Exception: + delivered_all = False + return delivered_all + + def _stats_loop(self) -> None: + """Periodically save stats/billing to DB and push local slices to cluster peers.""" + while not self._stats_stop.wait(_StatsCollector.SAVE_INTERVAL): + if self._stats is not None: + self._stats.save_to_db() + if self._billing is not None: + self._billing.save_to_db() + if self._accounts is not None: + self._accounts.save_to_db() + registry_log = getattr(self._contracts, "registry_log", None) + if registry_log is not None: + registry_log.save_to_db() + if self._cluster_peers and not self._hive_secret: + print( + "[tracker] WARNING: cluster peers configured without --hive-secret — " + "gossip pushes will be rejected (ADR-0017)", + flush=True, + ) + if self._stats is not None and self._cluster_peers: + self_url = self._cluster_self_url or f"http://{self._host}:{self.port}" + local_rpms = self._stats.get_local_rpms() + body = json.dumps({"tracker_url": self_url, "stats": local_rpms}).encode() + self._push_to_peers("/v1/stats/gossip", body) + if self._billing is not None and self._cluster_peers: + events, cursor = self._billing.events_since(self._billing_gossip_cursor) + if events: + body = json.dumps({"events": events}).encode() + # Only advance past events every peer has seen; unreachable + # peers get the full backlog on the next tick (idempotent + # via event-id dedupe on the receiving side). + if self._push_to_peers("/v1/billing/gossip", body): + self._billing_gossip_cursor = cursor + if self._accounts is not None and self._cluster_peers: + events, cursor = self._accounts.events_since(self._accounts_gossip_cursor) + if events: + body = json.dumps({"events": events}).encode() + if self._push_to_peers("/v1/accounts/gossip", body): + self._accounts_gossip_cursor = cursor + registry_log = getattr(self._contracts, "registry_log", None) + if registry_log is not None and self._cluster_peers: + events, cursor = registry_log.events_since(self._registry_gossip_cursor) + if events: + body = json.dumps({"events": events}).encode() + if self._push_to_peers("/v1/registry/gossip", body): + self._registry_gossip_cursor = cursor + + def stop(self) -> None: + if self._raft is not None: + self._raft.stop() + if self._gossip is not None: + self._gossip.stop() + if self._server is None: + return + self._rebalance_stop.set() + self._stats_stop.set() + self._deposit_stop.set() + self._settlement_stop.set() + self._hf_pricing_stop.set() + if self._stats is not None: + self._stats.save_to_db() + if self._billing is not None: + self._billing.save_to_db() + if self._accounts is not None: + self._accounts.save_to_db() + registry_log = getattr(self._contracts, "registry_log", None) + if registry_log is not None: + registry_log.save_to_db() + self._server.shutdown() + self._server.server_close() + if self._thread is not None: + self._thread.join(timeout=1) + if self._rebalance_thread is not None: + self._rebalance_thread.join(timeout=1) + if self._stats_thread is not None: + self._stats_thread.join(timeout=1) + if self._deposit_thread is not None: + self._deposit_thread.join(timeout=1) + self._deposit_thread = None + if self._settlement_thread is not None: + self._settlement_thread.join(timeout=1) + self._settlement_thread = None + if self._hf_pricing_thread is not None: + self._hf_pricing_thread.join(timeout=1) + self._hf_pricing_thread = None + self._server = None + self._thread = None + self._rebalance_thread = None + self._stats_thread = None + self._raft = None + self._gossip = None + self.port = None diff --git a/packages/validator/meshnet_validator/__init__.py b/packages/validator/meshnet_validator/__init__.py index bf6c96b..b55b104 100644 --- a/packages/validator/meshnet_validator/__init__.py +++ b/packages/validator/meshnet_validator/__init__.py @@ -1,552 +1,552 @@ -"""Optimistic fraud validator for completed inference requests.""" - -import json -import math -import random -import threading -import time -import urllib.request -from typing import Any - -from .audit import ( - ToplocAuditConfig, - ToplocProofClaim, - ToplocVerificationResult, - verify_activation_proofs, - verify_activation_proofs_detailed, -) -from .sampling import AdaptiveAuditSampler, AuditRateConfig -from .tripwire import detect_output_tripwire - -__version__ = "0.1.0" - - -class ValidatorProcess: - """Separate validator loop that samples completed requests and submits slashes.""" - - def __init__( - self, - *, - contracts: Any, - reference_node_url: str, - sample_rate: float = 0.05, - tolerance: float = 1e-6, - slash_amount: int = 100, - strike_threshold: int = 3, - random_seed: int | None = None, - webhook_url: str | None = None, - interval_seconds: float = 1.0, - billing: Any | None = None, - toploc_config: ToplocAuditConfig | None = None, - toploc_backend: Any | None = None, - audit_sampler: AdaptiveAuditSampler | None = None, - ) -> None: - if not 0.0 <= sample_rate <= 1.0: - raise ValueError("sample_rate must be between 0 and 1") - if tolerance < 0: - raise ValueError("tolerance must be non-negative") - if slash_amount <= 0: - raise ValueError("slash_amount must be positive") - if strike_threshold <= 0: - raise ValueError("strike_threshold must be positive") - if interval_seconds <= 0: - raise ValueError("interval_seconds must be positive") - - self._contracts = contracts - self._billing = billing - self._reference_node_url = reference_node_url.rstrip("/") - self._sample_rate = sample_rate - self._tolerance = tolerance - self._slash_amount = slash_amount - self._strike_threshold = strike_threshold - self._webhook_url = webhook_url - self._interval_seconds = interval_seconds - self._toploc_config = toploc_config or ToplocAuditConfig() - self._toploc_backend = toploc_backend - self._audit_sampler = audit_sampler - self._random = random.Random(random_seed) - self._last_event_index = -1 - self._running = False - self._thread: threading.Thread | None = None - self.sampled_count = 0 - - def validate_once(self) -> list[Any]: - """Run one validation cycle and return slash receipts submitted this cycle.""" - receipts: list[Any] = [] - events = self._contracts.validation.list_completed_inferences( - after_index=self._last_event_index, - ) - for event in events: - self._last_event_index = max(self._last_event_index, event.index) - if not self._should_sample(event): - continue - self.sampled_count += 1 - audit_result = self._validate_event(event) - if audit_result.ok: - self._record_clean_audit(event) - continue - receipts.extend(self._slash_node( - audit_result.culprit_node, - event.observed_output, - audit_result.reference_output, - reason=audit_result.reason, - )) - return receipts - - def _should_sample(self, event: Any) -> bool: - """ADR-0018 §1/§6-7: flat sample_rate stays the default; when an - AdaptiveAuditSampler is configured, the decision is reputation- and - tenure-weighted and budget-balanced against the fleet-wide target - instead of a uniform coin flip.""" - if self._audit_sampler is None: - return self._random.random() < self._sample_rate - - tripwire = detect_output_tripwire(_event_value(event, "observed_output") or "") - wallets = _route_wallets(event) - if not wallets: - return self._audit_sampler.should_audit( - completed_job_count=0, reputation=1.0, tripwire=tripwire, - ) - # A route is only as trustworthy as its least-trusted hop -- audit - # against whichever wallet on the route looks riskiest. - riskiest = min( - (self._contracts.registry.get_wallet(wallet) for wallet in wallets), - key=lambda wallet: wallet.reputation, - ) - return self._audit_sampler.should_audit( - completed_job_count=riskiest.completed_job_count, - reputation=riskiest.reputation, - tripwire=tripwire, - ) - - def start(self) -> None: - if self._running: - raise RuntimeError("ValidatorProcess is already running") - self._running = True - self._thread = threading.Thread(target=self._run_loop, daemon=True) - self._thread.start() - - def stop(self) -> None: - self._running = False - if self._thread is not None: - self._thread.join(timeout=2) - self._thread = None - - def _run_loop(self) -> None: - while self._running: - self.validate_once() - time.sleep(self._interval_seconds) - - def _run_reference(self, messages: list[dict]) -> str: - response = _post_json( - f"{self._reference_node_url}/v1/infer", - {"messages": messages}, - ) - text = response.get("text") - if not isinstance(text, str): - raise ValueError("reference node response did not contain text") - return text - - def _validate_event(self, event: Any) -> "_AuditResult": - event = self._event_with_on_demand_commitments(event) - hop_commitments = _hop_commitments_from_event(event) - if hop_commitments is not None and self._commitment_expired(event): - # ADR-0018 §3: the on-demand retention window has passed — nodes - # are no longer expected to hold the boundary activations needed - # to verify this commitment, so fall back to the text-only path. - hop_commitments = None - - if hop_commitments is None: - reference_output = self._run_reference(event.messages) - ok = _outputs_match(event.observed_output, reference_output, self._tolerance) - return _AuditResult( - ok=ok, - reference_output=reference_output, - reason="reference output diverged", - # Text comparison has no per-hop signal; the last hop is the - # best-effort guess (text-only fallback), never used when - # hop-boundary commitments make real bisection possible. - culprit_node=None if ok else _final_text_node(event.route_nodes), - ) - - if len(hop_commitments) == 1: - # Single-commitment route (AH-006 whole-route format, or a - # genuinely one-hop pipeline): reuse the original teacher-forced - # call so existing single-hop reference integrations keep working. - only = hop_commitments[0] - reference_activations_by_hop = [self._run_teacher_forced_prefill( - model=_event_value(event, "model"), - messages=_event_value(event, "messages"), - claimed_token_ids=only.token_ids, - claim=only.claim, - )] - else: - reference_activations_by_hop = self._run_teacher_forced_prefill_hops( - model=_event_value(event, "model"), - messages=_event_value(event, "messages"), - hop_commitments=hop_commitments, - ) - culprit_index = _first_divergent_hop( - hop_commitments, - reference_activations_by_hop, - config=self._toploc_config, - backend=self._toploc_backend, - ) - ok = culprit_index is None - return _AuditResult( - ok=ok, - reference_output=( - "TOPLOC activation proof accepted" - if ok - else f"TOPLOC activation proof mismatch at hop {culprit_index}" - ), - reason="TOPLOC activation proof mismatch", - culprit_node=None if ok else hop_commitments[culprit_index].node, - ) - - def _event_with_on_demand_commitments(self, event: Any) -> Any: - """Fetch missing per-hop TOPLOC commitments only after audit sampling. - - Tracker validation events deliberately carry ordinary route metadata, - not a pre-announced audit flag. When this validator samples an event, it - asks each hop for its short-lived boundary commitment and splices the - returned proof into a local event copy for the bisection verifier. - """ - route_nodes = _event_value(event, "route_nodes") or [] - if not isinstance(route_nodes, list) or not route_nodes: - return event - updated_nodes: list[dict] = [] - changed = False - for node in route_nodes: - if not isinstance(node, dict): - updated_nodes.append(node) - continue - updated = dict(node) - if _mapping_value(updated, "toploc_proof") is None: - commitment = self._fetch_hop_commitment(event, updated) - if commitment is not None: - updated.update(commitment) - changed = True - updated_nodes.append(updated) - if not changed: - return event - if isinstance(event, dict): - copied = dict(event) - else: - copied = dict(vars(event)) - copied["route_nodes"] = updated_nodes - return copied - - def _fetch_hop_commitment(self, event: Any, node: dict) -> dict[str, Any] | None: - endpoint = node.get("endpoint") - if not isinstance(endpoint, str) or not endpoint: - return None - try: - response = _post_json( - f"{endpoint.rstrip('/')}/v1/audit/toploc/commitment", - { - "session_id": _event_value(event, "session_id"), - "model": _event_value(event, "model"), - "messages": _event_value(event, "messages") or [], - "shard_start": node.get("shard_start"), - "shard_end": node.get("shard_end"), - }, - timeout=2.0, - ) - except (OSError, ValueError, json.JSONDecodeError): - return None - proof = response.get("toploc_proof") or response.get("activation_proof") - token_ids = response.get("claimed_token_ids") or response.get("output_token_ids") - if not isinstance(proof, dict): - return None - if not isinstance(token_ids, list) or not all(isinstance(token, int) for token in token_ids): - return None - return {"toploc_proof": proof, "claimed_token_ids": token_ids} - - def _commitment_expired(self, event: Any) -> bool: - ts = _event_value(event, "ts") - if ts is None: - return False - return (time.time() - float(ts)) > self._toploc_config.commitment_ttl_seconds - - def _run_teacher_forced_prefill( - self, - *, - model: str, - messages: list[dict], - claimed_token_ids: list[int], - claim: ToplocProofClaim, - ) -> list[Any]: - response = _post_json( - f"{self._reference_node_url}/v1/audit/toploc", - { - "model": model, - "messages": messages, - "claimed_token_ids": claimed_token_ids, - "dtype": claim.dtype, - "quantization": claim.quantization, - "decode_batching_size": claim.decode_batching_size, - "topk": claim.topk, - "skip_prefill": claim.skip_prefill, - }, - ) - activations = response.get("activations") - if not isinstance(activations, list): - raise ValueError("reference node audit response did not contain activations") - return activations - - def _run_teacher_forced_prefill_hops( - self, - *, - model: str, - messages: list[dict], - hop_commitments: list["_HopCommitment"], - ) -> list[list[Any]]: - """Teacher-force the claimed tokens once and collect reference - activations at every hop's boundary layer (ADR-0018 §4 / research §1.2: - one referee forward pass, compared at each cut-point).""" - reference_claim = hop_commitments[0].claim - response = _post_json( - f"{self._reference_node_url}/v1/audit/toploc", - { - "model": model, - "messages": messages, - "claimed_token_ids": hop_commitments[-1].token_ids, - "hop_boundaries": [hop.shard_end for hop in hop_commitments], - "dtype": reference_claim.dtype, - "quantization": reference_claim.quantization, - "decode_batching_size": reference_claim.decode_batching_size, - "topk": reference_claim.topk, - "skip_prefill": reference_claim.skip_prefill, - }, - ) - activations_by_hop = response.get("activations_by_hop") - if not isinstance(activations_by_hop, list) or len(activations_by_hop) != len(hop_commitments): - raise ValueError("reference node audit response did not contain per-hop activations") - return activations_by_hop - - def _record_clean_audit(self, event: Any) -> None: - """ADR-0018 §6: reputation derives only from tracker-verified audit - outcomes — a clean audit credits every node on the verified route.""" - for wallet_address in _route_wallets(event): - if self._contracts.registry.get_wallet(wallet_address).banned: - continue - self._contracts.registry.record_audit_outcome(wallet_address, passed=True) - - def _slash_node( - self, - node: dict | None, - observed_output: str, - reference_output: str, - *, - reason: str = "reference output diverged", - ) -> list[Any]: - receipts: list[Any] = [] - wallet_address = node.get("wallet_address") if node else None - if not wallet_address: - return receipts - if self._contracts.registry.get_wallet(wallet_address).banned: - return receipts - receipts.append(self._contracts.registry.submit_slash_proof( - wallet_address=wallet_address, - slash_amount=self._slash_amount, - strike_threshold=self._strike_threshold, - reason=( - f"{reason} " - f"(observed={observed_output!r}, reference={reference_output!r})" - ), - webhook_url=self._webhook_url, - )) - # ADR-0018 §6: reputation loss is separate from the strike/ban that - # submit_slash_proof already recorded above — never double-strike. - self._contracts.registry.record_audit_outcome(wallet_address, passed=False) - # ADR-0015: the pending balance is the collateral — forfeit it in the - # same validation cycle as the strike. - if self._billing is not None: - forfeit = self._billing.forfeit_pending(wallet_address, reason="fraud-divergence") - print( - f"[validator] forfeited pending balance of {wallet_address}: " - f"{forfeit['amount']:.6f} USDT (fraud-divergence)", - flush=True, - ) - return receipts - - -def _route_wallets(event: Any) -> list[str]: - """Unique wallet addresses across a route, in hop order.""" - route_nodes = _event_value(event, "route_nodes") or [] - seen: set[str] = set() - wallets: list[str] = [] - for node in route_nodes: - wallet_address = node.get("wallet_address") if isinstance(node, dict) else None - if wallet_address and wallet_address not in seen: - seen.add(wallet_address) - wallets.append(wallet_address) - return wallets - - -def _final_text_node(route_nodes: list[dict]) -> dict | None: - """Text-only fallback blame: when the audit has no per-hop fingerprints - to bisect (free-running text comparison only), guess the last hop. - Never used once hop-boundary commitments make real bisection possible.""" - if not route_nodes: - return None - return max(route_nodes, key=lambda node: int(node.get("shard_end", 0))) - - -class _AuditResult: - def __init__( - self, - *, - ok: bool, - reference_output: str, - reason: str, - culprit_node: dict | None = None, - ) -> None: - self.ok = ok - self.reference_output = reference_output - self.reason = reason - self.culprit_node = culprit_node - - -class _HopCommitment: - """One hop's on-demand TOPLOC commitment plus the route node it blames.""" - - def __init__(self, node: dict | None, claim: ToplocProofClaim, token_ids: list[int]) -> None: - self.node = node - self.claim = claim - self.token_ids = token_ids - self.shard_end = int(node.get("shard_end", 0)) if node else None - - -def _hop_commitments_from_event(event: Any) -> list[_HopCommitment] | None: - """Per-hop bisection commitments (ADR-0018 §3/§4): each route node reports - its own output-boundary fingerprint. Falls back to the AH-006 whole-route - commitment format (one fingerprint, no bisection) when hops don't carry - individual commitments.""" - route_nodes = _event_value(event, "route_nodes") or [] - per_hop_nodes = [ - node for node in route_nodes - if isinstance(node, dict) and _mapping_value(node, "toploc_proof") is not None - ] - if per_hop_nodes and len(per_hop_nodes) == len(route_nodes): - ordered = sorted(route_nodes, key=lambda node: int(node.get("shard_start", 0))) - default_token_ids = _event_value(event, "claimed_token_ids") - commitments = [] - for node in ordered: - token_ids = node.get("claimed_token_ids", default_token_ids) - if not isinstance(token_ids, list) or not all(isinstance(t, int) for t in token_ids): - raise ValueError("TOPLOC hop commitments must include claimed_token_ids") - commitments.append(_HopCommitment( - node, - ToplocProofClaim.from_mapping(node["toploc_proof"]), - token_ids, - )) - return commitments - - whole_route = _toploc_audit_from_event(event) - if whole_route is None: - return None - token_ids, claim = whole_route - return [_HopCommitment(route_nodes[-1] if route_nodes else None, claim, token_ids)] - - -def _first_divergent_hop( - hop_commitments: list[_HopCommitment], - reference_activations_by_hop: list[Any], - *, - config: ToplocAuditConfig, - backend: Any | None, -) -> int | None: - """First hop whose committed output fingerprint diverges from the - referee's independently-computed reference activations at that same - cut-point (research §1.2: no interactive game needed at hop granularity — - the referee checks every cut-point in one replay).""" - for index, commitment in enumerate(hop_commitments): - ok = verify_activation_proofs( - reference_activations_by_hop[index], - commitment.claim, - config=config, - backend=backend, - ) - if not ok: - return index - return None - - -def _toploc_audit_from_event(event: Any) -> tuple[list[int], ToplocProofClaim] | None: - audit = _event_mapping(event, "audit") - claim_data = ( - _event_mapping(event, "toploc_proof") - or _event_mapping(event, "activation_proof") - or _mapping_value(audit, "toploc_proof") - or _mapping_value(audit, "activation_proof") - or _mapping_value(audit, "toploc") - ) - if claim_data is None: - return None - token_ids = ( - _event_value(event, "claimed_token_ids") - or _event_value(event, "output_token_ids") - or _mapping_value(audit, "claimed_token_ids") - or _mapping_value(audit, "output_token_ids") - ) - if not isinstance(token_ids, list) or not all(isinstance(token, int) for token in token_ids): - raise ValueError("TOPLOC audit events must include claimed_token_ids") - return token_ids, ToplocProofClaim.from_mapping(claim_data) - - -def _event_value(event: Any, name: str) -> Any: - if hasattr(event, name): - return getattr(event, name) - if isinstance(event, dict): - return event.get(name) - return None - - -def _event_mapping(event: Any, name: str) -> dict[str, Any] | None: - value = _event_value(event, name) - return value if isinstance(value, dict) else None - - -def _mapping_value(mapping: dict[str, Any] | None, name: str) -> Any: - if mapping is None: - return None - return mapping.get(name) - - -def _outputs_match(observed: str, reference: str, tolerance: float) -> bool: - observed_float = _parse_float(observed) - reference_float = _parse_float(reference) - if observed_float is not None and reference_float is not None: - return math.isclose(observed_float, reference_float, rel_tol=tolerance, abs_tol=tolerance) - return observed == reference - - -def _parse_float(value: str) -> float | None: - try: - return float(value) - except ValueError: - return None - - -def _post_json(url: str, payload: dict, timeout: float = 5.0) -> 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, timeout=timeout) as response: - return json.loads(response.read()) - - -__all__ = [ - "ToplocAuditConfig", - "ToplocProofClaim", - "ValidatorProcess", - "AdaptiveAuditSampler", - "AuditRateConfig", - "detect_output_tripwire", -] +"""Optimistic fraud validator for completed inference requests.""" + +import json +import math +import random +import threading +import time +import urllib.request +from typing import Any + +from .audit import ( + ToplocAuditConfig, + ToplocProofClaim, + ToplocVerificationResult, + verify_activation_proofs, + verify_activation_proofs_detailed, +) +from .sampling import AdaptiveAuditSampler, AuditRateConfig +from .tripwire import detect_output_tripwire + +__version__ = "0.1.0" + + +class ValidatorProcess: + """Separate validator loop that samples completed requests and submits slashes.""" + + def __init__( + self, + *, + contracts: Any, + reference_node_url: str, + sample_rate: float = 0.05, + tolerance: float = 1e-6, + slash_amount: int = 100, + strike_threshold: int = 3, + random_seed: int | None = None, + webhook_url: str | None = None, + interval_seconds: float = 1.0, + billing: Any | None = None, + toploc_config: ToplocAuditConfig | None = None, + toploc_backend: Any | None = None, + audit_sampler: AdaptiveAuditSampler | None = None, + ) -> None: + if not 0.0 <= sample_rate <= 1.0: + raise ValueError("sample_rate must be between 0 and 1") + if tolerance < 0: + raise ValueError("tolerance must be non-negative") + if slash_amount <= 0: + raise ValueError("slash_amount must be positive") + if strike_threshold <= 0: + raise ValueError("strike_threshold must be positive") + if interval_seconds <= 0: + raise ValueError("interval_seconds must be positive") + + self._contracts = contracts + self._billing = billing + self._reference_node_url = reference_node_url.rstrip("/") + self._sample_rate = sample_rate + self._tolerance = tolerance + self._slash_amount = slash_amount + self._strike_threshold = strike_threshold + self._webhook_url = webhook_url + self._interval_seconds = interval_seconds + self._toploc_config = toploc_config or ToplocAuditConfig() + self._toploc_backend = toploc_backend + self._audit_sampler = audit_sampler + self._random = random.Random(random_seed) + self._last_event_index = -1 + self._running = False + self._thread: threading.Thread | None = None + self.sampled_count = 0 + + def validate_once(self) -> list[Any]: + """Run one validation cycle and return slash receipts submitted this cycle.""" + receipts: list[Any] = [] + events = self._contracts.validation.list_completed_inferences( + after_index=self._last_event_index, + ) + for event in events: + self._last_event_index = max(self._last_event_index, event.index) + if not self._should_sample(event): + continue + self.sampled_count += 1 + audit_result = self._validate_event(event) + if audit_result.ok: + self._record_clean_audit(event) + continue + receipts.extend(self._slash_node( + audit_result.culprit_node, + event.observed_output, + audit_result.reference_output, + reason=audit_result.reason, + )) + return receipts + + def _should_sample(self, event: Any) -> bool: + """ADR-0018 §1/§6-7: flat sample_rate stays the default; when an + AdaptiveAuditSampler is configured, the decision is reputation- and + tenure-weighted and budget-balanced against the fleet-wide target + instead of a uniform coin flip.""" + if self._audit_sampler is None: + return self._random.random() < self._sample_rate + + tripwire = detect_output_tripwire(_event_value(event, "observed_output") or "") + wallets = _route_wallets(event) + if not wallets: + return self._audit_sampler.should_audit( + completed_job_count=0, reputation=1.0, tripwire=tripwire, + ) + # A route is only as trustworthy as its least-trusted hop -- audit + # against whichever wallet on the route looks riskiest. + riskiest = min( + (self._contracts.registry.get_wallet(wallet) for wallet in wallets), + key=lambda wallet: wallet.reputation, + ) + return self._audit_sampler.should_audit( + completed_job_count=riskiest.completed_job_count, + reputation=riskiest.reputation, + tripwire=tripwire, + ) + + def start(self) -> None: + if self._running: + raise RuntimeError("ValidatorProcess is already running") + self._running = True + self._thread = threading.Thread(target=self._run_loop, daemon=True) + self._thread.start() + + def stop(self) -> None: + self._running = False + if self._thread is not None: + self._thread.join(timeout=2) + self._thread = None + + def _run_loop(self) -> None: + while self._running: + self.validate_once() + time.sleep(self._interval_seconds) + + def _run_reference(self, messages: list[dict]) -> str: + response = _post_json( + f"{self._reference_node_url}/v1/infer", + {"messages": messages}, + ) + text = response.get("text") + if not isinstance(text, str): + raise ValueError("reference node response did not contain text") + return text + + def _validate_event(self, event: Any) -> "_AuditResult": + event = self._event_with_on_demand_commitments(event) + hop_commitments = _hop_commitments_from_event(event) + if hop_commitments is not None and self._commitment_expired(event): + # ADR-0018 §3: the on-demand retention window has passed — nodes + # are no longer expected to hold the boundary activations needed + # to verify this commitment, so fall back to the text-only path. + hop_commitments = None + + if hop_commitments is None: + reference_output = self._run_reference(event.messages) + ok = _outputs_match(event.observed_output, reference_output, self._tolerance) + return _AuditResult( + ok=ok, + reference_output=reference_output, + reason="reference output diverged", + # Text comparison has no per-hop signal; the last hop is the + # best-effort guess (text-only fallback), never used when + # hop-boundary commitments make real bisection possible. + culprit_node=None if ok else _final_text_node(event.route_nodes), + ) + + if len(hop_commitments) == 1: + # Single-commitment route (AH-006 whole-route format, or a + # genuinely one-hop pipeline): reuse the original teacher-forced + # call so existing single-hop reference integrations keep working. + only = hop_commitments[0] + reference_activations_by_hop = [self._run_teacher_forced_prefill( + model=_event_value(event, "model"), + messages=_event_value(event, "messages"), + claimed_token_ids=only.token_ids, + claim=only.claim, + )] + else: + reference_activations_by_hop = self._run_teacher_forced_prefill_hops( + model=_event_value(event, "model"), + messages=_event_value(event, "messages"), + hop_commitments=hop_commitments, + ) + culprit_index = _first_divergent_hop( + hop_commitments, + reference_activations_by_hop, + config=self._toploc_config, + backend=self._toploc_backend, + ) + ok = culprit_index is None + return _AuditResult( + ok=ok, + reference_output=( + "TOPLOC activation proof accepted" + if ok + else f"TOPLOC activation proof mismatch at hop {culprit_index}" + ), + reason="TOPLOC activation proof mismatch", + culprit_node=None if ok else hop_commitments[culprit_index].node, + ) + + def _event_with_on_demand_commitments(self, event: Any) -> Any: + """Fetch missing per-hop TOPLOC commitments only after audit sampling. + + Tracker validation events deliberately carry ordinary route metadata, + not a pre-announced audit flag. When this validator samples an event, it + asks each hop for its short-lived boundary commitment and splices the + returned proof into a local event copy for the bisection verifier. + """ + route_nodes = _event_value(event, "route_nodes") or [] + if not isinstance(route_nodes, list) or not route_nodes: + return event + updated_nodes: list[dict] = [] + changed = False + for node in route_nodes: + if not isinstance(node, dict): + updated_nodes.append(node) + continue + updated = dict(node) + if _mapping_value(updated, "toploc_proof") is None: + commitment = self._fetch_hop_commitment(event, updated) + if commitment is not None: + updated.update(commitment) + changed = True + updated_nodes.append(updated) + if not changed: + return event + if isinstance(event, dict): + copied = dict(event) + else: + copied = dict(vars(event)) + copied["route_nodes"] = updated_nodes + return copied + + def _fetch_hop_commitment(self, event: Any, node: dict) -> dict[str, Any] | None: + endpoint = node.get("endpoint") + if not isinstance(endpoint, str) or not endpoint: + return None + try: + response = _post_json( + f"{endpoint.rstrip('/')}/v1/audit/toploc/commitment", + { + "session_id": _event_value(event, "session_id"), + "model": _event_value(event, "model"), + "messages": _event_value(event, "messages") or [], + "shard_start": node.get("shard_start"), + "shard_end": node.get("shard_end"), + }, + timeout=2.0, + ) + except (OSError, ValueError, json.JSONDecodeError): + return None + proof = response.get("toploc_proof") or response.get("activation_proof") + token_ids = response.get("claimed_token_ids") or response.get("output_token_ids") + if not isinstance(proof, dict): + return None + if not isinstance(token_ids, list) or not all(isinstance(token, int) for token in token_ids): + return None + return {"toploc_proof": proof, "claimed_token_ids": token_ids} + + def _commitment_expired(self, event: Any) -> bool: + ts = _event_value(event, "ts") + if ts is None: + return False + return (time.time() - float(ts)) > self._toploc_config.commitment_ttl_seconds + + def _run_teacher_forced_prefill( + self, + *, + model: str, + messages: list[dict], + claimed_token_ids: list[int], + claim: ToplocProofClaim, + ) -> list[Any]: + response = _post_json( + f"{self._reference_node_url}/v1/audit/toploc", + { + "model": model, + "messages": messages, + "claimed_token_ids": claimed_token_ids, + "dtype": claim.dtype, + "quantization": claim.quantization, + "decode_batching_size": claim.decode_batching_size, + "topk": claim.topk, + "skip_prefill": claim.skip_prefill, + }, + ) + activations = response.get("activations") + if not isinstance(activations, list): + raise ValueError("reference node audit response did not contain activations") + return activations + + def _run_teacher_forced_prefill_hops( + self, + *, + model: str, + messages: list[dict], + hop_commitments: list["_HopCommitment"], + ) -> list[list[Any]]: + """Teacher-force the claimed tokens once and collect reference + activations at every hop's boundary layer (ADR-0018 §4 / research §1.2: + one referee forward pass, compared at each cut-point).""" + reference_claim = hop_commitments[0].claim + response = _post_json( + f"{self._reference_node_url}/v1/audit/toploc", + { + "model": model, + "messages": messages, + "claimed_token_ids": hop_commitments[-1].token_ids, + "hop_boundaries": [hop.shard_end for hop in hop_commitments], + "dtype": reference_claim.dtype, + "quantization": reference_claim.quantization, + "decode_batching_size": reference_claim.decode_batching_size, + "topk": reference_claim.topk, + "skip_prefill": reference_claim.skip_prefill, + }, + ) + activations_by_hop = response.get("activations_by_hop") + if not isinstance(activations_by_hop, list) or len(activations_by_hop) != len(hop_commitments): + raise ValueError("reference node audit response did not contain per-hop activations") + return activations_by_hop + + def _record_clean_audit(self, event: Any) -> None: + """ADR-0018 §6: reputation derives only from tracker-verified audit + outcomes — a clean audit credits every node on the verified route.""" + for wallet_address in _route_wallets(event): + if self._contracts.registry.get_wallet(wallet_address).banned: + continue + self._contracts.registry.record_audit_outcome(wallet_address, passed=True) + + def _slash_node( + self, + node: dict | None, + observed_output: str, + reference_output: str, + *, + reason: str = "reference output diverged", + ) -> list[Any]: + receipts: list[Any] = [] + wallet_address = node.get("wallet_address") if node else None + if not wallet_address: + return receipts + if self._contracts.registry.get_wallet(wallet_address).banned: + return receipts + receipts.append(self._contracts.registry.submit_slash_proof( + wallet_address=wallet_address, + slash_amount=self._slash_amount, + strike_threshold=self._strike_threshold, + reason=( + f"{reason} " + f"(observed={observed_output!r}, reference={reference_output!r})" + ), + webhook_url=self._webhook_url, + )) + # ADR-0018 §6: reputation loss is separate from the strike/ban that + # submit_slash_proof already recorded above — never double-strike. + self._contracts.registry.record_audit_outcome(wallet_address, passed=False) + # ADR-0015: the pending balance is the collateral — forfeit it in the + # same validation cycle as the strike. + if self._billing is not None: + forfeit = self._billing.forfeit_pending(wallet_address, reason="fraud-divergence") + print( + f"[validator] forfeited pending balance of {wallet_address}: " + f"{forfeit['amount']:.6f} USDT (fraud-divergence)", + flush=True, + ) + return receipts + + +def _route_wallets(event: Any) -> list[str]: + """Unique wallet addresses across a route, in hop order.""" + route_nodes = _event_value(event, "route_nodes") or [] + seen: set[str] = set() + wallets: list[str] = [] + for node in route_nodes: + wallet_address = node.get("wallet_address") if isinstance(node, dict) else None + if wallet_address and wallet_address not in seen: + seen.add(wallet_address) + wallets.append(wallet_address) + return wallets + + +def _final_text_node(route_nodes: list[dict]) -> dict | None: + """Text-only fallback blame: when the audit has no per-hop fingerprints + to bisect (free-running text comparison only), guess the last hop. + Never used once hop-boundary commitments make real bisection possible.""" + if not route_nodes: + return None + return max(route_nodes, key=lambda node: int(node.get("shard_end", 0))) + + +class _AuditResult: + def __init__( + self, + *, + ok: bool, + reference_output: str, + reason: str, + culprit_node: dict | None = None, + ) -> None: + self.ok = ok + self.reference_output = reference_output + self.reason = reason + self.culprit_node = culprit_node + + +class _HopCommitment: + """One hop's on-demand TOPLOC commitment plus the route node it blames.""" + + def __init__(self, node: dict | None, claim: ToplocProofClaim, token_ids: list[int]) -> None: + self.node = node + self.claim = claim + self.token_ids = token_ids + self.shard_end = int(node.get("shard_end", 0)) if node else None + + +def _hop_commitments_from_event(event: Any) -> list[_HopCommitment] | None: + """Per-hop bisection commitments (ADR-0018 §3/§4): each route node reports + its own output-boundary fingerprint. Falls back to the AH-006 whole-route + commitment format (one fingerprint, no bisection) when hops don't carry + individual commitments.""" + route_nodes = _event_value(event, "route_nodes") or [] + per_hop_nodes = [ + node for node in route_nodes + if isinstance(node, dict) and _mapping_value(node, "toploc_proof") is not None + ] + if per_hop_nodes and len(per_hop_nodes) == len(route_nodes): + ordered = sorted(route_nodes, key=lambda node: int(node.get("shard_start", 0))) + default_token_ids = _event_value(event, "claimed_token_ids") + commitments = [] + for node in ordered: + token_ids = node.get("claimed_token_ids", default_token_ids) + if not isinstance(token_ids, list) or not all(isinstance(t, int) for t in token_ids): + raise ValueError("TOPLOC hop commitments must include claimed_token_ids") + commitments.append(_HopCommitment( + node, + ToplocProofClaim.from_mapping(node["toploc_proof"]), + token_ids, + )) + return commitments + + whole_route = _toploc_audit_from_event(event) + if whole_route is None: + return None + token_ids, claim = whole_route + return [_HopCommitment(route_nodes[-1] if route_nodes else None, claim, token_ids)] + + +def _first_divergent_hop( + hop_commitments: list[_HopCommitment], + reference_activations_by_hop: list[Any], + *, + config: ToplocAuditConfig, + backend: Any | None, +) -> int | None: + """First hop whose committed output fingerprint diverges from the + referee's independently-computed reference activations at that same + cut-point (research §1.2: no interactive game needed at hop granularity — + the referee checks every cut-point in one replay).""" + for index, commitment in enumerate(hop_commitments): + ok = verify_activation_proofs( + reference_activations_by_hop[index], + commitment.claim, + config=config, + backend=backend, + ) + if not ok: + return index + return None + + +def _toploc_audit_from_event(event: Any) -> tuple[list[int], ToplocProofClaim] | None: + audit = _event_mapping(event, "audit") + claim_data = ( + _event_mapping(event, "toploc_proof") + or _event_mapping(event, "activation_proof") + or _mapping_value(audit, "toploc_proof") + or _mapping_value(audit, "activation_proof") + or _mapping_value(audit, "toploc") + ) + if claim_data is None: + return None + token_ids = ( + _event_value(event, "claimed_token_ids") + or _event_value(event, "output_token_ids") + or _mapping_value(audit, "claimed_token_ids") + or _mapping_value(audit, "output_token_ids") + ) + if not isinstance(token_ids, list) or not all(isinstance(token, int) for token in token_ids): + raise ValueError("TOPLOC audit events must include claimed_token_ids") + return token_ids, ToplocProofClaim.from_mapping(claim_data) + + +def _event_value(event: Any, name: str) -> Any: + if hasattr(event, name): + return getattr(event, name) + if isinstance(event, dict): + return event.get(name) + return None + + +def _event_mapping(event: Any, name: str) -> dict[str, Any] | None: + value = _event_value(event, name) + return value if isinstance(value, dict) else None + + +def _mapping_value(mapping: dict[str, Any] | None, name: str) -> Any: + if mapping is None: + return None + return mapping.get(name) + + +def _outputs_match(observed: str, reference: str, tolerance: float) -> bool: + observed_float = _parse_float(observed) + reference_float = _parse_float(reference) + if observed_float is not None and reference_float is not None: + return math.isclose(observed_float, reference_float, rel_tol=tolerance, abs_tol=tolerance) + return observed == reference + + +def _parse_float(value: str) -> float | None: + try: + return float(value) + except ValueError: + return None + + +def _post_json(url: str, payload: dict, timeout: float = 5.0) -> 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, timeout=timeout) as response: + return json.loads(response.read()) + + +__all__ = [ + "ToplocAuditConfig", + "ToplocProofClaim", + "ValidatorProcess", + "AdaptiveAuditSampler", + "AuditRateConfig", + "detect_output_tripwire", +] diff --git a/tests/test_accounts.py b/tests/test_accounts.py index c90d375..62388ef 100644 --- a/tests/test_accounts.py +++ b/tests/test_accounts.py @@ -1,384 +1,384 @@ -"""Dashboard user accounts: registration, login, roles, API keys, usage. - -Unit tests for AccountStore plus HTTP integration on the tracker: -register/login/logout, per-account balance and usage, API-key lifecycle -(revoked keys rejected by the OpenAI proxy), and the admin listing. -""" - -import http.cookies -import json -import urllib.error -import urllib.request - -import pytest - -from meshnet_tracker.accounts import AccountStore -from meshnet_tracker.auth import sign_hive_request -from meshnet_tracker.billing import BillingLedger -from meshnet_tracker.server import TrackerServer - -HIVE_SECRET = "test-hive-secret" - - -# ---------------------------------------------------------------- unit tests - - -def test_first_account_is_admin_then_users(): - store = AccountStore() - first = store.register(email="admin@example.com", password="secret-123") - second = store.register(email="user@example.com", password="secret-123") - assert first["role"] == "admin" - assert second["role"] == "user" - - -def test_register_requires_email_or_wallet_and_password_length(): - store = AccountStore() - with pytest.raises(ValueError, match="email or a wallet"): - store.register(password="secret-123") - with pytest.raises(ValueError, match="invalid email"): - store.register(email="not-an-email", password="secret-123") - with pytest.raises(ValueError, match="at least 8"): - store.register(email="a@b.co", password="short") - - -def test_register_rejects_duplicate_identifiers(): - store = AccountStore() - store.register(email="dup@example.com", password="secret-123") - with pytest.raises(ValueError, match="already exists"): - store.register(email="DUP@example.com", password="other-secret") - - -def test_login_by_email_or_wallet(): - store = AccountStore() - account = store.register( - email="both@example.com", wallet="WalletXYZ", password="secret-123" - ) - assert store.verify_login("both@example.com", "secret-123")["account_id"] == account["account_id"] - assert store.verify_login("WalletXYZ", "secret-123")["account_id"] == account["account_id"] - assert store.verify_login("both@example.com", "wrong-password") is None - assert store.verify_login("nobody@example.com", "secret-123") is None - - -def test_sessions_resolve_and_destroy(): - store = AccountStore() - account = store.register(email="s@example.com", password="secret-123") - token = store.create_session(account["account_id"]) - assert store.session_account(token)["account_id"] == account["account_id"] - store.destroy_session(token) - assert store.session_account(token) is None - assert store.session_account("bogus") is None - - -def test_sessions_persist_across_restart(tmp_path): - db = str(tmp_path / "accounts.db") - store = AccountStore(db_path=db) - account = store.register(email="cookie@example.com", password="secret-123") - token = store.create_session(account["account_id"]) - store.save_to_db() - - reloaded = AccountStore(db_path=db) - assert reloaded.session_account(token)["account_id"] == account["account_id"] - - -def test_api_key_lifecycle(): - store = AccountStore() - account = store.register(email="k@example.com", password="secret-123") - other = store.register(email="other@example.com", password="secret-123") - key = store.create_api_key(account["account_id"]) - assert key.startswith("sk-mesh-") - assert store.keys_for(account["account_id"]) == [key] - # someone else's account cannot revoke it - assert store.revoke_api_key(other["account_id"], key) is False - assert store.revoke_api_key(account["account_id"], key) is True - assert store.keys_for(account["account_id"]) == [] - assert store.is_key_revoked(key) - - -def test_accounts_persist_across_restart(tmp_path): - db = str(tmp_path / "accounts.db") - store = AccountStore(db_path=db) - account = store.register(email="p@example.com", password="secret-123") - key = store.create_api_key(account["account_id"]) - store.save_to_db() - - reloaded = AccountStore(db_path=db) - assert reloaded.verify_login("p@example.com", "secret-123") is not None - assert reloaded.keys_for(account["account_id"]) == [key] - - -def test_account_events_replicate_and_dedupe(): - leader = AccountStore() - follower = AccountStore() - account = leader.register(email="r@example.com", password="secret-123") - key = leader.create_api_key(account["account_id"]) - leader.revoke_api_key(account["account_id"], key) - - events, cursor = leader.events_since(0) - assert follower.apply_events(events) == len(events) - assert follower.apply_events(events) == 0 # replay is a no-op - assert follower.verify_login("r@example.com", "secret-123") is not None - assert follower.is_key_revoked(key) - more, _ = leader.events_since(cursor) - assert more == [] - - -# ---------------------------------------------------------- HTTP integration - - -def _call(url, method="GET", body=None, token=None): - headers = {"Content-Type": "application/json"} - if token: - headers["Authorization"] = f"Bearer {token}" - data = json.dumps(body).encode() if body is not None else None - req = urllib.request.Request(url, data=data, headers=headers, method=method) - with urllib.request.urlopen(req) as r: - return json.loads(r.read()) - - -@pytest.fixture -def account_tracker(): - """Tracker with credit features pinned OFF (defaults are devnet-friendly 1.0).""" - ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) - tracker = TrackerServer( - billing=ledger, - accounts=AccountStore(), - hive_secret=HIVE_SECRET, - starting_credit=0.0, - devnet_topup_amount=0.0, - ) - port = tracker.start() - yield f"http://127.0.0.1:{port}", ledger - tracker.stop() - - -def test_register_login_and_account_view(account_tracker): - url, _ = account_tracker - reg = _call(f"{url}/v1/auth/register", "POST", - {"email": "admin@example.com", "password": "secret-123"}) - assert reg["account"]["role"] == "admin" - assert reg["api_key"].startswith("sk-mesh-") - assert reg["session_token"] - - login = _call(f"{url}/v1/auth/login", "POST", - {"identifier": "admin@example.com", "password": "secret-123"}) - me = _call(f"{url}/v1/account", token=login["session_token"]) - assert me["account"]["email"] == "admin@example.com" - assert me["api_keys"] == [reg["api_key"]] - assert me["total_balance"] == pytest.approx(0.0) - assert me["usage"]["requests"] == 0 - - -def test_login_sets_cookie_and_cookie_auth_survives_tracker_restart(tmp_path): - accounts_db = str(tmp_path / "accounts.db") - tracker = TrackerServer( - billing=BillingLedger(starting_credit=0.0, default_price_per_1k=0.02), - accounts_db=accounts_db, - starting_credit=0.0, - devnet_topup_amount=0.0, - ) - port = tracker.start() - url = f"http://127.0.0.1:{port}" - try: - _call(f"{url}/v1/auth/register", "POST", - {"email": "cookie-http@example.com", "password": "secret-123"}) - req = urllib.request.Request( - f"{url}/v1/auth/login", - data=json.dumps({ - "identifier": "cookie-http@example.com", - "password": "secret-123", - }).encode(), - headers={"Content-Type": "application/json"}, - method="POST", - ) - with urllib.request.urlopen(req) as r: - assert json.loads(r.read())["session_token"] - cookie_header = r.headers["Set-Cookie"] - finally: - tracker.stop() - - cookie = http.cookies.SimpleCookie(cookie_header) - session_cookie = cookie["meshnet_session"].OutputString() - - restarted = TrackerServer( - billing=BillingLedger(starting_credit=0.0, default_price_per_1k=0.02), - accounts_db=accounts_db, - starting_credit=0.0, - devnet_topup_amount=0.0, - ) - restarted_port = restarted.start() - restarted_url = f"http://127.0.0.1:{restarted_port}" - try: - req = urllib.request.Request( - f"{restarted_url}/v1/account", - headers={"Cookie": session_cookie}, - method="GET", - ) - with urllib.request.urlopen(req) as r: - me = json.loads(r.read()) - finally: - restarted.stop() - - assert me["account"]["email"] == "cookie-http@example.com" - - -def test_bad_credentials_and_missing_session_are_401(account_tracker): - url, _ = account_tracker - _call(f"{url}/v1/auth/register", "POST", - {"email": "a@example.com", "password": "secret-123"}) - with pytest.raises(urllib.error.HTTPError) as exc_info: - _call(f"{url}/v1/auth/login", "POST", - {"identifier": "a@example.com", "password": "wrong-pass"}) - assert exc_info.value.code == 401 - with pytest.raises(urllib.error.HTTPError) as exc_info: - _call(f"{url}/v1/account") - assert exc_info.value.code == 401 - - -def test_key_create_revoke_and_revoked_key_rejected_by_proxy(account_tracker): - url, _ = account_tracker - reg = _call(f"{url}/v1/auth/register", "POST", - {"email": "k@example.com", "password": "secret-123"}) - token = reg["session_token"] - - new_key = _call(f"{url}/v1/account/keys", "POST", {}, token=token)["api_key"] - me = _call(f"{url}/v1/account", token=token) - assert sorted(me["api_keys"]) == sorted([reg["api_key"], new_key]) - - _call(f"{url}/v1/account/keys/revoke", "POST", {"api_key": new_key}, token=token) - with pytest.raises(urllib.error.HTTPError) as exc_info: - _call(f"{url}/v1/chat/completions", "POST", - {"model": "any", "messages": []}, token=new_key) - assert exc_info.value.code == 401 - assert "revoked" in exc_info.value.read().decode() - - -def test_admin_listing_requires_admin_role(account_tracker): - url, _ = account_tracker - admin = _call(f"{url}/v1/auth/register", "POST", - {"email": "admin@example.com", "password": "secret-123"}) - user = _call(f"{url}/v1/auth/register", "POST", - {"wallet": "WalletUser1", "password": "secret-123"}) - - with pytest.raises(urllib.error.HTTPError) as exc_info: - _call(f"{url}/v1/admin/accounts", token=user["session_token"]) - assert exc_info.value.code == 403 - - listing = _call(f"{url}/v1/admin/accounts", token=admin["session_token"]) - accounts = listing["accounts"] - assert len(accounts) == 2 - assert accounts[0]["role"] == "admin" - assert accounts[1]["wallet"] == "WalletUser1" - assert "balances" in accounts[0] - - -def test_accounts_gossip_endpoint_applies_events(account_tracker): - url, _ = account_tracker - peer = AccountStore() - peer.register(email="remote@example.com", password="secret-123") - events, _ = peer.events_since(0) - body = json.dumps({"events": events}).encode() - req = urllib.request.Request( - f"{url}/v1/accounts/gossip", data=body, - headers={"Content-Type": "application/json", **sign_hive_request(HIVE_SECRET, body)}, - method="POST", - ) - with urllib.request.urlopen(req) as r: - result = json.loads(r.read()) - assert result["applied"] == len(events) - login = _call(f"{url}/v1/auth/login", "POST", - {"identifier": "remote@example.com", "password": "secret-123"}) - assert login["account"]["email"] == "remote@example.com" - - -def test_accounts_endpoints_404_when_disabled(): - tracker = TrackerServer() # no accounts, no billing - port = tracker.start() - try: - with pytest.raises(urllib.error.HTTPError) as exc_info: - _call(f"http://127.0.0.1:{port}/v1/auth/register", "POST", - {"email": "x@example.com", "password": "secret-123"}) - assert exc_info.value.code == 404 - finally: - tracker.stop() - - -# ------------------------------------------- US-039/US-040: credit and top-up - - -@pytest.fixture -def funded_tracker(): - """Tracker with Caller Credit and the devnet top-up faucet enabled.""" - ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) - tracker = TrackerServer( - billing=ledger, - accounts=AccountStore(), - hive_secret=HIVE_SECRET, - starting_credit=1.0, - devnet_topup_amount=10.0, - ) - port = tracker.start() - yield f"http://127.0.0.1:{port}", ledger - tracker.stop() - - -def test_caller_credit_granted_once_per_account(funded_tracker): - url, ledger = funded_tracker - reg = _call(f"{url}/v1/auth/register", "POST", - {"email": "c@example.com", "password": "secret-123"}) - token = reg["session_token"] - first_key = reg["api_key"] - assert ledger.get_client_balance(first_key) == pytest.approx(1.0) - - # A second key never re-grants — not even after revoking the first. - second = _call(f"{url}/v1/account/keys", "POST", {}, token=token) - assert second["caller_credit_granted"] is False - assert ledger.get_client_balance(second["api_key"]) == pytest.approx(0.0) - _call(f"{url}/v1/account/keys/revoke", "POST", {"api_key": first_key}, token=token) - third = _call(f"{url}/v1/account/keys", "POST", {}, token=token) - assert third["caller_credit_granted"] is False - assert ledger.get_client_balance(third["api_key"]) == pytest.approx(0.0) - - -def test_unknown_bearer_key_rejected_by_proxy(funded_tracker): - url, ledger = funded_tracker - with pytest.raises(urllib.error.HTTPError) as exc_info: - _call(f"{url}/v1/chat/completions", "POST", - {"model": "any", "messages": []}, token="sk-mesh-made-up-key") - assert exc_info.value.code == 401 - assert "unknown API key" in exc_info.value.read().decode() - # The invented key must not have become a billable client. - assert ledger.get_client_balance("sk-mesh-made-up-key") == pytest.approx(0.0) - - -def test_devnet_topup_credits_own_key_only(funded_tracker): - url, ledger = funded_tracker - owner = _call(f"{url}/v1/auth/register", "POST", - {"email": "own@example.com", "password": "secret-123"}) - other = _call(f"{url}/v1/auth/register", "POST", - {"email": "oth@example.com", "password": "secret-123"}) - - me = _call(f"{url}/v1/account", token=owner["session_token"]) - assert me["topup_amount"] == pytest.approx(10.0) - - result = _call(f"{url}/v1/account/topup", "POST", - {"api_key": owner["api_key"]}, token=owner["session_token"]) - assert result["credited"] == pytest.approx(10.0) - assert result["balance"] == pytest.approx(11.0) # 1.0 caller credit + 10.0 - - with pytest.raises(urllib.error.HTTPError) as exc_info: - _call(f"{url}/v1/account/topup", "POST", - {"api_key": owner["api_key"]}, token=other["session_token"]) - assert exc_info.value.code == 403 - assert ledger.get_client_balance(owner["api_key"]) == pytest.approx(11.0) - - -def test_topup_404_when_disabled(account_tracker): - url, _ = account_tracker - reg = _call(f"{url}/v1/auth/register", "POST", - {"email": "t@example.com", "password": "secret-123"}) - me = _call(f"{url}/v1/account", token=reg["session_token"]) - assert me["topup_amount"] == pytest.approx(0.0) - with pytest.raises(urllib.error.HTTPError) as exc_info: - _call(f"{url}/v1/account/topup", "POST", - {"api_key": reg["api_key"]}, token=reg["session_token"]) - assert exc_info.value.code == 404 +"""Dashboard user accounts: registration, login, roles, API keys, usage. + +Unit tests for AccountStore plus HTTP integration on the tracker: +register/login/logout, per-account balance and usage, API-key lifecycle +(revoked keys rejected by the OpenAI proxy), and the admin listing. +""" + +import http.cookies +import json +import urllib.error +import urllib.request + +import pytest + +from meshnet_tracker.accounts import AccountStore +from meshnet_tracker.auth import sign_hive_request +from meshnet_tracker.billing import BillingLedger +from meshnet_tracker.server import TrackerServer + +HIVE_SECRET = "test-hive-secret" + + +# ---------------------------------------------------------------- unit tests + + +def test_first_account_is_admin_then_users(): + store = AccountStore() + first = store.register(email="admin@example.com", password="secret-123") + second = store.register(email="user@example.com", password="secret-123") + assert first["role"] == "admin" + assert second["role"] == "user" + + +def test_register_requires_email_or_wallet_and_password_length(): + store = AccountStore() + with pytest.raises(ValueError, match="email or a wallet"): + store.register(password="secret-123") + with pytest.raises(ValueError, match="invalid email"): + store.register(email="not-an-email", password="secret-123") + with pytest.raises(ValueError, match="at least 8"): + store.register(email="a@b.co", password="short") + + +def test_register_rejects_duplicate_identifiers(): + store = AccountStore() + store.register(email="dup@example.com", password="secret-123") + with pytest.raises(ValueError, match="already exists"): + store.register(email="DUP@example.com", password="other-secret") + + +def test_login_by_email_or_wallet(): + store = AccountStore() + account = store.register( + email="both@example.com", wallet="WalletXYZ", password="secret-123" + ) + assert store.verify_login("both@example.com", "secret-123")["account_id"] == account["account_id"] + assert store.verify_login("WalletXYZ", "secret-123")["account_id"] == account["account_id"] + assert store.verify_login("both@example.com", "wrong-password") is None + assert store.verify_login("nobody@example.com", "secret-123") is None + + +def test_sessions_resolve_and_destroy(): + store = AccountStore() + account = store.register(email="s@example.com", password="secret-123") + token = store.create_session(account["account_id"]) + assert store.session_account(token)["account_id"] == account["account_id"] + store.destroy_session(token) + assert store.session_account(token) is None + assert store.session_account("bogus") is None + + +def test_sessions_persist_across_restart(tmp_path): + db = str(tmp_path / "accounts.db") + store = AccountStore(db_path=db) + account = store.register(email="cookie@example.com", password="secret-123") + token = store.create_session(account["account_id"]) + store.save_to_db() + + reloaded = AccountStore(db_path=db) + assert reloaded.session_account(token)["account_id"] == account["account_id"] + + +def test_api_key_lifecycle(): + store = AccountStore() + account = store.register(email="k@example.com", password="secret-123") + other = store.register(email="other@example.com", password="secret-123") + key = store.create_api_key(account["account_id"]) + assert key.startswith("sk-mesh-") + assert store.keys_for(account["account_id"]) == [key] + # someone else's account cannot revoke it + assert store.revoke_api_key(other["account_id"], key) is False + assert store.revoke_api_key(account["account_id"], key) is True + assert store.keys_for(account["account_id"]) == [] + assert store.is_key_revoked(key) + + +def test_accounts_persist_across_restart(tmp_path): + db = str(tmp_path / "accounts.db") + store = AccountStore(db_path=db) + account = store.register(email="p@example.com", password="secret-123") + key = store.create_api_key(account["account_id"]) + store.save_to_db() + + reloaded = AccountStore(db_path=db) + assert reloaded.verify_login("p@example.com", "secret-123") is not None + assert reloaded.keys_for(account["account_id"]) == [key] + + +def test_account_events_replicate_and_dedupe(): + leader = AccountStore() + follower = AccountStore() + account = leader.register(email="r@example.com", password="secret-123") + key = leader.create_api_key(account["account_id"]) + leader.revoke_api_key(account["account_id"], key) + + events, cursor = leader.events_since(0) + assert follower.apply_events(events) == len(events) + assert follower.apply_events(events) == 0 # replay is a no-op + assert follower.verify_login("r@example.com", "secret-123") is not None + assert follower.is_key_revoked(key) + more, _ = leader.events_since(cursor) + assert more == [] + + +# ---------------------------------------------------------- HTTP integration + + +def _call(url, method="GET", body=None, token=None): + headers = {"Content-Type": "application/json"} + if token: + headers["Authorization"] = f"Bearer {token}" + data = json.dumps(body).encode() if body is not None else None + req = urllib.request.Request(url, data=data, headers=headers, method=method) + with urllib.request.urlopen(req) as r: + return json.loads(r.read()) + + +@pytest.fixture +def account_tracker(): + """Tracker with credit features pinned OFF (defaults are devnet-friendly 1.0).""" + ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) + tracker = TrackerServer( + billing=ledger, + accounts=AccountStore(), + hive_secret=HIVE_SECRET, + starting_credit=0.0, + devnet_topup_amount=0.0, + ) + port = tracker.start() + yield f"http://127.0.0.1:{port}", ledger + tracker.stop() + + +def test_register_login_and_account_view(account_tracker): + url, _ = account_tracker + reg = _call(f"{url}/v1/auth/register", "POST", + {"email": "admin@example.com", "password": "secret-123"}) + assert reg["account"]["role"] == "admin" + assert reg["api_key"].startswith("sk-mesh-") + assert reg["session_token"] + + login = _call(f"{url}/v1/auth/login", "POST", + {"identifier": "admin@example.com", "password": "secret-123"}) + me = _call(f"{url}/v1/account", token=login["session_token"]) + assert me["account"]["email"] == "admin@example.com" + assert me["api_keys"] == [reg["api_key"]] + assert me["total_balance"] == pytest.approx(0.0) + assert me["usage"]["requests"] == 0 + + +def test_login_sets_cookie_and_cookie_auth_survives_tracker_restart(tmp_path): + accounts_db = str(tmp_path / "accounts.db") + tracker = TrackerServer( + billing=BillingLedger(starting_credit=0.0, default_price_per_1k=0.02), + accounts_db=accounts_db, + starting_credit=0.0, + devnet_topup_amount=0.0, + ) + port = tracker.start() + url = f"http://127.0.0.1:{port}" + try: + _call(f"{url}/v1/auth/register", "POST", + {"email": "cookie-http@example.com", "password": "secret-123"}) + req = urllib.request.Request( + f"{url}/v1/auth/login", + data=json.dumps({ + "identifier": "cookie-http@example.com", + "password": "secret-123", + }).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req) as r: + assert json.loads(r.read())["session_token"] + cookie_header = r.headers["Set-Cookie"] + finally: + tracker.stop() + + cookie = http.cookies.SimpleCookie(cookie_header) + session_cookie = cookie["meshnet_session"].OutputString() + + restarted = TrackerServer( + billing=BillingLedger(starting_credit=0.0, default_price_per_1k=0.02), + accounts_db=accounts_db, + starting_credit=0.0, + devnet_topup_amount=0.0, + ) + restarted_port = restarted.start() + restarted_url = f"http://127.0.0.1:{restarted_port}" + try: + req = urllib.request.Request( + f"{restarted_url}/v1/account", + headers={"Cookie": session_cookie}, + method="GET", + ) + with urllib.request.urlopen(req) as r: + me = json.loads(r.read()) + finally: + restarted.stop() + + assert me["account"]["email"] == "cookie-http@example.com" + + +def test_bad_credentials_and_missing_session_are_401(account_tracker): + url, _ = account_tracker + _call(f"{url}/v1/auth/register", "POST", + {"email": "a@example.com", "password": "secret-123"}) + with pytest.raises(urllib.error.HTTPError) as exc_info: + _call(f"{url}/v1/auth/login", "POST", + {"identifier": "a@example.com", "password": "wrong-pass"}) + assert exc_info.value.code == 401 + with pytest.raises(urllib.error.HTTPError) as exc_info: + _call(f"{url}/v1/account") + assert exc_info.value.code == 401 + + +def test_key_create_revoke_and_revoked_key_rejected_by_proxy(account_tracker): + url, _ = account_tracker + reg = _call(f"{url}/v1/auth/register", "POST", + {"email": "k@example.com", "password": "secret-123"}) + token = reg["session_token"] + + new_key = _call(f"{url}/v1/account/keys", "POST", {}, token=token)["api_key"] + me = _call(f"{url}/v1/account", token=token) + assert sorted(me["api_keys"]) == sorted([reg["api_key"], new_key]) + + _call(f"{url}/v1/account/keys/revoke", "POST", {"api_key": new_key}, token=token) + with pytest.raises(urllib.error.HTTPError) as exc_info: + _call(f"{url}/v1/chat/completions", "POST", + {"model": "any", "messages": []}, token=new_key) + assert exc_info.value.code == 401 + assert "revoked" in exc_info.value.read().decode() + + +def test_admin_listing_requires_admin_role(account_tracker): + url, _ = account_tracker + admin = _call(f"{url}/v1/auth/register", "POST", + {"email": "admin@example.com", "password": "secret-123"}) + user = _call(f"{url}/v1/auth/register", "POST", + {"wallet": "WalletUser1", "password": "secret-123"}) + + with pytest.raises(urllib.error.HTTPError) as exc_info: + _call(f"{url}/v1/admin/accounts", token=user["session_token"]) + assert exc_info.value.code == 403 + + listing = _call(f"{url}/v1/admin/accounts", token=admin["session_token"]) + accounts = listing["accounts"] + assert len(accounts) == 2 + assert accounts[0]["role"] == "admin" + assert accounts[1]["wallet"] == "WalletUser1" + assert "balances" in accounts[0] + + +def test_accounts_gossip_endpoint_applies_events(account_tracker): + url, _ = account_tracker + peer = AccountStore() + peer.register(email="remote@example.com", password="secret-123") + events, _ = peer.events_since(0) + body = json.dumps({"events": events}).encode() + req = urllib.request.Request( + f"{url}/v1/accounts/gossip", data=body, + headers={"Content-Type": "application/json", **sign_hive_request(HIVE_SECRET, body)}, + method="POST", + ) + with urllib.request.urlopen(req) as r: + result = json.loads(r.read()) + assert result["applied"] == len(events) + login = _call(f"{url}/v1/auth/login", "POST", + {"identifier": "remote@example.com", "password": "secret-123"}) + assert login["account"]["email"] == "remote@example.com" + + +def test_accounts_endpoints_404_when_disabled(): + tracker = TrackerServer() # no accounts, no billing + port = tracker.start() + try: + with pytest.raises(urllib.error.HTTPError) as exc_info: + _call(f"http://127.0.0.1:{port}/v1/auth/register", "POST", + {"email": "x@example.com", "password": "secret-123"}) + assert exc_info.value.code == 404 + finally: + tracker.stop() + + +# ------------------------------------------- US-039/US-040: credit and top-up + + +@pytest.fixture +def funded_tracker(): + """Tracker with Caller Credit and the devnet top-up faucet enabled.""" + ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) + tracker = TrackerServer( + billing=ledger, + accounts=AccountStore(), + hive_secret=HIVE_SECRET, + starting_credit=1.0, + devnet_topup_amount=10.0, + ) + port = tracker.start() + yield f"http://127.0.0.1:{port}", ledger + tracker.stop() + + +def test_caller_credit_granted_once_per_account(funded_tracker): + url, ledger = funded_tracker + reg = _call(f"{url}/v1/auth/register", "POST", + {"email": "c@example.com", "password": "secret-123"}) + token = reg["session_token"] + first_key = reg["api_key"] + assert ledger.get_client_balance(first_key) == pytest.approx(1.0) + + # A second key never re-grants — not even after revoking the first. + second = _call(f"{url}/v1/account/keys", "POST", {}, token=token) + assert second["caller_credit_granted"] is False + assert ledger.get_client_balance(second["api_key"]) == pytest.approx(0.0) + _call(f"{url}/v1/account/keys/revoke", "POST", {"api_key": first_key}, token=token) + third = _call(f"{url}/v1/account/keys", "POST", {}, token=token) + assert third["caller_credit_granted"] is False + assert ledger.get_client_balance(third["api_key"]) == pytest.approx(0.0) + + +def test_unknown_bearer_key_rejected_by_proxy(funded_tracker): + url, ledger = funded_tracker + with pytest.raises(urllib.error.HTTPError) as exc_info: + _call(f"{url}/v1/chat/completions", "POST", + {"model": "any", "messages": []}, token="sk-mesh-made-up-key") + assert exc_info.value.code == 401 + assert "unknown API key" in exc_info.value.read().decode() + # The invented key must not have become a billable client. + assert ledger.get_client_balance("sk-mesh-made-up-key") == pytest.approx(0.0) + + +def test_devnet_topup_credits_own_key_only(funded_tracker): + url, ledger = funded_tracker + owner = _call(f"{url}/v1/auth/register", "POST", + {"email": "own@example.com", "password": "secret-123"}) + other = _call(f"{url}/v1/auth/register", "POST", + {"email": "oth@example.com", "password": "secret-123"}) + + me = _call(f"{url}/v1/account", token=owner["session_token"]) + assert me["topup_amount"] == pytest.approx(10.0) + + result = _call(f"{url}/v1/account/topup", "POST", + {"api_key": owner["api_key"]}, token=owner["session_token"]) + assert result["credited"] == pytest.approx(10.0) + assert result["balance"] == pytest.approx(11.0) # 1.0 caller credit + 10.0 + + with pytest.raises(urllib.error.HTTPError) as exc_info: + _call(f"{url}/v1/account/topup", "POST", + {"api_key": owner["api_key"]}, token=other["session_token"]) + assert exc_info.value.code == 403 + assert ledger.get_client_balance(owner["api_key"]) == pytest.approx(11.0) + + +def test_topup_404_when_disabled(account_tracker): + url, _ = account_tracker + reg = _call(f"{url}/v1/auth/register", "POST", + {"email": "t@example.com", "password": "secret-123"}) + me = _call(f"{url}/v1/account", token=reg["session_token"]) + assert me["topup_amount"] == pytest.approx(0.0) + with pytest.raises(urllib.error.HTTPError) as exc_info: + _call(f"{url}/v1/account/topup", "POST", + {"api_key": reg["api_key"]}, token=reg["session_token"]) + assert exc_info.value.code == 404 diff --git a/tests/test_billing_ledger.py b/tests/test_billing_ledger.py index 362a8d8..40e0dff 100644 --- a/tests/test_billing_ledger.py +++ b/tests/test_billing_ledger.py @@ -1,674 +1,674 @@ -"""US-031: billing ledger — per-token pricing, 90/10 split, pending balances. - -Unit tests for BillingLedger math/persistence/replication, plus HTTP -integration: 401 without an API key, 402 for unfunded keys, billed 200 after -explicit credit, and 402 once the balance is exhausted. -""" - -import http.server -import json -import socketserver -import threading -import urllib.error -import urllib.request - -import pytest - -from meshnet_tracker.auth import sign_hive_request -from meshnet_tracker.billing import DEFAULT_STARTING_CREDIT, BillingLedger -from meshnet_tracker.server import ( - TrackerServer, - _billable_non_stream_tokens, - _billable_stream_tokens, - _estimate_prompt_tokens, - _observed_non_stream_completion_tokens, - _observed_stream_tokens, -) - -MODEL = "openai-community/gpt2" -HIVE_SECRET = "test-hive-secret" - - -# ---------------------------------------------------------------- unit tests - - -def test_charge_single_node_gets_90_percent(): - ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) - ledger.ensure_client("key-a") - event = ledger.charge_request("key-a", MODEL, total_tokens=1000, node_work=[("wallet-1", 12)]) - assert event["cost"] == pytest.approx(0.02) - assert ledger.get_client_balance("key-a") == pytest.approx(1.0 - 0.02) - assert ledger.get_node_pending("wallet-1") == pytest.approx(0.02 * 0.90) - assert ledger.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10) - - -def test_default_starting_credit_is_zero_and_fresh_key_is_unfunded(): - ledger = BillingLedger(default_price_per_1k=0.02) - - assert DEFAULT_STARTING_CREDIT == 0.0 - assert ledger.ensure_client("fresh-key") == pytest.approx(0.0) - assert ledger.get_client_balance("fresh-key") == pytest.approx(0.0) - assert ledger.has_funds("fresh-key") is False - assert ledger.events_since(0)[0] == [] - - -def test_charge_three_node_split_by_work_units(): - ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) - ledger.ensure_client("key-a") - ledger.charge_request( - "key-a", MODEL, total_tokens=1000, - node_work=[("wallet-1", 6), ("wallet-2", 3), ("wallet-3", 3)], - ) - pool = 0.02 * 0.90 - assert ledger.get_node_pending("wallet-1") == pytest.approx(pool * 6 / 12) - assert ledger.get_node_pending("wallet-2") == pytest.approx(pool * 3 / 12) - assert ledger.get_node_pending("wallet-3") == pytest.approx(pool * 3 / 12) - assert ledger.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10) - - -def test_walletless_node_share_accrues_to_protocol_cut(): - ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) - ledger.ensure_client("key-a") - ledger.charge_request( - "key-a", MODEL, total_tokens=1000, - node_work=[("wallet-1", 6), (None, 6)], - ) - pool = 0.02 * 0.90 - assert ledger.get_node_pending("wallet-1") == pytest.approx(pool / 2) - # walletless half of the pool + the 10% cut both land in protocol_cut - assert ledger.snapshot()["protocol_cut"] == pytest.approx(pool / 2 + 0.02 * 0.10) - - -def test_per_model_price_override(): - ledger = BillingLedger( - starting_credit=1.0, - default_price_per_1k=0.02, - prices={MODEL: 0.10}, - ) - ledger.ensure_client("key-a") - event = ledger.charge_request("key-a", MODEL, 500, [("wallet-1", 1)]) - assert event["cost"] == pytest.approx(0.10 * 500 / 1000) - event = ledger.charge_request("key-a", "other-model", 500, [("wallet-1", 1)]) - assert event["cost"] == pytest.approx(0.02 * 500 / 1000) - - -def test_non_stream_billable_tokens_cap_usage_by_request_bound(): - payload = {"usage": {"total_tokens": 1_000_000}} - request = {"messages": [{"role": "user", "content": "hello tracker"}], "max_tokens": 3} - - assert _estimate_prompt_tokens(request) == 2 - assert _billable_non_stream_tokens(payload, request) == 5 - - -def test_non_stream_billable_tokens_fallback_when_usage_missing(): - payload = {"choices": [{"message": {"role": "assistant", "content": "ok now"}}]} - request = {"messages": [{"role": "user", "content": "hello tracker"}], "max_tokens": 10} - - assert _observed_non_stream_completion_tokens(payload) == 2 - assert _billable_non_stream_tokens(payload, request) == 4 - - -def test_stream_billable_tokens_allow_usage_to_lower_observed_count_only(): - observed_payload = { - "choices": [{ - "index": 0, - "delta": {"content": "hello tracker"}, - "finish_reason": None, - }] - } - - assert _observed_stream_tokens(observed_payload) == 2 - assert _billable_stream_tokens(observed_tokens=2, reported_tokens=1_000_000) == 2 - assert _billable_stream_tokens(observed_tokens=2, reported_tokens=1) == 1 - - -def test_payout_and_forfeit_hooks(): - ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) - ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 1)]) - pending = ledger.get_node_pending("wallet-1") - ledger.settle_node_payout("wallet-1", pending, reference="tx-sig-1") - assert ledger.get_node_pending("wallet-1") == pytest.approx(0.0) - - ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 1)]) - cut_before = ledger.snapshot()["protocol_cut"] - forfeited = ledger.forfeit_pending("wallet-1")["amount"] - assert forfeited == pytest.approx(0.02 * 0.90) - assert ledger.get_node_pending("wallet-1") == pytest.approx(0.0) - assert ledger.snapshot()["protocol_cut"] == pytest.approx(cut_before + forfeited) - - -def test_restart_persistence(tmp_path): - db = str(tmp_path / "billing.db") - ledger = BillingLedger(db_path=db, starting_credit=1.0, default_price_per_1k=0.02) - ledger.credit_client("key-a", 5.0) - ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 12)]) - ledger.save_to_db() - - reloaded = BillingLedger(db_path=db, starting_credit=1.0, default_price_per_1k=0.02) - assert reloaded.get_client_balance("key-a") == pytest.approx( - ledger.get_client_balance("key-a") - ) - assert reloaded.get_node_pending("wallet-1") == pytest.approx(0.02 * 0.90) - assert reloaded.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10) - - -def test_tracker_enables_billing_with_default_db_when_requested(tmp_path, monkeypatch): - from meshnet_tracker.billing import DEFAULT_BILLING_DB_PATH - - monkeypatch.chdir(tmp_path) - tracker = TrackerServer(enable_billing=True) - port = tracker.start() - try: - # /v1/billing/summary is admin-gated now; just confirm the server is up. - health = json.loads( - urllib.request.urlopen(f"http://127.0.0.1:{port}/v1/health").read() - ) - assert health["status"] == "ok" - finally: - tracker.stop() - # enabling billing creates the ledger DB in the working directory - assert (tmp_path / DEFAULT_BILLING_DB_PATH).exists() - - -def test_event_replication_converges_and_dedupes(): - leader = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) - follower = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) - - leader.credit_client("key-a", 10.0) - leader.charge_request("key-a", MODEL, 2000, [("wallet-1", 8), ("wallet-2", 4)]) - - events, cursor = leader.events_since(0) - assert follower.apply_events(events) == len(events) - # replaying the same batch must be a no-op - assert follower.apply_events(events) == 0 - - assert follower.get_client_balance("key-a") == pytest.approx( - leader.get_client_balance("key-a") - ) - assert follower.get_node_pending("wallet-1") == pytest.approx( - leader.get_node_pending("wallet-1") - ) - assert follower.snapshot()["protocol_cut"] == pytest.approx( - leader.snapshot()["protocol_cut"] - ) - # incremental cursor: nothing new after full sync - more, _ = leader.events_since(cursor) - assert more == [] - - -# ---------------------------------------------------------- HTTP integration - - -class _UsageStubNode: - """Minimal head node returning a chat completion with real usage numbers.""" - - def __init__( - self, - total_tokens: int = 1000, - *, - stream_chunks: list[str] | None = None, - stream_usage_total: int | None = None, - ): - self.total_tokens = total_tokens - self.stream_chunks = stream_chunks or [] - self.stream_usage_total = stream_usage_total - self.request_count = 0 - outer = self - - class Handler(http.server.BaseHTTPRequestHandler): - def log_message(self, fmt, *args): - pass - - def do_POST(self): - outer.request_count += 1 - length = int(self.headers.get("Content-Length", 0)) - raw_body = self.rfile.read(length) - try: - request_body = json.loads(raw_body) - except json.JSONDecodeError: - request_body = {} - if request_body.get("stream"): - self.send_response(200) - self.send_header("Content-Type", "text/event-stream; charset=utf-8") - self.end_headers() - for chunk in outer.stream_chunks: - payload = json.dumps({ - "id": "chatcmpl-stub", - "object": "chat.completion.chunk", - "model": MODEL, - "choices": [{ - "index": 0, - "delta": {"content": chunk}, - "finish_reason": None, - }], - }).encode() - self.wfile.write(b"data: " + payload + b"\n\n") - if outer.stream_usage_total is not None: - usage_payload = json.dumps({ - "id": "chatcmpl-stub", - "object": "chat.completion.chunk", - "model": MODEL, - "choices": [], - "usage": { - "prompt_tokens": 0, - "completion_tokens": outer.stream_usage_total, - "total_tokens": outer.stream_usage_total, - }, - }).encode() - self.wfile.write(b"data: " + usage_payload + b"\n\n") - self.wfile.write(b"data: [DONE]\n\n") - return - body = json.dumps({ - "id": "chatcmpl-stub", - "object": "chat.completion", - "model": MODEL, - "choices": [{ - "index": 0, - "message": {"role": "assistant", "content": "ok"}, - "finish_reason": "stop", - }], - "usage": { - "prompt_tokens": 0, - "completion_tokens": outer.total_tokens, - "total_tokens": outer.total_tokens, - }, - }).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(body))) - self.end_headers() - self.wfile.write(body) - - self._server = socketserver.ThreadingTCPServer(("127.0.0.1", 0), Handler) - self._server.daemon_threads = True - self._thread: threading.Thread | None = None - - def start(self) -> int: - self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) - self._thread.start() - return self._server.server_address[1] - - def stop(self): - self._server.shutdown() - self._server.server_close() - - -@pytest.fixture -def billed_tracker(): - ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) - tracker = TrackerServer( - model_presets={ - MODEL: { - "layers_start": 0, - "layers_end": 11, - "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, - } - }, - billing=ledger, - hive_secret=HIVE_SECRET, - ) - port = tracker.start() - tracker_url = f"http://127.0.0.1:{port}" - - stub = _UsageStubNode(total_tokens=1000) - stub_port = stub.start() - reg = json.dumps({ - "endpoint": f"http://127.0.0.1:{stub_port}", - "shard_start": 0, - "shard_end": 11, - "model": MODEL, - "hardware_profile": {}, - "score": 1.0, - "tracker_mode": True, - "wallet_address": "wallet-head", - }).encode() - req = urllib.request.Request( - f"{tracker_url}/v1/nodes/register", - data=reg, - headers={"Content-Type": "application/json"}, - method="POST", - ) - with urllib.request.urlopen(req) as r: - r.read() - - yield tracker_url, ledger, stub - - stub.stop() - tracker.stop() - - -def _chat(tracker_url: str, api_key: str | None, **body_overrides): - body = { - "model": MODEL, - "messages": [{"role": "user", "content": "hi"}], - } - body.update(body_overrides) - data = json.dumps(body).encode() - headers = {"Content-Type": "application/json"} - if api_key: - headers["Authorization"] = f"Bearer {api_key}" - req = urllib.request.Request( - f"{tracker_url}/v1/chat/completions", - data=data, - headers=headers, - method="POST", - ) - with urllib.request.urlopen(req) as r: - if body.get("stream"): - return r.read().decode() - return json.loads(r.read()) - - -def test_proxy_chat_requires_api_key_when_billing_enabled(billed_tracker): - tracker_url, _, _ = billed_tracker - with pytest.raises(urllib.error.HTTPError) as exc_info: - _chat(tracker_url, api_key=None) - assert exc_info.value.code == 401 - - -def test_proxy_chat_402_for_fresh_key_before_routing(billed_tracker): - tracker_url, ledger, stub = billed_tracker - with pytest.raises(urllib.error.HTTPError) as exc_info: - _chat(tracker_url, api_key="fresh-client") - assert exc_info.value.code == 402 - assert ledger.get_client_balance("fresh-client") == pytest.approx(0.0) - assert stub.request_count == 0 - - -def test_proxy_chat_bills_credited_client_and_credits_node(billed_tracker): - tracker_url, ledger, _ = billed_tracker - ledger.credit_client("client-1", 0.03, note="admin-credit") - _chat(tracker_url, api_key="client-1") - # 1000 tokens at 0.02/1K = 0.02 USDT; explicit credit 0.03 - assert ledger.get_client_balance("client-1") == pytest.approx(0.03 - 0.02) - assert ledger.get_node_pending("wallet-head") == pytest.approx(0.02 * 0.90) - - # /v1/billing/summary is admin-gated now (ADR-0017); auth is covered in - # test_auth_boundary.py, so verify the numbers via the ledger snapshot. - summary = ledger.snapshot() - assert summary["node_pending"]["wallet-head"] == pytest.approx(0.02 * 0.90) - assert summary["protocol_cut"] == pytest.approx(0.02 * 0.10) - - stats = json.loads(urllib.request.urlopen(f"{tracker_url}/v1/stats").read()) - node_stats = next(iter(stats["nodes"].values()))["models"][MODEL] - assert node_stats["tokens_per_sec_last_hour"] is not None - assert node_stats["sample_count_last_hour"] == 1 - - -def test_proxy_chat_caps_inflated_non_streaming_usage_by_request_bounds(billed_tracker): - tracker_url, ledger, stub = billed_tracker - stub.total_tokens = 1_000_000 - ledger.credit_client("bounded-client", 100.0, note="admin-credit") - - _chat(tracker_url, api_key="bounded-client", max_tokens=2) - - # messages=[{"content": "hi"}] has a local prompt estimate of one token, so - # billable total is capped at max_tokens + prompt estimate, not node usage. - expected_tokens = 3 - assert ledger.get_client_balance("bounded-client") == pytest.approx(100.0 - 0.02 * expected_tokens / 1000) - events, _ = ledger.events_since(0) - charge = next(event for event in events if event["type"] == "charge") - assert charge["total_tokens"] == expected_tokens - - -def test_proxy_chat_caps_inflated_streaming_usage_by_observed_chunks(): - ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) - ledger.credit_client("stream-client", 1.0, note="admin-credit") - tracker = TrackerServer( - model_presets={ - MODEL: { - "layers_start": 0, - "layers_end": 11, - "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, - } - }, - billing=ledger, - ) - port = tracker.start() - tracker_url = f"http://127.0.0.1:{port}" - stub = _UsageStubNode(stream_chunks=["one", "two"], stream_usage_total=1_000_000) - stub_port = stub.start() - try: - reg = json.dumps({ - "endpoint": f"http://127.0.0.1:{stub_port}", - "shard_start": 0, - "shard_end": 11, - "model": MODEL, - "hardware_profile": {}, - "score": 1.0, - "tracker_mode": True, - "wallet_address": "wallet-head", - }).encode() - req = urllib.request.Request( - f"{tracker_url}/v1/nodes/register", - data=reg, - headers={"Content-Type": "application/json"}, - method="POST", - ) - with urllib.request.urlopen(req) as r: - r.read() - - _chat(tracker_url, api_key="stream-client", stream=True, max_tokens=2) - - events, _ = ledger.events_since(0) - charge = next(event for event in events if event["type"] == "charge") - assert charge["total_tokens"] == 2 - assert ledger.get_client_balance("stream-client") == pytest.approx(1.0 - 0.02 * 2 / 1000) - finally: - stub.stop() - tracker.stop() - - -def test_proxy_chat_splits_payout_by_tracker_assigned_route_span(): - ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) - ledger.credit_client("route-client", 1.0, note="admin-credit") - tracker = TrackerServer( - model_presets={ - MODEL: { - "total_layers": 12, - "bytes_per_layer": {"bfloat16": 1_000}, - } - }, - billing=ledger, - ) - port = tracker.start() - tracker_url = f"http://127.0.0.1:{port}" - head = _UsageStubNode(total_tokens=1000) - head_port = head.start() - try: - for endpoint, wallet, vram, bench in ( - (f"http://127.0.0.1:{head_port}", "wallet-head", 5_000, 2.0), - ("http://127.0.0.1:1", "wallet-tail", 10_000, 1.0), - ): - reg = json.dumps({ - "endpoint": endpoint, - "model": MODEL, - "vram_bytes": vram, - "ram_bytes": 10_000, - "quantizations": ["bfloat16"], - "benchmark_tokens_per_sec": bench, - "hardware_profile": {}, - "score": 1.0, - "tracker_mode": endpoint.endswith(str(head_port)), - "wallet_address": wallet, - "managed_assignment": True, - "shard_start": 0, - "shard_end": 999, - }).encode() - req = urllib.request.Request( - f"{tracker_url}/v1/nodes/register", - data=reg, - headers={"Content-Type": "application/json"}, - method="POST", - ) - with urllib.request.urlopen(req) as r: - r.read() - - _chat(tracker_url, api_key="route-client", max_tokens=1000) - - pool = 0.02 * 0.90 - assert ledger.get_node_pending("wallet-head") == pytest.approx(pool * 4 / 12) - assert ledger.get_node_pending("wallet-tail") == pytest.approx(pool * 8 / 12) - finally: - head.stop() - tracker.stop() - - -def test_proxy_chat_402_when_balance_exhausted(billed_tracker): - tracker_url, ledger, _ = billed_tracker - ledger.credit_client("client-2", 0.03, note="admin-credit") - _chat(tracker_url, api_key="client-2") # 0.03 -> 0.01 - _chat(tracker_url, api_key="client-2") # 0.01 -> -0.01 (post-pay drift) - assert ledger.get_client_balance("client-2") == pytest.approx(-0.01) - with pytest.raises(urllib.error.HTTPError) as exc_info: - _chat(tracker_url, api_key="client-2") - assert exc_info.value.code == 402 - # rejected before routing: nothing further was billed - assert ledger.get_client_balance("client-2") == pytest.approx(-0.01) - - -def test_proxy_chat_rejects_request_above_spend_cap_before_routing(): - ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) - ledger.credit_client("capped-client", 10.0, note="admin-credit") - tracker = TrackerServer( - model_presets={ - MODEL: { - "layers_start": 0, - "layers_end": 11, - "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, - } - }, - billing=ledger, - max_charge_per_request=0.01, - ) - port = tracker.start() - tracker_url = f"http://127.0.0.1:{port}" - stub = _UsageStubNode(total_tokens=1000) - stub_port = stub.start() - try: - reg = json.dumps({ - "endpoint": f"http://127.0.0.1:{stub_port}", - "shard_start": 0, - "shard_end": 11, - "model": MODEL, - "hardware_profile": {}, - "score": 1.0, - "tracker_mode": True, - "wallet_address": "wallet-head", - }).encode() - req = urllib.request.Request( - f"{tracker_url}/v1/nodes/register", - data=reg, - headers={"Content-Type": "application/json"}, - method="POST", - ) - with urllib.request.urlopen(req) as r: - r.read() - - with pytest.raises(urllib.error.HTTPError) as exc_info: - _chat( - tracker_url, - api_key="capped-client", - messages=[{"role": "user", "content": " ".join(["prompt"] * 1000)}], - max_tokens=1, - ) - body = json.loads(exc_info.value.read()) - assert exc_info.value.code == 402 - assert body["error"]["code"] == "spend_cap_exceeded" - assert "max_charge_per_request" in body["error"]["message"] - assert ledger.get_client_balance("capped-client") == pytest.approx(10.0) - assert stub.request_count == 0 - finally: - stub.stop() - tracker.stop() - - -def test_proxy_chat_records_validation_event_with_plain_route_metadata(): - class FakeRegistry: - def get_wallet(self, wallet_address): - return type("Wallet", (), {"banned": False})() - - class FakeValidation: - def __init__(self): - self.events = [] - - def record_completed_inference(self, **kwargs): - self.events.append(kwargs) - return kwargs - - class FakeContracts: - def __init__(self): - self.registry = FakeRegistry() - self.validation = FakeValidation() - - contracts = FakeContracts() - tracker = TrackerServer( - model_presets={ - MODEL: { - "layers_start": 0, - "layers_end": 11, - "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, - } - }, - contracts=contracts, - ) - port = tracker.start() - tracker_url = f"http://127.0.0.1:{port}" - stub = _UsageStubNode(total_tokens=1000) - stub_port = stub.start() - try: - reg = json.dumps({ - "endpoint": f"http://127.0.0.1:{stub_port}", - "shard_start": 0, - "shard_end": 11, - "model": MODEL, - "hardware_profile": {}, - "score": 1.0, - "tracker_mode": True, - "wallet_address": "wallet-head", - }).encode() - req = urllib.request.Request( - f"{tracker_url}/v1/nodes/register", - data=reg, - headers={"Content-Type": "application/json"}, - method="POST", - ) - with urllib.request.urlopen(req) as r: - r.read() - - _chat(tracker_url, api_key=None) - - assert len(contracts.validation.events) == 1 - event = contracts.validation.events[0] - assert event["model"] == MODEL - assert event["messages"] == [{"role": "user", "content": "hi"}] - assert event["observed_output"] == "ok" - assert event["route_nodes"] == [{ - "node_id": next(iter(tracker._registry)), - "endpoint": f"http://127.0.0.1:{stub_port}", - "wallet_address": "wallet-head", - "shard_start": 0, - "shard_end": 11, - }] - assert "toploc_proof" not in event["route_nodes"][0] - finally: - stub.stop() - tracker.stop() - - -def test_billing_gossip_endpoint_applies_events(billed_tracker): - tracker_url, ledger, _ = billed_tracker - peer = BillingLedger(starting_credit=0.03, default_price_per_1k=0.02) - peer.credit_client("remote-client", 7.0) - events, _ = peer.events_since(0) - body = json.dumps({"events": events}).encode() - req = urllib.request.Request( - f"{tracker_url}/v1/billing/gossip", - data=body, - headers={"Content-Type": "application/json", **sign_hive_request(HIVE_SECRET, body)}, - method="POST", - ) - with urllib.request.urlopen(req) as r: - assert json.loads(r.read())["applied"] == len(events) - # only the replicated credit event lands here. - assert ledger.get_client_balance("remote-client") == pytest.approx(7.0) +"""US-031: billing ledger — per-token pricing, 90/10 split, pending balances. + +Unit tests for BillingLedger math/persistence/replication, plus HTTP +integration: 401 without an API key, 402 for unfunded keys, billed 200 after +explicit credit, and 402 once the balance is exhausted. +""" + +import http.server +import json +import socketserver +import threading +import urllib.error +import urllib.request + +import pytest + +from meshnet_tracker.auth import sign_hive_request +from meshnet_tracker.billing import DEFAULT_STARTING_CREDIT, BillingLedger +from meshnet_tracker.server import ( + TrackerServer, + _billable_non_stream_tokens, + _billable_stream_tokens, + _estimate_prompt_tokens, + _observed_non_stream_completion_tokens, + _observed_stream_tokens, +) + +MODEL = "openai-community/gpt2" +HIVE_SECRET = "test-hive-secret" + + +# ---------------------------------------------------------------- unit tests + + +def test_charge_single_node_gets_90_percent(): + ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) + ledger.ensure_client("key-a") + event = ledger.charge_request("key-a", MODEL, total_tokens=1000, node_work=[("wallet-1", 12)]) + assert event["cost"] == pytest.approx(0.02) + assert ledger.get_client_balance("key-a") == pytest.approx(1.0 - 0.02) + assert ledger.get_node_pending("wallet-1") == pytest.approx(0.02 * 0.90) + assert ledger.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10) + + +def test_default_starting_credit_is_zero_and_fresh_key_is_unfunded(): + ledger = BillingLedger(default_price_per_1k=0.02) + + assert DEFAULT_STARTING_CREDIT == 0.0 + assert ledger.ensure_client("fresh-key") == pytest.approx(0.0) + assert ledger.get_client_balance("fresh-key") == pytest.approx(0.0) + assert ledger.has_funds("fresh-key") is False + assert ledger.events_since(0)[0] == [] + + +def test_charge_three_node_split_by_work_units(): + ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) + ledger.ensure_client("key-a") + ledger.charge_request( + "key-a", MODEL, total_tokens=1000, + node_work=[("wallet-1", 6), ("wallet-2", 3), ("wallet-3", 3)], + ) + pool = 0.02 * 0.90 + assert ledger.get_node_pending("wallet-1") == pytest.approx(pool * 6 / 12) + assert ledger.get_node_pending("wallet-2") == pytest.approx(pool * 3 / 12) + assert ledger.get_node_pending("wallet-3") == pytest.approx(pool * 3 / 12) + assert ledger.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10) + + +def test_walletless_node_share_accrues_to_protocol_cut(): + ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) + ledger.ensure_client("key-a") + ledger.charge_request( + "key-a", MODEL, total_tokens=1000, + node_work=[("wallet-1", 6), (None, 6)], + ) + pool = 0.02 * 0.90 + assert ledger.get_node_pending("wallet-1") == pytest.approx(pool / 2) + # walletless half of the pool + the 10% cut both land in protocol_cut + assert ledger.snapshot()["protocol_cut"] == pytest.approx(pool / 2 + 0.02 * 0.10) + + +def test_per_model_price_override(): + ledger = BillingLedger( + starting_credit=1.0, + default_price_per_1k=0.02, + prices={MODEL: 0.10}, + ) + ledger.ensure_client("key-a") + event = ledger.charge_request("key-a", MODEL, 500, [("wallet-1", 1)]) + assert event["cost"] == pytest.approx(0.10 * 500 / 1000) + event = ledger.charge_request("key-a", "other-model", 500, [("wallet-1", 1)]) + assert event["cost"] == pytest.approx(0.02 * 500 / 1000) + + +def test_non_stream_billable_tokens_cap_usage_by_request_bound(): + payload = {"usage": {"total_tokens": 1_000_000}} + request = {"messages": [{"role": "user", "content": "hello tracker"}], "max_tokens": 3} + + assert _estimate_prompt_tokens(request) == 2 + assert _billable_non_stream_tokens(payload, request) == 5 + + +def test_non_stream_billable_tokens_fallback_when_usage_missing(): + payload = {"choices": [{"message": {"role": "assistant", "content": "ok now"}}]} + request = {"messages": [{"role": "user", "content": "hello tracker"}], "max_tokens": 10} + + assert _observed_non_stream_completion_tokens(payload) == 2 + assert _billable_non_stream_tokens(payload, request) == 4 + + +def test_stream_billable_tokens_allow_usage_to_lower_observed_count_only(): + observed_payload = { + "choices": [{ + "index": 0, + "delta": {"content": "hello tracker"}, + "finish_reason": None, + }] + } + + assert _observed_stream_tokens(observed_payload) == 2 + assert _billable_stream_tokens(observed_tokens=2, reported_tokens=1_000_000) == 2 + assert _billable_stream_tokens(observed_tokens=2, reported_tokens=1) == 1 + + +def test_payout_and_forfeit_hooks(): + ledger = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) + ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 1)]) + pending = ledger.get_node_pending("wallet-1") + ledger.settle_node_payout("wallet-1", pending, reference="tx-sig-1") + assert ledger.get_node_pending("wallet-1") == pytest.approx(0.0) + + ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 1)]) + cut_before = ledger.snapshot()["protocol_cut"] + forfeited = ledger.forfeit_pending("wallet-1")["amount"] + assert forfeited == pytest.approx(0.02 * 0.90) + assert ledger.get_node_pending("wallet-1") == pytest.approx(0.0) + assert ledger.snapshot()["protocol_cut"] == pytest.approx(cut_before + forfeited) + + +def test_restart_persistence(tmp_path): + db = str(tmp_path / "billing.db") + ledger = BillingLedger(db_path=db, starting_credit=1.0, default_price_per_1k=0.02) + ledger.credit_client("key-a", 5.0) + ledger.charge_request("key-a", MODEL, 1000, [("wallet-1", 12)]) + ledger.save_to_db() + + reloaded = BillingLedger(db_path=db, starting_credit=1.0, default_price_per_1k=0.02) + assert reloaded.get_client_balance("key-a") == pytest.approx( + ledger.get_client_balance("key-a") + ) + assert reloaded.get_node_pending("wallet-1") == pytest.approx(0.02 * 0.90) + assert reloaded.snapshot()["protocol_cut"] == pytest.approx(0.02 * 0.10) + + +def test_tracker_enables_billing_with_default_db_when_requested(tmp_path, monkeypatch): + from meshnet_tracker.billing import DEFAULT_BILLING_DB_PATH + + monkeypatch.chdir(tmp_path) + tracker = TrackerServer(enable_billing=True) + port = tracker.start() + try: + # /v1/billing/summary is admin-gated now; just confirm the server is up. + health = json.loads( + urllib.request.urlopen(f"http://127.0.0.1:{port}/v1/health").read() + ) + assert health["status"] == "ok" + finally: + tracker.stop() + # enabling billing creates the ledger DB in the working directory + assert (tmp_path / DEFAULT_BILLING_DB_PATH).exists() + + +def test_event_replication_converges_and_dedupes(): + leader = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) + follower = BillingLedger(starting_credit=1.0, default_price_per_1k=0.02) + + leader.credit_client("key-a", 10.0) + leader.charge_request("key-a", MODEL, 2000, [("wallet-1", 8), ("wallet-2", 4)]) + + events, cursor = leader.events_since(0) + assert follower.apply_events(events) == len(events) + # replaying the same batch must be a no-op + assert follower.apply_events(events) == 0 + + assert follower.get_client_balance("key-a") == pytest.approx( + leader.get_client_balance("key-a") + ) + assert follower.get_node_pending("wallet-1") == pytest.approx( + leader.get_node_pending("wallet-1") + ) + assert follower.snapshot()["protocol_cut"] == pytest.approx( + leader.snapshot()["protocol_cut"] + ) + # incremental cursor: nothing new after full sync + more, _ = leader.events_since(cursor) + assert more == [] + + +# ---------------------------------------------------------- HTTP integration + + +class _UsageStubNode: + """Minimal head node returning a chat completion with real usage numbers.""" + + def __init__( + self, + total_tokens: int = 1000, + *, + stream_chunks: list[str] | None = None, + stream_usage_total: int | None = None, + ): + self.total_tokens = total_tokens + self.stream_chunks = stream_chunks or [] + self.stream_usage_total = stream_usage_total + self.request_count = 0 + outer = self + + class Handler(http.server.BaseHTTPRequestHandler): + def log_message(self, fmt, *args): + pass + + def do_POST(self): + outer.request_count += 1 + length = int(self.headers.get("Content-Length", 0)) + raw_body = self.rfile.read(length) + try: + request_body = json.loads(raw_body) + except json.JSONDecodeError: + request_body = {} + if request_body.get("stream"): + self.send_response(200) + self.send_header("Content-Type", "text/event-stream; charset=utf-8") + self.end_headers() + for chunk in outer.stream_chunks: + payload = json.dumps({ + "id": "chatcmpl-stub", + "object": "chat.completion.chunk", + "model": MODEL, + "choices": [{ + "index": 0, + "delta": {"content": chunk}, + "finish_reason": None, + }], + }).encode() + self.wfile.write(b"data: " + payload + b"\n\n") + if outer.stream_usage_total is not None: + usage_payload = json.dumps({ + "id": "chatcmpl-stub", + "object": "chat.completion.chunk", + "model": MODEL, + "choices": [], + "usage": { + "prompt_tokens": 0, + "completion_tokens": outer.stream_usage_total, + "total_tokens": outer.stream_usage_total, + }, + }).encode() + self.wfile.write(b"data: " + usage_payload + b"\n\n") + self.wfile.write(b"data: [DONE]\n\n") + return + body = json.dumps({ + "id": "chatcmpl-stub", + "object": "chat.completion", + "model": MODEL, + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + }], + "usage": { + "prompt_tokens": 0, + "completion_tokens": outer.total_tokens, + "total_tokens": outer.total_tokens, + }, + }).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + self._server = socketserver.ThreadingTCPServer(("127.0.0.1", 0), Handler) + self._server.daemon_threads = True + self._thread: threading.Thread | None = None + + def start(self) -> int: + self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + self._thread.start() + return self._server.server_address[1] + + def stop(self): + self._server.shutdown() + self._server.server_close() + + +@pytest.fixture +def billed_tracker(): + ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) + tracker = TrackerServer( + model_presets={ + MODEL: { + "layers_start": 0, + "layers_end": 11, + "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, + } + }, + billing=ledger, + hive_secret=HIVE_SECRET, + ) + port = tracker.start() + tracker_url = f"http://127.0.0.1:{port}" + + stub = _UsageStubNode(total_tokens=1000) + stub_port = stub.start() + reg = json.dumps({ + "endpoint": f"http://127.0.0.1:{stub_port}", + "shard_start": 0, + "shard_end": 11, + "model": MODEL, + "hardware_profile": {}, + "score": 1.0, + "tracker_mode": True, + "wallet_address": "wallet-head", + }).encode() + req = urllib.request.Request( + f"{tracker_url}/v1/nodes/register", + data=reg, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req) as r: + r.read() + + yield tracker_url, ledger, stub + + stub.stop() + tracker.stop() + + +def _chat(tracker_url: str, api_key: str | None, **body_overrides): + body = { + "model": MODEL, + "messages": [{"role": "user", "content": "hi"}], + } + body.update(body_overrides) + data = json.dumps(body).encode() + headers = {"Content-Type": "application/json"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + req = urllib.request.Request( + f"{tracker_url}/v1/chat/completions", + data=data, + headers=headers, + method="POST", + ) + with urllib.request.urlopen(req) as r: + if body.get("stream"): + return r.read().decode() + return json.loads(r.read()) + + +def test_proxy_chat_requires_api_key_when_billing_enabled(billed_tracker): + tracker_url, _, _ = billed_tracker + with pytest.raises(urllib.error.HTTPError) as exc_info: + _chat(tracker_url, api_key=None) + assert exc_info.value.code == 401 + + +def test_proxy_chat_402_for_fresh_key_before_routing(billed_tracker): + tracker_url, ledger, stub = billed_tracker + with pytest.raises(urllib.error.HTTPError) as exc_info: + _chat(tracker_url, api_key="fresh-client") + assert exc_info.value.code == 402 + assert ledger.get_client_balance("fresh-client") == pytest.approx(0.0) + assert stub.request_count == 0 + + +def test_proxy_chat_bills_credited_client_and_credits_node(billed_tracker): + tracker_url, ledger, _ = billed_tracker + ledger.credit_client("client-1", 0.03, note="admin-credit") + _chat(tracker_url, api_key="client-1") + # 1000 tokens at 0.02/1K = 0.02 USDT; explicit credit 0.03 + assert ledger.get_client_balance("client-1") == pytest.approx(0.03 - 0.02) + assert ledger.get_node_pending("wallet-head") == pytest.approx(0.02 * 0.90) + + # /v1/billing/summary is admin-gated now (ADR-0017); auth is covered in + # test_auth_boundary.py, so verify the numbers via the ledger snapshot. + summary = ledger.snapshot() + assert summary["node_pending"]["wallet-head"] == pytest.approx(0.02 * 0.90) + assert summary["protocol_cut"] == pytest.approx(0.02 * 0.10) + + stats = json.loads(urllib.request.urlopen(f"{tracker_url}/v1/stats").read()) + node_stats = next(iter(stats["nodes"].values()))["models"][MODEL] + assert node_stats["tokens_per_sec_last_hour"] is not None + assert node_stats["sample_count_last_hour"] == 1 + + +def test_proxy_chat_caps_inflated_non_streaming_usage_by_request_bounds(billed_tracker): + tracker_url, ledger, stub = billed_tracker + stub.total_tokens = 1_000_000 + ledger.credit_client("bounded-client", 100.0, note="admin-credit") + + _chat(tracker_url, api_key="bounded-client", max_tokens=2) + + # messages=[{"content": "hi"}] has a local prompt estimate of one token, so + # billable total is capped at max_tokens + prompt estimate, not node usage. + expected_tokens = 3 + assert ledger.get_client_balance("bounded-client") == pytest.approx(100.0 - 0.02 * expected_tokens / 1000) + events, _ = ledger.events_since(0) + charge = next(event for event in events if event["type"] == "charge") + assert charge["total_tokens"] == expected_tokens + + +def test_proxy_chat_caps_inflated_streaming_usage_by_observed_chunks(): + ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) + ledger.credit_client("stream-client", 1.0, note="admin-credit") + tracker = TrackerServer( + model_presets={ + MODEL: { + "layers_start": 0, + "layers_end": 11, + "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, + } + }, + billing=ledger, + ) + port = tracker.start() + tracker_url = f"http://127.0.0.1:{port}" + stub = _UsageStubNode(stream_chunks=["one", "two"], stream_usage_total=1_000_000) + stub_port = stub.start() + try: + reg = json.dumps({ + "endpoint": f"http://127.0.0.1:{stub_port}", + "shard_start": 0, + "shard_end": 11, + "model": MODEL, + "hardware_profile": {}, + "score": 1.0, + "tracker_mode": True, + "wallet_address": "wallet-head", + }).encode() + req = urllib.request.Request( + f"{tracker_url}/v1/nodes/register", + data=reg, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req) as r: + r.read() + + _chat(tracker_url, api_key="stream-client", stream=True, max_tokens=2) + + events, _ = ledger.events_since(0) + charge = next(event for event in events if event["type"] == "charge") + assert charge["total_tokens"] == 2 + assert ledger.get_client_balance("stream-client") == pytest.approx(1.0 - 0.02 * 2 / 1000) + finally: + stub.stop() + tracker.stop() + + +def test_proxy_chat_splits_payout_by_tracker_assigned_route_span(): + ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) + ledger.credit_client("route-client", 1.0, note="admin-credit") + tracker = TrackerServer( + model_presets={ + MODEL: { + "total_layers": 12, + "bytes_per_layer": {"bfloat16": 1_000}, + } + }, + billing=ledger, + ) + port = tracker.start() + tracker_url = f"http://127.0.0.1:{port}" + head = _UsageStubNode(total_tokens=1000) + head_port = head.start() + try: + for endpoint, wallet, vram, bench in ( + (f"http://127.0.0.1:{head_port}", "wallet-head", 5_000, 2.0), + ("http://127.0.0.1:1", "wallet-tail", 10_000, 1.0), + ): + reg = json.dumps({ + "endpoint": endpoint, + "model": MODEL, + "vram_bytes": vram, + "ram_bytes": 10_000, + "quantizations": ["bfloat16"], + "benchmark_tokens_per_sec": bench, + "hardware_profile": {}, + "score": 1.0, + "tracker_mode": endpoint.endswith(str(head_port)), + "wallet_address": wallet, + "managed_assignment": True, + "shard_start": 0, + "shard_end": 999, + }).encode() + req = urllib.request.Request( + f"{tracker_url}/v1/nodes/register", + data=reg, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req) as r: + r.read() + + _chat(tracker_url, api_key="route-client", max_tokens=1000) + + pool = 0.02 * 0.90 + assert ledger.get_node_pending("wallet-head") == pytest.approx(pool * 4 / 12) + assert ledger.get_node_pending("wallet-tail") == pytest.approx(pool * 8 / 12) + finally: + head.stop() + tracker.stop() + + +def test_proxy_chat_402_when_balance_exhausted(billed_tracker): + tracker_url, ledger, _ = billed_tracker + ledger.credit_client("client-2", 0.03, note="admin-credit") + _chat(tracker_url, api_key="client-2") # 0.03 -> 0.01 + _chat(tracker_url, api_key="client-2") # 0.01 -> -0.01 (post-pay drift) + assert ledger.get_client_balance("client-2") == pytest.approx(-0.01) + with pytest.raises(urllib.error.HTTPError) as exc_info: + _chat(tracker_url, api_key="client-2") + assert exc_info.value.code == 402 + # rejected before routing: nothing further was billed + assert ledger.get_client_balance("client-2") == pytest.approx(-0.01) + + +def test_proxy_chat_rejects_request_above_spend_cap_before_routing(): + ledger = BillingLedger(starting_credit=0.0, default_price_per_1k=0.02) + ledger.credit_client("capped-client", 10.0, note="admin-credit") + tracker = TrackerServer( + model_presets={ + MODEL: { + "layers_start": 0, + "layers_end": 11, + "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, + } + }, + billing=ledger, + max_charge_per_request=0.01, + ) + port = tracker.start() + tracker_url = f"http://127.0.0.1:{port}" + stub = _UsageStubNode(total_tokens=1000) + stub_port = stub.start() + try: + reg = json.dumps({ + "endpoint": f"http://127.0.0.1:{stub_port}", + "shard_start": 0, + "shard_end": 11, + "model": MODEL, + "hardware_profile": {}, + "score": 1.0, + "tracker_mode": True, + "wallet_address": "wallet-head", + }).encode() + req = urllib.request.Request( + f"{tracker_url}/v1/nodes/register", + data=reg, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req) as r: + r.read() + + with pytest.raises(urllib.error.HTTPError) as exc_info: + _chat( + tracker_url, + api_key="capped-client", + messages=[{"role": "user", "content": " ".join(["prompt"] * 1000)}], + max_tokens=1, + ) + body = json.loads(exc_info.value.read()) + assert exc_info.value.code == 402 + assert body["error"]["code"] == "spend_cap_exceeded" + assert "max_charge_per_request" in body["error"]["message"] + assert ledger.get_client_balance("capped-client") == pytest.approx(10.0) + assert stub.request_count == 0 + finally: + stub.stop() + tracker.stop() + + +def test_proxy_chat_records_validation_event_with_plain_route_metadata(): + class FakeRegistry: + def get_wallet(self, wallet_address): + return type("Wallet", (), {"banned": False})() + + class FakeValidation: + def __init__(self): + self.events = [] + + def record_completed_inference(self, **kwargs): + self.events.append(kwargs) + return kwargs + + class FakeContracts: + def __init__(self): + self.registry = FakeRegistry() + self.validation = FakeValidation() + + contracts = FakeContracts() + tracker = TrackerServer( + model_presets={ + MODEL: { + "layers_start": 0, + "layers_end": 11, + "bytes_per_layer": {"bfloat16": 30 * 1024 * 1024}, + } + }, + contracts=contracts, + ) + port = tracker.start() + tracker_url = f"http://127.0.0.1:{port}" + stub = _UsageStubNode(total_tokens=1000) + stub_port = stub.start() + try: + reg = json.dumps({ + "endpoint": f"http://127.0.0.1:{stub_port}", + "shard_start": 0, + "shard_end": 11, + "model": MODEL, + "hardware_profile": {}, + "score": 1.0, + "tracker_mode": True, + "wallet_address": "wallet-head", + }).encode() + req = urllib.request.Request( + f"{tracker_url}/v1/nodes/register", + data=reg, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req) as r: + r.read() + + _chat(tracker_url, api_key=None) + + assert len(contracts.validation.events) == 1 + event = contracts.validation.events[0] + assert event["model"] == MODEL + assert event["messages"] == [{"role": "user", "content": "hi"}] + assert event["observed_output"] == "ok" + assert event["route_nodes"] == [{ + "node_id": next(iter(tracker._registry)), + "endpoint": f"http://127.0.0.1:{stub_port}", + "wallet_address": "wallet-head", + "shard_start": 0, + "shard_end": 11, + }] + assert "toploc_proof" not in event["route_nodes"][0] + finally: + stub.stop() + tracker.stop() + + +def test_billing_gossip_endpoint_applies_events(billed_tracker): + tracker_url, ledger, _ = billed_tracker + peer = BillingLedger(starting_credit=0.03, default_price_per_1k=0.02) + peer.credit_client("remote-client", 7.0) + events, _ = peer.events_since(0) + body = json.dumps({"events": events}).encode() + req = urllib.request.Request( + f"{tracker_url}/v1/billing/gossip", + data=body, + headers={"Content-Type": "application/json", **sign_hive_request(HIVE_SECRET, body)}, + method="POST", + ) + with urllib.request.urlopen(req) as r: + assert json.loads(r.read())["applied"] == len(events) + # only the replicated credit event lands here. + assert ledger.get_client_balance("remote-client") == pytest.approx(7.0) diff --git a/tests/test_dashboard.py b/tests/test_dashboard.py index af8fb5b..a077842 100644 --- a/tests/test_dashboard.py +++ b/tests/test_dashboard.py @@ -29,6 +29,9 @@ def test_dashboard_served_with_all_panels(): for panel in PANELS: assert panel in html assert "