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 (inter-session and intra-session gaps)
  - 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 N sessions (or until target request count reached)
  3. Re-zero timestamps so first event starts at t=0
  4. Optionally compress time axis to increase load density

Usage:
    python scripts/sample_trace.py \\
        --input ~/ali-trace/trace-glm5.1-formatted/051315-051317.jsonl \\
        --output traces/sampled.jsonl \\
        --target-requests 5000 \\
        --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]],
    *,
    target_requests: int,
    seed: int,
    strategy: str = "random",
) -> list[str]:
    """Select sessions until target request count is reached."""
    all_sids = list(rows_by_session.keys())
    rng = random.Random(seed)

    if strategy == "random":
        rng.shuffle(all_sids)
    elif strategy == "sequential":
        pass  # keep file order
    else:
        raise ValueError(f"Unknown strategy: {strategy}")

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

    return selected


def build_output(
    rows_by_session: dict[str, list[dict]],
    selected: list[str],
    *,
    time_scale: float = 1.0,
) -> list[dict]:
    """Build output rows with re-zeroed timestamps."""
    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

    # Re-zero: subtract earliest timestamp
    t0 = float(out_rows[0]["timestamp"])
    for row in out_rows:
        row["timestamp"] = (float(row["timestamp"]) - t0) / time_scale

    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
    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
    starts_sorted = sorted(session_starts.values())
    deltas = [starts_sorted[i+1] - starts_sorted[i]
              for i in range(len(starts_sorted) - 1)]

    # hash_ids overlap: count unique hash_ids across all requests
    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"  Unique hash blocks: {len(all_hashes)}")
    if deltas:
        deltas.sort()
        p = lambda q: deltas[min(int(q * len(deltas)), len(deltas) - 1)]
        print(f"  Session arrival deltas (s): p10={p(0.1):.2f} p50={p(0.5):.2f} "
              f"p90={p(0.9):.2f} max={max(deltas):.2f}")


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("--target-requests", type=int, default=5000,
                   help="Target number of requests (stops after session that crosses it)")
    p.add_argument("--strategy", choices=["random", "sequential"], default="random",
                   help="Session selection strategy")
    p.add_argument("--time-scale", type=float, default=1.0,
                   help="Compress time axis by this factor (>1 = faster arrival)")
    p.add_argument("--seed", type=int, default=42)
    args = p.parse_args()

    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,
        target_requests=args.target_requests,
        seed=args.seed,
        strategy=args.strategy,
    )

    out_rows = build_output(
        rows_by_session, selected,
        time_scale=args.time_scale,
    )

    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()