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 _is_failed_request(row: RequestMetrics) -> bool: if row.error is not None: return True if row.finish_reason is not None: fr = str(row.finish_reason).lower() if "abort" in fr or "badrequest" in fr: return True return False def write_summary_json( path: Path, rows: list[RequestMetrics], *, trace_path: Path, router_url: str | None, ) -> None: successful = [row for row in rows if not _is_failed_request(row)] latencies = [row.latency_s for row in successful if row.latency_s is not None] ttfts = [row.ttft_s for row in successful if row.ttft_s is not None] tpots = [row.tpot_s for row in successful 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), "abort_count": sum( 1 for row in rows if row.error is None and row.finish_reason is not None and ( "abort" in str(row.finish_reason).lower() or "badrequest" in str(row.finish_reason).lower() ) ), "failure_count": sum(1 for row in rows if _is_failed_request(row)), "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]