Gahow Wang
c6b7c3471b
Three additions land together because B3's whole point is comparing
LMetric against meaningful controls.
- scripts/cache_aware_proxy.py: two new --policy values.
- load_only: pure min(num_requests) routing, no cache or affinity.
The B3 control that strips locality so the LMetric-vs-load gap is
legible.
- sticky: first turn goes to min-load, subsequent turns ALWAYS
return to the same instance, even under saturation. The B3
control that maxes out locality so the hot-spot cost is legible.
- scripts/build_capped_trace.py: per-session turn cap (default 8).
Generates the session-mass-equalized variant the TODO calls for so
that hot-spot index can be re-measured with the heavy-tail removed.
- scripts/b3_sweep.sh: orchestrates the 5-cell sweep.
- GPU_INDICES makes it easy to skip a dead GPU.
- EXTRA_VLLM_ARGS defaults to --enable-prompt-tokens-details so
usage.prompt_tokens_details.cached_tokens is populated. vLLM
0.18.1 omits the field by default and breaks the reuse-decomp
pipeline; the smoke run surfaced this.
- Trap kills EngineCore by name in addition to "vllm serve" — the
parent dies first but the child holds GPU memory. Was the root
cause of the 89 GB ghost on GPU 0 earlier today.
- Proxy readiness is a polling loop, not a fixed sleep.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
178 lines
6.1 KiB
Bash
Executable File
178 lines
6.1 KiB
Bash
Executable File
#!/usr/bin/env bash
|
||
# B3 routing sweep: 5 policies on 8x TP1 instances with full instrumentation.
|
||
#
|
||
# Policies:
|
||
# lmetric — cache-aware P_tokens × BS routing (main baseline)
|
||
# load_only — pure min-num_requests (B3 control: no cache)
|
||
# sticky — hard session affinity (B3 control: perfect locality)
|
||
# unified — hybrid affinity + LMetric fallback
|
||
# capped — lmetric on a per-session turn-capped trace
|
||
#
|
||
# Each policy run produces metrics.jsonl + breakdown.json + worker_state.json
|
||
# + run_window.json (start/end unix timestamps so the analyzer can slice the
|
||
# shared engine_*.jsonl by time).
|
||
|
||
set -euo pipefail
|
||
|
||
ROOT="${ROOT:-/home/admin/cpfs/wjh/agentic-kv}"
|
||
VENV="$ROOT/.venv/bin"
|
||
MODEL="${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}"
|
||
TRACE="${TRACE:-$ROOT/traces/w600_r0.0015_st30.jsonl}"
|
||
OUTDIR="${OUTDIR:-$ROOT/outputs/b3_sweep_$(date +%Y%m%d_%H%M%S)}"
|
||
PROXY_PORT="${PROXY_PORT:-9300}"
|
||
BASE_PORT="${BASE_PORT:-8000}"
|
||
# Space-separated list of GPU indices to use, one vLLM instance per index.
|
||
# Override via GPU_INDICES="1 2 3 4 5 6 7" when GPU 0 holds ghost memory.
|
||
GPU_INDICES="${GPU_INDICES:-0 1 2 3 4 5 6 7}"
|
||
POLICIES="${POLICIES:-lmetric load_only sticky unified}"
|
||
MAX_TURNS_CAP="${MAX_TURNS_CAP:-8}"
|
||
EXTRA_VLLM_ARGS="${EXTRA_VLLM_ARGS:---enable-prompt-tokens-details}"
|
||
|
||
# Derive N_INSTANCES from GPU_INDICES
|
||
N_INSTANCES=$(echo $GPU_INDICES | wc -w)
|
||
|
||
mkdir -p "$OUTDIR/engine_state" "$OUTDIR/logs"
|
||
echo "[b3_sweep] OUTDIR=$OUTDIR"
|
||
|
||
cleanup() {
|
||
pkill -9 -f "vllm serve" 2>/dev/null || true
|
||
# vLLM spawns an EngineCore child whose process name is
|
||
# "VLLM::EngineCor" — pkill -f "vllm serve" misses it and leaves
|
||
# the GPU memory locked by a dead-but-tracked-by-driver context.
|
||
pkill -9 -f "EngineCore" 2>/dev/null || true
|
||
pkill -9 -f cache_aware_proxy 2>/dev/null || true
|
||
sleep 3
|
||
}
|
||
trap cleanup EXIT
|
||
|
||
# 1) Launch one vLLM per GPU index in GPU_INDICES; each emits engine_<i>.jsonl
|
||
launch_vllm() {
|
||
echo "[b3_sweep] launching $N_INSTANCES vLLM instances on GPUs $GPU_INDICES ..."
|
||
local i=0
|
||
for gpu in $GPU_INDICES; do
|
||
local port=$((BASE_PORT + i))
|
||
local master=$((29500 + i))
|
||
local log="$OUTDIR/logs/vllm_inst_${i}_gpu${gpu}.log"
|
||
AGENTIC_STEP_LOG_PATH="$OUTDIR/engine_state/engine_${i}.jsonl" \
|
||
AGENTIC_WORKER_ID="engine_${i}" \
|
||
CUDA_VISIBLE_DEVICES=$gpu \
|
||
MASTER_PORT=$master \
|
||
nohup "$VENV/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 \
|
||
$EXTRA_VLLM_ARGS \
|
||
> "$log" 2>&1 &
|
||
disown
|
||
sleep 2
|
||
i=$((i + 1))
|
||
done
|
||
|
||
echo "[b3_sweep] waiting for vLLM health ..."
|
||
for i in $(seq 0 $((N_INSTANCES - 1))); do
|
||
local port=$((BASE_PORT + i))
|
||
local tries=0
|
||
while ! curl -sf "http://127.0.0.1:$port/health" >/dev/null 2>&1; do
|
||
tries=$((tries + 1))
|
||
if [ $tries -gt 90 ]; then
|
||
echo "[b3_sweep] FATAL: inst_$i (port $port) not healthy after 180s"
|
||
exit 1
|
||
fi
|
||
sleep 2
|
||
done
|
||
echo " inst_$i ready"
|
||
done
|
||
}
|
||
|
||
launch_proxy() {
|
||
local policy="$1"
|
||
local logfile="$2"
|
||
local combined_args=""
|
||
for i in $(seq 0 $((N_INSTANCES - 1))); do
|
||
combined_args="$combined_args http://127.0.0.1:$((BASE_PORT + i))"
|
||
done
|
||
nohup "$VENV/python" "$ROOT/scripts/cache_aware_proxy.py" \
|
||
--port "$PROXY_PORT" \
|
||
--combined $combined_args \
|
||
--policy "$policy" \
|
||
> "$logfile" 2>&1 &
|
||
disown
|
||
local tries=0
|
||
until curl -sf "http://127.0.0.1:$PROXY_PORT/stats" >/dev/null 2>&1; do
|
||
tries=$((tries + 1))
|
||
if [ $tries -gt 30 ]; then
|
||
echo "[b3_sweep] FATAL: proxy did not come up in 60s"
|
||
tail -30 "$logfile"
|
||
exit 1
|
||
fi
|
||
sleep 2
|
||
done
|
||
}
|
||
|
||
run_policy() {
|
||
local policy="$1"
|
||
local trace="$2"
|
||
local rundir="$OUTDIR/$policy"
|
||
mkdir -p "$rundir"
|
||
echo "[b3_sweep] === policy=$policy trace=$(basename "$trace") ==="
|
||
|
||
pkill -9 -f cache_aware_proxy 2>/dev/null || true
|
||
sleep 2
|
||
launch_proxy "$policy" "$rundir/proxy.log"
|
||
|
||
local t_start
|
||
t_start=$(date +%s.%N)
|
||
echo "{\"policy\": \"$policy\", \"trace\": \"$trace\", \"t_start_unix\": $t_start}" \
|
||
> "$rundir/run_window.json.partial"
|
||
|
||
PYTHONPATH="$ROOT" "$VENV/python" -m replayer \
|
||
--trace "$trace" \
|
||
--output "$rundir/metrics.jsonl" \
|
||
--endpoint "http://127.0.0.1:$PROXY_PORT" \
|
||
--model "$MODEL" \
|
||
2>&1 | tee "$rundir/replayer.log" | tail -3
|
||
|
||
local t_end
|
||
t_end=$(date +%s.%N)
|
||
python3 - "$rundir" "$policy" "$trace" "$t_start" "$t_end" <<'PY'
|
||
import json, sys
|
||
rundir, policy, trace, t_start, t_end = sys.argv[1:]
|
||
with open(f"{rundir}/run_window.json", "w") as f:
|
||
json.dump({
|
||
"policy": policy, "trace": trace,
|
||
"t_start_unix": float(t_start),
|
||
"t_end_unix": float(t_end),
|
||
}, f, indent=2)
|
||
PY
|
||
rm -f "$rundir/run_window.json.partial"
|
||
|
||
curl -s "http://127.0.0.1:$PROXY_PORT/breakdown" > "$rundir/breakdown.json"
|
||
curl -s "http://127.0.0.1:$PROXY_PORT/worker_state" > "$rundir/worker_state.json"
|
||
curl -s "http://127.0.0.1:$PROXY_PORT/stats" > "$rundir/stats.json"
|
||
echo "[b3_sweep] $policy done: $(wc -l < "$rundir/metrics.jsonl") metric rows"
|
||
}
|
||
|
||
# 2) Run each policy
|
||
launch_vllm
|
||
|
||
for policy in $POLICIES; do
|
||
run_policy "$policy" "$TRACE"
|
||
done
|
||
|
||
# 3) Capped variant on lmetric
|
||
echo "[b3_sweep] building capped trace (max_turns=$MAX_TURNS_CAP) ..."
|
||
CAPPED_TRACE="$OUTDIR/capped/trace.jsonl"
|
||
mkdir -p "$OUTDIR/capped"
|
||
"$VENV/python" "$ROOT/scripts/build_capped_trace.py" \
|
||
--input "$TRACE" \
|
||
--output "$CAPPED_TRACE" \
|
||
--max-turns "$MAX_TURNS_CAP" | tee "$OUTDIR/capped/build.log"
|
||
run_policy "capped" "$CAPPED_TRACE"
|
||
|
||
# 4) Snapshot final engine state file sizes for the analyzer
|
||
ls -l "$OUTDIR/engine_state/" > "$OUTDIR/engine_state_files.txt"
|
||
|
||
echo "[b3_sweep] sweep complete. OUTDIR=$OUTDIR"
|