asyncio compute/output pattens

async.py

import asyncio
from time import perf_counter, sleep


def compute(i):
    print('start compute', i)
    sleep(5)
    print('finish compute', i)
    return i


async def output_async(i):
    print('start io', i)
    sleep(5)
    print('finish io', i)


def output_sync(i):
    print('start io', i)
    sleep(5)
    print('finish io', i)


option = 1


async def main():
    result = 0
    background_tasks = set()

    for i in range(10):
        # Assumption: compute should block output and next compute (uses previous result)

        if option == 1:  # Option 1: Compute on separate thread, output on main thread - as required e.g. for cv2.imshow
            result = await asyncio.to_thread(compute, result + 1)
            task = asyncio.create_task(output_async(result))

        elif option == 2:  # Option 2: Both compute and output on separate threads
            result = await asyncio.to_thread(compute, result + 1)
            task = asyncio.create_task(asyncio.to_thread(output_sync, result))

        elif option == 3:  # Option 3: Compute on main thread, output on separate thread
            result = compute(result + 1)
            task = asyncio.create_task(asyncio.to_thread(output_sync, result))

        # Save the reference. See: https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
        background_tasks.add(task)
        task.add_done_callback(background_tasks.discard)  # Prevent explosion (for longer runs)

        if option == 3:
            await asyncio.sleep(0)  # Task only starts at an await following it (here only needed for option 3)

    # Await all background tasks to ensure they finish
    await asyncio.gather(*background_tasks)


start = perf_counter()
asyncio.run(main())
print(perf_counter() - start)