MB5 analysis: per-role KV split proves static-partition mismatch
aggregate_mb5.py: - Split the cluster KV timeline by role (P-pool vs D-pool) using a PID->role map parsed from vllm_logs filenames. The cluster average hid the result — 6P+2D/4P+4D look ~45% utilized but the decode pool is actually pegged at ~100% while prefill idles at ~30%. - Two-stage reduce/plot: --reduce-to (numpy-only, runs on the serving host over multi-GB snapshot dirs) dumps a compact JSON; --from-reduced (matplotlib) renders locally. matplotlib import is now lazy. - New plot_role_split figure + p/d peak/steady columns in the CSV. PD_DISAGG_RESULTS.md: consolidated writeup with figures inline. Verdict: no static P:D ratio beats 8C colocation. The binding constraint moves with the ratio (D-pool saturates at 6P+2D/4P+4D, P-pool jams at 2P+6D -> 91% request loss); 8C's shared pool stays elastic at 34% steady, 100% completion. PD wins TPOT (10-35x cleaner, the MB1 phase-isolation benefit is real) but loses TTFT and sheds load. Round-robin P routing also zeroes prefix-cache reuse; a session-affinity re-run of 6P+2D is in flight to test the fix. Figures (rep1): mb5_kv_timeline, mb5_role_split, mb5_peak_utilization, mb5_latency_compare + mb5_summary.csv. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
@@ -31,11 +31,11 @@ import json
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
# matplotlib is imported lazily inside the plot functions so the --reduce
|
||||
# path (numpy-only) can run on a serving host without matplotlib installed.
|
||||
|
||||
|
||||
def load_snapshots_for_run(snap_dir: Path) -> list[dict]:
|
||||
"""Merge all per-PID snapshot files in snap_dir, tag with pid, sort by t_unix."""
|
||||
@@ -57,19 +57,60 @@ def load_snapshots_for_run(snap_dir: Path) -> list[dict]:
|
||||
return out
|
||||
|
||||
|
||||
def cluster_timeline(snaps: list[dict], bin_size_s: float = 1.0) -> tuple[np.ndarray, ...]:
|
||||
def load_pid_roles(logs_dir: Path) -> dict[int, str]:
|
||||
"""Map EngineCore PID -> 'P' | 'D' | 'C' by parsing vllm_logs filenames.
|
||||
|
||||
File names look like vllm_idx{i}_gpu{g}_kv_{producer|consumer|both}.log and
|
||||
each contains '(EngineCore pid=NNNN)'. Returns {} if no logs found.
|
||||
"""
|
||||
role_map = {"producer": "P", "consumer": "D", "both": "C"}
|
||||
out: dict[int, str] = {}
|
||||
if not logs_dir.is_dir():
|
||||
return out
|
||||
for f in logs_dir.glob("vllm_idx*_kv_*.log"):
|
||||
role = None
|
||||
for key, short in role_map.items():
|
||||
if f.name.endswith(f"kv_{key}.log"):
|
||||
role = short
|
||||
break
|
||||
if role is None:
|
||||
continue
|
||||
with f.open(errors="ignore") as fh:
|
||||
for line in fh:
|
||||
if "EngineCore pid=" in line:
|
||||
try:
|
||||
pid = int(line.split("EngineCore pid=")[1].split(")")[0].split()[0])
|
||||
out[pid] = role
|
||||
break
|
||||
except (ValueError, IndexError):
|
||||
continue
|
||||
return out
|
||||
|
||||
|
||||
def cluster_timeline(snaps: list[dict], bin_size_s: float = 1.0,
|
||||
keep_pids: set | None = None,
|
||||
t0: float | None = None,
|
||||
n_bins: int | None = None) -> tuple[np.ndarray, ...]:
|
||||
"""Bin per-PID snapshots into a cluster-wide timeline.
|
||||
|
||||
For each time bin, sum used_blocks across PIDs that emitted a snapshot
|
||||
in that bin. PIDs without a sample in a bin carry their previous value
|
||||
forward (so a quiet PID doesn't artificially drop the total).
|
||||
|
||||
If keep_pids is given, only those PIDs are counted (used for per-role
|
||||
P-pool / D-pool splits); the pool ceiling is summed over the same subset.
|
||||
Pass a shared t0/n_bins so role-splits land on the same time grid.
|
||||
"""
|
||||
if keep_pids is not None:
|
||||
snaps = [s for s in snaps if s["pid"] in keep_pids]
|
||||
if not snaps:
|
||||
empty = np.array([], dtype=float)
|
||||
return empty, empty, empty, empty, empty
|
||||
t0 = snaps[0]["t_unix"]
|
||||
t_end = snaps[-1]["t_unix"]
|
||||
n_bins = max(1, int(np.ceil((t_end - t0) / bin_size_s)) + 1)
|
||||
if t0 is None:
|
||||
t0 = snaps[0]["t_unix"]
|
||||
if n_bins is None:
|
||||
t_end = snaps[-1]["t_unix"]
|
||||
n_bins = max(1, int(np.ceil((t_end - t0) / bin_size_s)) + 1)
|
||||
times = np.arange(n_bins) * bin_size_s
|
||||
|
||||
pids = sorted({s["pid"] for s in snaps})
|
||||
@@ -118,34 +159,60 @@ def load_summary(rundir: Path) -> dict | None:
|
||||
return json.loads(p.read_text())
|
||||
|
||||
|
||||
def _steady_median(arr: np.ndarray) -> float:
|
||||
n = len(arr)
|
||||
if n == 0:
|
||||
return 0.0
|
||||
if n >= 10:
|
||||
return float(np.median(arr[int(n * 0.1):int(n * 0.9)]))
|
||||
return float(np.median(arr))
|
||||
|
||||
|
||||
def per_run_metrics(snaps_dir: Path, rundir: Path) -> dict:
|
||||
snaps = load_snapshots_for_run(snaps_dir)
|
||||
times, total_used, pool_frac, total_waiting, total_running = cluster_timeline(snaps)
|
||||
summary = load_summary(rundir) or {}
|
||||
|
||||
# Trim the warmup/cooldown 10% to compute "steady-state" stats
|
||||
n = len(times)
|
||||
if n >= 10:
|
||||
lo, hi = int(n * 0.1), int(n * 0.9)
|
||||
frac_steady = pool_frac[lo:hi]
|
||||
wait_steady = total_waiting[lo:hi]
|
||||
# Establish a shared time grid (global t0 / n_bins) so the overall and
|
||||
# per-role timelines all line up on the same x axis.
|
||||
if snaps:
|
||||
t0 = snaps[0]["t_unix"]
|
||||
t_end = snaps[-1]["t_unix"]
|
||||
n_bins = max(1, int(np.ceil(t_end - t0)) + 1)
|
||||
else:
|
||||
frac_steady = pool_frac
|
||||
wait_steady = total_waiting
|
||||
t0, n_bins = None, None
|
||||
|
||||
return {
|
||||
"snaps": snaps,
|
||||
"times": times,
|
||||
"total_used": total_used,
|
||||
"pool_frac": pool_frac,
|
||||
"total_waiting": total_waiting,
|
||||
"total_running": total_running,
|
||||
times, total_used, pool_frac, total_waiting, total_running = cluster_timeline(
|
||||
snaps, t0=t0, n_bins=n_bins
|
||||
)
|
||||
n = len(times)
|
||||
|
||||
out = {
|
||||
"times": times.tolist(),
|
||||
"total_used": total_used.tolist(),
|
||||
"pool_frac": pool_frac.tolist(),
|
||||
"total_waiting": total_waiting.tolist(),
|
||||
"total_running": total_running.tolist(),
|
||||
"peak_pool_frac": float(pool_frac.max()) if n else 0.0,
|
||||
"steady_pool_frac": float(np.median(frac_steady)) if n else 0.0,
|
||||
"steady_pool_frac": _steady_median(pool_frac),
|
||||
"peak_waiting": int(total_waiting.max()) if n else 0,
|
||||
"summary": summary,
|
||||
}
|
||||
|
||||
# Per-role (P-pool vs D-pool) split for PD configs.
|
||||
roles = load_pid_roles(snaps_dir.parent / "vllm_logs")
|
||||
p_pids = {pid for pid, r in roles.items() if r == "P"}
|
||||
d_pids = {pid for pid, r in roles.items() if r == "D"}
|
||||
if p_pids and d_pids:
|
||||
for tag, subset in (("p", p_pids), ("d", d_pids)):
|
||||
_, _, frac, _, run = cluster_timeline(
|
||||
snaps, keep_pids=subset, t0=t0, n_bins=n_bins
|
||||
)
|
||||
out[f"{tag}_pool_frac"] = frac.tolist()
|
||||
out[f"{tag}_running"] = run.tolist()
|
||||
out[f"{tag}_peak_frac"] = float(frac.max()) if len(frac) else 0.0
|
||||
out[f"{tag}_steady_frac"] = _steady_median(frac)
|
||||
return out
|
||||
|
||||
|
||||
def collect_sweep(sweep_root: Path, tag: str, configs: list[str], reps: int) -> dict:
|
||||
"""Returns {config: [run_record_per_rep]}."""
|
||||
@@ -169,6 +236,10 @@ def collect_sweep(sweep_root: Path, tag: str, configs: list[str], reps: int) ->
|
||||
|
||||
|
||||
def plot_kv_timeline(sweep: dict, out: Path) -> None:
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
n_configs = len(sweep)
|
||||
if n_configs == 0:
|
||||
return
|
||||
@@ -177,8 +248,8 @@ def plot_kv_timeline(sweep: dict, out: Path) -> None:
|
||||
axes = [axes]
|
||||
for ax, (config, reps) in zip(axes, sweep.items()):
|
||||
for rep_data in reps:
|
||||
t = rep_data["times"]
|
||||
ax.plot(t, rep_data["pool_frac"] * 100, alpha=0.4, lw=1.0,
|
||||
t = np.asarray(rep_data["times"])
|
||||
ax.plot(t, np.asarray(rep_data["pool_frac"]) * 100, alpha=0.4, lw=1.0,
|
||||
label=f"rep{rep_data['rep']}")
|
||||
# bold median across reps (need to align times — use longest series)
|
||||
if reps:
|
||||
@@ -203,6 +274,10 @@ def plot_kv_timeline(sweep: dict, out: Path) -> None:
|
||||
|
||||
|
||||
def plot_peak_utilization(sweep: dict, out: Path) -> None:
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
configs = list(sweep.keys())
|
||||
peaks = [[r["peak_pool_frac"] * 100 for r in sweep[c]] for c in configs]
|
||||
steady = [[r["steady_pool_frac"] * 100 for r in sweep[c]] for c in configs]
|
||||
@@ -234,6 +309,10 @@ def plot_peak_utilization(sweep: dict, out: Path) -> None:
|
||||
|
||||
|
||||
def plot_latency_compare(sweep: dict, out: Path) -> None:
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
configs = list(sweep.keys())
|
||||
metrics = ["p50", "p90", "p99"]
|
||||
data = {m: [] for m in metrics}
|
||||
@@ -264,6 +343,50 @@ def plot_latency_compare(sweep: dict, out: Path) -> None:
|
||||
print(f"wrote {out}")
|
||||
|
||||
|
||||
def plot_role_split(sweep: dict, out: Path) -> None:
|
||||
"""For PD configs, show P-pool vs D-pool KV % over time (rep1) — exposes
|
||||
the imbalance that the cluster average hides. 8C (no role split) shows
|
||||
the overall cluster line for reference."""
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
n_configs = len(sweep)
|
||||
if n_configs == 0:
|
||||
return
|
||||
fig, axes = plt.subplots(n_configs, 1, figsize=(14, 2.6 * n_configs), sharex=True)
|
||||
if n_configs == 1:
|
||||
axes = [axes]
|
||||
for ax, (config, reps) in zip(axes, sweep.items()):
|
||||
if not reps:
|
||||
continue
|
||||
r = reps[0] # rep1
|
||||
t = np.asarray(r["times"])
|
||||
if "p_pool_frac" in r and "d_pool_frac" in r:
|
||||
ax.plot(t, np.asarray(r["p_pool_frac"]) * 100, color="#4c72b0",
|
||||
lw=1.5, label="P-pool (prefill)")
|
||||
ax.plot(t, np.asarray(r["d_pool_frac"]) * 100, color="#c44e52",
|
||||
lw=1.5, label="D-pool (decode)")
|
||||
ax.plot(t, np.asarray(r["pool_frac"]) * 100, color="#999",
|
||||
lw=1.0, ls=":", label="cluster avg")
|
||||
else:
|
||||
ax.plot(t, np.asarray(r["pool_frac"]) * 100, color="#222",
|
||||
lw=1.5, label="cluster (kv_both)")
|
||||
ax.axhline(90, color="#444", ls="--", alpha=0.5, lw=1)
|
||||
ax.set_ylim(0, 105)
|
||||
ax.set_ylabel(f"{config}\nKV pool (%)")
|
||||
ax.grid(True, alpha=0.3)
|
||||
ax.legend(loc="upper right", fontsize=8, ncol=3)
|
||||
axes[-1].set_xlabel("wall-clock since first snapshot (s)")
|
||||
fig.suptitle("MB5: per-role KV pool utilization (P-pool vs D-pool), rep1",
|
||||
fontsize=12)
|
||||
fig.tight_layout()
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
fig.savefig(out, dpi=120)
|
||||
plt.close(fig)
|
||||
print(f"wrote {out}")
|
||||
|
||||
|
||||
def write_summary_csv(sweep: dict, out: Path) -> None:
|
||||
rows = []
|
||||
for config, reps in sweep.items():
|
||||
@@ -279,6 +402,10 @@ def write_summary_csv(sweep: dict, out: Path) -> None:
|
||||
"wall_clock_s": s.get("wall_clock_s"),
|
||||
"peak_pool_frac": r["peak_pool_frac"],
|
||||
"steady_pool_frac": r["steady_pool_frac"],
|
||||
"p_pool_peak_frac": r.get("p_peak_frac"),
|
||||
"p_pool_steady_frac": r.get("p_steady_frac"),
|
||||
"d_pool_peak_frac": r.get("d_peak_frac"),
|
||||
"d_pool_steady_frac": r.get("d_steady_frac"),
|
||||
"peak_waiting": r["peak_waiting"],
|
||||
"latency_p50_s": lat.get("p50"),
|
||||
"latency_p90_s": lat.get("p90"),
|
||||
@@ -299,24 +426,55 @@ def write_summary_csv(sweep: dict, out: Path) -> None:
|
||||
print(f"wrote {out} ({len(rows)} rows)")
|
||||
|
||||
|
||||
def render_all(sweep: dict, out_dir: Path) -> None:
|
||||
plot_kv_timeline(sweep, out_dir / "mb5_kv_timeline.png")
|
||||
plot_role_split(sweep, out_dir / "mb5_role_split.png")
|
||||
plot_peak_utilization(sweep, out_dir / "mb5_peak_utilization.png")
|
||||
plot_latency_compare(sweep, out_dir / "mb5_latency_compare.png")
|
||||
write_summary_csv(sweep, out_dir / "mb5_summary.csv")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument("--sweep-root", type=Path, required=True,
|
||||
p = argparse.ArgumentParser(
|
||||
description="MB5 aggregate. Two-stage: --reduce (numpy-only, runs on "
|
||||
"a serving host) dumps a compact JSON; --from-reduced "
|
||||
"(needs matplotlib) renders figures locally. Or run "
|
||||
"directly (raw snapshots -> figures) when both the data "
|
||||
"and matplotlib are local."
|
||||
)
|
||||
p.add_argument("--sweep-root", type=Path,
|
||||
help="dir containing ${tag}_${config}_rep${N}/ subdirs")
|
||||
p.add_argument("--tag", required=True)
|
||||
p.add_argument("--tag")
|
||||
p.add_argument("--configs", default="8C 6P+2D 4P+4D 2P+6D",
|
||||
help="space-separated config names")
|
||||
p.add_argument("--reps", type=int, default=3)
|
||||
p.add_argument("--out-dir", type=Path, default=Path("figs/mb5"))
|
||||
p.add_argument("--reduce-to", type=Path,
|
||||
help="numpy-only: write reduced sweep JSON here and exit "
|
||||
"(no plotting, no matplotlib needed)")
|
||||
p.add_argument("--from-reduced", type=Path,
|
||||
help="load a reduced sweep JSON (from --reduce-to) and "
|
||||
"render figures into --out-dir")
|
||||
args = p.parse_args()
|
||||
|
||||
if args.from_reduced:
|
||||
sweep = json.loads(args.from_reduced.read_text())
|
||||
render_all(sweep, args.out_dir)
|
||||
return
|
||||
|
||||
if not (args.sweep_root and args.tag):
|
||||
p.error("--sweep-root and --tag are required unless --from-reduced is given")
|
||||
|
||||
configs = args.configs.split()
|
||||
sweep = collect_sweep(args.sweep_root, args.tag, configs, args.reps)
|
||||
|
||||
plot_kv_timeline(sweep, args.out_dir / "mb5_kv_timeline.png")
|
||||
plot_peak_utilization(sweep, args.out_dir / "mb5_peak_utilization.png")
|
||||
plot_latency_compare(sweep, args.out_dir / "mb5_latency_compare.png")
|
||||
write_summary_csv(sweep, args.out_dir / "mb5_summary.csv")
|
||||
if args.reduce_to:
|
||||
args.reduce_to.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.reduce_to.write_text(json.dumps(sweep))
|
||||
print(f"wrote reduced sweep -> {args.reduce_to}")
|
||||
return
|
||||
|
||||
render_all(sweep, args.out_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user