Extract dumper and prefill delayer tests common utils (#18857)

python/sglang/test/test_utils.py

@@ -2004,6 +2004,56 @@ async def send_concurrent_generate_requests_with_custom_params(
     return await asyncio.gather(*tasks)
 
 
+def run_distributed_test(func, world_size=2, backend="nccl", **kwargs):
+    """Spawn ``world_size`` processes, initialise torch.distributed in each,
+    run *func(rank, **kwargs)*, and propagate any worker exception to the caller.
+    """
+    import torch.multiprocessing as mp
+
+    ctx = mp.get_context("spawn")
+    result_queue = ctx.Queue()
+    port = find_available_port(29500)
+
+    processes = []
+    for rank in range(world_size):
+        p = ctx.Process(
+            target=_distributed_worker,
+            args=(rank, world_size, backend, port, 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 _distributed_worker(rank, world_size, backend, port, func, result_queue, kwargs):
+    import traceback
+
+    import torch.distributed as dist
+
+    if backend == "nccl":
+        torch.cuda.set_device(rank)
+    dist.init_process_group(
+        backend=backend,
+        init_method=f"tcp://127.0.0.1:{port}",
+        world_size=world_size,
+        rank=rank,
+    )
+    try:
+        func(rank, **kwargs)
+        result_queue.put(None)
+    except Exception as e:
+        result_queue.put(f"Rank {rank}: {e}\n{traceback.format_exc()}")
+    finally:
+        dist.destroy_process_group()
+
+
 class CustomTestCase(unittest.TestCase):
     def _callTestMethod(self, method):
         max_retry = envs.SGLANG_TEST_MAX_RETRY.get()

test/registered/debug_utils/test_dumper.py

@@ -1,37 +1,34 @@
 import os
-import tempfile
+import sys
 import time
-import unittest
 from pathlib import Path
 
+import pytest
 import requests
 import torch
 import torch.distributed as dist
-import torch.multiprocessing as mp
 
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
-from sglang.test.test_utils import CustomTestCase
+from sglang.test.test_utils import run_distributed_test
 
 register_cuda_ci(est_time=30, suite="nightly-2-gpu", nightly=True)
 register_amd_ci(est_time=60, suite="nightly-amd", nightly=True)
 
 
-class TestDumperPureFunctions(CustomTestCase):
+class TestDumperPureFunctions:
     def test_get_truncated_value(self):
         from sglang.srt.debug_utils.dumper import get_truncated_value
 
-        self.assertIsNone(get_truncated_value(None))
+        assert get_truncated_value(None) is None
-        self.assertEqual(get_truncated_value(42), 42)
+        assert get_truncated_value(42) == 42
-        self.assertEqual(
+        assert len(get_truncated_value((torch.randn(10), torch.randn(20)))) == 2
-            len(get_truncated_value((torch.randn(10), torch.randn(20)))), 2
+        assert get_truncated_value(torch.randn(10, 10)).shape == (10, 10)
-        )
+        assert get_truncated_value(torch.randn(100, 100)).shape == (5, 5)
-        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})
+        assert _obj_to_dict({"a": 1}) == {"a": 1}
 
         class Obj:
             x, y = 10, 20
@@ -40,111 +37,112 @@ class TestDumperPureFunctions(CustomTestCase):
                 pass
 
         result = _obj_to_dict(Obj())
-        self.assertEqual(result["x"], 10)
+        assert result["x"] == 10
-        self.assertNotIn("method", result)
+        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="]:
-            self.assertIn(key, info)
+            assert key in info
 
-        self.assertIn("value=42", get_tensor_info(42))
+        assert "value=42" in get_tensor_info(42)
-        self.assertIn("min=None", get_tensor_info(torch.tensor([])))
+        assert "min=None" in get_tensor_info(torch.tensor([]))
 
 
-class TestDumperDistributed(CustomTestCase):
+class TestDumperDistributed:
-    def test_basic(self):
+    def test_basic(self, tmp_path):
-        with tempfile.TemporaryDirectory(prefix="test_dumper_") as tmpdir:
+        run_distributed_test(self._test_basic_func, tmpdir=str(tmp_path))
-            _run_distributed_test(_test_basic_func, tmpdir=tmpdir)
+
+    @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}")
+
+        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_enable(self):
-        _run_distributed_test(_test_http_func)
+        run_distributed_test(self._test_http_func)
 
-    def test_filter(self):
+    @staticmethod
-        with tempfile.TemporaryDirectory(prefix="test_dumper_") as tmpdir:
+    def _test_http_func(rank):
-            _run_distributed_test(_test_filter_func, tmpdir=tmpdir)
+        os.environ["SGLANG_DUMPER_ENABLE"] = "0"
+        from sglang.srt.debug_utils.dumper import dumper
 
-    def test_write_disabled(self):
+        assert not dumper._enable
-        with tempfile.TemporaryDirectory(prefix="test_dumper_") as tmpdir:
+        dumper.on_forward_pass_start()
-            _run_distributed_test(_test_write_disabled_func, tmpdir=tmpdir)
 
+        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_basic_func(rank, tmpdir):
+    def test_filter(self, tmp_path):
-    os.environ["SGLANG_DUMPER_DIR"] = tmpdir
+        run_distributed_test(self._test_filter_func, tmpdir=str(tmp_path))
-    from sglang.srt.debug_utils.dumper import dumper
 
-    tensor = torch.randn(10, 10, device=f"cuda:{rank}")
+    @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()
+        dumper.on_forward_pass_start()
-    dumper.dump("tensor_a", tensor, arg=100)
+        dumper.dump("keep_this", torch.randn(5, device=f"cuda:{rank}"))
+        dumper.dump("skip_this", torch.randn(5, device=f"cuda:{rank}"))
+        dumper.dump("not_keep_this", torch.randn(5, device=f"cuda:{rank}"))
 
-    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:
+        filenames = _get_filenames(tmpdir)
-            time.sleep(0.1)
+        _assert_files(
-            requests.post(
+            filenames,
-                "http://localhost:40000/dumper", json={"enable": enable}
+            exist=["keep_this"],
-            ).raise_for_status()
+            not_exist=["skip_this", "not_keep_this"],
+        )
+
+    def test_write_disabled(self, tmp_path):
+        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()
+        dumper.dump("no_write", torch.randn(5, device=f"cuda:{rank}"))
+
         dist.barrier()
-        assert dumper._enable == enable
+        assert len(_get_filenames(tmpdir)) == 0
-
-
-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):
@@ -160,47 +158,5 @@ def _assert_files(filenames, *, exist=(), not_exist=()):
         ), 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()
+    sys.exit(pytest.main([__file__]))

test/registered/scheduler/test_prefill_delayer.py

@@ -3,7 +3,6 @@ import os
 import re
 import time
 import unittest
-from collections import defaultdict
 from dataclasses import dataclass
 from types import SimpleNamespace
 from typing import List, Optional
@@ -11,7 +10,6 @@ from typing import List, Optional
 import openai
 import requests
 import torch
-import torch.multiprocessing as mp
 
 from sglang.bench_serving import run_benchmark
 from sglang.srt.managers.prefill_delayer import PrefillDelayer
@@ -25,6 +23,7 @@ from sglang.test.test_utils import (
     CustomTestCase,
     get_benchmark_args,
     popen_launch_server,
+    run_distributed_test,
 )
 
 register_cuda_ci(
@@ -54,13 +53,8 @@ class NegotiateTestCase:
     expected_reason: str
 
 
-def _run_negotiate_test(rank, world_size, test_cases, results_queue, port):
+def _run_negotiate_test(rank, test_cases):
-    torch.distributed.init_process_group(
+    world_size = torch.distributed.get_world_size()
-        backend="gloo",
-        init_method=f"tcp://127.0.0.1:{port}",
-        world_size=world_size,
-        rank=rank,
-    )
     cpu_group = torch.distributed.new_group(backend="gloo")
 
     for case in test_cases:
@@ -83,9 +77,10 @@ def _run_negotiate_test(rank, world_size, test_cases, results_queue, port):
                 token_usage=call.token_usage[rank],
             )
 
-        results_queue.put((rank, case.name, result.output_allow, result.output_reason))
+        assert (result.output_allow, result.output_reason) == (
-
+            case.expected_allow,
-    torch.distributed.destroy_process_group()
+            case.expected_reason,
+        ), f"Case {case.name} rank {rank}"
 
 
 _NEGOTIATE_TEST_CASES = [
@@ -210,38 +205,12 @@ _NEGOTIATE_TEST_CASES = [
 
 class TestPrefillDelayerNegotiate(unittest.TestCase):
     def test_negotiate(self):
-        world_size = 4
+        run_distributed_test(
-        test_cases = _NEGOTIATE_TEST_CASES
+            _run_negotiate_test,
-
+            world_size=4,
-        ctx = mp.get_context("spawn")
+            backend="gloo",
-        results_queue = ctx.Queue()
+            test_cases=_NEGOTIATE_TEST_CASES,
-        port = 29500 + os.getpid() % 1000
+        )
-
-        processes = []
-        for rank in range(world_size):
-            p = ctx.Process(
-                target=_run_negotiate_test,
-                args=(rank, world_size, test_cases, results_queue, port),
-            )
-            p.start()
-            processes.append(p)
-
-        for p in processes:
-            p.join()
-
-        results = defaultdict(dict)
-        for _ in range(world_size * len(test_cases)):
-            rank, case_name, output_allow, output_reason = results_queue.get()
-            results[case_name][rank] = (output_allow, output_reason)
-
-        for case in test_cases:
-            for rank in range(world_size):
-                output_allow, output_reason = results[case.name][rank]
-                self.assertEqual(
-                    (output_allow, output_reason),
-                    (case.expected_allow, case.expected_reason),
-                    f"Case {case.name} rank {rank}",
-                )
 
 
 # ============================ E2E Tests ============================