diff --git a/python/sglang/srt/debug_utils/schedule_simulator/__init__.py b/python/sglang/srt/debug_utils/schedule_simulator/__init__.py index c406f765c..b9c815de6 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/__init__.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/__init__.py @@ -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", diff --git a/python/sglang/srt/debug_utils/schedule_simulator/__main__.py b/python/sglang/srt/debug_utils/schedule_simulator/__main__.py index c3a907581..aff33faaa 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/__main__.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/__main__.py @@ -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) diff --git a/python/sglang/srt/debug_utils/schedule_simulator/entrypoint.py b/python/sglang/srt/debug_utils/schedule_simulator/entrypoint.py index 1a831d5b6..32431ecea 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/entrypoint.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/entrypoint.py @@ -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 diff --git a/python/sglang/srt/debug_utils/schedule_simulator/gpu_state.py b/python/sglang/srt/debug_utils/schedule_simulator/gpu_state.py index 059d4751a..634d01fef 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/gpu_state.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/gpu_state.py @@ -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 diff --git a/python/sglang/srt/debug_utils/schedule_simulator/metrics.py b/python/sglang/srt/debug_utils/schedule_simulator/metrics.py index fdf9fe1d4..d640072af 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/metrics.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/metrics.py @@ -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} diff --git a/python/sglang/srt/debug_utils/schedule_simulator/simulator.py b/python/sglang/srt/debug_utils/schedule_simulator/simulator.py index e8cf78826..d1fb48015 100644 --- a/python/sglang/srt/debug_utils/schedule_simulator/simulator.py +++ b/python/sglang/srt/debug_utils/schedule_simulator/simulator.py @@ -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) diff --git a/test/registered/debug_utils/test_schedule_simulator.py b/test/registered/debug_utils/test_schedule_simulator.py index b17643ec5..c30f75ef9 100644 --- a/test/registered/debug_utils/test_schedule_simulator.py +++ b/test/registered/debug_utils/test_schedule_simulator.py @@ -5,7 +5,7 @@ import tempfile import unittest from sglang.srt.debug_utils.schedule_simulator import ( - AttentionBalancednessRecorder, + AttentionComputeBalancednessRecorder, BatchSizeBalancednessRecorder, FIFOScheduler, GPUState, @@ -16,14 +16,16 @@ from sglang.srt.debug_utils.schedule_simulator import ( Simulator, StepRecord, StickyRouter, + create_arg_parser, generate_gsp_requests, generate_random_requests, load_from_request_logger, + main, ) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase -register_cpu_ci(est_time=60, suite="default", nightly=True) +register_cpu_ci(est_time=120, suite="default", nightly=True) # ==================== Non-E2E Tests ==================== @@ -244,8 +246,8 @@ class TestMetrics(CustomTestCase): recorder.get_summary()["batch_size_balancedness_mean"], 0.75 ) - def test_attention_balancedness(self): - recorder = AttentionBalancednessRecorder() + def test_attention_compute_balancedness(self): + recorder = AttentionComputeBalancednessRecorder() gpu_states = [GPUState(gpu_id=i, max_total_tokens=10000) for i in range(2)] gpu_states[0].running_requests = [ SimRequest(request_id="r1", input_len=100, output_len=50) @@ -255,7 +257,7 @@ class TestMetrics(CustomTestCase): ] recorder.on_step_end(0, gpu_states) self.assertAlmostEqual( - recorder.get_summary()["attention_balancedness_mean"], 0.75 + recorder.get_summary()["attention_compute_balancedness_mean"], 0.75 ) def test_empty_history(self): @@ -408,7 +410,7 @@ class TestSimulator(CustomTestCase): scheduler=FIFOScheduler(), recorders=[ BatchSizeBalancednessRecorder(), - AttentionBalancednessRecorder(), + AttentionComputeBalancednessRecorder(), ], max_total_tokens=100, ) @@ -545,30 +547,13 @@ class TestCLI(CustomTestCase): self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}") self.assertIn("router=random", result.stdout) - def test_cli_sticky_router(self): - result = self._run_cli( - "--synth-gsp", - "--synth-gsp-num-groups", - "2", - "--synth-gsp-prompts-per-group", - "3", - "--synth-seed", - "42", - "--num-gpus", - "4", - "--router", - "sticky", - ) - self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}") - self.assertIn("router=sticky", result.stdout) - def test_e2e_sticky_router_group_locality(self): result = self._run_cli( "--synth-gsp", "--synth-gsp-num-groups", - "2", + "1", "--synth-gsp-prompts-per-group", - "2", + "4", "--synth-gsp-system-prompt-len", "10", "--synth-gsp-question-len", @@ -576,7 +561,7 @@ class TestCLI(CustomTestCase): "--synth-gsp-output-len", "2", "--synth-seed", - "123", + "42", "--num-gpus", "2", "--router", @@ -587,11 +572,8 @@ class TestCLI(CustomTestCase): "2", ) self.assertEqual(result.returncode, 0, f"CLI failed: {result.stderr}") - for expected in [ - "step=0 | GPU0[R=2:gsp0,gsp1 Q=0:-] | GPU1[R=2:gsp2,gsp3 Q=0:-]", - "step=1 | GPU0[R=0:- Q=0:-] | GPU1[R=0:- Q=0:-]", - ]: - self.assertIn(expected, result.stdout) + self.assertIn("R=4:", result.stdout) + self.assertIn("R=0:-", result.stdout) def test_cli_synthetic(self): result = self._run_cli( @@ -773,5 +755,91 @@ step=13 | GPU0[R=0:- Q=0:-]""", self.assertNotIn("R=2:", result.stdout) +class TestLargerScale(CustomTestCase): + def _run_main(self, *cli_args) -> SimulationResult: + parser = create_arg_parser() + args = parser.parse_args(cli_args) + return main(args) + + def _assert_in_range(self, value, lo, hi, name): + self.assertGreaterEqual(value, lo, f"{name}={value} < {lo}") + self.assertLessEqual(value, hi, f"{name}={value} > {hi}") + + def test_vanilla_workload_random_policy(self): + result = self._run_main( + "--synthetic", + "--synth-random-num-requests", + "2000", + "--synth-random-input-len", + "32000", + "--synth-random-output-len", + "2000", + "--synth-seed", + "42", + "--num-gpus", + "8", + "--router", + "random", + "--max-total-tokens", + "2000000", + "--stop-criteria", + "exist_no_pending", + ) + self._assert_in_range( + result.summary["attention_compute_balancedness_mean"], 0.95, 1.0, "attn" + ) + self._assert_in_range( + result.summary["batch_size_balancedness_mean"], 0.90, 0.98, "bs" + ) + self._assert_in_range(result.summary["avg_batch_size"], 127, 141, "avg_bs") + + def _run_gsp_workload(self, router: str) -> SimulationResult: + return self._run_main( + "--synth-gsp", + "--synth-gsp-num-groups", + "200", + "--synth-gsp-prompts-per-group", + "20", + "--synth-gsp-system-prompt-len", + "31000", + "--synth-gsp-question-len", + "1000", + "--synth-gsp-output-len", + "8000", + "--synth-seed", + "42", + "--num-gpus", + "8", + "--router", + router, + "--max-total-tokens", + "500000", + "--stop-criteria", + "exist_no_pending", + "--max-steps", + "1500", + ) + + def test_gsp_workload_random_policy(self): + result = self._run_gsp_workload("random") + self._assert_in_range( + result.summary["attention_compute_balancedness_mean"], 0.88, 0.98, "attn" + ) + self._assert_in_range( + result.summary["batch_size_balancedness_mean"], 0.88, 0.98, "bs" + ) + self._assert_in_range(result.summary["avg_batch_size"], 24, 27, "avg_bs") + + def test_gsp_workload_sticky_policy(self): + result = self._run_gsp_workload("sticky") + self._assert_in_range( + result.summary["attention_compute_balancedness_mean"], 0.63, 0.70, "attn" + ) + self._assert_in_range( + result.summary["batch_size_balancedness_mean"], 0.63, 0.71, "bs" + ) + self._assert_in_range(result.summary["avg_batch_size"], 33, 37, "avg_bs") + + if __name__ == "__main__": unittest.main()