NVIDIA Earth-2 修正扩散 NIM 部署指南#

使用本文档详细了解如何部署 NVIDIA Earth-2 修正扩散 (CorrDiff) NIM

重要提示

在您使用本文档之前,您必须满足所有前提条件

查看可用的 NIM 容器信息#

容器镜像标签可以使用以下命令通过 NGC CLI 查看,类似于 NGC 上的其他容器镜像。

ngc registry image info nvcr.io/nim/nvidia/corrdiff:1.0.0

拉取容器镜像#

使用以下命令之一拉取容器镜像

Docker#

docker pull nvcr.io/nim/nvidia/corrdiff:1.0.0

NGC CLI#

ngc registry image pull nvcr.io/nim/nvidia/corrdiff:1.0.0

运行容器#

如同 快速入门指南 中所示,您可以运行以下命令来启动 CorrDiff NIM

export LOCAL_NIM_CACHE=~/.cache/nim
export NGC_API_KEY=<NGC API Key>

docker run --rm --name corrdiff --runtime=nvidia --gpus all --shm-size 1g \
    -p 8000:8000 \
    -e NGC_API_KEY \
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -u $(id -u) \
    nvcr.io/nim/nvidia/corrdiff:1.0.0

以下是对命令及其选项以及 NIM 启动时可用的其他选项的概述

  • docker run — 这是从 Docker 镜像运行新容器的命令。

  • --rm — 此标志告诉 Docker 在容器退出时自动删除容器。此属性对于一次性运行或测试非常有用,因为它可防止容器被遗留。

  • --name corrdiff — 此标志为容器指定名称 “corrdiff”。

  • --runtime=nvidia — 此标志指定用于容器的运行时。在本例中,它设置为 “nvidia”,用于 GPU 加速。

  • --gpus all — 使 NIM 能够使用部署容器的机器上所有可用的 GPU。

  • --shm-size 1g — 为 Triton 分配主机内存以供使用。您可能需要根据部署情况增加此值。

  • -p 8000:8000 — 此标志将主机上的端口 8000 映射到容器中的端口 8000。这允许您从主机访问容器的服务。

  • -e NGC_API_KEY — 这会将 NGC_API_KEY 环境变量(以及父终端中设置的值)传递给容器。这用于 NVIDIA 的 NGC(NVIDIA GPU Cloud)服务的身份验证,包括在 NIM 缓存中不存在模型数据时下载模型数据。

  • -v <source>:<dest> — 将主机目录 /home/$USER/.cache/nim 作为 /home/nvs/.cache 挂载到容器中,以便它可以用于存储和重用下载的模型。

  • -u $(id -u) — 在 NIM 容器内使用与系统用户相同的用户,以避免在本地缓存目录中下载模型时出现权限不匹配。

检查 CorrDiff NIM 的状态#

您可以使用以下命令查看系统上正在运行的 Docker 容器

docker ps

如果 NIM 正在运行,这将返回如下所示的输出

CONTAINER ID   IMAGE                 COMMAND                  CREATED          STATUS          PORTS                                                                                                          NAMES
d114948j4f55   corrdiff   "/opt/nvidia/nvidia_…"   46 minutes ago   Up 46 minutes   6006/tcp, 0.0.0.0:8000->8000/tcp, :::8000->8000/tcp, 8888/tcp, 0.0.0.0:50051->50051/tcp, :::50051->50051/tcp   test1

输出的第一列是 Docker 容器 ID,这对于与容器交互非常有用。其余字段描述了容器正在运行的镜像、命令(在本例中为 NIM 服务器软件)、容器构建时间、其状态、运行时间、任何打开的端口以及最终名称(由启动命令给出)。

要检查 NIM 是否已完全启动且运行状况良好,请使用 health API,如果 NIM 准备就绪,它将返回 200 响应

curl -X 'GET' \
    'http://127.0.0.1:8000/v1/health/ready' \
    -H 'accept: application/json'

停止正在运行的容器#

如果由于某种原因您的容器进入了您无法与之交互的状态,或者由于任何原因必须停止它,您可以使用从 docker ps 命令获得的 ID 使用 kill 命令。这将立即终止正在运行的容器。

注意

任何正在进行的请求都将被取消,并且数据可能会丢失。

以下示例使用前面示例中的 ID 停止正在运行的容器。

docker kill d114948j4f55

模型检查点缓存#

在初始启动时,容器从 NGC 下载 CorrDiff 模型参数和支持数据。您可以通过使用缓存目录本地缓存模型权重来跳过后续运行中的此下载步骤,如下例所示。

# Create the cache directory on the host machine
export LOCAL_NIM_CACHE=~/.cache/nim
mkdir -p $LOCAL_NIM_CACHE

# Run the container with the cache directory mounted in the appropriate location
docker run --rm --runtime=nvidia --gpus all --shm-size 4g \
    -p 8000:8000 \
    -e NGC_API_KEY \
    -v $LOCAL_NIM_CACHE:/opt/nim/.cache \
    -u $(id -u) \
    nvcr.io/nim/nvidia/corrdiff:1.0.0

注意

缓存模型检查点可以为后续容器运行节省大量时间。

重要提示

关键是 LOCAL_NIM_CACHE 目录为 Docker 容器内的用户提供读/写访问权限。如果 NIM 在下载模型检查点时生成权限错误,则表示挂载的文件夹的权限不正确。

气隙部署(离线缓存)#

CorrDiff NIM 支持在气隙系统(也称为气墙、气隙或断开连接的网络)中提供模型服务。启动时,如果 NIM 检测到模型文件已下载,它将直接从缓存提供模型服务。用户可以通过挂载本地文件夹到缓存,在有权访问 NGC 服务器的硬件上启动 NIM,然后将缓存的模型文件移动到相应的气隙系统,来预缓存 NIM 的模型包。

从挂载模型文件夹的在线部署开始

export LOCAL_NIM_CACHE=<Local folder>
export NGC_API_KEY=<NGC API Key>

docker run --rm --runtime=nvidia --gpus all --shm-size 4g \
    -p 8000:8000 \
    -e NGC_API_KEY \
    -v ${LOCAL_NIM_CACHE}:/opt/nim/.cache \
    -u $(id -u) \
    -t nvcr.io/nim/nvidia/corrdiff:1.0.0

模型包现在应该在 LOCAL_NIM_CACHE 中。现在可以使用此缓存启动 NIM,而无需 NGC_API_KEY