run_evals_parallel.sh

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

# ==============================================================================
# Run BFCL evals across N checkpoints on two parallel GPU lanes.
#
#   Lane A: GPUs 0,1,2,3  (port 8000, container ytahtah-vllm-eval-A)
#   Lane B: GPUs 4,5,6,7  (port 8001, container ytahtah-vllm-eval-B)
#
# Checkpoints are split: ceil(N/2) to lane A, floor(N/2) to lane B
# (lane A gets the extra one when N is odd). Within a lane, checkpoints
# run sequentially; each one boots vLLM, does --runs evaluations, then
# the container is torn down before the next checkpoint.
#
# On a failure within a lane, the lane's vLLM container is force-killed
# and the lane continues with the next checkpoint.
#
# checkpoints.txt format (one per line, '#' comments and blank lines OK):
#   <name>|<model_path>
#   <name>|<model_path>|<bfcl_key>       # optional 3rd field overrides
#                                          --bfcl-model-key for this line only
# <name> is used as the output suffix (eval_results/eval_<name>_fc).
# <model_path> is an HF id or absolute local path.
#
# Usage:
#   ./run_evals_parallel.sh                     # checkpoints.txt, defaults
#   ./run_evals_parallel.sh --runs 3
#   ./run_evals_parallel.sh --checkpoints my_list.txt
# ==============================================================================

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
RUN_EVAL="${SCRIPT_DIR}/run_eval.sh"

# Defaults
RUNS=5
SERVE_NAME="allenai/Olmo-3-7B-Instruct-SFT"
BFCL_MODEL_KEY="allenai/Olmo-3-7B-Instruct-SFT-FC"
CHECKPOINTS_FILE="${SCRIPT_DIR}/checkpoints.txt"
TEST_CATEGORY="all"        # pass-through to run_eval.sh; matches its default
OUTPUT_SUFFIX=""           # appended to eval_results/eval_<name>_fc<suffix>

# Hardcoded lanes
LANE_A_GPUS="0,1,2,3"
LANE_A_PORT=8000
LANE_A_CONTAINER="ytahtah-vllm-eval-A"
LANE_B_GPUS="4,5,6,7"
LANE_B_PORT=8001
LANE_B_CONTAINER="ytahtah-vllm-eval-B"

# Per-lane cache dirs (avoid contention on shared compile/triton/tmp dirs)
LANE_A_VLLM_CACHE="/mnt/nfs/ytahtah/.cache/vllm_compile/lane_a"
LANE_A_TMP="/mnt/nfs/ytahtah/tmp_compile/lane_a"
LANE_A_TRITON_CACHE="/mnt/nfs/ytahtah/.triton_cache/lane_a"
LANE_B_VLLM_CACHE="/mnt/nfs/ytahtah/.cache/vllm_compile/lane_b"
LANE_B_TMP="/mnt/nfs/ytahtah/tmp_compile/lane_b"
LANE_B_TRITON_CACHE="/mnt/nfs/ytahtah/.triton_cache/lane_b"

# Per-lane BFCL_PROJECT_ROOT (where bfcl writes its result/, score/,
# .file_locks/). Without this, both lanes would share the same dirs and
# trample each other's intermediate state.
LANE_A_PROJECT_ROOT="${SCRIPT_DIR}/.workspaces/lane_a"
LANE_B_PROJECT_ROOT="${SCRIPT_DIR}/.workspaces/lane_b"

usage() {
    cat <<'USAGE'
Usage: ./run_evals_parallel.sh [options]

Options:
  --runs N                Number of eval rounds per checkpoint (default 5)
  --serve-name NAME       vLLM API name (default allenai/Olmo-3-7B-Instruct-SFT)
  --bfcl-model-key KEY    BFCL model key (default allenai/Olmo-3-7B-Instruct-SFT-FC)
  --checkpoints FILE      Checkpoints list (default ./checkpoints.txt)
  --test-category CATS    Comma-separated categories or collection name
                          (default 'all'; see category_mapping.py)
  --output-suffix STR     Appended to eval_results/eval_<name>_fc<suffix>.
                          Use to avoid clobbering a prior run (default '').
  -h, --help              Show this help
USAGE
    exit 0
}

while [[ $# -gt 0 ]]; do
    case $1 in
        --runs)            RUNS="$2"; shift 2 ;;
        --serve-name)      SERVE_NAME="$2"; shift 2 ;;
        --bfcl-model-key)  BFCL_MODEL_KEY="$2"; shift 2 ;;
        --checkpoints)     CHECKPOINTS_FILE="$2"; shift 2 ;;
        --test-category)   TEST_CATEGORY="$2"; shift 2 ;;
        --output-suffix)   OUTPUT_SUFFIX="$2"; shift 2 ;;
        -h|--help)         usage ;;
        *) echo "Unknown option: $1" >&2; exit 1 ;;
    esac
done

[[ -f "$CHECKPOINTS_FILE" ]] || { echo "Checkpoints file not found: $CHECKPOINTS_FILE" >&2; exit 1; }
[[ -x "$RUN_EVAL" ]]         || { echo "run_eval.sh not found or not executable: $RUN_EVAL" >&2; exit 1; }

# Refuse to start if the web-search driver's containers are live — both drivers
# share GPUs 0-7 and ports 8000/8001, so running them at once collides. This
# mirrors the guard run_web_search.sh has against ytahtah-vllm-eval-A/B.
for c in ytahtah-vllm-websearch-A ytahtah-vllm-websearch-B; do
    if docker ps --format '{{.Names}}' 2>/dev/null | grep -qx "$c"; then
        echo "ERROR: container '${c}' is running — run_web_search.sh appears active." >&2
        echo "  These scripts share GPUs and ports. Stop the web-search run first." >&2
        exit 1
    fi
done

# Parse checkpoints (skip comment/blank lines). sed strips trailing CR so a
# CRLF-saved checkpoints file doesn't poison the last field.
mapfile -t CHECKPOINTS < <(grep -vE '^[[:space:]]*(#|$)' "$CHECKPOINTS_FILE" | sed 's/\r$//')
N=${#CHECKPOINTS[@]}
(( N > 0 )) || { echo "No checkpoints in $CHECKPOINTS_FILE" >&2; exit 1; }

# Split: lane A gets ceil(N/2), lane B gets floor(N/2)
HALF=$(( (N + 1) / 2 ))
LANE_A_CKPTS=("${CHECKPOINTS[@]:0:$HALF}")
LANE_B_CKPTS=()
(( HALF < N )) && LANE_B_CKPTS=("${CHECKPOINTS[@]:$HALF}")

mkdir -p "${SCRIPT_DIR}/logs"
TS=$(date +%Y%m%d_%H%M%S)
LANE_A_LOG="${SCRIPT_DIR}/logs/lane_a_${TS}.log"
LANE_B_LOG="${SCRIPT_DIR}/logs/lane_b_${TS}.log"
# Per-lane status files (ok|name / fail|name per checkpoint) for the roll-up.
LANE_A_STATUS="${SCRIPT_DIR}/logs/lane_status_a_${TS}.tsv"
LANE_B_STATUS="${SCRIPT_DIR}/logs/lane_status_b_${TS}.tsv"
: > "$LANE_A_STATUS"
: > "$LANE_B_STATUS"

echo "================================================================"
echo "Parallel evaluation across 2 lanes"
echo "  N checkpoints      : ${N}"
echo "  Lane A (GPUs ${LANE_A_GPUS}) : ${#LANE_A_CKPTS[@]} checkpoints  -> ${LANE_A_LOG}"
echo "  Lane B (GPUs ${LANE_B_GPUS}) : ${#LANE_B_CKPTS[@]} checkpoints  -> ${LANE_B_LOG}"
echo "  Runs per checkpoint: ${RUNS}"
echo "  Serve name         : ${SERVE_NAME}"
echo "  BFCL model key     : ${BFCL_MODEL_KEY}"
echo "  Test categories    : ${TEST_CATEGORY}"
echo "  Output suffix      : '${OUTPUT_SUFFIX}'  (eval_results/eval_<name>_fc${OUTPUT_SUFFIX}/)"
echo "================================================================"

run_lane() {
    local lane="$1" gpus="$2" port="$3" container="$4"
    local vllm_cache="$5" tmp_dir="$6" triton_cache="$7" project_root="$8" status_file="$9"
    shift 9
    local ckpts=("$@")

    if [[ ${#ckpts[@]} -eq 0 ]]; then
        echo "[lane $lane] no checkpoints, skipping"
        return 0
    fi

    local idx=0 total=${#ckpts[@]} fail=0 ok=0
    for entry in "${ckpts[@]}"; do
        idx=$((idx + 1))
        # 3rd field is an optional BFCL key override. Empty -> use global default.
        IFS='|' read -r name model_path line_bfcl_key <<< "$entry"
        local effective_bfcl_key="${line_bfcl_key:-$BFCL_MODEL_KEY}"
        local output_dir="eval_results/eval_${name}_fc${OUTPUT_SUFFIX}"

        echo ""
        echo "──[lane $lane] [${idx}/${total}] ${name} (${model_path}) [bfcl-key=${effective_bfcl_key}]──"

        local -a args=(
            --model "$model_path"
            --bfcl-model-key "$effective_bfcl_key"
            --runs "$RUNS"
            --gpus "$gpus"
            --port "$port"
            --container-name "$container"
            --output-dir "$output_dir"
            --test-category "$TEST_CATEGORY"
            --vllm-cache-dir "$vllm_cache"
            --tmp-dir "$tmp_dir"
            --triton-cache-dir "$triton_cache"
            --project-root "$project_root"
        )
        # Local path -> need --serve-name; HF id -> default to model
        [[ "$model_path" == /* ]] && args+=(--serve-name "$SERVE_NAME")

        if "$RUN_EVAL" "${args[@]}"; then
            ok=$((ok + 1))
            echo "ok|${name}" >> "$status_file"
            echo "[lane $lane] ✓ ${name}"
        else
            fail=$((fail + 1))
            echo "fail|${name}" >> "$status_file"
            echo "[lane $lane] ✗ ${name} — killing container '${container}' before next checkpoint"
            docker rm -f "$container" >/dev/null 2>&1 || true
        fi
    done

    echo ""
    echo "[lane $lane] done: ${ok}/${total} ok, ${fail} failed"
    # Surface failure: any failed checkpoint makes this lane exit nonzero so the
    # driver's overall exit reflects it (no more silent green on total failure).
    (( fail > 0 )) && return 1
    return 0
}

# Always tear down lane subshells + vLLM containers when this script exits,
# whatever the cause: normal end, error, signal, lane crash. Idempotent so
# it's safe even when run_eval.sh already cleaned up its own container.
cleanup_on_exit() {
    local code=$?
    [[ -n "${PID_A:-}" ]] && kill "$PID_A" 2>/dev/null || true
    [[ -n "${PID_B:-}" ]] && kill "$PID_B" 2>/dev/null || true
    docker rm -f "$LANE_A_CONTAINER" "$LANE_B_CONTAINER" >/dev/null 2>&1 || true
    exit "$code"
}
on_interrupt() {
    echo ""
    echo "Interrupted (SIGINT/SIGTERM) — cleaning up..."
    exit 130   # triggers cleanup_on_exit via the EXIT trap
}
trap cleanup_on_exit EXIT
trap on_interrupt INT TERM

# Startup sweep: kill any leftover lane containers from a previous crashed run
docker rm -f "$LANE_A_CONTAINER" "$LANE_B_CONTAINER" >/dev/null 2>&1 || true

# Pre-create per-lane dirs so docker doesn't create them as root
mkdir -p "$LANE_A_VLLM_CACHE" "$LANE_A_TMP" "$LANE_A_TRITON_CACHE" \
         "$LANE_B_VLLM_CACHE" "$LANE_B_TMP" "$LANE_B_TRITON_CACHE" \
         "$LANE_A_PROJECT_ROOT" "$LANE_B_PROJECT_ROOT"

run_lane "A" "$LANE_A_GPUS" "$LANE_A_PORT" "$LANE_A_CONTAINER" \
    "$LANE_A_VLLM_CACHE" "$LANE_A_TMP" "$LANE_A_TRITON_CACHE" \
    "$LANE_A_PROJECT_ROOT" "$LANE_A_STATUS" \
    "${LANE_A_CKPTS[@]}" \
    > "$LANE_A_LOG" 2>&1 &
PID_A=$!

if [[ ${#LANE_B_CKPTS[@]} -gt 0 ]]; then
    run_lane "B" "$LANE_B_GPUS" "$LANE_B_PORT" "$LANE_B_CONTAINER" \
        "$LANE_B_VLLM_CACHE" "$LANE_B_TMP" "$LANE_B_TRITON_CACHE" \
        "$LANE_B_PROJECT_ROOT" "$LANE_B_STATUS" \
        "${LANE_B_CKPTS[@]}" \
        > "$LANE_B_LOG" 2>&1 &
    PID_B=$!
else
    PID_B=""
fi

echo ""
echo "Both lanes launched. Tail logs with:"
echo "  tail -f ${LANE_A_LOG} ${LANE_B_LOG}"
echo ""

wait "$PID_A"; A_STATUS=$?
B_STATUS=0
[[ -n "$PID_B" ]] && { wait "$PID_B"; B_STATUS=$?; }

# ── Roll-up: aggregate per-checkpoint outcomes from both lanes ───────────────
TOTAL_OK=0 TOTAL_FAIL=0
FAILED_NAMES=()
for sf in "$LANE_A_STATUS" "$LANE_B_STATUS"; do
    [[ -f "$sf" ]] || continue
    while IFS='|' read -r st nm; do
        case "$st" in
            ok)   TOTAL_OK=$((TOTAL_OK + 1)) ;;
            fail) TOTAL_FAIL=$((TOTAL_FAIL + 1)); FAILED_NAMES+=("$nm") ;;
        esac
    done < "$sf"
done

echo ""
echo "================================================================"
echo "All lanes complete."
echo "  Lane A exit: ${A_STATUS}  (log: ${LANE_A_LOG})"
echo "  Lane B exit: ${B_STATUS}  (log: ${LANE_B_LOG})"
echo "  Checkpoints: ${TOTAL_OK} ok, ${TOTAL_FAIL} failed"
if (( TOTAL_FAIL > 0 )); then
    echo "  ✗ FAILED checkpoints (need investigation / re-run):"
    for nm in "${FAILED_NAMES[@]}"; do echo "      - ${nm}"; done
fi
echo "================================================================"

# Overall exit is nonzero if ANY checkpoint failed or a lane wrapper errored —
# never silently green when a checkpoint didn't complete.
OVERALL=$(( A_STATUS | B_STATUS ))
(( TOTAL_FAIL > 0 )) && OVERALL=1
exit "$OVERALL"