diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 19a67a928..21a39f01b 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -1,6 +1,7 @@ import functools import json import os +import random import re import socket import threading @@ -17,11 +18,11 @@ from typing import Any, List, Literal, Optional, Union, get_args, get_type_hints import torch import torch.distributed as dist -# -------------------------------------- frozen config base ------------------------------------------ +# -------------------------------------- config base ------------------------------------------ @dataclass(frozen=True) -class _FrozenConfig(ABC): +class _BaseConfig(ABC): def __post_init__(self) -> None: self._verify_types() @@ -48,7 +49,7 @@ class _FrozenConfig(ABC): return f"{cls._env_prefix()}{field_name.upper()}" @classmethod - def from_env(cls) -> "_FrozenConfig": + def from_env(cls) -> "_BaseConfig": return cls( **{ f.name: cls._parse_env_field(cls._env_name(f.name), f.default) @@ -56,7 +57,7 @@ class _FrozenConfig(ABC): } ) - def with_defaults(self, **kwargs) -> "_FrozenConfig": + def with_defaults(self, **kwargs) -> "_BaseConfig": cls = type(self) actual = { key: value @@ -72,9 +73,9 @@ class _FrozenConfig(ABC): return next(a for a in args if a is not type(None)) return hint - @staticmethod - def _parse_env_field(env_name: str, default): - return _FrozenConfig._parse_env_value(os.getenv(env_name), default) + @classmethod + def _parse_env_field(cls, env_name: str, default): + return cls._parse_env_value(os.getenv(env_name), default) @staticmethod def _parse_env_value(raw, default): @@ -86,9 +87,42 @@ class _FrozenConfig(ABC): return int(raw) return raw + @classmethod + def from_kv_pairs(cls, pairs: Optional[List[str]]) -> "_BaseConfig": + return cls(**cls._kv_pairs_to_dict(pairs)) + + @classmethod + def _kv_pairs_to_dict(cls, pairs: Optional[List[str]]) -> dict: + if not pairs: + return {} + + missing = object() + defaults = {f.name: f.default for f in fields(cls)} + result: dict = {} + + for pair in pairs: + key, sep, value = pair.partition("=") + if not sep: + raise ValueError(f"Invalid config pair (missing '='): {pair!r}") + default = defaults.get(key, missing) + if default is missing: + raise ValueError( + f"Unknown config key {key!r}. Valid keys: {sorted(defaults)}" + ) + try: + result[key] = cls._parse_env_value(value, default) + except (ValueError, TypeError) as exc: + field_type = type(default).__name__ + raise TypeError(f"{key}: expected {field_type}, got {value!r}") from exc + + return result + + +_DEFAULT_EXP_NAME_PREFIX = "dump_" + @dataclass(frozen=True) -class _DumperConfig(_FrozenConfig): +class _DumperConfig(_BaseConfig): enable: bool = False filter: Optional[str] = None dir: str = "/tmp/dumper" @@ -431,7 +465,9 @@ class _Dumper: else: if not self._cleanup_previous_handled: self._cleanup_previous_handled = True - _cleanup_old_dumps(Path(self._config.dir)) + _cleanup_old_dumps( + Path(self._config.dir), exp_name=self._config.exp_name + ) path.parent.mkdir(parents=True, exist_ok=True) _torch_save(output_data, str(path)) @@ -454,7 +490,6 @@ class _Dumper: rpc_broadcast = _create_zmq_rpc_broadcast( self, - base_port=get_int_env_var("DUMPER_ZMQ_BASE_PORT", 16800), timeout_seconds=self._config.collective_timeout, ) @@ -652,7 +687,20 @@ def _collective_with_timeout(fn, operation_name: str, timeout_seconds: int = 60) def _get_default_exp_name(timeout_seconds: int = 60): rank = _get_rank() - object_list = [f"dump_{time.time()}" if rank == 0 else None] + now = time.time() + ms = int((now % 1) * 1000) + rand_suffix = random.randint(0, 999) + object_list = [ + ( + ( + f"{_DEFAULT_EXP_NAME_PREFIX}" + f"{time.strftime('%Y%m%d_%H%M%S', time.gmtime(now))}" + f"_{ms:03d}{rand_suffix:03d}" + ) + if rank == 0 + else None + ) + ] if dist.is_initialized(): _collective_with_timeout( @@ -664,14 +712,18 @@ def _get_default_exp_name(timeout_seconds: int = 60): return object_list[0] -def _cleanup_old_dumps(base_dir: Path) -> None: +def _cleanup_old_dumps(base_dir: Path, exp_name: Optional[str] = None) -> None: import shutil if _get_rank() == 0: - for entry in base_dir.glob("dump_*"): - if entry.is_dir(): - shutil.rmtree(entry) - print(f"[Dumper] Cleaned up {entry}") + targets = {entry for entry in base_dir.glob(f"{_DEFAULT_EXP_NAME_PREFIX}*")} + if exp_name: + targets.add(base_dir / exp_name) + targets = {d for d in targets if d.is_dir()} + + for entry in targets: + shutil.rmtree(entry) + print(f"[Dumper] Cleaned up {entry}") if dist.is_initialized(): dist.barrier() @@ -785,19 +837,19 @@ def _make_http_handler(*, prefix: str, target): def _create_zmq_rpc_broadcast( - handler, base_port: int, timeout_seconds: int = 60 + handler, timeout_seconds: int = 60 ) -> Optional["_ZmqRpcBroadcast"]: """A general-purpose minimal RPC to support broadcasting executions to multi processes""" import zmq rank = _get_rank() world_size = dist.get_world_size() if dist.is_initialized() else 1 - port = base_port + rank - local_addr = f"tcp://{_get_local_ip_by_remote()}:{port}" ctx = zmq.Context() sock = ctx.socket(zmq.REP) - sock.bind(f"tcp://*:{port}") + sock.bind("tcp://*:0") + bound_port = int(sock.getsockopt_string(zmq.LAST_ENDPOINT).rsplit(":", 1)[1]) + local_addr = f"tcp://{_get_local_ip_by_remote()}:{bound_port}" def serve_loop(): while True: diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index f215742a2..73029eb6f 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -118,6 +118,95 @@ class TestDumperConfig: assert d2.may_enable is True +class TestKvPairsParsing: + def test_from_kv_pairs_none_returns_defaults(self): + assert _DumperConfig.from_kv_pairs(None) == _DumperConfig() + + def test_from_kv_pairs_empty_returns_defaults(self): + assert _DumperConfig.from_kv_pairs([]) == _DumperConfig() + + def test_from_kv_pairs_bool_field(self): + cfg = _DumperConfig.from_kv_pairs(["enable=true"]) + assert cfg.enable is True + assert cfg.dir == "/tmp/dumper" + + def test_from_kv_pairs_bool_numeric(self): + assert _DumperConfig.from_kv_pairs(["enable=1"]).enable is True + assert _DumperConfig.from_kv_pairs(["enable=0"]).enable is False + + def test_from_kv_pairs_int_field(self): + cfg = _DumperConfig.from_kv_pairs(["collective_timeout=120"]) + assert cfg.collective_timeout == 120 + assert type(cfg.collective_timeout) is int + + def test_from_kv_pairs_int_field_zero_stays_int(self): + cfg = _DumperConfig.from_kv_pairs(["collective_timeout=0"]) + assert cfg.collective_timeout == 0 + assert type(cfg.collective_timeout) is int + + def test_from_kv_pairs_str_field_not_coerced(self): + cfg = _DumperConfig.from_kv_pairs(["server_port=0"]) + assert cfg.server_port == "0" + assert type(cfg.server_port) is str + + def test_from_kv_pairs_str_field_one_stays_str(self): + cfg = _DumperConfig.from_kv_pairs(["server_port=1"]) + assert cfg.server_port == "1" + assert type(cfg.server_port) is str + + def test_from_kv_pairs_optional_str_field(self): + cfg = _DumperConfig.from_kv_pairs(["filter=layer_id=[0-3]"]) + assert cfg.filter == "layer_id=[0-3]" + + def test_from_kv_pairs_optional_str_exp_name(self): + cfg = _DumperConfig.from_kv_pairs(["exp_name=my_experiment"]) + assert cfg.exp_name == "my_experiment" + + def test_from_kv_pairs_multiple_fields(self): + cfg = _DumperConfig.from_kv_pairs( + [ + "enable=true", + "dir=/my/dir", + "filter=name=foo", + "collective_timeout=30", + "enable_grad=1", + ] + ) + assert cfg.enable is True + assert cfg.dir == "/my/dir" + assert cfg.filter == "name=foo" + assert cfg.collective_timeout == 30 + assert cfg.enable_grad is True + + def test_from_kv_pairs_missing_equals_raises(self): + with pytest.raises(ValueError, match="missing '='"): + _DumperConfig.from_kv_pairs(["enable"]) + + def test_from_kv_pairs_unknown_key_raises(self): + with pytest.raises(ValueError, match="Unknown config key"): + _DumperConfig.from_kv_pairs(["nonexistent=true"]) + + def test_kv_pairs_to_dict_returns_only_explicit(self): + d = _DumperConfig._kv_pairs_to_dict(["enable=true", "dir=/x"]) + assert d == {"enable": True, "dir": "/x"} + assert "filter" not in d + assert "collective_timeout" not in d + + def test_kv_pairs_to_dict_none_returns_empty(self): + assert _DumperConfig._kv_pairs_to_dict(None) == {} + + def test_kv_pairs_to_dict_empty_returns_empty(self): + assert _DumperConfig._kv_pairs_to_dict([]) == {} + + def test_from_kv_pairs_value_with_equals_in_value(self): + cfg = _DumperConfig.from_kv_pairs(["filter=name=foo"]) + assert cfg.filter == "name=foo" + + def test_from_kv_pairs_type_validation_still_works(self): + with pytest.raises(TypeError, match="collective_timeout.*expected int"): + _DumperConfig.from_kv_pairs(["collective_timeout=not_a_number"]) + + class TestDumperPureFunctions: def test_get_truncated_value(self): assert get_truncated_value(None) is None @@ -452,13 +541,14 @@ class TestDumpDictFormat: def _make_test_dumper(tmp_path, **overrides) -> _Dumper: """Create a _Dumper for CPU testing without HTTP server or distributed.""" - config = _DumperConfig( + defaults = dict( enable=True, dir=str(tmp_path), exp_name="test", enable_http_server=False, - **overrides, ) + defaults.update(overrides) + config = _DumperConfig(**defaults) return _Dumper(config=config) @@ -795,6 +885,35 @@ class TestCleanup: assert not old_dir.exists() _assert_files(_get_filenames(tmp_path), exist=["new_tensor"]) + def test_cleanup_removes_exp_name_dir(self, tmp_path): + exp_name = "my_custom_exp" + old_exp_dir = tmp_path / exp_name + old_exp_dir.mkdir() + (old_exp_dir / "old_data.pt").touch() + + dumper = _make_test_dumper(tmp_path, exp_name=exp_name, cleanup_previous=True) + dumper.dump("new_tensor", torch.randn(3, 3)) + + assert not (tmp_path / exp_name / "old_data.pt").exists() + _assert_files(_get_filenames(tmp_path), exist=["new_tensor"]) + + def test_cleanup_removes_both_dump_prefix_and_exp_name(self, tmp_path): + old_dump = tmp_path / "dump_old" + old_dump.mkdir() + (old_dump / "dummy.pt").touch() + + exp_name = "custom_run" + old_exp = tmp_path / exp_name + old_exp.mkdir() + (old_exp / "stale.pt").touch() + + dumper = _make_test_dumper(tmp_path, exp_name=exp_name, cleanup_previous=True) + dumper.dump("new_tensor", torch.randn(3, 3)) + + assert not old_dump.exists() + assert not (tmp_path / exp_name / "stale.pt").exists() + _assert_files(_get_filenames(tmp_path), exist=["new_tensor"]) + def test_no_cleanup_by_default(self, tmp_path): old_dir = tmp_path / "dump_old" old_dir.mkdir() @@ -832,6 +951,68 @@ class TestReset: assert "dump_index=1" in post_file.name +def _dumper_worker(rank, http_port: int, stop_event): + """Minimal distributed dumper worker: configure, step (triggers ZMQ setup), then wait.""" + dumper.configure(enable=False, server_port=str(http_port)) + dumper.step() + stop_event.wait() + + +def _wait_for_dumper_http(url: str, timeout: float = 30) -> None: + deadline = time.time() + timeout + while time.time() < deadline: + try: + requests.post(f"{url}/dumper/configure", json={}, timeout=2) + return + except requests.ConnectionError: + time.sleep(0.5) + raise TimeoutError(f"Dumper HTTP server not reachable at {url}") + + +class TestZmqPortIsolation: + """Multiple independent dumper instances (each with 2 ranks) must not conflict on ZMQ ports.""" + + NUM_INSTANCES = 3 + + def test_concurrent_instances_no_port_conflict(self): + ports = [ + find_available_port(40000 + i * 1000) for i in range(self.NUM_INSTANCES) + ] + stop_events = [] + threads = [] + ctx = multiprocessing.get_context("spawn") + + for port in ports: + stop_event = ctx.Event() + stop_events.append(stop_event) + thread = threading.Thread( + target=run_distributed_test, + args=(_dumper_worker,), + kwargs={"http_port": port, "stop_event": stop_event}, + ) + thread.start() + threads.append(thread) + + try: + for port in ports: + _wait_for_dumper_http(f"http://127.0.0.1:{port}") + + for i, port in enumerate(ports): + resp = requests.post( + f"http://127.0.0.1:{port}/dumper/get_state", json={} + ) + resp.raise_for_status() + states = resp.json() + assert ( + len(states) == 2 + ), f"Instance {i} (port {port}): expected 2 ranks, got {len(states)}" + finally: + for event in stop_events: + event.set() + for thread in threads: + thread.join(timeout=10) + + class TestDumperHttp: """Test /dumper/* HTTP control — parametrized over standalone vs sglang server.""" @@ -843,12 +1024,12 @@ class TestDumperHttp: stop_event = multiprocessing.get_context("spawn").Event() thread = threading.Thread( target=run_distributed_test, - args=(TestDumperHttp._standalone_mode_worker,), + args=(_dumper_worker,), kwargs={"http_port": http_port, "stop_event": stop_event}, ) thread.start() try: - TestDumperHttp._wait_for_http(base_url) + _wait_for_dumper_http(base_url) yield base_url finally: stop_event.set() @@ -868,23 +1049,6 @@ class TestDumperHttp: finally: kill_process_tree(proc.pid) - @staticmethod - def _standalone_mode_worker(rank, http_port: int, stop_event): - dumper.configure(enable=False, server_port=str(http_port)) - dumper.step() - stop_event.wait() - - @staticmethod - def _wait_for_http(url: str, timeout: float = 30) -> None: - deadline = time.time() + timeout - while time.time() < deadline: - try: - requests.post(f"{url}/dumper/configure", json={}, timeout=2) - return - except requests.ConnectionError: - time.sleep(0.5) - raise TimeoutError(f"Standalone dumper HTTP server not reachable at {url}") - @staticmethod def _post(base_url: str, method: str, **kwargs) -> list[dict]: resp = requests.post(f"{base_url}/dumper/{method}", json=kwargs or None) @@ -997,6 +1161,15 @@ class _NonIntrusiveTestBase: def _make_dumper(tmp_path, **overrides) -> "_Dumper": return _make_test_dumper(tmp_path, non_intrusive_mode="all", **overrides) + def _run(self, tmp_path, inner_cls, **dumper_overrides): + d = self._make_dumper(tmp_path, **dumper_overrides) + model = self._wrap_as_outer(inner_cls) + d.register_non_intrusive_dumper(model) + x = torch.randn(2, 4) + with d.capture_output() as captured: + output = model(x) + return captured, x, output + class TestNonIntrusiveDumper(_NonIntrusiveTestBase): """Tests for mode='all' — hooks on every module, non_intrusive__ prefix.""" @@ -1011,13 +1184,7 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase): def forward(self, x): return self.relu(self.linear(x)) - d = self._make_dumper(tmp_path) - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - output = model(x) + captured, x, output = self._run(tmp_path, Inner) self._assert_captured_contains( captured, @@ -1103,13 +1270,7 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase): x = layer(x) return x - d = self._make_dumper(tmp_path) - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - model(x) + captured, x, output = self._run(tmp_path, Inner) self._assert_captured_contains( captured, @@ -1144,13 +1305,7 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase): a, b = self.split(x) return self.linear(a + b) - d = self._make_dumper(tmp_path) - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - model(x) + captured, x, output = self._run(tmp_path, Inner) assert "non_intrusive__model.split.output.0" in captured assert "non_intrusive__model.split.output.1" in captured @@ -1171,13 +1326,7 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase): def forward(self, x): return self.wrap(x)[0] - d = self._make_dumper(tmp_path) - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - model(x) + captured, x, output = self._run(tmp_path, Inner) assert "non_intrusive__model.wrap.output" in captured assert "non_intrusive__model.wrap.output.0" not in captured @@ -1196,13 +1345,7 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase): mask = torch.ones_like(x) return self.mul(x, mask) - d = self._make_dumper(tmp_path) - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - model(x) + captured, x, output = self._run(tmp_path, Inner) assert "non_intrusive__model.mul.inputs.0" in captured assert "non_intrusive__model.mul.inputs.1" in captured @@ -1221,13 +1364,7 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase): self.sink(x) return x - d = self._make_dumper(tmp_path) - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - model(x) + captured, x, output = self._run(tmp_path, Inner) assert "non_intrusive__model.sink.inputs.0" in captured assert not any( @@ -1248,13 +1385,7 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase): self.const(x) return x - d = self._make_dumper(tmp_path) - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - model(x) + captured, x, output = self._run(tmp_path, Inner) assert "non_intrusive__model.const.inputs.0" in captured assert not any( @@ -1287,15 +1418,9 @@ class TestNonIntrusiveDumper(_NonIntrusiveTestBase): def forward(self, x): return self.relu(self.linear(x)) - d = self._make_dumper( - tmp_path, filter="name=non_intrusive__model.linear.output" + captured, x, output = self._run( + tmp_path, Inner, filter="name=non_intrusive__model.linear.output" ) - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - model(x) assert "non_intrusive__model.linear.output" in captured assert "non_intrusive__model.relu.output" not in captured @@ -1422,38 +1547,37 @@ class TestNonIntrusiveDumperConfigMode(_NonIntrusiveTestBase): assert "non_intrusive__layer.output" in captured +class _LayerWithNumber(torch.nn.Module): + """Test helper: module with a ``layer_number`` attribute (Megatron style).""" + + def __init__(self, layer_number: int): + super().__init__() + self.layer_number = layer_number + self.linear = torch.nn.Linear(4, 4) + + def forward(self, x): + return self.linear(x) + + class TestNonIntrusiveLayerIdCtx(_NonIntrusiveTestBase): """Tests for automatic layer_id context injection via set_ctx.""" def test_layer_id_from_layer_number(self, tmp_path): """Megatron PP: layer_number (1-based global) -> layer_id = layer_number - 1.""" - class Layer(torch.nn.Module): - def __init__(self, layer_number: int): - super().__init__() - self.layer_number = layer_number - self.linear = torch.nn.Linear(4, 4) - - def forward(self, x): - return self.linear(x) - class Inner(torch.nn.Module): def __init__(self): super().__init__() - self.layers = torch.nn.ModuleList([Layer(10), Layer(11)]) + self.layers = torch.nn.ModuleList( + [_LayerWithNumber(10), _LayerWithNumber(11)] + ) def forward(self, x): for layer in self.layers: x = layer(x) return x - d = self._make_dumper(tmp_path) - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - model(x) + captured, x, output = self._run(tmp_path, Inner) layer0_key = "non_intrusive__model.layers.0.linear.output" layer1_key = "non_intrusive__model.layers.1.linear.output" @@ -1488,13 +1612,7 @@ class TestNonIntrusiveLayerIdCtx(_NonIntrusiveTestBase): x = layer(x) return x - d = self._make_dumper(tmp_path) - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - model(x) + captured, x, output = self._run(tmp_path, Inner) layer_key = "non_intrusive__model.layers.0.linear.output" assert layer_key in captured @@ -1515,13 +1633,7 @@ class TestNonIntrusiveLayerIdCtx(_NonIntrusiveTestBase): x = layer(x) return x - d = self._make_dumper(tmp_path) - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - model(x) + captured, x, output = self._run(tmp_path, Inner) assert len(captured) > 0 for key, entry in captured.items(): @@ -1530,32 +1642,19 @@ class TestNonIntrusiveLayerIdCtx(_NonIntrusiveTestBase): def test_filter_by_layer_id(self, tmp_path): """filter='layer_id=0' keeps only layer 0 dumps.""" - class Layer(torch.nn.Module): - def __init__(self, layer_number: int): - super().__init__() - self.layer_number = layer_number - self.linear = torch.nn.Linear(4, 4) - - def forward(self, x): - return self.linear(x) - class Inner(torch.nn.Module): def __init__(self): super().__init__() - self.layers = torch.nn.ModuleList([Layer(1), Layer(2)]) + self.layers = torch.nn.ModuleList( + [_LayerWithNumber(1), _LayerWithNumber(2)] + ) def forward(self, x): for layer in self.layers: x = layer(x) return x - d = self._make_dumper(tmp_path, filter="layer_id=0") - model = self._wrap_as_outer(Inner) - d.register_non_intrusive_dumper(model) - - x = torch.randn(2, 4) - with d.capture_output() as captured: - model(x) + captured, x, output = self._run(tmp_path, Inner, filter="layer_id=0") layer0_keys = [k for k in captured if "layers.0" in k] layer1_keys = [k for k in captured if "layers.1" in k]