OpenMPI手动终止时如何在上下文管理器捕获KeyboardInterrupt

当使用mpi4py运行并行代码时，通过python file.py执行时，键盘中断（Ctrl+C）会正常触发上下文管理器的__exit__方法；但用mpirun -np 1 file.py（单/多进程均如此）执行时，进程会直接终止，不会进入__exit__流程。即使手动添加SIGTERM信号处理也无效。

最小复现示例 #

from mpi4py import MPI

def f(i):
    # raise KeyboardInterrupt()
    return i**0.5
    
class ContextManager():
    def __init__(self,):
        return
    
    def __enter__(self,):
        return self
        
    def __exit__(self, exc_type, exc_value, traceback):
        comm = MPI.COMM_WORLD
        rank = comm.Get_rank()
        size = comm.Get_size()
        
        print("Exiting.")
        
        if rank == 0:
            print(rank, size)
            print(exc_type)
            print(exc_value)
            print(traceback)
        

if __name__ == "__main__":
    with ContextManager():
        for i in range(1_000_000):
            print(f(i))

测试现象 #

• 手动抛出KeyboardInterrupt时，无论用哪种方式运行，都会触发__exit__
• 按下Ctrl+C时：
  • python file.py：正常进入__exit__并输出内容
  • mpirun -np 1 file.py：进程直接终止，无任何输出

版本信息：

python: 3.10.14
mpirun: 4.1.4
mpi4py: 3.1.6
Ubuntu: 22.04.4 LTS

问题原因 #

MPI启动的进程会修改默认信号处理逻辑：当按下Ctrl+C时，mpirun会向所有进程发送SIGINT信号，而mpi4py初始化时会覆盖Python默认的SIGINT处理函数，导致信号无法被Python的上下文管理器捕获。之前尝试处理SIGTERM是无效的，因为进程实际收到的是SIGINT。

解决方案 #

重新接管SIGINT信号的处理，抛出Python能识别的KeyboardInterrupt异常，同时确保MPI进程同步退出。

修改后的代码 #

from mpi4py import MPI
import signal

def sigint_handler(signum, frame):
    # 抛出KeyboardInterrupt，让上下文管理器捕获异常
    raise KeyboardInterrupt("MPI process interrupted by Ctrl+C")

# 替换默认的SIGINT信号处理
signal.signal(signal.SIGINT, sigint_handler)

def f(i):
    return i**0.5

class ContextManager():
    def __init__(self):
        self.comm = MPI.COMM_WORLD
        self.rank = self.comm.Get_rank()
        self.size = self.comm.Get_size()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_value, traceback):
        print(f"Rank {self.rank}: Exiting.")
        if self.rank == 0:
            print(f"Rank 0 | Size: {self.size}")
            print(f"Exception Type: {exc_type}")
            print(f"Exception Value: {exc_value}")
        # 调用MPI Abort确保所有进程同步退出，避免挂起
        self.comm.Abort(0 if exc_type is None else 1)

if __name__ == "__main__":
    try:
        with ContextManager():
            for i in range(1_000_000):
                print(f"Rank {MPI.COMM_WORLD.Get_rank()}: {f(i)}")
    except KeyboardInterrupt:
        # 捕获中断，确保上下文管理器的__exit__正常执行
        pass

关键说明 #

1. 处理SIGINT信号：Ctrl+C发送的是SIGINT，而非SIGTERM，因此需要替换SIGINT的处理函数
2. 抛出KeyboardInterrupt：让Python的上下文管理器能捕获异常，触发__exit__流程
3. 调用comm.Abort()：确保所有MPI进程同步退出，避免部分进程因等待通信而挂起
4. 外层try-except：捕获抛出的中断异常，保证上下文管理器的__exit__方法完整执行

内容的提问来源于stack exchange，提问作者jared