512 lines
18 KiB
Python
512 lines
18 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import statistics
|
|
from collections import Counter, defaultdict
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
|
|
from agentic_pd_hybrid.trace import TraceRequest, load_trace
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ProfileConfig:
|
|
trace_path: Path
|
|
output_path: Path | None = None
|
|
metrics_path: Path | None = None
|
|
baseline_metrics_path: Path | None = None
|
|
candidate_metrics_path: Path | None = None
|
|
direct_max_uncached_tokens: int = 2048
|
|
|
|
|
|
def write_profile(config: ProfileConfig) -> dict[str, Any]:
|
|
report = build_profile(config)
|
|
if config.output_path is not None:
|
|
config.output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
with config.output_path.open("w", encoding="utf-8") as handle:
|
|
json.dump(report, handle, indent=2, sort_keys=True)
|
|
return report
|
|
|
|
|
|
def build_profile(config: ProfileConfig) -> dict[str, Any]:
|
|
requests = load_trace(config.trace_path)
|
|
features = _build_trace_features(
|
|
requests,
|
|
direct_max_uncached_tokens=config.direct_max_uncached_tokens,
|
|
)
|
|
report: dict[str, Any] = {
|
|
"trace_path": str(config.trace_path),
|
|
"direct_max_uncached_tokens": config.direct_max_uncached_tokens,
|
|
"trace_profile": _trace_profile(requests, features),
|
|
}
|
|
|
|
if config.metrics_path is not None:
|
|
metrics = _load_jsonl(config.metrics_path)
|
|
report["metrics_path"] = str(config.metrics_path)
|
|
report["metrics_profile"] = _metrics_profile(metrics, features)
|
|
|
|
if (
|
|
config.baseline_metrics_path is not None
|
|
and config.candidate_metrics_path is not None
|
|
):
|
|
baseline = _load_jsonl(config.baseline_metrics_path)
|
|
candidate = _load_jsonl(config.candidate_metrics_path)
|
|
report["baseline_metrics_path"] = str(config.baseline_metrics_path)
|
|
report["candidate_metrics_path"] = str(config.candidate_metrics_path)
|
|
report["baseline_profile"] = _metrics_profile(baseline, features)
|
|
report["candidate_profile"] = _metrics_profile(candidate, features)
|
|
report["paired_comparison"] = _paired_comparison(
|
|
baseline=baseline,
|
|
candidate=candidate,
|
|
features=features,
|
|
)
|
|
|
|
return report
|
|
|
|
|
|
def print_profile_summary(report: dict[str, Any]) -> None:
|
|
trace = report["trace_profile"]
|
|
print(
|
|
"trace: "
|
|
f"{trace['request_count']} requests, "
|
|
f"{trace['session_count']} sessions, "
|
|
f"{trace['multi_turn_session_count']} multi-turn sessions"
|
|
)
|
|
print(
|
|
"direct-eligible turns: "
|
|
f"{trace['direct_eligible_turn2plus_count']}/"
|
|
f"{trace['turn2plus_count']} "
|
|
f"({trace['direct_eligible_turn2plus_ratio']:.3f})"
|
|
)
|
|
append_stats = trace.get("append_input_tokens_stats")
|
|
output_stats = trace.get("output_tokens_stats")
|
|
if append_stats is not None:
|
|
print(
|
|
"append tokens: "
|
|
f"mean={append_stats['mean']:.1f} "
|
|
f"p50={append_stats['p50']:.1f} "
|
|
f"p90={append_stats['p90']:.1f} "
|
|
f"p99={append_stats['p99']:.1f}"
|
|
)
|
|
if output_stats is not None:
|
|
print(
|
|
"output tokens: "
|
|
f"mean={output_stats['mean']:.1f} "
|
|
f"p50={output_stats['p50']:.1f} "
|
|
f"p90={output_stats['p90']:.1f} "
|
|
f"p99={output_stats['p99']:.1f}"
|
|
)
|
|
|
|
comparison = report.get("paired_comparison")
|
|
if isinstance(comparison, dict):
|
|
overall = comparison.get("overall", {})
|
|
delta = overall.get("latency_delta_s_stats")
|
|
if delta is not None:
|
|
print(
|
|
"candidate - baseline E2E: "
|
|
f"mean={delta['mean']:.3f}s "
|
|
f"p50={delta['p50']:.3f}s "
|
|
f"p90={delta['p90']:.3f}s"
|
|
)
|
|
print(
|
|
"paired wins/losses: "
|
|
f"{overall.get('candidate_faster_count', 0)} faster, "
|
|
f"{overall.get('candidate_slower_count', 0)} slower, "
|
|
f"{overall.get('paired_count', 0)} paired"
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _TraceFeature:
|
|
request_id: str
|
|
session_id: str
|
|
turn_id: int
|
|
input_length: int
|
|
output_length: int
|
|
resident_tokens: int
|
|
append_input_tokens: int | None
|
|
inter_turn_gap_s: float | None
|
|
overlap_blocks_with_previous: int | None
|
|
overlap_ratio_with_previous: float | None
|
|
direct_eligible: bool
|
|
turn_class: str
|
|
append_bin: str
|
|
input_bin: str
|
|
output_bin: str
|
|
resident_bin: str
|
|
|
|
|
|
def _build_trace_features(
|
|
requests: list[TraceRequest],
|
|
*,
|
|
direct_max_uncached_tokens: int,
|
|
) -> dict[str, _TraceFeature]:
|
|
ordered_by_session: dict[str, list[TraceRequest]] = defaultdict(list)
|
|
for request in requests:
|
|
ordered_by_session[request.session_id].append(request)
|
|
|
|
previous_by_request_id: dict[str, TraceRequest | None] = {}
|
|
for session_requests in ordered_by_session.values():
|
|
ordered = sorted(
|
|
session_requests,
|
|
key=lambda request: (request.timestamp_s, request.turn_id, request.chat_id),
|
|
)
|
|
previous: TraceRequest | None = None
|
|
for request in ordered:
|
|
previous_by_request_id[request.request_id] = previous
|
|
previous = request
|
|
|
|
features: dict[str, _TraceFeature] = {}
|
|
for request in requests:
|
|
previous = previous_by_request_id.get(request.request_id)
|
|
append_input_tokens: int | None = None
|
|
inter_turn_gap_s: float | None = None
|
|
overlap_blocks: int | None = None
|
|
overlap_ratio: float | None = None
|
|
direct_eligible = False
|
|
|
|
if previous is not None:
|
|
append_input_tokens = request.input_length - (
|
|
previous.input_length + previous.output_length
|
|
)
|
|
inter_turn_gap_s = request.timestamp_s - previous.timestamp_s
|
|
previous_blocks = set(previous.hash_ids)
|
|
overlap_blocks = sum(1 for block in request.hash_ids if block in previous_blocks)
|
|
overlap_ratio = (
|
|
overlap_blocks / len(request.hash_ids) if request.hash_ids else 0.0
|
|
)
|
|
direct_eligible = (
|
|
append_input_tokens > 0
|
|
and append_input_tokens <= direct_max_uncached_tokens
|
|
and overlap_blocks > 0
|
|
)
|
|
|
|
features[request.request_id] = _TraceFeature(
|
|
request_id=request.request_id,
|
|
session_id=request.session_id,
|
|
turn_id=request.turn_id,
|
|
input_length=request.input_length,
|
|
output_length=request.output_length,
|
|
resident_tokens=request.input_length + request.output_length,
|
|
append_input_tokens=append_input_tokens,
|
|
inter_turn_gap_s=inter_turn_gap_s,
|
|
overlap_blocks_with_previous=overlap_blocks,
|
|
overlap_ratio_with_previous=overlap_ratio,
|
|
direct_eligible=direct_eligible,
|
|
turn_class="turn1" if request.turn_id <= 1 else "turn2plus",
|
|
append_bin=_token_bin(append_input_tokens),
|
|
input_bin=_token_bin(request.input_length),
|
|
output_bin=_token_bin(request.output_length),
|
|
resident_bin=_token_bin(request.input_length + request.output_length),
|
|
)
|
|
return features
|
|
|
|
|
|
def _trace_profile(
|
|
requests: list[TraceRequest],
|
|
features: dict[str, _TraceFeature],
|
|
) -> dict[str, Any]:
|
|
session_turns = Counter(request.session_id for request in requests)
|
|
turn2plus = [feature for feature in features.values() if feature.turn_id > 1]
|
|
direct_eligible = [feature for feature in turn2plus if feature.direct_eligible]
|
|
append_values = [
|
|
feature.append_input_tokens
|
|
for feature in turn2plus
|
|
if feature.append_input_tokens is not None
|
|
]
|
|
positive_append_values = [
|
|
value for value in append_values if value is not None and value > 0
|
|
]
|
|
overlap_ratios = [
|
|
feature.overlap_ratio_with_previous
|
|
for feature in turn2plus
|
|
if feature.overlap_ratio_with_previous is not None
|
|
]
|
|
gaps = [
|
|
feature.inter_turn_gap_s
|
|
for feature in turn2plus
|
|
if feature.inter_turn_gap_s is not None
|
|
]
|
|
return {
|
|
"request_count": len(requests),
|
|
"session_count": len(session_turns),
|
|
"multi_turn_session_count": sum(1 for turns in session_turns.values() if turns > 1),
|
|
"turn2plus_count": len(turn2plus),
|
|
"direct_eligible_turn2plus_count": len(direct_eligible),
|
|
"direct_eligible_turn2plus_ratio": (
|
|
len(direct_eligible) / len(turn2plus) if turn2plus else 0.0
|
|
),
|
|
"turn_count_distribution": dict(sorted(Counter(session_turns.values()).items())),
|
|
"request_type_distribution": dict(
|
|
sorted(Counter(request.request_type for request in requests).items())
|
|
),
|
|
"turn_id_distribution": dict(
|
|
sorted(Counter(request.turn_id for request in requests).items())
|
|
),
|
|
"append_bin_distribution": dict(
|
|
sorted(Counter(feature.append_bin for feature in turn2plus).items())
|
|
),
|
|
"input_bin_distribution": dict(
|
|
sorted(Counter(feature.input_bin for feature in features.values()).items())
|
|
),
|
|
"output_bin_distribution": dict(
|
|
sorted(Counter(feature.output_bin for feature in features.values()).items())
|
|
),
|
|
"resident_bin_distribution": dict(
|
|
sorted(Counter(feature.resident_bin for feature in features.values()).items())
|
|
),
|
|
"input_tokens_stats": _stats(
|
|
[float(request.input_length) for request in requests]
|
|
),
|
|
"output_tokens_stats": _stats(
|
|
[float(request.output_length) for request in requests]
|
|
),
|
|
"resident_tokens_stats": _stats(
|
|
[float(feature.resident_tokens) for feature in features.values()]
|
|
),
|
|
"append_input_tokens_stats": _stats(
|
|
[float(value) for value in append_values if value is not None]
|
|
),
|
|
"positive_append_input_tokens_stats": _stats(
|
|
[float(value) for value in positive_append_values]
|
|
),
|
|
"inter_turn_gap_s_stats": _stats([float(value) for value in gaps]),
|
|
"overlap_ratio_stats": _stats([float(value) for value in overlap_ratios]),
|
|
"non_positive_append_count": sum(
|
|
1 for value in append_values if value is not None and value <= 0
|
|
),
|
|
}
|
|
|
|
|
|
def _metrics_profile(
|
|
rows: list[dict[str, Any]],
|
|
features: dict[str, _TraceFeature],
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"request_count": len(rows),
|
|
"mechanism_distribution": dict(
|
|
sorted(Counter(str(row.get("mechanism_name")) for row in rows).items())
|
|
),
|
|
"execution_mode_distribution": dict(
|
|
sorted(Counter(str(row.get("execution_mode")) for row in rows).items())
|
|
),
|
|
"latency_s_stats": _stats(_numeric_values(rows, "latency_s")),
|
|
"ttft_s_stats": _stats(_numeric_values(rows, "ttft_s")),
|
|
"tpot_s_stats": _stats(_numeric_values(rows, "tpot_s")),
|
|
"cached_tokens_stats": _stats(_numeric_values(rows, "cached_tokens")),
|
|
"actual_kv_transfer_blocks_stats": _stats(
|
|
_numeric_values(rows, "actual_kv_transfer_blocks")
|
|
),
|
|
"session_reused_count": sum(1 for row in rows if row.get("session_reused")),
|
|
"session_reset_count": sum(1 for row in rows if row.get("session_reset")),
|
|
"error_count": sum(1 for row in rows if row.get("error") is not None),
|
|
"by_turn_class": _group_metrics(rows, features, lambda feature, _row: feature.turn_class),
|
|
"by_direct_eligible": _group_metrics(
|
|
rows,
|
|
features,
|
|
lambda feature, _row: "eligible" if feature.direct_eligible else "not_eligible",
|
|
),
|
|
"by_append_bin": _group_metrics(rows, features, lambda feature, _row: feature.append_bin),
|
|
"by_resident_bin": _group_metrics(
|
|
rows,
|
|
features,
|
|
lambda feature, _row: feature.resident_bin,
|
|
),
|
|
"by_execution_mode": _group_metrics(
|
|
rows,
|
|
features,
|
|
lambda _feature, row: str(row.get("execution_mode")),
|
|
),
|
|
}
|
|
|
|
|
|
def _paired_comparison(
|
|
*,
|
|
baseline: list[dict[str, Any]],
|
|
candidate: list[dict[str, Any]],
|
|
features: dict[str, _TraceFeature],
|
|
) -> dict[str, Any]:
|
|
baseline_by_id = {
|
|
str(row.get("request_id")): row
|
|
for row in baseline
|
|
if row.get("latency_s") is not None
|
|
}
|
|
candidate_by_id = {
|
|
str(row.get("request_id")): row
|
|
for row in candidate
|
|
if row.get("latency_s") is not None
|
|
}
|
|
paired_ids = sorted(set(baseline_by_id) & set(candidate_by_id))
|
|
pairs = [
|
|
(baseline_by_id[request_id], candidate_by_id[request_id], features.get(request_id))
|
|
for request_id in paired_ids
|
|
]
|
|
pairs = [pair for pair in pairs if pair[2] is not None]
|
|
|
|
return {
|
|
"overall": _delta_summary(pairs),
|
|
"by_turn_class": _group_deltas(
|
|
pairs,
|
|
lambda feature, _base, _cand: feature.turn_class,
|
|
),
|
|
"by_direct_eligible": _group_deltas(
|
|
pairs,
|
|
lambda feature, _base, _cand: (
|
|
"eligible" if feature.direct_eligible else "not_eligible"
|
|
),
|
|
),
|
|
"by_append_bin": _group_deltas(
|
|
pairs,
|
|
lambda feature, _base, _cand: feature.append_bin,
|
|
),
|
|
"by_resident_bin": _group_deltas(
|
|
pairs,
|
|
lambda feature, _base, _cand: feature.resident_bin,
|
|
),
|
|
"by_candidate_execution_mode": _group_deltas(
|
|
pairs,
|
|
lambda _feature, _base, cand: str(cand.get("execution_mode")),
|
|
),
|
|
}
|
|
|
|
|
|
def _group_metrics(
|
|
rows: list[dict[str, Any]],
|
|
features: dict[str, _TraceFeature],
|
|
key_fn,
|
|
) -> dict[str, Any]:
|
|
grouped: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
|
for row in rows:
|
|
feature = features.get(str(row.get("request_id")))
|
|
if feature is None:
|
|
continue
|
|
grouped[str(key_fn(feature, row))].append(row)
|
|
return {
|
|
key: {
|
|
"count": len(group_rows),
|
|
"latency_s_stats": _stats(_numeric_values(group_rows, "latency_s")),
|
|
"ttft_s_stats": _stats(_numeric_values(group_rows, "ttft_s")),
|
|
"tpot_s_stats": _stats(_numeric_values(group_rows, "tpot_s")),
|
|
"session_reused_count": sum(
|
|
1 for row in group_rows if row.get("session_reused")
|
|
),
|
|
"error_count": sum(1 for row in group_rows if row.get("error") is not None),
|
|
}
|
|
for key, group_rows in sorted(grouped.items())
|
|
}
|
|
|
|
|
|
def _group_deltas(
|
|
pairs: list[tuple[dict[str, Any], dict[str, Any], _TraceFeature | None]],
|
|
key_fn,
|
|
) -> dict[str, Any]:
|
|
grouped: dict[str, list[tuple[dict[str, Any], dict[str, Any], _TraceFeature | None]]] = (
|
|
defaultdict(list)
|
|
)
|
|
for base, cand, feature in pairs:
|
|
if feature is None:
|
|
continue
|
|
grouped[str(key_fn(feature, base, cand))].append((base, cand, feature))
|
|
return {key: _delta_summary(group_pairs) for key, group_pairs in sorted(grouped.items())}
|
|
|
|
|
|
def _delta_summary(
|
|
pairs: list[tuple[dict[str, Any], dict[str, Any], _TraceFeature | None]],
|
|
) -> dict[str, Any]:
|
|
latency_deltas = [
|
|
float(cand["latency_s"]) - float(base["latency_s"])
|
|
for base, cand, _feature in pairs
|
|
if base.get("latency_s") is not None and cand.get("latency_s") is not None
|
|
]
|
|
ttft_deltas = [
|
|
float(cand["ttft_s"]) - float(base["ttft_s"])
|
|
for base, cand, _feature in pairs
|
|
if base.get("ttft_s") is not None and cand.get("ttft_s") is not None
|
|
]
|
|
return {
|
|
"paired_count": len(latency_deltas),
|
|
"candidate_faster_count": sum(1 for delta in latency_deltas if delta < 0),
|
|
"candidate_slower_count": sum(1 for delta in latency_deltas if delta > 0),
|
|
"latency_delta_s_stats": _stats(latency_deltas),
|
|
"ttft_delta_s_stats": _stats(ttft_deltas),
|
|
"total_latency_delta_s": sum(latency_deltas),
|
|
"mean_baseline_latency_s": _mean(
|
|
[
|
|
float(base["latency_s"])
|
|
for base, cand, _feature in pairs
|
|
if base.get("latency_s") is not None and cand.get("latency_s") is not None
|
|
]
|
|
),
|
|
"mean_candidate_latency_s": _mean(
|
|
[
|
|
float(cand["latency_s"])
|
|
for base, cand, _feature in pairs
|
|
if base.get("latency_s") is not None and cand.get("latency_s") is not None
|
|
]
|
|
),
|
|
}
|
|
|
|
|
|
def _load_jsonl(path: Path) -> list[dict[str, Any]]:
|
|
rows: list[dict[str, Any]] = []
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
for line in handle:
|
|
if line.strip():
|
|
rows.append(json.loads(line))
|
|
return rows
|
|
|
|
|
|
def _numeric_values(rows: Iterable[dict[str, Any]], key: str) -> list[float]:
|
|
values: list[float] = []
|
|
for row in rows:
|
|
value = row.get(key)
|
|
if value is not None:
|
|
values.append(float(value))
|
|
return values
|
|
|
|
|
|
def _stats(values: list[float]) -> 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),
|
|
"min": clean[0],
|
|
"max": clean[-1],
|
|
}
|
|
|
|
|
|
def _percentile(sorted_values: list[float], percentile: float) -> float:
|
|
if len(sorted_values) == 1:
|
|
return sorted_values[0]
|
|
index = round((len(sorted_values) - 1) * percentile)
|
|
return sorted_values[index]
|
|
|
|
|
|
def _mean(values: list[float]) -> float | None:
|
|
if not values:
|
|
return None
|
|
return statistics.fmean(values)
|
|
|
|
|
|
def _token_bin(value: int | None) -> str:
|
|
if value is None:
|
|
return "none"
|
|
if value <= 0:
|
|
return "<=0"
|
|
if value <= 512:
|
|
return "1-512"
|
|
if value <= 2048:
|
|
return "513-2048"
|
|
if value <= 8192:
|
|
return "2049-8192"
|
|
if value <= 32768:
|
|
return "8193-32768"
|
|
return ">32768"
|