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/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 9a5a398..5ee3989 100644 --- a/QUICKSTART.md +++ b/QUICKSTART.md @@ -1,43 +1,43 @@ -# 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`). +# 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`). - **Linear-attention fast path (GPU only).** Qwen3.5/3.6 use hybrid linear-attention layers; without optional CUDA kernels, Transformers falls back to slower pure-PyTorch code and prints `The fast path is not available…` at startup. That warning is @@ -54,748 +54,748 @@ python3 -m venv .venv # NVIDIA (CUDA) pip install flash-linear-attention[cuda] causal-conv1d - # AMD (ROCm) — match your torch index, then: - pip install flash-linear-attention[rocm] causal-conv1d - ``` - - Restart the node after install; the warning should disappear. Expect the largest - gain on GPU nodes serving linear-attention layers (roughly three quarters of - Qwen3.6 layers); end-to-end chat speed still depends on the slowest hop in a - split route. -- `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 -``` + # AMD (ROCm) — match your torch index, then: + pip install flash-linear-attention[rocm] causal-conv1d + ``` + + Restart the node after install; the warning should disappear. Expect the largest + gain on GPU nodes serving linear-attention layers (roughly three quarters of + Qwen3.6 layers); end-to-end chat speed still depends on the slowest hop in a + split route. +- `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 +``` diff --git a/docs/INSTALL_WINDOWS.md b/docs/INSTALL_WINDOWS.md index c204cad..b7730dc 100644 --- a/docs/INSTALL_WINDOWS.md +++ b/docs/INSTALL_WINDOWS.md @@ -1,111 +1,111 @@ -# 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__)" -``` - +# 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 @@ -124,107 +124,107 @@ 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. + +## 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/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/pyproject.toml b/packages/node/pyproject.toml index 67f9314..3ba2d56 100644 --- a/packages/node/pyproject.toml +++ b/packages/node/pyproject.toml @@ -1,34 +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", +[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"] + "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 cd8fb26..0b40b55 100644 --- a/packages/tracker/meshnet_tracker/cli.py +++ b/packages/tracker/meshnet_tracker/cli.py @@ -1,419 +1,419 @@ -"""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, - DEFAULT_LOG_DIR, - 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 _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, - metavar="PATH", - help=( - "Directory for rotating tracker logs " - f"(default: {DEFAULT_LOG_DIR}; files: info.log, warning.log, error.log)" - ), - ) - common.add_argument( - "--log-max-bytes", - type=int, - default=DEFAULT_LOG_MAX_BYTES, - metavar="BYTES", - help=f"Rotate each tracker log file after this many bytes (default: {DEFAULT_LOG_MAX_BYTES})", - ) - common.add_argument( - "--log-backup-count", - type=int, - default=DEFAULT_LOG_BACKUP_COUNT, - metavar="N", - help=f"Number of rotated tracker log files to keep per level (default: {DEFAULT_LOG_BACKUP_COUNT})", - ) - common.add_argument( - "--no-file-logs", - 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() - - if args.command in {None, "start"}: - if not args.no_file_logs: - log_dir = configure_tracker_file_logging( - args.log_dir, - max_bytes=args.log_max_bytes, - backup_count=args.log_backup_count, - ) - 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, - 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() +"""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, + DEFAULT_LOG_DIR, + 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 _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, + metavar="PATH", + help=( + "Directory for rotating tracker logs " + f"(default: {DEFAULT_LOG_DIR}; files: info.log, warning.log, error.log)" + ), + ) + common.add_argument( + "--log-max-bytes", + type=int, + default=DEFAULT_LOG_MAX_BYTES, + metavar="BYTES", + help=f"Rotate each tracker log file after this many bytes (default: {DEFAULT_LOG_MAX_BYTES})", + ) + common.add_argument( + "--log-backup-count", + type=int, + default=DEFAULT_LOG_BACKUP_COUNT, + metavar="N", + help=f"Number of rotated tracker log files to keep per level (default: {DEFAULT_LOG_BACKUP_COUNT})", + ) + common.add_argument( + "--no-file-logs", + 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() + + if args.command in {None, "start"}: + if not args.no_file_logs: + log_dir = configure_tracker_file_logging( + args.log_dir, + max_bytes=args.log_max_bytes, + backup_count=args.log_backup_count, + ) + 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, + 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 e3ba6a7..40c6145 100644 --- a/packages/tracker/meshnet_tracker/dashboard.html +++ b/packages/tracker/meshnet_tracker/dashboard.html @@ -1,1566 +1,1566 @@ - - - - - -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
-
- - - + + + + + +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/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_real_model_backend.py b/tests/test_real_model_backend.py index 87caa67..38815fe 100644 --- a/tests/test_real_model_backend.py +++ b/tests/test_real_model_backend.py @@ -1,1051 +1,1051 @@ -"""US-012 tests for the real PyTorch node backend.""" - -import json -import os -from pathlib import Path -import sys -import threading -import time -import types -import urllib.request - -import pytest - -from meshnet_node.model_backend import ( - InsufficientVRAMError, - PartialModelLoadUnsupported, - TensorPayload, - TorchModelShard, - _call_layer, - _checkpoint_tensor_name_for_model, - _load_partial_model_from_snapshot, - _should_partial_materialize_shard, - _decoder_attention_mask, - _int_tensor_header, - build_quantization_config, - validate_quantization, -) -from meshnet_node.torch_server import TorchNodeServer - - -class _FakeBackend: - model_id = "fake-model" - total_layers = 12 - is_head = True - is_tail = False - - def encode_prompt(self, prompt: str) -> TensorPayload: - assert prompt == "The capital of France is" - return TensorPayload( - body=b"\x00" * (1 * 6 * 8 * 2), - shape=[1, 6, 8], - attention_mask_header=None, - position_ids_header=None, - ) - - def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): - assert shape == [1, 6, 8] - return TensorPayload( - body=body, - shape=shape, - attention_mask_header=attention_mask_header, - position_ids_header=position_ids_header, - ) - - -class _FakeTailBackend(_FakeBackend): - is_head = False - is_tail = True - - def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): - assert len(body) == 1 * 6 * 8 * 2 - return " Paris" - - -class _FakeFullBackend(_FakeBackend): - is_head = True - is_tail = True - - def generate_text( - self, - messages: list[dict], - max_new_tokens: int = 16, - temperature: float = 1.0, - top_p: float = 1.0, - ) -> str: - assert messages == [{"role": "user", "content": "What is 7 times 8?"}] - assert max_new_tokens == 7 - assert temperature == 1.0 - assert top_p == 1.0 - return "56" - - def count_prompt_tokens(self, messages: list[dict]) -> int: - assert messages == [{"role": "user", "content": "What is 7 times 8?"}] - return 8 - - def count_text_tokens(self, text: str) -> int: - assert text == "56" - return 1 - - -class _FakeChatTokenizer: - eos_token = "" - - def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False): - assert add_generation_prompt is True - assert tokenize is False - return "debug prompt" - - -class _FakePipelineHeadBackend(_FakeBackend): - tokenizer = _FakeChatTokenizer() - - def encode_prompt(self, prompt: str) -> TensorPayload: - assert prompt.startswith("debug prompt") - return TensorPayload( - body=b"\x00" * (1 * 6 * 8 * 2), - shape=[1, 6, 8], - attention_mask_header=None, - position_ids_header=None, - ) - - -class _FakePipelineTailBackend(_FakeTailBackend): - def __init__(self) -> None: - self.start_layers: list[int | None] = [] - - def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): - self.start_layers.append(start_layer) - assert len(body) == 1 * 6 * 8 * 2 - return " token" - - -class _BlockingStreamingTailBackend(_FakeTailBackend): - def __init__(self, second_token_release: threading.Event) -> None: - self._release = second_token_release - self.calls = 0 - - def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): - self.calls += 1 - if self.calls == 1: - return " first" - self._release.wait(timeout=3.0) - return " second" - - -def test_quantization_flag_validation(): - assert validate_quantization("bfloat16") == "bfloat16" - assert validate_quantization("int8") == "int8" - assert validate_quantization("nf4") == "nf4" - with pytest.raises(ValueError, match="quantization"): - validate_quantization("float32") - - -def test_node_package_declares_torch_dependency(): - pyproject = Path("packages/node/pyproject.toml").read_text(encoding="utf-8") - - assert '"torch>=' in pyproject - - -def test_bitsandbytes_configs_are_created_lazily(monkeypatch): - calls = [] - - class FakeBitsAndBytesConfig: - def __init__(self, **kwargs): - calls.append(kwargs) - - monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(bfloat16="bf16")) - monkeypatch.setitem( - sys.modules, - "transformers", - types.SimpleNamespace(BitsAndBytesConfig=FakeBitsAndBytesConfig), - ) - - assert build_quantization_config("bfloat16") is None - build_quantization_config("int8") - build_quantization_config("nf4") - - assert calls == [ - {"load_in_8bit": True}, - { - "load_in_4bit": True, - "bnb_4bit_quant_type": "nf4", - "bnb_4bit_compute_dtype": "bf16", - }, - ] - - -def test_head_forward_accepts_text_prompt_and_returns_bfloat16_activations(): - node = TorchNodeServer(backend=_FakeBackend()) - port = node.start() - try: - payload = json.dumps({"prompt": "The capital of France is"}).encode() - req = urllib.request.Request( - f"http://127.0.0.1:{port}/forward", - data=payload, - headers={"Content-Type": "application/json"}, - method="POST", - ) - with urllib.request.urlopen(req, timeout=5) as resp: - body = resp.read() - headers = {key.lower(): value for key, value in resp.headers.items()} - - assert len(body) == 1 * 6 * 8 * 2 - assert headers["x-meshnet-shape"] == "1,6,8" - assert headers["x-meshnet-dtype"] == "bfloat16" - assert headers["x-meshnet-wire"] == "2" - finally: - node.stop() - - -def test_tail_forward_returns_text_completion_from_binary_activations(): - node = TorchNodeServer(backend=_FakeTailBackend()) - port = node.start() - try: - req = urllib.request.Request( - f"http://127.0.0.1:{port}/forward", - data=b"\x00" * (1 * 6 * 8 * 2), - headers={ - "Content-Type": "application/octet-stream", - "X-Meshnet-Shape": "1,6,8", - "X-Meshnet-Dtype": "bfloat16", - "X-Meshnet-Session": "session-1", - "X-Meshnet-Chunk-Index": "0", - "X-Meshnet-Chunk-Total": "1", - "X-Meshnet-Hop-Index": "1", - }, - method="POST", - ) - with urllib.request.urlopen(req, timeout=5) as resp: - body = json.loads(resp.read()) - - assert body == {"text": " Paris"} - assert node.received_activations - assert node.forward_chunk_count == 1 - finally: - node.stop() - - -def test_full_model_chat_completion_uses_generation_not_single_token_decode(capsys): - node = TorchNodeServer(backend=_FakeFullBackend()) - port = node.start() - try: - payload = json.dumps({ - "model": "fake-model", - "messages": [{"role": "user", "content": "What is 7 times 8?"}], - "max_tokens": 7, - }).encode() - req = urllib.request.Request( - f"http://127.0.0.1:{port}/v1/chat/completions", - data=payload, - headers={ - "Content-Type": "application/json", - "X-Meshnet-Request-Id": "req-test-123", - }, - method="POST", - ) - with urllib.request.urlopen(req, timeout=5) as resp: - body = json.loads(resp.read()) - - assert body["choices"][0]["message"]["content"] == "56" - assert body["usage"] == {"prompt_tokens": 8, "completion_tokens": 1, "total_tokens": 9} - finally: - node.stop() - - out = capsys.readouterr().out - assert " [node] processing chat model='fake-model' stream=False max_tokens=7 request_id=req-test-123" in out - assert " [node] chat complete tokens=1 elapsed_s=" in out - - -def test_pipeline_hop_logs_are_suppressed_without_debug(capsys): - tail_backend = _FakePipelineTailBackend() - head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True) - tail = TorchNodeServer(backend=tail_backend) - head_port = head.start() - tail_port = tail.start() - try: - payload = json.dumps({ - "model": "fake-model", - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 1, - }).encode() - req = urllib.request.Request( - f"http://127.0.0.1:{head_port}/v1/chat/completions", - data=payload, - headers={ - "Content-Type": "application/json", - "X-Meshnet-Route": json.dumps([ - {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22}, - ]), - }, - method="POST", - ) - with urllib.request.urlopen(req, timeout=5) as resp: - body = json.loads(resp.read()) - finally: - head.stop() - tail.stop() - - out = capsys.readouterr().out - assert body["choices"][0]["message"]["content"] == " token" - assert tail_backend.start_layers == [22] - assert "pipeline hop 0:" not in out - assert "pipeline hop 0 returned text" not in out - - -def test_pipeline_hop_logs_are_enabled_with_debug(capsys): - head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True, debug=True) - tail = TorchNodeServer(backend=_FakePipelineTailBackend()) - head_port = head.start() - tail_port = tail.start() - try: - payload = json.dumps({ - "model": "fake-model", - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 1, - }).encode() - req = urllib.request.Request( - f"http://127.0.0.1:{head_port}/v1/chat/completions", - data=payload, - headers={ - "Content-Type": "application/json", - "X-Meshnet-Route": json.dumps([ - {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22}, - ]), - }, - method="POST", - ) - with urllib.request.urlopen(req, timeout=5) as resp: - json.loads(resp.read()) - finally: - head.stop() - tail.stop() - - out = capsys.readouterr().out - assert f" [node] pipeline hop 0: http://127.0.0.1:{tail_port} start_layer=22" in out - assert " [node] pipeline hop 0 returned text=' token'" in out - - -def test_split_shard_chat_streams_each_generated_token_incrementally(): - release_second = threading.Event() - head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True) - tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second)) - head_port = head.start() - tail_port = tail.start() - response = None - try: - payload = json.dumps({ - "model": "fake-model", - "messages": [{"role": "user", "content": "hello"}], - "stream": True, - "max_tokens": 2, - }).encode() - req = urllib.request.Request( - f"http://127.0.0.1:{head_port}/v1/chat/completions", - data=payload, - headers={ - "Content-Type": "application/json", - "X-Meshnet-Route": json.dumps([ - {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22}, - ]), - }, - method="POST", - ) - response = urllib.request.urlopen(req, timeout=5) - - first_token_line = "" - deadline = time.time() + 2.0 - while time.time() < deadline: - line = response.readline().decode() - if '"content": " first"' in line: - first_token_line = line - break - - assert first_token_line - assert not release_second.is_set() - release_second.set() - rest = response.read().decode() - finally: - release_second.set() - if response is not None: - response.close() - head.stop() - tail.stop() - - assert '"content": " second"' in rest - assert "data: [DONE]" in rest - - -def test_current_requests_snapshot_while_generating(): - release_second = threading.Event() - head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True) - tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second)) - head_port = head.start() - tail_port = tail.start() - response = None - try: - payload = json.dumps({ - "model": "fake-model", - "messages": [{"role": "user", "content": "hello"}], - "stream": True, - "max_tokens": 2, - }).encode() - req = urllib.request.Request( - f"http://127.0.0.1:{head_port}/v1/chat/completions", - data=payload, - headers={ - "Content-Type": "application/json", - "X-Meshnet-Request-Id": "req-live-1", - "X-Meshnet-Route": json.dumps([ - {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22}, - ]), - }, - method="POST", - ) - response = urllib.request.urlopen(req, timeout=5) - deadline = time.time() + 2.0 - while time.time() < deadline: - live = head.current_requests - if live and live[0]["request_id"] == "req-live-1" and live[0]["tokens"] >= 1: - break - time.sleep(0.02) - assert head.current_requests - snap = head.current_requests[0] - assert snap["request_id"] == "req-live-1" - assert snap["tokens"] >= 1 - assert snap["tokens_per_sec"] >= 0 - assert snap["routing_complete"] is True - release_second.set() - response.read() - finally: - release_second.set() - if response is not None: - response.close() - head.stop() - tail.stop() - - assert head.current_requests == [] - - -def test_distributed_generating_log_includes_tps(capsys): - head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True) - tail = TorchNodeServer(backend=_FakePipelineTailBackend()) - head_port = head.start() - tail_port = tail.start() - try: - payload = json.dumps({ - "model": "fake-model", - "messages": [{"role": "user", "content": "hello"}], - "max_tokens": 1, - }).encode() - req = urllib.request.Request( - f"http://127.0.0.1:{head_port}/v1/chat/completions", - data=payload, - headers={ - "Content-Type": "application/json", - "X-Meshnet-Route": json.dumps([ - {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22}, - ]), - }, - method="POST", - ) - with urllib.request.urlopen(req, timeout=5) as resp: - json.loads(resp.read()) - finally: - head.stop() - tail.stop() - - out = capsys.readouterr().out - assert "generating step=1/1" in out - assert " tps=" in out - assert "generation complete tokens=1" in out - assert out.count("generating step=1/1") == 1 - - -def test_int_tensor_header_serializes_torch_tensors(): - torch = pytest.importorskip("torch") - - header = _int_tensor_header(torch.tensor([[1, 2, 3]], dtype=torch.long)) - - assert header.startswith("1,3:") - - -def test_decoder_attention_mask_is_causal_float_mask(): - torch = pytest.importorskip("torch") - - hidden_states = torch.zeros((1, 3, 8), dtype=torch.bfloat16) - mask = _decoder_attention_mask(torch.ones((1, 3), dtype=torch.long), hidden_states, torch) - - assert mask.shape == (1, 1, 3, 3) - assert mask.dtype == torch.bfloat16 - assert mask[0, 0, 0, 1] < 0 - assert mask[0, 0, 2, 0] == 0 - - -def test_call_layer_passes_rotary_position_embeddings(): - class NeedsPositionEmbeddings: - def __call__(self, hidden_states, **kwargs): - assert kwargs["position_embeddings"] == "rotary" - return hidden_states - - assert _call_layer( - NeedsPositionEmbeddings(), - "hidden", - attention_mask=None, - position_ids="positions", - position_embeddings="rotary", - ) == "hidden" - - -def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapshot(tmp_path): - snapshot_dir = tmp_path / "snapshot" - snapshot_dir.mkdir() - (snapshot_dir / "config.json").write_text("{}") - (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}') - - assert _should_partial_materialize_shard( - str(snapshot_dir), - 4, - 7, - total_layers_hint=40, - uses_quantized_weights=False, - ) is True - assert _should_partial_materialize_shard( - str(snapshot_dir), - 0, - 39, - total_layers_hint=40, - uses_quantized_weights=False, - ) is True - assert _should_partial_materialize_shard( - str(snapshot_dir), - 4, - 7, - total_layers_hint=40, - uses_quantized_weights=True, - ) is False - assert _should_partial_materialize_shard( - "repo/model", - 4, - 7, - total_layers_hint=40, - uses_quantized_weights=False, - ) is False - - -def test_checkpoint_tensor_name_remapped_for_text_only_causal_lm(): - class TextOnlyModel: - def __init__(self): - self.model = types.SimpleNamespace(layers=[]) - - model = TextOnlyModel() - assert _checkpoint_tensor_name_for_model( - model, - "model.language_model.layers.0.mlp.gate.weight", - ) == "model.layers.0.mlp.gate.weight" - assert _checkpoint_tensor_name_for_model( - model, - "model.language_model.embed_tokens.weight", - ) == "model.embed_tokens.weight" - - -def test_checkpoint_tensor_name_kept_for_multimodal_backbone(): - class MultimodalModel: - def __init__(self): - self.model = types.SimpleNamespace(language_model=types.SimpleNamespace()) - - model = MultimodalModel() - name = "model.language_model.layers.0.mlp.gate.weight" - assert _checkpoint_tensor_name_for_model(model, name) == name - - -def test_partial_snapshot_loader_remaps_language_model_checkpoint_keys(tmp_path): - snapshot_dir = tmp_path / "snapshot" - snapshot_dir.mkdir() - (snapshot_dir / "config.json").write_text(json.dumps({ - "text_config": {"num_hidden_layers": 3}, - })) - (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({ - "weight_map": { - "model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors", - } - })) - (snapshot_dir / "shard-2.safetensors").write_bytes(b"stub") - - class FakeModule: - def __init__(self): - self.to_calls = [] - - def to(self, device): - self.to_calls.append(device) - return self - - class FakeModel: - def __init__(self): - self.model = types.SimpleNamespace( - layers=[FakeModule(), FakeModule(), FakeModule()], - rotary_emb=FakeModule(), - ) - - def tie_weights(self): - pass - - class AutoConfigStub: - @staticmethod - def from_pretrained(model_id): - return types.SimpleNamespace( - text_config=types.SimpleNamespace(num_hidden_layers=3), - get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3), - ) - - class AutoModelStub: - @staticmethod - def from_config(cfg, torch_dtype=None): - return FakeModel() - - set_calls = [] - - def fake_set_tensor(module, tensor_name, device, value=None, dtype=None): - set_calls.append(tensor_name) - - class FakeSafeOpen: - def __init__(self, filename, framework, device): - self.filename = Path(filename).name - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def get_tensor(self, tensor_name): - return tensor_name - - class UnusedContext: - def __enter__(self): - return None - - def __exit__(self, exc_type, exc, tb): - return False - - _load_partial_model_from_snapshot( - AutoConfigStub, - AutoModelStub, - types.SimpleNamespace(), - str(snapshot_dir), - 1, - 1, - "bf16", - "cpu:0", - init_empty_weights_fn=UnusedContext, - set_tensor_fn=fake_set_tensor, - safe_open_fn=FakeSafeOpen, - ) - - assert set_calls == ["model.layers.1.self_attn.q_proj.weight"] - - -def test_partial_snapshot_loader_skips_tensors_absent_from_causal_lm(tmp_path): - # Multimodal/MTP checkpoints (Qwen3.5/3.6-MoE) carry mtp.* and model.visual.* - # tensors that the text-only CausalLM never builds — they must be skipped, - # not assigned (assignment raises AttributeError: 'mtp' / 'visual'). - snapshot_dir = tmp_path / "snapshot" - snapshot_dir.mkdir() - (snapshot_dir / "config.json").write_text(json.dumps({ - "text_config": {"num_hidden_layers": 3}, - })) - (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({ - "weight_map": { - "model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors", - "mtp.layers.1.input_layernorm.weight": "shard-2.safetensors", - "model.visual.blocks.1.attn.qkv.weight": "shard-2.safetensors", - } - })) - (snapshot_dir / "shard-2.safetensors").write_bytes(b"stub") - - class FakeModule: - def to(self, device): - return self - - class FakeModel: - def __init__(self): - self.model = types.SimpleNamespace( - layers=[FakeModule(), FakeModule(), FakeModule()], - rotary_emb=FakeModule(), - ) - - def tie_weights(self): - pass - - def state_dict(self): - return {"model.layers.1.self_attn.q_proj.weight": None} - - class AutoConfigStub: - @staticmethod - def from_pretrained(model_id): - return types.SimpleNamespace( - text_config=types.SimpleNamespace(num_hidden_layers=3), - get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3), - ) - - class AutoModelStub: - @staticmethod - def from_config(cfg, torch_dtype=None): - return FakeModel() - - set_calls = [] - - def fake_set_tensor(module, tensor_name, device, value=None, dtype=None): - set_calls.append(tensor_name) - - class FakeSafeOpen: - def __init__(self, filename, framework, device): - pass - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def get_tensor(self, tensor_name): - return tensor_name - - class UnusedContext: - def __enter__(self): - return None - - def __exit__(self, exc_type, exc, tb): - return False - - _load_partial_model_from_snapshot( - AutoConfigStub, - AutoModelStub, - types.SimpleNamespace(), - str(snapshot_dir), - 1, - 1, - "bf16", - "cpu:0", - init_empty_weights_fn=UnusedContext, - set_tensor_fn=fake_set_tensor, - safe_open_fn=FakeSafeOpen, - ) - - assert set_calls == ["model.layers.1.self_attn.q_proj.weight"] - - -def test_partial_snapshot_loader_materializes_only_assigned_tensors(tmp_path): - snapshot_dir = tmp_path / "snapshot" - snapshot_dir.mkdir() - (snapshot_dir / "config.json").write_text("{}") - (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({ - "weight_map": { - "model.embed_tokens.weight": "shard-1.safetensors", - "model.layers.0.self_attn.q_proj.weight": "shard-1.safetensors", - "model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors", - "model.layers.2.self_attn.q_proj.weight": "shard-3.safetensors", - "model.norm.weight": "shard-3.safetensors", - "lm_head.weight": "shard-3.safetensors", - } - })) - for rel in ("shard-1.safetensors", "shard-2.safetensors", "shard-3.safetensors"): - (snapshot_dir / rel).write_bytes(b"stub") - - class FakeModule: - def __init__(self, name): - self.name = name - self.to_calls = [] - - def to(self, device): - self.to_calls.append(device) - return self - - class FakeModel: - def __init__(self): - self.model = types.SimpleNamespace( - embed_tokens=FakeModule("embed"), - layers=[FakeModule("layer0"), FakeModule("layer1"), FakeModule("layer2")], - rotary_emb=FakeModule("rotary"), - norm=FakeModule("norm"), - ) - self.lm_head = FakeModule("lm_head") - self.tie_weights_called = 0 - - def tie_weights(self): - self.tie_weights_called += 1 - - class AutoConfigStub: - @staticmethod - def from_pretrained(model_id): - assert model_id == str(snapshot_dir) - return types.SimpleNamespace(num_hidden_layers=3) - - class AutoModelStub: - @staticmethod - def from_config(cfg, torch_dtype=None): - assert cfg.num_hidden_layers == 3 - assert torch_dtype == "bf16" - return FakeModel() - - class EmptyWeights: - def __init__(self): - self.entered = 0 - self.exited = 0 - - def __call__(self): - return self - - def __enter__(self): - self.entered += 1 - return None - - def __exit__(self, exc_type, exc, tb): - self.exited += 1 - return False - - init_empty_weights = EmptyWeights() - set_calls = [] - - def fake_set_tensor(module, tensor_name, device, value=None, dtype=None): - set_calls.append((tensor_name, device, value, dtype)) - - tensors = { - "shard-1.safetensors": { - "model.embed_tokens.weight": "embed", - "model.layers.0.self_attn.q_proj.weight": "layer0", - }, - "shard-2.safetensors": { - "model.layers.1.self_attn.q_proj.weight": "layer1", - }, - "shard-3.safetensors": { - "model.layers.2.self_attn.q_proj.weight": "layer2", - "model.norm.weight": "norm", - "lm_head.weight": "lm_head", - }, - } - - class FakeSafeOpen: - def __init__(self, filename, framework, device): - assert framework == "pt" - assert device == "cpu" - self.filename = Path(filename).name - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def get_tensor(self, tensor_name): - return tensors[self.filename][tensor_name] - - model = _load_partial_model_from_snapshot( - AutoConfigStub, - AutoModelStub, - types.SimpleNamespace(), - str(snapshot_dir), - 1, - 1, - "bf16", - "cpu:0", - init_empty_weights_fn=init_empty_weights, - set_tensor_fn=fake_set_tensor, - safe_open_fn=FakeSafeOpen, - ) - - assert init_empty_weights.entered == 1 - assert init_empty_weights.exited == 1 - assert model.tie_weights_called == 1 - assert [call[0] for call in set_calls] == ["model.layers.1.self_attn.q_proj.weight"] - assert model.model.layers[1].to_calls == ["cpu:0"] - assert model.model.layers[0].to_calls == [] - assert model.model.layers[2].to_calls == [] - assert model.model.embed_tokens.to_calls == [] - assert model.model.norm.to_calls == [] - assert model.lm_head.to_calls == [] - assert model.model.rotary_emb.to_calls == ["cpu:0"] - - -def test_partial_snapshot_loader_requires_known_layer_count(tmp_path): - snapshot_dir = tmp_path / "snapshot" - snapshot_dir.mkdir() - (snapshot_dir / "config.json").write_text("{}") - (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({ - "weight_map": {"model.layers.0.self_attn.q_proj.weight": "shard.safetensors"} - })) - (snapshot_dir / "shard.safetensors").write_bytes(b"stub") - - class AutoConfigStub: - @staticmethod - def from_pretrained(model_id): - return types.SimpleNamespace() - - class AutoModelStub: - @staticmethod - def from_config(cfg, torch_dtype=None): - raise AssertionError("from_config should not run without a known layer count") - - class UnusedContext: - def __enter__(self): - return None - - def __exit__(self, exc_type, exc, tb): - return False - - with pytest.raises(PartialModelLoadUnsupported, match="num_hidden_layers"): - _load_partial_model_from_snapshot( - AutoConfigStub, - AutoModelStub, - types.SimpleNamespace(), - str(snapshot_dir), - 0, - 0, - "bf16", - "cpu:0", - init_empty_weights_fn=lambda: UnusedContext(), - set_tensor_fn=lambda *args, **kwargs: None, - safe_open_fn=lambda *args, **kwargs: None, - ) - - -def test_torch_model_shard_prefers_partial_loader_for_local_snapshot(tmp_path, monkeypatch): - import meshnet_node.model_backend as backend - - snapshot_dir = tmp_path / "snapshot" - snapshot_dir.mkdir() - (snapshot_dir / "config.json").write_text("{}") - (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}') - - class FakeModel: - def __init__(self): - self.model = types.SimpleNamespace( - layers=[object(), object(), object()], - embed_tokens=object(), - ) - self.config = types.SimpleNamespace(hidden_size=8) - self.eval_called = 0 - - def eval(self): - self.eval_called += 1 - - fake_model = FakeModel() - partial_calls = [] - - class AutoConfigStub: - @staticmethod - def from_pretrained(model_id, cache_dir=None): - return types.SimpleNamespace(num_hidden_layers=3, text_config=types.SimpleNamespace(dtype="torch.bfloat16")) - - class AutoModelStub: - @staticmethod - def from_pretrained(*args, **kwargs): - raise AssertionError("full model load should not run for partial local shards") - - class AutoTokenizerStub: - @staticmethod - def from_pretrained(model_id, cache_dir=None): - assert model_id == str(snapshot_dir) - return types.SimpleNamespace() - - monkeypatch.setitem( - sys.modules, - "torch", - types.SimpleNamespace( - cuda=types.SimpleNamespace(is_available=lambda: False), - device=lambda value: value, - bfloat16="bf16", - ), - ) - monkeypatch.setitem( - sys.modules, - "transformers", - types.SimpleNamespace( - AutoConfig=AutoConfigStub, - AutoModelForCausalLM=AutoModelStub, - AutoTokenizer=AutoTokenizerStub, - ), - ) - monkeypatch.setattr( - backend, - "_load_partial_model_from_snapshot", - lambda *args, **kwargs: partial_calls.append((args, kwargs)) or fake_model, - ) - - shard = TorchModelShard( - "repo/model", - 1, - 1, - quantization="auto", - cache_dir=snapshot_dir, - ) - - assert len(partial_calls) == 1 - assert shard.model is fake_model - assert fake_model.eval_called == 1 - assert shard.total_layers == 3 - assert shard.is_head is False - assert shard.is_tail is False - - -@pytest.mark.integration -def test_two_node_gpt2_completion_is_deterministic(): - if os.environ.get("CI"): - pytest.skip("GPT-2 integration test is skipped in CI") - torch = pytest.importorskip("torch") - pytest.importorskip("transformers") - pytest.importorskip("safetensors") - pytest.importorskip("accelerate") - pytest.importorskip("bitsandbytes") - if not torch.cuda.is_available(): - pytest.skip("GPT-2 integration test requires a CUDA GPU") - - head = TorchNodeServer( - model_id="openai-community/gpt2", - shard_start=0, - shard_end=6, - quantization="bfloat16", - ) - tail = TorchNodeServer( - model_id="openai-community/gpt2", - shard_start=6, - shard_end=12, - quantization="bfloat16", - ) - head_port = head.start() - tail_port = tail.start() - try: - prompt_req = urllib.request.Request( - f"http://127.0.0.1:{head_port}/forward", - data=json.dumps({"prompt": "The capital of France is"}).encode(), - headers={"Content-Type": "application/json"}, - method="POST", - ) - with urllib.request.urlopen(prompt_req, timeout=60) as resp: - activation = resp.read() - head_headers = resp.headers - - tail_req = urllib.request.Request( - f"http://127.0.0.1:{tail_port}/forward", - data=activation, - headers={ - "Content-Type": "application/octet-stream", - "X-Meshnet-Shape": head_headers["X-Meshnet-Shape"], - "X-Meshnet-Dtype": head_headers["X-Meshnet-Dtype"], - "X-Meshnet-Session": "gpt2-session", - "X-Meshnet-Chunk-Index": "0", - "X-Meshnet-Chunk-Total": "1", - "X-Meshnet-Hop-Index": "1", - "X-Meshnet-Attn-Mask": head_headers["X-Meshnet-Attn-Mask"], - "X-Meshnet-Position-Ids": head_headers["X-Meshnet-Position-Ids"], - }, - method="POST", - ) - with urllib.request.urlopen(tail_req, timeout=60) as resp: - body = json.loads(resp.read()) - - assert body["text"].strip() - assert body["text"] == " Paris" - finally: - head.stop() - tail.stop() +"""US-012 tests for the real PyTorch node backend.""" + +import json +import os +from pathlib import Path +import sys +import threading +import time +import types +import urllib.request + +import pytest + +from meshnet_node.model_backend import ( + InsufficientVRAMError, + PartialModelLoadUnsupported, + TensorPayload, + TorchModelShard, + _call_layer, + _checkpoint_tensor_name_for_model, + _load_partial_model_from_snapshot, + _should_partial_materialize_shard, + _decoder_attention_mask, + _int_tensor_header, + build_quantization_config, + validate_quantization, +) +from meshnet_node.torch_server import TorchNodeServer + + +class _FakeBackend: + model_id = "fake-model" + total_layers = 12 + is_head = True + is_tail = False + + def encode_prompt(self, prompt: str) -> TensorPayload: + assert prompt == "The capital of France is" + return TensorPayload( + body=b"\x00" * (1 * 6 * 8 * 2), + shape=[1, 6, 8], + attention_mask_header=None, + position_ids_header=None, + ) + + def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): + assert shape == [1, 6, 8] + return TensorPayload( + body=body, + shape=shape, + attention_mask_header=attention_mask_header, + position_ids_header=position_ids_header, + ) + + +class _FakeTailBackend(_FakeBackend): + is_head = False + is_tail = True + + def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): + assert len(body) == 1 * 6 * 8 * 2 + return " Paris" + + +class _FakeFullBackend(_FakeBackend): + is_head = True + is_tail = True + + def generate_text( + self, + messages: list[dict], + max_new_tokens: int = 16, + temperature: float = 1.0, + top_p: float = 1.0, + ) -> str: + assert messages == [{"role": "user", "content": "What is 7 times 8?"}] + assert max_new_tokens == 7 + assert temperature == 1.0 + assert top_p == 1.0 + return "56" + + def count_prompt_tokens(self, messages: list[dict]) -> int: + assert messages == [{"role": "user", "content": "What is 7 times 8?"}] + return 8 + + def count_text_tokens(self, text: str) -> int: + assert text == "56" + return 1 + + +class _FakeChatTokenizer: + eos_token = "" + + def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False): + assert add_generation_prompt is True + assert tokenize is False + return "debug prompt" + + +class _FakePipelineHeadBackend(_FakeBackend): + tokenizer = _FakeChatTokenizer() + + def encode_prompt(self, prompt: str) -> TensorPayload: + assert prompt.startswith("debug prompt") + return TensorPayload( + body=b"\x00" * (1 * 6 * 8 * 2), + shape=[1, 6, 8], + attention_mask_header=None, + position_ids_header=None, + ) + + +class _FakePipelineTailBackend(_FakeTailBackend): + def __init__(self) -> None: + self.start_layers: list[int | None] = [] + + def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): + self.start_layers.append(start_layer) + assert len(body) == 1 * 6 * 8 * 2 + return " token" + + +class _BlockingStreamingTailBackend(_FakeTailBackend): + def __init__(self, second_token_release: threading.Event) -> None: + self._release = second_token_release + self.calls = 0 + + def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None): + self.calls += 1 + if self.calls == 1: + return " first" + self._release.wait(timeout=3.0) + return " second" + + +def test_quantization_flag_validation(): + assert validate_quantization("bfloat16") == "bfloat16" + assert validate_quantization("int8") == "int8" + assert validate_quantization("nf4") == "nf4" + with pytest.raises(ValueError, match="quantization"): + validate_quantization("float32") + + +def test_node_package_declares_torch_dependency(): + pyproject = Path("packages/node/pyproject.toml").read_text(encoding="utf-8") + + assert '"torch>=' in pyproject + + +def test_bitsandbytes_configs_are_created_lazily(monkeypatch): + calls = [] + + class FakeBitsAndBytesConfig: + def __init__(self, **kwargs): + calls.append(kwargs) + + monkeypatch.setitem(sys.modules, "torch", types.SimpleNamespace(bfloat16="bf16")) + monkeypatch.setitem( + sys.modules, + "transformers", + types.SimpleNamespace(BitsAndBytesConfig=FakeBitsAndBytesConfig), + ) + + assert build_quantization_config("bfloat16") is None + build_quantization_config("int8") + build_quantization_config("nf4") + + assert calls == [ + {"load_in_8bit": True}, + { + "load_in_4bit": True, + "bnb_4bit_quant_type": "nf4", + "bnb_4bit_compute_dtype": "bf16", + }, + ] + + +def test_head_forward_accepts_text_prompt_and_returns_bfloat16_activations(): + node = TorchNodeServer(backend=_FakeBackend()) + port = node.start() + try: + payload = json.dumps({"prompt": "The capital of France is"}).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{port}/forward", + data=payload, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5) as resp: + body = resp.read() + headers = {key.lower(): value for key, value in resp.headers.items()} + + assert len(body) == 1 * 6 * 8 * 2 + assert headers["x-meshnet-shape"] == "1,6,8" + assert headers["x-meshnet-dtype"] == "bfloat16" + assert headers["x-meshnet-wire"] == "2" + finally: + node.stop() + + +def test_tail_forward_returns_text_completion_from_binary_activations(): + node = TorchNodeServer(backend=_FakeTailBackend()) + port = node.start() + try: + req = urllib.request.Request( + f"http://127.0.0.1:{port}/forward", + data=b"\x00" * (1 * 6 * 8 * 2), + headers={ + "Content-Type": "application/octet-stream", + "X-Meshnet-Shape": "1,6,8", + "X-Meshnet-Dtype": "bfloat16", + "X-Meshnet-Session": "session-1", + "X-Meshnet-Chunk-Index": "0", + "X-Meshnet-Chunk-Total": "1", + "X-Meshnet-Hop-Index": "1", + }, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5) as resp: + body = json.loads(resp.read()) + + assert body == {"text": " Paris"} + assert node.received_activations + assert node.forward_chunk_count == 1 + finally: + node.stop() + + +def test_full_model_chat_completion_uses_generation_not_single_token_decode(capsys): + node = TorchNodeServer(backend=_FakeFullBackend()) + port = node.start() + try: + payload = json.dumps({ + "model": "fake-model", + "messages": [{"role": "user", "content": "What is 7 times 8?"}], + "max_tokens": 7, + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{port}/v1/chat/completions", + data=payload, + headers={ + "Content-Type": "application/json", + "X-Meshnet-Request-Id": "req-test-123", + }, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5) as resp: + body = json.loads(resp.read()) + + assert body["choices"][0]["message"]["content"] == "56" + assert body["usage"] == {"prompt_tokens": 8, "completion_tokens": 1, "total_tokens": 9} + finally: + node.stop() + + out = capsys.readouterr().out + assert " [node] processing chat model='fake-model' stream=False max_tokens=7 request_id=req-test-123" in out + assert " [node] chat complete tokens=1 elapsed_s=" in out + + +def test_pipeline_hop_logs_are_suppressed_without_debug(capsys): + tail_backend = _FakePipelineTailBackend() + head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True) + tail = TorchNodeServer(backend=tail_backend) + head_port = head.start() + tail_port = tail.start() + try: + payload = json.dumps({ + "model": "fake-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 1, + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{head_port}/v1/chat/completions", + data=payload, + headers={ + "Content-Type": "application/json", + "X-Meshnet-Route": json.dumps([ + {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22}, + ]), + }, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5) as resp: + body = json.loads(resp.read()) + finally: + head.stop() + tail.stop() + + out = capsys.readouterr().out + assert body["choices"][0]["message"]["content"] == " token" + assert tail_backend.start_layers == [22] + assert "pipeline hop 0:" not in out + assert "pipeline hop 0 returned text" not in out + + +def test_pipeline_hop_logs_are_enabled_with_debug(capsys): + head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True, debug=True) + tail = TorchNodeServer(backend=_FakePipelineTailBackend()) + head_port = head.start() + tail_port = tail.start() + try: + payload = json.dumps({ + "model": "fake-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 1, + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{head_port}/v1/chat/completions", + data=payload, + headers={ + "Content-Type": "application/json", + "X-Meshnet-Route": json.dumps([ + {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22}, + ]), + }, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5) as resp: + json.loads(resp.read()) + finally: + head.stop() + tail.stop() + + out = capsys.readouterr().out + assert f" [node] pipeline hop 0: http://127.0.0.1:{tail_port} start_layer=22" in out + assert " [node] pipeline hop 0 returned text=' token'" in out + + +def test_split_shard_chat_streams_each_generated_token_incrementally(): + release_second = threading.Event() + head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True) + tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second)) + head_port = head.start() + tail_port = tail.start() + response = None + try: + payload = json.dumps({ + "model": "fake-model", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + "max_tokens": 2, + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{head_port}/v1/chat/completions", + data=payload, + headers={ + "Content-Type": "application/json", + "X-Meshnet-Route": json.dumps([ + {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22}, + ]), + }, + method="POST", + ) + response = urllib.request.urlopen(req, timeout=5) + + first_token_line = "" + deadline = time.time() + 2.0 + while time.time() < deadline: + line = response.readline().decode() + if '"content": " first"' in line: + first_token_line = line + break + + assert first_token_line + assert not release_second.is_set() + release_second.set() + rest = response.read().decode() + finally: + release_second.set() + if response is not None: + response.close() + head.stop() + tail.stop() + + assert '"content": " second"' in rest + assert "data: [DONE]" in rest + + +def test_current_requests_snapshot_while_generating(): + release_second = threading.Event() + head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True) + tail = TorchNodeServer(backend=_BlockingStreamingTailBackend(release_second)) + head_port = head.start() + tail_port = tail.start() + response = None + try: + payload = json.dumps({ + "model": "fake-model", + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + "max_tokens": 2, + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{head_port}/v1/chat/completions", + data=payload, + headers={ + "Content-Type": "application/json", + "X-Meshnet-Request-Id": "req-live-1", + "X-Meshnet-Route": json.dumps([ + {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22}, + ]), + }, + method="POST", + ) + response = urllib.request.urlopen(req, timeout=5) + deadline = time.time() + 2.0 + while time.time() < deadline: + live = head.current_requests + if live and live[0]["request_id"] == "req-live-1" and live[0]["tokens"] >= 1: + break + time.sleep(0.02) + assert head.current_requests + snap = head.current_requests[0] + assert snap["request_id"] == "req-live-1" + assert snap["tokens"] >= 1 + assert snap["tokens_per_sec"] >= 0 + assert snap["routing_complete"] is True + release_second.set() + response.read() + finally: + release_second.set() + if response is not None: + response.close() + head.stop() + tail.stop() + + assert head.current_requests == [] + + +def test_distributed_generating_log_includes_tps(capsys): + head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True) + tail = TorchNodeServer(backend=_FakePipelineTailBackend()) + head_port = head.start() + tail_port = tail.start() + try: + payload = json.dumps({ + "model": "fake-model", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 1, + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{head_port}/v1/chat/completions", + data=payload, + headers={ + "Content-Type": "application/json", + "X-Meshnet-Route": json.dumps([ + {"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22}, + ]), + }, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5) as resp: + json.loads(resp.read()) + finally: + head.stop() + tail.stop() + + out = capsys.readouterr().out + assert "generating step=1/1" in out + assert " tps=" in out + assert "generation complete tokens=1" in out + assert out.count("generating step=1/1") == 1 + + +def test_int_tensor_header_serializes_torch_tensors(): + torch = pytest.importorskip("torch") + + header = _int_tensor_header(torch.tensor([[1, 2, 3]], dtype=torch.long)) + + assert header.startswith("1,3:") + + +def test_decoder_attention_mask_is_causal_float_mask(): + torch = pytest.importorskip("torch") + + hidden_states = torch.zeros((1, 3, 8), dtype=torch.bfloat16) + mask = _decoder_attention_mask(torch.ones((1, 3), dtype=torch.long), hidden_states, torch) + + assert mask.shape == (1, 1, 3, 3) + assert mask.dtype == torch.bfloat16 + assert mask[0, 0, 0, 1] < 0 + assert mask[0, 0, 2, 0] == 0 + + +def test_call_layer_passes_rotary_position_embeddings(): + class NeedsPositionEmbeddings: + def __call__(self, hidden_states, **kwargs): + assert kwargs["position_embeddings"] == "rotary" + return hidden_states + + assert _call_layer( + NeedsPositionEmbeddings(), + "hidden", + attention_mask=None, + position_ids="positions", + position_embeddings="rotary", + ) == "hidden" + + +def test_partial_materialize_guard_requires_local_non_full_non_quantized_snapshot(tmp_path): + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "config.json").write_text("{}") + (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}') + + assert _should_partial_materialize_shard( + str(snapshot_dir), + 4, + 7, + total_layers_hint=40, + uses_quantized_weights=False, + ) is True + assert _should_partial_materialize_shard( + str(snapshot_dir), + 0, + 39, + total_layers_hint=40, + uses_quantized_weights=False, + ) is True + assert _should_partial_materialize_shard( + str(snapshot_dir), + 4, + 7, + total_layers_hint=40, + uses_quantized_weights=True, + ) is False + assert _should_partial_materialize_shard( + "repo/model", + 4, + 7, + total_layers_hint=40, + uses_quantized_weights=False, + ) is False + + +def test_checkpoint_tensor_name_remapped_for_text_only_causal_lm(): + class TextOnlyModel: + def __init__(self): + self.model = types.SimpleNamespace(layers=[]) + + model = TextOnlyModel() + assert _checkpoint_tensor_name_for_model( + model, + "model.language_model.layers.0.mlp.gate.weight", + ) == "model.layers.0.mlp.gate.weight" + assert _checkpoint_tensor_name_for_model( + model, + "model.language_model.embed_tokens.weight", + ) == "model.embed_tokens.weight" + + +def test_checkpoint_tensor_name_kept_for_multimodal_backbone(): + class MultimodalModel: + def __init__(self): + self.model = types.SimpleNamespace(language_model=types.SimpleNamespace()) + + model = MultimodalModel() + name = "model.language_model.layers.0.mlp.gate.weight" + assert _checkpoint_tensor_name_for_model(model, name) == name + + +def test_partial_snapshot_loader_remaps_language_model_checkpoint_keys(tmp_path): + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "config.json").write_text(json.dumps({ + "text_config": {"num_hidden_layers": 3}, + })) + (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({ + "weight_map": { + "model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors", + } + })) + (snapshot_dir / "shard-2.safetensors").write_bytes(b"stub") + + class FakeModule: + def __init__(self): + self.to_calls = [] + + def to(self, device): + self.to_calls.append(device) + return self + + class FakeModel: + def __init__(self): + self.model = types.SimpleNamespace( + layers=[FakeModule(), FakeModule(), FakeModule()], + rotary_emb=FakeModule(), + ) + + def tie_weights(self): + pass + + class AutoConfigStub: + @staticmethod + def from_pretrained(model_id): + return types.SimpleNamespace( + text_config=types.SimpleNamespace(num_hidden_layers=3), + get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3), + ) + + class AutoModelStub: + @staticmethod + def from_config(cfg, torch_dtype=None): + return FakeModel() + + set_calls = [] + + def fake_set_tensor(module, tensor_name, device, value=None, dtype=None): + set_calls.append(tensor_name) + + class FakeSafeOpen: + def __init__(self, filename, framework, device): + self.filename = Path(filename).name + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def get_tensor(self, tensor_name): + return tensor_name + + class UnusedContext: + def __enter__(self): + return None + + def __exit__(self, exc_type, exc, tb): + return False + + _load_partial_model_from_snapshot( + AutoConfigStub, + AutoModelStub, + types.SimpleNamespace(), + str(snapshot_dir), + 1, + 1, + "bf16", + "cpu:0", + init_empty_weights_fn=UnusedContext, + set_tensor_fn=fake_set_tensor, + safe_open_fn=FakeSafeOpen, + ) + + assert set_calls == ["model.layers.1.self_attn.q_proj.weight"] + + +def test_partial_snapshot_loader_skips_tensors_absent_from_causal_lm(tmp_path): + # Multimodal/MTP checkpoints (Qwen3.5/3.6-MoE) carry mtp.* and model.visual.* + # tensors that the text-only CausalLM never builds — they must be skipped, + # not assigned (assignment raises AttributeError: 'mtp' / 'visual'). + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "config.json").write_text(json.dumps({ + "text_config": {"num_hidden_layers": 3}, + })) + (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({ + "weight_map": { + "model.language_model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors", + "mtp.layers.1.input_layernorm.weight": "shard-2.safetensors", + "model.visual.blocks.1.attn.qkv.weight": "shard-2.safetensors", + } + })) + (snapshot_dir / "shard-2.safetensors").write_bytes(b"stub") + + class FakeModule: + def to(self, device): + return self + + class FakeModel: + def __init__(self): + self.model = types.SimpleNamespace( + layers=[FakeModule(), FakeModule(), FakeModule()], + rotary_emb=FakeModule(), + ) + + def tie_weights(self): + pass + + def state_dict(self): + return {"model.layers.1.self_attn.q_proj.weight": None} + + class AutoConfigStub: + @staticmethod + def from_pretrained(model_id): + return types.SimpleNamespace( + text_config=types.SimpleNamespace(num_hidden_layers=3), + get_text_config=lambda: types.SimpleNamespace(num_hidden_layers=3), + ) + + class AutoModelStub: + @staticmethod + def from_config(cfg, torch_dtype=None): + return FakeModel() + + set_calls = [] + + def fake_set_tensor(module, tensor_name, device, value=None, dtype=None): + set_calls.append(tensor_name) + + class FakeSafeOpen: + def __init__(self, filename, framework, device): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def get_tensor(self, tensor_name): + return tensor_name + + class UnusedContext: + def __enter__(self): + return None + + def __exit__(self, exc_type, exc, tb): + return False + + _load_partial_model_from_snapshot( + AutoConfigStub, + AutoModelStub, + types.SimpleNamespace(), + str(snapshot_dir), + 1, + 1, + "bf16", + "cpu:0", + init_empty_weights_fn=UnusedContext, + set_tensor_fn=fake_set_tensor, + safe_open_fn=FakeSafeOpen, + ) + + assert set_calls == ["model.layers.1.self_attn.q_proj.weight"] + + +def test_partial_snapshot_loader_materializes_only_assigned_tensors(tmp_path): + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "config.json").write_text("{}") + (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({ + "weight_map": { + "model.embed_tokens.weight": "shard-1.safetensors", + "model.layers.0.self_attn.q_proj.weight": "shard-1.safetensors", + "model.layers.1.self_attn.q_proj.weight": "shard-2.safetensors", + "model.layers.2.self_attn.q_proj.weight": "shard-3.safetensors", + "model.norm.weight": "shard-3.safetensors", + "lm_head.weight": "shard-3.safetensors", + } + })) + for rel in ("shard-1.safetensors", "shard-2.safetensors", "shard-3.safetensors"): + (snapshot_dir / rel).write_bytes(b"stub") + + class FakeModule: + def __init__(self, name): + self.name = name + self.to_calls = [] + + def to(self, device): + self.to_calls.append(device) + return self + + class FakeModel: + def __init__(self): + self.model = types.SimpleNamespace( + embed_tokens=FakeModule("embed"), + layers=[FakeModule("layer0"), FakeModule("layer1"), FakeModule("layer2")], + rotary_emb=FakeModule("rotary"), + norm=FakeModule("norm"), + ) + self.lm_head = FakeModule("lm_head") + self.tie_weights_called = 0 + + def tie_weights(self): + self.tie_weights_called += 1 + + class AutoConfigStub: + @staticmethod + def from_pretrained(model_id): + assert model_id == str(snapshot_dir) + return types.SimpleNamespace(num_hidden_layers=3) + + class AutoModelStub: + @staticmethod + def from_config(cfg, torch_dtype=None): + assert cfg.num_hidden_layers == 3 + assert torch_dtype == "bf16" + return FakeModel() + + class EmptyWeights: + def __init__(self): + self.entered = 0 + self.exited = 0 + + def __call__(self): + return self + + def __enter__(self): + self.entered += 1 + return None + + def __exit__(self, exc_type, exc, tb): + self.exited += 1 + return False + + init_empty_weights = EmptyWeights() + set_calls = [] + + def fake_set_tensor(module, tensor_name, device, value=None, dtype=None): + set_calls.append((tensor_name, device, value, dtype)) + + tensors = { + "shard-1.safetensors": { + "model.embed_tokens.weight": "embed", + "model.layers.0.self_attn.q_proj.weight": "layer0", + }, + "shard-2.safetensors": { + "model.layers.1.self_attn.q_proj.weight": "layer1", + }, + "shard-3.safetensors": { + "model.layers.2.self_attn.q_proj.weight": "layer2", + "model.norm.weight": "norm", + "lm_head.weight": "lm_head", + }, + } + + class FakeSafeOpen: + def __init__(self, filename, framework, device): + assert framework == "pt" + assert device == "cpu" + self.filename = Path(filename).name + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def get_tensor(self, tensor_name): + return tensors[self.filename][tensor_name] + + model = _load_partial_model_from_snapshot( + AutoConfigStub, + AutoModelStub, + types.SimpleNamespace(), + str(snapshot_dir), + 1, + 1, + "bf16", + "cpu:0", + init_empty_weights_fn=init_empty_weights, + set_tensor_fn=fake_set_tensor, + safe_open_fn=FakeSafeOpen, + ) + + assert init_empty_weights.entered == 1 + assert init_empty_weights.exited == 1 + assert model.tie_weights_called == 1 + assert [call[0] for call in set_calls] == ["model.layers.1.self_attn.q_proj.weight"] + assert model.model.layers[1].to_calls == ["cpu:0"] + assert model.model.layers[0].to_calls == [] + assert model.model.layers[2].to_calls == [] + assert model.model.embed_tokens.to_calls == [] + assert model.model.norm.to_calls == [] + assert model.lm_head.to_calls == [] + assert model.model.rotary_emb.to_calls == ["cpu:0"] + + +def test_partial_snapshot_loader_requires_known_layer_count(tmp_path): + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "config.json").write_text("{}") + (snapshot_dir / "model.safetensors.index.json").write_text(json.dumps({ + "weight_map": {"model.layers.0.self_attn.q_proj.weight": "shard.safetensors"} + })) + (snapshot_dir / "shard.safetensors").write_bytes(b"stub") + + class AutoConfigStub: + @staticmethod + def from_pretrained(model_id): + return types.SimpleNamespace() + + class AutoModelStub: + @staticmethod + def from_config(cfg, torch_dtype=None): + raise AssertionError("from_config should not run without a known layer count") + + class UnusedContext: + def __enter__(self): + return None + + def __exit__(self, exc_type, exc, tb): + return False + + with pytest.raises(PartialModelLoadUnsupported, match="num_hidden_layers"): + _load_partial_model_from_snapshot( + AutoConfigStub, + AutoModelStub, + types.SimpleNamespace(), + str(snapshot_dir), + 0, + 0, + "bf16", + "cpu:0", + init_empty_weights_fn=lambda: UnusedContext(), + set_tensor_fn=lambda *args, **kwargs: None, + safe_open_fn=lambda *args, **kwargs: None, + ) + + +def test_torch_model_shard_prefers_partial_loader_for_local_snapshot(tmp_path, monkeypatch): + import meshnet_node.model_backend as backend + + snapshot_dir = tmp_path / "snapshot" + snapshot_dir.mkdir() + (snapshot_dir / "config.json").write_text("{}") + (snapshot_dir / "model.safetensors.index.json").write_text('{"weight_map": {}}') + + class FakeModel: + def __init__(self): + self.model = types.SimpleNamespace( + layers=[object(), object(), object()], + embed_tokens=object(), + ) + self.config = types.SimpleNamespace(hidden_size=8) + self.eval_called = 0 + + def eval(self): + self.eval_called += 1 + + fake_model = FakeModel() + partial_calls = [] + + class AutoConfigStub: + @staticmethod + def from_pretrained(model_id, cache_dir=None): + return types.SimpleNamespace(num_hidden_layers=3, text_config=types.SimpleNamespace(dtype="torch.bfloat16")) + + class AutoModelStub: + @staticmethod + def from_pretrained(*args, **kwargs): + raise AssertionError("full model load should not run for partial local shards") + + class AutoTokenizerStub: + @staticmethod + def from_pretrained(model_id, cache_dir=None): + assert model_id == str(snapshot_dir) + return types.SimpleNamespace() + + monkeypatch.setitem( + sys.modules, + "torch", + types.SimpleNamespace( + cuda=types.SimpleNamespace(is_available=lambda: False), + device=lambda value: value, + bfloat16="bf16", + ), + ) + monkeypatch.setitem( + sys.modules, + "transformers", + types.SimpleNamespace( + AutoConfig=AutoConfigStub, + AutoModelForCausalLM=AutoModelStub, + AutoTokenizer=AutoTokenizerStub, + ), + ) + monkeypatch.setattr( + backend, + "_load_partial_model_from_snapshot", + lambda *args, **kwargs: partial_calls.append((args, kwargs)) or fake_model, + ) + + shard = TorchModelShard( + "repo/model", + 1, + 1, + quantization="auto", + cache_dir=snapshot_dir, + ) + + assert len(partial_calls) == 1 + assert shard.model is fake_model + assert fake_model.eval_called == 1 + assert shard.total_layers == 3 + assert shard.is_head is False + assert shard.is_tail is False + + +@pytest.mark.integration +def test_two_node_gpt2_completion_is_deterministic(): + if os.environ.get("CI"): + pytest.skip("GPT-2 integration test is skipped in CI") + torch = pytest.importorskip("torch") + pytest.importorskip("transformers") + pytest.importorskip("safetensors") + pytest.importorskip("accelerate") + pytest.importorskip("bitsandbytes") + if not torch.cuda.is_available(): + pytest.skip("GPT-2 integration test requires a CUDA GPU") + + head = TorchNodeServer( + model_id="openai-community/gpt2", + shard_start=0, + shard_end=6, + quantization="bfloat16", + ) + tail = TorchNodeServer( + model_id="openai-community/gpt2", + shard_start=6, + shard_end=12, + quantization="bfloat16", + ) + head_port = head.start() + tail_port = tail.start() + try: + prompt_req = urllib.request.Request( + f"http://127.0.0.1:{head_port}/forward", + data=json.dumps({"prompt": "The capital of France is"}).encode(), + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(prompt_req, timeout=60) as resp: + activation = resp.read() + head_headers = resp.headers + + tail_req = urllib.request.Request( + f"http://127.0.0.1:{tail_port}/forward", + data=activation, + headers={ + "Content-Type": "application/octet-stream", + "X-Meshnet-Shape": head_headers["X-Meshnet-Shape"], + "X-Meshnet-Dtype": head_headers["X-Meshnet-Dtype"], + "X-Meshnet-Session": "gpt2-session", + "X-Meshnet-Chunk-Index": "0", + "X-Meshnet-Chunk-Total": "1", + "X-Meshnet-Hop-Index": "1", + "X-Meshnet-Attn-Mask": head_headers["X-Meshnet-Attn-Mask"], + "X-Meshnet-Position-Ids": head_headers["X-Meshnet-Position-Ids"], + }, + method="POST", + ) + with urllib.request.urlopen(tail_req, timeout=60) as resp: + body = json.loads(resp.read()) + + assert body["text"].strip() + assert body["text"] == " Paris" + finally: + head.stop() + tail.stop()