feat(us-021): --route-timeout CLI flag for node tracker route lookup

Default 30s replaces the hardcoded 5s. Wired through TorchNodeServer →
_TorchHTTPServer → _get_remaining_route. Available on both the legacy
`start` subcommand and the default wizard path.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Dobromir Popov
2026-06-30 13:06:00 +03:00
parent c691e8d5d1
commit e1ba120912
3 changed files with 17 additions and 1 deletions

View File

@@ -37,6 +37,7 @@ class _TorchHTTPServer(http.server.HTTPServer):
backend: TorchModelShard,
tracker_mode: bool = False,
tracker_url: str | None = None,
route_timeout: float = 30.0,
):
super().__init__(addr, handler)
self.backend = backend
@@ -44,6 +45,7 @@ class _TorchHTTPServer(http.server.HTTPServer):
self.forward_chunk_count = 0
self.tracker_mode = tracker_mode
self.tracker_url = tracker_url
self.route_timeout = route_timeout
class _TorchHandler(http.server.BaseHTTPRequestHandler):
@@ -318,7 +320,7 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
route_model = getattr(server.backend, "model_id", None) or model
try:
url = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(route_model)}"
with urllib.request.urlopen(url, timeout=5.0) as r:
with urllib.request.urlopen(url, timeout=server.route_timeout) as r:
route_resp = json.loads(r.read())
route = route_resp.get("route", [])
own_port = server.server_address[1]
@@ -558,6 +560,7 @@ class TorchNodeServer:
backend: TorchModelShard | None = None,
tracker_mode: bool | None = None,
tracker_url: str | None = None,
route_timeout: float = 30.0,
) -> None:
self._host = host
self._requested_port = port
@@ -570,6 +573,7 @@ class TorchNodeServer:
# Auto-detect tracker mode: enabled when shard_start == 0 or explicitly set
self._tracker_mode = tracker_mode if tracker_mode is not None else (shard_start == 0)
self._tracker_url = tracker_url
self._route_timeout = route_timeout
self._server: _TorchHTTPServer | None = None
self._thread: threading.Thread | None = None
self.port: int | None = None
@@ -595,6 +599,7 @@ class TorchNodeServer:
self._backend,
self._tracker_mode,
self._tracker_url,
self._route_timeout,
)
self.port = self._server.server_address[1]
self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)