pp-simulator/web/app.py

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()