From 815715110216f54980b6aac7c1dc5653ffb4cf17 Mon Sep 17 00:00:00 2001 From: Dobromir Popov Date: Tue, 30 Jun 2026 18:02:25 +0300 Subject: [PATCH] feat(us-016): outbound relay client for NAT/internet pipeline hops 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 --- packages/node/meshnet_node/relay_bridge.py | 26 ++++- packages/node/meshnet_node/torch_server.py | 121 ++++++++++++++++----- packages/tracker/meshnet_tracker/server.py | 5 +- 3 files changed, 119 insertions(+), 33 deletions(-) diff --git a/packages/node/meshnet_node/relay_bridge.py b/packages/node/meshnet_node/relay_bridge.py index d7b63ea..4127160 100644 --- a/packages/node/meshnet_node/relay_bridge.py +++ b/packages/node/meshnet_node/relay_bridge.py @@ -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, diff --git a/packages/node/meshnet_node/torch_server.py b/packages/node/meshnet_node/torch_server.py index 036f17b..16e7ffd 100644 --- a/packages/node/meshnet_node/torch_server.py +++ b/packages/node/meshnet_node/torch_server.py @@ -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: diff --git a/packages/tracker/meshnet_tracker/server.py b/packages/tracker/meshnet_tracker/server.py index 7b92fd2..6d3a81b 100644 --- a/packages/tracker/meshnet_tracker/server.py +++ b/packages/tracker/meshnet_tracker/server.py @@ -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("/")]