如何使用 Python asyncio 限制並發？

Question

假設我們有一堆鏈接要下載，每個鏈接可能需要不同的時間來下載。 而且我只能使用最多 3 個連接進行下載。 現在，我想確保使用 asyncio 有效地執行此操作。

這是我想要實現的目標：在任何時間點，盡量確保我至少運行 3 次下載。

Connection 1: 1---------7---9---
Connection 2: 2---4----6-----
Connection 3: 3-----5---8-----

數字代表下載鏈接，而連字符代表等待下載。

這是我現在正在使用的代碼

from random import randint
import asyncio

count = 0


async def download(code, permit_download, no_concurrent, downloading_event):
    global count
    downloading_event.set()
    wait_time = randint(1, 3)
    print('downloading {} will take {} second(s)'.format(code, wait_time))
    await asyncio.sleep(wait_time)  # I/O, context will switch to main function
    print('downloaded {}'.format(code))
    count -= 1
    if count < no_concurrent and not permit_download.is_set():
        permit_download.set()


async def main(loop):
    global count
    permit_download = asyncio.Event()
    permit_download.set()
    downloading_event = asyncio.Event()
    no_concurrent = 3
    i = 0
    while i < 9:
        if permit_download.is_set():
            count += 1
            if count >= no_concurrent:
                permit_download.clear()
            loop.create_task(download(i, permit_download, no_concurrent, downloading_event))
            await downloading_event.wait()  # To force context to switch to download function
            downloading_event.clear()
            i += 1
        else:
            await permit_download.wait()
    await asyncio.sleep(9)

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    try:
        loop.run_until_complete(main(loop))
    finally:
        loop.close()

輸出符合預期：

downloading 0 will take 2 second(s)
downloading 1 will take 3 second(s)
downloading 2 will take 1 second(s)
downloaded 2
downloading 3 will take 2 second(s)
downloaded 0
downloading 4 will take 3 second(s)
downloaded 1
downloaded 3
downloading 5 will take 2 second(s)
downloading 6 will take 2 second(s)
downloaded 5
downloaded 6
downloaded 4
downloading 7 will take 1 second(s)
downloading 8 will take 1 second(s)
downloaded 7
downloaded 8

但這里是我的問題：

1. 目前，我只是等待 9 秒以保持主函數運行，直到下載完成。 有沒有一種有效的方法可以在退出main函數之前等待最后一次下載完成？ （我知道有asyncio.wait ，但我需要存儲所有任務引用才能使其工作）

2. 什么是完成此類任務的好圖書館？ 我知道 javascript 有很多異步庫，但是 Python 呢？

編輯： 2. 什么是處理常見異步模式的好庫？ （類似於async ）

Answer 1

如果我沒記錯的話，您正在尋找asyncio.Semaphore 。 用法示例：

import asyncio
from random import randint


async def download(code):
    wait_time = randint(1, 3)
    print('downloading {} will take {} second(s)'.format(code, wait_time))
    await asyncio.sleep(wait_time)  # I/O, context will switch to main function
    print('downloaded {}'.format(code))


sem = asyncio.Semaphore(3)


async def safe_download(i):
    async with sem:  # semaphore limits num of simultaneous downloads
        return await download(i)


async def main():
    tasks = [
        asyncio.ensure_future(safe_download(i))  # creating task starts coroutine
        for i
        in range(9)
    ]
    await asyncio.gather(*tasks)  # await moment all downloads done


if __name__ ==  '__main__':
    loop = asyncio.get_event_loop()
    try:
        loop.run_until_complete(main())
    finally:
        loop.run_until_complete(loop.shutdown_asyncgens())
        loop.close()

輸出：

downloading 0 will take 3 second(s)
downloading 1 will take 3 second(s)
downloading 2 will take 1 second(s)
downloaded 2
downloading 3 will take 3 second(s)
downloaded 1
downloaded 0
downloading 4 will take 2 second(s)
downloading 5 will take 1 second(s)
downloaded 5
downloaded 3
downloading 6 will take 3 second(s)
downloading 7 will take 1 second(s)
downloaded 4
downloading 8 will take 2 second(s)
downloaded 7
downloaded 8
downloaded 6

可以在此處找到使用aiohttp進行異步下載的示例。

Answer 2

我用了 Mikhails 的回答，最后得到了這個小寶石

async def gather_with_concurrency(n, *tasks):
    semaphore = asyncio.Semaphore(n)

    async def sem_task(task):
        async with semaphore:
            return await task
    return await asyncio.gather(*(sem_task(task) for task in tasks))

您將運行而不是正常收集

await gather_with_concurrency(100, *my_coroutines)

Answer 3

在閱讀本答案的其余部分之前，請注意使用asyncio.Semaphore限制並行任務數量的慣用方法是使用asyncio.Semaphore ，如Mikhail 的回答所示並在Andrei 的回答中優雅地抽象。 這個答案包含工作，但實現相同的更復雜的方法。 我留下答案是因為在某些情況下，這種方法可能比信號量更具優勢，特別是當要完成的工作非常大或無限時，並且您無法提前創建所有協程。 在這種情況下，第二個（基於隊列的）解決方案就是這個答案就是你想要的。 但是在大多數常規情況下，例如通過 aiohttp 並行下載，您應該改用信號量。

---

您基本上需要一個固定大小的下載任務池。 asyncio沒有預制的任務池，但很容易創建一個：只需保留一組任務，不要讓它增長超過限制。 盡管問題表明您不願意走這條路，但代碼最終變得更加優雅：

import asyncio
import random

async def download(code):
    wait_time = random.randint(1, 3)
    print('downloading {} will take {} second(s)'.format(code, wait_time))
    await asyncio.sleep(wait_time)  # I/O, context will switch to main function
    print('downloaded {}'.format(code))

async def main(loop):
    no_concurrent = 3
    dltasks = set()
    i = 0
    while i < 9:
        if len(dltasks) >= no_concurrent:
            # Wait for some download to finish before adding a new one
            _done, dltasks = await asyncio.wait(
                dltasks, return_when=asyncio.FIRST_COMPLETED)
        dltasks.add(loop.create_task(download(i)))
        i += 1
    # Wait for the remaining downloads to finish
    await asyncio.wait(dltasks)

另一種方法是創建固定數量的協程進行下載，就像固定大小的線程池一樣，並使用asyncio.Queue它們提供工作。 這消除了手動限制下載次數的需要，下載次數將自動受到調用download()的協程數量的限制：

# download() defined as above

async def download_worker(q):
    while True:
        code = await q.get()
        await download(code)
        q.task_done()

async def main(loop):
    q = asyncio.Queue()
    workers = [loop.create_task(download_worker(q)) for _ in range(3)]
    i = 0
    while i < 9:
        await q.put(i)
        i += 1
    await q.join()  # wait for all tasks to be processed
    for worker in workers:
        worker.cancel()
    await asyncio.gather(*workers, return_exceptions=True)

至於你的另一個問題，顯而易見的選擇是aiohttp 。

Answer 4

asyncio-pool 庫正是您所需要的。

https://pypi.org/project/asyncio-pool/

LIST_OF_URLS = ("http://www.google.com", "......")

pool = AioPool(size=3)
await pool.map(your_download_coroutine, LIST_OF_URLS)

Answer 5

使用信號量，你還可以創建一個裝飾器來包裝函數

import asyncio
from functools import wraps
def request_concurrency_limit_decorator(limit=3):
    # Bind the default event loop 
    sem = asyncio.Semaphore(limit)

    def executor(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            async with sem:
                return await func(*args, **kwargs)

        return wrapper

    return executor

然后，將裝飾器添加到原點下載功能。

@request_concurrency_limit_decorator(limit=...)
async def download(...):
    ...

現在你可以像以前一樣調用下載函數，但使用信號量來限制並發。

await download(...)

需要注意的是，在執行裝飾器函數時，創建的信號量綁定到默認的事件循環，因此不能調用asyncio.run來創建新的循環。 相反，調用asyncio.get_event_loop().run...以使用默認事件循環。

asyncio.Semaphore RuntimeError: Task got Future 附加到不同的循環

Answer 6

小更新：不再需要創建循環。 我調整了下面的代碼。 只是稍微清理一下。

# download(code) is the same

async def main():
    no_concurrent = 3
    dltasks = set()
    for i in range(9):
        if len(dltasks) >= no_concurrent:
            # Wait for some download to finish before adding a new one
            _done, dltasks = await asyncio.wait(dltasks, return_when=asyncio.FIRST_COMPLETED)
        dltasks.add(asyncio.create_task(download(i)))
    # Wait for the remaining downloads to finish
    await asyncio.wait(dltasks)

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

Answer 7

如果您有一個生成器來生成您的任務，那么可能有更多的任務可以同時容納在內存中。

經典的asyncio.Semaphore上下文管理器模式將所有任務同時asyncio.Semaphore到內存中。

我不喜歡asyncio.Queue模式。 您可以阻止它將所有任務預加載到內存中（通過設置maxsize=1 ），但它仍然需要樣板來定義、啟動和關閉工作協程（從隊列中消耗），並且您必須確保工作人員獲勝如果任務拋出異常，則不會失敗。 感覺不是 Pythonic，就像實現你自己的multiprocessing.pool 。

相反，這里有一個替代方案：

sem = asyncio.Semaphore(n := 5) # specify maximum concurrency

async def task_wrapper(args):
    try:
        await my_task(*args)
    finally:
        sem.release()

for args in my_generator: # may yield too many to list
    await sem.acquire() 
    asyncio.create_task(task_wrapper(args))

# wait for all tasks to complete
for i in range(n):
    await sem.acquire()

當有足夠多的活動任務時，這會暫停生成器，並讓事件循環清理已完成的任務。 請注意，對於較舊的 Python 版本，請將create_task替換為ensure_future 。