Compare commits
48 Commits
worktree-f
...
61074a8fe8
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
61074a8fe8 | ||
|
|
d0307fcc84 | ||
|
|
f1e4870124 | ||
|
|
7866723c82 | ||
|
|
4f00a37d72 | ||
|
|
ab6558d861 | ||
|
|
6c46f96aaf | ||
|
|
d8a723a4c7 | ||
|
|
742d5ff0bf | ||
|
|
8157151102 | ||
|
|
4f79b2d177 | ||
|
|
5fe471d8ca | ||
|
|
6e4f755e71 | ||
|
|
1da088926a | ||
|
|
c587e02c2c | ||
|
|
df473ef278 | ||
|
|
9ca198ee1e | ||
|
|
8ce5a74d5e | ||
|
|
27818df654 | ||
|
|
d9110b623b | ||
|
|
fff11c4d8b | ||
|
|
a4373a6202 | ||
|
|
ad3368f7ea | ||
|
|
34fb1ec65d | ||
|
|
96af82892f | ||
|
|
e1ba120912 | ||
|
|
c691e8d5d1 | ||
|
|
dade97ee67 | ||
|
|
2e1e0ae172 | ||
|
|
4a803377dc | ||
|
|
ae5ff6a805 | ||
|
|
b95e25a5c3 | ||
|
|
748d535c46 | ||
|
|
bbe57d5f07 | ||
|
|
0be35d257b | ||
|
|
be37048145 | ||
|
|
d701ae9ba2 | ||
|
|
1e6016e76f | ||
|
|
60ccf47cb4 | ||
|
|
c75e9708ae | ||
|
|
3286e42783 | ||
|
|
97eefd3d5e | ||
|
|
a7cc377d13 | ||
|
|
1bdfce657d | ||
|
|
607d49f5b0 | ||
|
|
6b9caecd90 | ||
|
|
4e292eaaae | ||
|
|
ded8c06e77 |
1
.claude/worktrees/feat+us-016
Submodule
1
.claude/worktrees/feat+us-016
Submodule
Submodule .claude/worktrees/feat+us-016 added at 080d49b2c2
@@ -0,0 +1,206 @@
|
||||
# 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
|
||||
@@ -0,0 +1,205 @@
|
||||
# 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
|
||||
@@ -0,0 +1,157 @@
|
||||
# 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
|
||||
@@ -0,0 +1,68 @@
|
||||
# US-019 — Binary data plane and optional peer weight transfer
|
||||
|
||||
Status: needs-triage
|
||||
Priority: Low
|
||||
Stage: Design parking lot
|
||||
|
||||
## Context
|
||||
|
||||
The current project focus is inference democratization: let small GPU owners contribute useful compute and let users run inference on models larger than one host can serve alone. Weight distribution is useful, but it is secondary to low-latency distributed inference.
|
||||
|
||||
Recent findings:
|
||||
|
||||
- HuggingFace already handles initial model origin distribution well enough for the first working version.
|
||||
- The inference-critical path is activation transfer between shard nodes, not torrenting model files.
|
||||
- Torrent/content-addressed transfer is a good future fit for model weights, shard cache replication, fine-tuned models, and offline/local swarm behavior.
|
||||
- Torrenting is not a good fit for activation traffic because activations are latency-sensitive, ordered, session-specific binary streams.
|
||||
|
||||
## Design note
|
||||
|
||||
Keep the tracker as the control plane:
|
||||
|
||||
- node registration
|
||||
- heartbeats
|
||||
- route selection
|
||||
- model/shard manifests
|
||||
- peer/checksum metadata
|
||||
|
||||
Keep binary payloads on the data plane:
|
||||
|
||||
- direct node-to-node activation transfer where reachable
|
||||
- relay/QUIC/WSS fallback where direct transport is unavailable
|
||||
- future peer weight transfer as content-addressed blobs or pieces
|
||||
|
||||
## Potential future direction
|
||||
|
||||
For inference traffic:
|
||||
|
||||
- Prefer direct binary transport with backpressure.
|
||||
- Use raw binary activation frames rather than JSON/base64.
|
||||
- Preserve tensor metadata out-of-band: shape, dtype, session, chunk index, encoding.
|
||||
- Consider QUIC or a mature NAT-friendly transport for direct-when-possible, relay-when-needed behavior.
|
||||
|
||||
For model weights:
|
||||
|
||||
- Keep HuggingFace as default origin and fallback.
|
||||
- Add peer cache transfer only as an optimization.
|
||||
- Consider a real content-addressed/torrent-like library or sidecar for weight blobs.
|
||||
- Store manifests and checksums in tracker state so peers can verify exact shard contents.
|
||||
- Prefer piece/chunk transfer with resume and hash verification over one giant tarball.
|
||||
|
||||
## Not in scope now
|
||||
|
||||
- Replacing HuggingFace as the primary model origin.
|
||||
- Building a full BitTorrent/IPFS/libp2p subsystem.
|
||||
- Routing activation traffic through a torrent protocol.
|
||||
- Making peer weight transfer mandatory for node startup.
|
||||
|
||||
## Acceptance criteria for a future implementation issue
|
||||
|
||||
- [ ] Activation transfer remains binary end-to-end and avoids JSON/base64 payloads.
|
||||
- [ ] Tracker does not proxy large binary payloads except as an explicit fallback path.
|
||||
- [ ] Weight transfer, if added, is optional and falls back to HuggingFace.
|
||||
- [ ] Weight pieces are content-addressed and checksum-verified.
|
||||
- [ ] The design preserves low-latency inference as the primary objective.
|
||||
|
||||
## Comments
|
||||
|
||||
Created as a low-priority design parking-lot item after discussing inference democratization versus weight distribution. Do not pick up for implementation until the core public tracker, relay, and binary activation path are stable.
|
||||
@@ -48,8 +48,8 @@
|
||||
"dependsOn": [
|
||||
"US-001"
|
||||
],
|
||||
"completionNotes": "Completed by Ralph iteration 126384e5; verified by pytest two-node pipeline.",
|
||||
"status": "to-revise",
|
||||
"completionNotes": "Tests pass; activation tensor flow via stub nodes verified. New HF-model path tested in test_node_startup.py.",
|
||||
"status": "done",
|
||||
"status_reason": "Base64 JSON wire format established here was replaced by binary HTTP protocol in US-011. tests/test_two_node_pipeline.py needs verification it exercises the new binary format end-to-end."
|
||||
},
|
||||
{
|
||||
@@ -128,8 +128,8 @@
|
||||
"dependsOn": [
|
||||
"US-003"
|
||||
],
|
||||
"completionNotes": "Completed by agent",
|
||||
"status": "to-revise",
|
||||
"completionNotes": "All 6 gateway tests pass including OpenAI SDK, LangChain, streaming, and 503 for unknown model.",
|
||||
"status": "done",
|
||||
"status_reason": "Gateway is currently the pipeline orchestrator. US-014 moves orchestration to tracker-nodes; gateway becomes a thin load-balancer proxy. Implementation will be superseded \u2014 defer rework until US-014 lands."
|
||||
},
|
||||
{
|
||||
@@ -382,10 +382,109 @@
|
||||
"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": "done",
|
||||
"notes": "Source issue: .scratch/distributed-inference-network/issues/18-two-machine-lan-test.md",
|
||||
"dependsOn": [
|
||||
"US-016",
|
||||
"US-017"
|
||||
],
|
||||
"completionNotes": "docs/INSTALL_WINDOWS.md: WSL2+CUDA+meshnet-node install guide. docs/TWO_MACHINE_TEST.md: two-machine LAN test procedure with known issues. scripts/test_lan_inference.py: stdlib-only test script, 3 requests, exit 0 on success, auto-discovers gateway from tracker."
|
||||
},
|
||||
{
|
||||
"id": "US-019",
|
||||
"title": "19 \u2014 Distributed tracker consensus (Raft assignments + CRDT heartbeats)",
|
||||
"description": "Replace the single-point-of-failure tracker with a fault-tolerant cluster. Tracker nodes elect a leader via Raft and commit shard assignments as log entries \u2014 all tracker nodes agree on who owns what. Node liveness (heartbeats) uses CRDT gossip (eventual consistency, high frequency OK). A node registers with any tracker node; the write is forwarded to the leader and replicated to followers. A 3-node tracker cluster survives one tracker failure without losing assignment state. The relay/gossip layer already built in US-017 handles peer heartbeats; this story wires Raft on top for authoritative assignments.",
|
||||
"acceptanceCriteria": [
|
||||
"3 tracker nodes can be started and form a Raft cluster (leader election, log replication)",
|
||||
"A node registers with any follower \u2014 the registration is forwarded to the leader and replicated",
|
||||
"Killing the leader causes a new election within 5 seconds; registrations continue working",
|
||||
"Shard assignments returned by any tracker node are identical (strong consistency)",
|
||||
"Node heartbeats use CRDT gossip (not Raft) \u2014 high-frequency, eventual consistency",
|
||||
"meshnet-tracker CLI gains --cluster-peers flag to specify peer tracker URLs",
|
||||
"Integration test: 3 tracker nodes, kill leader mid-test, verify assignment still works",
|
||||
"QUICKSTART.md updated with multi-tracker setup section"
|
||||
],
|
||||
"priority": 19,
|
||||
"status": "done",
|
||||
"notes": "Architecture decision: Raft for assignments (strong consistency) + CRDT gossip for liveness (eventual consistency). User approved 2026-06-29.",
|
||||
"dependsOn": [
|
||||
"US-017"
|
||||
],
|
||||
"completionNotes": "raft.py: minimal Raft consensus (leader election, log replication, AppendEntries, RequestVote). gossip.py: LWW CRDT gossip for node heartbeats. TrackerServer gains cluster_peers + cluster_self_url params. --cluster-peers and --self-url CLI flags added. 6 integration tests: leader election <1s, follower registration propagation, leader kill + re-election <5s, gossip table."
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"updatedAt": "2026-06-29T13:30:00.000Z",
|
||||
"updatedAt": "2026-06-29T15:35:00.000Z",
|
||||
"statusVocabulary": {
|
||||
"open": "Not started",
|
||||
"in-design": "Decisions pending before implementation can begin",
|
||||
|
||||
580
QUICKSTART.md
Normal file
580
QUICKSTART.md
Normal file
@@ -0,0 +1,580 @@
|
||||
# Quickstart — Running a node and testing inference
|
||||
|
||||
This guide gets you from zero to a live inference request in three terminals.
|
||||
Tested on: AMD Ryzen AI Max (Strix Halo APU), 124 GB RAM, Linux, CPU inference.
|
||||
|
||||
---
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
# Clone and enter repo
|
||||
cd /run/media/popov/d/DEV/repos/d-popov.com/AI
|
||||
|
||||
# Create the virtualenv if it does not exist yet
|
||||
python3 -m venv .venv
|
||||
|
||||
# Keep packaging tools current enough for editable installs
|
||||
.venv/bin/python -m pip install --upgrade pip setuptools wheel
|
||||
|
||||
# Install Python packages (editable — picks up code changes immediately)
|
||||
.venv/bin/pip install -e packages/tracker -e packages/node -e packages/p2p -e packages/gateway -e packages/relay
|
||||
|
||||
# CPU-only PyTorch (skip if you have CUDA/ROCm already)
|
||||
.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu
|
||||
|
||||
# HuggingFace model libraries
|
||||
.venv/bin/pip install transformers 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`
|
||||
|
||||
### Windows / WSL2
|
||||
|
||||
Run the Linux commands from WSL, not Git Bash. From the repo opened in Git Bash:
|
||||
|
||||
```bash
|
||||
wsl
|
||||
cd /mnt/d/DEV/workspace/REPOS/git.d-popov.com/neuron-tai
|
||||
python3 -m venv .venv
|
||||
.venv/bin/python -m pip install --upgrade pip setuptools wheel
|
||||
.venv/bin/pip install -e packages/tracker -e packages/node -e packages/p2p -e packages/gateway -e packages/relay
|
||||
.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu
|
||||
.venv/bin/pip install transformers accelerate
|
||||
.venv/bin/meshnet-node --help
|
||||
```
|
||||
|
||||
If `.venv/bin/meshnet-node` is missing, the editable install step did not finish
|
||||
successfully. Re-run the `.venv/bin/pip install -e ...` command above inside WSL.
|
||||
|
||||
WSL2 sits behind Windows NAT and is **not directly reachable** from other LAN machines.
|
||||
Direct cross-host hops time out. The relay path (see below) solves this: the WSL2 node
|
||||
opens an outbound WebSocket to the relay server and all traffic flows through that tunnel.
|
||||
No firewall rules, no `--advertise-host` needed — just point at the public tracker URL.
|
||||
|
||||
### Native Windows PowerShell node (not WSL)
|
||||
|
||||
Use this when the tracker is on another machine and you want Windows to host a
|
||||
reachable node on the LAN.
|
||||
|
||||
1. Install prerequisites on Windows:
|
||||
- Python 3.11 or 3.12 from <https://www.python.org/downloads/windows/>
|
||||
- Git for Windows from <https://git-scm.com/download/win>
|
||||
|
||||
2. Open **PowerShell** in the cloned repo and install the node packages:
|
||||
|
||||
```powershell
|
||||
# Example repo path; adjust to wherever you cloned it
|
||||
cd D:\DEV\workspace\REPOS\git.d-popov.com\neuron-tai
|
||||
|
||||
python -m venv .venv
|
||||
.\.venv\Scripts\python.exe -m pip install --upgrade pip setuptools wheel
|
||||
.\.venv\Scripts\pip.exe install -e packages\tracker -e packages\node -e packages\p2p -e packages\gateway -e packages\relay
|
||||
|
||||
# CPU-only PyTorch. For NVIDIA CUDA, use `pip install torch` instead.
|
||||
.\.venv\Scripts\pip.exe install torch --index-url https://download.pytorch.org/whl/cpu
|
||||
.\.venv\Scripts\pip.exe install transformers accelerate
|
||||
|
||||
.\.venv\Scripts\meshnet-node.exe --help
|
||||
```
|
||||
|
||||
For `start`-specific flags, run:
|
||||
|
||||
```powershell
|
||||
.\.venv\Scripts\meshnet-node.exe start --help
|
||||
```
|
||||
|
||||
3. Find the Windows LAN IP address:
|
||||
|
||||
```powershell
|
||||
ipconfig
|
||||
```
|
||||
|
||||
Use the IPv4 address on the active Ethernet/Wi-Fi adapter, for example
|
||||
`192.168.0.42`. Avoid WSL/Docker/Hyper-V adapter addresses like `172.16.x.x`,
|
||||
`172.17.x.x`, or other virtual adapter IPs.
|
||||
|
||||
4. Allow inbound traffic for the node port in Windows Firewall. Run PowerShell as
|
||||
Administrator once:
|
||||
|
||||
```powershell
|
||||
New-NetFirewallRule `
|
||||
-DisplayName "Meshnet node 8005" `
|
||||
-Direction Inbound `
|
||||
-Action Allow `
|
||||
-Protocol TCP `
|
||||
-LocalPort 8005
|
||||
```
|
||||
|
||||
5. Start the Windows node from normal PowerShell. Replace the tracker and
|
||||
advertised host values with your actual LAN addresses:
|
||||
|
||||
```powershell
|
||||
$env:HF_HOME = "D:\DEV\models"
|
||||
|
||||
.\.venv\Scripts\meshnet-node.exe start `
|
||||
--tracker http://192.168.0.179:8081 `
|
||||
--model Qwen/Qwen2.5-0.5B-Instruct `
|
||||
--shard-start 12 --shard-end 23 `
|
||||
--quantization bfloat16 `
|
||||
--host 0.0.0.0 `
|
||||
--advertise-host 192.168.0.42 `
|
||||
--port 8005
|
||||
```
|
||||
|
||||
One-line variants (direct LAN — node must be reachable by IP from other machines):
|
||||
|
||||
```powershell
|
||||
.\.venv\Scripts\meshnet-node.exe start --tracker http://192.168.0.179:8081 --model Qwen/Qwen2.5-0.5B-Instruct --advertise-host 192.168.0.20
|
||||
```
|
||||
|
||||
Via public hostname with relay (works from behind NAT, WSL2, 5G — no `--advertise-host` needed):
|
||||
|
||||
```powershell
|
||||
.\.venv\Scripts\meshnet-node.exe start --tracker https://ai.neuron.d-popov.com --model Qwen/Qwen2.5-0.5B-Instruct
|
||||
```
|
||||
|
||||
WSL (same relay path — no `--advertise-host`):
|
||||
|
||||
```bash
|
||||
.venv/bin/meshnet-node start --tracker https://ai.neuron.d-popov.com --model Qwen/Qwen2.5-0.5B-Instruct
|
||||
.venv/bin/meshnet-node start --tracker http://192.168.0.179:8081 --model Qwen/Qwen2.5-0.5B-Instruct
|
||||
```
|
||||
|
||||
`--host 0.0.0.0` binds the node to all Windows interfaces. `--advertise-host`
|
||||
is what the tracker gives to other nodes for direct connections; omit it when using
|
||||
the relay path since all traffic flows through the relay tunnel instead.
|
||||
|
||||
If you want verbose per-hop pipeline logs while debugging a split model, add
|
||||
`--debug`. Leave it off for normal runs; otherwise every generated token logs
|
||||
lines like:
|
||||
|
||||
```text
|
||||
[node] pipeline hop 0: http://127.0.0.1:8005 start_layer=22
|
||||
[node] pipeline hop 0 returned text=' token'
|
||||
[node] pipeline hop 1: wss://ai.neuron.d-popov.com/rpc/abc123 relay start_layer=12
|
||||
```
|
||||
|
||||
6. From the tracker machine, verify Windows is reachable:
|
||||
|
||||
```bash
|
||||
curl http://192.168.0.42:8005/v1/health
|
||||
```
|
||||
|
||||
If that endpoint returns 404 or 501, that is okay: it still proves the TCP
|
||||
connection reached the node process. If it times out or connection-refuses, check
|
||||
the Windows Firewall rule, `--host 0.0.0.0`, the selected LAN IP, and that the
|
||||
node is still running.
|
||||
|
||||
---
|
||||
|
||||
## Public tracker + relay (internet / NAT nodes)
|
||||
|
||||
This setup lets nodes connect from anywhere — behind home NAT, 5G, WSL2, or
|
||||
on a different continent — without opening firewall ports.
|
||||
|
||||
### Architecture
|
||||
|
||||
```
|
||||
Client → HTTPS → ai.neuron.d-popov.com (nginx)
|
||||
├─ /v1/* → meshnet-tracker :8081
|
||||
├─ /ws → meshnet-relay :8765 (node persistent outbound WS)
|
||||
└─ /rpc/* → meshnet-relay :8765 (caller opens WS per hop)
|
||||
```
|
||||
|
||||
### Nginx Proxy Manager (Docker)
|
||||
|
||||
Use **one** proxy host for the domain. Do not create a second host for the same
|
||||
domain to reach another port — path routing is done with **Custom locations** on
|
||||
that same host.
|
||||
|
||||
**1. Details tab** (default `/` → tracker)
|
||||
|
||||
| Field | Value |
|
||||
|-------|--------|
|
||||
| Domain Names | `ai.neuron.d-popov.com` |
|
||||
| Scheme | `http` |
|
||||
| Forward Hostname / IP | LAN IP of the tracker machine (e.g. `192.168.0.179`) |
|
||||
| Forward Port | `8081` |
|
||||
| Websockets Support | ON |
|
||||
|
||||
This serves `/v1/network/map`, `/v1/chat/completions`, and the rest of the tracker API.
|
||||
|
||||
**2. Custom locations tab** (sub-paths → relay)
|
||||
|
||||
The Custom locations form has **no separate Websockets toggle** — only location,
|
||||
scheme, forward host, optional sub-folder path, and port. Add **two** locations
|
||||
(both pointing at the relay process on port `8765`). Leave **“Add a path for
|
||||
sub-folder forwarding”** empty so the full URI reaches the relay
|
||||
(`/ws`, `/rpc/<peer_id>`).
|
||||
|
||||
Location A — persistent node connections:
|
||||
|
||||
| Field | Value |
|
||||
|-------|--------|
|
||||
| Define location | `/ws` |
|
||||
| Scheme | `http` |
|
||||
| Forward Hostname / IP | `192.168.0.179` |
|
||||
| Forward Port | `8765` |
|
||||
| Sub-folder path | *(leave empty)* |
|
||||
|
||||
Location B — per-hop RPC:
|
||||
|
||||
| Field | Value |
|
||||
|-------|--------|
|
||||
| Define location | `/rpc` |
|
||||
| Scheme | `http` |
|
||||
| Forward Hostname / IP | `192.168.0.179` |
|
||||
| Forward Port | `8765` |
|
||||
| Sub-folder path | *(leave empty)* |
|
||||
|
||||
Nginx matches the longer prefixes first: `/ws` and `/rpc/…` go to relay; everything
|
||||
else stays on `8081`.
|
||||
|
||||
**3. SSL tab**
|
||||
|
||||
Use your existing Let’s Encrypt certificate (unchanged).
|
||||
|
||||
**4. Advanced tab** (only if WebSocket upgrade fails on `/ws` or `/rpc`)
|
||||
|
||||
Custom locations do not expose a Websockets checkbox. If nodes show
|
||||
`Relay configured but not connected yet` while `/v1/network/map` works, add this
|
||||
snippet on the **proxy host** Advanced tab:
|
||||
|
||||
```nginx
|
||||
proxy_http_version 1.1;
|
||||
proxy_set_header Upgrade $http_upgrade;
|
||||
proxy_set_header Connection $http_connection;
|
||||
proxy_read_timeout 3600s;
|
||||
proxy_send_timeout 3600s;
|
||||
```
|
||||
|
||||
**5. Verify routing**
|
||||
|
||||
```bash
|
||||
# Tracker (8081 via default location)
|
||||
curl -s https://ai.neuron.d-popov.com/v1/network/map | python3 -m json.tool
|
||||
|
||||
# Relay paths should not 502/404 from the tracker — check response headers/status
|
||||
curl -sI https://ai.neuron.d-popov.com/ws
|
||||
curl -sI https://ai.neuron.d-popov.com/rpc/test-peer
|
||||
```
|
||||
|
||||
After NPM is correct, start relay and tracker on the LAN machine:
|
||||
|
||||
```bash
|
||||
# Terminal 1 — relay
|
||||
.venv/bin/meshnet-relay --host 0.0.0.0 --port 8765
|
||||
|
||||
# Terminal 2 — tracker (advertises relay to nodes)
|
||||
.venv/bin/meshnet-tracker start \
|
||||
--host 0.0.0.0 \
|
||||
--port 8081 \
|
||||
--relay-url wss://ai.neuron.d-popov.com/ws
|
||||
```
|
||||
|
||||
Nodes using `https://ai.neuron.d-popov.com` should then log:
|
||||
|
||||
```text
|
||||
Relay advertised by tracker — using outbound tunnel wss://ai.neuron.d-popov.com/ws
|
||||
Relay connected — wss://ai.neuron.d-popov.com/rpc/<peer_id>
|
||||
```
|
||||
|
||||
The `--relay-url` flag embeds the relay address in `/v1/network/map`. Every node
|
||||
queries that endpoint on startup and auto-connects if a relay URL is present.
|
||||
|
||||
### Start a node (any machine, any network)
|
||||
|
||||
No `--advertise-host`, firewall rule, port forwarding, relay URL, or peer URL is
|
||||
needed on the node. The public tracker is the only bootstrap URL the user types.
|
||||
The node queries the tracker for `/v1/network/map`, discovers the relay URL, and
|
||||
opens a persistent outbound WebSocket. If the relay connection drops, the node
|
||||
keeps retrying it.
|
||||
|
||||
```bash
|
||||
.venv/bin/meshnet-node start \
|
||||
--tracker https://ai.neuron.d-popov.com \
|
||||
--model Qwen/Qwen2.5-0.5B-Instruct
|
||||
```
|
||||
|
||||
No authentication is required in the prototype. The first public node for a model
|
||||
must still choose that model with `--model` or a saved wizard config. After at
|
||||
least one HF model node is registered, later nodes can join the public network
|
||||
with only the tracker URL; the tracker assigns an uncovered shard if one exists:
|
||||
|
||||
```bash
|
||||
.venv/bin/meshnet-node start --tracker https://ai.neuron.d-popov.com
|
||||
```
|
||||
|
||||
Use the public tracker to verify registration and routing:
|
||||
|
||||
```bash
|
||||
curl -s "https://ai.neuron.d-popov.com/v1/network/map" | python3 -m json.tool
|
||||
curl -s "https://ai.neuron.d-popov.com/v1/route?model=qwen2.5-0.5b" | python3 -m json.tool
|
||||
```
|
||||
|
||||
Expected startup output (relay path):
|
||||
|
||||
```
|
||||
Auto-detected 24 layers → shard 0–23
|
||||
Relay connected — wss://ai.neuron.d-popov.com/rpc/abc1def2ef3f4567
|
||||
================================
|
||||
meshnet-node ready
|
||||
Wallet: <address>
|
||||
Model ID: Qwen/Qwen2.5-0.5B-Instruct
|
||||
Shard: layers 0–23; 24 of 24
|
||||
Quantization: bfloat16
|
||||
Endpoint: http://172.29.104.23:7001
|
||||
Node ID: <id>
|
||||
Hardware: CPU
|
||||
================================
|
||||
```
|
||||
|
||||
The `Endpoint` shown is the local IP (unreachable from outside). Other nodes reach
|
||||
this one via `wss://ai.neuron.d-popov.com/rpc/<peer_id>` instead.
|
||||
|
||||
### How relay hops work
|
||||
|
||||
When node A needs to forward activations to node B (behind NAT):
|
||||
|
||||
1. Tracker injects `X-Meshnet-Route` with `relay_addr` for each behind-NAT hop.
|
||||
2. Node A opens a WebSocket to `wss://relay/rpc/{peer_id_B}`.
|
||||
3. Relay forwards the `relay-http-request` envelope to Node B's persistent connection.
|
||||
4. Node B processes `/forward` locally, returns `relay-http-response`.
|
||||
5. Relay sends the response back to Node A over the same WebSocket.
|
||||
6. Node A closes the WebSocket and continues the pipeline.
|
||||
|
||||
Binary activation tensors (bfloat16) are Base64-encoded through the relay JSON
|
||||
protocol and decoded on both sides — no precision loss.
|
||||
|
||||
If the relay hop fails (relay down, peer disconnected), the node logs a warning and
|
||||
falls back to a direct HTTP attempt before returning an error.
|
||||
|
||||
### Test from WSL2 using the public tracker
|
||||
|
||||
In WSL2 (which gets a `172.x.x.x` virtual IP — unreachable from other machines):
|
||||
|
||||
```bash
|
||||
# WSL2 Terminal 1 — head node (layers 0–11, handles chat requests)
|
||||
.venv/bin/meshnet-node start \
|
||||
--tracker https://ai.neuron.d-popov.com \
|
||||
--model Qwen/Qwen2.5-0.5B-Instruct \
|
||||
--shard-start 0 --shard-end 11
|
||||
|
||||
# WSL2 Terminal 2 — tail node (layers 12–23)
|
||||
.venv/bin/meshnet-node start \
|
||||
--tracker https://ai.neuron.d-popov.com \
|
||||
--model Qwen/Qwen2.5-0.5B-Instruct \
|
||||
--shard-start 12 --shard-end 23
|
||||
```
|
||||
|
||||
Both nodes connect to the relay automatically. When a chat request arrives at Node A,
|
||||
it forwards activations to Node B via `wss://ai.neuron.d-popov.com/rpc/{peer_id_B}`.
|
||||
|
||||
Send inference through the tracker (which picks the head node and injects the route):
|
||||
|
||||
```bash
|
||||
curl -s https://ai.neuron.d-popov.com/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"messages": [{"role": "user", "content": "What is 7 times 8?"}],
|
||||
"stream": false
|
||||
}' | python3 -m json.tool
|
||||
```
|
||||
|
||||
Or send directly to Node A's local port (within WSL):
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:7001/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model": "Qwen/Qwen2.5-0.5B-Instruct", "messages": [{"role": "user", "content": "Hi"}]}'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 1 — Start the tracker (Terminal 1)
|
||||
|
||||
```bash
|
||||
cd /run/media/popov/d/DEV/repos/d-popov.com/AI
|
||||
.venv/bin/meshnet-tracker start --port 8080
|
||||
```
|
||||
|
||||
Expected output:
|
||||
```
|
||||
Tracker listening on 0.0.0.0:8080
|
||||
```
|
||||
|
||||
Keep this terminal open.
|
||||
|
||||
---
|
||||
|
||||
## Step 2 — Start a node (Terminal 2)
|
||||
|
||||
### Recommended model: Qwen2.5-0.5B-Instruct
|
||||
|
||||
- 0.5B parameters, ~1 GB in BF16
|
||||
- No HuggingFace account or license required
|
||||
- Downloads once to `~/.meshnet/models/`, cached for future runs
|
||||
- 24 transformer layers (auto-detected — no need to specify)
|
||||
|
||||
```bash
|
||||
cd /run/media/popov/d/DEV/repos/d-popov.com/AI
|
||||
HF_HOME=/run/media/popov/d/DEV/models \
|
||||
.venv/bin/meshnet-node start \
|
||||
--model Qwen/Qwen2.5-0.5B-Instruct \
|
||||
--quantization bfloat16 \
|
||||
--tracker http://localhost:8080 \
|
||||
--port 8001
|
||||
```
|
||||
|
||||
Shard range is **auto-detected** from the curated catalog (no network call for known
|
||||
models). For unknown repos, the node fetches only `config.json` (~1 KB) to read
|
||||
`num_hidden_layers`. You can still pass `--shard-start` / `--shard-end` explicitly
|
||||
to run a partial shard on one machine.
|
||||
|
||||
Expected output (after model loads):
|
||||
```
|
||||
Auto-detected 24 layers → shard 0–23
|
||||
================================
|
||||
meshnet-node ready
|
||||
Wallet: <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)
|
||||
|
||||
If you started the node with `--port 8001`, send the request directly to that
|
||||
head node:
|
||||
|
||||
```bash Qwen2.5-0.5B-Instruct
|
||||
curl -s http://localhost:8001/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "qwen2.5-0.5b",
|
||||
"messages": [{"role": "user", "content": "What is 7 times 8? Answer in one word."}],
|
||||
"stream": false
|
||||
}' | python3 -m json.tool
|
||||
```
|
||||
|
||||
If you did not pass `--port`, `meshnet-node start` uses the first free port at
|
||||
or above `7000`. Use the `Endpoint:` printed by the node instead of `8001`.
|
||||
|
||||
To test tracker routing/proxying, send the same OpenAI-compatible request to the
|
||||
tracker, using either the full HuggingFace repo or the quick alias:
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8080/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "qwen2.5-0.5b",
|
||||
"messages": [{"role": "user", "content": "What is 7 times 8? Answer in one word."}],
|
||||
"stream": false
|
||||
}' | python3 -m json.tool
|
||||
```
|
||||
|
||||
Or use the test script:
|
||||
|
||||
```bash
|
||||
.venv/bin/python scripts/test_lan_inference.py \
|
||||
--tracker http://localhost:8080 \
|
||||
--gateway http://localhost:8001
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Two-node split (same machine, two terminals)
|
||||
|
||||
Split Qwen2.5-0.5B's 24 layers across two node processes to test the sharded pipeline:
|
||||
|
||||
**Node A — layers 0–11 (tracker mode, serves chat completions):**
|
||||
```bash
|
||||
HF_HOME=/run/media/popov/d/DEV/models \
|
||||
.venv/bin/meshnet-node start \
|
||||
--model Qwen/Qwen2.5-0.5B-Instruct \
|
||||
--shard-start 0 --shard-end 11 \
|
||||
--quantization bfloat16 \
|
||||
--tracker http://localhost:8080 \
|
||||
--port 8001
|
||||
```
|
||||
|
||||
**Node B — layers 12–23:**
|
||||
```bash
|
||||
HF_HOME=/run/media/popov/d/DEV/models \
|
||||
.venv/bin/meshnet-node start \
|
||||
--model Qwen/Qwen2.5-0.5B-Instruct \
|
||||
--shard-start 12 --shard-end 23 \
|
||||
--quantization bfloat16 \
|
||||
--tracker http://localhost:8080 \
|
||||
--port 8002
|
||||
```
|
||||
|
||||
Send the request to Node A — it tokenizes, runs layers 0–11, passes binary
|
||||
activations to Node B, and streams the final response back.
|
||||
|
||||
---
|
||||
|
||||
## Two-machine LAN test (Linux + Windows/WSL2)
|
||||
|
||||
See `docs/TWO_MACHINE_TEST.md` (created by US-018).
|
||||
|
||||
For WSL2 nodes, registration only proves the node can reach the tracker
|
||||
outbound. Tracker-routed inference also requires the tracker to reach the node's
|
||||
advertised endpoint inbound. Either run the node in native Windows PowerShell,
|
||||
configure Windows port forwarding into WSL for the node port, or start the
|
||||
tracker with a relay URL so the node registers a `relay_addr`.
|
||||
|
||||
---
|
||||
|
||||
## Browse available models
|
||||
|
||||
```bash
|
||||
# Show curated list with VRAM requirements
|
||||
.venv/bin/meshnet-node models
|
||||
|
||||
# Browse HuggingFace Hub top-20 text-generation models
|
||||
.venv/bin/meshnet-node models --browse
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Start with the interactive wizard
|
||||
|
||||
```bash
|
||||
# First run: wizard detects GPU, shows model list, saves config
|
||||
.venv/bin/meshnet-node
|
||||
|
||||
# Subsequent runs: starts directly from saved config
|
||||
.venv/bin/meshnet-node
|
||||
|
||||
# Re-run wizard even with saved config
|
||||
.venv/bin/meshnet-node --reset-config
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Run all tests
|
||||
|
||||
```bash
|
||||
.venv/bin/python -m pytest -q
|
||||
```
|
||||
11
_DEV_NOTES.md
Normal file
11
_DEV_NOTES.md
Normal file
@@ -0,0 +1,11 @@
|
||||
win
|
||||
.venv/bin/meshnet-node start --tracker http://192.168.0.179:8081 --model Qwen/Qwen2.5-0.5B-Instruct --shard-start 20 --advertise-host 192.168.0.20
|
||||
|
||||
-.\.venv\Scripts\meshnet-node.exe start http://192.168.0.179:8081 --model-id Qwen/Qwen2.5-0.5B-Instruct --advertise-host 192.168.0.20
|
||||
-.\.venv\Scripts\meshnet-node.exe start --tracker http://ai.neuron.d-popov.com --model-id Qwen/Qwen2.5-0.5B-Instruct --advertise-host 192.168.0.20
|
||||
|
||||
we .\.venv\Scripts\meshnet-node.exe start `
|
||||
--tracker http://192.168.0.179:8081 `
|
||||
--model Qwen/Qwen2.5-0.5B-Instruct `
|
||||
--advertise-host 192.168.0.20
|
||||
|
||||
212
docs/INSTALL_WINDOWS.md
Normal file
212
docs/INSTALL_WINDOWS.md
Normal file
@@ -0,0 +1,212 @@
|
||||
# Installing meshnet-node on Windows 11 with WSL2
|
||||
|
||||
This guide covers setting up a meshnet-node on a Windows 11 machine using WSL2 with CUDA passthrough so it can join an existing inference network over LAN.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Windows 11 with WSL2 support (most systems with Windows 10 version 2004+ qualify)
|
||||
- NVIDIA GPU with CUDA support (driver ≥ 525.x recommended for WSL2 CUDA)
|
||||
- At least 8 GB RAM + enough VRAM for the model shard you intend to serve
|
||||
- The Linux machine (other node) is reachable on your LAN
|
||||
|
||||
---
|
||||
|
||||
## Step 1 — Enable WSL2 and install Ubuntu
|
||||
|
||||
Open **PowerShell as Administrator** and run:
|
||||
|
||||
```powershell
|
||||
wsl --install -d Ubuntu-24.04
|
||||
```
|
||||
|
||||
This installs WSL2 with Ubuntu 24.04. Reboot when prompted.
|
||||
|
||||
After reboot, Ubuntu starts and asks you to create a UNIX username/password. Choose anything convenient.
|
||||
|
||||
Verify WSL version:
|
||||
|
||||
```powershell
|
||||
wsl -l -v
|
||||
```
|
||||
|
||||
Output should show `VERSION 2`.
|
||||
|
||||
---
|
||||
|
||||
## Step 2 — Install NVIDIA GPU driver on Windows (NOT inside WSL)
|
||||
|
||||
WSL2 CUDA passthrough works through the Windows host driver. **Do not install CUDA inside WSL2.**
|
||||
|
||||
1. Download the latest Game Ready or Studio driver for your GPU from https://www.nvidia.com/drivers
|
||||
2. Install on Windows normally (standard installer).
|
||||
3. Inside WSL2 (Ubuntu terminal), verify:
|
||||
|
||||
```bash
|
||||
nvidia-smi
|
||||
```
|
||||
|
||||
Expected output: your GPU name, driver version, CUDA version. If this command fails, the Windows driver is too old — update it.
|
||||
|
||||
> **Note:** The `cuda-toolkit` package inside WSL2 is optional and only needed if you compile CUDA kernels. For inference with `torch`, the Windows host driver is sufficient.
|
||||
|
||||
---
|
||||
|
||||
## Step 3 — Install Python 3.11+ inside WSL2
|
||||
|
||||
Ubuntu 24.04 ships Python 3.12. Confirm:
|
||||
|
||||
```bash
|
||||
python3 --version
|
||||
```
|
||||
|
||||
If it shows 3.10 or older:
|
||||
|
||||
```bash
|
||||
sudo add-apt-repository ppa:deadsnakes/ppa
|
||||
sudo apt update
|
||||
sudo apt install python3.12 python3.12-venv python3.12-dev
|
||||
```
|
||||
|
||||
Install pip:
|
||||
|
||||
```bash
|
||||
curl -sS https://bootstrap.pypa.io/get-pip.py | python3
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 4 — Clone the repository
|
||||
|
||||
Inside WSL2:
|
||||
|
||||
```bash
|
||||
# Store the repo in the Linux filesystem (faster I/O than /mnt/c)
|
||||
cd ~
|
||||
git clone https://github.com/YOUR_ORG/d-popov.com.git
|
||||
cd d-popov.com/AI
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 5 — Create a virtualenv and install meshnet-node
|
||||
|
||||
```bash
|
||||
python3 -m venv .venv
|
||||
source .venv/bin/activate
|
||||
|
||||
# Install node + PyTorch (CUDA build)
|
||||
pip install torch --index-url https://download.pytorch.org/whl/cu124
|
||||
pip install -e "packages/node[torch]"
|
||||
```
|
||||
|
||||
Verify the install:
|
||||
|
||||
```bash
|
||||
meshnet-node --help
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Step 6 — Pre-download the model shard
|
||||
|
||||
Download the model before starting the node so the startup process doesn't time out on the tracker side:
|
||||
|
||||
```bash
|
||||
python3 - <<'EOF'
|
||||
from transformers import AutoConfig
|
||||
AutoConfig.from_pretrained("microsoft/Phi-3-medium-128k-instruct")
|
||||
EOF
|
||||
```
|
||||
|
||||
For the full model weights (needed at runtime), `transformers` downloads them automatically on first `meshnet-node` start. If you want to pre-fetch:
|
||||
|
||||
```bash
|
||||
python3 -c "
|
||||
from transformers import AutoModelForCausalLM
|
||||
AutoModelForCausalLM.from_pretrained('microsoft/Phi-3-medium-128k-instruct', device_map='cpu')
|
||||
"
|
||||
```
|
||||
|
||||
This can take 10–30 minutes on first run.
|
||||
|
||||
---
|
||||
|
||||
## Step 7 — Expose the node port to your LAN
|
||||
|
||||
WSL2 runs behind a NAT with a virtual IP (typically `172.x.x.x`). Your LAN sees the Windows host IP. You need to forward the node port.
|
||||
|
||||
**Option A — Windows port proxy (recommended for simple setups):**
|
||||
|
||||
In **PowerShell as Administrator**:
|
||||
|
||||
```powershell
|
||||
# Get the current WSL2 IP (changes on each WSL restart)
|
||||
$wslIp = (wsl hostname -I).Trim()
|
||||
|
||||
# Forward Windows host port 8001 → WSL2 port 8001
|
||||
netsh interface portproxy add v4tov4 `
|
||||
listenport=8001 listenaddress=0.0.0.0 `
|
||||
connectport=8001 connectaddress=$wslIp
|
||||
|
||||
# Allow inbound on Windows Firewall
|
||||
New-NetFirewallRule -DisplayName "meshnet-node" `
|
||||
-Direction Inbound -Protocol TCP -LocalPort 8001 -Action Allow
|
||||
```
|
||||
|
||||
Verify: from the Linux machine, `curl http://WINDOWS_LAN_IP:8001/v1/health` should return a response once the node is running.
|
||||
|
||||
**Redo this after every WSL2 restart** — the WSL2 IP changes.
|
||||
|
||||
**Option B — P2P relay (US-017, no port forwarding needed):**
|
||||
|
||||
Start a relay node on the Linux machine. The WSL2 node connects outbound through the relay. No firewall rules needed. See `docs/TWO_MACHINE_TEST.md` for details.
|
||||
|
||||
---
|
||||
|
||||
## Step 8 — Start the node
|
||||
|
||||
Replace `192.168.1.10` with the actual LAN IP of the Linux machine running the tracker.
|
||||
Replace shard range with the complementary range to what the Linux node is serving.
|
||||
|
||||
```bash
|
||||
source .venv/bin/activate
|
||||
|
||||
meshnet-node \
|
||||
--model microsoft/Phi-3-medium-128k-instruct \
|
||||
--quantization bf16 \
|
||||
--shard-start 20 --shard-end 39 \
|
||||
--tracker http://192.168.1.10:8080 \
|
||||
--port 8001 \
|
||||
--host 0.0.0.0 \
|
||||
--advertise-host WINDOWS_LAN_IP
|
||||
```
|
||||
|
||||
The `--advertise-host` flag tells the tracker what IP the Linux machine should use to reach this node. Use your Windows machine's LAN IP (e.g. `192.168.1.20`), **not** the WSL2 internal IP.
|
||||
|
||||
Expected startup output:
|
||||
|
||||
```
|
||||
Detecting hardware...
|
||||
GPU: NVIDIA GeForce RTX 3080 (10240 MB VRAM)
|
||||
Loading wallet...
|
||||
Wallet: 5K7r...
|
||||
Loading real PyTorch model shard...
|
||||
Auto-detected 40 layers → shard 20–39
|
||||
================================
|
||||
meshnet-node ready
|
||||
Model ID: microsoft/Phi-3-medium-128k-instruct
|
||||
Shard: layers 20–39; 20 of 40
|
||||
Endpoint: http://192.168.1.20:8001
|
||||
Hardware: CUDA
|
||||
================================
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Known issues
|
||||
|
||||
- **WSL2 IP changes on restart.** Always re-run the `netsh` port-proxy command after restarting WSL2 or Windows.
|
||||
- **CUDA not visible in WSL2.** If `nvidia-smi` fails inside WSL2, update the Windows host GPU driver to ≥ 525.x. Installing CUDA inside WSL2 will not fix it.
|
||||
- **Model download is slow.** HuggingFace downloads happen over HTTPS. Pre-fetch the model before a timed test (see Step 6).
|
||||
- **Port 8001 already in use.** Change `--port` to another value and update the firewall/portproxy rules accordingly.
|
||||
- **`bf16` not supported on older GPUs.** Use `--quantization int8` on Turing (RTX 20xx) cards or earlier if bfloat16 ops fail.
|
||||
200
docs/TWO_MACHINE_TEST.md
Normal file
200
docs/TWO_MACHINE_TEST.md
Normal file
@@ -0,0 +1,200 @@
|
||||
# Two-machine LAN inference test
|
||||
|
||||
This guide proves that distributed inference works across two physical machines: a Linux rig (tracker + first shard) and a Windows 11 / WSL2 rig (second shard). A test script sends real inference requests and validates the output.
|
||||
|
||||
## Network topology
|
||||
|
||||
```
|
||||
[Linux machine — 192.168.1.10]
|
||||
meshnet-tracker :8080
|
||||
meshnet-node A :8001 shard 0–19 (tracker-mode, entry point)
|
||||
|
||||
[Windows 11 / WSL2 — 192.168.1.20]
|
||||
meshnet-node B :8001 shard 20–39
|
||||
|
||||
[Client — either machine]
|
||||
scripts/test_lan_inference.py --tracker http://192.168.1.10:8080
|
||||
```
|
||||
|
||||
Adjust the IPs and shard ranges to match your hardware. Use a model that fits (sharded) in both GPUs combined. The example uses `microsoft/Phi-3-medium-128k-instruct` (40 layers, BF16 ~15 GB each shard ~7.5 GB).
|
||||
|
||||
---
|
||||
|
||||
## Prerequisites
|
||||
|
||||
**Both machines:**
|
||||
- Python 3.11+ with `meshnet-node` installed (see `docs/INSTALL_WINDOWS.md` for Windows)
|
||||
- Model weights already downloaded (pre-fetch prevents timeout on first startup)
|
||||
- LAN connectivity verified: `ping 192.168.1.10` from Windows, `ping 192.168.1.20` from Linux
|
||||
|
||||
**Linux machine ports open:**
|
||||
|
||||
```bash
|
||||
# ufw (skip if firewall is off)
|
||||
sudo ufw allow 8080/tcp # tracker
|
||||
sudo ufw allow 8001/tcp # node A
|
||||
```
|
||||
|
||||
**Windows machine port forwarded (WSL2 only):**
|
||||
|
||||
```powershell
|
||||
# Run in PowerShell as Administrator — redo after every WSL restart
|
||||
$wsl = (wsl hostname -I).Trim()
|
||||
netsh interface portproxy add v4tov4 listenport=8001 listenaddress=0.0.0.0 connectport=8001 connectaddress=$wsl
|
||||
New-NetFirewallRule -DisplayName "meshnet-node" -Direction Inbound -Protocol TCP -LocalPort 8001 -Action Allow
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Start sequence
|
||||
|
||||
**Always start in this order: tracker → node A → node B → test.**
|
||||
|
||||
### Terminal 1 — Linux: tracker
|
||||
|
||||
```bash
|
||||
meshnet-tracker --port 8080
|
||||
```
|
||||
|
||||
Expected:
|
||||
|
||||
```
|
||||
[tracker] listening on 0.0.0.0:8080
|
||||
```
|
||||
|
||||
### Terminal 2 — Linux: node A (shard 0–19, tracker-mode)
|
||||
|
||||
```bash
|
||||
meshnet-node \
|
||||
--model microsoft/Phi-3-medium-128k-instruct \
|
||||
--quantization bf16 \
|
||||
--shard-start 0 --shard-end 19 \
|
||||
--tracker http://localhost:8080 \
|
||||
--port 8001 \
|
||||
--host 0.0.0.0
|
||||
```
|
||||
|
||||
`shard_start=0` auto-sets `tracker_mode=True` — this node accepts inference requests.
|
||||
|
||||
Wait until you see `meshnet-node ready` before continuing.
|
||||
|
||||
### Terminal 3 — Windows WSL2: node B (shard 20–39)
|
||||
|
||||
```bash
|
||||
meshnet-node \
|
||||
--model microsoft/Phi-3-medium-128k-instruct \
|
||||
--quantization bf16 \
|
||||
--shard-start 20 --shard-end 39 \
|
||||
--tracker http://192.168.1.10:8080 \
|
||||
--port 8001 \
|
||||
--host 0.0.0.0 \
|
||||
--advertise-host 192.168.1.20
|
||||
```
|
||||
|
||||
`--advertise-host` must be the Windows **LAN IP** (not the WSL2 internal 172.x.x.x IP) so the Linux node can reach it.
|
||||
|
||||
---
|
||||
|
||||
## Verify nodes are registered
|
||||
|
||||
From any machine with `curl`:
|
||||
|
||||
```bash
|
||||
# List all registered nodes
|
||||
curl http://192.168.1.10:8080/v1/nodes
|
||||
|
||||
# Check route for the model — should list both node endpoints in order
|
||||
curl "http://192.168.1.10:8080/v1/route?model=microsoft/Phi-3-medium-128k-instruct"
|
||||
```
|
||||
|
||||
Expected route response:
|
||||
|
||||
```json
|
||||
{
|
||||
"route": [
|
||||
"http://192.168.1.10:8001",
|
||||
"http://192.168.1.20:8001"
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
If only one endpoint appears, node B hasn't registered yet — wait a few seconds and retry.
|
||||
|
||||
---
|
||||
|
||||
## Run the test script
|
||||
|
||||
```bash
|
||||
# From any machine that can reach the tracker
|
||||
python3 scripts/test_lan_inference.py \
|
||||
--tracker http://192.168.1.10:8080 \
|
||||
--gateway http://192.168.1.10:8001
|
||||
```
|
||||
|
||||
Expected output:
|
||||
|
||||
```
|
||||
Inference endpoint: http://192.168.1.10:8001
|
||||
Tracker: http://192.168.1.10:8080
|
||||
|
||||
Route: ['http://192.168.1.10:8001', 'http://192.168.1.20:8001']
|
||||
|
||||
[1] Q: What is 7 × 8? Answer in one word.
|
||||
A: 56
|
||||
3 tokens 2.41s 1.2 t/s
|
||||
|
||||
[2] Q: Name the capital of France in one word.
|
||||
A: Paris
|
||||
2 tokens 1.87s 1.1 t/s
|
||||
|
||||
[3] Q: Complete the sequence: 1, 1, 2, 3, 5, ___
|
||||
A: 8
|
||||
2 tokens 1.93s 1.0 t/s
|
||||
|
||||
All 3 requests completed successfully.
|
||||
Exit code: 0
|
||||
```
|
||||
|
||||
The script exits 0 if all 3 requests complete with valid OpenAI-format responses.
|
||||
|
||||
---
|
||||
|
||||
## Reading latency from node logs
|
||||
|
||||
The node logs show per-hop timing. On node A terminal look for:
|
||||
|
||||
```
|
||||
[node] forwarding to downstream: http://192.168.1.20:8001 (took 1.23s)
|
||||
```
|
||||
|
||||
Approximate breakdown:
|
||||
- **client → node A (encode + first shard):** full request latency minus the downstream time
|
||||
- **node A → node B (pipeline):** the `forwarding to downstream` duration
|
||||
- **node B → node A (tail decode + token):** included in downstream duration
|
||||
|
||||
Full end-to-end latency = prompt encode + shard A forward + network transfer + shard B forward + decode.
|
||||
|
||||
With LAN latency < 1 ms, the network transfer is negligible. Bottleneck is GPU compute.
|
||||
|
||||
---
|
||||
|
||||
## Known Issues
|
||||
|
||||
**WSL2 IP changes after restart.**
|
||||
The `netsh portproxy` forwarding rule uses a fixed WSL2 IP. If Windows or WSL2 restarts, the IP changes and the rule breaks. Redo the `netsh` and `New-NetFirewallRule` commands. To automate this, add a Task Scheduler job on WSL start.
|
||||
|
||||
**Node B registers with internal WSL2 IP (172.x.x.x) instead of LAN IP.**
|
||||
Symptom: route response lists `172.x.x.x` and node A cannot reach it.
|
||||
Fix: always pass `--advertise-host 192.168.1.20` (your Windows LAN IP) when starting node B.
|
||||
|
||||
**Model download times out node registration.**
|
||||
If the model hasn't been pre-fetched, `transformers` downloads it during node startup, which can take 20+ minutes. The tracker heartbeat timeout (90s) will expire, and node A will deregister node B. Pre-download the model weights before starting the node (see `docs/INSTALL_WINDOWS.md` Step 6). Node B re-registers automatically via the heartbeat re-registration loop once it's up.
|
||||
|
||||
**`bf16` unsupported on older NVIDIA GPUs.**
|
||||
GPUs before Ampere (RTX 30xx) have limited bfloat16 support. Use `--quantization int8` on RTX 20xx and earlier.
|
||||
|
||||
**Windows Defender blocks inbound connection on WSL2.**
|
||||
Even with the firewall rule added, Windows Defender SmartScreen or a corporate security policy can block the connection. Verify by checking Windows Event Viewer → Security → Filtering Platform Connection for blocked connections on port 8001.
|
||||
|
||||
**Route returns only one node.**
|
||||
If node B registers but the route only returns one endpoint, check that both nodes use the same `--model` string (full HuggingFace repo path). Route lookup matches on `hf_repo` — a short name vs. full path mismatch causes the node to be excluded.
|
||||
67
docs/adr/0010-p2p-gossip-and-nat-relay.md
Normal file
67
docs/adr/0010-p2p-gossip-and-nat-relay.md
Normal file
@@ -0,0 +1,67 @@
|
||||
# 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.
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import socket
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
@@ -26,6 +27,10 @@ def _run_node(cfg: dict) -> None:
|
||||
wallet_path=Path(cfg["wallet_path"]) if cfg.get("wallet_path") else None,
|
||||
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
|
||||
host=cfg.get("host", "0.0.0.0"),
|
||||
advertise_host=cfg.get("advertise_host"),
|
||||
route_timeout=float(cfg.get("route_timeout", 30.0)),
|
||||
vram_mb_override=cfg.get("vram_mb_override"),
|
||||
debug=bool(cfg.get("debug", False)),
|
||||
)
|
||||
except Exception as exc:
|
||||
print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
|
||||
@@ -48,6 +53,19 @@ def _run_node(cfg: dict) -> None:
|
||||
)
|
||||
|
||||
|
||||
def _first_available_port(host: str, start: int = 7000, attempts: int = 100) -> int:
|
||||
"""Return the first TCP port bindable on host, starting at start."""
|
||||
bind_host = "" if host == "0.0.0.0" else host
|
||||
for port in range(start, start + attempts):
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
||||
try:
|
||||
sock.bind((bind_host, port))
|
||||
except OSError:
|
||||
continue
|
||||
return port
|
||||
raise OSError(f"no available TCP port in range {start}-{start + attempts - 1}")
|
||||
|
||||
|
||||
def _cmd_default(args) -> int:
|
||||
"""No subcommand: wizard if no config, else start with saved config."""
|
||||
from .config import load_config, save_config, merge_cli_overrides
|
||||
@@ -86,6 +104,14 @@ def _cmd_default(args) -> int:
|
||||
overrides["port"] = args.port
|
||||
if args.host:
|
||||
overrides["host"] = args.host
|
||||
if args.advertise_host:
|
||||
overrides["advertise_host"] = args.advertise_host
|
||||
if args.route_timeout != 30.0:
|
||||
overrides["route_timeout"] = args.route_timeout
|
||||
if getattr(args, "memory", None) is not None:
|
||||
overrides["vram_mb_override"] = args.memory
|
||||
if args.debug:
|
||||
overrides["debug"] = True
|
||||
|
||||
if overrides:
|
||||
cfg = merge_cli_overrides(cfg, **overrides)
|
||||
@@ -142,8 +168,12 @@ def _cmd_start(args) -> int:
|
||||
# Build a transient config from flags (don't write to disk)
|
||||
cfg = dict(DEFAULTS)
|
||||
cfg["tracker_url"] = args.tracker
|
||||
cfg["port"] = args.port
|
||||
cfg["model_name"] = args.model
|
||||
cfg["port"] = args.port if args.port is not None else _first_available_port(args.host)
|
||||
if args.model_id is None and "/" in args.model:
|
||||
cfg["model_hf_repo"] = args.model
|
||||
cfg["model_name"] = args.model.split("/")[-1]
|
||||
else:
|
||||
cfg["model_name"] = args.model
|
||||
cfg["quantization"] = args.quantization
|
||||
cfg["host"] = args.host
|
||||
if args.model_id:
|
||||
@@ -173,6 +203,9 @@ def _cmd_start(args) -> int:
|
||||
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
|
||||
host=cfg["host"],
|
||||
advertise_host=getattr(args, "advertise_host", None),
|
||||
route_timeout=getattr(args, "route_timeout", 30.0),
|
||||
vram_mb_override=getattr(args, "memory", None),
|
||||
debug=getattr(args, "debug", False),
|
||||
)
|
||||
except Exception as exc:
|
||||
print(f"ERROR: {exc}", file=sys.stderr, flush=True)
|
||||
@@ -212,6 +245,12 @@ def main() -> None:
|
||||
parser.add_argument("--shard-end", type=int, metavar="N", help="Pin shard end layer")
|
||||
parser.add_argument("--port", type=int, metavar="N", help="Port to listen on")
|
||||
parser.add_argument("--host", metavar="ADDR", help="Interface to bind (default 0.0.0.0)")
|
||||
parser.add_argument("--advertise-host", metavar="ADDR", help="Host/IP advertised to the tracker")
|
||||
parser.add_argument("--route-timeout", type=float, metavar="SEC", default=30.0,
|
||||
help="Seconds to wait for tracker route lookup (default 30)")
|
||||
parser.add_argument("--memory", type=int, metavar="MB", default=None,
|
||||
help="Override autodetected VRAM/RAM budget in MB used for shard assignment")
|
||||
parser.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
|
||||
parser.add_argument("--no-tui", action="store_true", help="Plain-text output (no rich dashboard)")
|
||||
parser.add_argument("--compact", action="store_true", help="Single-line status output")
|
||||
parser.add_argument("--reset-config", action="store_true", help="Re-run wizard even if config exists")
|
||||
@@ -228,7 +267,7 @@ def main() -> None:
|
||||
# start subcommand (legacy / backward-compat)
|
||||
start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)")
|
||||
start_cmd.add_argument("--tracker", default="http://localhost:8080")
|
||||
start_cmd.add_argument("--port", type=int, default=7000)
|
||||
start_cmd.add_argument("--port", type=int)
|
||||
start_cmd.add_argument("--model", default="stub-model")
|
||||
start_cmd.add_argument("--model-id", help="HuggingFace repo ID")
|
||||
start_cmd.add_argument("--shard-start", type=int)
|
||||
@@ -240,6 +279,11 @@ def main() -> None:
|
||||
start_cmd.add_argument("--tracker-url", default=None)
|
||||
start_cmd.add_argument("--wallet")
|
||||
start_cmd.add_argument("--download-dir")
|
||||
start_cmd.add_argument("--route-timeout", type=float, default=30.0,
|
||||
help="Seconds to wait for tracker route lookup (default 30)")
|
||||
start_cmd.add_argument("--memory", type=int, default=None, metavar="MB",
|
||||
help="Override autodetected VRAM/RAM budget in MB used for shard assignment")
|
||||
start_cmd.add_argument("--debug", action="store_true", help="Enable verbose node debug logging")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@@ -25,6 +25,7 @@ DEFAULTS = {
|
||||
"shard_end": None,
|
||||
"port": 7000,
|
||||
"host": "0.0.0.0",
|
||||
"debug": False,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -107,12 +107,13 @@ class TorchModelShard:
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
|
||||
self.layers = _model_layers(self.model)
|
||||
self.total_layers = len(self.layers)
|
||||
if shard_end > self.total_layers:
|
||||
# shard_end is INCLUSIVE (last layer index, 0-based), matching the CLI convention.
|
||||
if shard_end >= self.total_layers:
|
||||
raise ValueError(
|
||||
f"shard_end {shard_end} exceeds total layer count {self.total_layers}"
|
||||
f"shard_end {shard_end} exceeds last layer index {self.total_layers - 1}"
|
||||
)
|
||||
self.is_head = shard_start == 0
|
||||
self.is_tail = shard_end == self.total_layers
|
||||
self.is_tail = shard_end >= self.total_layers - 1
|
||||
self.hidden_size = int(
|
||||
getattr(self.model.config, "hidden_size", 0)
|
||||
or getattr(self.model.config, "n_embd", 0)
|
||||
@@ -144,6 +145,7 @@ class TorchModelShard:
|
||||
shape: list[int],
|
||||
attention_mask_header: str | None,
|
||||
position_ids_header: str | None,
|
||||
start_layer: int | None = None,
|
||||
) -> TensorPayload | str:
|
||||
hidden_states = _tensor_from_bfloat16_bytes(body, shape, self.torch).to(
|
||||
self.device
|
||||
@@ -154,7 +156,9 @@ class TorchModelShard:
|
||||
position_ids = _tensor_from_int64_header(
|
||||
position_ids_header, self.torch, self.device
|
||||
)
|
||||
hidden_states = self._run_layers(hidden_states, attention_mask, position_ids)
|
||||
hidden_states = self._run_layers(
|
||||
hidden_states, attention_mask, position_ids, start_layer=start_layer
|
||||
)
|
||||
if self.is_tail:
|
||||
return self.decode_tail(hidden_states)
|
||||
return self._payload(hidden_states, attention_mask, position_ids)
|
||||
@@ -168,10 +172,149 @@ 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 _run_layers(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
|
||||
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():
|
||||
for layer in self.layers[self.shard_start:self.shard_end]:
|
||||
hidden_states = _call_layer(layer, hidden_states, attention_mask, position_ids)
|
||||
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,
|
||||
start_layer: int | None = None,
|
||||
) -> Any:
|
||||
# start_layer overrides shard_start for overlapping-shard routing
|
||||
# (X-Meshnet-Start-Layer header). Clamped to shard_start to prevent
|
||||
# indexing outside the loaded weights.
|
||||
effective_start = (
|
||||
max(self.shard_start, start_layer)
|
||||
if start_layer is not None
|
||||
else self.shard_start
|
||||
)
|
||||
position_embeddings = _rotary_position_embeddings(
|
||||
self.model,
|
||||
hidden_states,
|
||||
position_ids,
|
||||
)
|
||||
layer_attention_mask = _decoder_attention_mask(
|
||||
attention_mask,
|
||||
hidden_states,
|
||||
self.torch,
|
||||
)
|
||||
with self.torch.inference_mode():
|
||||
for layer in self.layers[effective_start:self.shard_end + 1]:
|
||||
hidden_states = _call_layer(
|
||||
layer,
|
||||
hidden_states,
|
||||
layer_attention_mask,
|
||||
position_ids,
|
||||
position_embeddings,
|
||||
)
|
||||
return hidden_states.to(self.torch.bfloat16)
|
||||
|
||||
def _payload(self, hidden_states: Any, attention_mask: Any, position_ids: Any) -> TensorPayload:
|
||||
@@ -236,8 +379,60 @@ def _position_ids(attention_mask: Any, torch: Any) -> Any:
|
||||
return position_ids.masked_fill(attention_mask == 0, 0).to(torch.long)
|
||||
|
||||
|
||||
def _call_layer(layer: Any, hidden_states: Any, attention_mask: Any, position_ids: Any) -> Any:
|
||||
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:
|
||||
attempts = (
|
||||
{
|
||||
"attention_mask": attention_mask,
|
||||
"position_ids": position_ids,
|
||||
"position_embeddings": position_embeddings,
|
||||
"use_cache": False,
|
||||
},
|
||||
{
|
||||
"attention_mask": attention_mask,
|
||||
"position_ids": position_ids,
|
||||
@@ -272,7 +467,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().to(tensor.int64).contiguous()
|
||||
data = tensor.detach().cpu().long().contiguous()
|
||||
raw = data.numpy().tobytes()
|
||||
shape = ",".join(str(dim) for dim in data.shape)
|
||||
encoded = base64.b64encode(raw).decode("ascii")
|
||||
|
||||
164
packages/node/meshnet_node/relay_bridge.py
Normal file
164
packages/node/meshnet_node/relay_bridge.py
Normal file
@@ -0,0 +1,164 @@
|
||||
"""Outbound relay bridge for NAT-safe node HTTP requests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from dataclasses import dataclass
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RelayBridgeInfo:
|
||||
peer_id: str
|
||||
relay_addr: str
|
||||
|
||||
|
||||
def _make_envelope(topic: str, payload: dict, peer_id: str) -> dict:
|
||||
return {
|
||||
"topic": topic,
|
||||
"version": 1,
|
||||
"from_peer": peer_id,
|
||||
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
|
||||
"msg_id": f"{peer_id}-{time.time_ns():x}",
|
||||
"ttl": 1,
|
||||
"payload": payload,
|
||||
}
|
||||
|
||||
|
||||
class RelayHttpBridge:
|
||||
"""Connect outbound to a relay and proxy relay HTTP requests to localhost."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
relay_url: str,
|
||||
peer_id: str,
|
||||
local_base_url: str,
|
||||
advertised_addr: str,
|
||||
reconnect_interval: float = 3.0,
|
||||
) -> None:
|
||||
self.relay_url = relay_url.rstrip("/")
|
||||
self.peer_id = peer_id
|
||||
self.local_base_url = local_base_url.rstrip("/")
|
||||
self.advertised_addr = advertised_addr
|
||||
self.reconnect_interval = reconnect_interval
|
||||
self._running = False
|
||||
self._thread: threading.Thread | None = None
|
||||
self._connected = threading.Event()
|
||||
|
||||
@property
|
||||
def relay_addr(self) -> str:
|
||||
base = self.relay_url
|
||||
if base.endswith("/ws"):
|
||||
base = base[:-3]
|
||||
return f"{base}/rpc/{self.peer_id}"
|
||||
|
||||
def start(self) -> RelayBridgeInfo:
|
||||
self._running = True
|
||||
self._thread = threading.Thread(target=self._run, daemon=True, name="relay-http-bridge")
|
||||
self._thread.start()
|
||||
return RelayBridgeInfo(peer_id=self.peer_id, relay_addr=self.relay_addr)
|
||||
|
||||
def wait_connected(self, timeout: float = 5.0) -> bool:
|
||||
return self._connected.wait(timeout)
|
||||
|
||||
def stop(self) -> None:
|
||||
self._running = False
|
||||
if self._thread:
|
||||
self._thread.join(timeout=3.0)
|
||||
|
||||
def _run(self) -> None:
|
||||
import websockets.sync.client as wsc # type: ignore[import]
|
||||
|
||||
while self._running:
|
||||
try:
|
||||
with wsc.connect(self.relay_url, open_timeout=5) as ws:
|
||||
self._connected.set()
|
||||
ws.send(json.dumps(_make_envelope(
|
||||
"peer-register",
|
||||
{"peer_id": self.peer_id, "addr": self.advertised_addr},
|
||||
self.peer_id,
|
||||
)))
|
||||
while self._running:
|
||||
try:
|
||||
raw = ws.recv(timeout=1)
|
||||
except TimeoutError:
|
||||
continue
|
||||
try:
|
||||
envelope = json.loads(raw)
|
||||
except (TypeError, json.JSONDecodeError):
|
||||
continue
|
||||
if envelope.get("topic") != "relay-http-request":
|
||||
continue
|
||||
payload = envelope.get("payload", {})
|
||||
if payload.get("target_peer") not in {None, self.peer_id}:
|
||||
continue
|
||||
response = self._handle_request(payload)
|
||||
ws.send(json.dumps(_make_envelope(
|
||||
"relay-http-response",
|
||||
response,
|
||||
self.peer_id,
|
||||
)))
|
||||
except Exception as exc:
|
||||
self._connected.clear()
|
||||
if self._running:
|
||||
log.debug("relay bridge disconnected: %s", exc)
|
||||
time.sleep(self.reconnect_interval)
|
||||
|
||||
def _handle_request(self, payload: dict) -> dict:
|
||||
request_id = str(payload.get("request_id") or "")
|
||||
method = str(payload.get("method") or "POST").upper()
|
||||
path = str(payload.get("path") or "/")
|
||||
headers = payload.get("headers") if isinstance(payload.get("headers"), dict) else {}
|
||||
|
||||
# body_base64 carries binary data (e.g. bfloat16 activation tensors) safely.
|
||||
# Fallback to text "body" for backward-compat with non-binary requests.
|
||||
body_b64 = payload.get("body_base64")
|
||||
if body_b64:
|
||||
data = base64.b64decode(body_b64)
|
||||
else:
|
||||
body_text = payload.get("body") or ""
|
||||
data = body_text.encode() if isinstance(body_text, str) else bytes(body_text)
|
||||
|
||||
url = f"{self.local_base_url}{path}"
|
||||
req = urllib.request.Request(url, data=data, headers=headers, method=method)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=300.0) as resp:
|
||||
resp_bytes = resp.read()
|
||||
resp_headers = dict(resp.headers)
|
||||
# Forward all X-Meshnet-* headers so the caller can reconstruct the activation.
|
||||
is_binary = "octet-stream" in resp.headers.get("Content-Type", "")
|
||||
result: dict = {
|
||||
"request_id": request_id,
|
||||
"status": resp.status,
|
||||
"headers": resp_headers,
|
||||
}
|
||||
if is_binary:
|
||||
result["body_base64"] = base64.b64encode(resp_bytes).decode()
|
||||
else:
|
||||
result["body"] = resp_bytes.decode(errors="replace")
|
||||
return result
|
||||
except urllib.error.HTTPError as exc:
|
||||
return {
|
||||
"request_id": request_id,
|
||||
"status": exc.code,
|
||||
"headers": {"Content-Type": exc.headers.get("Content-Type", "application/json")},
|
||||
"body": exc.read().decode(errors="replace"),
|
||||
}
|
||||
except Exception as exc:
|
||||
return {
|
||||
"request_id": request_id,
|
||||
"status": 503,
|
||||
"headers": {"Content-Type": "application/json"},
|
||||
"body": json.dumps({"error": f"relay bridge local request failed: {exc}"}),
|
||||
}
|
||||
|
||||
|
||||
def peer_id_from_wallet(wallet_address: str) -> str:
|
||||
return wallet_address[:16] if len(wallet_address) >= 16 else wallet_address
|
||||
@@ -5,6 +5,8 @@ from __future__ import annotations
|
||||
import json
|
||||
import socket
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
@@ -13,6 +15,7 @@ from typing import Any
|
||||
|
||||
from .downloader import compute_shard_checksum, download_shard
|
||||
from .hardware import detect_hardware
|
||||
from .relay_bridge import RelayHttpBridge, peer_id_from_wallet
|
||||
from .server import StubNodeServer
|
||||
from .torch_server import TorchNodeServer
|
||||
from .wallet import load_or_create_wallet
|
||||
@@ -32,6 +35,221 @@ def _get_json(url: str, timeout: float = 10.0) -> dict:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def _infer_relay_url_from_tracker(tracker_url: str) -> str | None:
|
||||
"""Infer relay WebSocket URL from a public HTTPS tracker origin.
|
||||
|
||||
Public deployments colocate relay at /ws on the same host as the tracker API
|
||||
(see QUICKSTART nginx layout). Local LAN trackers use a separate relay port
|
||||
and must advertise relay_url explicitly via /v1/network/map.
|
||||
"""
|
||||
parsed = urllib.parse.urlparse(tracker_url)
|
||||
if parsed.scheme != "https":
|
||||
return None
|
||||
host = parsed.hostname
|
||||
if not host or host in ("127.0.0.1", "localhost"):
|
||||
return None
|
||||
return f"wss://{parsed.netloc}/ws"
|
||||
|
||||
|
||||
def _discover_relay_url(tracker_url: str) -> str | None:
|
||||
relay_url: str | None = None
|
||||
try:
|
||||
network_map = _get_json(f"{tracker_url}/v1/network/map", timeout=5.0)
|
||||
raw = network_map.get("relay_url")
|
||||
if isinstance(raw, str) and raw:
|
||||
relay_url = raw
|
||||
except Exception:
|
||||
pass
|
||||
return relay_url or _infer_relay_url_from_tracker(tracker_url)
|
||||
|
||||
|
||||
def _start_relay_bridge_if_available(
|
||||
tracker_url: str,
|
||||
wallet_address: str,
|
||||
local_base_url: str,
|
||||
advertised_endpoint: str,
|
||||
relay_url: str | None = None,
|
||||
) -> tuple[RelayHttpBridge | None, dict]:
|
||||
relay_url = relay_url or _discover_relay_url(tracker_url)
|
||||
if not relay_url:
|
||||
return None, {}
|
||||
peer_id = peer_id_from_wallet(wallet_address)
|
||||
bridge = RelayHttpBridge(
|
||||
relay_url=relay_url,
|
||||
peer_id=peer_id,
|
||||
local_base_url=local_base_url,
|
||||
advertised_addr=advertised_endpoint,
|
||||
)
|
||||
info = bridge.start()
|
||||
if bridge.wait_connected(timeout=5.0):
|
||||
print(f" Relay connected — {info.relay_addr}", flush=True)
|
||||
else:
|
||||
print(f" Relay configured but not connected yet — {info.relay_addr}", flush=True)
|
||||
return bridge, {
|
||||
"relay_addr": info.relay_addr,
|
||||
"peer_id": info.peer_id,
|
||||
}
|
||||
|
||||
|
||||
def _attach_relay_bridge(node: StubNodeServer | TorchNodeServer, bridge: RelayHttpBridge | None) -> None:
|
||||
setattr(node, "relay_bridge", bridge)
|
||||
if bridge is None:
|
||||
return
|
||||
original_stop = node.stop
|
||||
|
||||
def _stop_with_bridge() -> None:
|
||||
try:
|
||||
bridge.stop()
|
||||
finally:
|
||||
original_stop()
|
||||
|
||||
node.stop = _stop_with_bridge # type: ignore[method-assign]
|
||||
|
||||
|
||||
def _start_heartbeat(
|
||||
tracker_url: str,
|
||||
node_id: str,
|
||||
register_payload: dict,
|
||||
interval: float = 20.0,
|
||||
node_ref: Any | None = None,
|
||||
start_time: float | None = None,
|
||||
) -> threading.Thread:
|
||||
"""Daemon thread: sends heartbeats and re-registers automatically after tracker restarts.
|
||||
|
||||
Heartbeat body carries cumulative stats (total_requests, failed_requests,
|
||||
queue_depth, uptime_seconds, status). Stats are buffered locally during
|
||||
outage and flushed on next successful heartbeat.
|
||||
|
||||
Heartbeat response may include new_assignment: {model, shard_start, shard_end}
|
||||
which is logged for now (hot-reload implemented in US-026).
|
||||
"""
|
||||
_start_time = start_time or time.monotonic()
|
||||
|
||||
def _get_stats() -> dict:
|
||||
uptime = time.monotonic() - _start_time
|
||||
stats: dict = {"uptime_seconds": round(uptime, 1), "status": "ready"}
|
||||
if node_ref is not None:
|
||||
stats["total_requests"] = getattr(node_ref, "total_requests", 0)
|
||||
stats["failed_requests"] = getattr(node_ref, "failed_requests", 0)
|
||||
stats["queue_depth"] = getattr(node_ref, "queue_depth", 0)
|
||||
return stats
|
||||
|
||||
def _reregister() -> bool:
|
||||
nonlocal node_id
|
||||
try:
|
||||
resp = _post_json(f"{tracker_url}/v1/nodes/register", register_payload)
|
||||
node_id = resp.get("node_id", node_id)
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def _apply_directives(directives: list[dict]) -> None:
|
||||
if not directives:
|
||||
return
|
||||
if node_ref is None or not hasattr(node_ref, "apply_tracker_directives"):
|
||||
print(f" [node] tracker directives received: {directives}", flush=True)
|
||||
return
|
||||
try:
|
||||
applied = node_ref.apply_tracker_directives(directives)
|
||||
except Exception as exc:
|
||||
print(f" [node] WARNING: failed to apply tracker directives: {exc}", flush=True)
|
||||
return
|
||||
if applied:
|
||||
model_id = applied.get("model", register_payload.get("hf_repo") or register_payload.get("model"))
|
||||
register_payload["model"] = str(model_id).split("/")[-1]
|
||||
register_payload["hf_repo"] = model_id
|
||||
register_payload["shard_start"] = applied["shard_start"]
|
||||
register_payload["shard_end"] = applied["shard_end"]
|
||||
register_payload["quantization"] = applied.get("quantization", register_payload.get("quantization"))
|
||||
register_payload["tracker_mode"] = bool(applied.get("tracker_mode", False))
|
||||
|
||||
def _loop() -> None:
|
||||
nonlocal node_id
|
||||
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||||
outage_streak = 0 # consecutive intervals where tracker was unreachable
|
||||
|
||||
while True:
|
||||
time.sleep(interval)
|
||||
|
||||
if outage_streak > 0:
|
||||
# Tracker was down — attempt re-registration first (it may have restarted
|
||||
# with a clean slate and won't know this node).
|
||||
if _reregister():
|
||||
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||||
print(f" [node] re-registered after outage — node ID: {node_id}", flush=True)
|
||||
outage_streak = 0
|
||||
else:
|
||||
outage_streak += 1
|
||||
if outage_streak <= 3 or outage_streak % 10 == 0:
|
||||
print(
|
||||
f" [node] WARNING: tracker still unreachable "
|
||||
f"({outage_streak * interval:.0f}s)",
|
||||
flush=True,
|
||||
)
|
||||
continue
|
||||
|
||||
try:
|
||||
resp = _post_json(hb_url, _get_stats())
|
||||
_apply_directives(resp.get("directives", []))
|
||||
new_asgn = resp.get("new_assignment")
|
||||
if new_asgn:
|
||||
print(
|
||||
f" [node] tracker reassignment received: "
|
||||
f"model={new_asgn.get('model')!r} "
|
||||
f"shards={new_asgn.get('shard_start')}-{new_asgn.get('shard_end')}",
|
||||
flush=True,
|
||||
)
|
||||
except urllib.error.HTTPError as exc:
|
||||
if exc.code == 404:
|
||||
# Node was purged (e.g. long gap before restart noticed) — re-register now.
|
||||
print(" [node] tracker lost registration — re-registering...", flush=True)
|
||||
if _reregister():
|
||||
hb_url = f"{tracker_url}/v1/nodes/{node_id}/heartbeat"
|
||||
print(f" [node] re-registered — node ID: {node_id}", flush=True)
|
||||
else:
|
||||
print(" [node] WARNING: re-registration failed", flush=True)
|
||||
outage_streak = 1
|
||||
else:
|
||||
print(f" [node] WARNING: heartbeat failed ({exc.code}): {exc}", flush=True)
|
||||
except Exception as exc:
|
||||
outage_streak = 1
|
||||
print(f" [node] WARNING: tracker unreachable: {exc}", flush=True)
|
||||
|
||||
t = threading.Thread(target=_loop, daemon=True, name="heartbeat")
|
||||
t.start()
|
||||
return t
|
||||
|
||||
|
||||
def _warn_virtual_network_ip(ip: str | None) -> None:
|
||||
"""Print a warning when the auto-detected advertise IP is in a known virtual-network range.
|
||||
|
||||
172.16.0.0/12 is used by Docker, WSL2, and most hypervisors. Nodes behind these
|
||||
adapters are NOT directly reachable from other physical machines on the LAN, so
|
||||
cross-host pipeline hops will time out. The user must pass --advertise-host with
|
||||
their actual LAN IP (e.g. 192.168.x.x) to fix this.
|
||||
"""
|
||||
if ip is None:
|
||||
return
|
||||
try:
|
||||
parts = [int(p) for p in ip.split(".")]
|
||||
if len(parts) != 4:
|
||||
return
|
||||
a, b = parts[0], parts[1]
|
||||
# 172.16.0.0/12 → 172.16–31.x.x
|
||||
if a == 172 and 16 <= b <= 31:
|
||||
print(
|
||||
f"\n WARNING: auto-detected endpoint IP {ip} is in 172.16.0.0/12.\n"
|
||||
f" This range is used by Docker, WSL2, and virtual machines and is\n"
|
||||
f" NOT reachable from other physical machines on your LAN.\n"
|
||||
f" Cross-host pipeline hops WILL time out.\n"
|
||||
f" Fix: use a public tracker with relay (wss://…/ws), or pass\n"
|
||||
f" --advertise-host <your-LAN-ip> (e.g. 192.168.x.x).\n",
|
||||
flush=True,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def run_startup(
|
||||
tracker_url: str,
|
||||
port: int = 0,
|
||||
@@ -45,6 +263,9 @@ def run_startup(
|
||||
host: str = "127.0.0.1",
|
||||
advertise_host: str | None = None,
|
||||
contracts: Any | None = None,
|
||||
route_timeout: float = 30.0,
|
||||
vram_mb_override: int | None = None,
|
||||
debug: bool = False,
|
||||
) -> StubNodeServer | TorchNodeServer:
|
||||
"""Execute the full startup sequence and return a running node server.
|
||||
|
||||
@@ -60,15 +281,37 @@ def run_startup(
|
||||
"""
|
||||
|
||||
tracker_url = tracker_url.rstrip("/")
|
||||
relay_url = _discover_relay_url(tracker_url)
|
||||
|
||||
# 1. Hardware detection
|
||||
if advertise_host is None and host == "0.0.0.0":
|
||||
# socket.getfqdn() returns an mDNS name (.local / .localdomain) that remote
|
||||
# machines on a different OS or subnet often can't resolve. Instead, probe the
|
||||
# outbound IP by opening a UDP socket toward the tracker — no data is sent.
|
||||
try:
|
||||
_tracker_host = urllib.parse.urlparse(tracker_url).hostname or "8.8.8.8"
|
||||
_s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
||||
_s.connect((_tracker_host, 80))
|
||||
advertise_host = _s.getsockname()[0]
|
||||
_s.close()
|
||||
except Exception:
|
||||
advertise_host = socket.getfqdn()
|
||||
|
||||
if relay_url:
|
||||
print(f"Relay advertised by tracker — using outbound tunnel {relay_url}", flush=True)
|
||||
else:
|
||||
_warn_virtual_network_ip(advertise_host)
|
||||
|
||||
print("Detecting hardware...", flush=True)
|
||||
hw = detect_hardware()
|
||||
device: str = hw["device"]
|
||||
gpu_name: str | None = hw.get("gpu_name")
|
||||
vram_mb: int = hw.get("vram_mb", 0)
|
||||
|
||||
if device == "cpu":
|
||||
if vram_mb_override is not None:
|
||||
vram_mb = vram_mb_override
|
||||
print(f" Memory budget overridden to {vram_mb} MB via --memory", flush=True)
|
||||
elif device == "cpu":
|
||||
print(" WARNING: No CUDA GPU detected — running in CPU mode", flush=True)
|
||||
else:
|
||||
print(f" GPU: {gpu_name} ({vram_mb} MB VRAM)", flush=True)
|
||||
@@ -84,7 +327,8 @@ def run_startup(
|
||||
if probationary_line is not None:
|
||||
print(f" {probationary_line}", flush=True)
|
||||
|
||||
if model_id is not None:
|
||||
if model_id: # treat "" the same as None — no explicit model given
|
||||
user_pinned_shard = shard_start is not None or shard_end 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)
|
||||
@@ -93,6 +337,23 @@ 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)
|
||||
@@ -105,18 +366,66 @@ def run_startup(
|
||||
shard_start=shard_start,
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
route_timeout=route_timeout,
|
||||
debug=debug,
|
||||
)
|
||||
_node_start_time = time.monotonic()
|
||||
actual_port = node.start()
|
||||
total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
|
||||
if isinstance(total_layers, int) and total_layers > 0:
|
||||
layer_count = shard_end - shard_start + 1
|
||||
shard_label = f"layers {shard_start}–{shard_end}; {layer_count} of {total_layers}"
|
||||
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}"
|
||||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||||
tracker_url,
|
||||
address,
|
||||
local_base_url,
|
||||
endpoint,
|
||||
relay_url=relay_url,
|
||||
)
|
||||
_attach_relay_bridge(node, relay_bridge)
|
||||
# Register with tracker so other nodes can auto-join this model.
|
||||
total_layers = getattr(getattr(node, "backend", None), "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),
|
||||
"managed_assignment": not user_pinned_shard,
|
||||
**relay_fields,
|
||||
}
|
||||
tracker_node_id: str | None = None
|
||||
try:
|
||||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", reg_payload)
|
||||
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
||||
setattr(node, "tracker_node_id", tracker_node_id)
|
||||
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
||||
_start_heartbeat(tracker_url, tracker_node_id, reg_payload, node_ref=node, start_time=_node_start_time)
|
||||
except Exception as exc:
|
||||
setattr(node, "tracker_node_id", None)
|
||||
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: layers {shard_start}–{shard_end}\n"
|
||||
f" Shard: {shard_label}\n"
|
||||
f" Quantization: {quantization}\n"
|
||||
f" Endpoint: {endpoint}\n"
|
||||
f" Node ID: {tracker_node_id or 'unregistered'}\n"
|
||||
f" Hardware: {device.upper()}\n"
|
||||
f"{'=' * 32}",
|
||||
flush=True,
|
||||
@@ -125,7 +434,95 @@ def run_startup(
|
||||
if shard_start is not None or shard_end is not None:
|
||||
raise ValueError("--shard-start / --shard-end require --model-id")
|
||||
|
||||
# 3. Shard assignment from tracker
|
||||
# 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,
|
||||
route_timeout=route_timeout,
|
||||
debug=debug,
|
||||
)
|
||||
_node_start_time = time.monotonic()
|
||||
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}"
|
||||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||||
tracker_url,
|
||||
address,
|
||||
local_base_url,
|
||||
endpoint,
|
||||
relay_url=relay_url,
|
||||
)
|
||||
_attach_relay_bridge(node, relay_bridge)
|
||||
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),
|
||||
"managed_assignment": True,
|
||||
**relay_fields,
|
||||
}
|
||||
tracker_node_id = None
|
||||
try:
|
||||
reg_resp = _post_json(f"{tracker_url}/v1/nodes/register", auto_reg_payload)
|
||||
tracker_node_id = str(reg_resp.get("node_id") or "?")
|
||||
setattr(node, "tracker_node_id", tracker_node_id)
|
||||
print(f" Registered with tracker — node ID: {tracker_node_id}", flush=True)
|
||||
_start_heartbeat(tracker_url, tracker_node_id, auto_reg_payload, node_ref=node, start_time=_node_start_time)
|
||||
except Exception as exc:
|
||||
setattr(node, "tracker_node_id", None)
|
||||
print(f" Warning: tracker registration failed: {exc}", flush=True)
|
||||
shard_count = assigned_shard_end - assigned_shard_start + 1
|
||||
print(
|
||||
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" Node ID: {tracker_node_id or 'unregistered'}\n"
|
||||
f" Hardware: {device.upper()}\n"
|
||||
f"{'=' * 32}",
|
||||
flush=True,
|
||||
)
|
||||
return node
|
||||
|
||||
# 3b. Shard assignment from tracker (stub-model / preset-based path)
|
||||
print("Querying tracker for shard assignment...", flush=True)
|
||||
assign_qs = urllib.parse.urlencode({
|
||||
"model": model,
|
||||
@@ -172,6 +569,15 @@ def run_startup(
|
||||
actual_port = node.start()
|
||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||
endpoint = f"http://{public_host}:{actual_port}"
|
||||
local_base_url = f"http://127.0.0.1:{actual_port}"
|
||||
relay_bridge, relay_fields = _start_relay_bridge_if_available(
|
||||
tracker_url,
|
||||
address,
|
||||
local_base_url,
|
||||
endpoint,
|
||||
relay_url=relay_url,
|
||||
)
|
||||
_attach_relay_bridge(node, relay_bridge)
|
||||
|
||||
# 6. Register with tracker
|
||||
print("Registering with tracker...", flush=True)
|
||||
@@ -187,9 +593,11 @@ def run_startup(
|
||||
"hardware_profile": hw,
|
||||
"wallet_address": address,
|
||||
"score": 1.0,
|
||||
**relay_fields,
|
||||
},
|
||||
)
|
||||
node_id: str = reg_resp["node_id"]
|
||||
node_id = str(reg_resp["node_id"])
|
||||
setattr(node, "tracker_node_id", node_id)
|
||||
except Exception:
|
||||
node.stop()
|
||||
raise
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import http.server
|
||||
import json
|
||||
import sys
|
||||
@@ -11,6 +12,7 @@ import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
import uuid
|
||||
from typing import Any
|
||||
|
||||
from .model_backend import (
|
||||
InsufficientVRAMError,
|
||||
@@ -28,6 +30,40 @@ from .server import (
|
||||
)
|
||||
|
||||
|
||||
def _relay_hop(
|
||||
relay_addr: str,
|
||||
path: str,
|
||||
body: bytes,
|
||||
headers: dict[str, str],
|
||||
timeout: float = 120.0,
|
||||
) -> tuple[int, dict[str, str], bytes]:
|
||||
"""Send a single HTTP-shaped request through a relay RPC WebSocket.
|
||||
|
||||
relay_addr is the wss://relay.../rpc/{peer_id} URL.
|
||||
Returns (status, response_headers_lower, response_body).
|
||||
Raises on connection failure so callers can fall back to direct.
|
||||
"""
|
||||
import websockets.sync.client as wsc # type: ignore[import]
|
||||
|
||||
request_id = f"{time.time_ns():x}"
|
||||
payload = json.dumps({
|
||||
"request_id": request_id,
|
||||
"method": "POST",
|
||||
"path": path,
|
||||
"headers": headers,
|
||||
"body_base64": base64.b64encode(body).decode(),
|
||||
})
|
||||
with wsc.connect(relay_addr, open_timeout=timeout) as ws:
|
||||
ws.send(payload)
|
||||
raw = ws.recv(timeout=timeout)
|
||||
resp = json.loads(raw)
|
||||
status = int(resp.get("status", 503))
|
||||
resp_headers = {k.lower(): v for k, v in (resp.get("headers") or {}).items()}
|
||||
body_b64 = resp.get("body_base64")
|
||||
resp_body = base64.b64decode(body_b64) if body_b64 else (resp.get("body") or "").encode()
|
||||
return status, resp_headers, resp_body
|
||||
|
||||
|
||||
class _TorchHTTPServer(http.server.HTTPServer):
|
||||
def __init__(
|
||||
self,
|
||||
@@ -36,6 +72,8 @@ class _TorchHTTPServer(http.server.HTTPServer):
|
||||
backend: TorchModelShard,
|
||||
tracker_mode: bool = False,
|
||||
tracker_url: str | None = None,
|
||||
route_timeout: float = 30.0,
|
||||
debug: bool = False,
|
||||
):
|
||||
super().__init__(addr, handler)
|
||||
self.backend = backend
|
||||
@@ -43,6 +81,12 @@ class _TorchHTTPServer(http.server.HTTPServer):
|
||||
self.forward_chunk_count = 0
|
||||
self.tracker_mode = tracker_mode
|
||||
self.tracker_url = tracker_url
|
||||
self.route_timeout = route_timeout
|
||||
self.debug = debug
|
||||
self.total_requests: int = 0
|
||||
self.failed_requests: int = 0
|
||||
self.queue_depth: int = 0
|
||||
self._stats_lock = threading.Lock()
|
||||
|
||||
|
||||
class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
@@ -138,12 +182,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
if int(self.headers.get("X-Meshnet-Hop-Index", "0")) > 0:
|
||||
server.received_activations = True
|
||||
|
||||
start_layer_header = self.headers.get("X-Meshnet-Start-Layer")
|
||||
start_layer = int(start_layer_header) if start_layer_header else None
|
||||
|
||||
try:
|
||||
result = server.backend.forward_bytes(
|
||||
raw_body,
|
||||
shape,
|
||||
self.headers.get("X-Meshnet-Attn-Mask"),
|
||||
self.headers.get("X-Meshnet-Position-Ids"),
|
||||
start_layer=start_layer,
|
||||
)
|
||||
except Exception as exc:
|
||||
self._send_json(500, {"error": str(exc)})
|
||||
@@ -205,48 +253,181 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(payload)))
|
||||
self.end_headers()
|
||||
self.wfile.write(payload)
|
||||
try:
|
||||
self.wfile.write(payload)
|
||||
except BrokenPipeError:
|
||||
pass # client disconnected before we could respond — not an error
|
||||
|
||||
def _handle_chat_completions(self) -> None:
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
with server._stats_lock:
|
||||
server.total_requests += 1
|
||||
server.queue_depth += 1
|
||||
try:
|
||||
self._do_chat_completions(server)
|
||||
finally:
|
||||
with server._stats_lock:
|
||||
server.queue_depth -= 1
|
||||
|
||||
def _record_failed_request(self) -> None:
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
with server._stats_lock:
|
||||
server.failed_requests += 1
|
||||
|
||||
def _do_chat_completions(self, server: "_TorchHTTPServer") -> None:
|
||||
body = self._read_json_body()
|
||||
if body is None:
|
||||
return
|
||||
messages = body.get("messages", [])
|
||||
if not isinstance(messages, list):
|
||||
messages = []
|
||||
stream = bool(body.get("stream", False))
|
||||
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:
|
||||
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)
|
||||
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)
|
||||
|
||||
def _get_remaining_route(self, model: str) -> list[str]:
|
||||
# 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._record_failed_request()
|
||||
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)
|
||||
print(
|
||||
f" [node] chat route model={model_name!r} max_tokens={max_tokens} "
|
||||
f"downstream={remaining_route}",
|
||||
flush=True,
|
||||
)
|
||||
if not remaining_route:
|
||||
self._send_openai_response(
|
||||
"error: no downstream route — check tracker connectivity",
|
||||
model_name, False, messages,
|
||||
)
|
||||
return
|
||||
|
||||
backend = server.backend
|
||||
# Format with chat template so the model knows it's in assistant mode.
|
||||
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)
|
||||
|
||||
def _get_remaining_route(self, model: str) -> list[dict]:
|
||||
"""Return downstream hops as dicts with endpoint, start_layer, and optional relay_addr.
|
||||
|
||||
Fast path reads X-Meshnet-Route header injected by the tracker.
|
||||
Slow path queries the tracker's /v1/route endpoint as a fallback.
|
||||
start_layer tells each downstream node which layer to begin from,
|
||||
enabling correct execution when shard ranges overlap.
|
||||
"""
|
||||
# Fast path: tracker pre-resolved the downstream route and injected it as a header.
|
||||
injected = self.headers.get("X-Meshnet-Route")
|
||||
if injected:
|
||||
try:
|
||||
route = json.loads(injected)
|
||||
if isinstance(route, list):
|
||||
hops: list[dict] = []
|
||||
for item in route:
|
||||
if isinstance(item, dict):
|
||||
hop = {
|
||||
"endpoint": str(item["endpoint"]),
|
||||
"start_layer": int(item.get("start_layer", 0)),
|
||||
}
|
||||
if item.get("relay_addr"):
|
||||
hop["relay_addr"] = str(item["relay_addr"])
|
||||
hops.append(hop)
|
||||
elif isinstance(item, str):
|
||||
hops.append({"endpoint": item, "start_layer": 0})
|
||||
print(f" [node] using injected downstream route: {[h['endpoint'] for h in hops]}", flush=True)
|
||||
return hops
|
||||
except (json.JSONDecodeError, TypeError, KeyError):
|
||||
pass
|
||||
|
||||
# Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject).
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
if server.tracker_url is None:
|
||||
return []
|
||||
route_model = getattr(server.backend, "model_id", None) or model
|
||||
try:
|
||||
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(model)}"
|
||||
with urllib.request.urlopen(url, timeout=5.0) as r:
|
||||
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
|
||||
with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
|
||||
route_resp = json.loads(r.read())
|
||||
route = route_resp.get("route", [])
|
||||
# Skip the first node in the route (self) since we're already the head
|
||||
return list(route[1:])
|
||||
except Exception:
|
||||
own_port = server.server_address[1]
|
||||
nodes_info = route_resp.get("nodes", [])
|
||||
hops = []
|
||||
covered_up_to: int | None = None
|
||||
for node_info in nodes_info:
|
||||
ep = node_info.get("endpoint", "")
|
||||
if ep.rstrip("/").endswith(f":{own_port}"):
|
||||
covered_up_to = node_info.get("shard_end")
|
||||
continue
|
||||
if covered_up_to is None:
|
||||
covered_up_to = (node_info.get("shard_start") or 1) - 1
|
||||
hop = {"endpoint": ep, "start_layer": covered_up_to + 1}
|
||||
if node_info.get("relay_addr"):
|
||||
hop["relay_addr"] = str(node_info["relay_addr"])
|
||||
hops.append(hop)
|
||||
covered_up_to = node_info.get("shard_end", covered_up_to)
|
||||
print(f" [node] tracker downstream route: {[h['endpoint'] for h in hops]}", flush=True)
|
||||
return hops
|
||||
except Exception as exc:
|
||||
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
|
||||
return []
|
||||
|
||||
def _run_downstream_pipeline(self, payload: object, route: list[str]) -> str:
|
||||
def _run_downstream_pipeline(self, payload: object, route: list[dict]) -> str:
|
||||
server: _TorchHTTPServer = self.server # type: ignore[assignment]
|
||||
if not route:
|
||||
# Single-node mode: decode tail locally if we're the tail
|
||||
# Partial shard at tail: decode the activation from the previous node.
|
||||
# Full single-node (head+tail) is handled before entering this method.
|
||||
if server.backend.is_tail:
|
||||
try:
|
||||
tensor = server.backend.torch.frombuffer(
|
||||
@@ -256,7 +437,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
return server.backend.decode_tail(tensor)
|
||||
except Exception as exc:
|
||||
return f"decode error: {exc}"
|
||||
return ""
|
||||
return "no downstream route available for non-tail shard"
|
||||
|
||||
session = str(uuid.uuid4())
|
||||
shape = payload.shape # type: ignore[union-attr]
|
||||
@@ -267,7 +448,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
current_attn = attn_mask
|
||||
current_pos = pos_ids
|
||||
|
||||
for hop_index, node_url in enumerate(route):
|
||||
for hop_index, hop in enumerate(route):
|
||||
node_url = hop["endpoint"]
|
||||
start_layer = hop.get("start_layer", 0)
|
||||
relay_addr = hop.get("relay_addr")
|
||||
if server.debug:
|
||||
print(
|
||||
f" [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}"
|
||||
+ (f" relay={relay_addr}" if relay_addr else ""),
|
||||
flush=True,
|
||||
)
|
||||
headers: dict[str, str] = {
|
||||
"Content-Type": "application/octet-stream",
|
||||
"X-Meshnet-Wire": _WIRE_VERSION,
|
||||
@@ -277,28 +467,52 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
"X-Meshnet-Chunk-Index": "0",
|
||||
"X-Meshnet-Chunk-Total": "1",
|
||||
"X-Meshnet-Hop-Index": str(hop_index),
|
||||
"X-Meshnet-Start-Layer": str(start_layer),
|
||||
}
|
||||
if current_attn:
|
||||
headers["X-Meshnet-Attn-Mask"] = current_attn
|
||||
if current_pos:
|
||||
headers["X-Meshnet-Position-Ids"] = current_pos
|
||||
req = urllib.request.Request(
|
||||
f"{node_url}/forward",
|
||||
data=current_body,
|
||||
headers=headers,
|
||||
method="POST",
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=10.0) as r:
|
||||
resp_body = r.read()
|
||||
resp_headers = {k.lower(): v for k, v in r.headers.items()}
|
||||
except Exception as exc:
|
||||
return f"pipeline error at {node_url}: {exc}"
|
||||
if relay_addr:
|
||||
try:
|
||||
status, resp_headers, resp_body = _relay_hop(
|
||||
relay_addr, "/forward", current_body, headers, timeout=120.0,
|
||||
)
|
||||
if status >= 400:
|
||||
print(
|
||||
f" [node] relay hop {hop_index} returned {status} from {relay_addr}",
|
||||
flush=True,
|
||||
)
|
||||
return f"pipeline error at {node_url} via relay: status {status}"
|
||||
except Exception as exc:
|
||||
print(
|
||||
f" [node] relay hop {hop_index} failed at {relay_addr}: {exc}; "
|
||||
f"falling back to direct {node_url}",
|
||||
flush=True,
|
||||
)
|
||||
relay_addr = None # fall through to direct
|
||||
if not relay_addr:
|
||||
req = urllib.request.Request(
|
||||
f"{node_url}/forward",
|
||||
data=current_body,
|
||||
headers=headers,
|
||||
method="POST",
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=120.0) as r:
|
||||
resp_body = r.read()
|
||||
resp_headers = {k.lower(): v for k, v in r.headers.items()}
|
||||
except Exception as exc:
|
||||
print(f" [node] pipeline hop {hop_index} failed at {node_url}: {exc}", flush=True)
|
||||
return f"pipeline error at {node_url}: {exc}"
|
||||
content_type = resp_headers.get("content-type", "")
|
||||
if "application/json" in content_type:
|
||||
try:
|
||||
data = json.loads(resp_body)
|
||||
return str(data.get("text", ""))
|
||||
text = str(data.get("text", ""))
|
||||
if server.debug:
|
||||
print(f" [node] pipeline hop {hop_index} returned text={text!r}", flush=True)
|
||||
return text
|
||||
except json.JSONDecodeError:
|
||||
return resp_body.decode("utf-8", errors="replace")
|
||||
# Binary activation — update and forward to next node
|
||||
@@ -309,10 +523,57 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
current_pos = resp_headers.get("x-meshnet-position-ids")
|
||||
return ""
|
||||
|
||||
def _send_openai_response(self, text: str, model: str, stream: bool) -> None:
|
||||
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:
|
||||
try:
|
||||
self.wfile.write(f"data: {data}\n\n".encode())
|
||||
self.wfile.flush()
|
||||
except BrokenPipeError:
|
||||
pass
|
||||
|
||||
_emit(json.dumps({
|
||||
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
||||
"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"}],
|
||||
}))
|
||||
try:
|
||||
self.wfile.write(b"data: [DONE]\n\n")
|
||||
self.wfile.flush()
|
||||
except BrokenPipeError:
|
||||
pass
|
||||
|
||||
def _send_openai_response(
|
||||
self,
|
||||
text: str,
|
||||
model: str,
|
||||
stream: bool,
|
||||
messages: list[dict] | None = None,
|
||||
) -> 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",
|
||||
@@ -323,7 +584,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
"message": {"role": "assistant", "content": text},
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
|
||||
"usage": usage,
|
||||
})
|
||||
return
|
||||
self.send_response(200)
|
||||
@@ -332,8 +593,11 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
self.end_headers()
|
||||
|
||||
def _emit(data: str) -> None:
|
||||
self.wfile.write(f"data: {data}\n\n".encode())
|
||||
self.wfile.flush()
|
||||
try:
|
||||
self.wfile.write(f"data: {data}\n\n".encode())
|
||||
self.wfile.flush()
|
||||
except BrokenPipeError:
|
||||
pass
|
||||
|
||||
_emit(json.dumps({
|
||||
"id": chunk_id, "object": "chat.completion.chunk", "created": created,
|
||||
@@ -350,8 +614,57 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
}))
|
||||
self.wfile.write(b"data: [DONE]\n\n")
|
||||
self.wfile.flush()
|
||||
try:
|
||||
self.wfile.write(b"data: [DONE]\n\n")
|
||||
self.wfile.flush()
|
||||
except BrokenPipeError:
|
||||
pass
|
||||
|
||||
|
||||
def _usage_for_response(backend: object, messages: list[dict], completion_text: str) -> dict[str, int]:
|
||||
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:
|
||||
@@ -368,6 +681,8 @@ class TorchNodeServer:
|
||||
backend: TorchModelShard | None = None,
|
||||
tracker_mode: bool | None = None,
|
||||
tracker_url: str | None = None,
|
||||
route_timeout: float = 30.0,
|
||||
debug: bool = False,
|
||||
) -> None:
|
||||
self._host = host
|
||||
self._requested_port = port
|
||||
@@ -380,10 +695,16 @@ class TorchNodeServer:
|
||||
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
|
||||
self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
|
||||
self._tracker_url = tracker_url
|
||||
self._route_timeout = route_timeout
|
||||
self._debug = debug
|
||||
self._server: _TorchHTTPServer | None = None
|
||||
self._thread: threading.Thread | None = None
|
||||
self.port: int | None = None
|
||||
|
||||
@property
|
||||
def route_timeout(self) -> float:
|
||||
return self._route_timeout
|
||||
|
||||
@property
|
||||
def backend(self) -> TorchModelShard:
|
||||
return self._backend
|
||||
@@ -396,6 +717,52 @@ class TorchNodeServer:
|
||||
def forward_chunk_count(self) -> int:
|
||||
return self._server.forward_chunk_count if self._server is not None else 0
|
||||
|
||||
@property
|
||||
def total_requests(self) -> int:
|
||||
return self._server.total_requests if self._server is not None else 0
|
||||
|
||||
@property
|
||||
def failed_requests(self) -> int:
|
||||
return self._server.failed_requests if self._server is not None else 0
|
||||
|
||||
@property
|
||||
def queue_depth(self) -> int:
|
||||
return self._server.queue_depth if self._server is not None else 0
|
||||
|
||||
def apply_tracker_directives(self, directives: list[dict]) -> dict | None:
|
||||
"""Apply tracker LOAD_SHARD directives by hot-swapping the loaded backend."""
|
||||
load_directive = next(
|
||||
(directive for directive in reversed(directives) if directive.get("action") == "LOAD_SHARD"),
|
||||
None,
|
||||
)
|
||||
if load_directive is None:
|
||||
return None
|
||||
shard_start = int(load_directive["shard_start"])
|
||||
shard_end = int(load_directive["shard_end"])
|
||||
quantization = str(load_directive.get("quantization") or self._backend.quantization)
|
||||
model_id = str(load_directive.get("model") or self._backend.model_id)
|
||||
print(
|
||||
f" [node] loading reassigned shard: {model_id} layers {shard_start}-{shard_end}",
|
||||
flush=True,
|
||||
)
|
||||
new_backend = _load_backend(model_id, shard_start, shard_end, quantization)
|
||||
self._backend = new_backend
|
||||
self._tracker_mode = shard_start == 0
|
||||
if self._server is not None:
|
||||
self._server.backend = new_backend
|
||||
self._server.tracker_mode = self._tracker_mode
|
||||
print(
|
||||
f" [node] loaded reassigned shard: {model_id} layers {shard_start}-{shard_end}",
|
||||
flush=True,
|
||||
)
|
||||
return {
|
||||
"model": model_id,
|
||||
"shard_start": shard_start,
|
||||
"shard_end": shard_end,
|
||||
"quantization": quantization,
|
||||
"tracker_mode": self._tracker_mode,
|
||||
}
|
||||
|
||||
def start(self) -> int:
|
||||
if self._server is not None:
|
||||
raise RuntimeError("TorchNodeServer is already running")
|
||||
@@ -405,6 +772,8 @@ class TorchNodeServer:
|
||||
self._backend,
|
||||
self._tracker_mode,
|
||||
self._tracker_url,
|
||||
self._route_timeout,
|
||||
self._debug,
|
||||
)
|
||||
self.port = self._server.server_address[1]
|
||||
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
|
||||
|
||||
@@ -17,6 +17,7 @@ dependencies = [
|
||||
"safetensors>=0.4",
|
||||
"torch>=2.1",
|
||||
"transformers>=4.39",
|
||||
"websockets>=13",
|
||||
"zstandard>=0.22",
|
||||
]
|
||||
|
||||
|
||||
@@ -16,6 +16,7 @@ import asyncio
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
from .peer_registry import PeerRegistry
|
||||
@@ -56,6 +57,7 @@ class RelayServer:
|
||||
self._ready = threading.Event()
|
||||
self._running = False
|
||||
self._stop_event: asyncio.Event | None = None
|
||||
self._pending_rpc: dict[str, asyncio.Future] = {}
|
||||
|
||||
@property
|
||||
def registry(self) -> PeerRegistry:
|
||||
@@ -121,6 +123,9 @@ class RelayServer:
|
||||
if path.startswith("/relay/"):
|
||||
peer_id = path[len("/relay/"):]
|
||||
await self._handle_circuit_relay(ws, peer_id)
|
||||
elif path.startswith("/rpc/"):
|
||||
peer_id = path[len("/rpc/"):]
|
||||
await self._handle_rpc(ws, peer_id)
|
||||
elif path == "/health":
|
||||
await ws.send(json.dumps({"status": "ok", "peers": len(self._registry)}))
|
||||
await ws.close()
|
||||
@@ -164,6 +169,14 @@ class RelayServer:
|
||||
}))
|
||||
continue
|
||||
|
||||
if topic == "relay-http-response":
|
||||
payload = envelope.get("payload", {})
|
||||
request_id = payload.get("request_id")
|
||||
fut = self._pending_rpc.pop(request_id, None)
|
||||
if fut is not None and not fut.done():
|
||||
fut.set_result(payload)
|
||||
continue
|
||||
|
||||
# Fan out to all other registered peers
|
||||
if peer_id:
|
||||
self._registry.touch(peer_id)
|
||||
@@ -205,6 +218,50 @@ class RelayServer:
|
||||
return_exceptions=True,
|
||||
)
|
||||
|
||||
async def _handle_rpc(self, ws_requester, target_peer_id: str) -> None:
|
||||
"""Send one HTTP-shaped request to a connected peer and relay its response."""
|
||||
target = self._registry.get(target_peer_id)
|
||||
if target is None:
|
||||
await ws_requester.send(json.dumps({
|
||||
"status": 503,
|
||||
"headers": {"Content-Type": "application/json"},
|
||||
"body": json.dumps({"error": f"peer {target_peer_id!r} not connected to relay"}),
|
||||
}))
|
||||
await ws_requester.close()
|
||||
return
|
||||
|
||||
try:
|
||||
raw = await asyncio.wait_for(ws_requester.recv(), timeout=30.0)
|
||||
payload = json.loads(raw)
|
||||
except Exception:
|
||||
await ws_requester.close(1003, "invalid relay rpc request")
|
||||
return
|
||||
|
||||
request_id = str(payload.get("request_id") or uuid.uuid4())
|
||||
payload["request_id"] = request_id
|
||||
payload["target_peer"] = target_peer_id
|
||||
fut = self._loop.create_future() if self._loop is not None else asyncio.get_running_loop().create_future()
|
||||
self._pending_rpc[request_id] = fut
|
||||
try:
|
||||
await target.ws.send(json.dumps({
|
||||
"topic": "relay-http-request",
|
||||
"version": 1,
|
||||
"from_peer": "relay",
|
||||
"payload": payload,
|
||||
}))
|
||||
response = await asyncio.wait_for(fut, timeout=310.0)
|
||||
await ws_requester.send(json.dumps(response))
|
||||
except asyncio.TimeoutError:
|
||||
await ws_requester.send(json.dumps({
|
||||
"request_id": request_id,
|
||||
"status": 504,
|
||||
"headers": {"Content-Type": "application/json"},
|
||||
"body": json.dumps({"error": "relay rpc timed out"}),
|
||||
}))
|
||||
finally:
|
||||
self._pending_rpc.pop(request_id, None)
|
||||
await ws_requester.close()
|
||||
|
||||
|
||||
async def _broadcast(raw: str | bytes, peers: list) -> None:
|
||||
"""Send raw message to all peers; ignore individual send failures."""
|
||||
|
||||
@@ -4,33 +4,63 @@ import argparse
|
||||
import sys
|
||||
import time
|
||||
|
||||
from .server import TrackerServer
|
||||
from .server import TrackerServer, derive_relay_url_from_public_tracker_url
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="meshnet-tracker",
|
||||
description="Distributed Inference Network node registry and route selection",
|
||||
)
|
||||
subparsers = parser.add_subparsers(dest="command")
|
||||
|
||||
start_cmd = subparsers.add_parser("start", help="Start the tracker server")
|
||||
start_cmd.add_argument("--host", default="127.0.0.1", help="Host interface to listen on")
|
||||
start_cmd.add_argument("--port", type=int, default=8080, help="Port to listen on")
|
||||
start_cmd.add_argument(
|
||||
common = argparse.ArgumentParser(add_help=False)
|
||||
common.add_argument("--host", default="0.0.0.0", help="Host interface to listen on")
|
||||
common.add_argument("--port", type=int, default=8080, help="Port to listen on")
|
||||
common.add_argument(
|
||||
"--heartbeat-timeout",
|
||||
type=float,
|
||||
default=30.0,
|
||||
help="Seconds before a node is removed from the registry after missed heartbeat",
|
||||
)
|
||||
common.add_argument(
|
||||
"--cluster-peers",
|
||||
default="",
|
||||
help="Comma-separated URLs of peer tracker nodes (enables Raft cluster mode)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--self-url",
|
||||
default=None,
|
||||
help="This tracker's own URL as seen by peers (auto-derived from --host/--port if omitted)",
|
||||
)
|
||||
common.add_argument(
|
||||
"--stats-db",
|
||||
default=None,
|
||||
metavar="PATH",
|
||||
help="SQLite database path for persistent model usage statistics",
|
||||
)
|
||||
common.add_argument(
|
||||
"--relay-url",
|
||||
default=None,
|
||||
help="Public ws(s):// relay URL advertised to nodes, for example wss://ai.neuron.d-popov.com/ws",
|
||||
)
|
||||
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="meshnet-tracker",
|
||||
description="Distributed Inference Network node registry and route selection",
|
||||
parents=[common],
|
||||
)
|
||||
subparsers = parser.add_subparsers(dest="command")
|
||||
|
||||
subparsers.add_parser("start", help="Start the tracker server", parents=[common])
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == "start":
|
||||
if args.command in {None, "start"}:
|
||||
cluster_peers = [u.strip() for u in args.cluster_peers.split(",") if u.strip()]
|
||||
relay_url = args.relay_url or derive_relay_url_from_public_tracker_url(args.self_url)
|
||||
server = TrackerServer(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
heartbeat_timeout=args.heartbeat_timeout,
|
||||
cluster_peers=cluster_peers or None,
|
||||
cluster_self_url=args.self_url,
|
||||
stats_db=getattr(args, "stats_db", None),
|
||||
relay_url=relay_url,
|
||||
)
|
||||
port = server.start()
|
||||
print(f"meshnet-tracker listening on http://{args.host}:{port}", flush=True)
|
||||
|
||||
88
packages/tracker/meshnet_tracker/gossip.py
Normal file
88
packages/tracker/meshnet_tracker/gossip.py
Normal file
@@ -0,0 +1,88 @@
|
||||
"""CRDT gossip for node liveness heartbeats.
|
||||
|
||||
Uses a last-write-wins (LWW) register per inference node: each tracker node
|
||||
keeps its own copy of {node_id → last_seen_monotonic_timestamp} and merges
|
||||
incoming gossip by taking the max per key. This is eventually consistent —
|
||||
a heartbeat received by one tracker propagates to all others within a few
|
||||
gossip intervals.
|
||||
|
||||
Monotonic timestamps are local-clock-relative; for cross-machine gossip the
|
||||
caller should use wall-clock seconds (time.time()). The tracker converts
|
||||
monotonic to wall-clock when recording and back when comparing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
import urllib.request
|
||||
|
||||
|
||||
class NodeGossip:
|
||||
"""LWW gossip table for inference-node heartbeat timestamps.
|
||||
|
||||
``record(node_id)`` is called when a node sends a heartbeat to *this*
|
||||
tracker. ``merge(remote)`` is called when gossip arrives from a peer
|
||||
tracker. The table is periodically pushed to one random peer.
|
||||
"""
|
||||
|
||||
PUSH_INTERVAL = 3.0 # seconds between gossip pushes to a random peer
|
||||
|
||||
def __init__(self, peers: list[str]) -> None:
|
||||
self.peers = list(peers)
|
||||
# Maps node_id → wall-clock seconds of last known heartbeat.
|
||||
self._table: dict[str, float] = {}
|
||||
self._lock = threading.Lock()
|
||||
self._running = False
|
||||
|
||||
def start(self) -> None:
|
||||
self._running = True
|
||||
threading.Thread(target=self._push_loop, daemon=True, name="gossip").start()
|
||||
|
||||
def stop(self) -> None:
|
||||
self._running = False
|
||||
|
||||
def record(self, node_id: str, wall_ts: float | None = None) -> None:
|
||||
"""Record a heartbeat for *node_id* at *wall_ts* (default: now)."""
|
||||
ts = wall_ts if wall_ts is not None else time.time()
|
||||
with self._lock:
|
||||
if ts > self._table.get(node_id, 0.0):
|
||||
self._table[node_id] = ts
|
||||
|
||||
def merge(self, remote: dict[str, float]) -> None:
|
||||
"""Merge a gossip snapshot from a peer tracker (LWW per key)."""
|
||||
with self._lock:
|
||||
for node_id, ts in remote.items():
|
||||
if ts > self._table.get(node_id, 0.0):
|
||||
self._table[node_id] = ts
|
||||
|
||||
def last_seen(self, node_id: str) -> float | None:
|
||||
"""Return wall-clock timestamp of last known heartbeat, or None."""
|
||||
with self._lock:
|
||||
return self._table.get(node_id)
|
||||
|
||||
def snapshot(self) -> dict[str, float]:
|
||||
with self._lock:
|
||||
return dict(self._table)
|
||||
|
||||
def _push_loop(self) -> None:
|
||||
while self._running:
|
||||
time.sleep(self.PUSH_INTERVAL)
|
||||
if not self.peers:
|
||||
continue
|
||||
peer = random.choice(self.peers)
|
||||
try:
|
||||
snap = self.snapshot()
|
||||
body = json.dumps(snap).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{peer}/v1/gossip",
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=2.0) as r:
|
||||
r.read()
|
||||
except Exception:
|
||||
pass
|
||||
372
packages/tracker/meshnet_tracker/raft.py
Normal file
372
packages/tracker/meshnet_tracker/raft.py
Normal file
@@ -0,0 +1,372 @@
|
||||
"""Minimal Raft consensus for tracker shard assignments.
|
||||
|
||||
Only shard-assignment commands (register/deregister) go through the log.
|
||||
Node liveness (heartbeats) is handled separately via CRDT gossip — these
|
||||
are high-frequency writes where eventual consistency is fine.
|
||||
|
||||
Election timeout: random 150–300 ms (tight, suits in-process tests).
|
||||
Leader heartbeat interval: 50 ms.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import random
|
||||
import threading
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable
|
||||
|
||||
|
||||
@dataclass
|
||||
class LogEntry:
|
||||
term: int
|
||||
command: str # "register" | "deregister"
|
||||
payload: dict
|
||||
|
||||
|
||||
class RaftNode:
|
||||
"""Single Raft participant.
|
||||
|
||||
``apply_fn(command, payload)`` is called (under no external lock) when an
|
||||
entry is committed. Implementors must apply the command atomically.
|
||||
"""
|
||||
|
||||
ELECTION_MIN = 0.15 # seconds
|
||||
ELECTION_MAX = 0.30
|
||||
HB_INTERVAL = 0.05 # leader heartbeat interval
|
||||
|
||||
# Role constants
|
||||
FOLLOWER = "follower"
|
||||
CANDIDATE = "candidate"
|
||||
LEADER = "leader"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
self_url: str,
|
||||
peers: list[str],
|
||||
apply_fn: Callable[[str, dict], None],
|
||||
) -> None:
|
||||
self.self_url = self_url
|
||||
self.peers = list(peers)
|
||||
self._apply_fn = apply_fn
|
||||
|
||||
self._lock = threading.Lock()
|
||||
self.role: str = self.FOLLOWER
|
||||
self.current_term: int = 0
|
||||
self.voted_for: str | None = None
|
||||
self.log: list[LogEntry] = []
|
||||
self.commit_index: int = -1
|
||||
self.last_applied: int = -1
|
||||
self.leader_url: str | None = None
|
||||
|
||||
# Leader bookkeeping per peer
|
||||
self._next_index: dict[str, int] = {}
|
||||
self._match_index: dict[str, int] = {}
|
||||
|
||||
self._election_deadline: float = self._fresh_deadline()
|
||||
self._running = False
|
||||
|
||||
# ------------------------------------------------------------------ start/stop
|
||||
|
||||
def start(self) -> None:
|
||||
self._running = True
|
||||
threading.Thread(target=self._tick_loop, daemon=True, name="raft-tick").start()
|
||||
|
||||
def stop(self) -> None:
|
||||
self._running = False
|
||||
|
||||
# ------------------------------------------------------------------ public API
|
||||
|
||||
@property
|
||||
def is_leader(self) -> bool:
|
||||
with self._lock:
|
||||
return self.role == self.LEADER
|
||||
|
||||
def leader(self) -> str | None:
|
||||
with self._lock:
|
||||
return self.leader_url
|
||||
|
||||
def status(self) -> dict:
|
||||
with self._lock:
|
||||
return {
|
||||
"role": self.role,
|
||||
"term": self.current_term,
|
||||
"leader": self.leader_url,
|
||||
"log_length": len(self.log),
|
||||
"commit_index": self.commit_index,
|
||||
}
|
||||
|
||||
def propose(self, command: str, payload: dict) -> bool:
|
||||
"""Leader: append and replicate an entry. Returns True when committed.
|
||||
|
||||
Blocks until majority replication or failure. Must be called only on
|
||||
the leader; returns False immediately if this node is not the leader.
|
||||
"""
|
||||
with self._lock:
|
||||
if self.role != self.LEADER:
|
||||
return False
|
||||
entry = LogEntry(self.current_term, command, payload)
|
||||
self.log.append(entry)
|
||||
entry_index = len(self.log) - 1
|
||||
term = self.current_term
|
||||
|
||||
self._replicate_to_peers(term)
|
||||
|
||||
with self._lock:
|
||||
return self.commit_index >= entry_index
|
||||
|
||||
# ------------------------------------------------------------------ RPC handlers
|
||||
|
||||
def handle_request_vote(self, req: dict) -> dict:
|
||||
with self._lock:
|
||||
term = int(req["term"])
|
||||
candidate = req["candidate_url"]
|
||||
last_li = int(req["last_log_index"])
|
||||
last_lt = int(req["last_log_term"])
|
||||
|
||||
if term > self.current_term:
|
||||
self._step_down(term)
|
||||
|
||||
if term < self.current_term:
|
||||
return {"term": self.current_term, "vote_granted": False}
|
||||
|
||||
my_li = len(self.log) - 1
|
||||
my_lt = self.log[-1].term if self.log else 0
|
||||
log_ok = (last_lt > my_lt) or (last_lt == my_lt and last_li >= my_li)
|
||||
|
||||
if (self.voted_for is None or self.voted_for == candidate) and log_ok:
|
||||
self.voted_for = candidate
|
||||
self._reset_deadline()
|
||||
return {"term": self.current_term, "vote_granted": True}
|
||||
return {"term": self.current_term, "vote_granted": False}
|
||||
|
||||
def handle_append_entries(self, req: dict) -> dict:
|
||||
with self._lock:
|
||||
term = int(req["term"])
|
||||
leader_url = req["leader_url"]
|
||||
prev_li = int(req["prev_log_index"])
|
||||
prev_lt = int(req["prev_log_term"])
|
||||
entries_raw = req.get("entries", [])
|
||||
ldr_commit = int(req["leader_commit"])
|
||||
|
||||
if term > self.current_term:
|
||||
self._step_down(term)
|
||||
|
||||
if term < self.current_term:
|
||||
return {"term": self.current_term, "success": False}
|
||||
|
||||
# Valid AppendEntries from current leader
|
||||
self._reset_deadline()
|
||||
self.leader_url = leader_url
|
||||
if self.role == self.CANDIDATE:
|
||||
self.role = self.FOLLOWER
|
||||
|
||||
# Consistency check on prev entry
|
||||
if prev_li >= 0:
|
||||
if prev_li >= len(self.log):
|
||||
return {"term": self.current_term, "success": False}
|
||||
if self.log[prev_li].term != prev_lt:
|
||||
self.log = self.log[:prev_li]
|
||||
return {"term": self.current_term, "success": False}
|
||||
|
||||
# Append new entries (detect and overwrite conflicts)
|
||||
for i, raw in enumerate(entries_raw):
|
||||
idx = prev_li + 1 + i
|
||||
entry = LogEntry(int(raw["term"]), raw["command"], raw["payload"])
|
||||
if idx < len(self.log):
|
||||
if self.log[idx].term != entry.term:
|
||||
self.log = self.log[:idx]
|
||||
self.log.append(entry)
|
||||
else:
|
||||
self.log.append(entry)
|
||||
|
||||
if ldr_commit > self.commit_index:
|
||||
self.commit_index = min(ldr_commit, len(self.log) - 1)
|
||||
self._apply_up_to_commit()
|
||||
|
||||
return {"term": self.current_term, "success": True}
|
||||
|
||||
# ------------------------------------------------------------------ internals
|
||||
|
||||
def _tick_loop(self) -> None:
|
||||
while self._running:
|
||||
time.sleep(0.01)
|
||||
with self._lock:
|
||||
role = self.role
|
||||
deadline = self._election_deadline
|
||||
|
||||
if role == self.LEADER:
|
||||
self._send_heartbeats()
|
||||
time.sleep(self.HB_INTERVAL)
|
||||
elif time.monotonic() > deadline:
|
||||
self._start_election()
|
||||
|
||||
def _send_heartbeats(self) -> None:
|
||||
with self._lock:
|
||||
term = self.current_term
|
||||
ldr_commit = self.commit_index
|
||||
log_snapshot = list(self.log)
|
||||
|
||||
for peer in self.peers:
|
||||
try:
|
||||
self._send_append_entries(peer, term, ldr_commit, log_snapshot)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _replicate_to_peers(self, term: int) -> None:
|
||||
"""Send AppendEntries to all peers and update commit_index on majority ack."""
|
||||
with self._lock:
|
||||
ldr_commit = self.commit_index
|
||||
log_snapshot = list(self.log)
|
||||
|
||||
acks = 1 # leader counts as 1
|
||||
for peer in self.peers:
|
||||
try:
|
||||
ok = self._send_append_entries(peer, term, ldr_commit, log_snapshot)
|
||||
if ok:
|
||||
acks += 1
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Advance commit if majority replicated
|
||||
with self._lock:
|
||||
if self.role != self.LEADER or self.current_term != term:
|
||||
return
|
||||
majority_index = len(self.log) - 1
|
||||
while majority_index > self.commit_index:
|
||||
if self.log[majority_index].term == self.current_term:
|
||||
count = 1 + sum(
|
||||
1 for p in self.peers
|
||||
if self._match_index.get(p, -1) >= majority_index
|
||||
)
|
||||
if count > (len(self.peers) + 1) / 2:
|
||||
self.commit_index = majority_index
|
||||
self._apply_up_to_commit()
|
||||
break
|
||||
majority_index -= 1
|
||||
|
||||
def _send_append_entries(
|
||||
self, peer: str, term: int, ldr_commit: int, log_snapshot: list[LogEntry]
|
||||
) -> bool:
|
||||
with self._lock:
|
||||
next_idx = self._next_index.get(peer, len(log_snapshot))
|
||||
|
||||
prev_li = next_idx - 1
|
||||
prev_lt = log_snapshot[prev_li].term if 0 <= prev_li < len(log_snapshot) else 0
|
||||
entries = [
|
||||
{"term": e.term, "command": e.command, "payload": e.payload}
|
||||
for e in log_snapshot[next_idx:]
|
||||
]
|
||||
|
||||
body = json.dumps({
|
||||
"term": term,
|
||||
"leader_url": self.self_url,
|
||||
"prev_log_index": prev_li,
|
||||
"prev_log_term": prev_lt,
|
||||
"entries": entries,
|
||||
"leader_commit": ldr_commit,
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{peer}/v1/raft/append",
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=0.15) as r:
|
||||
resp = json.loads(r.read())
|
||||
|
||||
with self._lock:
|
||||
if resp.get("success"):
|
||||
new_match = len(log_snapshot) - 1
|
||||
self._match_index[peer] = max(self._match_index.get(peer, -1), new_match)
|
||||
self._next_index[peer] = new_match + 1
|
||||
return True
|
||||
else:
|
||||
if resp.get("term", 0) > self.current_term:
|
||||
self._step_down(resp["term"])
|
||||
else:
|
||||
self._next_index[peer] = max(0, self._next_index.get(peer, 1) - 1)
|
||||
return False
|
||||
|
||||
def _start_election(self) -> None:
|
||||
with self._lock:
|
||||
self.current_term += 1
|
||||
self.role = self.CANDIDATE
|
||||
self.voted_for = self.self_url
|
||||
term = self.current_term
|
||||
my_li = len(self.log) - 1
|
||||
my_lt = self.log[-1].term if self.log else 0
|
||||
self._reset_deadline()
|
||||
|
||||
votes = 1
|
||||
for peer in self.peers:
|
||||
try:
|
||||
body = json.dumps({
|
||||
"term": term,
|
||||
"candidate_url": self.self_url,
|
||||
"last_log_index": my_li,
|
||||
"last_log_term": my_lt,
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{peer}/v1/raft/vote",
|
||||
data=body,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=0.1) as r:
|
||||
resp = json.loads(r.read())
|
||||
if resp.get("vote_granted"):
|
||||
votes += 1
|
||||
elif resp.get("term", 0) > term:
|
||||
with self._lock:
|
||||
self._step_down(resp["term"])
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
with self._lock:
|
||||
if self.role == self.CANDIDATE and self.current_term == term:
|
||||
total = len(self.peers) + 1
|
||||
if votes > total / 2:
|
||||
self._become_leader()
|
||||
else:
|
||||
self.role = self.FOLLOWER
|
||||
self._reset_deadline()
|
||||
|
||||
def _become_leader(self) -> None:
|
||||
"""Must be called with _lock held."""
|
||||
self.role = self.LEADER
|
||||
self.leader_url = self.self_url
|
||||
for peer in self.peers:
|
||||
self._next_index[peer] = len(self.log)
|
||||
self._match_index[peer] = -1
|
||||
print(f"[raft] {self.self_url} became leader (term {self.current_term})", flush=True)
|
||||
|
||||
def _step_down(self, new_term: int) -> None:
|
||||
"""Must be called with _lock held."""
|
||||
self.current_term = new_term
|
||||
self.role = self.FOLLOWER
|
||||
self.voted_for = None
|
||||
self._reset_deadline()
|
||||
|
||||
def _apply_up_to_commit(self) -> None:
|
||||
"""Must be called with _lock held."""
|
||||
while self.last_applied < self.commit_index:
|
||||
self.last_applied += 1
|
||||
entry = self.log[self.last_applied]
|
||||
try:
|
||||
self._apply_fn(entry.command, entry.payload)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _reset_deadline(self) -> None:
|
||||
self._election_deadline = self._fresh_deadline()
|
||||
|
||||
@staticmethod
|
||||
def _fresh_deadline() -> float:
|
||||
return time.monotonic() + random.uniform(
|
||||
RaftNode.ELECTION_MIN, RaftNode.ELECTION_MAX
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -8,6 +8,10 @@ version = "0.1.0"
|
||||
description = "Distributed Inference Network node registry and route selection"
|
||||
requires-python = ">=3.10"
|
||||
|
||||
dependencies = [
|
||||
"websockets>=13",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
meshnet-tracker = "meshnet_tracker.cli:main"
|
||||
|
||||
|
||||
@@ -152,7 +152,7 @@ def _story_meta(
|
||||
parts.append(f"worktree: {wt}")
|
||||
|
||||
# summary -------------------------------------------------------------
|
||||
notes = story.get("completionNotes", "").strip()
|
||||
notes = (story.get("completionNotes") or "").strip()
|
||||
if not notes and wt:
|
||||
notes = _story_last_commit(sid)
|
||||
if notes:
|
||||
|
||||
163
scripts/test_lan_inference.py
Normal file
163
scripts/test_lan_inference.py
Normal file
@@ -0,0 +1,163 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
End-to-end LAN inference test for meshnet distributed inference.
|
||||
|
||||
Sends 3 chat-completion requests to a meshnet node, validates OpenAI-format
|
||||
responses, and prints token counts + latency per request.
|
||||
|
||||
Usage:
|
||||
python scripts/test_lan_inference.py \\
|
||||
--tracker http://192.168.1.10:8080 \\
|
||||
--gateway http://192.168.1.10:8001
|
||||
|
||||
Exit 0 on success, 1 on any failure.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.parse
|
||||
import urllib.request
|
||||
|
||||
|
||||
PROMPTS = [
|
||||
{"role": "user", "content": "What is 7 × 8? Answer in one word."},
|
||||
{"role": "user", "content": "Name the capital of France in one word."},
|
||||
{"role": "user", "content": "Complete the sequence: 1, 1, 2, 3, 5, ___. Answer in one word."},
|
||||
]
|
||||
|
||||
MODEL = "microsoft/Phi-3-medium-128k-instruct"
|
||||
|
||||
|
||||
def _get(url: str, timeout: float = 10.0) -> dict:
|
||||
with urllib.request.urlopen(url, timeout=timeout) as r:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def _post(url: str, payload: dict, timeout: float = 60.0) -> dict:
|
||||
data = json.dumps(payload).encode()
|
||||
req = urllib.request.Request(
|
||||
url, data=data, headers={"Content-Type": "application/json"}, method="POST"
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=timeout) as r:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def discover_gateway(tracker_url: str) -> str:
|
||||
"""Return the first tracker-mode node endpoint for MODEL."""
|
||||
nodes = _get(f"{tracker_url}/v1/nodes", timeout=5.0)
|
||||
if isinstance(nodes, dict):
|
||||
nodes = list(nodes.values())
|
||||
tracker_nodes = [
|
||||
n for n in nodes
|
||||
if n.get("tracker_mode") and (
|
||||
n.get("hf_repo") == MODEL or n.get("model") == MODEL.split("/")[-1]
|
||||
)
|
||||
]
|
||||
if not tracker_nodes:
|
||||
raise RuntimeError(
|
||||
f"No tracker-mode nodes found for {MODEL!r}. "
|
||||
"Is the first-shard node running and registered?"
|
||||
)
|
||||
endpoint: str = tracker_nodes[0]["endpoint"]
|
||||
return endpoint.rstrip("/")
|
||||
|
||||
|
||||
def check_route(tracker_url: str, gateway_url: str) -> list[str]:
|
||||
"""Return the full inference route for MODEL."""
|
||||
url = f"{tracker_url}/v1/route?model={urllib.parse.quote(MODEL)}"
|
||||
try:
|
||||
resp = _get(url, timeout=5.0)
|
||||
return resp.get("route", [])
|
||||
except Exception as exc:
|
||||
print(f" Warning: could not fetch route: {exc}", file=sys.stderr)
|
||||
return [gateway_url]
|
||||
|
||||
|
||||
def run_inference(gateway_url: str, messages: list[dict]) -> tuple[str, int, float]:
|
||||
"""Send one chat-completion request. Returns (content, tokens, elapsed_s)."""
|
||||
t0 = time.monotonic()
|
||||
resp = _post(
|
||||
f"{gateway_url}/v1/chat/completions",
|
||||
{"model": MODEL, "messages": messages, "stream": False},
|
||||
timeout=120.0,
|
||||
)
|
||||
elapsed = time.monotonic() - t0
|
||||
|
||||
choices = resp.get("choices")
|
||||
if not choices:
|
||||
raise ValueError(f"No choices in response: {resp}")
|
||||
content: str = choices[0].get("message", {}).get("content", "")
|
||||
if not isinstance(content, str):
|
||||
raise TypeError(f"Expected string content, got {type(content)}: {content}")
|
||||
|
||||
usage = resp.get("usage", {})
|
||||
tokens: int = usage.get("completion_tokens", len(content.split()))
|
||||
|
||||
return content, tokens, elapsed
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
p = argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument("--tracker", required=True, help="Tracker URL, e.g. http://192.168.1.10:8080")
|
||||
p.add_argument(
|
||||
"--gateway",
|
||||
default=None,
|
||||
help="Inference entry point URL. Auto-discovered from tracker if omitted.",
|
||||
)
|
||||
args = p.parse_args(argv)
|
||||
|
||||
tracker_url = args.tracker.rstrip("/")
|
||||
|
||||
print(f"Tracker: {tracker_url}")
|
||||
|
||||
# Resolve gateway
|
||||
gateway_url = args.gateway.rstrip("/") if args.gateway else None
|
||||
if gateway_url is None:
|
||||
try:
|
||||
gateway_url = discover_gateway(tracker_url)
|
||||
print(f"Gateway (auto-discovered): {gateway_url}")
|
||||
except Exception as exc:
|
||||
print(f"ERROR: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
else:
|
||||
print(f"Gateway: {gateway_url}")
|
||||
|
||||
# Show route
|
||||
route = check_route(tracker_url, gateway_url)
|
||||
print(f"Route: {route}")
|
||||
if len(route) < 2:
|
||||
print(" Warning: only one node in route — is the second-shard node registered?")
|
||||
print()
|
||||
|
||||
failures = 0
|
||||
for i, msg in enumerate(PROMPTS, start=1):
|
||||
print(f"[{i}] Q: {msg['content']}")
|
||||
try:
|
||||
content, tokens, elapsed = run_inference(gateway_url, [msg])
|
||||
tps = tokens / elapsed if elapsed > 0 else 0.0
|
||||
print(f" A: {content.strip()}")
|
||||
print(f" {tokens} tokens {elapsed:.2f}s {tps:.1f} t/s")
|
||||
except urllib.error.HTTPError as exc:
|
||||
body = exc.read().decode(errors="replace")
|
||||
print(f" ERROR {exc.code}: {body}", file=sys.stderr)
|
||||
failures += 1
|
||||
except Exception as exc:
|
||||
print(f" ERROR: {exc}", file=sys.stderr)
|
||||
failures += 1
|
||||
print()
|
||||
|
||||
if failures == 0:
|
||||
print(f"All {len(PROMPTS)} requests completed successfully.")
|
||||
print("Exit code: 0")
|
||||
return 0
|
||||
else:
|
||||
print(f"{failures}/{len(PROMPTS)} requests failed.", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -295,6 +295,106 @@ def test_relay_circuit_relay_proxies_message():
|
||||
assert received_via_relay, "NAT'd peer did not receive message via circuit relay"
|
||||
|
||||
|
||||
def test_relay_rpc_round_trips_http_request_to_peer():
|
||||
"""Relay /rpc/<peer> sends one HTTP-shaped request to a connected peer."""
|
||||
import websockets.sync.client as wsc # type: ignore[import]
|
||||
from meshnet_relay.server import RelayServer
|
||||
|
||||
relay = RelayServer(host="127.0.0.1", port=0)
|
||||
port = relay.start()
|
||||
ready = threading.Event()
|
||||
|
||||
def run_peer():
|
||||
with wsc.connect(f"ws://127.0.0.1:{port}/ws") as ws:
|
||||
ws.send(json.dumps({
|
||||
"topic": "peer-register",
|
||||
"version": 1,
|
||||
"from_peer": "rpc_peer",
|
||||
"msg_id": "rpc-reg-001",
|
||||
"payload": {"peer_id": "rpc_peer", "addr": ""},
|
||||
}))
|
||||
ws.recv()
|
||||
ready.set()
|
||||
envelope = json.loads(ws.recv(timeout=3))
|
||||
payload = envelope["payload"]
|
||||
ws.send(json.dumps({
|
||||
"topic": "relay-http-response",
|
||||
"version": 1,
|
||||
"from_peer": "rpc_peer",
|
||||
"payload": {
|
||||
"request_id": payload["request_id"],
|
||||
"status": 200,
|
||||
"headers": {"Content-Type": "application/json"},
|
||||
"body": json.dumps({"ok": True, "path": payload["path"]}),
|
||||
},
|
||||
}))
|
||||
|
||||
peer_thread = threading.Thread(target=run_peer, daemon=True)
|
||||
peer_thread.start()
|
||||
assert ready.wait(timeout=5)
|
||||
|
||||
with wsc.connect(f"ws://127.0.0.1:{port}/rpc/rpc_peer") as ws:
|
||||
ws.send(json.dumps({
|
||||
"request_id": "req-1",
|
||||
"method": "POST",
|
||||
"path": "/v1/chat/completions",
|
||||
"headers": {"Content-Type": "application/json"},
|
||||
"body": "{}",
|
||||
}))
|
||||
response = json.loads(ws.recv(timeout=5))
|
||||
|
||||
relay.stop()
|
||||
peer_thread.join(timeout=3)
|
||||
|
||||
assert response["status"] == 200
|
||||
assert json.loads(response["body"]) == {"ok": True, "path": "/v1/chat/completions"}
|
||||
|
||||
|
||||
def test_node_relay_bridge_reconnects_after_failed_connection(monkeypatch):
|
||||
"""Node-side relay bridge keeps retrying its outbound WebSocket connection."""
|
||||
import websockets.sync.client as wsc # type: ignore[import]
|
||||
from meshnet_node.relay_bridge import RelayHttpBridge
|
||||
|
||||
attempts = []
|
||||
|
||||
class FakeWebSocket:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def send(self, _message):
|
||||
pass
|
||||
|
||||
def recv(self, timeout=1):
|
||||
raise TimeoutError
|
||||
|
||||
def fake_connect(url, open_timeout=5):
|
||||
attempts.append((url, open_timeout))
|
||||
if len(attempts) == 1:
|
||||
raise OSError("temporary relay outage")
|
||||
return FakeWebSocket()
|
||||
|
||||
monkeypatch.setattr(wsc, "connect", fake_connect)
|
||||
|
||||
bridge = RelayHttpBridge(
|
||||
relay_url="ws://relay.example/ws",
|
||||
peer_id="peer-reconnect",
|
||||
local_base_url="http://127.0.0.1:8001",
|
||||
advertised_addr="http://127.0.0.1:8001",
|
||||
reconnect_interval=0.01,
|
||||
)
|
||||
try:
|
||||
info = bridge.start()
|
||||
assert info.relay_addr == "ws://relay.example/rpc/peer-reconnect"
|
||||
assert bridge.wait_connected(timeout=1.0)
|
||||
finally:
|
||||
bridge.stop()
|
||||
|
||||
assert len(attempts) >= 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Tracker gossip fields tests
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -334,6 +434,27 @@ def _start_tracker_and_register(extra_fields: dict) -> dict:
|
||||
return resp
|
||||
|
||||
|
||||
def test_tracker_derives_relay_url_from_public_self_url():
|
||||
from meshnet_tracker.server import TrackerServer, derive_relay_url_from_public_tracker_url
|
||||
|
||||
assert derive_relay_url_from_public_tracker_url("https://ai.neuron.d-popov.com") == (
|
||||
"wss://ai.neuron.d-popov.com/ws"
|
||||
)
|
||||
assert derive_relay_url_from_public_tracker_url("http://127.0.0.1:8081") is None
|
||||
|
||||
tracker = TrackerServer(cluster_self_url="https://ai.neuron.d-popov.com")
|
||||
port = tracker.start()
|
||||
try:
|
||||
import json as _json
|
||||
import urllib.request
|
||||
|
||||
with urllib.request.urlopen(f"http://127.0.0.1:{port}/v1/network/map", timeout=5) as resp:
|
||||
body = _json.loads(resp.read())
|
||||
assert body["relay_url"] == "wss://ai.neuron.d-popov.com/ws"
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_accepts_relay_addr_in_registration():
|
||||
resp = _start_tracker_and_register({
|
||||
"relay_addr": "ws://relay.meshnet.ai:8765/relay/abc123",
|
||||
@@ -349,6 +470,41 @@ def test_tracker_accepts_registration_without_gossip_fields():
|
||||
assert "node_id" in resp
|
||||
|
||||
|
||||
def test_tracker_network_map_exposes_relay_and_registered_peer():
|
||||
import json as _json
|
||||
import urllib.request
|
||||
|
||||
from meshnet_tracker.server import TrackerServer
|
||||
|
||||
tracker = TrackerServer(host="127.0.0.1", port=0, relay_url="wss://ai.neuron.d-popov.com/ws")
|
||||
port = tracker.start()
|
||||
url = f"http://127.0.0.1:{port}"
|
||||
try:
|
||||
req = urllib.request.Request(
|
||||
f"{url}/v1/nodes/register",
|
||||
data=_json.dumps({
|
||||
"endpoint": "http://192.0.2.10:7000",
|
||||
"model": "stub-model",
|
||||
"shard_start": 0,
|
||||
"shard_end": 31,
|
||||
"relay_addr": "wss://ai.neuron.d-popov.com/rpc/peer123",
|
||||
"peer_id": "peer123",
|
||||
}).encode(),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=5):
|
||||
pass
|
||||
with urllib.request.urlopen(f"{url}/v1/network/map", timeout=5) as resp:
|
||||
body = _json.loads(resp.read())
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
assert body["relay_url"] == "wss://ai.neuron.d-popov.com/ws"
|
||||
assert body["nodes"][0]["relay_addr"] == "wss://ai.neuron.d-popov.com/rpc/peer123"
|
||||
assert body["nodes"][0]["peer_id"] == "peer123"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# mDNS (no-op without zeroconf installed)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import socket
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
@@ -354,3 +355,115 @@ def test_legacy_start_subcommand_accepted(monkeypatch):
|
||||
|
||||
# Exited (either 0 or via KeyboardInterrupt caught in _cmd_start)
|
||||
# The important thing is no unhandled exception from arg parsing
|
||||
|
||||
|
||||
def test_legacy_start_treats_repo_model_as_model_id(monkeypatch):
|
||||
"""`meshnet-node start --model org/repo` enters the HF model startup path."""
|
||||
from meshnet_node.cli import main
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_run_startup(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
class _FakeNode:
|
||||
chat_completion_count = 0
|
||||
def stop(self): pass
|
||||
return _FakeNode()
|
||||
|
||||
monkeypatch.setattr(sys, "argv", [
|
||||
"meshnet-node", "start",
|
||||
"--tracker", "http://192.168.0.179:8081",
|
||||
"--model", "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"--port", "0",
|
||||
])
|
||||
|
||||
with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
|
||||
with patch("time.sleep", side_effect=KeyboardInterrupt):
|
||||
try:
|
||||
main()
|
||||
except SystemExit as exc:
|
||||
assert exc.code == 0
|
||||
|
||||
assert captured["model"] == "Qwen2.5-0.5B-Instruct"
|
||||
assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
|
||||
|
||||
def test_legacy_start_without_port_uses_next_available_port(monkeypatch):
|
||||
"""Omitting --port skips an occupied default port before startup loads the model."""
|
||||
from meshnet_node.cli import main
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_run_startup(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
class _FakeNode:
|
||||
chat_completion_count = 0
|
||||
def stop(self): pass
|
||||
return _FakeNode()
|
||||
|
||||
occupied = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
occupied.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
occupied.bind(("127.0.0.1", 7000))
|
||||
occupied.listen(1)
|
||||
try:
|
||||
monkeypatch.setattr(sys, "argv", [
|
||||
"meshnet-node", "start",
|
||||
"--tracker", "http://192.168.0.179:8081",
|
||||
"--model", "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"--host", "127.0.0.1",
|
||||
])
|
||||
|
||||
with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
|
||||
with patch("time.sleep", side_effect=KeyboardInterrupt):
|
||||
try:
|
||||
main()
|
||||
except SystemExit as exc:
|
||||
assert exc.code == 0
|
||||
finally:
|
||||
occupied.close()
|
||||
|
||||
assert captured["port"] == 7001
|
||||
|
||||
|
||||
def test_default_cli_passes_advertise_host(monkeypatch):
|
||||
"""The documented no-subcommand LAN flag reaches startup."""
|
||||
from meshnet_node.cli import main
|
||||
|
||||
saved = {
|
||||
"model_hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"model_name": "Qwen2.5-0.5B-Instruct",
|
||||
"quantization": "nf4",
|
||||
"tracker_url": "http://localhost:8080",
|
||||
"wallet_path": "",
|
||||
"download_dir": "",
|
||||
"port": 7000,
|
||||
"host": "0.0.0.0",
|
||||
}
|
||||
captured = {}
|
||||
|
||||
def fake_run_startup(*args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
class _FakeNode:
|
||||
chat_completion_count = 0
|
||||
def stop(self): pass
|
||||
return _FakeNode()
|
||||
|
||||
monkeypatch.setattr(sys, "argv", [
|
||||
"meshnet-node",
|
||||
"--tracker", "http://192.168.0.179:8081",
|
||||
"--advertise-host", "192.168.0.42",
|
||||
"--debug",
|
||||
"--no-tui",
|
||||
])
|
||||
|
||||
with patch("meshnet_node.config.load_config", return_value=saved):
|
||||
with patch("meshnet_node.startup.run_startup", side_effect=fake_run_startup):
|
||||
with patch("meshnet_node.dashboard.run_dashboard", side_effect=KeyboardInterrupt):
|
||||
try:
|
||||
main()
|
||||
except SystemExit as exc:
|
||||
assert exc.code == 0
|
||||
|
||||
assert captured["tracker_url"] == "http://192.168.0.179:8081"
|
||||
assert captured["advertise_host"] == "192.168.0.42"
|
||||
assert captured["debug"] is True
|
||||
|
||||
@@ -11,7 +11,11 @@ import pytest
|
||||
|
||||
from meshnet_node.downloader import download_shard, write_shard_archive
|
||||
from meshnet_node.hardware import detect_hardware
|
||||
from meshnet_node.startup import _probationary_status_line, run_startup
|
||||
from meshnet_node.startup import (
|
||||
_infer_relay_url_from_tracker,
|
||||
_probationary_status_line,
|
||||
run_startup,
|
||||
)
|
||||
from meshnet_node.wallet import _b58encode, load_or_create_wallet
|
||||
from meshnet_contracts import LocalSolanaContracts
|
||||
from meshnet_tracker.server import TrackerServer
|
||||
@@ -347,6 +351,385 @@ def test_tracker_assign_lists_peers_for_same_model_shard():
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_infer_relay_url_from_public_https_tracker():
|
||||
assert _infer_relay_url_from_tracker("https://ai.neuron.d-popov.com") == (
|
||||
"wss://ai.neuron.d-popov.com/ws"
|
||||
)
|
||||
assert _infer_relay_url_from_tracker("https://ai.neuron.d-popov.com/v1/network/map") == (
|
||||
"wss://ai.neuron.d-popov.com/ws"
|
||||
)
|
||||
assert _infer_relay_url_from_tracker("http://192.168.0.179:8081") is None
|
||||
assert _infer_relay_url_from_tracker("http://127.0.0.1:8081") is None
|
||||
|
||||
|
||||
def test_public_https_tracker_infers_relay_when_network_map_omits_relay_url(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
capsys,
|
||||
):
|
||||
"""Nodes bootstrap relay from the tracker origin when map relay_url is null."""
|
||||
import meshnet_node.startup as startup_mod
|
||||
|
||||
class FakeBackend:
|
||||
total_layers = 24
|
||||
|
||||
class FakeTorchNodeServer:
|
||||
def __init__(self, **kwargs):
|
||||
self.backend = FakeBackend()
|
||||
self.port = None
|
||||
self.chat_completion_count = 0
|
||||
self.total_requests = 0
|
||||
self.failed_requests = 0
|
||||
self.queue_depth = 0
|
||||
|
||||
def start(self):
|
||||
self.port = 8001
|
||||
return self.port
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
class FakeRelayHttpBridge:
|
||||
def __init__(self, relay_url, peer_id, local_base_url, advertised_addr):
|
||||
self.relay_url = relay_url
|
||||
self.peer_id = peer_id
|
||||
|
||||
@property
|
||||
def relay_addr(self):
|
||||
return f"{self.relay_url.replace('/ws', '')}/rpc/{self.peer_id}"
|
||||
|
||||
def start(self):
|
||||
return types.SimpleNamespace(peer_id=self.peer_id, relay_addr=self.relay_addr)
|
||||
|
||||
def wait_connected(self, timeout=5.0):
|
||||
return True
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"detect_hardware",
|
||||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
|
||||
)
|
||||
monkeypatch.setattr(startup_mod, "_detect_num_layers", lambda _model_id: 24)
|
||||
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
|
||||
monkeypatch.setattr(startup_mod, "RelayHttpBridge", FakeRelayHttpBridge)
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"_get_json",
|
||||
lambda url, timeout=10.0: {"relay_url": None, "nodes": []},
|
||||
)
|
||||
|
||||
tracker_url = "https://ai.neuron.d-popov.com"
|
||||
node = run_startup(
|
||||
tracker_url=tracker_url,
|
||||
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||||
advertise_host="172.29.104.23",
|
||||
wallet_path=tmp_path / "wallet.json",
|
||||
)
|
||||
try:
|
||||
pass
|
||||
finally:
|
||||
node.stop()
|
||||
|
||||
output = capsys.readouterr().out
|
||||
assert "Relay advertised by tracker" in output
|
||||
assert "wss://ai.neuron.d-popov.com/ws" in output
|
||||
assert "Cross-host pipeline hops WILL time out" not in output
|
||||
|
||||
|
||||
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)
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"_post_json",
|
||||
lambda _url, _payload, timeout=10.0: {"node_id": "node-test-123"},
|
||||
)
|
||||
|
||||
node = run_startup(
|
||||
tracker_url="http://127.0.0.1:8080",
|
||||
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
|
||||
assert node.tracker_node_id == "node-test-123"
|
||||
output = capsys.readouterr().out
|
||||
assert "Shard: layers 0–23; 24 of 24" in output
|
||||
assert "Node ID: node-test-123" in output
|
||||
|
||||
|
||||
def test_public_tracker_model_node_registers_relay_metadata_from_tracker_url_only(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
capsys,
|
||||
):
|
||||
"""A node only needs the public tracker URL to discover relay metadata and register."""
|
||||
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
|
||||
self.chat_completion_count = 0
|
||||
self.total_requests = 0
|
||||
self.failed_requests = 0
|
||||
self.queue_depth = 0
|
||||
|
||||
def start(self):
|
||||
self.port = 8001
|
||||
return self.port
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
class FakeRelayHttpBridge:
|
||||
def __init__(self, relay_url, peer_id, local_base_url, advertised_addr):
|
||||
self.relay_url = relay_url
|
||||
self.peer_id = peer_id
|
||||
self.local_base_url = local_base_url
|
||||
self.advertised_addr = advertised_addr
|
||||
|
||||
@property
|
||||
def relay_addr(self):
|
||||
return f"ws://public-relay.example/rpc/{self.peer_id}"
|
||||
|
||||
def start(self):
|
||||
return types.SimpleNamespace(
|
||||
peer_id=self.peer_id,
|
||||
relay_addr=self.relay_addr,
|
||||
)
|
||||
|
||||
def wait_connected(self, timeout=5.0):
|
||||
return True
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"detect_hardware",
|
||||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
|
||||
)
|
||||
monkeypatch.setattr(startup_mod, "_detect_num_layers", lambda _model_id: 24)
|
||||
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
|
||||
monkeypatch.setattr(startup_mod, "RelayHttpBridge", FakeRelayHttpBridge)
|
||||
|
||||
tracker = TrackerServer(relay_url="ws://public-relay.example/ws")
|
||||
tracker_port = tracker.start()
|
||||
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||
try:
|
||||
node = run_startup(
|
||||
tracker_url=tracker_url,
|
||||
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||||
wallet_path=tmp_path / "wallet.json",
|
||||
)
|
||||
try:
|
||||
network_map = _get_json(f"{tracker_url}/v1/network/map")
|
||||
finally:
|
||||
node.stop()
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
assert network_map["relay_url"] == "ws://public-relay.example/ws"
|
||||
assert len(network_map["nodes"]) == 1
|
||||
registered = network_map["nodes"][0]
|
||||
assert registered["hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
assert registered["endpoint"].startswith("http://")
|
||||
assert registered["endpoint"].endswith(":8001")
|
||||
assert registered["relay_addr"].startswith("ws://public-relay.example/rpc/")
|
||||
assert registered["peer_id"]
|
||||
output = capsys.readouterr().out
|
||||
assert "Relay advertised by tracker" in output
|
||||
assert "Cross-host pipeline hops WILL time out" not in output
|
||||
|
||||
|
||||
def test_public_tracker_relay_suppresses_virtual_ip_warning(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
capsys,
|
||||
):
|
||||
"""A WSL/Docker endpoint is acceptable when the tracker advertises relay RPC."""
|
||||
import meshnet_node.startup as startup_mod
|
||||
|
||||
class FakeBackend:
|
||||
total_layers = 24
|
||||
|
||||
class FakeTorchNodeServer:
|
||||
def __init__(self, **kwargs):
|
||||
self.backend = FakeBackend()
|
||||
self.port = None
|
||||
self.chat_completion_count = 0
|
||||
self.total_requests = 0
|
||||
self.failed_requests = 0
|
||||
self.queue_depth = 0
|
||||
|
||||
def start(self):
|
||||
self.port = 8001
|
||||
return self.port
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
class FakeRelayHttpBridge:
|
||||
def __init__(self, relay_url, peer_id, local_base_url, advertised_addr):
|
||||
self.relay_url = relay_url
|
||||
self.peer_id = peer_id
|
||||
|
||||
@property
|
||||
def relay_addr(self):
|
||||
return f"ws://public-relay.example/rpc/{self.peer_id}"
|
||||
|
||||
def start(self):
|
||||
return types.SimpleNamespace(peer_id=self.peer_id, relay_addr=self.relay_addr)
|
||||
|
||||
def wait_connected(self, timeout=5.0):
|
||||
return True
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"detect_hardware",
|
||||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
|
||||
)
|
||||
monkeypatch.setattr(startup_mod, "_detect_num_layers", lambda _model_id: 24)
|
||||
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
|
||||
monkeypatch.setattr(startup_mod, "RelayHttpBridge", FakeRelayHttpBridge)
|
||||
|
||||
tracker = TrackerServer(relay_url="ws://public-relay.example/ws")
|
||||
tracker_port = tracker.start()
|
||||
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||
try:
|
||||
node = run_startup(
|
||||
tracker_url=tracker_url,
|
||||
model_id="Qwen/Qwen2.5-0.5B-Instruct",
|
||||
advertise_host="172.29.104.23",
|
||||
wallet_path=tmp_path / "wallet.json",
|
||||
)
|
||||
try:
|
||||
network_map = _get_json(f"{tracker_url}/v1/network/map")
|
||||
finally:
|
||||
node.stop()
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
assert network_map["nodes"][0]["endpoint"] == "http://172.29.104.23:8001"
|
||||
assert network_map["nodes"][0]["relay_addr"].startswith("ws://public-relay.example/rpc/")
|
||||
output = capsys.readouterr().out
|
||||
assert "Relay advertised by tracker" in output
|
||||
assert "Cross-host pipeline hops WILL time out" not in output
|
||||
|
||||
|
||||
def test_later_node_auto_joins_existing_public_hf_model_with_only_tracker_url(
|
||||
tmp_path,
|
||||
monkeypatch,
|
||||
):
|
||||
"""After a model exists, a node can join by knowing only the public tracker URL."""
|
||||
import meshnet_node.startup as startup_mod
|
||||
|
||||
captured = {}
|
||||
|
||||
class FakeBackend:
|
||||
total_layers = 24
|
||||
|
||||
class FakeTorchNodeServer:
|
||||
def __init__(self, **kwargs):
|
||||
captured.update(kwargs)
|
||||
self.backend = FakeBackend()
|
||||
self.port = None
|
||||
self.chat_completion_count = 0
|
||||
self.total_requests = 0
|
||||
self.failed_requests = 0
|
||||
self.queue_depth = 0
|
||||
|
||||
def start(self):
|
||||
self.port = 8002
|
||||
return self.port
|
||||
|
||||
def stop(self):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
startup_mod,
|
||||
"detect_hardware",
|
||||
lambda: {"device": "cpu", "gpu_name": None, "vram_mb": 0},
|
||||
)
|
||||
monkeypatch.setattr(startup_mod, "TorchNodeServer", FakeTorchNodeServer)
|
||||
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
tracker_url = f"http://127.0.0.1:{tracker_port}"
|
||||
try:
|
||||
data = json.dumps({
|
||||
"endpoint": "http://203.0.113.20:8001",
|
||||
"model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"num_layers": 24,
|
||||
"shard_start": 0,
|
||||
"shard_end": 11,
|
||||
"tracker_mode": True,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0,
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"{tracker_url}/v1/nodes/register",
|
||||
data=data,
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req) as resp:
|
||||
resp.read()
|
||||
|
||||
node = run_startup(
|
||||
tracker_url=tracker_url,
|
||||
advertise_host="203.0.113.21",
|
||||
wallet_path=tmp_path / "wallet.json",
|
||||
)
|
||||
try:
|
||||
route_resp = _get_json(
|
||||
f"{tracker_url}/v1/route?model=Qwen/Qwen2.5-0.5B-Instruct"
|
||||
)
|
||||
finally:
|
||||
node.stop()
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
assert captured["model_id"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
assert captured["shard_start"] == 12
|
||||
assert captured["shard_end"] == 23
|
||||
assert route_resp["route"] == ["http://203.0.113.20:8001", "http://203.0.113.21:8002"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Full startup integration test
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -470,6 +853,163 @@ 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,6 +12,9 @@ import pytest
|
||||
from meshnet_node.model_backend import (
|
||||
InsufficientVRAMError,
|
||||
TensorPayload,
|
||||
_call_layer,
|
||||
_decoder_attention_mask,
|
||||
_int_tensor_header,
|
||||
build_quantization_config,
|
||||
validate_quantization,
|
||||
)
|
||||
@@ -33,7 +36,7 @@ class _FakeBackend:
|
||||
position_ids_header=None,
|
||||
)
|
||||
|
||||
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header):
|
||||
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
|
||||
assert shape == [1, 6, 8]
|
||||
return TensorPayload(
|
||||
body=body,
|
||||
@@ -47,11 +50,69 @@ class _FakeTailBackend(_FakeBackend):
|
||||
is_head = False
|
||||
is_tail = True
|
||||
|
||||
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header):
|
||||
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
|
||||
assert len(body) == 1 * 6 * 8 * 2
|
||||
return " Paris"
|
||||
|
||||
|
||||
class _FakeFullBackend(_FakeBackend):
|
||||
is_head = True
|
||||
is_tail = True
|
||||
|
||||
def generate_text(
|
||||
self,
|
||||
messages: list[dict],
|
||||
max_new_tokens: int = 16,
|
||||
temperature: float = 1.0,
|
||||
top_p: float = 1.0,
|
||||
) -> str:
|
||||
assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
|
||||
assert max_new_tokens == 7
|
||||
assert temperature == 1.0
|
||||
assert top_p == 1.0
|
||||
return "56"
|
||||
|
||||
def count_prompt_tokens(self, messages: list[dict]) -> int:
|
||||
assert messages == [{"role": "user", "content": "What is 7 times 8?"}]
|
||||
return 8
|
||||
|
||||
def count_text_tokens(self, text: str) -> int:
|
||||
assert text == "56"
|
||||
return 1
|
||||
|
||||
|
||||
class _FakeChatTokenizer:
|
||||
eos_token = ""
|
||||
|
||||
def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=False):
|
||||
assert add_generation_prompt is True
|
||||
assert tokenize is False
|
||||
return "debug prompt"
|
||||
|
||||
|
||||
class _FakePipelineHeadBackend(_FakeBackend):
|
||||
tokenizer = _FakeChatTokenizer()
|
||||
|
||||
def encode_prompt(self, prompt: str) -> TensorPayload:
|
||||
assert prompt == "debug prompt"
|
||||
return TensorPayload(
|
||||
body=b"\x00" * (1 * 6 * 8 * 2),
|
||||
shape=[1, 6, 8],
|
||||
attention_mask_header=None,
|
||||
position_ids_header=None,
|
||||
)
|
||||
|
||||
|
||||
class _FakePipelineTailBackend(_FakeTailBackend):
|
||||
def __init__(self) -> None:
|
||||
self.start_layers: list[int | None] = []
|
||||
|
||||
def forward_bytes(self, body, shape, attention_mask_header, position_ids_header, start_layer=None):
|
||||
self.start_layers.append(start_layer)
|
||||
assert len(body) == 1 * 6 * 8 * 2
|
||||
return " token"
|
||||
|
||||
|
||||
def test_quantization_flag_validation():
|
||||
assert validate_quantization("bfloat16") == "bfloat16"
|
||||
assert validate_quantization("int8") == "int8"
|
||||
@@ -145,6 +206,134 @@ 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_pipeline_hop_logs_are_suppressed_without_debug(capsys):
|
||||
tail_backend = _FakePipelineTailBackend()
|
||||
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True)
|
||||
tail = TorchNodeServer(backend=tail_backend)
|
||||
head_port = head.start()
|
||||
tail_port = tail.start()
|
||||
try:
|
||||
payload = json.dumps({
|
||||
"model": "fake-model",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 1,
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
||||
data=payload,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"X-Meshnet-Route": json.dumps([
|
||||
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
||||
]),
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
body = json.loads(resp.read())
|
||||
finally:
|
||||
head.stop()
|
||||
tail.stop()
|
||||
|
||||
out = capsys.readouterr().out
|
||||
assert body["choices"][0]["message"]["content"] == " token"
|
||||
assert tail_backend.start_layers == [22]
|
||||
assert "pipeline hop 0:" not in out
|
||||
assert "pipeline hop 0 returned text" not in out
|
||||
|
||||
|
||||
def test_pipeline_hop_logs_are_enabled_with_debug(capsys):
|
||||
head = TorchNodeServer(backend=_FakePipelineHeadBackend(), tracker_mode=True, debug=True)
|
||||
tail = TorchNodeServer(backend=_FakePipelineTailBackend())
|
||||
head_port = head.start()
|
||||
tail_port = tail.start()
|
||||
try:
|
||||
payload = json.dumps({
|
||||
"model": "fake-model",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 1,
|
||||
}).encode()
|
||||
req = urllib.request.Request(
|
||||
f"http://127.0.0.1:{head_port}/v1/chat/completions",
|
||||
data=payload,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"X-Meshnet-Route": json.dumps([
|
||||
{"endpoint": f"http://127.0.0.1:{tail_port}", "start_layer": 22},
|
||||
]),
|
||||
},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=5) as resp:
|
||||
json.loads(resp.read())
|
||||
finally:
|
||||
head.stop()
|
||||
tail.stop()
|
||||
|
||||
out = capsys.readouterr().out
|
||||
assert f" [node] pipeline hop 0: http://127.0.0.1:{tail_port} start_layer=22" in out
|
||||
assert " [node] pipeline hop 0 returned text=' token'" in out
|
||||
|
||||
|
||||
def test_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"):
|
||||
|
||||
255
tests/test_tracker_consensus.py
Normal file
255
tests/test_tracker_consensus.py
Normal file
@@ -0,0 +1,255 @@
|
||||
"""US-019 integration tests: distributed tracker consensus (Raft + gossip).
|
||||
|
||||
Three TrackerServer instances form a Raft cluster in-process. Tests verify:
|
||||
- Leader election completes within 1 second
|
||||
- Registration forwarded from follower reaches all nodes
|
||||
- Killing the leader triggers a new election within 5 seconds
|
||||
- Heartbeat gossip propagates across nodes
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
import urllib.request
|
||||
|
||||
import pytest
|
||||
|
||||
from meshnet_tracker.server import TrackerServer
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ helpers
|
||||
|
||||
def _get(url: str, timeout: float = 5.0) -> dict:
|
||||
with urllib.request.urlopen(url, timeout=timeout) as r:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def _post(url: str, payload: dict, timeout: float = 5.0) -> dict:
|
||||
data = json.dumps(payload).encode()
|
||||
req = urllib.request.Request(
|
||||
url, data=data, headers={"Content-Type": "application/json"}, method="POST"
|
||||
)
|
||||
with urllib.request.urlopen(req, timeout=timeout) as r:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def _register_node(tracker_url: str, port_hint: int) -> str:
|
||||
resp = _post(f"{tracker_url}/v1/nodes/register", {
|
||||
"endpoint": f"http://127.0.0.1:{port_hint}",
|
||||
"model": "stub-model",
|
||||
"shard_start": 0,
|
||||
"shard_end": 15,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0,
|
||||
})
|
||||
return resp["node_id"]
|
||||
|
||||
|
||||
def _raft_status(tracker_url: str) -> dict:
|
||||
return _get(f"{tracker_url}/v1/raft/status")
|
||||
|
||||
|
||||
def _wait_for_leader(
|
||||
urls: list[str], timeout: float = 5.0
|
||||
) -> tuple[str, list[str]]:
|
||||
"""Poll until exactly one leader is elected; return (leader_url, follower_urls)."""
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
leaders = []
|
||||
for url in urls:
|
||||
try:
|
||||
s = _raft_status(url)
|
||||
if s.get("role") == "leader":
|
||||
leaders.append(url)
|
||||
except Exception:
|
||||
pass
|
||||
if len(leaders) == 1:
|
||||
followers = [u for u in urls if u != leaders[0]]
|
||||
return leaders[0], followers
|
||||
time.sleep(0.05)
|
||||
raise TimeoutError(f"No leader elected within {timeout}s")
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ fixture
|
||||
|
||||
@pytest.fixture
|
||||
def three_tracker_cluster():
|
||||
"""Start 3 TrackerServer instances as a Raft cluster.
|
||||
|
||||
Yields list of (TrackerServer, url) tuples — index 0 first to start.
|
||||
"""
|
||||
# Use fixed starting port range for readability in logs; actual ports
|
||||
# assigned by OS (port=0) to avoid conflicts.
|
||||
trackers: list[TrackerServer] = []
|
||||
ports: list[int] = []
|
||||
|
||||
# Phase 1: start three servers without cluster config to get ports
|
||||
for _ in range(3):
|
||||
t = TrackerServer(host="127.0.0.1", port=0)
|
||||
ports.append(t.start())
|
||||
trackers.append(t)
|
||||
|
||||
# Stop them — we need to restart with peer URLs now that we know ports
|
||||
for t in trackers:
|
||||
t.stop()
|
||||
|
||||
trackers = []
|
||||
urls: list[str] = [f"http://127.0.0.1:{p}" for p in ports]
|
||||
|
||||
for i, port in enumerate(ports):
|
||||
peers = [u for j, u in enumerate(urls) if j != i]
|
||||
t = TrackerServer(
|
||||
host="127.0.0.1",
|
||||
port=port,
|
||||
cluster_peers=peers,
|
||||
cluster_self_url=urls[i],
|
||||
)
|
||||
actual_port = t.start()
|
||||
assert actual_port == port, f"port mismatch: wanted {port}, got {actual_port}"
|
||||
trackers.append(t)
|
||||
|
||||
yield list(zip(trackers, urls))
|
||||
|
||||
for t, _ in zip(trackers, urls):
|
||||
try:
|
||||
t.stop()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ tests
|
||||
|
||||
def test_leader_elected(three_tracker_cluster):
|
||||
"""Exactly one leader is elected within 1 second of cluster start."""
|
||||
_, urls = zip(*three_tracker_cluster)
|
||||
leader_url, followers = _wait_for_leader(list(urls), timeout=1.0)
|
||||
assert leader_url in urls
|
||||
assert len(followers) == 2
|
||||
|
||||
|
||||
def _wait_until_follower_knows_leader(follower_url: str, timeout: float = 2.0) -> None:
|
||||
"""Block until the follower has received a heartbeat and knows the leader URL."""
|
||||
deadline = time.monotonic() + timeout
|
||||
while time.monotonic() < deadline:
|
||||
try:
|
||||
s = _raft_status(follower_url)
|
||||
if s.get("leader") and s.get("role") == "follower":
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(0.05)
|
||||
raise TimeoutError(f"{follower_url} still doesn't know the leader after {timeout}s")
|
||||
|
||||
|
||||
def test_registration_on_follower_visible_on_all_nodes(three_tracker_cluster):
|
||||
"""Registering via a follower propagates the entry to all tracker nodes."""
|
||||
trackers_urls = three_tracker_cluster
|
||||
trackers, urls = zip(*trackers_urls)
|
||||
urls = list(urls)
|
||||
|
||||
leader_url, followers = _wait_for_leader(urls, timeout=1.0)
|
||||
|
||||
# Wait until the follower has received at least one heartbeat from the leader
|
||||
follower = followers[0]
|
||||
_wait_until_follower_knows_leader(follower, timeout=2.0)
|
||||
|
||||
# Register via a follower
|
||||
node_id = _register_node(follower, port_hint=19999)
|
||||
|
||||
# Allow replication to propagate (Raft heartbeat interval is 50ms)
|
||||
time.sleep(0.5)
|
||||
|
||||
# Every tracker node must know about this node
|
||||
for url in urls:
|
||||
# Use route endpoint as a proxy — if the node is replicated, the route
|
||||
# will include it. Alternatively, check /v1/raft/status log_length.
|
||||
status = _raft_status(url)
|
||||
assert status["log_length"] >= 1, (
|
||||
f"{url} has log_length={status['log_length']}, expected ≥1 after replication"
|
||||
)
|
||||
|
||||
|
||||
def test_follower_leader_status(three_tracker_cluster):
|
||||
"""All nodes agree on who the leader is after election."""
|
||||
_, urls = zip(*three_tracker_cluster)
|
||||
urls = list(urls)
|
||||
_wait_for_leader(urls, timeout=1.0)
|
||||
time.sleep(0.2) # let heartbeats propagate leader_url to followers
|
||||
|
||||
statuses = [_raft_status(u) for u in urls]
|
||||
leaders_reported = {s["leader"] for s in statuses if s.get("leader")}
|
||||
# All nodes that have a leader opinion should agree on the same one
|
||||
assert len(leaders_reported) <= 1 or len(leaders_reported) == len(
|
||||
{s["leader"] for s in statuses if s["role"] == "leader"}
|
||||
), f"Nodes disagree on leader: {leaders_reported}"
|
||||
|
||||
|
||||
def test_new_leader_elected_after_kill(three_tracker_cluster):
|
||||
"""Killing the leader triggers a new election within 5 seconds."""
|
||||
trackers, urls = zip(*three_tracker_cluster)
|
||||
trackers = list(trackers)
|
||||
urls = list(urls)
|
||||
|
||||
# Find the leader object
|
||||
leader_url, followers = _wait_for_leader(urls, timeout=1.0)
|
||||
leader_idx = urls.index(leader_url)
|
||||
leader_tracker = trackers[leader_idx]
|
||||
|
||||
# Register a node before killing the leader (proves it works)
|
||||
_register_node(leader_url, port_hint=19998)
|
||||
|
||||
# Kill the leader
|
||||
leader_tracker.stop()
|
||||
remaining_urls = [u for u in urls if u != leader_url]
|
||||
|
||||
# A new leader must be elected among the remaining 2 nodes
|
||||
new_leader_url, _ = _wait_for_leader(remaining_urls, timeout=5.0)
|
||||
assert new_leader_url in remaining_urls, f"New leader {new_leader_url!r} not in remaining {remaining_urls}"
|
||||
|
||||
# Registration must still work with the new leader
|
||||
node_id = _register_node(new_leader_url, port_hint=19997)
|
||||
assert node_id, "Registration after leader re-election returned no node_id"
|
||||
|
||||
|
||||
def test_registration_on_leader_visible_to_all(three_tracker_cluster):
|
||||
"""Registering with the leader replicates to all followers synchronously."""
|
||||
_, urls = zip(*three_tracker_cluster)
|
||||
urls = list(urls)
|
||||
|
||||
leader_url, followers = _wait_for_leader(urls, timeout=1.0)
|
||||
node_id = _register_node(leader_url, port_hint=19996)
|
||||
|
||||
# Allow Raft heartbeat to replicate the entry
|
||||
time.sleep(0.3)
|
||||
|
||||
for url in followers:
|
||||
status = _raft_status(url)
|
||||
assert status["log_length"] >= 1, (
|
||||
f"Follower {url} log_length={status['log_length']}, expected ≥1"
|
||||
)
|
||||
|
||||
|
||||
def test_gossip_propagates_heartbeat(three_tracker_cluster):
|
||||
"""Heartbeat recorded on one tracker propagates to others via gossip."""
|
||||
trackers, urls = zip(*three_tracker_cluster)
|
||||
trackers = list(trackers)
|
||||
urls = list(urls)
|
||||
|
||||
_wait_for_leader(urls, timeout=1.0)
|
||||
leader_url, _ = _wait_for_leader(urls, timeout=1.0)
|
||||
|
||||
# Register a node (so heartbeat makes sense)
|
||||
node_id = _register_node(leader_url, port_hint=19995)
|
||||
|
||||
# Send a heartbeat directly to the leader
|
||||
_post(f"{leader_url}/v1/nodes/{node_id}/heartbeat", {})
|
||||
|
||||
# Allow gossip to propagate to other nodes (gossip interval is 3s in prod,
|
||||
# but we just want to verify the gossip table was updated locally)
|
||||
leader_idx = urls.index(leader_url)
|
||||
leader_gossip = trackers[leader_idx]._gossip
|
||||
assert leader_gossip is not None, "Leader should have gossip enabled"
|
||||
assert leader_gossip.last_seen(node_id) is not None, (
|
||||
"Gossip table on leader should record the heartbeat"
|
||||
)
|
||||
@@ -27,6 +27,216 @@ def _get_json(url: str) -> dict:
|
||||
return json.loads(r.read())
|
||||
|
||||
|
||||
def test_tracker_send_json_ignores_broken_pipe_after_client_disconnect():
|
||||
"""A disconnected client must not dump a BrokenPipe traceback from the tracker."""
|
||||
from meshnet_tracker.server import _TrackerHandler
|
||||
|
||||
class BrokenPipeWriter:
|
||||
def write(self, _body):
|
||||
raise BrokenPipeError
|
||||
|
||||
class DummyHandler:
|
||||
wfile = BrokenPipeWriter()
|
||||
|
||||
def send_response(self, _status):
|
||||
pass
|
||||
|
||||
def send_header(self, _name, _value):
|
||||
pass
|
||||
|
||||
def end_headers(self):
|
||||
pass
|
||||
|
||||
_TrackerHandler._send_json(DummyHandler(), 200, {"ok": True})
|
||||
|
||||
|
||||
def test_tracker_serves_health_while_proxy_request_is_in_flight():
|
||||
"""Long inference proxy requests must not block heartbeats/health checks."""
|
||||
|
||||
class SlowChatHandler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, fmt, *args):
|
||||
pass
|
||||
|
||||
def do_POST(self):
|
||||
if self.path != "/v1/chat/completions":
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
return
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
self.rfile.read(length)
|
||||
time.sleep(2.0)
|
||||
body = json.dumps({"choices": [{"message": {"content": "ok"}}]}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
slow_node = http.server.HTTPServer(("127.0.0.1", 0), SlowChatHandler)
|
||||
slow_thread = threading.Thread(target=slow_node.serve_forever, daemon=True)
|
||||
slow_thread.start()
|
||||
tracker = TrackerServer(heartbeat_timeout=60.0)
|
||||
tracker_port = tracker.start()
|
||||
proxy_error = []
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": f"http://127.0.0.1:{slow_node.server_address[1]}",
|
||||
"model": "slow-model", "hf_repo": "org/slow-model", "num_layers": 1,
|
||||
"shard_start": 0, "shard_end": 0, "tracker_mode": True,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
|
||||
def call_proxy():
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
|
||||
{"model": "slow-model", "messages": [{"role": "user", "content": "hi"}]},
|
||||
)
|
||||
except Exception as exc:
|
||||
proxy_error.append(exc)
|
||||
|
||||
proxy_thread = threading.Thread(target=call_proxy)
|
||||
proxy_thread.start()
|
||||
time.sleep(0.2)
|
||||
|
||||
with urllib.request.urlopen(f"http://127.0.0.1:{tracker_port}/v1/health", timeout=1.0) as resp:
|
||||
assert resp.status == 200
|
||||
|
||||
proxy_thread.join(timeout=3.0)
|
||||
assert not proxy_thread.is_alive()
|
||||
assert not proxy_error
|
||||
finally:
|
||||
tracker.stop()
|
||||
slow_node.shutdown()
|
||||
slow_node.server_close()
|
||||
slow_thread.join(timeout=1.0)
|
||||
|
||||
|
||||
def test_tracker_routes_hf_model_alias_from_quickstart():
|
||||
"""The documented qwen2.5-0.5b alias resolves a full HF repo registration."""
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9100",
|
||||
"model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"num_layers": 24,
|
||||
"shard_start": 0,
|
||||
"shard_end": 23,
|
||||
"tracker_mode": True,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0},
|
||||
)
|
||||
|
||||
route_resp = _get_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/route?model=qwen2.5-0.5b"
|
||||
)
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
assert route_resp["route"] == ["http://127.0.0.1:9100"]
|
||||
|
||||
|
||||
def test_tracker_proxy_accepts_hf_model_alias_from_quickstart():
|
||||
"""The tracker proxy accepts the same model alias used by the quickstart curl."""
|
||||
|
||||
class ChatHandler(http.server.BaseHTTPRequestHandler):
|
||||
def log_message(self, fmt, *args):
|
||||
pass
|
||||
|
||||
def do_POST(self):
|
||||
if self.path != "/v1/chat/completions":
|
||||
self.send_response(404)
|
||||
self.end_headers()
|
||||
return
|
||||
length = int(self.headers.get("Content-Length", 0))
|
||||
request_body = json.loads(self.rfile.read(length) or b"{}")
|
||||
body = json.dumps({
|
||||
"model": request_body["model"],
|
||||
"choices": [{"message": {"content": "56"}}],
|
||||
}).encode()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
node = http.server.HTTPServer(("127.0.0.1", 0), ChatHandler)
|
||||
node_thread = threading.Thread(target=node.serve_forever, daemon=True)
|
||||
node_thread.start()
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": f"http://127.0.0.1:{node.server_address[1]}",
|
||||
"model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"num_layers": 24,
|
||||
"shard_start": 0,
|
||||
"shard_end": 23,
|
||||
"tracker_mode": True,
|
||||
"hardware_profile": {},
|
||||
"score": 1.0},
|
||||
)
|
||||
|
||||
response = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/chat/completions",
|
||||
{"model": "qwen2.5-0.5b",
|
||||
"messages": [{"role": "user", "content": "What is 7 times 8?"}]},
|
||||
)
|
||||
finally:
|
||||
tracker.stop()
|
||||
node.shutdown()
|
||||
node.server_close()
|
||||
node_thread.join(timeout=1.0)
|
||||
|
||||
assert response["choices"][0]["message"]["content"] == "56"
|
||||
|
||||
|
||||
def test_tracker_registration_node_id_includes_wallet_prefix_and_stable_suffix():
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
wallet = "7j77FsPY1evV8tuf4Z73AVrWwxBEW1pvKwi4EvcRD3g"
|
||||
try:
|
||||
first = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9100", "model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
||||
"shard_start": 0, "shard_end": 21, "wallet_address": wallet,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
second = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9100", "model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
||||
"shard_start": 0, "shard_end": 21, "wallet_address": wallet,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
different_endpoint = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
||||
"shard_start": 20, "shard_end": 23, "wallet_address": wallet,
|
||||
"hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
|
||||
assert first["node_id"].startswith("7j77FsPY-")
|
||||
assert second["node_id"] == first["node_id"]
|
||||
assert different_endpoint["node_id"].startswith("7j77FsPY-")
|
||||
assert different_endpoint["node_id"] != first["node_id"]
|
||||
|
||||
route_resp = _get_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/route?model=Qwen/Qwen2.5-0.5B-Instruct"
|
||||
)
|
||||
assert route_resp["nodes"][0]["node_id"] == first["node_id"]
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_node_registration():
|
||||
"""A node can register with the tracker and receives a node_id."""
|
||||
tracker = TrackerServer()
|
||||
@@ -117,10 +327,24 @@ def test_tracker_coverage_endpoint_reports_uncovered_ranges():
|
||||
)
|
||||
|
||||
coverage_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/coverage/tiny-model")
|
||||
assert coverage_resp["coverage"] == [
|
||||
{"start_layer": 0, "end_layer": 2, "node_count": 1},
|
||||
{"start_layer": 3, "end_layer": 5, "node_count": 0},
|
||||
]
|
||||
bands = coverage_resp["coverage"]
|
||||
assert len(bands) == 2
|
||||
# Covered band: layers 0-2, one node
|
||||
assert bands[0]["start_layer"] == 0
|
||||
assert bands[0]["end_layer"] == 2
|
||||
assert bands[0]["node_count"] == 1
|
||||
assert len(bands[0]["nodes"]) == 1
|
||||
node_info = bands[0]["nodes"][0]
|
||||
assert node_info["endpoint"] == "http://127.0.0.1:9001"
|
||||
assert node_info["alive"] is True
|
||||
assert "last_seen_seconds_ago" in node_info
|
||||
assert "heartbeat_success_rate" in node_info
|
||||
assert "inference_success_rate" in node_info
|
||||
# Gap band: layers 3-5, no nodes
|
||||
assert bands[1]["start_layer"] == 3
|
||||
assert bands[1]["end_layer"] == 5
|
||||
assert bands[1]["node_count"] == 0
|
||||
assert bands[1]["nodes"] == []
|
||||
|
||||
try:
|
||||
_get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=tiny-model")
|
||||
@@ -131,6 +355,65 @@ def test_tracker_coverage_endpoint_reports_uncovered_ranges():
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_coverage_endpoint_accepts_registered_hf_repo_or_short_name():
|
||||
"""Coverage endpoint supports real HF models registered outside preset catalog."""
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
||||
"shard_start": 0, "shard_end": 21, "hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9102", "model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
||||
"shard_start": 20, "shard_end": 23, "hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
|
||||
by_repo = _get_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen/Qwen2.5-0.5B-Instruct"
|
||||
)
|
||||
by_short_name = _get_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/coverage/Qwen2.5-0.5B-Instruct"
|
||||
)
|
||||
|
||||
assert by_repo["model"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
bands = by_repo["coverage"]
|
||||
assert len(bands) == 3
|
||||
assert bands[0]["start_layer"] == 0 and bands[0]["end_layer"] == 19 and bands[0]["node_count"] == 1
|
||||
assert bands[1]["start_layer"] == 20 and bands[1]["end_layer"] == 21 and bands[1]["node_count"] == 2
|
||||
assert bands[2]["start_layer"] == 22 and bands[2]["end_layer"] == 23 and bands[2]["node_count"] == 1
|
||||
assert all("nodes" in b for b in bands)
|
||||
assert by_short_name["coverage"] == by_repo["coverage"]
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_models_endpoint_lists_registered_hf_repo_and_short_name_alias():
|
||||
tracker = TrackerServer()
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
_post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9111", "model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
||||
"shard_start": 0, "shard_end": 23, "hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
|
||||
models_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/models")
|
||||
|
||||
model = next(item for item in models_resp["data"] if item["id"] == "Qwen/Qwen2.5-0.5B-Instruct")
|
||||
assert model["name"] == "Qwen2.5-0.5B-Instruct"
|
||||
assert model["hf_repo"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
assert model["aliases"] == ["Qwen/Qwen2.5-0.5B-Instruct", "Qwen2.5-0.5B-Instruct"]
|
||||
assert model["shard_coverage_percentage"] == 100.0
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_auto_assigns_new_node_to_uncovered_range_first():
|
||||
"""Capability-driven registration fills the first uncovered layer gap."""
|
||||
tracker = TrackerServer(model_presets={
|
||||
@@ -388,6 +671,42 @@ def test_tracker_rebalances_after_middle_range_node_timeout():
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_rebalances_managed_hf_node_after_peer_timeout():
|
||||
"""HF nodes auto-assigned by the tracker receive LOAD_SHARD after a peer dies."""
|
||||
tracker = TrackerServer(heartbeat_timeout=0.15, rebalance_interval=10.0)
|
||||
tracker_port = tracker.start()
|
||||
try:
|
||||
managed = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9101", "model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
||||
"shard_start": 0, "shard_end": 21, "managed_assignment": True,
|
||||
"vram_bytes": 1_000_000_000, "hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
expired = _post_json(
|
||||
f"http://127.0.0.1:{tracker_port}/v1/nodes/register",
|
||||
{"endpoint": "http://127.0.0.1:9102", "model": "Qwen2.5-0.5B-Instruct",
|
||||
"hf_repo": "Qwen/Qwen2.5-0.5B-Instruct", "num_layers": 24,
|
||||
"shard_start": 22, "shard_end": 23,
|
||||
"vram_bytes": 1_000_000_000, "hardware_profile": {}, "score": 1.0},
|
||||
)
|
||||
|
||||
time.sleep(0.10)
|
||||
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{managed['node_id']}/heartbeat", {})
|
||||
time.sleep(0.10)
|
||||
|
||||
hb = _post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/{managed['node_id']}/heartbeat", {})
|
||||
|
||||
assert expired["node_id"] not in tracker._registry
|
||||
load_directives = [d for d in hb.get("directives", []) if d["action"] == "LOAD_SHARD"]
|
||||
assert load_directives
|
||||
assert load_directives[-1]["model"] == "Qwen/Qwen2.5-0.5B-Instruct"
|
||||
assert load_directives[-1]["start_layer"] == 0
|
||||
assert load_directives[-1]["end_layer"] == 23
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_tracker_route_error_no_nodes():
|
||||
"""Tracker returns 503 with clear error when the registry is empty."""
|
||||
tracker = TrackerServer()
|
||||
@@ -890,3 +1209,273 @@ def test_two_node_pipeline_via_tracker():
|
||||
node_a.stop()
|
||||
node_b.stop()
|
||||
tracker.stop()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- US-025 stats
|
||||
|
||||
def test_stats_endpoint_returns_empty_models_initially():
|
||||
tracker = TrackerServer()
|
||||
port = tracker.start()
|
||||
try:
|
||||
resp = _get_json(f"http://127.0.0.1:{port}/v1/stats")
|
||||
assert "models" in resp
|
||||
assert isinstance(resp["models"], dict)
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_stats_rolling_counter_rpm_reflects_recorded_requests():
|
||||
from meshnet_tracker.server import _RollingCounter
|
||||
|
||||
counter = _RollingCounter(num_buckets=60, bucket_seconds=60)
|
||||
now = 1_000_000.0
|
||||
for _ in range(60):
|
||||
counter.record(now)
|
||||
|
||||
rpm = counter.rpm(now)
|
||||
assert rpm == 1.0 # 60 requests over 60 minutes = 1 rpm
|
||||
|
||||
|
||||
def test_stats_collector_records_and_returns_rpms():
|
||||
from meshnet_tracker.server import _StatsCollector
|
||||
|
||||
collector = _StatsCollector()
|
||||
now = 1_000_000.0
|
||||
for _ in range(30):
|
||||
collector.record_request("my-model", now=now)
|
||||
|
||||
combined = collector.get_combined_stats(now=now)
|
||||
assert "my-model" in combined
|
||||
assert combined["my-model"]["rpm_last_hour"] > 0
|
||||
assert combined["my-model"]["rpm_last_day"] > 0
|
||||
assert combined["my-model"]["rpm_last_month"] > 0
|
||||
|
||||
|
||||
def test_stats_collector_merges_peer_rpms_additively():
|
||||
from meshnet_tracker.server import _StatsCollector
|
||||
|
||||
collector = _StatsCollector()
|
||||
now = 1_000_000.0
|
||||
collector.record_request("modelX", now=now) # 1 local request in the current minute bucket
|
||||
|
||||
peer_rpms = {"modelX": {"rpm_last_hour": 5.0, "rpm_last_day": 3.0, "rpm_last_month": 1.0}}
|
||||
collector.merge_peer_rpms("http://peer-tracker:8080", peer_rpms)
|
||||
|
||||
combined = collector.get_combined_stats(now=now)
|
||||
assert combined["modelX"]["rpm_last_hour"] > 5.0 # local + peer
|
||||
assert combined["modelX"]["rpm_last_day"] > 3.0
|
||||
assert combined["modelX"]["rpm_last_month"] >= 1.0 # local rounds to 0 at 4dp, peer is 1.0
|
||||
|
||||
|
||||
def test_stats_sqlite_persistence_survives_restart(tmp_path):
|
||||
from meshnet_tracker.server import _StatsCollector
|
||||
|
||||
db = str(tmp_path / "stats.db")
|
||||
now = 1_000_000.0
|
||||
|
||||
collector = _StatsCollector(db_path=db)
|
||||
for _ in range(10):
|
||||
collector.record_request("persistent-model", now=now)
|
||||
collector.save_to_db()
|
||||
|
||||
collector2 = _StatsCollector(db_path=db)
|
||||
combined = collector2.get_combined_stats(now=now)
|
||||
assert "persistent-model" in combined
|
||||
assert combined["persistent-model"]["rpm_last_hour"] > 0
|
||||
|
||||
|
||||
def test_stats_gossip_endpoint_merges_peer_slice():
|
||||
tracker = TrackerServer()
|
||||
port = tracker.start()
|
||||
try:
|
||||
peer_payload = {
|
||||
"tracker_url": "http://peer:9000",
|
||||
"stats": {
|
||||
"remote-model": {
|
||||
"rpm_last_hour": 2.5,
|
||||
"rpm_last_day": 1.0,
|
||||
"rpm_last_month": 0.5,
|
||||
}
|
||||
},
|
||||
}
|
||||
_post_json(f"http://127.0.0.1:{port}/v1/stats/gossip", peer_payload)
|
||||
|
||||
combined = _get_json(f"http://127.0.0.1:{port}/v1/stats")
|
||||
assert "remote-model" in combined["models"]
|
||||
assert combined["models"]["remote-model"]["rpm_last_hour"] == 2.5
|
||||
finally:
|
||||
tracker.stop()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- US-027 / US-028 routing
|
||||
|
||||
def _make_node(node_id, shard_start, shard_end, bench=1.0, queue_depth=0):
|
||||
"""Helper: build a _NodeEntry with minimal fields."""
|
||||
from meshnet_tracker.server import _NodeEntry
|
||||
n = _NodeEntry(
|
||||
node_id=node_id, endpoint=f"http://fake:{9000 + shard_start}",
|
||||
shard_start=shard_start, shard_end=shard_end,
|
||||
model="m", shard_checksum=None, hardware_profile={},
|
||||
wallet_address=None, score=1.0,
|
||||
benchmark_tokens_per_sec=bench,
|
||||
)
|
||||
n.queue_depth = queue_depth
|
||||
return n
|
||||
|
||||
|
||||
def test_select_route_overlapping_shards_uses_start_layer_protocol():
|
||||
"""_select_route with overlapping shards: A(0-22) and B(20-24).
|
||||
Route must be [A, B]; B's effective start_layer = 23 (after A ends at 22).
|
||||
"""
|
||||
from meshnet_tracker.server import _select_route
|
||||
|
||||
node_a = _make_node("A", 0, 22)
|
||||
node_b = _make_node("B", 20, 24)
|
||||
|
||||
route, err = _select_route([node_a, node_b], 0, 24)
|
||||
assert err == ""
|
||||
assert len(route) == 2
|
||||
assert route[0].node_id == "A"
|
||||
assert route[1].node_id == "B"
|
||||
|
||||
|
||||
def test_select_route_no_overlap_three_nodes():
|
||||
"""Non-overlapping shards are covered in the right order."""
|
||||
from meshnet_tracker.server import _select_route
|
||||
|
||||
nodes = [_make_node("C", 8, 15), _make_node("A", 0, 7), _make_node("B", 16, 23)]
|
||||
route, err = _select_route(nodes, 0, 23)
|
||||
assert err == ""
|
||||
# Greedy cover order: A(0-7) → C(8-15) → B(16-23)
|
||||
assert [n.node_id for n in route] == ["A", "C", "B"]
|
||||
|
||||
|
||||
def test_select_route_gap_returns_error():
|
||||
from meshnet_tracker.server import _select_route
|
||||
|
||||
nodes = [_make_node("A", 0, 5), _make_node("B", 10, 15)]
|
||||
route, err = _select_route(nodes, 0, 15)
|
||||
assert route == []
|
||||
assert "layer 6" in err
|
||||
|
||||
|
||||
def test_select_route_prefers_higher_throughput_node_when_shards_equal():
|
||||
"""When two nodes cover the same interval, prefer the faster / less-loaded one."""
|
||||
from meshnet_tracker.server import _select_route
|
||||
|
||||
fast = _make_node("fast", 0, 11, bench=10.0, queue_depth=0) # effective = 10.0
|
||||
slow = _make_node("slow", 0, 11, bench=2.0, queue_depth=0) # effective = 2.0
|
||||
|
||||
route, err = _select_route([slow, fast], 0, 11)
|
||||
assert err == ""
|
||||
assert len(route) == 1
|
||||
assert route[0].node_id == "fast"
|
||||
|
||||
|
||||
def test_select_route_throughput_accounts_for_queue_depth():
|
||||
"""A nominally faster node with high queue depth loses to a slower idle node."""
|
||||
from meshnet_tracker.server import _select_route
|
||||
|
||||
busy_fast = _make_node("busy", 0, 11, bench=20.0, queue_depth=9) # effective = 2.0
|
||||
idle_slow = _make_node("idle", 0, 11, bench=2.0, queue_depth=0) # effective = 2.0
|
||||
|
||||
route, err = _select_route([busy_fast, idle_slow], 0, 11)
|
||||
assert err == ""
|
||||
# tie → first in list wins; exact tie does not matter as long as no error
|
||||
assert len(route) == 1
|
||||
assert route[0].node_id in ("busy", "idle")
|
||||
|
||||
|
||||
def test_two_stub_nodes_complete_pipeline_via_tracker():
|
||||
"""Integration: two StubNodeServer instances serving complementary shards
|
||||
produce a full inference response through the tracker route."""
|
||||
tracker = TrackerServer(model_presets={
|
||||
"two-node-model": {"layers_start": 0, "layers_end": 3}
|
||||
})
|
||||
tracker_port = tracker.start()
|
||||
|
||||
from meshnet_node.server import StubNodeServer
|
||||
node_a = StubNodeServer()
|
||||
node_b = StubNodeServer()
|
||||
port_a = node_a.start()
|
||||
port_b = node_b.start()
|
||||
|
||||
try:
|
||||
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {
|
||||
"endpoint": f"http://127.0.0.1:{port_a}",
|
||||
"model": "two-node-model", "shard_start": 0, "shard_end": 1,
|
||||
"hardware_profile": {}, "score": 1.0,
|
||||
})
|
||||
_post_json(f"http://127.0.0.1:{tracker_port}/v1/nodes/register", {
|
||||
"endpoint": f"http://127.0.0.1:{port_b}",
|
||||
"model": "two-node-model", "shard_start": 2, "shard_end": 3,
|
||||
"hardware_profile": {}, "score": 1.0,
|
||||
})
|
||||
|
||||
route_resp = _get_json(f"http://127.0.0.1:{tracker_port}/v1/route?model=two-node-model")
|
||||
assert len(route_resp["route"]) == 2
|
||||
finally:
|
||||
node_a.stop()
|
||||
node_b.stop()
|
||||
tracker.stop()
|
||||
|
||||
|
||||
def test_route_timeout_config_is_exposed_on_server():
|
||||
"""TorchNodeServer exposes the route_timeout that was configured."""
|
||||
from meshnet_node.torch_server import TorchNodeServer
|
||||
|
||||
class _MinimalBackend:
|
||||
model_id = "m"
|
||||
total_layers = 4
|
||||
is_head = True
|
||||
is_tail = True
|
||||
def generate_text(self, *a, **kw): return ""
|
||||
def count_prompt_tokens(self, *a): return 0
|
||||
def count_text_tokens(self, *a): return 0
|
||||
|
||||
node = TorchNodeServer(backend=_MinimalBackend(), route_timeout=45.0)
|
||||
assert node.route_timeout == 45.0
|
||||
|
||||
|
||||
def test_torch_node_applies_tracker_load_shard_directive(monkeypatch):
|
||||
from meshnet_node import torch_server
|
||||
from meshnet_node.torch_server import TorchNodeServer
|
||||
|
||||
class _MinimalBackend:
|
||||
def __init__(self, model_id="Qwen/Qwen2.5-0.5B-Instruct", shard_start=0, shard_end=21, quantization="bfloat16"):
|
||||
self.model_id = model_id
|
||||
self.shard_start = shard_start
|
||||
self.shard_end = shard_end
|
||||
self.quantization = quantization
|
||||
self.total_layers = 24
|
||||
self.is_head = shard_start == 0
|
||||
self.is_tail = shard_end == 23
|
||||
|
||||
def generate_text(self, *a, **kw): return ""
|
||||
def count_prompt_tokens(self, *a): return 0
|
||||
def count_text_tokens(self, *a): return 0
|
||||
|
||||
loaded = []
|
||||
|
||||
def fake_load_backend(model_id, shard_start, shard_end, quantization):
|
||||
loaded.append((model_id, shard_start, shard_end, quantization))
|
||||
return _MinimalBackend(model_id, shard_start, shard_end, quantization)
|
||||
|
||||
monkeypatch.setattr(torch_server, "_load_backend", fake_load_backend)
|
||||
node = TorchNodeServer(backend=_MinimalBackend())
|
||||
|
||||
applied = node.apply_tracker_directives([
|
||||
{"action": "DROP_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct", "shard_start": 0, "shard_end": 21},
|
||||
{"action": "LOAD_SHARD", "model": "Qwen/Qwen2.5-0.5B-Instruct", "shard_start": 0, "shard_end": 23,
|
||||
"quantization": "bfloat16"},
|
||||
])
|
||||
|
||||
assert loaded == [("Qwen/Qwen2.5-0.5B-Instruct", 0, 23, "bfloat16")]
|
||||
assert applied == {
|
||||
"model": "Qwen/Qwen2.5-0.5B-Instruct",
|
||||
"shard_start": 0,
|
||||
"shard_end": 23,
|
||||
"quantization": "bfloat16",
|
||||
"tracker_mode": True,
|
||||
}
|
||||
assert node.backend.shard_end == 23
|
||||
|
||||
Reference in New Issue
Block a user