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

relay preparing for public internet

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-09 08:01:22 +02:00
parent 9ec4ca9ce1
commit 6ba8546c55
8 changed files with 344 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
packages/node/meshnet_node/torch_server.py
View File

@@ -169,6 +169,35 @@ class _TorchHTTPServer(http.server.HTTPServer):
self.queue_depth: int = 0
self._stats_lock = threading.Lock()
self._active_requests: dict[str, dict[str, Any]] = {}
self._decode_log: dict[str, dict[str, float]] = {}
def note_decode_step(
self, session: str, now: float | None = None,
) -> int | None:
"""Count one decode forward; return the cumulative step count when a
log line is due (first step of a session, then every 5s), else None."""
if now is None:
now = time.monotonic()
with self._stats_lock:
rec = self._decode_log.get(session)
if rec is None:
if len(self._decode_log) >= 64:
stale = [
sid for sid, r in self._decode_log.items()
if now - r["seen"] > 600.0
]
for sid in stale:
del self._decode_log[sid]
while len(self._decode_log) >= 64:
self._decode_log.pop(next(iter(self._decode_log)))
self._decode_log[session] = {"steps": 1.0, "logged": now, "seen": now}
return 1
rec["steps"] += 1
rec["seen"] = now
if now - rec["logged"] >= 5.0:
rec["logged"] = now
return int(rec["steps"])
return None
def snapshot_current_requests(self) -> list[dict[str, Any]]:
"""In-flight request snapshots for tracker heartbeats."""
@@ -346,23 +375,34 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
hop_index = int(self.headers.get("X-Meshnet-Hop-Index", "0"))
if hop_index > 0:
server.received_activations = True
if chunk_index_value == 0:
shard_start = getattr(server.backend, "shard_start", "?")
shard_end = getattr(server.backend, "shard_end", "?")
print(
f" [node] forward hop={hop_index} "
f"layers={shard_start}-{shard_end} "
f"session={session[:8]}{self._request_log_suffix()}",
flush=True,
)
start_layer_header = self.headers.get("X-Meshnet-Start-Layer")
start_layer = int(start_layer_header) if start_layer_header else None
# Session KV-cache protocol: prefill establishes per-session state on
# this node's layer range; decode reuses it. Absent header = legacy
# stateless call (also the signature fake backends implement).
cache_mode = self.headers.get("X-Meshnet-Cache")
if chunk_index_value == 0:
shard_start = getattr(server.backend, "shard_start", "?")
shard_end = getattr(server.backend, "shard_end", "?")
if cache_mode == "decode":
# One decode forward arrives per generated token — log a
# periodic per-session summary instead of one line per token.
steps = server.note_decode_step(session)
if steps is not None:
print(
f" [node] decoding layers={shard_start}-{shard_end} "
f"session={session[:8]} steps={steps}"
f"{self._request_log_suffix()}",
flush=True,
)
else:
print(
f" [node] forward hop={hop_index} "
f"layers={shard_start}-{shard_end} "
f"session={session[:8]}{self._request_log_suffix()}",
flush=True,
)
start_layer_header = self.headers.get("X-Meshnet-Start-Layer")
start_layer = int(start_layer_header) if start_layer_header else None
forward_kwargs: dict[str, object] = {}
if cache_mode in ("prefill", "decode"):
past_len_header = self.headers.get("X-Meshnet-Past-Len")
@@ -585,12 +625,14 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
session_id = str(uuid.uuid4())
use_kv = bool(getattr(backend, "supports_kv_cache", False))
# EOS detection by id must work on the stateless path too: the tail
# returns token_id regardless of caching, and EOS usually decodes to
# "" (skip_special_tokens), so the text comparison never fires.
eos_ids: set[int] = set()
if use_kv:
try:
eos_ids = set(backend.eos_token_ids())
except Exception:
eos_ids = set()
try:
eos_ids = set(backend.eos_token_ids())
except Exception:
eos_ids = set()
stream_emit = None
if stream:

26
packages/tracker/meshnet_tracker/cli.py
View File

@@ -103,6 +103,28 @@ def main() -> None:
default=None,
help="Public ws(s):// relay URL advertised to nodes, for example wss://ai.neuron.d-popov.com/ws",
)
common.add_argument(
"--embedded-relay",
action="store_true",
help="Run the relay WebSocket server in this tracker process (still uses meshnet_relay.RelayServer)",
)
common.add_argument(
"--relay-host",
default="127.0.0.1",
help="Bind address for --embedded-relay (default: 127.0.0.1; use 0.0.0.0 only when exposing the relay port directly)",
)
common.add_argument(
"--relay-port",
type=int,
default=8765,
help="Bind port for --embedded-relay (default: 8765)",
)
common.add_argument(
"--relay-max-peers",
type=int,
default=500,
help="Maximum WebSocket peers accepted by --embedded-relay",
)
common.add_argument(
"--billing-db",
default=DEFAULT_BILLING_DB_PATH,
@@ -378,6 +400,10 @@ def main() -> None:
cluster_self_url=args.self_url,
stats_db=getattr(args, "stats_db", None),
relay_url=relay_url,
embedded_relay=args.embedded_relay,
embedded_relay_host=args.relay_host,
embedded_relay_port=args.relay_port,
embedded_relay_max_peers=args.relay_max_peers,
enable_billing=not args.no_billing,
billing_db=None if args.no_billing else args.billing_db,
max_charge_per_request=args.max_charge_per_request,

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

@@ -101,6 +101,12 @@ def derive_relay_url_from_public_tracker_url(url: str | None) -> str | None:
return f"wss://{parsed.netloc}/ws"
def _local_relay_url(host: str, port: int) -> str:
"""Return a node-consumable ws:// URL for an embedded relay bind address."""
advertise_host = "127.0.0.1" if host in {"", "0.0.0.0", "::"} else host
return f"ws://{advertise_host}:{port}/ws"
def _load_model_presets_from_data() -> dict[str, dict]:
"""Load recommended model presets from package JSON data."""
try:
@@ -2459,6 +2465,7 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
hf_pricing_log: "HfPricingLog | None" = None,
models_dir: Path | None = None,
route_stats: "RouteStatsStore | None" = None,
relay_status: dict | None = None,
) -> None:
super().__init__(addr, handler)
self.registry = registry
@@ -2496,6 +2503,9 @@ class _TrackerHTTPServer(socketserver.ThreadingMixIn, http.server.HTTPServer):
self.active_proxies_lock = threading.Lock()
self.route_stats: RouteStatsStore = route_stats or RouteStatsStore()
self.route_rng = random.Random()
self.relay_status = dict(
relay_status or {"mode": "external" if self.relay_url else "off"}
)
class _TrackerHandler(http.server.BaseHTTPRequestHandler):
@@ -2930,6 +2940,7 @@ class _TrackerHandler(http.server.BaseHTTPRequestHandler):
self._send_json(200, {
"relay_url": server.relay_url,
"relay": dict(server.relay_status),
"pool": _pool_summary(nodes),
"memory_pool": memory_pool,
"recommended_models": [
@@ -5779,6 +5790,10 @@ class TrackerServer:
cluster_self_url: str | None = None,
stats_db: str | None = None,
relay_url: str | None = None,
embedded_relay: bool = False,
embedded_relay_host: str = "127.0.0.1",
embedded_relay_port: int = 8765,
embedded_relay_max_peers: int = 500,
billing: BillingLedger | None = None,
enable_billing: bool = False,
billing_db: str | None = None,
@@ -5821,6 +5836,12 @@ class TrackerServer:
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._embedded_relay_enabled = bool(embedded_relay)
self._embedded_relay_host = embedded_relay_host
self._embedded_relay_port = int(embedded_relay_port)
self._embedded_relay_max_peers = int(embedded_relay_max_peers)
self._embedded_relay: Any | None = None
self._embedded_relay_actual_port: int | None = None
self._registry: dict[str, _NodeEntry] = {}
self._lock = threading.Lock()
self._server: _TrackerHTTPServer | None = None
@@ -5924,13 +5945,59 @@ class TrackerServer:
self._route_stats = RouteStatsStore(routing_config)
self.port: int | None = None
def _start_embedded_relay(self) -> dict:
"""Start the shared RelayServer class in-process for tracker+relay deployments."""
if not self._embedded_relay_enabled:
return {"mode": "external" if self._relay_url else "off"}
if self._embedded_relay is not None:
return self._embedded_relay_status()
try:
from meshnet_relay.server import RelayServer # type: ignore[import-not-found]
except Exception as exc: # pragma: no cover - only hit in broken installs
raise RuntimeError(
"embedded relay requires the meshnet-relay package to be installed"
) from exc
relay = RelayServer(
host=self._embedded_relay_host,
port=self._embedded_relay_port,
max_peers=self._embedded_relay_max_peers,
)
actual_port = relay.start()
self._embedded_relay = relay
self._embedded_relay_actual_port = actual_port
return self._embedded_relay_status()
def _embedded_relay_status(self) -> dict:
advertised_url = (
self._relay_url
or derive_relay_url_from_public_tracker_url(self._cluster_self_url)
or _local_relay_url(self._embedded_relay_host, self._embedded_relay_actual_port or self._embedded_relay_port)
)
return {
"mode": "embedded",
"url": advertised_url,
"bind_host": self._embedded_relay_host,
"bind_port": self._embedded_relay_actual_port or self._embedded_relay_port,
"max_peers": self._embedded_relay_max_peers,
}
def start(self) -> int:
if self._server is not None:
raise RuntimeError("TrackerServer is already running")
relay_status = self._start_embedded_relay()
effective_relay_url = (
self._relay_url
or derive_relay_url_from_public_tracker_url(self._cluster_self_url)
or (relay_status.get("url") if relay_status.get("mode") == "embedded" else None)
)
http_relay_status = (
{**relay_status, "url": effective_relay_url}
if relay_status.get("mode") == "embedded"
else {"mode": "external", "url": effective_relay_url}
if effective_relay_url
else {"mode": "off"}
)
# Start HTTP server first so we know our port
@@ -5960,6 +6027,7 @@ class TrackerServer:
hf_pricing_log=self._hf_pricing_log,
models_dir=self._models_dir,
route_stats=self._route_stats,
relay_status=http_relay_status,
)
self.port = self._server.server_address[1]
@@ -6275,6 +6343,11 @@ class TrackerServer:
self._raft.stop()
if self._gossip is not None:
self._gossip.stop()
embedded_relay = self._embedded_relay
self._embedded_relay = None
self._embedded_relay_actual_port = None
if embedded_relay is not None:
embedded_relay.stop()
if self._server is None:
return
self._rebalance_stop.set()