Failed to de-serialize the local function

[superstar54]

from aiida_workgraph import WorkGraph, task
import pathlib
from aiida import orm, load_profile

import pathlib

load_profile()

@task.graph_builder(outputs = [{"name": "eigvals", "from": "cat_task.eigvals"}])
def query_and_diag(matrix_pk: orm.Int):
    from aiida_shell.data import PickledData
    wg = WorkGraph()

    # In an .ipynb at least defining it here works 
    def myparser(dirpath: pathlib.Path) -> dict[str, orm.Data]:
        x = 1
        return {'eigvals': orm.Int(x)}

    wg.add_task("ShellJob", name="cat_task",
                        command="cat",
                        arguments=[],
                        parser=myparser,
                        outputs=[],
                        parser_outputs=[{"name": "eigvals"}],
                    )
    return wg
i = 5
wg = WorkGraph("processing_data")
query_and_diag_task = wg.add_task(query_and_diag, name=f"query_and_diag_pk{i}", matrix_pk=orm.Int(i))
wg.run()

It gave this error about "!!python/name:__main__.myparser ''\n"

In workgraph, we use aiida.orm.utils.serialize.serialize (yaml serialize) to serialize the input data of a workgraph and save it to the node.base.extras. However, the yaml seralizer can not handle the local function properly. If the function is defined locally, it typically has a name like <function __main__.outer_function.<locals>.myparser>. However, the YAML serialization system is not designed to handle local functions correctly, as shown below:

In [20]: wg.tasks['cat_task'].inputs['parser'].value
Out[20]: <function __main__.query_and_diag.<locals>.myparser(dirpath: pathlib.Path) -> dict[str, aiida.orm.nodes.data.data.Data]>

In [21]: serialize(wg.tasks['cat_task'].inputs['parser'].value)
Out[21]: "!!python/name:__main__.myparser ''\n"

After serialization, it change it to __main__

• Loss of Context: Since the function is local, it’s not available directly in the __main__ module, but the serializer still tries to serialize it as if it were a top-level function. This leads to a mismatch during deserialization, where Python cannot find myparser in __main__.

Solutions

use pickle to serialize it and then use yaml, we need to use yaml because we want to save it into node.base.extrase using the aiida serialize