agentic-kvc/scripts/sample_trace.py

"""Sample sessions from the full cluster-scale trace to fit a single machine.

Preserves:
  - Complete session structure (all turns within a session kept together)
  - Original arrival timing (re-zeroed to t=0 but NOT compressed)
  - hash_ids for KV cache reuse patterns
  - Request type distribution

Sampling strategy:
  1. Group requests by session (derived from parent_chat_id chains)
  2. Randomly sample a fraction of sessions (--sample-ratio)
     OR sample until target request count (--target-requests)
  3. Re-zero timestamps so first event starts at t=0
  4. The resulting QPS is proportional to the sample ratio,
     preserving the production arrival pattern

Usage:
    # Sample 1.6% of sessions (e.g., 8 GPUs / 500 cluster GPUs)
    python scripts/sample_trace.py \\
        --input ~/ali-trace/trace-glm5.1-formatted/051315-051317.jsonl \\
        --output traces/sampled_ratio016.jsonl \\
        --sample-ratio 0.016 --seed 42

    # Sample by request count (legacy)
    python scripts/sample_trace.py \\
        --input ... --output ... --target-requests 1000 --seed 42
"""

from __future__ import annotations

import argparse
import collections
import json
import random
import sys
from pathlib import Path


def load_raw_rows(path: Path) -> dict[str, list[dict]]:
    """Load trace, group rows by resolved session_id. Preserve file order."""
    chat_to_session: dict[int, str] = {}
    rows_by_session: dict[str, list[dict]] = collections.OrderedDict()

    with path.open("r", encoding="utf-8") as fh:
        for line in fh:
            row = json.loads(line)
            cid = int(row["chat_id"])
            pid = int(row["parent_chat_id"])

            if "session_id" in row:
                sid = str(row["session_id"])
            elif pid < 0:
                sid = str(cid)
            else:
                sid = chat_to_session.get(pid, str(pid))
            chat_to_session[cid] = sid

            row["_session_id"] = sid
            rows_by_session.setdefault(sid, []).append(row)

    return rows_by_session


def sample_sessions(
    rows_by_session: dict[str, list[dict]],
    *,
    sample_ratio: float | None = None,
    target_requests: int | None = None,
    seed: int,
) -> list[str]:
    """Select sessions by ratio or until target request count."""
    all_sids = list(rows_by_session.keys())
    rng = random.Random(seed)
    rng.shuffle(all_sids)

    if sample_ratio is not None:
        n_select = max(1, int(len(all_sids) * sample_ratio))
        return all_sids[:n_select]

    if target_requests is not None:
        selected = []
        total = 0
        for sid in all_sids:
            selected.append(sid)
            total += len(rows_by_session[sid])
            if total >= target_requests:
                break
        return selected

    raise ValueError("Must specify --sample-ratio or --target-requests")


def build_output(
    rows_by_session: dict[str, list[dict]],
    selected: list[str],
) -> list[dict]:
    """Build output rows with re-zeroed timestamps (no time compression)."""
    out_rows = []
    for sid in selected:
        for row in rows_by_session[sid]:
            out = {k: v for k, v in row.items() if not k.startswith("_")}
            out["session_id"] = sid
            out_rows.append(out)

    out_rows.sort(key=lambda r: float(r["timestamp"]))

    if not out_rows:
        return out_rows

    t0 = float(out_rows[0]["timestamp"])
    for row in out_rows:
        row["timestamp"] = float(row["timestamp"]) - t0

    return out_rows


def print_summary(
    rows_by_session: dict[str, list[dict]],
    selected: list[str],
    out_rows: list[dict],
) -> None:
    n_sessions = len(selected)
    n_requests = len(out_rows)
    turns_per_session = [len(rows_by_session[s]) for s in selected]
    multi_turn = sum(1 for t in turns_per_session if t > 1)

    input_lens = [r["input_length"] for r in out_rows]
    output_lens = [r["output_length"] for r in out_rows]

    span_s = float(out_rows[-1]["timestamp"]) if out_rows else 0
    qps = n_requests / span_s if span_s > 0 else 0

    session_starts = {}
    for r in out_rows:
        sid = r["session_id"]
        ts = float(r["timestamp"])
        if sid not in session_starts:
            session_starts[sid] = ts

    # hash_ids overlap
    all_hashes = set()
    for r in out_rows:
        all_hashes.update(r.get("hash_ids", []))

    print(f"Sampled: {n_sessions} sessions, {n_requests} requests")
    print(f"  Multi-turn sessions: {multi_turn} ({multi_turn/n_sessions*100:.1f}%)")
    print(f"  Turns/session: min={min(turns_per_session)} max={max(turns_per_session)} "
          f"avg={sum(turns_per_session)/len(turns_per_session):.1f}")
    print(f"  Input length: min={min(input_lens)} max={max(input_lens)} "
          f"avg={sum(input_lens)/len(input_lens):.0f}")
    print(f"  Output length: min={min(output_lens)} max={max(output_lens)} "
          f"avg={sum(output_lens)/len(output_lens):.0f}")
    print(f"  Trace span: {span_s:.1f}s ({span_s/60:.1f} min)")
    print(f"  QPS: {qps:.2f} req/s")
    print(f"  Unique hash blocks: {len(all_hashes)}")


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__,
                                formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("--input", type=Path, required=True,
                   help="Path to the full trace JSONL file")
    p.add_argument("--output", type=Path, required=True,
                   help="Path to write sampled trace JSONL")
    p.add_argument("--sample-ratio", type=float, default=None,
                   help="Fraction of sessions to sample (e.g. 0.016 for 8/500 GPU ratio)")
    p.add_argument("--target-requests", type=int, default=None,
                   help="Target number of requests (legacy, stops after session that crosses it)")
    p.add_argument("--seed", type=int, default=42)
    args = p.parse_args()

    if args.sample_ratio is None and args.target_requests is None:
        p.error("Must specify --sample-ratio or --target-requests")

    print(f"Loading trace from {args.input} ...")
    rows_by_session = load_raw_rows(args.input)
    total_sessions = len(rows_by_session)
    total_requests = sum(len(v) for v in rows_by_session.values())
    print(f"Full trace: {total_sessions} sessions, {total_requests} requests")

    selected = sample_sessions(
        rows_by_session,
        sample_ratio=args.sample_ratio,
        target_requests=args.target_requests,
        seed=args.seed,
    )

    out_rows = build_output(rows_by_session, selected)

    print_summary(rows_by_session, selected, out_rows)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as fh:
        for row in out_rows:
            fh.write(json.dumps(row, ensure_ascii=False) + "\n")
    print(f"\nWrote {len(out_rows)} rows to {args.output}")


if __name__ == "__main__":
    main()