Skip to navigation

Prerequisites

View as Markdown
docker run --rm --gpus all ubuntu nvidia-smi

Example output:

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 590.44.01    Driver Version: 590.44.01    CUDA Version: 13.1     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  Off  | 00000000:01:00.0 Off |                  N/A |
| 41%   30C    P8     1W / 260W |   2244MiB / 11264MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

The CUDA Version reported by nvidia-smi is the CUDA version bundled with the installed driver, not the one used by the container. The OpenFold2 container ships with a CUDA 13.2 runtime and can run on systems that meet the documented minimum driver requirement for OpenFold2, through CUDA minor version compatibility.

For more information on enumerating multi-GPU systems, refer to the NVIDIA Container Toolkit’s GPU Enumeration Docs.

NGC (NVIDIA GPU Cloud) Account

  1. Create an account on NGC

  2. Generate an API Key

  3. Export your NGC API key so the following commands can reference it (replace <Enter your NGC key here> with the key from step 2):

    export NGC_API_KEY=<Enter your NGC key here>

  4. Log in to Docker with your NGC API key:

    docker login nvcr.io --username='$oauthtoken' --password=${NGC_API_KEY}

NGC CLI Tool

  1. Download the NGC CLI tool for your OS.

Use NGC CLI version 3.41.1 or newer. Here is the command to install this on AMD64 Linux in your home directory:

wget --content-disposition https://api.ngc.nvidia.com/v2/resources/nvidia/ngc-apps/ngc_cli/versions/3.41.3/files/ngccli_linux.zip -O ~/ngccli_linux.zip && \
unzip ~/ngccli_linux.zip -d ~/ngc && \
chmod u+x ~/ngc/ngc-cli/ngc && \
echo "export PATH=\"\$PATH:~/ngc/ngc-cli\"" >> ~/.bash_profile && source ~/.bash_profile
  1. Set up your NGC CLI Tool locally (you need your API key for this):
ngc config set

After you enter your API key, you might see multiple options for the org and team. Select as desired or press Enter to accept the default.

  1. Log in to NGC

You need to log in to NGC through Docker and set the NGC_API_KEY environment variable to pull images:

docker login nvcr.io
Username: $oauthtoken
Password: <Enter your NGC key here>

Then, set the relevant environment variables in your shell. You need to set the NGC_API_KEY variable:

export NGC_API_KEY=<Enter your NGC key here>
  1. Set up your NIM cache

The NIM cache allows you to download models and store previously downloaded models so that you do not need to download them again later when you run the NIM again. The NIM cache must be readable and writable by the NIM, so in addition to creating the directory, the permissions on this directory need to be set to globally readable and writable. The NIM cache directory can be set up as follows:

## Create the NIM cache directory
mkdir -p ~/.cache/nim
## Set the NIM cache directory permissions to the correct values
chmod -R 777 ~/.cache/nim

Now, you should be able to pull the container and download the model using the environment variables. To get started, refer to Getting Started.

Running on Slurm With Enroot

This section describes how to run the OpenFold2 NIM on a Slurm-based HPC cluster using Enroot to run the NGC container.

NGC_API_KEY: You must export your NGC API key before running the Slurm script so that Enroot can pull the image from NGC. Then run the script (for example, bash openfold2_srun.sh).

Example:

export NGC_API_KEY=<your_NGC_API_key>
bash openfold2_srun.sh

Environment

  • Enroot: Version 3.4.1 or newer is supported.

1. Check That Enroot Is Available

On the login node, verify Enroot is installed:

enroot version
which enroot

If Enroot is not installed, contact your cluster administrator. Typical path: /usr/bin/enroot.

Create the Enroot config directory if it does not exist:

mkdir -p ~/.config/enroot

2. Create and Add NGC Credentials

Create the credentials file so Enroot can pull images from NGC. Replace YOUR_NGC_API_KEY with your NGC API key:

echo 'machine nvcr.io login $oauthtoken password YOUR_NGC_API_KEY' > ~/.config/enroot/.credentials

Restrict permissions and confirm the file content:

chmod 600 ~/.config/enroot/.credentials
cat ~/.config/enroot/.credentials

If you prefer not to echo the key on the command line, use vi ~/.config/enroot/.credentials or another editor.

3. Configure and Run the Slurm Job Script

Copy the script below to a file, such as openfold2_srun.sh. Set the variables at the top for your cluster and paths:

  • ACCOUNT: Slurm account name
  • PARTITION: Slurm partition (for example, interactive or gpu)
  • ROOT_DIR: Your working directory on the cluster (for example, Lustre or NFS path)
  • NGC_IMAGE: NGC image with tag; use # between registry and image path for Enroot (refer to the example below)
  • NGC_API_KEY: Must be exported in your environment before running the script (for example, export NGC_API_KEY=<your_key>)

Use a shared filesystem path for ROOT_DIR and CACHE_PATH so that the SQSH image and cache are visible from the node where the job runs. Use a node-local or large-quota path for TMP_PATH if /tmp has strict quotas (for example, Lustre scratch).

Save the following as openfold2_srun.sh and run it with bash openfold2_srun.sh:

#!/bin/bash
# --- Slurm job configuration (customize for your cluster) ---
ACCOUNT="your_slurm_account"
PARTITION="interactive"
GPUS_PER_NODE=1
MEMORY="32G"
TIME="04:00:00"
# --- Paths (customize for your cluster) ---
# Root directory on shared storage (for example, Lustre/NFS)
ROOT_DIR="/path/to/your/workspace"
# NIM cache on shared storage
CACHE_PATH="${ROOT_DIR}/.nim_cache"
WORKING_CACHE_DIR="/opt/nim/.cache"
# Project/data mount inside container
WORKING_PATH="${ROOT_DIR}"
MOUNT_PATH="/workspace/data"
# Temp directory (use shared storage if node /tmp has quota limits)
TMP_PATH="${ROOT_DIR}/tmp"
MOUNT_TMP_PATH="${MOUNT_PATH}/tmp"
# --- NGC image (use # between registry and image path for Enroot) ---
NGC_IMAGE="docker://nvcr.io#nim/openfold/openfold2:2.6.0"
CONTAINER_NAME="openfold2"
ACTIVE_SQSH_PATH="${ROOT_DIR}/openfold2.sqsh"
# --- Helpers ---
handle_error() {
    echo "Error: $1"
    exit 1
}
mkdir -p "$(dirname "${ACTIVE_SQSH_PATH}")" || handle_error "Failed to create directory for SQSH file"
mkdir -p "${CACHE_PATH}" || handle_error "Failed to create cache directory"
mkdir -p "${TMP_PATH}" || handle_error "Failed to create tmp directory"
echo "Step 1: Requesting interactive resources..."
srun --account="${ACCOUNT}" \
     --partition="${PARTITION}" \
     --gpus-per-node="${GPUS_PER_NODE}" \
     --mem="${MEMORY}" \
     --time="${TIME}" \
     --export=ALL \
     -o /dev/tty -e /dev/tty \
     bash -c "
    echo \"Loading environment...\"
    # Step 2: Import Docker image (if not present)
    echo \"Step 2: Importing Docker image\"
    if [ ! -f \"${ACTIVE_SQSH_PATH}\" ]; then
        echo \"Importing from NGC...\"
        mkdir -p \"\$(dirname \"${ACTIVE_SQSH_PATH}\")\" || { echo \"Failed to create directory\"; exit 1; }
        enroot import -o \"${ACTIVE_SQSH_PATH}\" ${NGC_IMAGE} || { echo \"Failed to import Docker image\"; exit 1; }
    else
        echo \"Docker image already exists, skipping import.\"
    fi
    # Step 3: Create Enroot container
    echo \"Step 3: Creating Enroot container...\"
    if enroot list 2>/dev/null | grep -q \"${CONTAINER_NAME}\"; then
        echo \"Removing existing container...\"
        enroot remove -f ${CONTAINER_NAME} || true
    fi
    enroot create --name ${CONTAINER_NAME} \"${ACTIVE_SQSH_PATH}\" || { echo \"Failed to create container\"; exit 1; }
    # Step 4: Start container and NIM server
    NODE=\$(hostname)
    echo \"\$NODE\" > \"${WORKING_PATH}/.openfold2_node\"
    echo \"Step 4: OpenFold2 container ready.\"
    echo \"=====================================\"
    echo \"Working directory in container: ${MOUNT_PATH}\"
    echo \"\"
    echo \"From another terminal (login node or your machine), run:\"
    echo \"  ssh -L 8000:localhost:8000 \$NODE\"
    echo \"Then keep that SSH session open and call the API (see Step 4 in docs).\"
    echo \"Node name saved to: ${WORKING_PATH}/.openfold2_node\"
    echo \"Type exit to leave the container.\"
    echo \"=====================================\"
    cat > \"\${TMPDIR:-/tmp}/rc.local\" << RCEOF
#!/bin/sh
export TMPDIR=${MOUNT_TMP_PATH}
export TEMP=${MOUNT_TMP_PATH}
export TMP=${MOUNT_TMP_PATH}
export HOME=${MOUNT_PATH}
export XDG_CACHE_HOME=${WORKING_CACHE_DIR}
export XDG_DATA_HOME=${MOUNT_PATH}/.local/share
mkdir -p ${MOUNT_TMP_PATH} ${MOUNT_PATH}/.local/share
/opt/nim/start_server.sh &
exec /bin/bash
RCEOF
    chmod +x \"\${TMPDIR:-/tmp}/rc.local\"
    enroot start \\
      --mount ${CACHE_PATH}:${WORKING_CACHE_DIR} \\
      --mount ${WORKING_PATH}:${MOUNT_PATH} \\
      --mount \"\${TMPDIR:-/tmp}/rc.local:/etc/rc.local\" \\
      -e NGC_API_KEY \\
      -e TMPDIR=${MOUNT_TMP_PATH} \\
      -e TEMP=${MOUNT_TMP_PATH} \\
      -e TMP=${MOUNT_TMP_PATH} \\
      -e HOME=${MOUNT_PATH} \\
      -e XDG_CACHE_HOME=${WORKING_CACHE_DIR} \\
      -e XDG_DATA_HOME=${MOUNT_PATH}/.local/share \\
      ${CONTAINER_NAME} || { echo \"Failed to start container\"; exit 1; }
    echo \"=====================================\"
    echo \"Workflow completed.\"
"

Submit the job (interactive):

export NGC_API_KEY=<your_NGC_API_key>
bash openfold2_srun.sh

Alternatively, for a batch job, wrap the same script body in an sbatch script and set ACCOUNT, PARTITION, and other Slurm directives as needed. Ensure NGC_API_KEY is exported or passed into the job environment.

4. Call the API From Your Machine

After the container is running on a compute node, the script prints the node name and saves it to ROOT_DIR/.openfold2_node.

  1. From your laptop or the login node, create an SSH tunnel to that node:

    ssh -L 8000:localhost:8000 $(cat /path/to/your/workspace/.openfold2_node)

    Or use the node name directly if you know it:

    ssh -L 8000:localhost:8000 <compute-node-name>

  2. In another terminal (with the tunnel still open), run a test request:

    curl -s -X POST 'http://localhost:8000/biology/openfold/openfold2/predict-structure-from-msa-and-template' \
      -H "content-type: application/json" \
      -H "NVCF-POLL-SECONDS: 300" \
      -d '{"sequence":"GGSKENEISHHAKEIERLQKEIERHKQSIKKLKQSEQSNPPPNPEGTRQARRNRRRRWRERQRQKENEISHHAKEIERLQKEIERHKQSIKKLKQSEC","selected_models":[1,2]}' \
      --max-time 300

If the request succeeds, you get a JSON response containing the prediction result.