手动终止 OpenMPI 运行时在上下文管理器中捕获键盘中断

Question

我正在使用

mpi4py

并行运行代码。我注意到，如果我运行代码并执行键盘中断，如果我将代码作为

__exit__

运行，我的上下文管理器

python file.py

将会运行，但当作为

mpirun -np 1 file.py

运行时则不会运行（这只是一个进程，但是它产生的结果与我使用更多进程运行代码相同）。

如何让终止的 MPI 运行使代码像正常的 python 进程一样退出并进入上下文管理器的

__exit__

进程？

最小可重现示例：

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))

如果我取消注释中断，无论代码是通过

python file.py

还是

mpirun -np 1 file.py

运行，上下文管理器

__exit__

都会运行并显示打印输出。

如果我运行如上所示的代码（中断被注释）并在运行过程中点击

ctrl+C

，那么：

• 如果我运行代码为

  python file.py

  ，则输入

  __exit__

  并显示打印输出
• 如果我按照

  mpirun -np 1 python file.py

  运行代码，则代码会直接终止并且永远不会进入

  __exit__

  （即，没有打印输出）。

版本：

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

Answer 1

我认为您可能需要将整个

with

块放入

try/except

块中。

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

我认为这是因为使用Python的某些实现，

KeyboardInterrupt

（SIGINT）只是终止程序，而不退出

with

块并调用

__exit__()

函数。 对于

try/except

块，它会调用

__exit__()

，因为它正在捕获

KeyboardInterrupt

并退出调用

with

函数的

__exit__()

块。然而，它可能不会打印回溯......