Asyncio Log Task Exceptions With Done Callback Function

February 11, 2024 Python Asyncio

You can log task exceptions automatically in asyncio by defining a general done callback that checks the task for an exception, and if present, logs it.

This general done callback function can then be added routinely to all tasks created in the program to ensure that task exceptions are always logged as soon as each task is done.

In this tutorial, you will discover how to automatically log task exceptions using a done callback function.

Let's get started.

Need To Log Task Exceptions Automatically

Asyncio tasks can fail with exceptions.

If an exception is raised within a coroutine or task, it will cause the task to terminate.

If the exception is not explicitly retrieved, it will be reported or logged automatically when the asyncio event loop is terminated.

For example, you can learn more about the logging of never-retrieved exceptions in the tutorial:

If the program runs for a long time, such as hours or days, these faults may not be seen for a long time and therefore not acted upon.

As such, we need a way to automatically log task exceptions as soon as they cause a task to terminate.

How can we automatically log task exceptions in asyncio?

How to Log Task Exceptions With a Done Callback

We can automatically log task exceptions as soon as they terminate a task using a done callback function.

This can be achieved by defining a general done callback function that can be added to all tasks created in an asyncio program.

This general done callback function can check the task to see if an exception was raised, and if so retrieve it and log it.

The Python standard library logging module will automatically log the stack traces of raised exceptions when logging them as an exception-level message within a try-except block.

Therefore, our general done callback function can cause the exception that terminated the task to be re-raised and logged within a try-exception block, ensuring the trace of the exception is also logged.

The helper_done_callback() done callback function below implements this.

# callback func called for all tasks
def helper_done_callback(task):
    # check if the task had an exception
    if task.exception():
        try:
            # re-raise exception
            task.result()
        except Exception as e:
            logging.exception(f'Task {task} failed with exception {e}')

This done callback can then be added to all asyncio tasks created in a program.

For example:

...
# add custom callback to task
task.add_done_callback(helper_done_callback)

This will run the done callback for all tasks and only log exceptions from those tasks that failed with an unhandled exception.

Importantly, the exception is logged as soon as the task is done.

We can make this easier by creating a helper function that creates tasks for us and always adds the done callback function.

The create_task_helper() below implements this.

# helper for creating all tasks
def create_task_helper(coroutine):
    # wrap and schedule the task
    task = asyncio.create_task(coroutine)
    # add custom callback to task
    task.add_done_callback(helper_done_callback)
    # return the task that was created
    return task

We can then create tasks in our program by calling this helper function and passing in a created coroutine object.

For example:

...
# schedule a task
task = create_task_helper(work())

The create_task_helper() function could be updated to support an asyncio.TaskGroup, if needed.

Now that we know how we can automatically log task exceptions using a done callback, let's look at some worked examples.

Example of Never-Retrieved Task Exception

We will look at exploring the automatic logging of an exception with a done callback function, let's start with a coroutine that fails with an exception that is not retrieved.

We can define a task that reports a message, sleeps a moment, then fails with an exception. The coroutine can be run as a background task and the exception never-retrieved, causing it to be reported by the event loop when the program is exited.

The complete example is listed below.

# SuperFastPython.com
# example of a never-retrieved exception
import logging
import asyncio

# task that does work and logs
async def work():
    # log a message
    logging.info(f'Task is starting')
    # simulate doing work
    await asyncio.sleep(1)
    # fail with an exception
    raise RuntimeError('Something bad happened')
    # log a message
    logging.info(f'Task is done')

# main coroutine
async def main():
    # log a message
    logging.info(f'Main is starting')
    # schedule a task
    task = asyncio.create_task(work())
    # wait around
    await asyncio.sleep(2)
    # log a message
    logging.info(f'Main is done')

# prepare the logger
logging.basicConfig(level=logging.DEBUG)
# start the event loop
asyncio.run(main())

Running the example first configures the logging infrastructure then starts the event loop and runs the main() coroutine.

The main() coroutine runs and reports a log message, then creates and starts an asyncio.Task for the work() coroutine to run in the background, then suspends and sleeps for two seconds.

The work() coroutine runs and logs a message. It then sleeps for one second. Once it resumes it fails with an exception.

The main() coroutine resumes and logs a final message and the program exits.

We can see that when the event loop is closed, an ERROR message is logged indicating that a task exception was never logged, followed by the details of the task and a trace of the exception.

This highlights that never-retrieved task exceptions are logged automatically by the event loop after the event loop is closed.

DEBUG:asyncio:Using selector: KqueueSelector
INFO:root:Main is starting
INFO:root:Task is starting
INFO:root:Main is done
ERROR:asyncio:Task exception was never retrieved
future: <Task finished name='Task-2' coro=<work() done, defined at ...:7> exception=RuntimeError('Something bad happened')>
Traceback (most recent call last):
  File "...", line 13, in work
    raise RuntimeError('Something bad happened')
RuntimeError: Something bad happened

Next, let's look at how we can log task errors automatically using a done callback function.

Example of Logging Task Exception With Done Callback

We can explore an example of automatically logging task exceptions using a done callback function.

In this case, can update the above example to systematically create tasks using our create_task_helper() function and add our helper_done_callback() done callback function to each task that is created.

# callback func called for all tasks
def helper_done_callback(task):
    # check if the task had an exception
    if task.exception():
        try:
            # re-raise exception
            task.result()
        except Exception as e:
            logging.exception(f'Task {task} failed with exception {e}')

# helper for creating all tasks
def create_task_helper(coroutine):
    # wrap and schedule the task
    task = asyncio.create_task(coroutine)
    # add custom callback to task
    task.add_done_callback(helper_done_callback)
    # return the task that was created
    return task

This will ensure that any tasks created will automatically log their exceptions via the general done callback function.

To achieve this, all we need to do is update the main() coroutine to create the work() background task using our create_task_helper() helper function.

...
# schedule a task
task = create_task_helper(work())

The updated main() coroutine with these changes is listed below.

# main coroutine
async def main():
    # log a message
    logging.info(f'Main is starting')
    # schedule a task
    task = create_task_helper(work())
    # wait around
    await asyncio.sleep(2)
    # log a message
    logging.info(f'Main is done')

Tying this together, the complete example is listed below.

# SuperFastPython.com
# example of logging exception with done callback function
import logging
import asyncio

# callback func called for all tasks
def helper_done_callback(task):
    # check if the task had an exception
    if task.exception():
        try:
            # re-raise exception
            task.result()
        except Exception as e:
            logging.exception(f'Task {task} failed with exception {e}')

# helper for creating all tasks
def create_task_helper(coroutine):
    # wrap and schedule the task
    task = asyncio.create_task(coroutine)
    # add custom callback to task
    task.add_done_callback(helper_done_callback)
    # return the task that was created
    return task

# task that does work and logs
async def work():
    # log a message
    logging.info(f'Task is starting')
    # simulate doing work
    await asyncio.sleep(1)
    # fail with an exception
    raise RuntimeError('Something bad happened')
    # log a message
    logging.info(f'Task is done')

# main coroutine
async def main():
    # log a message
    logging.info(f'Main is starting')
    # schedule a task
    task = create_task_helper(work())
    # wait around
    await asyncio.sleep(2)
    # log a message
    logging.info(f'Main is done')

# prepare the logger
logging.basicConfig(level=logging.DEBUG)
# start the event loop
asyncio.run(main())

Running the example first configures the logging infrastructure then starts the event loop and runs the main() coroutine.

The main() coroutine runs and reports a log message, then creates and starts an asyncio.Task for the work() coroutine to run in the background using our helper function.

This creates the task as per normal, then adds our custom done callback function to log task exceptions.

The main() coroutine then suspends and sleeps for two seconds.

The work() coroutine runs and logs a message. It then sleeps for one second. Once it resumes it fails with an exception.

The work() task is done, and then the done callback function runs. It checks if the task has an exception, which it does. It then attempts to retrieve the result of the task, which re-raises the exception. The exception is handled and logged using logging.exception(), ensuring the trace of the exception is also logged.

The main() coroutine resumes and logs a final message and the program exits.

This highlights how we can use a done callback function to automatically log task exceptions as soon as the task is done using a done callback function.

DEBUG:asyncio:Using selector: KqueueSelector
INFO:root:Main is starting
INFO:root:Task is starting
ERROR:root:Task <Task finished name='Task-2' coro=<work() done, defined at ...:26> exception=RuntimeError('Something bad happened')> failed with exception Something bad happened
Traceback (most recent call last):
  File "...", line 12, in helper_done_callback
    task.result()
  File "...", line 32, in work
    raise RuntimeError('Something bad happened')
RuntimeError: Something bad happened
INFO:root:Main is done

Takeaways

You now know how to automatically log task exceptions using a done callback function.