[router] Refactor router and policy traits with dependency injection (#7987)

Co-authored-by: Jin Pan <jpan236@wisc.edu>
Co-authored-by: Keru Yang <rukeyang@gmail.com>
Co-authored-by: Yingyi Huang <yingyihuang2000@outlook.com>
Co-authored-by: Philip Zhu <phlipzhux@gmail.com>
This commit is contained in:
Simo Lin
2025-07-18 14:24:24 -07:00
committed by GitHub
parent 1f76fc8747
commit c8f31042a8
24 changed files with 3190 additions and 1944 deletions

View File

@@ -164,56 +164,47 @@ class TestLaunchRouter(unittest.TestCase):
"""Test that policy validation works correctly for PD and regular modes."""
from sglang_router.launch_router import RouterArgs, launch_router
# Test 1: PowerOfTwo is only valid in PD mode
# Test 1: PowerOfTwo requires at least 2 workers
args = self.create_router_args(
pd_disaggregation=False,
policy="power_of_two",
worker_urls=["http://localhost:8000"],
worker_urls=["http://localhost:8000"], # Only 1 worker
)
# Should raise error
with self.assertRaises(ValueError) as cm:
launch_router(args)
self.assertIn(
"PowerOfTwo policy is only supported in PD disaggregated mode",
"Power-of-two policy requires at least 2 workers",
str(cm.exception),
)
# Test 2: RoundRobin is not valid in PD mode
# Test 2: PowerOfTwo with sufficient workers should succeed
args = self.create_router_args(
pd_disaggregation=True,
policy="round_robin",
prefill=[["http://prefill1:8080", "9000"]],
decode=[["http://decode1:8081"]],
worker_urls=[],
pd_disaggregation=False,
policy="power_of_two",
worker_urls=["http://localhost:8000", "http://localhost:8001"], # 2 workers
)
# This should not raise an error (validation passes)
# Should raise error
with self.assertRaises(ValueError) as cm:
launch_router(args)
self.assertIn(
"RoundRobin policy is not supported in PD disaggregated mode",
str(cm.exception),
)
# Test 3: Valid combinations should not raise errors
# Test 3: All policies now work in both modes
# Regular mode with RoundRobin
args = self.create_router_args(
pd_disaggregation=False,
policy="round_robin",
worker_urls=["http://localhost:8000"],
)
# This should not raise (though it may fail to connect)
# This should not raise validation error
# PD mode with PowerOfTwo
# PD mode with RoundRobin (now supported!)
args = self.create_router_args(
pd_disaggregation=True,
policy="power_of_two",
policy="round_robin",
prefill=[["http://prefill1:8080", "9000"]],
decode=[["http://decode1:8081"]],
worker_urls=[],
)
# This should not raise (though it may fail to connect)
# This should not raise validation error
def test_pd_service_discovery_args_parsing(self):
"""Test PD service discovery CLI argument parsing."""