minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

test/disaggregation: dp-attention keeps total_tokens e2e, simpler LB algos -> unit tests (#26617)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-05-28 16:50:46 -07:00
committed by GitHub
parent 3e255fd493
commit 407106300e
3 changed files with 289 additions and 177 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/disaggregation/test_disaggregation_dp_attention.py
+7 -86
View File
@@ -20,9 +20,15 @@ register_cuda_ci(est_time=443, stage="base-c", runner_config="8-gpu-h20")
class TestDisaggregationDPAttention(PDDisaggregationServerBase):
"""PD-disagg + DP-attention e2e on `total_tokens` LB — the most complex
dispatch (token accounting + tie-break + estimated_tokens). Simpler
algorithms are unit-tested in
test/registered/unit/managers/test_data_parallel_controller.py.
"""
PREFILL_DP_SIZE = 4
DECODE_DP_SIZE = 4
LOAD_BALANCE_METHOD = "auto"
LOAD_BALANCE_METHOD = "total_tokens"
@classmethod
def setUpClass(cls):
@@ -107,11 +113,6 @@ class TestDisaggregationDPAttention(PDDisaggregationServerBase):
self.assertGreater(metrics["score"], 0.60)
class TestDisaggregationDPAttentionRoundRobin(TestDisaggregationDPAttention):
LOAD_BALANCE_METHOD = "round_robin"
# TODO: add a balancedness metric
def test_bench_serving(self):
args = get_benchmark_args(
base_url=f"http://{self.base_host}:{self.lb_port}",
@@ -129,85 +130,5 @@ class TestDisaggregationDPAttentionRoundRobin(TestDisaggregationDPAttention):
self.assertEqual(result["completed"], 1000)
class TestDisaggregationDPAttentionTotalRequests(TestDisaggregationDPAttention):
LOAD_BALANCE_METHOD = "total_requests"
test_gsm8k = unittest.skip(
"Covered by base class; this class targets total_requests path."
)(TestDisaggregationDPAttention.test_gsm8k)
def test_bench_serving(self):
args = get_benchmark_args(
base_url=f"http://{self.base_host}:{self.lb_port}",
dataset_name="random",
tokenizer=self.model,
num_prompts=256,
random_input_len=2048,
random_output_len=512,
request_rate=float("inf"),
max_concurrency=128,
)
result = run_benchmark(args)
self.assertEqual(result["completed"], 256)
class TestDisaggregationDPAttentionTotalTokens(TestDisaggregationDPAttention):
LOAD_BALANCE_METHOD = "total_tokens"
test_gsm8k = unittest.skip(
"Covered by base class; this class targets total_tokens path."
)(TestDisaggregationDPAttention.test_gsm8k)
def test_bench_serving(self):
args = get_benchmark_args(
base_url=f"http://{self.base_host}:{self.lb_port}",
dataset_name="random",
tokenizer=self.model,
num_prompts=256,
random_input_len=2048,
random_output_len=512,
request_rate=float("inf"),
max_concurrency=128,
)
result = run_benchmark(args)
self.assertEqual(result["completed"], 256)
@unittest.skip(
"Skip this test until new testing logic in mini-lb has been updated in docker image."
)
class TestDisaggregationDPAttentionExternalRouting(TestDisaggregationDPAttention):
"""Test external DP rank assignment via mini-lb --test-external-dp-routing.
NOTE: In PD disaggregation the response comes from the decode server,
so meta_info["dp_rank"] reflects the decode-side DP rank. Prefill DP
rank correctness is verified implicitly — if the wrong prefill DP
worker were used, KV transfer would fail and the request would error.
The mini-lb internally verifies meta_info["dp_rank"] matches the
assigned decode dp_rank; a mismatch returns HTTP 500.
"""
@classmethod
def launch_lb(cls):
from sglang.test.test_utils import popen_with_error_check
lb_command = [
"python3",
"-m",
"sglang_router.launch_router",
"--pd-disaggregation",
"--mini-lb",
"--test-external-dp-routing",
"--prefill",
cls.prefill_url,
"--decode",
cls.decode_url,
"--host",
cls.base_host,
"--port",
cls.lb_port,
]
cls.process_lb = popen_with_error_check(lb_command)
cls.wait_server_ready(cls.lb_url + "/health", process=cls.process_lb)
if __name__ == "__main__":
unittest.main()
test/registered/unit/managers/test_data_parallel_controller.py
+282
View File
@@ -0,0 +1,282 @@
"""DPBudget + DataParallelController dispatch tests.
`total_tokens` (the most complex algorithm) is exercised end-to-end in
test/registered/disaggregation/test_disaggregation_dp_attention.py; its
tie-break on `total_requests` transitively covers that state.
Fragility: scheduler tests bypass `DataParallelController.__init__` via
`__new__` and inject only the attrs the schedulers read (`workers`,
`status`, `round_robin_counter`, `dp_budget`). Update `_make_controller`
if a scheduler starts reading another attr. `maybe_external_dp_rank_routing`
is exercised as the real method, no mock.
"""
import dataclasses
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
maybe_stub_sgl_kernel()
from sglang.srt.managers.data_parallel_controller import (
DataParallelController,
DPBudget,
LoadBalanceMethod,
)
from sglang.srt.managers.io_struct import GetLoadsReqOutput, WatchLoadUpdateReq
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
_BASE_LOAD = GetLoadsReqOutput(
dp_rank=0,
timestamp=0.0,
num_running_reqs=0,
num_waiting_reqs=0,
num_used_tokens=0,
num_total_tokens=0,
max_total_num_tokens=4096,
token_usage=0.0,
gen_throughput=0.0,
cache_hit_rate=0.0,
utilization=0.0,
max_running_requests=128,
)
def _load(**overrides) -> GetLoadsReqOutput:
return dataclasses.replace(_BASE_LOAD, **overrides)
def _make_controller(dp_size: int) -> DataParallelController:
"""Bypass __init__; inject only the attrs dispatch methods read."""
ctl = DataParallelController.__new__(DataParallelController)
ctl.workers = [MagicMock(name=f"worker_{i}") for i in range(dp_size)]
ctl.status = [True] * dp_size
ctl.round_robin_counter = 0
ctl.dp_budget = DPBudget(dp_size=dp_size)
return ctl
def _req(routed_dp_rank=None, bootstrap_room=None, input_ids=None):
"""Req stand-in; SimpleNamespace avoids pinning to the Req dataclass schema."""
return SimpleNamespace(
routed_dp_rank=routed_dp_rank,
bootstrap_room=bootstrap_room,
input_ids=input_ids or [],
)
class TestDPBudgetUpdateBudget(CustomTestCase):
def test_maps_running_plus_waiting_to_total_requests(self):
budget = DPBudget(dp_size=2)
budget.update_budget(
WatchLoadUpdateReq(
loads=[
_load(dp_rank=0, num_running_reqs=3, num_waiting_reqs=2),
_load(dp_rank=1, num_running_reqs=5, num_waiting_reqs=1),
]
)
)
self.assertEqual(budget.total_requests, [5, 6])
def test_maps_num_total_tokens_not_num_used_tokens(self):
# Reads num_total_tokens (used + pending prefill), NOT num_used_tokens.
# A silent swap here would break DP balance for long-prompt workloads.
budget = DPBudget(dp_size=2)
budget.update_budget(
WatchLoadUpdateReq(
loads=[
_load(dp_rank=0, num_used_tokens=100, num_total_tokens=150),
_load(dp_rank=1, num_used_tokens=80, num_total_tokens=80),
]
)
)
self.assertEqual(budget.total_tokens, [150, 80])
def test_partial_update_only_affects_reported_rank(self):
budget = DPBudget(dp_size=3)
budget.total_requests = [10, 20, 30]
budget.total_tokens = [100, 200, 300]
budget.update_budget(
WatchLoadUpdateReq(
loads=[
_load(
dp_rank=1,
num_running_reqs=1,
num_waiting_reqs=1,
num_total_tokens=50,
)
]
)
)
self.assertEqual(budget.total_requests, [10, 2, 30])
self.assertEqual(budget.total_tokens, [100, 50, 300])
class TestDPBudgetDispatch(CustomTestCase):
"""DPBudget.dispatch picks a rank from current state and updates counters."""
def test_total_requests_dispatch_picks_min_and_increments(self):
budget = DPBudget(dp_size=3)
budget.total_requests = [4, 2, 7]
rank = budget.dispatch(LoadBalanceMethod.TOTAL_REQUESTS)
self.assertEqual(rank, 1)
self.assertEqual(
budget.total_requests[1],
3,
"dispatch should increment chosen worker's request count",
)
def test_total_tokens_dispatch_applies_estimated_tokens(self):
budget = DPBudget(dp_size=3)
budget.total_tokens = [100, 50, 200]
budget.total_requests = [0, 0, 0]
rank = budget.dispatch(LoadBalanceMethod.TOTAL_TOKENS, estimated_tokens=30)
self.assertEqual(rank, 1, "should pick worker with min total_tokens")
self.assertEqual(
budget.total_tokens[1],
80,
"dispatch should add estimated_tokens to chosen worker",
)
self.assertEqual(
budget.total_requests[1],
1,
"dispatch should also increment request count",
)
def test_total_tokens_tie_breaks_on_total_requests(self):
budget = DPBudget(dp_size=3)
budget.total_tokens = [50, 50, 50]
budget.total_requests = [4, 2, 7]
rank = budget.dispatch(LoadBalanceMethod.TOTAL_TOKENS, estimated_tokens=10)
self.assertEqual(
rank, 1, "tie on total_tokens should fall back to min total_requests"
)
def test_dispatch_returns_none_for_methods_not_handled(self):
"""Round-robin and follow_bootstrap_room dispatch elsewhere; DPBudget
only handles the load-aware variants."""
budget = DPBudget(dp_size=3)
self.assertIsNone(budget.dispatch(LoadBalanceMethod.ROUND_ROBIN))
self.assertIsNone(budget.dispatch(LoadBalanceMethod.FOLLOW_BOOTSTRAP_ROOM))
class TestRoundRobinScheduler(CustomTestCase):
def test_cycles_through_active_workers_in_order(self):
ctl = _make_controller(dp_size=4)
for _ in range(8):
ctl.round_robin_scheduler(_req())
# 8 reqs across 4 active workers — 2 each, in round-robin order
for i, worker in enumerate(ctl.workers):
self.assertEqual(worker.send_pyobj.call_count, 2, f"worker {i} call count")
def test_first_dispatch_picks_worker_zero(self):
ctl = _make_controller(dp_size=4)
ctl.round_robin_scheduler(_req())
ctl.workers[0].send_pyobj.assert_called_once()
for i in (1, 2, 3):
ctl.workers[i].send_pyobj.assert_not_called()
def test_skips_inactive_workers(self):
ctl = _make_controller(dp_size=4)
ctl.status[1] = False
ctl.status[3] = False
for _ in range(6):
ctl.round_robin_scheduler(_req())
# Only workers 0 and 2 are active — should split 6 reqs evenly
self.assertEqual(ctl.workers[0].send_pyobj.call_count, 3)
ctl.workers[1].send_pyobj.assert_not_called()
self.assertEqual(ctl.workers[2].send_pyobj.call_count, 3)
ctl.workers[3].send_pyobj.assert_not_called()
def test_routed_dp_rank_bypasses_counter(self):
"""External dp-rank routing must not advance the counter."""
ctl = _make_controller(dp_size=4)
ctl.round_robin_scheduler(_req(routed_dp_rank=2))
ctl.workers[2].send_pyobj.assert_called_once()
self.assertEqual(
ctl.round_robin_counter,
0,
"external routing must not advance the round-robin counter",
)
# Subsequent round-robin req still lands on worker 0
ctl.round_robin_scheduler(_req())
ctl.workers[0].send_pyobj.assert_called_once()
class TestFollowBootstrapRoomScheduler(CustomTestCase):
def test_dispatches_by_bootstrap_room_modulo(self):
ctl = _make_controller(dp_size=4)
for room, expected_rank in [
(0, 0),
(1, 1),
(4, 0),
(5, 1),
(100, 0),
(101, 1),
]:
ctl.follow_bootstrap_room_scheduler(_req(bootstrap_room=room))
ctl.workers[expected_rank].send_pyobj.assert_called()
def test_requires_bootstrap_room(self):
ctl = _make_controller(dp_size=4)
with self.assertRaises(AssertionError):
ctl.follow_bootstrap_room_scheduler(_req(bootstrap_room=None))
def test_routed_dp_rank_bypasses_bootstrap_room(self):
ctl = _make_controller(dp_size=4)
ctl.follow_bootstrap_room_scheduler(_req(routed_dp_rank=3, bootstrap_room=1))
ctl.workers[3].send_pyobj.assert_called_once()
ctl.workers[1].send_pyobj.assert_not_called()
class TestTotalRequestsScheduler(CustomTestCase):
def test_dispatches_to_min_request_worker(self):
ctl = _make_controller(dp_size=4)
ctl.dp_budget.total_requests = [5, 3, 1, 4]
ctl.total_requests_scheduler(_req())
ctl.workers[2].send_pyobj.assert_called_once()
for i in (0, 1, 3):
ctl.workers[i].send_pyobj.assert_not_called()
self.assertEqual(
ctl.dp_budget.total_requests[2],
2,
"DPBudget must record the dispatch by incrementing the counter",
)
def test_routed_dp_rank_bypasses_budget(self):
ctl = _make_controller(dp_size=4)
ctl.dp_budget.total_requests = [5, 3, 1, 4]
ctl.total_requests_scheduler(_req(routed_dp_rank=0))
ctl.workers[0].send_pyobj.assert_called_once()
# DPBudget must not be touched when bypassed
self.assertEqual(
ctl.dp_budget.total_requests,
[5, 3, 1, 4],
"external routing must not mutate DPBudget state",
)
class TestStatusAwarenessInconsistency(CustomTestCase):
"""Document a divergence: ``round_robin_scheduler`` skips workers whose
``status`` is False, but ``total_requests_scheduler`` /
``total_tokens_scheduler`` route purely by DPBudget — they do NOT
consult ``self.status``. If a future change unifies this behaviour,
this test will fail and force a reviewer to confirm intent."""
def test_total_requests_ignores_status(self):
ctl = _make_controller(dp_size=4)
# Worker 2 is the global minimum AND marked inactive.
ctl.dp_budget.total_requests = [5, 3, 1, 4]
ctl.status[2] = False
ctl.total_requests_scheduler(_req())
# Current behaviour: still dispatches to the inactive worker.
ctl.workers[2].send_pyobj.assert_called_once()
if __name__ == "__main__":
unittest.main()
test/registered/unit/managers/test_dp_budget.py
-91
View File
@@ -1,91 +0,0 @@
"""Unit tests for DPBudget — field mapping regression guard.
This PR changed DPBudget.update_budget to read num_running_reqs +
num_waiting_reqs and num_total_tokens from the new GetLoadsReqOutput.
These tests lock in that mapping. Pre-existing dispatch logic is not
retested here — it's covered by DP balance integration tests.
"""
import dataclasses
import unittest
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
maybe_stub_sgl_kernel()
from sglang.srt.managers.data_parallel_controller import DPBudget
from sglang.srt.managers.io_struct import GetLoadsReqOutput, WatchLoadUpdateReq
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
_BASE_LOAD = GetLoadsReqOutput(
dp_rank=0,
timestamp=0.0,
num_running_reqs=0,
num_waiting_reqs=0,
num_used_tokens=0,
num_total_tokens=0,
max_total_num_tokens=4096,
token_usage=0.0,
gen_throughput=0.0,
cache_hit_rate=0.0,
utilization=0.0,
max_running_requests=128,
)
def _load(**overrides) -> GetLoadsReqOutput:
return dataclasses.replace(_BASE_LOAD, **overrides)
class TestDPBudgetUpdateBudget(CustomTestCase):
def test_maps_running_plus_waiting_to_total_requests(self):
budget = DPBudget(dp_size=2)
budget.update_budget(
WatchLoadUpdateReq(
loads=[
_load(dp_rank=0, num_running_reqs=3, num_waiting_reqs=2),
_load(dp_rank=1, num_running_reqs=5, num_waiting_reqs=1),
]
)
)
self.assertEqual(budget.total_requests, [5, 6])
def test_maps_num_total_tokens_not_num_used_tokens(self):
# Reads num_total_tokens (used + pending prefill), NOT num_used_tokens.
# A silent swap here would break DP balance for long-prompt workloads.
budget = DPBudget(dp_size=2)
budget.update_budget(
WatchLoadUpdateReq(
loads=[
_load(dp_rank=0, num_used_tokens=100, num_total_tokens=150),
_load(dp_rank=1, num_used_tokens=80, num_total_tokens=80),
]
)
)
self.assertEqual(budget.total_tokens, [150, 80])
def test_partial_update_only_affects_reported_rank(self):
budget = DPBudget(dp_size=3)
budget.total_requests = [10, 20, 30]
budget.total_tokens = [100, 200, 300]
budget.update_budget(
WatchLoadUpdateReq(
loads=[
_load(
dp_rank=1,
num_running_reqs=1,
num_waiting_reqs=1,
num_total_tokens=50,
)
]
)
)
self.assertEqual(budget.total_requests, [10, 2, 30])
self.assertEqual(budget.total_tokens, [100, 50, 300])
if __name__ == "__main__":
unittest.main()