File size: 3,173 Bytes
e79d917 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | #!/usr/bin/env bash
# Ordered, resumable run of the whole model-3 pipeline.
#
# bash scripts/queue.sh # full pipeline
# bash scripts/queue.sh embed # embeddings only
# STAGE_GPU_TRAIN=0 STAGE_GPU_TEST=1 bash scripts/queue.sh
#
# Every stage is resumable: re-running skips work that is already complete, so a
# kill (shared GPU, disk, anything) costs at most one 250k-row shard.
#
# Check the GPUs are actually free before launching:
# nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv
set -euo pipefail
cd "$(dirname "$0")/.."
ROOT=$(pwd)
PY=${PY:-python3}
MODEL=${M3_MODEL:-e5-small}
GPU_TRAIN=${STAGE_GPU_TRAIN:-0}
GPU_TEST=${STAGE_GPU_TEST:-1}
LOGS=${M3_LOGS:-$ROOT/logs}
mkdir -p "$LOGS"
STAGE=${1:-all}
export M3_MODEL="$MODEL"
export PYTHONPATH="$ROOT/src:${PYTHONPATH:-}"
# Never let a stray HF cache land on /home (it hit 96% full and killed a run).
export HF_HOME=${HF_HOME:-/scratch/user4/mtp/cache/huggingface}
export TOKENIZERS_PARALLELISM=false # fork deadlock with DataLoader workers
ts() { date +%H:%M:%S; }
run() { echo "== [$(ts)] $*"; "$@"; }
echo "== model-3 queue | model=$MODEL | train-gpu=$GPU_TRAIN test-gpu=$GPU_TEST"
echo " work dir: ${M3_WORK:-$ROOT/work}"
if [[ "$STAGE" == "all" || "$STAGE" == "embed" ]]; then
# All four parts are persisted (train pool, train S1, test pool, test S1).
# At e5-small speed train S1 costs ~6 min, and caching it means the training
# sample can be scaled past 200k later without re-encoding anything.
# Order follows the pipeline dependency: train parts first so run_train_eval.py
# can start while the test corpora are still encoding.
if [[ "${SINGLE_GPU:-0}" == "1" ]]; then
for spec in "train pool" "train s1" "test pool" "test s1"; do
set -- $spec
run $PY src/embed.py --split "$1" --part "$2" --gpu "$GPU_TRAIN" \
2>&1 | tee "$LOGS/embed_$1_$2.log"
done
else
run $PY src/embed.py --split train --part all --gpu "$GPU_TRAIN" \
2>&1 | tee "$LOGS/embed_train.log" &
P1=$!
run $PY src/embed.py --split test --part all --gpu "$GPU_TEST" \
2>&1 | tee "$LOGS/embed_test.log" &
P2=$!
wait $P1 $P2
fi
fi
if [[ "$STAGE" == "all" || "$STAGE" == "train" ]]; then
# CUDA_VISIBLE_DEVICES, not just --gpu: cross_encoder.py falls back to a bare
# "cuda" (= cuda:0), so --gpu 1 alone still put the 568M reranker on GPU 0 --
# the shared card. Masking the device makes every default land correctly.
run env CUDA_VISIBLE_DEVICES="${TRAIN_GPUS:-$GPU_TRAIN}" $PY src/run_train_eval.py --gpu 0 2>&1 | tee "$LOGS/train_eval.log"
fi
if [[ "$STAGE" == "all" || "$STAGE" == "test" ]]; then
ART="${M3_WORK:-$ROOT/work}/results_A"
run env CUDA_VISIBLE_DEVICES="$GPU_TEST" $PY src/run_test.py --gpu 0 --artifacts "$ART" 2>&1 | tee "$LOGS/run_test.log"
VAL=$(cd "$ROOT/../dataset" && pwd)/utils/validate_submission.py
run $PY "$VAL" \
--matching "$ROOT/output/matching_results.tsv" \
--candidate "$ROOT/output/candidate_pairs.tsv" \
--test-dir "$(cd "$ROOT/../dataset/dataset" && pwd)/test" \
2>&1 | tee "$LOGS/validate.log"
fi
echo "== [$(ts)] queue finished"
|