MB5 driver: launcher, orchestrator, KV-pool timeline plotter

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<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 <noreply@anthropic.com>
This commit is contained in:
2026-05-27 23:02:57 +08:00
parent a4f5dd56aa
commit e9abd70c8d
4 changed files with 475 additions and 2 deletions

120
microbench/fresh_setup/mb5_run.sh Executable file
View File

@@ -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}"