chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Tiny support sticky routing algorithm in schedule simulator (#16355)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-04 08:36:52 +08:00
committed by GitHub
parent 216ea910ad
commit 8ca9597057
5 changed files with 124 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

85
test/registered/debug_utils/test_schedule_simulator.py
View File

@@ -15,6 +15,7 @@ from sglang.srt.debug_utils.schedule_simulator import (
SimulationResult,
Simulator,
StepRecord,
StickyRouter,
generate_gsp_requests,
generate_random_requests,
load_from_request_logger,
@@ -154,6 +155,42 @@ class TestRouters(CustomTestCase):
results = [router.route(req, gpu_states) for _ in range(100)]
self.assertTrue(all(0 <= r < 4 for r in results))
def test_sticky_router_same_group_same_gpu(self):
router = StickyRouter(num_gpus=4)
gpu_states = [GPUState(gpu_id=i, max_total_tokens=10000) for i in range(4)]
reqs = [
SimRequest(request_id=f"r{i}", input_len=100, output_len=50, group_id="g0")
for i in range(10)
]
results = [router.route(req, gpu_states) for req in reqs]
self.assertEqual(len(set(results)), 1)
def test_sticky_router_no_group_fallback(self):
router = StickyRouter(num_gpus=4)
gpu_states = [GPUState(gpu_id=i, max_total_tokens=10000) for i in range(4)]
reqs = [
SimRequest(request_id=f"r{i}", input_len=100, output_len=50)
for i in range(100)
]
results = [router.route(req, gpu_states) for req in reqs]
self.assertTrue(all(0 <= r < 4 for r in results))
def test_sticky_router_multiple_groups(self):
router = StickyRouter(num_gpus=4)
gpu_states = [GPUState(gpu_id=i, max_total_tokens=10000) for i in range(4)]
for group_id in ["g0", "g1", "g2"]:
reqs = [
SimRequest(
request_id=f"{group_id}_r{i}",
input_len=100,
output_len=50,
group_id=group_id,
)
for i in range(5)
]
results = [router.route(req, gpu_states) for req in reqs]
self.assertEqual(len(set(results)), 1)
class TestFIFOScheduler(CustomTestCase):
def test_runs_pending_requests(self):
@@ -508,6 +545,54 @@ class TestCLI(CustomTestCase):
self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}")
self.assertIn("router=random", result.stdout)
def test_cli_sticky_router(self):
result = self._run_cli(
"--synth-gsp",
"--synth-gsp-num-groups",
"2",
"--synth-gsp-prompts-per-group",
"3",
"--synth-seed",
"42",
"--num-gpus",
"4",
"--router",
"sticky",
)
self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}")
self.assertIn("router=sticky", result.stdout)
def test_e2e_sticky_router_group_locality(self):
result = self._run_cli(
"--synth-gsp",
"--synth-gsp-num-groups",
"2",
"--synth-gsp-prompts-per-group",
"2",
"--synth-gsp-system-prompt-len",
"10",
"--synth-gsp-question-len",
"10",
"--synth-gsp-output-len",
"2",
"--synth-seed",
"123",
"--num-gpus",
"2",
"--router",
"sticky",
"--max-total-tokens",
"1000",
"--log-level",
"2",
)
self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}")
for expected in [
"step=0 | GPU0[R=2:gsp0,gsp1 Q=0:-] | GPU1[R=2:gsp2,gsp3 Q=0:-]",
"step=1 | GPU0[R=0:- Q=0:-] | GPU1[R=0:- Q=0:-]",
]:
self.assertIn(expected, result.stdout)
def test_cli_synthetic(self):
result = self._run_cli(
"--synthetic",