How to pass multiprocessing.Pool instance to apply_async callback function?

Question

Here is my prime factorization program,i added a callback function in pool.apply_async(findK, args=(N,begin,end)) ,a message prompt out prime factorization is over when factorization is over,it works fine.

import math
import multiprocessing 

def findK(N,begin,end):
    for k in range(begin,end):
        if N% k == 0:
            print(N,"=" ,k ,"*", N/k)
            return True
    return False


def prompt(result):
    if result:
        print("prime factorization is over")


def mainFun(N,process_num):
    pool = multiprocessing.Pool(process_num)
    for i in range(process_num):
        if i ==0 :
            begin =2
        else:
            begin = int(math.sqrt(N)/process_num*i)+1
        end = int(math.sqrt(N)/process_num*(i+1))
        pool.apply_async(findK, args=(N,begin,end) , callback = prompt)    
    pool.close()
    pool.join()    

if __name__ == "__main__":
    N = 684568031001583853
    process_num = 16
    mainFun(N,process_num)

Now i want to change the callback function in apply_async,to change prompt into a shutdown function to kill all other process.

def prompt(result):
    if result:
        pool.terminate()

The pool instance is not defined in prompt scope or passed into prompt.
pool.terminate() can't work in prompt function.
How to pass multiprocessing.Pool instance to apply_async'callback function ?
(I have made it done in class format,just to add a class method and call self.pool.terminate can kill all other process, how to do the job in function format?)

if not set pool as global variable, can pool be passed into callback function?

Answer 1

Passing extra arguments to the callback function is not supported. Yet you have plenty of elegant ways to workaround that.

You can encapsulate your pool logic into an object:

class Executor:
    def __init__(self, process_num):
        self.pool = multiprocessing.Pool(process_num)

    def prompt(self, result):
        if result:
            print("prime factorization is over")
            self.pool.terminate()

    def schedule(self, function, args):
        self.pool.apply_async(function, args=args, callback=self.prompt)

    def wait(self):
        self.pool.close()
        self.pool.join() 


def main(N,process_num):
    executor = Executor(process_num)
    for i in range(process_num):
        ...
        executor.schedule(findK, (N,begin,end))   
    executor.wait()

Or you can use the concurrent.futures.Executor implementation which returns a Future object. You just append the pool to the Future object before setting the callback.

def prompt(future):
    if future.result():
        print("prime factorization is over")
        future.pool_executor.shutdown(wait=False)

def main(N,process_num):
    executor = concurrent.futures.ProcessPoolExecutor(max_workers=process_num)
    for i in range(process_num):
        ...
        future = executor.submit(findK, N,begin,end)
        future.pool_executor = executor
        future.add_done_callback(prompt)

Answer 2

You can simply define a local close function as a callback:

import math
import multiprocessing 


def findK(N, begin, end):
    for k in range(begin, end):
        if N % k == 0:
            print(N, "=", k, "*", N / k)
            return True
    return False


def mainFun(N, process_num):
    pool = multiprocessing.Pool(process_num)

    def close(result):
        if result:
            print("prime factorization is over")
            pool.terminate()
    for i in range(process_num):
        if i == 0:
            begin = 2
        else:
            begin = int(math.sqrt(N) / process_num * i) + 1
        end = int(math.sqrt(N) / process_num * (i + 1))
        pool.apply_async(findK, args=(N, begin, end), callback=close)
    pool.close()
    pool.join()


if __name__ == "__main__":
    N = 684568031001583853
    process_num = 16
    mainFun(N, process_num)

You can also use a partial function from functool , with

import functools

def close_pool(pool, results):
    if result:
        pool.terminate()

def mainFun(N, process_num):
    pool = multiprocessing.Pool(process_num)

    close = funtools.partial(close_pool, pool)
....

Answer 3

You need to have pool end up in prompt 's environment. One possibility is to move pool into the global scope (though this isn't really best-practice). This appears to work:

import math
import multiprocessing 

pool = None

def findK(N,begin,end):
    for k in range(begin,end):
        if N% k == 0:
            print(N,"=" ,k ,"*", N/k)
            return True
    return False


def prompt(result):
    if result:
        print("prime factorization is over")
        pool.terminate()


def mainFun(N,process_num):
    global pool
    pool = multiprocessing.Pool(process_num)
    for i in range(process_num):
        if i ==0 :
            begin =2
        else:
            begin = int(math.sqrt(N)/process_num*i)+1
        end = int(math.sqrt(N)/process_num*(i+1))
        pool.apply_async(findK, args=(N,begin,end) , callback = prompt)    
    pool.close()
    pool.join()    

if __name__ == "__main__":
    N = 684568031001583853
    process_num = 16
    mainFun(N,process_num)