Switch to multi-stage build for image

Also don't download Torch models

More complicated file (and will take longer on first run) but image is 16GB -> 9GB.

Author: jsilter

.dockerignore

@@ -1,9 +1,14 @@
 # Ignore hidden files and directories
-.*
+**/.*
 
 # Ignore git directories (redundant)
 .git
 .gitattributes
 
 # Ignore generated user results
 results/user_inference/*
+
+# Do not include PDBBind files
+# (quite likely they are in the same directory)
+**/pdbbind_sequences.fasta
+**/PDBBind_processed/

docker/cuda-ubuntu22.04/Dockerfile

@@ -1,47 +1,54 @@
-FROM nvidia/cuda:11.7.1-devel-ubuntu22.04
+# Stage 1: Build Environment Setup
+FROM nvidia/cuda:11.7.1-devel-ubuntu22.04 as builder
 
-RUN apt-get update -y
-RUN apt-get install -y sudo wget curl nano git && rm -rf /var/lib/apt/lists/*
+RUN apt-get update -y && apt-get install -y wget curl git tar bzip2 && rm -rf /var/lib/apt/lists/*
 
 # Create a user
 ENV APPUSER="appuser"
 ENV HOME=/home/$APPUSER
 RUN useradd -m -u 1000 $APPUSER
-
 USER $APPUSER
 WORKDIR $HOME
 
-RUN mkdir -p ~/.cache/torch/hub/checkpoints
-# Download ESM models
-ENV ESM_MODEL="esm2_t33_650M_UR50D"
-ENV CONTACT_MODEL="$ESM_MODEL-contact-regression"
-RUN curl -o ~/.cache/torch/hub/checkpoints/$ESM_MODEL.pt https://dl.fbaipublicfiles.com/fair-esm/models/$ESM_MODEL.pt
-RUN curl -o ~/.cache/torch/hub/checkpoints/$CONTACT_MODEL.pt https://dl.fbaipublicfiles.com/fair-esm/regression/$CONTACT_MODEL.pt
+ENV ENV_NAME="diffdock-pocket"
 
-# Build is significantly faster with micromamba, rather than conda
+# Install micromamba
 RUN curl -Ls https://micro.mamba.pm/api/micromamba/linux-64/latest | tar -xj bin/micromamba
-
-# Set the environment variables
-ENV MAMBA_ROOT_PREFIX=$HOME/micromamba
 ENV PATH=$HOME/bin:$HOME/.local/bin:$PATH
 
-RUN ~/bin/micromamba shell init -s bash --root-prefix $MAMBA_ROOT_PREFIX  # this writes to your .bashrc file
-RUN ~/bin/micromamba --version
-
-# Create conda environment first
-# This is time-consuming and don't want to reproduce it often
+# Copy and create Conda environment
 ENV ENV_FILE_NAME=environment.yml
-ENV ENV_NAME="diffdock-pocket"
 COPY --chown=$APPUSER:$APPUSER ./$ENV_FILE_NAME .
+RUN ~/bin/micromamba env create --file $ENV_FILE_NAME && ~/bin/micromamba clean -afy --quiet
 
-RUN ~/bin/micromamba env create --file $ENV_FILE_NAME && ~/bin/micromamba clean -afy
-
+# Copy application code
 COPY --chown=$APPUSER:$APPUSER . $HOME/DiffDock-Pocket
 
+# Stage 2: Runtime Environment
+FROM nvidia/cuda:11.7.1-runtime-ubuntu22.04
+
+# Create user and setup environment
+ENV APPUSER="appuser"
+ENV HOME=/home/$APPUSER
+RUN useradd -m -u 1000 $APPUSER
+USER $APPUSER
+WORKDIR $HOME
+
+ENV ENV_NAME="diffdock-pocket"
+
+# Copy the Conda environment and application code from the builder stage
+COPY --from=builder --chown=$APPUSER:$APPUSER $HOME/micromamba $HOME/micromamba
+COPY --from=builder --chown=$APPUSER:$APPUSER $HOME/bin $HOME/bin
+COPY --from=builder --chown=$APPUSER:$APPUSER $HOME/DiffDock-Pocket $HOME/DiffDock-Pocket
 WORKDIR $HOME/DiffDock-Pocket
 
-# Expose default streamlit and gradio ports
+# Set the environment variables
+ENV MAMBA_ROOT_PREFIX=$HOME/micromamba
+ENV PATH=$HOME/bin:$HOME/.local/bin:$PATH
+RUN micromamba shell init -s bash --root-prefix $MAMBA_ROOT_PREFIX
+
+# Expose ports for streamlit and gradio
 EXPOSE 7860 8501
 
-# Default command just prints the device (CPU or GPU), also ensure using correct python
+# Default command
 CMD ["sh", "-c", "micromamba run -n ${ENV_NAME} python utils/print_device.py"]