Как ограничить количество одновременных операций чтения / записи с помощью файлов AI?

Question

Моя программа одновременно загружала около 10 миллионов фрагментов данных с помощью aiohttp, а затем записывала данные примерно в 4000 файлов на диск.

Я использую библиотеку aiofiles, потому что хочу, чтобы моя программа выполняла и другие действия при чтении / записи файла.

Но меня беспокоит, что если программа попытается записать все 4000 файлов одновременно, жесткий диск не сможет выполнить все операции записи так быстро.

Можно ли ограничить количество одновременных записей с помощью aiofiles (или другой библиотеки)? Aiofiles уже это делает?

Спасибо.

тестовый код:

import aiofiles
import asyncio


async def write_to_disk(fname):
    async with aiofiles.open(fname, "w+") as f:
        await f.write("asdf")


async def main():
    tasks = [asyncio.create_task(write_to_disk("%d.txt" % i)) 
             for i in range(10)]
    await asyncio.gather(*tasks)


asyncio.run(main())

Answer 1

Вы можете использовать asyncio.Semaphoreчтобы ограничить количество одновременных задач. Просто заставьте свойwrite_to_disk функция для получения семафора перед записью:

import aiofiles
import asyncio


async def write_to_disk(fname, sema):
    # Edit to address comment: acquire semaphore after opening file
    async with aiofiles.open(fname, "w+") as f, sema:
        print("Writing", fname)
        await f.write("asdf")
        print("Done writing", fname)


async def main():
    sema = asyncio.Semaphore(3)  # Allow 3 concurrent writers
    tasks = [asyncio.create_task(write_to_disk("%d.txt" % i, sema)) for i in range(10)]
    await asyncio.gather(*tasks)


asyncio.run(main())

Обратите внимание на sema = asyncio.Semaphore(3) линия, а также добавление sema, в async with.

Выход:

"""
Writing 1.txt
Writing 0.txt
Writing 2.txt
Done writing 1.txt
Done writing 0.txt
Done writing 2.txt
Writing 3.txt
Writing 4.txt
Writing 5.txt
Done writing 3.txt
Done writing 4.txt
Done writing 5.txt
Writing 6.txt
Writing 7.txt
Writing 8.txt
Done writing 6.txt
Done writing 7.txt
Done writing 8.txt
Writing 9.txt
Done writing 9.txt
"""