less verbose node hop if not in debug mode
This commit is contained in:
@@ -38,6 +38,7 @@ class _TorchHTTPServer(http.server.HTTPServer):
|
||||
tracker_mode: bool = False,
|
||||
tracker_url: str | None = None,
|
||||
route_timeout: float = 30.0,
|
||||
debug: bool = False,
|
||||
):
|
||||
super().__init__(addr, handler)
|
||||
self.backend = backend
|
||||
@@ -46,6 +47,7 @@ class _TorchHTTPServer(http.server.HTTPServer):
|
||||
self.tracker_mode = tracker_mode
|
||||
self.tracker_url = tracker_url
|
||||
self.route_timeout = route_timeout
|
||||
self.debug = debug
|
||||
self.total_requests: int = 0
|
||||
self.failed_requests: int = 0
|
||||
self.queue_depth: int = 0
|
||||
@@ -406,7 +408,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
current_pos = pos_ids
|
||||
|
||||
for hop_index, (node_url, start_layer) in enumerate(route):
|
||||
print(f" [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}", flush=True)
|
||||
if server.debug:
|
||||
print(f" [node] pipeline hop {hop_index}: {node_url} start_layer={start_layer}", flush=True)
|
||||
headers: dict[str, str] = {
|
||||
"Content-Type": "application/octet-stream",
|
||||
"X-Meshnet-Wire": _WIRE_VERSION,
|
||||
@@ -440,7 +443,8 @@ class _TorchHandler(http.server.BaseHTTPRequestHandler):
|
||||
try:
|
||||
data = json.loads(resp_body)
|
||||
text = str(data.get("text", ""))
|
||||
print(f" [node] pipeline hop {hop_index} returned text={text!r}", flush=True)
|
||||
if server.debug:
|
||||
print(f" [node] pipeline hop {hop_index} returned text={text!r}", flush=True)
|
||||
return text
|
||||
except json.JSONDecodeError:
|
||||
return resp_body.decode("utf-8", errors="replace")
|
||||
@@ -611,6 +615,7 @@ class TorchNodeServer:
|
||||
tracker_mode: bool | None = None,
|
||||
tracker_url: str | None = None,
|
||||
route_timeout: float = 30.0,
|
||||
debug: bool = False,
|
||||
) -> None:
|
||||
self._host = host
|
||||
self._requested_port = port
|
||||
@@ -624,6 +629,7 @@ class TorchNodeServer:
|
||||
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._debug = debug
|
||||
self._server: _TorchHTTPServer | None = None
|
||||
self._thread: threading.Thread | None = None
|
||||
self.port: int | None = None
|
||||
@@ -666,6 +672,7 @@ class TorchNodeServer:
|
||||
self._tracker_mode,
|
||||
self._tracker_url,
|
||||
self._route_timeout,
|
||||
self._debug,
|
||||
)
|
||||
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