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

packages/tracker/meshnet_tracker/server.py

@@ -29,6 +29,8 @@ import sqlite3
 import threading
 import time
 import urllib.parse
+import urllib.request
+import uuid
 from typing import Any
 
 from .gossip import NodeGossip
@@ -507,6 +509,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,
@@ -701,6 +732,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,
         stats: "_StatsCollector | None" = None,
@@ -712,6 +744,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
         self.stats: _StatsCollector | None = stats
@@ -785,6 +818,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/"):
@@ -930,12 +965,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:
@@ -1059,6 +1120,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}",
@@ -1100,6 +1180,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()
@@ -1686,6 +1780,8 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
                     "node_id": e.node_id,
                     "endpoint": e.endpoint,
                     "start_layer": start,
+                    "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,
@@ -1746,6 +1842,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,
@@ -1780,6 +1878,7 @@ class TrackerServer:
         cluster_peers: list[str] | None = None,
         cluster_self_url: str | None = None,
         stats_db: str | None = None,
+        relay_url: str | None = None,
     ) -> None:
         self._host = host
         self._requested_port = port
@@ -1792,6 +1891,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
@@ -1819,6 +1919,7 @@ class TrackerServer:
             self._model_presets,
             self._contracts,
             self._minimum_stake,
+            relay_url=self._relay_url,
             stats=self._stats,
         )
         self.port = self._server.server_address[1]