Agentic workload PD separation analysis with trace-driven benchmarks
Systematic study of prefill-decode disaggregation for agentic LLM workloads using production GLM-5.1 coder trace (2.1M requests, 71B input tokens). Key findings: - Cache-aware routing improves TPOT p90 by 15% and APC from 20.8% to 44.7% without PD separation, matching PD-Sep's decode isolation benefit - PD separation adds +72% TTFT overhead (KV transfer) with no TPOT gain when using the same cache-aware scheduler - Prefill remains compute-bound even at 95% KV cache reuse (AI >1000x vs decode AI <2), but absolute FLOPs drop 71% from cache hits - For agentic MoE workloads, cache-aware routing > PD separation Infrastructure: - Trace sampler preserving session structure + hash_ids for prefix sharing - Async trace replayer with streaming TTFT/TPOT/E2E measurement - Unified cache-aware + token-level load-balanced global scheduler proxy supporting both PD-colocated and PD-disaggregated (Mooncake/RDMA) modes - vLLM 0.18.1 scheduler patch for KV transfer abort race condition - Roofline analysis tool for prefill/decode compute characterization Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
204
scripts/sample_trace.py
Normal file
204
scripts/sample_trace.py
Normal file
@@ -0,0 +1,204 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user