Support creating tasks from builders

[GeigerJ2]

Currently, only nested dictionaries can be used to populate the parameters of a task. Support setting also from a builder (similar to submit in AiiDA, where you can provide either the class and the input parameters, or the builder).

As noted by @rikigigi, who had to first convert the builder data to a nested dictionary.

[superstar54]

Hi @GeigerJ2 , thanks for taking care of this issue. Setting task inputs from the builder is now supported in the set_from_builder method, but creating tasks from the builder has not been implemented yet.

I think we want to support syntax similar to the following code:

from aiida.calculations.arithmetic.add import ArithmeticAddCalculation
from aiida_workgraph import WorkGraph
from aiida import orm

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

wg = WorkGraph("test_set_inputs_from_builder")
# use builder to create task
add1 = wg.add_task(builder, name="add1")

[GeigerJ2]

Sorry for the delay, been busy with other tasks. Starting to look into this now again.

[GeigerJ2]

Just writing down some thoughts while working on this. On one hand, this achieves successful execution of the WG for the ArithmeticAddCalculation:

from aiida.calculations.arithmetic.add import ArithmeticAddCalculation
from aiida_workgraph import WorkGraph
from aiida import orm, load_profile
from aiida_workgraph.utils import get_dict_from_builder

load_profile()

builder = ArithmeticAddCalculation.get_builder()

builder.code = orm.load_code('add@localhost')
builder.x = orm.Int(2)
builder.y = orm.Int(3)

builder_dict = get_dict_from_builder(builder)

wg = WorkGraph('test-builder')
wg.add_task(builder.process_class, name="add_builder", **builder_dict)

wg.run()

While that's good to know, we want to take the necessary steps to make it work inside the implementation such that the user doesn't have to worry, and can directly just pass the ProcessBuilder, as written by @superstar54.

Now, when I check the implementation, I see that WorkGraph.add_task() calls self.tasks._new() implemented in TaskCollection, so one could put the code there, and re-use build_task_from_callable function from decorator.py, like so:

        if isinstance(identifier, ProcessBuilder):
            from aiida_workgraph.utils import get_dict_from_builder
            kwargs = {**kwargs, **get_dict_from_builder(identifier)}
            identifier = build_task_from_callable(identifier.process_class)
        return super()._new(identifier, name, uuid, **kwargs)

(see here).

rather than adding some additional function (which is something I also considered, e.g., a build_task_from_builder function).

This seems to me like a pretty easy and straightforward implementation. Will now test it also with more complex entities, such as the MultiplyAddWorkchain. And I still need to think about how one would handle the linking of input/output Sockets from other tasks, but I'm somewhat hoping this will also work out of the box, due to the re-use of existing functionality. Feel free to share any thoughts at this point already, @superstar54. Thanks!

[superstar54]

Hi @GeigerJ2 , LGTM, please go ahead.

And I still need to think about how one would handle the linking of input/output Sockets from other tasks, but I'm somewhat hoping this will also work out of the box,

yes, this should work. But the user should be cuation that one can not pass socket to the builder itself directly, i.e., the following code will raise an error:

builder.y = another_task.outputs.result

They always need to add the task first, and then link the socket. Maybe you could mention this in the docs.

[GeigerJ2]

Thanks, @superstar54, for getting back here so quickly. I'll finalize this, add tests, and documentation one of those days 🔥
I think the issue you mentioned makes sense, as the builder.<something> AiiDA core typical syntax anyway doesn't allow for the functionality of linking non-existing entities (i.e., Sockets), so one has to go the WorkGraph route anyway.