How to Create Task in Done Callback

April 14, 2024 Python Asyncio

We can schedule asyncio tasks in the event loop from done callback functions.

Done callback functions are regular Python functions executed after a task is done. Because they are regular Python functions, they cannot await tasks or coroutines.

Nevertheless, we can create and schedule new tasks from done callback functions, and this may require using a mechanism to inform the main() coroutine that there are more tasks to complete before terminating and exiting the event loop.

In this tutorial, you will discover how to create and schedule a new asyncio task in a done callback function.

Let's get started.

Asyncio Callback Functions

We can add a done callback function to tasks in asyncio.

This callback gets triggered when the task is completed, whether it's successfully done, canceled, or raises an exception.

This is helpful in order to perform some action or cleanup after the task is completed, such as logging, resource cleanup, or initiating another action after the task finishes its execution.

This can be achieved using the add_done_callback() method and specifying the function name to be called when the task is completed.

For example:

...
# add a callback to a task
task.add_done_callback(callback)

You can learn more about how to use done callback functions in the tutorial:

How to Use Asyncio Task Done Callback Functions

Cannot Await in a Done Callback Function

The done callback function is a regular Python function, it is not a coroutine.

This means that we cannot await in a done callback function.

For example:

# custom done callback function
def callback(task):
	# await another coroutine
	await coro()

This will result in a SyntaxError exception.

We get a SyntaxError because we cannot use the "await" expression outside of a coroutine, such as in a regular Python function.

You can learn more about this syntax error in the tutorial:

SyntaxError: 'await' outside function

How to Create Tasks in Callback Functions

Although we cannot await in a done callback function, we can create and schedule new asyncio tasks.

This can be achieved via the asyncio.create_task() function.

We can call this function and pass it a coroutine object.

This will wrap the coroutine in an asyncio.Task and schedule it for execution in the current asyncio event loop.

For example:

# custom done callback function
def callback(task):
	# create and schedule another task
	other = asyncio.create_task(coro())

This can be helpful if we want to chain tasks together and pass results from one task to another.

For example:

# custom done callback function
def callback(task):
	# get the result from the done task
	result = task.result()
	# create and schedule another task and pass the result
	other = asyncio.create_task(coro(result))

You can see an example of using this approach to chain coroutines in the tutorial:

Asyncio Coroutine Chaining

Now that we know how to create an asyncio task in a done callback function, let's look at some worked examples.

Example Cannot Await in a Callback

We can explore an example of a SyntaxError when attempting to await in a done callback function.

In this case, we will define a done callback function that awaits another coroutine, asyncio.sleep() in this case. We will then create a task and add our custom done callback to the task.

Firstly, we can define a simple task. In this case, our task will report a message, sleep for 2 seconds, then report a final message.

The work() coroutine below implements this.

# coroutine to perform some work
async def work():
    # report a message
    print('Work is starting...')
    # suspend a moment
    await asyncio.sleep(2)
    # report a message
    print('Work is done')

Next, we can define our custom callback function.

The function takes the task instance as an argument, and then awaits the asyncio.sleep() coroutine.

# callback function to schedule another task
def callback(task):
    # await another coroutine
    await asyncio.sleep(1)

Finally, we can define the main() coroutine.

It first reports a message, then creates a task from our work() coroutine and adds our callback() done callback function to the task.

It then awaits the task and reports a final message.

# main coroutine
async def main():
    # report a message
    print('Main is starting...')
    # create the task
    work_task = asyncio.create_task(work())
    # add the done callback
    work_task.add_done_callback(callback)
    # wait for the task to be done
    await work_task
    # report a final message
    print('Main Done.')

Tying this together, the complete example is listed below.

# SuperFastPython.com
# example of cannot await within callback function
import asyncio

# coroutine to perform some work
async def work():
    # report a message
    print('Work is starting...')
    # suspend a moment
    await asyncio.sleep(2)
    # report a message
    print('Work is done')

# callback function to schedule another task
def callback(task):
    # await another coroutine
    await asyncio.sleep(1)

# main coroutine
async def main():
    # report a message
    print('Main is starting...')
    # create the task
    work_task = asyncio.create_task(work())
    # add the done callback
    work_task.add_done_callback(callback)
    # wait for the task to be done
    await work_task
    # report a final message
    print('Main Done.')

# start the event loop
asyncio.run(main())

Running the example fails before it is able to run.

A SyntaxError is raised that reports that we attempted to use an "await" expression outside of a coroutine

This highlights that we cannot directly await a coroutine or a task within a done callback function.

SyntaxError: 'await' outside async function

Next, let's explore an example of creating and scheduling a new asyncio task from a done callback function.

Example of Creating a Task in a Callback

We can explore an example of scheduling a new asyncio task from a custom done callback function.

In this case, we will define a second coroutine called more_work() and schedule it as a task from a done callback added to the first task. This can be achieved by updating the above example.

Firstly, we can define the new more_work() that reports a message, sleeps, and reports a final message, much like the work() coroutine.

# coroutine to perform some more work
async def more_work():
    # report a message
    print('More work is starting...')
    # suspend a moment
    await asyncio.sleep(2)
    # report a message
    print('More work is done')

Next, we can define a done callback function that creates and schedules the more_work() as a new asyncio task.

# callback function to schedule another task
def callback(task):
    # schedule another task
    task = asyncio.create_task(more_work())

The main() coroutine then needs to be updated to wait for all tasks in the event loop to be completed.

This can be achieved using a loop where each iteration it gets a list of all running tasks via the asyncio.all_tasks() function, removes the current task from the set by first getting it via the asyncio.current_task(), and then awaiting all remaining tasks. If there are no remaining tasks, then the loop is exited.

This ensures that the work() task is completed and the new more_work() task is completed before the main() coroutine exits.

...
# wait for all tasks to be done
while True:
    # get all tasks
    tasks = asyncio.all_tasks()
    # remove this task
    tasks.remove(asyncio.current_task())
    # check for no more tasks
    if not tasks:
        break
    # wait on remaining tasks
    for task in tasks:
        await task

If the main coroutine does not wait for all of the tasks to be done, it will exit and in turn the asyncio event loop will terminate and cancel all running tasks.

You can learn more about this special property of the main() coroutine in the tutorial:

What is the Main Coroutine

Tying this together, the updated main() coroutine with these changes is listed below.

# main coroutine
async def main():
    # report a message
    print('Main is starting...')
    # create the task
    work_task = asyncio.create_task(work())
    # add the done callback
    work_task.add_done_callback(callback)
    # wait for all tasks to be done
    while True:
        # get all tasks
        tasks = asyncio.all_tasks()
        # remove this task
        tasks.remove(asyncio.current_task())
        # check for no more tasks
        if not tasks:
            break
        # wait on remaining tasks
        for task in tasks:
            await task
    # report a final message
    print('Main Done.')

Tying this together, the complete example of scheduling a task from a done callback function is listed below.

# SuperFastPython.com
# example of scheduling a task from a done callback function
import asyncio

# coroutine to perform some more work
async def more_work():
    # report a message
    print('More work is starting...')
    # suspend a moment
    await asyncio.sleep(2)
    # report a message
    print('More work is done')

# coroutine to perform some work
async def work():
    # report a message
    print('Work is starting...')
    # suspend a moment
    await asyncio.sleep(2)
    # report a message
    print('Work is done')

# callback function to schedule another task
def callback(task):
    # schedule another task
    task = asyncio.create_task(more_work())

# main coroutine
async def main():
    # report a message
    print('Main is starting...')
    # create the task
    work_task = asyncio.create_task(work())
    # add the done callback
    work_task.add_done_callback(callback)
    # wait for all tasks to be done
    while True:
        # get all tasks
        tasks = asyncio.all_tasks()
        # remove this task
        tasks.remove(asyncio.current_task())
        # check for no more tasks
        if not tasks:
            break
        # wait on remaining tasks
        for task in tasks:
            await task
    # report a final message
    print('Main Done.')

# start the event loop
asyncio.run(main())

Running the example first creates the main() coroutine and creates and starts the asyncio event loop to run it.

The main() coroutine runs and reports a message. It then creates the work() task and adds the callback() done callback. The main() coroutine then suspends and waits for all running tasks to be done.

The work() task runs, reporting a message, sleeping for 2 seconds, and then reporting a final message before terminating.

The callback() function is then executed automatically and the more_work() coroutine is then created and scheduled as a task.

The more_work() then runs, reports a message, sleeps for 2 seconds, then reports a final message before terminating.

The main() coroutine then determines that there are no other tasks running and then exits, closing the event loop.

This highlights how we can create and schedule a new asyncio task from a done callback function, and have the main() coroutine wait for the new tasks to complete.

Main is starting...
Work is starting...
Work is done
More work is starting...
More work is done
Main Done.

Next, let's look at a simpler way to notify the main() coroutine that all scheduled tasks are done.

Example of Simplified Waiting For Pipeline of Tasks

We can explore a simplified version of the pipeline defined above.

In the above example, the main coroutine had to loop and introspect the internal status of the asyncio event loop in order to determine if there were other tasks running. This may be an overly complicated solution.

We can simplify the example by using an asyncio.Event.

An event is a coroutine-safe boolean variable that can be set and cleared by different coroutines and allows one coroutine to suspend and wait for it to be set.

You can learn more about how to use the asyncio.Event in the tutorial:

How to Use an Asyncio Event in Python

The main() coroutine can declare and define an event as a global variable, and then wait on it after issuing the work() coroutine.

For example:

...
# declare a shared event
global event
# define shared event
event = asyncio.Event()
...
# wait for all tasks to be done
await event.wait()

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

# main coroutine
async def main():
    # declare a shared event
    global event
    # define shared event
    event = asyncio.Event()
    # report a message
    print('Main is starting...')
    # create the task
    work_task = asyncio.create_task(work())
    # add the done callback
    work_task.add_done_callback(callback)
    # wait for all tasks to be done
    await event.wait()
    # report a final message
    print('Main Done.')

The more_work() coroutine can then declare the event as a global variable and set it.

...
# declare shared event
global event
# mark the pipeline done
event.set()

The updated more_work() coroutine with this change is listed below.

# coroutine to perform some more work
async def more_work():
    # report a message
    print('More work is starting...')
    # suspend a moment
    await asyncio.sleep(2)
    # report a message
    print('More work is done')
    # declare shared event
    global event
    # mark the pipeline done
    event.set()

This will have the effect of signaling to the asyncio event loop that all tasks are done.

A variation on this approach would be to define a second done callback function that sets the event. This is left as an exercise.

Tying this together, the complete example of using an asyncio.Event to signal that all tasks are done is listed below.

# SuperFastPython.com
# example of scheduling a task from a done callback function
import asyncio

# coroutine to perform some more work
async def more_work():
    # report a message
    print('More work is starting...')
    # suspend a moment
    await asyncio.sleep(2)
    # report a message
    print('More work is done')
    # declare shared event
    global event
    # mark the pipeline done
    event.set()

# coroutine to perform some work
async def work():
    # report a message
    print('Work is starting...')
    # suspend a moment
    await asyncio.sleep(2)
    # report a message
    print('Work is done')

# callback function to schedule another task
def callback(task):
    # schedule another task
    task = asyncio.create_task(more_work())

# main coroutine
async def main():
    # declare a shared event
    global event
    # define shared event
    event = asyncio.Event()
    # report a message
    print('Main is starting...')
    # create the task
    work_task = asyncio.create_task(work())
    # add the done callback
    work_task.add_done_callback(callback)
    # wait for all tasks to be done
    await event.wait()
    # report a final message
    print('Main Done.')

# start the event loop
asyncio.run(main())

Running the example first creates the main() coroutine and creates and starts the asyncio event loop to run it.

The main() coroutine runs and declares the asyncio.Event as a global variable, then defines it. It then reports a message.

Next, the main() coroutine then creates the work() task and adds the callback() done callback.

The main() coroutine then suspends and waits for the asyncio.Event to be set.

The work() task runs, reporting a message, sleeping for 2 seconds, and then reporting a final message before terminating.

The callback() function is then executed automatically and the more_work() coroutine is then created and scheduled as a task.

The more_work() then runs, reports a message, sleeps for 2 seconds, then reports a final message. It then sets the event and terminates.

The main() coroutine resumes, reports a final message then terminates, closing the event loop.

This highlights how we can use an asyncio.Event to signal to the main() coroutine that all tasks scheduled in a done callback function are completed.

Main is starting...
Work is starting...
Work is done
More work is starting...
More work is done
Main Done.

Takeaways

You now know how to create and schedule a new asyncio task in a done callback function.

If you enjoyed this tutorial, you will love my book: Python Asyncio Jump-Start. It covers everything you need to master the topic with hands-on examples and clear explanations.