generation_slurm.sh

#!/bin/bash
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=2
#SBATCH --gres=gpu:1
#SBATCH --gpus-per-task=1
#SBATCH --time=03:00:00
#SBATCH --mem-per-cpu=49152

export OMP_NUM_THREADS=$SLURM_CPUS_PER_TASK

# Load modules
module purge
module load release/24.10 CUDA/12.8.0
module load GCCcore/13.2.0 Python/3.11.5

cd ~/DissBench

# Load environment variables from file
echo "=== Loading environment variables from .env file..."
set -a # automatically export all variables
source .env
set +a

echo "=== Loading venv at $VENV_PATH..."
source $VENV_PATH/bin/activate

# Set additional environment variables that aren't in .env
echo "=== Setting additional environment variables..."
export PIP_CACHE_DIR=/tmp/pip_cache
export TORCH_EXTENSIONS_DIR=/tmp/torch_extensions
export NVCC_CACHE_DIR=/tmp/nvcc_cache
export HF_METRICS_CACHE=/tmp/hf_metrics_cache
export TMPDIR=/tmp
export XDG_CACHE_HOME=/tmp
export VLLM_WORKER_MULTIPROC_METHOD=spawn

echo "=== Environment variables set:"
# Benchmark Specific
echo "DATASET_PATH=$DATASET_PATH"
echo "OUTPUT_DIR=$OUTPUT_DIR"

# General
echo "HF_HOME=$HF_HOME"
echo "HF_HUB_CACHE=$HF_HUB_CACHE"
echo "PIP_CACHE_DIR=$PIP_CACHE_DIR"
echo "TORCH_EXTENSIONS_DIR=$TORCH_EXTENSIONS_DIR"
echo "TMPDIR=$TMPDIR"
echo "NVCC_CACHE_DIR=$NVCC_CACHE_DIR"
echo "HF_METRICS_CACHE=$HF_METRICS_CACHE"
echo "XDG_CACHE_HOME=$XDG_CACHE_HOME"
echo "VLLM_WORKER_MULTIPROC_METHOD=$VLLM_WORKER_MULTIPROC_METHOD"

echo "=== Running the Benchmark script..."
srun ./run_benchmark.sh

exit 0