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

python/sglang/srt/debug_utils/schedule_simulator/__init__.py

@@ -15,6 +15,7 @@ from sglang.srt.debug_utils.schedule_simulator.routers import (
     RandomRouter,
     RoundRobinRouter,
     RouterPolicy,
+    StickyRouter,
 )
 from sglang.srt.debug_utils.schedule_simulator.schedulers import (
     FIFOScheduler,
@@ -34,6 +35,7 @@ __all__ = [
     "RouterPolicy",
     "RandomRouter",
     "RoundRobinRouter",
+    "StickyRouter",
     "SchedulerPolicy",
     "FIFOScheduler",
     "MetricRecorder",

python/sglang/srt/debug_utils/schedule_simulator/entrypoint.py

@@ -20,6 +20,7 @@ from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
 from sglang.srt.debug_utils.schedule_simulator.routers import (
     RandomRouter,
     RoundRobinRouter,
+    StickyRouter,
 )
 from sglang.srt.debug_utils.schedule_simulator.schedulers import FIFOScheduler
 from sglang.srt.debug_utils.schedule_simulator.simulator import Simulator
@@ -62,7 +63,10 @@ def create_arg_parser() -> argparse.ArgumentParser:
 
     parser.add_argument("--num-gpus", type=int, default=8)
     parser.add_argument(
-        "--router", type=str, choices=["random", "round_robin"], default="round_robin"
+        "--router",
+        type=str,
+        choices=["random", "round_robin", "sticky"],
+        default="round_robin",
     )
     parser.add_argument("--scheduler", type=str, choices=["fifo"], default="fifo")
     parser.add_argument("--max-total-tokens", type=int, default=100000)
@@ -97,11 +101,13 @@ def _load_requests(args: argparse.Namespace) -> List[SimRequest]:
     return requests
 
 
-def _create_router(name: str):
+def _create_router(name: str, num_gpus: int):
     if name == "random":
         return RandomRouter()
     if name == "round_robin":
         return RoundRobinRouter()
+    if name == "sticky":
+        return StickyRouter(num_gpus)
     raise ValueError(f"Unknown router: {name}")
 
 
@@ -113,7 +119,7 @@ def _create_scheduler(name: str):
 
 def main(args: argparse.Namespace) -> pl.DataFrame:
     requests = _load_requests(args)
-    router = _create_router(args.router)
+    router = _create_router(args.router, args.num_gpus)
     scheduler = _create_scheduler(args.scheduler)
 
     sim = Simulator(

python/sglang/srt/debug_utils/schedule_simulator/routers/__init__.py

@@ -3,5 +3,6 @@ from sglang.srt.debug_utils.schedule_simulator.routers.random_router import Rand
 from sglang.srt.debug_utils.schedule_simulator.routers.round_robin_router import (
     RoundRobinRouter,
 )
+from sglang.srt.debug_utils.schedule_simulator.routers.sticky_router import StickyRouter
 
-__all__ = ["RouterPolicy", "RandomRouter", "RoundRobinRouter"]
+__all__ = ["RouterPolicy", "RandomRouter", "RoundRobinRouter", "StickyRouter"]

python/sglang/srt/debug_utils/schedule_simulator/routers/sticky_router.py

@@ -0,0 +1,26 @@
+import random
+from collections import defaultdict
+from typing import List
+
+from sglang.srt.debug_utils.schedule_simulator.gpu_state import GPUState
+from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
+from sglang.srt.debug_utils.schedule_simulator.routers.base import RouterPolicy
+
+
+class StickyRouter(RouterPolicy):
+    def __init__(self, num_gpus: int):
+        self._num_gpus = num_gpus
+        self._group_to_gpu = defaultdict(self._assign_gpu)
+
+    def _assign_gpu(self) -> int:
+        return random.randint(0, self._num_gpus - 1)
+
+    def route(
+        self,
+        incoming_request: SimRequest,
+        gpu_states: List[GPUState],
+    ) -> int:
+        group_id = incoming_request.group_id
+        if group_id is None:
+            return random.randint(0, len(gpu_states) - 1)
+        return self._group_to_gpu[group_id]