Add missing dumper tests (#18859)
This commit is contained in:
@@ -51,6 +51,43 @@ class TestDumperPureFunctions:
|
|||||||
assert "min=None" in get_tensor_info(torch.tensor([]))
|
assert "min=None" in get_tensor_info(torch.tensor([]))
|
||||||
|
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
_torch_save(tensor, path)
|
||||||
|
|
||||||
|
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")
|
||||||
|
|
||||||
|
path = str(tmp_path / "b.pt")
|
||||||
|
param = BadParam(torch.randn(4))
|
||||||
|
|
||||||
|
_torch_save(param, path)
|
||||||
|
|
||||||
|
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)
|
||||||
|
|
||||||
|
captured = capsys.readouterr()
|
||||||
|
assert "[Dumper] Observe error=" in captured.out
|
||||||
|
assert "skip the tensor" in captured.out
|
||||||
|
|
||||||
|
|
||||||
class TestDumperDistributed:
|
class TestDumperDistributed:
|
||||||
def test_basic(self, tmp_path):
|
def test_basic(self, tmp_path):
|
||||||
run_distributed_test(self._test_basic_func, tmpdir=str(tmp_path))
|
run_distributed_test(self._test_basic_func, tmpdir=str(tmp_path))
|
||||||
@@ -107,6 +144,26 @@ class TestDumperDistributed:
|
|||||||
dist.barrier()
|
dist.barrier()
|
||||||
assert dumper._enable == enable
|
assert dumper._enable == enable
|
||||||
|
|
||||||
|
def test_file_content_correctness(self, 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()
|
||||||
|
|
||||||
|
dumper.on_forward_pass_start()
|
||||||
|
dumper.dump("content_check", tensor)
|
||||||
|
|
||||||
|
dist.barrier()
|
||||||
|
path = _find_dump_file(tmpdir, rank=rank, name="content_check")
|
||||||
|
loaded = torch.load(path, map_location="cpu", weights_only=True)
|
||||||
|
assert torch.equal(loaded, tensor.cpu())
|
||||||
|
|
||||||
|
|
||||||
|
class TestDumperFileWriteControl:
|
||||||
def test_filter(self, tmp_path):
|
def test_filter(self, tmp_path):
|
||||||
run_distributed_test(self._test_filter_func, tmpdir=str(tmp_path))
|
run_distributed_test(self._test_filter_func, tmpdir=str(tmp_path))
|
||||||
|
|
||||||
@@ -144,6 +201,20 @@ class TestDumperDistributed:
|
|||||||
dist.barrier()
|
dist.barrier()
|
||||||
assert len(_get_filenames(tmpdir)) == 0
|
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))
|
||||||
|
|
||||||
|
@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()
|
||||||
|
dumper.dump("no_save_tensor", torch.randn(5, device=f"cuda:{rank}"), save=False)
|
||||||
|
|
||||||
|
dist.barrier()
|
||||||
|
assert len(_get_filenames(tmpdir)) == 0
|
||||||
|
|
||||||
|
|
||||||
def _get_filenames(tmpdir):
|
def _get_filenames(tmpdir):
|
||||||
return {f.name for f in Path(tmpdir).glob("sglang_dump_*/*.pt")}
|
return {f.name for f in Path(tmpdir).glob("sglang_dump_*/*.pt")}
|
||||||
@@ -158,5 +229,17 @@ def _assert_files(filenames, *, exist=(), not_exist=()):
|
|||||||
), f"{p} should not exist in {filenames}"
|
), f"{p} should not exist in {filenames}"
|
||||||
|
|
||||||
|
|
||||||
|
def _find_dump_file(tmpdir, *, rank: int, name: str) -> Path:
|
||||||
|
matches = [
|
||||||
|
f
|
||||||
|
for f in Path(tmpdir).glob("sglang_dump_*/*.pt")
|
||||||
|
if f"rank={rank}" in f.name and name in f.name
|
||||||
|
]
|
||||||
|
assert (
|
||||||
|
len(matches) == 1
|
||||||
|
), f"Expected 1 file matching rank={rank} name={name}, got {matches}"
|
||||||
|
return matches[0]
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
sys.exit(pytest.main([__file__]))
|
sys.exit(pytest.main([__file__]))
|
||||||
|
|||||||
Reference in New Issue
Block a user