WorkGraph

WorkGraph is a collection of tasks and links.

Create workgraph

First, create an empty workgraph:

from aiida_workgraph import WorkGraph, task

wg = WorkGraph(name='my_first_workgraph')

Define and use tasks

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


# Add tasks to the workgraph
add1 = wg.add_task(add, name='add1')
add2 = wg.add_task(add, name='add2')

Add a link between tasks:

wg.add_link(add1.outputs.result, add2.inputs.x)

# Visualize the graph
wg

Note

By creating a link to add2.inputs.x, you can no longer set this task’s input directly. It can only be changed through changes to add1.outputs.result.

Execute the workgraph

With the graph defined, you can now execute it. You provide the inputs for the tasks.

from aiida import load_profile

load_profile()
wg.run(inputs={'add1': {'x': 1, 'y': 2}, 'add2': {'y': 3}})

{}

Graph-level inputs and outputs

As workflows grow, managing inputs for many tasks can become cumbersome. WorkGraph allows you to define graph-level inputs and outputs to create a cleaner, more user-friendly interface for your complex logic.

This lets you:

• Reuse a single input across multiple tasks.

• Hide internal complexity and only expose essential inputs.

• Collect and rename key results as named workflow outputs.

wg = WorkGraph('graph_inputs_outputs')

# Define graph-level input
wg.inputs.x = 2

# Add tasks using the graph-level input
wg.add_task(add, 'add1', x=wg.inputs.x, y=3)
wg.add_task(add, 'add2', x=wg.inputs.x, y=wg.tasks.add1.outputs.result)

# Define graph-level outputs to expose selected task results
wg.outputs.sum1 = wg.tasks.add1.outputs.result
wg.outputs.sum2 = wg.tasks.add2.outputs.result

# Run the WorkGraph
wg.run()

# Verify the final output
assert wg.outputs.sum2.value == 2 + (2 + 3)

# Visualize the graph with inputs and outputs
wg

Note

As when explicitly linking tasks, graph-level inputs create a link that cannot be overridden by setting task-level inputs. In this case, this means you cannot set wg.tasks.add1.inputs.x or wg.tasks.add2.inputs.x, in addition to being unable to set wg.tasks.add2.inputs.y due to the link to add1’s result.

Context variables

Context variables (ctx) are used to store and pass intermediate data within a workflow that isn’t directly an input or output of a task. This is especially useful for workflows with conditional logic (if/else) or loops, where you need to manage state between steps.

wg = WorkGraph(name='context_example')
# Setting the ``ctx`` attribute of the WorkGraph directly, on initialization
wg.ctx = {'x': 2, 'data.y': 3}
wg.add_task(add, 'add1', x=wg.ctx.x, y=wg.ctx.data.y)
# Assign the result of a task to a context variable
wg.ctx.sum = wg.tasks.add1.outputs.result
# Use the context variable in another task
wg.add_task(add, 'add2', x=wg.ctx.x, y=wg.ctx.sum)

Context variables can be nested, allowing you to organize complex data structures. For example, you can store multiple results in a structured way:

wg.ctx.data = {
    'sum1': wg.tasks.add1.outputs.result,
    'sum2': wg.tasks.add2.outputs.result,
}

WorkGraph engine

The WorkGraph engine operates on a dataflow programming model. Once submitted, the engine continuously monitors the tasks in the graph. A task is executed only when all of its inputs are available. This means:

• Tasks with no inputs are executed first.

• A task starts only after all the upstream tasks linked to its inputs are finished.