使用asyncio並行化生成器

Question

我的應用程序從慢速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() ，緩沖結果然后將它們寫入文件是不行的 。我需要一個可以超出主存的大數據集的解決方案，這就是為什么我一直在到目前為止使用發電機。然而，可以安全地假設消費者總是比所有生產者組合更快）

Answer 1

從python 3.6和異步生成器開始 ，需要進行很少的更改才能使代碼與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