#!/bin/bash
# Train one ledger run on the GPU cluster. Placeholders __LIKE_THIS__ are filled in
# once the cluster facts are known (see slurm/README.md).
#
#   sbatch slurm/train.sbatch <run_id>      # run created with: python tools/ledger.py new ...
#
#SBATCH --job-name=aarelc-train
#SBATCH --partition=__PARTITION__          # mx
#SBATCH --account=__ACCOUNT__
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=32G
#SBATCH --time=__WALLTIME__                 # e.g. 08:00:00
#SBATCH --output=__SHARED_FS__/AareLC_train/runs/slurm-%j.out
set -euo pipefail

RUN_ID="${1:?usage: sbatch slurm/train.sbatch <run_id>}"
REPO="${AARELC_TRAIN_REPO:-__SHARED_FS__/AareLC_train}"   # a checkout on the shared filesystem

# --- environment: keep exactly one of the three ---------------------------------
# module load __MODULES__                                   # e.g. cuda/12 python/3.12
# source __VENV__/bin/activate                               # venv with ultralytics + opencv
# APPTAINER="apptainer exec --nv __CONTAINER__"             # container with CUDA + ultralytics
APPTAINER="${APPTAINER:-}"

cd "$REPO"
git fetch -q --tags
git pull -q --ff-only
DATASET_TAG="$(python -c "import json;print(json.load(open('runs/$RUN_ID/run.json'))['dataset'])")"
git checkout -q "$DATASET_TAG" -- data/datasets                # the dataset exactly as tagged
git lfs pull -q -I "data/pool/images/**,data/gold/**"          # only the labelled pool and the gold set

python tools/ledger.py set "$RUN_ID" queued --kv "job=$SLURM_JOB_ID" "host=$(hostname)" "partition=$SLURM_JOB_PARTITION"
$APPTAINER python tools/train_run.py "$RUN_ID" --device 0 --workers "$SLURM_CPUS_PER_TASK" 2>&1 | tee "runs/$RUN_ID/train.log"
STATUS=${PIPESTATUS[0]}

# --- report back: commit the run, the registry entry and, on a pass, the promotion --
git add "runs/$RUN_ID" models
git commit -q -m "run $RUN_ID: $(python -c "import json;print(json.load(open('runs/$RUN_ID/run.json'))['status'])")" || true
# git push -q origin HEAD:main                                 # needs a deploy key on the cluster (__DEPLOY_KEY__)
exit "$STATUS"