From e9abd70c8d18e836c866c9dc480616a9849ce011 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Wed, 27 May 2026 23:02:57 +0800 Subject: [PATCH] MB5 driver: launcher, orchestrator, KV-pool timeline plotter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three new files to drive the PD ratio sweep + per-request KV occupancy capture, plus a deploy.sh update so the patched replayer rides along to the fresh-venv host. mb5_launch.sh One script handles all four configs we plan to sweep: CONFIG=8C / 6P+2D / 4P+4D / 2P+6D - For 8C: 8 vLLM instances with kv_role=kv_both on GPU 0-7. Replayer talks to them via the existing comma-separated round-robin in replayer/replay.py — no proxy. - For PD configs: kv_role=kv_producer for the P pool (with VLLM_MOONCAKE_BOOTSTRAP_PORT) + kv_role=kv_consumer for the D pool, routed by the official vLLM example third_party/vllm/examples/online_serving/disaggregated_serving/ mooncake_connector/mooncake_connector_proxy.py — no policy choice made by us, per user instruction to use the standard recipe. - Applies instrument_kv_snapshot.py before launching so every EngineCore writes its per-step KV snapshot to $RUN_ROOT/kv_snapshots/mb5_kv_snapshot_pid.jsonl - Reverts the patch on stop. - Emits ENDPOINTS= line on stdout for the orchestrator to read. mb5_run.sh For each CONFIG × rep: launch, replay w600 trace via the existing replayer, capture wall-clock, tear down, cool down 10 s. Defaults: CONFIGS="8C 6P+2D 4P+4D 2P+6D" REPS=3 TRACE=traces/w600_r0.0015_st30.jsonl All artefacts go under $FRESH_ROOT/mb5_runs/$RUN_TAG_${config}_rep${rep}/ (vllm_logs/, kv_snapshots/, replay_metrics.jsonl, wall_clock_s.txt). plot_kv_pool_timeline.py Reads one or more mb5_kv_snapshot_pid*.jsonl files and renders a stacked-area chart per file: x = wall-clock since first snapshot y = KV block count, stacked by per-request contribution overlay: pool-total ceiling, 90% line, waiting-queue depth subplot Bands are colored by a deterministic hash of request_id so individual requests are visually tractable across the run. This is the figure the user asked for — turns headline "PD-disagg is 10× worse" into a system-level picture of *where* the KV pool is blocked, when, and by which requests. deploy.sh Also tar-syncs the local replayer/ dir to /home/admin/cpfs/wjh/agentic-kv-fresh/replayer/ so mb5_run.sh can `python -m replayer` against the patched (trace_span_s/amplification) version, not the older copy under /home/admin/cpfs/wjh/agentic-kv/. Co-Authored-By: Claude Opus 4.7 --- microbench/fresh_setup/deploy.sh | 15 +- microbench/fresh_setup/mb5_launch.sh | 201 ++++++++++++++++++ microbench/fresh_setup/mb5_run.sh | 120 +++++++++++ .../fresh_setup/plot_kv_pool_timeline.py | 141 ++++++++++++ 4 files changed, 475 insertions(+), 2 deletions(-) create mode 100755 microbench/fresh_setup/mb5_launch.sh create mode 100755 microbench/fresh_setup/mb5_run.sh create mode 100644 microbench/fresh_setup/plot_kv_pool_timeline.py diff --git a/microbench/fresh_setup/deploy.sh b/microbench/fresh_setup/deploy.sh index 9a3c874..7cba4ff 100755 --- a/microbench/fresh_setup/deploy.sh +++ b/microbench/fresh_setup/deploy.sh @@ -12,11 +12,22 @@ SRC="${SCRIPT_DIR}/" DEST_HOST="${1:-dash1}" DEST="/home/admin/cpfs/wjh/agentic-kv-fresh/scripts/" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" +REPLAYER_SRC="${REPO_ROOT}/replayer" +REPLAYER_DEST="/home/admin/cpfs/wjh/agentic-kv-fresh/replayer/" + echo "[deploy] syncing ${SRC} -> ${DEST_HOST}:${DEST}" -ssh "${DEST_HOST}" "mkdir -p ${DEST}" +ssh "${DEST_HOST}" "mkdir -p ${DEST} ${REPLAYER_DEST}" # dash1/2 don't have rsync; use tar over ssh. tar -C "${SCRIPT_DIR}" --exclude='__pycache__' --exclude='*.pyc' \ --exclude='deploy.sh' -czf - . \ | ssh "${DEST_HOST}" "cd ${DEST} && tar -xzf -" + +if [ -d "${REPLAYER_SRC}" ]; then + echo "[deploy] syncing ${REPLAYER_SRC}/ -> ${DEST_HOST}:${REPLAYER_DEST}" + tar -C "${REPLAYER_SRC}" --exclude='__pycache__' --exclude='*.pyc' -czf - . \ + | ssh "${DEST_HOST}" "cd ${REPLAYER_DEST} && tar -xzf -" +fi + echo "[deploy] done" -ssh "${DEST_HOST}" "ls -la ${DEST}" +ssh "${DEST_HOST}" "ls -la ${DEST}; echo '---replayer---'; ls -la ${REPLAYER_DEST}" diff --git a/microbench/fresh_setup/mb5_launch.sh b/microbench/fresh_setup/mb5_launch.sh new file mode 100755 index 0000000..03edfd2 --- /dev/null +++ b/microbench/fresh_setup/mb5_launch.sh @@ -0,0 +1,201 @@ +#!/usr/bin/env bash +# Launch vLLM instances + (optional) proxy for one MB5 config. +# +# CONFIG=8C : 8 kv_both instances on GPU 0-7 (replayer +# talks directly via round-robin) +# CONFIG=6P+2D : 6 kv_producer (GPU 0-5) + 2 kv_consumer (GPU 6-7) +# CONFIG=4P+4D : 4 producer (GPU 0-3) + 4 consumer (GPU 4-7) +# CONFIG=2P+6D : 2 producer (GPU 0-1) + 6 consumer (GPU 2-7) +# +# All configs use the fresh venv (vanilla vLLM 0.18.1 + Mooncake 0.3.11), +# kv_both/kv_producer/kv_consumer roles, MB5 scheduler instrumentation +# applied. PD configs are launched per the official vLLM example +# (run_mooncake_connector.sh) — round-robin P / round-robin D via +# mooncake_connector_proxy.py on PROXY_PORT. +# +# Usage: +# CONFIG=4P+4D RUN_LABEL=run1 bash mb5_launch.sh start +# bash mb5_launch.sh status +# bash mb5_launch.sh stop +# +# After "start", grep "ENDPOINTS=" from this script's output to get the +# URL(s) the replayer should target. + +set -eo pipefail + +FRESH_ROOT="/home/admin/cpfs/wjh/agentic-kv-fresh" +VENV="${FRESH_ROOT}/.venv" +MODEL="${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +INSTRUMENT="${SCRIPT_DIR}/instrument_kv_snapshot.py" +PROXY_SRC="${SCRIPT_DIR}/../../third_party/vllm/examples/online_serving/disaggregated_serving/mooncake_connector/mooncake_connector_proxy.py" + +CONFIG="${CONFIG:-8C}" +RUN_LABEL="${RUN_LABEL:-default}" + +# All artefacts for this run live here +RUN_ROOT="${FRESH_ROOT}/mb5_runs/${RUN_LABEL}_${CONFIG}" +LOGS_DIR="${RUN_ROOT}/vllm_logs" +SNAPSHOT_DIR="${RUN_ROOT}/kv_snapshots" + +PROXY_PORT="${PROXY_PORT:-8100}" +BASE_HTTP=8000 +BASE_BP=8998 +BASE_MASTER=29500 + +stop_all() { + pkill -9 -f "mooncake_connector_proxy.py" 2>/dev/null || true + pkill -9 -f "vllm serve" 2>/dev/null || true + pkill -9 -f "EngineCore" 2>/dev/null || true + sleep 3 +} + +case "${1:-start}" in + stop) + stop_all + python "${INSTRUMENT}" --revert --venv "${VENV}" 2>/dev/null || true + exit 0;; + status) + ports=() + case "${CONFIG}" in + 8C) for i in 0 1 2 3 4 5 6 7; do ports+=( $((BASE_HTTP+i)) ); done ;; + *) ports=( ${PROXY_PORT} ) ;; + esac + for p in "${ports[@]}"; do + if curl -sf "http://127.0.0.1:${p}/health" >/dev/null 2>&1 \ + || curl -sf "http://127.0.0.1:${p}/v1/models" >/dev/null 2>&1; then + echo "port ${p}: UP" + else + echo "port ${p}: DOWN" + fi + done + exit 0;; + start) ;; + *) echo "Unknown command: $1"; exit 1;; +esac + +# --- parse CONFIG into (prefill_gpus, decode_gpus) ---------------- +case "${CONFIG}" in + 8C) ROLES="combined"; P_GPUS=""; D_GPUS=""; COMBINED_GPUS="0,1,2,3,4,5,6,7" ;; + 6P+2D) ROLES="pd"; P_GPUS="0,1,2,3,4,5"; D_GPUS="6,7" ;; + 4P+4D) ROLES="pd"; P_GPUS="0,1,2,3"; D_GPUS="4,5,6,7" ;; + 2P+6D) ROLES="pd"; P_GPUS="0,1"; D_GPUS="2,3,4,5,6,7" ;; + *) echo "Unknown CONFIG=${CONFIG} (expected: 8C, 6P+2D, 4P+4D, 2P+6D)"; exit 1;; +esac + +stop_all +mkdir -p "${LOGS_DIR}" "${SNAPSHOT_DIR}" +source "${VENV}/bin/activate" + +# Apply MB5 patch (idempotent — affects entire shared cpfs venv). +python "${INSTRUMENT}" --apply --venv "${VENV}" + +launch_vllm() { + # $1 idx, $2 gpu, $3 port, $4 role, $5 bp_or_dash + local idx="$1" gpu="$2" port="$3" role="$4" bp="$5" + local cfg="{\"kv_connector\":\"MooncakeConnector\",\"kv_role\":\"${role}\"}" + local master=$((BASE_MASTER+idx)) + local env_bp_arg="" + if [ -n "${bp}" ] && [ "${bp}" != "-" ]; then + env_bp_arg="VLLM_MOONCAKE_BOOTSTRAP_PORT=${bp}" + fi + echo "[mb5] launching idx=${idx} gpu=${gpu} port=${port} role=${role} bp=${bp:-none}" + PYTHONHASHSEED=42 \ + MB5_LOG_DIR="${SNAPSHOT_DIR}" \ + CUDA_VISIBLE_DEVICES="${gpu}" \ + MASTER_PORT="${master}" \ + ${env_bp_arg} \ + nohup vllm serve "${MODEL}" \ + --host 0.0.0.0 --port "${port}" \ + --tensor-parallel-size 1 \ + --trust-remote-code --enable-prefix-caching \ + --dtype auto --gpu-memory-utilization 0.9 \ + --max-model-len 200000 \ + --max-num-batched-tokens 8192 \ + --kv-transfer-config "${cfg}" \ + --enable-prompt-tokens-details \ + > "${LOGS_DIR}/vllm_idx${idx}_gpu${gpu}_${role}.log" 2>&1 & + disown +} + +idx=0 +proxy_args=() +ENDPOINTS="" + +case "${ROLES}" in + combined) + IFS=',' read -ra GPUS <<< "${COMBINED_GPUS}" + for gpu in "${GPUS[@]}"; do + port=$((BASE_HTTP+idx)) + bp=$((BASE_BP+idx)) + launch_vllm "${idx}" "${gpu}" "${port}" "kv_both" "${bp}" + ENDPOINTS+="${ENDPOINTS:+,}http://127.0.0.1:${port}" + idx=$((idx+1)) + sleep 1 + done + ;; + pd) + # producers + IFS=',' read -ra PG <<< "${P_GPUS}" + for gpu in "${PG[@]}"; do + port=$((BASE_HTTP+idx)) + bp=$((BASE_BP+idx)) + launch_vllm "${idx}" "${gpu}" "${port}" "kv_producer" "${bp}" + proxy_args+=( --prefill "http://127.0.0.1:${port}" "${bp}" ) + idx=$((idx+1)) + sleep 1 + done + # consumers + IFS=',' read -ra DG <<< "${D_GPUS}" + for gpu in "${DG[@]}"; do + port=$((BASE_HTTP+idx)) + launch_vllm "${idx}" "${gpu}" "${port}" "kv_consumer" "-" + proxy_args+=( --decode "http://127.0.0.1:${port}" ) + idx=$((idx+1)) + sleep 1 + done + ;; +esac + +echo "[mb5] waiting for all vllm /v1/models endpoints..." +all_ports=() +for ((p=0; p/dev/null 2>&1; do + tries=$((tries+1)) + if [ ${tries} -gt 300 ]; then + echo "[mb5] FATAL port ${port} did not come up in 10 min" + tail -40 "${LOGS_DIR}/vllm_idx"*"_gpu"*".log" || true + exit 1 + fi + sleep 2 + done + echo " port=${port} ready" +done + +if [ "${ROLES}" = "pd" ]; then + echo "[mb5] launching mooncake_connector_proxy on ${PROXY_PORT}" + nohup python "${PROXY_SRC}" "${proxy_args[@]}" --port "${PROXY_PORT}" --host 0.0.0.0 \ + > "${LOGS_DIR}/proxy.log" 2>&1 & + disown + # wait for proxy + tries=0 + while ! curl -sf "http://127.0.0.1:${PROXY_PORT}/v1/models" >/dev/null 2>&1; do + tries=$((tries+1)) + if [ ${tries} -gt 60 ]; then + echo "[mb5] FATAL proxy did not come up in 2 min" + tail -40 "${LOGS_DIR}/proxy.log" || true + exit 1 + fi + sleep 2 + done + echo " proxy port=${PROXY_PORT} ready" + ENDPOINTS="http://127.0.0.1:${PROXY_PORT}" +fi + +echo "[mb5] CONFIG=${CONFIG} RUN_LABEL=${RUN_LABEL} UP" +echo "ENDPOINTS=${ENDPOINTS}" +echo "RUN_ROOT=${RUN_ROOT}" +echo "SNAPSHOT_DIR=${SNAPSHOT_DIR}" +echo "VLLM_LOGS=${LOGS_DIR}" diff --git a/microbench/fresh_setup/mb5_run.sh b/microbench/fresh_setup/mb5_run.sh new file mode 100755 index 0000000..8e32d39 --- /dev/null +++ b/microbench/fresh_setup/mb5_run.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +# Orchestrator for MB5: for each CONFIG × rep, bring up the stack, run a +# trace replay against it, collect KV snapshots and replayer metrics, +# tear down. +# +# Designed to be run on dash1 (or any host with cpfs mounted at +# /home/admin/cpfs/wjh/). +# +# Env vars (with defaults): +# CONFIGS : space-separated MB5 configs (default: "8C 6P+2D 4P+4D 2P+6D") +# REPS : reps per config (default: 3) +# TRACE : trace JSONL path +# (default: /home/admin/cpfs/wjh/agentic-kv/traces/w600_r0.0015_st30.jsonl) +# RUN_TAG : output root tag (default: $(date +%Y%m%d_%H%M%S)) +# REQUEST_LIMIT : optional, cap replay requests (default: none) + +set -eo pipefail + +FRESH_ROOT="/home/admin/cpfs/wjh/agentic-kv-fresh" +VENV="${FRESH_ROOT}/.venv" +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +LAUNCH="${SCRIPT_DIR}/mb5_launch.sh" +REPLAYER_DIR="${FRESH_ROOT}/replayer" + +CONFIGS="${CONFIGS:-8C 6P+2D 4P+4D 2P+6D}" +REPS="${REPS:-3}" +TRACE="${TRACE:-/home/admin/cpfs/wjh/agentic-kv/traces/w600_r0.0015_st30.jsonl}" +RUN_TAG="${RUN_TAG:-$(date +%Y%m%d_%H%M%S)}" +MODEL_NAME="${MODEL_NAME:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}" +REQUEST_LIMIT_ARG="" +if [ -n "${REQUEST_LIMIT:-}" ]; then + REQUEST_LIMIT_ARG="--request-limit ${REQUEST_LIMIT}" +fi + +OUT_ROOT="${FRESH_ROOT}/mb5_runs/${RUN_TAG}" +mkdir -p "${OUT_ROOT}" +echo "[mb5-run] RUN_TAG=${RUN_TAG}" +echo "[mb5-run] OUT_ROOT=${OUT_ROOT}" +echo "[mb5-run] CONFIGS=${CONFIGS}" +echo "[mb5-run] REPS=${REPS}" +echo "[mb5-run] TRACE=${TRACE}" + +run_one() { + local config="$1" rep="$2" + local label="${RUN_TAG}_${config}_rep${rep}" + local rundir="${FRESH_ROOT}/mb5_runs/${label}" + echo "" + echo "======== ${config} rep${rep} ========" + + # Launch + if ! CONFIG="${config}" RUN_LABEL="${RUN_TAG}_${config}_rep${rep}" \ + bash "${LAUNCH}" start > "${OUT_ROOT}/${config}_rep${rep}_launch.log" 2>&1; then + echo "[mb5-run] LAUNCH FAILED for ${config} rep${rep}; see ${OUT_ROOT}/${config}_rep${rep}_launch.log" + return 1 + fi + + # Extract ENDPOINTS line emitted by mb5_launch.sh + local endpoints + endpoints=$(grep "^ENDPOINTS=" "${OUT_ROOT}/${config}_rep${rep}_launch.log" | tail -1 | cut -d= -f2-) + if [ -z "${endpoints}" ]; then + echo "[mb5-run] ERROR: no ENDPOINTS in launch log" + bash "${LAUNCH}" stop > /dev/null 2>&1 || true + return 1 + fi + echo "[mb5-run] endpoints: ${endpoints}" + + # Replay + source "${VENV}/bin/activate" + local replay_out="${rundir}/replay_metrics.jsonl" + mkdir -p "$(dirname "${replay_out}")" + local t0 + t0=$(date +%s.%N) + if ! PYTHONPATH="${FRESH_ROOT}" python -m replayer \ + --endpoint-url "${endpoints}" \ + --trace-path "${TRACE}" \ + --output-path "${replay_out}" \ + --model-name "${MODEL_NAME}" \ + ${REQUEST_LIMIT_ARG} \ + > "${OUT_ROOT}/${config}_rep${rep}_replay.log" 2>&1; then + local t1=$(date +%s.%N) + echo "[mb5-run] REPLAY FAILED after $(echo "$t1 - $t0" | bc) s; see ${OUT_ROOT}/${config}_rep${rep}_replay.log" + bash "${LAUNCH}" stop > /dev/null 2>&1 || true + return 1 + fi + local t1=$(date +%s.%N) + local wall_clock_s=$(echo "$t1 - $t0" | bc) + echo "[mb5-run] replay done in ${wall_clock_s}s" + echo "${wall_clock_s}" > "${rundir}/wall_clock_s.txt" + + # Stop launch (cleans up vllm + proxy; reverts patch on last call) + CONFIG="${config}" RUN_LABEL="${RUN_TAG}_${config}_rep${rep}" \ + bash "${LAUNCH}" stop > "${OUT_ROOT}/${config}_rep${rep}_stop.log" 2>&1 || true + + sleep 10 # cooldown so GPUs settle before next config + echo "[mb5-run] DONE ${config} rep${rep}" +} + +# Quick check that the launch script and replayer are reachable +if [ ! -f "${LAUNCH}" ]; then echo "missing ${LAUNCH}"; exit 1; fi +if [ ! -d "${REPLAYER_DIR}" ]; then echo "missing ${REPLAYER_DIR}"; exit 1; fi +if [ ! -f "${TRACE}" ]; then echo "missing trace ${TRACE}"; exit 1; fi + +# Iterate +failures=0 +for config in ${CONFIGS}; do + for ((rep=1; rep<=REPS; rep++)); do + if ! run_one "${config}" "${rep}"; then + failures=$((failures+1)) + fi + done +done + +# Final patch revert (defensive — mb5_launch.sh stop also reverts) +python "${SCRIPT_DIR}/instrument_kv_snapshot.py" --revert --venv "${VENV}" 2>/dev/null || true + +echo "" +echo "======== ALL CONFIGS DONE ========" +echo "failures: ${failures}" +echo "results under: ${FRESH_ROOT}/mb5_runs/${RUN_TAG}_*" +echo "to plot: python plot_kv_pool_timeline.py --run-tag ${RUN_TAG}" diff --git a/microbench/fresh_setup/plot_kv_pool_timeline.py b/microbench/fresh_setup/plot_kv_pool_timeline.py new file mode 100644 index 0000000..3dc2304 --- /dev/null +++ b/microbench/fresh_setup/plot_kv_pool_timeline.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +"""Plot per-instance KV-pool composition over time from MB5 snapshots. + +Input: a directory of mb5_kv_snapshot_pid.jsonl files (one per +EngineCore PID) — typically MB5_LOG_DIR set during a run. + +For each snapshot file, builds a stacked-area chart: + x axis = wall-clock time since first snapshot + y axis = KV block count (total pool size = ceiling line) + stacked bands = per-request block usage; each request gets one + band colored by a hash-based palette so the eye + can track individual requests across the run. + +Also overlays: + - free_blocks (white area at the top) + - 90% capacity line (red dashed) + - waiting-queue depth (separate small subplot beneath) +""" +from __future__ import annotations + +import argparse +import hashlib +import json +from collections import defaultdict +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np + + +def load_snapshots(path: Path) -> list[dict]: + out = [] + with path.open() as f: + for line in f: + line = line.strip() + if not line: + continue + try: + out.append(json.loads(line)) + except json.JSONDecodeError: + continue + return out + + +def req_color(req_id: str) -> str: + """Deterministic color per request_id.""" + h = hashlib.md5(req_id.encode()).hexdigest() + return f"#{h[:6]}" + + +def plot_one_instance(snaps: list[dict], out: Path, title: str) -> None: + if not snaps: + return + snaps = sorted(snaps, key=lambda s: s["t_unix"]) + t0 = snaps[0]["t_unix"] + times = [s["t_unix"] - t0 for s in snaps] + total_blocks = snaps[0]["total_blocks"] + + # Build the request × time block-count matrix + all_req_ids: list[str] = [] + req_first_seen: dict[str, float] = {} + for s in snaps: + for r in s.get("running", []): + rid = r["req_id"] + if rid not in req_first_seen: + req_first_seen[rid] = s["t_unix"] - t0 + all_req_ids.append(rid) + # Sort by first-seen time so the band order follows arrival + all_req_ids.sort(key=lambda r: req_first_seen[r]) + + matrix = np.zeros((len(all_req_ids), len(times)), dtype=np.int64) + req_to_row = {r: i for i, r in enumerate(all_req_ids)} + for j, s in enumerate(snaps): + for r in s.get("running", []): + i = req_to_row[r["req_id"]] + matrix[i, j] = r.get("n_blocks", 0) + + fig, (ax1, ax2) = plt.subplots( + 2, 1, figsize=(13, 6), + sharex=True, + gridspec_kw={"height_ratios": [4, 1]}, + ) + + colors = [req_color(r) for r in all_req_ids] + ax1.stackplot(times, matrix, colors=colors, linewidth=0) + ax1.axhline(total_blocks, color="#444", lw=1.5, ls="-", + label=f"pool total = {total_blocks} blocks") + ax1.axhline(total_blocks * 0.9, color="#c44e52", lw=1.2, ls="--", alpha=0.7, + label="90% capacity") + ax1.set_ylabel("KV blocks (per-request stacked)") + ax1.set_ylim(0, total_blocks * 1.05) + ax1.set_title(title) + ax1.legend(loc="upper right", fontsize=9, framealpha=0.95) + ax1.grid(True, alpha=0.3) + + # Waiting queue depth subplot + wait_lens = [len(s.get("waiting", [])) for s in snaps] + ax2.fill_between(times, 0, wait_lens, color="#c44e52", alpha=0.55, + label="waiting requests") + ax2.set_ylabel("queue\ndepth") + ax2.set_xlabel("wall-clock since first snapshot (s)") + ax2.set_ylim(0, max(max(wait_lens, default=0), 1) * 1.2 + 1) + ax2.grid(True, alpha=0.3) + ax2.legend(loc="upper right", fontsize=8) + + out.parent.mkdir(parents=True, exist_ok=True) + fig.tight_layout() + fig.savefig(out, dpi=120) + plt.close(fig) + print(f"wrote {out} (n_snapshots={len(snaps)}, n_requests={len(all_req_ids)})") + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("--snapshot-dir", type=Path, required=True, + help="dir containing mb5_kv_snapshot_pid*.jsonl files") + p.add_argument("--out-dir", type=Path, required=True) + p.add_argument("--label", default="", + help="prefix for output filenames + figure title") + args = p.parse_args() + + files = sorted(args.snapshot_dir.glob("mb5_kv_snapshot_pid*.jsonl")) + if not files: + print(f"[plot] no snapshot files in {args.snapshot_dir}") + return + + for f in files: + pid = f.stem.replace("mb5_kv_snapshot_pid", "") + snaps = load_snapshots(f) + if not snaps: + print(f"[plot] {f.name}: empty") + continue + out = args.out_dir / f"{args.label}_pid{pid}.png" + title = f"{args.label} pid={pid} (n_snap={len(snaps)})" + plot_one_instance(snaps, out, title) + + +if __name__ == "__main__": + main()