Сбой быстро с MPI4PY

Question

Мне бы хотелось, чтобы при запуске сценария MPI с mpi4py было следующее поведение: когда какой-либо процесс выдает исключение, 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/ я прочитал:

mpi4py инициализирует / завершает MPI для вас. Инициализация происходит во время импорта, и финализация, когда завершается процесс Python (для этого я использую вызов C-API Py_AtExit ()). Поскольку 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 может быть запущен как модуль, исправляющий эту проблему (внутренне, вызывая Abort(), как говорит jcgiret):

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