#!/usr/bin/env bash
# B300 native-topology matrix: topo x precision x task.
# Topology axes: tp / ulysses / replicas / instance batching / quantization.
# 32 requests per task: 2 tasks x 4 short-edge resolutions x 8 prompts (same gauge as 6000D).
set -Eeuo pipefail

TOTAL_GPUS=${TOTAL_GPUS:-8}
NUM_INFERENCE_STEPS=${NUM_INFERENCE_STEPS:-20}
DURATION_SECONDS=${DURATION_SECONDS:-5}
# name|replicas|tp|ulysses|batching
TOPO_LIST=${TOPO_LIST:-"u8x1|1|1|8|1 u8x1_b2|1|1|8|2 u8x1_b4|1|1|8|4 u4x2|2|1|4|1 u4x2_b2|2|1|4|2 tp1x8|8|1|1|1 tp1x8_b2|8|1|1|2 share2x8|16|1|1|1"}
QUANT_LIST=${QUANT_LIST:-"bf16 fp8"}
TASKS=${TASKS:-"fl2va ref2va"}
RESOLUTIONS=${RESOLUTIONS:-"480,720,768,1080"}
REQUESTS_PER_RESOLUTION=${REQUESTS_PER_RESOLUTION:-8}
REQUESTS_PER_TASK=$((REQUESTS_PER_RESOLUTION * 4))
GPU_MODE=${GPU_MODE:-partition}   # partition | share (share requires tp*ulysses==1)
SKIP_ON_FAIL=${SKIP_ON_FAIL:-1}   # 1: record phase failure and continue (exploration mode)

BASE_PORT=${BASE_PORT:-30010}
PORT_STRIDE=${PORT_STRIDE:-10}
MASTER_PORT_BASE=${MASTER_PORT_BASE:-31000}
SCHEDULER_PORT_BASE=${SCHEDULER_PORT_BASE:-32000}
HOST=${HOST:-127.0.0.1}
MODEL=${MODEL:-/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}
PYTHON=${PYTHON:-/root/.miniconda3/envs/sglang/bin/python}
SGLANG_BIN=${SGLANG_BIN:-/root/.miniconda3/envs/sglang/bin/sglang}
CLIENT_SCRIPT=${CLIENT_SCRIPT:-/data/wxy/sskj-h3/throughput/sglang-base-b300/scripts/minimax_h3_b300_bench.py}
SERVER_START_TIMEOUT=${SERVER_START_TIMEOUT:-1800}
RUN_ID=${RUN_ID:-b300-matrix-$(date '+%Y%m%d-%H%M%S')}
RESULT_ROOT=${RESULT_ROOT:-/data/wxy/results/minimax_h3_b300_matrix/$RUN_ID}

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 "$SGLANG_BIN" ]] || die "sglang not executable: $SGLANG_BIN"
[[ -f "$CLIENT_SCRIPT" ]] || die "client script missing: $CLIENT_SCRIPT"
[[ -f "$REFERENCE_IMAGE" ]] || die "reference image missing: $REFERENCE_IMAGE"
mkdir -p "$RESULT_ROOT"
SUMMARY_TSV="$RESULT_ROOT/summary.tsv"
printf 'topo\tprec\treplicas\ttp\tulysses\tbatching\tinflight\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[@]} server(s)"
  for pid in "${SERVER_PIDS[@]}"; do kill -INT "$pid" 2>/dev/null || true; done
  deadline=$((SECONDS + 120))
  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 SIGINT; terminating process group"
      kill -TERM -- "-$pid" 2>/dev/null || kill -TERM "$pid" 2>/dev/null || true
      sleep 5
      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 -100 "$log_file" >&2 || true; return 1; fi
    sleep 5
  done
  tail -100 "$log_file" >&2 || true
  return 1
}

gpu_csv_for() {
  # $1=replica_index $2=replicas $3=gpus_per_instance -> prints CUDA_VISIBLE_DEVICES csv
  local replica=$1 replicas=$2 k=$3 gpu_csv="" offset gpu
  if [[ "$GPU_MODE" == share ]]; then
    printf '%s' "$((replica % TOTAL_GPUS))"
    return 0
  fi
  for ((offset=0; offset<k; offset++)); do
    gpu=$((replica * k + offset)); [[ -z "$gpu_csv" ]] && gpu_csv="$gpu" || gpu_csv+=",$gpu"
  done
  printf '%s' "$gpu_csv"
}

start_servers() {
  local topo=$1 prec=$2 replicas=$3 tp=$4 ulysses=$5 batching=$6 variant=$7 phase_dir=$8
  local k=$((tp * ulysses)) replica port master_port scheduler_port gpu_csv server_dir server_log candidate enc_flag
  [[ -n "$batching" && "$batching" -gt 1 ]] && enc_flag="--encoder-parallel dp" || enc_flag=""
  SERVER_PIDS=()
  for ((replica=0; replica<replicas; replica++)); do
    port=$((BASE_PORT + replica * PORT_STRIDE))
    master_port=$((MASTER_PORT_BASE + replica * PORT_STRIDE))
    scheduler_port=$((SCHEDULER_PORT_BASE + replica * PORT_STRIDE))
    for candidate in "$port" "$((port+1))" "$master_port" "$scheduler_port"; do
      port_is_open "$candidate" && die "port already in use: $candidate"
    done
    gpu_csv=$(gpu_csv_for "$replica" "$replicas" "$k")
    server_dir="$phase_dir/server_${replica}_port${port}"; mkdir -p "$server_dir/outputs"
    server_log="$server_dir/server.log"; printf '%s\n' "$gpu_csv" > "$server_dir/cuda_visible_devices.txt"
    log "starting topo=$topo prec=$prec variant=$variant replica=$replica GPUs=$gpu_csv port=$port batching=$batching$([[ -n "$enc_flag" ]] && echo " [encoder-dp]")"
    CUDA_VISIBLE_DEVICES="$gpu_csv" PYTHONUNBUFFERED=1 TOKENIZERS_PARALLELISM=false \
      SGLANG_USE_RUNAI_MODEL_STREAMER=false setsid "$SGLANG_BIN" serve \
      --model-path "$MODEL" --model-variant "$variant" --backend sglang --performance-mode speed \
      --num-gpus "$k" --tp-size "$tp" --ulysses-degree "$ulysses" --use-fsdp-inference false \
      --enable-torch-compile false --batching-max-size "$batching" --batching-delay-ms 0 \
      $([[ "$prec" == fp8 ]] && echo --quantization fp8) $enc_flag \
      --warmup-resolutions 1344x768 \
      --host 0.0.0.0 --port "$port" --master-port "$master_port" --scheduler-port "$scheduler_port" \
      --output-path "$server_dir/outputs" >"$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" || {
      log "WARN: server failed startup: topo=$topo prec=$prec variant=$variant replica=$replica"
      return 1
    }
    log "variant=$variant replica=$replica healthy port=$port"
  done
}

run_clients() {
  local topo=$1 prec=$2 replicas=$3 tp=$4 ulysses=$5 batching=$6 task=$7 phase_dir=$8
  local replica port client_dir in_flight failed=0
  in_flight=$batching   # in-flight jobs per instance aligns with server batching ceiling
  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"
    "$PYTHON" "$CLIENT_SCRIPT" run --host "$HOST" --port "$port" --model "$MODEL" --task "$task" \
      --reference-image "$REFERENCE_IMAGE" --prompt-file "$PROMPT_FILE" --resolutions "$RESOLUTIONS" \
      --requests-per-resolution "$REQUESTS_PER_RESOLUTION" --replica-index "$replica" --num-replicas "$replicas" \
      --in-flight "$in_flight" \
      --num-inference-steps "$NUM_INFERENCE_STEPS" --warmup-requests 1 --warmup-inference-steps 5 \
      --duration-seconds "$DURATION_SECONDS" --aspect-ratio 16:9 --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)) in_flight=$in_flight"
  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" --topo "$topo" --prec "$prec" --tp "$tp" --ulysses "$ulysses" \
    --replicas "$replicas" --batching "$batching" --in-flight "$in_flight" \
    --expected-requests "$REQUESTS_PER_TASK" >> "$SUMMARY_TSV" || failed=1
  return "$failed"
}

if [[ "${DRY_RUN:-0}" == 1 ]]; then
  echo "===== B300 matrix plan (DRY_RUN) ====="
  for topo_spec in $TOPO_LIST; do
    IFS='|' read -r name replicas tp ulysses batching <<< "$topo_spec"
    for prec in $QUANT_LIST; do
      for task in $TASKS; do
        echo "  topo=$name prec=$prec task=$task replicas=$replicas tp=$tp ulysses=$ulysses batching=$batching gpu_mode=$GPU_MODE"
      done
    done
  done
  echo "RESULT_ROOT=$RESULT_ROOT"
  exit 0
fi

for topo_spec in $TOPO_LIST; do
  IFS='|' read -r name replicas tp ulysses batching <<< "$topo_spec"
  k=$((tp * ulysses))
  topo_skip=0
  if [[ "$GPU_MODE" == share ]]; then
    ((k == 1)) || { log "WARN: skip topo $name: GPU_MODE=share requires tp*ulysses==1"; topo_skip=1; }
    ((replicas % TOTAL_GPUS == 0)) || { log "WARN: skip topo $name: replicas=$replicas not multiple of TOTAL_GPUS=$TOTAL_GPUS"; topo_skip=1; }
  else
    ((replicas * k == TOTAL_GPUS)) || { log "WARN: skip topo $name: replicas*K=$((replicas*k)) != TOTAL_GPUS=$TOTAL_GPUS (multi-instance-per-GPU topo needs GPU_MODE=share)"; topo_skip=1; }
  fi
  ((REQUESTS_PER_TASK % replicas == 0)) || { log "WARN: skip topo $name: requests/task not divisible by replicas=$replicas"; topo_skip=1; }
  ((topo_skip == 1)) && continue
  for prec in $QUANT_LIST; do
    for task in $TASKS; do
      [[ "$task" == fl2va ]] && variant=FL2VA || variant=Ref2VA
      phase_dir="$RESULT_ROOT/${name}_${prec}/${task}"; mkdir -p "$phase_dir"
      log "===== topo=$name prec=$prec task=$task replicas=$replicas tp=$tp ulysses=$ulysses batching=$batching ====="
      phase_failed=0
      if start_servers "$name" "$prec" "$replicas" "$tp" "$ulysses" "$batching" "$variant" "$phase_dir"; then
        run_clients "$name" "$prec" "$replicas" "$tp" "$ulysses" "$batching" "$task" "$phase_dir" || phase_failed=1
      else
        phase_failed=1
      fi
      stop_servers
      if ((phase_failed == 1)); then
        if [[ "$SKIP_ON_FAIL" == 1 ]]; then
          log "WARN: phase failed (topo=$name prec=$prec task=$task); skipping and continuing"
        else
          die "phase failed: topo=$name prec=$prec task=$task; inspect $phase_dir"
        fi
      fi
    done
  done
done

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