HOOOS

Slurm 调度下 MPI 作业的 NVIDIA MPS 动态启停与自动配置方案

0 10 HPC基建指南 SlurmNVIDIA MPSMPI高性能计算
Apple

在利用 Slurm 调度器运行 MPI 多机多卡作业时，若多个 MPI 进程（Ranks）需要共享同一张 GPU 卡，默认情况下会因为 CUDA Context 切换开销巨大而导致显卡利用率低下。NVIDIA MPS（Multi-Process Service，多进程服务）通过允许多个 CUDA 进程共享同一个 GPU 上下文，能显著提升 MPI 作业的并发性能。

需要明确的是，MPS 服务应当在实际运行 MPI 任务的 GPU 计算节点（Compute Nodes）上启动，而非 Slurm 的主控制节点（Controller Node）。

在生产环境中，如果全局常开 MPS，会导致不同用户的作业相互干扰、显存无法物理隔离甚至引发安全越权问题。因此，实现 “作业级动态启停与配置” 是最优解。以下提供两种主流的生产级部署方案。

方案一：利用 Slurm 原生 GRES MPS 机制（推荐）

自 Slurm 19.05 版本起，系统已经原生支持将 MPS 作为一种通用的通用资源（GRES）进行管理。Slurm 会在作业启动时自动拉起 nvidia-cuda-mps-control 守护进程，并在作业结束时自动将其销毁。

1. 配置 slurm.conf

在控制节点和计算节点的 slurm.conf 中，将 mps 加入到通用资源类型中：

GresTypes=gpu,mps

2. 配置 gres.conf

在每个 GPU 计算节点上，定义 GPU 与 MPS 的对应关系。例如，节点有 2 张 A100 GPU，每张卡虚拟出 100 个 MPS 份额（即代表显卡计算资源的百分比）：

Name=gpu File=/dev/nvidia0
Name=mps Count=100 File=/dev/nvidia0

Name=gpu File=/dev/nvidia1
Name=mps Count=100 File=/dev/nvidia1

3. 用户提交 MPI 作业

用户在提交作业时，只需通过 --gres 参数同时申请 GPU 和对应的 MPS 比例。

# 申请 1 张 GPU 和 50% 的 GPU 算力，运行 4 个 MPI 进程
srun --gres=gpu:1,mps:50 -n 4 --mpi=pmix ./mpi_gpu_application

工作原理： Slurm 会在分配给该作业的计算节点上自动创建专属于该 Job 的 MPS 目录，启动对应份额的守护进程，并将环境变量 CUDA_VISIBLE_DEVICES 和 CUDA_MPS_ACTIVE_THREAD_PERCENTAGE=50 自动注入到 MPI 进程中。

方案二：基于 Slurm Prolog/Epilog 脚本动态启停（深度定制方案）

如果使用的是老版本 Slurm，或者需要实现更严苛的多租户安全隔离（例如：强制要求 MPS 守护进程以作业属主普通用户身份运行，严禁以 root 运行），则需要利用 Slurm 的节点级钩子脚本（Prolog/Epilog）进行动态干预。

1. 核心技术陷阱与解决机制

  • 安全风险：如果以 root 运行 MPS Control，普通用户可以通过共享管道向其发送控制命令，存在越权风险。因此必须使用 su - $SLURM_JOB_USER 以作业用户身份拉起。
  • 变量透传问题：Prolog 脚本运行在 root 权限下，其设置的临时环境变量（如 CUDA_MPS_PIPE_DIRECTORY）无法直接传递给用户的作业环境。必须借助 TaskProlog 机制进行变量注入。

2. 编写 Prolog 脚本（启动 MPS）

在计算节点创建 /etc/slurm/prolog.d/mps_start（确保权限为 700，属主为 root）：

#!/bin/bash
# 路径：/etc/slurm/prolog.d/mps_start

# 检测作业是否申请了 GPU
if [ -z "$SLURM_JOB_GPUS" ]; then
    exit 0
fi

# 定义该作业独占的 MPS 管道与日志目录（引入 JobID 防止多作业冲突）
MPS_DIR="/tmp/nvidia-mps_${SLURM_JOB_ID}"
MPS_LOG="/tmp/nvidia-log_${SLURM_JOB_ID}"

mkdir -p "$MPS_DIR" "$MPS_LOG"
chown -R "${SLURM_JOB_UID}:${SLURM_JOB_GID}" "$MPS_DIR" "$MPS_LOG"

# 将分配给该作业的实际 GPU 物理卡号转换为 MPS 可识别的格式
# 注：Slurm 提供给 Prolog 的变量为 SLURM_JOB_GPUS=0,1 等
export CUDA_VISIBLE_DEVICES=$SLURM_JOB_GPUS

# 以作业用户的身份启动 MPS 守护进程
su - "$SLURM_JOB_USER" -c "
    export CUDA_MPS_PIPE_DIRECTORY=$MPS_DIR; \
    export CUDA_MPS_LOG_DIRECTORY=$MPS_LOG; \
    export CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES; \
    nvidia-cuda-mps-control -d
"

exit 0

3. 编写 TaskProlog 脚本（注入环境变量）

为了让 MPI 进程在运行时能够寻址到正确的 MPS 管道，必须在任务拉起前向其注入环境变量。

在 slurm.conf 中配置：

TaskProlog=/etc/slurm/task_prolog

编写 /etc/slurm/task_prolog：

#!/bin/bash
# 路径：/etc/slurm/task_prolog

# 如果当前作业存在对应的 MPS 专属管道目录，则进行注入
MPS_DIR="/tmp/nvidia-mps_${SLURM_JOB_ID}"
if [ -d "$MPS_DIR" ]; then
    echo "export CUDA_MPS_PIPE_DIRECTORY=$MPS_DIR"
    echo "export CUDA_MPS_LOG_DIRECTORY=/tmp/nvidia-log_${SLURM_JOB_ID}"
fi

注：Slurm 允许通过 TaskProlog 的 stdout 打印 export VAR=VAL 来动态修改任务的环境变量。

4. 编写 Epilog 脚本（停止与清理）

当 MPI 作业结束或被 scancel 强行终止时，必须确保清理 MPS 进程及临时管道文件，防止残留污染后续作业。

创建 /etc/slurm/epilog.d/mps_stop：

#!/bin/bash
# 路径：/etc/slurm/epilog.d/mps_stop

MPS_DIR="/tmp/nvidia-mps_${SLURM_JOB_ID}"

if [ -d "$MPS_DIR" ]; then
    # 以作业用户身份优雅关闭该 Job 的 MPS Control 守护进程
    su - "$SLURM_JOB_USER" -c "
        export CUDA_MPS_PIPE_DIRECTORY=$MPS_DIR; \
        echo quit | nvidia-cuda-mps-control
    "
    # 强制清理残留文件
    rm -rf "$MPS_DIR"
    rm -rf "/tmp/nvidia-log_${SLURM_JOB_ID}"
fi

exit 0

MPI 配合 MPS 运行的关键最佳实践

在配置好上述自动启停机制后，要使 MPI 作业充分压榨 GPU 算力，还需注意以下运行细节：

  1. 显卡模式设置为独占进程模式（Exclusive Process）
    在计算节点上，建议将 GPU 设置为独占进程模式，以确保非 MPS 作业不会误入此 GPU。
    nvidia-smi -c EXCLUSIVE_PROCESS
  2. 设置合理的 CPU-GPU 亲和性
    MPI 任务调度中，CPU 核心与 GPU 物理位置的绑定至关重要。启动 srun 时，必须配合 --cpu-bind 确保 MPI rank 与对应的 NUMA 节点对齐，否则 PCIe 跨 Socket 通信会严重抵消 MPS 带来的性能红利。
    srun --cpu-bind=cores --gres=gpu:1,mps:100 -n 8 --mpi=pmix ./mpi_app
  3. 限制单个进程的显存/算力上限
    如果不加限制，某个失控的 MPI Rank 可能会吃掉整张显卡的全部显存。可以在 Job 提交脚本中设置以下环境变量限制单个 Rank 的资源占比：
    # 限制单个 MPI 进程最多使用 GPU 30% 的线程资源
export CUDA_MPS_ACTIVE_THREAD_PERCENTAGE=30

点评评价

captcha
健康