Created
September 26, 2014 20:58
-
-
Save benanne/02c1dbafe966d2736cf4 to your computer and use it in GitHub Desktop.
Running a (slow) generator in a separate process
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import multiprocessing as mp | |
def buffered_gen_mp(source_gen, buffer_size=2): | |
""" | |
Generator that runs a slow source generator in a separate process. | |
buffer_size: the maximal number of items to pre-generate (length of the buffer) | |
""" | |
if buffer_size < 2: | |
raise RuntimeError("Minimal buffer size is 2!") | |
buffer = mp.Queue(maxsize=buffer_size - 1) | |
# the effective buffer size is one less, because the generation process | |
# will generate one extra element and block until there is room in the buffer. | |
def _buffered_generation_process(source_gen, buffer): | |
for data in source_gen: | |
buffer.put(data, block=True) | |
buffer.put(None) # sentinel: signal the end of the iterator | |
buffer.close() # unfortunately this does not suffice as a signal: if buffer.get() | |
# was called and subsequently the buffer is closed, it will block forever. | |
process = mp.Process(target=_buffered_generation_process, args=(source_gen, buffer)) | |
process.start() | |
for data in iter(buffer.get, None): | |
yield data |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment