[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:
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user