chore: import upstream snapshot with attribution

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

@@ -0,0 +1,8 @@
+from sglang.srt.debug_utils.schedule_simulator.routers.base import RouterPolicy
+from sglang.srt.debug_utils.schedule_simulator.routers.random_router import RandomRouter
+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", "StickyRouter"]

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

@@ -0,0 +1,8 @@
+from abc import ABC, abstractmethod
+
+from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
+
+
+class RouterPolicy(ABC):
+    @abstractmethod
+    def route(self, incoming_request: SimRequest) -> int: ...

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

@@ -0,0 +1,12 @@
+import random
+
+from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
+from sglang.srt.debug_utils.schedule_simulator.routers.base import RouterPolicy
+
+
+class RandomRouter(RouterPolicy):
+    def __init__(self, num_gpus: int):
+        self._num_gpus = num_gpus
+
+    def route(self, incoming_request: SimRequest) -> int:
+        return random.randint(0, self._num_gpus - 1)

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

@@ -0,0 +1,13 @@
+from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
+from sglang.srt.debug_utils.schedule_simulator.routers.base import RouterPolicy
+
+
+class RoundRobinRouter(RouterPolicy):
+    def __init__(self, num_gpus: int):
+        self._num_gpus = num_gpus
+        self._counter = 0
+
+    def route(self, incoming_request: SimRequest) -> int:
+        gpu_id = self._counter % self._num_gpus
+        self._counter += 1
+        return gpu_id

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

@@ -0,0 +1,20 @@
+import random
+from collections import defaultdict
+
+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) -> int:
+        group_id = incoming_request.group_id
+        if group_id is None:
+            return random.randint(0, self._num_gpus - 1)
+        return self._group_to_gpu[group_id]