Run async functions as tasks

Introduction

The awaitable decorator allows for the integration of asyncio within tasks, letting users control asynchronous functions.

We’ll temporarily set the AiiDA log level to REPORT, so that we can inspect the execution order of the workgraph.

from aiida_workgraph.utils.logging import set_aiida_loglevel

set_aiida_loglevel("REPORT")
import asyncio

from aiida import load_profile
from aiida_workgraph import WorkGraph, task

load_profile()
Profile<uuid='fa705ae5c8fb4460b0d168fea18da213' name='presto'>

Use the @task.awaitable decorator on an async function to make it non‑blocking:

@task.awaitable
async def awaitable_task(x, y):
    await asyncio.sleep(0.5)
    return x + y


with WorkGraph("AwaitableGraph") as wg:
    wg.inputs = dict.fromkeys(["x", "y"])
    awaitable_task(x=wg.inputs.x, y=wg.inputs.y)
    wg.inputs.x + wg.inputs.y

wg.run(
    inputs={
        "x": 1,
        "y": 2,
    },
)
07/23/2025 07:46:15 AM <2576> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [391|WorkGraphEngine|continue_workgraph]: tasks ready to run: awaitable_task,op_add
07/23/2025 07:46:16 AM <2576> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [391|WorkGraphEngine|update_task_state]: Task: op_add, type: PyFunction, finished.
07/23/2025 07:46:16 AM <2576> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [391|WorkGraphEngine|continue_workgraph]: tasks ready to run:
07/23/2025 07:46:16 AM <2576> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [391|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: awaitable_task
07/23/2025 07:46:16 AM <2576> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [391|WorkGraphEngine|update_normal_task_state]: Task: awaitable_task finished.
07/23/2025 07:46:16 AM <2576> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [391|WorkGraphEngine|update_task_state]: Task: awaitable_task finished.
07/23/2025 07:46:16 AM <2576> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [391|WorkGraphEngine|continue_workgraph]: tasks ready to run:
07/23/2025 07:46:16 AM <2576> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [391|WorkGraphEngine|finalize]: Finalize workgraph.

Note the timestamps. The addition task runs while the awaitable task sleeps. As the above tasks are functional, they would block one another. The awaitable decorator allows them to run concurrently.

Notes on asyncio integration

The awaitable task lets the workgraph enter a WAITING state, yielding control to the asyncio event loop. This enables other tasks to run concurrently, though long-running calculations may delay the execution of awaitable tasks.

Summary

In this section, we’ve explored the awaitable decorator for integrating asynchronous functions within tasks.

set_aiida_loglevel("ERROR")

Total running time of the script: (0 minutes 2.338 seconds)

Gallery generated by Sphinx-Gallery