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

Add routing key based schedule policy (#16840)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-10 10:16:23 +08:00
committed by GitHub
parent 15da3061bc
commit 1f0ea4f958
5 changed files with 293 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
python/sglang/srt/managers/schedule_policy.py
View File

@@ -1,6 +1,12 @@
from __future__ import annotations
import logging
from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor
from sglang.srt.utils import get_bool_env_var
_ROUTING_KEY_POLICY_DEBUG_LOG = get_bool_env_var("SGLANG_ROUTING_KEY_POLICY_DEBUG_LOG")
logger = logging.getLogger(__name__)
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
@@ -19,7 +25,7 @@ from sglang.srt.managers.prefill_delayer import PrefillDelayerSinglePassExecutor
import os
import random
from collections import defaultdict
from collections import Counter, defaultdict
from contextlib import contextmanager
from enum import Enum, auto
from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
@@ -78,6 +84,7 @@ class CacheAgnosticPolicy(Enum):
FCFS = "fcfs" # first come first serve
LOF = "lof" # longest output first
RANDOM = "random"
ROUTING_KEY = "routing-key" # prioritize by routing key frequency in running batch
class SchedulePolicy:
@@ -101,7 +108,9 @@ class SchedulePolicy:
# It is used to find the matching prefix for in-batch prefix caching.
self.waiting_queue_radix_tree = RadixCache.create_simulated()
def calc_priority(self, waiting_queue: List[Req]) -> bool:
def calc_priority(
self, waiting_queue: List[Req], running_batch: Optional[ScheduleBatch] = None
) -> bool:
if self.policy == CacheAgnosticPolicy.FCFS:
if self.enable_priority_scheduling:
SchedulePolicy._sort_by_priority_and_fcfs(
@@ -136,6 +145,9 @@ class SchedulePolicy:
)
elif policy == CacheAgnosticPolicy.RANDOM:
SchedulePolicy._sort_randomly(waiting_queue)
elif policy == CacheAgnosticPolicy.ROUTING_KEY:
if running_batch is not None:
SchedulePolicy._sort_by_routing_key(waiting_queue, running_batch)
else:
raise ValueError(f"Unknown CacheAgnostic Policy: {policy=}")
return prefix_computed
@@ -288,6 +300,39 @@ class SchedulePolicy:
)
)
@staticmethod
def _sort_by_routing_key(
waiting_queue: List[Req], running_batch: ScheduleBatch
) -> None:
"""Sorts waiting queue by routing key frequency in running batch."""
routing_key_counts = Counter(
r.routing_key for r in running_batch.reqs if r.routing_key
)
if _ROUTING_KEY_POLICY_DEBUG_LOG:
waiting_keys_before = [r.routing_key for r in waiting_queue]
logger.info(
f"routing_key_counts={dict(routing_key_counts)}, "
f"waiting_keys_before={waiting_keys_before}"
)
if not routing_key_counts:
return
def sort_key(req: Req):
key = req.routing_key
if key and key in routing_key_counts:
count = routing_key_counts[key]
return (0, -count, key)
else:
return (1, 0, key or "")
waiting_queue.sort(key=sort_key)
if _ROUTING_KEY_POLICY_DEBUG_LOG:
waiting_keys_after = [r.routing_key for r in waiting_queue]
logger.info(f"waiting_keys_after={waiting_keys_after}")
@staticmethod
def _calc_weight(cur_node: TreeNode, node_to_weight: Dict[TreeNode, int]) -> None:
for child in cur_node.children.values():

3
python/sglang/srt/managers/scheduler.py
View File

@@ -1433,6 +1433,7 @@ class Scheduler(
metrics_collector=(
self.metrics_collector if self.enable_metrics else None
),
routing_key=recv_req.routing_key,
http_worker_ipc=recv_req.http_worker_ipc,
dllm_config=self.dllm_config,
)
@@ -1907,7 +1908,7 @@ class Scheduler(
self.tree_cache.check_hicache_events()
# Get priority queue
self.policy.calc_priority(self.waiting_queue)
self.policy.calc_priority(self.waiting_queue, self.running_batch)
if TEST_RETRACT and running_bs > TEST_RETRACT_NO_PREFILL_BS:
# If we are testing retraction and the running batch size exceeds

10
python/sglang/srt/server_args.py
View File

@@ -2837,7 +2837,15 @@ class ServerArgs:
"--schedule-policy",
type=str,
default=ServerArgs.schedule_policy,
choices=["lpm", "random", "fcfs", "dfs-weight", "lof", "priority"],
choices=[
"lpm",
"random",
"fcfs",
"dfs-weight",
"lof",
"priority",
"routing-key",
],
help="The scheduling policy of the requests.",
)
parser.add_argument(

114
test/manual/test_schedule_policy.py
View File

@@ -1,6 +1,6 @@
import unittest
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
from sglang.srt.managers.schedule_policy import (
CacheAgnosticPolicy,
CacheAwarePolicy,
@@ -197,6 +197,118 @@ class TestSchedulePolicy(CustomTestCase):
self.assertEqual(waiting_queue[1].rid, 2)
self.assertEqual(waiting_queue[2].rid, 3)
def test_calc_priority_routing_key_scheduling(self):
"""Test routing-key policy: prioritize by routing key frequency in running batch."""
tree_cache = RadixCache.create_simulated()
running_reqs = [
Req("r1", "a", [1], SamplingParams(), routing_key="key_a"),
Req("r2", "b", [2], SamplingParams(), routing_key="key_a"),
Req("r3", "c", [3], SamplingParams(), routing_key="key_b"),
]
running_batch = ScheduleBatch(reqs=running_reqs)
waiting_queue = [
Req("w1", "d", [4], SamplingParams(), routing_key="key_b"),
Req("w2", "e", [5], SamplingParams(), routing_key="key_a"),
Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
]
policy = SchedulePolicy(
policy="routing-key",
tree_cache=tree_cache,
enable_hierarchical_cache=False,
enable_priority_scheduling=False,
schedule_low_priority_values_first=False,
)
policy.calc_priority(waiting_queue, running_batch)
self.assertEqual(waiting_queue[0].rid, "w2")
self.assertEqual(waiting_queue[1].rid, "w1")
self.assertEqual(waiting_queue[2].rid, "w3")
def test_calc_priority_routing_key_tie_break_by_lexicographic_order(self):
"""Test routing-key policy: tie-break by lexicographic order."""
tree_cache = RadixCache.create_simulated()
running_reqs = [
Req("r1", "a", [1], SamplingParams(), routing_key="key_b"),
Req("r2", "b", [2], SamplingParams(), routing_key="key_a"),
]
running_batch = ScheduleBatch(reqs=running_reqs)
waiting_queue = [
Req("w1", "d", [4], SamplingParams(), routing_key="key_b"),
Req("w2", "e", [5], SamplingParams(), routing_key="key_a"),
]
policy = SchedulePolicy(
policy="routing-key",
tree_cache=tree_cache,
enable_hierarchical_cache=False,
enable_priority_scheduling=False,
schedule_low_priority_values_first=False,
)
policy.calc_priority(waiting_queue, running_batch)
self.assertEqual(waiting_queue[0].rid, "w2")
self.assertEqual(waiting_queue[1].rid, "w1")
def test_calc_priority_routing_key_no_match_deprioritized(self):
"""Test routing-key policy: requests without matching routing keys are deprioritized."""
tree_cache = RadixCache.create_simulated()
running_reqs = [
Req("r1", "a", [1], SamplingParams(), routing_key="key_a"),
Req("r2", "b", [2], SamplingParams(), routing_key="key_b"),
Req("r3", "c", [3], SamplingParams(), routing_key="key_c"),
]
running_batch = ScheduleBatch(reqs=running_reqs)
waiting_queue = [
Req("w1", "d", [4], SamplingParams(), routing_key="key_d"),
Req("w2", "e", [5], SamplingParams(), routing_key="key_e"),
Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
]
policy = SchedulePolicy(
policy="routing-key",
tree_cache=tree_cache,
enable_hierarchical_cache=False,
enable_priority_scheduling=False,
schedule_low_priority_values_first=False,
)
policy.calc_priority(waiting_queue, running_batch)
self.assertEqual(waiting_queue[0].rid, "w3")
self.assertEqual(waiting_queue[1].rid, "w1")
self.assertEqual(waiting_queue[2].rid, "w2")
def test_calc_priority_routing_key_empty_running_batch(self):
"""Test routing-key policy: empty running batch keeps original order."""
tree_cache = RadixCache.create_simulated()
running_batch = ScheduleBatch(reqs=[])
waiting_queue = [
Req("w1", "d", [4], SamplingParams(), routing_key="key_a"),
Req("w2", "e", [5], SamplingParams(), routing_key="key_b"),
Req("w3", "f", [6], SamplingParams(), routing_key="key_c"),
]
policy = SchedulePolicy(
policy="routing-key",
tree_cache=tree_cache,
enable_hierarchical_cache=False,
enable_priority_scheduling=False,
schedule_low_priority_values_first=False,
)
policy.calc_priority(waiting_queue, running_batch)
self.assertEqual(waiting_queue[0].rid, "w1")
self.assertEqual(waiting_queue[1].rid, "w2")
self.assertEqual(waiting_queue[2].rid, "w3")
if __name__ == "__main__":
unittest.main()

122
test/registered/scheduler/test_routing_key_scheduling.py Normal file
View File

@@ -0,0 +1,122 @@
import asyncio
import os
import time
import unittest
import aiohttp
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
STDERR_FILENAME,
STDOUT_FILENAME,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(est_time=120, suite="nightly-1-gpu", nightly=True)
class TestRoutingKeyScheduling(CustomTestCase):
@classmethod
def setUpClass(cls):
os.environ["SGLANG_ROUTING_KEY_POLICY_DEBUG_LOG"] = "1"
cls.model = "Qwen/Qwen3-0.6B"
cls.base_url = DEFAULT_URL_FOR_TEST
cls.stdout = open(STDOUT_FILENAME, "w")
cls.stderr = open(STDERR_FILENAME, "w")
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=(
"--max-running-requests",
"3",
"--schedule-policy",
"routing-key",
),
return_stdout_stderr=(cls.stdout, cls.stderr),
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
cls.stdout.close()
cls.stderr.close()
os.remove(STDOUT_FILENAME)
os.remove(STDERR_FILENAME)
def test_routing_key_scheduling_order(self):
"""Verify requests with matching routing keys are prioritized.
Test strategy:
1. First send 2 long-running key_a requests to occupy running batch
2. Then send 10 key_a and 10 key_b short requests concurrently
3. With max_running_requests=3, key_a requests should be prioritized
because running batch has 2 key_a requests
4. Verify key_a requests finish before key_b requests on average
"""
asyncio.run(self._test_routing_key_scheduling_order())
async def _test_routing_key_scheduling_order(self):
long_running_tasks = [
asyncio.create_task(self._send_chat_request("key_a", 20000)),
asyncio.create_task(self._send_chat_request("key_a", 20000)),
]
await asyncio.sleep(2.0)
short_tasks = []
for _ in range(10):
short_tasks.append(
asyncio.create_task(self._send_chat_request("key_a", 10))
)
short_tasks.append(
asyncio.create_task(self._send_chat_request("key_b", 10))
)
all_short_results = await asyncio.gather(*short_tasks)
await asyncio.gather(*long_running_tasks)
key_a_latencies = [lat for key, lat in all_short_results if key == "key_a"]
key_b_latencies = [lat for key, lat in all_short_results if key == "key_b"]
avg_key_a = sum(key_a_latencies) / len(key_a_latencies)
avg_key_b = sum(key_b_latencies) / len(key_b_latencies)
print(f"Average key_a latency: {avg_key_a:.3f}s")
print(f"Average key_b latency: {avg_key_b:.3f}s")
self.assertLess(
avg_key_a,
avg_key_b,
f"key_a requests (avg={avg_key_a:.3f}s) should finish before key_b (avg={avg_key_b:.3f}s)",
)
async def _send_chat_request(self, routing_key: str, max_tokens: int):
payload = {
"model": self.model,
"messages": [{"role": "user", "content": "What is 1+1?"}],
"max_tokens": max_tokens,
"temperature": 0,
}
headers = {"x-smg-routing-key": routing_key}
start_time = time.perf_counter()
async with aiohttp.ClientSession() as session:
async with session.post(
f"{self.base_url}/v1/chat/completions",
json=payload,
headers=headers,
) as resp:
await resp.json()
latency = time.perf_counter() - start_time
return routing_key, latency
if __name__ == "__main__":
unittest.main()