chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Flip dumper to disable by default and refactor environment handling (#18878)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-16 13:29:32 +08:00
committed by GitHub
parent 5ddc84e33e
commit 02816abc0d
3 changed files with 117 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
python/sglang/srt/debug_utils/dumper.py
View File

@@ -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 ------------------------------------------

16
python/sglang/srt/environ.py
View File

@@ -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:

70
test/registered/debug_utils/test_dumper.py
View File

@@ -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")