aiidateam / plumpy

A python workflows library that supports writing Processes with a well defined set of inputs and outputs that can be strung together.
https://plumpy.readthedocs.io
Other
8 stars 17 forks source link

ProcessListener added to Process are not persistent #273

Closed rikigigi closed 11 months ago

rikigigi commented 1 year ago

I would expect that after reloading a checkpoint, the listener that I added to the Process is persisted. This can be useful in the AiiDA daemon. For example, I may add a listener that connects to a remote URL to notify when the WorkChain has finished. I apologize if I have misunderstood the intended functionality.

Below you can find an example of code that does not load the listener after a checkpoint is loaded.

Thank you.

import plumpy

persister = plumpy.InMemoryPersister()

class TestListener(plumpy.ProcessListener):

    def on_process_excepted(self, process, reason):
        print('process excepted')

    def on_process_killed(self, process, reason):
        print('process killed')

    def on_process_finished(self, process, output):
        print('process finished')

class SimpleWorkChain(plumpy.WorkChain):
    @classmethod
    def define(cls, spec):
        super().define(spec)
        spec.outline(
            cls.step1,
            cls.step2,
        )

    def step1(self):
        print('step1')
        persister.save_checkpoint(self,'step1')

    def step2(self):
        print('step2')
        persister.save_checkpoint(self,'step2')

workchain = SimpleWorkChain()
workchain.add_process_listener(TestListener())
output = workchain.execute()

print('reload persister checkpoint:')
workchain_checkpoint = persister.load_checkpoint(workchain.pid, 'step1').unbundle()
workchain_checkpoint.execute()
sphuber commented 1 year ago

I had a look through the source code and it seems that the event listeners are simply never persisted, and so they are also not reconstructed when the process is reloaded from a checkpoint. This functionality is not actually used in aiida-core so that's why it has gone unnoticed. Would be happy to accept a PR. You essentially need to update the Process.save_instance_state and Process.load_instance_state to serialize and deserialize the listeners into the Savable.

rikigigi commented 11 months ago

fixed by https://github.com/aiidateam/plumpy/pull/277 and https://github.com/aiidateam/plumpy/pull/279

and tested in the combination

plumpy @ git+https://github.com/aiidateam/plumpy@01dc91dc7767ff56a76c4d62cddeda832243b69d
aiida-core @ git+https://github.com/aiidateam/aiida-core@8b8887e559e02eadac832a89f7012872040e1cbc