How to link dynamic ports

[ahkole]

Is there any recommended procedure for linking dynamic ports in a namespace? For example, say that I have a workgraph wg1 which has dynamic output namespace scf_outputs and I want to link the port scf_outputs.zero to the input of another task (I know zero will exist in the namespace for the particular case I am trying to construct). Using scf_outputs.zero does not work, because the names of the dynamic ports are not in the OutputSpec of wg1. Right now I am using a workfunction to still link them,

from aiida.engine import ExitCode
from aiida_workgraph import task


@task.workfunction
def get_dynamic_socket(
        name: str,
        **inputs,
):
    try:
        aiida_name = name.value
    except AttributeError:
        aiida_name = name
    try:
        result = inputs
        for k in aiida_name.split('.'):
            result = result[k]
        return result
    except KeyError:
        return ExitCode(300, f'No socket {aiida_name} in dynamic namespace')

This seems to work, but I was wondering if there is a recommended, perhaps better way of linking these kind of dynamic ports?

[ahkole]

Maybe an additional question. Would my approach also work with task instead of task.workfunction? And do you still keep the data provenance of linking the correct output Data node as the input Data node of the next task?

[superstar54]

You can define fixed fields in the dynamic namespace.

# out1 and out2 always exist
dynamic(int, out1=int, out2=int])

You can define the output as optional if it only exists in a particular case:

dynamic(t.Any, out1=t.Annotated[t.Any, meta(required=False)])

Here is an example:

from aiida_workgraph import task, dynamic, meta
import typing as t
from aiida import load_profile

load_profile()

@task
def generate_dynamic_outputs_with_fixed_field(
    n: int,
) -> t.Annotated[
    dict,
    dynamic(t.Any, known_output=t.Annotated[t.Any, meta(required=False)]),
]:
    """Generate a dict of square numbers with a fixed field 'known_output'."""
    results = {f'square_{i}': i**2 for i in range(n)}
    results['known_output'] = n
    return results

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


@task.graph
def SquareNumbersGenerator(n: int):
    out1 = generate_dynamic_outputs_with_fixed_field(n=n)
    return add(out1.known_output, 2).result


wg = SquareNumbersGenerator.build(n=5)
results = wg.run()
print("result: ", results)

The doc misses this part 😢

[ahkole]

Hi, thanks for the suggestions. In my case, the dynamic port is not necessarily always present, but its name is taken from the name of a dynamic input port (so since I am supplying this name in the input I know it will be present in the output). The examples are therefore not directly applicable to my case, but it's good to know they exist.

[superstar54]

Could you give a simple example to show the tasks and links?

[ahkole]

Here is a simplified example that is representative of what I'm trying to do,

@task.graph
def run_scf(
    scf_inputs: t.Annotated[
        dict,
        tSiestaBaseWorkChain.inputs,
    ],
) -> t.Annotated[
    dict,
    namespace(
        scf_outputs=dynamic(tSiestaBaseWorkChain.outputs),
    )
]:
    # Run scf
    ports = {}
    for k, v in scf_inputs.items():
        ports[k] = tSiestaBaseWorkChain(**v)

    return {'scf_outputs': ports}


# Link output port to new task
with WorkGraph('test') as wg:
    task1 = run_scf.build(
        scf_inputs = {'zero': scf_inputs},
    )
    task2 = tSiestaBaseWorkChain(parent_calc_folder=task1.scf_outputs.zero.remote_folder, **scf_inputs2)

Here the link is from task1.scf_outputs.zero.remote_folder (which should exist at runtime) to the input of task2. This example is of course greatly simplified for illustration but it gives the general idea where I am using a label from a dynamic input port to label a dynamic output port. The above approach does not work because the label zero is not in the OutputSpec of run_scf, so I am trying to fix it with the utility function get_dynamic_socket I showed above (which does work but gives a meaningless/superfluous node in the provenance graph).