NVIDIA Earth-2 FourCastNet NIM 部署指南#

使用本文档了解如何部署 NVIDIA Earth-2 FourCastNet NIM 的详细信息。

重要提示

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

查看可用的 NIM 容器信息#

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

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

拉取容器镜像#

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

Docker#

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

NGC CLI#

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

运行容器#

正如快速入门指南中所述,您可以运行以下命令来启动 FourCastNet NIM

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

docker run --rm --name fourcastnet --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/fourcastnet:1.0.0

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

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

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

  • --name fourcastnet — 此标志为容器命名为 “fourcastnet”。

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

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

  • --shm-size 4g — 为 Triton 分配主机内存以供使用。您可能需要根据部署和推理请求大小来增加此值。建议最小值为 4 GB。

  • -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 作为 /opt/nim/.cache 挂载到容器中,以便它可以用于存储和重用下载的模型。

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

检查 FourCastNet NIM 的状态#

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

docker ps

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

CONTAINER ID   IMAGE                 COMMAND                  CREATED          STATUS          PORTS                                                                                                          NAMES
d114948j4f55   fourcastnet   "/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'

停止正在运行的容器#

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

注意

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

以下示例使用上一个示例中的 ID 停止正在运行的容器。

docker kill d114948j4f55

模型检查点缓存#

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

# 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/fourcastnet:1.0.0

注意

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

重要提示

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

气隙部署(离线缓存)#

FourCastNet 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/fourcastnet:1.0.0

验证模型包现在是否在 LOCAL_NIM_CACHE 中。现在可以使用此缓存启动 NIM,而无需 NGC_API_KEY

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

气隙部署(本地模型)#

也可以通过使用 NGC CLI 工具 手动下载模型检查点并将这些文件挂载到容器中,以气隙方式运行 NIM。此模式也可用于运行 NIM 未分发的其他 SFNO 检查点。首先将模型文件下载到本地计算机。例如,考虑 FourCastNet NIM 的默认模型 earth2-sfno-era5-73ch:1.0.0-finetuned。使用 NGC CLI 将此模型下载到文件

ngc registry model download-version nim/nvidia/earth2-sfno-era5-73ch:1.0.0-finetuned

下载的文件夹现在将充当 LOCAL_NIM_CACHE。将源模型包挂载到 /opt/nim/workspace/ 目录并禁用模型下载。现在可以使用此缓存启动 NIM,而无需 NGC_API_KEY

export LOCAL_NIM_CACHE=${PWD}/earth2-sfno-era5-73ch_v1.0.0-finetuned/

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