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:
@@ -26,6 +26,7 @@ def _run_node(cfg: dict) -> 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,
|
||||
host=cfg.get("host", "0.0.0.0"),
|
||||
route_timeout=float(cfg.get("route_timeout", 30.0)),
|
||||
)
|
||||
except Exception as exc:
|
||||
print(f"\nERROR: {exc}", file=sys.stderr, flush=True)
|
||||
@@ -86,6 +87,8 @@ def _cmd_default(args) -> int:
|
||||
overrides["port"] = args.port
|
||||
if args.host:
|
||||
overrides["host"] = args.host
|
||||
if args.route_timeout != 30.0:
|
||||
overrides["route_timeout"] = args.route_timeout
|
||||
|
||||
if 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,
|
||||
host=cfg["host"],
|
||||
advertise_host=getattr(args, "advertise_host", None),
|
||||
route_timeout=getattr(args, "route_timeout", 30.0),
|
||||
)
|
||||
except Exception as exc:
|
||||
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("--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("--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("--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")
|
||||
@@ -240,6 +246,8 @@ def main() -> None:
|
||||
start_cmd.add_argument("--tracker-url", default=None)
|
||||
start_cmd.add_argument("--wallet")
|
||||
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()
|
||||
|
||||
|
||||
@@ -111,6 +111,7 @@ def run_startup(
|
||||
host: str = "127.0.0.1",
|
||||
advertise_host: str | None = None,
|
||||
contracts: Any | None = None,
|
||||
route_timeout: float = 30.0,
|
||||
) -> StubNodeServer | TorchNodeServer:
|
||||
"""Execute the full startup sequence and return a running node server.
|
||||
|
||||
@@ -202,6 +203,7 @@ def run_startup(
|
||||
shard_end=shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
route_timeout=route_timeout,
|
||||
)
|
||||
actual_port = node.start()
|
||||
total_layers = getattr(getattr(node, "backend", None), "total_layers", None)
|
||||
@@ -285,6 +287,7 @@ def run_startup(
|
||||
shard_end=assigned_shard_end,
|
||||
quantization=quantization,
|
||||
tracker_url=tracker_url,
|
||||
route_timeout=route_timeout,
|
||||
)
|
||||
actual_port = node.start()
|
||||
public_host = advertise_host or (socket.getfqdn() if host == "0.0.0.0" else host)
|
||||
|
||||
@@ -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