入门#

先决条件#

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

示例输出

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

注意

有关枚举多 GPU 系统的更多信息，请参阅 NVIDIA Container Toolkit 的GPU 枚举文档

NGC (NVIDIA GPU Cloud) 帐户#

  1. 在 NGC 上创建帐户

  2. 生成 API 密钥

  3. 使用您的 NGC API 密钥进行 Docker 登录，使用 docker login nvcr.io --username='$oauthtoken' --password=${NGC_CLI_API_KEY}

NGC CLI 工具#

  1. 为您的操作系统下载 NGC CLI 工具。

重要提示

使用 NGC CLI 版本 3.41.1 或更高版本。以下是在 AMD64 Linux 家庭目录中安装此版本的命令

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. 在本地设置您的 NGC CLI 工具（您将需要您的 API 密钥！）

ngc config set

注意

输入 API 密钥后，您可能会看到组织和团队的多个选项。根据需要选择或按 Enter 键接受默认值。

模型特定要求#

以下是 DiffDock NIM 的特定要求。

硬件#

  • 支持的 GPU 型号

  • Hopper GPU (H100)

  • Ampere GPU（例如，A100 和 A6000，详细信息请参见此处）

  • Ada GPU（例如，L40S，详细信息请参见此处）

  • Volta GPU（例如，V100，详细信息请参见此处）

  • 最低 GPU 内存 (GB)：16

满足上述要求后，您将使用快速入门指南来拉取 NIM 容器和模型，执行健康检查，然后运行推理。

软件#

  • 最低驱动程序版本：535.104.05

启动 DiffDock NIM#

  1. 拉取 NIM 容器。

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

  1. 运行容器。

注意

环境变量 NGC_API_KEY 必须在您的本地环境中定义且有效，以确保以下命令可以继续执行。有关个人 API 密钥设置的信息，请参阅此网站。

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

  1. 打开一个新的终端，使用以下命令检查 API 的状态，直到它返回 true。这可能需要几分钟。

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

运行推理#

  1. 打开一个新的终端，保持当前终端打开并运行服务。

注意

打开一个新的终端，保持当前终端打开并运行服务。

注意

“sed” 命令用于将多行文本文件转换为单行以进行 JSON 编码。

  1. 准备 JSON 格式的 post-data。此步骤需要在最常见的 bash shell 环境中启动（Linux）。用户可以使用命令 echo $0 验证当前会话是否为 bash。如果不是，请在此步骤之前运行命令 /bin/bash。

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

  1. 运行推理并保存到 output.json。

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

  1. 输出文件 output.json 是 JSON 格式的内容，其中包含预测的对接姿势（配体原子的坐标），结构如下

字段

类型

描述

status

str

报告此请求的“成功”或“失败”

ligand_positions

字符串列表

作为生成的姿势的 SDF 格式文本列表

position_confidence

浮点数列表

生成的姿势的置信度分数列表

trajectory

字符串列表

PDB 格式文本列表，作为生成姿势的扩散轨迹（可选）

导出生成的姿势#

  1. 本节中提供了一个简单的 Python 脚本，用于将推理结果（配体的对接姿势）导出到名为 output 的文件夹中。创建一个新的空白文件，将其命名为 dump_output.py，并将以下内容复制到其中。

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])

  1. 运行以下命令以启动 Python 脚本。

python3 dump_output.py

  1. 列出 output 文件夹中的内容。

$ ls output

rank01_confidence_-0.98.sdf

停止容器#

完成端点测试后，您可以通过在新终端中运行 docker stop diffdock-nim 来关闭容器。

在本页中