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      |
|=============================================================================|
+-----------------------------------------------------------------------------+

Note

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. Login to Docker with your NGC API key using docker login nvcr.io --username='$oauthtoken' --password=<your personal NGC key>

NGC CLI Tool

1. Download the NGC CLI tool for your OS.

Important

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

Note

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:

  GPU	GPU Memory (GB)	Precision	# of GPUs	Disk Space (GB)	CPU RAM (GB)
  H100	80	BF16	1	40	8
  A100	40	BF16	1	40	8
  A6000	48	BF16	1	40	8
  A10G	24	BF16	1	40	8
  L40S	48	BF16	1	40	8

Software

• Minimum Driver version: 535.104.05

Launch DiffDock NIM

1. Pull the NIM container.

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

2. Run container.

Note

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.

export NGC_API_KEY=<your personal NGC key>

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.1.0

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.

Note

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

Note

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
details	str	detailed information for success or failure
protein	str	protein PDB text from the input
ligand	str	ligand molecule description from the input
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.