Kieren Anderson
Kieren Anderson

Reputation: 528

Python asyncio: exit program immediately if a task throws exception

I have a variable number of long running tasks created with asyncio.create_task(). I would like to gracefully exit my program immediately if any of these tasks raises an exception. I cannot asyncio.wait() the tasks since not all tasks are created at the same time (some running tasks create new tasks at varying times as required).

My thinking is to use Task.add_done_callback() to be notified when tasks complete, but the Python documentation suggests this should be used for low-level callback code only, is there an alternative?

Upvotes: 4

Views: 1681

Answers (1)

user4815162342
user4815162342

Reputation: 154846

If you control the code that creates the task, you can wrap the task coroutine to catch exceptions and handle the failure as needed. For example, if the task is created like this:

# let some_coro() run unattended
asyncio.create_task(some_coro())

you could modify it as follows:

async def run_guarded(aw):
    try:
        await aw
    except:
        import traceback, sys
        traceback.print_exc()
        sys.exit(1)

asyncio.create_task(run_guarded(some_coro()))

If you can't modify the place(s) where tasks are spawned, you will need to use add_done_callback. Don't worry about it being "low-level", that warning is just to prevent the overuse of callbacks in general, and nudge users towards awaiting where possible. This kind of use case is what add_done_callback is good for.

def check_error(t):
    if not t.cancelled() and t.exception() is not None:
        import traceback, sys
        e = t.exception()
        traceback.print_exception(None, e, e.__traceback__)
        sys.exit(1)

# task previously created with create_task()
task.add_done_callback(check_error)

Upvotes: 4

Related Questions