Python: TaskGroup, Timeouts and Cancellation

Description

Starting concurrent work in asyncio is easy. Ending it well is the hard part: when one of three calls fails, what happens to the other two? When a call takes too long, how do you stop it? When a user closes the page, does the work behind it stop too? Python 3.11 added two tools that answer these questions - TaskGroup and asyncio.timeout - and both depend on one mechanism: cancellation. This lesson covers all three, and the habits that keep cancellation working.

Here's the kind of bug we mean. A dashboard fetches a profile (slow) and orders (fast) with asyncio.gather. The orders service goes down, and the request fails after 100 ms - good. But the profile fetch keeps running in the background for its full duration, holding a database connection, for every failed request. During an outage the connection pool fills up with work nobody will ever read. By the end of the lesson you'll know why gather does this, and what to use instead.

We'll go from simple to more technical:

1. gather and how it handles errors.
2. TaskGroup: concurrent work that ends cleanly.
3. Catching errors from a group with except*.
4. Timeouts.
5. Cancellation - and the one rule you must never break.

gather and its errors

try:
    await asyncio.gather(fetch_profile(uid), fetch_orders(uid))
except ValueError as e:
    ...   # orders failed - but fetch_profile is still running

gather runs its arguments concurrently. When one of them raises:

• The first exception is passed straight to you, and the others are not cancelled. In a test, the 0.3 s profile fetch wasn't done when gather raised after 0.1 s - and it finished anyway, 0.2 s later, with nobody waiting for it.
• With return_exceptions=True, every call runs to the end and failures come back as values in the results list: ['a', ValueError('orders service down')]. You then check each result yourself.

Neither mode stops the remaining work when something fails. For that you need a TaskGroup.

TaskGroup

async with asyncio.TaskGroup() as tg:
    profile = tg.create_task(fetch_profile(uid))
    orders = tg.create_task(fetch_orders(uid))
page = render(profile.result(), orders.result())   # both are done here

A TaskGroup (Python 3.11+) follows two simple rules:

• Nobody leaves early. The async with block doesn't end until every task started in it has finished. No task is forgotten in the background, and no exception goes unnoticed.
• One fails, all stop. When a task raises, the group cancels the other tasks, waits for them to stop, and then raises. In a test, a failure after 0.1 s ended the whole group at 0.1 s, with the 0.3 s profile fetch cancelled rather than finished.

That's structured concurrency: concurrent work lives inside a block, like a function call, and is always finished - one way or another - when the block ends. Prefer TaskGroup to gather in new code. (Java's equivalent is StructuredTaskScope.)

Catching errors from a group: except*

A TaskGroup raises an ExceptionGroup - a bundle of exceptions - because several tasks can fail at the same time. It does so even when only one task failed. That means a plain except ValueError doesn't catch it:

try:
    async with asyncio.TaskGroup() as tg:
        tg.create_task(fetch_orders(uid))
except ValueError:
    ...   # never runs - the error is an ExceptionGroup

Python 3.11 added except* for this (PEP 654). It matches the exceptions of that type inside the group:

try:
    async with asyncio.TaskGroup() as tg:
        tg.create_task(fetch_profile(uid))
        tg.create_task(fetch_orders(uid))
except* ValueError as eg:
    log.warning("dashboard failed: %s", eg.exceptions)   # a tuple of every ValueError

If two tasks raised ValueError, eg.exceptions holds both.

Timeouts

Every call over a network needs a time limit. asyncio has two ways to set one:

async with asyncio.timeout(2.0):          # Python 3.11+: a deadline for everything in the block
    user = await get_user(uid)
    orders = await get_orders(uid)

data = await asyncio.wait_for(fetch(url), 2.0)   # a time limit for one call

Both cancel the work when time runs out and raise TimeoutError. Since 3.11, asyncio.TimeoutError is the same class as the built-in TimeoutError, so one except TimeoutError: covers both. asyncio.timeout is usually nicer: one deadline can cover several awaits, and it can be moved while it runs with reschedule().

Cancellation

TaskGroup, timeouts and server shutdown all stop work the same way: task.cancel(). That makes asyncio.CancelledError get raised inside the task, at the await where it's paused. The exception travels up through the task's code, running finally blocks and async with exits on the way, and the task ends as "cancelled".

This only works if the exception is allowed to travel. So:

• CancelledError is a BaseException, not an Exception (since Python 3.8). An except Exception: block lets it through - which is exactly what you want.
• Never swallow it. A coroutine that catches CancelledError - or uses a bare except: or except BaseException: - and carries on breaks everything built on cancellation. In a test, a coroutine that caught it and returned a default value ran under asyncio.timeout(0.1): the block ended normally, with no TimeoutError at all. The timeout silently stopped working.

If you need to clean up when cancelled, use finally, or catch it, clean up, and raise it again:

try:
    await upload(chunk)
finally:
    await release_slot()          # runs on success, failure and cancellation

Some work must finish even if the caller gives up - a payment that's already been sent, a write that's half done. Wrap it in asyncio.shield(): cancelling or timing out the waiter doesn't cancel the shielded task. In a test, wait_for(asyncio.shield(payment), 0.05) raised TimeoutError, the payment task wasn't cancelled, and it finished with 'paid' afterwards. Keep a reference to the shielded task, and make sure something eventually awaits it.

Which tool for which job?

Common mistakes

Check yourself

Try to answer before opening each one. Recalling beats re-reading.

1. One of two calls in gather fails. What happens to the other one?

2. Why doesn't except ValueError catch a failure from inside a TaskGroup?

3. A coroutine catches asyncio.CancelledError and returns a default. What happens when it runs under asyncio.timeout(0.1) and is too slow?

Key takeaways

• gather doesn't stop the other calls when one fails; TaskGroup does, and never leaves tasks behind.
• A TaskGroup raises an ExceptionGroup - catch it with except*.
• asyncio.timeout and wait_for cancel the work and raise the built-in TimeoutError.
• CancelledError is a BaseException: never swallow it; clean up in finally.
• shield the work that must finish.