social-data-attribution-w2/scripts/slurm/analysis/multiseed/eval_array.sbatch

#!/usr/bin/env bash
# Throttled array eval: each array task reads one line from the manifest TSV
# (label \t adapter \t tasks \t outdir), merges the LoRA into base on node-local /tmp,
# evaluates, and CLEANS UP the merged model on exit so disk never accumulates.
#SBATCH --job-name=oevarr
set -uo pipefail
cd "$HOME/dev/data-attribution"
export PATH=.venv/bin:$PATH
export PYTHONPATH=src:olmes
export HF_HOME=$HOME/scratch/hf_cache
export HF_HUB_OFFLINE=${HFOFFLINE:-1}   # set HFOFFLINE=0 to allow fetching eval datasets the maintenance reboot evicted
export TRANSFORMERS_OFFLINE=1
export TOKENIZERS_PARALLELISM=false
export PYTHONUNBUFFERED=1
export VLLM_WORKER_MULTIPROC_METHOD=spawn

MANIFEST="${MANIFEST:?need MANIFEST}"
BASE_LOCAL="${BASE_LOCAL:?need BASE_LOCAL}"
LINE=$(sed -n "${SLURM_ARRAY_TASK_ID}p" "$MANIFEST")
LABEL=$(echo "$LINE" | cut -f1)
ADAPTER=$(echo "$LINE" | cut -f2)
TASKS=$(echo "$LINE" | cut -f3)
OUTDIR=$(echo "$LINE" | cut -f4)
echo "task=${SLURM_ARRAY_TASK_ID} label=$LABEL adapter=$ADAPTER"
# start clean: stale partial requests/predictions from a prior failed run cause
# oe_eval to "resume" with empty predictions -> IndexError. Always start fresh.
rm -rf "$OUTDIR"
mkdir -p "$OUTDIR"

# unique node-local merged dir; clean on exit so /tmp never accumulates across array tasks
MERGED_DIR="/tmp/merged_${SLURM_JOB_ID}_${SLURM_ARRAY_TASK_ID}"
VLLM_CACHE="/tmp/vllm_${SLURM_JOB_ID}_${SLURM_ARRAY_TASK_ID}"
mkdir -p "$MERGED_DIR" "$VLLM_CACHE"
export VLLM_CACHE_ROOT="$VLLM_CACHE"
cleanup() { rm -rf "$MERGED_DIR" "$VLLM_CACHE"; }
trap cleanup EXIT

python -c "
import logging; logging.basicConfig(level=logging.INFO)
from unlearning.eval_harness import merge_adapter
merge_adapter('${BASE_LOCAL}', '${ADAPTER}', '${MERGED_DIR}')
" || { echo 'MERGE FAILED'; exit 1; }

python -m oe_eval.launch \
    --model "$MERGED_DIR" \
    --model-type vllm \
    --model-args "{\"trust_remote_code\": true, \"max_length\": 8192, \"tokenizer\": \"${BASE_LOCAL}\"}" \
    --task ${TASKS} \
    --output-dir "$OUTDIR" --gpus 1
rc=$?
# prune verbose prompt-replay dumps (regenerable; not needed for gamma) to avoid scratch-quota fill
rm -f "$OUTDIR"/*requests*.jsonl 2>/dev/null
echo "done rc=$rc $(date)"
exit $rc