283 lines
10 KiB
Python
283 lines
10 KiB
Python
"""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 unittest
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock
|
|
|
|
import msgspec.structs
|
|
|
|
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.load_snapshot import LoadSnapshot
|
|
|
|
register_cpu_ci(est_time=11, suite="base-a-test-cpu")
|
|
|
|
|
|
_BASE_LOAD = msgspec.structs.replace(
|
|
LoadSnapshot(dp_rank=0),
|
|
max_total_num_tokens=4096,
|
|
max_running_requests=128,
|
|
)
|
|
|
|
|
|
def _load(**overrides) -> LoadSnapshot:
|
|
return msgspec.structs.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(
|
|
[
|
|
_load(dp_rank=0, timestamp=1.0, num_running_reqs=3, num_waiting_reqs=2),
|
|
_load(dp_rank=1, timestamp=1.0, 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):
|
|
budget = DPBudget(dp_size=2)
|
|
budget.update_budget(
|
|
[
|
|
_load(
|
|
dp_rank=0, timestamp=1.0, num_used_tokens=100, num_total_tokens=150
|
|
),
|
|
_load(
|
|
dp_rank=1, timestamp=1.0, 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.update_budget(
|
|
[
|
|
_load(
|
|
dp_rank=0, timestamp=1.0, num_running_reqs=10, num_total_tokens=100
|
|
),
|
|
_load(
|
|
dp_rank=1, timestamp=1.0, num_running_reqs=20, num_total_tokens=200
|
|
),
|
|
_load(
|
|
dp_rank=2, timestamp=1.0, num_running_reqs=30, num_total_tokens=300
|
|
),
|
|
]
|
|
)
|
|
budget.update_budget(
|
|
[
|
|
_load(
|
|
dp_rank=1,
|
|
timestamp=2.0,
|
|
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()
|