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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user