Compare commits
3 Commits
d701ae9ba2
...
worktree-f
| Author | SHA1 | Date | |
|---|---|---|---|
|
080d49b2c2 | ||
|
|
a2258d3df4 | ||
|
|
65f3ee6a85 |
Submodule .claude/worktrees/feat+us-016 deleted from 080d49b2c2
@@ -1,206 +0,0 @@
|
||||
# US-016 — Mining-style node startup CLI + live dashboard
|
||||
|
||||
## Goal
|
||||
|
||||
Replace the bare flag-driven `meshnet-node start` with a wizard-guided first-run experience modelled on GPU mining clients (like PhoenixMiner, lolMiner, etc.). After the wizard, the terminal switches to a live status dashboard showing real-time node health and earnings.
|
||||
|
||||
## Wizard flow (first run only)
|
||||
|
||||
```
|
||||
╔══════════════════════════════════════════════════════════╗
|
||||
║ meshnet-node v0.1.0 ║
|
||||
║ Distributed AI Inference — Node Setup ║
|
||||
╚══════════════════════════════════════════════════════════╝
|
||||
|
||||
Detecting hardware...
|
||||
GPU 0: NVIDIA RTX 4090 24 GB VRAM ✓
|
||||
GPU 1: NVIDIA RTX 3090 24 GB VRAM ✓
|
||||
|
||||
Select a model to serve:
|
||||
|
||||
# Model Layers NF4 INT8 BF16
|
||||
1 Llama-3-70B-Instruct 80 ✓18GB ✓40GB ✗80GB
|
||||
2 Qwen-2.5-72B-Instruct 80 ✓19GB ✗41GB ✗81GB
|
||||
3 Mixtral-8x7B-Instruct-v0.1 32 ✓ 7GB ✓14GB ✓27GB
|
||||
4 Phi-3-medium-128k-instruct 40 ✓ 4GB ✓ 8GB ✓15GB
|
||||
5 [Browse HuggingFace…]
|
||||
|
||||
Enter number [1]: _
|
||||
|
||||
Quantization [nf4/int8/bf16] (nf4 recommended for 24GB): _
|
||||
|
||||
Download directory [~/.meshnet/models]: _
|
||||
|
||||
Tracker URL [http://localhost:8080]: _
|
||||
|
||||
Wallet path [~/.config/meshnet/wallet.json] (new wallet will be created): _
|
||||
|
||||
Config saved to ~/.config/meshnet/config.json
|
||||
Starting node…
|
||||
```
|
||||
|
||||
Second run with existing config:
|
||||
|
||||
```
|
||||
meshnet-node
|
||||
Reading config from ~/.config/meshnet/config.json
|
||||
Model: Llama-3-70B-Instruct Quant: nf4 Shard: layers 0–15
|
||||
Tracker: http://192.168.1.10:8080
|
||||
Starting…
|
||||
```
|
||||
|
||||
## Live dashboard (once running)
|
||||
|
||||
Renders every 2 seconds using `rich.live`. Fallback: plain-text status line if `rich` is unavailable or terminal is not a TTY (important for WSL2 / SSH).
|
||||
|
||||
```
|
||||
meshnet-node Llama-3-70B-Instruct [nf4] shard 0–15/80 up 00:03:22
|
||||
|
||||
GPU 0 RTX 4090 GPU ████████░░ 73% VRAM 18.2/24.0 GB 45°C
|
||||
GPU 1 RTX 3090 GPU ███░░░░░░░ 28% VRAM 8.7/24.0 GB 38°C
|
||||
|
||||
Tokens/sec ▁▂▃▄▅▆▇█ 42.3 t/s (EMA 30s)
|
||||
Requests 1,247 served 3 active
|
||||
Peers 8 connected (tracker: ✓ relay: ✓)
|
||||
TAI earned 0.00 TAI (payments active after US-006)
|
||||
Uptime 00:03:22
|
||||
|
||||
[q] quit [r] reset stats [c] compact view
|
||||
```
|
||||
|
||||
Compact mode (`--compact` or pressing `c`) shows a single status line:
|
||||
```
|
||||
[43t/s VRAM18.2GB req1247 peers8 up3m22s]
|
||||
```
|
||||
|
||||
## Implementation notes
|
||||
|
||||
### Hardware detection
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
def detect_gpus() -> list[dict]:
|
||||
gpus = []
|
||||
if torch.cuda.is_available():
|
||||
for i in range(torch.cuda.device_count()):
|
||||
props = torch.cuda.get_device_properties(i)
|
||||
gpus.append({
|
||||
"index": i,
|
||||
"name": props.name,
|
||||
"vram_gb": props.total_memory / 1e9,
|
||||
"backend": "cuda"
|
||||
})
|
||||
# ROCm / Apple Silicon stubs for later
|
||||
return gpus
|
||||
```
|
||||
|
||||
### Curated model list
|
||||
|
||||
`packages/node/meshnet_node/model_catalog.py` — a hardcoded list of `ModelPreset` dataclasses:
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class ModelPreset:
|
||||
name: str # display name
|
||||
hf_repo: str # HuggingFace repo ID
|
||||
num_layers: int
|
||||
vram_gb: dict # {"nf4": 18, "int8": 40, "bf16": 80}
|
||||
description: str # one-line description
|
||||
```
|
||||
|
||||
Initial list (expand over time):
|
||||
- `meta-llama/Meta-Llama-3-70B-Instruct` — 80L, NF4 18GB, INT8 40GB, BF16 80GB
|
||||
- `Qwen/Qwen2.5-72B-Instruct` — 80L, NF4 19GB, INT8 41GB, BF16 81GB
|
||||
- `mistralai/Mixtral-8x7B-Instruct-v0.1` — 32L, NF4 7GB, INT8 14GB, BF16 27GB
|
||||
- `microsoft/Phi-3-medium-128k-instruct` — 40L, NF4 4GB, INT8 8GB, BF16 15GB
|
||||
- `google/gemma-2-27b-it` — 46L, NF4 10GB, INT8 20GB, BF16 40GB
|
||||
|
||||
### HuggingFace Browse
|
||||
|
||||
```python
|
||||
from huggingface_hub import list_models
|
||||
|
||||
def browse_hf(top_n=20) -> list[dict]:
|
||||
models = list_models(
|
||||
pipeline_tag="text-generation",
|
||||
library="transformers",
|
||||
sort="downloads",
|
||||
direction=-1,
|
||||
limit=top_n,
|
||||
cardData=True,
|
||||
)
|
||||
return [{"repo": m.modelId, "downloads": m.downloads} for m in models]
|
||||
```
|
||||
|
||||
### Persistent config
|
||||
|
||||
`~/.config/meshnet/config.json`:
|
||||
```json
|
||||
{
|
||||
"model_hf_repo": "meta-llama/Meta-Llama-3-70B-Instruct",
|
||||
"quantization": "nf4",
|
||||
"download_dir": "~/.meshnet/models",
|
||||
"tracker_url": "http://192.168.1.10:8080",
|
||||
"wallet_path": "~/.config/meshnet/wallet.json",
|
||||
"shard_start": null,
|
||||
"shard_end": null,
|
||||
"updatedAt": "2026-06-29T..."
|
||||
}
|
||||
```
|
||||
|
||||
`shard_start`/`shard_end`: null means tracker auto-assigns. User can pin a range for dedicated partial-model nodes.
|
||||
|
||||
### CLI flags
|
||||
|
||||
All wizard answers are overridable without re-running the wizard:
|
||||
|
||||
```
|
||||
meshnet-node [start]
|
||||
--model <hf-repo-id> # e.g. meta-llama/Meta-Llama-3-70B-Instruct
|
||||
--quantization [bf16|int8|nf4]
|
||||
--download-dir <path>
|
||||
--tracker <url>
|
||||
--wallet <path>
|
||||
--shard-start <int> # pin shard range (optional)
|
||||
--shard-end <int>
|
||||
--reset-config # ignore saved config, re-run wizard
|
||||
--no-tui # plain-text output (for CI / headless)
|
||||
--compact # single-line status instead of full dashboard
|
||||
|
||||
meshnet-node models # list curated models and exit
|
||||
meshnet-node models --browse # list HF Hub top-20 and exit
|
||||
meshnet-node config # print current config and exit
|
||||
```
|
||||
|
||||
### WSL2 / non-TTY fallback
|
||||
|
||||
```python
|
||||
import sys, os
|
||||
|
||||
def is_interactive_tty() -> bool:
|
||||
return sys.stdout.isatty() and os.environ.get("TERM") not in ("dumb", "")
|
||||
|
||||
if not is_interactive_tty():
|
||||
# fall back to plain-text periodic status
|
||||
run_plain_status_loop(node)
|
||||
else:
|
||||
run_rich_dashboard(node)
|
||||
```
|
||||
|
||||
Do NOT use `termios`, `fcntl`, or `/dev/tty` — these break in Windows cmd.exe and some WSL2 terminal emulators.
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- `meshnet-node` with no args and no config → wizard starts
|
||||
- Wizard detects GPU and marks `[too large]` for models that exceed available VRAM
|
||||
- `meshnet-node models` prints curated list and exits
|
||||
- `meshnet-node models --browse` calls HF Hub API, prints top-20, exits
|
||||
- Second run (config exists) → skips wizard, starts immediately
|
||||
- `--reset-config` re-runs wizard even with config present
|
||||
- All wizard inputs override-able via CLI flags
|
||||
- Live rich dashboard renders and updates every 2s when running in a TTY
|
||||
- Falls back to plain-text when not a TTY (CI / WSL2 without TERM set)
|
||||
- Ctrl-C prints a clean summary line and exits 0
|
||||
- `python -m pytest` passes from repo root
|
||||
- Commit only this story's changes
|
||||
@@ -1,205 +0,0 @@
|
||||
# US-017 — P2P gossip, NAT-traversal relay node, and SSL/TLS
|
||||
|
||||
## Goal
|
||||
|
||||
Nodes must work behind NAT (home routers, cloud VMs without public IPs) and must communicate securely. Implement:
|
||||
|
||||
1. **SSL/TLS everywhere** — all HTTP between nodes/tracker is HTTPS; all WebSocket gossip is WSS
|
||||
2. **mDNS peer discovery** — nodes on the same LAN find each other automatically (no config)
|
||||
3. **WebSocket gossip PubSub** — nodes propagate join/leave/coverage-update events in near-real-time
|
||||
4. **Circuit relay node** — team-run public relay (`packages/relay`) that enables NAT traversal and bootstraps new nodes joining from the internet
|
||||
|
||||
Architecture is designed to migrate to libp2p GossipSub + Kademlia DHT without breaking the message schema (topic names and payload formats are stable contracts).
|
||||
|
||||
## Gossip protocol
|
||||
|
||||
### Transport
|
||||
|
||||
WebSocket (`wss://`) using the `websockets` Python library. Each node maintains persistent WSS connections to:
|
||||
- The relay node (always, bootstraps peer list)
|
||||
- Up to 8 direct peers (Kademlia-style target fanout; peers discovered via mDNS + relay peer list)
|
||||
|
||||
### Topics
|
||||
|
||||
All messages are JSON with an envelope:
|
||||
|
||||
```json
|
||||
{
|
||||
"topic": "node-join",
|
||||
"version": 1,
|
||||
"from_peer": "<peer_id>",
|
||||
"timestamp": "<iso8601>",
|
||||
"payload": { ... }
|
||||
}
|
||||
```
|
||||
|
||||
| Topic | Direction | Payload |
|
||||
|-------|-----------|---------|
|
||||
| `node-join` | broadcast | `{peer_id, addr, models: [{model_preset, shard_start, shard_end}], vram_gb, quant}` |
|
||||
| `node-leave` | broadcast | `{peer_id, reason}` |
|
||||
| `coverage-update` | broadcast | `{model_preset, coverage: [{start, end, count}]}` |
|
||||
| `heartbeat` | peer→relay | `{peer_id, addr, uptime_s, tokens_per_sec}` |
|
||||
| `peer-list` | relay→peer | `{peers: [{peer_id, addr}]}` |
|
||||
| `relay-announce` | relay→all | `{relay_id, relay_url, capacity}` |
|
||||
|
||||
Gossip fanout: each node re-broadcasts received messages to all its peers (simple flooding with `seen_ids` dedup, TTL=3 hops). Migration to GossipSub mesh routing is a later ADR.
|
||||
|
||||
### Peer ID
|
||||
|
||||
`peer_id = sha256(public_key)[:16].hex()` — generated on first run, stored in `~/.config/meshnet/identity.json`. The same keypair is used for TLS client certificates (mTLS) in future work.
|
||||
|
||||
## mDNS LAN discovery
|
||||
|
||||
Use Python `zeroconf` library. Service type: `_meshnet._tcp.local.`
|
||||
|
||||
```python
|
||||
from zeroconf import ServiceInfo, Zeroconf
|
||||
|
||||
info = ServiceInfo(
|
||||
"_meshnet._tcp.local.",
|
||||
f"{peer_id}._meshnet._tcp.local.",
|
||||
addresses=[socket.inet_aton(local_ip)],
|
||||
port=node_port,
|
||||
properties={"peer_id": peer_id, "version": "1"},
|
||||
)
|
||||
zc = Zeroconf()
|
||||
zc.register_service(info)
|
||||
```
|
||||
|
||||
On startup, nodes also browse for `_meshnet._tcp.local.` to discover existing nodes. mDNS is LAN-only (does not traverse routers), which is correct for LAN discovery.
|
||||
|
||||
## NAT traversal: circuit relay
|
||||
|
||||
### How it works
|
||||
|
||||
1. Node A (behind NAT) cannot accept inbound TCP connections
|
||||
2. Node A connects outbound to the public relay via WSS
|
||||
3. Node A tells the tracker: `"effective_addr": "wss://relay.meshnet.ai/relay/{peer_id_A}"`
|
||||
4. Node B (wants to call A) connects to the relay at the above URL
|
||||
5. Relay proxies the TCP stream between A and B
|
||||
|
||||
Hole-punching (direct connection via STUN) is attempted first (future work). Relay is the fallback.
|
||||
|
||||
### meshnet-relay
|
||||
|
||||
`packages/relay/meshnet_relay/server.py` — a standalone aiohttp server:
|
||||
|
||||
```
|
||||
GET /health → {status: ok}
|
||||
GET /v1/peers → [{peer_id, addr, last_seen}]
|
||||
POST /v1/gossip → receive a gossip message, fan out to connected peers
|
||||
WSS /ws → persistent gossip connection (subscribe to all topics)
|
||||
WSS /relay/{peer_id} → circuit relay proxy to that peer_id
|
||||
GET /v1/relay/capacity → {connected_peers: N, max_peers: 500}
|
||||
```
|
||||
|
||||
CLI:
|
||||
|
||||
```
|
||||
meshnet-relay [--port 8443] [--cert path/to/cert.pem] [--key path/to/key.pem]
|
||||
[--tracker-url http://...] [--max-peers 500]
|
||||
```
|
||||
|
||||
The relay can optionally proxy to the tracker (so `relay.meshnet.ai` is the single internet-visible endpoint).
|
||||
|
||||
## SSL/TLS setup
|
||||
|
||||
### Node certificate (self-signed, auto-generated)
|
||||
|
||||
On first run, `meshnet-node` generates a self-signed RSA-2048 cert valid for 10 years:
|
||||
|
||||
```python
|
||||
from cryptography import x509
|
||||
from cryptography.hazmat.primitives import hashes, serialization
|
||||
from cryptography.hazmat.primitives.asymmetric import rsa
|
||||
```
|
||||
|
||||
Cert saved to `~/.config/meshnet/node_cert.pem` + `node_key.pem`. Fingerprint stored in config and shared with tracker via heartbeat. Nodes connecting to each other validate the fingerprint (TOFU — trust on first use), not the CA chain.
|
||||
|
||||
### Relay certificate
|
||||
|
||||
The relay uses a real Let's Encrypt cert (cert-bot or acme.sh). The relay cert is pinned in `packages/p2p/relay_bootstrap.json`:
|
||||
|
||||
```json
|
||||
{
|
||||
"relays": [
|
||||
{
|
||||
"url": "wss://relay.meshnet.ai:8443",
|
||||
"cert_fingerprint": "sha256:<hex>",
|
||||
"operator": "meshnet-team"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### All HTTP switched to HTTPS
|
||||
|
||||
`meshnet-node` starts an HTTPS server using `ssl.SSLContext`. `meshnet-tracker` similarly. All outbound `httpx` / `aiohttp` calls use TLS verification against pinned fingerprints (not the system CA store — too many corporate proxies break this).
|
||||
|
||||
## Tracker changes
|
||||
|
||||
Heartbeat payload gains new fields:
|
||||
|
||||
```json
|
||||
{
|
||||
"peer_id": "a1b2c3d4e5f6a1b2",
|
||||
"effective_addr": "https://192.168.1.42:8001",
|
||||
"relay_addr": "wss://relay.meshnet.ai:8443/relay/a1b2c3d4e5f6a1b2",
|
||||
"cert_fingerprint": "sha256:...",
|
||||
"gossip_peers": ["peer_id_1", "peer_id_2"]
|
||||
}
|
||||
```
|
||||
|
||||
Tracker uses `effective_addr` (direct) or `relay_addr` (fallback) when building inference routes.
|
||||
|
||||
## Integration test
|
||||
|
||||
```
|
||||
tests/test_gossip_and_relay.py
|
||||
|
||||
scenario:
|
||||
1. Start a local relay (localhost:18443)
|
||||
2. Start node A (no inbound port — simulate NAT by binding to 127.0.0.1 only)
|
||||
3. Start node B (public-reachable on localhost)
|
||||
4. Both register with relay; relay peer-list includes both
|
||||
5. Node B sends a gossip node-join message
|
||||
6. Assert node A receives it within 500ms
|
||||
7. Start tracker; confirm tracker's node registry includes node A via relay_addr
|
||||
8. Send inference request; assert it routes through relay to node A
|
||||
```
|
||||
|
||||
## Package layout
|
||||
|
||||
```
|
||||
packages/relay/
|
||||
pyproject.toml
|
||||
meshnet_relay/
|
||||
__init__.py
|
||||
server.py # aiohttp relay + gossip hub + circuit relay proxy
|
||||
cli.py # meshnet-relay entrypoint
|
||||
peer_registry.py # in-memory {peer_id: {addr, last_seen, ...}}
|
||||
circuit_relay.py # WSS proxy between two peers
|
||||
|
||||
packages/p2p/
|
||||
meshnet_p2p/
|
||||
gossip.py # GossipClient — connect to relay + peers, pub/sub
|
||||
mdns.py # ZeroconfDiscovery — mDNS announce + browse
|
||||
identity.py # PeerIdentity — generate/load peer_id + keypair
|
||||
tls.py # cert generation, fingerprint, SSLContext helpers
|
||||
|
||||
packages/node/meshnet_node/
|
||||
gossip_integration.py # wires GossipClient into node lifecycle
|
||||
```
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- All node↔node and node↔tracker HTTP uses HTTPS; self-signed cert auto-generated on first run
|
||||
- `cert_fingerprint` included in heartbeat; tracker stores and logs it
|
||||
- mDNS: two nodes on the same LAN discover each other without manual tracker URL (test with two localhost processes using different mDNS names)
|
||||
- Relay: `meshnet-relay` starts, accepts WSS connections, fans out gossip messages to all connected peers
|
||||
- Circuit relay: node A (127.0.0.1-only) can receive a gossip message via the relay from node B
|
||||
- Tracker routes inference to node A using `relay_addr` when direct addr not reachable
|
||||
- `relay_bootstrap.json` exists in `packages/p2p/` with at least one entry (localhost for tests)
|
||||
- ADR-0010 documents the gossip architecture and libp2p migration path
|
||||
- `python -m pytest` passes from repo root
|
||||
- Commit only this story's changes
|
||||
@@ -1,157 +0,0 @@
|
||||
# US-018 — End-to-end two-machine LAN inference test
|
||||
|
||||
## Goal
|
||||
|
||||
Run real distributed inference across two physical machines: the Linux rig and a Windows 11 rig running WSL2. Document every setup step, firewall rule, and gotcha so this is repeatable. The test script exits 0 with token output and timing, proving the network works.
|
||||
|
||||
## Network topology for LAN test
|
||||
|
||||
```
|
||||
[Linux machine] [Windows 11 / WSL2]
|
||||
meshnet-tracker :8080 meshnet-node (shard B)
|
||||
meshnet-node :8001 (shard A, tracker-mode)
|
||||
meshnet-gateway :8000 (optional, for OpenAI-compat)
|
||||
|
||||
Client (either machine):
|
||||
scripts/test_lan_inference.py --tracker http://192.168.1.10:8080
|
||||
```
|
||||
|
||||
The Linux machine runs the tracker + the first-shard node (tracker-mode). The Windows/WSL2 machine runs the second-shard node. A small model (e.g. Phi-3-medium at BF16, fits on one GPU each) is split across both.
|
||||
|
||||
## WSL2 setup (Windows side)
|
||||
|
||||
`docs/INSTALL_WINDOWS.md` covers:
|
||||
|
||||
1. Enable WSL2: `wsl --install -d Ubuntu-24.04`
|
||||
2. CUDA in WSL2: install NVIDIA driver on Windows (NOT inside WSL); WSL2 gets CUDA automatically
|
||||
- Verify: `nvidia-smi` inside WSL2 should show GPU
|
||||
3. Install Python 3.11+ and pip inside WSL2
|
||||
4. `pip install -e packages/node packages/p2p` (clone repo first)
|
||||
5. Firewall: Windows Defender must allow inbound WSL2 → LAN on node port
|
||||
- PowerShell: `New-NetFirewallRule -DisplayName "meshnet-node" -Direction Inbound -Protocol TCP -LocalPort 8001 -Action Allow`
|
||||
6. WSL2 IP: WSL2 has its own NAT'd IP (172.x.x.x); to expose to LAN, either:
|
||||
- Option A: `netsh interface portproxy add v4tov4 listenport=8001 listenaddress=0.0.0.0 connectport=8001 connectaddress=$(wsl hostname -I)`
|
||||
- Option B: use the relay node (US-017) — no port forwarding needed
|
||||
|
||||
## Linux setup
|
||||
|
||||
Standard install (already done after US-016). Firewall:
|
||||
|
||||
```bash
|
||||
# If using ufw
|
||||
sudo ufw allow 8080/tcp # tracker
|
||||
sudo ufw allow 8001/tcp # node
|
||||
sudo ufw allow 8000/tcp # gateway (optional)
|
||||
```
|
||||
|
||||
## Model split
|
||||
|
||||
For the test, use a model that has enough layers to split meaningfully but fits comfortably in memory. Phi-3-medium-128k-instruct (40 layers, BF16 15GB) works on a single 24GB GPU on each machine:
|
||||
|
||||
- Linux node: layers 0–19 (tracker-mode, owns tokenizer + embed_tokens)
|
||||
- Windows/WSL2 node: layers 20–39
|
||||
|
||||
Start sequence:
|
||||
```bash
|
||||
# Terminal 1 (Linux) — tracker
|
||||
meshnet-tracker --port 8080
|
||||
|
||||
# Terminal 2 (Linux) — first-shard node (tracker-mode auto-detected because shard_start=0)
|
||||
meshnet-node --model microsoft/Phi-3-medium-128k-instruct \
|
||||
--quantization bf16 \
|
||||
--shard-start 0 --shard-end 19 \
|
||||
--tracker http://localhost:8080 \
|
||||
--port 8001
|
||||
|
||||
# Terminal 3 (Windows WSL2) — second-shard node
|
||||
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
|
||||
```
|
||||
|
||||
## Test script
|
||||
|
||||
`scripts/test_lan_inference.py`:
|
||||
|
||||
```python
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
End-to-end LAN inference test.
|
||||
Usage: python scripts/test_lan_inference.py --tracker http://192.168.1.10:8080
|
||||
"""
|
||||
import argparse, time, httpx, json
|
||||
|
||||
MESSAGES = [
|
||||
{"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, ___"},
|
||||
]
|
||||
|
||||
def run_test(tracker_url: str, gateway_url: str | None):
|
||||
# Discover inference entry point via tracker if gateway not given
|
||||
if not gateway_url:
|
||||
r = httpx.get(f"{tracker_url}/v1/tracker-nodes/phi-3-medium", timeout=5)
|
||||
r.raise_for_status()
|
||||
nodes = r.json()
|
||||
assert nodes, "No tracker-mode nodes registered — is the first-shard node running?"
|
||||
gateway_url = nodes[0]["url"]
|
||||
|
||||
print(f"Inference endpoint: {gateway_url}")
|
||||
print(f"Tracker: {tracker_url}")
|
||||
print()
|
||||
|
||||
for i, msg in enumerate(MESSAGES):
|
||||
t0 = time.monotonic()
|
||||
r = httpx.post(
|
||||
f"{gateway_url}/v1/chat/completions",
|
||||
json={"model": "phi-3-medium", "messages": [msg], "stream": False},
|
||||
timeout=60,
|
||||
)
|
||||
r.raise_for_status()
|
||||
data = r.json()
|
||||
elapsed = time.monotonic() - t0
|
||||
|
||||
content = data["choices"][0]["message"]["content"]
|
||||
tokens = data["usage"]["completion_tokens"]
|
||||
tps = tokens / elapsed if elapsed > 0 else 0
|
||||
|
||||
print(f"[{i+1}] Q: {msg['content']}")
|
||||
print(f" A: {content}")
|
||||
print(f" {tokens} tokens {elapsed:.2f}s {tps:.1f} t/s")
|
||||
print()
|
||||
|
||||
print("✓ All 3 requests completed successfully")
|
||||
|
||||
if __name__ == "__main__":
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--tracker", required=True)
|
||||
p.add_argument("--gateway", default=None)
|
||||
args = p.parse_args()
|
||||
run_test(args.tracker, args.gateway)
|
||||
```
|
||||
|
||||
## Docs: TWO_MACHINE_TEST.md
|
||||
|
||||
`docs/TWO_MACHINE_TEST.md` must cover:
|
||||
|
||||
1. Prerequisites (models downloaded on both machines, same model ID, complementary shard ranges)
|
||||
2. Start order: tracker first, then nodes, then test script
|
||||
3. How to verify nodes are registered: `GET /v1/nodes` on tracker
|
||||
4. How to verify coverage: `GET /v1/coverage/phi-3-medium` — all 40 layers must show node_count ≥ 1
|
||||
5. How to run the test script
|
||||
6. Expected output
|
||||
7. Latency breakdown: how to read per-hop latency from node logs
|
||||
8. **Known Issues** section — updated during actual test run with real gotchas
|
||||
|
||||
## Acceptance criteria
|
||||
|
||||
- `docs/INSTALL_WINDOWS.md` covers WSL2 + CUDA + meshnet-node install end-to-end
|
||||
- `docs/TWO_MACHINE_TEST.md` covers the full two-machine setup and test procedure
|
||||
- `scripts/test_lan_inference.py` exists and is executable
|
||||
- When run against a real two-machine LAN setup: script exits 0, prints 3 valid answers with timing
|
||||
- Coverage map shows 100% coverage (no gap) after both nodes register
|
||||
- Known Issues section in TWO_MACHINE_TEST.md contains at least the issues encountered during this test run
|
||||
- No new pytest failures from repo root (this story adds docs + a script, not new Python packages)
|
||||
- Commit only this story's changes
|
||||
@@ -382,107 +382,10 @@
|
||||
"notes": "Source issue: .scratch/distributed-inference-network/issues/15-ralph-agent-agnostic-status-aware.md",
|
||||
"dependsOn": [],
|
||||
"completionNotes": "Implemented by agent: status-aware helpers (_is_done, _needs_attention, _is_active, _is_in_design), 6-bucket _story_sets, attention dashboard section, _review_report Attention Required block, auto --include-revise, set-agent subcommand with persistent agent-config.json, _run_openrouter stub, custom agent support, list-parallel subcommand, and auto --parallel N worktree orchestration. All 65 tests pass."
|
||||
},
|
||||
{
|
||||
"id": "US-016",
|
||||
"title": "16 \u2014 Mining-style node startup CLI + live dashboard",
|
||||
"description": "Replace the bare flag-driven node CLI with a wizard-guided first-run experience (like a GPU mining client) followed by a live terminal dashboard once the node is running. On first run, the wizard auto-detects GPU VRAM, presents a curated list of compatible models with VRAM requirements at each quantization level, lets the user pick a download location, and writes a persistent config file so subsequent starts are one command. Once the node is running, the wizard gives way to a rich live status panel showing: GPU temp + VRAM used, tokens/sec, requests served, peers connected, TAI earned (stub until US-006 is live). A Browse HuggingFace option calls the HF Hub API so users can load any HF model beyond the curated list.",
|
||||
"acceptanceCriteria": [
|
||||
" with no args and no config file enters the interactive setup wizard",
|
||||
"Wizard step 1: auto-detect GPU(s) via torch.cuda / torch.version.hip; print GPU name + total VRAM",
|
||||
"Wizard step 2: show curated model list (name, HF repo, layers, VRAM@NF4/INT8/BF16); mark models that do NOT fit available VRAM as [too large]",
|
||||
"Wizard step 3: offer [B] Browse HuggingFace \u2014 calls HF Hub API (huggingface_hub.list_models filtered by pipeline_tag=text-generation, sorted by downloads, top 20) and lets user enter a custom HF repo ID",
|
||||
"Wizard step 4: prompt for download directory (default ~/.meshnet/models/); validate writable; show estimated disk usage for chosen model+quantization",
|
||||
"Wizard step 5: prompt for tracker URL (default http://localhost:8080); validate connection",
|
||||
"Wizard writes ~/.config/meshnet/config.json; second run skips wizard and starts directly",
|
||||
"All wizard values overridable via CLI flags: --model, --download-dir, --quantization [bf16|int8|nf4], --tracker, --wallet, --reset-config",
|
||||
"Once node is running, wizard clears and a live dashboard renders every 2s (rich.live): GPU util%, VRAM used/total, tokens/sec (EMA), requests served, TAI earned (stub 0.0), peers connected, uptime, current model/shard range",
|
||||
"Dashboard exits cleanly on Ctrl-C with a summary line",
|
||||
"Works inside WSL2 (no termios/ioctl calls that fail on Windows terminal; fall back to plain-text status if rich is not available)",
|
||||
" passes from repo root",
|
||||
"Commit only this story changes"
|
||||
],
|
||||
"priority": 16,
|
||||
"status": "done",
|
||||
"notes": "Source issue: .scratch/distributed-inference-network/issues/16-mining-cli-ux.md",
|
||||
"dependsOn": [
|
||||
"US-004",
|
||||
"US-012"
|
||||
],
|
||||
"completionNotes": "Implemented: mining-style wizard with GPU detection, curated model list (7 models with NF4/INT8/BF16 VRAM requirements), HF Hub browse, persistent config, rich live dashboard with plain-text WSL2 fallback. 19 tests, 97 passed total."
|
||||
},
|
||||
{
|
||||
"id": "US-017",
|
||||
"title": "17 \u2014 P2P gossip, NAT-traversal relay node, and SSL/TLS",
|
||||
"description": "Add a gossip layer so nodes discover each other and propagate coverage-map changes without polling the tracker continuously. Introduce a publicly-hosted relay node (run by the team) that solves NAT traversal using circuit relay (Petals-style) and serves as the bootstrap peer list for new nodes. Encrypt all node-to-node and node-to-tracker communications with TLS. Gossip protocol: WebSocket-based PubSub over wss://, topics: node-join / node-leave / coverage-update / heartbeat. Peer discovery: mDNS (zeroconf) for LAN, public relay bootstrap list for internet. NAT traversal: relay node acts as TCP-level circuit relay when direct connection fails (hole-punching first, relay second). Architecture is designed to migrate to libp2p GossipSub + Kademlia DHT in a future story without breaking the message schema.",
|
||||
"acceptanceCriteria": [
|
||||
"All HTTP between nodes and tracker uses HTTPS (TLS 1.3); self-signed cert generated on first run and fingerprint pinned in config; relay node uses Let's Encrypt",
|
||||
"Nodes broadcast node-join / node-leave events over wss:// to known peers within 1s of registration",
|
||||
"mDNS peer discovery (Python zeroconf) finds other meshnet nodes on the same LAN segment without manual tracker URL entry",
|
||||
"Public relay bootstrap list (hardcoded relay URL + ) is consulted when no LAN peers found",
|
||||
"Relay node is a standalone meshnet package () with CLI: starts a WebSocket relay server + circuit relay + optional tracker proxy",
|
||||
"When a node behind NAT cannot accept inbound connections, the relay forwards its traffic; node advertises relay address (relay_url/node_id) to tracker as its effective endpoint",
|
||||
"Tracker accepts both direct node URLs and relay-proxied URLs in heartbeat payloads",
|
||||
"Integration test: two nodes in separate processes on localhost (simulating NAT) communicate via a local relay process; inference request routes correctly",
|
||||
"ADR-0010 documents the gossip protocol, relay architecture, and migration path to libp2p",
|
||||
" passes from repo root",
|
||||
"Commit only this story changes"
|
||||
],
|
||||
"priority": 17,
|
||||
"status": "done",
|
||||
"notes": "Source issue: .scratch/distributed-inference-network/issues/17-p2p-gossip-relay-ssl.md",
|
||||
"dependsOn": [
|
||||
"US-013",
|
||||
"US-014"
|
||||
],
|
||||
"completionNotes": "Implemented: packages/p2p (identity, TLS cert+fingerprint, GossipClient WSS PubSub, MdnsDiscovery with zeroconf optional), packages/relay (RelayServer gossip hub + circuit relay proxy, meshnet-relay CLI), tracker extended with relay_addr/cert_fingerprint/peer_id, relay_bootstrap.json, ADR-0010. 18 new tests; 115 total passed."
|
||||
},
|
||||
{
|
||||
"id": "US-018",
|
||||
"title": "18 \u2014 End-to-end two-machine LAN inference test",
|
||||
"description": "Prove the network works across two real machines: the Linux rig (this machine) and a Windows 11 rig running WSL2. One machine runs the tracker + first-shard node (inference entry point). The other machine runs a second-shard node. A client sends a real inference request and receives a streamed response. This story is primarily a test plan + setup guide + test execution script; it produces documented evidence (logs, timing, token output) that real distributed inference works. It also surfaces any real-world issues (port forwarding, CUDA driver version mismatches, WSL2 CUDA passthrough, model download paths) that need fixing.",
|
||||
"acceptanceCriteria": [
|
||||
"docs/INSTALL_WINDOWS.md exists: step-by-step WSL2 + CUDA + meshnet-node install on Windows 11",
|
||||
"docs/TWO_MACHINE_TEST.md exists: how to start tracker on machine A, node on machine B, run inference, interpret output",
|
||||
"A test script scripts/test_lan_inference.py: given --tracker-url, --gateway-url, sends 3 chat completion requests, asserts valid OpenAI format, prints token count + latency + which nodes served each request",
|
||||
"Both machines can reach each other on LAN (documented: firewall rules, port list)",
|
||||
"At least one successful inference recorded: the test script exits 0 with output showing tokens generated and node IDs",
|
||||
"Latency breakdown logged: gateway\u2192node-A, node-A\u2192node-B, node-B\u2192gateway (approximate, from server logs)",
|
||||
"Known issues during test documented in docs/TWO_MACHINE_TEST.md under a Known Issues section",
|
||||
"Commit only this story changes"
|
||||
],
|
||||
"priority": 18,
|
||||
"status": "open",
|
||||
"notes": "Source issue: .scratch/distributed-inference-network/issues/18-two-machine-lan-test.md",
|
||||
"dependsOn": [
|
||||
"US-016",
|
||||
"US-017"
|
||||
],
|
||||
"completionNotes": null
|
||||
},
|
||||
{
|
||||
"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.",
|
||||
"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",
|
||||
"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",
|
||||
"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",
|
||||
"notes": "Architecture decision: Raft for assignments (strong consistency) + CRDT gossip for liveness (eventual consistency). User approved 2026-06-29.",
|
||||
"dependsOn": ["US-017"],
|
||||
"completionNotes": null
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"updatedAt": "2026-06-29T15:35:00.000Z",
|
||||
"updatedAt": "2026-06-29T13:30:00.000Z",
|
||||
"statusVocabulary": {
|
||||
"open": "Not started",
|
||||
"in-design": "Decisions pending before implementation can begin",
|
||||
|
||||
199
QUICKSTART.md
199
QUICKSTART.md
@@ -1,199 +0,0 @@
|
||||
# Quickstart — Running a node and testing inference
|
||||
|
||||
This guide gets you from zero to a live inference request in three terminals.
|
||||
Tested on: AMD Ryzen AI Max (Strix Halo APU), 124 GB RAM, Linux, CPU inference.
|
||||
|
||||
---
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
# Clone and enter repo
|
||||
cd /run/media/popov/d/DEV/repos/d-popov.com/AI
|
||||
|
||||
# Install Python packages (editable — picks up code changes immediately)
|
||||
.venv/bin/pip install -e packages/tracker packages/node packages/p2p packages/relay
|
||||
|
||||
# CPU-only PyTorch (skip if you have CUDA/ROCm already)
|
||||
.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu
|
||||
|
||||
# HuggingFace model libraries
|
||||
.venv/bin/pip install transformers accelerate
|
||||
```
|
||||
|
||||
> **NVIDIA GPU (CUDA):** replace the torch line with `pip install torch` (default index).
|
||||
> **AMD GPU (ROCm):** `pip install torch --index-url https://download.pytorch.org/whl/rocm6.2`
|
||||
|
||||
---
|
||||
|
||||
## Step 1 — Start the tracker (Terminal 1)
|
||||
|
||||
```bash
|
||||
cd /run/media/popov/d/DEV/repos/d-popov.com/AI
|
||||
.venv/bin/meshnet-tracker start --port 8080
|
||||
```
|
||||
|
||||
Expected output:
|
||||
```
|
||||
Tracker listening on 0.0.0.0:8080
|
||||
```
|
||||
|
||||
Keep this terminal open.
|
||||
|
||||
---
|
||||
|
||||
## Step 2 — Start a node (Terminal 2)
|
||||
|
||||
### Recommended model: Qwen2.5-0.5B-Instruct
|
||||
|
||||
- 0.5B parameters, ~1 GB in BF16
|
||||
- No HuggingFace account or license required
|
||||
- Downloads once to `~/.meshnet/models/`, cached for future runs
|
||||
- 24 transformer layers (auto-detected — no need to specify)
|
||||
|
||||
```bash
|
||||
cd /run/media/popov/d/DEV/repos/d-popov.com/AI
|
||||
HF_HOME=/run/media/popov/d/DEV/models \
|
||||
.venv/bin/meshnet-node start \
|
||||
--model-id Qwen/Qwen2.5-0.5B-Instruct \
|
||||
--quantization bfloat16 \
|
||||
--tracker http://localhost:8080 \
|
||||
--port 8001
|
||||
```
|
||||
|
||||
Shard range is **auto-detected** from the curated catalog (no network call for known
|
||||
models). For unknown repos, the node fetches only `config.json` (~1 KB) to read
|
||||
`num_hidden_layers`. You can still pass `--shard-start` / `--shard-end` explicitly
|
||||
to run a partial shard on one machine.
|
||||
|
||||
Expected output (after model loads):
|
||||
```
|
||||
Auto-detected 24 layers → shard 0–23
|
||||
================================
|
||||
meshnet-node ready
|
||||
Wallet: <address>
|
||||
Model ID: Qwen/Qwen2.5-0.5B-Instruct
|
||||
Shard: layers 0–23
|
||||
Quantization: bfloat16
|
||||
Endpoint: http://<host>:8001
|
||||
Hardware: CPU
|
||||
================================
|
||||
```
|
||||
|
||||
### Other model options (all CPU-friendly)
|
||||
|
||||
| Model | HF repo | Layers | BF16 size | Notes |
|
||||
|-------|---------|--------|-----------|-------|
|
||||
| Qwen2.5-0.5B | `Qwen/Qwen2.5-0.5B-Instruct` | 24 | ~1 GB | Fastest, no gating |
|
||||
| Qwen2.5-1.5B | `Qwen/Qwen2.5-1.5B-Instruct` | 28 | ~3 GB | Better quality |
|
||||
| Phi-3-mini | `microsoft/Phi-3-mini-4k-instruct` | 32 | ~7.5 GB | Best CPU quality |
|
||||
| Llama-3.2-1B | `meta-llama/Llama-3.2-1B-Instruct` | 16 | ~2 GB | Requires HF login |
|
||||
| Llama-3.2-3B | `meta-llama/Llama-3.2-3B-Instruct` | 28 | ~6 GB | Requires HF login |
|
||||
|
||||
For gated models (Llama), run `huggingface-cli login` first.
|
||||
|
||||
---
|
||||
|
||||
## Step 3 — Send an inference request (Terminal 3)
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8001/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "qwen2.5-0.5b",
|
||||
"messages": [{"role": "user", "content": "What is 7 times 8? Answer in one word."}],
|
||||
"stream": false
|
||||
}' | python3 -m json.tool
|
||||
```
|
||||
|
||||
Or use the test script:
|
||||
|
||||
```bash
|
||||
.venv/bin/python scripts/test_lan_inference.py \
|
||||
--tracker http://localhost:8080 \
|
||||
--gateway http://localhost:8001
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Two-node split (same machine, two terminals)
|
||||
|
||||
Split Qwen2.5-0.5B's 24 layers across two node processes to test the sharded pipeline:
|
||||
|
||||
**Node A — layers 0–11 (tracker mode, serves chat completions):**
|
||||
```bash
|
||||
HF_HOME=/run/media/popov/d/DEV/models \
|
||||
.venv/bin/meshnet-node start \
|
||||
--model-id Qwen/Qwen2.5-0.5B-Instruct \
|
||||
--shard-start 0 --shard-end 11 \
|
||||
--quantization bfloat16 \
|
||||
--tracker http://localhost:8080 \
|
||||
--port 8001
|
||||
```
|
||||
|
||||
**Node B — layers 12–23:**
|
||||
```bash
|
||||
HF_HOME=/run/media/popov/d/DEV/models \
|
||||
.venv/bin/meshnet-node start \
|
||||
--model-id Qwen/Qwen2.5-0.5B-Instruct \
|
||||
--shard-start 12 --shard-end 23 \
|
||||
--quantization bfloat16 \
|
||||
--tracker http://localhost:8080 \
|
||||
--port 8002
|
||||
```
|
||||
|
||||
Send the request to Node A — it tokenizes, runs layers 0–13, passes binary
|
||||
activations to Node B, and streams the final response back.
|
||||
|
||||
---
|
||||
|
||||
## Two-machine LAN test (Linux + Windows/WSL2)
|
||||
|
||||
See `docs/TWO_MACHINE_TEST.md` (created by US-018).
|
||||
|
||||
---
|
||||
|
||||
## Browse available models
|
||||
|
||||
```bash
|
||||
# Show curated list with VRAM requirements
|
||||
.venv/bin/meshnet-node models
|
||||
|
||||
# Browse HuggingFace Hub top-20 text-generation models
|
||||
.venv/bin/meshnet-node models --browse
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Start with the interactive wizard
|
||||
|
||||
```bash
|
||||
# First run: wizard detects GPU, shows model list, saves config
|
||||
.venv/bin/meshnet-node
|
||||
|
||||
# Subsequent runs: starts directly from saved config
|
||||
.venv/bin/meshnet-node
|
||||
|
||||
# Re-run wizard even with saved config
|
||||
.venv/bin/meshnet-node --reset-config
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Start the relay node (for NAT traversal)
|
||||
|
||||
```bash
|
||||
.venv/bin/pip install -e packages/relay
|
||||
.venv/bin/meshnet-relay --port 8765
|
||||
```
|
||||
|
||||
Nodes behind NAT connect to the relay and advertise their relay address to the
|
||||
tracker. See `docs/adr/0010-p2p-gossip-and-nat-relay.md`.
|
||||
|
||||
---
|
||||
|
||||
## Run all tests
|
||||
|
||||
```bash
|
||||
.venv/bin/python -m pytest -q
|
||||
```
|
||||
@@ -11,6 +11,7 @@ _packages = [
|
||||
"packages/sdk",
|
||||
"packages/contracts",
|
||||
"packages/p2p",
|
||||
"packages/relay",
|
||||
"packages/validator",
|
||||
]
|
||||
|
||||
|
||||
@@ -1,67 +0,0 @@
|
||||
# ADR-0010: P2P gossip, NAT-traversal relay, and TLS
|
||||
|
||||
## Status: Accepted
|
||||
|
||||
## Context
|
||||
|
||||
All node-to-node and node-to-tracker communication in the prototype is plain HTTP over a LAN or direct-IP internet connection. This has three problems:
|
||||
|
||||
1. **NAT blocking**: Most home and cloud nodes cannot accept inbound TCP connections.
|
||||
2. **No encryption**: Activations and heartbeats are in plaintext.
|
||||
3. **Polling overhead**: Nodes poll the tracker for coverage changes every 30s. This is slow to react to node churn and does not scale past a few hundred nodes.
|
||||
|
||||
The reference implementation (Petals) solves this with libp2p — GossipSub for pub/sub and Kademlia DHT for peer discovery. We adopt the same goals but start with simpler, more stable building blocks that can be swapped for libp2p later without changing the message schema.
|
||||
|
||||
## Decisions
|
||||
|
||||
### 1. TLS everywhere
|
||||
|
||||
All HTTP between nodes, tracker, and gateway uses HTTPS (TLS 1.3). Self-signed certificates are auto-generated on first node start and stored in `~/.config/meshnet/`. The certificate fingerprint is included in every heartbeat and gossip envelope. Nodes use TOFU (trust on first use) — they accept a peer's cert on first contact and pin the fingerprint; connections from the same peer with a different fingerprint are rejected.
|
||||
|
||||
The relay node uses a real CA-signed certificate (Let's Encrypt) because it is the internet-facing bootstrap point.
|
||||
|
||||
### 2. mDNS for LAN peer discovery
|
||||
|
||||
Python `zeroconf` library. Service type: `_meshnet._tcp.local.`. A node announces itself on startup and browses for existing peers. This is zero-config discovery for home and lab networks. mDNS does not traverse routers, which is correct — LAN discovery should not bleed into the internet.
|
||||
|
||||
### 3. WebSocket PubSub for gossip
|
||||
|
||||
Each node maintains persistent WSS connections to the relay and up to 8 direct peers. Messages use a stable JSON envelope with a `topic`, `version`, `from_peer`, and `payload`. Topics: `node-join`, `node-leave`, `coverage-update`, `heartbeat`, `peer-list`, `relay-announce`.
|
||||
|
||||
Simple flooding with `seen_ids` dedup and TTL=3 is good enough for the prototype. The message schema is stable; the fanout mechanism can be replaced with GossipSub mesh routing without changing the schema.
|
||||
|
||||
### 4. Circuit relay node for NAT traversal
|
||||
|
||||
A team-operated public relay (`packages/relay`, CLI: `meshnet-relay`) is the internet bootstrap point. A node behind NAT:
|
||||
|
||||
1. Connects outbound to the relay via WSS
|
||||
2. Advertises `relay_addr = wss://relay.meshnet.ai:8443/relay/{peer_id}` to the tracker
|
||||
3. Other nodes proxy connections through the relay when the direct addr is not reachable
|
||||
|
||||
Hole-punching (STUN + simultaneous TCP open) is deferred to a future story. Circuit relay is the reliable fallback.
|
||||
|
||||
The relay is stateless in terms of inference — it only proxies bytes. It does not decrypt activations.
|
||||
|
||||
### 5. Bootstrap peer list
|
||||
|
||||
`packages/p2p/relay_bootstrap.json` contains the team-operated relay endpoints with their TLS fingerprints. New nodes load this file on startup to find their first peer. The list is bundled with the package and updated via pip upgrades.
|
||||
|
||||
### Migration path to libp2p
|
||||
|
||||
When the network has enough volume to justify the complexity:
|
||||
|
||||
1. Replace the WebSocket gossip layer with libp2p GossipSub (same topics and payload schemas, different transport)
|
||||
2. Replace mDNS + relay peer list with Kademlia DHT
|
||||
3. Replace circuit relay with libp2p circuit relay v2
|
||||
|
||||
The gossip envelope schema (`topic`, `version`, `from_peer`, `payload`) is the stable contract. As long as messages on the wire are identical, the transport layer can be swapped without touching node business logic.
|
||||
|
||||
## Alternatives rejected
|
||||
|
||||
**libp2p from the start**: `py-libp2p` is experimental and not production-ready. A Go libp2p sidecar is operationally complex. The benefits of real libp2p (mesh routing, Kademlia DHT, hole-punching) are not needed until we have hundreds of nodes.
|
||||
|
||||
**NATS**: Stable and fast but requires a central NATS server. Adds operational dependency and contradicts the P2P goal.
|
||||
|
||||
**ZeroMQ**: No NAT traversal built in. Requires manual topology management.
|
||||
|
||||
**No gossip (keep polling)**: Does not scale; slow to react to node churn; misses the relay/NAT requirement.
|
||||
@@ -107,13 +107,12 @@ class TorchModelShard:
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
|
||||
self.layers = _model_layers(self.model)
|
||||
self.total_layers = len(self.layers)
|
||||
# shard_end is INCLUSIVE (last layer index, 0-based), matching the CLI convention.
|
||||
if shard_end >= self.total_layers:
|
||||
if shard_end > self.total_layers:
|
||||
raise ValueError(
|
||||
f"shard_end {shard_end} exceeds last layer index {self.total_layers - 1}"
|
||||
f"shard_end {shard_end} exceeds total layer count {self.total_layers}"
|
||||
)
|
||||
self.is_head = shard_start == 0
|
||||
self.is_tail = shard_end >= self.total_layers - 1
|
||||
self.is_tail = shard_end == self.total_layers
|
||||
self.hidden_size = int(
|
||||
getattr(self.model.config, "hidden_size", 0)
|
||||
or getattr(self.model.config, "n_embd", 0)
|
||||
@@ -169,135 +168,10 @@ class TorchModelShard:
|
||||
token_id = int(self.torch.argmax(logits[:, -1, :], dim=-1)[0].item())
|
||||
return self.tokenizer.decode([token_id], skip_special_tokens=True)
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 256,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
) -> str:
|
||||
"""Autoregressive generation using HF generate() — single-node (head+tail) mode."""
|
||||
if not self.is_head or not self.is_tail:
|
||||
raise ModelBackendError("local generation requires a full head+tail shard")
|
||||
encoded = self._encode_messages(messages)
|
||||
input_ids = encoded["input_ids"].to(self.device)
|
||||
attention_mask = encoded.get("attention_mask")
|
||||
if attention_mask is not None:
|
||||
attention_mask = attention_mask.to(self.device)
|
||||
pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None)
|
||||
do_sample = temperature != 1.0 or top_p != 1.0
|
||||
with self.torch.inference_mode():
|
||||
generated = self.model.generate(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
max_new_tokens=max(1, int(max_new_tokens)),
|
||||
do_sample=do_sample,
|
||||
temperature=temperature if do_sample else None,
|
||||
top_p=top_p if do_sample else None,
|
||||
pad_token_id=pad_token_id,
|
||||
)
|
||||
new_tokens = generated[0, input_ids.shape[-1]:]
|
||||
return self.tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
|
||||
|
||||
def generate_text_streaming(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 256,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
):
|
||||
"""Yield decoded token strings one at a time using HF TextIteratorStreamer."""
|
||||
if not self.is_head or not self.is_tail:
|
||||
raise ModelBackendError("streaming generation requires a full head+tail shard")
|
||||
import threading
|
||||
try:
|
||||
from transformers import TextIteratorStreamer # type: ignore[import]
|
||||
except ImportError:
|
||||
yield self.generate_text(messages, max_new_tokens, temperature, top_p)
|
||||
return
|
||||
|
||||
encoded = self._encode_messages(messages)
|
||||
input_ids = encoded["input_ids"].to(self.device)
|
||||
attention_mask = encoded.get("attention_mask")
|
||||
if attention_mask is not None:
|
||||
attention_mask = attention_mask.to(self.device)
|
||||
pad_token_id = getattr(self.tokenizer, "pad_token_id", None) or getattr(self.tokenizer, "eos_token_id", None)
|
||||
do_sample = temperature != 1.0 or top_p != 1.0
|
||||
|
||||
streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True)
|
||||
gen_kwargs = dict(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
max_new_tokens=max(1, int(max_new_tokens)),
|
||||
do_sample=do_sample,
|
||||
temperature=temperature if do_sample else None,
|
||||
top_p=top_p if do_sample else None,
|
||||
pad_token_id=pad_token_id,
|
||||
streamer=streamer,
|
||||
)
|
||||
t = threading.Thread(target=self.model.generate, kwargs=gen_kwargs, daemon=True)
|
||||
t.start()
|
||||
for token_text in streamer:
|
||||
yield token_text
|
||||
t.join()
|
||||
|
||||
def count_prompt_tokens(self, messages: list[dict]) -> int:
|
||||
"""Return tokenizer-backed prompt token count for OpenAI usage metadata."""
|
||||
encoded = self._encode_messages(messages)
|
||||
input_ids = encoded["input_ids"]
|
||||
return int(input_ids.shape[-1])
|
||||
|
||||
def count_text_tokens(self, text: str) -> int:
|
||||
"""Return tokenizer-backed completion token count for OpenAI usage metadata."""
|
||||
try:
|
||||
encoded = self.tokenizer(
|
||||
text,
|
||||
return_tensors="pt",
|
||||
add_special_tokens=False,
|
||||
)
|
||||
except TypeError:
|
||||
encoded = self.tokenizer(text, return_tensors="pt")
|
||||
return int(encoded["input_ids"].shape[-1])
|
||||
|
||||
def _encode_messages(self, messages: list[dict]) -> dict:
|
||||
"""Format messages with chat template (if available) and tokenize."""
|
||||
if hasattr(self.tokenizer, "apply_chat_template"):
|
||||
try:
|
||||
prompt_str = self.tokenizer.apply_chat_template(
|
||||
messages,
|
||||
add_generation_prompt=True,
|
||||
tokenize=False,
|
||||
)
|
||||
return dict(self.tokenizer(prompt_str, return_tensors="pt"))
|
||||
except Exception:
|
||||
pass
|
||||
prompt = " ".join(
|
||||
str(m.get("content", ""))
|
||||
for m in messages
|
||||
if isinstance(m, dict) and m.get("role") == "user"
|
||||
)
|
||||
return dict(self.tokenizer(prompt, return_tensors="pt"))
|
||||
|
||||
def _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
|
||||
position_embeddings = _rotary_position_embeddings(
|
||||
self.model,
|
||||
hidden_states,
|
||||
position_ids,
|
||||
)
|
||||
layer_attention_mask = _decoder_attention_mask(
|
||||
attention_mask,
|
||||
hidden_states,
|
||||
self.torch,
|
||||
)
|
||||
with self.torch.inference_mode():
|
||||
for layer in self.layers[self.shard_start:self.shard_end + 1]:
|
||||
hidden_states = _call_layer(
|
||||
layer,
|
||||
hidden_states,
|
||||
layer_attention_mask,
|
||||
position_ids,
|
||||
position_embeddings,
|
||||
)
|
||||
for layer in self.layers[self.shard_start:self.shard_end]:
|
||||
hidden_states = _call_layer(layer, hidden_states, attention_mask, position_ids)
|
||||
return hidden_states.to(self.torch.bfloat16)
|
||||
|
||||
def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload:
|
||||
@@ -362,60 +236,8 @@ def _position_ids(attention_mask: Any, torch: Any) -> Any:
|
||||
return position_ids.masked_fill(attention_mask == 0, 0).to(torch.long)
|
||||
|
||||
|
||||
def _decoder_attention_mask(attention_mask: Any, hidden_states: Any, torch: Any) -> Any:
|
||||
"""Build a causal additive mask for decoder layers called outside model.forward."""
|
||||
if attention_mask is None:
|
||||
return None
|
||||
if len(getattr(attention_mask, "shape", ())) != 2:
|
||||
return attention_mask
|
||||
batch_size, seq_len = attention_mask.shape
|
||||
if seq_len <= 1:
|
||||
return None if bool(attention_mask.all()) else attention_mask.to(hidden_states.dtype)
|
||||
|
||||
min_value = torch.finfo(hidden_states.dtype).min
|
||||
causal = torch.full(
|
||||
(seq_len, seq_len),
|
||||
min_value,
|
||||
dtype=hidden_states.dtype,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
causal = torch.triu(causal, diagonal=1)
|
||||
causal = causal[None, None, :, :].expand(batch_size, 1, seq_len, seq_len).clone()
|
||||
|
||||
padding = attention_mask.to(device=hidden_states.device)
|
||||
if not bool(padding.all()):
|
||||
causal = causal.masked_fill(padding[:, None, None, :] == 0, min_value)
|
||||
return causal
|
||||
|
||||
|
||||
def _rotary_position_embeddings(model: Any, hidden_states: Any, position_ids: Any) -> Any | None:
|
||||
"""Return model-level rotary embeddings required by newer HF decoder layers."""
|
||||
if position_ids is None:
|
||||
return None
|
||||
rotary = None
|
||||
if hasattr(model, "model") and hasattr(model.model, "rotary_emb"):
|
||||
rotary = model.model.rotary_emb
|
||||
elif hasattr(model, "transformer") and hasattr(model.transformer, "rotary_emb"):
|
||||
rotary = model.transformer.rotary_emb
|
||||
if rotary is None:
|
||||
return None
|
||||
return rotary(hidden_states, position_ids)
|
||||
|
||||
|
||||
def _call_layer(
|
||||
layer: Any,
|
||||
hidden_states: Any,
|
||||
attention_mask: Any,
|
||||
position_ids: Any,
|
||||
position_embeddings: Any | None = None,
|
||||
) -> Any:
|
||||
def _call_layer(layer: Any, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
|
||||
attempts = (
|
||||
{
|
||||
"attention_mask": attention_mask,
|
||||
"position_ids": position_ids,
|
||||
"position_embeddings": position_embeddings,
|
||||
"use_cache": False,
|
||||
},
|
||||
{
|
||||
"attention_mask": attention_mask,
|
||||
"position_ids": position_ids,
|
||||
@@ -450,7 +272,7 @@ def _tensor_from_bfloat16_bytes(body: bytes, shape: list[int], torch: Any) -> An
|
||||
|
||||
|
||||
def _int_tensor_header(tensor: Any) -> str:
|
||||
data = tensor.detach().cpu().long().contiguous()
|
||||
data = tensor.detach().cpu().to(tensor.int64).contiguous()
|
||||
raw = data.numpy().tobytes()
|
||||
shape = ",".join(str(dim) for dim in data.shape)
|
||||
encoded = base64.b64encode(raw).decode("ascii")
|
||||
|
||||
@@ -5,8 +5,6 @@ from __future__ import annotations
|
||||
import json
|
||||
import socket
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
@@ -34,40 +32,6 @@ def _get_json(url: str, timeout: float = 10.0) -> dict:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def _start_heartbeat(
|
||||
tracker_url: str,
|
||||
node_id: str,
|
||||
register_payload: dict,
|
||||
interval: float = 20.0,
|
||||
) -> threading.Thread:
|
||||
"""Daemon thread: sends heartbeats; re-registers automatically on 404 (tracker restart)."""
|
||||
def _loop() -> None:
|
||||
nonlocal node_id
|
||||
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||||
while True:
|
||||
time.sleep(interval)
|
||||
try:
|
||||
_post_json(hb_url, {})
|
||||
except urllib.error.HTTPError as exc:
|
||||
if exc.code == 404:
|
||||
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)
|
||||
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)
|
||||
except Exception as exc:
|
||||
print(f" [node] WARNING: heartbeat failed: {exc}", flush=True)
|
||||
|
||||
t = threading.Thread(target=_loop, daemon=True, name="heartbeat")
|
||||
t.start()
|
||||
return t
|
||||
|
||||
|
||||
def run_startup(
|
||||
tracker_url: str,
|
||||
port: int = 0,
|
||||
@@ -120,7 +84,7 @@ def run_startup(
|
||||
if probationary_line is not None:
|
||||
print(f" {probationary_line}", flush=True)
|
||||
|
||||
if model_id: # treat "" the same as None — no explicit model given
|
||||
if model_id is not None:
|
||||
# Auto-detect shard range from model config if not explicitly provided
|
||||
if shard_start is None or shard_end is None:
|
||||
detected = _detect_num_layers(model_id)
|
||||
@@ -129,23 +93,6 @@ def run_startup(
|
||||
f"Could not read num_hidden_layers from {model_id} config. "
|
||||
"Pass --shard-start and --shard-end explicitly."
|
||||
)
|
||||
# When no explicit shard range given, ask the tracker if there's a gap for this model.
|
||||
if shard_start is None and shard_end is None:
|
||||
try:
|
||||
qs = urllib.parse.urlencode({
|
||||
"device": device, "vram_mb": vram_mb, "hf_repo": model_id,
|
||||
})
|
||||
net_asgn = _get_json(f"{tracker_url}/v1/network/assign?{qs}", timeout=5.0)
|
||||
if net_asgn.get("hf_repo") == model_id and net_asgn.get("gap_found"):
|
||||
shard_start = net_asgn["shard_start"]
|
||||
shard_end = net_asgn["shard_end"]
|
||||
print(
|
||||
f" Tracker found uncovered shard: "
|
||||
f"layers {shard_start}–{shard_end} (of {detected})",
|
||||
flush=True,
|
||||
)
|
||||
except Exception:
|
||||
pass # No other nodes registered yet — default to full model below
|
||||
shard_start = shard_start if shard_start is not None else 0
|
||||
shard_end = shard_end if shard_end is not None else detected - 1
|
||||
print(f" Auto-detected {detected} layers → shard {shard_start}–{shard_end}", flush=True)
|
||||
@@ -158,46 +105,16 @@ def run_startup(
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
)
|
||||
actual_port = node.start()
|
||||
total_layers = getattr(node.backend, "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}"
|
||||
else:
|
||||
shard_label = f"layers {shard_start}–{shard_end}"
|
||||
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)
|
||||
reg_payload = {
|
||||
"endpoint": endpoint,
|
||||
"model": model_id.split("/")[-1],
|
||||
"hf_repo": model_id,
|
||||
"num_layers": total_layers,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (shard_start == 0),
|
||||
}
|
||||
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)
|
||||
except Exception as exc:
|
||||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||||
|
||||
print(
|
||||
f"\n{'=' * 32}\n"
|
||||
f"meshnet-node ready\n"
|
||||
f" Wallet: {address}\n"
|
||||
f" Model ID: {model_id}\n"
|
||||
f" Shard: {shard_label}\n"
|
||||
f" Shard: layers {shard_start}–{shard_end}\n"
|
||||
f" Quantization: {quantization}\n"
|
||||
f" Endpoint: {endpoint}\n"
|
||||
f" Hardware: {device.upper()}\n"
|
||||
@@ -208,77 +125,7 @@ def run_startup(
|
||||
if shard_start is not None or shard_end is not None:
|
||||
raise ValueError("--shard-start / --shard-end require --model-id")
|
||||
|
||||
# 3a. Auto-join: query tracker for network-wide HF model assignment.
|
||||
print("Querying tracker for network assignment...", flush=True)
|
||||
assign_qs = urllib.parse.urlencode({"device": device, "vram_mb": vram_mb})
|
||||
net_assignment: dict = {}
|
||||
try:
|
||||
net_assignment = _get_json(f"{tracker_url}/v1/network/assign?{assign_qs}")
|
||||
except Exception as exc:
|
||||
print(f" (auto-join unavailable: {exc})", flush=True)
|
||||
assigned_hf_repo: str | None = net_assignment.get("hf_repo")
|
||||
_gap_found: bool = bool(net_assignment.get("gap_found", False))
|
||||
|
||||
if assigned_hf_repo and _gap_found:
|
||||
assigned_shard_start: int = net_assignment["shard_start"]
|
||||
assigned_shard_end: int = net_assignment["shard_end"]
|
||||
assigned_num_layers: int = net_assignment["num_layers"]
|
||||
print(
|
||||
f" Assigned: {assigned_hf_repo} "
|
||||
f"layers {assigned_shard_start}–{assigned_shard_end} "
|
||||
f"(of {assigned_num_layers})",
|
||||
flush=True,
|
||||
)
|
||||
print("Loading real PyTorch model shard...", flush=True)
|
||||
node = TorchNodeServer(
|
||||
host=host,
|
||||
port=port,
|
||||
model_id=assigned_hf_repo,
|
||||
shard_start=assigned_shard_start,
|
||||
shard_end=assigned_shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
)
|
||||
actual_port = node.start()
|
||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||
endpoint = f"http://{public_host}:{actual_port}"
|
||||
auto_reg_payload = {
|
||||
"endpoint": endpoint,
|
||||
"model": assigned_hf_repo.split("/")[-1],
|
||||
"hf_repo": assigned_hf_repo,
|
||||
"num_layers": assigned_num_layers,
|
||||
"shard_start": assigned_shard_start,
|
||||
"shard_end": assigned_shard_end,
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"quantization": quantization,
|
||||
"score": 1.0,
|
||||
"tracker_mode": (assigned_shard_start == 0),
|
||||
}
|
||||
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)
|
||||
except Exception as exc:
|
||||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||||
shard_count = assigned_shard_end - assigned_shard_start + 1
|
||||
print(
|
||||
f"\n{'=' * 32}\n"
|
||||
f"meshnet-node ready (auto-joined)\n"
|
||||
f" Wallet: {address}\n"
|
||||
f" Model ID: {assigned_hf_repo}\n"
|
||||
f" Shard: layers {assigned_shard_start}–{assigned_shard_end} "
|
||||
f"({shard_count} of {assigned_num_layers})\n"
|
||||
f" Quantization: {quantization}\n"
|
||||
f" Endpoint: {endpoint}\n"
|
||||
f" Hardware: {device.upper()}\n"
|
||||
f"{'=' * 32}",
|
||||
flush=True,
|
||||
)
|
||||
return node
|
||||
|
||||
# 3b. Shard assignment from tracker (stub-model / preset-based path)
|
||||
# 3. Shard assignment from tracker
|
||||
print("Querying tracker for shard assignment...", flush=True)
|
||||
assign_qs = urllib.parse.urlencode({
|
||||
"model": model,
|
||||
|
||||
@@ -11,7 +11,6 @@ import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from .model_backend import (
|
||||
InsufficientVRAMError,
|
||||
@@ -214,106 +213,40 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
if body is None:
|
||||
return
|
||||
messages = body.get("messages", [])
|
||||
if not isinstance(messages, list):
|
||||
messages = []
|
||||
stream = bool(body.get("stream", False))
|
||||
model_name = str(body.get("model", ""))
|
||||
max_tokens = int(body.get("max_tokens") or body.get("max_new_tokens") or 256)
|
||||
temperature = float(body.get("temperature") or 1.0)
|
||||
top_p = float(body.get("top_p") or 1.0)
|
||||
|
||||
# Fast path: this node owns the complete model — use HF generate() with KV cache.
|
||||
# Avoids the single-token-per-forward-pass limitation of the distributed path.
|
||||
if server.backend.is_head and server.backend.is_tail:
|
||||
try:
|
||||
if stream:
|
||||
self._stream_openai_response(
|
||||
server.backend.generate_text_streaming(messages, max_tokens, temperature, top_p),
|
||||
model_name,
|
||||
)
|
||||
else:
|
||||
text = server.backend.generate_text(messages, max_tokens, temperature, top_p)
|
||||
self._send_openai_response(text, model_name, False, messages)
|
||||
except Exception as exc:
|
||||
self._send_json(500, {"error": f"generation failed: {exc}"})
|
||||
return
|
||||
|
||||
# Distributed path: autoregressive generation across shards.
|
||||
# We do N single-step forward passes (no cross-node KV cache), which is slow
|
||||
# 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)
|
||||
if not remaining_route:
|
||||
self._send_openai_response(
|
||||
"error: no downstream route — check tracker connectivity",
|
||||
model_name, False, messages,
|
||||
)
|
||||
return
|
||||
|
||||
backend = server.backend
|
||||
# Format with chat template so the model knows it's in assistant mode.
|
||||
model = str(body.get("model", ""))
|
||||
prompt = " ".join(
|
||||
str(m.get("content", ""))
|
||||
for m in messages
|
||||
if isinstance(m, dict) and m.get("role") == "user"
|
||||
)
|
||||
try:
|
||||
if hasattr(backend.tokenizer, "apply_chat_template"):
|
||||
prompt_text: str = backend.tokenizer.apply_chat_template(
|
||||
messages, add_generation_prompt=True, tokenize=False,
|
||||
)
|
||||
else:
|
||||
raise AttributeError("no apply_chat_template")
|
||||
except Exception:
|
||||
prompt_text = " ".join(
|
||||
str(m.get("content", ""))
|
||||
for m in messages
|
||||
if isinstance(m, dict) and m.get("role") == "user"
|
||||
)
|
||||
|
||||
eos_token: str = getattr(backend.tokenizer, "eos_token", "") or ""
|
||||
generated: list[str] = []
|
||||
current_text = prompt_text
|
||||
|
||||
for _ in range(max_tokens):
|
||||
try:
|
||||
payload = backend.encode_prompt(current_text)
|
||||
except Exception as exc:
|
||||
print(f" [node] distributed encode error: {exc}", flush=True)
|
||||
break
|
||||
token_str = self._run_downstream_pipeline(payload, remaining_route)
|
||||
if not token_str:
|
||||
break
|
||||
# Stop on error responses or EOS.
|
||||
if token_str.startswith(("pipeline error", "decode error", "no downstream", "error:")):
|
||||
break
|
||||
if eos_token and token_str == eos_token:
|
||||
break
|
||||
generated.append(token_str)
|
||||
current_text = current_text + token_str
|
||||
|
||||
result_text = "".join(generated)
|
||||
self._send_openai_response(result_text, model_name, stream, messages)
|
||||
payload = server.backend.encode_prompt(prompt)
|
||||
except Exception as exc:
|
||||
self._send_json(500, {"error": f"encode_prompt failed: {exc}"})
|
||||
return
|
||||
remaining_route = self._get_remaining_route(model)
|
||||
result_text = self._run_downstream_pipeline(payload, remaining_route)
|
||||
self._send_openai_response(result_text, model, stream)
|
||||
|
||||
def _get_remaining_route(self, model: str) -> list[str]:
|
||||
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)}"
|
||||
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(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}")]
|
||||
except Exception as exc:
|
||||
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
|
||||
# Skip the first node in the route (self) since we're already the head
|
||||
return list(route[1:])
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
def _run_downstream_pipeline(self, payload: object, route: list[str]) -> str:
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
if not route:
|
||||
# Partial shard at tail: decode the activation from the previous node.
|
||||
# Full single-node (head+tail) is handled before entering this method.
|
||||
# Single-node mode: decode tail locally if we're the tail
|
||||
if server.backend.is_tail:
|
||||
try:
|
||||
tensor = server.backend.torch.frombuffer(
|
||||
@@ -323,7 +256,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
return server.backend.decode_tail(tensor)
|
||||
except Exception as exc:
|
||||
return f"decode error: {exc}"
|
||||
return "no downstream route available for non-tail shard"
|
||||
return ""
|
||||
|
||||
session = str(uuid.uuid4())
|
||||
shape = payload.shape # type: ignore[union-attr]
|
||||
@@ -376,51 +309,10 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
current_pos = resp_headers.get("x-meshnet-position-ids")
|
||||
return ""
|
||||
|
||||
def _stream_openai_response(self, token_iter, model: str) -> None:
|
||||
"""Stream tokens from an iterator as SSE chunks."""
|
||||
chunk_id = "chatcmpl-node"
|
||||
created = int(time.time())
|
||||
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()
|
||||
|
||||
def _emit(data: str) -> None:
|
||||
self.wfile.write(f"data: {data}\n\n".encode())
|
||||
self.wfile.flush()
|
||||
|
||||
_emit(json.dumps({
|
||||
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
|
||||
}))
|
||||
for token_text in token_iter:
|
||||
if not token_text:
|
||||
continue
|
||||
_emit(json.dumps({
|
||||
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": {"content": token_text}, "finish_reason": None}],
|
||||
}))
|
||||
_emit(json.dumps({
|
||||
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
}))
|
||||
self.wfile.write(b"data: [DONE]\n\n")
|
||||
self.wfile.flush()
|
||||
|
||||
def _send_openai_response(
|
||||
self,
|
||||
text: str,
|
||||
model: str,
|
||||
stream: bool,
|
||||
messages: list[dict] | None = None,
|
||||
) -> None:
|
||||
def _send_openai_response(self, text: str, model: str, stream: bool) -> None:
|
||||
chunk_id = "chatcmpl-node"
|
||||
created = int(time.time())
|
||||
if not stream:
|
||||
usage = _usage_for_response(self.server.backend, messages or [], text) # type: ignore[attr-defined]
|
||||
self._send_json(200, {
|
||||
"id": chunk_id,
|
||||
"object": "chat.completion",
|
||||
@@ -431,7 +323,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
"message": {"role": "assistant", "content": text},
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
"usage": usage,
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
||||
})
|
||||
return
|
||||
self.send_response(200)
|
||||
@@ -462,52 +354,6 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
self.wfile.flush()
|
||||
|
||||
|
||||
def _usage_for_response(backend: object, messages: list[dict], completion_text: str) -> dict[str, int]:
|
||||
prompt_tokens = _backend_token_count(
|
||||
backend,
|
||||
"count_prompt_tokens",
|
||||
messages,
|
||||
fallback=_fallback_message_token_count(messages),
|
||||
)
|
||||
completion_tokens = _backend_token_count(
|
||||
backend,
|
||||
"count_text_tokens",
|
||||
completion_text,
|
||||
fallback=_fallback_text_token_count(completion_text),
|
||||
)
|
||||
return {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
}
|
||||
|
||||
|
||||
def _backend_token_count(backend: object, method_name: str, value: object, fallback: int) -> int:
|
||||
method: Any = getattr(backend, method_name, None)
|
||||
if callable(method):
|
||||
try:
|
||||
return max(0, int(method(value)))
|
||||
except Exception:
|
||||
pass
|
||||
return max(0, int(fallback))
|
||||
|
||||
|
||||
def _fallback_message_token_count(messages: list[dict]) -> int:
|
||||
text = " ".join(
|
||||
str(message.get("content", ""))
|
||||
for message in messages
|
||||
if isinstance(message, dict)
|
||||
)
|
||||
return _fallback_text_token_count(text)
|
||||
|
||||
|
||||
def _fallback_text_token_count(text: str) -> int:
|
||||
parts = text.split()
|
||||
if parts:
|
||||
return len(parts)
|
||||
return 1 if text else 0
|
||||
|
||||
|
||||
class TorchNodeServer:
|
||||
"""HTTP server backed by a HuggingFace causal language model shard."""
|
||||
|
||||
|
||||
@@ -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*"]
|
||||
@@ -52,10 +52,11 @@ DEFAULT_BENCHMARK_TOKENS_PER_SEC = 1.0
|
||||
class _NodeEntry:
|
||||
__slots__ = (
|
||||
"node_id", "endpoint", "shard_start", "shard_end",
|
||||
"model", "hf_repo", "num_layers", "shard_checksum", "hardware_profile", "wallet_address",
|
||||
"model", "shard_checksum", "hardware_profile", "wallet_address",
|
||||
"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__(
|
||||
@@ -76,8 +77,9 @@ class _NodeEntry:
|
||||
quantization: str | None = None,
|
||||
managed_assignment: bool = False,
|
||||
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
|
||||
@@ -95,8 +97,9 @@ class _NodeEntry:
|
||||
self.quantization = quantization
|
||||
self.managed_assignment = managed_assignment
|
||||
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()
|
||||
|
||||
@@ -258,12 +261,7 @@ def _purge_expired_nodes_locked(server: "_TrackerHTTPServer") -> list[str]:
|
||||
if (now - entry.last_heartbeat) > server.heartbeat_timeout
|
||||
]
|
||||
for node_id in expired_ids:
|
||||
entry = server.registry.pop(node_id)
|
||||
print(
|
||||
f"[tracker] node expired: {node_id[:8]} {entry.endpoint} "
|
||||
f"(no heartbeat for >{server.heartbeat_timeout:.0f}s)",
|
||||
flush=True,
|
||||
)
|
||||
del server.registry[node_id]
|
||||
if expired_ids:
|
||||
_rebalance_all_locked(server)
|
||||
return expired_ids
|
||||
@@ -435,8 +433,6 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
self._handle_routes(parsed)
|
||||
elif parsed.path == "/v1/nodes/assign":
|
||||
self._handle_assign(parsed)
|
||||
elif parsed.path == "/v1/network/assign":
|
||||
self._handle_network_assign(parsed)
|
||||
elif parsed.path == "/v1/models":
|
||||
self._handle_models()
|
||||
elif parsed.path.startswith("/v1/coverage/"):
|
||||
@@ -618,18 +614,9 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
return
|
||||
|
||||
tracker_mode = bool(body.get("tracker_mode", False))
|
||||
hf_repo = body.get("hf_repo")
|
||||
if hf_repo is not None and not isinstance(hf_repo, str):
|
||||
self._send_json(400, {"error": "hf_repo must be a string"})
|
||||
return
|
||||
num_layers_body = body.get("num_layers")
|
||||
num_layers: int | None = None
|
||||
if num_layers_body is not None:
|
||||
try:
|
||||
num_layers = int(num_layers_body)
|
||||
except (TypeError, ValueError):
|
||||
self._send_json(400, {"error": "num_layers must be an integer"})
|
||||
return
|
||||
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())
|
||||
entry = _NodeEntry(
|
||||
@@ -649,23 +636,12 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
quantization=quantization,
|
||||
managed_assignment=not explicit_shard,
|
||||
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()
|
||||
# Dedup: if this endpoint is already registered, remove the old entry first.
|
||||
stale_ids = [
|
||||
eid for eid, e in server.registry.items()
|
||||
if e.endpoint == entry.endpoint.rstrip("/")
|
||||
]
|
||||
for eid in stale_ids:
|
||||
old = server.registry.pop(eid)
|
||||
print(
|
||||
f"[tracker] node re-registered: replaced {eid[:8]} with {node_id[:8]}"
|
||||
f" {old.endpoint}",
|
||||
flush=True,
|
||||
)
|
||||
server.registry[node_id] = entry
|
||||
if entry.managed_assignment:
|
||||
_rebalance_model_locked(server, model)
|
||||
@@ -673,13 +649,6 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
if assignment_directive is not None:
|
||||
entry.pending_directives.clear()
|
||||
|
||||
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}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
payload = {"node_id": node_id}
|
||||
if assignment_directive is not None:
|
||||
payload["directive"] = assignment_directive
|
||||
@@ -687,6 +656,13 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
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,13 +670,15 @@ 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 directives:
|
||||
self._send_json(200, {"directives": directives})
|
||||
else:
|
||||
@@ -799,111 +777,6 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
**({"hf_repo": preset["hf_repo"]} if "hf_repo" in preset else {}),
|
||||
})
|
||||
|
||||
def _handle_network_assign(self, parsed: urllib.parse.ParseResult):
|
||||
"""Assign a new node to fill the biggest uncovered shard gap across HF-model nodes.
|
||||
|
||||
Query params:
|
||||
vram_mb — integer VRAM in MB (0 = CPU-only node)
|
||||
device — "cuda" | "cpu"
|
||||
hf_repo — optional; if set, restrict search to this repo only
|
||||
|
||||
Returns:
|
||||
{hf_repo, shard_start, shard_end, num_layers, gap_found}
|
||||
gap_found=true means a real uncovered gap was assigned; false means redundancy.
|
||||
"""
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
params = urllib.parse.parse_qs(parsed.query)
|
||||
try:
|
||||
vram_mb = int(params.get("vram_mb", ["0"])[0])
|
||||
except ValueError:
|
||||
vram_mb = 0
|
||||
device = params.get("device", ["cpu"])[0]
|
||||
filter_repo = params.get("hf_repo", [None])[0] # optional repo filter
|
||||
|
||||
with server.lock:
|
||||
self._purge_expired_nodes()
|
||||
all_nodes = list(server.registry.values())
|
||||
|
||||
# Collect only nodes that registered a real HF model (have hf_repo + shard bounds).
|
||||
hf_nodes = [
|
||||
n for n in all_nodes
|
||||
if n.hf_repo
|
||||
and n.shard_start is not None
|
||||
and n.shard_end is not None
|
||||
and n.num_layers is not None
|
||||
and (filter_repo is None or n.hf_repo == filter_repo)
|
||||
]
|
||||
|
||||
if not hf_nodes:
|
||||
msg = (
|
||||
f"no HF-model nodes registered for {filter_repo!r}"
|
||||
if filter_repo
|
||||
else "no HF-model nodes registered; cannot assign shards"
|
||||
)
|
||||
self._send_json(503, {"error": msg})
|
||||
return
|
||||
|
||||
# Group by hf_repo; pick the one with the largest total_layers and biggest gap.
|
||||
from collections import defaultdict
|
||||
repo_groups: dict = defaultdict(list)
|
||||
repo_layers: dict = {}
|
||||
for n in hf_nodes:
|
||||
repo_groups[n.hf_repo].append(n)
|
||||
# Use the largest num_layers seen for this repo.
|
||||
if n.hf_repo not in repo_layers or n.num_layers > repo_layers[n.hf_repo]:
|
||||
repo_layers[n.hf_repo] = n.num_layers
|
||||
|
||||
# Pick the repo where the gap is largest (most work to do).
|
||||
best_repo = None
|
||||
best_gap_size = -1
|
||||
best_gap_start = 0
|
||||
best_num_layers = 0
|
||||
|
||||
for repo, nodes in repo_groups.items():
|
||||
total = repo_layers[repo]
|
||||
covered = sorted(
|
||||
[(n.shard_start, n.shard_end) for n in nodes],
|
||||
key=lambda t: t[0],
|
||||
)
|
||||
# Walk from 0 to find first uncovered layer.
|
||||
gap_start = 0
|
||||
for s, e in covered:
|
||||
if s <= gap_start:
|
||||
gap_start = max(gap_start, e + 1)
|
||||
else:
|
||||
break
|
||||
gap_size = max(0, (total - 1) - gap_start + 1) # layers remaining uncovered
|
||||
if gap_size > best_gap_size:
|
||||
best_gap_size = gap_size
|
||||
best_gap_start = gap_start
|
||||
best_repo = repo
|
||||
best_num_layers = total
|
||||
|
||||
gap_found = best_gap_size > 0
|
||||
if not gap_found:
|
||||
# All shards are covered — still assign to the model with most nodes for redundancy.
|
||||
best_repo = max(repo_groups, key=lambda r: len(repo_groups[r]))
|
||||
best_gap_start = 0
|
||||
best_num_layers = repo_layers[best_repo]
|
||||
|
||||
# Capacity: CPU nodes get at most half the layers; CUDA nodes based on VRAM.
|
||||
total_l = best_num_layers
|
||||
if device == "cuda" and vram_mb >= 8192:
|
||||
max_layers = total_l
|
||||
else:
|
||||
max_layers = max(1, total_l // 2)
|
||||
|
||||
shard_start = best_gap_start
|
||||
shard_end = min(total_l - 1, shard_start + max_layers - 1)
|
||||
|
||||
self._send_json(200, {
|
||||
"hf_repo": best_repo,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"num_layers": total_l,
|
||||
"gap_found": gap_found,
|
||||
})
|
||||
|
||||
def _handle_route(self, parsed: urllib.parse.ParseResult):
|
||||
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
|
||||
params = urllib.parse.parse_qs(parsed.query)
|
||||
@@ -914,28 +787,15 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
|
||||
|
||||
model = model_list[0]
|
||||
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()
|
||||
if preset is not None:
|
||||
# Preset-based routing (stub-model system).
|
||||
alive = [node for node in server.registry.values() if node.model == model]
|
||||
required_start, required_end = _preset_layer_bounds(preset)
|
||||
else:
|
||||
# HF model routing: match by hf_repo (full) or model short name.
|
||||
alive = [
|
||||
node for node in server.registry.values()
|
||||
if (node.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 alive:
|
||||
self._send_json(404, {"error": f"no nodes registered for model {model!r}"})
|
||||
return
|
||||
required_start = 0
|
||||
required_end = max(n.num_layers for n in alive) - 1 # type: ignore[type-var]
|
||||
|
||||
alive = [node for node in server.registry.values() if node.model == model]
|
||||
if server.contracts is not None:
|
||||
alive = [
|
||||
node for node in alive
|
||||
@@ -1037,7 +897,7 @@ class TrackerServer:
|
||||
self,
|
||||
host: str = "127.0.0.1",
|
||||
port: int = 0,
|
||||
heartbeat_timeout: float = 90.0,
|
||||
heartbeat_timeout: float = 30.0,
|
||||
rebalance_interval: float = 30.0,
|
||||
model_presets: dict | None = None,
|
||||
contracts: Any | None = None,
|
||||
|
||||
@@ -152,7 +152,7 @@ def _story_meta(
|
||||
parts.append(f"worktree: {wt}")
|
||||
|
||||
# summary -------------------------------------------------------------
|
||||
notes = (story.get("completionNotes") or "").strip()
|
||||
notes = story.get("completionNotes", "").strip()
|
||||
if not notes and wt:
|
||||
notes = _story_last_commit(sid)
|
||||
if notes:
|
||||
|
||||
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
|
||||
@@ -347,43 +347,6 @@ def test_tracker_assign_lists_peers_for_same_model_shard():
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_real_model_startup_summary_shows_total_layers(tmp_path, monkeypatch, capsys):
|
||||
"""Real-model startup summary prints the shard range plus total model layers."""
|
||||
import meshnet_node.startup as startup_mod
|
||||
|
||||
class FakeBackend:
|
||||
total_layers = 24
|
||||
|
||||
class FakeTorchNodeServer:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
self.backend = FakeBackend()
|
||||
self.port = None
|
||||
|
||||
def start(self):
|
||||
self.port = 8001
|
||||
return self.port
|
||||
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"detect_hardware",
|
||||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
|
||||
)
|
||||
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
|
||||
|
||||
node = run_startup(
|
||||
tracker_url="http://127.0.0.1:8080",
|
||||
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||||
shard_start=0,
|
||||
shard_end=23,
|
||||
wallet_path=tmp_path / "wallet.json",
|
||||
)
|
||||
|
||||
assert node.backend.total_layers == 24
|
||||
output = capsys.readouterr().out
|
||||
assert "Shard: layers 0–23; 24 of 24" in output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Full startup integration test
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -507,163 +470,6 @@ def test_second_node_downloads_same_shard_from_peer_without_huggingface(
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_network_assign_gap_found_field():
|
||||
"""network/assign sets gap_found=True when a real gap exists, False when fully covered."""
|
||||
import json as _json
|
||||
import urllib.request as _ur
|
||||
|
||||
tracker = TrackerServer()
|
||||
port = tracker.start()
|
||||
try:
|
||||
# Register a node covering only layers 0-11 of a 24-layer model.
|
||||
data = _json.dumps({
|
||||
"endpoint": "http://127.0.0.1:9200",
|
||||
"model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"num_layers": 24,
|
||||
"shard_start": 0,
|
||||
"shard_end": 11,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0,
|
||||
}).encode()
|
||||
req = _ur.Request(
|
||||
f"http://127.0.0.1:{port}/v1/nodes/register",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with _ur.urlopen(req) as r:
|
||||
r.read()
|
||||
|
||||
# A new node should be told there is a gap (layers 12-23).
|
||||
resp = _get_json(
|
||||
f"http://127.0.0.1:{port}/v1/network/assign?device=cpu&vram_mb=0"
|
||||
"&hf_repo=Qwen/Qwen2.5-0.5B-Instruct"
|
||||
)
|
||||
assert resp["gap_found"] is True
|
||||
assert resp["shard_start"] == 12, f"expected gap at 12, got {resp['shard_start']}"
|
||||
assert resp["shard_end"] == 23
|
||||
|
||||
# Register the second node covering the gap.
|
||||
data2 = _json.dumps({
|
||||
"endpoint": "http://127.0.0.1:9201",
|
||||
"model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"num_layers": 24,
|
||||
"shard_start": 12,
|
||||
"shard_end": 23,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0,
|
||||
}).encode()
|
||||
req2 = _ur.Request(
|
||||
f"http://127.0.0.1:{port}/v1/nodes/register",
|
||||
data=data2,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with _ur.urlopen(req2) as r:
|
||||
r.read()
|
||||
|
||||
# Now fully covered — gap_found should be False.
|
||||
resp2 = _get_json(
|
||||
f"http://127.0.0.1:{port}/v1/network/assign?device=cpu&vram_mb=0"
|
||||
"&hf_repo=Qwen/Qwen2.5-0.5B-Instruct"
|
||||
)
|
||||
assert resp2["gap_found"] is False
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_route_finds_hf_model_across_two_nodes():
|
||||
"""Tracker /v1/route returns ordered route for HF model even without a preset."""
|
||||
import json as _json
|
||||
import urllib.request as _ur
|
||||
|
||||
tracker = TrackerServer()
|
||||
port = tracker.start()
|
||||
try:
|
||||
def register(endpoint, shard_start, shard_end):
|
||||
data = _json.dumps({
|
||||
"endpoint": endpoint,
|
||||
"model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"num_layers": 24,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0,
|
||||
}).encode()
|
||||
req = _ur.Request(
|
||||
f"http://127.0.0.1:{port}/v1/nodes/register",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with _ur.urlopen(req) as r:
|
||||
r.read()
|
||||
|
||||
register("http://127.0.0.1:9300", 0, 11)
|
||||
register("http://127.0.0.1:9301", 12, 23)
|
||||
|
||||
# Route by hf_repo (full identifier).
|
||||
resp = _get_json(
|
||||
f"http://127.0.0.1:{port}/v1/route?model=Qwen/Qwen2.5-0.5B-Instruct"
|
||||
)
|
||||
assert resp["route"] == ["http://127.0.0.1:9300", "http://127.0.0.1:9301"]
|
||||
|
||||
# Route also works by short model name.
|
||||
resp2 = _get_json(
|
||||
f"http://127.0.0.1:{port}/v1/route?model=Qwen2.5-0.5B-Instruct"
|
||||
)
|
||||
assert resp2["route"] == ["http://127.0.0.1:9300", "http://127.0.0.1:9301"]
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_register_deduplicates_same_endpoint():
|
||||
"""Re-registering the same endpoint replaces the old entry, not duplicates it."""
|
||||
import json as _json
|
||||
import urllib.request as _ur
|
||||
|
||||
tracker = TrackerServer()
|
||||
port = tracker.start()
|
||||
try:
|
||||
def register(shard_start, shard_end):
|
||||
data = _json.dumps({
|
||||
"endpoint": "http://127.0.0.1:9400",
|
||||
"model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"num_layers": 24,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0,
|
||||
}).encode()
|
||||
req = _ur.Request(
|
||||
f"http://127.0.0.1:{port}/v1/nodes/register",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with _ur.urlopen(req) as r:
|
||||
return _json.loads(r.read())
|
||||
|
||||
register(0, 23) # initial full-model registration
|
||||
register(12, 23) # re-register with corrected shard range
|
||||
|
||||
# After re-register, tracker should see only one node at 12-23 for this endpoint.
|
||||
# If both were still registered, the gap scan would find no gap (0-23 still covers).
|
||||
# With dedup, the old 0-23 is gone and a real gap 0-11 exists.
|
||||
assign_resp = _get_json(
|
||||
f"http://127.0.0.1:{port}/v1/network/assign?device=cpu&vram_mb=0"
|
||||
"&hf_repo=Qwen/Qwen2.5-0.5B-Instruct"
|
||||
)
|
||||
assert assign_resp["gap_found"] is True
|
||||
assert assign_resp["shard_start"] == 0, "old 0-23 entry should have been replaced"
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_startup_cpu_fallback(tmp_path, monkeypatch):
|
||||
"""Node starts with CPU warning when no GPU is detected."""
|
||||
import meshnet_node.startup as startup_mod
|
||||
|
||||
@@ -12,9 +12,6 @@ import pytest
|
||||
from meshnet_node.model_backend import (
|
||||
InsufficientVRAMError,
|
||||
TensorPayload,
|
||||
_call_layer,
|
||||
_decoder_attention_mask,
|
||||
_int_tensor_header,
|
||||
build_quantization_config,
|
||||
validate_quantization,
|
||||
)
|
||||
@@ -55,32 +52,6 @@ class _FakeTailBackend(_FakeBackend):
|
||||
return " Paris"
|
||||
|
||||
|
||||
class _FakeFullBackend(_FakeBackend):
|
||||
is_head = True
|
||||
is_tail = True
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 16,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
) -> str:
|
||||
assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
|
||||
assert max_new_tokens == 7
|
||||
assert temperature == 1.0
|
||||
assert top_p == 1.0
|
||||
return "56"
|
||||
|
||||
def count_prompt_tokens(self, messages: list[dict]) -> int:
|
||||
assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
|
||||
return 8
|
||||
|
||||
def count_text_tokens(self, text: str) -> int:
|
||||
assert text == "56"
|
||||
return 1
|
||||
|
||||
|
||||
def test_quantization_flag_validation():
|
||||
assert validate_quantization("bfloat16") == "bfloat16"
|
||||
assert validate_quantization("int8") == "int8"
|
||||
@@ -174,65 +145,6 @@ def test_tail_forward_returns_text_completion_from_binary_activations():
|
||||
node.stop()
|
||||
|
||||
|
||||
def test_full_model_chat_completion_uses_generation_not_single_token_decode():
|
||||
node = TorchNodeServer(backend=_FakeFullBackend())
|
||||
port = node.start()
|
||||
try:
|
||||
payload = json.dumps({
|
||||
"model": "fake-model",
|
||||
"messages": [{"role": "user", "content": "What is 7 times 8?"}],
|
||||
"max_tokens": 7,
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{port}/v1/chat/completions",
|
||||
data=payload,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
body = json.loads(resp.read())
|
||||
|
||||
assert body["choices"][0]["message"]["content"] == "56"
|
||||
assert body["usage"] == {"prompt_tokens": 8, "completion_tokens": 1, "total_tokens": 9}
|
||||
finally:
|
||||
node.stop()
|
||||
|
||||
|
||||
def test_int_tensor_header_serializes_torch_tensors():
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
header = _int_tensor_header(torch.tensor([[1, 2, 3]], dtype=torch.long))
|
||||
|
||||
assert header.startswith("1,3:")
|
||||
|
||||
|
||||
def test_decoder_attention_mask_is_causal_float_mask():
|
||||
torch = pytest.importorskip("torch")
|
||||
|
||||
hidden_states = torch.zeros((1, 3, 8), dtype=torch.bfloat16)
|
||||
mask = _decoder_attention_mask(torch.ones((1, 3), dtype=torch.long), hidden_states, torch)
|
||||
|
||||
assert mask.shape == (1, 1, 3, 3)
|
||||
assert mask.dtype == torch.bfloat16
|
||||
assert mask[0, 0, 0, 1] < 0
|
||||
assert mask[0, 0, 2, 0] == 0
|
||||
|
||||
|
||||
def test_call_layer_passes_rotary_position_embeddings():
|
||||
class NeedsPositionEmbeddings:
|
||||
def __call__(self, hidden_states, **kwargs):
|
||||
assert kwargs["position_embeddings"] == "rotary"
|
||||
return hidden_states
|
||||
|
||||
assert _call_layer(
|
||||
NeedsPositionEmbeddings(),
|
||||
"hidden",
|
||||
attention_mask=None,
|
||||
position_ids="positions",
|
||||
position_embeddings="rotary",
|
||||
) == "hidden"
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
def test_two_node_gpt2_completion_is_deterministic():
|
||||
if os.environ.get("CI"):
|
||||
|
||||
Reference in New Issue
Block a user