使用MPI4PY快速失败

Andre 发表于 Dev

其他

我希望使用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()

如何获得所需的行为（如果任何进程失败，则迅速失败）？

谢谢！

吉吉雷特

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

mpi4py为您初始化/完成MPI。初始化在导入时进行，在Python进程即将完成时进行初始化（我正在使用Py_AtExit（）C-API调用来完成此操作）。由于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__

本文收集自互联网，转载请注明来源。

如有侵权，请联系 [email protected] 删除。