Skip to content

← Concurrency with Threads and asyncio step 49 of 56

Hard End-to-End

Map Over a Stream with at Most N Calls at Once

You have a stream of user IDs and an async fetch to call for each one. bounded_map(fetch, ids, concurrency=10) runs up to 10 calls at once and yields results as calls finish. It reads a new ID only when a call slot is free. That is the backpressure. If your loop stops taking results, bounded_map stops reading IDs.

async with aclosing(bounded_map(fetch, ids(), concurrency=10)) as results:
    async for user in results:
        ...

The previous article's semaphore limits the calls running at once, but its example creates every task up front. Here you start a call only when a slot is free, so you never hold more than concurrency items.

What to write

bounded_map(fn, items, *, concurrency) is a plain def. If concurrency < 1, raise ValueError(f"concurrency must be at least 1, got {concurrency}") straight away, before you read an item or call fn. Otherwise return the async generator _bounded_map(fn, items, concurrency).

_bounded_map repeats these steps:

  1. Fill the free slots. While fewer than concurrency calls are running, read the next item with await anext(items) and start a task for it. Stop reading when the source runs out. Number the tasks in the order you start them.
  2. Stop when idle. If no call is running, return.
  3. Wait for a call to finish. await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) returns as soon as at least one task has finished. Its first value is the set of done tasks, which can hold more than one.
  4. Check the whole done set. Sort it by task number, and stop tracking those tasks. Look at every outcome before you yield anything from the set.
  5. Raise or yield. If any task in the set failed, raise as described below, and yield none of the set's results. Otherwise yield its results in task-number order, then go back to step 1. Don't refill between yields from the same set.

So results come out one done set at a time, in start order within a set.

Give each call its own task. Wrap the call in a small coroutine, async def invoke(item: T) -> R: return await fn(item), and pass that to asyncio.create_task. fn can return any awaitable, and create_task accepts only coroutines. An error that fn raises then becomes that task's failure too.

Failures and cleanup

  • Collect the failures of the done set in task-number order. A cancelled task counts as a failure, with its CancelledError.
  • Raise BaseExceptionGroup("bounded_map failed", failures). When every failure is an Exception, Python makes it an ExceptionGroup.
  • In a finally, cancel every task still running and await them with asyncio.gather(..., return_exceptions=True). This runs on a failure, when the caller closes the generator, and when the caller is cancelled. A cancellation from outside propagates after the cleanup, as a plain CancelledError.
  • An error from reading items also propagates after the cleanup. bounded_map did not open items, so closing it is the caller's job.

Track the tasks yourself, in a dictionary from task to number. A TaskGroup would stay open across your yield. The article on yield and TaskGroups showed why that goes wrong.

What the driver reports

Keep _drive and solve as they are. The items are indices into hops. Call i waits hops[i] loop turns, then returns i * 100 + hops[i]. At index fail_at it raises RuntimeError(f"boom at {i}") instead. A missing, negative or too-large fail_at never fails. After each result, the consumer waits consumer_lag more loop turns.

  • results: the values you yielded.
  • max_inflight: the most calls running at the same time. Calls with zero hops finish without overlapping, so this can be 1 with free slots.
  • max_lookahead: the most items read but not yet delivered. It never goes over concurrency, even with a slow consumer.
  • pulled: how many items you read.
  • pending_tasks: other tasks still running at the end. It must be 0.
  • error: "ValueError: <message>", or "ExceptionGroup[1]: RuntimeError: boom at <i>" for one failed call, or "".

For hops=[24, 8, 16] and concurrency=3, results is [108, 216, 24]. With five calls of 4 hops each, concurrency=2 and fail_at=2, calls 2 and 3 finish in the same done set. That set holds a failure, so results stops at [4, 104].