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

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

@@ -28,6 +28,8 @@ import socketserver
import threading
import time
import urllib.parse
import urllib.request
import uuid
from typing import Any
from .gossip import NodeGossip
@@ -238,6 +240,35 @@ def _coverage_gaps(coverage: list[dict]) -> list[tuple[int, int]]:
]
def _relay_http_request(
relay_addr: str,
path: str,
body: bytes,
headers: dict[str, str],
timeout: float = 310.0,
) -> dict | None:
"""Send an HTTP-shaped request through a relay RPC WebSocket."""
try:
import websockets.sync.client as wsc # type: ignore[import]
request_id = str(uuid.uuid4())
with wsc.connect(relay_addr, open_timeout=10, close_timeout=5) as ws:
ws.send(json.dumps({
"request_id": request_id,
"method": "POST",
"path": path,
"headers": headers,
"body": body.decode(errors="replace"),
}))
raw = ws.recv(timeout=timeout)
response = json.loads(raw)
if response.get("request_id") not in {None, request_id}:
return None
return response
except Exception:
return None
def _nodes_and_bounds_for_model(
server: "_TrackerHTTPServer",
model: str,
@@ -432,6 +463,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
model_presets: dict,
contracts: Any | None,
minimum_stake: int,
relay_url: str | None = None,
raft: "RaftNode | None" = None,
gossip: "NodeGossip | None" = None,
) -> None:
@@ -442,6 +474,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
self.model_presets = model_presets
self.contracts = contracts
self.minimum_stake = minimum_stake
self.relay_url = relay_url.rstrip("/") if relay_url else None
self.raft = raft
self.gossip = gossip
@@ -511,6 +544,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self._handle_assign(parsed)
elif parsed.path == "/v1/network/assign":
self._handle_network_assign(parsed)
elif parsed.path == "/v1/network/map":
self._handle_network_map()
elif parsed.path == "/v1/models":
self._handle_models()
elif parsed.path.startswith("/v1/coverage/"):
@@ -653,12 +688,38 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
{
"node_id": node.node_id,
"endpoint": node.endpoint,
"relay_addr": node.relay_addr,
"peer_id": node.peer_id,
"benchmark_tokens_per_sec": node.benchmark_tokens_per_sec,
}
for node in tracker_nodes
],
})
def _handle_network_map(self) -> None:
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
with server.lock:
self._purge_expired_nodes()
nodes = list(server.registry.values())
self._send_json(200, {
"relay_url": server.relay_url,
"nodes": [
{
"node_id": node.node_id,
"endpoint": node.endpoint,
"relay_addr": node.relay_addr,
"peer_id": node.peer_id,
"model": node.model,
"hf_repo": node.hf_repo,
"shard_start": node.shard_start,
"shard_end": node.shard_end,
"tracker_mode": node.tracker_mode,
"last_heartbeat": node.last_heartbeat,
}
for node in nodes
],
})
# ---------------------------------------------------------------- OpenAI proxy
def _handle_proxy_chat(self) -> None:
@@ -771,6 +832,25 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
pass
return
except Exception as exc:
if node.relay_addr:
print(
f"[tracker] direct proxy failed {request_id}: {target_url}: {exc}; "
f"trying relay {node.relay_addr}",
flush=True,
)
relayed = _relay_http_request(
node.relay_addr,
path="/v1/chat/completions",
body=raw_body,
headers={
"Content-Type": "application/json",
"X-Meshnet-Route": downstream_urls,
**({"Authorization": auth} if auth else {}),
},
)
if relayed is not None:
self._send_relayed_response(relayed)
return
print(f"[tracker] proxy failed {request_id}: {target_url}: {exc}", flush=True)
self._send_json(503, {"error": {
"message": f"upstream node unreachable: {exc}",
@@ -812,6 +892,20 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
except BrokenPipeError:
pass
def _send_relayed_response(self, response: dict) -> None:
status = int(response.get("status", 503))
headers = response.get("headers") if isinstance(response.get("headers"), dict) else {}
body_text = response.get("body") or ""
body = body_text.encode() if isinstance(body_text, str) else bytes(body_text)
self.send_response(status)
self.send_header("Content-Type", headers.get("Content-Type", "application/json"))
self.send_header("Content-Length", str(len(body)))
self.end_headers()
try:
self.wfile.write(body)
except BrokenPipeError:
pass
def _handle_register(self):
server: _TrackerHTTPServer = self.server # type: ignore[assignment]
body = self._read_json_body()
@@ -1340,6 +1434,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
{
"node_id": e.node_id,
"endpoint": e.endpoint,
"relay_addr": e.relay_addr,
"peer_id": e.peer_id,
"wallet_address": e.wallet_address,
"shard_start": e.shard_start,
"shard_end": e.shard_end,
@@ -1400,6 +1496,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
{
"node_id": e.node_id,
"endpoint": e.endpoint,
"relay_addr": e.relay_addr,
"peer_id": e.peer_id,
"wallet_address": e.wallet_address,
"shard_start": e.shard_start,
"shard_end": e.shard_end,
@@ -1433,6 +1531,7 @@ class TrackerServer:
minimum_stake: int = 0,
cluster_peers: list[str] | None = None,
cluster_self_url: str | None = None,
relay_url: str | None = None,
) -> None:
self._host = host
self._requested_port = port
@@ -1445,6 +1544,7 @@ class TrackerServer:
self._minimum_stake = minimum_stake
self._cluster_peers: list[str] = list(cluster_peers) if cluster_peers else []
self._cluster_self_url: str | None = cluster_self_url
self._relay_url = relay_url
self._registry: dict[str, _NodeEntry] = {}
self._lock = threading.Lock()
self._server: _TrackerHTTPServer | None = None
@@ -1469,6 +1569,7 @@ class TrackerServer:
self._model_presets,
self._contracts,
self._minimum_stake,
self._relay_url,
)
self.port = self._server.server_address[1]