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

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

Browse Source
This commit is contained in:
Dobromir Popov
2026-06-30 17:27:12 +03:00
parent 8ce5a74d5e c587e02c2c
commit 1da088926a
13 changed files with 764 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
packages/node/meshnet_node/cli.py
View File

@@ -3,6 +3,7 @@
from __future__ import annotations
import argparse
import socket
import sys
import time
from pathlib import Path
@@ -26,6 +27,7 @@ def _run_node(cfg: dict) -> None:
wallet_path=Path(cfg["wallet_path"]) if cfg.get("wallet_path") else None,
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
host=cfg.get("host", "0.0.0.0"),
advertise_host=cfg.get("advertise_host"),
route_timeout=float(cfg.get("route_timeout", 30.0)),
vram_mb_override=cfg.get("vram_mb_override"),
)
@@ -50,6 +52,19 @@ def _run_node(cfg: dict) -> None:
)
def _first_available_port(host: str, start: int = 7000, attempts: int = 100) -> int:
"""Return the first TCP port bindable on host, starting at start."""
bind_host = "" if host == "0.0.0.0" else host
for port in range(start, start + attempts):
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
try:
sock.bind((bind_host, port))
except OSError:
continue
return port
raise OSError(f"no available TCP port in range {start}-{start + attempts - 1}")
def _cmd_default(args) -> int:
"""No subcommand: wizard if no config, else start with saved config."""
from .config import load_config, save_config, merge_cli_overrides
@@ -88,6 +103,8 @@ def _cmd_default(args) -> int:
overrides["port"] = args.port
if args.host:
overrides["host"] = args.host
if args.advertise_host:
overrides["advertise_host"] = args.advertise_host
if args.route_timeout != 30.0:
overrides["route_timeout"] = args.route_timeout
if getattr(args, "memory", None) is not None:
@@ -148,8 +165,12 @@ def _cmd_start(args) -> int:
# Build a transient config from flags (don't write to disk)
cfg = dict(DEFAULTS)
cfg["tracker_url"] = args.tracker
cfg["port"] = args.port
cfg["model_name"] = args.model
cfg["port"] = args.port if args.port is not None else _first_available_port(args.host)
if args.model_id is None and "/" in args.model:
cfg["model_hf_repo"] = args.model
cfg["model_name"] = args.model.split("/")[-1]
else:
cfg["model_name"] = args.model
cfg["quantization"] = args.quantization
cfg["host"] = args.host
if args.model_id:
@@ -220,6 +241,7 @@ def main() -> None:
parser.add_argument("--shard-end", type=int, metavar="N", help="Pin shard end layer")
parser.add_argument("--port", type=int, metavar="N", help="Port to listen on")
parser.add_argument("--host", metavar="ADDR", help="Interface to bind (default 0.0.0.0)")
parser.add_argument("--advertise-host", metavar="ADDR", help="Host/IP advertised to the tracker")
parser.add_argument("--route-timeout", type=float, metavar="SEC", default=30.0,
help="Seconds to wait for tracker route lookup (default 30)")
parser.add_argument("--memory", type=int, metavar="MB", default=None,
@@ -240,7 +262,7 @@ def main() -> None:
# start subcommand (legacy / backward-compat)
start_cmd = subparsers.add_parser("start", help="Start node (legacy flags)")
start_cmd.add_argument("--tracker", default="http://localhost:8080")
start_cmd.add_argument("--port", type=int, default=7000)
start_cmd.add_argument("--port", type=int)
start_cmd.add_argument("--model", default="stub-model")
start_cmd.add_argument("--model-id", help="HuggingFace repo ID")
start_cmd.add_argument("--shard-start", type=int)

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,
@@ -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"])

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",
]