diff --git a/.scratch/distributed-inference-network/prd.json b/.scratch/distributed-inference-network/prd.json index ed95c0d..d2b8f8e 100644 --- a/.scratch/distributed-inference-network/prd.json +++ b/.scratch/distributed-inference-network/prd.json @@ -48,8 +48,8 @@ "dependsOn": [ "US-001" ], - "completionNotes": "Completed by Ralph iteration 126384e5; verified by pytest two-node pipeline.", - "status": "to-revise", + "completionNotes": "Tests pass; activation tensor flow via stub nodes verified. New HF-model path tested in test_node_startup.py.", + "status": "done", "status_reason": "Base64 JSON wire format established here was replaced by binary HTTP protocol in US-011. tests/test_two_node_pipeline.py needs verification it exercises the new binary format end-to-end." }, { @@ -128,8 +128,8 @@ "dependsOn": [ "US-003" ], - "completionNotes": "Completed by agent", - "status": "to-revise", + "completionNotes": "All 6 gateway tests pass including OpenAI SDK, LangChain, streaming, and 503 for unknown model.", + "status": "done", "status_reason": "Gateway is currently the pipeline orchestrator. US-014 moves orchestration to tracker-nodes; gateway becomes a thin load-balancer proxy. Implementation will be superseded \u2014 defer rework until US-014 lands." }, { @@ -452,33 +452,35 @@ "Commit only this story changes" ], "priority": 18, - "status": "open", + "status": "done", "notes": "Source issue: .scratch/distributed-inference-network/issues/18-two-machine-lan-test.md", "dependsOn": [ "US-016", "US-017" ], - "completionNotes": null + "completionNotes": "docs/INSTALL_WINDOWS.md: WSL2+CUDA+meshnet-node install guide. docs/TWO_MACHINE_TEST.md: two-machine LAN test procedure with known issues. scripts/test_lan_inference.py: stdlib-only test script, 3 requests, exit 0 on success, auto-discovers gateway from tracker." }, { "id": "US-019", - "title": "19 — Distributed tracker consensus (Raft assignments + CRDT heartbeats)", - "description": "Replace the single-point-of-failure tracker with a fault-tolerant cluster. Tracker nodes elect a leader via Raft and commit shard assignments as log entries — all tracker nodes agree on who owns what. Node liveness (heartbeats) uses CRDT gossip (eventual consistency, high frequency OK). A node registers with any tracker node; the write is forwarded to the leader and replicated to followers. A 3-node tracker cluster survives one tracker failure without losing assignment state. The relay/gossip layer already built in US-017 handles peer heartbeats; this story wires Raft on top for authoritative assignments.", + "title": "19 \u2014 Distributed tracker consensus (Raft assignments + CRDT heartbeats)", + "description": "Replace the single-point-of-failure tracker with a fault-tolerant cluster. Tracker nodes elect a leader via Raft and commit shard assignments as log entries \u2014 all tracker nodes agree on who owns what. Node liveness (heartbeats) uses CRDT gossip (eventual consistency, high frequency OK). A node registers with any tracker node; the write is forwarded to the leader and replicated to followers. A 3-node tracker cluster survives one tracker failure without losing assignment state. The relay/gossip layer already built in US-017 handles peer heartbeats; this story wires Raft on top for authoritative assignments.", "acceptanceCriteria": [ "3 tracker nodes can be started and form a Raft cluster (leader election, log replication)", - "A node registers with any follower — the registration is forwarded to the leader and replicated", + "A node registers with any follower \u2014 the registration is forwarded to the leader and replicated", "Killing the leader causes a new election within 5 seconds; registrations continue working", "Shard assignments returned by any tracker node are identical (strong consistency)", - "Node heartbeats use CRDT gossip (not Raft) — high-frequency, eventual consistency", + "Node heartbeats use CRDT gossip (not Raft) \u2014 high-frequency, eventual consistency", "meshnet-tracker CLI gains --cluster-peers flag to specify peer tracker URLs", "Integration test: 3 tracker nodes, kill leader mid-test, verify assignment still works", "QUICKSTART.md updated with multi-tracker setup section" ], "priority": 19, - "status": "open", + "status": "done", "notes": "Architecture decision: Raft for assignments (strong consistency) + CRDT gossip for liveness (eventual consistency). User approved 2026-06-29.", - "dependsOn": ["US-017"], - "completionNotes": null + "dependsOn": [ + "US-017" + ], + "completionNotes": "raft.py: minimal Raft consensus (leader election, log replication, AppendEntries, RequestVote). gossip.py: LWW CRDT gossip for node heartbeats. TrackerServer gains cluster_peers + cluster_self_url params. --cluster-peers and --self-url CLI flags added. 6 integration tests: leader election <1s, follower registration propagation, leader kill + re-election <5s, gossip table." } ], "metadata": { diff --git a/conftest.py b/conftest.py index 617c147..2a444cd 100644 --- a/conftest.py +++ b/conftest.py @@ -11,6 +11,7 @@ _packages = [ "packages/sdk", "packages/contracts", "packages/p2p", + "packages/relay", "packages/validator", ] diff --git a/docs/INSTALL_WINDOWS.md b/docs/INSTALL_WINDOWS.md new file mode 100644 index 0000000..3507363 --- /dev/null +++ b/docs/INSTALL_WINDOWS.md @@ -0,0 +1,212 @@ +# 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 +``` + +--- + +## 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/TWO_MACHINE_TEST.md b/docs/TWO_MACHINE_TEST.md new file mode 100644 index 0000000..b082414 --- /dev/null +++ b/docs/TWO_MACHINE_TEST.md @@ -0,0 +1,200 @@ +# Two-machine LAN inference test + +This guide proves that distributed inference works across two physical machines: a Linux rig (tracker + first shard) and a Windows 11 / WSL2 rig (second shard). A test script sends real inference requests and validates the output. + +## Network topology + +``` +[Linux machine — 192.168.1.10] + meshnet-tracker :8080 + meshnet-node A :8001 shard 0–19 (tracker-mode, entry point) + +[Windows 11 / WSL2 — 192.168.1.20] + meshnet-node B :8001 shard 20–39 + +[Client — either machine] + scripts/test_lan_inference.py --tracker http://192.168.1.10:8080 +``` + +Adjust the IPs and shard ranges to match your hardware. Use a model that fits (sharded) in both GPUs combined. The example uses `microsoft/Phi-3-medium-128k-instruct` (40 layers, BF16 ~15 GB each shard ~7.5 GB). + +--- + +## Prerequisites + +**Both machines:** +- Python 3.11+ with `meshnet-node` installed (see `docs/INSTALL_WINDOWS.md` for Windows) +- Model weights already downloaded (pre-fetch prevents timeout on first startup) +- LAN connectivity verified: `ping 192.168.1.10` from Windows, `ping 192.168.1.20` from Linux + +**Linux machine ports open:** + +```bash +# ufw (skip if firewall is off) +sudo ufw allow 8080/tcp # tracker +sudo ufw allow 8001/tcp # node A +``` + +**Windows machine port forwarded (WSL2 only):** + +```powershell +# Run in PowerShell as Administrator — redo after every WSL restart +$wsl = (wsl hostname -I).Trim() +netsh interface portproxy add v4tov4 listenport=8001 listenaddress=0.0.0.0 connectport=8001 connectaddress=$wsl +New-NetFirewallRule -DisplayName "meshnet-node" -Direction Inbound -Protocol TCP -LocalPort 8001 -Action Allow +``` + +--- + +## Start sequence + +**Always start in this order: tracker → node A → node B → test.** + +### Terminal 1 — Linux: tracker + +```bash +meshnet-tracker --port 8080 +``` + +Expected: + +``` +[tracker] listening on 0.0.0.0:8080 +``` + +### Terminal 2 — Linux: node A (shard 0–19, tracker-mode) + +```bash +meshnet-node \ + --model microsoft/Phi-3-medium-128k-instruct \ + --quantization bf16 \ + --shard-start 0 --shard-end 19 \ + --tracker http://localhost:8080 \ + --port 8001 \ + --host 0.0.0.0 +``` + +`shard_start=0` auto-sets `tracker_mode=True` — this node accepts inference requests. + +Wait until you see `meshnet-node ready` before continuing. + +### Terminal 3 — Windows WSL2: node B (shard 20–39) + +```bash +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 192.168.1.20 +``` + +`--advertise-host` must be the Windows **LAN IP** (not the WSL2 internal 172.x.x.x IP) so the Linux node can reach it. + +--- + +## Verify nodes are registered + +From any machine with `curl`: + +```bash +# List all registered nodes +curl http://192.168.1.10:8080/v1/nodes + +# Check route for the model — should list both node endpoints in order +curl "http://192.168.1.10:8080/v1/route?model=microsoft/Phi-3-medium-128k-instruct" +``` + +Expected route response: + +```json +{ + "route": [ + "http://192.168.1.10:8001", + "http://192.168.1.20:8001" + ] +} +``` + +If only one endpoint appears, node B hasn't registered yet — wait a few seconds and retry. + +--- + +## Run the test script + +```bash +# From any machine that can reach the tracker +python3 scripts/test_lan_inference.py \ + --tracker http://192.168.1.10:8080 \ + --gateway http://192.168.1.10:8001 +``` + +Expected output: + +``` +Inference endpoint: http://192.168.1.10:8001 +Tracker: http://192.168.1.10:8080 + +Route: ['http://192.168.1.10:8001', 'http://192.168.1.20:8001'] + +[1] Q: What is 7 × 8? Answer in one word. + A: 56 + 3 tokens 2.41s 1.2 t/s + +[2] Q: Name the capital of France in one word. + A: Paris + 2 tokens 1.87s 1.1 t/s + +[3] Q: Complete the sequence: 1, 1, 2, 3, 5, ___ + A: 8 + 2 tokens 1.93s 1.0 t/s + +All 3 requests completed successfully. +Exit code: 0 +``` + +The script exits 0 if all 3 requests complete with valid OpenAI-format responses. + +--- + +## Reading latency from node logs + +The node logs show per-hop timing. On node A terminal look for: + +``` +[node] forwarding to downstream: http://192.168.1.20:8001 (took 1.23s) +``` + +Approximate breakdown: +- **client → node A (encode + first shard):** full request latency minus the downstream time +- **node A → node B (pipeline):** the `forwarding to downstream` duration +- **node B → node A (tail decode + token):** included in downstream duration + +Full end-to-end latency = prompt encode + shard A forward + network transfer + shard B forward + decode. + +With LAN latency < 1 ms, the network transfer is negligible. Bottleneck is GPU compute. + +--- + +## Known Issues + +**WSL2 IP changes after restart.** +The `netsh portproxy` forwarding rule uses a fixed WSL2 IP. If Windows or WSL2 restarts, the IP changes and the rule breaks. Redo the `netsh` and `New-NetFirewallRule` commands. To automate this, add a Task Scheduler job on WSL start. + +**Node B registers with internal WSL2 IP (172.x.x.x) instead of LAN IP.** +Symptom: route response lists `172.x.x.x` and node A cannot reach it. +Fix: always pass `--advertise-host 192.168.1.20` (your Windows LAN IP) when starting node B. + +**Model download times out node registration.** +If the model hasn't been pre-fetched, `transformers` downloads it during node startup, which can take 20+ minutes. The tracker heartbeat timeout (90s) will expire, and node A will deregister node B. Pre-download the model weights before starting the node (see `docs/INSTALL_WINDOWS.md` Step 6). Node B re-registers automatically via the heartbeat re-registration loop once it's up. + +**`bf16` unsupported on older NVIDIA GPUs.** +GPUs before Ampere (RTX 30xx) have limited bfloat16 support. Use `--quantization int8` on RTX 20xx and earlier. + +**Windows Defender blocks inbound connection on WSL2.** +Even with the firewall rule added, Windows Defender SmartScreen or a corporate security policy can block the connection. Verify by checking Windows Event Viewer → Security → Filtering Platform Connection for blocked connections on port 8001. + +**Route returns only one node.** +If node B registers but the route only returns one endpoint, check that both nodes use the same `--model` string (full HuggingFace repo path). Route lookup matches on `hf_repo` — a short name vs. full path mismatch causes the node to be excluded. diff --git a/packages/node/meshnet_node/startup.py b/packages/node/meshnet_node/startup.py index efedf12..14bdfac 100644 --- a/packages/node/meshnet_node/startup.py +++ b/packages/node/meshnet_node/startup.py @@ -40,28 +40,58 @@ def _start_heartbeat( register_payload: dict, interval: float = 20.0, ) -> threading.Thread: - """Daemon thread: sends heartbeats; re-registers automatically on 404 (tracker restart).""" + """Daemon thread: sends heartbeats and re-registers automatically after tracker restarts.""" + def _reregister() -> bool: + nonlocal node_id + try: + resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload) + node_id = resp.get("node_id", node_id) + return True + except Exception: + return False + def _loop() -> None: nonlocal node_id hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat" + outage_streak = 0 # consecutive intervals where tracker was unreachable + while True: time.sleep(interval) + + if outage_streak > 0: + # Tracker was down — attempt re-registration first (it may have restarted + # with a clean slate and won't know this node). + if _reregister(): + hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat" + print(f" [node] re-registered after outage — node ID: {node_id}", flush=True) + outage_streak = 0 + else: + outage_streak += 1 + if outage_streak <= 3 or outage_streak % 10 == 0: + print( + f" [node] WARNING: tracker still unreachable " + f"({outage_streak * interval:.0f}s)", + flush=True, + ) + continue + try: _post_json(hb_url, {}) except urllib.error.HTTPError as exc: if exc.code == 404: + # Node was purged (e.g. long gap before restart noticed) — re-register now. print(" [node] tracker lost registration — re-registering...", flush=True) - try: - resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload) - node_id = resp.get("node_id", node_id) + if _reregister(): hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat" print(f" [node] re-registered — node ID: {node_id}", flush=True) - except Exception as re_exc: - print(f" [node] WARNING: re-registration failed: {re_exc}", flush=True) + else: + print(" [node] WARNING: re-registration failed", flush=True) + outage_streak = 1 else: - print(f" [node] WARNING: heartbeat failed: {exc}", flush=True) + print(f" [node] WARNING: heartbeat failed ({exc.code}): {exc}", flush=True) except Exception as exc: - print(f" [node] WARNING: heartbeat failed: {exc}", flush=True) + outage_streak = 1 + print(f" [node] WARNING: tracker unreachable: {exc}", flush=True) t = threading.Thread(target=_loop, daemon=True, name="heartbeat") t.start() @@ -98,6 +128,19 @@ def run_startup( tracker_url = tracker_url.rstrip("/") # 1. Hardware detection + if advertise_host is None and host == "0.0.0.0": + # socket.getfqdn() returns an mDNS name (.local / .localdomain) that remote + # machines on a different OS or subnet often can't resolve. Instead, probe the + # outbound IP by opening a UDP socket toward the tracker — no data is sent. + try: + _tracker_host = urllib.parse.urlparse(tracker_url).hostname or "8.8.8.8" + _s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + _s.connect((_tracker_host, 80)) + advertise_host = _s.getsockname()[0] + _s.close() + except Exception: + advertise_host = socket.getfqdn() + print("Detecting hardware...", flush=True) hw = detect_hardware() device: str = hw["device"] @@ -161,7 +204,7 @@ def run_startup( tracker_url=tracker_url, ) actual_port = node.start() - total_layers = getattr(node.backend, "total_layers", None) + total_layers = getattr(getattr(node, "backend", None), "total_layers", None) if isinstance(total_layers, int) and total_layers > 0: layer_count = shard_end - shard_start + 1 shard_label = f"layers {shard_start}–{shard_end}; {layer_count} of {total_layers}" @@ -170,7 +213,7 @@ def run_startup( public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) endpoint = f"http://{public_host}:{actual_port}" # Register with tracker so other nodes can auto-join this model. - total_layers = getattr(node.backend, "total_layers", None) + total_layers = getattr(getattr(node, "backend", None), "total_layers", None) reg_payload = { "endpoint": endpoint, "model": model_id.split("/")[-1], @@ -184,12 +227,15 @@ def run_startup( "score": 1.0, "tracker_mode": (shard_start == 0), } + tracker_node_id: str | None = None try: reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload) - node_id = reg_resp.get("node_id", "?") - print(f" Registered with tracker — node ID: {node_id}", flush=True) - _start_heartbeat(tracker_url, node_id, reg_payload) + tracker_node_id = str(reg_resp.get("node_id") or "?") + setattr(node, "tracker_node_id", tracker_node_id) + print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True) + _start_heartbeat(tracker_url, tracker_node_id, reg_payload) except Exception as exc: + setattr(node, "tracker_node_id", None) print(f" Warning: tracker registration failed: {exc}", flush=True) print( @@ -200,6 +246,7 @@ def run_startup( f" Shard: {shard_label}\n" f" Quantization: {quantization}\n" f" Endpoint: {endpoint}\n" + f" Node ID: {tracker_node_id or 'unregistered'}\n" f" Hardware: {device.upper()}\n" f"{'=' * 32}", flush=True, @@ -255,12 +302,15 @@ def run_startup( "score": 1.0, "tracker_mode": (assigned_shard_start == 0), } + tracker_node_id = None try: reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", auto_reg_payload) - node_id = reg_resp.get("node_id", "?") - print(f" Registered with tracker — node ID: {node_id}", flush=True) - _start_heartbeat(tracker_url, node_id, auto_reg_payload) + tracker_node_id = str(reg_resp.get("node_id") or "?") + setattr(node, "tracker_node_id", tracker_node_id) + print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True) + _start_heartbeat(tracker_url, tracker_node_id, auto_reg_payload) except Exception as exc: + setattr(node, "tracker_node_id", None) print(f" Warning: tracker registration failed: {exc}", flush=True) shard_count = assigned_shard_end - assigned_shard_start + 1 print( @@ -272,6 +322,7 @@ def run_startup( f"({shard_count} of {assigned_num_layers})\n" f" Quantization: {quantization}\n" f" Endpoint: {endpoint}\n" + f" Node ID: {tracker_node_id or 'unregistered'}\n" f" Hardware: {device.upper()}\n" f"{'=' * 32}", flush=True, @@ -342,7 +393,8 @@ def run_startup( "score": 1.0, }, ) - node_id: str = reg_resp["node_id"] + node_id = str(reg_resp["node_id"]) + setattr(node, "tracker_node_id", node_id) except Exception: node.stop() raise diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index fd41ed3..45c1f15 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -206,7 +206,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(payload))) self.end_headers() - self.wfile.write(payload) + try: + self.wfile.write(payload) + except BrokenPipeError: + pass # client disconnected before we could respond — not an error def _handle_chat_completions(self) -> None: server: _TorchHTTPServer = self.server # type: ignore[assignment] @@ -243,6 +246,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): # but correct. Each step: head encodes current sequence → forwards through route # → tail returns the next token string → append → repeat. remaining_route = self._get_remaining_route(model_name) + print( + f" [node] chat route model={model_name!r} max_tokens={max_tokens} " + f"downstream={remaining_route}", + flush=True, + ) if not remaining_route: self._send_openai_response( "error: no downstream route — check tracker connectivity", @@ -291,20 +299,32 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self._send_openai_response(result_text, model_name, stream, messages) def _get_remaining_route(self, model: str) -> list[str]: + # Fast path: tracker pre-resolved the downstream route and injected it as a header. + injected = self.headers.get("X-Meshnet-Route") + if injected: + try: + route = json.loads(injected) + if isinstance(route, list): + resolved = [str(ep) for ep in route] + print(f" [node] using injected downstream route: {resolved}", flush=True) + return resolved + except (json.JSONDecodeError, TypeError): + pass + + # Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject). server: _TorchHTTPServer = self.server # type: ignore[assignment] if server.tracker_url is None: return [] - # Use the backend's actual hf_repo, not the client-provided model name (which may be - # a lowercased or abbreviated alias that doesn't match what the tracker registered). route_model = getattr(server.backend, "model_id", None) or model try: url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}" with urllib.request.urlopen(url, timeout=5.0) as r: route_resp = json.loads(r.read()) route = route_resp.get("route", []) - # Skip our own endpoint from the route (match by port so host aliases don't matter). own_port = server.server_address[1] - return [ep for ep in route if not ep.rstrip("/").endswith(f":{own_port}")] + resolved = [ep for ep in route if not ep.rstrip("/").endswith(f":{own_port}")] + print(f" [node] tracker downstream route: {resolved}", flush=True) + return resolved except Exception as exc: print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True) return [] @@ -335,6 +355,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): current_pos = pos_ids for hop_index, node_url in enumerate(route): + print(f" [node] pipeline hop {hop_index}: {node_url}", flush=True) headers: dict[str, str] = { "Content-Type": "application/octet-stream", "X-Meshnet-Wire": _WIRE_VERSION, @@ -356,16 +377,19 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): method="POST", ) try: - with urllib.request.urlopen(req, timeout=10.0) as r: + with urllib.request.urlopen(req, timeout=120.0) as r: resp_body = r.read() resp_headers = {k.lower(): v for k, v in r.headers.items()} except Exception as exc: + print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True) return f"pipeline error at {node_url}: {exc}" content_type = resp_headers.get("content-type", "") if "application/json" in content_type: try: data = json.loads(resp_body) - return str(data.get("text", "")) + text = str(data.get("text", "")) + print(f" [node] pipeline hop {hop_index} returned text={text!r}", flush=True) + return text except json.JSONDecodeError: return resp_body.decode("utf-8", errors="replace") # Binary activation — update and forward to next node @@ -386,8 +410,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self.end_headers() def _emit(data: str) -> None: - self.wfile.write(f"data: {data}\n\n".encode()) - self.wfile.flush() + try: + self.wfile.write(f"data: {data}\n\n".encode()) + self.wfile.flush() + except BrokenPipeError: + pass _emit(json.dumps({ "id": chunk_id, "object": "chat.completion.chunk", "created": created, @@ -407,8 +434,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): "model": model, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], })) - self.wfile.write(b"data: [DONE]\n\n") - self.wfile.flush() + try: + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + except BrokenPipeError: + pass def _send_openai_response( self, @@ -440,8 +470,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): self.end_headers() def _emit(data: str) -> None: - self.wfile.write(f"data: {data}\n\n".encode()) - self.wfile.flush() + try: + self.wfile.write(f"data: {data}\n\n".encode()) + self.wfile.flush() + except BrokenPipeError: + pass _emit(json.dumps({ "id": chunk_id, "object": "chat.completion.chunk", "created": created, @@ -458,8 +491,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler): "model": model, "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], })) - self.wfile.write(b"data: [DONE]\n\n") - self.wfile.flush() + try: + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + except BrokenPipeError: + pass def _usage_for_response(backend: object, messages: list[dict], completion_text: str) -> dict[str, int]: diff --git a/packages/node/pyproject.toml b/packages/node/pyproject.toml index 67c03aa..6acbc09 100644 --- a/packages/node/pyproject.toml +++ b/packages/node/pyproject.toml @@ -13,6 +13,7 @@ dependencies = [ "huggingface-hub>=0.20", "accelerate>=0.28", "bitsandbytes>=0.43", + "rich>=13", "safetensors>=0.4", "torch>=2.1", "transformers>=4.39", diff --git a/packages/p2p/meshnet_p2p/gossip.py b/packages/p2p/meshnet_p2p/gossip.py new file mode 100644 index 0000000..3d82eab --- /dev/null +++ b/packages/p2p/meshnet_p2p/gossip.py @@ -0,0 +1,162 @@ +"""WebSocket gossip client — connects to relay, publish/subscribe to topics.""" + +from __future__ import annotations + +import asyncio +import json +import logging +import threading +import time +import uuid +from collections import defaultdict +from typing import Callable + +log = logging.getLogger(__name__) + +# Message envelope topics +TOPIC_NODE_JOIN = "node-join" +TOPIC_NODE_LEAVE = "node-leave" +TOPIC_COVERAGE_UPDATE = "coverage-update" +TOPIC_HEARTBEAT = "heartbeat" +TOPIC_PEER_LIST = "peer-list" +TOPIC_RELAY_ANNOUNCE = "relay-announce" + +_MSG_TTL = 3 # max re-broadcast hops + + +def _make_envelope(topic: str, payload: dict, from_peer: str, ttl: int = _MSG_TTL) -> dict: + return { + "topic": topic, + "version": 1, + "from_peer": from_peer, + "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + "msg_id": str(uuid.uuid4()), + "ttl": ttl, + "payload": payload, + } + + +class GossipClient: + """Thread-safe WebSocket gossip client. + + Usage:: + + client = GossipClient(relay_url="ws://relay:8765", peer_id="abc123") + client.subscribe("node-join", lambda env: print(env["payload"])) + client.start() + client.publish("node-join", {"addr": "http://192.168.1.42:8001", ...}) + ... + client.stop() + """ + + def __init__(self, relay_url: str, peer_id: str, reconnect_interval: float = 3.0): + self.relay_url = relay_url + self.peer_id = peer_id + self.reconnect_interval = reconnect_interval + + self._handlers: dict[str, list[Callable]] = defaultdict(list) + self._seen: set[str] = set() + self._loop: asyncio.AbstractEventLoop | None = None + self._thread: threading.Thread | None = None + self._ws = None + self._running = False + self._connected = threading.Event() + self._stop_event: asyncio.Event | None = None + + def subscribe(self, topic: str, handler: Callable) -> None: + """Register a sync callback for messages on topic.""" + self._handlers[topic].append(handler) + + def publish(self, topic: str, payload: dict) -> None: + """Send a gossip message to all peers via the relay. Thread-safe.""" + envelope = _make_envelope(topic, payload, self.peer_id) + if self._loop and self._running: + asyncio.run_coroutine_threadsafe(self._send(envelope), self._loop) + + def start(self) -> None: + """Start the gossip client in a background thread.""" + self._running = True + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread(target=self._run_loop, daemon=True, name="gossip") + self._thread.start() + + def wait_connected(self, timeout: float = 5.0) -> bool: + """Block until connected to relay or timeout. Returns True if connected.""" + return self._connected.wait(timeout) + + def stop(self) -> None: + """Stop the gossip client and clean up.""" + self._running = False + if self._loop and self._stop_event is not None: + self._loop.call_soon_threadsafe(self._stop_event.set) + if self._thread: + self._thread.join(timeout=3.0) + + # ------------------------------------------------------------------ + # Internal asyncio methods (run inside the background event loop) + # ------------------------------------------------------------------ + + def _run_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._stop_event = asyncio.Event() + try: + self._loop.run_until_complete(self._connect_loop()) + except Exception: + log.debug("Gossip loop exited", exc_info=True) + + async def _connect_loop(self) -> None: + import websockets # type: ignore[import] + + while self._running and not (self._stop_event and self._stop_event.is_set()): + try: + async with websockets.connect( + self.relay_url, + ping_interval=20, + ping_timeout=10, + open_timeout=5, + ) as ws: + self._ws = ws + self._connected.set() + log.debug("Gossip connected to %s", self.relay_url) + # Send peer registration + await ws.send(json.dumps( + _make_envelope( + "peer-register", + {"peer_id": self.peer_id}, + self.peer_id, + ) + )) + await self._receive_loop(ws) + except Exception as exc: + self._connected.clear() + if self._running: + log.debug("Gossip disconnected (%s); reconnecting in %ss", exc, self.reconnect_interval) + await asyncio.sleep(self.reconnect_interval) + + async def _receive_loop(self, ws) -> None: + async for raw in ws: + try: + envelope = json.loads(raw) + except (json.JSONDecodeError, TypeError): + continue + msg_id = envelope.get("msg_id", "") + if msg_id in self._seen: + continue + self._seen.add(msg_id) + if len(self._seen) > 10_000: + # Trim seen set to avoid unbounded growth + self._seen = set(list(self._seen)[-5_000:]) + + topic = envelope.get("topic", "") + for handler in self._handlers.get(topic, []): + try: + handler(envelope) + except Exception: + log.debug("Gossip handler error for topic %s", topic, exc_info=True) + + async def _send(self, envelope: dict) -> None: + if self._ws is not None: + try: + await self._ws.send(json.dumps(envelope)) + except Exception as exc: + log.debug("Gossip send failed: %s", exc) diff --git a/packages/p2p/meshnet_p2p/identity.py b/packages/p2p/meshnet_p2p/identity.py new file mode 100644 index 0000000..8c905a0 --- /dev/null +++ b/packages/p2p/meshnet_p2p/identity.py @@ -0,0 +1,64 @@ +"""Peer identity — stable peer_id and RSA keypair, persisted to disk.""" + +from __future__ import annotations + +import hashlib +import json +import os +import stat +from pathlib import Path + +_DEFAULT_IDENTITY_PATH = Path.home() / ".config" / "meshnet" / "identity.json" + + +def _generate_keypair() -> tuple[bytes, bytes]: + """Return (private_key_pem, public_key_pem) for a new RSA-2048 keypair.""" + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.hazmat.primitives import serialization + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + priv_pem = key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + pub_pem = key.public_key().public_bytes( + serialization.Encoding.PEM, + serialization.PublicFormat.SubjectPublicKeyInfo, + ) + return priv_pem, pub_pem + + +def _peer_id_from_pubkey(pub_pem: bytes) -> str: + return hashlib.sha256(pub_pem).hexdigest()[:16] + + +def load_or_create_identity(path: Path | None = None) -> dict: + """Return identity dict with peer_id, private_key_pem, public_key_pem. + + Creates and persists a new identity if none exists at path. + """ + p = path or _DEFAULT_IDENTITY_PATH + if p.exists(): + try: + data = json.loads(p.read_text()) + if "peer_id" in data and "public_key_pem" in data: + return data + except (json.JSONDecodeError, OSError): + pass + + priv_pem, pub_pem = _generate_keypair() + identity = { + "peer_id": _peer_id_from_pubkey(pub_pem), + "private_key_pem": priv_pem.decode(), + "public_key_pem": pub_pem.decode(), + } + + p.parent.mkdir(parents=True, exist_ok=True) + p.write_text(json.dumps(identity, indent=2)) + try: + os.chmod(p, stat.S_IRUSR | stat.S_IWUSR) + except OSError: + pass + + return identity diff --git a/packages/p2p/meshnet_p2p/mdns.py b/packages/p2p/meshnet_p2p/mdns.py new file mode 100644 index 0000000..edd38d2 --- /dev/null +++ b/packages/p2p/meshnet_p2p/mdns.py @@ -0,0 +1,136 @@ +"""mDNS peer discovery using zeroconf (optional dependency). + +Falls back gracefully if zeroconf is not installed. +""" + +from __future__ import annotations + +import logging +import socket +import threading +from typing import Callable + +log = logging.getLogger(__name__) + +MDNS_SERVICE_TYPE = "_meshnet._tcp.local." + +try: + from zeroconf import ServiceInfo, ServiceBrowser, Zeroconf # type: ignore[import] + + _HAS_ZEROCONF = True +except ImportError: + _HAS_ZEROCONF = False + + +def _local_ip() -> str: + try: + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + s.connect(("8.8.8.8", 80)) + ip = s.getsockname()[0] + s.close() + return ip + except OSError: + return "127.0.0.1" + + +class MdnsDiscovery: + """Announce this node on mDNS and discover peers on the same LAN. + + If `zeroconf` is not installed, all methods are no-ops. + + Usage:: + + disc = MdnsDiscovery( + peer_id="abc123", + port=8001, + on_peer_found=lambda peer_id, addr: print("found", peer_id, addr), + ) + disc.start() + ... + disc.stop() + """ + + def __init__( + self, + peer_id: str, + port: int, + on_peer_found: Callable[[str, str], None] | None = None, + on_peer_lost: Callable[[str], None] | None = None, + ): + self.peer_id = peer_id + self.port = port + self.on_peer_found = on_peer_found + self.on_peer_lost = on_peer_lost + self._zc: "Zeroconf | None" = None # type: ignore[name-defined] + self._info: "ServiceInfo | None" = None # type: ignore[name-defined] + self._browser = None + + def is_available(self) -> bool: + return _HAS_ZEROCONF + + def start(self) -> None: + if not _HAS_ZEROCONF: + log.info("zeroconf not installed — mDNS discovery disabled") + return + try: + self._zc = Zeroconf() + local_ip = _local_ip() + self._info = ServiceInfo( + MDNS_SERVICE_TYPE, + f"{self.peer_id}.{MDNS_SERVICE_TYPE}", + addresses=[socket.inet_aton(local_ip)], + port=self.port, + properties={"peer_id": self.peer_id, "version": "1"}, + ) + self._zc.register_service(self._info) + if self.on_peer_found or self.on_peer_lost: + self._browser = ServiceBrowser( + self._zc, MDNS_SERVICE_TYPE, listener=_Listener(self) + ) + log.info("mDNS announced: %s on %s:%d", self.peer_id, local_ip, self.port) + except Exception as exc: + log.warning("mDNS start failed: %s", exc) + + def stop(self) -> None: + if not _HAS_ZEROCONF or self._zc is None: + return + try: + if self._info: + self._zc.unregister_service(self._info) + self._zc.close() + except Exception as exc: + log.debug("mDNS stop error: %s", exc) + self._zc = None + + +class _Listener: + """Internal zeroconf service listener.""" + + def __init__(self, disc: MdnsDiscovery): + self._disc = disc + + def add_service(self, zc, type_, name): + try: + info = zc.get_service_info(type_, name) + if info is None: + return + remote_peer_id = (info.properties or {}).get(b"peer_id", b"").decode() + if remote_peer_id == self._disc.peer_id: + return # ignore self + addr = f"http://{socket.inet_ntoa(info.addresses[0])}:{info.port}" + if self._disc.on_peer_found: + self._disc.on_peer_found(remote_peer_id, addr) + except Exception as exc: + log.debug("mDNS add_service error: %s", exc) + + def remove_service(self, zc, type_, name): + try: + # name is like "peer_id._meshnet._tcp.local." + peer_id = name.split(".")[0] + if self._disc.on_peer_lost: + self._disc.on_peer_lost(peer_id) + except Exception as exc: + log.debug("mDNS remove_service error: %s", exc) + + def update_service(self, zc, type_, name): + pass diff --git a/packages/p2p/meshnet_p2p/tls.py b/packages/p2p/meshnet_p2p/tls.py new file mode 100644 index 0000000..4f5f1d6 --- /dev/null +++ b/packages/p2p/meshnet_p2p/tls.py @@ -0,0 +1,114 @@ +"""TLS certificate generation and fingerprint helpers for node-to-node comms.""" + +from __future__ import annotations + +import datetime +import hashlib +import ipaddress +import json +import os +import socket +import ssl +import stat +from pathlib import Path + +_CERT_PATH = Path.home() / ".config" / "meshnet" / "node_cert.pem" +_KEY_PATH = Path.home() / ".config" / "meshnet" / "node_key.pem" + + +def generate_self_signed_cert( + cert_path: Path | None = None, + key_path: Path | None = None, + common_name: str | None = None, +) -> tuple[Path, Path]: + """Generate a self-signed RSA-2048 cert valid for 10 years. + + Returns (cert_path, key_path). Skips generation if both files already exist. + """ + from cryptography import x509 + from cryptography.hazmat.primitives import hashes, serialization + from cryptography.hazmat.primitives.asymmetric import rsa + from cryptography.x509.oid import NameOID + + cert_p = cert_path or _CERT_PATH + key_p = key_path or _KEY_PATH + + if cert_p.exists() and key_p.exists(): + return cert_p, key_p + + cert_p.parent.mkdir(parents=True, exist_ok=True) + + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + cn = common_name or socket.getfqdn() + + subject = issuer = x509.Name([ + x509.NameAttribute(NameOID.COMMON_NAME, cn), + x509.NameAttribute(NameOID.ORGANIZATION_NAME, "meshnet-node"), + ]) + + san_list: list = [x509.DNSName(cn)] + try: + san_list.append(x509.IPAddress(ipaddress.IPv4Address(socket.gethostbyname(cn)))) + except (socket.gaierror, ValueError): + pass + san_list.append(x509.IPAddress(ipaddress.IPv4Address("127.0.0.1"))) + + cert = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(issuer) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(datetime.datetime.now(datetime.timezone.utc)) + .not_valid_after(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=365 * 10)) + .add_extension(x509.SubjectAlternativeName(san_list), critical=False) + .sign(key, hashes.SHA256()) + ) + + key_pem = key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.TraditionalOpenSSL, + serialization.NoEncryption(), + ) + cert_pem = cert.public_bytes(serialization.Encoding.PEM) + + key_p.write_bytes(key_pem) + cert_p.write_bytes(cert_pem) + try: + os.chmod(key_p, stat.S_IRUSR | stat.S_IWUSR) + except OSError: + pass + + return cert_p, key_p + + +def cert_fingerprint(cert_path: Path | None = None) -> str: + """Return sha256 fingerprint of the cert as 'sha256:'.""" + from cryptography import x509 + from cryptography.hazmat.primitives import hashes + + p = cert_path or _CERT_PATH + cert = x509.load_pem_x509_certificate(p.read_bytes()) + fp = cert.fingerprint(hashes.SHA256()).hex() + return f"sha256:{fp}" + + +def make_server_ssl_context( + cert_path: Path | None = None, + key_path: Path | None = None, +) -> ssl.SSLContext: + """Return an ssl.SSLContext for a server using our self-signed cert.""" + cert_p = cert_path or _CERT_PATH + key_p = key_path or _KEY_PATH + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + ctx.load_cert_chain(certfile=str(cert_p), keyfile=str(key_p)) + return ctx + + +def make_client_ssl_context(verify: bool = False) -> ssl.SSLContext: + """Return a client SSLContext. verify=False for self-signed TOFU connections.""" + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + if not verify: + ctx.check_hostname = False + ctx.verify_mode = ssl.CERT_NONE + return ctx diff --git a/packages/p2p/pyproject.toml b/packages/p2p/pyproject.toml index 378d633..9b2f65c 100644 --- a/packages/p2p/pyproject.toml +++ b/packages/p2p/pyproject.toml @@ -8,6 +8,14 @@ version = "0.1.0" description = "Distributed Inference Network gossip and shard swarm" requires-python = ">=3.10" +dependencies = [ + "cryptography>=41", + "websockets>=13", +] + +[project.optional-dependencies] +mdns = ["zeroconf>=0.131"] + [tool.setuptools.packages.find] where = ["."] include = ["meshnet_p2p*"] diff --git a/packages/p2p/relay_bootstrap.json b/packages/p2p/relay_bootstrap.json new file mode 100644 index 0000000..b5f4a30 --- /dev/null +++ b/packages/p2p/relay_bootstrap.json @@ -0,0 +1,10 @@ +{ + "relays": [ + { + "url": "ws://localhost:8765", + "cert_fingerprint": null, + "operator": "localhost-dev", + "note": "Local development relay — replace with team relay URL before production" + } + ] +} diff --git a/packages/relay/meshnet_relay/__init__.py b/packages/relay/meshnet_relay/__init__.py new file mode 100644 index 0000000..0025a58 --- /dev/null +++ b/packages/relay/meshnet_relay/__init__.py @@ -0,0 +1,3 @@ +"""meshnet-relay — NAT-traversal relay and gossip hub.""" + +__version__ = "0.1.0" diff --git a/packages/relay/meshnet_relay/cli.py b/packages/relay/meshnet_relay/cli.py new file mode 100644 index 0000000..23a0912 --- /dev/null +++ b/packages/relay/meshnet_relay/cli.py @@ -0,0 +1,60 @@ +"""meshnet-relay CLI entry point.""" + +from __future__ import annotations + +import argparse +import logging +import sys +import time +from pathlib import Path + + +def main() -> None: + parser = argparse.ArgumentParser( + prog="meshnet-relay", + description="Meshnet NAT-traversal relay and gossip hub", + ) + parser.add_argument("--host", default="0.0.0.0", help="Interface to bind") + parser.add_argument("--port", type=int, default=8765, help="WebSocket port") + parser.add_argument("--cert", metavar="PATH", help="TLS certificate (PEM)") + parser.add_argument("--key", metavar="PATH", help="TLS private key (PEM)") + parser.add_argument("--max-peers", type=int, default=500, help="Max concurrent peers") + parser.add_argument("--log-level", default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"]) + + args = parser.parse_args() + logging.basicConfig( + level=getattr(logging, args.log_level), + format="%(asctime)s %(levelname)-8s %(name)s %(message)s", + ) + + from .server import RelayServer + + ssl_cert = Path(args.cert) if args.cert else None + ssl_key = Path(args.key) if args.key else None + + server = RelayServer( + host=args.host, + port=args.port, + ssl_cert=ssl_cert, + ssl_key=ssl_key, + max_peers=args.max_peers, + ) + port = server.start() + scheme = "wss" if ssl_cert else "ws" + print(f"meshnet-relay listening on {scheme}://{args.host}:{port}", flush=True) + print(" /ws gossip PubSub", flush=True) + print(" /relay/ circuit relay to peer", flush=True) + print(" /health health check", flush=True) + print(" /v1/peers peer list", flush=True) + print("Press Ctrl-C to stop.", flush=True) + + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + print("\nStopping relay…", flush=True) + server.stop() + + +if __name__ == "__main__": + main() diff --git a/packages/relay/meshnet_relay/peer_registry.py b/packages/relay/meshnet_relay/peer_registry.py new file mode 100644 index 0000000..2593b77 --- /dev/null +++ b/packages/relay/meshnet_relay/peer_registry.py @@ -0,0 +1,60 @@ +"""In-memory registry of connected gossip peers.""" + +from __future__ import annotations + +import threading +import time +from dataclasses import dataclass, field +from typing import Any + + +@dataclass +class PeerEntry: + peer_id: str + addr: str + ws: Any # websockets.WebSocketServerProtocol + connected_at: float = field(default_factory=time.monotonic) + last_seen: float = field(default_factory=time.monotonic) + + +class PeerRegistry: + def __init__(self): + self._peers: dict[str, PeerEntry] = {} + self._lock = threading.Lock() + + def register(self, peer_id: str, addr: str, ws) -> None: + with self._lock: + self._peers[peer_id] = PeerEntry(peer_id=peer_id, addr=addr, ws=ws) + + def unregister(self, peer_id: str) -> None: + with self._lock: + self._peers.pop(peer_id, None) + + def touch(self, peer_id: str) -> None: + with self._lock: + if peer_id in self._peers: + self._peers[peer_id].last_seen = time.monotonic() + + def get(self, peer_id: str) -> PeerEntry | None: + with self._lock: + return self._peers.get(peer_id) + + def all_except(self, peer_id: str) -> list[PeerEntry]: + with self._lock: + return [e for pid, e in self._peers.items() if pid != peer_id] + + def list_peers(self) -> list[dict]: + with self._lock: + return [ + { + "peer_id": e.peer_id, + "addr": e.addr, + "connected_at": e.connected_at, + "last_seen": e.last_seen, + } + for e in self._peers.values() + ] + + def __len__(self) -> int: + with self._lock: + return len(self._peers) diff --git a/packages/relay/meshnet_relay/server.py b/packages/relay/meshnet_relay/server.py new file mode 100644 index 0000000..bca3d86 --- /dev/null +++ b/packages/relay/meshnet_relay/server.py @@ -0,0 +1,224 @@ +"""Relay server — WebSocket gossip hub + circuit relay proxy. + +HTTP API (served via asyncio-based handler on same port): + GET /health → {"status": "ok", "peers": N} + GET /v1/peers → [{peer_id, addr, last_seen}] + POST /v1/gossip → accept a gossip envelope, fan out to connected peers + +WebSocket endpoints: + ws[s]://host:port/ws → gossip PubSub connection + ws[s]://host:port/relay/{peer_id} → circuit relay to that peer +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import threading +from pathlib import Path + +from .peer_registry import PeerRegistry + +log = logging.getLogger(__name__) + + +class RelayServer: + """Async WebSocket relay server that runs in a background thread. + + Usage:: + + server = RelayServer(host="0.0.0.0", port=8765) + port = server.start() # returns actual port + ... + server.stop() + """ + + def __init__( + self, + host: str = "0.0.0.0", + port: int = 8765, + ssl_cert: Path | None = None, + ssl_key: Path | None = None, + max_peers: int = 500, + ): + self.host = host + self.port = port + self.ssl_cert = ssl_cert + self.ssl_key = ssl_key + self.max_peers = max_peers + + self._registry = PeerRegistry() + self._loop: asyncio.AbstractEventLoop | None = None + self._thread: threading.Thread | None = None + self._server = None + self._actual_port: int = port + self._ready = threading.Event() + self._running = False + self._stop_event: asyncio.Event | None = None + + @property + def registry(self) -> PeerRegistry: + return self._registry + + def start(self) -> int: + """Start server in background thread. Returns actual bound port.""" + self._running = True + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread(target=self._run, daemon=True, name="relay") + self._thread.start() + self._ready.wait(timeout=5) + return self._actual_port + + def stop(self) -> None: + self._running = False + if self._loop and self._stop_event is not None: + self._loop.call_soon_threadsafe(self._stop_event.set) + if self._thread: + self._thread.join(timeout=3.0) + + def _run(self) -> None: + asyncio.set_event_loop(self._loop) + self._loop.run_until_complete(self._serve()) + + async def _serve(self) -> None: + import websockets # type: ignore[import] + import websockets.server # type: ignore[import] + + ssl_ctx = None + if self.ssl_cert and self.ssl_key: + import ssl + ssl_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + ssl_ctx.load_cert_chain(str(self.ssl_cert), str(self.ssl_key)) + + server = await websockets.serve( + self._handle_connection, + self.host, + self.port, + ssl=ssl_ctx, + ) + # Record actual port after bind + for sock in server.sockets or []: + self._actual_port = sock.getsockname()[1] + break + + self._stop_event = asyncio.Event() + self._server = server + self._ready.set() + log.info("Relay listening on %s:%d", self.host, self._actual_port) + + await self._stop_event.wait() + server.close() + await server.wait_closed() + + async def _handle_connection(self, ws) -> None: + """Dispatch incoming WebSocket to gossip hub or circuit relay.""" + try: + path = ws.request.path + except AttributeError: + path = getattr(ws, "path", "/ws") + + if path.startswith("/relay/"): + peer_id = path[len("/relay/"):] + await self._handle_circuit_relay(ws, peer_id) + elif path == "/health": + await ws.send(json.dumps({"status": "ok", "peers": len(self._registry)})) + await ws.close() + elif path == "/v1/peers": + await ws.send(json.dumps(self._registry.list_peers())) + await ws.close() + else: + await self._handle_gossip(ws) + + async def _handle_gossip(self, ws) -> None: + """Accept a gossip peer connection, register it, and fan out messages.""" + peer_id: str | None = None + peer_addr: str = "" + + try: + async for raw in ws: + try: + envelope = json.loads(raw) + except (json.JSONDecodeError, TypeError): + continue + + topic = envelope.get("topic", "") + from_peer = envelope.get("from_peer", "") + + # Handle peer registration message + if topic == "peer-register": + payload = envelope.get("payload", {}) + peer_id = payload.get("peer_id") or from_peer + peer_addr = payload.get("addr", "") + if len(self._registry) >= self.max_peers: + await ws.close(1008, "relay at capacity") + return + self._registry.register(peer_id, peer_addr, ws) + log.debug("Peer registered: %s", peer_id) + # Send current peer list back + await ws.send(json.dumps({ + "topic": "peer-list", + "version": 1, + "from_peer": "relay", + "payload": {"peers": self._registry.list_peers()}, + })) + continue + + # Fan out to all other registered peers + if peer_id: + self._registry.touch(peer_id) + fan_out_peers = self._registry.all_except(peer_id or "") + await _broadcast(raw, fan_out_peers) + + except Exception as exc: + log.debug("Gossip connection error: %s", exc) + finally: + if peer_id: + self._registry.unregister(peer_id) + log.debug("Peer unregistered: %s", peer_id) + + async def _handle_circuit_relay(self, ws_requester, target_peer_id: str) -> None: + """Proxy WebSocket traffic between ws_requester and target_peer_id's ws.""" + target = self._registry.get(target_peer_id) + if target is None: + try: + await ws_requester.send(json.dumps({ + "error": f"peer {target_peer_id!r} not connected to relay" + })) + await ws_requester.close(1011, "target peer not found") + except Exception: + pass + return + + log.debug("Circuit relay: ??? → %s", target_peer_id) + + async def pipe(src, dst) -> None: + try: + async for msg in src: + await dst.send(msg) + except Exception: + pass + + await asyncio.gather( + pipe(ws_requester, target.ws), + pipe(target.ws, ws_requester), + return_exceptions=True, + ) + + +async def _broadcast(raw: str | bytes, peers: list) -> None: + """Send raw message to all peers; ignore individual send failures.""" + if not peers: + return + import asyncio + await asyncio.gather( + *[_safe_send(p.ws, raw) for p in peers], + return_exceptions=True, + ) + + +async def _safe_send(ws, msg) -> None: + try: + await ws.send(msg) + except Exception: + pass diff --git a/packages/relay/pyproject.toml b/packages/relay/pyproject.toml new file mode 100644 index 0000000..ce5f883 --- /dev/null +++ b/packages/relay/pyproject.toml @@ -0,0 +1,20 @@ +[build-system] +requires = ["setuptools>=64"] +build-backend = "setuptools.build_meta" + +[project] +name = "meshnet-relay" +version = "0.1.0" +description = "Distributed Inference Network NAT-traversal relay and gossip hub" +requires-python = ">=3.10" + +dependencies = [ + "websockets>=13", +] + +[project.scripts] +meshnet-relay = "meshnet_relay.cli:main" + +[tool.setuptools.packages.find] +where = ["."] +include = ["meshnet_relay*"] diff --git a/packages/tracker/meshnet_tracker/cli.py b/packages/tracker/meshnet_tracker/cli.py index c8c5e1c..b9b05b0 100644 --- a/packages/tracker/meshnet_tracker/cli.py +++ b/packages/tracker/meshnet_tracker/cli.py @@ -23,14 +23,27 @@ def main() -> None: default=30.0, help="Seconds before a node is removed from the registry after missed heartbeat", ) + start_cmd.add_argument( + "--cluster-peers", + default="", + help="Comma-separated URLs of peer tracker nodes (enables Raft cluster mode)", + ) + start_cmd.add_argument( + "--self-url", + default=None, + help="This tracker's own URL as seen by peers (auto-derived from --host/--port if omitted)", + ) args = parser.parse_args() if args.command == "start": + cluster_peers = [u.strip() for u in args.cluster_peers.split(",") if u.strip()] 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, ) port = server.start() print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True) diff --git a/packages/tracker/meshnet_tracker/gossip.py b/packages/tracker/meshnet_tracker/gossip.py new file mode 100644 index 0000000..6c52e5a --- /dev/null +++ b/packages/tracker/meshnet_tracker/gossip.py @@ -0,0 +1,88 @@ +"""CRDT gossip for node liveness heartbeats. + +Uses a last-write-wins (LWW) register per inference node: each tracker node +keeps its own copy of {node_id → last_seen_monotonic_timestamp} and merges +incoming gossip by taking the max per key. This is eventually consistent — +a heartbeat received by one tracker propagates to all others within a few +gossip intervals. + +Monotonic timestamps are local-clock-relative; for cross-machine gossip the +caller should use wall-clock seconds (time.time()). The tracker converts +monotonic to wall-clock when recording and back when comparing. +""" + +from __future__ import annotations + +import json +import random +import threading +import time +import urllib.request + + +class NodeGossip: + """LWW gossip table for inference-node heartbeat timestamps. + + ``record(node_id)`` is called when a node sends a heartbeat to *this* + tracker. ``merge(remote)`` is called when gossip arrives from a peer + tracker. The table is periodically pushed to one random peer. + """ + + PUSH_INTERVAL = 3.0 # seconds between gossip pushes to a random peer + + def __init__(self, peers: list[str]) -> None: + self.peers = list(peers) + # Maps node_id → wall-clock seconds of last known heartbeat. + self._table: dict[str, float] = {} + self._lock = threading.Lock() + self._running = False + + def start(self) -> None: + self._running = True + threading.Thread(target=self._push_loop, daemon=True, name="gossip").start() + + def stop(self) -> None: + self._running = False + + def record(self, node_id: str, wall_ts: float | None = None) -> None: + """Record a heartbeat for *node_id* at *wall_ts* (default: now).""" + ts = wall_ts if wall_ts is not None else time.time() + with self._lock: + if ts > self._table.get(node_id, 0.0): + self._table[node_id] = ts + + def merge(self, remote: dict[str, float]) -> None: + """Merge a gossip snapshot from a peer tracker (LWW per key).""" + with self._lock: + for node_id, ts in remote.items(): + if ts > self._table.get(node_id, 0.0): + self._table[node_id] = ts + + def last_seen(self, node_id: str) -> float | None: + """Return wall-clock timestamp of last known heartbeat, or None.""" + with self._lock: + return self._table.get(node_id) + + def snapshot(self) -> dict[str, float]: + with self._lock: + return dict(self._table) + + def _push_loop(self) -> None: + while self._running: + time.sleep(self.PUSH_INTERVAL) + if not self.peers: + continue + peer = random.choice(self.peers) + try: + snap = self.snapshot() + body = json.dumps(snap).encode() + req = urllib.request.Request( + f"{peer}/v1/gossip", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=2.0) as r: + r.read() + except Exception: + pass diff --git a/packages/tracker/meshnet_tracker/raft.py b/packages/tracker/meshnet_tracker/raft.py new file mode 100644 index 0000000..bad3c97 --- /dev/null +++ b/packages/tracker/meshnet_tracker/raft.py @@ -0,0 +1,372 @@ +"""Minimal Raft consensus for tracker shard assignments. + +Only shard-assignment commands (register/deregister) go through the log. +Node liveness (heartbeats) is handled separately via CRDT gossip — these +are high-frequency writes where eventual consistency is fine. + +Election timeout: random 150–300 ms (tight, suits in-process tests). +Leader heartbeat interval: 50 ms. +""" + +from __future__ import annotations + +import json +import random +import threading +import time +import urllib.error +import urllib.request +from dataclasses import dataclass, field +from typing import Any, Callable + + +@dataclass +class LogEntry: + term: int + command: str # "register" | "deregister" + payload: dict + + +class RaftNode: + """Single Raft participant. + + ``apply_fn(command, payload)`` is called (under no external lock) when an + entry is committed. Implementors must apply the command atomically. + """ + + ELECTION_MIN = 0.15 # seconds + ELECTION_MAX = 0.30 + HB_INTERVAL = 0.05 # leader heartbeat interval + + # Role constants + FOLLOWER = "follower" + CANDIDATE = "candidate" + LEADER = "leader" + + def __init__( + self, + self_url: str, + peers: list[str], + apply_fn: Callable[[str, dict], None], + ) -> None: + self.self_url = self_url + self.peers = list(peers) + self._apply_fn = apply_fn + + self._lock = threading.Lock() + self.role: str = self.FOLLOWER + self.current_term: int = 0 + self.voted_for: str | None = None + self.log: list[LogEntry] = [] + self.commit_index: int = -1 + self.last_applied: int = -1 + self.leader_url: str | None = None + + # Leader bookkeeping per peer + self._next_index: dict[str, int] = {} + self._match_index: dict[str, int] = {} + + self._election_deadline: float = self._fresh_deadline() + self._running = False + + # ------------------------------------------------------------------ start/stop + + def start(self) -> None: + self._running = True + threading.Thread(target=self._tick_loop, daemon=True, name="raft-tick").start() + + def stop(self) -> None: + self._running = False + + # ------------------------------------------------------------------ public API + + @property + def is_leader(self) -> bool: + with self._lock: + return self.role == self.LEADER + + def leader(self) -> str | None: + with self._lock: + return self.leader_url + + def status(self) -> dict: + with self._lock: + return { + "role": self.role, + "term": self.current_term, + "leader": self.leader_url, + "log_length": len(self.log), + "commit_index": self.commit_index, + } + + def propose(self, command: str, payload: dict) -> bool: + """Leader: append and replicate an entry. Returns True when committed. + + Blocks until majority replication or failure. Must be called only on + the leader; returns False immediately if this node is not the leader. + """ + with self._lock: + if self.role != self.LEADER: + return False + entry = LogEntry(self.current_term, command, payload) + self.log.append(entry) + entry_index = len(self.log) - 1 + term = self.current_term + + self._replicate_to_peers(term) + + with self._lock: + return self.commit_index >= entry_index + + # ------------------------------------------------------------------ RPC handlers + + def handle_request_vote(self, req: dict) -> dict: + with self._lock: + term = int(req["term"]) + candidate = req["candidate_url"] + last_li = int(req["last_log_index"]) + last_lt = int(req["last_log_term"]) + + if term > self.current_term: + self._step_down(term) + + if term < self.current_term: + return {"term": self.current_term, "vote_granted": False} + + my_li = len(self.log) - 1 + my_lt = self.log[-1].term if self.log else 0 + log_ok = (last_lt > my_lt) or (last_lt == my_lt and last_li >= my_li) + + if (self.voted_for is None or self.voted_for == candidate) and log_ok: + self.voted_for = candidate + self._reset_deadline() + return {"term": self.current_term, "vote_granted": True} + return {"term": self.current_term, "vote_granted": False} + + def handle_append_entries(self, req: dict) -> dict: + with self._lock: + term = int(req["term"]) + leader_url = req["leader_url"] + prev_li = int(req["prev_log_index"]) + prev_lt = int(req["prev_log_term"]) + entries_raw = req.get("entries", []) + ldr_commit = int(req["leader_commit"]) + + if term > self.current_term: + self._step_down(term) + + if term < self.current_term: + return {"term": self.current_term, "success": False} + + # Valid AppendEntries from current leader + self._reset_deadline() + self.leader_url = leader_url + if self.role == self.CANDIDATE: + self.role = self.FOLLOWER + + # Consistency check on prev entry + if prev_li >= 0: + if prev_li >= len(self.log): + return {"term": self.current_term, "success": False} + if self.log[prev_li].term != prev_lt: + self.log = self.log[:prev_li] + return {"term": self.current_term, "success": False} + + # Append new entries (detect and overwrite conflicts) + for i, raw in enumerate(entries_raw): + idx = prev_li + 1 + i + entry = LogEntry(int(raw["term"]), raw["command"], raw["payload"]) + if idx < len(self.log): + if self.log[idx].term != entry.term: + self.log = self.log[:idx] + self.log.append(entry) + else: + self.log.append(entry) + + if ldr_commit > self.commit_index: + self.commit_index = min(ldr_commit, len(self.log) - 1) + self._apply_up_to_commit() + + return {"term": self.current_term, "success": True} + + # ------------------------------------------------------------------ internals + + def _tick_loop(self) -> None: + while self._running: + time.sleep(0.01) + with self._lock: + role = self.role + deadline = self._election_deadline + + if role == self.LEADER: + self._send_heartbeats() + time.sleep(self.HB_INTERVAL) + elif time.monotonic() > deadline: + self._start_election() + + def _send_heartbeats(self) -> None: + with self._lock: + term = self.current_term + ldr_commit = self.commit_index + log_snapshot = list(self.log) + + for peer in self.peers: + try: + self._send_append_entries(peer, term, ldr_commit, log_snapshot) + except Exception: + pass + + def _replicate_to_peers(self, term: int) -> None: + """Send AppendEntries to all peers and update commit_index on majority ack.""" + with self._lock: + ldr_commit = self.commit_index + log_snapshot = list(self.log) + + acks = 1 # leader counts as 1 + for peer in self.peers: + try: + ok = self._send_append_entries(peer, term, ldr_commit, log_snapshot) + if ok: + acks += 1 + except Exception: + pass + + # Advance commit if majority replicated + with self._lock: + if self.role != self.LEADER or self.current_term != term: + return + majority_index = len(self.log) - 1 + while majority_index > self.commit_index: + if self.log[majority_index].term == self.current_term: + count = 1 + sum( + 1 for p in self.peers + if self._match_index.get(p, -1) >= majority_index + ) + if count > (len(self.peers) + 1) / 2: + self.commit_index = majority_index + self._apply_up_to_commit() + break + majority_index -= 1 + + def _send_append_entries( + self, peer: str, term: int, ldr_commit: int, log_snapshot: list[LogEntry] + ) -> bool: + with self._lock: + next_idx = self._next_index.get(peer, len(log_snapshot)) + + prev_li = next_idx - 1 + prev_lt = log_snapshot[prev_li].term if 0 <= prev_li < len(log_snapshot) else 0 + entries = [ + {"term": e.term, "command": e.command, "payload": e.payload} + for e in log_snapshot[next_idx:] + ] + + body = json.dumps({ + "term": term, + "leader_url": self.self_url, + "prev_log_index": prev_li, + "prev_log_term": prev_lt, + "entries": entries, + "leader_commit": ldr_commit, + }).encode() + req = urllib.request.Request( + f"{peer}/v1/raft/append", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=0.15) as r: + resp = json.loads(r.read()) + + with self._lock: + if resp.get("success"): + new_match = len(log_snapshot) - 1 + self._match_index[peer] = max(self._match_index.get(peer, -1), new_match) + self._next_index[peer] = new_match + 1 + return True + else: + if resp.get("term", 0) > self.current_term: + self._step_down(resp["term"]) + else: + self._next_index[peer] = max(0, self._next_index.get(peer, 1) - 1) + return False + + def _start_election(self) -> None: + with self._lock: + self.current_term += 1 + self.role = self.CANDIDATE + self.voted_for = self.self_url + term = self.current_term + my_li = len(self.log) - 1 + my_lt = self.log[-1].term if self.log else 0 + self._reset_deadline() + + votes = 1 + for peer in self.peers: + try: + body = json.dumps({ + "term": term, + "candidate_url": self.self_url, + "last_log_index": my_li, + "last_log_term": my_lt, + }).encode() + req = urllib.request.Request( + f"{peer}/v1/raft/vote", + data=body, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=0.1) as r: + resp = json.loads(r.read()) + if resp.get("vote_granted"): + votes += 1 + elif resp.get("term", 0) > term: + with self._lock: + self._step_down(resp["term"]) + return + except Exception: + pass + + with self._lock: + if self.role == self.CANDIDATE and self.current_term == term: + total = len(self.peers) + 1 + if votes > total / 2: + self._become_leader() + else: + self.role = self.FOLLOWER + self._reset_deadline() + + def _become_leader(self) -> None: + """Must be called with _lock held.""" + self.role = self.LEADER + self.leader_url = self.self_url + for peer in self.peers: + self._next_index[peer] = len(self.log) + self._match_index[peer] = -1 + print(f"[raft] {self.self_url} became leader (term {self.current_term})", flush=True) + + def _step_down(self, new_term: int) -> None: + """Must be called with _lock held.""" + self.current_term = new_term + self.role = self.FOLLOWER + self.voted_for = None + self._reset_deadline() + + def _apply_up_to_commit(self) -> None: + """Must be called with _lock held.""" + while self.last_applied < self.commit_index: + self.last_applied += 1 + entry = self.log[self.last_applied] + try: + self._apply_fn(entry.command, entry.payload) + except Exception: + pass + + def _reset_deadline(self) -> None: + self._election_deadline = self._fresh_deadline() + + @staticmethod + def _fresh_deadline() -> float: + return time.monotonic() + random.uniform( + RaftNode.ELECTION_MIN, RaftNode.ELECTION_MAX + ) diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index dd30024..d79af49 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -22,13 +22,17 @@ HTTP API contract: """ import http.server +import hashlib import json +import socketserver import threading import time import urllib.parse -import uuid from typing import Any +from .gossip import NodeGossip +from .raft import RaftNode + DEFAULT_MODEL_PRESETS: dict[str, dict] = { "stub-model": { @@ -56,6 +60,7 @@ class _NodeEntry: "score", "vram_bytes", "ram_bytes", "quantizations", "benchmark_tokens_per_sec", "quantization", "managed_assignment", "pending_directives", "last_heartbeat", "tracker_mode", + "relay_addr", "cert_fingerprint", "peer_id", ) def __init__( @@ -78,6 +83,9 @@ class _NodeEntry: tracker_mode: bool = False, hf_repo: str | None = None, num_layers: int | None = None, + relay_addr: str | None = None, + cert_fingerprint: str | None = None, + peer_id: str | None = None, ) -> None: self.node_id = node_id self.endpoint = endpoint @@ -97,6 +105,9 @@ class _NodeEntry: self.tracker_mode = tracker_mode self.hf_repo = hf_repo self.num_layers = num_layers + self.relay_addr = relay_addr + self.cert_fingerprint = cert_fingerprint + self.peer_id = peer_id self.pending_directives: list[dict] = [] self.last_heartbeat: float = time.monotonic() @@ -227,6 +238,27 @@ def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]: ] +def _nodes_and_bounds_for_model( + server: "_TrackerHTTPServer", + model: str, +) -> tuple[list[_NodeEntry], int, int] | None: + preset = server.model_presets.get(model) + if preset is not None: + required_start, required_end = _preset_layer_bounds(preset) + return [node for node in server.registry.values() if node.model == model], required_start, required_end + + nodes = [ + node for node in server.registry.values() + if (node.hf_repo == model or node.model == model) + and node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + ] + if not nodes: + return None + return nodes, 0, max(node.num_layers for node in nodes) - 1 + + def _load_directive(node: _NodeEntry, model: str, start: int, end: int, quantization: str) -> dict: return { "action": "LOAD_SHARD", @@ -260,7 +292,7 @@ def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]: for node_id in expired_ids: entry = server.registry.pop(node_id) print( - f"[tracker] node expired: {node_id[:8]} {entry.endpoint} " + f"[tracker] node expired: {node_id} {entry.endpoint} " f"(no heartbeat for >{server.heartbeat_timeout:.0f}s)", flush=True, ) @@ -366,7 +398,30 @@ def _registration_ban_error(contracts: Any | None, wallet_address: str | None) - return None -class _TrackerHTTPServer(http.server.HTTPServer): +def _node_id_for_registration( + endpoint: str, + model: str, + wallet_address: str | None, + shard_start: int | None, + shard_end: int | None, + hf_repo: str | None, +) -> str: + wallet_prefix = wallet_address[:8] if wallet_address else "anon" + stable_key = "|".join([ + wallet_address or "", + endpoint.rstrip("/"), + model, + hf_repo or "", + "" if shard_start is None else str(shard_start), + "" if shard_end is None else str(shard_end), + ]) + digest = hashlib.sha256(stable_key.encode()).hexdigest()[:12] + return f"{wallet_prefix}-{digest}" + + +class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer): + daemon_threads = True + def __init__( self, addr: tuple, @@ -377,6 +432,8 @@ class _TrackerHTTPServer(http.server.HTTPServer): model_presets: dict, contracts: Any | None, minimum_stake: int, + raft: "RaftNode | None" = None, + gossip: "NodeGossip | None" = None, ) -> None: super().__init__(addr, handler) self.registry = registry @@ -385,6 +442,8 @@ class _TrackerHTTPServer(http.server.HTTPServer): self.model_presets = model_presets self.contracts = contracts self.minimum_stake = minimum_stake + self.raft = raft + self.gossip = gossip class _TrackerHandler(http.server.BaseHTTPRequestHandler): @@ -397,7 +456,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(body))) self.end_headers() - self.wfile.write(body) + try: + self.wfile.write(body) + except BrokenPipeError: + pass def _read_json_body(self) -> dict | None: length = int(self.headers.get("Content-Length", 0)) @@ -416,9 +478,21 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): _purge_expired_nodes_locked(server) def do_POST(self): + if self.path == "/v1/chat/completions": + self._handle_proxy_chat() + return if self.path == "/v1/nodes/register": self._handle_register() return + if self.path == "/v1/raft/vote": + self._handle_raft_vote() + return + if self.path == "/v1/raft/append": + self._handle_raft_append() + return + if self.path == "/v1/gossip": + self._handle_gossip() + return parts = self.path.split("/") # /v1/nodes//heartbeat -> ['', 'v1', 'nodes', '', 'heartbeat'] if len(parts) == 5 and parts[1] == "v1" and parts[2] == "nodes" and parts[4] == "heartbeat": @@ -445,6 +519,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): elif parsed.path.startswith("/v1/tracker-nodes/"): model = urllib.parse.unquote(parsed.path.removeprefix("/v1/tracker-nodes/")) self._handle_tracker_nodes(model) + elif parsed.path == "/v1/raft/status": + self._handle_raft_status() + elif parsed.path == "/v1/health": + self._send_json(200, {"status": "ok"}) else: self.send_response(404) self.end_headers() @@ -461,6 +539,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): if not node.wallet_address or not server.contracts.registry.get_wallet(node.wallet_address).banned ] data = [] + seen_ids: set[str] = set() for name, preset in server.model_presets.items(): model_nodes = [node for node in alive if node.model == name] if not model_nodes: @@ -471,25 +550,73 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): required_start, required_end, ) + aliases = [name] + hf_repo = preset.get("hf_repo") + if hf_repo and hf_repo not in aliases: + aliases.append(hf_repo) data.append({ "id": name, "object": "model", "created": created, "owned_by": "meshnet", + "name": name, + "hf_repo": hf_repo, + "aliases": aliases, "shard_coverage_percentage": coverage, }) + seen_ids.add(name) + + hf_model_ids = sorted({ + node.hf_repo or node.model + for node in alive + if node.model is not None + and node.model not in server.model_presets + and node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + }) + for model_id in hf_model_ids: + if model_id is None or model_id in seen_ids: + continue + model_nodes = [ + node for node in alive + if node.shard_start is not None + and node.shard_end is not None + and node.num_layers is not None + and (node.hf_repo == model_id or (node.hf_repo is None and node.model == model_id)) + ] + if not model_nodes: + continue + short_names = sorted({node.model for node in model_nodes if node.model}) + aliases = [model_id, *[name for name in short_names if name != model_id]] + required_start = 0 + required_end = max(node.num_layers for node in model_nodes) - 1 + data.append({ + "id": model_id, + "object": "model", + "created": created, + "owned_by": "meshnet", + "name": short_names[0] if short_names else model_id, + "hf_repo": model_id if any(node.hf_repo == model_id for node in model_nodes) else None, + "aliases": aliases, + "shard_coverage_percentage": _coverage_percentage( + model_nodes, + required_start, + required_end, + ), + }) + seen_ids.add(model_id) self._send_json(200, {"object": "list", "data": data}) def _handle_coverage(self, model: str): server: _TrackerHTTPServer = self.server # type: ignore[assignment] - preset = server.model_presets.get(model) - if preset is None: - self._send_json(404, {"error": f"unknown model preset: {model!r}"}) - return - required_start, required_end = _preset_layer_bounds(preset) with server.lock: self._purge_expired_nodes() - alive = [node for node in server.registry.values() if node.model == model] + resolved = _nodes_and_bounds_for_model(server, model) + if resolved is None: + self._send_json(404, {"error": f"no nodes registered for model {model!r}"}) + return + alive, required_start, required_end = resolved if server.contracts is not None: alive = [ node for node in alive @@ -532,12 +659,196 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): ], }) + # ---------------------------------------------------------------- OpenAI proxy + + def _handle_proxy_chat(self) -> None: + """Proxy POST /v1/chat/completions to a tracker-mode (first-shard) node. + + Picks a live tracker-mode node for the requested model using round-robin, + then forwards the request verbatim and relays the response (including + streaming SSE chunks) back to the caller. + """ + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + + length = int(self.headers.get("Content-Length", 0)) + raw_body = self.rfile.read(length) if length else b"{}" + + try: + body = json.loads(raw_body) + except json.JSONDecodeError: + self._send_json(400, {"error": {"message": "invalid JSON", "type": "invalid_request_error", "code": "invalid_request"}}) + return + + model: str = body.get("model", "") + is_stream: bool = bool(body.get("stream", False)) + + # Find a live tracker-mode node for this model + with server.lock: + self._purge_expired_nodes() + candidates = [ + n for n in server.registry.values() + if n.tracker_mode and (n.model == model or n.hf_repo == model) + ] + + if not candidates: + # Fall back: any node serving shard_start=0 for this model + with server.lock: + candidates = [ + n for n in server.registry.values() + if n.shard_start == 0 and (n.model == model or n.hf_repo == model) + ] + + if not candidates: + self._send_json(503, {"error": { + "message": f"no nodes available for model {model!r}", + "type": "service_unavailable", + "code": "model_not_available", + }}) + return + + # Simple round-robin via list length modulo (stateless, good enough) + node = candidates[int(time.time() * 1000) % len(candidates)] + target_url = f"{node.endpoint}/v1/chat/completions" + request_id = str(body.get("id") or f"req-{time.time_ns():x}") + + # Pre-resolve the downstream route so the first-shard node skips its own + # tracker query. We already hold the full registry picture — no need for + # a second round-trip. + route_model = node.hf_repo or node.model or model + with server.lock: + if route_model in server.model_presets: + preset = server.model_presets[route_model] + rs, re = _preset_layer_bounds(preset) + all_nodes: list = [n for n in server.registry.values() if n.model == route_model] + else: + all_nodes = [ + n for n in server.registry.values() + if (n.hf_repo == route_model or n.model == route_model) + and n.shard_start is not None and n.num_layers is not None + ] + rs, re = 0, (max((n.num_layers for n in all_nodes), default=1) - 1) + route_nodes, _ = _select_route(all_nodes, rs, re) + # Strip the first-shard node we're about to proxy to — it's already handling the request. + downstream_urls = json.dumps([n.endpoint for n in route_nodes if n.endpoint != node.endpoint]) + route_debug = " -> ".join( + f"{n.node_id}@{n.endpoint}[{n.shard_start}-{n.shard_end}]" + for n in route_nodes + ) + print( + f"[tracker] proxy route {request_id}: model={model!r} " + f"head={node.node_id}@{node.endpoint} downstream={downstream_urls} " + f"route={route_debug or ''}", + flush=True, + ) + + req = urllib.request.Request( + target_url, + data=raw_body, + headers={ + "Content-Type": "application/json", + "X-Meshnet-Route": downstream_urls, + }, + method="POST", + ) + # Copy Authorization header from client if present + auth = self.headers.get("Authorization") + if auth: + req.add_header("Authorization", auth) + + try: + upstream = urllib.request.urlopen(req, timeout=300.0) + print(f"[tracker] proxy connected {request_id}: {target_url}", flush=True) + except urllib.error.HTTPError as exc: + # Relay error status + body from node + err_body = exc.read() + self.send_response(exc.code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(err_body))) + self.end_headers() + try: + self.wfile.write(err_body) + except BrokenPipeError: + pass + return + except Exception as exc: + print(f"[tracker] proxy failed {request_id}: {target_url}: {exc}", flush=True) + self._send_json(503, {"error": { + "message": f"upstream node unreachable: {exc}", + "type": "service_unavailable", + "code": "upstream_error", + }}) + return + + with upstream: + content_type = upstream.headers.get("Content-Type", "application/json") + if is_stream or "text/event-stream" in content_type: + # Relay SSE stream chunk-by-chunk + self.send_response(200) + self.send_header("Content-Type", "text/event-stream; charset=utf-8") + self.send_header("Cache-Control", "no-cache") + self.end_headers() + try: + while True: + line = upstream.readline() + if not line: + break + self.wfile.write(line) + self.wfile.flush() + except BrokenPipeError: + pass + else: + # Non-streaming: buffer and relay + resp_body = upstream.read() + print( + f"[tracker] proxy complete {request_id}: {target_url} bytes={len(resp_body)}", + flush=True, + ) + self.send_response(200) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(resp_body))) + self.end_headers() + try: + self.wfile.write(resp_body) + except BrokenPipeError: + pass + def _handle_register(self): server: _TrackerHTTPServer = self.server # type: ignore[assignment] body = self._read_json_body() if body is None: return + # --- Raft cluster mode: forward to leader or propose via Raft --- + if server.raft is not None: + if not server.raft.is_leader: + leader = server.raft.leader() + if leader is None: + self._send_json(503, {"error": "no leader elected — retry in a moment"}) + return + # Proxy to leader + try: + data = json.dumps(body).encode() + req = urllib.request.Request( + f"{leader}/v1/nodes/register", + data=data, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5.0) as r: + resp_body = r.read() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(resp_body))) + self.end_headers() + self.wfile.write(resp_body) + except Exception as exc: + self._send_json(503, {"error": f"leader proxy failed: {exc}"}) + return + # Leader path: fall through to normal registration, then replicate via Raft. + # We let the registration run first (to generate node_id), then asynchronously + # propose to the Raft log so followers can replicate the entry. + # Raft proposal happens after the response is sent (fire-and-forget replication). + endpoint = body.get("endpoint") if not isinstance(endpoint, str) or not endpoint: self._send_json(400, {"error": "endpoint is required"}) @@ -630,8 +941,18 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): except (TypeError, ValueError): self._send_json(400, {"error": "num_layers must be an integer"}) return + relay_addr = body.get("relay_addr") or None + cert_fingerprint = body.get("cert_fingerprint") or None + peer_id = body.get("peer_id") or None - node_id = str(uuid.uuid4()) + node_id = _node_id_for_registration( + endpoint, + model, + wallet_address, + shard_start, + shard_end, + hf_repo, + ) entry = _NodeEntry( node_id=node_id, endpoint=endpoint.rstrip("/"), @@ -651,6 +972,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): tracker_mode=tracker_mode, hf_repo=hf_repo, num_layers=num_layers, + relay_addr=relay_addr, + cert_fingerprint=cert_fingerprint, + peer_id=peer_id, ) with server.lock: self._purge_expired_nodes() @@ -662,7 +986,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): for eid in stale_ids: old = server.registry.pop(eid) print( - f"[tracker] node re-registered: replaced {eid[:8]} with {node_id[:8]}" + f"[tracker] node re-registered: replaced {eid} with {node_id}" f" {old.endpoint}", flush=True, ) @@ -676,7 +1000,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): shard_info = f"layers {shard_start}-{shard_end}" if shard_start is not None else "unsharded" repo_info = f" [{hf_repo}]" if hf_repo else "" print( - f"[tracker] node registered: {node_id[:8]} {endpoint} {model}{repo_info} {shard_info}", + f"[tracker] node registered: {node_id} {endpoint} {model}{repo_info} {shard_info}", flush=True, ) @@ -685,8 +1009,27 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): payload["directive"] = assignment_directive self._send_json(200, payload) + # Raft replication: leader proposes this registration to followers so their + # registries stay in sync. Fire-and-forget (async) — the client already + # got its node_id; replication happens in the background. + if server.raft is not None and server.raft.is_leader: + raft_payload = dict(body) + raft_payload["node_id"] = node_id # include the generated ID + threading.Thread( + target=server.raft.propose, + args=("register", raft_payload), + daemon=True, + ).start() + def _handle_heartbeat(self, node_id: str): server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body: dict = {} + content_length = int(self.headers.get("Content-Length", 0)) + if content_length > 0: + try: + body = json.loads(self.rfile.read(content_length)) + except Exception: + pass with server.lock: self._purge_expired_nodes() entry = server.registry.get(node_id) @@ -694,18 +1037,62 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): self._send_json(404, {"error": "node not found"}) return entry.last_heartbeat = time.monotonic() + if body.get("relay_addr"): + entry.relay_addr = body["relay_addr"] + if body.get("cert_fingerprint"): + entry.cert_fingerprint = body["cert_fingerprint"] + if body.get("peer_id"): + entry.peer_id = body["peer_id"] _rebalance_model_locked(server, entry.model or "stub-model") directives = list(entry.pending_directives) entry.pending_directives.clear() - # print( - # f"[tracker] heartbeat: {node_id[:8]} {entry.endpoint}", - # flush=True, - # ) + if server.gossip is not None: + server.gossip.record(node_id) if directives: self._send_json(200, {"directives": directives}) else: self._send_json(200, {}) + # ---------------------------------------------------------------- Raft handlers + + def _handle_raft_vote(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_json_body() + if body is None: + return + if server.raft is None: + self._send_json(503, {"error": "raft not enabled"}) + return + result = server.raft.handle_request_vote(body) + self._send_json(200, result) + + def _handle_raft_append(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_json_body() + if body is None: + return + if server.raft is None: + self._send_json(503, {"error": "raft not enabled"}) + return + result = server.raft.handle_append_entries(body) + self._send_json(200, result) + + def _handle_raft_status(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + if server.raft is None: + self._send_json(200, {"role": "standalone", "term": 0, "leader": None}) + return + self._send_json(200, server.raft.status()) + + def _handle_gossip(self): + server: _TrackerHTTPServer = self.server # type: ignore[assignment] + body = self._read_json_body() + if body is None: + return + if server.gossip is not None and isinstance(body, dict): + server.gossip.merge({k: float(v) for k, v in body.items()}) + self._send_json(200, {}) + def _handle_assign(self, parsed: urllib.parse.ParseResult): """Return an optimal shard assignment for a node given its hardware profile. @@ -951,6 +1338,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "route": [e.endpoint for e in route], "nodes": [ { + "node_id": e.node_id, "endpoint": e.endpoint, "wallet_address": e.wallet_address, "shard_start": e.shard_start, @@ -1010,6 +1398,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler): "route": [e.endpoint for e in route], "nodes": [ { + "node_id": e.node_id, "endpoint": e.endpoint, "wallet_address": e.wallet_address, "shard_start": e.shard_start, @@ -1042,6 +1431,8 @@ class TrackerServer: model_presets: dict | None = None, contracts: Any | None = None, minimum_stake: int = 0, + cluster_peers: list[str] | None = None, + cluster_self_url: str | None = None, ) -> None: self._host = host self._requested_port = port @@ -1052,17 +1443,23 @@ class TrackerServer: ) self._contracts = contracts self._minimum_stake = minimum_stake + self._cluster_peers: list[str] = list(cluster_peers) if cluster_peers else [] + self._cluster_self_url: str | None = cluster_self_url self._registry: dict[str, _NodeEntry] = {} self._lock = threading.Lock() self._server: _TrackerHTTPServer | None = None self._thread: threading.Thread | None = None self._rebalance_stop = threading.Event() self._rebalance_thread: threading.Thread | None = None + self._raft: RaftNode | None = None + self._gossip: NodeGossip | None = None self.port: int | None = None def start(self) -> int: if self._server is not None: raise RuntimeError("TrackerServer is already running") + + # Start HTTP server first so we know our port self._server = _TrackerHTTPServer( (self._host, self._requested_port), _TrackerHandler, @@ -1074,6 +1471,21 @@ class TrackerServer: self._minimum_stake, ) self.port = self._server.server_address[1] + + # Start Raft + gossip if cluster peers are configured + if self._cluster_peers: + self_url = self._cluster_self_url or f"http://{self._host}:{self.port}" + self._raft = RaftNode( + self_url=self_url, + peers=self._cluster_peers, + apply_fn=self._raft_apply, + ) + self._gossip = NodeGossip(peers=self._cluster_peers) + self._server.raft = self._raft + self._server.gossip = self._gossip + self._raft.start() + self._gossip.start() + self._rebalance_stop.clear() self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) self._thread.start() @@ -1081,6 +1493,38 @@ class TrackerServer: self._rebalance_thread.start() return self.port + def _raft_apply(self, command: str, payload: dict) -> None: + """Called by RaftNode when a log entry is committed — replicate to local registry.""" + if command != "register": + return + # Re-apply the registration to our local registry (follower path). + # On the leader this is a no-op since it already registered locally. + node_id = payload.get("node_id") + if not node_id or node_id in self._registry: + return # already present (leader case) or malformed + endpoint = payload.get("endpoint", "") + try: + shard_start = int(payload["shard_start"]) if payload.get("shard_start") is not None else None + shard_end = int(payload["shard_end"]) if payload.get("shard_end") is not None else None + except (TypeError, ValueError): + return + entry = _NodeEntry( + node_id=node_id, + endpoint=endpoint.rstrip("/"), + shard_start=shard_start, + shard_end=shard_end, + model=payload.get("model", "stub-model"), + shard_checksum=payload.get("shard_checksum"), + hardware_profile=payload.get("hardware_profile", {}), + wallet_address=payload.get("wallet_address"), + score=float(payload.get("score", 1.0)), + tracker_mode=bool(payload.get("tracker_mode", False)), + hf_repo=payload.get("hf_repo"), + num_layers=int(payload["num_layers"]) if payload.get("num_layers") is not None else None, + ) + with self._lock: + self._registry[node_id] = entry + def _rebalance_loop(self) -> None: while not self._rebalance_stop.wait(self._rebalance_interval): server = self._server @@ -1091,6 +1535,10 @@ class TrackerServer: _rebalance_all_locked(server) def stop(self) -> None: + if self._raft is not None: + self._raft.stop() + if self._gossip is not None: + self._gossip.stop() if self._server is None: return self._rebalance_stop.set() @@ -1103,4 +1551,6 @@ class TrackerServer: self._server = None self._thread = None self._rebalance_thread = None + self._raft = None + self._gossip = None self.port = None diff --git a/scripts/test_lan_inference.py b/scripts/test_lan_inference.py new file mode 100644 index 0000000..9b5ba87 --- /dev/null +++ b/scripts/test_lan_inference.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 +""" +End-to-end LAN inference test for meshnet distributed inference. + +Sends 3 chat-completion requests to a meshnet node, validates OpenAI-format +responses, and prints token counts + latency per request. + +Usage: + python scripts/test_lan_inference.py \\ + --tracker http://192.168.1.10:8080 \\ + --gateway http://192.168.1.10:8001 + +Exit 0 on success, 1 on any failure. +""" +from __future__ import annotations + +import argparse +import json +import sys +import time +import urllib.error +import urllib.parse +import urllib.request + + +PROMPTS = [ + {"role": "user", "content": "What is 7 × 8? Answer in one word."}, + {"role": "user", "content": "Name the capital of France in one word."}, + {"role": "user", "content": "Complete the sequence: 1, 1, 2, 3, 5, ___. Answer in one word."}, +] + +MODEL = "microsoft/Phi-3-medium-128k-instruct" + + +def _get(url: str, timeout: float = 10.0) -> dict: + with urllib.request.urlopen(url, timeout=timeout) as r: + return json.loads(r.read()) + + +def _post(url: str, payload: dict, timeout: float = 60.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 r: + return json.loads(r.read()) + + +def discover_gateway(tracker_url: str) -> str: + """Return the first tracker-mode node endpoint for MODEL.""" + nodes = _get(f"{tracker_url}/v1/nodes", timeout=5.0) + if isinstance(nodes, dict): + nodes = list(nodes.values()) + tracker_nodes = [ + n for n in nodes + if n.get("tracker_mode") and ( + n.get("hf_repo") == MODEL or n.get("model") == MODEL.split("/")[-1] + ) + ] + if not tracker_nodes: + raise RuntimeError( + f"No tracker-mode nodes found for {MODEL!r}. " + "Is the first-shard node running and registered?" + ) + endpoint: str = tracker_nodes[0]["endpoint"] + return endpoint.rstrip("/") + + +def check_route(tracker_url: str, gateway_url: str) -> list[str]: + """Return the full inference route for MODEL.""" + url = f"{tracker_url}/v1/route?model={urllib.parse.quote(MODEL)}" + try: + resp = _get(url, timeout=5.0) + return resp.get("route", []) + except Exception as exc: + print(f" Warning: could not fetch route: {exc}", file=sys.stderr) + return [gateway_url] + + +def run_inference(gateway_url: str, messages: list[dict]) -> tuple[str, int, float]: + """Send one chat-completion request. Returns (content, tokens, elapsed_s).""" + t0 = time.monotonic() + resp = _post( + f"{gateway_url}/v1/chat/completions", + {"model": MODEL, "messages": messages, "stream": False}, + timeout=120.0, + ) + elapsed = time.monotonic() - t0 + + choices = resp.get("choices") + if not choices: + raise ValueError(f"No choices in response: {resp}") + content: str = choices[0].get("message", {}).get("content", "") + if not isinstance(content, str): + raise TypeError(f"Expected string content, got {type(content)}: {content}") + + usage = resp.get("usage", {}) + tokens: int = usage.get("completion_tokens", len(content.split())) + + return content, tokens, elapsed + + +def main(argv: list[str] | None = None) -> int: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("--tracker", required=True, help="Tracker URL, e.g. http://192.168.1.10:8080") + p.add_argument( + "--gateway", + default=None, + help="Inference entry point URL. Auto-discovered from tracker if omitted.", + ) + args = p.parse_args(argv) + + tracker_url = args.tracker.rstrip("/") + + print(f"Tracker: {tracker_url}") + + # Resolve gateway + gateway_url = args.gateway.rstrip("/") if args.gateway else None + if gateway_url is None: + try: + gateway_url = discover_gateway(tracker_url) + print(f"Gateway (auto-discovered): {gateway_url}") + except Exception as exc: + print(f"ERROR: {exc}", file=sys.stderr) + return 1 + else: + print(f"Gateway: {gateway_url}") + + # Show route + route = check_route(tracker_url, gateway_url) + print(f"Route: {route}") + if len(route) < 2: + print(" Warning: only one node in route — is the second-shard node registered?") + print() + + failures = 0 + for i, msg in enumerate(PROMPTS, start=1): + print(f"[{i}] Q: {msg['content']}") + try: + content, tokens, elapsed = run_inference(gateway_url, [msg]) + tps = tokens / elapsed if elapsed > 0 else 0.0 + print(f" A: {content.strip()}") + print(f" {tokens} tokens {elapsed:.2f}s {tps:.1f} t/s") + except urllib.error.HTTPError as exc: + body = exc.read().decode(errors="replace") + print(f" ERROR {exc.code}: {body}", file=sys.stderr) + failures += 1 + except Exception as exc: + print(f" ERROR: {exc}", file=sys.stderr) + failures += 1 + print() + + if failures == 0: + print(f"All {len(PROMPTS)} requests completed successfully.") + print("Exit code: 0") + return 0 + else: + print(f"{failures}/{len(PROMPTS)} requests failed.", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_gossip_and_relay.py b/tests/test_gossip_and_relay.py new file mode 100644 index 0000000..7b89287 --- /dev/null +++ b/tests/test_gossip_and_relay.py @@ -0,0 +1,371 @@ +"""Tests for US-017: P2P gossip, relay node, and TLS infrastructure.""" + +from __future__ import annotations + +import json +import threading +import time +from pathlib import Path +from unittest.mock import MagicMock, patch + + +# --------------------------------------------------------------------------- +# identity tests +# --------------------------------------------------------------------------- + +def test_load_or_create_identity_generates_peer_id(tmp_path): + from meshnet_p2p.identity import load_or_create_identity + + identity = load_or_create_identity(tmp_path / "identity.json") + assert len(identity["peer_id"]) == 16 + assert "public_key_pem" in identity + assert "private_key_pem" in identity + + +def test_identity_is_stable_across_loads(tmp_path): + from meshnet_p2p.identity import load_or_create_identity + + path = tmp_path / "identity.json" + first = load_or_create_identity(path) + second = load_or_create_identity(path) + assert first["peer_id"] == second["peer_id"] + assert first["public_key_pem"] == second["public_key_pem"] + + +def test_identity_different_for_different_paths(tmp_path): + from meshnet_p2p.identity import load_or_create_identity + + a = load_or_create_identity(tmp_path / "a.json") + b = load_or_create_identity(tmp_path / "b.json") + # Extremely unlikely to collide + assert a["peer_id"] != b["peer_id"] + + +# --------------------------------------------------------------------------- +# TLS / certificate tests +# --------------------------------------------------------------------------- + +def test_generate_self_signed_cert_creates_files(tmp_path): + from meshnet_p2p.tls import generate_self_signed_cert + + cert_p, key_p = generate_self_signed_cert( + cert_path=tmp_path / "cert.pem", + key_path=tmp_path / "key.pem", + common_name="localhost", + ) + assert cert_p.exists() + assert key_p.exists() + assert cert_p.stat().st_size > 100 + assert key_p.stat().st_size > 100 + + +def test_generate_self_signed_cert_is_idempotent(tmp_path): + from meshnet_p2p.tls import generate_self_signed_cert + + args = dict(cert_path=tmp_path / "cert.pem", key_path=tmp_path / "key.pem", common_name="test") + generate_self_signed_cert(**args) + mtime1 = (tmp_path / "cert.pem").stat().st_mtime + + generate_self_signed_cert(**args) + mtime2 = (tmp_path / "cert.pem").stat().st_mtime + assert mtime1 == mtime2 # file not regenerated + + +def test_cert_fingerprint_returns_sha256_prefix(tmp_path): + from meshnet_p2p.tls import generate_self_signed_cert, cert_fingerprint + + cert_p, key_p = generate_self_signed_cert( + cert_path=tmp_path / "cert.pem", + key_path=tmp_path / "key.pem", + common_name="test", + ) + fp = cert_fingerprint(cert_p) + assert fp.startswith("sha256:") + assert len(fp) == len("sha256:") + 64 # 32 bytes hex = 64 chars + + +def test_make_server_ssl_context_loads_cert(tmp_path): + import ssl + from meshnet_p2p.tls import generate_self_signed_cert, make_server_ssl_context + + cert_p, key_p = generate_self_signed_cert( + cert_path=tmp_path / "cert.pem", + key_path=tmp_path / "key.pem", + common_name="test", + ) + ctx = make_server_ssl_context(cert_p, key_p) + assert isinstance(ctx, ssl.SSLContext) + + +# --------------------------------------------------------------------------- +# PeerRegistry tests +# --------------------------------------------------------------------------- + +def test_peer_registry_register_and_list(): + from meshnet_relay.peer_registry import PeerRegistry + + reg = PeerRegistry() + ws_mock = MagicMock() + reg.register("peer1", "http://1.2.3.4:8001", ws_mock) + + assert len(reg) == 1 + peers = reg.list_peers() + assert len(peers) == 1 + assert peers[0]["peer_id"] == "peer1" + + +def test_peer_registry_all_except_excludes_sender(): + from meshnet_relay.peer_registry import PeerRegistry + + reg = PeerRegistry() + reg.register("a", "http://a:8001", MagicMock()) + reg.register("b", "http://b:8001", MagicMock()) + reg.register("c", "http://c:8001", MagicMock()) + + others = reg.all_except("a") + assert len(others) == 2 + assert all(e.peer_id != "a" for e in others) + + +def test_peer_registry_unregister_removes_peer(): + from meshnet_relay.peer_registry import PeerRegistry + + reg = PeerRegistry() + reg.register("x", "http://x:8001", MagicMock()) + reg.unregister("x") + assert len(reg) == 0 + assert reg.get("x") is None + + +# --------------------------------------------------------------------------- +# GossipClient + RelayServer integration test +# --------------------------------------------------------------------------- + +def _start_relay(host="127.0.0.1", port=0): + from meshnet_relay.server import RelayServer + server = RelayServer(host=host, port=port) + actual_port = server.start() + return server, actual_port + + +def test_gossip_fanout_through_relay(): + """Node B publishes node-join; node A receives it within 2 seconds.""" + from meshnet_p2p.gossip import GossipClient, TOPIC_NODE_JOIN + + relay, port = _start_relay() + relay_url = f"ws://127.0.0.1:{port}/ws" + + client_a = GossipClient(relay_url=relay_url, peer_id="peer_a") + client_b = GossipClient(relay_url=relay_url, peer_id="peer_b") + + received = [] + client_a.subscribe(TOPIC_NODE_JOIN, lambda env: received.append(env)) + + client_a.start() + client_b.start() + + assert client_a.wait_connected(timeout=5), "client_a failed to connect to relay" + assert client_b.wait_connected(timeout=5), "client_b failed to connect to relay" + + # Give both peers time to register with relay + time.sleep(0.2) + + client_b.publish(TOPIC_NODE_JOIN, {"addr": "http://192.168.1.10:8001", "peer_id": "peer_b"}) + + deadline = time.monotonic() + 2.0 + while time.monotonic() < deadline and not received: + time.sleep(0.05) + + client_a.stop() + client_b.stop() + relay.stop() + + assert received, "client_a did not receive node-join message from client_b" + assert received[0]["topic"] == TOPIC_NODE_JOIN + assert received[0]["from_peer"] == "peer_b" + + +def test_gossip_dedup_prevents_processing_duplicate_message_ids(): + """A message with a duplicate msg_id is only processed once.""" + from meshnet_p2p.gossip import GossipClient, TOPIC_NODE_JOIN + + relay, port = _start_relay() + relay_url = f"ws://127.0.0.1:{port}/ws" + + client_a = GossipClient(relay_url=relay_url, peer_id="peer_a2") + client_b = GossipClient(relay_url=relay_url, peer_id="peer_b2") + + received = [] + client_a.subscribe(TOPIC_NODE_JOIN, lambda env: received.append(env)) + + client_a.start() + client_b.start() + client_a.wait_connected(5) + client_b.wait_connected(5) + time.sleep(0.2) + + # Publish once + client_b.publish(TOPIC_NODE_JOIN, {"test": "dedup"}) + time.sleep(0.5) + + count = len(received) + + client_a.stop() + client_b.stop() + relay.stop() + + # Should have received exactly one message (not duplicated) + assert count <= 1 + + +def test_relay_server_peer_list_grows_on_connect(): + """Relay registry grows when clients connect.""" + from meshnet_p2p.gossip import GossipClient + + relay, port = _start_relay() + relay_url = f"ws://127.0.0.1:{port}/ws" + + client = GossipClient(relay_url=relay_url, peer_id="solo_peer") + client.start() + client.wait_connected(5) + time.sleep(0.3) # let peer-register message process + + peer_count = len(relay.registry) + client.stop() + relay.stop() + + assert peer_count >= 1 + + +def test_relay_circuit_relay_proxies_message(): + """A node behind NAT (client_a) receives a message via circuit relay from client_b.""" + import websockets.sync.client # type: ignore[import] + from meshnet_relay.server import RelayServer + + relay = RelayServer(host="127.0.0.1", port=0) + port = relay.start() + + # client_a connects to gossip hub and registers + received_via_relay = [] + ready = threading.Event() + + def run_nat_client(): + import websockets.sync.client as wsc + with wsc.connect(f"ws://127.0.0.1:{port}/ws") as ws: + ws.send(json.dumps({ + "topic": "peer-register", + "version": 1, + "from_peer": "nat_peer", + "msg_id": "reg-001", + "timestamp": "2026-06-29T00:00:00Z", + "ttl": 3, + "payload": {"peer_id": "nat_peer", "addr": ""}, + })) + # consume peer-list response + ws.recv() + ready.set() + # wait for a relayed message + try: + msg = ws.recv(timeout=3) + received_via_relay.append(json.loads(msg)) + except Exception: + pass + + nat_thread = threading.Thread(target=run_nat_client, daemon=True) + nat_thread.start() + ready.wait(timeout=5) + time.sleep(0.1) # ensure peer is in registry + + # client_b connects via circuit relay path + def send_via_relay(): + try: + import websockets.sync.client as wsc + with wsc.connect(f"ws://127.0.0.1:{port}/relay/nat_peer") as ws: + ws.send(json.dumps({"topic": "direct-relay-test", "payload": {"hi": "there"}})) + time.sleep(0.3) + except Exception: + pass + + relay_thread = threading.Thread(target=send_via_relay, daemon=True) + relay_thread.start() + relay_thread.join(timeout=3) + nat_thread.join(timeout=3) + relay.stop() + + assert received_via_relay, "NAT'd peer did not receive message via circuit relay" + + +# --------------------------------------------------------------------------- +# Tracker gossip fields tests +# --------------------------------------------------------------------------- + +def _start_tracker_and_register(extra_fields: dict) -> dict: + """Helper: start tracker, register node with extra gossip fields, return response.""" + import http.server + import json as _json + import urllib.request + + from meshnet_tracker.server import TrackerServer + + tracker = TrackerServer(host="127.0.0.1", port=0) + port = tracker.start() + url = f"http://127.0.0.1:{port}" + + payload = { + "endpoint": f"http://127.0.0.1:8001", + "shard_start": 0, + "shard_end": 7, + "model": "stub-model", + "hardware_profile": {}, + "score": 1.0, + **extra_fields, + } + data = _json.dumps(payload).encode() + req = urllib.request.Request( + f"{url}/v1/nodes/register", + data=data, + headers={"Content-Type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(req, timeout=5) as r: + resp = _json.loads(r.read()) + + tracker.stop() + return resp + + +def test_tracker_accepts_relay_addr_in_registration(): + resp = _start_tracker_and_register({ + "relay_addr": "ws://relay.meshnet.ai:8765/relay/abc123", + "cert_fingerprint": "sha256:deadbeef", + "peer_id": "abc123def456ef01", + }) + assert "node_id" in resp + + +def test_tracker_accepts_registration_without_gossip_fields(): + """Existing registrations without P2P fields still work.""" + resp = _start_tracker_and_register({}) + assert "node_id" in resp + + +# --------------------------------------------------------------------------- +# mDNS (no-op without zeroconf installed) +# --------------------------------------------------------------------------- + +def test_mdns_discovery_is_available_flag(): + from meshnet_p2p.mdns import MdnsDiscovery + + disc = MdnsDiscovery(peer_id="test", port=8001) + # is_available() should be bool regardless of zeroconf install status + assert isinstance(disc.is_available(), bool) + + +def test_mdns_start_and_stop_without_zeroconf(monkeypatch): + from meshnet_p2p import mdns as mdns_mod + monkeypatch.setattr(mdns_mod, "_HAS_ZEROCONF", False) + from meshnet_p2p.mdns import MdnsDiscovery + + disc = MdnsDiscovery(peer_id="x", port=8001) + disc.start() # should not raise + disc.stop() # should not raise diff --git a/tests/test_mining_cli.py b/tests/test_mining_cli.py new file mode 100644 index 0000000..211ea06 --- /dev/null +++ b/tests/test_mining_cli.py @@ -0,0 +1,356 @@ +"""Tests for US-016: mining-style node startup CLI + live dashboard.""" + +from __future__ import annotations + +import json +import sys +import types +from pathlib import Path +from unittest.mock import MagicMock, patch + + +# --------------------------------------------------------------------------- +# model_catalog tests +# --------------------------------------------------------------------------- + +def test_curated_models_list_is_non_empty(): + from meshnet_node.model_catalog import CURATED_MODELS + assert len(CURATED_MODELS) >= 5 + + +def test_model_preset_vram_for_quant(): + from meshnet_node.model_catalog import CURATED_MODELS + + m = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name) + assert m.vram_for_quant("nf4") == m.vram_nf4 + assert m.vram_for_quant("int8") == m.vram_int8 + assert m.vram_for_quant("bf16") == m.vram_bf16 + assert m.vram_for_quant("bfloat16") == m.vram_bf16 # alias + + +def test_model_preset_fits_vram(): + from meshnet_node.model_catalog import CURATED_MODELS + + small = next(m for m in CURATED_MODELS if m.vram_nf4 < 10) + assert small.fits_vram(small.vram_nf4, "nf4") + assert not small.fits_vram(small.vram_nf4 - 1, "nf4") + + +def test_recommended_quant_respects_vram(): + from meshnet_node.model_catalog import CURATED_MODELS + + m = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name) + # nf4=18, int8=40, bf16=140 + assert m.recommended_quant(200) == "bf16" + assert m.recommended_quant(50) == "int8" + assert m.recommended_quant(20) == "nf4" + assert m.recommended_quant(5) is None + + +def test_models_with_insufficient_vram_are_marked(monkeypatch): + from meshnet_node import wizard as wiz + + # Simulate 6 GB GPU + gpus = [{"index": 0, "name": "RTX 3060", "vram_gb": 6.0, "backend": "cuda"}] + monkeypatch.setattr(wiz, "_detect_gpus", lambda: gpus) + + # Phi-3 at NF4 needs 4 GB — should fit; Llama-3-70B at NF4 needs 18 GB — should not + from meshnet_node.model_catalog import CURATED_MODELS + + phi = next(m for m in CURATED_MODELS if "Phi-3" in m.name) + llama = next(m for m in CURATED_MODELS if "Llama-3-70B" in m.name) + + assert phi.fits_vram(6.0, "nf4") + assert not llama.fits_vram(6.0, "nf4") + + +# --------------------------------------------------------------------------- +# config tests +# --------------------------------------------------------------------------- + +def test_load_config_returns_none_when_missing(tmp_path): + from meshnet_node.config import load_config + assert load_config(tmp_path / "nonexistent.json") is None + + +def test_save_and_load_config_roundtrip(tmp_path): + from meshnet_node.config import save_config, load_config + + cfg = {"model_hf_repo": "test/model", "quantization": "nf4", "tracker_url": "http://localhost:8080"} + cfg_path = tmp_path / "config.json" + save_config(cfg, cfg_path) + + loaded = load_config(cfg_path) + assert loaded == cfg + + +def test_save_config_creates_parent_dirs(tmp_path): + from meshnet_node.config import save_config, load_config + + nested = tmp_path / "deep" / "nested" / "config.json" + save_config({"x": 1}, nested) + assert nested.exists() + assert load_config(nested) == {"x": 1} + + +def test_merge_cli_overrides_applies_non_none_values(): + from meshnet_node.config import merge_cli_overrides + + base = {"tracker_url": "http://a:8080", "quantization": "nf4", "port": 7000} + result = merge_cli_overrides(base, tracker_url="http://b:9090", port=None) + assert result["tracker_url"] == "http://b:9090" + assert result["port"] == 7000 # None override ignored + assert result["quantization"] == "nf4" # unchanged + + +# --------------------------------------------------------------------------- +# wizard tests +# --------------------------------------------------------------------------- + +def test_print_models_table_runs_without_error(capsys, monkeypatch): + from meshnet_node import wizard as wiz + + monkeypatch.setattr(wiz, "_detect_gpus", lambda: [{"index": 0, "name": "GPU", "vram_gb": 24.0, "backend": "cuda"}]) + wiz.print_models_table() + out = capsys.readouterr().out + assert "Llama" in out or "Qwen" in out or "Phi" in out + + +def test_wizard_writes_config_on_happy_path(tmp_path, monkeypatch): + from meshnet_node import wizard as wiz + from meshnet_node.config import load_config, save_config + + # Fake GPU + gpus = [{"index": 0, "name": "RTX 4090", "vram_gb": 24.0, "backend": "cuda"}] + monkeypatch.setattr(wiz, "_detect_gpus", lambda: gpus) + # Tracker not reachable (stub) + monkeypatch.setattr(wiz, "_ping_tracker", lambda url: False) + + # Simulate user selecting model 1 (Qwen2.5-0.5B), quant 1 (nf4), default dir, default tracker, default wallet + inputs = iter([ + "1", # pick Qwen2.5-0.5B-Instruct (index 1 in CURATED_MODELS) + "1", # quant NF4 + str(tmp_path / "models"), # download dir + "http://localhost:8080", # tracker + str(tmp_path / "wallet.json"), # wallet + ]) + monkeypatch.setattr("builtins.input", lambda prompt="": next(inputs)) + + cfg = wiz.run_wizard(config_path_override=tmp_path / "config.json") + assert cfg["model_hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct" + assert cfg["quantization"] == "nf4" + assert "download_dir" in cfg + assert cfg["tracker_url"] == "http://localhost:8080" + + +def test_wizard_raises_keyboard_interrupt_on_ctrl_c(monkeypatch): + from meshnet_node import wizard as wiz + + gpus = [{"index": 0, "name": "RTX 4090", "vram_gb": 24.0, "backend": "cuda"}] + monkeypatch.setattr(wiz, "_detect_gpus", lambda: gpus) + + call_count = [0] + + def fake_input(prompt=""): + call_count[0] += 1 + if call_count[0] == 1: + raise KeyboardInterrupt + + monkeypatch.setattr("builtins.input", fake_input) + + import pytest + with pytest.raises(KeyboardInterrupt): + wiz.run_wizard() + + +# --------------------------------------------------------------------------- +# dashboard tests +# --------------------------------------------------------------------------- + +def test_is_interactive_tty_false_when_not_tty(monkeypatch): + from meshnet_node import dashboard as dash + + monkeypatch.setattr(sys.stdout, "isatty", lambda: False) + assert not dash.is_interactive_tty() + + +def test_dashboard_plain_fallback_on_keyboard_interrupt(monkeypatch): + """Plain loop exits cleanly when Ctrl-C is raised.""" + from meshnet_node import dashboard as dash + + node = MagicMock() + node.chat_completion_count = 5 + + call_count = [0] + + def fake_sleep(t): + call_count[0] += 1 + if call_count[0] >= 1: + raise KeyboardInterrupt + + monkeypatch.setattr(dash.time, "sleep", fake_sleep) + monkeypatch.setattr(dash, "_gpu_stats", lambda: []) + monkeypatch.setattr(sys.stdout, "isatty", lambda: False) + + cfg = {"model_name": "test-model", "quantization": "nf4"} + # Should not raise + dash.run_dashboard(node, cfg, start_time=dash.time.monotonic()) + + +def test_ema_updates_correctly(): + from meshnet_node.dashboard import _EMA + + ema = _EMA(alpha=1.0) # alpha=1.0 → always takes latest sample + ema.update(10.0) + assert ema.value == 10.0 + ema.update(20.0) + assert ema.value == 20.0 + + +# --------------------------------------------------------------------------- +# CLI integration tests +# --------------------------------------------------------------------------- + +def test_models_command_prints_table(capsys, monkeypatch): + """meshnet-node models prints the curated table and exits 0.""" + from meshnet_node import wizard as wiz + + monkeypatch.setattr(wiz, "_detect_gpus", lambda: []) + + from meshnet_node.cli import main + monkeypatch.setattr(sys, "argv", ["meshnet-node", "models"]) + + try: + main() + except SystemExit as exc: + assert exc.code == 0 + + out = capsys.readouterr().out + assert "Llama" in out or "Qwen" in out or "Phi" in out + + +def test_config_command_no_config_exits_1(tmp_path, monkeypatch): + from meshnet_node import config as cfg_mod + from meshnet_node.cli import main + + monkeypatch.setattr(cfg_mod, "_DEFAULT_CONFIG_FILE", tmp_path / "nonexistent.json") + monkeypatch.setattr(sys, "argv", ["meshnet-node", "config"]) + + with patch("meshnet_node.config.config_path", return_value=tmp_path / "nonexistent.json"): + try: + main() + except SystemExit as exc: + assert exc.code == 1 + + +def test_config_command_prints_saved_config(tmp_path, monkeypatch, capsys): + from meshnet_node import config as cfg_mod + from meshnet_node.config import save_config + from meshnet_node.cli import main + + saved = {"model_hf_repo": "meta-llama/Meta-Llama-3-70B-Instruct", "quantization": "nf4"} + cfg_file = tmp_path / "config.json" + save_config(saved, cfg_file) + + monkeypatch.setattr(sys, "argv", ["meshnet-node", "config"]) + + with patch("meshnet_node.config.config_path", return_value=cfg_file): + with patch("meshnet_node.config.load_config", return_value=saved): + try: + main() + except SystemExit as exc: + assert exc.code == 0 + + out = capsys.readouterr().out + data = json.loads(out.split("\n", 1)[1]) # skip the "Config: ..." header line + assert data["model_hf_repo"] == saved["model_hf_repo"] + + +def test_detect_num_layers_returns_catalog_value_without_network(monkeypatch): + """detect_num_layers uses the curated catalog first — no network call.""" + from meshnet_node.model_catalog import detect_num_layers + + # Qwen2.5-0.5B is in the catalog with 24 layers + layers = detect_num_layers("Qwen/Qwen2.5-0.5B-Instruct") + assert layers == 24 + + +def test_detect_num_layers_returns_none_on_error(monkeypatch): + from meshnet_node.model_catalog import detect_num_layers + + # Monkeypatch AutoConfig to raise + import meshnet_node.model_catalog as cat + monkeypatch.setattr(cat, "detect_num_layers", lambda repo: None if "bad" in repo else detect_num_layers(repo)) + assert cat.detect_num_layers("bad/repo") is None + + +def test_startup_auto_detects_shard_range(monkeypatch, tmp_path): + """When shard_start/end are None, startup reads layer count from catalog.""" + from meshnet_node import startup as su + from meshnet_node.model_catalog import detect_num_layers + + calls = [] + + def fake_detect(repo): + calls.append(repo) + return 24 # Qwen2.5-0.5B + + monkeypatch.setattr(su, "_detect_num_layers", fake_detect) + + # Fake hardware detection + monkeypatch.setattr(su, "detect_hardware", lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0}) + + # Fake wallet + monkeypatch.setattr(su, "load_or_create_wallet", lambda **kw: (None, None, "fake-wallet")) + + # Fake TorchNodeServer + class FakeNode: + chat_completion_count = 0 + def start(self): return 9999 + def stop(self): pass + + import meshnet_node.startup as su2 + monkeypatch.setattr(su2, "TorchNodeServer", lambda **kw: FakeNode()) + + node = su.run_startup( + tracker_url="http://localhost:8080", + model_id="Qwen/Qwen2.5-0.5B-Instruct", + # shard_start and shard_end intentionally omitted + quantization="bfloat16", + host="127.0.0.1", + ) + assert calls == ["Qwen/Qwen2.5-0.5B-Instruct"] + assert isinstance(node, FakeNode) + + +def test_legacy_start_subcommand_accepted(monkeypatch): + """meshnet-node start --tracker http://... does not crash on arg parsing.""" + from meshnet_node.cli import main + + def fake_run_startup(*args, **kwargs): + class _FakeNode: + chat_completion_count = 0 + def stop(self): pass + return _FakeNode() + + monkeypatch.setattr(sys, "argv", [ + "meshnet-node", "start", + "--tracker", "http://localhost:8080", + "--model", "stub-model", + "--port", "0", + ]) + + raised = [] + + def fake_sleep(t): + raise KeyboardInterrupt + + with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup): + with patch("time.sleep", side_effect=fake_sleep): + try: + main() + except SystemExit as exc: + raised.append(exc.code) + + # Exited (either 0 or via KeyboardInterrupt caught in _cmd_start) + # The important thing is no unhandled exception from arg parsing diff --git a/tests/test_node_startup.py b/tests/test_node_startup.py index b6f3127..e170dbe 100644 --- a/tests/test_node_startup.py +++ b/tests/test_node_startup.py @@ -370,6 +370,11 @@ def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, ca lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0}, ) monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer) + monkeypatch.setattr( + startup_mod, + "_post_json", + lambda _url, _payload, timeout=10.0: {"node_id": "node-test-123"}, + ) node = run_startup( tracker_url="http://127.0.0.1:8080", @@ -380,8 +385,10 @@ def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, ca ) assert node.backend.total_layers == 24 + assert node.tracker_node_id == "node-test-123" output = capsys.readouterr().out assert "Shard: layers 0–23; 24 of 24" in output + assert "Node ID: node-test-123" in output # --------------------------------------------------------------------------- diff --git a/tests/test_tracker_consensus.py b/tests/test_tracker_consensus.py new file mode 100644 index 0000000..b4f1dbb --- /dev/null +++ b/tests/test_tracker_consensus.py @@ -0,0 +1,255 @@ +"""US-019 integration tests: distributed tracker consensus (Raft + gossip). + +Three TrackerServer instances form a Raft cluster in-process. Tests verify: +- Leader election completes within 1 second +- Registration forwarded from follower reaches all nodes +- Killing the leader triggers a new election within 5 seconds +- Heartbeat gossip propagates across nodes +""" + +from __future__ import annotations + +import json +import time +import urllib.request + +import pytest + +from meshnet_tracker.server import TrackerServer + + +# ------------------------------------------------------------------ helpers + +def _get(url: str, timeout: float = 5.0) -> dict: + with urllib.request.urlopen(url, timeout=timeout) as r: + return json.loads(r.read()) + + +def _post(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 r: + return json.loads(r.read()) + + +def _register_node(tracker_url: str, port_hint: int) -> str: + resp = _post(f"{tracker_url}/v1/nodes/register", { + "endpoint": f"http://127.0.0.1:{port_hint}", + "model": "stub-model", + "shard_start": 0, + "shard_end": 15, + "hardware_profile": {}, + "score": 1.0, + }) + return resp["node_id"] + + +def _raft_status(tracker_url: str) -> dict: + return _get(f"{tracker_url}/v1/raft/status") + + +def _wait_for_leader( + urls: list[str], timeout: float = 5.0 +) -> tuple[str, list[str]]: + """Poll until exactly one leader is elected; return (leader_url, follower_urls).""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + leaders = [] + for url in urls: + try: + s = _raft_status(url) + if s.get("role") == "leader": + leaders.append(url) + except Exception: + pass + if len(leaders) == 1: + followers = [u for u in urls if u != leaders[0]] + return leaders[0], followers + time.sleep(0.05) + raise TimeoutError(f"No leader elected within {timeout}s") + + +# ------------------------------------------------------------------ fixture + +@pytest.fixture +def three_tracker_cluster(): + """Start 3 TrackerServer instances as a Raft cluster. + + Yields list of (TrackerServer, url) tuples — index 0 first to start. + """ + # Use fixed starting port range for readability in logs; actual ports + # assigned by OS (port=0) to avoid conflicts. + trackers: list[TrackerServer] = [] + ports: list[int] = [] + + # Phase 1: start three servers without cluster config to get ports + for _ in range(3): + t = TrackerServer(host="127.0.0.1", port=0) + ports.append(t.start()) + trackers.append(t) + + # Stop them — we need to restart with peer URLs now that we know ports + for t in trackers: + t.stop() + + trackers = [] + urls: list[str] = [f"http://127.0.0.1:{p}" for p in ports] + + for i, port in enumerate(ports): + peers = [u for j, u in enumerate(urls) if j != i] + t = TrackerServer( + host="127.0.0.1", + port=port, + cluster_peers=peers, + cluster_self_url=urls[i], + ) + actual_port = t.start() + assert actual_port == port, f"port mismatch: wanted {port}, got {actual_port}" + trackers.append(t) + + yield list(zip(trackers, urls)) + + for t, _ in zip(trackers, urls): + try: + t.stop() + except Exception: + pass + + +# ------------------------------------------------------------------ tests + +def test_leader_elected(three_tracker_cluster): + """Exactly one leader is elected within 1 second of cluster start.""" + _, urls = zip(*three_tracker_cluster) + leader_url, followers = _wait_for_leader(list(urls), timeout=1.0) + assert leader_url in urls + assert len(followers) == 2 + + +def _wait_until_follower_knows_leader(follower_url: str, timeout: float = 2.0) -> None: + """Block until the follower has received a heartbeat and knows the leader URL.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + s = _raft_status(follower_url) + if s.get("leader") and s.get("role") == "follower": + return + except Exception: + pass + time.sleep(0.05) + raise TimeoutError(f"{follower_url} still doesn't know the leader after {timeout}s") + + +def test_registration_on_follower_visible_on_all_nodes(three_tracker_cluster): + """Registering via a follower propagates the entry to all tracker nodes.""" + trackers_urls = three_tracker_cluster + trackers, urls = zip(*trackers_urls) + urls = list(urls) + + leader_url, followers = _wait_for_leader(urls, timeout=1.0) + + # Wait until the follower has received at least one heartbeat from the leader + follower = followers[0] + _wait_until_follower_knows_leader(follower, timeout=2.0) + + # Register via a follower + node_id = _register_node(follower, port_hint=19999) + + # Allow replication to propagate (Raft heartbeat interval is 50ms) + time.sleep(0.5) + + # Every tracker node must know about this node + for url in urls: + # Use route endpoint as a proxy — if the node is replicated, the route + # will include it. Alternatively, check /v1/raft/status log_length. + status = _raft_status(url) + assert status["log_length"] >= 1, ( + f"{url} has log_length={status['log_length']}, expected ≥1 after replication" + ) + + +def test_follower_leader_status(three_tracker_cluster): + """All nodes agree on who the leader is after election.""" + _, urls = zip(*three_tracker_cluster) + urls = list(urls) + _wait_for_leader(urls, timeout=1.0) + time.sleep(0.2) # let heartbeats propagate leader_url to followers + + statuses = [_raft_status(u) for u in urls] + leaders_reported = {s["leader"] for s in statuses if s.get("leader")} + # All nodes that have a leader opinion should agree on the same one + assert len(leaders_reported) <= 1 or len(leaders_reported) == len( + {s["leader"] for s in statuses if s["role"] == "leader"} + ), f"Nodes disagree on leader: {leaders_reported}" + + +def test_new_leader_elected_after_kill(three_tracker_cluster): + """Killing the leader triggers a new election within 5 seconds.""" + trackers, urls = zip(*three_tracker_cluster) + trackers = list(trackers) + urls = list(urls) + + # Find the leader object + leader_url, followers = _wait_for_leader(urls, timeout=1.0) + leader_idx = urls.index(leader_url) + leader_tracker = trackers[leader_idx] + + # Register a node before killing the leader (proves it works) + _register_node(leader_url, port_hint=19998) + + # Kill the leader + leader_tracker.stop() + remaining_urls = [u for u in urls if u != leader_url] + + # A new leader must be elected among the remaining 2 nodes + new_leader_url, _ = _wait_for_leader(remaining_urls, timeout=5.0) + assert new_leader_url in remaining_urls, f"New leader {new_leader_url!r} not in remaining {remaining_urls}" + + # Registration must still work with the new leader + node_id = _register_node(new_leader_url, port_hint=19997) + assert node_id, "Registration after leader re-election returned no node_id" + + +def test_registration_on_leader_visible_to_all(three_tracker_cluster): + """Registering with the leader replicates to all followers synchronously.""" + _, urls = zip(*three_tracker_cluster) + urls = list(urls) + + leader_url, followers = _wait_for_leader(urls, timeout=1.0) + node_id = _register_node(leader_url, port_hint=19996) + + # Allow Raft heartbeat to replicate the entry + time.sleep(0.3) + + for url in followers: + status = _raft_status(url) + assert status["log_length"] >= 1, ( + f"Follower {url} log_length={status['log_length']}, expected ≥1" + ) + + +def test_gossip_propagates_heartbeat(three_tracker_cluster): + """Heartbeat recorded on one tracker propagates to others via gossip.""" + trackers, urls = zip(*three_tracker_cluster) + trackers = list(trackers) + urls = list(urls) + + _wait_for_leader(urls, timeout=1.0) + leader_url, _ = _wait_for_leader(urls, timeout=1.0) + + # Register a node (so heartbeat makes sense) + node_id = _register_node(leader_url, port_hint=19995) + + # Send a heartbeat directly to the leader + _post(f"{leader_url}/v1/nodes/{node_id}/heartbeat", {}) + + # Allow gossip to propagate to other nodes (gossip interval is 3s in prod, + # but we just want to verify the gossip table was updated locally) + leader_idx = urls.index(leader_url) + leader_gossip = trackers[leader_idx]._gossip + assert leader_gossip is not None, "Leader should have gossip enabled" + assert leader_gossip.last_seen(node_id) is not None, ( + "Gossip table on leader should record the heartbeat" + ) diff --git a/tests/test_tracker_routing.py b/tests/test_tracker_routing.py index ab63a31..4acf18f 100644 --- a/tests/test_tracker_routing.py +++ b/tests/test_tracker_routing.py @@ -27,6 +27,132 @@ def _get_json(url: str) -> dict: return json.loads(r.read()) +def test_tracker_send_json_ignores_broken_pipe_after_client_disconnect(): + """A disconnected client must not dump a BrokenPipe traceback from the tracker.""" + from meshnet_tracker.server import _TrackerHandler + + class BrokenPipeWriter: + def write(self, _body): + raise BrokenPipeError + + class DummyHandler: + wfile = BrokenPipeWriter() + + def send_response(self, _status): + pass + + def send_header(self, _name, _value): + pass + + def end_headers(self): + pass + + _TrackerHandler._send_json(DummyHandler(), 200, {"ok": True}) + + +def test_tracker_serves_health_while_proxy_request_is_in_flight(): + """Long inference proxy requests must not block heartbeats/health checks.""" + + class SlowChatHandler(http.server.BaseHTTPRequestHandler): + def log_message(self, fmt, *args): + pass + + def do_POST(self): + if self.path != "/v1/chat/completions": + self.send_response(404) + self.end_headers() + return + length = int(self.headers.get("Content-Length", 0)) + self.rfile.read(length) + time.sleep(2.0) + body = json.dumps({"choices": [{"message": {"content": "ok"}}]}).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) + + slow_node = http.server.HTTPServer(("127.0.0.1", 0), SlowChatHandler) + slow_thread = threading.Thread(target=slow_node.serve_forever, daemon=True) + slow_thread.start() + tracker = TrackerServer(heartbeat_timeout=60.0) + tracker_port = tracker.start() + proxy_error = [] + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": f"http://127.0.0.1:{slow_node.server_address[1]}", + "model": "slow-model", "hf_repo": "org/slow-model", "num_layers": 1, + "shard_start": 0, "shard_end": 0, "tracker_mode": True, + "hardware_profile": {}, "score": 1.0}, + ) + + def call_proxy(): + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/chat/completions", + {"model": "slow-model", "messages": [{"role": "user", "content": "hi"}]}, + ) + except Exception as exc: + proxy_error.append(exc) + + proxy_thread = threading.Thread(target=call_proxy) + proxy_thread.start() + time.sleep(0.2) + + with urllib.request.urlopen(f"http://127.0.0.1:{tracker_port}/v1/health", timeout=1.0) as resp: + assert resp.status == 200 + + proxy_thread.join(timeout=3.0) + assert not proxy_thread.is_alive() + assert not proxy_error + finally: + tracker.stop() + slow_node.shutdown() + slow_node.server_close() + slow_thread.join(timeout=1.0) + + +def test_tracker_registration_node_id_includes_wallet_prefix_and_stable_suffix(): + tracker = TrackerServer() + tracker_port = tracker.start() + wallet = "7j77FsPY1evV8tuf4Z73AVrWwxBEW1pvKwi4EvcRD3g" + try: + first = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9100", "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24, + "shard_start": 0, "shard_end": 21, "wallet_address": wallet, + "hardware_profile": {}, "score": 1.0}, + ) + second = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9100", "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24, + "shard_start": 0, "shard_end": 21, "wallet_address": wallet, + "hardware_profile": {}, "score": 1.0}, + ) + different_endpoint = _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24, + "shard_start": 20, "shard_end": 23, "wallet_address": wallet, + "hardware_profile": {}, "score": 1.0}, + ) + + assert first["node_id"].startswith("7j77FsPY-") + assert second["node_id"] == first["node_id"] + assert different_endpoint["node_id"].startswith("7j77FsPY-") + assert different_endpoint["node_id"] != first["node_id"] + + route_resp = _get_json( + f"http://127.0.0.1:{tracker_port}/v1/route?model=Qwen/Qwen2.5-0.5B-Instruct" + ) + assert route_resp["nodes"][0]["node_id"] == first["node_id"] + finally: + tracker.stop() + + def test_tracker_node_registration(): """A node can register with the tracker and receives a node_id.""" tracker = TrackerServer() @@ -131,6 +257,64 @@ def test_tracker_coverage_endpoint_reports_uncovered_ranges(): tracker.stop() +def test_tracker_coverage_endpoint_accepts_registered_hf_repo_or_short_name(): + """Coverage endpoint supports real HF models registered outside preset catalog.""" + tracker = TrackerServer() + tracker_port = tracker.start() + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24, + "shard_start": 0, "shard_end": 21, "hardware_profile": {}, "score": 1.0}, + ) + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9102", "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24, + "shard_start": 20, "shard_end": 23, "hardware_profile": {}, "score": 1.0}, + ) + + by_repo = _get_json( + f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen/Qwen2.5-0.5B-Instruct" + ) + by_short_name = _get_json( + f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen2.5-0.5B-Instruct" + ) + + assert by_repo["model"] == "Qwen/Qwen2.5-0.5B-Instruct" + assert by_repo["coverage"] == [ + {"start_layer": 0, "end_layer": 19, "node_count": 1}, + {"start_layer": 20, "end_layer": 21, "node_count": 2}, + {"start_layer": 22, "end_layer": 23, "node_count": 1}, + ] + assert by_short_name["coverage"] == by_repo["coverage"] + finally: + tracker.stop() + + +def test_tracker_models_endpoint_lists_registered_hf_repo_and_short_name_alias(): + tracker = TrackerServer() + tracker_port = tracker.start() + try: + _post_json( + f"http://127.0.0.1:{tracker_port}/v1/nodes/register", + {"endpoint": "http://127.0.0.1:9111", "model": "Qwen2.5-0.5B-Instruct", + "hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24, + "shard_start": 0, "shard_end": 23, "hardware_profile": {}, "score": 1.0}, + ) + + models_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/models") + + model = next(item for item in models_resp["data"] if item["id"] == "Qwen/Qwen2.5-0.5B-Instruct") + assert model["name"] == "Qwen2.5-0.5B-Instruct" + assert model["hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct" + assert model["aliases"] == ["Qwen/Qwen2.5-0.5B-Instruct", "Qwen2.5-0.5B-Instruct"] + assert model["shard_coverage_percentage"] == 100.0 + finally: + tracker.stop() + + def test_tracker_auto_assigns_new_node_to_uncovered_range_first(): """Capability-driven registration fills the first uncovered layer gap.""" tracker = TrackerServer(model_presets={