用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。初始化发生在导入时，当process即将完成时(我正在使用Py_AtExit() call调用来完成)。由于MPI_Finalize()是集体的，并且在大多数MPI内隐式中可能是阻塞的，所以会出现死锁。

一个解决方案是覆盖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可以作为一个模块运行，解决这个问题(在内部，如jcgiret所说，调用Abort() )：

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