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

routing

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-08 18:48:50 +02:00
parent 29db25108f
commit e44abc910d
7 changed files with 191 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
packages/node/meshnet_node/startup.py
View File

@@ -829,6 +829,7 @@ def run_startup(
shard_label = _format_shard_label(shard_start, shard_end, total_layers)
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
endpoint = f"http://{public_host}:{actual_port}"
node.set_advertised_endpoint(endpoint)
local_base_url = f"http://127.0.0.1:{actual_port}"
relay_bridge, relay_fields = _start_relay_bridge_if_available(
tracker_url,
@@ -977,6 +978,7 @@ 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}"
node.set_advertised_endpoint(endpoint)
local_base_url = f"http://127.0.0.1:{actual_port}"
relay_bridge, relay_fields = _start_relay_bridge_if_available(
tracker_url,
@@ -1154,6 +1156,7 @@ def run_startup(
shard_label = f"{shard_label} (pinned)"
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
endpoint = f"http://{public_host}:{actual_port}"
node.set_advertised_endpoint(endpoint)
local_base_url = f"http://127.0.0.1:{actual_port}"
relay_bridge, relay_fields = _start_relay_bridge_if_available(
tracker_url,

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

@@ -31,6 +31,52 @@ from .server import (
)
def _endpoint_key(url: str) -> str:
"""Normalize http(s) endpoints for host:port comparison."""
parsed = urllib.parse.urlparse(url.rstrip("/"))
host = (parsed.hostname or "").lower()
if not host:
return url.rstrip("/").lower()
port = parsed.port
if port is None:
port = 443 if parsed.scheme == "https" else 80
return f"{host}:{port}"
def _own_endpoint_key(server: _TorchHTTPServer) -> str:
advertised = getattr(server, "advertised_endpoint", None)
if advertised:
return _endpoint_key(advertised)
host, port = server.server_address
return _endpoint_key(f"http://{host}:{port}")
def _clamp_downstream_hops(
hops: list[dict],
backend: TorchModelShard | None,
) -> list[dict]:
"""Ensure downstream start_layer continues after this shard's layers."""
if not hops or backend is None:
return hops
shard_end = getattr(backend, "shard_end", None)
if shard_end is None:
return hops
min_start = int(shard_end) + 1
clamped: list[dict] = []
for hop in hops:
adjusted = dict(hop)
if int(adjusted.get("start_layer", 0)) < min_start:
adjusted["start_layer"] = min_start
clamped.append(adjusted)
return clamped
def _format_downstream_route(hops: list[dict]) -> str:
return ", ".join(
f"{h['endpoint']}@{h.get('start_layer', 0)}" for h in hops
)
def _write_progress_line(state: list[bool], message: str, *, final: bool = False) -> None:
"""Rewrite one in-place progress line (\\r) or finish with a newline."""
if final:
@@ -104,6 +150,7 @@ class _TorchHTTPServer(http.server.HTTPServer):
self.route_timeout = route_timeout
self.debug = debug
self.max_loaded_shards = max(1, max_loaded_shards)
self.advertised_endpoint: str | None = None
self.total_requests: int = 0
self.failed_requests: int = 0
self.queue_depth: int = 0
@@ -572,6 +619,9 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
start_layer tells each downstream node which layer to begin from,
enabling correct execution when shard ranges overlap.
"""
server: _TorchHTTPServer = self.server # type: ignore[assignment]
active_backend = backend or server.backend
# Fast path: tracker pre-resolved the downstream route and injected it as a header.
injected = self.headers.get("X-Meshnet-Route")
if injected:
@@ -590,14 +640,16 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
hops.append(hop)
elif isinstance(item, str):
hops.append({"endpoint": item, "start_layer": 0})
print(f" [node] using injected downstream route: {[h['endpoint'] for h in hops]}", flush=True)
hops = _clamp_downstream_hops(hops, active_backend)
print(
f" [node] using injected downstream route: {_format_downstream_route(hops)}",
flush=True,
)
return hops
except (json.JSONDecodeError, TypeError, KeyError):
pass
# Slow path: query the tracker (direct node-to-node calls, or tracker didn't inject).
server: _TorchHTTPServer = self.server # type: ignore[assignment]
active_backend = backend or server.backend
if server.tracker_url is None:
return []
route_model = getattr(active_backend, "model_id", None) or model
@@ -605,23 +657,31 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
route_resp = json.loads(r.read())
own_port = server.server_address[1]
own_key = _own_endpoint_key(server)
nodes_info = route_resp.get("nodes", [])
hops = []
covered_up_to: int | None = None
hops: list[dict] = []
passed_self = False
for node_info in nodes_info:
ep = node_info.get("endpoint", "")
if ep.rstrip("/").endswith(f":{own_port}"):
covered_up_to = node_info.get("shard_end")
if not ep:
continue
if covered_up_to is None:
covered_up_to = (node_info.get("shard_start") or 1) - 1
hop = {"endpoint": ep, "start_layer": covered_up_to + 1}
if _endpoint_key(ep) == own_key:
passed_self = True
continue
if not passed_self:
continue
hop = {
"endpoint": ep,
"start_layer": int(node_info.get("start_layer", 0)),
}
if node_info.get("relay_addr"):
hop["relay_addr"] = str(node_info["relay_addr"])
hops.append(hop)
covered_up_to = node_info.get("shard_end", covered_up_to)
print(f" [node] tracker downstream route: {[h['endpoint'] for h in hops]}", flush=True)
hops = _clamp_downstream_hops(hops, active_backend)
print(
f" [node] tracker downstream route: {_format_downstream_route(hops)}",
flush=True,
)
return hops
except Exception as exc:
print(f" [node] WARNING: route lookup failed for {route_model!r}: {exc}", flush=True)
@@ -1039,6 +1099,11 @@ class TorchNodeServer:
self._thread.start()
return self.port
def set_advertised_endpoint(self, endpoint: str) -> None:
"""Set the LAN-facing endpoint used for route self-detection."""
if self._server is not None:
self._server.advertised_endpoint = endpoint
def stop(self) -> None:
if self._server is None:
return