Gahow Wang
e23128ad65
scripts/b2_interference.py is the controlled microbench. It runs two coroutines against the open proxy bypass (direct vLLM endpoints): - decode_load: continuous short-prompt requests at fixed QPS into a designated decode instance, to keep it decode-saturated. - prefill_injections: N large one-token requests at fixed interval, pointed at either the same instance (same-worker variant) or a paired one (different-worker control). Each cell (variant × prefill_size) gets its own metrics.jsonl plus a run_window.json containing t_start_unix/t_end_unix. The shared engine_*.jsonl from the scheduler patch is sliced by that window in the aggregator. analysis/characterization/b2_sweep_analysis.py walks the cell tree, slices the per-worker step log by each cell's window, runs the A5 interference_index() against the slice, and emits a single b2_sweep_summary.json with one row per cell. This is what feeds the "interference vs uncached prefill size" figure. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
118 lines
4.3 KiB
Python
118 lines
4.3 KiB
Python
"""Aggregate B2 microbench cells into a single interference-index sweep summary.
|
||
|
||
Per cell (variant × prefill_size):
|
||
- read metrics.jsonl + run_window.json
|
||
- slice the shared engine_*.jsonl by run window
|
||
- run interference_index() against the slice
|
||
- record (variant, prefill_size, n_overlap, n_clean, tpot_p90_*, idx)
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
from collections import defaultdict
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from analysis.characterization.joined_analysis import (
|
||
_percentile,
|
||
_vllm_rid_matches,
|
||
interference_index,
|
||
load_engine_state,
|
||
load_jsonl,
|
||
write_json,
|
||
)
|
||
|
||
|
||
def _slice_engine_state(
|
||
engine_state_by_worker: dict[str, list[dict]],
|
||
t_start: float,
|
||
t_end: float,
|
||
) -> dict[str, list[dict]]:
|
||
sliced: dict[str, list[dict]] = {}
|
||
for worker, steps in engine_state_by_worker.items():
|
||
sliced[worker] = [s for s in steps
|
||
if t_start <= (s.get("t_unix") or 0.0) <= t_end]
|
||
return sliced
|
||
|
||
|
||
def _to_joined_shape(metrics_rows: list[dict], variant: str) -> list[dict]:
|
||
"""Adapt B2 metric rows to what interference_index expects."""
|
||
joined: list[dict] = []
|
||
for r in metrics_rows:
|
||
if r.get("workload") != "decode":
|
||
continue
|
||
joined.append({
|
||
"request_id": r["request_id"],
|
||
"tpot_s": r.get("tpot_s"),
|
||
"ttft_s": r.get("ttft_s"),
|
||
"latency_s": r.get("latency_s"),
|
||
"endpoint_url": r.get("endpoint"),
|
||
"routed_to": r.get("endpoint"),
|
||
"t_first_token_unix": (
|
||
(r["t_dispatch_unix"] + r["ttft_s"])
|
||
if r.get("ttft_s") is not None
|
||
and r.get("t_dispatch_unix") is not None else None
|
||
),
|
||
"t_finish_unix": r.get("t_finish_unix"),
|
||
"error": r.get("error"),
|
||
})
|
||
return joined
|
||
|
||
|
||
def main() -> None:
|
||
p = argparse.ArgumentParser(description="B2 sweep aggregation")
|
||
p.add_argument("--sweep-dir", type=Path, required=True,
|
||
help="Top-level dir produced by scripts/b2_interference.py")
|
||
p.add_argument("--engine-state-dir", type=Path, required=True)
|
||
p.add_argument("--worker-map", type=str, required=True,
|
||
help="URL=worker_id pairs, comma-separated")
|
||
p.add_argument("--out", type=Path, default=None)
|
||
args = p.parse_args()
|
||
|
||
worker_map = {}
|
||
for entry in args.worker_map.split(","):
|
||
url, _, wid = entry.strip().partition("=")
|
||
if url and wid:
|
||
worker_map[url] = wid
|
||
|
||
engine_state = load_engine_state(args.engine_state_dir)
|
||
rows: list[dict] = []
|
||
for variant_dir in sorted(args.sweep_dir.glob("*/")):
|
||
if variant_dir.name in ("logs",):
|
||
continue
|
||
for cell_dir in sorted(variant_dir.glob("p*/")):
|
||
window_path = cell_dir / "run_window.json"
|
||
metrics_path = cell_dir / "metrics.jsonl"
|
||
if not window_path.exists() or not metrics_path.exists():
|
||
continue
|
||
window = json.loads(window_path.read_text())
|
||
metrics_rows = load_jsonl(metrics_path)
|
||
joined = _to_joined_shape(metrics_rows, variant_dir.name)
|
||
sliced = _slice_engine_state(
|
||
engine_state, window["t_start_unix"], window["t_end_unix"],
|
||
)
|
||
idx = interference_index(joined, sliced, worker_map)
|
||
rows.append({
|
||
"variant": variant_dir.name,
|
||
"prefill_size": int(window["prefill_size"]),
|
||
"decode_endpoint": window["decode_endpoint"],
|
||
"prefill_endpoint": window["prefill_endpoint"],
|
||
"n_decode_requests": sum(1 for r in metrics_rows
|
||
if r.get("workload") == "decode"
|
||
and r.get("error") is None),
|
||
"n_prefill_injections": sum(1 for r in metrics_rows
|
||
if r.get("workload") == "prefill"
|
||
and r.get("error") is None),
|
||
**idx,
|
||
})
|
||
summary = {"rows": rows}
|
||
out_path = args.out or args.sweep_dir / "b2_sweep_summary.json"
|
||
write_json(out_path, summary)
|
||
print(json.dumps(rows, indent=2))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|