#!/usr/bin/env bash
# Cache-DiT + Larry LoRA over 8 Feishu Ref2VA cases, TP4 x 2 replicas.
# Resumable per case.
set -Eeuo pipefail
MODEL=${MODEL:-/data/hf_models/MiniMax-H3}
REPO_ROOT=/data/wxy/sskj-h3
INPUT_ROOT=$REPO_ROOT/throughput/sglang-base/inputs/ref2va-feishu-20260831
RECORDS=${RECORDS:-$REPO_ROOT/throughput/sglang-lora/records_rewritten.json}
ASSETS_ROOT=$INPUT_ROOT/assets
PYTHON=/root/.miniconda3/envs/sglang-lora/bin/python
SGLANG_BIN=/root/.miniconda3/envs/sglang-lora/bin/sglang
CACHE_PATCH=/data/ymrl/h3_cache_patch
CLIENT_SCRIPT=$REPO_ROOT/throughput/sglang-base/scripts/ref2va_feishu_bench.py
MEDIA_BIN_DIR=/root/.miniconda3/envs/deploy/bin
LORA_PATH=/data/hf_models/MiniMax-H3-Turbo-Lora
LORA_WEIGHT_NAME=minimax_h3_turbo_v4_step600_ema.safetensors
LORA_NICKNAME=h3-turbo-v4-600
LORA_SCALE=1.0
LORA_MERGE_MODE=merge
BASE_PORT=34010; PORT_STRIDE=10
MASTER_PORT_BASE=35000; SCHEDULER_PORT_BASE=36000
HOST=127.0.0.1
TP_SIZE=4; REPLICAS=2
NUM_INFERENCE_STEPS=9
SHORT_EDGE=768; ASPECT_RATIO=9:16; DURATION_SECONDS=15
FLOW_SHIFT=12.0; AUDIO_FLOW_SHIFT=3.0
CLIENT_WARMUP_STEPS=5
CACHE_FN=2; CACHE_BN=0; CACHE_WARMUP=2
MAX_ATTEMPTS=3; SERVER_START_TIMEOUT=1800
RUN_ID=${RUN_ID:-ref2va-feishu-cache-lora-tp4x2-3tier-$(date +%Y%m%d-%H%M%S)}
RESULT_ROOT=${RESULT_ROOT:-$REPO_ROOT/throughput/sglang-lora/results/$RUN_ID}
PROFILES=(
  "quality_rdt003_mc2 0.03 2 quality"
  "balanced_rdt005_mc3 0.05 3 balanced"
  "throughput_rdt010_mc4 0.10 4 throughput"
)
declare -a SERVER_PIDS=()
declare -a CLIENT_PIDS=()
CURRENT_PROFILE_DIR=""
mkdir -p "$RESULT_ROOT"
MASTER_LOG="$RESULT_ROOT/matrix.log"
exec > >(tee -a "$MASTER_LOG") 2>&1
log() { printf "[%s] %s\n" "$(date +%F_%T)" "$*"; }
die() { log "ERROR: $*" >&2; exit 1; }
[[ -x "$PYTHON" ]] || die "python missing"
[[ -x "$SGLANG_BIN" ]] || die "sglang missing"
[[ -f "$MEDIA_BIN_DIR/ffprobe" ]] || die "ffprobe missing"
[[ -f "$CLIENT_SCRIPT" ]] || die "client missing"
[[ -f "$CACHE_PATCH/sitecustomize.py" ]] || die "patch missing"
[[ -f "$LORA_PATH/$LORA_WEIGHT_NAME" ]] || die "lora missing"
[[ -f "$RECORDS" ]] || die "records missing"

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 deadline
  ((${#SERVER_PIDS[@]})) || return 0
  log "stopping ${#SERVER_PIDS[@]} server(s)"
  for pid in "${SERVER_PIDS[@]}"; do kill -INT -- "-$pid" 2>/dev/null || kill -INT "$pid" 2>/dev/null || true; done
  deadline=$((SECONDS + 180))
  while ((SECONDS < deadline)); do
    local 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
      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 -160 "$log_file" >&2 || true; return 1; fi
    sleep 5
  done
  tail -160 "$log_file" >&2 || true
  return 1
}
start_servers() {
  local tag=$1 rdt=$2 mc=$3 attempt=$4
  local replica port master_port scheduler_port first_gpu gpu_csv offset gpu server_dir server_log candidate
  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 busy: $candidate"
    done
    first_gpu=$((replica * TP_SIZE)); gpu_csv=""
    for ((offset=0; offset<TP_SIZE; offset++)); do
      gpu=$((first_gpu + offset)); [[ -z "$gpu_csv" ]] && gpu_csv="$gpu" || gpu_csv+=",$gpu"
    done
    server_dir="$CURRENT_PROFILE_DIR/server_attempt_${attempt}_${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 cache+lora TP$TP_SIZE profile=$tag attempt=$attempt replica=$replica GPUs=$gpu_csv port=$port RDT=$rdt MC=$mc"
    CUDA_VISIBLE_DEVICES="$gpu_csv" \
      PATH="$MEDIA_BIN_DIR:$PATH" \
      PYTHONPATH="$CACHE_PATCH${PYTHONPATH:+:$PYTHONPATH}" \
      PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
      PYTHONUNBUFFERED=1 TOKENIZERS_PARALLELISM=false SGLANG_USE_RUNAI_MODEL_STREAMER=false \
      SGLANG_CACHE_DIT_ENABLED=true SGLANG_CACHE_DIT_FN="$CACHE_FN" SGLANG_CACHE_DIT_BN="$CACHE_BN" \
      SGLANG_CACHE_DIT_RDT="$rdt" SGLANG_CACHE_DIT_MC="$mc" SGLANG_CACHE_DIT_WARMUP="$CACHE_WARMUP" \
      SGLANG_CACHE_DIT_TAYLORSEER=false SGLANG_CACHE_DIT_SCM_PRESET=none SGLANG_CACHE_DIT_SCM_POLICY=dynamic \
      lora_path="" setsid "$PYTHON" "$CACHE_PATCH/../scripts-nonexistent" 2>/dev/null; true
    CUDA_VISIBLE_DEVICES="$gpu_csv" \
      PATH="$MEDIA_BIN_DIR:$PATH" \
      PYTHONPATH="$CACHE_PATCH${PYTHONPATH:+:$PYTHONPATH}" \
      PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
      PYTHONUNBUFFERED=1 TOKENIZERS_PARALLELISM=false SGLANG_USE_RUNAI_MODEL_STREAMER=false \
      SGLANG_CACHE_DIT_ENABLED=true SGLANG_CACHE_DIT_FN="$CACHE_FN" SGLANG_CACHE_DIT_BN="$CACHE_BN" \
      SGLANG_CACHE_DIT_RDT="$rdt" SGLANG_CACHE_DIT_MC="$mc" SGLANG_CACHE_DIT_WARMUP="$CACHE_WARMUP" \
      SGLANG_CACHE_DIT_TAYLORSEER=false SGLANG_CACHE_DIT_SCM_PRESET=none SGLANG_CACHE_DIT_SCM_POLICY=dynamic \
      setsid "$SGLANG_BIN" serve --model-path "$MODEL" --model-variant Ref2VA --backend sglang \
      --performance-mode speed --num-gpus "$TP_SIZE" --tp-size "$TP_SIZE" --ulysses-degree 1 \
      --use-fsdp-inference false --enable-torch-compile false --batching-max-size 1 --batching-delay-ms 0 \
      --lora-path "$LORA_PATH" --lora-weight-name "$LORA_WEIGHT_NAME" --lora-nickname "$LORA_NICKNAME" \
      --lora-scale "$LORA_SCALE" --lora-merge-mode "$LORA_MERGE_MODE" \
      --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="$CURRENT_PROFILE_DIR/server_attempt_${attempt}_${replica}_port${port}/server.log"
    wait_healthy "$port" "${SERVER_PIDS[$replica]}" "$server_log" || return 1
    log "profile=$tag attempt=$attempt replica=$replica healthy port=$port"
  done
}
case_successful() {
  local case_number=$1 batch replica path
  batch=$((case_number <= 4 ? 1 : 2))
  replica=$(((case_number - 1) % REPLICAS))
  path="$CURRENT_PROFILE_DIR/batch_$(printf "%02d" "$batch")/client_${replica}_case$(printf "%02d" "$case_number")_port$((BASE_PORT + replica * PORT_STRIDE))/result.json"
  "$PYTHON" - "$path" <<"PY"
import json, pathlib, sys
p = pathlib.Path(sys.argv[1])
try:
    row = json.loads(p.read_text())
    video = pathlib.Path(str(row.get("file_path") or ""))
    ok = row.get("success") is True and video.is_file() and video.stat().st_size > 0
except Exception:
    ok = False
raise SystemExit(0 if ok else 1)
PY
}
run_batch_pending() {
  local tag=$1 batch=$2 first_case=$3 warmup_steps=$4
  local replica port case_number client_dir pid launched=0 failed=0
  CLIENT_PIDS=()
  for ((replica=0; replica<REPLICAS; replica++)); do
    case_number=$((first_case + replica))
    if case_successful "$case_number"; then
      log "profile=$tag batch=$batch case=$case_number already done; skip"
      continue
    fi
    port=$((BASE_PORT + replica * PORT_STRIDE))
    client_dir="$CURRENT_PROFILE_DIR/batch_$(printf "%02d" "$batch")/client_${replica}_case$(printf "%02d" "$case_number")_port${port}"
    mkdir -p "$client_dir"
    "$PYTHON" "$CLIENT_SCRIPT" run --host "$HOST" --port "$port" --replica-index "$replica" \
      --case-number "$case_number" --records "$RECORDS" --assets-root "$ASSETS_ROOT" \
      --model "$MODEL" --num-inference-steps "$NUM_INFERENCE_STEPS" --warmup-steps "$warmup_steps" \
      --short-edge "$SHORT_EDGE" --aspect-ratio "$ASPECT_RATIO" --duration-seconds "$DURATION_SECONDS" \
      --flow-shift "$FLOW_SHIFT" --audio-flow-shift "$AUDIO_FLOW_SHIFT" \
      --output "$client_dir/result.json" >"$client_dir/client.log" 2>&1 &
    pid=$!
    CLIENT_PIDS+=("$pid")
    launched=$((launched + 1))
    log "profile=$tag batch=$batch started case=$case_number replica=$replica port=$port"
  done
  ((launched > 0)) || return 0
  for pid in "${CLIENT_PIDS[@]}"; do wait "$pid" || failed=1; done
  CLIENT_PIDS=()
  ((failed == 0)) || { log "profile=$tag batch=$batch had failed requests"; return 1; }
}
completed_count() {
  local case_number count=0
  for case_number in {1..8}; do case_successful "$case_number" && count=$((count + 1)); done
  printf "%s\n" "$count"
}
run_profile() {
  local tag=$1 rdt=$2 mc=$3 objective=$4 attempt done_count pass_failed
  CURRENT_PROFILE_DIR="$RESULT_ROOT/$tag"
  mkdir -p "$CURRENT_PROFILE_DIR"
  [[ -f "$CURRENT_PROFILE_DIR/PROFILE_COMPLETED" ]] && { log "profile=$tag already done"; return 0; }
  {
    printf "profile=%s\nobjective=%s\n" "$tag" "$objective"
    printf "rdt=%s\nmc=%s\nfn=%s\nbn=%s\ncache_warmup=%s\nlora_merge_mode=%s\ntp_size=%s\nreplicas=%s\n" \
      "$rdt" "$mc" "$CACHE_FN" "$CACHE_BN" "$CACHE_WARMUP" "$LORA_MERGE_MODE" "$TP_SIZE" "$REPLICAS"
    printf "num_inference_steps=%s\nshort_edge=%s\naspect_ratio=%s\nduration_seconds=%s\n" \
      "$NUM_INFERENCE_STEPS" "$SHORT_EDGE" "$ASPECT_RATIO" "$DURATION_SECONDS"
  } >"$CURRENT_PROFILE_DIR/run_config.env"
  done_count=$(completed_count)
  log "===== profile=$tag RDT=$rdt MC=$mc existing_success=$done_count/8 ====="
  for ((attempt=1; attempt<=MAX_ATTEMPTS && done_count<8; attempt++)); do
    start_servers "$tag" "$rdt" "$mc" "$attempt" || { stop_servers; sleep 10; continue; }
    pass_failed=0
    launched=0
    # REPLICAS covers subset of the 8 cases per server pass; loop 4 chunks.
    for ((slice=0; slice<8; slice+=REPLICAS)); do
      for ((r=0; r<REPLICAS; r++)); do
        case_number=$((slice + r + 1))
        if case_successful "$case_number"; then
          log "profile=$tag case=$case_number already done; skip"
          continue
        fi
        batch=$(( case_number <= 4 ? 1 : 2 ))
        port=$((BASE_PORT + r * PORT_STRIDE))
        client_dir="$CURRENT_PROFILE_DIR/batch_$(printf "%02d" "$batch")/client_${r}_case$(printf "%02d" "$case_number")_port${port}"
        mkdir -p "$client_dir"
        "$PYTHON" "$CLIENT_SCRIPT" run --host "$HOST" --port "$port" --replica-index "$r" \
          --case-number "$case_number" --records "$RECORDS" --assets-root "$ASSETS_ROOT" \
          --model "$MODEL" --num-inference-steps "$NUM_INFERENCE_STEPS" --warmup-steps 0 \
          --short-edge "$SHORT_EDGE" --aspect-ratio "$ASPECT_RATIO" --duration-seconds "$DURATION_SECONDS" \
          --flow-shift "$FLOW_SHIFT" --audio-flow-shift "$AUDIO_FLOW_SHIFT" \
          --output "$client_dir/result.json" >"$client_dir/client.log" 2>&1 &
        pid=$!
        CLIENT_PIDS+=("$pid")
        launched=$((launched + 1))
        log "profile=$tag started case=$case_number replica=$r port=$port"
      done
      for pid in "${CLIENT_PIDS[@]}"; do wait "$pid" || pass_failed=1; done
      CLIENT_PIDS=()
    done
    stop_servers
    done_count=$(completed_count)
    log "profile=$tag attempt=$attempt success=$done_count/8 pass_failed=$pass_failed"
  done
  ((done_count == 8)) || die "profile=$tag incomplete: $done_count/8"
  "${CLIENT_SCRIPT%%bench.py}custom_summary" 2>/dev/null || true
  printf "%s\n" "$(date -Is)" >"$CURRENT_PROFILE_DIR/PROFILE_COMPLETED"
  log "profile=$tag complete 8/8"
}
"$PYTHON" "$CLIENT_SCRIPT" summarize --help >/dev/null 2>&1 || true
for profile in "${PROFILES[@]}"; do
  read -r tag rdt mc objective <<< "$profile"
  run_profile "$tag" "$rdt" "$mc" "$objective"
done
trap - EXIT INT TERM
log "CACHE+LORA TP4x2 ALL PROFILES COMPLETE: $RESULT_ROOT"