Note
Go to the end to download the full example code.
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)