From 407106300e4895fc2acb30540833ef5d225ae715 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 28 May 2026 16:50:46 -0700 Subject: [PATCH] test/disaggregation: dp-attention keeps total_tokens e2e, simpler LB algos -> unit tests (#26617) --- .../test_disaggregation_dp_attention.py | 93 +----- .../managers/test_data_parallel_controller.py | 282 ++++++++++++++++++ .../unit/managers/test_dp_budget.py | 91 ------ 3 files changed, 289 insertions(+), 177 deletions(-) create mode 100644 test/registered/unit/managers/test_data_parallel_controller.py delete mode 100644 test/registered/unit/managers/test_dp_budget.py diff --git a/test/registered/disaggregation/test_disaggregation_dp_attention.py b/test/registered/disaggregation/test_disaggregation_dp_attention.py index d7fba1c27..a9662130c 100644 --- a/test/registered/disaggregation/test_disaggregation_dp_attention.py +++ b/test/registered/disaggregation/test_disaggregation_dp_attention.py @@ -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() diff --git a/test/registered/unit/managers/test_data_parallel_controller.py b/test/registered/unit/managers/test_data_parallel_controller.py new file mode 100644 index 000000000..23a738bcb --- /dev/null +++ b/test/registered/unit/managers/test_data_parallel_controller.py @@ -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() diff --git a/test/registered/unit/managers/test_dp_budget.py b/test/registered/unit/managers/test_dp_budget.py deleted file mode 100644 index 59d35c294..000000000 --- a/test/registered/unit/managers/test_dp_budget.py +++ /dev/null @@ -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()