Starmap combined with tqdm?
It's not possible with starmap()
, but it's possible with a patch adding Pool.istarmap()
. It's based on the code for imap()
. All you have to do, is create the istarmap.py
-file and import the module to apply the patch before you make your regular multiprocessing-imports.
Python <3.8
# istarmap.py for Python <3.8
import multiprocessing.pool as mpp
def istarmap(self, func, iterable, chunksize=1):
"""starmap-version of imap
"""
if self._state != mpp.RUN:
raise ValueError("Pool not running")
if chunksize < 1:
raise ValueError(
"Chunksize must be 1+, not {0:n}".format(
chunksize))
task_batches = mpp.Pool._get_tasks(func, iterable, chunksize)
result = mpp.IMapIterator(self._cache)
self._taskqueue.put(
(
self._guarded_task_generation(result._job,
mpp.starmapstar,
task_batches),
result._set_length
))
return (item for chunk in result for item in chunk)
mpp.Pool.istarmap = istarmap
Python 3.8+
# istarmap.py for Python 3.8+
import multiprocessing.pool as mpp
def istarmap(self, func, iterable, chunksize=1):
"""starmap-version of imap
"""
self._check_running()
if chunksize < 1:
raise ValueError(
"Chunksize must be 1+, not {0:n}".format(
chunksize))
task_batches = mpp.Pool._get_tasks(func, iterable, chunksize)
result = mpp.IMapIterator(self)
self._taskqueue.put(
(
self._guarded_task_generation(result._job,
mpp.starmapstar,
task_batches),
result._set_length
))
return (item for chunk in result for item in chunk)
mpp.Pool.istarmap = istarmap
Then in your script:
import istarmap # import to apply patch
from multiprocessing import Pool
import tqdm
def foo(a, b):
for _ in range(int(50e6)):
pass
return a, b
if __name__ == '__main__':
with Pool(4) as pool:
iterable = [(i, 'x') for i in range(10)]
for _ in tqdm.tqdm(pool.istarmap(foo, iterable),
total=len(iterable)):
pass
The simplest way would probably be to apply tqdm() around the inputs, rather than the mapping function. For example:
inputs = zip(param1, param2, param3)
with mp.Pool(8) as pool:
results = pool.starmap(my_function, tqdm.tqdm(inputs, total=len(param1)))