How to break parallelization in joblib after finding wanted value?

Viewed 207

Suppose I have a function like this:

def f(angle):
    if angle % 5==0:
        return angle 

    return None 

Then I use joblib.Parallel and joblib.delayed in the following way:

from joblib import Parallel, delayed
results = Parallel(n_jobs=4)(delayed(f)(angle) for angle in range(360))

This gives me all values that I want and they will be different from "None" only when angle is divisible by 5. Then i do the following

for x in results:
   if x is not None:
       return x 

With this I return first result that is not None. But, what if we need to do some hard calculations that takes too much time? Then we would want this to be finished during parallelization. I don't want to brute force all parallelization in order to find value that I want. I would like if I could obtain this value immidiately during parallelization, after which parallelization stops and gives me the value i want. What is the best way to accomplish that using joblib library?

1 Answers

Two years later, but...

From what I can see, joblib doesn't have any method to stop mid-process. So I believe you will simply have to resort to using the multiprocessing library instead:

from time import sleep
from multiprocessing import Pool


def f(angle):
    sleep(5)  # very expensive computations here

    if angle % 5 == 0:
        return angle

    return None


def main():
    p = Pool(4)

    for x in p.imap_unordered(f, range(360)):
        if x is not None:
            p.terminate()
            return x

Unfortunetely, imap() and imap_unordered() can only map a single argument into the function. So if f() required multiple arguments you'd need to create an iterable object that yields multiple results:

from time import sleep
from multiprocessing import Pool


def f(args):
    angle = args[0]
    angle2 = args[1]

    sleep(5)  # very expensive computations here

    if angle * angle2 % 5 == 0:
        return angle

    return None


class iterable(object):
    def __init__(self, x, y):
        self.i = 0
        self.x = range(x)
        self.y = range(y)
        self.sz1 = len(self.x)
        self.sz2 = len(self.y)

    def __iter__(self):
        return self

    def __next__(self):
        self.i += 1

        if self.i >= self.sz1 or self.i >= self.sz2:
            raise StopIteration
        else:
            return self.x[self.i - 1], self.y[self.i - 1]


def main():
    p = Pool(4)

    for x in p.imap_unordered(f, iterable(360, 360)):
        if x is not None:
            p.terminate()
            return x
Related