Compute the outputs of a generator in parallel

Question

I am processing an unknown "length" of generator object. I have to keep things "lazy" because of memory management. The processing is compute heavy, so writing it multiproc style is the solution (or at least it seems for me).

I have solved this problem of multiproc on generator object with a combination of monkey patch and a bounded Queue.

What really itches me is the monkey patch... Do you think this is fine to apply imap() on a generator object? How would you do this?

I would like to underline that the focus is to compute the outputs of a generator in parallel. From the perspective of this "minimal example" :

process_line, process_line_init, process_first_n_line

are the functions I am most interested about your opinion.

import multiprocessing as mp
import psutil
import queue 
from typing import Any, Dict, Iterable, Set

def yield_n_line(n: int)-> Iterable[Dict[str, str]]:
    for i in range(n):
        yield {'body': "Never try to 'put' without a timeout sec declared"}

def get_unique_words(x: Dict[str, str])-> Set[str]:
    return set(x['body'].split())

def process_line(x:Dict[str, str])-> Set[str]:
    try:
        process_line.q.put(x, block=True, timeout=2)
    except queue.Full:
        pass
    return get_unique_words(x)

def process_line_init(q: mp.Queue)-> None:
    process_line.q = q

def process_first_n_line(number_of_lines: int)-> Any:
    n_line = yield_n_line(number_of_lines)
    
    if psutil.cpu_count(logical=False) > 4:
        cpu_count = psutil.cpu_count(logical=False)-2
    else:
        cpu_count = psutil.cpu_count(logical=False)
    q = mp.Queue(maxsize=8000)
    p = mp.Pool(cpu_count, process_line_init, [q])
    results = p.imap(process_line, n_line)
    for _ in range(number_of_lines):
        try:
            q.get(timeout=2)
        except queue.Empty:
            q.close()
            q.join_thread()
        yield results.next()
    p.close()
    p.terminate()
    p.join()
    pass

def yield_uniqueword_chunks(
    n_line: int = 10_000_000,
    chunksize: int = 1_787_000)-> Iterable[Set[str]]:
    chunk = set()
    for result in process_first_n_line(n_line):
        chunk.update(result)
        if len(chunk) > chunksize:
            yield chunk
            chunk = set()
    yield chunk

def main()-> None:
    for chunk in yield_uniqueword_chunks(
        n_line=1000, #Number of total comments to process
        chunksize=200 #number of unique words in a chunk (around 32MB)
        ):
        print(chunk)
        #export(chunk)
    
if __name__ == "__main__":
    main()

Answer 1

Here's a condensed illustration of how to achieve your stated purpose, namely to compute the outputs of a generator in parallel. I offer it because I could not understand the purpose of most of the complexity in your current code. I suspect there are issues that you have not explained to us or that I failed to infer (if so, this answer might be off the mark). In any case, sometimes it helps to look at a problem in a simplified form to see whether your actual use case might work just fine within those parameters.

import multiprocessing as mp

def main():
    for chunk in process_data(1000, 50):
        print()
        print(len(chunk), chunk)
    
def process_data(n, chunksize):
    # Set up the Pool using a context manager.
    # This relieves you of the hassle of join/close.
    with mp.Pool() as p:
        s = set()
        # Just iterate over the results as they come in.
        # No need to check for empty queues, etc.
        for res in p.imap(worker, source_gen(n)):
            s.update(res)
            if len(s) >= chunksize:
                yield s
                s = set()

def source_gen(n):
    # A data source that will yield N values.
    # To get the data in this example:
    # curl 'https://www.gutenberg.org/files/2600/2600-0.txt' -o war-peace.txt
    with open('war-peace.txt') as fh:
        for i, line in enumerate(fh):
            yield line
            if i >= n:
                break

def worker(line):
    # A single-argument worker function.
    # If you need multiple args, bundle them in tuple/dict/etc.
    return [word.lower() for word in line.split()]

if __name__ == "__main__":
    main()

This example illustrates a garden-variety type of parallelization:

• The data source is running in a single process (the parent).

• Downstream computation is running in multiple child processes.

• And final aggregation/reporting is running in a single process (also the parent).

Behind the scenes, Pool.imap() is using pickling and queues to shuttle data between processes. If your real use case requires parallelizing data generation and/or reporting, a different approach is needed.