feat: add agentic pd hybrid benchmark prototype

src/agentic_pd_hybrid/metrics.py

@@ -0,0 +1,165 @@
+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
+    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
+
+    @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,
+        error: 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,
+            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,
+        )
+
+
+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)
+
+    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())),
+        "error_count": sum(1 for row in rows if row.error 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]