vllm-omni/vllm_omni/metrics/stats.py

from __future__ import annotations

import time
from collections import defaultdict
from collections.abc import Callable
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any

from vllm.logger import init_logger

from vllm_omni.metrics.utils import _build_field_defs, _build_row, _format_table

logger = init_logger(__name__)


@dataclass
class StageStats:
    total_token: int = 0
    total_gen_time_ms: float = 0.0

    @property
    def avg_tokens_per_s(self) -> float:
        return (self.total_token * 1000.0 / self.total_gen_time_ms) if self.total_gen_time_ms > 0 else 0.0


@dataclass
class StageRequestStats:
    batch_id: int
    batch_size: int
    num_tokens_in: int
    num_tokens_out: int
    stage_gen_time_ms: float
    rx_transfer_bytes: int
    rx_decode_time_ms: float
    rx_in_flight_time_ms: float
    stage_stats: StageStats
    stage_id: int | None = None
    final_output_type: str | None = None
    request_id: str | None = None
    postprocess_time_ms: float = 0.0
    diffusion_metrics: dict[str, int] = None
    audio_generated_frames: int = 0
    pipeline_timings: dict[str, float] | None = None

    @property
    def rx_mbps(self) -> float:
        return (
            (float(self.rx_transfer_bytes) * 8.0) / (max(float(self.rx_decode_time_ms), 1e-6) * 1000.0)
            if self.rx_transfer_bytes > 0
            else 0.0
        )

    @property
    def tokens_per_s(self) -> float:
        return (self.num_tokens_out * 1000.0 / self.stage_gen_time_ms) if (self.stage_gen_time_ms > 0) else 0.0


@dataclass
class TransferEdgeStats:
    from_stage: int
    to_stage: int
    request_id: str
    size_bytes: int
    tx_time_ms: float
    used_shm: bool = False
    rx_decode_time_ms: float = 0.0
    in_flight_time_ms: float = 0.0

    @property
    def total_time_ms(self) -> float:
        return float(self.tx_time_ms) + float(self.rx_decode_time_ms) + float(self.in_flight_time_ms)


@dataclass
class RequestE2EStats:
    request_id: str
    e2e_total_ms: float
    e2e_total_tokens: int
    transfers_total_time_ms: float
    transfers_total_bytes: int

    @property
    def e2e_tpt(self) -> float:
        return (self.e2e_total_ms / self.e2e_total_tokens) if self.e2e_total_tokens > 0 else 0.0


# === Field Configuration ===
# Fields requiring unit conversion:  original_field_name -> (display_name, transform_fn)
FIELD_TRANSFORMS: dict[str, tuple[str, Callable[[Any], Any]]] = {
    "rx_transfer_bytes": ("rx_transfer_kbytes", lambda v: v / 1024.0),
    "size_bytes": ("size_kbytes", lambda v: v / 1024.0),
    "transfers_total_bytes": ("transfers_total_kbytes", lambda v: v / 1024.0),
}

# Fields to exclude from table display for each event type
STAGE_EXCLUDE = {
    "stage_stats",
    "stage_id",
    "request_id",
    "rx_transfer_bytes",
    "rx_decode_time_ms",
    "rx_in_flight_time_ms",
    "final_output_type",
    "pipeline_timings",
}
TRANSFER_EXCLUDE = {"from_stage", "to_stage", "request_id", "used_shm"}
E2E_EXCLUDE = {"request_id"}

# Decide the order of overall summary fields, or None for auto
OVERALL_FIELDS: list[str] | None = None
STAGE_FIELDS = _build_field_defs(StageRequestStats, STAGE_EXCLUDE, FIELD_TRANSFORMS)
TRANSFER_FIELDS = _build_field_defs(TransferEdgeStats, TRANSFER_EXCLUDE, FIELD_TRANSFORMS)
E2E_FIELDS = _build_field_defs(RequestE2EStats, E2E_EXCLUDE, FIELD_TRANSFORMS)


class OrchestratorAggregator:
    def __init__(
        self,
        num_stages: int,
        log_stats: bool,
        wall_start_ts: float,
        final_stage_id_for_e2e: dict[str, int] | int,
    ) -> None:
        self.num_stages = int(num_stages)
        self.log_stats = bool(log_stats)
        self.final_stage_id_for_e2e = final_stage_id_for_e2e
        self.init_run_state(wall_start_ts)
        self.stage_events: dict[str, list[StageRequestStats]] = {}
        self.transfer_events: dict[
            tuple[int, int, str], TransferEdgeStats
        ] = {}  # Key: (from_stage, to_stage, request_id)
        self.e2e_events: list[RequestE2EStats] = []

    def init_run_state(self, wall_start_ts: float) -> None:
        # Per-run aggregates and timing state
        self.stage_total_tokens = [0 for _ in range(self.num_stages)]
        self.e2e_total_ms = 0.0
        self.e2e_total_tokens = 0
        self.e2e_count = 0
        self.e2e_done = set()
        self.wall_start_ts = float(wall_start_ts)
        self.last_finish_ts = float(wall_start_ts)
        self.stage_first_ts = [None for _ in range(self.num_stages)]
        self.stage_last_ts = [None for _ in range(self.num_stages)]
        self.accumulated_gen_time_ms: defaultdict[str, defaultdict[int, float]] = defaultdict(
            lambda: defaultdict(float)
        )  # {request_id: {stage_id:accumulated_gen_time_ms}}
        self.diffusion_metrics: defaultdict[str, defaultdict[str, float]] = defaultdict(
            lambda: defaultdict(float)
        )  # {request_id: {diffusion_metrics_key: accumulated_metrics_data}}

    def _get_or_create_transfer_event(
        self,
        from_stage: int,
        to_stage: int,
        request_id: str,
    ) -> TransferEdgeStats:
        key = (from_stage, to_stage, request_id)
        evt = self.transfer_events.get(key)
        if evt is None:
            evt = TransferEdgeStats(
                from_stage=from_stage,
                to_stage=to_stage,
                request_id=request_id,
                size_bytes=0,
                tx_time_ms=0.0,
                used_shm=False,
                rx_decode_time_ms=0.0,
                in_flight_time_ms=0.0,
            )
            self.transfer_events[key] = evt
        return evt

    def record_transfer_tx(
        self,
        from_stage: int,
        to_stage: int,
        request_id: Any,
        size_bytes: int,
        tx_time_ms: float,
        used_shm: bool,
    ) -> TransferEdgeStats | None:
        try:
            evt = self._get_or_create_transfer_event(
                int(from_stage),
                int(to_stage),
                str(request_id),
            )
            # Accumulate tx metrics
            evt.size_bytes += int(size_bytes)
            evt.tx_time_ms += float(tx_time_ms)
            evt.used_shm = evt.used_shm or bool(used_shm)
            return evt
        except Exception:
            return None

    def record_transfer_rx(
        self,
        stats: StageRequestStats,
    ) -> TransferEdgeStats | None:
        try:
            if stats.stage_id is None or stats.stage_id <= 0:
                return None
            from_stage = int(stats.stage_id) - 1
            to_stage = int(stats.stage_id)
            rid_key = str(stats.request_id)
            evt = self._get_or_create_transfer_event(from_stage, to_stage, rid_key)
            # Accumulate rx metrics
            if evt.size_bytes == 0:
                # size_bytes has been recorded in tx phase
                evt.size_bytes = int(stats.rx_transfer_bytes)
            evt.rx_decode_time_ms += float(stats.rx_decode_time_ms)
            evt.in_flight_time_ms += float(stats.rx_in_flight_time_ms)
            return evt
        except Exception:
            return None

    def record_audio_generated_frames(
        self,
        output_to_yield: Any,
        stage_id: int,
        request_id: str,
    ) -> None:
        try:
            if (
                output_to_yield.final_output_type == "audio"
                and (multimodal_output := output_to_yield.multimodal_output.get("audio")) is not None
                and len(multimodal_output) > 0
            ):
                nframes = sum(
                    int(t.shape[0]) if t.ndim > 0 else 1
                    for t in (multimodal_output if isinstance(multimodal_output, list) else [multimodal_output])
                )
                stage_events_for_req = self.stage_events.get(request_id, [])
                if stage_events_for_req:
                    for stage_event in stage_events_for_req:
                        if stage_event.stage_id == stage_id:
                            stage_event.audio_generated_frames += nframes
                            break
                else:
                    logger.warning(
                        "Failed to record audio generated frames for request %s at stage %s: no stage event found",
                        request_id,
                        stage_id,
                    )
        except Exception:
            logger.debug(
                "Failed to record audio frames for request %s",
                request_id,
                exc_info=True,
            )

    def process_stage_metrics(
        self,
        *,
        result: dict[str, Any],
        stage_type: str,
        stage_id: int,
        req_id: str,
        engine_outputs: Any,
        finished: bool,
        final_output_type: str | None,
        output_to_yield: Any | None,
    ) -> None:
        """Process and record stage metrics.

        Args:
            result: Result dict containing metrics from stage
            stage_type: Type of the stage (e.g., 'llm', 'diffusion')
            stage_id: Stage identifier
            req_id: Request identifier
            engine_outputs: Engine output object
            finished: Whether stage processing is finished
            final_output_type: Type of final output (e.g., 'text', 'audio')
            output_to_yield: Output object to attach metrics to
        """
        try:
            _m: StageRequestStats | None = result.get("metrics")

            # 1. Accumulate metrics from stage stats
            if _m is not None:
                self.accumulated_gen_time_ms[req_id][stage_id] += _m.stage_gen_time_ms
                self.accumulate_diffusion_metrics(stage_type, req_id, engine_outputs)
                if finished:
                    self.on_stage_metrics(stage_id, req_id, _m, final_output_type)

            # 2. No output to yield, nothing more to do
            if output_to_yield is None:
                return

            # 3. Not finished yet — empty metrics, skip audio recording
            if not finished:
                output_to_yield.metrics = {}
                return

            # 4. Finished with output: assign text metrics if available
            output_to_yield.metrics = {}
            stage_event = next(
                (evt for evt in reversed(self.stage_events.get(req_id, [])) if evt.stage_id == stage_id),
                None,
            )
            if stage_event is not None and stage_event.final_output_type == "text":
                output_to_yield.metrics = {
                    "num_tokens_in": stage_event.num_tokens_in,
                    "num_tokens_out": stage_event.num_tokens_out,
                    "stage_id": stage_event.stage_id,
                    "final_output_type": stage_event.final_output_type,
                }

            # 5. Finished: record audio generated frames
            self.record_audio_generated_frames(output_to_yield, stage_id, req_id)

        except Exception:
            logger.exception(
                "Failed to process metrics for stage %s, req %s",
                stage_id,
                req_id,
            )

    def _as_stage_request_stats(
        self,
        stage_id: int,
        req_id: str,
        metrics: StageRequestStats,
        final_output_type: str | None = None,
    ) -> StageRequestStats:
        "Convert dict to StageRequestStats if needed."
        stats = metrics
        stats.stage_id = stage_id
        stats.request_id = req_id
        stats.final_output_type = final_output_type
        stats.diffusion_metrics = (
            {k: int(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
            if req_id in self.diffusion_metrics
            else None
        )
        return stats

    def on_stage_metrics(
        self,
        stage_id: int,
        req_id: Any,
        metrics: StageRequestStats,
        final_output_type: str | None = None,
    ) -> None:
        stats = self._as_stage_request_stats(stage_id, req_id, metrics, final_output_type)
        self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_out)
        if stats.stage_id == 0:
            self.stage_total_tokens[stats.stage_id] += int(stats.num_tokens_in)
        self.stage_events.setdefault(str(stats.request_id), []).append(stats)

        self.record_transfer_rx(stats)

    def record_stage_postprocess_time(self, stage_id: int, req_id: Any, postproc_time_ms: float) -> None:
        if req_id in self.stage_events:
            for stats in self.stage_events[req_id]:
                if stats.stage_id == stage_id:
                    stats.postprocess_time_ms = float(postproc_time_ms)
                    break
        else:
            logger.warning(
                "Failed to record postprocess time for request %s at stage %s: no stage event found",
                req_id,
                stage_id,
            )

    @contextmanager
    def stage_postprocess_timer(self, stage_id: int, req_id: Any):
        """Context manager for measuring and recording stage postprocessing time.

        Usage:
            with metrics.stage_postprocess_timer(stage_id, request_id):
                next_inputs = next_stage.process_engine_inputs(...)
        """
        _t0 = time.perf_counter()
        try:
            yield
        finally:
            _postproc_ms = (time.perf_counter() - _t0) * 1000.0
            self.record_stage_postprocess_time(stage_id, req_id, _postproc_ms)

    def accumulate_diffusion_metrics(self, stage_type: str, req_id: Any, engine_outputs: Any) -> None:
        """Accumulate diffusion metrics for a request.

        Handles extraction and accumulation of diffusion stage metrics.

        Args:
            req_id: Request ID
            engine_outputs: Engine output object containing metrics
        """
        if stage_type != "diffusion":
            return
        engine_output = engine_outputs[0] if isinstance(engine_outputs, list) and engine_outputs else engine_outputs
        diffusion_metrics: dict = getattr(engine_output, "metrics", {})
        if isinstance(diffusion_metrics, list):
            diffusion_metrics = diffusion_metrics[0]
        if diffusion_metrics:
            for key, value in diffusion_metrics.items():
                self.diffusion_metrics[req_id][key] += value

    def on_forward(
        self,
        from_stage: int,
        to_stage: int,
        req_id: Any,
        size_bytes: int,
        tx_ms: float,
        used_shm: bool,
    ) -> None:
        # Mark first input time for the destination stage if not set
        if self.stage_first_ts[to_stage] is None:
            self.stage_first_ts[to_stage] = time.time()
        self.record_transfer_tx(
            from_stage=from_stage,
            to_stage=to_stage,
            request_id=req_id,
            size_bytes=size_bytes,
            tx_time_ms=tx_ms,
            used_shm=used_shm,
        )

    def on_finalize_request(
        self,
        stage_id: int,
        req_id: Any,
        req_start_ts: float,
    ) -> None:
        rid_key = str(req_id)
        if rid_key in self.e2e_done:
            return  # Already finalized
        _t0 = float(req_start_ts)
        _t1 = time.time()
        # Update last output time for this stage
        prev_last = self.stage_last_ts[stage_id]
        self.stage_last_ts[stage_id] = _t1 if prev_last is None else max(prev_last, _t1)
        self.last_finish_ts = max(self.last_finish_ts, _t1)
        e2e_ms = (_t1 - _t0) * 1000.0

        # Sum tokens from all stages for this request
        # Include input tokens from stage 0 + output tokens from all stages
        total_tokens = 0
        if rid_key in self.stage_events:
            for evt in self.stage_events[rid_key]:
                if evt.stage_id == 0:
                    total_tokens += int(evt.num_tokens_in)
                total_tokens += int(evt.num_tokens_out)

        self.e2e_total_ms += e2e_ms
        self.e2e_total_tokens += total_tokens
        self.e2e_count += 1
        self.e2e_done.add(rid_key)
        per_req_record = RequestE2EStats(
            request_id=rid_key,
            e2e_total_ms=e2e_ms,
            e2e_total_tokens=total_tokens,
            transfers_total_time_ms=float(
                sum(evt.total_time_ms for evt in self.transfer_events.values() if evt.request_id == rid_key)
            ),
            transfers_total_bytes=int(
                sum(evt.size_bytes for evt in self.transfer_events.values() if evt.request_id == rid_key)
            ),
        )
        self.e2e_events.append(per_req_record)

    def build_and_log_summary(self) -> dict[str, Any]:
        if not self.log_stats:
            return {}
        wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
        e2e_avg_req = (wall_time_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
        e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / wall_time_ms) if wall_time_ms > 0 else 0.0

        if isinstance(self.final_stage_id_for_e2e, int):
            final_stage_id_map: dict[str, int] = {"*": int(self.final_stage_id_for_e2e)}
        else:
            final_stage_id_map = self.final_stage_id_for_e2e

        stage_wall_time_ms = [
            ((self.stage_last_ts[i] - self.stage_first_ts[i]) * 1000.0)
            if (self.stage_first_ts[i] is not None and self.stage_last_ts[i] is not None)
            else 0.0
            for i in range(self.num_stages)
        ]

        overall_summary = {
            "e2e_requests": int(self.e2e_count),
            "e2e_wall_time_ms": float(wall_time_ms),
            "e2e_total_tokens": int(self.e2e_total_tokens),
            "e2e_avg_time_per_request_ms": float(e2e_avg_req),
            "e2e_avg_tokens_per_s": float(e2e_avg_tok),
        }
        # Add average input preprocess time across all requests
        preprocess_times = []
        for _rid, evts in self.stage_events.items():
            for evt in evts:
                if evt.pipeline_timings and "preprocess_ms" in evt.pipeline_timings:
                    preprocess_times.append(evt.pipeline_timings["preprocess_ms"])
                    break  # only once per request
        if preprocess_times:
            overall_summary["input_preprocess_time_ms"] = sum(preprocess_times) / len(preprocess_times)
        # Add stage_wall_time_ms as separate fields for each stage
        for idx, wall_time in enumerate(stage_wall_time_ms):
            overall_summary[f"e2e_stage_{idx}_wall_time_ms"] = wall_time

        # Print overall summary
        # filter out all-zero fields for logging
        overall_fields = []
        for k in OVERALL_FIELDS or list(overall_summary.keys()):
            v = overall_summary.get(k, None)
            if v not in (0, 0.0, 0.000, None, ""):
                overall_fields.append(k)
        if overall_fields:
            logger.info(
                "\n%s",
                _format_table("Overall Summary", overall_summary, overall_fields),
            )

        all_request_ids = sorted(set(self.stage_events.keys()) | {e.request_id for e in self.e2e_events})

        result_stage_table = []
        result_trans_table = []
        result_e2e_table = []

        for rid in all_request_ids:
            # === E2E table (single column) ===
            e2e_evt = next((e for e in self.e2e_events if e.request_id == rid), None)
            if e2e_evt:
                e2e_data = _build_row(e2e_evt, E2E_FIELDS)
                result_e2e_table.append({"request_id": rid, **e2e_data})

                # filter out all-zero fields for logging
                nonzero_e2e_fields = set()
                for k, v in e2e_data.items():
                    if v not in (0, 0.000, None, ""):
                        nonzero_e2e_fields.add(k)
                value_fields_e2e = sorted(nonzero_e2e_fields)

                if value_fields_e2e:
                    logger.info(
                        "\n%s",
                        _format_table(
                            f"RequestE2EStats [request_id={rid}]",
                            e2e_data,
                            value_fields=value_fields_e2e,
                        ),
                    )

            # === [OmniTiming] concise per-request summary ===
            stage_evts = sorted(
                self.stage_events.get(rid, []),
                key=lambda e: e.stage_id if e.stage_id is not None else -1,
            )
            pt = {}
            if stage_evts:
                pt = stage_evts[-1].pipeline_timings or {}
            if pt or e2e_evt:
                parts = [f"req={rid}"]
                if e2e_evt:
                    parts.append(f"total={e2e_evt.e2e_total_ms / 1000.0:.2f}s")
                if "preprocess_ms" in pt:
                    parts.append(f"preprocess={pt['preprocess_ms'] / 1000.0:.2f}s")
                if e2e_evt:
                    engine_ms = e2e_evt.e2e_total_ms - pt.get("preprocess_ms", 0.0)
                    parts.append(f"engine={engine_ms / 1000.0:.2f}s")
                stage_parts = []
                for evt in stage_evts:
                    sid = evt.stage_id if evt.stage_id is not None else "?"
                    t = evt.stage_gen_time_ms / 1000.0
                    stage_parts.append(f"{sid}:{t:.2f}s")
                if stage_parts:
                    parts.append(f"stages=[{','.join(stage_parts)}]")
                transfer_parts = []
                for te in self.transfer_events.values():
                    if te.request_id == rid:
                        transfer_parts.append(f"{te.from_stage}->{te.to_stage}={te.tx_time_ms:.2f}ms")
                if transfer_parts:
                    parts.append(f"transfers=[{','.join(transfer_parts)}]")
                if "ar2diffusion_ms" in pt:
                    parts.append(f"ar2diffusion={pt['ar2diffusion_ms']:.2f}ms")
                logger.info("[OmniTiming] %s", " ".join(parts))

            # === Stage table (columns = stage_id) ===
            # if any stage has diffusion_metrics, remove postprocess_time_ms field
            # because it is already included in diffusion_metrics
            local_exclude = STAGE_EXCLUDE.copy()
            has_diffusion_metrics = any(getattr(evt, "diffusion_metrics", None) for evt in stage_evts)
            if has_diffusion_metrics:
                local_exclude.add("postprocess_time_ms")
            local_stage_fields = _build_field_defs(StageRequestStats, local_exclude, FIELD_TRANSFORMS)

            # if diffusion_metrics is present, expand it into multiple columns
            # then remove diffusion_metrics from the table
            stage_rows = []
            for evt in stage_evts:
                row = {"stage_id": evt.stage_id, **_build_row(evt, local_stage_fields)}
                if evt.diffusion_metrics:
                    row.update(evt.diffusion_metrics)
                row.pop("diffusion_metrics", None)  # Remove the dict itself
                stage_rows.append(row)

            result_stage_table.append({"request_id": rid, "stages": stage_rows})

            if stage_rows:
                # filter out all-zero fields for logging
                all_value_fields = set()
                for row in stage_rows:
                    for k in row.keys():
                        if k != "stage_id":
                            all_value_fields.add(k)
                value_fields_list = []
                for field in sorted(all_value_fields):
                    all_zero = True
                    for row in stage_rows:
                        v = row.get(field, None)
                        if v not in (0, 0.0, 0.000, None, ""):
                            all_zero = False
                            break
                    if not all_zero:
                        value_fields_list.append(field)

                if value_fields_list:
                    logger.info(
                        "\n%s",
                        _format_table(
                            f"StageRequestStats [request_id={rid}]",
                            stage_rows,
                            column_key="stage_id",
                            value_fields=value_fields_list,
                        ),
                    )

            # === Transfer table (columns = edge) ===
            transfer_evts = sorted(
                [e for e in self.transfer_events.values() if e.request_id == rid],
                key=lambda e: (e.from_stage, e.to_stage),
            )
            transfer_rows = [
                {"edge": f"{evt.from_stage}->{evt.to_stage}", **_build_row(evt, TRANSFER_FIELDS)}
                for evt in transfer_evts
            ]
            result_trans_table.append({"request_id": rid, "transfers": transfer_rows})

            if transfer_rows:
                # filter out all-zero fields for logging
                all_value_fields = set()
                for row in transfer_rows:
                    for k in row.keys():
                        if k != "edge":
                            all_value_fields.add(k)
                value_fields_list = []
                for field in sorted(all_value_fields):
                    all_zero = True
                    for row in transfer_rows:
                        v = row.get(field, None)
                        if v not in (0, 0.0, 0.000, None, ""):
                            all_zero = False
                            break
                    if not all_zero:
                        value_fields_list.append(field)

                if value_fields_list:
                    logger.info(
                        "\n%s",
                        _format_table(
                            f"TransferEdgeStats [request_id={rid}]",
                            transfer_rows,
                            column_key="edge",
                            value_fields=value_fields_list,
                        ),
                    )

        return {
            "final_stage_id": final_stage_id_map,
            "overall_summary": overall_summary,
            "stage_table": result_stage_table,
            "trans_table": result_trans_table,
            "e2e_table": result_e2e_table,
        }