运行时配置 NIM#

AlphaFold2-Multimer NIM 允许在运行时进行多项自定义，从而可以

  • 调整 NIM 的性能。

  • 充分利用您独特的硬件。

  • 修改运行时变量（例如 NIM 缓存的位置）。

下面，我们将详细介绍各种 AlphaFold2-Multimer NIM 特定的运行时配置选项。

简要回顾：启动 AlphaFold2-Multimer NIM#

让我们首先回顾一下如何启动 NIM

export LOCAL_NIM_CACHE=~/.cache/nim

docker run -it --rm --name alphafold2-multimer --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_CLI_API_KEY \
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2-multimer:1.0.0

请记住

  • -p 定义与 NIM 的服务端点交互和查询的端口。

  • -e 选项定义环境变量，这些变量在运行时传递到 NIM 容器中。

  • --rm 在容器终止时（即通过 CTRL+C 或其他终止信号）停止并移除容器。

  • -it 允许直接在 CLI 中与容器终端交互，例如 CTRL+C。

注意

如果在下载 AlphaFold2 模型后，docker run 命令由于文件权限错误而失败，请运行以下命令：sudo chmod -R 777 $LOCAL_NIM_CACHE。之后，使用 docker run 命令重新运行 NIM。

配置 NIM 缓存位置#

运行 NIM 容器会自动将 AlphaFold2 模型下载到容器文件系统中的默认路径：/opt/nim/.cache。但是，NIM 模型缓存的位置可以在本地文件系统和 NIM 内部进行更改。当需要更大的驱动器来存储 MSA 数据库（MSA 数据库占 AlphaFold2 模型大小的很大一部分，约为 613 GB），或者当用户想要自定义其在 NIM 容器中的工作区以向 NIM 添加其他功能（例如，修改后的入口点或挂载的插件）时，这非常有用。

要更改主机文件系统上的 NIM 缓存位置，请更改 NIM 启动命令中的源挂载位置，我们将其称为环境变量 LOCAL_NIM_CACHE

export LOCAL_NIM_CACHE=/mount/largedisk/nim/.cache

docker run -it --rm --name alphafold2-multimer --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_CLI_API_KEY \
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2-multimer:1.0.0

这只会影响容器外部的位置，即运行容器的主机上。要更改容器内部挂载的缓存位置，您必须

  1. 确保挂载新的目标路径，即 -v $LOCAL_NIM_CACHE:$NIM_CACHE_PATH。

  2. 在 NIM_CACHE_PATH 环境变量中提供新路径。

这是一个示例，我们将容器内的缓存移动到 /workdir/cache/af2/

export LOCAL_NIM_CACHE=/mount/largedisk/nim/.cache

docker run -it --rm --name alphafold2-multimer --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_CLI_API_KEY \
    -e NIM_CACHE_PATH=/workdir/cache/af2 \ ## Note: we must set this environment variable...
    -v $LOCAL_NIM_CACHE:/workdir/cache/af2 \ ## and provide the destination mount.
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2-multimer:1.0.0

Air-Gapped NIM 部署#

为了下载、检测和验证缓存模型的内容，NIM 需要连接到 NGC 的互联网连接。在 离线 或 air-gapped 环境中运行 NIM 时，这是不可能的。

为了解决这些限制，请配置 NIM 以忽略启动时的模型下载和验证错误，并将本地下载的 AF2 模型挂载到离线 NIM 服务器。工作流程如下

  1. 将模型下载到本地磁盘，该磁盘可以挂载到 NIM 的离线主机。

  2. 在运行 NIM 时，将缓存模型挂载到您的离线主机。

  3. 设置 NIM_IGNORE_MODEL_DOWNLOAD_FAIL=true，以便容器尝试模型下载和检测，但会优雅地失败，而不会终止。

要将模型下载到本地磁盘，请运行以下命令

ngc registry model download-version nim/deepmind/alphafold2-data:1.1.0

启动 NIM 容器，并从 $LOCAL_NIM_CACHE 中指定的路径挂载缓存模型

docker run -it --rm --name alphafold2-multimer --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_CLI_API_KEY \
    -e NIM_IGNORE_MODEL_DOWNLOAD_FAIL=true \
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2-multimer:1.0.0

为 NIM 查询/请求使用备用端口#

如果您有其他 HTTP 或 gRPC 服务器正在运行（例如，其他 NIM），您可能需要使用另一个端口用于您的 NIM，从而使 8000 端口可用于这些其他服务器。这可以通过以下方式完成

  1. 通过设置备用端口来更改暴露的端口，即 -p $HOST_PORT:$CONTAINER_PORT。

  2. 将 NIM_HTTP_API_PORT 环境变量更新为新的 $CONTAINER_PORT。

这是一个示例，我们将 NIM 设置为（对称地）在端口 7979 上运行

export LOCAL_NIM_CACHE=/mount/largedisk/nim/.cache

docker run -it --rm --name alphafold2-multimer --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_CLI_API_KEY \
    -e NIM_HTTP_API_PORT=7979 \ ## We must set the NIM_HTTP_API_PORT environment variable...
    -e NIM_CACHE_PATH=/workdir/cache/af2 \
    -v $LOCAL_NIM_CACHE:/workdir/cache/af2 \
    -p 7979:7979 \ ## as well as forward the port to host.
    nvcr.io/nim/deepmind/alphafold2-multimer:1.0.0

调整 MSA Runner 和线程#

AlphaFold2-Multimer NIM 在许多方面都进行了优化。其中一项优化包括能够根据您的机器上可用的 CPU 核心数量、可用 GPU 数量和可用 SSD 的速度来调整 MSA 进程，以获得更好的性能。（注意：我们不建议在使用硬盘驱动器进行存储的机器上运行 AlphaFold2-Multimer NIM，因为这会显著降低 NIM 的性能。）

AlphaFold2-Multimer NIM 中 MSA 的可调并行性有两个维度，由两个不同的环境变量设置

  1. NIM_PARALLEL_MSA_RUNNERS：控制 MSA 进程的粗粒度并行性。

  2. NIM_PARALLEL_THREADS_PER_MSA：控制每个 MSA 进程使用的线程数。

提示

NIM_PARALLEL_MSA_RUNNERS 和 NIM_PARALLEL_THREADS_PER_MSA 的乘积应 ≥ NIM 可用的 CPU 核心数。

例如，在具有 48 个 CPU 核心的机器上，合适的设置是

NIM_PARALLEL_MSA_RUNNERS=4
NIM_PARALLEL_THREADS_PER_MSA=12

这将分配 4 个 runner，每个 runner 有 12 个线程。

注意

NVIDIA 不建议使用超过 5 个进程 runner，并且当每个 MSA 的线程数是 CPU 中线程数的因子时，NIM 性能最佳。NIM_PARALLEL_MSA_RUNNERS=3 和 NIM_PARALLEL_THREADS_PER_MSA=8 的默认值适用于大多数机器。

要在运行时修改这些变量，请在 NIM 启动命令中将它们作为环境变量传递。例如，在 64 核机器上，当运行许多 MSA 数据库时，我们可能会将 NIM_PARALLEL_MSA_RUNNERS 增加到 6，并将 NIM_PARALLEL_THREADS_PER_MSA 减少到 10，以便我们使我们的工作负载很好地适应我们的机器

docker run -it --rm --name alphafold2-multimer --runtime=nvidia \
    -e CUDA_VISIBLE_DEVICES=0 \
    -e NGC_CLI_API_KEY \
    -e NIM_PARALLEL_MSA_RUNNERS=6 \ ## We set the number of parallel runners to 6, perhaps because we have used many MSA databases...
    -e NIM_PARALLEL_THREADS_PER_MSA=10 \ ## and we set the number of threads per MSA to 12, to distribute alignment workloads...
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -p 8000:8000 \
    nvcr.io/nim/deepmind/alphafold2-multimer:1.0.0
在本页中