1"""Load raw per-attempt records into a tidy DataFrame with derived columns (scenario costs, method labels)."""2from __future__ import annotations3import glob, json4import pandas as pd5import numpy as np6from ..util import ROOT, read_jsonl78TEXT_METHODS = ["A0", "A1", "A2", "A3", "A4", "A5", "A6", "A7", "A8"]9CACHE_METHODS = ["B0", "B1", "B2", "B2b", "B3", "B4", "B5", "B6", "B9", "C0", "C2", "C3"]10LABELS = {"A0": "A0 no message", "A1": "A1 full text", "A2": "A2 head/tail", "A3": "A3 extractive", "A4": "A4 summary", "A5": "A5 structured evidence", "A6": "A6 refs+retrieval",11 "A7": "A7 zstd text", "A8": "A8 LLMLingua-2", "B0": "B0 native cache", "B1": "B1 serialized reuse", "B2": "B2 packed quant", "B3": "B3 sink+recent", "B4": "B4 random",12 "B5": "B5 attention (SnapKV-style)", "B6": "B6 pyramid", "B9": "B9 select+quant", "C0": "C0 identity (neg. control)", "C2": "C2 linear mapper", "C3": "C3 mapper+select"}1314def load_experiments(patterns: list[str], last_attempt_only: bool = True) -> pd.DataFrame:15 """Load records. By default keep the LAST attempt per trial (retries supersede); all attempts are kept when16 last_attempt_only=False (used for the failure/attempt tables)."""17 rows = []18 for pat in patterns:19 for p in glob.glob(str(ROOT / "runs/raw" / pat)):20 for r in read_jsonl(p):21 rows.append(flatten(r))22 df = pd.DataFrame(rows)23 if last_attempt_only and len(df):24 df = df.sort_values("timestamp_utc").drop_duplicates("trial_id", keep="last")25 return df2627def _scenario_costs(r: dict) -> dict:28 sp = r.get("timing_spans") or {}29 m = r["method_id"]30 enc = sp.get("encode_ms", 0.0) or 0.031 recv = (sp.get("receiver_prefill_ms", 0.0) or 0.0) + (sp.get("ttft_ms", 0.0) or 0.0) + (sp.get("decode_ms", 0.0) or 0.0)32 if m.startswith("B") or m.startswith("C"):33 doc_prefill = sp.get("doc_prefill_ms", 0.0) or 0.034 compress = (sp.get("compress_ms", 0.0) or 0.0) + (sp.get("serialize_ms", 0.0) or 0.0)35 load = sp.get("load_ms", 0.0) or 0.036 # cold: document must be read & cached, compressed, loaded, then query processed37 cold = doc_prefill + compress + load + recv38 # warm (cache already built & compressed once, reused across N queries): per-query = load + query39 warm_query = load + recv40 reuse_n = max(1, r.get("reuse_count") or 1)41 amortized = (doc_prefill + compress) / reuse_n + warm_query42 return {"cost_cold_ms": cold, "cost_warm_query_ms": warm_query, "cost_amortized_ms": amortized, "cost_build_ms": doc_prefill + compress, "receiver_ms": recv}43 else:44 reuse_n = max(1, r.get("reuse_count") or 1)45 cold = enc + recv46 warm_query = recv # message already built47 amortized = enc / reuse_n + recv48 return {"cost_cold_ms": cold, "cost_warm_query_ms": warm_query, "cost_amortized_ms": amortized, "cost_build_ms": enc, "receiver_ms": recv}4950def flatten(r: dict) -> dict:51 sc = r.get("task_scores") or {}52 cc = r.get("compression_config") or {}53 out = {k: r.get(k) for k in ["experiment_id", "trial_id", "attempt_id", "timestamp_utc", "run_group_id", "task_family", "document_id", "question_id", "split", "cluster_id", "method_id",54 "question_access", "topology", "handoff_index", "reuse_count", "status", "error_code", "serialized_payload_bytes", "message_utf8_bytes",55 "cache_tensor_bytes", "fetched_bytes", "total_wall_ms", "ttft_ms", "peak_prefill_memory", "peak_decode_memory", "prediction_artifact", "protocol_hash"]}56 out["sender"] = (r.get("sender_model_revision") or "").split("@")[0].split("/")[-1]57 out["receiver"] = (r.get("receiver_model_revision") or "").split("@")[0].split("/")[-1]58 out["pair"] = f"{out['sender']}->{out['receiver']}" if out["sender"] else out["receiver"]59 out["correct"] = sc.get("correct", np.nan); out["f1"] = sc.get("f1", np.nan); out["em"] = sc.get("em", np.nan); out["pass_rate"] = sc.get("pass_rate", np.nan); out["abstained"] = sc.get("abstained", np.nan)60 out["quality"] = sc.get("f1") if r["task_family"].startswith("public_") else sc.get("correct")61 if r["status"] != "success":62 out["quality_operational"] = 0.063 else:64 out["quality_operational"] = out["quality"]65 out["budget_tokens"] = cc.get("budget_tokens"); out["bits"] = cc.get("bits"); out["keep_frac"] = cc.get("keep_frac"); out["codec"] = cc.get("codec")66 out["condition"] = r["run_group_id"].split(":", 1)[1] if ":" in (r.get("run_group_id") or "") else r["method_id"]67 out["label"] = LABELS.get(r["method_id"], r["method_id"])68 out["qtype"] = r["question_id"].split("_", 1)[1] if "_" in r["question_id"] else r["task_family"]69 ev = r.get("evidence_artifact") or {}70 if isinstance(ev, dict) and "ledger" in ev:71 L = ev["ledger"]; out["retention_rate"] = L["retention_rate"]; out["unsupported"] = L["unsupported_numbers"] + L["unsupported_identifiers"]; out["distorted"] = L["distorted_sentences"]; out["msg_bytes_hop"] = L["message_bytes"]72 out.update(_scenario_costs(r))73 sp = r.get("timing_spans") or {}74 for k in ["encode_ms", "doc_prefill_ms", "compress_ms", "serialize_ms", "load_ms", "receiver_prefill_ms", "query_prefill_ms", "decode_ms", "sender_prefill_ms"]:75 out["span_" + k] = sp.get(k)76 cost = r.get("cost_with_source") or {}77 out["cost_usd"] = cost.get("usd") if isinstance(cost, dict) else None78 return out79