如何在越来越多的任务上使用 asyncio.wait？

Question

在以下代码中，创建了一个任务A并将其添加到一组任务tasks中。
然后我使用await asyncio.wait(tasks)等待任务完成。

但这没有考虑在任务A中创建的任务B1 （递归 function 调用）。

然后下面的代码不会等待B1完成：看下面的结果，任务B1永远不会完成。

我认为原因是当在线（**）评估tasks时，它仍然只有一个元素。

问题：如何让await asyncio.wait(tasks)在不断发展/增长的任务集上工作？

import asyncio

tasks = set()
i = 0

async def mytask(s):
    global i
    print('mytask %s starting' % s)
    await asyncio.sleep(1)
    if i < 4:    # limit number of tasks
        print('mytask %s creating new task' % s)
        i += 1
        tasks.add(asyncio.create_task(mytask('B%i' % i)))
    print('mytask %s len tasks:' % s, len(tasks))
    await asyncio.sleep(0.5)
    print('mystak %s finished' % s)

async def main():
    print('main starting')
    tasks.add(asyncio.create_task(mytask('A')))
    print('len tasks:', len(tasks))
    await asyncio.wait(tasks)            # (**)
    # await asyncio.sleep(10)
    print('main finished')

asyncio.run(main())

结果：

main starting
len tasks: 1
mytask A starting
mytask A creating new task
mytask A len tasks: 2
mytask B1 starting       # <--- mytask B1 will never complete!
mystak A finished
main finished

如果我们将 (**) 行替换为await asyncio.sleep(10) ，当然，所有任务都会完成：

main starting
len tasks: 1
mytask A starting
mytask A creating new task
mytask A len tasks: 2
mytask B1 starting
mystak A finished
mytask B1 creating new task
mytask B1 len tasks: 3
mytask B2 starting
mystak B1 finished
mytask B2 creating new task
mytask B2 len tasks: 4
mytask B3 starting
mystak B2 finished
mytask B3 creating new task
mytask B3 len tasks: 5
mytask B4 starting
mystak B3 finished
mytask B4 len tasks: 5
mystak B4 finished
main finished

Answer 1

要开始直接回答您的问题，您可以等待一组带有循环的动态任务，例如：

while tasks:
    prev_tasks = tasks.copy()
    # use gather() so exceptions are propagated rather than discarded
    await asyncio.gather(*tasks)
    tasks.difference_update(prev_tasks)

但你可能不需要这样做。 相反，您可以让每个任务等待它创建的子任务以及它自己的工作。 这样，您甚至不需要拥有一组全局任务，也无需担心在main()中等待所有任务：

import asyncio

i = 0

async def mytask(s):
    global i
    print('mytask %s starting' % s)
    await asyncio.sleep(1)
    if i < 4:    # limit number of tasks
        print('mytask %s creating new task' % s)
        i += 1
        task = asyncio.create_task(mytask('B%i' % i))
    else:
        task = None
    print('mytask %s len tasks:' % s, i)
    await asyncio.sleep(0.5)   # our actual work
    print('mystak %s finished' % s)
    # after doing the work, wait for the child task if we created one
    if task is not None:
        await task

async def main():
    print('main starting')
    await mytask('A')
    # await asyncio.sleep(10)
    print('main finished')

asyncio.run(main())

Answer 2

使用一个元素的列表调用asyncio.wait ，直到它的开始/处理没有 B 任务。 在您的情况下，最简单的解决方案就是在 A 中等待 B 任务，但 A 在 B 完成之前不会返回。

如果这不适合您，您可以使用某种忙等待 - 检查任务长度的无限循环：

import asyncio

tasks = set()


async def mytask(s):
    print('mytask %s starting' % s)
    await asyncio.sleep(1)
    print('mytask %s create new task' % s)
    tasks.add(asyncio.create_task(mytask('B')))
    print('mytask %s len tasks:' % s, len(tasks))
    await asyncio.sleep(0.5)
    print('mystak %s finished' % s)

async def main():
    print('main starting')
    tasks.add(asyncio.create_task(mytask('A')))
    print('len tasks:', len(tasks))
    while True:
        if all([task.done() for task in tasks]): break
        await asyncio.wait(tasks)
    print('main finished')

asyncio.run(main())

请记住，忙碌的等待经常被过度使用。 此外，它看起来像是在 asyncio 的任务调度器之上实现了一个任务调度器（它在引擎盖下也有while True ）。

另一种解决方案是run_forever循环而不是asyncio.run 。 对于工作人员长时间运行的应用程序来说感觉很好。

asyncio.ensure_future(main())
asyncio.get_event_loop().run_forever()

您还可以重构代码以利用asyncio.Queue并使用Queue.join等待所有已处理的项目。