queue_bfcl.sh

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

# =============================================================================
# queue_bfcl.sh — unified launcher for BFCL evals.
#
# Replaces queue_bfcl_after_gpus_free.sh and queue_websearch_after_gpus_free.sh.
#
#   --phases {both|main|ws}        which BFCL phases to run (default both)
#   --no-wait                      skip the GPU-free gate; start immediately
#
# Phases:
#   main = run_evals_parallel.sh (5x, --main-categories) over --checkpoints-main
#   ws   = run_web_search.sh    (--ws-runs x)            over --checkpoints-ws
#
# GPU GATING — by default the script waits until ALL GPUs are free (idle for
# STABLE_SECONDS continuous), re-checked immediately before EACH phase, then
# claims them. It waits no matter WHAT occupies the GPUs — another user's job,
# a stray container, anything — and never kills anything. --no-wait skips it.
#
# SAFETY — the per-phase re-check matters because a foreign job can appear
# during a multi-hour phase and leave the GPUs busy for the next one. The child
# drivers exit nonzero and print a failed-checkpoint roll-up when any checkpoint
# fails; this script surfaces that and exits nonzero too — never a silent green.
#
# Run under tmux and leave it:
#   tmux new -s bfcl ; ./queue_bfcl.sh [opts] ; # Ctrl-b d
# =============================================================================

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"

# Defaults
PHASES="both"
NO_WAIT=false           # --no-wait skips the GPU-free gate and starts immediately
POLL_SECONDS=30
STABLE_SECONDS=120
MEM_FREE_THRESHOLD_MIB=2000
MAIN_CATEGORIES="single_turn,multi_turn,memory,format_sensitivity"
WS_RUNS=5
CHECKPOINTS_MAIN="${SCRIPT_DIR}/checkpoints.txt"
CHECKPOINTS_WS="${SCRIPT_DIR}/checkpoints_websearch.txt"

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

  --phases {both|main|ws}        phases to run (default both)
  --no-wait                      skip the GPU-free gate; start immediately
                                 (default: wait until all GPUs are free)
  --ws-runs N                    web-search runs per checkpoint (default 5)
  --main-categories CATS         main eval categories
                                 (default single_turn,multi_turn,memory,format_sensitivity)
  --checkpoints-main FILE        main eval checkpoints (default checkpoints.txt)
  --checkpoints-ws FILE          web-search checkpoints (default checkpoints_websearch.txt)
  --stable-seconds N             continuous-free window before a phase (default 120)
  --poll-seconds N               GPU poll interval (default 30)
  -h, --help                     show this help
USAGE
    exit 0
}

while [[ $# -gt 0 ]]; do
    case "$1" in
        --no-wait)          NO_WAIT=true; shift ;;
        --phases)           PHASES="$2"; shift 2 ;;
        --ws-runs)          WS_RUNS="$2"; shift 2 ;;
        --main-categories)  MAIN_CATEGORIES="$2"; shift 2 ;;
        --checkpoints-main) CHECKPOINTS_MAIN="$2"; shift 2 ;;
        --checkpoints-ws)   CHECKPOINTS_WS="$2"; shift 2 ;;
        --stable-seconds)   STABLE_SECONDS="$2"; shift 2 ;;
        --poll-seconds)     POLL_SECONDS="$2"; shift 2 ;;
        -h|--help)          usage ;;
        *) echo "Unknown option: $1" >&2; exit 1 ;;
    esac
done

case "$PHASES" in both|main|ws) ;; *) echo "ERROR: --phases must be both|main|ws" >&2; exit 1 ;; esac

mkdir -p "${SCRIPT_DIR}/logs"
LOG="${SCRIPT_DIR}/logs/queued_bfcl_$(date +%Y%m%d_%H%M%S).log"
log() { printf '[%s] %s\n' "$(date '+%Y-%m-%d %H:%M:%S')" "$*" | tee -a "$LOG"; }

# Echo a single status token describing GPU occupancy:
#   "FREE maxmem=<n>MiB" | "BUSY procs=<n>" | "BUSY maxmem=<n>MiB" | "UNKNOWN <reason>"
# Two signals, busy if EITHER says busy:
#   (a) compute-process count. NOTE: on some nodes `--query-compute-apps` lists
#       only YOUR processes (permission-restricted), so this can read 0 even
#       when a foreign job is present — it is a best-effort signal, not a
#       guarantee. (On this node it does show other users' procs.)
#   (b) per-GPU used memory vs threshold. This is the RELIABLE signal: a
#       foreign job's VRAM shows up regardless of process visibility.
# Only "FREE ..." counts as free; anything else (incl. UNKNOWN) means busy.
gpu_state() {
    local proc_out rc nproc mem_out maxmem=0 m
    proc_out=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null)
    rc=$?
    if (( rc != 0 )); then echo "UNKNOWN nvidia-smi(compute-apps)_rc=${rc}"; return; fi
    nproc=$(printf '%s\n' "$proc_out" | grep -c '[0-9]')
    if (( nproc > 0 )); then echo "BUSY procs=${nproc}"; return; fi
    mem_out=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null)
    rc=$?
    if (( rc != 0 )); then echo "UNKNOWN nvidia-smi(memory)_rc=${rc}"; return; fi
    while read -r m; do
        m="${m//[[:space:]]/}"
        [[ -z "$m" ]] && continue
        [[ "$m" =~ ^[0-9]+$ ]] || { echo "UNKNOWN mem_parse='${m}'"; return; }
        (( m > maxmem )) && maxmem=$m
    done <<< "$mem_out"
    if (( maxmem > MEM_FREE_THRESHOLD_MIB )); then echo "BUSY maxmem=${maxmem}MiB"; return; fi
    echo "FREE maxmem=${maxmem}MiB"
}

# Block until all GPUs have been continuously FREE for STABLE_SECONDS.
wait_for_gpus_free() {
    local label="$1" free_since="" now state free_for
    log "── Gate: waiting for GPUs free before ${label} (need ${STABLE_SECONDS}s continuous) ──"
    while true; do
        now=$(date +%s); state=$(gpu_state)
        if [[ "$state" == FREE* ]]; then
            if [[ -z "$free_since" ]]; then
                free_since=$now
                log "GPUs FREE (${state}) — starting ${STABLE_SECONDS}s timer [${label}]"
            fi
            free_for=$(( now - free_since ))
            log "GPUs free for ${free_for}s / ${STABLE_SECONDS}s (${state}) [${label}]"
            if (( free_for >= STABLE_SECONDS )); then
                log "GPUs free ${STABLE_SECONDS}s continuous — proceeding to ${label}"
                return 0
            fi
        else
            if [[ -n "$free_since" ]]; then
                log "GPUs busy again (${state}) — resetting timer [${label}]"
            else
                log "GPUs busy (${state}) [${label}]"
            fi
            free_since=""
        fi
        sleep "$POLL_SECONDS"
    done
}

# ── Cleanup: kill the in-flight child driver and free GPUs on exit/interrupt ──
# Ctrl-C reaches the driver via the process group (it has its own trap), but a
# `kill` to just this PID would orphan it — so we kill the tracked child and,
# as a backstop, force-remove any of our lane containers that survive.
QUEUE_CHILD_PID=""
cleanup_queue() {
    local code=$?
    [[ -n "$QUEUE_CHILD_PID" ]] && kill "$QUEUE_CHILD_PID" 2>/dev/null || true
    docker rm -f ytahtah-vllm-eval-A ytahtah-vllm-eval-B ytahtah-vllm-websearch-A ytahtah-vllm-websearch-B >/dev/null 2>&1 || true
    exit "$code"
}
trap cleanup_queue EXIT
trap 'exit 130' INT TERM

# Run a child driver in the background + wait, so the trap can kill it. Returns
# the driver's exit code (nonzero if any checkpoint failed, per the drivers).
run_driver() {
    "$@" >> "$LOG" 2>&1 &
    QUEUE_CHILD_PID=$!
    wait "$QUEUE_CHILD_PID"; local rc=$?
    QUEUE_CHILD_PID=""
    return "$rc"
}

# GPU-free re-check before a phase (skipped only with --no-wait).
phase_gate() {
    if [[ "$NO_WAIT" == "true" ]]; then
        log "--no-wait — skipping GPU-free check before $1"
        return 0
    fi
    wait_for_gpus_free "$1"
    sleep 15   # settle margin before claiming the GPUs
}

log "queue_bfcl.sh starting: phases=${PHASES} no_wait=${NO_WAIT}"
log "  poll=${POLL_SECONDS}s stable=${STABLE_SECONDS}s mem_thresh=${MEM_FREE_THRESHOLD_MIB}MiB"
log "  main: 5 runs × [${MAIN_CATEGORIES}] over ${CHECKPOINTS_MAIN}"
log "  ws  : ${WS_RUNS} runs × web_search over ${CHECKPOINTS_WS}"
log "  combined log: ${LOG}"
log "NOTE: never kills foreign GPU processes — only waits for them to end."

RUN_MAIN=false; RUN_WS=false
[[ "$PHASES" == "both" || "$PHASES" == "main" ]] && RUN_MAIN=true
[[ "$PHASES" == "both" || "$PHASES" == "ws" ]] && RUN_WS=true
OVERALL_FAIL=0

# ── Phase: main eval ─────────────────────────────────────────────────────────
if $RUN_MAIN; then
    phase_gate "MAIN eval"
    log "launching MAIN eval (run_evals_parallel.sh) — per-lane logs under logs/lane_*"
    run_driver "${SCRIPT_DIR}/run_evals_parallel.sh" \
        --test-category "$MAIN_CATEGORIES" \
        --checkpoints "$CHECKPOINTS_MAIN"
    main_status=$?
    log "MAIN eval exited ${main_status}"
    if (( main_status != 0 )); then
        OVERALL_FAIL=1
        log "⚠ MAIN eval reported FAILURES (exit ${main_status}) — see the roll-up above / lane logs"
    fi
fi

# ── Phase: web search ────────────────────────────────────────────────────────
if $RUN_WS; then
    $RUN_MAIN && sleep 15   # let main-eval containers tear down before re-gating
    phase_gate "WEB SEARCH eval"
    log "launching WEB SEARCH eval (run_web_search.sh, ${WS_RUNS} runs) — logs/websearch_lane_*"
    run_driver "${SCRIPT_DIR}/run_web_search.sh" \
        --runs "$WS_RUNS" \
        --checkpoints "$CHECKPOINTS_WS"
    ws_status=$?
    log "WEB SEARCH eval exited ${ws_status}"
    if (( ws_status != 0 )); then
        OVERALL_FAIL=1
        log "⚠ WEB SEARCH reported FAILURES (exit ${ws_status}) — see the roll-up above / lane logs"
    fi
fi

if (( OVERALL_FAIL != 0 )); then
    log "queue_bfcl.sh COMPLETE WITH FAILURES — investigate the roll-ups / lane logs above."
else
    log "queue_bfcl.sh complete — all requested phases succeeded."
fi
exit "$OVERALL_FAIL"