#!/bin/bash
#SBATCH --job-name=${DOVLA_JOB_NAME:-dovla_scaling}
#SBATCH --partition=${DOVLA_PARTITION:-gpu}
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=${DOVLA_CPUS_PER_TASK:-8}
#SBATCH --gres=gpu:${DOVLA_GPUS_PER_TASK:-1}
#SBATCH --mem=${DOVLA_MEM:-64G}
#SBATCH --time=${DOVLA_TIME:-24:00:00}
#SBATCH --output=${DOVLA_LOG_DIR:-logs/slurm}/%x_%j.out
#SBATCH --error=${DOVLA_LOG_DIR:-logs/slurm}/%x_%j.err

set -euo pipefail

PROJECT_DIR="${PROJECT_DIR:-$PWD}"
VENV_PATH="${VENV_PATH:-$PROJECT_DIR/.venv}"
BACKEND="${BACKEND:-toy}"
TASKS="${TASKS:-builtins}"
OUT_DIR="${OUT_DIR:-$PROJECT_DIR/runs/scaling_toy}"
TOTAL_RECORDS="${TOTAL_RECORDS:-4096}"
K_VALUES="${K_VALUES:-1,2,4,8,16,32}"
EPOCHS="${EPOCHS:-3}"
SEED="${SEED:-0}"
DEVICE="${DEVICE:-auto}"

mkdir -p "${DOVLA_LOG_DIR:-logs/slurm}" "$OUT_DIR"
cd "$PROJECT_DIR"

if [ -f "$VENV_PATH/bin/activate" ]; then
  # shellcheck disable=SC1091
  source "$VENV_PATH/bin/activate"
fi

python scripts/run_scaling.py \
  --backend "$BACKEND" \
  --tasks "$TASKS" \
  --out "$OUT_DIR" \
  --total-records "$TOTAL_RECORDS" \
  --k-values "$K_VALUES" \
  --epochs "$EPOCHS" \
  --seed "$SEED" \
  --device "$DEVICE"