#!/usr/bin/env bash
# MiniMax-H3 vLLM-Omni benchmark on one 8-GPU RTX 6000D host.
# The 8/4/2 labels mean GPUs per model instance. Per the official RTX PRO
# 6000 recipe, the internal layouts are TP2xUSP4, TP2xUSP2, and TP2xUSP1.
set -Eeuo pipefail

TOTAL_GPUS=${TOTAL_GPUS:-8}
GPU_MATRIX=${GPU_MATRIX:-"8 4 2"}
TASKS=${TASKS:-"fl2va ref2va"}
RESOLUTION_MAP=${RESOLUTION_MAP:-"480:864x480,720:1280x736,768:1344x768,1080:1920x1088"}
REQUESTS_PER_RESOLUTION=${REQUESTS_PER_RESOLUTION:-8}
REQUESTS_PER_TASK=$((REQUESTS_PER_RESOLUTION * 4))
TOTAL_REQUESTS_PER_CONFIG=$((REQUESTS_PER_TASK * 2))
NUM_INFERENCE_STEPS=${NUM_INFERENCE_STEPS:-20}
DURATION_SECONDS=${DURATION_SECONDS:-5}
WARMUP_REQUESTS=${WARMUP_REQUESTS:-2}
WARMUP_INFERENCE_STEPS=${WARMUP_INFERENCE_STEPS:-5}

HOST=${HOST:-127.0.0.1}
BASE_PORT=${BASE_PORT:-30010}
PORT_STRIDE=${PORT_STRIDE:-100}
VLLM_INTERNAL_PORT_BASE=${VLLM_INTERNAL_PORT_BASE:-35000}
OMNI_MASTER_PORT_BASE=${OMNI_MASTER_PORT_BASE:-36000}
SERVER_START_TIMEOUT=${SERVER_START_TIMEOUT:-1800}
VIDEO_REQUEST_TIMEOUT=${VIDEO_REQUEST_TIMEOUT:-14400}

MODEL_ROOT=${MODEL_ROOT:-/data/hf_models/MiniMax-H3}
REFERENCE_IMAGE=${REFERENCE_IMAGE:-/data/wxy/sskj-MiniMax-H3/assets/reference_images/landscape_mountain_lake.jpg}
PROMPT_FILE=${PROMPT_FILE:-/root/.cache/sglang/vbench_subject_consistency.txt}
VLLM_ENV=${VLLM_ENV:-/root/.miniconda3/envs/deploy}
PYTHON=${PYTHON:-$VLLM_ENV/bin/python}
VLLM_BIN=${VLLM_BIN:-$VLLM_ENV/bin/vllm}
CLIENT_SCRIPT=${CLIENT_SCRIPT:-/data/wxy/minimax_h3_vllm_bench.py}
RUN_ID=${RUN_ID:-vllm-balanced64-$(date '+%Y%m%d-%H%M%S')}
RESULT_ROOT=${RESULT_ROOT:-/data/wxy/results/minimax_h3_vllm_matrix/$RUN_ID}

export PATH="$VLLM_ENV/bin:$PATH"
export LD_LIBRARY_PATH="$VLLM_ENV/lib:${LD_LIBRARY_PATH:-}"

declare -a SERVER_PIDS=()
declare -a CLIENT_PIDS=()

log() { printf '[%s] %s\n' "$(date '+%F %T')" "$*"; }
die() { log "ERROR: $*" >&2; exit 1; }

[[ -x "$PYTHON" ]] || die "python not executable: $PYTHON"
[[ -x "$VLLM_BIN" ]] || die "vllm not executable: $VLLM_BIN"
[[ -f "$CLIENT_SCRIPT" ]] || die "client script missing: $CLIENT_SCRIPT"
[[ -d "$MODEL_ROOT/FL2VA" ]] || die "FL2VA checkpoint missing under $MODEL_ROOT"
[[ -d "$MODEL_ROOT/Ref2VA" ]] || die "Ref2VA checkpoint missing under $MODEL_ROOT"
[[ -f "$REFERENCE_IMAGE" ]] || die "reference image missing: $REFERENCE_IMAGE"
[[ -s "$PROMPT_FILE" ]] || die "VBench prompt file missing or empty: $PROMPT_FILE"

mkdir -p "$RESULT_ROOT"
SUMMARY_TSV="$RESULT_ROOT/summary.tsv"
printf 'gpus_per_instance\tdit_tp\tusp\treplicas\ttask\texpected\trecorded\tcompleted\tfailed\tmachine_qps\tlatency_mean_s\tlatency_p95_s\tmachine_wall_s\n' > "$SUMMARY_TSV"

port_is_open() {
  "$PYTHON" - "$HOST" "$1" <<'PY'
import socket, sys
s = socket.socket(); s.settimeout(0.5)
try: s.connect((sys.argv[1], int(sys.argv[2])))
except OSError: raise SystemExit(1)
else: raise SystemExit(0)
finally: s.close()
PY
}

stop_servers() {
  local pid alive deadline
  ((${#SERVER_PIDS[@]})) || return 0
  log "gracefully stopping ${#SERVER_PIDS[@]} vLLM-Omni server(s)"
  for pid in "${SERVER_PIDS[@]}"; do
    kill -TERM -- "-$pid" 2>/dev/null || kill -TERM "$pid" 2>/dev/null || true
  done
  deadline=$((SECONDS + 180))
  while ((SECONDS < deadline)); do
    alive=0
    for pid in "${SERVER_PIDS[@]}"; do kill -0 "$pid" 2>/dev/null && alive=1; done
    ((alive == 0)) && break
    sleep 2
  done
  for pid in "${SERVER_PIDS[@]}"; do
    if kill -0 "$pid" 2>/dev/null; then
      log "server pid=$pid did not exit after SIGTERM; killing process group"
      kill -KILL -- "-$pid" 2>/dev/null || kill -KILL "$pid" 2>/dev/null || true
    fi
    wait "$pid" 2>/dev/null || true
  done
  SERVER_PIDS=()
}

cleanup() {
  local rc=$? pid
  trap - EXIT INT TERM
  for pid in "${CLIENT_PIDS[@]}"; do kill -TERM "$pid" 2>/dev/null || true; done
  stop_servers
  exit "$rc"
}
trap cleanup EXIT INT TERM

wait_healthy() {
  local port=$1 pid=$2 log_file=$3 deadline=$((SECONDS + SERVER_START_TIMEOUT))
  while ((SECONDS < deadline)); do
    curl -fsS --max-time 5 "http://${HOST}:${port}/health" >/dev/null 2>&1 && return 0
    if ! kill -0 "$pid" 2>/dev/null; then
      tail -120 "$log_file" >&2 || true
      return 1
    fi
    sleep 5
  done
  tail -120 "$log_file" >&2 || true
  return 1
}

parallel_layout() {
  local gpus_per_instance=$1
  case "$gpus_per_instance" in
    2) printf '2 1 2 2\n' ;;
    4) printf '2 2 4 4\n' ;;
    8) printf '2 4 8 8\n' ;;
    *) die "supported GPUs per instance are 2, 4, and 8; got $gpus_per_instance" ;;
  esac
}

start_servers() {
  local gpus_per_instance=$1 replicas=$2 variant=$3 phase_dir=$4
  local dit_tp usp text_tp vae_pp replica port internal_port omni_port
  local first_gpu offset gpu gpu_csv server_dir server_log rpc_dir candidate
  read -r dit_tp usp text_tp vae_pp < <(parallel_layout "$gpus_per_instance")
  SERVER_PIDS=()
  for ((replica=0; replica<replicas; replica++)); do
    port=$((BASE_PORT + replica * PORT_STRIDE))
    internal_port=$((VLLM_INTERNAL_PORT_BASE + replica * PORT_STRIDE))
    omni_port=$((OMNI_MASTER_PORT_BASE + replica * PORT_STRIDE))
    for candidate in "$port" "$internal_port" "$omni_port"; do
      port_is_open "$candidate" && die "port already in use: $candidate"
    done
    first_gpu=$((replica * gpus_per_instance)); gpu_csv=""
    for ((offset=0; offset<gpus_per_instance; offset++)); do
      gpu=$((first_gpu + offset))
      [[ -z "$gpu_csv" ]] && gpu_csv="$gpu" || gpu_csv+=",$gpu"
    done
    server_dir="$phase_dir/server_${replica}_port${port}"
    # Unix-domain socket paths are limited to roughly 107 characters on Linux.
    # Keep vLLM's RPC base independent of the deliberately descriptive result path.
    rpc_dir=$(mktemp -d /tmp/vh3rpc.XXXXXX)
    mkdir -p "$server_dir/outputs"
    printf '%s\n' "$rpc_dir" > "$server_dir/vllm_rpc_base_path.txt"
    server_log="$server_dir/server.log"
    printf '%s\n' "$gpu_csv" > "$server_dir/cuda_visible_devices.txt"
    log "starting variant=$variant replica=$replica GPUs=$gpu_csv port=$port TP=$dit_tp USP=$usp"
    CUDA_VISIBLE_DEVICES="$gpu_csv" \
      VLLM_WORKER_MULTIPROC_METHOD=spawn \
      VLLM_OMNI_VIDEO_SYNC_TIMEOUT="$VIDEO_REQUEST_TIMEOUT" \
      VLLM_PORT="$internal_port" \
      VLLM_RPC_BASE_PATH="$rpc_dir" \
      PYTHONUNBUFFERED=1 \
      setsid "$VLLM_BIN" serve "$MODEL_ROOT/$variant" \
        --omni \
        --host 0.0.0.0 \
        --port "$port" \
        --trust-remote-code \
        --num-gpus "$gpus_per_instance" \
        --tensor-parallel-size "$dit_tp" \
        --usp "$usp" \
        --ring 1 \
        --text-encoder-tp-size "$text_tp" \
        --vae-patch-parallel-size "$vae_pp" \
        --vae-parallel-mode tile \
        --vae-use-tiling \
        --diffusion-attention-backend CUDNN_ATTN \
        --enable-diffusion-pipeline-profiler \
        --omni-master-address 127.0.0.1 \
        --omni-master-port "$omni_port" \
        >"$server_log" 2>&1 &
    SERVER_PIDS+=("$!")
  done
  for ((replica=0; replica<replicas; replica++)); do
    port=$((BASE_PORT + replica * PORT_STRIDE))
    server_log="$phase_dir/server_${replica}_port${port}/server.log"
    wait_healthy "$port" "${SERVER_PIDS[$replica]}" "$server_log" \
      || die "$variant replica=$replica failed startup"
    log "variant=$variant replica=$replica healthy port=$port"
  done
}

run_clients() {
  local gpus_per_instance=$1 replicas=$2 task=$3 phase_dir=$4
  local dit_tp usp text_tp vae_pp replica port client_dir failed=0
  read -r dit_tp usp text_tp vae_pp < <(parallel_layout "$gpus_per_instance")
  CLIENT_PIDS=()
  for ((replica=0; replica<replicas; replica++)); do
    port=$((BASE_PORT + replica * PORT_STRIDE))
    client_dir="$phase_dir/client_${replica}_port${port}"
    mkdir -p "$client_dir/videos"
    "$PYTHON" "$CLIENT_SCRIPT" run \
      --host "$HOST" --port "$port" --task "$task" \
      --reference-image "$REFERENCE_IMAGE" --prompt-file "$PROMPT_FILE" \
      --resolution-map "$RESOLUTION_MAP" --requests-per-resolution "$REQUESTS_PER_RESOLUTION" \
      --replica-index "$replica" --num-replicas "$replicas" \
      --gpus-per-instance "$gpus_per_instance" --dit-tp "$dit_tp" --usp "$usp" \
      --num-inference-steps "$NUM_INFERENCE_STEPS" \
      --warmup-requests "$WARMUP_REQUESTS" --warmup-inference-steps "$WARMUP_INFERENCE_STEPS" \
      --duration-seconds "$DURATION_SECONDS" --aspect-ratio 16:9 \
      --request-timeout "$VIDEO_REQUEST_TIMEOUT" \
      --video-dir "$client_dir/videos" --output "$client_dir/results.jsonl" \
      >"$client_dir/client.log" 2>&1 &
    CLIENT_PIDS+=("$!")
    log "started task=$task client=$replica port=$port requests=$((REQUESTS_PER_TASK / replicas))"
  done
  for ((replica=0; replica<replicas; replica++)); do
    wait "${CLIENT_PIDS[$replica]}" || failed=1
  done
  CLIENT_PIDS=()
  "$PYTHON" "$CLIENT_SCRIPT" summarize \
    --input-dir "$phase_dir" --output "$phase_dir/summary.json" \
    --task "$task" --gpus-per-instance "$gpus_per_instance" \
    --dit-tp "$dit_tp" --usp "$usp" --replicas "$replicas" \
    --expected-requests "$REQUESTS_PER_TASK" >> "$SUMMARY_TSV" || failed=1
  return "$failed"
}

read -r -a GPU_VALUES <<< "$GPU_MATRIX"
read -r -a TASK_VALUES <<< "$TASKS"
for gpus_per_instance in "${GPU_VALUES[@]}"; do
  ((TOTAL_GPUS % gpus_per_instance == 0)) \
    || die "TOTAL_GPUS=$TOTAL_GPUS not divisible by gpus_per_instance=$gpus_per_instance"
  replicas=$((TOTAL_GPUS / gpus_per_instance))
  ((REQUESTS_PER_RESOLUTION % replicas == 0)) \
    || die "REQUESTS_PER_RESOLUTION=$REQUESTS_PER_RESOLUTION not divisible by replicas=$replicas"
  read -r dit_tp usp text_tp vae_pp < <(parallel_layout "$gpus_per_instance")
  log "===== GPUs/instance=$gpus_per_instance TP=$dit_tp USP=$usp replicas=$replicas total_requests=$TOTAL_REQUESTS_PER_CONFIG ====="
  for task in "${TASK_VALUES[@]}"; do
    [[ "$task" == fl2va ]] && variant=FL2VA || variant=Ref2VA
    phase_dir="$RESULT_ROOT/gpus${gpus_per_instance}_tp${dit_tp}_usp${usp}_replicas${replicas}/${task}"
    mkdir -p "$phase_dir"
    start_servers "$gpus_per_instance" "$replicas" "$variant" "$phase_dir"
    phase_failed=0
    run_clients "$gpus_per_instance" "$replicas" "$task" "$phase_dir" || phase_failed=1
    stop_servers
    ((phase_failed == 0)) \
      || die "GPUs/instance=$gpus_per_instance task=$task failed; inspect $phase_dir"
  done
done

trap - EXIT INT TERM
log "vLLM-Omni matrix complete: $SUMMARY_TSV"