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:
@@ -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}"
|
||||
|
||||
201
microbench/fresh_setup/mb5_launch.sh
Executable file
201
microbench/fresh_setup/mb5_launch.sh
Executable file
@@ -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<idx; p++)); do all_ports+=( $((BASE_HTTP+p)) ); done
|
||||
for port in "${all_ports[@]}"; do
|
||||
tries=0
|
||||
while ! curl -sf "http://127.0.0.1:${port}/v1/models" >/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}"
|
||||
120
microbench/fresh_setup/mb5_run.sh
Executable file
120
microbench/fresh_setup/mb5_run.sh
Executable 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}"
|
||||
141
microbench/fresh_setup/plot_kv_pool_timeline.py
Normal file
141
microbench/fresh_setup/plot_kv_pool_timeline.py
Normal file
@@ -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<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()
|
||||
Reference in New Issue
Block a user