Sockets

In WorkGraph, a Socket is a fundamental concept that defines the connection points for data to flow between tasks. Think of them as the inputs and outputs of each step in your workflow. They are analogous to AiiDA’s Ports but with an important extension: sockets are designed to manage not just the data type but also the links between tasks in the graph.

This guide will walk you through how to define, customize, and organize sockets for your tasks.

from aiida.manage import load_profile
from aiida_workgraph import task, WorkGraph, spec
from aiida import orm
from typing import Any

# Load the AiiDA profile to interact with the database
load_profile()

Profile<uuid='2c6af82ef2b241408a4c27c7c32d5490' name='presto'>

Automatic Socket Creation

The easiest way to create sockets is to let WorkGraph do it for you. When you decorate a Python function as a @task, WorkGraph automatically inspects its signature:

• Function arguments become input sockets.

• The return value becomes a single output socket named result.

Beside, there are also some built-in sockets from the WorkGraph, like _wait and _outputs.

@task()
def multiply(x, y):
    """A simple task to multiply two numbers."""
    return x * y


# Let's inspect the automatically generated sockets
wg = WorkGraph()
task1 = wg.add_task(multiply, x=3, y=4)
print('Input sockets: ', task1.get_input_names())
print('Output sockets: ', task1.get_output_names())

Input sockets:  ['x', 'y', 'metadata', 'function_data', 'process_label', 'function_inputs', '_wait']
Output sockets:  ['result', '_outputs', '_wait']

Customizing Output Sockets

When a function needs to return multiple outputs, you can define custom output sockets with the outputs argument in the @task.decorator. The task will map returned values to the named sockets in two ways:

• The function returns a dictionary: The keys of the returned dictionary must match the socket names.

@task
def add_and_subtract(x, y) -> spec.namespace(sum=Any, difference=Any):
    """Return the sum and difference of two numbers in a dict."""
    return {
        'sum': x + y,
        'difference': x - y,
    }


# Inspect the new input and output sockets
task2 = wg.add_task(add_and_subtract, x=3, y=4)
print('Input sockets: ', task2.get_input_names())
print('Output sockets: ', task2.get_output_names())

Input sockets:  ['x', 'y', 'metadata', 'function_data', 'process_label', 'function_inputs', '_wait']
Output sockets:  ['sum', 'difference', '_outputs', '_wait']

• The function returns a tuple: The elements of the tuple are mapped to the sockets in the order they are declared in the outputs definition.

Note

Be sure that the number of elements in the returned tuple matches the number of defined output sockets.

@task
def add_and_subtract(x, y) -> spec.namespace(sum=Any, difference=Any):
    """Return the sum and difference of two numbers as a tuple."""
    return x + y, x - y


# Inspect the new input and output sockets
task2 = wg.add_task(add_and_subtract, x=3, y=4)
print('Input sockets: ', task2.get_input_names())
print('Output sockets: ', task2.get_output_names())

Input sockets:  ['x', 'y', 'metadata', 'function_data', 'process_label', 'function_inputs', '_wait']
Output sockets:  ['sum', 'difference', '_outputs', '_wait']

The identifier is used to indicates the data type. The data type tell the code how to display the port in the GUI, validate the data, and serialize data into database. We use workgraph.Any for any data type.

Warning

The data validation feature is experimental. The API is subject to change in future releases. We welcome your feedback on its functionality.

For convenience, we also support using a list of strings as the output definition, which is equivalent to using workgraph.Any for each output socket.

@task(outputs=['sum', 'difference'])
def add_and_subtract(x, y):
    """This task returns both the sum and difference of two numbers."""
    return {'sum': x + y, 'difference': x - y}

Socket Types from Python Type Hints

To ensure data integrity and leverage AiiDA’s data provenance, sockets have types. WorkGraph can automatically assign AiiDA-compatible socket types based on standard Python type hints in your function signature. This is the recommended modern approach.

Python Type Hint	AiiDA Socket Type
int	orm.Int
float	orm.Float
str	orm.Str
bool	orm.Bool
list	orm.List
dict	orm.Dict

@task.calcfunction()
def add_typed(x: orm.Int, y: orm.Float) -> orm.Float:
    """A task with typed sockets to add an Int and a Float."""
    return orm.Float(x.value + y.value)


# When you call this task, WorkGraph expects AiiDA data types
task3 = wg.add_task(add_typed, x=orm.Int(3), y=orm.Int(4))
print('Input x: ', task3.inputs.x)

Input x:  SocketInt(name='x', value=uuid: a0c332e1-e929-4a77-a160-741decaa9f77 (unstored) value: 3)

Default Values as Socket Properties

A socket can have a property, which is a default value used when no data is linked to an input. The most Pythonic way to define a property is by using a default argument in your function definition.

@task
def power(base: int, exponent: int = 2):
    """Calculates base to the power of exponent, defaulting to square."""
    return base**exponent


# The `exponent` socket now has a default value of 2.
# We only need to provide the `base`.
with WorkGraph() as wg:
    outputs1 = power(base=3)
    # We can still override the default value by providing an input.
    outputs2 = power(base=3, exponent=3)
    wg.run()
    print(f'3 to the default power of 2 is: {outputs1.result.value}')
    print(f'3 to the power of 3 is: {outputs2.result.value}')

3 to the default power of 2 is: uuid: f3752e44-11d6-4298-9032-8901f3b4f2b5 (pk: 83) value: 9
3 to the power of 3 is: uuid: cc8e5a16-0e01-4607-9e78-d152450c727f (pk: 85) value: 27

Organizing Sockets with Namespaces

As workflows grow, you might have many related inputs or outputs. To keep them organized and avoid name clashes, you can group them into a namespace.

Simple Output Namespace

If a task has multiple outputs (as defined in the outputs decorator argument), WorkGraph automatically groups them into a SocketNamespace.

with WorkGraph('simple_namespace_example') as wg:
    # The `add_and_subtract` task from before has two outputs.
    # The `outputs` object becomes a namespace.
    outputs = add_and_subtract(x=10, y=4)
    wg.run()

    # You can access the results like attributes of an object:
    print('Accessing outputs from a namespace:')
    print(f'  Sum: {outputs.sum.value}')
    print(f'  Difference: {outputs.difference.value}')

Accessing outputs from a namespace:
  Sum: uuid: 2268c4dc-af62-4c67-b249-e259a907ce8c (pk: 90) value: 14
  Difference: uuid: 89e11f99-4e27-4648-9112-9eff4bec553e (pk: 91) value: 6

Nested Namespaces

For more complex data structures, you can define nested namespaces. This allows you to create a hierarchical organization for your sockets.

out = spec.namespace(
    normal=spec.namespace(sum=Any, product=Any),
    squared=spec.namespace(sum=Any, product=Any),
)


@task(outputs=out)
def advanced_math(x, y):
    """A task with a nested output structure."""
    return {
        'normal': {'sum': x + y, 'product': x * y},
        'squared': {'sum': x**2 + y**2, 'product': x**2 * y**2},
    }


with WorkGraph('nested_namespace_example') as wg:
    outputs = advanced_math(x=2, y=3)
    wg.run()
    print('\nAccessing outputs from a nested namespace:')
    print(f'  Normal sum: {outputs.normal.sum.value}')
    print(f'  Squared product: {outputs.squared.product.value}')
    assert outputs.normal.sum.value == 5
    assert outputs.squared.product.value == 36

Accessing outputs from a nested namespace:
  Normal sum: uuid: fe25de2e-9241-4980-b2d2-24360fd5958b (pk: 96) value: 5
  Squared product: uuid: eb3a4de4-0067-432e-80fd-0e2de70f4126 (pk: 99) value: 36

Dynamic namespaces

In many workflows, you might not know the exact number of outputs a task will generate until it runs. A dynamic namespace is a powerful feature in AiiDA-WorkGraph that allows a task to emit an arbitrary number of named output ports at runtime.

To create one, define a namespace socket and set its metadata to {“dynamic”: True}.

@task
def generate_squares(n: int) -> spec.namespace(squares=spec.dynamic(int)):
    """Generates a dynamic number of square values."""
    # The return dictionary must match the output socket names.
    # The value for the "squares" dynamic namespace is another dictionary,
    # where each key-value pair will become a separate AiiDA node.
    return {'squares': {f'n_{i}': i**2 for i in range(n)}}

This approach offers two significant advantages:

1. Fine-grained provenance: Each key-value pair (“n_0”: 0, “n_1”: 1, etc.) is stored as an individual AiiDA node. This gives you a complete and transparent graph, allowing you to trace the history and reuse of every single result.

2. Database integrity: It avoids serializing large, complex Python objects into a single, opaque file. Instead, each result is a native, queryable entry in the AiiDA database.

To consume these dynamic outputs, we can create a task that accepts an arbitrary number of keyword arguments.

@task
def sum_all(**data):
    """A task to sum all values from a dictionary of inputs."""
    # The **data syntax gathers all keyword arguments into a dictionary named 'data'.
    return sum(data.values())

Now, let’s build the WorkGraph.

with WorkGraph('dynamic_namespace_example') as wg:
    # The first task generates a dynamic set of outputs under the "squares" namespace.
    dynamic_outputs = generate_squares(n=4)

    # The entire "squares" namespace is then linked to the next task.
    total = sum_all(data=dynamic_outputs.squares)

    wg.run()

    # You can access the individual dynamic outputs
    print('\nIndividual dynamic outputs:')
    for i in range(4):
        print(f'  n_{i}: {dynamic_outputs.squares[f"n_{i}"].value}')

    print(f'\nSum of all dynamic outputs: {total.result.value}')
    assert total.result.value == 14

Individual dynamic outputs:
  n_0: uuid: e6a01a9e-f7ac-4b25-84e3-214fa17ea34c (pk: 103) value: 0
  n_1: uuid: 112e8fbc-aa8a-4bb6-9a5b-d885a84ccc9c (pk: 104) value: 1
  n_2: uuid: 0f247769-6c81-4698-93ea-d24abad104d8 (pk: 105) value: 4
  n_3: uuid: 3d95a88c-67bb-41d8-8c83-052babcf5160 (pk: 106) value: 9

Sum of all dynamic outputs: uuid: adb04098-40b8-4bb8-820a-1fe85b3b1c99 (pk: 108) value: 14

Important

Connecting sockets: why `data=…` is essential

When linking a dynamic namespace to a task that accepts keyword arguments (**kwargs), you must explicitly name the input argument. In our example:

# CORRECT: Connect the 'squares' output socket to the 'data' input socket.
sum_all(data=dynamic_outputs.squares)

This is because you are not passing values directly, but rather connecting the sockets of the graph. dynamic_outputs.squares is a reference to the output socket that will hold all the dynamic results. The task sum_all has a corresponding input socket named data (derived from the **data signature). The line above tells WorkGraph to wire the entire dynamic namespace from the output socket to this specific input socket.

Using Python’s argument unpacking syntax will fail:

# INCORRECT: This will raise an error.
sum_all(**dynamic_outputs.squares)

At the time the graph is built, dynamic_outputs.squares is a socket object, not a Python dictionary with concrete values. The ** operator attempts to unpack a dictionary immediately, which cannot be done on a socket reference that represents future results. You must name the argument to establish the connection for the engine to resolve later.