#!/usr/bin/env bash
# Orchestrator for the Qwen3-235B-A22B TP=8 benchmark on NVIDIA RTX 6000D.
# Cross-machine counterpart of experiments/p800/qwen3_235b_p800_sglang_tp8/run_bench.sh.
#
# Flow: write metadata -> start server (start_server.sh) -> bench_serving
# (isl=2048 osl=2048 c=16 n=160, inside the container) -> parse -> report.md.
# The bench client runs via `docker exec` into the running server container;
# /data is mounted same-path so results + dataset are reachable inside.
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"
# 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"
RUNTIME_BASE="${RUNTIME_BASE:-${SCRIPT_DIR}/runtime}"
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} image=${DOCKER_IMAGE}"
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 NVIDIA RTX 6000D (isl=osl=2048, c=16) - cross-machine vs p800.2"

"${PYTHON:-python3}" - "$METADATA_JSON" "$TP" "$CONTEXT_LENGTH" "$PORT" "$CONTAINER_NAME" "$DOCKER_IMAGE" "$MEM_FRACTION_STATIC" "$DTYPE" <<'PY'
import json, sys
p, tp, ctx, port, cont, img, memfrac, dtype = sys.argv[1:9]
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,
    "dtype": dtype,
    "mem_fraction_static": float(memfrac),
    "max_running_requests": 16,
    "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

stop_server() {
  if [[ -z "${SKIP_MANAGE_SERVER:-}" ]]; then
    log "stopping server (kill docker run PID -> --rm removes container)"
    if [[ -f "${RUNTIME_BASE}/${EXPERIMENT_NAME}.pid" ]]; then
      local pid
      pid="$(cat "${RUNTIME_BASE}/${EXPERIMENT_NAME}.pid")"
      kill "$pid" 2>/dev/null || true
      sleep 3
      kill -9 "$pid" 2>/dev/null || true
    fi
    docker rm -f "$CONTAINER_NAME" >/dev/null 2>&1 || true
  fi
}
on_exit() {
  local code=$?
  stop_server
  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"
  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" python3 -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"