run_vae.slurm

#!/bin/bash
#SBATCH --job-name=vae-sweep
#SBATCH --partition=gpu
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=4
#SBATCH --mem=32G
#SBATCH --time=24:00:00
#SBATCH --output=slurm/%x_%A_%a.out
#SBATCH --error=slurm/%x_%A_%a.err

set -euo pipefail

# Adjust these for your cluster environment
# module load cuda/11.8
# source ~/miniconda3/etc/profile.d/conda.sh
# conda activate vae

mkdir -p slurm
# -----------------------------------------------------------------------------
# Activate virtual environment (SLURM-safe)
# -----------------------------------------------------------------------------

echo "Activating virtual environment..."

if [[ -z "${SLURM_SUBMIT_DIR:-}" ]]; then
    echo "ERROR: SLURM_SUBMIT_DIR is not set"
    exit 1
fi

echo "SLURM_SUBMIT_DIR = ${SLURM_SUBMIT_DIR}"

# breaching/ is the submit directory
BASE_DIR="${SLURM_SUBMIT_DIR}"

# project root is one level above
PROJECT_ROOT="$(cd "${BASE_DIR}/.." && pwd)"

echo "BASE_DIR     = ${BASE_DIR}"
echo "PROJECT_ROOT = ${PROJECT_ROOT}"

cd "${PROJECT_ROOT}"

if [[ ! -f bin/activate ]]; then
    echo "ERROR: Virtual environment not found at ${PROJECT_ROOT}/bin/activate"
    exit 1
fi

source bin/activate

# -----------------------------------------------------------------------------

BASE_CONFIG="configs/vae.yaml"
WORKDIR="${SLURM_TMPDIR:-/tmp}/vae_${SLURM_JOB_ID}"
mkdir -p "$WORKDIR"

MODELS=("VanillaVAE")
LATENTS=(64 128)
LRS=(1e-3 5e-4 1e-4)
KLD_WEIGHTS=(2.5e-4 1e-3)
SEEDS=(1265)

for model in "${MODELS[@]}"; do
  for latent in "${LATENTS[@]}"; do
    for lr in "${LRS[@]}"; do
      for kld in "${KLD_WEIGHTS[@]}"; do
        for seed in "${SEEDS[@]}"; do
          safe_lr=${lr//./p}
          safe_kld=${kld//./p}
          run_name="${model}_ld${latent}_lr${safe_lr}_kld${safe_kld}_s${seed}"

          CONFIG_OUT="$WORKDIR/config_${run_name}.yaml"

          export BASE_CONFIG model latent lr kld seed run_name CONFIG_OUT

          python - <<'PY'
                        import os
                        import yaml

                        base = os.environ["BASE_CONFIG"]
                        out = os.environ["CONFIG_OUT"]
                        model = os.environ["model"]
                        latent = int(os.environ["latent"])
                        lr = float(os.environ["lr"])
                        kld = float(os.environ["kld"])
                        seed = int(os.environ["seed"])
                        run_name = os.environ["run_name"]

                        with open(base, "r") as f:
                            cfg = yaml.safe_load(f)

                        cfg["model_params"]["name"] = model
                        cfg["model_params"]["latent_dim"] = latent
                        cfg["exp_params"]["LR"] = lr
                        cfg["exp_params"]["kld_weight"] = kld
                        cfg["exp_params"]["manual_seed"] = seed

                        # Use the GPU allocated by Slurm
                        cfg["trainer_params"]["gpus"] = [0]

                        # Make logdir unique per run
                        cfg["logging_params"]["name"] = run_name

                        with open(out, "w") as f:
                            yaml.safe_dump(cfg, f, sort_keys=False)

                        print(f"Wrote config: {out}")
PY

          srun python run.py -c "$CONFIG_OUT"
        done
      done
    done
  done
done