#!/bin/bash --login
#SBATCH --job-name=boltz2
#SBATCH --output=af_output/%x.d%j/%x.o%j
#SBATCH --error=af_output/%x.d%j/%x.e%j
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=8
#SBATCH --time=0-02:00:00
#SBATCH --mem=40G
#SBATCH --partition=gpu
#SBATCH --qos=gpu
#SBATCH --gres=gpu:a100_40g:1
##SBATCH --account=<account>

shopt -q login_shell || exit 1
test -n "$SLURM_NODELIST" || exit 1
test $SLURM_NNODES -eq 1 || exit 1

module load apptainer
module load boltz2

export NGC_API_KEY="INSERT API KEY HERE"
export NIM_HTTP_API_PORT=$(hpc-find-free-tcp4-port 2>/dev/null || echo 8000)
export TMPDIR=$HOME/nim-cache
export APPTAINERENV_NGC_API_KEY=$NGC_API_KEY
export APPTAINERENV_LOCAL_NIM_CACHE=/nim-cache
export APPTAINERENV_NIM_CACHE=/nim-cache
export APPTAINERENV_NIM_WORKSPACE=/nim-cache/workspace
export APPTAINERENV_NIM_HTTP_API_PORT=$NIM_HTTP_API_PORT

mkdir -p $HOME/nim-cache/workspace

apptainer instance run \
  --nv \
  --bind $HOME/nim-cache:/nim-cache \
  --bind $HOME/nim-cache:/opt/nim/.cache \
  --bind $HOME/nim-cache/workspace:/opt/nim/workspace \
  $BOLTZ2_CONTAINER $SLURM_JOB_ID

apptainer instance list

# Waiting for the boltz2 server to start
until curl -sSf http://localhost:$NIM_HTTP_API_PORT/health/ready; do
    echo "Server not ready, waiting 10 seconds..."
    sleep 10
done
echo "Boltz2 server ready!"

# Testing boltz2
module load miniconda3
source "$CONDA_PREFIX/etc/profile.d/conda.sh"
conda activate boltz2-env
python $HOME/boltz2.py