agentic-pd-hybrid/src/agentic_pd_hybrid/metrics.py

from __future__ import annotations

import json
import statistics
from collections import Counter
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from agentic_pd_hybrid.policies import RoutingDecision
from agentic_pd_hybrid.trace import TraceRequest


@dataclass(frozen=True)
class RequestMetrics:
    request_id: str
    session_id: str
    turn_id: int
    mechanism_name: str
    execution_mode: str
    trace_timestamp_s: float
    input_length: int
    output_length: int
    request_type: str
    policy_name: str
    assigned_prefill_node: str
    assigned_decode_node: str
    assigned_decode_index: int
    inflight_decode_load_at_assignment: int
    reuse_expected: bool
    reuse_observed: bool
    observed_overlap_blocks: int
    kv_transfer_blocks: int
    actual_kv_transfer_blocks: int
    cached_tokens: int
    prefill_request_priority: int | None
    decode_request_priority: int | None
    re_prefill_required: bool
    effective_input_length: int | None
    session_reused: bool
    session_reset: bool
    latency_s: float | None
    ttft_s: float | None
    tpot_s: float | None
    error: str | None = None
    actual_output_tokens: int | None = None
    requested_output_tokens: int | None = None
    finish_reason: str | None = None

    @classmethod
    def from_decision(
        cls,
        request: TraceRequest,
        decision: RoutingDecision,
        *,
        mechanism_name: str,
        execution_mode: str,
        actual_kv_transfer_blocks: int,
        effective_input_length: int | None,
        cached_tokens: int,
        session_reused: bool,
        session_reset: bool,
        latency_s: float | None,
        ttft_s: float | None,
        tpot_s: float | None,
        prefill_request_priority: int | None = None,
        decode_request_priority: int | None = None,
        error: str | None = None,
        actual_output_tokens: int | None = None,
        requested_output_tokens: int | None = None,
        finish_reason: str | None = None,
    ) -> "RequestMetrics":
        return cls(
            request_id=request.request_id,
            session_id=request.session_id,
            turn_id=request.turn_id,
            mechanism_name=mechanism_name,
            execution_mode=execution_mode,
            trace_timestamp_s=request.timestamp_s,
            input_length=request.input_length,
            output_length=request.output_length,
            request_type=request.request_type,
            policy_name=decision.policy_name,
            assigned_prefill_node=decision.prefill_worker_id,
            assigned_decode_node=decision.decode_worker_id,
            assigned_decode_index=decision.decode_worker_index,
            inflight_decode_load_at_assignment=decision.inflight_decode_load,
            reuse_expected=decision.reuse_expected,
            reuse_observed=decision.observed_reuse,
            observed_overlap_blocks=decision.observed_overlap_blocks,
            kv_transfer_blocks=decision.kv_transfer_blocks,
            actual_kv_transfer_blocks=actual_kv_transfer_blocks,
            cached_tokens=cached_tokens,
            prefill_request_priority=prefill_request_priority,
            decode_request_priority=decode_request_priority,
            re_prefill_required=decision.re_prefill_required,
            effective_input_length=effective_input_length,
            session_reused=session_reused,
            session_reset=session_reset,
            latency_s=latency_s,
            ttft_s=ttft_s,
            tpot_s=tpot_s,
            error=error,
            actual_output_tokens=actual_output_tokens,
            requested_output_tokens=requested_output_tokens,
            finish_reason=finish_reason,
        )


def write_metrics_jsonl(path: Path, rows: list[RequestMetrics]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(asdict(row), sort_keys=True) + "\n")


def write_summary_json(
    path: Path,
    rows: list[RequestMetrics],
    *,
    trace_path: Path,
    router_url: str | None,
) -> None:
    latencies = [row.latency_s for row in rows if row.latency_s is not None]
    ttfts = [row.ttft_s for row in rows if row.ttft_s is not None]
    tpots = [row.tpot_s for row in rows if row.tpot_s is not None]
    per_decode_load = Counter(row.assigned_decode_node for row in rows)
    per_prefill_load = Counter(row.assigned_prefill_node for row in rows)
    prefill_priorities = Counter(
        row.prefill_request_priority
        for row in rows
        if row.prefill_request_priority is not None
    )
    decode_priorities = Counter(
        row.decode_request_priority
        for row in rows
        if row.decode_request_priority is not None
    )

    summary: dict[str, Any] = {
        "trace_path": str(trace_path),
        "router_url": router_url,
        "request_count": len(rows),
        "mechanisms": dict(sorted(Counter(row.mechanism_name for row in rows).items())),
        "execution_modes": dict(sorted(Counter(row.execution_mode for row in rows).items())),
        "latency_stats_s": _stats(latencies),
        "ttft_stats_s": _stats(ttfts),
        "tpot_stats_s": _stats(tpots),
        "reuse_expected_count": sum(1 for row in rows if row.reuse_expected),
        "reuse_observed_count": sum(1 for row in rows if row.reuse_observed),
        "re_prefill_count": sum(1 for row in rows if row.re_prefill_required),
        "cache_hit_request_count": sum(1 for row in rows if row.cached_tokens > 0),
        "total_cached_tokens": sum(row.cached_tokens for row in rows),
        "cached_tokens_stats": _stats([float(row.cached_tokens) for row in rows]),
        "session_reused_count": sum(1 for row in rows if row.session_reused),
        "session_reset_count": sum(1 for row in rows if row.session_reset),
        "total_kv_transfer_blocks": sum(row.kv_transfer_blocks for row in rows),
        "total_actual_kv_transfer_blocks": sum(
            row.actual_kv_transfer_blocks for row in rows
        ),
        "per_decode_load": dict(sorted(per_decode_load.items())),
        "per_prefill_load": dict(sorted(per_prefill_load.items())),
        "prefill_request_priorities": {
            str(key): value for key, value in sorted(prefill_priorities.items())
        },
        "decode_request_priorities": {
            str(key): value for key, value in sorted(decode_priorities.items())
        },
        "error_count": sum(1 for row in rows if row.error is not None),
        "truncated_request_count": sum(
            1
            for row in rows
            if row.actual_output_tokens is not None
            and row.requested_output_tokens is not None
            and row.requested_output_tokens > 1
            and row.actual_output_tokens < row.requested_output_tokens * 0.5
        ),
        "actual_output_tokens_stats": _stats(
            [float(row.actual_output_tokens) for row in rows if row.actual_output_tokens is not None]
        ),
    }
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        json.dump(summary, handle, indent=2, sort_keys=True)


def _stats(values: list[float | None]) -> dict[str, float] | None:
    clean = [value for value in values if value is not None]
    if not clean:
        return None
    clean.sort()
    return {
        "count": float(len(clean)),
        "mean": statistics.fmean(clean),
        "p50": _percentile(clean, 0.50),
        "p90": _percentile(clean, 0.90),
        "p99": _percentile(clean, 0.99),
    }


def _percentile(sorted_values: list[float], percentile: float) -> float:
    if not sorted_values:
        raise ValueError("sorted_values must not be empty")
    if len(sorted_values) == 1:
        return sorted_values[0]
    index = round((len(sorted_values) - 1) * percentile)
    return sorted_values[index]