Add routing key based schedule policy (#16840)

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():

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

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(

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()

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()