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

Enhance schedule simulator and add several e2e scenarios (#16357)

Browse Source
This commit is contained in:
fzyzcjy
2026-01-04 10:02:38 +08:00
committed by GitHub
parent 8ca9597057
commit 24e116ef6a
7 changed files with 169 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/debug_utils/schedule_simulator/__init__.py
View File

@@ -6,7 +6,8 @@ from sglang.srt.debug_utils.schedule_simulator.data_source import (
from sglang.srt.debug_utils.schedule_simulator.entrypoint import create_arg_parser, main
from sglang.srt.debug_utils.schedule_simulator.gpu_state import GPUState, StepRecord
from sglang.srt.debug_utils.schedule_simulator.metrics import (
AttentionBalancednessRecorder,
AttentionComputeBalancednessRecorder,
AvgBatchSizeRecorder,
BatchSizeBalancednessRecorder,
MetricRecorder,
)
@@ -40,7 +41,8 @@ __all__ = [
"FIFOScheduler",
"MetricRecorder",
"BatchSizeBalancednessRecorder",
"AttentionBalancednessRecorder",
"AttentionComputeBalancednessRecorder",
"AvgBatchSizeRecorder",
"load_from_request_logger",
"generate_random_requests",
"generate_gsp_requests",

4
python/sglang/srt/debug_utils/schedule_simulator/__main__.py
View File

@@ -3,6 +3,4 @@ from sglang.srt.debug_utils.schedule_simulator.entrypoint import create_arg_pars
if __name__ == "__main__":
parser = create_arg_parser()
args = parser.parse_args()
df = main(args)
print(f"\nDataFrame shape: {df.shape}")
print(df.head(20))
main(args)

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

@@ -1,10 +1,8 @@
import argparse
import json
from dataclasses import asdict
import random
from typing import List
import polars as pl
from sglang.srt.debug_utils.schedule_simulator.data_source.data_loader import (
load_from_request_logger,
)
@@ -13,7 +11,8 @@ from sglang.srt.debug_utils.schedule_simulator.data_source.data_synthesis import
generate_random_requests,
)
from sglang.srt.debug_utils.schedule_simulator.metrics import (
AttentionBalancednessRecorder,
AttentionComputeBalancednessRecorder,
AvgBatchSizeRecorder,
BatchSizeBalancednessRecorder,
)
from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
@@ -23,7 +22,10 @@ from sglang.srt.debug_utils.schedule_simulator.routers import (
StickyRouter,
)
from sglang.srt.debug_utils.schedule_simulator.schedulers import FIFOScheduler
from sglang.srt.debug_utils.schedule_simulator.simulator import Simulator
from sglang.srt.debug_utils.schedule_simulator.simulator import (
SimulationResult,
Simulator,
)
def create_arg_parser() -> argparse.ArgumentParser:
@@ -70,6 +72,14 @@ def create_arg_parser() -> argparse.ArgumentParser:
)
parser.add_argument("--scheduler", type=str, choices=["fifo"], default="fifo")
parser.add_argument("--max-total-tokens", type=int, default=100000)
parser.add_argument(
"--stop-criteria",
type=str,
choices=["all_done", "exist_no_pending"],
default="all_done",
help="all_done: run until all requests complete; exist_no_pending: stop when any GPU has no pending requests",
)
parser.add_argument("--max-steps", type=int, default=None)
parser.add_argument("--output", type=str, default=None)
parser.add_argument("--log-level", type=int, choices=[0, 1, 2], default=0)
@@ -117,7 +127,9 @@ def _create_scheduler(name: str):
raise ValueError(f"Unknown scheduler: {name}")
def main(args: argparse.Namespace) -> pl.DataFrame:
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)
scheduler = _create_scheduler(args.scheduler)
@@ -126,9 +138,15 @@ def main(args: argparse.Namespace) -> pl.DataFrame:
num_gpus=args.num_gpus,
router=router,
scheduler=scheduler,
recorders=[BatchSizeBalancednessRecorder(), AttentionBalancednessRecorder()],
recorders=[
BatchSizeBalancednessRecorder(),
AttentionComputeBalancednessRecorder(),
AvgBatchSizeRecorder(),
],
log_level=args.log_level,
max_total_tokens=args.max_total_tokens,
stop_criteria=args.stop_criteria,
max_steps=args.max_steps,
)
print(
@@ -136,8 +154,6 @@ def main(args: argparse.Namespace) -> pl.DataFrame:
)
result = sim.run(requests)
df = pl.DataFrame([asdict(r) for r in result.step_records])
print("\n=== Summary ===")
for key, value in result.summary.items():
print(f"{key}: {value:.4f}" if isinstance(value, float) else f"{key}: {value}")
@@ -147,4 +163,4 @@ def main(args: argparse.Namespace) -> pl.DataFrame:
json.dump(result.summary, f, indent=2)
print(f"\nSummary saved to {args.output}")
return df
return result

3
python/sglang/srt/debug_utils/schedule_simulator/gpu_state.py
View File

@@ -25,6 +25,9 @@ class GPUState:
def batch_size(self) -> int:
return len(self.running_requests)
def total_attention_compute(self) -> int:
return sum(req.seq_len() for req in self.running_requests)
def total_seq_len(self, extra_reqs: Optional[List[SimRequest]] = None) -> int:
seen_groups = set()
total = 0

19
python/sglang/srt/debug_utils/schedule_simulator/metrics.py
View File

@@ -39,7 +39,22 @@ def BatchSizeBalancednessRecorder() -> BalancednessRecorder:
return BalancednessRecorder("batch_size_balancedness", lambda gpu: gpu.batch_size())
def AttentionBalancednessRecorder() -> BalancednessRecorder:
def AttentionComputeBalancednessRecorder() -> BalancednessRecorder:
return BalancednessRecorder(
"attention_balancedness", lambda gpu: gpu.total_seq_len()
"attention_compute_balancedness", lambda gpu: gpu.total_attention_compute()
)
class AvgBatchSizeRecorder(MetricRecorder):
def __init__(self):
self._total_running = 0
self._num_records = 0
def on_step_end(self, step: int, gpu_states: List[GPUState]) -> None:
for gpu in gpu_states:
self._total_running += gpu.batch_size()
self._num_records += 1
def get_summary(self) -> Dict[str, Any]:
avg = self._total_running / self._num_records if self._num_records else 0.0
return {"avg_batch_size": avg}

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

@@ -23,6 +23,8 @@ class Simulator:
recorders: Optional[List[MetricRecorder]] = None,
log_level: int = 0,
max_total_tokens: int = 100000,
stop_criteria: str = "all_done",
max_steps: Optional[int] = None,
):
self.num_gpus = num_gpus
self.router = router
@@ -30,6 +32,8 @@ class Simulator:
self.recorders = recorders or []
self.log_level = log_level
self.max_total_tokens = max_total_tokens
self.stop_criteria = stop_criteria
self.max_steps = max_steps
self.gpu_states: List[GPUState] = []
self.step = 0
@@ -42,10 +46,12 @@ class Simulator:
step_records: List[StepRecord] = []
incoming_requests = list(requests)
while self._has_work(incoming_requests):
while True:
self._route_requests(incoming_requests)
incoming_requests.clear()
self._schedule_all_gpus()
if self._should_stop():
break
self._execute_step()
step_records.extend(
gpu.get_step_record(self.step) for gpu in self.gpu_states
@@ -56,10 +62,16 @@ class Simulator:
return SimulationResult(step_records=step_records, summary=self._get_summary())
def _has_work(self, incoming_requests: List[SimRequest]) -> bool:
return bool(incoming_requests) or any(
gpu.pending_requests or gpu.running_requests for gpu in self.gpu_states
)
def _should_stop(self) -> bool:
if self.max_steps is not None and self.step >= self.max_steps:
return True
if self.stop_criteria == "exist_no_pending":
return any(not gpu.pending_requests for gpu in self.gpu_states)
if self.stop_criteria == "all_done":
return not any(
gpu.pending_requests or gpu.running_requests for gpu in self.gpu_states
)
raise ValueError(f"Unknown stop criteria: {self.stop_criteria}")
def _route_requests(self, incoming_requests: List[SimRequest]) -> None:
for req in incoming_requests:
@@ -79,12 +91,12 @@ class Simulator:
gpu.execute_step()
def _log_step(self) -> None:
if self.log_level == 0:
if self.log_level == 0 and self.step % 100 != 0:
return
parts = [f"step={self.step:<4}"]
for gpu in self.gpu_states:
r, q = len(gpu.running_requests), len(gpu.pending_requests)
if self.log_level == 1:
if self.log_level <= 1:
parts.append(f"GPU{gpu.gpu_id}[R={r:<3} Q={q:<3}]")
else:
run_ids = _format_ids(gpu.running_requests)