Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai
This commit is contained in:
Dobromir Popov
@@ -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": {
|
||||
|
||||
@@ -11,6 +11,7 @@ _packages = [
|
||||
"packages/sdk",
|
||||
"packages/contracts",
|
||||
"packages/p2p",
|
||||
"packages/relay",
|
||||
"packages/validator",
|
||||
]
|
||||
|
||||
|
||||
212
docs/INSTALL_WINDOWS.md
Normal file
212
docs/INSTALL_WINDOWS.md
Normal file
@@ -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.
|
||||
200
docs/TWO_MACHINE_TEST.md
Normal file
200
docs/TWO_MACHINE_TEST.md
Normal file
@@ -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.
|
||||
@@ -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(f" [node] WARNING: heartbeat failed: {exc}", flush=True)
|
||||
print(" [node] WARNING: re-registration failed", flush=True)
|
||||
outage_streak = 1
|
||||
else:
|
||||
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
|
||||
|
||||
@@ -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()
|
||||
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:
|
||||
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"}],
|
||||
}))
|
||||
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:
|
||||
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"}],
|
||||
}))
|
||||
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]:
|
||||
|
||||
@@ -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",
|
||||
|
||||
162
packages/p2p/meshnet_p2p/gossip.py
Normal file
162
packages/p2p/meshnet_p2p/gossip.py
Normal file
@@ -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)
|
||||
64
packages/p2p/meshnet_p2p/identity.py
Normal file
64
packages/p2p/meshnet_p2p/identity.py
Normal file
@@ -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
|
||||
136
packages/p2p/meshnet_p2p/mdns.py
Normal file
136
packages/p2p/meshnet_p2p/mdns.py
Normal file
@@ -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
|
||||
114
packages/p2p/meshnet_p2p/tls.py
Normal file
114
packages/p2p/meshnet_p2p/tls.py
Normal file
@@ -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:<hex>'."""
|
||||
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
|
||||
@@ -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*"]
|
||||
|
||||
10
packages/p2p/relay_bootstrap.json
Normal file
10
packages/p2p/relay_bootstrap.json
Normal file
@@ -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"
|
||||
}
|
||||
]
|
||||
}
|
||||
3
packages/relay/meshnet_relay/__init__.py
Normal file
3
packages/relay/meshnet_relay/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""meshnet-relay — NAT-traversal relay and gossip hub."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
60
packages/relay/meshnet_relay/cli.py
Normal file
60
packages/relay/meshnet_relay/cli.py
Normal file
@@ -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/<id> 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()
|
||||
60
packages/relay/meshnet_relay/peer_registry.py
Normal file
60
packages/relay/meshnet_relay/peer_registry.py
Normal file
@@ -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)
|
||||
224
packages/relay/meshnet_relay/server.py
Normal file
224
packages/relay/meshnet_relay/server.py
Normal file
@@ -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
|
||||
20
packages/relay/pyproject.toml
Normal file
20
packages/relay/pyproject.toml
Normal file
@@ -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*"]
|
||||
@@ -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)
|
||||
|
||||
88
packages/tracker/meshnet_tracker/gossip.py
Normal file
88
packages/tracker/meshnet_tracker/gossip.py
Normal file
@@ -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
|
||||
372
packages/tracker/meshnet_tracker/raft.py
Normal file
372
packages/tracker/meshnet_tracker/raft.py
Normal file
@@ -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
|
||||
)
|
||||
@@ -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()
|
||||
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/<node_id>/heartbeat -> ['', 'v1', 'nodes', '<id>', '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 '<empty>'}",
|
||||
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
|
||||
|
||||
163
scripts/test_lan_inference.py
Normal file
163
scripts/test_lan_inference.py
Normal file
@@ -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())
|
||||
371
tests/test_gossip_and_relay.py
Normal file
371
tests/test_gossip_and_relay.py
Normal file
@@ -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
|
||||
356
tests/test_mining_cli.py
Normal file
356
tests/test_mining_cli.py
Normal file
@@ -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
|
||||
@@ -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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
255
tests/test_tracker_consensus.py
Normal file
255
tests/test_tracker_consensus.py
Normal file
@@ -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"
|
||||
)
|
||||
@@ -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={
|
||||
|
||||
Reference in New Issue
Block a user