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

@@ -26,6 +26,7 @@ def _run_node(cfg: dict) -> None:
wallet_path=Path(cfg["wallet_path"]) if cfg.get("wallet_path") else 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, cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
host=cfg.get("host", "0.0.0.0"), host=cfg.get("host", "0.0.0.0"),
route_timeout=float(cfg.get("route_timeout", 30.0)),
) )
except Exception as exc: except Exception as exc:
print(f"\nERROR: {exc}", file=sys.stderr, flush=True) print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
@@ -86,6 +87,8 @@ def _cmd_default(args) -> int:
overrides["port"] = args.port overrides["port"] = args.port
if args.host: if args.host:
overrides["host"] = args.host overrides["host"] = args.host
if args.route_timeout != 30.0:
overrides["route_timeout"] = args.route_timeout
if overrides: if overrides:
cfg = merge_cli_overrides(cfg, **overrides) cfg = merge_cli_overrides(cfg, **overrides)
@@ -173,6 +176,7 @@ def _cmd_start(args) -> int:
cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None, cache_dir=Path(cfg["download_dir"]) if cfg.get("download_dir") else None,
host=cfg["host"], host=cfg["host"],
advertise_host=getattr(args, "advertise_host", None), advertise_host=getattr(args, "advertise_host", None),
route_timeout=getattr(args, "route_timeout", 30.0),
) )
except Exception as exc: except Exception as exc:
print(f"ERROR: {exc}", file=sys.stderr, flush=True) print(f"ERROR: {exc}", file=sys.stderr, flush=True)
@@ -212,6 +216,8 @@ def main() -> None:
parser.add_argument("--shard-end", type=int, metavar="N", help="Pin shard end layer") 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("--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("--host", metavar="ADDR", help="Interface to bind (default 0.0.0.0)")
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("--no-tui", action="store_true", help="Plain-text output (no rich dashboard)") parser.add_argument("--no-tui", action="store_true", help="Plain-text output (no rich dashboard)")
parser.add_argument("--compact", action="store_true", help="Single-line status output") parser.add_argument("--compact", action="store_true", help="Single-line status output")
parser.add_argument("--reset-config", action="store_true", help="Re-run wizard even if config exists") parser.add_argument("--reset-config", action="store_true", help="Re-run wizard even if config exists")
@@ -240,6 +246,8 @@ def main() -> None:
start_cmd.add_argument("--tracker-url", default=None) start_cmd.add_argument("--tracker-url", default=None)
start_cmd.add_argument("--wallet") start_cmd.add_argument("--wallet")
start_cmd.add_argument("--download-dir") start_cmd.add_argument("--download-dir")
start_cmd.add_argument("--route-timeout", type=float, default=30.0,
help="Seconds to wait for tracker route lookup (default 30)")
args = parser.parse_args() args = parser.parse_args()

View File

@@ -111,6 +111,7 @@ def run_startup(
host: str = "127.0.0.1", host: str = "127.0.0.1",
advertise_host: str | None = None, advertise_host: str | None = None,
contracts: Any | None = None, contracts: Any | None = None,
route_timeout: float = 30.0,
) -> StubNodeServer | TorchNodeServer: ) -> StubNodeServer | TorchNodeServer:
"""Execute the full startup sequence and return a running node server. """Execute the full startup sequence and return a running node server.
@@ -202,6 +203,7 @@ def run_startup(
shard_end=shard_end, shard_end=shard_end,
quantization=quantization, quantization=quantization,
tracker_url=tracker_url, tracker_url=tracker_url,
route_timeout=route_timeout,
) )
actual_port = node.start() actual_port = node.start()
total_layers = getattr(getattr(node, "backend", None), "total_layers", None) total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
@@ -285,6 +287,7 @@ def run_startup(
shard_end=assigned_shard_end, shard_end=assigned_shard_end,
quantization=quantization, quantization=quantization,
tracker_url=tracker_url, tracker_url=tracker_url,
route_timeout=route_timeout,
) )
actual_port = node.start() actual_port = node.start()
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host) public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)

View File

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