gahow/agentic-kvc
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

MB2: per-stage instrumentation patch + launcher integration

Browse Source 
Per-stage breakdown of "step 2" (the B-side do_remote_prefill) requires
vLLM/mooncake-internal timing — we cannot infer it from black-box HTTP
E2E. This commit adds the four pieces to do that breakdown:

instrument_mooncake.py
  apply / revert / check patches on mooncake_connector_v1.py to emit
  structured JSONL transfer events at two key sites:

    send_blocks (P-side, on batch_transfer_sync_write):
      {event, remote_session, total_bytes, duration_s, t_start_unix,
       ret, tp_rank, t_log_unix}
    receive_kv (D-side, on the ZMQ-driven pull request):
      {event, path, local_req_ids, remote_req_ids, duration_s,
       t_start_unix, tp_rank, t_log_unix}

  All injected code is bracketed by `# MB2_INSTRUMENT_START/END` so the
  --revert pass is a single regex scan. Apply-revert round-trip
  validated on dash1 (PATCHED → py_compile ok → revert → CLEAN → ok).

start_vllm_pair.sh (updated)
  - Picks up instrument_mooncake.py via SCRIPT_DIR.
  - On `start`: applies patch before launching the two vLLM instances.
  - On `stop` (or trap exit): reverts patch.
  - Sets per-instance MB2_LOG_DIR = $FRESH_ROOT/mb2_transfer_logs/{A,B}/
    so send-side and receive-side events land in cleanly separated dirs.

deploy.sh
  tar-over-ssh sync of microbench/fresh_setup/ → cpfs
  /home/admin/cpfs/wjh/agentic-kv-fresh/scripts/ so dash1 / dash2 see
  the same scripts (dash{1,2} don't have rsync; tar pipe works).

The mb2_kv_transfer.py client still uses black-box E2E timing — the
next commit will teach it to ingest the per-instance JSONL logs to
produce the 4-way breakdown (queueing / setup / transfer / decode).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-27 18:12:44 +08:00
parent 7437422618
commit efdcf3c555
3 changed files with 259 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
microbench/fresh_setup/start_vllm_pair.sh
View File

@@ -30,10 +30,17 @@ BP_B=8999
MASTER_A=29500
MASTER_B=29501
MB2_LOG_ROOT="${FRESH_ROOT}/mb2_transfer_logs"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
INSTRUMENT="${SCRIPT_DIR}/instrument_mooncake.py"
stop_all() {
pkill -9 -f "vllm serve" 2>/dev/null || true
pkill -9 -f "EngineCore" 2>/dev/null || true
sleep 2
if [[ -f "${INSTRUMENT}" ]]; then
python "${INSTRUMENT}" --revert --venv "${VENV}" 2>/dev/null || true
fi
}
case "${1:-start}" in
@@ -61,6 +68,15 @@ stop_all
source "${VENV}/bin/activate"
if [[ -f "${INSTRUMENT}" ]]; then
echo "[mb2] applying instrumentation patch"
python "${INSTRUMENT}" --apply --venv "${VENV}"
else
echo "[mb2] WARN instrument_mooncake.py not found at ${INSTRUMENT}; transfer logs will be absent"
fi
mkdir -p "${MB2_LOG_ROOT}/A" "${MB2_LOG_ROOT}/B"
launch() {
local idx="$1" gpu="$2" port="$3" bp="$4" master="$5"
echo "[mb2] launching instance ${idx} on GPU ${gpu}, port ${port}, bp ${bp}"
@@ -68,6 +84,7 @@ launch() {
VLLM_MOONCAKE_BOOTSTRAP_PORT="${bp}" \
CUDA_VISIBLE_DEVICES="${gpu}" \
MASTER_PORT="${master}" \
MB2_LOG_DIR="${MB2_LOG_ROOT}/${idx}" \
nohup vllm serve "${MODEL}" \
--host 0.0.0.0 --port "${port}" \
--tensor-parallel-size 1 \
@@ -100,6 +117,6 @@ for port in "${PORT_A}" "${PORT_B}"; do
done
echo "[mb2] both instances UP"
echo " A: 127.0.0.1:${PORT_A} (GPU ${GPU_A}, bp ${BP_A})"
echo " B: 127.0.0.1:${PORT_B} (GPU ${GPU_B}, bp ${BP_B})"
echo " logs: ${LOGS_DIR}"
echo " A: 127.0.0.1:${PORT_A} (GPU ${GPU_A}, bp ${BP_A}, log_dir ${MB2_LOG_ROOT}/A)"
echo " B: 127.0.0.1:${PORT_B} (GPU ${GPU_B}, bp ${BP_B}, log_dir ${MB2_LOG_ROOT}/B)"
echo " vllm stdout: ${LOGS_DIR}"