#!/bin/bash
# A/B GPU utilization test: Combined vs PD-Sep
# Each run: clean start → warm up → benchmark with GPU monitoring → collect
set -euo pipefail

PROJECT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
VENV="$PROJECT_DIR/.venv/bin"
VLLM="$VENV/vllm"
PYTHON="$VENV/python"
MODEL="${MODEL_PATH:-$HOME/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}"
TRACE="$PROJECT_DIR/traces/sampled_1000req_seed42.jsonl"

# Use 200 requests for faster iteration
REQ_LIMIT=200
MAX_SESSIONS=8
MAX_CONCURRENT=16
TIME_SCALE=20  # 2x faster to reduce wall time
REQUEST_TIMEOUT=300

cleanup() {
    pkill -9 -f "gpu_monitor\|replayer\|cache_aware\|uvicorn" 2>/dev/null || true
    pkill -9 -f "vllm" 2>/dev/null || true
    sleep 5
    fuser /dev/nvidia* 2>/dev/null | tr " " "\n" | sort -u | xargs -r kill -9 2>/dev/null || true
    sleep 10
}

wait_server() {
    local port=$1
    timeout 600 bash -c "until curl -s localhost:$port/v1/models >/dev/null 2>&1; do sleep 5; done"
}

run_with_monitor() {
    local tag=$1
    local endpoint=$2
    local outdir="$PROJECT_DIR/outputs/gpu_ab_$tag"
    mkdir -p "$outdir"

    # Start GPU monitor
    bash "$PROJECT_DIR/scripts/gpu_monitor.sh" "$outdir/gpu_util.csv" 5 &
    local MON_PID=$!

    # Warm up: send 3 requests
    for i in 1 2 3; do
        curl -s -m 60 "$endpoint/v1/completions" -X POST \
            -H "Content-Type: application/json" \
            -d "{\"model\":\"$MODEL\",\"prompt\":[100,200,300],\"max_tokens\":5,\"temperature\":0}" > /dev/null 2>&1
    done
    sleep 5

    # Run benchmark
    $PYTHON -m replayer \
        --trace "$TRACE" \
        --output "$outdir/metrics.jsonl" \
        --endpoint "$endpoint" \
        --model "$MODEL" \
        --time-scale $TIME_SCALE \
        --max-inflight-sessions $MAX_SESSIONS \
        --concurrency-limit $MAX_CONCURRENT \
        --request-timeout $REQUEST_TIMEOUT \
        --request-limit $REQ_LIMIT \
        -v 2>&1 | tail -5

    # Stop monitor
    kill $MON_PID 2>/dev/null
    wait $MON_PID 2>/dev/null || true

    echo "  GPU data: $outdir/gpu_util.csv ($(wc -l < "$outdir/gpu_util.csv") samples)"
}

###############################################
# Test A: Combined TP=1 DP=8 (cache-aware)
###############################################
echo "================================================================"
echo "  TEST A: Combined TP=1 DP=8 + cache-aware scheduler"
echo "================================================================"
cleanup

for i in $(seq 0 7); do
    port=$((8000+i))
    mport=$((29500+i))
    MASTER_PORT=$mport CUDA_VISIBLE_DEVICES=$i $VLLM serve "$MODEL" \
        --host 0.0.0.0 --port $port --tensor-parallel-size 1 \
        --trust-remote-code --enable-prefix-caching --enforce-eager \
        --dtype auto --gpu-memory-utilization 0.9 --max-model-len 200000 \
        > /tmp/combined_$i.log 2>&1 &
    sleep 1
done

for i in $(seq 0 7); do wait_server $((8000+i)); done
echo "  All 8 instances ready"

$PYTHON "$PROJECT_DIR/scripts/cache_aware_proxy.py" \
    --combined http://127.0.0.1:8000 http://127.0.0.1:8001 http://127.0.0.1:8002 http://127.0.0.1:8003 \
               http://127.0.0.1:8004 http://127.0.0.1:8005 http://127.0.0.1:8006 http://127.0.0.1:8007 \
    --port 9090 > /tmp/proxy_combined.log 2>&1 &
sleep 10

echo "  Running benchmark..."
run_with_monitor "combined" "http://localhost:9090"

###############################################
# Test B: PD-Sep TP=1 4P+4D (cache-aware)
###############################################
echo ""
echo "================================================================"
echo "  TEST B: PD-Sep TP=1 4P+4D + cache-aware scheduler (Mooncake)"
echo "================================================================"
cleanup

# 4 prefill
for i in 0 1 2 3; do
    bp=$((8998+i))
    port=$((8010+i))
    mport=$((29500+i))
    MASTER_PORT=$mport VLLM_MOONCAKE_BOOTSTRAP_PORT=$bp CUDA_VISIBLE_DEVICES=$i \
    $VLLM serve "$MODEL" --host 0.0.0.0 --port $port --tensor-parallel-size 1 \
        --trust-remote-code --enable-prefix-caching --enforce-eager \
        --dtype auto --gpu-memory-utilization 0.9 --max-model-len 200000 \
        --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer"}' \
        > /tmp/prefill_$i.log 2>&1 &
    sleep 2
done

# 4 decode
for i in 0 1 2 3; do
    gpu=$((4+i))
    port=$((8020+i))
    mport=$((29510+i))
    MASTER_PORT=$mport CUDA_VISIBLE_DEVICES=$gpu \
    $VLLM serve "$MODEL" --host 0.0.0.0 --port $port --tensor-parallel-size 1 \
        --trust-remote-code --enable-prefix-caching --enforce-eager \
        --dtype auto --gpu-memory-utilization 0.9 --max-model-len 200000 \
        --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_consumer","kv_load_failure_policy":"recompute"}' \
        > /tmp/decode_$i.log 2>&1 &
    sleep 2
done

for i in 0 1 2 3; do wait_server $((8010+i)); done
for i in 0 1 2 3; do wait_server $((8020+i)); done
echo "  All 8 instances ready"

# Wait for bootstrap
for bp in 8998 8999 9000 9001; do
    timeout 120 bash -c "until curl -s localhost:$bp/query >/dev/null 2>&1; do sleep 2; done"
done

$PYTHON "$PROJECT_DIR/scripts/cache_aware_proxy.py" \
    --prefill http://127.0.0.1:8010 8998 --prefill http://127.0.0.1:8011 8999 \
    --prefill http://127.0.0.1:8012 9000 --prefill http://127.0.0.1:8013 9001 \
    --decode http://127.0.0.1:8020 --decode http://127.0.0.1:8021 \
    --decode http://127.0.0.1:8022 --decode http://127.0.0.1:8023 \
    --port 9090 > /tmp/proxy_pdsep.log 2>&1 &
sleep 15

echo "  Running benchmark..."
run_with_monitor "pdsep" "http://localhost:9090"

###############################################
# Analyze
###############################################
echo ""
echo "================================================================"
echo "  ANALYSIS"
echo "================================================================"
cleanup

$PYTHON << 'PYEOF'
import csv, json, statistics

def analyze_gpu(path, label, gpu_groups=None):
    """gpu_groups: dict of group_name -> list of gpu indices"""
    rows = []
    with open(path) as f:
        reader = csv.DictReader(f)
        for r in reader:
            rows.append(r)

    if not rows:
        print(f"  {label}: no GPU data")
        return

    # Group by GPU
    by_gpu = {}
    for r in rows:
        g = int(r["gpu"])
        by_gpu.setdefault(g, []).append(float(r["util_pct"]))

    if gpu_groups:
        for gname, indices in gpu_groups.items():
            vals = []
            for i in indices:
                vals.extend(by_gpu.get(i, []))
            if vals:
                print(f"  {label} {gname}: mean={statistics.fmean(vals):.1f}% p50={sorted(vals)[len(vals)//2]:.0f}% p90={sorted(vals)[int(0.9*len(vals))]:.0f}% max={max(vals):.0f}%")
    else:
        all_vals = []
        for vals in by_gpu.values():
            all_vals.extend(vals)
        if all_vals:
            print(f"  {label} all GPUs: mean={statistics.fmean(all_vals):.1f}% p50={sorted(all_vals)[len(all_vals)//2]:.0f}% p90={sorted(all_vals)[int(0.9*len(all_vals))]:.0f}%")

    # Per-GPU breakdown
    for g in sorted(by_gpu.keys()):
        vals = by_gpu[g]
        print(f"    GPU {g}: mean={statistics.fmean(vals):.1f}% samples={len(vals)}")

def analyze_metrics(path, label):
    rows = [json.loads(l) for l in open(path)]
    ok = [r for r in rows if not r.get("error")]
    print(f"  {label}: {len(ok)}/{len(rows)} OK")
    if ok:
        ttfts = sorted([r["ttft_s"] for r in ok if r.get("ttft_s")])
        tpots = sorted([r["tpot_s"] for r in ok if r.get("tpot_s") and r["tpot_s"]>0])
        p = lambda v,q: v[min(int(q*len(v)),len(v)-1)]
        if ttfts: print(f"    TTFT p50={p(ttfts,.5):.3f} p90={p(ttfts,.9):.3f}")
        if tpots: print(f"    TPOT p50={p(tpots,.5):.3f} p90={p(tpots,.9):.3f}")

import os
for tag, groups in [
    ("combined", None),
    ("pdsep", {"Prefill(0-3)": [0,1,2,3], "Decode(4-7)": [4,5,6,7]}),
]:
    d = f"outputs/gpu_ab_{tag}"
    if os.path.exists(f"{d}/gpu_util.csv"):
        analyze_gpu(f"{d}/gpu_util.csv", tag.upper(), groups)
    if os.path.exists(f"{d}/metrics.jsonl"):
        analyze_metrics(f"{d}/metrics.jsonl", tag.upper())
    print()
PYEOF