如何实现Python多进程程序的快捷键触发优雅终止？

当然可以实现！我来给你分享两种实用的方案，既能响应自定义快捷键，又能让你的多进程程序优雅地完成当前正在运行的任务后再退出，不会像Ctrl+C那样直接粗暴中断。

方案一：用第三方库监听快捷键（直观易用） #

我们可以用pynput库来监听键盘事件，捕获Q或Ctrl+Q按键，然后设置一个终止标志，让主进程在合适的时机关闭进程池，等待当前任务完成后退出。

首先安装依赖：

pip install pynput

完整代码示例：

import multiprocessing as mp
from pynput import keyboard
import time

# 全局终止标志，仅主进程修改和读取
should_terminate = False

def on_press(key):
    global should_terminate
    try:
        # 监听普通Q键（支持大小写，可根据需求调整）
        if key.char in ('q', 'Q'):
            print("\n收到终止信号，将在当前任务完成后退出...")
            should_terminate = True
            return False  # 停止监听
    except AttributeError:
        # 监听Ctrl+Q组合键（更准确的实现可结合on_release事件）
        if key == keyboard.Key.ctrl_l or key == keyboard.Key.ctrl_r:
            pass

# 替换成你的实际业务函数
def f(x, y):
    print(f"开始处理任务 ({x}, {y})")
    time.sleep(2)  # 模拟耗时操作
    result = x + y
    print(f"完成任务 ({x}, {y})，结果: {result}")
    return result

if __name__ == "__main__":
    # 启动键盘监听线程
    listener = keyboard.Listener(on_press=on_press)
    listener.start()

    # 你的输入列表
    L = [(1,2), (3,4), (5,6), (7,8), (9,10)]
    # 进程数
    N = 2

    with mp.Pool(processes=N) as p:
        # 用imap_unordered代替starmap，支持迭代获取结果并中途检查终止标志
        for result in p.imap_unordered(f, L):
            if should_terminate:
                print("正在等待剩余运行中任务完成...")
                p.close()  # 禁止提交新任务
                p.join()   # 等待所有正在运行的任务完成
                break

    print("程序已优雅退出")

代码说明： #

• 用pynput启动后台线程监听键盘，按下Q键时设置终止标志。
• 选择imap_unordered而非starmap：starmap是阻塞式的，会一次性提交所有任务并等待全部完成，无法中途检查终止信号；而imap_unordered会逐个返回任务结果，让我们有机会插入终止检查逻辑。
• p.close()禁止进程池接受新任务，p.join()等待当前运行的子进程完成，完美实现“当前迭代结束后终止”的需求。

方案二：用系统信号实现（无需第三方库） #

如果你的环境无法安装第三方库，可以用Python的signal模块监听自定义系统信号，比如SIGUSR1，通过终端发送信号触发终止。

完整代码示例：

import multiprocessing as mp
import signal
import time

should_terminate = False

def handle_terminate_signal(signum, frame):
    global should_terminate
    print("\n收到终止信号，将在当前任务完成后退出...")
    should_terminate = True

# 你的业务函数
def f(x, y):
    print(f"开始处理任务 ({x}, {y})")
    time.sleep(2)
    result = x + y
    print(f"完成任务 ({x}, {y})，结果: {result}")
    return result

if __name__ == "__main__":
    # 注册SIGUSR1信号的处理函数
    signal.signal(signal.SIGUSR1, handle_terminate_signal)
    print(f"当前主进程ID: {mp.current_process().pid}")
    print("要终止程序，请在另一个终端执行：kill -SIGUSR1 主进程ID")

    L = [(1,2), (3,4), (5,6), (7,8), (9,10)]
    N = 2

    with mp.Pool(processes=N) as p:
        for result in p.imap_unordered(f, L):
            if should_terminate:
                print("正在等待剩余运行中任务完成...")
                p.close()
                p.join()
                break

    print("程序已优雅退出")

使用方法： #

1. 运行程序后，记下终端输出的主进程ID。
2. 打开另一个终端，执行kill -SIGUSR1 <主进程ID>，程序就会触发终止逻辑，等待当前任务完成后退出。

关键注意事项 #

• 避免使用starmap：它是阻塞式的，一旦提交所有任务就会一直等到全部完成，无法中途检查终止信号，必须用imap或imap_unordered这类迭代式方法。
• 终止标志的安全性：这里的全局变量should_terminate只有主进程修改和读取，子进程不会操作它，因此是线程安全的。若需要子进程修改标志，要使用mp.Value这类多进程共享内存对象。
• Ctrl+C的替代方案：如果想用Ctrl+C实现优雅终止，也可以捕获SIGINT信号并执行类似的进程池关闭逻辑，但需要注意多进程下信号传递的问题，确保子进程不会被直接中断。

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