#!/usr/bin/env bash
# v3 trace replay with FULL migration instrumentation:
#   - instrument_dst_migration.py  : dst lifecycle -> T_kv_pull
#   - instrument_mooncake.py        : connector internals (send_blocks RDMA,
#                                      receive_kv window, ready_wait)
# Goal: decompose the 87% T_kv_pull into RDMA-actual vs control-plane to
# explain why effective bandwidth is far below the ~10 GB/s wire rate.
#
# Usage: bash run_v3_full_breakdown.sh

set -uo pipefail

PROJ_DIR="${PROJ_DIR:-/home/admin/cpfs/wjh/agentic-kv}"
TRACE="${TRACE:-$PROJ_DIR/traces/w600_r0.0015_st30.jsonl}"
DATE="$(date +%Y%m%d_%H%M)"
OUTROOT="${OUTROOT:-$PROJ_DIR/outputs/b3_v3_fullbreak_${DATE}}"
PYTHON="$PROJ_DIR/.venv/bin/python"
VENV="$PROJ_DIR/.venv"
VLLM_ROOT="${VLLM_ROOT:-$VENV/lib/python3.12/site-packages/vllm}"
DR_FIX="$PROJ_DIR/microbench/connector_tax/cache_sweep/apply_direct_read_fix.py"
DM_INSTR="$PROJ_DIR/microbench/fresh_setup/instrument_dst_migration.py"
MC_INSTR="$PROJ_DIR/microbench/fresh_setup/instrument_mooncake.py"
ANALYZE_DST="$PROJ_DIR/microbench/connector_tax/cache_sweep/analyze_dst_migration.py"
ANALYZE_XFER="$PROJ_DIR/microbench/connector_tax/cache_sweep/analyze_transfer_decomp.py"

mkdir -p "$OUTROOT"
DST_LOG_DIR="$OUTROOT/dst_mig_log"
XFER_LOG_DIR="$OUTROOT/xfer_log"
mkdir -p "$DST_LOG_DIR" "$XFER_LOG_DIR"

echo "=== unified_v3 + FULL migration breakdown ==="
echo "Out      : $OUTROOT"
echo "DST logs : $DST_LOG_DIR"
echo "XFER logs: $XFER_LOG_DIR"
echo ""

cleanup_all() {
    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
    "$PYTHON" "$DR_FIX"   --revert --vllm-root "$VLLM_ROOT" 2>/dev/null || true
    "$PYTHON" "$DM_INSTR" --revert --venv "$VENV" 2>/dev/null || true
    "$PYTHON" "$MC_INSTR" --revert --venv "$VLLM_ROOT/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py" 2>/dev/null || true
}
trap cleanup_all EXIT
cleanup_all

echo "[0a] DR_FIX"
"$PYTHON" "$DR_FIX" --apply --vllm-root "$VLLM_ROOT"
echo "[0b] DST migration instrument"
"$PYTHON" "$DM_INSTR" --apply --venv "$VENV"
echo "[0c] Mooncake transfer instrument"
"$PYTHON" "$MC_INSTR" --apply --venv "$VLLM_ROOT/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"
"$PYTHON" "$DM_INSTR" --check --venv "$VENV"
"$PYTHON" "$MC_INSTR" --check --venv "$VLLM_ROOT/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"

cfg_dir="$OUTROOT/unified_v3"
mkdir -p "$cfg_dir"

export VLLM_MOONCAKE_DISABLE_DIRECT_READ_SYNC=1
export DM_LOG_DIR="$DST_LOG_DIR"
export MB2_LOG_DIR="$XFER_LOG_DIR"

echo ""
echo "====== unified_v3 ; DM_LOG_DIR + MB2_LOG_DIR set ======"
bash "$PROJ_DIR/scripts/b3_isolated_policy.sh" "unified_v3" "$TRACE" "$cfg_dir" \
    2>&1 | tee "$cfg_dir/orchestrator.log" | tail -25

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

echo ""
echo "[Z] revert all instruments"
"$PYTHON" "$DR_FIX"   --revert --vllm-root "$VLLM_ROOT"
"$PYTHON" "$DM_INSTR" --revert --venv "$VENV"
"$PYTHON" "$MC_INSTR" --revert --venv "$VLLM_ROOT/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"

echo ""
echo "[analyze 1] dst-side T_kv_pull breakdown"
"$PYTHON" "$ANALYZE_DST" \
    --proxy-breakdown "$cfg_dir/breakdown.json" \
    --dst-log-dir "$DST_LOG_DIR" \
    --output "$cfg_dir/dst_migration_breakdown.csv" \
    --plot "$cfg_dir/dst_migration_breakdown.png" \
    2>&1 | tee "$cfg_dir/dst_migration_breakdown.txt" || echo "(dst analyze failed)"

echo ""
echo "[analyze 2] transfer decomposition: RDMA-actual vs control-plane"
"$PYTHON" "$ANALYZE_XFER" \
    --proxy-breakdown "$cfg_dir/breakdown.json" \
    --dst-log-dir "$DST_LOG_DIR" \
    --xfer-log-dir "$XFER_LOG_DIR" \
    2>&1 | tee "$cfg_dir/transfer_decomp.txt" || echo "(xfer analyze failed)"

echo ""
echo "Done. Artifacts in $cfg_dir/"
echo "  dst_migration_breakdown.{csv,png,txt}"
echo "  transfer_decomp.txt"
echo "  raw: $DST_LOG_DIR/  $XFER_LOG_DIR/"
