custom output parsing with running on specific computer leads to pickling error
bilke opened this issue · 11 comments
bilke commented
First thanks for this helpful plugin!
When using custom output parsing in combination with running on a specific computer:
results, node = launch_shell_job(
"echo",
arguments="some output",
parser=parser,
metadata={"options": {"computer": load_computer('some-computer')}}, # <-- Added this line to the example
)
print(results["string"].value)I get the following error:
.../site-packages/aiida/orm/entities.py", line 223, in __getstate__
raise InvalidOperation('pickling of AiiDA ORM instances is not supported.')
aiida.common.exceptions.InvalidOperation: pickling of AiiDA ORM instances is not supported.
Is this expected? Can I do anything about it?
sphuber commented
Thanks for the report @bilke . I cannot reproduce this problem. Could you please report what version of aiida-shell and aiida-core you have installed?
bilke commented
➜ pip list
Package Version
------------------ --------
aiida-core 2.5.1
aiida-shell 0.6.0
...I think I have tracked it down a bit more, this works:
...
metadata={"options": {"computer": load_computer('some-computer')}}
....This does not work:
some_computer = load_computer("some-computer")
...
metadata={"options": {"computer": some-computer}}
....And also this does not work:
some_computer = load_computer("some-computer")
...
metadata={"options": {"computer": load_computer("some-computer")}}
....And this does not work too (ie. doing a second shell job):
results, node = launch_shell_job(
"echo",
arguments="some output",
parser=parser,
metadata={"options": {"computer": load_computer("some-computer")}},
)
print(results["string"].value)
results, node = launch_shell_job(
"echo",
arguments="some output",
parser=parser,
metadata={"options": {"computer": load_computer("some-computer")}},
)
print(results["string"].value)sphuber commented
Could you please share the complete example script and the full stack trace of the exception you get?
bilke commented
This is the full script:
from aiida.orm import load_computer
from aiida_shell import launch_shell_job
envinf4 = load_computer("envinf4")
def parser(self, dirpath):
from aiida.orm import Str
return {"string": Str((dirpath / "stdout").read_text().strip())}
results, node = launch_shell_job(
"echo",
arguments="some output",
parser=parser,
metadata={"options": {"computer": envinf4}},
)
print(results["string"].value)And here is the stack trace:
➜ verdi run test.shell.py
Traceback (most recent call last):
File "/my-workdir/.direnv/python-3.11/bin/verdi", line 8, in <module>
sys.exit(verdi())
^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/click/core.py", line 1157, in __call__
return self.main(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/click/core.py", line 1078, in main
rv = self.invoke(ctx)
^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/click/core.py", line 1688, in invoke
return _process_result(sub_ctx.command.invoke(sub_ctx))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/click/core.py", line 1434, in invoke
return ctx.invoke(self.callback, **ctx.params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/click/core.py", line 783, in invoke
return __callback(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/cmdline/utils/decorators.py", line 80, in wrapper
return wrapped(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/cmdline/commands/cmd_run.py", line 119, in run
exec(compile(handle.read(), str(filepath), 'exec', dont_inherit=True), globals_dict)
File "test.shell.py", line 13, in <module>
results, node = launch_shell_job(
^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida_shell/launch.py", line 84, in launch_shell_job
results, node = launch.run_get_node(ShellJob, **inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/launch.py", line 64, in run_get_node
return runner.run_get_node(process, inputs, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/runners.py", line 288, in run_get_node
result, node = self._run(process, inputs, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/runners.py", line 242, in _run
process_inited = self.instantiate_process(process, **inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/runners.py", line 171, in instantiate_process
return instantiate_process(self, process, **inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/utils.py", line 83, in instantiate_process
process = process_class(runner=runner, inputs=inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/plumpy/base/state_machine.py", line 193, in __call__
inst = super().__call__(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/processes/calcjobs/calcjob.py", line 191, in __init__
super().__init__(*args, **kwargs)
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/processes/process.py", line 168, in __init__
inputs=self.spec().inputs.serialize(inputs),
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/processes/ports.py", line 270, in serialize
result[name] = port.serialize(value)
^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/processes/ports.py", line 126, in serialize
return self._serializer(value)
^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida_shell/calculations/shell.py", line 125, in serialize_parser
return PickledData(value)
^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida_shell/data/pickled.py", line 41, in __init__
pickled = self.get_pickler()(obj)
^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 280, in dumps
dump(obj, file, protocol, byref, fmode, recurse, **kwds)#, strictio)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 252, in dump
Pickler(file, protocol, **_kwds).dump(obj)
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 420, in dump
StockPickler.dump(self, obj)
File "/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 487, in dump
self.save(obj)
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 414, in save
StockPickler.save(self, obj, save_persistent_id)
File "/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 560, in save
f(self, obj) # Call unbound method with explicit self
^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 1985, in save_function
_save_with_postproc(pickler, (_create_function, (
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 1112, in _save_with_postproc
pickler._batch_setitems(iter(source.items()))
File "/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 998, in _batch_setitems
save(v)
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 414, in save
StockPickler.save(self, obj, save_persistent_id)
File "/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 578, in save
rv = reduce(self.proto)
^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/orm/entities.py", line 223, in __getstate__
raise InvalidOperation('pickling of AiiDA ORM instances is not supported.')
aiida.common.exceptions.InvalidOperation: pickling of AiiDA ORM instances is not supported.
sphuber commented
Thanks a lot. I can now reproduce it. I have tracked down what the problem is, but I cannot quite understand why it is happening or how to fix it. Essentially, whenever the scope contains any variables that are AiiDA ORM instances, i.e. a Computer instance or Node, the pickling of the parser function tries to pickle that variable as well, and pickling of AiiDA instances is not supported and so raises.
In your example, you assign the return value of load_computer (which is an instance of aiida.orm.computers.Computer and cannot be pickled) to the variable computer. When the code now tries to pickle the parser function, somehow that variable gets included in the scope, triggering the exception. If you do not assign load_computer to a variable, but directly add it to the metadata.options dictionary, the variable is not in the scope and the code works.
Another example to showcase the problem:
from aiida_shell import launch_shell_job
from aiida.orm import load_computer, Int
def parser(self, dirpath):
from aiida.orm import Str
return {'string': Str((dirpath / 'stdout').read_text().strip())}
not_used = Int(1)
results, node = launch_shell_job(
'echo',
arguments='some output',
parser=parser,
metadata={'options': {'computer': load_computer('localhost')}}
)Note how just assigning an AiiDA node Int(1) to some variable not_used, this will cause the code to fail. Just because Int(1) is now in the interpreter's scope somehow.
I am using dill as a more powerful pickling library as the built-in pickle module, and I am not sure why it is doing this. Or if this is even standard behavior for pickling. So I am also not sure how to fix this, or whether it can even be fixed.
As a pretty shitty workaround for the time being, you should just be sure not to assign any AiiDA ORM instances to variables.
bilke commented
Thanks for your explanation! But please note that this also does not work (having two launch_shell_job()-calls with load_computer() (and having no variable assignment):
from aiida.orm import load_computer
from aiida_shell import launch_shell_job
def parser(self, dirpath):
from aiida.orm import Str
return {"string": Str((dirpath / "stdout").read_text().strip())}
results, node = launch_shell_job(
"echo",
arguments="some output",
parser=parser,
metadata={"options": {"computer": load_computer("envinf4")}},
)
print(results["string"].value)
results, node = launch_shell_job(
"echo",
arguments="some output",
parser=parser,
metadata={"options": {"computer": load_computer("envinf4")}},
)
print(results["string"].value)Outputs:
➜ verdi run test.shell.py
some output
Traceback (most recent call last):
File "/my-workdir/.direnv/python-3.11/bin/verdi", line 8, in <module>
sys.exit(verdi())
^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/click/core.py", line 1157, in __call__
return self.main(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/click/core.py", line 1078, in main
rv = self.invoke(ctx)
^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/click/core.py", line 1688, in invoke
return _process_result(sub_ctx.command.invoke(sub_ctx))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/click/core.py", line 1434, in invoke
return ctx.invoke(self.callback, **ctx.params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/click/core.py", line 783, in invoke
return __callback(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/cmdline/utils/decorators.py", line 80, in wrapper
return wrapped(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/cmdline/commands/cmd_run.py", line 119, in run
exec(compile(handle.read(), str(filepath), 'exec', dont_inherit=True), globals_dict)
File "test.shell.py", line 21, in <module>
results, node = launch_shell_job(
^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida_shell/launch.py", line 84, in launch_shell_job
results, node = launch.run_get_node(ShellJob, **inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/launch.py", line 64, in run_get_node
return runner.run_get_node(process, inputs, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/runners.py", line 288, in run_get_node
result, node = self._run(process, inputs, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/runners.py", line 242, in _run
process_inited = self.instantiate_process(process, **inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/runners.py", line 171, in instantiate_process
return instantiate_process(self, process, **inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/utils.py", line 83, in instantiate_process
process = process_class(runner=runner, inputs=inputs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/plumpy/base/state_machine.py", line 193, in __call__
inst = super().__call__(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/processes/calcjobs/calcjob.py", line 191, in __init__
super().__init__(*args, **kwargs)
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/processes/process.py", line 168, in __init__
inputs=self.spec().inputs.serialize(inputs),
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/processes/ports.py", line 270, in serialize
result[name] = port.serialize(value)
^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/engine/processes/ports.py", line 126, in serialize
return self._serializer(value)
^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida_shell/calculations/shell.py", line 125, in serialize_parser
return PickledData(value)
^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida_shell/data/pickled.py", line 41, in __init__
pickled = self.get_pickler()(obj)
^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 280, in dumps
dump(obj, file, protocol, byref, fmode, recurse, **kwds)#, strictio)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 252, in dump
Pickler(file, protocol, **_kwds).dump(obj)
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 420, in dump
StockPickler.dump(self, obj)
File "/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 487, in dump
self.save(obj)
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 414, in save
StockPickler.save(self, obj, save_persistent_id)
File "/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 560, in save
f(self, obj) # Call unbound method with explicit self
^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 1985, in save_function
_save_with_postproc(pickler, (_create_function, (
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 1112, in _save_with_postproc
pickler._batch_setitems(iter(source.items()))
File "/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 998, in _batch_setitems
save(v)
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 414, in save
StockPickler.save(self, obj, save_persistent_id)
File "/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 560, in save
f(self, obj) # Call unbound method with explicit self
^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 1217, in save_module_dict
StockPickler.save_dict(pickler, obj)
File "/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 972, in save_dict
self._batch_setitems(obj.items())
File "/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 998, in _batch_setitems
save(v)
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/dill/_dill.py", line 414, in save
StockPickler.save(self, obj, save_persistent_id)
File "/opt/homebrew/Cellar/python@3.11/3.11.8/Frameworks/Python.framework/Versions/3.11/lib/python3.11/pickle.py", line 578, in save
rv = reduce(self.proto)
^^^^^^^^^^^^^^^^^^
File "/my-workdir/.direnv/python-3.11/lib/python3.11/site-packages/aiida/orm/entities.py", line 223, in __getstate__
raise InvalidOperation('pickling of AiiDA ORM instances is not supported.')
aiida.common.exceptions.InvalidOperation: pickling of AiiDA ORM instances is not supported.
sphuber commented
Thanks for your explanation! But please note that this also does not work (having two
launch_shell_job()-calls withload_computer()(and having no variable assignment):
Well, there are variable assignments, albeit a bit hidden 😄 The node variable is an AiiDA ORM instance, and the results is a dictionary that also contains them.
If you were to add
del node
del resultsjust before invoking the second launch_shell_job, it would work again.
sphuber commented
I think I may have found a bit of a fix in this branch. Could you please clone the repo and check out the branch and install it:
git clone https://github.com/sphuber/aiida-shell
cd aiida-shell
git checkout fix/082/dill-pickling-recurse
pip install -e .
and then run your example script again.
bilke commented
Thank you so much Sebastiaan, it works!
sphuber commented
I need to clean up the change a bit and add some tests. Then I will merge it and make a release soon so you can just install from PyPI again without having to clone.
sphuber commented
FYI: just released v0.7.0 on PyPI with the fix: https://pypi.org/project/aiida-shell/