Making a tqdm progress bar for asyncio

Question

Am attempting a tqdm progress bar with asyncio tasks gathered.

Want the progress bar to be progressively updated upon completion of a task. Tried the code:

import asyncio
import tqdm
import random

async def factorial(name, number):
    f = 1
    for i in range(2, number+1):
        await asyncio.sleep(random.random())
        f *= i
    print(f"Task {name}: factorial {number} = {f}")

async def tq(flen):
    for _ in tqdm.tqdm(range(flen)):
        await asyncio.sleep(0.1)

async def main():
    # Schedule the three concurrently

    flist = [factorial("A", 2),
        factorial("B", 3),
        factorial("C", 4)]

    await asyncio.gather(*flist, tq(len(flist)))

asyncio.run(main())

...but this simply completes the tqdm bar and then processes factorials.

Is there a way to make the progress bar move upon completion of each asyncio task?

Answer 1

As of tqdm version 4.48.0, it is possible to use tqdm.asyncio.tqdm.as_completed()

import tqdm.asyncio
...
for f in tqdm.asyncio.tqdm.as_completed(flist):
    await f

Answer 2

Now, I'm not particularly familiar with asyncho, though I've used tqdm with some success for multiprocesses in python. The following change to your code seems to update the progress bar and print the result at the same time, which might be enough to get you started.

responses = [await f
                 for f in tqdm.tqdm(asyncio.as_completed(flist), total=len(flist))]

The above should replace await asyncio.gather(*flist, tq(len(flist))) in your main definition.

For more information, the above was inspired from asyncio aiohttp progress bar with tqdm

To only print the bar once and update it, I've done the following, which updates the description of the progress bar to include your message:

import asyncio
import tqdm


async def factorial(name, number):
    f = 1
    for i in range(2, number + 1):
        await asyncio.sleep(1)
        f *= i
    return f"Task {name}: factorial {number} = {f}"

async def tq(flen):
    for _ in tqdm.tqdm(range(flen)):
        await asyncio.sleep(0.1)


async def main():
    # Schedule the three concurrently

    flist = [factorial("A", 2),
             factorial("B", 3),
             factorial("C", 4)]

    pbar = tqdm.tqdm(total=len(flist))
    for f in asyncio.as_completed(flist):
        value = await f
        pbar.set_description(value)
        pbar.update()

if __name__ == '__main__':
    asyncio.run(main())

Answer 3

Made a couple of small changes to Dragos' code in pbar format and used tqdm.write() to get almost what I want, as follows:

import asyncio
import random

import tqdm


async def factorial(name, number):
    f = 1
    for i in range(2, number + 1):
        await asyncio.sleep(random.random())
        f *= i
    return f"Task {name}: factorial {number} = {f}"

async def tq(flen):
    for _ in tqdm.tqdm(range(flen)):
        await asyncio.sleep(0.1)


async def main():

    flist = [factorial("A", 2),
             factorial("B", 3),
             factorial("C", 4)]

    pbar = tqdm.tqdm(total=len(flist), position=0, ncols=90)
    for f in asyncio.as_completed(flist):
        value = await f
        pbar.set_description(desc=value, refresh=True)
        tqdm.tqdm.write(value)
        pbar.update()

if __name__ == '__main__':
    asyncio.run(main())

Answer 4

As Danferno said in his comment there is an easier solution to this now where you just substitute asyncio gather with the gather function provided by tqdm. Your code can be changed to this:

import asyncio
from tqdm.asyncio import tqdm
import random

async def factorial(name, number):
    f = 1
    for i in range(2, number+1):
        await asyncio.sleep(random.random())
        f *= i
    print(f"Task {name}: factorial {number} = {f}")


async def main():
    # Schedule the three concurrently

    flist = [factorial("A", 2),
        factorial("B", 3),
        factorial("C", 4)]

    await tqdm.gather(*flist)

asyncio.run(main())

Answer 5

Here is an async wrapper around TQDM to return an ordered result:

import asyncio
from typing import Any, Coroutine, Iterable, List, Tuple

from tqdm import tqdm


async def aprogress(tasks: Iterable[Coroutine], **pbar_kws: Any) -> List[Any]:
    """Runs async tasks with a progress bar and returns an ordered result."""

    if not tasks:
        return []

    async def tup(idx: int, task: Coroutine) -> Tuple[int, Any]:
        """Returns the index and result of a task."""
        return idx, await task

    _tasks = [tup(i, t) for i, t in enumerate(tasks)]
    pbar = tqdm(asyncio.as_completed(_tasks), total=len(_tasks), **pbar_kws)
    res = [await t for t in pbar]
    return [r[1] for r in sorted(res, key=lambda r: r[0])]


if __name__ == "__main__":

    import random

    async def test(idx: int) -> Tuple[int, int]:
        sleep = random.randint(0, 5)
        await asyncio.sleep(sleep)
        return idx, sleep

    _tasks = [test(i) for i in range(10)]
    _res = asyncio.run(aprogress(_tasks, desc="pbar test"))
    print(_res)

Full source