Support for Semaphore / Parallel Execution / asyncio.FIRST_COMPLETE ?

[patrickvonplaten]

Amazing library! I definitely share your pain with error handling in asyncio!

Question: Can tinyio support asyncio.FIRST_COMPLETED / parallel execution / semaphore use cases?

Background

A common use case many people have with asyncio (especially when it comes to data loading in ML training runs) is to combine asyncio with ThreadPoolExecutors and/or asyncio's Semaphore.

More specifically:

• You want to lazily load large datasets for training. A Python iterator is a common abstraction for this.
• Iterators are then chained together and filtering / transformation methods are added. The final pipeline allows you to go from "load data iterator" -> .... -> sample batch for training iterator
• Now many intermediate transformation (that are blocking and can be very slow & heavy) as well as async functions might be called (requests to server etc...)
  => A common approach currently to solve this is to combine ThreadPoolExecutor with asyncio.Semaphore logic

Example:

E.g. the following:

import asyncio
from concurrent.futures import ThreadPoolExecutor
import time
import random

# Example slow function
def slow_transform(x):
    time.sleep(random.uniform(0.5, 2.0))  # simulate slow IO-bound work
    return x * x

# Exemplary async wrapper for the synchronous slow_transform function
# This way of wrapping block -> async and back is very annoying and makes error management difficult
async def async_slow_transform(x, semaphore, executor):
    async with semaphore:
        loop = asyncio.get_running_loop()
        result = await loop.run_in_executor(executor, slow_transform, x)
        return result

# Main coroutine that processes the iterator
async def process_iterator(iterator, max_concurrent=16):
    semaphore = asyncio.Semaphore(max_concurrent)
    executor = ThreadPoolExecutor(max_workers=max_concurrent)

    tasks = set()

    for item in iterator:
        task = asyncio.create_task(async_slow_transform(item, semaphore, executor))
        tasks.add(task)

        if len(tasks) >= max_concurrent:
            # Wait for the first task to finish and yield its result immediately
            done, tasks = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
            for completed_task in done:
                yield completed_task.result()

    # Yield remaining tasks
    for task in asyncio.as_completed(tasks):
        yield await task

    executor.shutdown()

# Just an example
async def main():
    iterator = range(1, 51)  # Example iterator
    async for result in process_iterator(iterator):
        print("Result:", result)

# Run the main coroutine
if __name__ == "__main__":
    asyncio.run(main())

=> Do you think tinyio could simplify the above workflow significantly and improve error handling?

[patrick-kidger]

Thanks! I'm glad you like this. So the short answer I think is yes, we can absolutely support these things. I got totally nerd-sniped by this, so settle in.

Our strategy is going to be to implement each of the pieces you note on top of the core tinyio abstractions. After that, we'll bring them together to implement your process_iterator in just a few lines of code :)

Step one: Semaphore

First things first, here is how we could implement Semaphore:

import contextlib

class Semaphore:
    def __init__(self, value: int):
        self._value = value

    def __call__(self):
        while self._value == 0:
            yield
        self._value -= 1
        return _close(self, [False])

@contextlib.contextmanager
def _close(semaphore: Semaphore, cell):
    if cell[0]:
        raise RuntimeError("Use a new `semaphore()` call in each `with (yield semaphore())`, do not re-use it.")
    cell[0] = True
    try:
        yield
    finally:
        semaphore._value += 1

And here's a standalone demo usage of it:

import time
import tinyio

def async_foo(semaphore):
    with (yield semaphore()):
        yield tinyio.sleep(1)

def semaphore_demo():
    semaphore = Semaphore(2)
    start = time.time()
    yield [async_foo(semaphore), async_foo(semaphore), async_foo(semaphore)]
    duration = time.time() - start
    print(f"Semaphore demo ran in {duration} seconds.")

loop = tinyio.Loop()
loop.run(semaphore_demo())
# Semaphore demo ran in 2.0001420974731445 seconds.

Step two: ThreadPool

Next up, a thread pool:

class ThreadPool:
    def __init__(self, value: int):
        self._semaphore = Semaphore(value)

    def run_in_thread(self, fn, /, *args, **kwargs):
        with (yield self._semaphore()):
            out = yield tinyio.run_in_thread(fn, *args, **kwargs)
        return out

And a standalone demo of a threadpool:

def blocking_foo():
    time.sleep(1)
    return 1

def thread_pool_demo():
    pool = ThreadPool(2)
    start = time.time()
    yield [
        pool.run_in_thread(blocking_foo),
        pool.run_in_thread(blocking_foo),
        pool.run_in_thread(blocking_foo),
    ]
    duration = time.time() - start
    print(f"Thread pool demo ran in {duration} seconds.")

loop = tinyio.Loop()
loop.run(thread_pool_demo())
# Thread pool demo ran in 2.022390127182007 seconds.

Step three: done callbacks

Okay, a quick helper function: done callbacks. This actually needs one new primitive from the tinyio loop, namely the ability to schedule work in the background. I've implemented this in #3. Running on top of that PR, then we begin by first implementing a done-callback:

from collections.abc import Callable
from typing import Any

def _with_callback(coro, success_callback, error_callback):
    try:
        yield
        out = yield coro
    except BaseException as e:
        if error_callback is not None:
            error_callback(e)
        raise
    else:
        if success_callback is not None:
            success_callback(out)

def add_done_callback(
    coro: tinyio.Coro,
    success_callback: Callable[[Any], None] | None,
    error_callback: Callable[[BaseException], None] | None,
):
    coro2 = _with_callback(coro, success_callback, error_callback)
    # prime it, so that if an error is raised it will definitely be inside the try/except block.
    next(coro2)
    yield {coro2}

Demo usage of done callbacks:

def _sleep(x):
    yield tinyio.sleep(x)
    return x

def done_callback_demo():
    yield add_done_callback(_sleep(2), lambda *a: print("Done callback demo: ", *a), None)
    yield add_done_callback(_sleep(1), lambda *a: print("Done callback demo: ", *a), None)

loop = tinyio.Loop()
loop.run(done_callback_demo())
# Done callback demo: 1
# Done callback demo: 2

Step four: AsAvailable

Now, on to implementing a loop that yields the result of coroutines as they become available. This is our equivalent of repeatdly iterating with asyncio.FIRST_COMPLETED.

import queue

class AsAvailable:
    def __init__(self, coros: list[tinyio.Coro]):
        self._coros = list(coros)
        self._count = 0
        self._queue = queue.SimpleQueue()
        self._started = False

    def done(self):
        return self._count == len(self._coros)

    def get(self):
        if not self._started:
            self._started = True
            yield {add_done_callback(coro, self._queue.put, None) for coro in self._coros}
        while self._queue.empty():
            yield
        self._count += 1
        return self._queue.get()

Demo usage of as-available:

def as_available_demo():
    iterator = AsAvailable([_sleep(7), _sleep(2), _sleep(4)])
    while not iterator.done():
        out = yield iterator.get()
        print(f"As available demo: {out}")

loop = tinyio.Loop()
loop.run(as_available_demo())

Step five: bringing it all together!

Treating each of these as library pieces, here is how we might implement an example like yours (which as a side-note does not require both a semaphore and a threadpool, as the latter already uses a semaphore internally):

import random

def slow_transform(x):
    time.sleep(random.uniform(0.5, 2.0))  # simulate slow IO-bound work
    return x * x

def process_iterator(iterator, max_concurrent=16):
    pool = ThreadPool(max_concurrent)
    return AsAvailable([pool.run_in_thread(slow_transform, item) for item in iterator])

def main():
    iterator = range(1, 51)
    async_iterator = process_iterator(iterator)
    while not async_iterator.done():
        out = yield async_iterator.get()
        print(f"Result: {out}")

loop = tinyio.Loop()
loop.run(main())

---

This was a fun challenge! 😁

[patrickvonplaten]

Wow super cool! I'll give that a try and report back asap

[patrick-kidger]

I've also just upstreamed these building blocks into tinyio in #4 :)