#!/usr/bin/env bash
set -euo pipefail

# Launch one source-partitioned 2k evaluation per retrieval GPU.  Every worker
# uses the same four Qwen endpoints but owns exactly one BGE encoder/reranker,
# avoiding cross-process CUDA contention.
ROOT="${ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
UV="${UV:-/home/clouduser/miniconda3/envs/vllm-nightly/bin/uv}"
SERVER_RUN_DIR="${SERVER_RUN_DIR:?Set to artifacts/servers/<server-run>}"
RUN_GROUP="${RUN_GROUP:-eval2000_$(date -u +%Y%m%d_%H%M%S)}"
CONCURRENCY="${CONCURRENCY:-16}"
EXPERIMENT="${EXPERIMENT:-configs/experiments/eval_2000.yaml}"
MAX_EXAMPLES="${MAX_EXAMPLES:-}"
# Learned BGE retrieval needs a GPU distinct from 32k-context Qwen replicas.
# One H100 can safely host the three source-partitioned BGE workers below.
RETRIEVAL_GPUS="${RETRIEVAL_GPUS:-5,5,5}"
PARALLEL_DATASETS="${PARALLEL_DATASETS:-0}"
cd "$ROOT"

# shellcheck disable=SC1090
source "$SERVER_RUN_DIR/endpoints.env"
if [[ -n "${MODEL_BASE_URLS_OVERRIDE:-}" ]]; then
  export MODEL_BASE_URLS="$MODEL_BASE_URLS_OVERRIDE"
fi
export MODEL_API_KEY="${MODEL_API_KEY:-EMPTY}"

datasets=(hotpotqa 2wikimultihopqa musique)
IFS=, read -r -a gpus <<< "$RETRIEVAL_GPUS"
if (( ${#gpus[@]} != ${#datasets[@]} )); then
  echo "RETRIEVAL_GPUS must provide exactly ${#datasets[@]} comma-separated GPU IDs" >&2
  exit 2
fi
run_dataset() {
  local dataset="$1"
  local gpu="$2"
  local log="artifacts/eval_logs/${RUN_GROUP}_${dataset}.log"
  local max_examples_args=()
  if [[ -n "$MAX_EXAMPLES" ]]; then
    max_examples_args=(--max-examples "$MAX_EXAMPLES")
  fi
  mkdir -p "$(dirname "$log")"
  # Three independent PyTorch/FAISS workers on one GPU can exceed the host's
  # BLAS memory-map/thread limit.  One source at a time is reproducible and
  # still lets the 16 benchmark workers use all Qwen replicas.
  OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 NUMEXPR_NUM_THREADS=1 \
    PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
    CUDA_VISIBLE_DEVICES="$gpu" "$UV" run agentic-search benchmark run \
      --examples "data/processed/benchmarks/${dataset}.jsonl" \
      --corpus data/processed/multihop_context_union.jsonl \
      --experiment "$EXPERIMENT" \
      "${max_examples_args[@]}" \
      --concurrency "$CONCURRENCY" \
      --no-cache \
      --run-id "${RUN_GROUP}_${dataset}" \
      >"$log" 2>&1
}

pids=()
for index in "${!datasets[@]}"; do
  dataset="${datasets[$index]}"
  gpu="${gpus[$index]}"
  if [[ "$PARALLEL_DATASETS" == "1" ]]; then
    run_dataset "$dataset" "$gpu" &
    pids+=("$!")
  else
    run_dataset "$dataset" "$gpu"
  fi
done
if (( ${#pids[@]} )); then
  printf '%s\n' "${pids[@]}" > "artifacts/eval_logs/${RUN_GROUP}.pids"
fi
printf 'run_group=%s\nlogs=artifacts/eval_logs\n' "$RUN_GROUP"
if (( ${#pids[@]} )); then
  wait "${pids[@]}"
fi