Skip to content

Commit

Permalink
add pmap. Add rate
Browse files Browse the repository at this point in the history
pmap maps using the input function using multiple processes.
  • Loading branch information
esheldon committed Oct 6, 2024
1 parent e2eca31 commit 66db3c7
Showing 1 changed file with 59 additions and 8 deletions.
67 changes: 59 additions & 8 deletions esutil/pbar.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,13 @@
https://github.com/noamraph/tqdm
"""
__all__ = ['PBar', 'prange']
__all__ = ['PBar', 'pbar', 'prange']

import sys
import time


def PBar(iterable, desc='', total=None, leave=True, file=sys.stderr,
def pbar(iterable, desc='', total=None, leave=True, file=sys.stderr,
mininterval=0.5, miniters=1, n_bars=20):
"""
Get an iterable object, and return an iterator which acts exactly like the
Expand Down Expand Up @@ -90,19 +90,60 @@ def PBar(iterable, desc='', total=None, leave=True, file=sys.stderr,
file.write('\n')


PBar = pbar


def prange(*args, **kwargs):
"""
A shortcut for writing PBar(range(...))
A shortcut for writing pbar(range(...))
Parameters
----------
Same args as for range. Extra keywords are sent to
Pbar
e.g.
import time
from pbar import prange
for i in prange(20):
print(i)
time.sleep(0.1)
"""
return PBar(range(*args), **kwargs)
return pbar(range(*args), **kwargs)


def pmap(fn, iterable, chunksize=1, nproc=1, **kw):
"""
Execute the function on the inputs using multiple processes, while showing
a progress bar. The result is equivalent to doing
list(map(fn, iterable))
Parameters
----------
fn: function
The function to execute
iterable: iterable data
The data over which to iterate
chunksize: int, optional
Default 1. It is often must faster to send large
chunks of data rather than 1.
nproc: int, optional
Number of processes to use, default 1
**kw:
Additional keyword arguments for the progress bar.
See pbar for details
Returns
-------
An list of data, the equivalent of
list(map(fn, iterable))
"""
from concurrent.futures import ProcessPoolExecutor

with ProcessPoolExecutor(max_workers=nproc) as ex:
res = list(pbar(ex.map(fn, iterable, chunksize=chunksize), **kw))

return res


def format_interval(t):
Expand Down Expand Up @@ -131,10 +172,20 @@ def format_meter(n, total, elapsed, n_bars=20):

percentage = '%3d%%' % (frac * 100)

if elapsed > 0:
it_per_second = n / elapsed # iterations per second
if it_per_second > 1:
rate_str = f'{it_per_second:.3g} it/s'
else:
second_per_it = elapsed / n
rate_str = f'{second_per_it:.3g} s/it'
else:
rate_str = '---'

left_str = format_interval(elapsed / n * (total-n)) if n else '?'

return '|%s| %d/%d %s [elapsed: %s left: %s]' % (
bar, n, total, percentage, elapsed_str, left_str)
return '|%s| %d/%d %s [elapsed: %s left: %s %s]' % (
bar, n, total, percentage, elapsed_str, left_str, rate_str)

else:
return '%d [elapsed: %s]' % (n, elapsed_str)
Expand Down

0 comments on commit 66db3c7

Please sign in to comment.