先决条件自动拓扑选择使用 topology/block 插件 UltraServer 拓扑最佳实践限制创建 IMEX 序言脚本部署 prolog 脚本验证 IMEX 设置问题排查

配置 UltraServer 实例

NVIDIA UltraServer 实例使用 NVLink 跨实例连接多个 GPU。 GPU-to-GPU 通过 NVLink 进行通信要求使用任务中所有计算节点的 IP 地址配置 NV IDIA IMEX 服务。本页向您介绍如何创建 Slurm 序言脚本以自动为每个任务配置 IMEX，以及如何将其部署到您的 PCS 集群中。 AWS

先决条件

装有 Slurm 25.05 或更高版本的 AWS PCS 集群。
使用带有容量块的 UltraServer 实例的计算节点组。有关更多信息，请参阅将 AWS PCS 计算节点组配置为使用容量块。
安装了 NVIDIA 驱动程序、CUDA 和 IMEX 软件包的 AMI。必须在 AMI 上启用 IMEX 服务 (nvidia-imex)。
prolog 脚本必须在集群的 Slur Prolog m 设置中配置的路径下在所有计算节点上可用。您可以使用共享文件系统或将脚本包含在自定义 AMI 中。

自动拓扑选择

当您使用 UltraServer 计算节点组创建 AWS PCS 集群时，系统会检查容量块，识别 UltraServer 类型，并使用相应的拓扑插件配置 Slurm。此过程自动运行，不需要任何配置。

AWS PCS 通过动态生成的topology.yaml文件管理拓扑。随着群集通过增加或移除计算节点组而演变， AWS PCS 会不断协调拓扑配置以反映当前群集状态。

分组拓扑模型均匀的高带宽通信域，其中所有 GPU 都参与单个高速 NVLink 域，延迟接近均匀。该topology/block插件根据定义的网络拓扑分配节点，确保将任务放置在同一拓扑结构中， UltraServer 并且 GPU-to-GPU 通信使用 NVLink 而不是回退到 EFA 网络。

使用 topology/block 插件

提交到 UltraServer 队列的任务由 Slurm 调度器自动放置在同一个区块 (UltraServer) 中。在使用下一个区块之前，将区块中的所有节点分配给任务。作业必须使用该--exclusive标志，因为 IMEX 要求每个任务都有一个专用 IMEX 域，而重新启动 IMEX 服务会中断在同一节点上运行的任何其他作业。有关该 topology/block 插件的更多信息，请参阅 Slurm 文档中的拓扑指南和 topology.yaml 参考文档。

提交任务时，您可以使用带sbatch和srun命令的--segment=N参数来指定要在同一块中组合在一起的节点数。分段的大小必须小于或等于规划区块大小。使用--segment并不能保证区段会被放置在不同的方块上。

以下是在具有 18 个节点的队列上分配区块的示例场景 UltraServers （例如，有 18 个p6e-gb200.36xlarge实例u-p6e-gb200x72的 a）：


# Allocate a whole block (single UltraServer)
sbatch -N18 --exclusive --partition=ultra job.sh

# Allocate two full blocks (two UltraServers)
sbatch -N36 --exclusive --partition=ultra job.sh

# Allocate 24 nodes: full block is filled before 2nd block is used
sbatch -N24 --exclusive --partition=ultra job.sh

# Allocate 24 nodes as segments of 12, segments can fit on blocks with available space
sbatch -N24 --segment=12 --exclusive --partition=ultra job.sh

UltraServer 拓扑最佳实践

为了使 AWS PCS 中的架 UltraServer 构实现最佳性能：

使用分段来提高可用性：使用值小于完整区块大小的--segment选项，允许在节点已耗尽或不可用的区块上调度任务。这使用绝对的 NVLink 本地性来缩短任务调度时间。有关更多信息，请参阅使用 Topology/Block 插件获得对节点调度的更多控制权。
考虑任务大小和块大小：如果BlockSizes=18，最多有 18 个实例的作业始终在单个实例上运行 UltraServer。请求超过 18 个节点的任务跨越多个节点 UltraServers，每个区块内有 NVLink 连接，区块之间有 EFA 网络。
分段权衡：使用会--segment禁用高阶区块大小调度。否--segment则，在使用下一个方块之前，将填充一个完整的方块。使用--segment，可以将分段放置在具有足够可用节点的任何区块上。

限制

AWS PCS 要求将 UltraServer 计算节点组置于专用队列中，与非UltraServer 计算节点组分开。
每个计算节点组只能与一个 UltraServer容量块关联。
当队列包含多个 UltraServer 计算节点组时，每个计算节点组的实例数之间必须有 2 的次方比（例如，9 和 18，或 2 和 8）。
UltraServer 队列中的任务必须使用该--exclusive标志来确保每个任务都有专用 IMEX 域。
由 prolog 脚本创建的 IMEX 域仅包含分配给任务的节点。由于 AWS PCS 为任务动态分配节点，因此无法预先创建 IMEX 域，必须在任务启动时进行配置。跨多个任务 UltraServers 不会在不同节点之间建立 NVLink 连接， UltraServers 因为 IMEX 域基于 IP 地址，无法跨不同的 NVLink 架构进行桥接。
在服务运行期间，nodes_config.cfg无法修改 IMEX。如果更改配置，IMEX 将永久阻止与不匹配节点的通信。唯一的恢复方法是停止所有节点上的 IMEX 并使用一致的配置重新启动。

创建 IMEX 序言脚本

以下91_nvidia_imex_prolog.sh脚本在 Slurm 任务启动时在计算节点上配置 NVIDIA IMEX。脚本名称使用遵循 SchedMD 命名约定的91前缀进行序言执行顺序。

此脚本将执行以下操作：

检查任务是否为独占任务。IMEX 配置需要独占任务，以防止中断同一节点上其他正在运行的作业。
检查实例类型是否为支持的 UltraServer 类型（p6e-gb200或p6e-gb300）。
检查 IMEX 服务是否已启用。
如果 IMEX 默认频道设备不存在，则创建该设备。
将所有任务节点的私有 IP 地址写入/etc/nvidia-imex/nodes_config.cfg。
重新启动 IMEX 服务。

重要

这个 prolog 脚本要求提交带有--exclusive标志的作业。当其他作业在同一节点上运行时重新启动 IMEX 将中断这些作业。

注意

此序言脚本支持p6e-gb200和p6e-gb300实例类型。要使其适应其他实例类型，请更新脚本中的ALLOWED_INSTANCE_TYPES模式。

将以下脚本另存为91_nvidia_imex_prolog.sh：


#!/usr/bin/env bash

# This prolog script configures NVIDIA IMEX on compute nodes involved in the job execution.
#
# In particular:
# - Checks whether the job is executed exclusively.
#   If not, it exits immediately because it requires jobs to be executed exclusively.
# - Checks if it is running on a supported UltraServer instance type.
#   If not, it exits immediately because IMEX must be configured only on those instance types.
# - Checks if the IMEX service is enabled.
#   If not, it exits immediately because IMEX must be enabled to get configured.
# - Creates the IMEX default channel.
#   For more information about IMEX channels, see
#   https://docs.nvidia.com/multi-node-nvlink-systems/imex-guide/imexchannels.html
# - Writes the private IP addresses of compute nodes into /etc/nvidia-imex/nodes_config.cfg.
# - Restarts the IMEX system service.
#
# REQUIREMENTS:
#  - This prolog assumes to be run only with exclusive jobs.

LOG_FILE_PATH="/var/log/nvidia-imex-prolog.log"
SCONTROL_CMD="scontrol"
IMEX_START_TIMEOUT=60
IMEX_STOP_TIMEOUT=15
ALLOWED_INSTANCE_TYPES="^(p6e-gb200|p6e-gb300)"
IMEX_SERVICE="nvidia-imex"
IMEX_NODES_CONFIG="/etc/nvidia-imex/nodes_config.cfg"

function info() {
  echo "$(date "+%Y-%m-%dT%H:%M:%S.%3N") [INFO] [PID:$$] [JOB:${SLURM_JOB_ID}] $1"
}

function warn() {
  echo "$(date "+%Y-%m-%dT%H:%M:%S.%3N") [WARN] [PID:$$] [JOB:${SLURM_JOB_ID}] $1"
}

function error() {
  echo "$(date "+%Y-%m-%dT%H:%M:%S.%3N") [ERROR] [PID:$$] [JOB:${SLURM_JOB_ID}] $1"
}

function prolog_end() {
    info "PROLOG End JobId=${SLURM_JOB_ID}: $0"
    info "----------------"
    exit 0
}

function get_instance_type() {
  local token=$(curl -X PUT -s "http://169.254.169.254/latest/api/token" -H "X-aws-ec2-metadata-token-ttl-seconds: 21600")
  curl -s -H "X-aws-ec2-metadata-token: ${token}" http://169.254.169.254/latest/meta-data/instance-type
}

function return_if_unsupported_instance_type() {
  local instance_type=$(get_instance_type)

  if [[ ! ${instance_type} =~ ${ALLOWED_INSTANCE_TYPES} ]]; then
    info "Skipping IMEX configuration because instance type ${instance_type} does not support it"
    prolog_end
  fi
}

function return_if_imex_disabled() {
  if ! systemctl is-enabled "${IMEX_SERVICE}" &>/dev/null; then
    warn "Skipping IMEX configuration because system service ${IMEX_SERVICE} is not enabled"
    prolog_end
  fi
}

function return_if_job_is_not_exclusive() {
  if [[ "${SLURM_JOB_OVERSUBSCRIBE}" =~ ^(NO|TOPO)$  ]]; then
    info "Job is exclusive, proceeding with IMEX configuration"
  else
    info "Skipping IMEX configuration because the job is not exclusive"
    prolog_end
  fi
}

function get_ips_from_node_names() {
  local _nodes=$1
  ${SCONTROL_CMD} -ao show node "${_nodes}" | sed 's/^.* NodeAddr=\([^ ]*\).*/\1/'
}

function reload_imex() {
  info "Stopping IMEX"
  timeout ${IMEX_STOP_TIMEOUT} systemctl stop ${IMEX_SERVICE}
  pkill -9 ${IMEX_SERVICE}

  info "Restarting IMEX"
  timeout ${IMEX_START_TIMEOUT} systemctl start ${IMEX_SERVICE}
}

function create_default_imex_channel() {
  info "Creating IMEX default channel"
  MAJOR_NUMBER=$(cat /proc/devices | grep nvidia-caps-imex-channels | cut -d' ' -f1)
  if [ ! -d "/dev/nvidia-caps-imex-channels" ]; then
    sudo mkdir /dev/nvidia-caps-imex-channels
  fi

  # Then check and create device node
  if [ ! -e "/dev/nvidia-caps-imex-channels/channel0" ]; then
    sudo mknod /dev/nvidia-caps-imex-channels/channel0 c $MAJOR_NUMBER 0
    info "IMEX default channel created"
  else
    info "IMEX default channel already exists"
  fi
}

{
  info "PROLOG Start JobId=${SLURM_JOB_ID}: $0"

  return_if_job_is_not_exclusive
  return_if_unsupported_instance_type
  return_if_imex_disabled

  create_default_imex_channel

  IPS_FROM_CR=$(get_ips_from_node_names "${SLURM_NODELIST}")

  info "Node Names: ${SLURM_NODELIST}"
  info "Node IPs: ${IPS_FROM_CR}"
  info "IMEX Nodes Config: ${IMEX_NODES_CONFIG}"

  info "Updating IMEX nodes config ${IMEX_NODES_CONFIG}"
  echo "${IPS_FROM_CR}" > "${IMEX_NODES_CONFIG}"
  reload_imex

  prolog_end

} 2>&1 | tee -a "${LOG_FILE_PATH}" | logger -t "91_nvidia_imex_prolog"

部署 prolog 脚本

要部署 prolog 脚本，请在集群上配置 Slurm Prolog 和PrologFlags自定义设置。您可以在创建集群或更新现有集群时进行设置。

有关配置自定义 Slurm 设置的更多信息，请参阅。的自定义 Slurm 设置 AWS PCS 集群有关更新集群的过程，请参阅更新一个 AWS PCS 集群。

在集群上设置以下自定义 Slurm 设置：

Prolog-例如/shared/scripts/prolog.d/91_nvidia_imex_prolog.sh，设置为计算节点上可用 prolog 脚本的路径。
PrologFlags-设置为Alloc,NoHold。当它首次分配给任务时，该Alloc标志会在每个节点上运行序言。该NoHold标志可防止在序言失败时暂停作业。

例使用 prolog 设置更新集群 AWS CLI


aws pcs update-cluster \
  --cluster-identifier <cluster-id> \
  --slurm-configuration '{
    "slurmCustomSettings": [
      {"parameterName": "Prolog", "parameterValue": "/shared/scripts/prolog.d/91_nvidia_imex_prolog.sh"},
      {"parameterName": "PrologFlags", "parameterValue": "Alloc,NoHold"}
    ]
  }'

验证 IMEX 设置

当您提交 Slurm 作业时，prolog 脚本会自动运行。提交测试任务以验证所有节点上的 IMEX 配置是否正确。

创建名为的作业脚本check-imex.sh：


#!/bin/bash
#SBATCH --job-name=nvidia-imex-status
#SBATCH --ntasks-per-node=1
#SBATCH --exclusive
#SBATCH --output=imex-status-%j.out
#SBATCH --error=imex-status-%j.err

srun bash -c "/usr/bin/nvidia-imex-ctl -N > imex_result_\${SLURM_JOB_ID}_\$(hostname).out 2>&1"

将任务提交到您的 UltraServer 队列中：


sbatch -p queue-name -N node-count check-imex.sh

检查输出文件。成功的 IMEX 配置会在连接表中显示所有处于C（已连接）READY状态的节点，并Domain State: UP且：


Connectivity Table Legend:
I - Invalid - Node wasn't reachable, no connection status available
N - Never Connected
R - Recovering - Connection was lost, but clean up has not yet been triggered.
D - Disconnected - Connection was lost, and clean up has been triggreed.
A - Authenticating - If GSSAPI enabled, client has initiated mutual authentication.
!V! - Version mismatch, communication disabled.
!M! - Node map mismatch, communication disabled.
C - Connected - Ready for operation

5/4/2026 06:08:10.580
Nodes:
Node #0   - 172.31.48.81    - READY                - Version: 570.172
Node #1   - 172.31.48.98    - READY                - Version: 570.172
Node #2   - 172.31.48.221   - READY                - Version: 570.172
Node #3   - 172.31.49.228   - READY                - Version: 570.172
Node #4   - 172.31.50.39    - READY                - Version: 570.172
Node #5   - 172.31.50.44    - READY                - Version: 570.172
Node #6   - 172.31.51.66    - READY                - Version: 570.172
Node #7   - 172.31.51.157   - READY                - Version: 570.172
Node #8   - 172.31.52.239   - READY                - Version: 570.172
Node #9   - 172.31.53.80    - READY                - Version: 570.172
Node #10  - 172.31.54.95    - READY                - Version: 570.172
Node #11  - 172.31.54.183   - READY                - Version: 570.172
Node #12  - 172.31.54.203   - READY                - Version: 570.172
Node #13  - 172.31.54.241   - READY                - Version: 570.172
Node #14  - 172.31.55.59    - READY                - Version: 570.172
Node #15  - 172.31.55.187   - READY                - Version: 570.172
Node #16  - 172.31.55.197   - READY                - Version: 570.172
Node #17  - 172.31.56.47    - READY                - Version: 570.172

 Nodes From\To  0   1   2   3   4   5   6   7   8   9   10  11  12  13  14  15  16  17
       0        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
       1        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
       2        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
       3        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
       4        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
       5        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
       6        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
       7        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
       8        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
       9        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
      10        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
      11        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
      12        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
      13        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
      14        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
      15        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
      16        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C
      17        C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C   C

Domain State: UP

问题排查

节点配置映射不匹配: 如果您!M!在连接表或 IMEX 日志Node configuration mapping mismatch中看到，该nodes_config.cfg文件在所有节点上都不一致。IMEX 永久阻止与不匹配节点的通信。验证 prolog 脚本是否在所有节点上运行，然后重新启动作业。
找不到 IMEX 频道设备: 如果 prolog 脚本无法创建 IMEX 频道设备，请验证 NVIDIA 驱动程序是否已安装且/proc/devices其中包含。nvidia-caps-imex-channels
Prolog 脚本未运行: 验证是否已在群集上配置Prolog和PrologFlags设置。检查脚本文件是否存在于计算节点上并具有执行权限 (chmod 0755)。

Javascript 在您的浏览器中被禁用或不可用。

要使用 Amazon Web Services 文档，必须启用 Javascript。请参阅浏览器的帮助页面以了解相关说明。

文档惯例

将计算节点组配置为使用容量块

带有 AWS PCS 的容量块的常见问题解答