diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index c199225ef..3bb2c371f 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -28,28 +28,53 @@ class _Dumper: from dumper import dumper ``` - Disable at startup and enable via HTTP: - 1. `SGLANG_DUMPER_ENABLE=0 python ...` + Then run the program: + `SGLANG_DUMPER_ENABLE=1 python ...` + + Alternatively, disable at startup and enable via HTTP: + 1. `python ...` 2. `curl -X POST http://localhost:40000/dumper -d '{"enable": true}'` Related: `sglang.srt.debug_utils.dump_comparator` for dump comparison """ - def __init__(self): - # Do not import `sglang` to make this file standalone - self._enable = bool(int(os.environ.get("SGLANG_DUMPER_ENABLE", "1"))) + def __init__( + self, + *, + enable: bool, + base_dir: Path, + filter: Optional[str] = None, + enable_write_file: bool = True, + partial_name: Optional[str] = None, + enable_http_server: bool = True, + ): + # Config + self._enable = enable # TODO (1) support filtering kv instead of name only (2) allow HTTP req change it - self._filter = os.environ.get("SGLANG_DUMPER_FILTER") - self._base_dir = Path(os.environ.get("SGLANG_DUMPER_DIR", "/tmp")) - self._enable_write_file = bool( - int(os.environ.get("SGLANG_DUMPER_WRITE_FILE", "1")) - ) - self._partial_name: Optional[str] = None + self._filter = filter + self._base_dir = base_dir + self._enable_write_file = enable_write_file + + # States + self._partial_name = partial_name self._dump_index = 0 self._forward_pass_id = 0 self._global_ctx = {} self._override_enable = None - self._http_server_handled = False + self._http_server_handled = not enable_http_server + + @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_write_file=get_bool_env_var("SGLANG_DUMPER_WRITE_FILE", "1"), + partial_name=_get_str_env_var("SGLANG_DUMPER_PARTIAL_NAME"), + enable_http_server=get_bool_env_var( + "SGLANG_ENABLE_DUMPER_HTTP_SERVER", "1" + ), + ) def on_forward_pass_start(self): """This should be called on all ranks.""" @@ -193,8 +218,8 @@ def _obj_to_dict(obj): def _start_maybe_http_server(dumper): - http_port = int(os.environ.get("SGLANG_DUMPER_SERVER_PORT", "40000")) - zmq_base_port = int(os.environ.get("SGLANG_DUMPER_ZMQ_BASE_PORT", "16800")) + http_port = get_int_env_var("SGLANG_DUMPER_SERVER_PORT", 40000) + zmq_base_port = get_int_env_var("SGLANG_DUMPER_ZMQ_BASE_PORT", 16800) if http_port <= 0: return @@ -321,6 +346,30 @@ 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(): + return default + try: + return int(value) + except ValueError: + return default + + def _get_local_ip_by_remote() -> Optional[str]: # try ipv4 s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) @@ -353,7 +402,7 @@ def _get_local_ip_by_remote() -> Optional[str]: # -------------------------------------- singleton ------------------------------------------ -dumper = _Dumper() +dumper = _Dumper.from_env() # -------------------------------------- other utility functions ------------------------------------------ diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 65f81d037..593fd03b5 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -7,11 +7,17 @@ from typing import Any @contextmanager -def temp_set_env(**env_vars: dict[str, Any]): - """Temporarily set non-sglang environment variables, e.g. OPENAI_API_KEY""" - for key in env_vars: - if key.startswith("SGLANG_") or key.startswith("SGL_"): - raise ValueError("temp_set_env should not be used for sglang env vars") +def temp_set_env(*, allow_sglang: bool = False, **env_vars: Any): + """Temporarily set environment variables, restoring originals on exit. + + By default, SGLANG_*/SGL_* keys are rejected — use ``Envs`` descriptors + for those. Pass ``allow_sglang=True`` only for special env vars that + intentionally bypass ``environ.py`` (e.g. ``SGLANG_DUMPER_*``). + """ + if not allow_sglang: + for key in env_vars: + if key.startswith("SGLANG_") or key.startswith("SGL_"): + raise ValueError("temp_set_env should not be used for sglang env vars") backup = {key: os.environ.get(key) for key in env_vars} try: diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index 69e874598..04a0da461 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -1,4 +1,3 @@ -import os import sys import time from pathlib import Path @@ -8,6 +7,13 @@ import requests import torch import torch.distributed as dist +from sglang.srt.debug_utils.dumper import ( + _obj_to_dict, + _torch_save, + get_tensor_info, + get_truncated_value, +) +from sglang.srt.environ import temp_set_env from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import run_distributed_test @@ -17,8 +23,6 @@ register_amd_ci(est_time=60, suite="nightly-amd", nightly=True) class TestDumperPureFunctions: def test_get_truncated_value(self): - from sglang.srt.debug_utils.dumper import get_truncated_value - assert get_truncated_value(None) is None assert get_truncated_value(42) == 42 assert len(get_truncated_value((torch.randn(10), torch.randn(20)))) == 2 @@ -26,8 +30,6 @@ class TestDumperPureFunctions: assert get_truncated_value(torch.randn(100, 100)).shape == (5, 5) def test_obj_to_dict(self): - from sglang.srt.debug_utils.dumper import _obj_to_dict - assert _obj_to_dict({"a": 1}) == {"a": 1} class Obj: @@ -41,8 +43,6 @@ class TestDumperPureFunctions: assert "method" not in result def test_get_tensor_info(self): - from sglang.srt.debug_utils.dumper import get_tensor_info - info = get_tensor_info(torch.randn(10, 10)) for key in ["shape=", "dtype=", "min=", "max=", "mean="]: assert key in info @@ -53,8 +53,6 @@ class TestDumperPureFunctions: class TestTorchSave: def test_normal(self, tmp_path): - from sglang.srt.debug_utils.dumper import _torch_save - path = str(tmp_path / "a.pt") tensor = torch.randn(3, 3) @@ -63,8 +61,6 @@ class TestTorchSave: assert torch.equal(torch.load(path, weights_only=True), tensor) def test_parameter_fallback(self, tmp_path): - from sglang.srt.debug_utils.dumper import _torch_save - class BadParam(torch.nn.Parameter): def __reduce_ex__(self, protocol): raise RuntimeError("not pickleable") @@ -77,8 +73,6 @@ class TestTorchSave: assert torch.equal(torch.load(path, weights_only=True), param.data) def test_silent_skip(self, tmp_path, capsys): - from sglang.srt.debug_utils.dumper import _torch_save - path = str(tmp_path / "c.pt") _torch_save({"fn": lambda: None}, path) @@ -90,11 +84,15 @@ class TestTorchSave: class TestDumperDistributed: def test_basic(self, tmp_path): - run_distributed_test(self._test_basic_func, tmpdir=str(tmp_path)) + with temp_set_env( + allow_sglang=True, + SGLANG_DUMPER_ENABLE="1", + SGLANG_DUMPER_DIR=str(tmp_path), + ): + run_distributed_test(self._test_basic_func, tmpdir=str(tmp_path)) @staticmethod def _test_basic_func(rank, tmpdir): - os.environ["SGLANG_DUMPER_DIR"] = tmpdir from sglang.srt.debug_utils.dumper import dumper tensor = torch.randn(10, 10, device=f"cuda:{rank}") @@ -124,11 +122,11 @@ class TestDumperDistributed: ) def test_http_enable(self): - run_distributed_test(self._test_http_func) + with temp_set_env(allow_sglang=True, SGLANG_DUMPER_ENABLE="0"): + run_distributed_test(self._test_http_func) @staticmethod def _test_http_func(rank): - os.environ["SGLANG_DUMPER_ENABLE"] = "0" from sglang.srt.debug_utils.dumper import dumper assert not dumper._enable @@ -145,11 +143,15 @@ class TestDumperDistributed: assert dumper._enable == enable def test_file_content_correctness(self, tmp_path): - run_distributed_test(self._test_file_content_func, tmpdir=str(tmp_path)) + with temp_set_env( + allow_sglang=True, + SGLANG_DUMPER_ENABLE="1", + SGLANG_DUMPER_DIR=str(tmp_path), + ): + run_distributed_test(self._test_file_content_func, tmpdir=str(tmp_path)) @staticmethod def _test_file_content_func(rank, tmpdir): - os.environ["SGLANG_DUMPER_DIR"] = tmpdir from sglang.srt.debug_utils.dumper import dumper tensor = torch.arange(12, device=f"cuda:{rank}").reshape(3, 4).float() @@ -165,12 +167,16 @@ class TestDumperDistributed: class TestDumperFileWriteControl: def test_filter(self, tmp_path): - run_distributed_test(self._test_filter_func, tmpdir=str(tmp_path)) + with temp_set_env( + allow_sglang=True, + SGLANG_DUMPER_ENABLE="1", + SGLANG_DUMPER_DIR=str(tmp_path), + SGLANG_DUMPER_FILTER="^keep", + ): + run_distributed_test(self._test_filter_func, tmpdir=str(tmp_path)) @staticmethod def _test_filter_func(rank, tmpdir): - os.environ["SGLANG_DUMPER_DIR"] = tmpdir - os.environ["SGLANG_DUMPER_FILTER"] = "^keep" from sglang.srt.debug_utils.dumper import dumper dumper.on_forward_pass_start() @@ -187,12 +193,16 @@ class TestDumperFileWriteControl: ) def test_write_disabled(self, tmp_path): - run_distributed_test(self._test_write_disabled_func, tmpdir=str(tmp_path)) + with temp_set_env( + allow_sglang=True, + SGLANG_DUMPER_ENABLE="1", + SGLANG_DUMPER_DIR=str(tmp_path), + SGLANG_DUMPER_WRITE_FILE="0", + ): + run_distributed_test(self._test_write_disabled_func, tmpdir=str(tmp_path)) @staticmethod def _test_write_disabled_func(rank, tmpdir): - os.environ["SGLANG_DUMPER_DIR"] = tmpdir - os.environ["SGLANG_DUMPER_WRITE_FILE"] = "0" from sglang.srt.debug_utils.dumper import dumper dumper.on_forward_pass_start() @@ -202,11 +212,15 @@ class TestDumperFileWriteControl: assert len(_get_filenames(tmpdir)) == 0 def test_save_false(self, tmp_path): - run_distributed_test(self._test_save_false_func, tmpdir=str(tmp_path)) + with temp_set_env( + allow_sglang=True, + SGLANG_DUMPER_ENABLE="1", + SGLANG_DUMPER_DIR=str(tmp_path), + ): + run_distributed_test(self._test_save_false_func, tmpdir=str(tmp_path)) @staticmethod def _test_save_false_func(rank, tmpdir): - os.environ["SGLANG_DUMPER_DIR"] = tmpdir from sglang.srt.debug_utils.dumper import dumper dumper.on_forward_pass_start() @@ -229,7 +243,7 @@ def _assert_files(filenames, *, exist=(), not_exist=()): ), f"{p} should not exist in {filenames}" -def _find_dump_file(tmpdir, *, rank: int, name: str) -> Path: +def _find_dump_file(tmpdir, *, rank: int = 0, name: str) -> Path: matches = [ f for f in Path(tmpdir).glob("sglang_dump_*/*.pt")