Q:

Python thread/process pool on infinite iterator?

I have an iterator function that yields an infinite stream of integers:
def all_ints(start=0):
  yield start
  yield all_ints(start+1)

I want to have a pool of threads or processes do calculations on these up to $POOLSIZE at a time. Each process will possibly save a result to some shared data structure so I do not need the return value from the process/thread function. It seems to me this use of the python3 Pool would achieve this:
# dummy example functions
def check_prime(n):
  return n % 2 == 0

def store_prime(p):
    ''' synchronize, write to some shared structure'''
    pass

p = Pool()

for n in all_ints():
    p.apply_async(check_prime, (n,), callback=store_prime)

But when I run this I get a python process that just continually uses more memory (and not from the iterator, that can run for days). I would expect this behavior if I was storing the results of all the apply_async calls, but I am not.
What am I doing wrong here? Or is there another API from the thread pool I should be using?

A:

I think you are looking for Pool.imap_unordered, which uses the pooled processes to apply a function to the elements yielded by an iterator. Its parameter chunksize allows you to specify how many items from the iterator are passed to the pool in each step.
Also I would avoid using any shared memory structures for IPC. Just let the "expensive" function sent to the pool return the information you need, and process it in the main process.
Here's an example (where I abort after 200'000 results; if you remove that part, you'll see the processes happily work in a fixed amount of RAM "forever"):
from multiprocessing import Pool
from math import sqrt
import itertools
import time

def check_prime(n): 
    if n == 2: return (n, True)
    if n % 2 == 0 or n < 2: return (n, False)
    for i in range(3, int(sqrt(n))+1, 2):
        if n % i == 0: return (n, False)
    return (n, True)    

def main():
    L = 200000   # limit for performance timing 
    p = Pool()
    n_primes = 0
    before = time.time()
    for (n, is_prime) in p.imap_unordered(check_prime, itertools.count(1), 1000):
        if is_prime:
            n_primes += 1
            if n_primes >= L: 
                break
    print("Computed %d primes in %.1fms" % (n_primes, (time.time()-before)*1000.0))
if __name__ == "__main__":
    main()

Output on my Intel Core i5 (2 Core, 4 Threads):
Computed 200000 primes in 15167.9ms

Output if I change it to Pool(1), so using just 1 subprocess:
Computed 200000 primes in 37909.2ms

HTH!

