dstack项目中MPI并行计算任务配置的优化方案

背景介绍

在分布式计算领域，MPI(Message Passing Interface)是一种广泛使用的并行编程标准。dstack作为一个开源的机器学习基础设施项目，近期对其MPI任务配置进行了优化，旨在简化用户在使用mpirun等工具时的配置复杂度。

原有问题分析

在优化前，用户在使用dstack运行MPI任务时需要编写大量自定义逻辑，包括但不限于：

1. 节点启动顺序控制
2. 主从节点同步机制
3. 任务终止条件判断
4. 节点间通信配置

这些复杂的配置不仅增加了用户的学习成本，也容易引入错误，降低了开发效率。

优化方案设计

dstack团队经过讨论，决定采用通用化而非特定化的解决方案，主要引入了两个关键配置项：

1. 节点启动顺序控制(startup_order)

该配置项支持两种模式：

• workers-first：工作节点优先启动
• master-first：主节点优先启动

对于MPI任务，通常需要选择workers-first模式，确保所有工作节点就绪后再启动主节点。

2. 任务停止条件(stop_criteria)

该配置项同样支持两种模式：

• all-done：所有节点任务完成后停止
• master-done：主节点完成后即停止

MPI任务通常选择master-done模式，主节点完成任务后即可终止整个运行。

优化后的配置示例

优化后，一个典型的NCCL测试任务配置简化为：

type: task
name: nccl-tests

nodes: 2
startup_order: workers-first
stop_criteria: master-done

image: dstackai/efa
env:
  - NCCL_DEBUG=INFO
commands:
  - |
    if [ ${DSTACK_NODE_RANK} -eq 0 ]; then
      cd /root/nccl-tests/build
      : > hostfile
      for ip in ${DSTACK_NODES_IPS}; do
        echo "${ip} slots=${DSTACK_GPUS_PER_NODE}" >> hostfile
      done
      MPIRUN='mpirun --allow-run-as-root --hostfile hostfile'
      ${MPIRUN} \
        -n ${DSTACK_GPUS_NUM} -N ${DSTACK_GPUS_PER_NODE} \
        --mca btl_tcp_if_exclude lo,docker0 \
        --bind-to none \
        ./all_reduce_perf -b 8 -e 8G -f 2 -g 1
    else
      sleep infinity
    fi

技术实现细节

1. 节点启动顺序保证：系统会严格按照配置的启动顺序初始化节点，确保依赖关系。

2. 自动主机文件生成：系统自动生成包含所有节点信息的主机文件，并通过环境变量DSTACK_NODES_IPS和DSTACK_HOSTFILE提供给用户。

3. 资源信息注入：系统自动注入节点GPU信息(DSTACK_GPUS_PER_NODE和DSTACK_GPUS_NUM)，简化资源配置。

4. 优雅终止机制：根据配置的停止条件，系统会智能判断何时终止整个任务运行。

最佳实践建议

1. 对于MPI类任务，推荐组合使用workers-first和master-done。

2. 工作节点应保持运行状态直到主节点完成任务，可通过sleep infinity实现。

3. 充分利用系统提供的环境变量简化配置，如节点IP列表和GPU信息。

4. 主节点负责协调任务，工作节点只需保持运行状态即可。

未来展望

虽然当前方案已经大大简化了MPI任务的配置，但团队仍在考虑进一步优化：

1. 内置MPI任务模板，进一步减少样板代码
2. 更智能的故障恢复机制
3. 更细粒度的节点控制选项

这种通用化的设计思路不仅适用于MPI任务，也为未来支持其他类型的分布式计算框架提供了良好的扩展基础。