Support simple schedule simulator (#16344)

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

@@ -0,0 +1,45 @@
+from sglang.srt.debug_utils.schedule_simulator.data_source import (
+    generate_random_requests,
+    load_from_request_logger,
+)
+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,
+    BatchSizeBalancednessRecorder,
+    MetricRecorder,
+)
+from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
+from sglang.srt.debug_utils.schedule_simulator.routers import (
+    RandomRouter,
+    RoundRobinRouter,
+    RouterPolicy,
+)
+from sglang.srt.debug_utils.schedule_simulator.schedulers import (
+    FIFOScheduler,
+    SchedulerPolicy,
+)
+from sglang.srt.debug_utils.schedule_simulator.simulator import (
+    SimulationResult,
+    Simulator,
+)
+
+__all__ = [
+    "SimRequest",
+    "GPUState",
+    "Simulator",
+    "SimulationResult",
+    "StepRecord",
+    "RouterPolicy",
+    "RandomRouter",
+    "RoundRobinRouter",
+    "SchedulerPolicy",
+    "FIFOScheduler",
+    "MetricRecorder",
+    "BatchSizeBalancednessRecorder",
+    "AttentionBalancednessRecorder",
+    "load_from_request_logger",
+    "generate_random_requests",
+    "create_arg_parser",
+    "main",
+]

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

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

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

@@ -0,0 +1,8 @@
+from sglang.srt.debug_utils.schedule_simulator.data_source.data_loader import (
+    load_from_request_logger,
+)
+from sglang.srt.debug_utils.schedule_simulator.data_source.data_synthesis import (
+    generate_random_requests,
+)
+
+__all__ = ["load_from_request_logger", "generate_random_requests"]

python/sglang/srt/debug_utils/schedule_simulator/data_source/data_loader.py

@@ -0,0 +1,34 @@
+import json
+from pathlib import Path
+from typing import List, Union
+
+from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
+
+
+def load_from_request_logger(file_path: Union[str, Path]) -> List[SimRequest]:
+    requests = []
+    file_path = Path(file_path)
+
+    with file_path.open(encoding="utf-8") as f:
+        for line_num, line in enumerate(f):
+            line = line.strip()
+            if not line or not line.startswith("{"):
+                continue
+
+            data = json.loads(line)
+
+            if data.get("event") != "request.finished":
+                continue
+
+            rid = data.get("rid", f"req_{line_num}")
+            meta_info = data["out"]["meta_info"]
+
+            requests.append(
+                SimRequest(
+                    request_id=rid,
+                    input_len=meta_info["prompt_tokens"],
+                    output_len=meta_info["completion_tokens"],
+                )
+            )
+
+    return requests

python/sglang/srt/debug_utils/schedule_simulator/data_source/data_synthesis.py

@@ -0,0 +1,34 @@
+import random
+from typing import List, Optional
+
+from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
+
+
+def generate_random_requests(
+    num_requests: int,
+    input_len: int,
+    output_len: int,
+    range_ratio: float = 1.0,
+    seed: Optional[int] = None,
+) -> List[SimRequest]:
+    if seed is not None:
+        random.seed(seed)
+
+    requests = []
+    for i in range(num_requests):
+        isl = _random_len(input_len, range_ratio)
+        osl = _random_len(output_len, range_ratio)
+        requests.append(
+            SimRequest(
+                request_id=f"syn{i}",
+                input_len=isl,
+                output_len=osl,
+            )
+        )
+
+    return requests
+
+
+def _random_len(full_len: int, range_ratio: float) -> int:
+    min_len = max(int(full_len * range_ratio), 1)
+    return random.randint(min_len, full_len)

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

@@ -0,0 +1,123 @@
+import argparse
+import json
+from dataclasses import asdict
+from typing import List
+
+import polars as pl
+
+from sglang.srt.debug_utils.schedule_simulator.data_source.data_loader import (
+    load_from_request_logger,
+)
+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,
+    BatchSizeBalancednessRecorder,
+)
+from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
+from sglang.srt.debug_utils.schedule_simulator.routers import (
+    RandomRouter,
+    RoundRobinRouter,
+)
+from sglang.srt.debug_utils.schedule_simulator.schedulers import FIFOScheduler
+from sglang.srt.debug_utils.schedule_simulator.simulator import Simulator
+
+
+def create_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        description="Schedule Simulator for analyzing request scheduling across GPUs"
+    )
+
+    data_group = parser.add_mutually_exclusive_group(required=True)
+    data_group.add_argument(
+        "--input", type=str, help="Path to request_logger JSON file"
+    )
+    data_group.add_argument(
+        "--synthetic", action="store_true", help="Use synthetic data generation"
+    )
+
+    parser.add_argument("--synth-num-requests", type=int, default=1000)
+    parser.add_argument("--synth-input-len", type=int, default=1024)
+    parser.add_argument("--synth-output-len", type=int, default=256)
+    parser.add_argument("--synth-range-ratio", type=float, default=1.0)
+    parser.add_argument("--synth-seed", type=int, default=None)
+
+    parser.add_argument("--num-gpus", type=int, default=8)
+    parser.add_argument(
+        "--router", type=str, choices=["random", "round_robin"], default="round_robin"
+    )
+    parser.add_argument("--scheduler", type=str, choices=["fifo"], default="fifo")
+    parser.add_argument("--max-total-tokens", type=int, default=100000)
+    parser.add_argument("--output", type=str, default=None)
+    parser.add_argument("--log-level", type=int, choices=[0, 1, 2], default=0)
+
+    return parser
+
+
+def _load_requests(args: argparse.Namespace) -> List[SimRequest]:
+    if args.input:
+        requests = load_from_request_logger(args.input)
+        print(f"Loaded {len(requests)} requests from {args.input}")
+    else:
+        requests = generate_random_requests(
+            num_requests=args.synth_num_requests,
+            input_len=args.synth_input_len,
+            output_len=args.synth_output_len,
+            range_ratio=args.synth_range_ratio,
+            seed=args.synth_seed,
+        )
+        print(
+            f"Generated {len(requests)} synthetic requests "
+            f"(synth_input_len={args.synth_input_len}, "
+            f"synth_output_len={args.synth_output_len}, "
+            f"synth_range_ratio={args.synth_range_ratio})"
+        )
+    return requests
+
+
+def _create_router(name: str):
+    if name == "random":
+        return RandomRouter()
+    if name == "round_robin":
+        return RoundRobinRouter()
+    raise ValueError(f"Unknown router: {name}")
+
+
+def _create_scheduler(name: str):
+    if name == "fifo":
+        return FIFOScheduler()
+    raise ValueError(f"Unknown scheduler: {name}")
+
+
+def main(args: argparse.Namespace) -> pl.DataFrame:
+    requests = _load_requests(args)
+    router = _create_router(args.router)
+    scheduler = _create_scheduler(args.scheduler)
+
+    sim = Simulator(
+        num_gpus=args.num_gpus,
+        router=router,
+        scheduler=scheduler,
+        recorders=[BatchSizeBalancednessRecorder(), AttentionBalancednessRecorder()],
+        log_level=args.log_level,
+        max_total_tokens=args.max_total_tokens,
+    )
+
+    print(
+        f"Running simulation with {args.num_gpus} GPUs, router={args.router}, scheduler={args.scheduler}"
+    )
+    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}")
+
+    if args.output:
+        with open(args.output, "w") as f:
+            json.dump(result.summary, f, indent=2)
+        print(f"\nSummary saved to {args.output}")
+
+    return df

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

@@ -0,0 +1,60 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
+
+
+@dataclass
+class StepRecord:
+    step: int
+    gpu_id: int
+    running_count: int
+    pending_count: int
+    total_seq_len: int
+    running_req_ids: List[str] = field(default_factory=list)
+    pending_req_ids: List[str] = field(default_factory=list)
+
+
+@dataclass
+class GPUState:
+    gpu_id: int
+    max_total_tokens: int
+    pending_requests: List[SimRequest] = field(default_factory=list)
+    running_requests: List[SimRequest] = field(default_factory=list)
+
+    def batch_size(self) -> int:
+        return len(self.running_requests)
+
+    def total_seq_len(self) -> int:
+        return sum(req.seq_len() for req in self.running_requests)
+
+    def is_valid(self) -> bool:
+        return self.total_seq_len() <= self.max_total_tokens
+
+    def start_request(self, req: SimRequest) -> None:
+        assert req in self.pending_requests
+        self.pending_requests.remove(req)
+        self.running_requests.append(req)
+
+    def evict_request(self, req: SimRequest) -> None:
+        assert req in self.running_requests
+        self.running_requests.remove(req)
+        self.pending_requests.insert(0, req)
+
+    def execute_step(self) -> None:
+        for req in self.running_requests:
+            req.decoded_tokens += 1
+        self.running_requests = [
+            r for r in self.running_requests if not r.is_finished()
+        ]
+
+    def get_step_record(self, step: int) -> StepRecord:
+        return StepRecord(
+            step=step,
+            gpu_id=self.gpu_id,
+            running_count=len(self.running_requests),
+            pending_count=len(self.pending_requests),
+            total_seq_len=self.total_seq_len(),
+            running_req_ids=[r.request_id for r in self.running_requests],
+            pending_req_ids=[r.request_id for r in self.pending_requests],
+        )

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

@@ -0,0 +1,45 @@
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, List
+
+from sglang.srt.debug_utils.schedule_simulator.gpu_state import GPUState
+
+
+class MetricRecorder(ABC):
+    @abstractmethod
+    def on_step_end(self, step: int, gpu_states: List[GPUState]) -> None: ...
+
+    @abstractmethod
+    def get_summary(self) -> Dict[str, Any]: ...
+
+
+class BalancednessRecorder(MetricRecorder):
+    def __init__(self, name: str, value_fn: Callable[[GPUState], float]):
+        self._name = name
+        self._value_fn = value_fn
+        self._history: List[float] = []
+
+    def on_step_end(self, step: int, gpu_states: List[GPUState]) -> None:
+        values = [self._value_fn(gpu) for gpu in gpu_states]
+        max_val = max(values) if values else 0
+        mean_val = sum(values) / len(values) if values else 0
+        balancedness = mean_val / max_val if max_val > 0 else 1.0
+        self._history.append(balancedness)
+
+    def get_summary(self) -> Dict[str, Any]:
+        if not self._history:
+            return {f"{self._name}_mean": 0.0}
+        return {
+            f"{self._name}_mean": sum(self._history) / len(self._history),
+            f"{self._name}_min": min(self._history),
+            f"{self._name}_max": max(self._history),
+        }
+
+
+def BatchSizeBalancednessRecorder() -> BalancednessRecorder:
+    return BalancednessRecorder("batch_size_balancedness", lambda gpu: gpu.batch_size())
+
+
+def AttentionBalancednessRecorder() -> BalancednessRecorder:
+    return BalancednessRecorder(
+        "attention_balancedness", lambda gpu: gpu.total_seq_len()
+    )

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

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class SimRequest:
+    request_id: str
+    input_len: int
+    output_len: int
+    decoded_tokens: int = 0
+
+    def seq_len(self) -> int:
+        return self.input_len + self.decoded_tokens
+
+    def is_finished(self) -> bool:
+        return self.decoded_tokens >= self.output_len

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

@@ -0,0 +1,7 @@
+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,
+)
+
+__all__ = ["RouterPolicy", "RandomRouter", "RoundRobinRouter"]

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

@@ -0,0 +1,14 @@
+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: ...

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

@@ -0,0 +1,15 @@
+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)

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

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

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

@@ -0,0 +1,6 @@
+from sglang.srt.debug_utils.schedule_simulator.schedulers.base import SchedulerPolicy
+from sglang.srt.debug_utils.schedule_simulator.schedulers.fifo_scheduler import (
+    FIFOScheduler,
+)
+
+__all__ = ["SchedulerPolicy", "FIFOScheduler"]

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

@@ -0,0 +1,10 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sglang.srt.debug_utils.schedule_simulator.gpu_state import GPUState
+
+
+class SchedulerPolicy(ABC):
+    @abstractmethod
+    def schedule(self, gpu_state: "GPUState") -> None: ...

python/sglang/srt/debug_utils/schedule_simulator/schedulers/fifo_scheduler.py

@@ -0,0 +1,16 @@
+from typing import TYPE_CHECKING
+
+from sglang.srt.debug_utils.schedule_simulator.schedulers.base import SchedulerPolicy
+
+if TYPE_CHECKING:
+    from sglang.srt.debug_utils.schedule_simulator.gpu_state import GPUState
+
+
+class FIFOScheduler(SchedulerPolicy):
+    def schedule(self, gpu_state: "GPUState") -> None:
+        while not gpu_state.is_valid() and gpu_state.running_requests:
+            gpu_state.evict_request(gpu_state.running_requests[-1])
+
+        for req in list(gpu_state.pending_requests):
+            if gpu_state.total_seq_len() + req.seq_len() <= gpu_state.max_total_tokens:
+                gpu_state.start_request(req)

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

@@ -0,0 +1,109 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+from sglang.srt.debug_utils.schedule_simulator.gpu_state import GPUState, StepRecord
+from sglang.srt.debug_utils.schedule_simulator.metrics import MetricRecorder
+from sglang.srt.debug_utils.schedule_simulator.request import SimRequest
+from sglang.srt.debug_utils.schedule_simulator.routers.base import RouterPolicy
+from sglang.srt.debug_utils.schedule_simulator.schedulers.base import SchedulerPolicy
+
+
+@dataclass
+class SimulationResult:
+    step_records: List[StepRecord]
+    summary: Dict[str, Any]
+
+
+class Simulator:
+    def __init__(
+        self,
+        num_gpus: int,
+        router: RouterPolicy,
+        scheduler: SchedulerPolicy,
+        recorders: Optional[List[MetricRecorder]] = None,
+        log_level: int = 0,
+        max_total_tokens: int = 100000,
+    ):
+        self.num_gpus = num_gpus
+        self.router = router
+        self.scheduler = scheduler
+        self.recorders = recorders or []
+        self.log_level = log_level
+        self.max_total_tokens = max_total_tokens
+        self.gpu_states: List[GPUState] = []
+        self.step = 0
+
+    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)
+        ]
+        self.step = 0
+        step_records: List[StepRecord] = []
+        incoming_requests = list(requests)
+
+        while self._has_work(incoming_requests):
+            self._route_requests(incoming_requests)
+            incoming_requests.clear()
+            self._schedule_all_gpus()
+            self._execute_step()
+            step_records.extend(
+                gpu.get_step_record(self.step) for gpu in self.gpu_states
+            )
+            self._log_step()
+            self._record_metrics()
+            self.step += 1
+
+        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 _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)
+
+    def _schedule_all_gpus(self) -> None:
+        for gpu in self.gpu_states:
+            self.scheduler.schedule(gpu)
+            assert gpu.is_valid(), (
+                f"GPU{gpu.gpu_id} invalid after scheduling "
+                f"({gpu.total_seq_len()=}, {gpu.max_total_tokens=})"
+            )
+
+    def _execute_step(self) -> None:
+        for gpu in self.gpu_states:
+            gpu.execute_step()
+
+    def _log_step(self) -> None:
+        if self.log_level == 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:
+                parts.append(f"GPU{gpu.gpu_id}[R={r:<3} Q={q:<3}]")
+            else:
+                run_ids = _format_ids(gpu.running_requests)
+                queue_ids = _format_ids(gpu.pending_requests)
+                parts.append(f"GPU{gpu.gpu_id}[R={r}:{run_ids} Q={q}:{queue_ids}]")
+        print(" | ".join(parts))
+
+    def _record_metrics(self) -> None:
+        for recorder in self.recorders:
+            recorder.on_step_end(self.step, self.gpu_states)
+
+    def _get_summary(self) -> Dict[str, Any]:
+        return {k: v for r in self.recorders for k, v in r.get_summary().items()}
+
+
+def _format_ids(requests: List[SimRequest], limit: int = 5) -> str:
+    if not requests:
+        return "-"
+    ids = ",".join(r.request_id for r in requests[:limit])
+    if len(requests) > limit:
+        ids += f"...+{len(requests) - limit}"
+    return ids