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

feat(us-016): outbound relay client for NAT/internet pipeline hops

Browse Source 
Nodes behind NAT (5G, WSL2, home routers) can now participate in
distributed pipeline inference over the internet via the relay server.

- torch_server: add module-level _relay_hop() that opens a WebSocket
  to relay.../rpc/{peer_id}, sends the binary activation with
  body_base64 encoding, and returns (status, headers, body)
- torch_server: _get_remaining_route returns list[dict] (was list[tuple])
  preserving relay_addr from injected X-Meshnet-Route header and
  from /v1/route slow-path node info
- torch_server: _run_downstream_pipeline dispatches via _relay_hop
  when hop has relay_addr; falls back to direct HTTP on relay error
- tracker server: downstream_hops dicts include relay_addr when node
  has one registered, so head node knows how to reach each peer
- relay_bridge: binary bodies (bfloat16 activations) use body_base64;
  response preserves all X-Meshnet-* headers

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-06-30 18:02:25 +03:00
parent 4f79b2d177
commit 8157151102
3 changed files with 119 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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 {}
url = f"{self.local_base_url}{path}"
# 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:
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,

95
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,6 +473,25 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
headers["X-Meshnet-Attn-Mask"] = current_attn
if current_pos:
headers["X-Meshnet-Position-Ids"] = current_pos
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,

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

@@ -1074,7 +1074,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("/")]