Python多进程池中的键盘中断

Question

我该如何在使用Python的多进程池时处理键盘中断事件呢？下面是一个简单的例子：

from multiprocessing import Pool
from time import sleep
from sys import exit

def slowly_square(i):
    sleep(1)
    return i*i

def go():
    pool = Pool(8)
    try:
        results = pool.map(slowly_square, range(40))
    except KeyboardInterrupt:
        # **** THIS PART NEVER EXECUTES. ****
        pool.terminate()
        print "You cancelled the program!"
        sys.exit(1)
    print "\nFinally, here are the results: ", results

if __name__ == "__main__":
    go()

当我运行上面的代码时，按下^C会触发KeyboardInterrupt，但是此时程序就会卡住，我必须手动结束它。

我希望能够随时按下^C，让所有进程都能优雅地退出。

Answer 1

出于某些原因，只有从基础的 Exception 类继承的异常才能正常处理。作为一种变通方法，你可以把你的 KeyboardInterrupt 重新抛出为一个 Exception 实例：

from multiprocessing import Pool
import time

class KeyboardInterruptError(Exception): pass

def f(x):
    try:
        time.sleep(x)
        return x
    except KeyboardInterrupt:
        raise KeyboardInterruptError()

def main():
    p = Pool(processes=4)
    try:
        print 'starting the pool map'
        print p.map(f, range(10))
        p.close()
        print 'pool map complete'
    except KeyboardInterrupt:
        print 'got ^C while pool mapping, terminating the pool'
        p.terminate()
        print 'pool is terminated'
    except Exception, e:
        print 'got exception: %r, terminating the pool' % (e,)
        p.terminate()
        print 'pool is terminated'
    finally:
        print 'joining pool processes'
        p.join()
        print 'join complete'
    print 'the end'

if __name__ == '__main__':
    main()

通常你会得到以下输出：

staring the pool map
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
pool map complete
joining pool processes
join complete
the end

所以如果你按下 ^C，你会得到：

staring the pool map
got ^C while pool mapping, terminating the pool
pool is terminated
joining pool processes
join complete
the end

Answer 2

根据我最近发现的情况，最好的解决办法是让工作进程完全忽略 SIGINT 信号，并把所有的清理代码放在父进程里。这样做可以解决无论是空闲还是忙碌的工作进程的问题，而且在子进程里不需要处理错误代码。

import signal

...

def init_worker():
    signal.signal(signal.SIGINT, signal.SIG_IGN)

...

def main()
    pool = multiprocessing.Pool(size, init_worker)

    ...

    except KeyboardInterrupt:
        pool.terminate()
        pool.join()

详细的解释和完整的示例代码可以在 http://noswap.com/blog/python-multiprocessing-keyboardinterrupt/ 和 http://github.com/jreese/multiprocessing-keyboardinterrupt 找到。

Answer 3

这是一个Python的bug。当在使用threading.Condition.wait()等待某个条件时，KeyboardInterrupt（也就是你按下Ctrl+C时的中断信号）不会被发送。下面是重现这个问题的方法：

import threading
cond = threading.Condition(threading.Lock())
cond.acquire()
cond.wait(None)
print "done"

KeyboardInterrupt这个异常不会被处理，直到wait()返回，而它实际上是不会返回的，所以中断就不会发生。理论上，KeyboardInterrupt应该能够打断条件的等待。

需要注意的是，如果指定了超时时间，这种情况就不会发生；比如使用cond.wait(1)时，按下中断键会立即生效。所以，一个解决办法就是指定一个超时时间。要做到这一点，可以把

    results = pool.map(slowly_square, range(40))

替换成

    results = pool.map_async(slowly_square, range(40)).get(9999999)

或者类似的写法。