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