From 21db2affb46ee2a1f558c496b345f880123ed9af Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Thu, 28 May 2026 17:28:48 +0800 Subject: [PATCH] Trace runner (run_v3_trace.sh) + concurrent mb7 correctness test --- .../connector_tax/layerwise/mb7_layerwise.py | 26 ++++++ microbench/connector_tax/layerwise/run_mb7.sh | 3 +- .../connector_tax/layerwise/run_v3_trace.sh | 93 +++++++++++++++++++ 3 files changed, 121 insertions(+), 1 deletion(-) create mode 100755 microbench/connector_tax/layerwise/run_v3_trace.sh diff --git a/microbench/connector_tax/layerwise/mb7_layerwise.py b/microbench/connector_tax/layerwise/mb7_layerwise.py index de436af..ee40131 100644 --- a/microbench/connector_tax/layerwise/mb7_layerwise.py +++ b/microbench/connector_tax/layerwise/mb7_layerwise.py @@ -193,6 +193,29 @@ async def main_async(a): print(f"[mb7] busy: A_run={na} B_run={nb}") break + # --- concurrent correctness mode: fire N transfers at once ---------- + if a.concurrent > 1 and a.mode == "layerwise": + print(f"[mb7] CONCURRENT correctness: {a.concurrent} simultaneous " + f"transfers per size (src=A stresses concurrent producing)") + all_ok = True + for sz in sizes: + tasks = [ + asyncio.create_task(measure_layerwise( + client, A, B, src_eid, src_bp_addr, + synth_prompt(sz * 1000 + j, sz), sz * 1000 + j)) + for j in range(a.concurrent) + ] + rows = await asyncio.gather(*tasks) + oks = [r["cached"] >= int(sz * 0.9) for r in rows] + all_ok = all_ok and all(oks) + print(f" sz={sz:>6} x{a.concurrent}: cached=" + f"{[r['cached'] for r in rows]} correct={oks}") + print(f"[mb7] CONCURRENT correctness: " + f"{'ALL PASS' if all_ok else 'FAILURE'}") + if loader: + await loader.stop() + return + results = [] for sz in sizes: for rep in range(a.repeats): @@ -259,6 +282,9 @@ def main(): p.add_argument("--repeats", type=int, default=3) p.add_argument("--bg-load", type=int, default=0, help="N concurrent background decode streams across A+B") + p.add_argument("--concurrent", type=int, default=1, + help="layerwise: fire N simultaneous transfers to test " + "concurrent-producing correctness") p.add_argument("--out", default="mb7_result.json") args = p.parse_args() asyncio.run(main_async(args)) diff --git a/microbench/connector_tax/layerwise/run_mb7.sh b/microbench/connector_tax/layerwise/run_mb7.sh index f8328c6..4829585 100644 --- a/microbench/connector_tax/layerwise/run_mb7.sh +++ b/microbench/connector_tax/layerwise/run_mb7.sh @@ -23,6 +23,7 @@ GPUS=(${GPUS:-0 1}) SIZES="${SIZES:-8192,16384,32768}" REPEATS="${REPEATS:-3}" BG_LOAD="${BG_LOAD:-0}" +CONCURRENT="${CONCURRENT:-1}" MAX_BATCHED="${MAX_BATCHED:-40960}" # >= max prompt => no chunked prefill DATE="$(date +%Y%m%d_%H%M)" OUTDIR="${OUTDIR:-$PROJ_DIR/outputs/mb7_${MODE}_${DATE}}" @@ -102,7 +103,7 @@ echo "[run] mb7 --mode $MODE" --src-port "${PORTS[0]}" --dst-port "${PORTS[1]}" \ --src-bp "${BPS[0]}" --dst-bp "${BPS[1]}" \ --sizes "$SIZES" --repeats "$REPEATS" --bg-load "$BG_LOAD" \ - --out "$OUTDIR/mb7_result.json" \ + --concurrent "$CONCURRENT" --out "$OUTDIR/mb7_result.json" \ 2>&1 | tee "$OUTDIR/mb7_run.txt" echo "[done] $OUTDIR" diff --git a/microbench/connector_tax/layerwise/run_v3_trace.sh b/microbench/connector_tax/layerwise/run_v3_trace.sh new file mode 100755 index 0000000..f8db56f --- /dev/null +++ b/microbench/connector_tax/layerwise/run_v3_trace.sh @@ -0,0 +1,93 @@ +#!/usr/bin/env bash +# Full 1200-req v3 trace, two modes (MODE env), for layer-wise re-profile. +# MODE=baseline : stock connector + stock proxy (post-hoc transfer) +# MODE=layerwise : LAYERWISE connector + write-mode proxy (overlapped) +# Both: unified_v3 routing + DR-fix. Connector & proxy restored from backup +# on exit. Output-equivalence/correctness gate = success rate + migrated-req +# TTFT distribution (byte-level KV correctness already validated on mb7). +# +# Usage (on dash0): MODE=baseline bash run_v3_trace.sh +# MODE=layerwise bash run_v3_trace.sh + +set -uo pipefail +MODE="${MODE:-baseline}" +PROJ_DIR="${PROJ_DIR:-/home/admin/cpfs/wjh/agentic-kv}" +VENV="$PROJ_DIR/.venv" +VLLM_ROOT="$VENV/lib/python3.12/site-packages/vllm" +TRACE="${TRACE:-$PROJ_DIR/traces/w600_r0.0015_st30.jsonl}" +DATE="$(date +%Y%m%d_%H%M)" +OUTROOT="${OUTROOT:-$PROJ_DIR/outputs/v3trace_${MODE}_${DATE}}" +PYTHON="$VENV/bin/python" +DR_FIX="$PROJ_DIR/microbench/connector_tax/cache_sweep/apply_direct_read_fix.py" +MC_FILE="$VLLM_ROOT/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py" +PROXY_FILE="$PROJ_DIR/scripts/cache_aware_proxy.py" +LW_CONN="${LW_CONN:-/tmp/mooncake_connector.LAYERWISE.py}" +WM_PROXY="${WM_PROXY:-/tmp/cache_aware_proxy.WRITEMODE.py}" + +mkdir -p "$OUTROOT" +cfg_dir="$OUTROOT/unified_v3"; mkdir -p "$cfg_dir" + +# Backups (connector backup already exists as .ORIG_BACKUP; make proxy one). +[ -f "$MC_FILE.ORIG_BACKUP" ] || cp "$MC_FILE" "$MC_FILE.ORIG_BACKUP" +[ -f "$PROXY_FILE.ORIG_BACKUP" ] || cp "$PROXY_FILE" "$PROXY_FILE.ORIG_BACKUP" + +restore() { + cp -f "$MC_FILE.ORIG_BACKUP" "$MC_FILE" + cp -f "$PROXY_FILE.ORIG_BACKUP" "$PROXY_FILE" + "$PYTHON" "$DR_FIX" --revert --vllm-root "$VLLM_ROOT" 2>/dev/null || true + echo "[restore] connector+proxy reset to ORIG, DR-fix reverted" +} +cleanup() { + pkill -9 -f cache_aware_proxy 2>/dev/null || true + pkill -9 -f "vllm serve" 2>/dev/null || true + pkill -9 -f "EngineCore" 2>/dev/null || true + sleep 5 + restore +} +trap cleanup EXIT +pkill -9 -f "vllm serve" 2>/dev/null || true; sleep 3 +restore # start from clean + +echo "=== v3 trace ($MODE) -> $OUTROOT ===" +if [ "$MODE" = "layerwise" ]; then + cp -f "$LW_CONN" "$MC_FILE" + cp -f "$WM_PROXY" "$PROXY_FILE" + "$PYTHON" -c "import ast; ast.parse(open('$MC_FILE').read()); ast.parse(open('$PROXY_FILE').read()); print('[deploy] LAYERWISE conn + WRITEMODE proxy AST OK')" || exit 1 + export MOONCAKE_LAYERWISE=1 + export EAR_WRITE_MODE=1 +fi + +echo "[DR-fix] apply" +"$PYTHON" "$DR_FIX" --apply --vllm-root "$VLLM_ROOT" +export VLLM_MOONCAKE_DISABLE_DIRECT_READ_SYNC=1 + +echo "[run] unified_v3 (MOONCAKE_LAYERWISE=${MOONCAKE_LAYERWISE:-0} EAR_WRITE_MODE=${EAR_WRITE_MODE:-0})" +bash "$PROJ_DIR/scripts/b3_isolated_policy.sh" "unified_v3" "$TRACE" "$cfg_dir" \ + 2>&1 | tee "$cfg_dir/orchestrator.log" | tail -20 + +pkill -9 -f cache_aware_proxy 2>/dev/null || true +pkill -9 -f "vllm serve" 2>/dev/null || true +sleep 5 + +echo "[stats] $MODE" +"$PYTHON" - "$cfg_dir" << 'PYEOF' +import json, sys, statistics +d = sys.argv[1] +ms = [json.loads(l) for l in open(f"{d}/metrics.jsonl")] +ok = [m for m in ms if not m.get("error")] +ttft = sorted(m["ttft_s"] for m in ok if m.get("ttft_s") is not None) +def p(q): return ttft[min(len(ttft)-1, int(q*len(ttft)))] if ttft else 0 +print(f" requests: {len(ms)} success: {len(ok)} ({len(ok)/max(1,len(ms))*100:.1f}%)") +print(f" TTFT s : p50={p(.5):.2f} p90={p(.9):.2f} p99={p(.99):.2f}") +# migrated reqs from proxy breakdown +try: + bd = json.load(open(f"{d}/breakdown.json")) + mig = [x for x in bd if x.get("route_class") == "PD_SEP_V2"] + mids = {x["request_id"] for x in mig} + mt = sorted(m["ttft_s"] for m in ok if m["request_id"] in mids and m.get("ttft_s")) + print(f" migrations: {len(mig)} migrated-req TTFT: " + f"p50={mt[len(mt)//2]:.2f} p90={mt[int(len(mt)*.9)]:.2f} max={mt[-1]:.2f}" if mt else f" migrations: {len(mig)}") +except Exception as e: + print(f" (breakdown parse: {e})") +PYEOF +echo "[done] $cfg_dir (metrics.jsonl, breakdown.json)"