#!/bin/bash
# Exp (a) 4th tier: remote global-KV-store hit over RDMA (Mooncake).
# Two kv_both MooncakeConnector instances (GPU0=src, GPU1=dst). For each prefix
# length: src prefills+caches the KV, dst serves the request by PULLING that KV
# over RDMA (do_remote_prefill) instead of recomputing -> that pull time is the
# remote-store hit latency. Mirrors the Mooncake-Store blog mechanism.
set -uo pipefail
cd /home/admin/cpfs/wjh/agentic-kv
PY=.venv/bin/python
MODEL=/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct
OUT=v2/exp_a_tier_latency/results
mkdir -p "$OUT"
PIDS=()

launch() {  # $1 gpu, $2 http port, $3 bootstrap port, $4 master port
    VLLM_MOONCAKE_BOOTSTRAP_PORT=$3 MASTER_PORT=$4 CUDA_VISIBLE_DEVICES=$1 VLLM_LOGGING_LEVEL=WARNING \
    $PY -m vllm.entrypoints.openai.api_server --model "$MODEL" \
        --host 0.0.0.0 --port $2 --tensor-parallel-size 1 --trust-remote-code \
        --enable-prefix-caching --enforce-eager --dtype auto --max-model-len 70000 \
        --gpu-memory-utilization 0.9 \
        --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_both"}' \
        > "$OUT/vllm_rdma_$2.log" 2>&1 &
    PIDS+=($!)
}
teardown() {
    for p in "${PIDS[@]:-}"; do kill -TERM "$p" 2>/dev/null; done
    sleep 6
    for p in $(pgrep -f "VLLM::EngineCore"); do kill -9 "$p" 2>/dev/null; done
    sleep 3
}
trap teardown EXIT

echo ">>> launch 2 kv_both instances (GPU0:8000/bp8998, GPU1:8001/bp8999)"
launch 0 8000 8998 29550
launch 1 8001 8999 29551
for port in 8000 8001; do
    echo -n "  wait health $port..."
    timeout 900 bash -c "until curl -sf http://127.0.0.1:$port/health >/dev/null 2>&1; do sleep 5; done" \
        && echo " ok" || { echo " FAIL"; tail -25 "$OUT/vllm_rdma_$port.log"; exit 1; }
done
for bp in 8998 8999; do
    timeout 180 bash -c "until curl -s http://127.0.0.1:$bp/query >/dev/null 2>&1; do sleep 2; done"
done
echo "  bootstrap ports ready."
sleep 3

$PY microbench/fresh_setup/mb2_kv_transfer.py \
    --src-host 127.0.0.1 --dst-host 127.0.0.1 \
    --src-port 8000 --dst-port 8001 --src-bp 8998 --dst-bp 8999 \
    --sizes 1024,2048,4096,8192,16384,32768,65536 --repeats 11 \
    --label rdma-intra-node --out "$OUT/rdma.json"

echo "=== exp (a) RDMA tier DONE ==="