sphuber / aiida-shell

AiiDA plugin that makes running shell commands easy.
MIT License
14 stars 7 forks source link

`PickledData`: Set `recurse=True` for `dill.dumps` #83

Closed sphuber closed 7 months ago

sphuber commented 7 months ago

Fixes #82

Fix InvalidOperation raised by aiida-core when pickling

To support inline custom parser functions being submitted to and executed by the daemon, they are pickled using the dill library through the PickledData plugin. For certain cases, the script would hit the exception InvalidOperation raised by aiida-core in aiida.orm.entities.Entity.__getstate__ which is called when the instance is attempted to be pickled by the code, which is not supported for AiiDA's ORM entities.

The immediate cause seems to be that by default, the entire global dictionary of the Python interpreter is pickled by dill and so if that contains an AiiDA entity, the exception is hit. The entity does not even have to be part of the parser function that is being pickled but can just be part of the global scope. For example, the following would hit the exception

from aiida_shell import launch_shell_job
from aiida.orm import Int

def parser(self, dirpath):
    pass

some_node = Int(1)

results, node = launch_shell_job(
    'echo',
    arguments='some output',
    parser=parser,
)

Merely defining the node Int and assigning it to a variable would be sufficient to cause the exception to be raised. Removing it from the scope, either by commenting out the line or deleting the variable from memory before calling launch_shell_job would fix the problem.

Passing the option recurse=True to dill.dumps appears to solve the problem. Although it is not quite understood how and why this works from the description of the option in dill's source code:

If *recurse=True*, then objects referred to in the global dictionary
are recursively traced and pickled, instead of the default behavior
of attempting to store the entire global dictionary. This is needed
for functions defined via *exec()*.

Weirdly enough, the behavior could not be reproduced in a unit test run through pytest. Maybe the global dictionary is handled differently in this environment. Unfortunately therefore no unit test could be added.