Add prefixed _log helper in dumper (#24506)

This commit is contained in:
fzyzcjy
2026-05-06 16:54:20 +08:00
committed by GitHub
parent 23172224b9
commit 61104d7d0a
2 changed files with 36 additions and 22 deletions
+23 -21
View File
@@ -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
+13 -1
View File
@@ -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()