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

Support multiple engines for router simulation in schedule simulator (#16368)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-04 11:46:05 +08:00
committed by GitHub
parent e797f0c570
commit d7aa0ce72f
7 changed files with 72 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
python/sglang/srt/debug_utils/schedule_simulator/entrypoint.py
View File

@@ -63,7 +63,8 @@ def create_arg_parser() -> argparse.ArgumentParser:
parser.add_argument("--synth-gsp-output-len", type=int, default=256)
parser.add_argument("--synth-gsp-range-ratio", type=float, default=1.0)
parser.add_argument("--num-gpus", type=int, default=8)
parser.add_argument("--num-gpus-per-engine", type=int, default=8)
parser.add_argument("--num-engines", type=int, default=1)
parser.add_argument(
"--router",
type=str,
@@ -111,13 +112,13 @@ def _load_requests(args: argparse.Namespace) -> List[SimRequest]:
return requests
def _create_router(name: str, num_gpus: int):
def _create_router(name: str, total_gpus: int):
if name == "random":
return RandomRouter()
return RandomRouter(total_gpus)
if name == "round_robin":
return RoundRobinRouter()
return RoundRobinRouter(total_gpus)
if name == "sticky":
return StickyRouter(num_gpus)
return StickyRouter(total_gpus)
raise ValueError(f"Unknown router: {name}")
@@ -131,11 +132,12 @@ def main(args: argparse.Namespace) -> SimulationResult:
if args.synth_seed is not None:
random.seed(args.synth_seed)
requests = _load_requests(args)
router = _create_router(args.router, args.num_gpus)
total_gpus = args.num_gpus_per_engine * args.num_engines
router = _create_router(args.router, total_gpus)
scheduler = _create_scheduler(args.scheduler)
sim = Simulator(
num_gpus=args.num_gpus,
num_gpus_per_engine=args.num_gpus_per_engine,
router=router,
scheduler=scheduler,
recorders=[
@@ -150,7 +152,7 @@ def main(args: argparse.Namespace) -> SimulationResult:
)
print(
f"Running simulation with {args.num_gpus} GPUs, router={args.router}, scheduler={args.scheduler}"
f"Running simulation with {args.num_gpus_per_engine} GPUs/engine x {args.num_engines} engines, router={args.router}, scheduler={args.scheduler}"
)
result = sim.run(requests)

8
python/sglang/srt/debug_utils/schedule_simulator/routers/base.py
View File

@@ -1,14 +1,8 @@
from abc import ABC, abstractmethod
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
class RouterPolicy(ABC):
@abstractmethod
def route(
self,
incoming_request: SimRequest,
gpu_states: List[GPUState],
) -> int: ...
def route(self, incoming_request: SimRequest) -> int: ...

13
python/sglang/srt/debug_utils/schedule_simulator/routers/random_router.py
View File

@@ -1,15 +1,12 @@
import random
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 RandomRouter(RouterPolicy):
def route(
self,
incoming_request: SimRequest,
gpu_states: List[GPUState],
) -> int:
return random.randint(0, len(gpu_states) - 1)
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)

14
python/sglang/srt/debug_utils/schedule_simulator/routers/round_robin_router.py
View File

@@ -1,19 +1,13 @@
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 RoundRobinRouter(RouterPolicy):
def __init__(self):
def __init__(self, num_gpus: int):
self._num_gpus = num_gpus
self._counter = 0
def route(
self,
incoming_request: SimRequest,
gpu_states: List[GPUState],
) -> int:
gpu_id = self._counter % len(gpu_states)
def route(self, incoming_request: SimRequest) -> int:
gpu_id = self._counter % self._num_gpus
self._counter += 1
return gpu_id

10
python/sglang/srt/debug_utils/schedule_simulator/routers/sticky_router.py
View File

@@ -1,8 +1,6 @@
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
@@ -15,12 +13,8 @@ class StickyRouter(RouterPolicy):
def _assign_gpu(self) -> int:
return random.randint(0, self._num_gpus - 1)
def route(
self,
incoming_request: SimRequest,
gpu_states: List[GPUState],
) -> int:
def route(self, incoming_request: SimRequest) -> int:
group_id = incoming_request.group_id
if group_id is None:
return random.randint(0, len(gpu_states) - 1)
return random.randint(0, self._num_gpus - 1)
return self._group_to_gpu[group_id]

11
python/sglang/srt/debug_utils/schedule_simulator/simulator.py
View File

@@ -17,7 +17,7 @@ class SimulationResult:
class Simulator:
def __init__(
self,
num_gpus: int,
num_gpus_per_engine: int,
router: RouterPolicy,
scheduler: SchedulerPolicy,
recorders: Optional[List[MetricRecorder]] = None,
@@ -26,7 +26,7 @@ class Simulator:
stop_criteria: str = "all_done",
max_steps: Optional[int] = None,
):
self.num_gpus = num_gpus
self.num_gpus_per_engine = num_gpus_per_engine
self.router = router
self.scheduler = scheduler
self.recorders = recorders or []
@@ -40,7 +40,7 @@ class Simulator:
def run(self, requests: List[SimRequest]) -> SimulationResult:
self.gpu_states = [
GPUState(gpu_id=i, max_total_tokens=self.max_total_tokens)
for i in range(self.num_gpus)
for i in range(self.num_gpus_per_engine)
]
self.step = 0
step_records: List[StepRecord] = []
@@ -75,8 +75,9 @@ class Simulator:
def _route_requests(self, incoming_requests: List[SimRequest]) -> None:
for req in incoming_requests:
gpu_id = self.router.route(req, self.gpu_states)
self.gpu_states[gpu_id].pending_requests.append(req)
gpu_id = self.router.route(req)
if gpu_id < self.num_gpus_per_engine:
self.gpu_states[gpu_id].pending_requests.append(req)
def _schedule_all_gpus(self) -> None:
for gpu in self.gpu_states: