#!/usr/bin/env bash
set -Eeuo pipefail

# Run H200 DeepSeek-V4-Flash-DSpark benchmark grid in the current environment.
# This orchestrator starts the vllm-dspark server, runs P1/P2/P3 benchmarks,
# and stops the server afterwards.
#
# Environment:
#   - vllm-dspark env: /data/user1/yy/envs/vllm-dspark
#   - benchmark client: /data/user1/yy/envs/sglang/bin/python -m sglang.bench_serving
#   - model: /data/models/DeepSeek-V4-Flash-DSpark
#   - server port: 30004
#
# Usage:
#   bash run_dspark_benchmark_grid.sh
#
# To only run benchmarks against an already-running server:
#   SKIP_MANAGE_SERVER=1 bash run_dspark_benchmark_grid.sh

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
START_SCRIPT="${SCRIPT_DIR}/../start_dsv4_dspark_8card.sh"
RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}"
RESULT_ROOT="${RESULT_ROOT:-/data/user1/yy/bench_results/dspark_grid_${RUN_ID}}"
LOG_DIR="${RESULT_ROOT}/logs"
PORT="${PORT:-30004}"
PID_FILE="/data/user1/yy/dsv4_flash_dspark.pid"

mkdir -p "${LOG_DIR}"

log() {
  echo "[$(date --iso-8601=seconds)] $*" | tee -a "${LOG_DIR}/orchestrator.log"
}

is_server_healthy() {
  curl --fail --silent --show-error --max-time 5 "http://127.0.0.1:${PORT}/health" >/dev/null 2>&1
}

stop_server() {
  if [[ -f "${PID_FILE}" ]]; then
    local pid
    pid="$(cat "${PID_FILE}")"
    if kill -0 "${pid}" 2>/dev/null; then
      log "stopping existing dspark server pid=${pid}"
      kill "${pid}" 2>/dev/null || true
      sleep 5
      kill -9 "${pid}" 2>/dev/null || true
    fi
    rm -f "${PID_FILE}"
  fi
  # Fallback: kill any lingering vllm serve processes for this model
  pkill -9 -f "vllm serve.*DeepSeek-V4-Flash-DSpark" 2>/dev/null || true
  sleep 2
}

start_server() {
  log "starting dspark server with ${START_SCRIPT}"
  if [[ ! -x "${START_SCRIPT}" ]]; then
    log "error: start script not found or not executable: ${START_SCRIPT}"
    exit 1
  fi
  bash "${START_SCRIPT}" >> "${LOG_DIR}/server.outer.log" 2>&1
  if ! is_server_healthy; then
    log "error: dspark server failed to become healthy"
    exit 1
  fi
  log "dspark server is healthy"
}

run_phase() {
  local name="$1"
  local script="$2"
  local log_file="${LOG_DIR}/${name}.log"

  log "running ${name}: ${script}"
  RUN_ID="${RUN_ID}" RESULT_ROOT="${RESULT_ROOT}" bash "${script}" 2>&1 | tee "${log_file}"
  log "finished ${name}"
}

on_exit() {
  local code="$?"
  log "orchestrator exiting with code=${code}"
  if [[ -z "${SKIP_MANAGE_SERVER:-}" ]]; then
    stop_server
  fi
  exit "${code}"
}
trap on_exit EXIT

main() {
  log "run_id=${RUN_ID}"
  log "result_root=${RESULT_ROOT}"
  log "script_dir=${SCRIPT_DIR}"

  if [[ -n "${SKIP_MANAGE_SERVER:-}" ]]; then
    log "SKIP_MANAGE_SERVER is set, assuming server is already running"
    if ! is_server_healthy; then
      log "error: no healthy server found at port ${PORT}"
      exit 1
    fi
  else
    stop_server
    start_server
  fi

  log "===== DSpark BENCHMARK START ====="
  run_phase "dspark_p1" "${SCRIPT_DIR}/bench_dspark_p1.sh"
  run_phase "dspark_p2" "${SCRIPT_DIR}/bench_dspark_p2.sh"
  run_phase "dspark_p3" "${SCRIPT_DIR}/bench_dspark_p3.sh"
  log "===== DSpark BENCHMARK DONE ====="

  log "all results saved to ${RESULT_ROOT}"
}

main "$@"