Add hint of duration for each iteration in tqdm

I have a list of tasks that each take a different amount of time. Let's say, I have 3 tasks, with durations close to 1x, 5x, 10*x. My tqdm code is something like:

from tqdm import tqdm

def create_task(n):
    def fib(x):
        if x == 1 or x == 0:
            return 1
        return fib(x - 1) + fib(x - 2)
    return lambda: fib(n)

n = 1
tasks = [create_task(n), create_task(5*n), create_task(10*n)]
for task in tqdm(tasks):
    task.run()

The problem is that tqdm thinks each iteration takes the same amount of time. As the first takes approximately 1/10 of the time, the ETA is unreliable.

My question: is it possible to somehow add a hint to tqdm to inform how much each iteration takes compared to the first? Something like informing the duration weights of each iteration...

Thanks!

Solution

There are two standard usages for tqdm progress bars: iterable-based, and manual. Manually updating a progress bar allows you to specify progress bar weights.

Consider the following code:

def my_func(x):
    """ Sleep for x / 5 seconds. """
    duration = x / 5
    time.sleep(duration)

in_vals = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

# Ex. 1: unweighted, updated automatically (iterable-based)
for x in tqdm(in_vals, position=0):
    my_func(x)

# Ex. 2: unweighted, updated manually
# functionally identical to example 1
with tqdm(total=len(in_vals), position=0) as pbar:
    for x in in_vals:
        my_func(x)
        pbar.update(n=1)

# Ex. 3: weighted, updated manually
with tqdm(total=sum(in_vals), position=0) as pbar:
    for x in in_vals:
        my_func(x)
        pbar.update(n=x)

The first two examples give a progress bar with an inaccurate estimated duration that grows longer as the iterations progress. This is because tqdm doesn't know that later iterations will take longer to complete than earlier iterations.

To fix this problem and get an accurate ETA from the start, we specify the total argument in tqdm() as the sum of all weights (example 3). To update the progress bar after completing an iteration with weight w, call pbar.update(n=w).

For the OP's use case, a tqdm progress bar with duration weights would look something like this:

def fib(x):
    time.sleep(.01) # added this so that progress bar isn't completed instantly
    if x == 1 or x == 0:
        return 1
    return fib(x - 1) + fib(x - 2)

in_vals = [1, 5, 10]

with tqdm(total=sum(in_vals), unit='calc', position=0) as pbar:
    for num in in_vals:
        fib(num)
        pbar.update(num)