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,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 ------------------------------------------