Using the Context

Introduction

In AiiDA workflows (both, traditional WorkChain s, as well as the WorkGraph), the Context (typically represented by ctx), is an internal container that can hold data shared between different tasks. It’s particularly useful for more complex workflows.

When to use

1. Many-to-many relationships where direct linking becomes unwieldy

2. Conditional workflows where links depend on runtime conditions

3. Graph builders that need to expose selective internal state

#

Passing data to the context

There are three ways to set data in the Context: - Setting the ctx attribute of the WorkGraph directly - Using the update_ctx method of the WorkGraph - Using the workgraph.set_context task

So let’s have a look how to use each of them:

from aiida_workgraph import WorkGraph, task
from aiida import load_profile
from aiida.orm import Int

load_profile()

wg1 = WorkGraph(name="context_1")

# 1. Setting the ``ctx`` attribute of the WorkGraph directly, on initialization
wg1.ctx = {"x": Int(2), "data.y": Int(3)}

# 2. Using the ``update_ctx`` method
wg2 = WorkGraph(name="context_2")


@task.calcfunction()
def add(x, y):
    return x + y


add1 = wg2.add_task(add, "add1", x=2, y=3)

# set result of add1 to ctx.sum
wg2.update_ctx({"sum": add1.outputs.result})

# 3. Using the ``workgraph.set_context`` task to set either a task result (socket) or a resolved value to the ctx
wg3 = WorkGraph(name="context_3")
add1 = wg3.add_task(add, "add1", x=2, y=3)
wg3.add_task(
    "workgraph.set_context", name="set_ctx1", key="sum", value=add1.outputs.result
)

SetContext(name='set_ctx1', properties=[], inputs=['context', 'key', 'value', '_wait'], outputs=['_wait'])

Nested context keys

To organize the context data in a hierarchical structure, the keys may contain dots .` that create nesting Here is an example, to group the results of multipl add tasks to ctx.sum:

wg = WorkGraph(name="ctx_nested")
add1 = wg.add_task(add, "add1", x=1, y=2)
add2 = wg.add_task(add, "add2", x=3, y=4)

wg.update_ctx({"sum.add1": add1.outputs.result})
wg.update_ctx({"sum.add2": add2.outputs.result})

# Or, alternatively:
# wg.update_ctx({
#     "sum": {
#         "add1": add1.outputs.result,
#         "add2": add2.outputs.result
#     }
# })

print(wg.ctx.sum)

TaskSocketNamespace(name='sum', sockets=['add1', 'add2'])

Use data from the context

Also for accessing data from the context, there are different approaches:

1. One can use elements from the wg.ctx.x, directly, e.g., as inputs for other tasks

add2 = wg1.add_task(add, "add2", x=wg1.ctx.x, y=3)

# 2. Similarly, nested context keys can be accessed, such as ``wg2.ctx.sum.add1``
print(wg.ctx.sum.add1)
# Gives: SocketAny(name='add1', value=None)

# 3. One can use the `get_context` task to get the data from ctx. **This task will be shown in the GUI**
#
wg3.add_task("workgraph.get_context", name="get_ctx1", key="sum.add1")

wg.show()
wg.to_html()

#
# 2. One can export the data from context to the graph builder outputs.
#
@task.graph()
def internal_add(x, y):
    outputs1 = add(x=x, y=y)
    outputs2 = add(x=outputs1.sum, y=5)
    return outputs2.sum


# Usage in a main WorkGraph
wg_main = WorkGraph("main")
builder_task = wg_main.add_task(internal_add, x=10, y=20)
final_task = wg_main.add_task(add, x=builder_task.outputs.result, y=100)

# In this example, the context can serve as a bridge between the internal workings of a subgraph and its external
# interface. This allows selectively exposing internal results as clean, named outputs.
#

SocketAny(name='add1', value=None)
--------------------------------------------------------------------------------
WorkGraph: ctx_nested, PK: None, State: CREATED
--------------------------------------------------------------------------------
Tasks:
Name       PK    State
---------  ----  -------
add1             PLANNED
add2             PLANNED
graph_ctx        PLANNED
--------------------------------------------------------------------------------

Context Tasks in the GUI

This example shows how context operations appear in the workflow visualization when using explicit context tasks instead of direct context access.

Note

Context tasks make data flow visible in the GUI, which is helpful for debugging and understanding workflow execution. Use them when you want explicit visibility of context operations.

wg = WorkGraph(name="context_gui_demo")

# Set initial context values
wg.ctx = {"x": 2, "multiplier": 10}

# Use context tasks - these appear as nodes in the GUI
get_x = wg.add_task("workgraph.get_context", name="get_x", key="x")
get_multiplier = wg.add_task(
    "workgraph.get_context", name="get_multiplier", key="multiplier"
)

# Perform calculation using context values
add1 = wg.add_task(add, "add1", x=get_x.outputs.result, y=get_multiplier.outputs.result)

# Store result back to context - also appears in GUI
wg.add_task(
    "workgraph.set_context",
    name="store_result",
    key="final_result",
    value=add1.outputs.result,
)

wg.to_html()
wg.show()

--------------------------------------------------------------------------------
WorkGraph: context_gui_demo, PK: None, State: CREATED
--------------------------------------------------------------------------------
Tasks:
Name            PK    State
--------------  ----  -------
graph_ctx             PLANNED
get_x                 PLANNED
get_multiplier        PLANNED
add1                  PLANNED
store_result          PLANNED
--------------------------------------------------------------------------------

As shown in the GUI, the get_context and set_context tasks appear as visible nodes in the workflow graph. This makes the data flow through context explicit and traceable, unlike direct context access via wg.ctx.x which is invisible in the visualization.

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

State of WorkGraph         : FINISHED
Result of add1            : uuid: 590bde7f-bb27-4189-97db-0f30d631edca (pk: 985) value: 12

# Generate node graph from the AiiDA process
from aiida_workgraph.utils import generate_node_graph

generate_node_graph(wg.pk)

Note

If you pass data from one task to another task through the context, you may need to use wait to wait for

the data to be ready. See waiting_on.