#!/usr/bin/env bash
# Fully automated baseline-vs-optimized benchmark for zai-org/GLM-5.2 (ROCm).
# Starts the baseline server, benchmarks it, tears it down, starts the
# optimized server, benchmarks it, tears it down, then prints a
# Metric | Baseline | Optimized | Delta table. No manual steps.
#
# Usage: ./run_benchmark.sh [baseline|optimized|both]   (default: both)
# Env overrides: MAX_CONCURRENCY, SEED, SERVER_READY_TIMEOUT
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODEL="zai-org/GLM-5.2"
HOST="127.0.0.1"
PORT="8000"
MAX_CONCURRENCY="${MAX_CONCURRENCY:-512}"
SEED="${SEED:-0}"
SERVER_READY_TIMEOUT="${SERVER_READY_TIMEOUT:-3600}"
MODE="${1:-both}"

CURRENT_SERVER_PID=""
stop_current_server() {
  [ -n "$CURRENT_SERVER_PID" ] || return 0
  kill -TERM -- "-$CURRENT_SERVER_PID" 2>/dev/null || kill -TERM "$CURRENT_SERVER_PID" 2>/dev/null || true
  for _ in $(seq 1 30); do
    kill -0 "$CURRENT_SERVER_PID" 2>/dev/null || { CURRENT_SERVER_PID=""; return 0; }
    sleep 1
  done
  kill -KILL -- "-$CURRENT_SERVER_PID" 2>/dev/null || kill -KILL "$CURRENT_SERVER_PID" 2>/dev/null || true
  CURRENT_SERVER_PID=""
}
trap stop_current_server EXIT INT TERM

wait_for_server() {
  python3 - "$HOST" "$PORT" "$SERVER_READY_TIMEOUT" <<'PY'
import sys, time, urllib.request
host, port, timeout = sys.argv[1], sys.argv[2], int(sys.argv[3])
url = f"http://{host}:{port}/health"
deadline = time.time() + timeout
while time.time() < deadline:
    try:
        urllib.request.urlopen(url, timeout=5)
        sys.exit(0)
    except Exception:
        time.sleep(10)
sys.exit(1)
PY
}

run_one() {
  local tag="$1"
  echo "=== [$tag] starting server (log: server_${tag}.log) ==="
  setsid "$SCRIPT_DIR/launch_server.sh" "$tag" >"server_${tag}.log" 2>&1 &
  CURRENT_SERVER_PID=$!

  echo "=== [$tag] waiting for server (up to ${SERVER_READY_TIMEOUT}s) ==="
  if ! wait_for_server; then
    echo "[$tag] server did not become healthy in time -- see server_${tag}.log" >&2
    stop_current_server
    exit 1
  fi

  echo "=== [$tag] benchmarking ==="
  local args=(
    --backend openai
    --host "$HOST" --port "$PORT"
    --model "$MODEL"
    --dataset-name random
    --random-input-len 2048
    --random-output-len 256
    --random-range-ratio 0.8
    --request-rate inf
    --max-concurrency "$MAX_CONCURRENCY"
    --seed "$SEED"
  )
  vllm bench serve "${args[@]}" --num-prompts 10 >/dev/null   # warm-up, discarded
  vllm bench serve "${args[@]}" --num-prompts 1000 \
    --save-result --result-filename "benchmark_${tag}.json"

  echo "=== [$tag] stopping server ==="
  stop_current_server
}

case "$MODE" in
  baseline|optimized) run_one "$MODE" ;;
  both) run_one baseline; run_one optimized ;;
  *)
    echo "Usage: $0 [baseline|optimized|both]" >&2
    exit 1
    ;;
esac

python3 - "$MODE" <<'PY'
import json, os, sys

mode = sys.argv[1]
tags = ["baseline", "optimized"] if mode == "both" else [mode]

data = {}
for tag in tags:
    path = f"benchmark_{tag}.json"
    if os.path.exists(path):
        with open(path) as f:
            data[tag] = json.load(f)

metrics = [
    ("Request throughput (req/s)", "request_throughput"),
    ("Output tok/s", "output_throughput"),
    ("Total tok/s", "total_token_throughput"),
    ("TTFT mean (ms)", "mean_ttft_ms"),
    ("TTFT median (ms)", "median_ttft_ms"),
    ("TTFT p99 (ms)", "p99_ttft_ms"),
    ("TPOT mean (ms)", "mean_tpot_ms"),
    ("TPOT median (ms)", "median_tpot_ms"),
    ("TPOT p99 (ms)", "p99_tpot_ms"),
    ("ITL mean (ms)", "mean_itl_ms"),
    ("ITL median (ms)", "median_itl_ms"),
    ("ITL p99 (ms)", "p99_itl_ms"),
    ("Duration (s)", "duration"),
]

def fmt(v):
    return f"{v:.1f}" if isinstance(v, (int, float)) else "n/a"

print()
if mode == "both":
    print("| Metric | Baseline | Optimized | Delta |")
    print("|---|---|---|---|")
    for name, key in metrics:
        b, o = data["baseline"].get(key), data["optimized"].get(key)
        delta = f"{(o - b) / b * 100:+.1f}%" if isinstance(b, (int, float)) and isinstance(o, (int, float)) and b else "n/a"
        print(f"| {name} | {fmt(b)} | {fmt(o)} | {delta} |")
else:
    print(f"| Metric | {mode.capitalize()} |")
    print("|---|---|")
    for name, key in metrics:
        print(f"| {name} | {fmt(data[mode].get(key))} |")
PY