Trace runner (run_v3_trace.sh) + concurrent mb7 correctness test

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))

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"

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)"