Getting Started

Prerequisites

• Install Docker
• Install the NVIDIA Container Toolkit
• Verify your container runtime supports NVIDIA GPUs by running

docker run --rm --runtime=nvidia --gpus all ubuntu nvidia-smi

Example output:

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.78.01    Driver Version: 525.78.01    CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| 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      |
|=============================================================================|
+-----------------------------------------------------------------------------+

NGC (NVIDIA GPU Cloud) Account

1. Create an account on NGC
2. Generate an API Key
3. Docker log in with your NGC API key using docker login nvcr.io --username='$oauthtoken' --password=${NGC_CLI_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

2. Set up your NGC CLI Tool locally (You’ll need your API key for this!):

ngc config set

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

Model Specific Requirements

The following are specific requirements for DiffDock NIM.

Hardware

• Supported GPU models:

• Hopper GPUs (H100)
• Ampere GPUs (for example, A100 and A6000, details can be found here)
• Ada GPUs (for example, L40S, details can be found here)
• Volta GPUs (for example, V100, details can be here)

• Minimum GPU memory (GB): 16

Once the above requirements are met, you will use the Quickstart Guide to pull the NIM container and model, perform a health check, and then run inference.

Software

• Minimum Driver version: 535.104.05

Launch DiffDock NIM

1. Pull the NIM container.

docker pull nvcr.io/nim/mit/diffdock:2.0.1

2. Run container.

The environment variable NGC_API_KEY must be defined and valid in your local environment to ensure the following command will proceed. For information regarding personal API key setup, refer to this website.

docker run --rm -it --name diffdock-nim \
  --runtime=nvidia -e NVIDIA_VISIBLE_DEVICES=0 \
  --shm-size=2G \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  -e NGC_API_KEY=$NGC_API_KEY \
  -p 8000:8000 \
  nvcr.io/nim/mit/diffdock:2.0.1

3. Open a new terminal, use the following command to check the status of API until it returns true. This can take a couple of minutes.

curl localhost:8000/v1/health/ready
...
true

Run Inference

1. Open a new terminal, leaving the current terminal open with the launched service.

Open a new terminal, leaving the current terminal open with the launched service.

The “sed” command is used to convert the multi-line text file into a single line for JSON encoding.

2. Prepare JSON formatted post-data. This step requires being launched in the most common bash shell environment in Linux. Users can verify if the current session is bash by using the command echo $0. If not, run the command /bin/bash before this step.

protein_bytes=`curl https://files.rcsb.org/download/8G43.pdb | grep -E '^ATOM' | sed -z 's/\n/\\\n/g'`; \
ligand_bytes=`curl https://files.rcsb.org/ligands/download/ZU6_ideal.sdf | sed -z 's/\n/\\\n/g'`; \
echo "{
    \"ligand\": \"${ligand_bytes}\",
    \"ligand_file_type\": \"sdf\",
    \"protein\": \"${protein_bytes}\",
    \"num_poses\": 1,
    \"time_divisions\": 20,
    \"steps\": 18,
    \"save_trajectory\": false,
    \"is_staged\": false
}" > diffdock.json

3. Run Inference and save to output.json.

curl --header "Content-Type: application/json" \
    --request POST \
    --data @diffdock.json \
    --output output.json \
    http://localhost:8000/molecular-docking/diffdock/generate

4. The output file output.json is a JSON formatted content with predicted docking poses (coordinates of ligand atoms) with the structure below:

Field	Type	Description
status	str	report ‘success’ or ‘fail’ for this request
ligand_positions	list of str	list of SDF formatted text as the generated poses
position_confidence	list of float	list of confidence scores for the generated poses
trajectory	list of str	list of PDB formatted text as the diffusion trajectories for the generated poses (optional)

Dump Generated Poses

1. A simple Python script provided in this section is used to dump the inference results (docked poses of ligands) into a folder named as output. Create a new blank file, name it as dump_output.py and copy the content below into it.

import json
import os
import shutil
def dump_one(folder, ligand_positions, position_confidence):
    os.makedirs(folder, exist_ok=True)
    for i, c in enumerate(position_confidence):
        with open('%s/rank%02d_confidence_%0.2f.sdf' % (folder, i+1, c), 'w') as f:
            f.write(ligand_positions[i])
shutil.rmtree('output', ignore_errors=True)
os.makedirs('output', exist_ok=True)
with open('output.json') as f:
    data = json.load(f)
if type(data['status']) == str:
    dump_one('output', data['ligand_positions'], data['position_confidence'])
else:
    for i in range(len(data['status'])):
        dump_one('output/ligand%d' % (i+1), data['ligand_positions'][i], data['position_confidence'][i])

2. Run the command below to launch the Python script.

python3 dump_output.py

3. List the content in the output folder.

$ ls output
rank01_confidence_-0.98.sdf

Stopping the Container

When you’re done testing the endpoint, you can bring down the container by running docker stop diffdock-nim in a new terminal.