如何通过KeyboardInterrupt优雅终止multiprocessing.Pool.map函数？

我尝试用KeyboardInterrupt优雅终止multiprocessing.Pool.map()函数，相关代码如下：

import multiprocessing
import time


def func(num):
    try:
        while True:
            print(f'Process num {num}')
            time.sleep(30)
    except KeyboardInterrupt:
        print(f'Process {num} was interrupted')
    finally:
        print('Cleaning up the thread')


def main():
    pool = multiprocessing.Pool()

    try:
        pool.map(func, range(9))
    except KeyboardInterrupt:
        print('Main process was interrupted')
        pool.Close()

    finally:
        print('Cleaning up main')


if __name__ == '__main__':
    main()

我的设备有8个逻辑核心，当传入range(8)时，按Ctrl+C后能正常终止所有进程，输出符合预期；但传入range(9)时，第一次按Ctrl+C仅终止正在运行的8个进程，Pool.map仍会生成新的进程（进程8），需要再次按Ctrl+C才会触发异常。请问有没有办法仅按一次Ctrl+C就停止Pool.map()？

问题根源 #

当Pool.map处理的任务数超过进程池大小（你的设备是8）时，第一次Ctrl+C会终止当前正在运行的8个进程，但主进程此时还在等待新的进程启动来处理剩余任务（第9个），直到新进程启动后再次收到信号才会触发主进程的KeyboardInterrupt。

解决方法 #

方法1：用pool.terminate()替代pool.close()

pool.close()仅阻止新任务提交，但不会终止待启动的任务；pool.terminate()会立即终止所有进程，无论任务是否完成。修改主进程的异常处理逻辑：

def main():
    pool = multiprocessing.Pool()

    try:
        pool.map(func, range(9))
    except KeyboardInterrupt:
        print('Main process was interrupted')
        pool.terminate()  # 替换为terminate，立即终止所有进程
    finally:
        pool.join()  # 等待所有子进程退出后再执行清理
        print('Cleaning up main')

方法2：注册全局信号处理函数

在主进程中捕获SIGINT信号，一旦收到就立即终止进程池：

import multiprocessing
import time
import signal

def func(num):
    try:
        while True:
            print(f'Process num {num}')
            time.sleep(30)
    except KeyboardInterrupt:
        print(f'Process {num} was interrupted')
    finally:
        print('Cleaning up the thread')

def handle_interrupt(signum, frame):
    print('Received interrupt signal, terminating pool')
    global pool
    pool.terminate()
    raise KeyboardInterrupt  # 重新抛出异常，让主进程进入finally块

def main():
    global pool
    pool = multiprocessing.Pool()
    signal.signal(signal.SIGINT, handle_interrupt)  # 注册信号处理

    try:
        pool.map(func, range(9))
    except KeyboardInterrupt:
        print('Main process was interrupted')
    finally:
        pool.join()
        print('Cleaning up main')

if __name__ == '__main__':
    main()

注意事项 #

• pool.terminate()会强制终止子进程，可能导致子进程的finally块无法执行。如果需要子进程完成清理，可在信号处理中先给子进程发送中断信号，短暂等待后再终止进程池，但这种方式可靠性有限。
• pool.join()必须在pool.close()或pool.terminate()之后调用，确保主进程等待所有子进程完全退出。

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