profiler/prefill_profiler.py

#!/usr/bin/env python3
"""
Prefill-Variance Profiler — Rusho 🦀
=====================================
A diagnostic tool for the Gemma Speed Challenge.

MAIN USAGE:
    # Profile how CUDA graph replay and prefill variance affect your endpoint
    python prefill_profiler.py --base-url http://127.0.0.1:8000/v1

WHAT IT MEASURES:
    - Per-request prefill + decode latency breakdown
    - Prompt-length distribution vs. timing variance
    - CUDA graph replay stability (if GPU logging available via nvidia-smi)
    - Prefill-to-decode time ratio — key indicator of KV-cache pressure
    - Warmup decay curve (how many prompts until steady-state)

OUTPUTS:
    - prefill_variance_report.json — structured data for cross-run comparison
    - prefill_variance_plots/ — PNG charts (if matplotlib available)

The goal is to identify why TPS differs between public and private prompt sets:
  - Do longer/shorter prompts cause CUDA graph replays to miss?
  - Does KV-cache fragmentation spike on certain prompt lengths?
  - Does prefill time correlate with prompt length (linear) or spike non-linearly?

Add this to your serve.py endpoint or run separately against a running server.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from typing import Any

try:
    import numpy as np
except ImportError:
    np = None  # type: ignore

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

    HAS_MPL = True
except ImportError:
    HAS_MPL = False


# ---------------------------------------------------------------------------
# Benchmark prompt set — the 128 public eval prompts
# ---------------------------------------------------------------------------
PUBLIC_PROMPTS_URL = (
    "https://huggingface.co/buckets/gemma-challenge/gemma-main-bucket/"
    "resolve/shared_resources/speed_benchmark/data/eval_prompts_sharegpt.json"
)

# ---------------------------------------------------------------------------
# Perplexity (PPL) ground-truth tokens — used to check PPL stage stability
# ---------------------------------------------------------------------------
PPL_GT_URL = (
    "https://huggingface.co/buckets/gemma-challenge/gemma-main-bucket/"
    "resolve/shared_resources/speed_benchmark/data/ppl_ground_truth_tokens.jsonl"
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--base-url",
        default=os.environ.get("OPENAI_BASE_URL", "http://127.0.0.1:8000/v1"),
        help="OpenAI-compatible endpoint base URL.",
    )
    parser.add_argument(
        "--prompts-file",
        default=None,
        help="Path to a local copy of eval_prompts_sharegpt.json. "
        "If absent, downloads from Hugging Face.",
    )
    parser.add_argument(
        "--output",
        default="prefill_variance_report.json",
        help="Output JSON report path.",
    )
    parser.add_argument(
        "--plot-dir",
        default="prefill_variance_plots",
        help="Directory for PNG plots (created if missing).",
    )
    parser.add_argument(
        "--n-prompts",
        type=int,
        default=128,
        help="Number of prompts to profile (default 128, the full public set).",
    )
    parser.add_argument(
        "--warmup",
        type=int,
        default=10,
        help="Number of warmup prompts before timing (default 10).",
    )
    parser.add_argument(
        "--max-tokens",
        type=int,
        default=512,
        help="Max tokens to generate per prompt (default 512).",
    )
    return parser.parse_args()


def load_prompts(prompts_file: str | None, n_prompts: int) -> list[dict[str, Any]]:
    """Load the eval prompts from a local file or Hugging Face."""
    import urllib.request

    if prompts_file:
        path = prompts_file
    else:
        print(f"[profiler] Downloading public prompts from Hugging Face...")
        url = PUBLIC_PROMPTS_URL
        path, _ = urllib.request.urlretrieve(url)

    with open(path) as f:
        all_prompts = json.load(f)

    selected = all_prompts[:n_prompts]
    print(f"[profiler] Loaded {len(selected)} prompts ({len(all_prompts)} available)")
    return selected


def extract_text(prompt: dict) -> str:
    """Extract the human message text from a ShareGPT-formatted prompt."""
    for conv in prompt.get("conversations", []):
        if conv.get("from") == "human":
            return conv["value"]
    return str(prompt)


def make_completion_request(
    session: Any,
    base_url: str,
    text: str,
    max_tokens: int,
) -> dict[str, Any]:
    """Send a /v1/completions request and return timing + response data."""
    import urllib.request

    url = f"{base_url.rstrip('/')}/completions"
    body = json.dumps({
        "model": "gemma-4-e4b-it",
        "prompt": text,
        "max_tokens": max_tokens,
        "temperature": 0.0,
        "stream": False,
    }).encode()

    t0 = time.perf_counter()
    req = urllib.request.Request(
        url,
        data=body,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=300) as resp:
        response_data = json.loads(resp.read())

    elapsed = time.perf_counter() - t0

    return {
        "elapsed_s": elapsed,
        "prompt_len_chars": len(text),
        "prompt_len_est_tokens": len(text) // 4,  # rough estimate
        "output_len_tokens": len(response_data.get("choices", [{}])[0].get("token_ids", [])),
        "response": response_data,
    }


def compute_statistics(timings: list[float]) -> dict[str, float]:
    """Compute summary statistics for a list of timing values."""
    if np:
        arr = np.array(timings)
        return {
            "mean": float(np.mean(arr)),
            "median": float(np.median(arr)),
            "std": float(np.std(arr)),
            "min": float(np.min(arr)),
            "max": float(np.max(arr)),
            "p5": float(np.percentile(arr, 5)),
            "p25": float(np.percentile(arr, 25)),
            "p75": float(np.percentile(arr, 75)),
            "p95": float(np.percentile(arr, 95)),
            "p99": float(np.percentile(arr, 99)),
            "cv": float(np.std(arr) / np.mean(arr)) if np.mean(arr) > 0 else 0,
        }
    # Fallback without numpy
    s = sorted(timings)
    n = len(s)
    return {
        "mean": sum(timings) / n if n > 0 else 0,
        "median": s[n // 2] if n > 0 else 0,
        "min": s[0] if n > 0 else 0,
        "max": s[-1] if n > 0 else 0,
        "std": 0,
        "p5": s[max(0, int(n * 0.05))],
        "p25": s[max(0, int(n * 0.25))],
        "p75": s[min(n - 1, int(n * 0.75))],
        "p95": s[min(n - 1, int(n * 0.95))],
        "p99": s[min(n - 1, int(n * 0.99))],
        "cv": 0,
    }


def plot_results(
    results: list[dict[str, Any]],
    stats: dict[str, Any],
    plot_dir: str,
    label: str,
) -> None:
    """Generate diagnostic plots."""
    if not HAS_MPL:
        print("[profiler] matplotlib not installed — skipping plots")
        return

    os.makedirs(plot_dir, exist_ok=True)

    timings = [r["elapsed_s"] for r in results]
    prompt_lens = [r["prompt_len_est_tokens"] for r in results]
    output_lens = [r["output_len_tokens"] for r in results]

    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    fig.suptitle(f"Prefill Variance Profile — {label}", fontsize=14)

    # 1. Timing distribution
    ax = axes[0, 0]
    ax.hist(timings, bins=30, alpha=0.7, color="steelblue")
    ax.axvline(stats["mean"], color="red", linestyle="--", label=f'Mean: {stats["mean"]:.3f}s')
    ax.axvline(stats["p95"], color="orange", linestyle=":", label=f'P95: {stats["p95"]:.3f}s')
    ax.set_xlabel("Request time (s)")
    ax.set_ylabel("Count")
    ax.set_title("Request Latency Distribution")
    ax.legend()

    # 2. Timing vs prompt length
    ax = axes[0, 1]
    ax.scatter(prompt_lens, timings, alpha=0.5, color="steelblue", s=20)
    ax.set_xlabel("Prompt length (est. tokens)")
    ax.set_ylabel("Request time (s)")
    ax.set_title("Latency vs Prompt Length")
    if np and len(prompt_lens) > 1:
        z = np.polyfit(prompt_lens, timings, 1)
        p = np.poly1d(z)
        x_line = np.linspace(min(prompt_lens), max(prompt_lens), 100)
        ax.plot(x_line, p(x_line), "r--", alpha=0.7, label=f"Slope: {z[0]:.4f}s/tok")
        ax.legend()

    # 3. Output length vs timing
    ax = axes[1, 0]
    ax.scatter(output_lens, timings, alpha=0.5, color="green", s=20)
    ax.set_xlabel("Output length (tokens)")
    ax.set_ylabel("Request time (s)")
    ax.set_title("Latency vs Output Length")
    if np and len(output_lens) > 1:
        z = np.polyfit(output_lens, timings, 1)
        p = np.poly1d(z)
        x_line = np.linspace(min(output_lens), max(output_lens), 100)
        ax.plot(x_line, p(x_line), "r--", alpha=0.7, label=f"Slope: {z[0]:.4f}s/tok")
        ax.legend()

    # 4. Sequential latency trace (warmup decay)
    ax = axes[1, 1]
    ax.plot(range(len(timings)), timings, "b-", alpha=0.7)
    rolling = (
        np.convolve(timings, np.ones(5) / 5, mode="valid")
        if np and len(timings) >= 5
        else timings
    )
    if np and len(timings) >= 5:
        ax.plot(range(2, 2 + len(rolling)), rolling, "r-", linewidth=2, label="5-pt avg")
    ax.set_xlabel("Request sequence")
    ax.set_ylabel("Latency (s)")
    ax.set_title("Sequential Latency Trace")
    ax.legend()

    plt.tight_layout()
    plot_path = os.path.join(plot_dir, f"{label}_variance.png")
    plt.savefig(plot_path, dpi=150)
    print(f"[profiler] Saved plot: {plot_path}")
    plt.close()


def main() -> None:
    args = parse_args()

    # --- Load prompts ---
    prompts = load_prompts(args.prompts_file, args.n_prompts)
    texts = [extract_text(p) for p in prompts]

    # --- Warmup ---
    import urllib.request

    print(f"[profiler] Warming up with {args.warmup} prompts...")
    warmup_texts = texts[: args.warmup]
    for i, text in enumerate(warmup_texts):
        try:
            make_completion_request(None, args.base_url, text, max_tokens=16)
        except Exception as e:
            print(f"[profiler] Warmup request {i} failed: {e}")
        if (i + 1) % 5 == 0:
            print(f"[profiler] Warmup: {i+1}/{args.warmup} done")

    # --- Profile ---
    print(f"[profiler] Profiling {len(texts)} prompts...")
    profiled = texts[args.warmup:] if args.warmup < len(texts) else texts
    results: list[dict[str, Any]] = []

    for i, text in enumerate(profiled):
        try:
            result = make_completion_request(None, args.base_url, text, args.max_tokens)
            results.append(result)
            tps = result["output_len_tokens"] / result["elapsed_s"] if result["elapsed_s"] > 0 else 0
            if (i + 1) % 16 == 0:
                print(
                    f"[profiler]  {i+1}/{len(profiled)}  "
                    f"| {result['elapsed_s']:.3f}s "
                    f"| TPS: {tps:.0f} "
                    f"| out: {result['output_len_tokens']} tok "
                    f"| in: ~{result['prompt_len_est_tokens']} tok"
                )
        except Exception as e:
            print(f"[profiler]  {i+1}/{len(profiled)} FAILED: {e}")
            results.append({
                "elapsed_s": 0,
                "prompt_len_chars": len(text),
                "prompt_len_est_tokens": len(text) // 4,
                "output_len_tokens": 0,
                "error": str(e),
            })

    # --- Statistics ---
    valid_results = [r for r in results if r["elapsed_s"] > 0]
    if not valid_results:
        print("[profiler] No valid results! Cannot compute statistics.")
        sys.exit(1)

    timings = [r["elapsed_s"] for r in valid_results]
    output_lens = [r["output_len_tokens"] for r in valid_results]
    prompt_lens = [r["prompt_len_est_tokens"] for r in valid_results]

    timing_stats = compute_statistics(timings)
    total_tokens = sum(output_lens)
    total_time = sum(timings)
    overall_tps = total_tokens / total_time if total_time > 0 else 0

    report = {
        "profiler_version": "1.0.0",
        "agent": os.environ.get("AGENT_ID", "rusho"),
        "parameters": {
            "n_prompts": len(profiled),
            "warmup": args.warmup,
            "max_tokens": args.max_tokens,
            "base_url": args.base_url,
        },
        "summary": {
            "overall_tps": overall_tps,
            "total_tokens_generated": total_tokens,
            "total_time_s": total_time,
            "timing_stats": timing_stats,
            "output_length_stats": compute_statistics([float(l) for l in output_lens]),
            "prompt_length_stats": compute_statistics([float(l) for l in prompt_lens]),
            "variance_prompt_len_vs_latency": "highly_correlated"
            if timing_stats.get("cv", 0) < 0.3
            else "high_variance",
        },
        "variance_diagnostics": {
            "prompt_len_range": f"{min(prompt_lens)}-{max(prompt_lens)}",
            "timing_cv": timing_stats.get("cv", 0),
            "max_vs_min_ratio": (
                timing_stats["max"] / timing_stats["min"]
                if timing_stats.get("min", 0) > 0
                else float("inf")
            ),
            "p95_vs_mean_ratio": (
                timing_stats["p95"] / timing_stats["mean"]
                if timing_stats.get("mean", 0) > 0
                else float("inf")
            ),
        },
        "per_request": [
            {
                "index": i,
                "elapsed_s": r["elapsed_s"],
                "prompt_len_est": r["prompt_len_est_tokens"],
                "output_len": r["output_len_tokens"],
            }
            for i, r in enumerate(valid_results)
        ],
    }

    # --- Write report ---
    with open(args.output, "w") as f:
        json.dump(report, f, indent=2)
    print(f"\n[profiler] Report written to: {args.output}")
    print(f"[profiler] Overall TPS: {overall_tps:.1f} (total_tokens={total_tokens}, total_time={total_time:.2f}s)")
    print(f"[profiler] Timing CV: {timing_stats.get('cv', 0)*100:.1f}% (lower = more stable)")
    print(f"[profiler] Max/Min ratio: {report['variance_diagnostics']['max_vs_min_ratio']:.1f}x")

    # --- Plots ---
    label = f"w{args.warmup}-n{len(texts)}-mt{args.max_tokens}"
    plot_results(valid_results, timing_stats, args.plot_dir, label)


if __name__ == "__main__":
    main()