Keyboard Interrupts with python's multiprocessing Pool

Viewed 89979

How can I handle KeyboardInterrupt events with python's multiprocessing Pools? Here is a simple example:

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

When running the code above, the KeyboardInterrupt gets raised when I press ^C, but the process simply hangs at that point and I have to kill it externally.

I want to be able to press ^C at any time and cause all of the processes to exit gracefully.

11 Answers

Many of these answers are old and/or they do not seem to work with later versions of Python (I am running 3.8.5) on Windows if you are executing a method such as Pool.map, which blocks until all the submitted tasks have completed. The following is my solution.

  1. Issue a call to signal.signal(signal.SIGINT, signal.SIG_IGN) in the main process to ignore Ctrl-C altogether.
  2. The processing pool will be initialized with a pool initializer that will initialize each processor thusly: Global variable ctrl_c_entered will be set to False and a a call to signal.signal(signal.SIGINT, signal.SIG_IGN) will be issued to initially ignore Ctrl-C. The return value from this call will be saved; this is the original, default handler that when re-established allows handing of KyboardInterrupt exceptions.
  3. A decorator, handle_ctrl_c, can be used to decorate multiprocessing functions and methods that should exit immediately on Ctrl-C being entered. This decorator will test to see if the global ctrl_c_entered flag is set and if so ,not even bother to run the function/method and instead will return a KeyboardInterrupt exception instance. Otherwise a try/catch handler for a KeyboardInterrupt will be established and the decorated function/method will be invoked. If Ctrl-C is entered, global ctrl_c_entered will be set to True and a KeyboardInterrupt exception instance will be returned. In any event, before returning the decorator will re-establish the SIG_IGN handler.

In essence all submitted tasks will be allowed to start but will immediately terminate with a return value of a KeyBoardInterrupt exception once a Ctrl-C has been entered. The main process can test the return values for the presence of such a return value to detect whether a Ctrl-C was entered.

from multiprocessing import Pool
import signal
from time import sleep
from functools import wraps

def handle_ctrl_c(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        global ctrl_c_entered
        if not ctrl_c_entered:
            signal.signal(signal.SIGINT, default_sigint_handler) # the default
            try:
                return func(*args, **kwargs)
            except KeyboardInterrupt:
                ctrl_c_entered = True
                return KeyboardInterrupt()
            finally:
                signal.signal(signal.SIGINT, pool_ctrl_c_handler)
        else:
            return KeyboardInterrupt()
    return wrapper

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

def pool_ctrl_c_handler(*args, **kwargs):
    global ctrl_c_entered
    ctrl_c_entered = True

def init_pool():
    # set global variable for each process in the pool:
    global ctrl_c_entered
    global default_sigint_handler
    ctrl_c_entered = False
    default_sigint_handler = signal.signal(signal.SIGINT, pool_ctrl_c_handler)

def main():
    signal.signal(signal.SIGINT, signal.SIG_IGN)
    pool = Pool(initializer=init_pool)
    results = pool.map(slowly_square, range(10))
    if any(map(lambda x: isinstance(x, KeyboardInterrupt), results)):
        print('Ctrl-C was entered.')
    print(results)
    pool.close()
    pool.join()

if __name__ == '__main__':
    main()

Prints:

Ctrl-C was entered.
[0, 1, 4, 9, 16, 25, 36, 49, KeyboardInterrupt(), KeyboardInterrupt()]

You can try using the apply_async method of a Pool object, like this:

import multiprocessing
import time
from datetime import datetime


def test_func(x):
    time.sleep(2)
    return x**2


def apply_multiprocessing(input_list, input_function):
    pool_size = 5
    pool = multiprocessing.Pool(processes=pool_size, maxtasksperchild=10)

    try:
        jobs = {}
        for value in input_list:
            jobs[value] = pool.apply_async(input_function, [value])

        results = {}
        for value, result in jobs.items():
            try:
                results[value] = result.get()
            except KeyboardInterrupt:
                print "Interrupted by user"
                pool.terminate()
                break
            except Exception as e:
                results[value] = e
        return results
    except Exception:
        raise
    finally:
        pool.close()
        pool.join()


if __name__ == "__main__":
    iterations = range(100)
    t0 = datetime.now()
    results1 = apply_multiprocessing(iterations, test_func)
    t1 = datetime.now()
    print results1
    print "Multi: {}".format(t1 - t0)

    t2 = datetime.now()
    results2 = {i: test_func(i) for i in iterations}
    t3 = datetime.now()
    print results2
    print "Non-multi: {}".format(t3 - t2)

Output:

100
Multiprocessing run time: 0:00:41.131000
100
Non-multiprocessing run time: 0:03:20.688000

An advantage of this method is that results processed before interruption will be returned in the results dictionary:

>>> apply_multiprocessing(range(100), test_func)
Interrupted by user
{0: 0, 1: 1, 2: 4, 3: 9, 4: 16, 5: 25}
Related