Spaces:
Sleeping
Sleeping
| """Aggregation of one benchmark run's `operational.jsonl` + `quality.jsonl` into `summary.json`. | |
| Scope note: this benchmark drives only the retrieval critical path | |
| (`_get_chunks()` -> `rag.retrieval` -> its embedding/chroma/result-prepare | |
| children), not a full backend request spanning memory/context/LLM/citation | |
| subsystems. So "per-request subsystem shares" here are computed relative to | |
| `retrieval_ms` (this workload's own measured parent operation), not | |
| `backend_operation_ms` -- that full-request denominator only applies to a | |
| workload that drives a whole chat/lesson turn, which this one deliberately | |
| does not (see plan preamble: "does not introduce independent application | |
| timers" / drives only the retrieval call under test). | |
| Per the workload protocol, cold and warm rows are never merged into one | |
| figure -- every stratum lives under its own `cache_state` key. Within each | |
| cache_state, rows are additionally stratified by `pipeline_version` and | |
| `query_category`, per the task-4 brief's Step 4. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import math | |
| from collections import defaultdict | |
| from pathlib import Path | |
| from typing import Any, Iterable | |
| # Retrieval stages this workload actually measures, in critical-path order. | |
| _RETRIEVAL_STAGES: tuple[str, ...] = ( | |
| "embedding.query", | |
| "chroma.collection_lookup", | |
| "chroma.collection_count", | |
| "chroma.vector_search", | |
| "rag.result_prepare", | |
| ) | |
| def _percentiles(values: list[float]) -> dict[str, float | int | None]: | |
| if not values: | |
| return {"p50": None, "p90": None, "p95": None, "max": None, "count": 0} | |
| ordered = sorted(values) | |
| n = len(ordered) | |
| def _pct(p: float) -> float: | |
| # Nearest-rank percentile over the sorted sample -- simple, deterministic, | |
| # and stable on the small sample sizes this benchmark produces (10-ish | |
| # warm repetitions per query), unlike interpolation methods that can | |
| # imply more precision than a 10-sample distribution actually supports. | |
| rank = max(0, min(n - 1, math.ceil(p * n) - 1)) | |
| return ordered[rank] | |
| return { | |
| "p50": _pct(0.50), | |
| "p90": _pct(0.90), | |
| "p95": _pct(0.95), | |
| "max": ordered[-1], | |
| "count": n, | |
| } | |
| def _mean(values: list[float]) -> float | None: | |
| return sum(values) / len(values) if values else None | |
| def load_jsonl(path: Path) -> list[dict[str, Any]]: | |
| path = Path(path) | |
| if not path.exists(): | |
| return [] | |
| 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 _experiment(row: dict[str, Any]) -> dict[str, Any]: | |
| return row.get("experiment") or {} | |
| def cache_state_of(row: dict[str, Any]) -> str: | |
| return str(_experiment(row).get("cache_state") or "unknown") | |
| def pipeline_version_of(row: dict[str, Any]) -> str: | |
| return str(_experiment(row).get("pipeline_version") or "unknown") | |
| def query_category_of(row: dict[str, Any]) -> str: | |
| return str(_experiment(row).get("query_category") or "unknown") | |
| def _retrieval_ms_of(row: dict[str, Any]) -> float | None: | |
| return (row.get("stage_durations_ms") or {}).get("retrieval_ms") | |
| def _stage_ms_of(row: dict[str, Any], stage: str) -> float | None: | |
| return (row.get("stage_durations_ms") or {}).get(stage) | |
| def _retrieval_status_of(row: dict[str, Any]) -> str | None: | |
| return (row.get("retrieval") or {}).get("status") | |
| def _retrieval_outcome_counts(rows: Iterable[dict[str, Any]]) -> dict[str, Any]: | |
| counts = {"success": 0, "success_empty": 0, "error_fallback": 0, "unknown": 0} | |
| for row in rows: | |
| status = _retrieval_status_of(row) or "unknown" | |
| counts[status] = counts.get(status, 0) + 1 | |
| total = sum(counts.values()) | |
| rates = {f"{k}_rate": (v / total if total else None) for k, v in counts.items()} | |
| return {"counts": counts, "total": total, **rates} | |
| def _operational_summary(rows: list[dict[str, Any]]) -> dict[str, Any]: | |
| retrieval_ms = [ms for r in rows if (ms := _retrieval_ms_of(r)) is not None] | |
| stage_summaries: dict[str, Any] = {} | |
| for stage in _RETRIEVAL_STAGES: | |
| values = [ms for r in rows if (ms := _stage_ms_of(r, stage)) is not None] | |
| stage_summaries[stage] = _percentiles(values) | |
| return {"retrieval_ms": _percentiles(retrieval_ms), "stage_ms": stage_summaries} | |
| def _retrieval_stage_shares(rows: list[dict[str, Any]]) -> dict[str, Any]: | |
| """Per-request stage_ms / retrieval_ms, then aggregated -- never | |
| p95(stage) / p95(retrieval_ms), per the plan preamble's explicit warning | |
| against dividing already-aggregated percentiles. | |
| Named ``retrieval_stage_shares`` (not ``subsystem_shares``) because this | |
| workload only ever drives the retrieval critical path (see module | |
| docstring): every share here is a *retrieval stage's* fraction of | |
| `retrieval_ms`, not a share of the plan's canonical cross-subsystem | |
| denominator (memory/retrieval/context/llm/citation over | |
| `backend_operation_ms`). The two are easy to confuse by name alone, so | |
| this field is named for exactly what it scopes over. | |
| """ | |
| shares: dict[str, list[float]] = defaultdict(list) | |
| for row in rows: | |
| retrieval_ms = _retrieval_ms_of(row) | |
| if not retrieval_ms: | |
| continue | |
| for stage in _RETRIEVAL_STAGES: | |
| value = _stage_ms_of(row, stage) | |
| if value is None: | |
| continue | |
| shares[stage].append(value / retrieval_ms) | |
| return {stage: {"mean": _mean(vals), **_percentiles(vals)} for stage, vals in shares.items()} | |
| def _quality_summary(rows: list[dict[str, Any]]) -> dict[str, Any]: | |
| n = len(rows) | |
| if n == 0: | |
| return { | |
| "count": 0, | |
| "document_hit_rate": None, | |
| "section_hit_rate": None, | |
| "section_hit_applicable_count": 0, | |
| "mrr": None, | |
| "citation_validity_rate": None, | |
| "citation_validity_applicable_count": 0, | |
| "empty_result_rate": None, | |
| "retrieval_error_rate": None, | |
| } | |
| section_applicable = [r for r in rows if r.get("section_hit") is not None] | |
| citation_applicable = [r for r in rows if r.get("citation_validity") is not None] | |
| return { | |
| "count": n, | |
| "document_hit_rate": sum(1 for r in rows if r.get("document_hit")) / n, | |
| "section_hit_rate": ( | |
| sum(1 for r in section_applicable if r.get("section_hit")) / len(section_applicable) | |
| if section_applicable | |
| else None | |
| ), | |
| "section_hit_applicable_count": len(section_applicable), | |
| "mrr": _mean([r.get("reciprocal_rank", 0.0) for r in rows]), | |
| "citation_validity_rate": ( | |
| sum(1 for r in citation_applicable if r.get("citation_validity")) / len(citation_applicable) | |
| if citation_applicable | |
| else None | |
| ), | |
| "citation_validity_applicable_count": len(citation_applicable), | |
| "empty_result_rate": sum(1 for r in rows if r.get("empty_result")) / n, | |
| "retrieval_error_rate": sum(1 for r in rows if r.get("retrieval_error")) / n, | |
| } | |
| def _stratum_summary(operational_rows: list[dict[str, Any]], quality_rows: list[dict[str, Any]]) -> dict[str, Any]: | |
| return { | |
| "request_count": len(operational_rows), | |
| "operational": _operational_summary(operational_rows), | |
| "retrieval_outcomes": _retrieval_outcome_counts(operational_rows), | |
| "retrieval_stage_shares": _retrieval_stage_shares(operational_rows), | |
| "quality": _quality_summary(quality_rows), | |
| } | |
| def summarize_run(operational_rows: list[dict[str, Any]], quality_rows: list[dict[str, Any]]) -> dict[str, Any]: | |
| """Build the nested `cache_state -> pipeline_version -> {overall, by_category}` summary. | |
| Cold and warm strata are always separate top-level keys -- never merged -- | |
| per the workload protocol. | |
| """ | |
| quality_by_trace = {r["trace_id"]: r for r in quality_rows if r.get("trace_id")} | |
| by_cache: dict[str, list[dict[str, Any]]] = defaultdict(list) | |
| for row in operational_rows: | |
| by_cache[cache_state_of(row)].append(row) | |
| summary: dict[str, Any] = {"cache_states": {}} | |
| for cache_state, op_rows in by_cache.items(): | |
| by_pipeline: dict[str, list[dict[str, Any]]] = defaultdict(list) | |
| for row in op_rows: | |
| by_pipeline[pipeline_version_of(row)].append(row) | |
| pipeline_summary: dict[str, Any] = {} | |
| for pipeline_version, p_rows in by_pipeline.items(): | |
| p_quality = [quality_by_trace[r["trace_id"]] for r in p_rows if r.get("trace_id") in quality_by_trace] | |
| overall = _stratum_summary(p_rows, p_quality) | |
| by_category: dict[str, Any] = {} | |
| categories = {query_category_of(r) for r in p_rows} | |
| for category in sorted(categories): | |
| cat_op_rows = [r for r in p_rows if query_category_of(r) == category] | |
| cat_quality = [ | |
| quality_by_trace[r["trace_id"]] for r in cat_op_rows if r.get("trace_id") in quality_by_trace | |
| ] | |
| by_category[category] = _stratum_summary(cat_op_rows, cat_quality) | |
| pipeline_summary[pipeline_version] = {"overall": overall, "by_category": by_category} | |
| summary["cache_states"][cache_state] = pipeline_summary | |
| return summary | |
| def summarize_run_dir(output_dir: Path) -> dict[str, Any]: | |
| output_dir = Path(output_dir) | |
| operational_rows = load_jsonl(output_dir / "operational.jsonl") | |
| quality_rows = load_jsonl(output_dir / "quality.jsonl") | |
| summary = summarize_run(operational_rows, quality_rows) | |
| manifest_path = output_dir / "manifest.json" | |
| if manifest_path.exists(): | |
| manifest = json.loads(manifest_path.read_text(encoding="utf-8")) | |
| summary["run_id"] = manifest.get("run_id") | |
| summary["experiment_id"] = manifest.get("experiment_id") | |
| summary["smoke"] = manifest.get("smoke") | |
| # Carried over so a reader of summary.json alone -- not manifest.json, | |
| # not the source -- still sees the cold-cache-state caveat. | |
| summary["cache_state_semantics"] = manifest.get("cache_state_semantics") | |
| return summary | |
| def write_summary(summary: dict[str, Any], path: Path) -> None: | |
| Path(path).write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") | |