From 1f0ea4f958c3fd85bbbd4c7ce96745dcc6a4fe1f Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sat, 10 Jan 2026 10:16:23 +0800 Subject: [PATCH] Add routing key based schedule policy (#16840) --- python/sglang/srt/managers/schedule_policy.py | 49 ++++++- python/sglang/srt/managers/scheduler.py | 3 +- python/sglang/srt/server_args.py | 10 +- test/manual/test_schedule_policy.py | 114 +++++++++++++++- .../scheduler/test_routing_key_scheduling.py | 122 ++++++++++++++++++ 5 files changed, 293 insertions(+), 5 deletions(-) create mode 100644 test/registered/scheduler/test_routing_key_scheduling.py diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index bb0197956..0698a59a9 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -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(): diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 6d9b00b8a..ce466dc94 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 4e0768fec..3aefcae35 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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( diff --git a/test/manual/test_schedule_policy.py b/test/manual/test_schedule_policy.py index 747a247e8..aa9ecc86c 100644 --- a/test/manual/test_schedule_policy.py +++ b/test/manual/test_schedule_policy.py @@ -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() diff --git a/test/registered/scheduler/test_routing_key_scheduling.py b/test/registered/scheduler/test_routing_key_scheduling.py new file mode 100644 index 000000000..be4a7d30e --- /dev/null +++ b/test/registered/scheduler/test_routing_key_scheduling.py @@ -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()