Tiny fix single-gpu dumper and add tests for dumper (#16285)

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:

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()

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()

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()