parmap: simplify generator termination

parmap.py

@@ -14,22 +14,19 @@ def chunkme(X, chunksize):
   if len(chunk):
     yield chunk
 
-def producer_fun(X, q_in, q_out, cont, nprocs, chunksize):
-  sent = 0
+def producer_fun(X, q_in, cont, nprocs, chunksize):
   for i, x in enumerate(chunkme(X, chunksize)):
     cont.acquire()
     q_in.put((i, x))
-    sent += 1
 
   for _ in range(nprocs):
     q_in.put((None, None))
 
-  q_out.put((None, sent))
-
 def worker_fun(f, q_in, q_out):
   while True:
     i, chunk = q_in.get()
     if i is None:
+      q_out.put((None, None))
       break
     q_out.put((i, [ f(x) for x in chunk ]))
 
@@ -53,7 +50,7 @@ def parmap(f, X, nprocs = None, chunksize = 1, chunks_in_flight = None,
 
     proc.append(multiprocessing.Process(
       target = producer_fun,
-      args = (X, q_in, q_out, cont, nprocs, chunksize)
+      args = (X, q_in, cont, nprocs, chunksize)
     ))
 
     for p in proc:
@@ -63,12 +60,12 @@ def parmap(f, X, nprocs = None, chunksize = 1, chunks_in_flight = None,
     def get_chunk():
       ret = {}
       chunk_index = 0
-      jobs = None
+      running_workers = nprocs
 
-      while jobs is None or chunk_index < jobs:
+      while running_workers > 0:
         i, val = q_out.get()
         if i is None:
-          jobs = val
+          running_workers -= 1
           continue
 
         if not inOrder: