from __future__ import annotations import argparse import json import os import sys from http import HTTPStatus from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import Any ROOT = Path(__file__).resolve().parents[1] WORKSPACE = ROOT.parent STATIC_DIR = Path(__file__).resolve().parent / "static" os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib") os.environ.setdefault("XDG_CACHE_HOME", "/tmp") sys.path.insert(0, str(ROOT)) sys.path.insert(0, str(WORKSPACE / "pp-scheduler")) from pp_scheduler import ( # noqa: E402 BigMac, BigMacVPP, DistTrainEncoder1F1B, DistTrain1F1B, GPipe, OpType, Sandwich, SandwichVPP, Schedule1F1B, UnifiedBigMacVPP, VPP1F1BSchedule, ) from pp_simulator import OpDurationSpec, PipelineSimulator # noqa: E402 from pp_simulator.duration import DurationSampler # noqa: E402 def _max_microbatches_from_env() -> int | None: raw_value = os.environ.get("PP_SIMULATOR_MAX_MICROBATCHES", "").strip() if not raw_value: return None max_microbatches = int(raw_value) if max_microbatches <= 0: raise ValueError("PP_SIMULATOR_MAX_MICROBATCHES must be positive when set") return max_microbatches SCHEDULERS = { "UnifiedBigMacVPP": UnifiedBigMacVPP, "DistTrain1F1B": DistTrain1F1B, "DistTrainEncoder1F1B": DistTrainEncoder1F1B, "BigMacVPP": BigMacVPP, "SandwichVPP": SandwichVPP, "VPP1F1BSchedule": VPP1F1BSchedule, "BigMac": BigMac, "Sandwich": Sandwich, "Schedule1F1B": Schedule1F1B, "GPipe": GPipe, } SCHEDULER_OPTIONS = [ { "name": "UnifiedBigMacVPP", "uses_vpp": True, "uses_ve_forward_limit": True, "vpp_mode": "vpp", "include_encoder": True, "include_generator": True, "op_types": ["F", "B", "VF", "VB", "GF", "GB"], "description": "Unified BigMac VPP", }, { "name": "DistTrain1F1B", "uses_vpp": False, "uses_ve_forward_limit": False, "vpp_mode": "no_vpp", "include_encoder": True, "include_generator": True, "op_types": ["F", "B", "VF", "VB", "GF", "GB"], "description": "DistTrain 1F1B", }, { "name": "DistTrainEncoder1F1B", "uses_vpp": False, "uses_ve_forward_limit": False, "vpp_mode": "no_vpp", "include_encoder": True, "include_generator": False, "op_types": ["F", "B", "VF", "VB"], "description": "DistTrain encoder + LLM 1F1B", }, { "name": "BigMacVPP", "uses_vpp": True, "uses_ve_forward_limit": True, "vpp_mode": "vpp", "include_encoder": True, "include_generator": False, "op_types": ["F", "B", "VF", "VB"], "description": "BigMac VPP", }, { "name": "SandwichVPP", "uses_vpp": True, "uses_ve_forward_limit": False, "vpp_mode": "vpp", "include_encoder": True, "include_generator": False, "op_types": ["F", "B", "VF", "VB"], "description": "Sandwich VPP", }, { "name": "VPP1F1BSchedule", "uses_vpp": True, "uses_ve_forward_limit": False, "vpp_mode": "vpp", "include_encoder": False, "include_generator": False, "op_types": ["F", "B"], "description": "VPP 1F1B", }, { "name": "BigMac", "uses_vpp": False, "uses_ve_forward_limit": True, "vpp_mode": "no_vpp", "include_encoder": True, "include_generator": False, "op_types": ["F", "B", "VF", "VB"], "description": "BigMac", }, { "name": "Sandwich", "uses_vpp": False, "uses_ve_forward_limit": False, "vpp_mode": "no_vpp", "include_encoder": True, "include_generator": False, "op_types": ["F", "B", "VF", "VB"], "description": "Sandwich", }, { "name": "Schedule1F1B", "uses_vpp": False, "uses_ve_forward_limit": False, "vpp_mode": "no_vpp", "include_encoder": False, "include_generator": False, "op_types": ["F", "B"], "description": "1F1B", }, { "name": "GPipe", "uses_vpp": False, "uses_ve_forward_limit": False, "vpp_mode": "no_vpp", "include_encoder": False, "include_generator": False, "op_types": ["F", "B"], "description": "GPipe", }, ] DEFAULT_DURATION_SPECS = { "F": {"mean": 1.0, "variance": 0.0}, "B": {"mean": 2.0, "variance": 0.0}, "VF": {"mean": 0.8, "variance": 0.01}, "VB": {"mean": 0.9, "variance": 0.01}, "GF": {"mean": 0.6, "variance": 0.01}, "GB": {"mean": 0.7, "variance": 0.01}, } MAX_MICROBATCHES = _max_microbatches_from_env() class SimulatorRequestHandler(SimpleHTTPRequestHandler): def __init__(self, *args, **kwargs): super().__init__(*args, directory=str(STATIC_DIR), **kwargs) def log_message(self, format: str, *args: Any) -> None: print(f"[pp-simulator] {self.address_string()} - {format % args}") def do_GET(self) -> None: if self.path == "/api/schedulers": self._send_json( { "schedulers": SCHEDULER_OPTIONS, "duration_specs": DEFAULT_DURATION_SPECS, "limits": {"max_microbatches": MAX_MICROBATCHES}, } ) return if self.path == "/": self.path = "/index.html" super().do_GET() def do_POST(self) -> None: if self.path not in {"/api/simulate", "/api/compare"}: self.send_error(HTTPStatus.NOT_FOUND, "Unknown endpoint") return try: payload = self._read_json() if self.path == "/api/compare": response = run_comparison(payload) else: response = run_simulation(payload) except Exception as exc: # Keep API errors JSON-shaped for the UI. self._send_json({"error": str(exc)}, status=HTTPStatus.BAD_REQUEST) return self._send_json(response) def _read_json(self) -> dict[str, Any]: length = int(self.headers.get("Content-Length", "0")) if length <= 0: return {} body = self.rfile.read(length) return json.loads(body.decode("utf-8")) def _send_json(self, payload: dict[str, Any], *, status: HTTPStatus = HTTPStatus.OK) -> None: data = json.dumps(payload, sort_keys=True).encode("utf-8") self.send_response(status) self.send_header("Content-Type", "application/json; charset=utf-8") self.send_header("Cache-Control", "no-store") self.send_header("Content-Length", str(len(data))) self.end_headers() self.wfile.write(data) def end_headers(self) -> None: self.send_header("Cache-Control", "no-store") super().end_headers() def run_simulation(payload: dict[str, Any]) -> dict[str, Any]: scheduler_name = str(payload.get("scheduler", "UnifiedBigMacVPP")) scheduler_cls = SCHEDULERS.get(scheduler_name) if scheduler_cls is None: raise ValueError(f"Unsupported scheduler: {scheduler_name}") pp_size = _positive_int(payload.get("pp_size", 4), "pp_size") vpp_size = _positive_int(payload.get("vpp_size", 2), "vpp_size") num_microbatches = _positive_int( payload.get("num_microbatches", 8), "num_microbatches", max_value=MAX_MICROBATCHES, ) ve_forward_limit = _positive_int(payload.get("ve_forward_limit", 3), "ve_forward_limit") seed = payload.get("seed", 7) seed = None if seed in ("", None) else int(seed) scheduler = _build_scheduler( scheduler_name, scheduler_cls, pp_size=pp_size, vpp_size=vpp_size, num_microbatches=num_microbatches, ve_forward_limit=ve_forward_limit, ) scheduler.generate_schedule() duration_specs = _duration_specs(payload.get("duration_specs", DEFAULT_DURATION_SPECS)) simulator = PipelineSimulator.from_scheduler(scheduler) result = simulator.simulate(duration_specs, seed=seed) simulator.validate_result(result) return { "simulation": result.to_dict(), "chrome_trace": result.to_chrome_trace_dict(), "perfetto_trace": result.to_chrome_trace_dict(perfetto_compat=True), } def run_comparison(payload: dict[str, Any]) -> dict[str, Any]: workload = payload.get("workload", payload) if not isinstance(workload, dict): raise ValueError("workload must be an object") scheduler_a_name = str(payload.get("scheduler_a", "")) scheduler_b_name = str(payload.get("scheduler_b", "")) scheduler_a_cls = SCHEDULERS.get(scheduler_a_name) scheduler_b_cls = SCHEDULERS.get(scheduler_b_name) if scheduler_a_cls is None: raise ValueError(f"Unsupported scheduler_a: {scheduler_a_name}") if scheduler_b_cls is None: raise ValueError(f"Unsupported scheduler_b: {scheduler_b_name}") pp_size = _positive_int(workload.get("pp_size", 4), "pp_size") vpp_size = _positive_int(workload.get("vpp_size", 2), "vpp_size") num_microbatches = _positive_int( workload.get("num_microbatches", 8), "num_microbatches", max_value=MAX_MICROBATCHES, ) ve_forward_limit = _positive_int(workload.get("ve_forward_limit", 3), "ve_forward_limit") include_encoder = _bool(workload.get("include_encoder", False)) include_generator = _bool(workload.get("include_generator", False)) scale_llm_by_pp_split = True allow_cross_vpp = True seed = payload.get("seed", 7) seed = None if seed in ("", None) else int(seed) shape = { "include_encoder": include_encoder, "include_generator": include_generator, "allow_cross_vpp": allow_cross_vpp, } _validate_scheduler_shape(scheduler_a_name, shape) _validate_scheduler_shape(scheduler_b_name, shape) scheduler_a = _build_compare_scheduler( scheduler_a_name, scheduler_a_cls, stage_budget=pp_size, vpp_size=vpp_size, num_microbatches=num_microbatches, ve_forward_limit=ve_forward_limit, include_encoder=include_encoder, include_generator=include_generator, ) scheduler_b = _build_compare_scheduler( scheduler_b_name, scheduler_b_cls, stage_budget=pp_size, vpp_size=vpp_size, num_microbatches=num_microbatches, ve_forward_limit=ve_forward_limit, include_encoder=include_encoder, include_generator=include_generator, ) scheduler_a.generate_schedule() scheduler_b.generate_schedule() duration_specs = _duration_specs(payload.get("duration_specs", DEFAULT_DURATION_SPECS)) simulator_a = PipelineSimulator.from_scheduler(scheduler_a) simulator_b = PipelineSimulator.from_scheduler(scheduler_b) normalize_llm_chunks = allow_cross_vpp and scale_llm_by_pp_split shared_table = _shared_duration_table( [simulator_a, simulator_b], duration_specs, seed=seed, normalize_llm_chunks=normalize_llm_chunks, ) result_a = simulator_a.simulate( duration_specs, seed=seed, duration_overrides=_duration_overrides( simulator_a, shared_table, stage_budget=pp_size, scale_llm_by_pp_split=scale_llm_by_pp_split, normalize_llm_chunks=normalize_llm_chunks, ), ) result_b = simulator_b.simulate( duration_specs, seed=seed, duration_overrides=_duration_overrides( simulator_b, shared_table, stage_budget=pp_size, scale_llm_by_pp_split=scale_llm_by_pp_split, normalize_llm_chunks=normalize_llm_chunks, ), ) simulator_a.validate_result(result_a) simulator_b.validate_result(result_b) makespan_a = float(result_a.summary["makespan"]) makespan_b = float(result_b.summary["makespan"]) return { "result_a": result_a.to_dict(), "result_b": result_b.to_dict(), "comparison": { "makespan_a": makespan_a, "makespan_b": makespan_b, "makespan_delta": makespan_b - makespan_a, "speedup_a_over_b": makespan_b / makespan_a if makespan_a > 0 else None, "shared_duration_count": len(shared_table), "workload": { "pp_size": pp_size, "vpp_size": vpp_size, "num_microbatches": num_microbatches, "ve_forward_limit": ve_forward_limit, "include_encoder": include_encoder, "include_generator": include_generator, "scale_llm_by_pp_split": scale_llm_by_pp_split, "allow_cross_vpp": allow_cross_vpp, }, }, } def _build_scheduler( scheduler_name: str, scheduler_cls, *, pp_size: int, vpp_size: int, num_microbatches: int, ve_forward_limit: int, ): if scheduler_name in {"UnifiedBigMacVPP", "BigMacVPP"}: return scheduler_cls( pp_size=pp_size, vpp_size=vpp_size, num_microbatches=num_microbatches, ve_forward_limit=ve_forward_limit, ) if scheduler_name in {"SandwichVPP", "VPP1F1BSchedule"}: return scheduler_cls(pp_size=pp_size, vpp_size=vpp_size, num_microbatches=num_microbatches) if scheduler_name in {"DistTrain1F1B", "DistTrainEncoder1F1B"}: return scheduler_cls(pp_size=pp_size, num_microbatches=num_microbatches) if scheduler_name == "BigMac": return scheduler_cls( pp_size=pp_size, num_microbatches=num_microbatches, ve_forward_limit=ve_forward_limit, check_ve_forward_limit=False, ) return scheduler_cls(pp_size=pp_size, num_microbatches=num_microbatches) def _build_compare_scheduler( scheduler_name: str, scheduler_cls, *, stage_budget: int, vpp_size: int, num_microbatches: int, ve_forward_limit: int, include_encoder: bool, include_generator: bool, ): pp_size = stage_budget if scheduler_name in {"DistTrain1F1B", "DistTrainEncoder1F1B"}: extra_stages = int(include_encoder) + int(include_generator) pp_size = stage_budget - extra_stages if pp_size <= 0: raise ValueError( f"DistTrain needs stage budget > {extra_stages}, got {stage_budget}" ) return _build_scheduler( scheduler_name, scheduler_cls, pp_size=pp_size, vpp_size=vpp_size, num_microbatches=num_microbatches, ve_forward_limit=ve_forward_limit, ) def _validate_scheduler_shape(scheduler_name: str, shape: dict[str, Any]) -> None: option = _scheduler_option(scheduler_name) keys = ["include_encoder", "include_generator"] for key in keys: if option.get(key) != shape[key]: raise ValueError( f"{scheduler_name} is not compatible with workload {key}={shape[key]!r}" ) def _scheduler_option(scheduler_name: str) -> dict[str, Any]: for option in SCHEDULER_OPTIONS: if option["name"] == scheduler_name: return option raise ValueError(f"Unsupported scheduler: {scheduler_name}") def _shared_duration_table( simulators: list[PipelineSimulator], duration_specs: dict[Any, OpDurationSpec], *, seed: int | None, normalize_llm_chunks: bool = False, ) -> dict[tuple[int, str, int | None, int | None], float]: sampler = DurationSampler(duration_specs, seed=seed) table: dict[tuple[int, str, int | None, int | None], float] = {} for simulator in simulators: for op in simulator.plan.ops: key = _duration_key(op, normalize_llm_chunks=normalize_llm_chunks) if key not in table: table[key] = sampler.sample(op.op_type, fallback_duration=op.base_duration) return table def _duration_overrides( simulator: PipelineSimulator, shared_table: dict[tuple[int, str, int | None, int | None], float], *, stage_budget: int, scale_llm_by_pp_split: bool, normalize_llm_chunks: bool = False, ) -> dict[str, float]: return { op.id: shared_table[_duration_key( op, normalize_llm_chunks=normalize_llm_chunks, )] * _duration_scale( simulator, op, stage_budget=stage_budget, scale_llm_by_pp_split=scale_llm_by_pp_split, ) for op in simulator.plan.ops } def _duration_scale( simulator: PipelineSimulator, op, *, stage_budget: int, scale_llm_by_pp_split: bool, ) -> float: if not scale_llm_by_pp_split or op.op_type not in {"F", "B"}: return 1.0 layout = simulator.plan.pipeline_layout or {} llm_pp_size = int(layout.get("llm_pp_size", stage_budget)) effective_llm_partitions = llm_pp_size * max(1, int(simulator.plan.vpp_size)) if effective_llm_partitions <= 0: return 1.0 return float(stage_budget) / float(effective_llm_partitions) def _duration_key( op, *, normalize_llm_chunks: bool = False, ) -> tuple[int, str, int | None, int | None]: chunk_id = 0 if op.chunk_id is None else int(op.chunk_id) if normalize_llm_chunks and op.op_type in {"F", "B"}: chunk_id = 0 return (int(op.rank), str(op.op_type), op.microbatch_id, chunk_id) def _duration_specs(raw_specs: Any) -> dict[Any, OpDurationSpec]: if not isinstance(raw_specs, dict): raise ValueError("duration_specs must be an object") specs = {} for op_value, default_spec in DEFAULT_DURATION_SPECS.items(): raw_spec = raw_specs.get(op_value, default_spec) if not isinstance(raw_spec, dict): raise ValueError(f"duration spec for {op_value} must be an object") specs[_op_type(op_value)] = OpDurationSpec( mean=float(raw_spec.get("mean", default_spec["mean"])), variance=float(raw_spec.get("variance", default_spec["variance"])), min_value=float(raw_spec.get("min_value", 0.0)), ) return specs def _op_type(op_value: str): for op_type in OpType: if op_type.value == op_value: return op_type return op_value def _positive_int(value: Any, name: str, *, max_value: int | None = None) -> int: number = int(value) if number <= 0: raise ValueError(f"{name} must be positive") if max_value is not None and number > max_value: raise ValueError(f"{name} must be <= {max_value} for this deployment") return number def _bool(value: Any) -> bool: if isinstance(value, bool): return value if isinstance(value, str): return value.lower() in {"1", "true", "yes", "on"} return bool(value) def main() -> None: parser = argparse.ArgumentParser(description="Run the PP Simulator web UI") parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--port", type=int, default=8765) args = parser.parse_args() server = ThreadingHTTPServer((args.host, args.port), SimulatorRequestHandler) url = f"http://{args.host}:{args.port}" print(f"PP Simulator web UI: {url}") try: server.serve_forever() except KeyboardInterrupt: print("\nStopping PP Simulator web UI") finally: server.server_close() if __name__ == "__main__": main()