Python Multiprocess Pool. How to exit the script w

问题:

mp.set_start_method('spawn')
total_count = Counter(0)
pool = mp.Pool(initializer=init, initargs=(total_count,), processes=num_proc)    

pool.map(part_crack_helper, product(seed_str, repeat=4))
pool.close()
pool.join()

So I have a pool of worker process that does some work. It just needs to find one solution. Therefore, when one of the worker processes finds the solution, I want to stop everything.

One way I thought of was just calling sys.exit(). However, that doesn't seem like it's working properly since other processes are running.

One other way was to check for the return value of each process calls (the return value of part_crack_helper function) and call terminate on that process. However, I don't know how to do that when using that map function.

How should I achieve this?

回答1:

You can use callbacks from Pool.apply_async.

Something like this should do the job for you.

from multiprocessing import Pool


def part_crack_helper(args):
    solution = do_job(args)
    if solution:
        return True
    else:
        return False


class Worker():
    def __init__(self, workers, initializer, initargs):
        self.pool = Pool(processes=workers, initializer=initializer, initargs=initargs)

    def callback(self, result):
        if result:
            print "Solution found! Yay!"
            self.pool.terminate()

    def do_job(self):
        for args in product(seed_str, repeat=4):
            self.pool.apply_async(part_crack_helper, args=args, callback=self.callback)
        self.pool.close()
        self.pool.join()
        print "good bye"


w = Worker(num_proc, init, [total_count])
w.do_job()