使用Python的multiprocessing Pool imap时的KeyboardInterrupt

Question

根据这个问题中的一些建议 使用Python的多进程池处理键盘中断，我写了一个测试代码，叫做test.py。

from multiprocessing import Pool
from multiprocessing.pool import IMapIterator
import fileinput

def wrapper(func):
  def wrap(self, timeout=None):
  # Note: the timeout of 1 googol seconds introduces a rather subtle
  # bug for Python scripts intended to run many times the age of the universe.
  return func(self, timeout=timeout if timeout is not None else 1e100)
return wrap
IMapIterator.next = wrapper(IMapIterator.next)

def main():
    p = Pool(1)
    try:
        for _ in p.imap(worker, fileinput.input(), 1):
            pass
    except KeyboardInterrupt:
        p.terminate()

def worker(line):
    print line

if __name__ == '__main__':
    main()

当我在终端运行python test.py时：

python test.py

我仍然收到了键盘中断的错误信息，而且我必须按两次Ctrl+C才能停止这个脚本。我该如何处理键盘中断，让它不打印错误信息，并且只用一次Ctrl+C就能优雅地停止它呢？

^CProcess PoolWorker-1:
Traceback (most recent call last):
File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/multiprocessing/process.py", line 258, in _bootstrap
    self.run()
File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/multiprocessing/process.py", line 114, in run
    self._target(*self._args, **self._kwargs)
File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/multiprocessing/pool.py", line 102, in worker
    task = get()
File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/multiprocessing/queues.py", line 376, in get
    return recv()
KeyboardInterrupt
^CTraceback (most recent call last):
File "test.py", line 27, in <module>
    main()
File "test.py", line 21, in main
    p.terminate()
File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/multiprocessing/pool.py", line 452, in terminate
    self._worker_handler._state = TERMINATE
File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/multiprocessing/util.py", line 201, in __call__
    res = self._callback(*self._args, **self._kwargs)
File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/multiprocessing/pool.py", line 504, in _terminate_pool
    if threading.current_thread() is not task_handler:
File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/threading.py", line 958, in join
    self.__block.wait(delay)
File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/threading.py", line 358, in wait
    _sleep(delay)
KeyboardInterrupt

Answer 1

看看标准库里包含的信号模块。你可以在主进程中注册一个信号处理器。

from signal import *

def siginthndlr(sig, frame):
    '''do what you need here'''
    print "Keyboard interrupt catched"

signal(SIGINT, siginthndlr) #Register SIGINT handler function

这样的话，它会优雅地结束工作进程，然后再结束主进程。