Enhance configure and env parsing in dumper (#19034)

python/sglang/srt/debug_utils/dumper.py

@@ -4,16 +4,106 @@ import re
 import socket
 import threading
 import time
+from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from copy import deepcopy
+from dataclasses import dataclass, fields, replace
 from functools import cached_property
 from http.server import BaseHTTPRequestHandler, HTTPServer
 from pathlib import Path
-from typing import List, Optional
+from typing import List, Optional, Self, get_args, get_type_hints
 
 import torch
 import torch.distributed as dist
 
+# -------------------------------------- frozen config base ------------------------------------------
+
+
+@dataclass(frozen=True)
+class _FrozenConfig(ABC):
+    def __post_init__(self) -> None:
+        self._verify_types()
+
+    def _verify_types(self) -> None:
+        hints = get_type_hints(type(self))
+        cls_name = type(self).__name__
+        for f in fields(self):
+            value = getattr(self, f.name)
+            if value is None:
+                continue
+            expected = self._unwrap_type(hints[f.name])
+            if not isinstance(value, expected):
+                raise TypeError(
+                    f"{cls_name}.{f.name}: expected {expected.__name__}, "
+                    f"got {type(value).__name__}"
+                )
+
+    @classmethod
+    @abstractmethod
+    def _env_prefix(cls) -> str: ...
+
+    @classmethod
+    def _env_name(cls, field_name: str) -> str:
+        return f"{cls._env_prefix()}{field_name.upper()}"
+
+    @classmethod
+    def from_env(cls) -> Self:
+        return cls(
+            **{
+                f.name: cls._parse_env_field(cls._env_name(f.name), f.default)
+                for f in fields(cls)
+            }
+        )
+
+    def with_defaults(self, **kwargs) -> Self:
+        cls = type(self)
+        actual = {
+            key: value
+            for key, value in kwargs.items()
+            if os.getenv(cls._env_name(key)) is None
+        }
+        return replace(self, **actual) if actual else self
+
+    @staticmethod
+    def _unwrap_type(hint) -> type:
+        args = get_args(hint)
+        if args:
+            return next(a for a in args if a is not type(None))
+        return hint
+
+    @staticmethod
+    def _parse_env_field(env_name: str, default):
+        raw = os.getenv(env_name)
+        if raw is None or not raw.strip():
+            return default
+        if isinstance(default, bool):
+            return raw.lower() in ("true", "1")
+        if isinstance(default, int):
+            return int(raw)
+        return raw
+
+
+@dataclass(frozen=True)
+class _DumperConfig(_FrozenConfig):
+    enable: bool = False
+    filter: Optional[str] = None
+    dir: str = "/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 _env_prefix(cls) -> str:
+        return "SGLANG_DUMPER_"
+
+
 # -------------------------------------- dumper core ------------------------------------------
 
 
@@ -44,66 +134,16 @@ class _Dumper:
     Related: `sglang.srt.debug_utils.dump_comparator` for dump comparison
     """
 
-    def __init__(
+    def __init__(self, *, config: _DumperConfig):
-        self,
+        self._config = config
-        *,
+
-        enable: bool,
+        self._http_server_handled = not config.enable_http_server
-        base_dir: Path,
+        self._cleanup_previous_handled = not config.cleanup_previous
-        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
-        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._dump_index = 0
         self._forward_pass_id = 0
-        self._global_ctx = {}
+        self._global_ctx: dict = {}
-        self._override_enable = 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,
-        )
 
     def on_forward_pass_start(self):
         """This should be called on all ranks."""
@@ -111,7 +151,7 @@ class _Dumper:
         # 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
@@ -126,20 +166,19 @@ class _Dumper:
         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:
+        if self._config.partial_name is None:
-            self._partial_name = _get_partial_name(
+            name = _get_partial_name(timeout_seconds=self._config.collective_timeout)
-                timeout_seconds=self._collective_timeout
+            self.configure(partial_name=name)
-            )
+            print(f"[Dumper] Choose partial_name={name}")
-            print(f"[Dumper] Choose partial_name={self._partial_name}")
 
     def set_ctx(self, **kwargs):
         """
         Example:
 
-        dumper.override_enable(self.layer_id <= 3)
+        dumper.configure_default(filter='layer_id=[0-3]')
         dumper.set_ctx(layer_id=self.layer_id)
         ...
         dumper.set_ctx(layer_id=None)
@@ -157,8 +196,11 @@ class _Dumper:
         finally:
             self._captured_output_data = None
 
-    def override_enable(self, value: bool):
+    def configure(self, **kwargs) -> None:
-        self._override_enable = value
+        self._config = replace(self._config, **kwargs)
+
+    def configure_default(self, **kwargs) -> None:
+        self._config = self._config.with_defaults(**kwargs)
 
     def dump_dict(self, name_prefix, data, save: bool = True, **kwargs):
         data = _obj_to_dict(data)
@@ -171,9 +213,9 @@ class _Dumper:
             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",
         )
@@ -191,8 +233,8 @@ class _Dumper:
                 value=param,
                 extra_kwargs=kwargs,
                 save=save,
-                enable_value=self._enable_model_value,
+                enable_value=self._config.enable_model_value,
-                enable_curr_grad=self._enable_model_grad,
+                enable_curr_grad=self._config.enable_model_grad,
                 enable_future_grad=False,
                 value_tag="Dumper.ParamValue",
                 grad_tag="Dumper.ParamGrad",
@@ -213,11 +255,13 @@ class _Dumper:
     ) -> None:
         self._ensure_http_server()
 
-        if not (self._enable and (self._override_enable is not False)):
+        if not self._config.enable:
             return
 
         tags = dict(name=name, **extra_kwargs, **self._global_ctx)
-        if (f := self._filter) is not None and re.search(f, _format_tags(tags)) is None:
+        if (f := self._config.filter) is not None and re.search(
+            f, _format_tags(tags)
+        ) is None:
             return
 
         if not (enable_value or enable_curr_grad or enable_future_grad):
@@ -307,9 +351,13 @@ class _Dumper:
             **tags,
         )
         full_filename = _format_tags(full_kwargs) + ".pt"
-        path = self._base_dir / f"sglang_dump_{self._partial_name}" / full_filename
+        path = (
+            Path(self._config.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)} "
@@ -321,7 +369,7 @@ class _Dumper:
             )
 
         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),
@@ -331,9 +379,9 @@ class _Dumper:
                 output_data["value"] = _deepcopy_or_clone(output_data["value"])
                 self._captured_output_data[tags["name"]] = output_data
             else:
-                if self._pending_cleanup:
+                if not self._cleanup_previous_handled:
-                    self._pending_cleanup = False
+                    self._cleanup_previous_handled = True
-                    _cleanup_old_dumps(self._base_dir)
+                    _cleanup_old_dumps(Path(self._config.dir))
 
                 path.parent.mkdir(parents=True, exist_ok=True)
                 _torch_save(output_data, str(path))
@@ -608,7 +656,7 @@ class _DumperRpcHandler:
 
     def set_enable(self, enable: bool):
         print(f"[DumperRpcHandler] set_enable {enable=}")
-        self._dumper._enable = enable
+        self._dumper.configure(enable=enable)
 
 
 # -------------------------------------- zmq rpc ------------------------------------------
@@ -694,20 +742,6 @@ class _ZmqRpcHandle:
 # --------------------------------- copied code (avoid dependency) --------------------------------------
 
 
-def get_bool_env_var(name: str, default: str = "false") -> bool:
-    value = os.getenv(name, default)
-    value = value.lower()
-    truthy_values = ("true", "1")
-    return value in truthy_values
-
-
-def _get_str_env_var(name: str, default: Optional[str] = None) -> Optional[str]:
-    value = os.getenv(name)
-    if value is None or not value.strip():
-        return default
-    return value
-
-
 def get_int_env_var(name: str, default: int = 0) -> int:
     value = os.getenv(name)
     if value is None or not value.strip():
@@ -750,7 +784,7 @@ def _get_local_ip_by_remote() -> Optional[str]:
 # -------------------------------------- singleton ------------------------------------------
 
 
-dumper = _Dumper.from_env()
+dumper = _Dumper(config=_DumperConfig.from_env())
 
 
 # -------------------------------------- other utility functions ------------------------------------------

test/registered/debug_utils/test_dump_comparator.py

@@ -89,7 +89,7 @@ class TestEndToEnd(CustomTestCase):
         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)
@@ -97,15 +97,18 @@ class TestEndToEnd(CustomTestCase):
 
             dump_dirs = []
             for d, tensor in [(d1, baseline_tensor), (d2, target_tensor)]:
-                with _with_env("SGLANG_DUMPER_DIR", d), _with_env(
+                dumper = _Dumper(
-                    "SGLANG_DUMPER_SERVER_PORT", "-1"
+                    config=_DumperConfig(
-                ):
+                        enable=True,
-                    dumper = _Dumper()
+                        dir=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._partial_name}")
+                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 @@ from sglang.srt.debug_utils.dumper import (
     _collect_sglang_parallel_info,
     _collective_with_timeout,
     _Dumper,
+    _DumperConfig,
     _format_tags,
     _materialize_value,
     _obj_to_dict,
@@ -41,6 +42,57 @@ def _capture_stdout():
         sys.stdout = old_stdout
 
 
+class TestDumperConfig:
+    def test_from_env_defaults_match_dataclass_defaults(self):
+        assert _DumperConfig.from_env() == _DumperConfig()
+
+    def test_from_env_bool(self):
+        with temp_set_env(allow_sglang=True, SGLANG_DUMPER_ENABLE="1"):
+            assert _DumperConfig.from_env().enable is True
+        with temp_set_env(allow_sglang=True, SGLANG_DUMPER_ENABLE="false"):
+            assert _DumperConfig.from_env().enable is False
+
+    def test_from_env_str(self):
+        with temp_set_env(allow_sglang=True, SGLANG_DUMPER_FILTER="layer_id=0"):
+            assert _DumperConfig.from_env().filter == "layer_id=0"
+
+    def test_from_env_dir(self):
+        with temp_set_env(allow_sglang=True, SGLANG_DUMPER_DIR="/my/dir"):
+            assert _DumperConfig.from_env().dir == "/my/dir"
+
+    def test_from_env_int(self):
+        with temp_set_env(allow_sglang=True, SGLANG_DUMPER_COLLECTIVE_TIMEOUT="120"):
+            assert _DumperConfig.from_env().collective_timeout == 120
+
+    def test_configure_overrides(self):
+        d = _make_test_dumper("/tmp")
+        d.configure(enable=False)
+        assert d._config.enable is False
+        d.configure(enable=True)
+        assert d._config.enable is True
+
+    def test_type_validation(self):
+        with pytest.raises(TypeError, match="enable.*expected bool.*got str"):
+            _DumperConfig(enable="yes")
+        with pytest.raises(
+            TypeError, match="collective_timeout.*expected int.*got str"
+        ):
+            _DumperConfig(collective_timeout="abc")
+        with pytest.raises(TypeError, match="filter.*expected str.*got int"):
+            _DumperConfig(filter=123)
+
+    def test_configure_default_skips_when_env_set(self):
+        with temp_set_env(allow_sglang=True, SGLANG_DUMPER_FILTER="from_env"):
+            d = _Dumper(config=_DumperConfig.from_env())
+            d.configure_default(filter="from_code")
+            assert d._config.filter == "from_env"
+
+    def test_configure_default_applies_when_no_env(self):
+        d = _Dumper(config=_DumperConfig.from_env())
+        d.configure_default(filter="from_code")
+        assert d._config.filter == "from_code"
+
+
 class TestDumperPureFunctions:
     def test_get_truncated_value(self):
         assert get_truncated_value(None) is None
@@ -154,9 +206,9 @@ class TestDumperDistributed:
         dumper.set_ctx(ctx_arg=None)
 
         dumper.on_forward_pass_start()
-        dumper.override_enable(False)
+        dumper.configure(filter=r"^$")
         dumper.dump("tensor_skip", tensor)
-        dumper.override_enable(True)
+        dumper.configure(filter=None)
 
         dumper.on_forward_pass_start()
         dumper.dump_dict("obj", {"a": torch.randn(3, device=f"cuda:{rank}"), "b": 42})
@@ -176,11 +228,11 @@ class TestDumperDistributed:
     @staticmethod
     def _test_collective_timeout_func(rank):
         dumper = _Dumper(
+            config=_DumperConfig(
                 enable=True,
-            base_dir=Path("/tmp"),
-            partial_name=None,
-            enable_http_server=False,
                 collective_timeout=3,
+                enable_http_server=False,
+            ),
         )
 
         with _capture_stdout() as captured:
@@ -203,7 +255,7 @@ class TestDumperDistributed:
     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]:
@@ -214,7 +266,7 @@ class TestDumperDistributed:
                     "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(
@@ -406,15 +458,16 @@ class TestDumpDictFormat:
         assert torch.equal(raw["value"], tensor)
 
 
-def _make_test_dumper(tmp_path: Path, **overrides) -> _Dumper:
+def _make_test_dumper(tmp_path, **overrides) -> _Dumper:
     """Create a _Dumper for CPU testing without HTTP server or distributed."""
-    defaults: dict = dict(
+    config = _DumperConfig(
         enable=True,
-        base_dir=tmp_path,
+        dir=str(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
 
@@ -487,7 +540,7 @@ class TestSaveValue:
 
 class TestStaticMetadata:
     def test_static_meta_contains_world_info(self):
-        dumper = _make_test_dumper(Path("/tmp"))
+        dumper = _make_test_dumper("/tmp")
         meta = dumper._static_meta
         assert "world_rank" in meta
         assert "world_size" in meta
@@ -495,7 +548,7 @@ class TestStaticMetadata:
         assert meta["world_size"] == 1
 
     def test_static_meta_caching(self):
-        dumper = _make_test_dumper(Path("/tmp"))
+        dumper = _make_test_dumper("/tmp")
         meta1 = dumper._static_meta
         meta2 = dumper._static_meta
         assert meta1 is meta2