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