chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Support captured dump output and console output control in dumper (#19017)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-20 12:26:24 +08:00
committed by GitHub
parent fc1500adc6
commit 261bca3c58
2 changed files with 128 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
python/sglang/srt/debug_utils/dumper.py
View File

@@ -4,6 +4,7 @@ import re
import socket
import threading
import time
from contextlib import contextmanager
from copy import deepcopy
from functools import cached_property
from http.server import BaseHTTPRequestHandler, HTTPServer
@@ -49,7 +50,8 @@ class _Dumper:
enable: bool,
base_dir: Path,
filter: Optional[str] = None,
enable_write_file: bool = True,
enable_output_file: bool = True,
enable_output_console: bool = True,
enable_value: bool = True,
enable_grad: bool = False,
enable_model_value: bool = True,
@@ -64,7 +66,8 @@ class _Dumper:
# TODO (1) support filtering kv instead of name only (2) allow HTTP req change it
self._filter = filter
self._base_dir = base_dir
self._enable_write_file = enable_write_file
self._enable_output_file = enable_output_file
self._enable_output_console = enable_output_console
self._enable_value = enable_value
self._enable_grad = enable_grad
self._enable_model_value = enable_model_value
@@ -77,6 +80,7 @@ class _Dumper:
self._forward_pass_id = 0
self._global_ctx = {}
self._override_enable = None
self._captured_output_data: Optional[dict] = None
self._http_server_handled = not enable_http_server
self._pending_cleanup = cleanup_previous
@@ -86,7 +90,8 @@ class _Dumper:
enable=get_bool_env_var("SGLANG_DUMPER_ENABLE", "0"),
base_dir=Path(_get_str_env_var("SGLANG_DUMPER_DIR", "/tmp")),
filter=_get_str_env_var("SGLANG_DUMPER_FILTER"),
enable_write_file=get_bool_env_var("SGLANG_DUMPER_WRITE_FILE", "1"),
enable_output_file=get_bool_env_var("SGLANG_DUMPER_OUTPUT_FILE", "1"),
enable_output_console=get_bool_env_var("SGLANG_DUMPER_OUTPUT_CONSOLE", "1"),
enable_value=get_bool_env_var("SGLANG_DUMPER_ENABLE_VALUE", "1"),
enable_grad=get_bool_env_var("SGLANG_DUMPER_ENABLE_GRAD", "0"),
enable_model_value=get_bool_env_var(
@@ -144,6 +149,15 @@ class _Dumper:
k: v for k, v in (self._global_ctx | kwargs).items() if v is not None
}
@contextmanager
def capture_output(self):
assert self._captured_output_data is None
self._captured_output_data = {}
try:
yield self._captured_output_data
finally:
self._captured_output_data = None
def override_enable(self, value: bool):
self._override_enable = value
@@ -294,27 +308,34 @@ class _Dumper:
full_filename = "___".join(f"{k}={v}" for k, v in full_kwargs.items()) + ".pt"
path = self._base_dir / f"sglang_dump_{self._partial_name}" / full_filename
print(
f"[{tag}] [{rank}, {time.time()}] {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} "
f"device={value.device if isinstance(value, torch.Tensor) else None} "
f"id={id(value)} "
f"sample_value={get_truncated_value(value)}"
)
if self._enable_output_console:
print(
f"[{tag}] [{rank}, {time.time()}] {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} "
f"device={value.device if isinstance(value, torch.Tensor) else None} "
f"id={id(value)} "
f"sample_value={get_truncated_value(value)}"
)
if self._enable_write_file and save:
if self._pending_cleanup:
self._pending_cleanup = False
_cleanup_old_dumps(self._base_dir)
path.parent.mkdir(parents=True, exist_ok=True)
capturing = self._captured_output_data is not None
if save and (self._enable_output_file or capturing):
output_data = {
"value": value.data if isinstance(value, torch.nn.Parameter) else value,
"meta": dict(**full_kwargs, **self._static_meta),
}
_torch_save(output_data, str(path))
if capturing:
output_data["value"] = _deepcopy_or_clone(output_data["value"])
self._captured_output_data[name] = output_data
else:
if self._pending_cleanup:
self._pending_cleanup = False
_cleanup_old_dumps(self._base_dir)
path.parent.mkdir(parents=True, exist_ok=True)
_torch_save(output_data, str(path))
@cached_property
def _static_meta(self) -> dict:
@@ -420,6 +441,12 @@ def _materialize_value(value):
return value
def _deepcopy_or_clone(x):
if isinstance(x, torch.Tensor):
return x.clone()
return deepcopy(x)
# -------------------------------------- static meta ------------------------------------------

101
test/registered/debug_utils/test_dumper.py
View File

@@ -269,25 +269,6 @@ class TestDumperFileWriteControl:
not_exist=["skip_this", "not_keep_this"],
)
def test_write_disabled(self, tmp_path):
with temp_set_env(
allow_sglang=True,
SGLANG_DUMPER_ENABLE="1",
SGLANG_DUMPER_DIR=str(tmp_path),
SGLANG_DUMPER_WRITE_FILE="0",
):
run_distributed_test(self._test_write_disabled_func, tmpdir=str(tmp_path))
@staticmethod
def _test_write_disabled_func(rank, tmpdir):
from sglang.srt.debug_utils.dumper import dumper
dumper.on_forward_pass_start()
dumper.dump("no_write", torch.randn(5, device=f"cuda:{rank}"))
dist.barrier()
assert len(_get_filenames(tmpdir)) == 0
def test_save_false(self, tmp_path):
with temp_set_env(
allow_sglang=True,
@@ -307,6 +288,88 @@ class TestDumperFileWriteControl:
assert len(_get_filenames(tmpdir)) == 0
class TestOutputControl:
def test_file_enabled_by_default(self, tmp_path):
d = _make_test_dumper(tmp_path)
d.dump("file_on", torch.randn(3, 3))
_assert_files(_get_filenames(tmp_path), exist=["file_on"])
def test_file_disabled(self, tmp_path, capsys):
d = _make_test_dumper(tmp_path, enable_output_file=False)
d.dump("file_off", torch.randn(3, 3))
assert len(_get_filenames(tmp_path)) == 0
assert "file_off" in capsys.readouterr().out
def test_console_enabled_by_default(self, tmp_path, capsys):
d = _make_test_dumper(tmp_path)
d.dump("console_on", torch.randn(3, 3))
captured = capsys.readouterr()
assert "[Dumper.Value]" in captured.out
assert "console_on" in captured.out
def test_console_disabled(self, tmp_path, capsys):
d = _make_test_dumper(tmp_path, enable_output_console=False)
d.dump("console_off", torch.randn(3, 3))
assert "console_off" not in capsys.readouterr().out
_assert_files(_get_filenames(tmp_path), exist=["console_off"])
def test_capture_output_basic(self, tmp_path):
d = _make_test_dumper(tmp_path)
tensor = torch.randn(4, 4)
with d.capture_output() as captured:
d.dump("cap_basic", tensor)
assert "cap_basic" in captured
assert set(captured["cap_basic"].keys()) == {"value", "meta"}
assert torch.equal(captured["cap_basic"]["value"], tensor)
assert captured["cap_basic"]["meta"]["name"] == "cap_basic"
def test_capture_output_no_file(self, tmp_path):
d = _make_test_dumper(tmp_path)
with d.capture_output() as captured:
d.dump("cap_no_file", torch.randn(3, 3))
assert "cap_no_file" in captured
assert len(_get_filenames(tmp_path)) == 0
def test_capture_output_multiple(self, tmp_path):
d = _make_test_dumper(tmp_path)
with d.capture_output() as captured:
d.dump("first", torch.randn(2, 2))
d.dump("second", torch.randn(3, 3))
assert set(captured.keys()) == {"first", "second"}
assert captured["first"]["value"].shape == (2, 2)
assert captured["second"]["value"].shape == (3, 3)
def test_capture_output_value_cloned(self, tmp_path):
d = _make_test_dumper(tmp_path)
tensor = torch.zeros(3, 3)
with d.capture_output() as captured:
d.dump("clone_check", tensor)
tensor.fill_(999.0)
assert torch.equal(captured["clone_check"]["value"], torch.zeros(3, 3))
def test_capture_output_respects_filter(self, tmp_path):
d = _make_test_dumper(tmp_path, filter="^keep")
with d.capture_output() as captured:
d.dump("keep_this", torch.randn(3, 3))
d.dump("skip_this", torch.randn(3, 3))
assert "keep_this" in captured
assert "skip_this" not in captured
class TestDumpDictFormat:
"""Verify that dump files use the dict output format: {"value": ..., "meta": {...}}."""