refactor(dumper): extract _DumperConfig frozen dataclass

Move all 13 _Dumper init parameters into a @DataClass(frozen=True)
_DumperConfig. _Dumper.__init__ now takes a single config parameter.

- Runtime mutations (enable via HTTP, lazy partial_name) use
  dataclasses.replace() to swap the frozen config
- _DumperConfig.from_env() centralizes env var parsing with defaults
  matching field defaults (verified by new UT)
- Rename _pending_cleanup to _cleanup_previous_handled for consistency
  with _http_server_handled

Author: fzyzcjy

python/sglang/srt/debug_utils/dumper.py

@@ -6,6 +6,7 @@
 import time
 from contextlib import contextmanager
 from copy import deepcopy
+from dataclasses import dataclass, replace
 from functools import cached_property
 from http.server import BaseHTTPRequestHandler, HTTPServer
 from pathlib import Path
@@ -17,6 +18,47 @@
 # -------------------------------------- dumper core ------------------------------------------
 
 
+@dataclass(frozen=True)
+class _DumperConfig:
+    enable: bool = False
+    filter: Optional[str] = None
+    base_dir: Path = Path("/tmp")
+    enable_output_file: bool = True
+    enable_output_console: bool = True
+    enable_value: bool = True
+    enable_grad: bool = False
+    enable_model_value: bool = True
+    enable_model_grad: bool = True
+    partial_name: Optional[str] = None
+    enable_http_server: bool = True
+    cleanup_previous: bool = False
+    collective_timeout: int = 60
+
+    @classmethod
+    def from_env(cls) -> "_DumperConfig":
+        return cls(
+            enable=get_bool_env_var("SGLANG_DUMPER_ENABLE", "0"),
+            filter=_get_str_env_var("SGLANG_DUMPER_FILTER"),
+            base_dir=Path(_get_str_env_var("SGLANG_DUMPER_DIR", "/tmp")),
+            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(
+                "SGLANG_DUMPER_ENABLE_MODEL_VALUE", "1"
+            ),
+            enable_model_grad=get_bool_env_var("SGLANG_DUMPER_ENABLE_MODEL_GRAD", "1"),
+            partial_name=_get_str_env_var("SGLANG_DUMPER_PARTIAL_NAME"),
+            enable_http_server=get_bool_env_var(
+                "SGLANG_ENABLE_DUMPER_HTTP_SERVER", "1"
+            ),
+            cleanup_previous=get_bool_env_var("SGLANG_DUMPER_CLEANUP_PREVIOUS", "0"),
+            collective_timeout=60,
+        )
+
+
 class _Dumper:
     """Utility to dump tensors, which can be useful when comparison checking models.
 
@@ -44,75 +86,30 @@ class _Dumper:
     Related: `sglang.srt.debug_utils.dump_comparator` for dump comparison
     """
 
-    def __init__(
-        self,
-        *,
-        enable: bool,
-        base_dir: Path,
-        filter: Optional[str] = None,
-        enable_output_file: bool = True,
-        enable_output_console: bool = True,
-        enable_value: bool = True,
-        enable_grad: bool = False,
-        enable_model_value: bool = True,
-        enable_model_grad: bool = True,
-        partial_name: Optional[str] = None,
-        enable_http_server: bool = True,
-        cleanup_previous: bool = False,
-        collective_timeout: int = 60,
-    ):
-        # Config
-        self._enable = enable
+    def __init__(self, *, config: _DumperConfig):
         # 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_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
-        self._enable_model_grad = enable_model_grad
-        self._collective_timeout = collective_timeout
-
-        # States
-        self._partial_name = partial_name
+        self._config = config
+
+        self._http_server_handled = not config.enable_http_server
+        self._cleanup_previous_handled = not config.cleanup_previous
+
         self._dump_index = 0
         self._forward_pass_id = 0
-        self._global_ctx = {}
-        self._override_enable = None
+        self._global_ctx: dict = {}
+        self._override_enable: Optional[bool] = None
         self._captured_output_data: Optional[dict] = None
-        self._http_server_handled = not enable_http_server
-        self._pending_cleanup = cleanup_previous
 
     @classmethod
     def from_env(cls) -> "_Dumper":
-        return cls(
-            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_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(
-                "SGLANG_DUMPER_ENABLE_MODEL_VALUE", "1"
-            ),
-            enable_model_grad=get_bool_env_var("SGLANG_DUMPER_ENABLE_MODEL_GRAD", "1"),
-            partial_name=_get_str_env_var("SGLANG_DUMPER_PARTIAL_NAME"),
-            enable_http_server=get_bool_env_var(
-                "SGLANG_ENABLE_DUMPER_HTTP_SERVER", "1"
-            ),
-            cleanup_previous=get_bool_env_var("SGLANG_DUMPER_CLEANUP_PREVIOUS", "0"),
-            collective_timeout=60,
-        )
+        return cls(config=_DumperConfig.from_env())
 
     def on_forward_pass_start(self):
         """This should be called on all ranks."""
 
         # Even if SGLANG_DUMPER_ENABLE=0, users may want to use HTTP endpoint to enable it
         self._ensure_http_server()
 
-        if not self._enable:
+        if not self._config.enable:
             return
 
         # Users may want to `dump` only on some ranks, thus determine name here
@@ -127,14 +124,15 @@ def _ensure_http_server(self):
         if self._http_server_handled:
             return
         self._http_server_handled = True
-        _start_maybe_http_server(self, timeout_seconds=self._collective_timeout)
+        _start_maybe_http_server(self, timeout_seconds=self._config.collective_timeout)
 
     def _ensure_partial_name(self):
-        if self._partial_name is None:
-            self._partial_name = _get_partial_name(
-                timeout_seconds=self._collective_timeout
+        if self._config.partial_name is None:
+            name = _get_partial_name(
+                timeout_seconds=self._config.collective_timeout
             )
-            print(f"[Dumper] Choose partial_name={self._partial_name}")
+            self._config = replace(self._config, partial_name=name)
+            print(f"[Dumper] Choose partial_name={name}")
 
     def set_ctx(self, **kwargs):
         """
@@ -172,9 +170,9 @@ def dump(self, name: str, value, save: bool = True, **kwargs) -> None:
             value=value,
             extra_kwargs=kwargs,
             save=save,
-            enable_value=self._enable_value,
+            enable_value=self._config.enable_value,
             enable_curr_grad=False,
-            enable_future_grad=self._enable_grad,
+            enable_future_grad=self._config.enable_grad,
             value_tag="Dumper.Value",
             grad_tag="Dumper.Grad",
         )
@@ -192,8 +190,8 @@ def dump_model(
                 value=param,
                 extra_kwargs=kwargs,
                 save=save,
-                enable_value=self._enable_model_value,
-                enable_curr_grad=self._enable_model_grad,
+                enable_value=self._config.enable_model_value,
+                enable_curr_grad=self._config.enable_model_grad,
                 enable_future_grad=False,
                 value_tag="Dumper.ParamValue",
                 grad_tag="Dumper.ParamGrad",
@@ -214,9 +212,9 @@ def _dump_inner(
     ) -> None:
         self._ensure_http_server()
 
-        if not (self._enable and (self._override_enable is not False)):
+        if not (self._config.enable and (self._override_enable is not False)):
             return
-        if (f := self._filter) is not None and re.search(f, name) is None:
+        if (f := self._config.filter) is not None and re.search(f, name) is None:
             return
         if not (enable_value or enable_curr_grad or enable_future_grad):
             return
@@ -306,9 +304,9 @@ def _dump_single(
             **self._global_ctx,
         )
         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
+        path = self._config.base_dir / f"sglang_dump_{self._config.partial_name}" / full_filename
 
-        if self._enable_output_console:
+        if self._config.enable_output_console:
             print(
                 f"[{tag}] [{rank}, {time.time()}] {path} "
                 f"type={type(value)} "
@@ -320,7 +318,7 @@ def _dump_single(
             )
 
         capturing = self._captured_output_data is not None
-        if save and (self._enable_output_file or capturing):
+        if save and (self._config.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),
@@ -330,9 +328,9 @@ def _dump_single(
                 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)
+                if not self._cleanup_previous_handled:
+                    self._cleanup_previous_handled = True
+                    _cleanup_old_dumps(self._config.base_dir)
 
                 path.parent.mkdir(parents=True, exist_ok=True)
                 _torch_save(output_data, str(path))
@@ -603,7 +601,7 @@ def __init__(self, dumper):
 
     def set_enable(self, enable: bool):
         print(f"[DumperRpcHandler] set_enable {enable=}")
-        self._dumper._enable = enable
+        self._dumper._config = replace(self._dumper._config, enable=enable)
 
 
 # -------------------------------------- zmq rpc ------------------------------------------

test/registered/debug_utils/test_dump_comparator.py

@@ -89,23 +89,22 @@ def test_main(self):
         from argparse import Namespace
 
         from sglang.srt.debug_utils.dump_comparator import main
-        from sglang.srt.debug_utils.dumper import _Dumper
+        from sglang.srt.debug_utils.dumper import _Dumper, _DumperConfig
 
         with tempfile.TemporaryDirectory() as d1, tempfile.TemporaryDirectory() as d2:
             baseline_tensor = torch.randn(10, 10)
             target_tensor = baseline_tensor + torch.randn(10, 10) * 0.01
 
             dump_dirs = []
             for d, tensor in [(d1, baseline_tensor), (d2, target_tensor)]:
-                with _with_env("SGLANG_DUMPER_DIR", d), _with_env(
-                    "SGLANG_DUMPER_SERVER_PORT", "-1"
-                ):
-                    dumper = _Dumper()
-                    dumper.on_forward_pass_start()
-                    dumper.dump("tensor_a", tensor)
-                    dumper.on_forward_pass_start()
-                    dumper.dump("tensor_b", tensor * 2)
-                    dump_dirs.append(Path(d) / f"sglang_dump_{dumper._partial_name}")
+                dumper = _Dumper(config=_DumperConfig(
+                    enable=True, base_dir=Path(d), enable_http_server=False,
+                ))
+                dumper.on_forward_pass_start()
+                dumper.dump("tensor_a", tensor)
+                dumper.on_forward_pass_start()
+                dumper.dump("tensor_b", tensor * 2)
+                dump_dirs.append(Path(d) / f"sglang_dump_{dumper._config.partial_name}")
 
             args = Namespace(
                 baseline_path=str(dump_dirs[0]),

test/registered/debug_utils/test_dumper.py

@@ -15,6 +15,7 @@
     _collect_sglang_parallel_info,
     _collective_with_timeout,
     _Dumper,
+    _DumperConfig,
     _materialize_value,
     _obj_to_dict,
     _torch_save,
@@ -40,6 +41,11 @@ def _capture_stdout():
         sys.stdout = old_stdout
 
 
+class TestDumperConfig:
+    def test_from_env_defaults_match_dataclass_defaults(self):
+        assert _DumperConfig.from_env() == _DumperConfig()
+
+
 class TestDumperPureFunctions:
     def test_get_truncated_value(self):
         assert get_truncated_value(None) is None
@@ -175,11 +181,9 @@ def test_collective_timeout(self):
     @staticmethod
     def _test_collective_timeout_func(rank):
         dumper = _Dumper(
-            enable=True,
-            base_dir=Path("/tmp"),
-            partial_name=None,
-            enable_http_server=False,
-            collective_timeout=3,
+            config=_DumperConfig(
+                enable=True, collective_timeout=3, enable_http_server=False,
+            ),
         )
 
         with _capture_stdout() as captured:
@@ -202,7 +206,7 @@ def test_http_enable(self):
     def _test_http_func(rank):
         from sglang.srt.debug_utils.dumper import dumper
 
-        assert not dumper._enable
+        assert not dumper._config.enable
         dumper.on_forward_pass_start()
 
         for enable in [True, False]:
@@ -213,7 +217,7 @@ def _test_http_func(rank):
                     "http://localhost:40000/dumper", json={"enable": enable}
                 ).raise_for_status()
             dist.barrier()
-            assert dumper._enable == enable
+            assert dumper._config.enable == enable
 
     def test_file_content_correctness(self, tmp_path):
         with temp_set_env(
@@ -407,13 +411,11 @@ def test_dict_format_with_context(self, tmp_path):
 
 def _make_test_dumper(tmp_path: Path, **overrides) -> _Dumper:
     """Create a _Dumper for CPU testing without HTTP server or distributed."""
-    defaults: dict = dict(
-        enable=True,
-        base_dir=tmp_path,
-        partial_name="test",
-        enable_http_server=False,
+    config = _DumperConfig(
+        enable=True, base_dir=tmp_path, partial_name="test",
+        enable_http_server=False, **overrides,
     )
-    d = _Dumper(**{**defaults, **overrides})
+    d = _Dumper(config=config)
     d.on_forward_pass_start()
     return d