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

packages/node/meshnet_node/startup.py

@@ -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,
@@ -248,6 +300,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 = {
@@ -262,6 +322,7 @@ def run_startup(
             "quantization": quantization,
             "score": 1.0,
             "tracker_mode": (shard_start == 0),
+            **relay_fields,
         }
         tracker_node_id: str | None = None
         try:
@@ -327,6 +388,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],
@@ -339,6 +408,7 @@ def run_startup(
             "quantization": quantization,
             "score": 1.0,
             "tracker_mode": (assigned_shard_start == 0),
+            **relay_fields,
         }
         tracker_node_id = None
         try:
@@ -414,6 +484,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)
@@ -429,6 +507,7 @@ def run_startup(
                 "hardware_profile": hw,
                 "wallet_address": address,
                 "score": 1.0,
+                **relay_fields,
             },
         )
         node_id = str(reg_resp["node_id"])