Note

Go to the end to download the full example code.

Quick Start

Introduction

To run this tutorial, you need to install aiida-workgraph. Open a terminal and run:

pip install aiida-workgraph

If you haven’t configured an AiiDA profile, you can run the following

verdi presto

Load the AiiDA profile.

from aiida import load_profile

load_profile()

Profile<uuid='90da34ae855c481f9f23e4e2526238f1' name='presto'>

First workflow

Suppose we want to calculate (x + y) * z in two steps:

add x and y
then multiply the result with z.

Create task

Task is the basic building block of the WorkGraph. A task has inputs, outputs and an executor.

from aiida_workgraph import task

# define add task
@task()
def add(x, y):
    return x + y


# define multiply task
@task()
def multiply(x, y):
    return x * y

Create the workflow

Three steps:

create an empty WorkGraph
add tasks: add and multiply.
link the output of the add task to the x input of the multiply task.

from aiida_workgraph import WorkGraph

wg = WorkGraph("add_multiply_workflow")
add_task = wg.add_task(add, name="add1")
# link the output of the `add` task to one of the `x` input of the `multiply` task.
wg.add_task(multiply, name="multiply1", x=add_task.outputs.result)

# export the workgraph to html file so that it can be visualized in a browser
wg.to_html()
# comment out the following line to visualize the workgraph in jupyter-notebook
# wg

Run and view results

from aiida.orm import Int

# ------------------------- Run the calculation -------------------
wg.run(
    inputs={"add1": {"x": Int(2), "y": Int(3)}, "multiply1": {"y": Int(4)}},
)

print("State of WorkGraph:   {}".format(wg.state))
print("Result of add      : {}".format(wg.tasks.add1.outputs.result.value))
print("Result of multiply : {}".format(wg.tasks.multiply1.outputs.result.value))

06/03/2025 01:42:37 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [107|WorkGraphEngine|continue_workgraph]: tasks ready to run: add1
06/03/2025 01:42:38 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [107|WorkGraphEngine|update_task_state]: Task: add1, type: PyFunction, finished.
06/03/2025 01:42:38 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [107|WorkGraphEngine|continue_workgraph]: tasks ready to run: multiply1
06/03/2025 01:42:38 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [107|WorkGraphEngine|update_task_state]: Task: multiply1, type: PyFunction, finished.
06/03/2025 01:42:38 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [107|WorkGraphEngine|continue_workgraph]: tasks ready to run:
06/03/2025 01:42:38 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [107|WorkGraphEngine|finalize]: Finalize workgraph.
State of WorkGraph:   FINISHED
Result of add      : uuid: fab7b3b5-2b0e-4e40-8e82-5ebc3fbe61eb (pk: 111) value: 5
Result of multiply : uuid: 25b65ba9-1373-470c-95af-6a1e352553d3 (pk: 115) value: 20

CalcJob and WorkChain

One can use AiiDA components (CalcJob and WorkChain) direclty in the WorkGraph. The inputs and outputs of the task is automatically generated based on the input and output port of the AiiDA components.

Here is an example of using the ArithmeticAddCalculation Calcjob inside the workgraph.

from aiida_workgraph import WorkGraph
from aiida.calculations.arithmetic.add import ArithmeticAddCalculation
from aiida.orm import Int, InstalledCode, load_computer, load_code
from aiida.common.exceptions import NotExistent


try:
    code = load_code("add@localhost")  # The computer label can also be omitted here
except NotExistent:
    code = InstalledCode(
        computer=load_computer("localhost"),
        filepath_executable="/bin/bash",
        label="add",
        default_calc_job_plugin="core.arithmetic.add",
    ).store()

wg = WorkGraph("test_add_multiply")
add1 = wg.add_task(ArithmeticAddCalculation, name="add1", x=Int(2), y=Int(3), code=code)
add2 = wg.add_task(ArithmeticAddCalculation, name="add2", y=Int(3), code=code)
wg.add_link(wg.tasks.add1.outputs.sum, wg.tasks.add2.inputs.x)
wg.to_html()

Run the workgraph and wait for the result.

wg.run()
print("\nResult of task add1: {}".format(wg.tasks.add2.outputs.sum.value))

06/03/2025 01:42:39 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [120|WorkGraphEngine|continue_workgraph]: tasks ready to run: add1
06/03/2025 01:42:39 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [120|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 121
06/03/2025 01:42:40 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [120|WorkGraphEngine|update_task_state]: Task: add1, type: CALCJOB, finished.
06/03/2025 01:42:41 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [120|WorkGraphEngine|continue_workgraph]: tasks ready to run: add2
06/03/2025 01:42:41 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [120|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 125
06/03/2025 01:42:42 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [120|WorkGraphEngine|update_task_state]: Task: add2, type: CALCJOB, finished.
06/03/2025 01:42:42 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [120|WorkGraphEngine|continue_workgraph]: tasks ready to run:
06/03/2025 01:42:42 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [120|WorkGraphEngine|finalize]: Finalize workgraph.

Result of task add1: uuid: 2d3f29dd-7ee6-4225-ba12-e7fe483562eb (pk: 128) value: 8

One can also create task from an AiiDA process builder directly.

from aiida.calculations.arithmetic.add import ArithmeticAddCalculation

builder = ArithmeticAddCalculation.get_builder()
builder.code = code
builder.x = Int(2)
builder.y = Int(3)

wg = WorkGraph("test_set_inputs_from_builder")
add1 = wg.add_task(builder, name="add1")

Graph builder

A WorkGraph is a group of tasks. One can treat a WorkGraph as a single task, and expose the inputs and outputs of the WorkGraph. This allow you to write:

nested workflows
dynamic workflow based on the input values. For example, if you want to use if and for to create the tasks, or repeat a calculation until it converges.

The Graph Builder allow user to create a dynamic workflow based on the input value, as well as nested workflows. Here is an example of nested workflow:

from aiida_workgraph import WorkGraph, task

# define add task
@task()
def add(x, y):
    return x + y


# define multiply task
@task()
def multiply(x, y):
    return x * y


# use task.graph_builder decorator, expose the "result" output of "multiply" task
# as the "multiply" output of the `WorkGraph`.
@task.graph_builder(outputs=[{"name": "multiply"}])
def add_multiply(x, y, z):
    # Create a WorkGraph
    wg = WorkGraph()
    wg.add_task(add, name="add", x=x, y=y)
    wg.add_task(multiply, name="multiply", x=z)
    wg.add_link(wg.tasks.add.outputs.result, wg.tasks.multiply.inputs.y)
    wg.outputs.multiply = wg.tasks.multiply.outputs.result
    # don't forget to return the `wg`
    return wg

Use this graph builder inside a WorkGraph:

from aiida_workgraph import WorkGraph
from aiida.orm import Int

wg = WorkGraph("test_graph_build")
# create a task using the graph builder
add_multiply1 = wg.add_task(add_multiply, x=Int(2), y=Int(3), z=Int(4))
add_multiply2 = wg.add_task(add_multiply, x=Int(2), y=Int(3))
# link the output of a task to the input of another task
wg.add_link(add_multiply1.outputs.multiply, add_multiply2.inputs.z)
wg.run()
print("WorkGraph state: ", wg.state)

06/03/2025 01:42:43 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [134|WorkGraphEngine|continue_workgraph]: tasks ready to run: add_multiply1
06/03/2025 01:42:43 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [134|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 135
06/03/2025 01:42:44 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [135|WorkGraphEngine|continue_workgraph]: tasks ready to run: add
06/03/2025 01:42:44 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [135|WorkGraphEngine|update_task_state]: Task: add, type: PyFunction, finished.
06/03/2025 01:42:44 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [135|WorkGraphEngine|continue_workgraph]: tasks ready to run: multiply
06/03/2025 01:42:44 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [135|WorkGraphEngine|update_task_state]: Task: multiply, type: PyFunction, finished.
06/03/2025 01:42:44 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [135|WorkGraphEngine|continue_workgraph]: tasks ready to run:
06/03/2025 01:42:44 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [135|WorkGraphEngine|finalize]: Finalize workgraph.
06/03/2025 01:42:45 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [134|WorkGraphEngine|update_task_state]: Task: add_multiply1, type: graph_builder, finished.
06/03/2025 01:42:45 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [134|WorkGraphEngine|continue_workgraph]: tasks ready to run: add_multiply2
06/03/2025 01:42:45 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [134|WorkGraphEngine|on_wait]: Process status: Waiting for child processes: 144
06/03/2025 01:42:45 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [144|WorkGraphEngine|continue_workgraph]: tasks ready to run: add
06/03/2025 01:42:46 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [144|WorkGraphEngine|update_task_state]: Task: add, type: PyFunction, finished.
06/03/2025 01:42:46 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [144|WorkGraphEngine|continue_workgraph]: tasks ready to run: multiply
06/03/2025 01:42:46 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [144|WorkGraphEngine|update_task_state]: Task: multiply, type: PyFunction, finished.
06/03/2025 01:42:46 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [144|WorkGraphEngine|continue_workgraph]: tasks ready to run:
06/03/2025 01:42:46 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [144|WorkGraphEngine|finalize]: Finalize workgraph.
06/03/2025 01:42:46 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [134|WorkGraphEngine|update_task_state]: Task: add_multiply2, type: graph_builder, finished.
06/03/2025 01:42:46 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [134|WorkGraphEngine|continue_workgraph]: tasks ready to run:
06/03/2025 01:42:46 PM <2571> aiida.orm.nodes.process.workflow.workchain.WorkChainNode: [REPORT] [134|WorkGraphEngine|finalize]: Finalize workgraph.
WorkGraph state:  FINISHED

Get the result of the tasks:

print("Result of task add_multiply1: {}".format(add_multiply1.outputs.multiply.value))

Result of task add_multiply1: uuid: 0f2c3112-9108-40da-ba97-2f26ed655c30 (pk: 143) value: 20

Start the web server

WorkGraph also provides a web GUI, where you can view and manage the workgraph. To use the web ui, first install the web ui package:

pip install aiida-workgraph-web-ui

Open a terminal, and run:

workgraph web start

Then visit the page http://127.0.0.1:8000/workgraph, you can view the workgraph later from here. You should find all the submited workgraph, e.g., the first_workflow WorkGraph. Please click the pk and view the workgraph.

What’s Next

Concepts	A brief introduction of WorkGraph’s main concepts.
Tutorials	Real-world examples in computational materials science and more.
HowTo	Advanced topics and tips, e.g flow control using `if`, `for`, `while` and `context`.

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

Gallery generated by Sphinx-Gallery