使用 MPI4PY 快速失败

Question

在使用 mpi4py 运行 MPI 脚本时，我希望以下行为：当任何进程引发异常时，mpirun(及其生成的进程(应立即退出并显示非零错误代码。但相反，我发现即使一个或多个进程引发异常，执行也会继续。

我正在使用 mpi4py 3.0.0 和 OpenMPI 2.1.2。我正在运行此脚本 mpirun --verbose -mca orte_abort_on_non_zero_status 1 -n 4 python my_script.py .我预计这会在睡眠被击中之前立即结束，但相反，进程与秩 ！= 0 睡眠：

import time
import mpi4py
def main():
    import mpi4py.MPI
    mpi_comm = mpi4py.MPI.COMM_WORLD
    if mpi_comm.rank == 0:
        raise ValueError('Failure')

    print('{} continuing to execute'.format(mpi_comm.rank))
    time.sleep(10)
    print('{} exiting'.format(mpi_comm.rank)

if __name__ == '__main__':
    main()

如何获得我想要的行为(如果任何进程失败，则快速失败(？

谢谢！

Answer 1

这似乎是 mpi4py 的已知问题。从 https://groups.google.com/forum/#!topic/mpi4py/RovYzJ8qkbc，我读到：

mpi4py 为您初始化/完成 MPI。初始化发生在 导入时间，以及 Python 进程即将完成时的完成 finalize(我正在使用 Py_AtExit(( C-API 调用来执行此操作(。如 MPI_Finalize(( 是集体的，在大多数 MPI impls 中可能是阻塞的， 你得到僵局。

解决方案是覆盖sys.excepthook并显式调用其中的MPI.COMM_WORLD.Abort。

这是您修改的代码：

import sys
import time
import mpi4py.MPI
mpi_comm = mpi4py.MPI.COMM_WORLD
def mpiabort_excepthook(type, value, traceback):
    mpi_comm.Abort()
    sys.__excepthook__(type, value, traceback)
def main():
    if mpi_comm.rank == 0:
        raise ValueError('Failure')

    print('{} continuing to execute'.format(mpi_comm.rank))
    time.sleep(10)
    print('{} exiting'.format(mpi_comm.rank))
if __name__ == "__main__":
    sys.excepthook = mpiabort_excepthook
    main()
    sys.excepthook = sys.__excepthook__

Answer 2

事实证明，mpi4py可以作为解决此问题的模块运行(内部通过调用Abort((，就像jcgiret说的那样(：

mpirun --verbose -mca orte_abort_on_non_zero_status 1 -n 4 python -m mpi4py my_script.py