如何将 asyncio 与很长的任务列表（生成器）一起使用

Question

我有一个小程序，它加载一个相当大的 CSV（超过 800MB，分块，使用pandas.read_csv 来限制内存使用）并“在野外”对服务器执行一些 API 调用，最后构建一个结果然后将对象存储在数据库中。

我已尽可能为网络请求添加缓存，但即便如此，代码仍需要 10 多个小时才能完成。当我使用 PySpy 分析代码时，大部分代码都在等待网络请求。

我尝试将其转换为使用 asyncio 来加快速度，并设法让代码在输入文件的一小部分上工作。然而，对于完整的文件，内存使用变得令人望而却步。

这是我尝试过的：

import pandas as pd
import httpx

async def process_item(item, client):
    # send a few requests with httpx session
    # process results
    await save_results_to_db(res)

async def get_items_from_csv():
    # loads the heavy CSV file
    for chunk in pd.read_csv(filename, ...):
        for row in chunk.itertuples():
            item = item_from_row(row)
            yield item

async def main():
    async with httpx.AsyncClient() as client:
        tasks = []
        for item in get_items_from_csv():
            tasks.append(process_item(item, client))
        await asyncio.gather(*tasks)

asyncio.run(main())

有没有办法避免创建tasks 列表，它会变成一个包含超过 150 万个项目的非常重的对象？另一个缺点是在读取整个文件之前似乎没有处理任何任务，这并不理想。 我使用的是 python 3.7，但如果需要可以轻松升级到 3.8。

Answer 1

我认为您在这里寻找的不是批量运行，而是运行 N 工作人员，它们同时从队列中拉出任务。

N = 10  # scale based on the processing power and memory you have

async def main():
    async with httpx.AsyncClient() as client:
        tasks = asyncio.Queue()
        for item in get_items_from_csv():
            tasks.put_nowait(process_item(item, client))

        async def worker():
            while not tasks.empty():
                await tasks.get_nowait()
            # for a server
            # while task := await tasks.get():
            #     await task

        await asyncio.gather(*[worker() for _ in range(N)])

我使用了asyncio.Queue，但您也可以只使用collections.deque，因为所有任务都在启动工作程序之前添加到队列中。前者在运行长时间运行的进程（例如服务器）中运行的工作人员时特别有用，其中项目可能会异步排队。