popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'master' of https://git.d-popov.com/popov/neuron-tai

Browse Source
This commit is contained in:
Dobromir Popov
2026-06-30 19:28:43 +02:00
parent f1e4870124 7866723c82
commit d0307fcc84
8 changed files with 590 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files

156
QUICKSTART.md
View File

@@ -48,13 +48,10 @@ python3 -m venv .venv
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 is still useful for local development, but do not rely on it for the
"another machine connects back to this node" LAN case. WSL2 commonly sits behind
Windows NAT/port-proxy behavior and may not accept inbound traffic from other LAN
machines without extra host networking setup. We intentionally leave that unfixed
because it is useful for testing NAT/relay scenarios. If you just want to bring up
a Windows node that other machines can reach directly, run the node in native
Windows PowerShell instead.
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)
@@ -126,16 +123,21 @@ $env:HF_HOME = "D:\DEV\models"
--port 8005
```
One-line variants:
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
.\.venv\Scripts\meshnet-node.exe start --tracker http://ai.neuron.d-popov.com --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
```
`--host 0.0.0.0` binds the node to all Windows interfaces. `--advertise-host`
is what the tracker gives to other nodes, so it must be the Windows LAN IP that
the tracker and peer nodes can actually reach.
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
@@ -144,6 +146,7 @@ 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:
@@ -152,32 +155,47 @@ lines like:
curl http://192.168.0.42:8005/v1/health
```
If that endpoint returns 404, 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.
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 + WSS relay
---
For internet nodes, expose one public HTTPS host and proxy these paths:
## Public tracker + relay (internet / NAT nodes)
```text
/v1/* -> meshnet-tracker, for registration, heartbeats, routing, and OpenAI requests
/ws -> meshnet-relay, for outbound node gossip/bridge connections
/rpc/* -> meshnet-relay, for tracker-to-node relay requests
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)
```
Start the tracker with the public relay URL it should advertise:
### Start the relay and tracker (server side)
```bash
# Terminal 1 — relay (WebSocket hub)
.venv/bin/meshnet-relay --host 0.0.0.0 --port 8765
# Terminal 2 — tracker (advertises relay URL to nodes)
.venv/bin/meshnet-tracker start \
--host 0.0.0.0 \
--port 8081 \
--relay-url wss://ai.neuron.d-popov.com/ws
```
Then a node only needs the public tracker address:
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` needed. The node discovers the relay URL from the tracker
and opens a persistent outbound WebSocket:
```bash
.venv/bin/meshnet-node start \
@@ -201,6 +219,84 @@ 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)
@@ -338,7 +434,7 @@ HF_HOME=/run/media/popov/d/DEV/models \
--port 8002
```
Send the request to Node A — it tokenizes, runs layers 0–13, passes binary
Send the request to Node A — it tokenizes, runs layers 0–11, passes binary
activations to Node B, and streams the final response back.
---
@@ -382,18 +478,6 @@ tracker with a relay URL so the node registers a `relay_addr`.
---
## Start the relay node (for NAT traversal)
```bash
.venv/bin/pip install -e packages/relay
.venv/bin/meshnet-relay --port 8765
```
Nodes behind NAT connect to the relay and advertise their relay address to the
tracker. See `docs/adr/0010-p2p-gossip-and-nat-relay.md`.
---
## Run all tests
```bash

26
packages/node/meshnet_node/relay_bridge.py
View File

@@ -2,6 +2,7 @@
from __future__ import annotations
import base64
import json
import logging
import threading
@@ -114,20 +115,35 @@ class RelayHttpBridge:
request_id = str(payload.get("request_id") or "")
method = str(payload.get("method") or "POST").upper()
path = str(payload.get("path") or "/")
body_text = payload.get("body") 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}"
data = body_text.encode() if isinstance(body_text, str) else bytes(body_text)
req = urllib.request.Request(url, data=data, headers=headers, method=method)
try:
with urllib.request.urlopen(req, timeout=300.0) as resp:
return {
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": {"Content-Type": resp.headers.get("Content-Type", "application/json")},
"body": resp.read().decode(errors="replace"),
"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,

24
packages/node/meshnet_node/startup.py
View File

@@ -123,6 +123,26 @@ def _start_heartbeat(
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"
@@ -150,6 +170,7 @@ def _start_heartbeat(
try:
resp = _post_json(hb_url, _get_stats())
_apply_directives(resp.get("directives", []))
new_asgn = resp.get("new_assignment")
if new_asgn:
print(
@@ -282,6 +303,7 @@ def run_startup(
print(f" {probationary_line}", flush=True)
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)
@@ -355,6 +377,7 @@ def run_startup(
"quantization": quantization,
"score": 1.0,
"tracker_mode": (shard_start == 0),
"managed_assignment": not user_pinned_shard,
**relay_fields,
}
tracker_node_id: str | None = None
@@ -442,6 +465,7 @@ def run_startup(
"quantization": quantization,
"score": 1.0,
"tracker_mode": (assigned_shard_start == 0),
"managed_assignment": True,
**relay_fields,
}
tracker_node_id = None

155
packages/node/meshnet_node/torch_server.py
View File

@@ -2,6 +2,7 @@
from __future__ import annotations
import base64
import http.server
import json
import sys
@@ -29,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,
@@ -326,8 +361,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
result_text = "".join(generated)
self._send_openai_response(result_text, model_name, stream, messages)
def _get_remaining_route(self, model: str) -> list[tuple[str, int]]:
"""Return downstream hops as (endpoint, start_layer) pairs.
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.
@@ -340,13 +375,19 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
try:
route = json.loads(injected)
if isinstance(route, list):
hops: list[tuple[str, int]] = []
hops: list[dict] = []
for item in route:
if isinstance(item, dict):
hops.append((str(item["endpoint"]), int(item.get("start_layer", 0))))
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((item, 0)) # backward-compat: plain string, no start_layer
print(f" [node] using injected downstream route: {[ep for ep, _ in hops]}", flush=True)
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
@@ -362,7 +403,6 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
route_resp = json.loads(r.read())
own_port = server.server_address[1]
nodes_info = route_resp.get("nodes", [])
# nodes_info is ordered; find own node and compute start_layers post-hoc
hops = []
covered_up_to: int | None = None
for node_info in nodes_info:
@@ -371,18 +411,19 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
covered_up_to = node_info.get("shard_end")
continue
if covered_up_to is None:
# Own node not found yet; use node's shard_start as fallback
covered_up_to = (node_info.get("shard_start") or 1) - 1
start_l = covered_up_to + 1
hops.append((ep, start_l))
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: {[ep for ep, _ in hops]}", flush=True)
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[tuple[str, int]]) -> str:
def _run_downstream_pipeline(self, payload: object, route: list[dict]) -> str:
server: _TorchHTTPServer = self.server # type: ignore[assignment]
if not route:
# Partial shard at tail: decode the activation from the previous node.
@@ -407,9 +448,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
current_attn = attn_mask
current_pos = pos_ids
for hop_index, (node_url, start_layer) 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}", flush=True)
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,
@@ -425,19 +473,38 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
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=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}"
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:
@@ -662,6 +729,40 @@ class TorchNodeServer:
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")

118
packages/tracker/meshnet_tracker/server.py
View File

@@ -706,9 +706,109 @@ def _rebalance_model_locked(server: "_TrackerHTTPServer", model: str) -> None:
)
def _hf_rebalance_preset(nodes: list[_NodeEntry]) -> dict:
total_layers = max(node.num_layers or 0 for node in nodes)
return {
"layers_start": 0,
"layers_end": total_layers - 1,
"bytes_per_layer": {"bfloat16": 30 * 1024 * 1024, "int8": 15 * 1024 * 1024, "nf4": 8 * 1024 * 1024},
}
def _rebalance_hf_model_locked(server: "_TrackerHTTPServer", hf_repo: str) -> None:
model_nodes = [
node for node in server.registry.values()
if node.hf_repo == hf_repo
and node.num_layers is not None
]
managed_nodes = [node for node in model_nodes if node.managed_assignment]
if not model_nodes or not managed_nodes:
return
preset = _hf_rebalance_preset(model_nodes)
required_start, required_end = _preset_layer_bounds(preset)
total_layers = required_end - required_start + 1
if total_layers <= 0:
return
previous_ranges = {
node.node_id: (node.shard_start, node.shard_end, node.quantization)
for node in managed_nodes
}
for node in managed_nodes:
node.shard_start = None
node.shard_end = None
managed_nodes.sort(
key=lambda node: (
-node.benchmark_tokens_per_sec,
-_node_layer_capacity(node, preset),
node.node_id,
)
)
base_nodes = [node for node in model_nodes if not node.managed_assignment]
coverage = _coverage_map(base_nodes, required_start, required_end)
gaps = _coverage_gaps(coverage)
if not gaps:
gaps = [(required_start, required_end)]
eligible_nodes = [
node for node in managed_nodes
if _node_layer_capacity(node, preset) > 0
]
node_index = 0
for gap_start, gap_end in gaps:
cursor = gap_start
while cursor <= gap_end and node_index < len(eligible_nodes):
node = eligible_nodes[node_index]
remaining_layers = gap_end - cursor + 1
remaining_nodes_after = len(eligible_nodes) - node_index - 1
capacity = min(
_node_layer_capacity(node, preset),
total_layers,
max(1, remaining_layers - remaining_nodes_after),
)
if capacity <= 0:
node_index += 1
continue
quantization = _node_quantization(node, preset)
node.quantization = quantization
node.shard_start = cursor
node.shard_end = min(gap_end, cursor + capacity - 1)
cursor = node.shard_end + 1
node_index += 1
for node in managed_nodes:
previous_start, previous_end, previous_quantization = previous_ranges[node.node_id]
current_range = (node.shard_start, node.shard_end, node.quantization)
if node.shard_start is None or node.shard_end is None or current_range == previous_ranges[node.node_id]:
continue
if previous_start is not None and previous_end is not None:
node.pending_directives.append(
_drop_directive(
node,
hf_repo,
previous_start,
previous_end,
previous_quantization or _node_quantization(node, preset),
)
)
node.pending_directives.append(
_load_directive(
node,
hf_repo,
node.shard_start,
node.shard_end,
node.quantization or _node_quantization(node, preset),
)
)
def _rebalance_all_locked(server: "_TrackerHTTPServer") -> None:
for model in list(server.model_presets):
_rebalance_model_locked(server, model)
for hf_repo in sorted({node.hf_repo for node in server.registry.values() if node.hf_repo}):
_rebalance_hf_model_locked(server, hf_repo)
def _registration_ban_error(contracts: Any | None, wallet_address: str | None) -> str | None:
@@ -1095,7 +1195,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
covered_up_to = rs - 1
route_hops: list[dict] = []
for rn in route_nodes:
route_hops.append({"endpoint": rn.endpoint, "start_layer": covered_up_to + 1})
hop: dict = {"endpoint": rn.endpoint, "start_layer": covered_up_to + 1}
if rn.relay_addr:
hop["relay_addr"] = rn.relay_addr
route_hops.append(hop)
covered_up_to = rn.shard_end if rn.shard_end is not None else covered_up_to
# Strip the first-shard node we're about to proxy to — it's already handling the request.
downstream_hops = [h for h in route_hops if h["endpoint"].rstrip("/") != node.endpoint.rstrip("/")]
@@ -1332,6 +1435,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
return
tracker_mode = bool(body.get("tracker_mode", False))
managed_assignment = bool(body.get("managed_assignment", False))
hf_repo = body.get("hf_repo")
if hf_repo is not None and not isinstance(hf_repo, str):
self._send_json(400, {"error": "hf_repo must be a string"})
@@ -1371,7 +1475,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
quantizations=quantizations,
benchmark_tokens_per_sec=benchmark_tokens_per_sec,
quantization=quantization,
managed_assignment=not explicit_shard,
managed_assignment=managed_assignment or not explicit_shard,
tracker_mode=tracker_mode,
hf_repo=hf_repo,
num_layers=num_layers,
@@ -1395,7 +1499,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
)
server.registry[node_id] = entry
if entry.managed_assignment:
_rebalance_model_locked(server, model)
if entry.hf_repo:
_rebalance_hf_model_locked(server, entry.hf_repo)
else:
_rebalance_model_locked(server, model)
assignment_directive = entry.pending_directives[-1] if entry.pending_directives else None
if assignment_directive is not None:
entry.pending_directives.clear()
@@ -1459,7 +1566,10 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
entry.uptime_seconds = float(body["uptime_seconds"])
if "status" in body and body["status"] in ("ready", "loading"):
entry.status = body["status"]
_rebalance_model_locked(server, entry.model or "stub-model")
if entry.hf_repo:
_rebalance_hf_model_locked(server, entry.hf_repo)
else:
_rebalance_model_locked(server, entry.model or "stub-model")
directives = list(entry.pending_directives)
entry.pending_directives.clear()
new_assignment = entry.pending_new_assignment

2
tests/test_mining_cli.py
View File

@@ -452,6 +452,7 @@ def test_default_cli_passes_advertise_host(monkeypatch):
"meshnet-node",
"--tracker", "http://192.168.0.179:8081",
"--advertise-host", "192.168.0.42",
"--debug",
"--no-tui",
])
@@ -465,3 +466,4 @@ def test_default_cli_passes_advertise_host(monkeypatch):
assert captured["tracker_url"] == "http://192.168.0.179:8081"
assert captured["advertise_host"] == "192.168.0.42"
assert captured["debug"] is True

101
tests/test_real_model_backend.py
View File

@@ -81,6 +81,38 @@ class _FakeFullBackend(_FakeBackend):
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"
@@ -198,6 +230,75 @@ def test_full_model_chat_completion_uses_generation_not_single_token_decode():
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")

80
tests/test_tracker_routing.py
View File

@@ -671,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()
@@ -1399,3 +1435,47 @@ def test_route_timeout_config_is_exposed_on_server():
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