chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
fzyzcjy
2026-02-15 18:33:23 +08:00
committed by GitHub
parent b992828ad2
commit 4c7f986c6b
3 changed files with 160 additions and 185 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
python/sglang/test/test_utils.py
View File

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

238
test/registered/debug_utils/test_dumper.py
View File

@@ -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))
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))
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
assert get_truncated_value(torch.randn(10, 10)).shape == (10, 10)
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
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)
self.assertNotIn("method", result)
assert result["x"] == 10
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))
self.assertIn("min=None", get_tensor_info(torch.tensor([])))
assert "value=42" in get_tensor_info(42)
assert "min=None" in 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)
class TestDumperDistributed:
def test_basic(self, 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}")
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):
with tempfile.TemporaryDirectory(prefix="test_dumper_") as tmpdir:
_run_distributed_test(_test_filter_func, tmpdir=tmpdir)
@staticmethod
def _test_http_func(rank):
os.environ["SGLANG_DUMPER_ENABLE"] = "0"
from sglang.srt.debug_utils.dumper import dumper
def test_write_disabled(self):
with tempfile.TemporaryDirectory(prefix="test_dumper_") as tmpdir:
_run_distributed_test(_test_write_disabled_func, tmpdir=tmpdir)
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_basic_func(rank, tmpdir):
os.environ["SGLANG_DUMPER_DIR"] = tmpdir
from sglang.srt.debug_utils.dumper import dumper
def test_filter(self, tmp_path):
run_distributed_test(self._test_filter_func, tmpdir=str(tmp_path))
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.dump("tensor_a", tensor, arg=100)
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}"))
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:
time.sleep(0.1)
requests.post(
"http://localhost:40000/dumper", json={"enable": enable}
).raise_for_status()
filenames = _get_filenames(tmpdir)
_assert_files(
filenames,
exist=["keep_this"],
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
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
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__]))

57
test/registered/scheduler/test_prefill_delayer.py
View File

@@ -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):
torch.distributed.init_process_group(
backend="gloo",
init_method=f"tcp://127.0.0.1:{port}",
world_size=world_size,
rank=rank,
)
def _run_negotiate_test(rank, test_cases):
world_size = torch.distributed.get_world_size()
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))
torch.distributed.destroy_process_group()
assert (result.output_allow, result.output_reason) == (
case.expected_allow,
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
test_cases = _NEGOTIATE_TEST_CASES
ctx = mp.get_context("spawn")
results_queue = ctx.Queue()
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}",
)
run_distributed_test(
_run_negotiate_test,
world_size=4,
backend="gloo",
test_cases=_NEGOTIATE_TEST_CASES,
)
# ============================ E2E Tests ============================