Python中如何將Tqdm與Asyncio結(jié)合使用呢

更新時間：2023年05月10日 11:27:30 作者：冷凍工廠

這篇文章主要和大家詳細(xì)介紹了在Python中如何將Tqdm與Asyncio結(jié)合使用呢，文中的示例代碼講解詳細(xì)，感興趣的小伙伴可以跟隨小編一起學(xué)習(xí)一下

簡介

困擾

在 Python 中使用并發(fā)編程來提高效率對于數(shù)據(jù)科學(xué)家來說并不罕見。在后臺觀察各種子進(jìn)程或并發(fā)線程以保持我的計算或 IO 綁定任務(wù)的順序總是令人滿意的。

但是還有一點困擾我的是，當(dāng)我在后臺并發(fā)處理成百上千個文件或者執(zhí)行成百上千個進(jìn)程時，我總是擔(dān)心會不會有幾個任務(wù)偷偷掛了，整個代碼永遠(yuǎn)跑不完。我也很難知道代碼現(xiàn)在在哪里執(zhí)行。

最糟糕的是，當(dāng)我看著一個空白屏幕時，很難說出我的代碼需要多長時間才能執(zhí)行或 ETA 是多少。這對我安排工作日程的能力非常不利。

因此，我想要一種方法讓我知道代碼執(zhí)行到了哪里。

已有方法

比較傳統(tǒng)的做法是任務(wù)之間共享一塊內(nèi)存區(qū)域，在這塊內(nèi)存區(qū)域放一個計數(shù)器，當(dāng)一個任務(wù)結(jié)束的時候讓這個計數(shù)器+1，然后用一個線程不停的打印這個計數(shù)器的值。

這從來都不是一個好的解決方案：一方面，我需要在你現(xiàn)有的業(yè)務(wù)邏輯中添加一段用于計數(shù)的代碼，這違反了“低耦合，高內(nèi)聚”的原則。另一方面，由于線程安全問題，我必須非常小心鎖定機(jī)制，這會導(dǎo)致不必要的性能問題。

tqdm

有一天，我發(fā)現(xiàn)了 tqdm 庫，它使用進(jìn)度條來可視化我的代碼進(jìn)度。我可以使用進(jìn)度條來可視化我的 asyncio 任務(wù)的完成和預(yù)計到達(dá)時間嗎？

那么本文我把這個方法分享給大家，讓每個程序員都有機(jī)會監(jiān)控自己并發(fā)任務(wù)的進(jìn)度。

異步

在我們開始之前，我希望您了解一些 Python asyncio 的背景知識。我的文章描述了asyncio的一些常用API的用法，這將有助于我們更好地理解tqdm的設(shè)計：

tqdm 概述

如官方網(wǎng)站所述，tqdm 是一個顯示循環(huán)進(jìn)度條的工具。它使用簡單、高度可定制并且占用資源少。

一個典型的用法是將一個可迭代對象傳遞給 tqdm 構(gòu)造函數(shù)，然后你會得到一個如下所示的進(jìn)度條：

from time import sleep
from tqdm import tqdm


def main():
    for _ in tqdm(range(100)):
        # do something in the loop
        sleep(0.1)


if __name__ == "__main__":
    main()

或者您可以在讀取文件時手動瀏覽并更新進(jìn)度條的進(jìn)度：

import os
from tqdm import tqdm


def main():
    filename = "../data/large-dataset"
    with (tqdm(total=os.path.getsize(filename)) as bar,
            open(filename, "r", encoding="utf-8") as f):
        for line in f:
            bar.update(len(line))


if __name__ == "__main__":
    main()

將 tqdm 與異步集成

總體而言，tqdm 非常易于使用。但是，GitHub 上需要更多關(guān)于將 tqdm 與 asyncio 集成的信息。所以我深入研究了源代碼，看看 tqdm 是否支持 asyncio。

幸運的是，最新版本的 tqdm 提供了包 tqdm.asyncio，它提供了類 tqdm_asyncio。

tqdm_asyncio 類有兩個相關(guān)的方法。一個是 tqdm_asyncio.as_completed。從源碼可以看出，它是對asyncio.as_completed的包裝：

@classmethod
    def as_completed(cls, fs, *, loop=None, timeout=None, total=None, **tqdm_kwargs):
        """
        Wrapper for `asyncio.as_completed`.
        """
        if total is None:
            total = len(fs)
        kwargs = {}
        if version_info[:2] < (3, 10):
            kwargs['loop'] = loop
        yield from cls(asyncio.as_completed(fs, timeout=timeout, **kwargs),
                       total=total, **tqdm_kwargs)

另一個是 tqdm_asyncio.gather ，從源代碼可以看出，它基于模擬 asyncio.gather 功能的 tqdm_asyncio.as_completed 的實現(xiàn)：

@classmethod
    async def gather(cls, *fs, loop=None, timeout=None, total=None, **tqdm_kwargs):
        """
        Wrapper for `asyncio.gather`.
        """
        async def wrap_awaitable(i, f):
            return i, await f

        ifs = [wrap_awaitable(i, f) for i, f in enumerate(fs)]
        res = [await f for f in cls.as_completed(ifs, loop=loop, timeout=timeout,
                                                 total=total, **tqdm_kwargs)]
        return [i for _, i in sorted(res)]

所以，接下來，我將描述這兩個API的用法。在開始之前，我們還需要做一些準(zhǔn)備工作。在這里，我寫了一個簡單的方法來模擬一個隨機(jī)休眠時間的并發(fā)任務(wù)：

import asyncio
import random

from tqdm.asyncio import tqdm_asyncio


class AsyncException(Exception):
    def __int__(self, message):
        super.__init__(self, message)


async def some_coro(simu_exception=False):
    delay = round(random.uniform(1.0, 5.0), 2)

    # We will simulate throwing an exception if simu_exception is True
    if delay > 4 and simu_exception:
        raise AsyncException("something wrong!")

    await asyncio.sleep(delay)

    return delay

緊接著，我們將創(chuàng)建 2000 個并發(fā)任務(wù)，然后使用 tqdm_asyncio.gather 而不是熟悉的 asyncio.gather 方法來查看進(jìn)度條是否正常工作：

async def main():
    tasks = []
    for _ in range(2000):
        tasks.append(some_coro())
    await tqdm_asyncio.gather(*tasks)

    print(f"All tasks done.")


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

或者讓我們用 tqdm_asyncio.as_completed 替換 tqdm_asyncio.gather 并重試：

async def main():
    tasks = []
    for _ in range(2000):
        tasks.append(some_coro())

    for done in tqdm_asyncio.as_completed(tasks):
        await done

    print(f"The tqdm_asyncio.as_completed also works fine.")


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