Infinite recursion calling fetch_value() on returned AiiDA Nodes

[louisponet]

Hi,

I was trying to use a PyFunction to construct the inputs to another CalcJob. For this I need to be able to return preexisting AiiDA data nodes such as e.g. KPointsData. This leads to an infinite recursion when calling fetch_value().

MWE to reproduce:

import aiida_dynamic_workflows as flows
from aiida.orm import Int

cluster_env = flows.engine.execution_environment(
    "aiida-dev",   # conda environment
    "localhost")

@flows.step(returns=["v1"])
def test(nk: int) -> int:
    return Int(nk)

z = flows.engine.apply(test, nk=4)
t = aiida.engine.run(z.on(cluster_env))
t['return_values']['kpoints'].fetch_value()

[jbweston]

Thanks for the report.

Functions decorated with flows.step are not meant to return AiiDA datatypes directly. The function typically runs on an external computer, which does not necessarily have access to the Aiida database, and so it is not possible to instantiate nodes directly there (additionally the execution environment may not even have Aiida installed).

I would imagine that this behavior is due to aiida.orm.Int not playing well with (cloud)pickle.

[jbweston]

Can confirm:

[jbweston]

@louisponet IMO this may be an AiiDA issue (OTOH I don't think any guarantees are made about cloudpickling AiiDA dataypes)

[jbweston]

But this raises a good point that we should be able to control more carefully how datatypes returned from steps are encoded.

Currently if you

@flows.step(returns=["v1"])
def test(nk: int) -> int:
    return nk

z = flows.engine.apply(test, nk=4)
t = aiida.engine.run(z.on(cluster_env))

then the Data node produced will be of type PyRemoteData, which is not particularly informative.

Indeed this feature (encoding types more richly) is something that we have been discussing recently.

Over the coming weeks we will add issues/projects here as we migrate to GitHub; I'll be sure to ping you on any relevant issues.

[louisponet]

Right, understood. I wonder if we should try to come up with a sensible conversion between bare python objects to aiida nodes. That would be useful on many fronts, but also needs to be well specified so that it does not become messy.

I have done some explorational work with dict2node functions w.r.t going from json to aiida nodes.

Thanks! I'm looking forward to working more with this :D.

[jbweston]

I wonder if we should try to come up with a sensible conversion between bare python objects to aiida nodes.

Sounds like a good idea; additionally different datatypes may have a less opaque serialization format than pickle, and that could be pluggable too.

I think the main design decision is how signal to AiiDA what type was returned (recall that the Python function actually runs in a cluster environment with no access to AiiDA directly). We could use the function's return type annotation for this, however we would also need to define what should happen if the function decides not to conform to its type signature.

[louisponet]

Ah right, I see that's indeed not very straightforward. The solution you'd be looking for would not change anything to the current syntax etc right?
One could imagine things like annotations that signal what desired AiiDA node a return value would be converted to etc, but that would maybe not be the most elegant solution.

[jbweston]

I have opened aiidateam/aiida-core#5285 to track the reported bug upstream.

I'll leave this issue open here until that one is resolved (and further discussion about feature requests should happen in a separate issue).