#!/usr/bin/env bash
# Orchestrator for the Qwen3-235B-A22B TP=8 benchmark on Kunlun P800.
#
# Flow: write metadata -> start server (start_server.sh) -> warmup+bench_serving
# (isl=2048 osl=2048 c=16 n=160) -> parse results -> report.md.
#
# The bench client (sglang.bench_serving) runs INSIDE the container; results are
# written to a /data1 path that is mounted same-path into the container, so the
# host-side parse_results.py can read them directly.
set -Eeuo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
EXPERIMENT_NAME="$(basename "$SCRIPT_DIR")"

# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../../scripts/common/lib.sh"
# Source config.env BEFORE platform.sh so our CONTAINER_NAME default wins
# over the platform's "sglang-dsv4-flash" default (avoids name collisions
# with the dsv4 experiments, which would docker rm -f this container).
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/config.env"
# shellcheck source=/dev/null
source "${SCRIPT_DIR}/../../../scripts/common/platform.sh"

RUN_ID="${RUN_ID:-$(date '+%Y%m%d-%H%M%S')}"
RESULT_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}"
RAW_DIR="${RESULT_ROOT}/raw_outputs"
LOG_DIR="${RESULT_ROOT}/logs"
mkdir -p "$RAW_DIR" "$LOG_DIR"

ensure_result_root "$RESULT_ROOT"
log_init "${LOG_DIR}/orchestrator.log"

log "experiment=${EXPERIMENT_NAME}"
log "run_id=${RUN_ID}"
log "result_root=${RESULT_ROOT}"
log "platform=${PLATFORM} chip=${CHIP} hardware=${HARDWARE}"
log "model=${MODEL_PATH} tp=${TP} backend=${BACKEND}"
log "scenarios=${SCENARIOS[*]}"

# --- metadata --------------------------------------------------------------
METADATA_JSON="${RESULT_ROOT}/results.json"
write_metadata_json \
  "$METADATA_JSON" \
  "$EXPERIMENT_NAME" \
  "$RUN_ID" \
  "$MODEL_PATH" \
  "$BACKEND" \
  "$ENGINE" \
  "$HARDWARE" \
  "$ACCELERATOR" \
  "$CHIP" \
  "experiments/${EXPERIMENT_NAME}/run_bench.sh" \
  "" \
  "Qwen3-235B-A22B SGLang TP=8 benchmark on Kunlun P800 (isl=osl=2048, c=16)"

"${PYTHON:-python3}" - "$METADATA_JSON" "$TP" "$CONTEXT_LENGTH" "$PORT" "$CONTAINER_NAME" "$DOCKER_IMAGE" "$MEM_FRACTION_STATIC" "$DTYPE" "$ATTENTION_BACKEND" <<'PY'
import json, sys
p, tp, ctx, port, cont, img, memfrac, dtype, attn = sys.argv[1:10]
with open(p) as f:
    d = json.load(f)
d["config"] = {
    "tp": int(tp),
    "context_length": int(ctx),
    "port": int(port),
    "container": cont,
    "docker_image": img,
    "attention_backend": attn,
    "dtype": dtype,
    "mem_fraction_static": float(memfrac),
    "max_running_requests": 16,
    "page_size": 64,
    "benchmark": {"input_len": 2048, "output_len": 2048, "concurrency": 16, "num_prompts": 160},
}
with open(p, "w") as f:
    json.dump(d, f, indent=2, ensure_ascii=False)
PY

# --- start server (unless caller manages it themselves) --------------------
if [[ -n "${SKIP_MANAGE_SERVER:-}" ]]; then
  log "SKIP_MANAGE_SERVER set; assuming server already running on port ${PORT}"
  if ! health_check 127.0.0.1 "$PORT" 30; then
    log "ERROR: no healthy server at port ${PORT}"
    exit 1
  fi
else
  log "starting server via start_server.sh"
  bash "${SCRIPT_DIR}/start_server.sh" >> "${LOG_DIR}/start_server.outer.log" 2>&1
fi

on_exit() {
  local code=$?
  if [[ -z "${SKIP_MANAGE_SERVER:-}" ]]; then
    log "stopping sglang inside ${CONTAINER_NAME} (container kept alive)"
    docker exec "$CONTAINER_NAME" bash -lc 'pkill -9 -f sglang.launch_server 2>/dev/null || true' || true
  fi
  log "orchestrator exiting with code=${code}"
  exit "$code"
}
trap on_exit EXIT

# --- run scenarios ---------------------------------------------------------
log "===== BENCHMARK START ====="

for scenario in "${SCENARIOS[@]}"; do
  scenario="${scenario//[()]/}"
  read -r concurrency input_len output_len num_prompts <<< "$scenario"

  today="$(date '+%m%d')"
  # Filename convention expected by parse_results.py:
  # {chip}_{engine}_{MMDD}_{concurrency}_{input_len}_{output_len}.jsonl
  output_file="${RAW_DIR}/${CHIP}_${ENGINE}_${today}_${concurrency}_${input_len}_${output_len}.jsonl"
  # parse_results.py reads the human-readable summary from this exact path.
  summary_log="${LOG_DIR}/sglang_c${concurrency}_i${input_len}_o${output_len}.log"

  log "running scenario: c=${concurrency} i=${input_len} o=${output_len} n=${num_prompts}"
  log "output_file=${output_file}"

  docker exec "$CONTAINER_NAME" bash -lc "
    source /root/miniconda/etc/profile.d/conda.sh
    conda activate python310_torch25_cuda
    python -m sglang.bench_serving \
      --backend ${BACKEND} \
      --host 127.0.0.1 \
      --port ${PORT} \
      --model ${MODEL_PATH} \
      --dataset-name random \
      --dataset-path ${DATASET_PATH} \
      --random-input-len ${input_len} \
      --random-output-len ${output_len} \
      --random-range-ratio ${RANDOM_RANGE_RATIO} \
      --num-prompts ${num_prompts} \
      --warmup-requests ${WARMUP} \
      --max-concurrency ${concurrency} \
      --request-rate ${REQUEST_RATE} \
      --output-file ${output_file} \
      --output-details \
      --seed ${SEED}
  " > "$summary_log" 2>&1 || {
    log "ERROR: scenario c=${concurrency} i=${input_len} o=${output_len} failed; see ${summary_log}"
    tail -60 "$summary_log" 2>/dev/null || true
    continue
  }

  log "finished scenario: output=${output_file}"
  log "----- summary -----"
  tail -40 "$summary_log" | sed 's/^/  /'
done

log "===== BENCHMARK DONE ====="
log "parsing results..."
"${PYTHON:-python3}" "${SCRIPT_DIR}/parse_results.py" "$RESULT_ROOT" || log "WARNING: parser failed"

log "all results saved to ${RESULT_ROOT}"
log "report: ${RESULT_ROOT}/report.md"