Simple way to parallelize embarrassingly parallelizable generator

Viewed 182

I have a generator (or, a list of generators). Let's call them gens

Each generator in gens is a complicated function that returns the next value of a complicated procedure. Fortunately, they are all independent of one another.

I want to call gen.__next__() for each element gen in gens, and return the resulting values in a list. However, multiprocessing is unhappy with pickling generators.

Is there a fast, simple way to do this in Python? I would like it such that gens of length m is mapped to n cores locally on my machine, where n could be larger or smaller than m. Each generator should run on a separate core.

If this is possible, can someone provide a minimal example?

4 Answers

You can't pickle generators. Read more about it here.

There is a blog post which explains it in much more detail. Referring a quote from it:

Let’s ignore that problem for a moment and look what we would need to do to pickle a generator. Since a generator is essentially a souped-up function, we would need to save its bytecode, which is not guarantee to be backward-compatible between Python’s versions, and its frame, which holds the state of the generator such as local variables, closures and the instruction pointer. And this latter is rather cumbersome to accomplish, since it basically requires to make the whole interpreter picklable. So, any support for pickling generators would require a large number of changes to CPython’s core.

Now if an object unsupported by pickle (e.g., a file handle, a socket, a database connection, etc) occurs in the local variables of a generator, then that generator could not be pickled automatically, regardless of any pickle support for generators we might implement. So in that case, you would still need to provide custom getstate and setstate methods. This problem renders any pickling support for generators rather limited.

He also suggests a solution, to use simple iterators.

the best solution to this problem to the rewrite the generators as simple iterators (i.e., one with a __next__ method). Iterators are easy and efficient space-wise to pickle because their state is explicit. You would still need to handle objects representing some external state explicitly however; you cannot get around this.

Another offered solution (which I haven't tried) suggests this

  1. Con­vert the gen­er­a­tor to a class in which the gen­er­a­tor code is the __iter__ method

  2. Add __getstate__ and __setstate__ meth­ods to the class, to han­dling pick­ling. Remem­ber that you can’t pickle file objects. So __setstate__ will have to re-open files, as necessary.

If your subtasks are truly parallel (do not rely on any shared state), you can do this with multiprocesing.Pool().

Take a look at https://docs.python.org/3/library/multiprocessing.html#multiprocessing.pool.Pool

This requires you to make the arguments of pool.map() serializable. You can't pass in a generator to your worker, but you can achieve something similar by defining your generator inside the target function, and pass in initialization arguments to the multiprocessing library:

import multiprocessing as mp
import time

def worker(value):
    # The generator is defined inside the multiprocessed function
    def gen():
        for k in range(value):
            time.sleep(1) # Simulate long running task
            yield k

    # Execute the generator
    for x in gen():
        print(x)
        # Do something with x?
        pass


pool = mp.Pool()
pool.map(worker, [2, 5, 2])

pool.join() # Wait for all the work to be finished.
pool.close() # Clean up system resources

The output will be:

0
0
0
1
1
1
2
3
4

Note that this solution only really works if you build your generators, then use them only once, as their final state is lost at the end of the worker function.

Keep in mind that anytime you want to use multiprocessing, you have to use for serializable objects due to the limitations of inter-process communication; this can often prove limiting.

If your process is not CPU bound but instead I/O bound (disk access, network access, etc), you'll have a much easier time using threads.

You don't need to pickle the generator, just send an index of the generator to the processing pool.

M = len(gens)
N = multiprocessing.cpu_count()

def proc(gen_idx):
    return [r for r in gens[gen_idx]()]

if __name__ == "__main__":
    with multiprocessing.Pool(N) as p:
        for r in p.imap_unordered(proc, range(M)):
            print(r)

Note that I don't call/initialize the generator until within the processing function.

Using imap_unordered will allow you to process the results as each generator completes.

It's quite easy to implement, just dont block the threads sincronusly, just constantly loop thru the states and join them on complition. This template shuld be good enough to give an idea, self.done alwais needs to be set last on thread complition and las on thread reuse.

import threading as th
import random
import time

class Gen_thread(th.Thread):

    def is_done(self):
        return self.done

    def get_result(self):
        return self.work_result

    def __init__(self, *args, **kwargs):
        self.g_id = kwargs['id']
        self.kwargs = kwargs
        self.args = args
        self.work_result = None
        self.done = False
        th.Thread.__init__(self)

    def run(self):
        # time.sleep(*self.args) to pass variables
        time.sleep(random.randint(1, 4))
        self.work_result = 'Thread {0} done'.format(self.g_id + 1)
        self.done = True

class Gens(object):

    def __init__(self, n):
        self.n_needed = 0
        self.n_done = 0
        self.n_loop = n

        self.workers_tmp = None
        self.workers = []

    def __iter__(self):
        return self

    def __next__(self):
        if self.n_needed == 0:
            for w in range(self.n_loop):
                self.workers.append(Gen_thread(id=w))
                self.workers[w].start()
                self.n_needed += 1

        while self.n_done != self.n_needed:
            for w in range(self.n_loop):
                if self.workers[w].is_done():
                    self.workers[w].join()
                    self.workers_tmp = self.workers[w].get_result()
                    self.workers.pop(w)
                    self.n_loop -= 1
                    self.n_done += 1
                    return self.workers_tmp

        raise StopIteration()


if __name__ == '__main__':
    for gen in Gens(4):
        print(gen)
Related