From 61104d7d0a7041bbb749578b10e165b86d4b1084 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Wed, 6 May 2026 16:54:20 +0800 Subject: [PATCH] Add prefixed _log helper in dumper (#24506) --- python/sglang/srt/debug_utils/dumper.py | 44 +++++++++++----------- test/registered/debug_utils/test_dumper.py | 14 ++++++- 2 files changed, 36 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index c2715f19e..9a9020add 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -222,7 +222,7 @@ class _Dumper: self._ensure_exp_name() self._state.step += 1 - print(f"[Dumper] [{time.time()}] step={self._state.step}") + _log(f"step={self._state.step}") def dump( self, @@ -340,7 +340,7 @@ class _Dumper: from sglang.srt.debug_utils.source_patcher import apply_patches_from_config yaml_content: str = Path(config_path).read_text() - print(f"[source_patcher] loading config from {config_path}") + _log(f"[source_patcher] loading config from {config_path}") apply_patches_from_config( yaml_content, extra_imports=["from sglang.srt.debug_utils.dumper import dumper"], @@ -521,8 +521,8 @@ class _Dumper: path = Path(self._config.dir) / self._config.exp_name / full_filename if self._config.enable_output_console: - print( - f"[{tag}] [{rank}, {time.time()}] {path} " + _log( + f"[{tag}] {path} " f"type={type(value)} " f"shape={value.shape if isinstance(value, torch.Tensor) else None} " f"dtype={value.dtype if isinstance(value, torch.Tensor) else None} " @@ -570,7 +570,7 @@ class _Dumper: timeout_seconds=self._config.collective_timeout ) self.configure(exp_name=name) - print(f"[Dumper] Choose exp_name={name}") + _log(f"Choose exp_name={name}") # -------------------------------------- hook dumper ------------------------------------------ @@ -754,11 +754,11 @@ def _torch_save(value, path: str): if "not pickleable" in str(e): stripped = _strip_parameter(value) if stripped is not value: - print(f"[Dumper] Observe error={e} and try pickling .data") + _log(f"Observe error={e} and try pickling .data") return _torch_save(stripped, path) raise except Exception as e: - print(f"[Dumper] Observe error={e} when saving data, skip the tensor") + _log(f"Observe error={e} when saving data, skip the tensor") def _map_tensor(value, fn: Callable[[torch.Tensor], torch.Tensor]): @@ -792,11 +792,10 @@ def _collective_with_timeout(fn, operation_name: str, timeout_seconds: int = 60) def watchdog(): if not completed.wait(timeout=timeout_seconds): - print( - f"\n[Dumper] WARNING: '{operation_name}' has not completed after " + _log( + f"WARNING: '{operation_name}' has not completed after " f"{timeout_seconds}s. This usually means not all ranks are " - f"participating in this collective operation.\n", - flush=True, + f"participating in this collective operation." ) thread = threading.Thread(target=watchdog, daemon=True) @@ -845,7 +844,7 @@ def _cleanup_old_dumps(base_dir: Path, exp_name: Optional[str] = None) -> None: for entry in targets: shutil.rmtree(entry) - print(f"[Dumper] Cleaned up {entry}") + _log(f"Cleaned up {entry}") if dist.is_initialized(): _collective_with_timeout( @@ -868,6 +867,11 @@ def _get_world_size(): return 1 +def _log(msg: str) -> None: + """Print a log line tagged with the current rank and wall-clock time.""" + print(f"[Dumper, rank={_get_rank()}, t={time.time():.3f}] {msg}", flush=True) + + def _obj_to_dict(obj): if isinstance(obj, dict): return obj @@ -962,12 +966,10 @@ class _DumperHttpManager: self._rpc_broadcast = rpc_broadcast if http_port == "reuse": - print( - "[Dumper] Standalone HTTP server disabled, reusing existing ports" - ) + _log("Standalone HTTP server disabled, reusing existing ports") else: _start_http_server(prefix="/dumper/", target=self, http_port=http_port) - print(f"[Dumper] HTTP server started on port {http_port}") + _log(f"HTTP server started on port {http_port}") # ------------------------------- public --------------------------------- @@ -1008,7 +1010,7 @@ def _make_http_handler(*, prefix: str, target): method = self.path[len(prefix) :] try: req_body = self._get_request_body() - print(f"[Dumper#{_get_rank()}] HTTP {self.path} {req_body=}") + _log(f"HTTP {self.path} {req_body=}") result = target.handle_request(method=method, body=req_body) resp_body = json.dumps(result).encode() self.send_response(200) @@ -1053,13 +1055,13 @@ def _create_zmq_rpc_broadcast( result = getattr(handler, req["method"])(*req["args"], **req["kwargs"]) resp = {"result": result, "error": None} except Exception as e: - print(f"[Dumper.ZmqRpc] error inside handler: {e}") + _log(f"[ZmqRpc] error inside handler: {e}") resp = {"result": None, "error": str(e)} sock.send_pyobj(resp) thread = threading.Thread(target=serve_loop, daemon=True) thread.start() - print(f"[Dumper.ZmqRpc] rank={rank} server started at {local_addr}") + _log(f"[ZmqRpc] server started at {local_addr}") if dist.is_initialized(): all_addresses = [None] * world_size @@ -1070,7 +1072,7 @@ def _create_zmq_rpc_broadcast( ) else: all_addresses = [local_addr] - print(f"[Dumper.ZmqRpc] rank={rank} all_addresses={all_addresses}") + _log(f"[ZmqRpc] all_addresses={all_addresses}") if rank == 0: handles = [] @@ -1166,7 +1168,7 @@ def _get_local_ip_by_remote() -> Optional[str]: s.connect(("2001:4860:4860::8888", 80)) # Doesn't need to be reachable return s.getsockname()[0] except Exception: - print("Can not get local ip by remote") + _log("Can not get local ip by remote") return None diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index d555009e8..a180fe9a7 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -20,6 +20,7 @@ from sglang.srt.debug_utils.dumper import ( _Dumper, _format_tags, _get_default_exp_name, + _log, _map_tensor, _materialize_value, _MegatronPlugin, @@ -377,10 +378,21 @@ class TestTorchSave: _torch_save({"fn": lambda: None}, path) captured = capsys.readouterr() - assert "[Dumper] Observe error=" in captured.out + assert "[Dumper, rank=" in captured.out + assert "Observe error=" in captured.out assert "skip the tensor" in captured.out +class TestLog: + def test_log_format(self): + with _capture_stdout() as captured: + _log("hello") + out = captured.getvalue() + assert "hello" in out, out + assert "[Dumper, rank=" in out, out + assert ", t=" in out, out + + class TestCollectiveTimeout: def test_watchdog_fires_on_timeout(self): block_event = threading.Event()