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