使用asyncio并行化生成器

时间:2017-10-17 11:29:21

标签: python multithreading generator python-asyncio

我的应用程序从慢速i / o源读取数据,进行一些处理,然后将其写入本地文件。我用这样的生成器实现了这个:

import time

def io_task(x):
    print("requesting data for input %s" % x)
    time.sleep(1)   # this simulates a blocking I/O task
    return 2*x

def producer(xs):
    for x in xs:
        yield io_task(x)

def consumer(xs):
    with open('output.txt', 'w') as fp:
        for x in xs:
            print("writing %s" % x)
            fp.write(str(x) + '\n')

data = [1,2,3,4,5]
consumer(producer(data))

现在我想在asyncio的帮助下并行化这项任务,但我似乎无法弄清楚如何。对我来说,主要问题是通过生成器直接从生产者向消费者提供数据,同时让asyncio向io_task(x)发出多个并行请求。此外,整个async def@asyncio.coroutine这件事令我困惑。

有人可以告诉我如何使用此示例代码中的asyncio构建一个最小的工作示例吗?

(注意:只需调用io_task() 即可,缓冲结果然后将其写入文件。我需要一个适用于大数据集的解决方案可以超过主记忆,这就是我到目前为止一直在使用发电机的原因。然而,假设消费者总是比所有生产商的总和更快,这是安全的。

1 个答案:

答案 0 :(得分:6)

自python 3.6和asynchronous generators以来,需要进行很少的更改才能使代码与asyncio兼容。

io_task函数成为协程:

async def io_task(x):
    await asyncio.sleep(1)
    return 2*x

producer生成器成为异步生成器:

async def producer(xs):
    for x in xs:
        yield await io_task(x)

consumer函数成为协程并使用aiofiles,异步上下文管理和异步迭代:

async def consumer(xs):
    async with aiofiles.open('output.txt', 'w') as fp:
        async for x in xs:
            await fp.write(str(x) + '\n')

主协程运行在一个事件循环中:

data = [1,2,3,4,5]
main = consumer(producer(data))
loop = asyncio.get_event_loop()
loop.run_until_complete(main)
loop.close()

此外,您可以考虑使用aiostream来管理生产者和消费者之间的一些处理操作。

编辑:使用as_completed可以在生产者端轻松地同时运行不同的I / O任务:

async def producer(xs):
    coros = [io_task(x) for x in xs]
    for future in asyncio.as_completed(coros):
        yield await future