diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 4bdc82257..07d42f50c 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/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__( - 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 - 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 + def __init__(self, *, config: _DumperConfig): + self._config = config + + self._http_server_handled = not config.enable_http_server + self._cleanup_previous_handled = not config.cleanup_previous - # States - self._partial_name = partial_name self._dump_index = 0 self._forward_pass_id = 0 - self._global_ctx = {} - self._override_enable = None + self._global_ctx: dict = {} 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: - self._partial_name = _get_partial_name( - timeout_seconds=self._collective_timeout - ) - print(f"[Dumper] Choose partial_name={self._partial_name}") + if self._config.partial_name is None: + name = _get_partial_name(timeout_seconds=self._config.collective_timeout) + self.configure(partial_name=name) + print(f"[Dumper] Choose partial_name={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): - self._override_enable = value + def configure(self, **kwargs) -> None: + 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_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", @@ -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: - self._pending_cleanup = False - _cleanup_old_dumps(self._base_dir) + if not self._cleanup_previous_handled: + self._cleanup_previous_handled = True + _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 ------------------------------------------ diff --git a/test/registered/debug_utils/test_dump_comparator.py b/test/registered/debug_utils/test_dump_comparator.py index 5b73dd48f..a302cb61f 100644 --- a/test/registered/debug_utils/test_dump_comparator.py +++ b/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( - "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, + 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._config.partial_name}") args = Namespace( baseline_path=str(dump_dirs[0]), diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 5f65e075d..69db0ca40 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/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( - 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: @@ -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