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

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

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

test/registered/unit/managers/test_dp_budget.py

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