Согласованное утверждение для параллельных прогонов в python?

Question

Мне кажется, что в коде python, который работает параллельно, утверждение, которое не удалось, по крайней мере, одному процессору, должно прервать все процессоры, так что:

1) сообщение об ошибке ясно видно (с трассировка стека)

2) оставшиеся процессоры не ждут навсегда.

Однако это не то, что делает стандартное утверждение.

Этот вопрос уже задан в python script running with mpirun not stopping if assert on processor 0 fails , но я не удовлетворен ответом. Там предлагается использовать функцию comm.Abort(), но это только соответствует пункту 2) выше.

Итак, мне было интересно: существует ли стандартная функция «assert» для параллельных кодов (например, с mpi4py), или я должен написать свои собственные утверждения для этой цели?

Спасибо!

Edit - вот моя попытка (в классе, но может быть снаружи), которые, безусловно, могут быть улучшены:

import mpi4py.MPI as mpi 
import traceback 

class My_code(): 

    def __init__(self, some_parameter=None): 

     self.current_com = mpi.COMM_WORLD 
     self.rank = self.current_com.rank 
     self.nb_procs = self.current_com.size 

     self.my_assert(some_parameter is not None) 
     self.parameter = some_parameter 
     print "Ok, parameter set to " + repr(self.parameter) 

    # some class functions here... 

    def my_assert(self, assertion): 
     """ 
     this is a try for an assert function that kills 
     every process in a parallel run 
     """ 
     if not assertion: 
      print 'Traceback (most recent call last):' 
      for line in traceback.format_stack()[:-1]: 
       print(line.strip()) 
      print 'AssertionError' 
      if self.nb_procs == 1: 
       exit() 
      else: 
       self.current_com.Abort() 

Answer 1

Я думаю, что следующий фрагмент кода отвечает на вопрос. Это вытекает из обсуждения, указываемой Dan D.

import mpi4py.MPI as mpi 
import sys 


# put this somewhere but before calling the asserts 
sys_excepthook = sys.excepthook 
def mpi_excepthook(type, value, traceback): 
    sys_excepthook(type, value, traceback) 
    if mpi.COMM_WORLD.size > 1: 
     mpi.COMM_WORLD.Abort(1) 
sys.excepthook = mpi_excepthook 

# example: 
if mpi.COMM_WORLD.rank == 0: 
    # with sys.excepthook redefined as above this will kill every processor 
    # otherwise this would only kill processor 0 
    assert 1==0   

# assume here we have a lot of print messages 
for i in range(50): 
    print "rank = ", mpi.COMM_WORLD.rank 

# with std asserts the code would be stuck here 
# and the error message from the failed assert above would hardly be visible 
mpi.COMM_WORLD.Barrier()