#!/usr/bin/env bash
# MB7 launcher (runs on dash0). Two 2-instance modes selected by MODE env:
#   MODE=baseline  : restore stock connector, no layerwise env
#   MODE=layerwise : deploy mooncake_connector.LAYERWISE.py + MOONCAKE_LAYERWISE=1
#
# Chunked prefill is DISABLED (max-num-batched-tokens >= max prompt) so the
# producer prefill is a single forward and save_kv_layer fires once per layer
# in order — the layer-wise counter assumes this.
#
# The connector is always restored from .ORIG_BACKUP on exit.
#
# Usage (on dash0):
#   MODE=baseline  bash run_mb7.sh
#   MODE=layerwise bash run_mb7.sh

set -uo pipefail

MODE="${MODE:-baseline}"
PROJ_DIR="${PROJ_DIR:-/home/admin/cpfs/wjh/agentic-kv}"
VENV="${VENV:-$PROJ_DIR/.venv}"
MODEL="${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}"
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}}"
PYTHON="$VENV/bin/python"
MC_FILE="$VENV/lib/python3.12/site-packages/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"
LW_SRC="${LW_SRC:-/tmp/mooncake_connector.LAYERWISE.py}"
DRIVER="$PROJ_DIR/microbench/connector_tax/layerwise/mb7_layerwise.py"

mkdir -p "$OUTDIR/logs"
PORTS=(8000 8001); BPS=(8998 8999)

echo "=== MB7 ($MODE) ==="
echo "Out: $OUTDIR ; connector: $MC_FILE"

restore_connector() {
    if [ -f "$MC_FILE.ORIG_BACKUP" ]; then
        cp -f "$MC_FILE.ORIG_BACKUP" "$MC_FILE"
        echo "[restore] connector reset to ORIG"
    fi
}
cleanup() {
    pkill -9 -f "vllm serve" 2>/dev/null || true
    pkill -9 -f "EngineCore" 2>/dev/null || true
    sleep 4
    restore_connector
}
trap cleanup EXIT
pkill -9 -f "vllm serve" 2>/dev/null || true; sleep 3

# Deploy the connector for the chosen mode.
if [ "$MODE" = "layerwise" ]; then
    if [ ! -f "$LW_SRC" ]; then echo "FATAL: $LW_SRC not found (scp it first)"; exit 1; fi
    cp -f "$LW_SRC" "$MC_FILE"
    "$PYTHON" -c "import ast; ast.parse(open('$MC_FILE').read()); print('[deploy] LAYERWISE connector AST OK')" || exit 1
    LW_ENV="MOONCAKE_LAYERWISE=1"
else
    restore_connector
    LW_ENV=""
fi

echo "[launch] 2 instances (max-num-batched-tokens=$MAX_BATCHED, chunked-prefill off)"
i=0
for gpu in "${GPUS[@]:0:2}"; do
    port=${PORTS[$i]}; bp=${BPS[$i]}; master=$((29700 + i))
    env $LW_ENV \
        PYTHONHASHSEED=42 VLLM_MOONCAKE_BOOTSTRAP_PORT=$bp \
        CUDA_VISIBLE_DEVICES=$gpu MASTER_PORT=$master \
        nohup "$VENV/bin/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 \
        --max-num-batched-tokens "$MAX_BATCHED" \
        --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_both"}' \
        --enable-prompt-tokens-details \
        > "$OUTDIR/logs/vllm_${i}_gpu${gpu}.log" 2>&1 &
    disown; sleep 2; i=$((i + 1))
done

echo "[health] waiting ..."
for i in 0 1; do
    port=${PORTS[$i]}; tries=0
    while ! curl -sf "http://127.0.0.1:$port/health" >/dev/null 2>&1; do
        tries=$((tries + 1)); [ $tries -gt 180 ] && { echo "FATAL inst_$i"; exit 1; }
        sleep 2
    done
    echo "  inst_$i ready"
done
for i in 0 1; do
    bp=${BPS[$i]}; tries=0
    while ! curl -sf "http://127.0.0.1:$bp/query" >/dev/null 2>&1; do
        tries=$((tries+1)); [ $tries -gt 60 ] && { echo "WARN bp $bp"; break; }; sleep 2
    done
done

echo "[run] mb7 --mode $MODE"
"$PYTHON" "$DRIVER" --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" \
    --concurrent "$CONCURRENT" --out "$OUTDIR/mb7_result.json" \
    2>&1 | tee "$OUTDIR/mb7_run.txt"

echo "[done] $OUTDIR"
# grep layerwise transfer logs from the producer (gpu0) for sanity
if [ "$MODE" = "layerwise" ]; then
    echo "=== producer layerwise log lines ==="
    grep -i "layerwise" "$OUTDIR/logs/vllm_0_gpu${GPUS[0]}.log" | tail -10 || true
fi