快速入门指南#

注意

本页假设您已安装并设置了前提软件 (Docker、NGC CLI、NGC 注册表访问)。

启动 NIM#

  1. 导出 NGC_API_KEY 变量。

export NGC_API_KEY=<your personal NGC key>

  1. NIM 容器自动下载模型。为了节省时间和带宽，建议提供本地缓存目录。这样 NIM 将能够重用已下载的模型。执行以下命令设置缓存目录。

export LOCAL_NIM_CACHE=~/.cache/nim
mkdir -p "$LOCAL_NIM_CACHE"
sudo chmod 0777 -R "$LOCAL_NIM_CACHE"

  1. 使用以下命令运行 NIM 容器。

docker run -it \
    --runtime=nvidia \
    --gpus='"device=0"' \
    -p 8000:8000 \
    -e NGC_API_KEY \
    -v "$LOCAL_NIM_CACHE":/opt/nim/.cache \
    nvcr.io/nim/ipd/rfdiffusion:2

此命令将启动 NIM 容器并暴露端口 8000，供用户与 NIM 交互。

  1. 打开一个新的终端，保持运行已启动服务的终端打开。在新终端中，等待健康检查端点返回 {"status":"ready"} 后再继续。这可能需要几分钟时间。您可以使用以下命令查询健康检查。

curl http://127.0.0.1:8000/v1/health/ready

Python 客户端示例#

  1. 将以下 Python 示例保存到名为 nim_client.py 的文件中。

#!/usr/bin/env python3
import requests
import os
import json
from pathlib import Path

def get_reduced_pdb():
    pdb = Path("1R42.pdb")
    if not pdb.exists():
        pdb.write_text(requests.get(f"https://files.rcsb.org/download/{pdb}").text)
    lines = filter(lambda line: line.startswith("ATOM"), pdb.read_text().split("\n"))
    return "\n".join(list(lines)[:400])

r = requests.post(
    url="http://127.0.0.1:8000/biology/ipd/rfdiffusion/generate",
    json={
        "input_pdb": get_reduced_pdb(),
        "contigs": "A20-60/0 50-100",
        "hotspot_res": ["A50","A51","A52","A53","A54"],
        "diffusion_steps": 15,
    },
)
print(r, "Saving to output.pdb:\n", r.text[:200], "...")
Path("output.pdb").write_text(json.loads(r.text)["output_pdb"])

  1. 执行示例。

chmod +x nim_client.py

./nim_client.py

  1. 该示例将结果以 PDB 格式保存到 output.pdb 文件中。您可以使用以下命令快速查看该文件。

less output.pdb

Shell 客户端示例#

  1. 将以下 Shell 示例保存到名为 nim_client.sh 的文件中。

#!/usr/bin/env bash
set -e

URL=http://127.0.0.1:8000/biology/ipd/rfdiffusion/generate

if [ ! -e 1R42.pdb ]; then curl -O https://files.rcsb.org/download/1R42.pdb; fi

pdb=$(cat 1R42.pdb | grep ^ATOM | head -n 400 | awk '{printf "%s\\n", $0}')

request='{
 "input_pdb": "'"$pdb"'",
 "contigs": "A20-60/0 50-100",
 "hotspot_res": ["A50","A51","A52","A53","A54"],
 "diffusion_steps": 15
}'
curl -H 'Content-Type: application/json' \
     -d "$request" "$URL"

  1. 执行示例。该示例以 JSON 格式将结果显示在终端中。

chmod +x nim_client.sh

./nim_client.sh
本页目录