From 749736ba952e45191fb6719f6f4509c3cbaa512a Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 2 Jan 2026 13:30:09 +0800 Subject: [PATCH] Tiny fix single-gpu dumper and add tests for dumper (#16285) --- python/sglang/srt/debug_utils/dumper.py | 11 +- .../debug_utils/test_dump_comparator.py | 136 ++++++++++++ .../debug_utils/test_dump_loader.py | 67 ++++++ test/registered/debug_utils/test_dumper.py | 205 ++++++++++++++++++ 4 files changed, 415 insertions(+), 4 deletions(-) create mode 100644 test/registered/debug_utils/test_dump_comparator.py create mode 100644 test/registered/debug_utils/test_dump_loader.py create mode 100644 test/registered/debug_utils/test_dumper.py diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 481f23ecf..c199225ef 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -250,8 +250,8 @@ class _DumperRpcHandler: def _create_zmq_rpc_handles(handler, base_port: int) -> Optional[List["_ZmqRpcHandle"]]: import zmq - rank = dist.get_rank() - world_size = dist.get_world_size() + 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}" @@ -274,8 +274,11 @@ def _create_zmq_rpc_handles(handler, base_port: int) -> Optional[List["_ZmqRpcHa thread.start() print(f"[Dumper.ZmqRpc] rank={rank} server started at {local_addr}") - all_addresses = [None] * world_size - dist.all_gather_object(all_addresses, local_addr) + if dist.is_initialized(): + all_addresses = [None] * world_size + dist.all_gather_object(all_addresses, local_addr) + else: + all_addresses = [local_addr] print(f"[Dumper.ZmqRpc] rank={rank} all_addresses={all_addresses}") if rank == 0: diff --git a/test/registered/debug_utils/test_dump_comparator.py b/test/registered/debug_utils/test_dump_comparator.py new file mode 100644 index 000000000..5b73dd48f --- /dev/null +++ b/test/registered/debug_utils/test_dump_comparator.py @@ -0,0 +1,136 @@ +import os +import tempfile +import unittest +from contextlib import contextmanager +from pathlib import Path + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=60, suite="default", nightly=True) + + +class TestDumpComparator(CustomTestCase): + def test_calc_rel_diff(self): + from sglang.srt.debug_utils.dump_comparator import _calc_rel_diff + + x = torch.randn(10, 10) + self.assertAlmostEqual(_calc_rel_diff(x, x).item(), 0.0, places=5) + self.assertAlmostEqual( + _calc_rel_diff(torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])).item(), + 1.0, + places=5, + ) + + def test_argmax_coord(self): + from sglang.srt.debug_utils.dump_comparator import _argmax_coord + + x = torch.zeros(2, 3, 4) + x[1, 2, 3] = 10.0 + self.assertEqual(_argmax_coord(x), (1, 2, 3)) + + def test_try_unify_shape(self): + from sglang.srt.debug_utils.dump_comparator import _try_unify_shape + + target = torch.Size([3, 4]) + self.assertEqual( + _try_unify_shape(torch.randn(1, 1, 3, 4), target).shape, target + ) + self.assertEqual( + _try_unify_shape(torch.randn(2, 3, 4), target).shape, (2, 3, 4) + ) + + def test_compute_smaller_dtype(self): + from sglang.srt.debug_utils.dump_comparator import _compute_smaller_dtype + + self.assertEqual( + _compute_smaller_dtype(torch.float32, torch.bfloat16), torch.bfloat16 + ) + self.assertIsNone(_compute_smaller_dtype(torch.float32, torch.float32)) + + def test_einops_pattern(self): + from sglang.srt.debug_utils.dump_comparator import ( + _get_einops_dim_index, + _split_einops_pattern, + ) + + self.assertEqual(_split_einops_pattern("a (b c) d"), ["a", "(b c)", "d"]) + self.assertEqual(_get_einops_dim_index("a b c", "b"), 1) + + def test_load_object(self): + from sglang.srt.debug_utils.dump_comparator import _load_object + + with tempfile.TemporaryDirectory() as tmpdir: + path = Path(tmpdir) / "tensor.pt" + torch.save(torch.randn(5, 5), path) + self.assertEqual(_load_object(path).shape, (5, 5)) + + torch.save({"dict": 1}, path) + self.assertIsNone(_load_object(path)) + + self.assertIsNone(_load_object("/nonexistent.pt")) + + def test_compute_and_print_diff(self): + from sglang.srt.debug_utils.dump_comparator import _compute_and_print_diff + + x = torch.ones(10, 10) + self.assertAlmostEqual( + _compute_and_print_diff(x, x, 1e-3)["max_abs_diff"], 0.0, places=5 + ) + self.assertAlmostEqual( + _compute_and_print_diff(x, x + 0.5, 1e-3)["max_abs_diff"], 0.5, places=4 + ) + + +class TestEndToEnd(CustomTestCase): + def test_main(self): + from argparse import Namespace + + from sglang.srt.debug_utils.dump_comparator import main + from sglang.srt.debug_utils.dumper import _Dumper + + with tempfile.TemporaryDirectory() as d1, tempfile.TemporaryDirectory() as d2: + baseline_tensor = torch.randn(10, 10) + target_tensor = baseline_tensor + torch.randn(10, 10) * 0.01 + + dump_dirs = [] + for d, tensor in [(d1, baseline_tensor), (d2, target_tensor)]: + with _with_env("SGLANG_DUMPER_DIR", d), _with_env( + "SGLANG_DUMPER_SERVER_PORT", "-1" + ): + dumper = _Dumper() + dumper.on_forward_pass_start() + dumper.dump("tensor_a", tensor) + dumper.on_forward_pass_start() + dumper.dump("tensor_b", tensor * 2) + dump_dirs.append(Path(d) / f"sglang_dump_{dumper._partial_name}") + + args = Namespace( + baseline_path=str(dump_dirs[0]), + target_path=str(dump_dirs[1]), + start_id=1, + end_id=2, + baseline_start_id=1, + diff_threshold=1e-3, + filter=None, + ) + main(args) + + +@contextmanager +def _with_env(name: str, value: str): + old = os.environ.get(name) + os.environ[name] = value + try: + yield + finally: + if old is None: + os.environ.pop(name, None) + else: + os.environ[name] = old + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/debug_utils/test_dump_loader.py b/test/registered/debug_utils/test_dump_loader.py new file mode 100644 index 000000000..b704f62f9 --- /dev/null +++ b/test/registered/debug_utils/test_dump_loader.py @@ -0,0 +1,67 @@ +import tempfile +import unittest +from pathlib import Path + +import polars as pl +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=30, suite="default", nightly=True) + + +class TestDumpLoader(CustomTestCase): + def test_read_meta(self): + from sglang.srt.debug_utils.dump_loader import read_meta + + with tempfile.TemporaryDirectory() as tmpdir: + for fn in [ + "forward_pass_id=1___rank=0___dump_index=1___name=a.pt", + "forward_pass_id=2___rank=0___dump_index=2___name=b.pt", + ]: + torch.save(torch.randn(5), Path(tmpdir) / fn) + + df = read_meta(tmpdir) + self.assertEqual(len(df), 2) + self.assertTrue( + all(c in df.columns for c in ["forward_pass_id", "rank", "name"]) + ) + + def test_find_row(self): + from sglang.srt.debug_utils.dump_loader import find_row + + df = pl.DataFrame({"id": [1, 2], "name": ["a", "b"], "file": ["f1", "f2"]}) + self.assertEqual(find_row(df, {"id": 2})["file"], "f2") + self.assertIsNone(find_row(df, {"id": 999})) + + df_dup = pl.DataFrame({"id": [1, 1], "file": ["f1", "f2"]}) + self.assertIsNone(find_row(df_dup, {"id": 1})) + + def test_cast_to_polars_dtype(self): + from sglang.srt.debug_utils.dump_loader import _cast_to_polars_dtype + + self.assertEqual(_cast_to_polars_dtype("42", pl.Int64), 42) + self.assertEqual(_cast_to_polars_dtype("3.14", pl.Float64), 3.14) + + def test_add_duplicate_index(self): + from sglang.srt.debug_utils.dump_loader import _add_duplicate_index + + df = pl.DataFrame( + { + "name": ["a", "a", "b"], + "dump_index": [1, 2, 3], + "filename": ["f1", "f2", "f3"], + } + ) + result = _add_duplicate_index(df) + self.assertEqual( + result.filter(pl.col("name") == "a") + .sort("dump_index")["duplicate_index"] + .to_list(), + [0, 1], + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py new file mode 100644 index 000000000..6eb52b800 --- /dev/null +++ b/test/registered/debug_utils/test_dumper.py @@ -0,0 +1,205 @@ +import os +import tempfile +import time +import unittest +from pathlib import Path + +import requests +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, suite="nightly-2-gpu", nightly=True) + + +class TestDumperPureFunctions(CustomTestCase): + def test_get_truncated_value(self): + from sglang.srt.debug_utils.dumper import get_truncated_value + + self.assertIsNone(get_truncated_value(None)) + self.assertEqual(get_truncated_value(42), 42) + self.assertEqual( + len(get_truncated_value((torch.randn(10), torch.randn(20)))), 2 + ) + self.assertEqual(get_truncated_value(torch.randn(10, 10)).shape, (10, 10)) + self.assertEqual(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 + + self.assertEqual(_obj_to_dict({"a": 1}), {"a": 1}) + + class Obj: + x, y = 10, 20 + + def method(self): + pass + + result = _obj_to_dict(Obj()) + self.assertEqual(result["x"], 10) + self.assertNotIn("method", 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="]: + self.assertIn(key, info) + + self.assertIn("value=42", get_tensor_info(42)) + self.assertIn("min=None", get_tensor_info(torch.tensor([]))) + + +class TestDumperDistributed(CustomTestCase): + def test_basic(self): + with tempfile.TemporaryDirectory(prefix="test_dumper_") as tmpdir: + _run_distributed_test(_test_basic_func, tmpdir=tmpdir) + + def test_http_enable(self): + _run_distributed_test(_test_http_func) + + def test_filter(self): + with tempfile.TemporaryDirectory(prefix="test_dumper_") as tmpdir: + _run_distributed_test(_test_filter_func, tmpdir=tmpdir) + + def test_write_disabled(self): + with tempfile.TemporaryDirectory(prefix="test_dumper_") as tmpdir: + _run_distributed_test(_test_write_disabled_func, tmpdir=tmpdir) + + +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}") + + dumper.on_forward_pass_start() + dumper.dump("tensor_a", tensor, arg=100) + + dumper.on_forward_pass_start() + dumper.set_ctx(ctx_arg=200) + dumper.dump("tensor_b", tensor) + dumper.set_ctx(ctx_arg=None) + + dumper.on_forward_pass_start() + dumper.override_enable(False) + dumper.dump("tensor_skip", tensor) + dumper.override_enable(True) + + dumper.on_forward_pass_start() + dumper.dump_dict("obj", {"a": torch.randn(3, device=f"cuda:{rank}"), "b": 42}) + + dist.barrier() + filenames = _get_filenames(tmpdir) + + _assert_files( + filenames, + exist=["tensor_a", "tensor_b", "arg=100", "ctx_arg=200", "obj_a", "obj_b"], + not_exist=["tensor_skip"], + ) + + +def _test_http_func(rank): + os.environ["SGLANG_DUMPER_ENABLE"] = "0" + from sglang.srt.debug_utils.dumper import dumper + + assert not dumper._enable + dumper.on_forward_pass_start() + + for enable in [True, False]: + dist.barrier() + if rank == 0: + time.sleep(0.1) + requests.post( + "http://localhost:40000/dumper", json={"enable": enable} + ).raise_for_status() + dist.barrier() + assert dumper._enable == enable + + +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() + dumper.dump("keep_this", torch.randn(5, device=f"cuda:{rank}")) + dumper.dump("skip_this", torch.randn(5, device=f"cuda:{rank}")) + + dist.barrier() + filenames = _get_filenames(tmpdir) + _assert_files(filenames, exist=["keep_this"], not_exist=["skip_this"]) + + +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() + dumper.dump("no_write", torch.randn(5, device=f"cuda:{rank}")) + + dist.barrier() + assert len(_get_filenames(tmpdir)) == 0 + + +def _get_filenames(tmpdir): + return {f.name for f in Path(tmpdir).glob("sglang_dump_*/*.pt")} + + +def _assert_files(filenames, *, exist=(), not_exist=()): + for p in exist: + assert any(p in f for f in filenames), f"{p} not found in {filenames}" + for p in not_exist: + assert not any( + p in f for f in filenames + ), f"{p} should not exist in {filenames}" + + +def _run_distributed_test(func, world_size=2, **kwargs): + ctx = mp.get_context("spawn") + result_queue = ctx.Queue() + processes = [] + + for rank in range(world_size): + p = ctx.Process( + target=_run_worker, args=(rank, world_size, func, result_queue, kwargs) + ) + p.start() + processes.append(p) + + for p in processes: + p.join() + + errors = [result_queue.get() for _ in range(world_size)] + errors = [e for e in errors if e] + if errors: + raise AssertionError("\n".join(errors)) + + +def _run_worker(rank, world_size, func, result_queue, kwargs): + os.environ.update( + MASTER_ADDR="localhost", + MASTER_PORT="29500", + RANK=str(rank), + WORLD_SIZE=str(world_size), + ) + torch.cuda.set_device(rank) + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + + try: + func(rank, **kwargs) + result_queue.put(None) + except Exception as e: + import traceback + + result_queue.put(f"Rank {rank}: {e}\n{traceback.format_exc()}") + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + unittest.main()