From 261bca3c58c617def40aa226b92be137167397d4 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 20 Feb 2026 12:26:24 +0800 Subject: [PATCH] Support captured dump output and console output control in dumper (#19017) --- python/sglang/srt/debug_utils/dumper.py | 65 +++++++++---- test/registered/debug_utils/test_dumper.py | 101 +++++++++++++++++---- 2 files changed, 128 insertions(+), 38 deletions(-) diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index d63dd0094..ae98491fe 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -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 ------------------------------------------ diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index ed7dad97a..be0102ccb 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -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": {...}}."""