search

sphuber / aiida-shell

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

custom output parsing with running on specific computer leads to pickling error #82

Closed bilke closed 6 months ago

bilke commented 6 months ago

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 6 months ago

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 6 months ago 
➜ 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 6 months ago

Could you please share the complete example script and the full stack trace of the exception you get?

bilke commented 6 months ago

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 6 months ago

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 6 months ago

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 6 months ago

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):

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 results

just before invoking the second launch_shell_job, it would work again.

sphuber commented 6 months ago

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 6 months ago

Thank you so much Sebastiaan, it works!

sphuber commented 6 months ago

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 6 months ago

FYI: just released v0.7.0 on PyPI with the fix: https://pypi.org/project/aiida-shell/