Tratando interrupções de teclado no pool de multiprocessamento do Python

No módulo de multiprocessamento do Python, a classe Pool fornece uma maneira conveniente de distribuir tarefas em vários processos. No entanto, lidar com eventos KeyboardInterrupt em Pools pode ser desafiador, conforme demonstrado pelo trecho de código:

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

Ao executar este código, pressionar Ctrl C não acionará o processo de limpeza, deixando os subprocessos em execução indefinidamente. Para resolver esse problema, considere a seguinte solução alternativa:

O comportamento observado no código é consequência de um bug do Python. KeyboardInterrupt não é enviado ao aguardar uma condição em threading.Condition.wait(). Como Pool.map() usa uma condição wait internamente, a interrupção nunca é recebida.

Uma solução é usar Pool.map_async(), que permite especificar um timeout. Ao definir um tempo limite suficientemente longo (por exemplo, 9999999), podemos garantir que a interrupção será acionada dentro de um tempo razoável.

Portanto, substitua:

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

com:

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

Esta solução alternativa fornece uma maneira de lidar normalmente com eventos KeyboardInterrupt em pools de multiprocessamento, permitindo o encerramento de todos os subprocessos quando o usuário cancela o programa.