#!/usr/bin/env bash
# Trace replay for unified_v3 WITHOUT affinity rotation.
#
# This is the #2 follow-up to cache_miss_audit: prior run showed v3 with
# rotation hits 9.5% cache on post-migration next turn vs 80.6% for unified.
# Hypothesis: rotation destroys prefix cache locality. Test by keeping the
# session affinity on prefill_host even after migration (i.e., the same
# behavior as unified for the post-migration write), so only the *current*
# turn's decode is migrated.
#
# Applies CT_DR_FIX (Mooncake DR sync disabled).
# Output: $PROJ_DIR/outputs/b3_v3_norot_${DATE}/unified_v3/

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_norot_${DATE}}"
PYTHON="$PROJ_DIR/.venv/bin/python"
DR_FIX_SCRIPT="$PROJ_DIR/microbench/connector_tax/cache_sweep/apply_direct_read_fix.py"
VLLM_ROOT="${VLLM_ROOT:-$PROJ_DIR/.venv/lib/python3.12/site-packages/vllm}"

mkdir -p "$OUTROOT"
echo "=== unified_v3 (no rotation) trace replay ==="
echo "Trace : $TRACE"
echo "Out   : $OUTROOT"
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_SCRIPT" --revert --vllm-root "$VLLM_ROOT" 2>/dev/null || true
}
trap cleanup_all EXIT
cleanup_all

echo "[stage 0] applying CT_DR_FIX (env-gated)"
"$PYTHON" "$DR_FIX_SCRIPT" --apply --vllm-root "$VLLM_ROOT"

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

export VLLM_MOONCAKE_DISABLE_DIRECT_READ_SYNC=1
export EXTRA_PROXY_ARGS="--v3-rotate-affinity 0"

echo ""
echo "====== unified_v3 (no rotation) ; DR_SYNC_DISABLED=1 ======"
bash "$PROJ_DIR/scripts/b3_isolated_policy.sh" "unified_v3" "$TRACE" "$cfg_dir" \
    2>&1 | tee "$cfg_dir/orchestrator.log" | tail -30

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 "[stage Z] reverting CT_DR_FIX"
"$PYTHON" "$DR_FIX_SCRIPT" --revert --vllm-root "$VLLM_ROOT"

echo ""
echo "Done. Artifacts: $OUTROOT/unified_v3/metrics.jsonl"
