def check_process_against_snapshots(loop, proc_class, snapshots):
"""
Take the series of snapshots from a Process that executed and run it
forward from each one. Check that the subsequent snapshots match.
This will only check up to the STARTED state because from that state back
they should of course differ.
Return True if they match, False otherwise.
:param loop: The event loop
:param proc_class: The process class to check
:type proc_class: :class:`Process`
:param snapshots: The snapshots taken from from an execution of that
process
:return: True if snapshots match False otherwise
:rtype: bool
"""
for i, bundle in zip(list(range(0, len(snapshots))), snapshots):
loaded = bundle.unbundle(plumpy.LoadSaveContext(loop=loop))
saver = ProcessSaver(loaded)
saver.capture()
# Now check going backwards until running that the saved states match
j = 1
while True:
if j >= min(len(snapshots), len(saver.snapshots)):
break
compare_dictionaries(
snapshots[-j], saver.snapshots[-j], snapshots[-j], saver.snapshots[-j], exclude={'exception'}
)
j += 1
return True
def save_checkpoint(self, process, tag=None):
"""Persist a Process instance.
:param process: :class:`aiida.engine.Process`
:param tag: optional checkpoint identifier to allow distinguishing multiple checkpoints for the same process
:raises: :class:`plumpy.PersistenceError` Raised if there was a problem saving the checkpoint
"""
LOGGER.debug('Persisting process<%d>', process.pid)
if tag is not None:
raise NotImplementedError('Checkpoint tags not supported yet')
try:
bundle = plumpy.Bundle(
process, plumpy.LoadSaveContext(loader=get_object_loader()))
except ImportError:
# Couldn't create the bundle
raise plumpy.PersistenceError(
f"Failed to create a bundle for '{process}': {traceback.format_exc()}"
)
try:
process.node.set_checkpoint(serialize.serialize(bundle))
except Exception:
raise plumpy.PersistenceError(
f"Failed to store a checkpoint for '{process}': {traceback.format_exc()}"
)
return bundle
def create_daemon_runner(self, loop=None):
"""Create and return a new daemon runner.
This is used by workers when the daemon is running and in testing.
:param loop: the (optional) tornado event loop to use
:type loop: :class:`tornado.ioloop.IOLoop`
:return: a runner configured to work in the daemon configuration
:rtype: :class:`aiida.engine.runners.Runner`
"""
import plumpy
from aiida.engine import persistence
from aiida.manage.external import rmq
runner = self.create_runner(rmq_submit=True, loop=loop)
runner_loop = runner.loop
# Listen for incoming launch requests
task_receiver = rmq.ProcessLauncher(
loop=runner_loop,
persister=self.get_persister(),
load_context=plumpy.LoadSaveContext(runner=runner),
loader=persistence.get_object_loader())
def callback(*args, **kwargs):
return plumpy.create_task(functools.partial(
task_receiver, *args, **kwargs),
loop=runner_loop)
runner.communicator.add_task_subscriber(callback)
return runner
def _save_round_trip_with_loader(self, savable):
"""
Do a round trip:
1) Save `savables` state
2) Recreate from the saved state
3) Save the state of the recreated `Savable`
4) Compare the two saved states (they should match)
:type savable: :class:`plumpy.Savable`
"""
object_loader = plumpy.get_object_loader()
saved_state1 = savable.save(plumpy.LoadSaveContext(object_loader))
loaded = savable.recreate_from(saved_state1)
saved_state2 = loaded.save(plumpy.LoadSaveContext(object_loader))
saved_state3 = loaded.save()
self.assertDictEqual(saved_state1, saved_state2)
self.assertNotEqual(saved_state1, saved_state3)
def test_bundle_load_context(self):
""" Check that the loop from the load context is used """
proc = test_utils.DummyProcess(loop=self.loop)
bundle = plumpy.Bundle(proc)
loop2 = plumpy.new_event_loop()
proc2 = bundle.unbundle(plumpy.LoadSaveContext(loop=loop2))
self.assertIs(loop2, proc2.loop())
def test_save_load_roundtrip(self):
"""
Test the plumpy.PicklePersister by taking a dummpy process, saving a checkpoint
and recreating it from the same checkpoint
"""
process = ProcessWithCheckpoint()
persister = plumpy.InMemoryPersister()
persister.save_checkpoint(process)
bundle = persister.load_checkpoint(process.pid)
load_context = plumpy.LoadSaveContext(loop=self.loop)
recreated = bundle.unbundle(load_context)
def test_save_load_roundtrip(self):
"""
Test the plumpy.PicklePersister by taking a dummpy process, saving a checkpoint
and recreating it from the same checkpoint
"""
loop = asyncio.get_event_loop()
process = ProcessWithCheckpoint()
with tempfile.TemporaryDirectory() as directory:
persister = plumpy.PicklePersister(directory)
persister.save_checkpoint(process)
bundle = persister.load_checkpoint(process.pid)
load_context = plumpy.LoadSaveContext(loop=loop)
recreated = bundle.unbundle(load_context)
def test_continue(self):
persister = plumpy.InMemoryPersister()
load_context = plumpy.LoadSaveContext(loop=self.loop)
launcher = plumpy.ProcessLauncher(persister=persister,
load_context=load_context)
process = test_utils.DummyProcess(loop=self.loop)
pid = process.pid
persister.save_checkpoint(process)
del process
process = None
result = yield launcher._continue(
None,
**plumpy.create_continue_body(pid)[process_comms.TASK_ARGS])
self.assertEqual(test_utils.DummyProcess.EXPECTED_OUTPUTS, result)
async def test_continue():
persister = plumpy.InMemoryPersister()
load_context = plumpy.LoadSaveContext()
launcher = plumpy.ProcessLauncher(persister=persister,
load_context=load_context)
process = utils.DummyProcess()
pid = process.pid
persister.save_checkpoint(process)
del process
process = None
result = await launcher._continue(
None,
**plumpy.create_continue_body(pid)[process_comms.TASK_ARGS])
assert result == utils.DummyProcess.EXPECTED_OUTPUTS
async def async_test():
await utils.run_until_waiting(proc)
saved_state = plumpy.Bundle(proc)
# Run the process to the end
proc.resume()
result1 = await proc.future()
# Load from saved state and run again
loader = plumpy.get_object_loader()
proc2 = saved_state.unbundle(plumpy.LoadSaveContext(loader))
asyncio.ensure_future(proc2.step_until_terminated())
proc2.resume()
result2 = await proc2.future()
# Check results match
self.assertEqual(result1, result2)
def run_async(proc):
yield test_utils.run_until_waiting(proc)
saved_state = plumpy.Bundle(proc)
# Run the process to the end
proc.resume()
result1 = yield proc.future()
# Load from saved state and run again
proc2 = saved_state.unbundle(
plumpy.LoadSaveContext(loop=self.loop))
self.loop.add_callback(proc2.step_until_terminated)
proc2.resume()
result2 = yield proc2.future()
# Check results match
self.assertEqual(result1, result2)
def _setup_rmq(self,
url,
prefix=None,
task_prefetch_count=None,
testing_mode=False):
super(DaemonRunner, self)._setup_rmq(url, prefix, task_prefetch_count,
testing_mode)
# Create a context for loading new processes
load_context = plumpy.LoadSaveContext(runner=self)
# Listen for incoming launch requests
task_receiver = rmq.ProcessLauncher(
loop=self.loop,
persister=self.persister,
load_context=load_context,
loader=persistence.get_object_loader())
self.communicator.add_task_subscriber(task_receiver)
def test_wait_save_continue(self):
""" Test that process saved while in WAITING state restarts correctly when loaded """
proc = test_utils.WaitForSignalProcess()
self.loop.add_callback(proc.step_until_terminated)
yield test_utils.run_until_waiting(proc)
saved_state = plumpy.Bundle(proc)
# Run the process to the end
proc.resume()
result1 = yield proc.future()
# Load from saved state and run again
proc2 = saved_state.unbundle(plumpy.LoadSaveContext(loop=self.loop))
self.loop.add_callback(proc2.step_until_terminated)
proc2.resume()
result2 = yield proc2.future()
# Check results match
self.assertEqual(result1, result2)
def test_save_future(self):
"""
test `SavableFuture` is initialized with the event loop of process
"""
loop = asyncio.new_event_loop()
nsync_comeback = SavePauseProc(loop=loop)
bundle = plumpy.Bundle(nsync_comeback)
# if loop is not specified will use the default asyncio loop
proc_unbundled = bundle.unbundle(plumpy.LoadSaveContext(loop=loop))
async def async_test():
await utils.run_until_paused(proc_unbundled)
# here the future should be a SavableFuture in process loop
proc_unbundled.play()
await proc_unbundled.future()
self.assertListEqual([SavePauseProc.run.__name__, SavePauseProc.step2.__name__], proc_unbundled.steps_ran)
loop.create_task(proc_unbundled.step_until_terminated())
loop.run_until_complete(async_test())
def save_checkpoint(self, process, tag=None):
"""
Persist a Process instance
:param process: :class:`aiida.work.Process`
:param tag: optional checkpoint identifier to allow distinguishing multiple checkpoints for the same process
:raises: :class:`plumpy.PersistenceError` Raised if there was a problem saving the checkpoint
"""
LOGGER.debug('Persisting process<%d>', process.pid)
if tag is not None:
raise NotImplementedError('Checkpoint tags not supported yet')
try:
bundle = plumpy.Bundle(process, plumpy.LoadSaveContext(loader=get_object_loader()))
except ValueError:
# Couldn't create the bundle
raise plumpy.PersistenceError("Failed to create a bundle for '{}':{}".format(
process, traceback.format_exc()))
else:
calc = process.calc
calc.set_checkpoint(yaml.dump(bundle))
return bundle
