pp-simulator/pp_simulator/simulator.py

from __future__ import annotations

import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Mapping

from .dependencies import DependencyGraph
from .duration import DurationSampler, OpDurationSpec
from .plan import Plan, PlanOp


CHROME_TRACE_CHUNK_COLORS = [
    "thread_state_running",
    "rail_response",
    "rail_animation",
    "rail_idle",
    "rail_load",
    "good",
    "bad",
    "terrible",
]


@dataclass(frozen=True)
class SimulatedOp:
    id: str
    rank: int
    index: int
    original_index: int
    op_type: str
    microbatch_id: int | None
    chunk_id: int | None
    label: str
    duration: float
    start_time: float
    end_time: float
    wait_time: float
    deps: list[str]
    dep_reasons: dict[str, list[str]]

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "rank": self.rank,
            "index": self.index,
            "original_index": self.original_index,
            "op_type": self.op_type,
            "microbatch_id": self.microbatch_id,
            "chunk_id": self.chunk_id,
            "label": self.label,
            "duration": self.duration,
            "start_time": self.start_time,
            "end_time": self.end_time,
            "wait_time": self.wait_time,
            "deps": self.deps,
            "dep_reasons": self.dep_reasons,
        }


@dataclass(frozen=True)
class SimulationResult:
    metadata: dict[str, Any]
    ops: list[SimulatedOp]
    summary: dict[str, Any]

    def to_dict(self) -> dict[str, Any]:
        return {
            "metadata": self.metadata,
            "summary": self.summary,
            "ops": [op.to_dict() for op in self.ops],
        }

    def to_json(self, path: str | Path | None = None, *, indent: int = 2) -> str:
        payload = json.dumps(self.to_dict(), indent=indent, sort_keys=True)
        if path is not None:
            Path(path).write_text(payload + "\n", encoding="utf-8")
        return payload

    def to_chrome_trace_dict(
        self,
        *,
        time_unit_scale: float = 1000.0,
        perfetto_compat: bool = False,
    ) -> dict[str, Any]:
        """Return a Chrome Trace JSON payload.

        Chrome Trace timestamps are conventionally microseconds. Simulator time
        units are abstract, so the default maps one simulator unit to 1000 us.
        """

        events: list[dict[str, Any]] = []
        ranks = sorted({op.rank for op in self.ops})
        pid = 1
        tid_by_rank = {rank: rank + 1 for rank in ranks}

        events.append(
            {
                "name": "process_name",
                "ph": "M",
                "pid": pid,
                "tid": 0,
                "args": {"name": self.metadata.get("scheduler", "PP Simulator")},
            }
        )

        for rank in ranks:
            events.append(
                {
                    "name": "thread_name",
                    "ph": "M",
                    "pid": pid,
                    "tid": tid_by_rank[rank],
                    "args": {"name": f"PP Rank {rank}"},
                }
            )
            events.append(
                {
                    "name": "thread_sort_index",
                    "ph": "M",
                    "pid": pid,
                    "tid": tid_by_rank[rank],
                    "args": {"sort_index": rank},
                }
            )

        slice_events = []
        for op in self.ops:
            pp_size = int(self.metadata.get("pp_size", 1))
            trace_category = _trace_category(op)
            slice_events.append(
                {
                    "name": _trace_event_name(op, perfetto_compat=perfetto_compat),
                    "cat": trace_category,
                    "cname": _trace_color_name(op),
                    "ph": "X",
                    "pid": pid,
                    "tid": tid_by_rank[op.rank],
                    "ts": op.start_time * time_unit_scale,
                    "dur": op.duration * time_unit_scale,
                    "args": {
                        "id": op.id,
                        "label": op.label,
                        "op_type": op.op_type,
                        "microbatch_id": op.microbatch_id,
                        "batch_id": (
                            op.microbatch_id // pp_size
                            if op.microbatch_id is not None and pp_size > 0
                            else None
                        ),
                        "chunk_id": op.chunk_id,
                        "rank": op.rank,
                        "pp_rank": op.rank,
                        "index": op.index,
                        "original_index": op.original_index,
                        "start_time": op.start_time,
                        "end_time": op.end_time,
                        "duration": op.duration,
                        "wait_time": op.wait_time,
                        "visual_category": trace_category,
                        "deps": op.deps,
                        "dep_reasons": op.dep_reasons,
                    },
                }
            )
        events.extend(sorted(slice_events, key=lambda event: (event["ts"], event["tid"], event["name"])))

        return {
            "displayTimeUnit": "ms",
            "metadata": self.metadata,
            "summary": self.summary,
            "traceEvents": events,
        }

    def to_chrome_trace(
        self,
        path: str | Path | None = None,
        *,
        indent: int = 2,
        time_unit_scale: float = 1000.0,
        perfetto_compat: bool = False,
    ) -> str:
        payload = json.dumps(
            self.to_chrome_trace_dict(
                time_unit_scale=time_unit_scale,
                perfetto_compat=perfetto_compat,
            ),
            indent=indent,
            sort_keys=True,
        )
        if path is not None:
            Path(path).write_text(payload + "\n", encoding="utf-8")
        return payload


@dataclass(frozen=True)
class MonteCarloReport:
    metadata: dict[str, Any]
    statistics: dict[str, Any]
    trial_summaries: list[dict[str, Any]]

    def to_dict(self) -> dict[str, Any]:
        return {
            "metadata": self.metadata,
            "statistics": self.statistics,
            "trial_summaries": self.trial_summaries,
        }

    def to_json(self, path: str | Path | None = None, *, indent: int = 2) -> str:
        payload = json.dumps(self.to_dict(), indent=indent, sort_keys=True)
        if path is not None:
            Path(path).write_text(payload + "\n", encoding="utf-8")
        return payload


def _trace_category(op: SimulatedOp) -> str:
    if op.chunk_id is None:
        return f"{op.op_type}/chunk_none"
    return f"{op.op_type}/chunk_{_chunk_label(op.chunk_id)}"


def _trace_event_name(op: SimulatedOp, *, perfetto_compat: bool) -> str:
    if not perfetto_compat:
        return op.label

    chunk = "none" if op.chunk_id is None else _chunk_label(op.chunk_id)
    mb = "none" if op.microbatch_id is None else str(op.microbatch_id)
    return f"rank{op.rank}/{op.op_type}/chunk_{chunk}/mb{mb}"


def _trace_color_name(op: SimulatedOp) -> str:
    if op.chunk_id is None:
        return CHROME_TRACE_CHUNK_COLORS[0]
    return CHROME_TRACE_CHUNK_COLORS[int(op.chunk_id) % len(CHROME_TRACE_CHUNK_COLORS)]


def _chunk_label(chunk_id: int) -> str:
    chunk_id = int(chunk_id)
    if chunk_id < 0:
        return str(chunk_id)

    letters = []
    value = chunk_id
    while True:
        letters.append(chr(ord("a") + (value % 26)))
        value = value // 26 - 1
        if value < 0:
            break
    return "".join(reversed(letters))


class PipelineSimulator:
    def __init__(self, plan: Plan):
        self.plan = plan
        self.graph = DependencyGraph(plan)
        self._ops_by_id = {op.id: op for op in plan.ops}

    @classmethod
    def from_scheduler(cls, scheduler: Any) -> "PipelineSimulator":
        return cls(Plan.from_scheduler(scheduler))

    def simulate(
        self,
        duration_specs: Mapping[Any, OpDurationSpec | Mapping[str, float]],
        *,
        seed: int | None = None,
        default_spec: OpDurationSpec | None = None,
        duration_overrides: Mapping[str, float] | None = None,
    ) -> SimulationResult:
        if duration_overrides is None:
            sampler = DurationSampler(duration_specs, default_spec=default_spec, seed=seed)
            sampled_durations = {
                op.id: sampler.sample(op.op_type, fallback_duration=op.base_duration)
                for op in self.plan.ops
            }
        else:
            sampled_durations = {
                op.id: max(0.0, float(duration_overrides[op.id]))
                for op in self.plan.ops
            }

        rank_free_time = {rank: 0.0 for rank in self.plan.ops_by_rank}
        simulated_by_id: dict[str, SimulatedOp] = {}

        for op_id in self.graph.topological_order:
            op = self._ops_by_id[op_id]
            dep_ids = sorted(self.graph.dependencies[op.id])
            dependency_ready_time = max(
                (simulated_by_id[dep_id].end_time for dep_id in dep_ids),
                default=0.0,
            )
            previous_rank_time = rank_free_time[op.rank]
            start_time = max(previous_rank_time, dependency_ready_time)
            duration = sampled_durations[op.id]
            end_time = start_time + duration

            simulated_by_id[op.id] = SimulatedOp(
                id=op.id,
                rank=op.rank,
                index=op.index,
                original_index=op.original_index,
                op_type=op.op_type,
                microbatch_id=op.microbatch_id,
                chunk_id=op.chunk_id,
                label=op.label,
                duration=duration,
                start_time=start_time,
                end_time=end_time,
                wait_time=max(0.0, start_time - previous_rank_time),
                deps=dep_ids,
                dep_reasons={
                    dep_id: sorted(self.graph.dependency_reasons[op.id][dep_id])
                    for dep_id in dep_ids
                },
            )
            rank_free_time[op.rank] = end_time

        ops = [
            simulated_by_id[op.id]
            for rank in sorted(self.plan.ops_by_rank)
            for op in self.plan.ops_by_rank[rank]
        ]
        summary = self._build_summary(ops)

        return SimulationResult(
            metadata={
                "scheduler": self.plan.scheduler_name,
                "pp_size": self.plan.pp_size,
                "vpp_size": self.plan.vpp_size,
                "num_microbatches": self.plan.num_microbatches,
                "seed": seed,
                "op_count": len(ops),
                "pipeline_layout": self.plan.pipeline_layout,
            },
            ops=ops,
            summary=summary,
        )

    def monte_carlo(
        self,
        duration_specs: Mapping[Any, OpDurationSpec | Mapping[str, float]],
        *,
        num_trials: int,
        seed: int | None = None,
        default_spec: OpDurationSpec | None = None,
        validate: bool = False,
    ) -> MonteCarloReport:
        if num_trials <= 0:
            raise ValueError(f"num_trials must be positive, got {num_trials}")

        trial_summaries: list[dict[str, Any]] = []
        for trial_index in range(num_trials):
            trial_seed = None if seed is None else seed + trial_index
            result = self.simulate(
                duration_specs,
                seed=trial_seed,
                default_spec=default_spec,
            )
            if validate:
                self.validate_result(result)
            trial_summaries.append(
                {
                    "trial_index": trial_index,
                    "seed": trial_seed,
                    "summary": result.summary,
                }
            )

        return MonteCarloReport(
            metadata={
                "scheduler": self.plan.scheduler_name,
                "pp_size": self.plan.pp_size,
                "vpp_size": self.plan.vpp_size,
                "num_microbatches": self.plan.num_microbatches,
                "op_count": len(self.plan.ops),
                "num_trials": num_trials,
                "seed": seed,
                "pipeline_layout": self.plan.pipeline_layout,
            },
            statistics=_build_monte_carlo_statistics(trial_summaries),
            trial_summaries=trial_summaries,
        )

    def validate_result(self, result: SimulationResult, *, tolerance: float = 1e-9) -> None:
        by_id = {op.id: op for op in result.ops}
        for op in result.ops:
            for dep_id in op.deps:
                if op.start_time + tolerance < by_id[dep_id].end_time:
                    raise AssertionError(f"{op.id} starts before dependency {dep_id} ends")

        for rank in sorted(self.plan.ops_by_rank):
            rank_ops = [op for op in result.ops if op.rank == rank]
            for previous, current in zip(rank_ops, rank_ops[1:]):
                if current.start_time + tolerance < previous.end_time:
                    raise AssertionError(f"{current.id} overlaps previous rank op {previous.id}")

    def _build_summary(self, ops: list[SimulatedOp]) -> dict[str, Any]:
        makespan = max((op.end_time for op in ops), default=0.0)
        rank_compute_time: dict[int, float] = {rank: 0.0 for rank in self.plan.ops_by_rank}
        rank_wait_time: dict[int, float] = {rank: 0.0 for rank in self.plan.ops_by_rank}
        op_type_time: dict[str, float] = {}

        for op in ops:
            rank_compute_time[op.rank] += op.duration
            rank_wait_time[op.rank] += op.wait_time
            op_type_time[op.op_type] = op_type_time.get(op.op_type, 0.0) + op.duration

        return {
            "makespan": makespan,
            "rank_compute_time": {str(rank): value for rank, value in rank_compute_time.items()},
            "rank_wait_time": {str(rank): value for rank, value in rank_wait_time.items()},
            "rank_utilization": {
                str(rank): (value / makespan if makespan > 0 else 0.0)
                for rank, value in rank_compute_time.items()
            },
            "op_type_time": dict(sorted(op_type_time.items())),
            "total_wait_time": sum(rank_wait_time.values()),
        }


def _build_monte_carlo_statistics(trial_summaries: list[dict[str, Any]]) -> dict[str, Any]:
    summaries = [trial["summary"] for trial in trial_summaries]
    return {
        "makespan": _summarize_values(summary["makespan"] for summary in summaries),
        "total_wait_time": _summarize_values(summary["total_wait_time"] for summary in summaries),
        "rank_compute_time": _summarize_nested_metric(summaries, "rank_compute_time"),
        "rank_wait_time": _summarize_nested_metric(summaries, "rank_wait_time"),
        "rank_utilization": _summarize_nested_metric(summaries, "rank_utilization"),
        "op_type_time": _summarize_nested_metric(summaries, "op_type_time"),
    }


def _summarize_nested_metric(summaries: list[dict[str, Any]], metric_name: str) -> dict[str, Any]:
    keys = sorted({key for summary in summaries for key in summary[metric_name]})
    return {
        key: _summarize_values(summary[metric_name].get(key, 0.0) for summary in summaries)
        for key in keys
    }


def _summarize_values(values: Any) -> dict[str, float]:
    sorted_values = sorted(float(value) for value in values)
    if not sorted_values:
        return {
            "count": 0,
            "mean": 0.0,
            "std": 0.0,
            "min": 0.0,
            "max": 0.0,
            "p50": 0.0,
            "p90": 0.0,
            "p95": 0.0,
            "p99": 0.0,
        }

    count = len(sorted_values)
    mean = sum(sorted_values) / count
    variance = sum((value - mean) ** 2 for value in sorted_values) / count
    return {
        "count": count,
        "mean": mean,
        "std": math.sqrt(variance),
        "min": sorted_values[0],
        "max": sorted_values[-1],
        "p50": _percentile(sorted_values, 50),
        "p90": _percentile(sorted_values, 90),
        "p95": _percentile(sorted_values, 95),
        "p99": _percentile(sorted_values, 99),
    }


def _percentile(sorted_values: list[float], percentile: float) -> float:
    if len(sorted_values) == 1:
        return sorted_values[0]

    position = (len(sorted_values) - 1) * percentile / 100.0
    lower = int(math.floor(position))
    upper = int(math.ceil(position))
    if lower == upper:
        return sorted_values[lower]

    weight = position - lower
    return sorted_values[lower] * (1.0 - weight) + sorted_values[upper] * weight