如何限制Python异步IO的并发性?

如果我没记错的话你在找 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 可以找到 here . 请注意 aiohttp公司 具有内置的等效信号量,您可以看到 here . 默认限制为100个连接。

我用了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, *tasks)

在阅读本答案的其余部分之前,请注意,限制异步IO并行任务数量的惯用方法是 asyncio.Semaphore ,如所示 Mikhail's answer 优雅地抽象在 Andrei's answer . 这个答案包含了工作,但实现同样目标的方式要复杂一些。我留下答案是因为在某些情况下,这种方法可能比信号量有优势,特别是当要做的工作非常大或没有边界,并且您无法提前创建所有协同路由时。在这种情况下,第二个(基于队列的)解决方案是这个答案是您想要的。但在大多数常规情况下,例如通过AIOHTTPP并行下载,应该使用信号量。

你基本上需要一个固定的尺寸 水塘 下载任务数。 asyncio 没有预先创建的任务池,但创建一个任务池很容易:只需保留一组任务,不允许其增长超过限制。虽然问题表明您不愿意走这条路,但代码最终要优雅得多:

import asyncio, 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 .

asyncio池库正好满足您的需要。

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)

使用信号量,您还可以创建一个装饰器来包装函数

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

然后,将decorator添加到origin下载函数中。

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

现在可以像以前一样调用下载函数,但使用信号量来限制并发性。

await download(...)

应该注意的是,在执行decorator函数时,创建的信号量被绑定到默认事件循环,因此您不能调用 asyncio.run 创建新循环。相反,请致电 asyncio.get_event_loop().run... 使用默认事件循环。

asyncio.Semaphore RuntimeError: Task got Future attached to a different loop

如果您有一个生成器来生成任务,那么可能会有更多的任务无法同时放入内存。

经典 asyncio.Semaphore 上下文管理器模式同时将所有任务分配到内存中。

我不喜欢 asyncio.Queue 图案你 可以 防止它将所有任务预加载到内存中(通过设置 maxsize=1 ),但它仍然需要样板文件来定义、启动和关闭工作者协同路由(使用que),并且您必须确保在任务引发异常时,工作者不会失败。感觉很不和谐,好像实现了你自己的 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 .

小更新: 不再需要创建循环。我调整了下面的代码。只是稍微清理一下。

# 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())