agentic-kvc/analysis/characterization/elastic_migration_v2/render_figures.py

"""Render PNG figures for the elastic_migration_v2 section.

Inputs in ./data/ :
- b3_policy_comparison.json
- breakdown_unified.json, breakdown_unified_kv_both.json,
  breakdown_unified_v2.json, breakdown_unified_v2_strict.json
- per_worker_<policy>.json for each of the four

Outputs in ./figures/ :
- fig_kv_both_overhead.png      — three-way latency bars (plain vs kv_both vs v2)
- fig_v2_trigger_funnel.png     — request count per fall-through reason
- fig_v2_predicted_vs_actual.png — cost-model migrate prediction vs realized TTFT
- fig_three_way_hotspot.png      — per-worker TTFT p90 grouped bars
"""

from __future__ import annotations

import json
from collections import Counter
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

ROOT = Path(__file__).parent
DATA = ROOT / "data"
OUT = ROOT / "figures"
OUT.mkdir(parents=True, exist_ok=True)


def _load(name: str):
    return json.loads((DATA / name).read_text())


POLICY_COLORS = {
    "unified":            "#2ca02c",
    "unified_kv_both":    "#9467bd",
    "unified_v2":         "#d62728",
    "unified_v2_strict":  "#ff7f0e",
}


def fig_kv_both_overhead():
    comp = _load("b3_policy_comparison.json")
    by = {r["policy"]: r for r in comp["rows"]}
    pols = ["unified", "unified_kv_both", "unified_v2"]
    metrics = [
        ("TTFT p90 (s)",   lambda r: r["ttft_p90_s"]),
        ("TPOT p90 (ms)",  lambda r: r["tpot_p90_s"] * 1000),
        ("E2E p90 (s)",    lambda r: r["e2e_p90_s"]),
        ("hotspot index",  lambda r: r["hotspot_index_ttft_p90"]),
    ]
    fig, axes = plt.subplots(1, 4, figsize=(14, 4))
    for ax, (label, fn) in zip(axes, metrics):
        vals = [fn(by[p]) for p in pols]
        bars = ax.bar(pols, vals,
                       color=[POLICY_COLORS[p] for p in pols],
                       edgecolor="black", linewidth=0.5)
        ax.set_title(label)
        ax.tick_params(axis="x", rotation=20, labelsize=9)
        for b, v in zip(bars, vals):
            ax.text(b.get_x() + b.get_width() / 2, v,
                     f"{v:.2f}" if v < 100 else f"{v:.0f}",
                     ha="center", va="bottom", fontsize=9)
        ax.grid(alpha=0.3, axis="y")
        # delta annotation
        baseline = vals[0]
        for i, v in enumerate(vals):
            if i == 0:
                continue
            pct = (v - baseline) / baseline * 100
            ax.text(i, v * 0.5, f"{pct:+.0f}%", ha="center",
                     fontsize=10, fontweight="bold",
                     color="darkred" if pct > 0 else "darkgreen")
    fig.suptitle(
        "kv_both adds ~45% to TTFT p90 even without PD-sep firing.\n"
        "v2's PD-sep barely recovers the gap (and overshoots TTFT p99)."
    )
    fig.tight_layout()
    fig.savefig(OUT / "fig_kv_both_overhead.png", dpi=120)
    plt.close(fig)


def _bucket_reasons(data):
    """Collapse v2_reason strings into the funnel buckets."""
    buckets = Counter()
    for r in data:
        if r.get("v2_pd_sep") is True:
            buckets["PD-sep TRIGGERED"] += 1
            continue
        reason = (r.get("v2_reason") or "no_v2_reason").split(" (")[0]
        if reason.startswith("local_cost"):
            reason = "cost_benefit not enough margin"
        buckets[reason] += 1
    return buckets


def fig_v2_trigger_funnel():
    strict = _load("breakdown_unified_v2_strict.json")
    relaxed = _load("breakdown_unified_v2.json")
    bs = _bucket_reasons(strict)
    br = _bucket_reasons(relaxed)
    order = [
        "new_local_below_threshold",
        "chosen_no_active_decode",
        "chosen_few_decodes",
        "src_cache_below_threshold",
        "src_not_meaningfully_more_cache",
        "cost_benefit not enough margin",
        "PD-sep TRIGGERED",
    ]
    labels = [k for k in order if k in bs or k in br]
    strict_vals = [bs.get(k, 0) for k in labels]
    relaxed_vals = [br.get(k, 0) for k in labels]

    x = range(len(labels))
    width = 0.4
    fig, ax = plt.subplots(figsize=(11, 5))
    ax.bar([i - width / 2 for i in x], strict_vals, width,
            label=f"v2.0 strict (PD-sep={bs['PD-sep TRIGGERED']}/{sum(bs.values())} "
                  f"= {bs['PD-sep TRIGGERED']*100/sum(bs.values()):.2f}%)",
            color="#ff7f0e", edgecolor="black", linewidth=0.5)
    ax.bar([i + width / 2 for i in x], relaxed_vals, width,
            label=f"v2.1 relaxed (PD-sep={br['PD-sep TRIGGERED']}/{sum(br.values())} "
                  f"= {br['PD-sep TRIGGERED']*100/sum(br.values()):.2f}%)",
            color="#d62728", edgecolor="black", linewidth=0.5)
    ax.set_xticks(list(x))
    ax.set_xticklabels(labels, rotation=20, ha="right", fontsize=9)
    ax.set_ylabel("request count")
    ax.set_yscale("log")
    ax.set_title(
        "Why v2 rarely PD-seps: 88-76% of requests have new_local < threshold\n"
        "(intra-session cache already hot). Relaxing thresholds barely helps."
    )
    ax.legend()
    ax.grid(alpha=0.3, axis="y", which="both")
    for i, (s, r) in enumerate(zip(strict_vals, relaxed_vals)):
        if s > 0:
            ax.text(i - width / 2, s * 1.05, str(s), ha="center", fontsize=8)
        if r > 0:
            ax.text(i + width / 2, r * 1.05, str(r), ha="center", fontsize=8)
    fig.tight_layout()
    fig.savefig(OUT / "fig_v2_trigger_funnel.png", dpi=120)
    plt.close(fig)


def fig_v2_predicted_vs_actual():
    """For each PD-sep'd request, plot model-predicted migrate cost
    vs realized TTFT. Should sit near y=x if model is calibrated; sits
    far above if mechanism is more expensive than modeled."""
    relaxed = _load("breakdown_unified_v2.json")
    triggered = [r for r in relaxed if r.get("v2_pd_sep") is True]
    if not triggered:
        return
    predicted = []
    actual = []
    sizes = []
    rids = []
    for r in triggered:
        cm = r.get("v2_cost_migrate_s")
        t0 = r.get("t_proxy_recv")
        t_first = r.get("t_first_token")
        if cm is None or t0 is None or t_first is None:
            continue
        ttft = t_first - t0
        predicted.append(cm)
        actual.append(ttft)
        sizes.append(r.get("input_length", 0))
        rids.append(r.get("request_id", "?"))

    fig, ax = plt.subplots(figsize=(7, 5))
    ax.scatter(predicted, actual,
                s=[max(100, sz / 100) for sz in sizes],
                color="#d62728", edgecolors="black", alpha=0.75)
    for p, a, sz, rid in zip(predicted, actual, sizes, rids):
        ax.annotate(f"input={sz}",
                     (p, a), xytext=(8, 6), textcoords="offset points",
                     fontsize=9)
    # y=x reference + 10x line + 20x line
    lo = 0.5
    hi = max(50, max(actual) * 1.2)
    ax.plot([lo, hi], [lo, hi], "k--", alpha=0.5, label="y = x (calibrated)")
    ax.plot([lo, hi], [lo * 10, hi * 10], color="gray", linestyle=":",
             alpha=0.4, label="10x")
    ax.plot([lo, hi], [lo * 20, hi * 20], color="lightgray", linestyle=":",
             alpha=0.4, label="20x")
    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlim(lo, hi)
    ax.set_ylim(lo, hi)
    ax.set_xlabel("Cost model: predicted migrate cost (s)")
    ax.set_ylabel("Realized TTFT (s)")
    ax.set_title(
        "All 5 PD-sep triggered requests in v2.1 sit far above y=x.\n"
        "Real transfer cost ~10-20x what the calibrated model predicted."
    )
    ax.grid(alpha=0.3, which="both")
    ax.legend(loc="lower right")
    fig.tight_layout()
    fig.savefig(OUT / "fig_v2_predicted_vs_actual.png", dpi=120)
    plt.close(fig)


def fig_three_way_hotspot():
    pols = ["unified", "unified_kv_both", "unified_v2"]
    per_worker = {p: _load(f"per_worker_{p}.json") for p in pols}
    workers = sorted(per_worker["unified"]["per_worker_ttft_p90_s"].keys())

    x = range(len(workers))
    width = 0.27
    fig, ax = plt.subplots(figsize=(11, 5))
    for i, p in enumerate(pols):
        d = per_worker[p]["per_worker_ttft_p90_s"]
        vals = [d[w] for w in workers]
        offset = (i - 1) * width
        ax.bar([j + offset for j in x], vals, width,
                label=f"{p} (hotspot={per_worker[p]['hotspot_index_ttft_p90']:.2f})",
                color=POLICY_COLORS[p], edgecolor="black", linewidth=0.4)
    short = [w.replace("http://127.0.0.1:", ":") for w in workers]
    ax.set_xticks(list(x))
    ax.set_xticklabels(short, rotation=0, fontsize=9)
    ax.set_ylabel("worker TTFT p90 (s)")
    ax.set_title(
        "Per-worker TTFT p90 distribution. kv_both alone makes the hot worker hotter\n"
        "(unified→kv_both: 37.7s→43.5s peak); v2's 5 PD-sep triggers nudge it back."
    )
    ax.legend(loc="upper left", fontsize=9)
    ax.grid(alpha=0.3, axis="y")
    fig.tight_layout()
    fig.savefig(OUT / "fig_three_way_hotspot.png", dpi=120)
    plt.close(fig)


def main():
    fig_kv_both_overhead()
    fig_v2_trigger_funnel()
    fig_v2_predicted_vs_actual()
    fig_three_way_hotspot()
    print(f"wrote 4 figures to {OUT}")


if __name__ == "__main__":
    main()