如何在Python程序中获取mpiexec -n参数指定的进程数

首先明确结论：直接使用sys模块、click库无法直接拿到-n参数的值，os.environ可以实现需求，更标准的方案是通过mpi4py库获取。

原因说明 #

你运行的mpiexec -n 6 python cuda.py命令中，-n 6是mpiexec本身的启动参数，不会传递给后面的cuda.py Python脚本，所以sys.argv和专门解析命令行参数的click库，都只能拿到cuda.py相关的传入参数，看不到-n 6的存在。

可行实现方案 #

方案1：通过os.environ读取环境变量 #

几乎所有主流MPI实现都会把总进程数（也就是-n后面的数值）写入环境变量，直接读取对应变量即可：

import os

# 兼容不同MPI实现的环境变量
mpi_size = os.getenv("OMPI_COMM_WORLD_SIZE")  # 适配OpenMPI环境
if not mpi_size:
    mpi_size = os.getenv("PMI_SIZE")  # 适配MPICH/Intel MPI/MS-MPI环境

if mpi_size:
    mpi_size = int(mpi_size)
print(f"mpiexec设置的-n参数值为: {mpi_size}")

方案2：通过mpi4py库获取（最稳定通用） #

如果你的程序本身就是MPI并行程序，最标准且跨MPI实现的方案是直接用mpi4py的接口获取，不需要关心环境变量差异：
首先安装依赖：pip install mpi4py
代码示例：

from mpi4py import MPI

comm = MPI.COMM_WORLD
mpi_size = comm.Get_size()
print(f"mpiexec设置的总进程数（-n参数值）为: {mpi_size}")

注意：如果是单进程场景下运行上述代码，拿到的数值是1，符合预期。

内容的提问来源于stack exchange，提问作者stapmoshun