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

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,6 +308,7 @@ 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
 
+        if self._enable_output_console:
             print(
                 f"[{tag}] [{rank}, {time.time()}] {path} "
                 f"type={type(value)} "
@@ -304,16 +319,22 @@ class _Dumper:
                 f"sample_value={get_truncated_value(value)}"
             )
 
-        if self._enable_write_file and save:
+        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),
+            }
+
+            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)
-            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))
 
     @cached_property
@@ -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 ------------------------------------------
 
 

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": {...}}."""