Flip dumper to disable by default and refactor environment handling (#18878)
This commit is contained in:
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user