build_container.slurm

#!/bin/bash
4#SBATCH --job-name=rfdiffusion_x86_build
#SBATCH --output=rfdiffusion_x86_build_%j.out
#SBATCH --error=rfdiffusion_x86_build_%j.err
#SBATCH --partition=YOUR_PARTITION             # Specify the partition/queue name
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --cpus-per-task=64                     # Adjust as needed
#SBATCH --time=04:00:00                        # Adjust as needed
#SBATCH --mem=256G                             # Adjust as needed
#SBATCH --account=YOUR_ACCOUNT               # Specify your allocation account
#SBATCH --gres=gpu:1                             # RFdiffusion build doesn't strictly need GPU, but good for consistency if runtime uses it
#SBATCH --gpus-per-task=1

# module load singularity # Uncomment and adjust if needed for your environment

MODEL_NAME="rfdiffusion"
ARCH="x86"
DEFINITION_FILE_RELATIVE_PATH="${MODEL_NAME}/${MODEL_NAME}_${ARCH}.def"

if [ -z "$1" ]; then
    echo "Error: Please provide the target build directory."
    echo "Usage: sbatch $0 /path/to/your/build/directory"
    exit 1
fi

BUILD_DIR=$1

if [ ! -f "${DEFINITION_FILE_RELATIVE_PATH}" ]; then
    echo "Error: Definition file ${DEFINITION_FILE_RELATIVE_PATH} not found from repository root."
    echo "Ensure the file exists and sbatch is run from the repository root."
    exit 1
fi

if [ ! -d "${BUILD_DIR}" ]; then
    echo "Build directory '${BUILD_DIR}' does not exist. Creating it."
    mkdir -p "${BUILD_DIR}"
    if [ $? -ne 0 ]; then
        echo "Error: Could not create build directory '${BUILD_DIR}'."
        exit 1
    fi
fi

DEFINITION_FILE_FOR_SINGULARITY="${DEFINITION_FILE_RELATIVE_PATH}"
SIF_NAME="${MODEL_NAME}_${ARCH}.sif"

echo "======================================================================="
echo "Starting Singularity container build for: ${MODEL_NAME} (${ARCH})"
echo "Using definition file: ${DEFINITION_FILE_FOR_SINGULARITY}"
echo "Target build directory: ${BUILD_DIR}"
echo "Output SIF name: ${SIF_NAME}"
echo "sbatch execution directory: $(pwd)"
echo "======================================================================="

# The SIF file will be created in BUILD_DIR.
# The definition file path is relative to the sbatch execution directory (assumed repo root).
singularity build --force "${BUILD_DIR}/${SIF_NAME}" "${DEFINITION_FILE_FOR_SINGULARITY}"

BUILD_EXIT_CODE=$?

if [ ${BUILD_EXIT_CODE} -eq 0 ]; then
    echo "======================================================================="
    echo "Successfully built ${SIF_NAME} in ${BUILD_DIR}"
    echo "Full path: ${BUILD_DIR}/${SIF_NAME}"
    echo "======================================================================="
else
    echo "======================================================================="
    echo "Error building ${SIF_NAME}. Exit code: ${BUILD_EXIT_CODE}"
    echo "Check SLURM output files in $(pwd) and any logs in ${BUILD_DIR}"
    echo "======================================================================="
    exit 1
fi

exit 0