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

Add relay-backed public node registration

Browse Source
This commit is contained in:
Dobromir Popov
2026-06-30 15:01:49 +02:00
parent a4373a6202
commit 9ca198ee1e
9 changed files with 516 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

148
packages/node/meshnet_node/relay_bridge.py Normal file
View File

@@ -0,0 +1,148 @@
"""Outbound relay bridge for NAT-safe node HTTP requests."""
from __future__ import annotations
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 "/")
body_text = payload.get("body") or ""
headers = payload.get("headers") if isinstance(payload.get("headers"), dict) else {}
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 {
"request_id": request_id,
"status": resp.status,
"headers": {"Content-Type": resp.headers.get("Content-Type", "application/json")},
"body": resp.read().decode(errors="replace"),
}
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

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

@@ -15,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
@@ -34,6 +35,57 @@ def _get_json(url: str, timeout: float = 10.0) -> dict:
return json.loads(r.read())
def _discover_relay_url(tracker_url: str) -> str | None:
try:
network_map = _get_json(f"{tracker_url}/v1/network/map", timeout=5.0)
except Exception:
return None
relay_url = network_map.get("relay_url")
return relay_url if isinstance(relay_url, str) and relay_url else None
def _start_relay_bridge_if_available(
tracker_url: str,
wallet_address: str,
local_base_url: str,
advertised_endpoint: str,
) -> tuple[RelayHttpBridge | None, dict]:
relay_url = _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,
@@ -212,6 +264,14 @@ def run_startup(
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,
)
_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 = {
@@ -226,6 +286,7 @@ def run_startup(
"quantization": quantization,
"score": 1.0,
"tracker_mode": (shard_start == 0),
**relay_fields,
}
tracker_node_id: str | None = None
try:
@@ -289,6 +350,14 @@ 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,
)
_attach_relay_bridge(node, relay_bridge)
auto_reg_payload = {
"endpoint": endpoint,
"model": assigned_hf_repo.split("/")[-1],
@@ -301,6 +370,7 @@ def run_startup(
"quantization": quantization,
"score": 1.0,
"tracker_mode": (assigned_shard_start == 0),
**relay_fields,
}
tracker_node_id = None
try:
@@ -376,6 +446,14 @@ 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,
)
_attach_relay_bridge(node, relay_bridge)
# 6. Register with tracker
print("Registering with tracker...", flush=True)
@@ -391,6 +469,7 @@ def run_startup(
"hardware_profile": hw,
"wallet_address": address,
"score": 1.0,
**relay_fields,
},
)
node_id = str(reg_resp["node_id"])

1
packages/node/pyproject.toml
View File

@@ -17,6 +17,7 @@ dependencies = [
"safetensors>=0.4",
"torch>=2.1",
"transformers>=4.39",
"websockets>=13",
"zstandard>=0.22",
]