async def run_async(workchain):
# run the original workchain until paused
await run_until_paused(workchain)
self.assertTrue(workchain.ctx.s1)
self.assertFalse(workchain.ctx.s2)
# Now bundle the workchain
bundle = plumpy.Bundle(workchain)
# Need to close the process before recreating a new instance
workchain.close()
# Load from saved state
workchain2 = bundle.unbundle()
self.assertTrue(workchain2.ctx.s1)
self.assertFalse(workchain2.ctx.s2)
# check bundling again creates the same saved state
bundle2 = plumpy.Bundle(workchain2)
self.assertDictEqual(bundle, bundle2)
# run the loaded workchain to completion
runner.schedule(workchain2)
workchain2.play()
await workchain2.future()
self.assertTrue(workchain2.ctx.s1)
self.assertTrue(workchain2.ctx.s2)
# ensure the original paused workchain future is finalised
# to avoid warnings
workchain.future().set_result(None)
def _check_round_trip(self, proc1):
bundle1 = plumpy.Bundle(proc1)
proc2 = bundle1.unbundle()
bundle2 = plumpy.Bundle(proc2)
self.assertEqual(proc1.pid, proc2.pid)
self.assertDictEqual(bundle1, bundle2)
def test_bundle_yaml(self):
bundle = plumpy.Bundle(Save1())
represent = yaml.dump({'bundle': bundle})
bundle_loaded = yaml.load(represent, Loader=yaml.Loader)['bundle']
self.assertIsInstance(bundle_loaded, plumpy.Bundle)
self.assertDictEqual(bundle_loaded, Save1().save())
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 test_running_save_instance_state(self):
nsync_comeback = SavePauseProc()
self.loop.add_callback(nsync_comeback.step_until_terminated)
yield test_utils.run_until_paused(nsync_comeback)
# Create a checkpoint
bundle = plumpy.Bundle(nsync_comeback)
self.assertListEqual([SavePauseProc.run.__name__],
nsync_comeback.steps_ran)
nsync_comeback.play()
yield nsync_comeback.future()
self.assertListEqual(
[SavePauseProc.run.__name__, SavePauseProc.step2.__name__],
nsync_comeback.steps_ran)
proc_unbundled = bundle.unbundle()
# At bundle time the Process was paused, the future of which will be persisted to the bundle.
# As a result the process, recreated from that bundle, will also be paused and will have to be played
proc_unbundled.play()
self.assertEqual(0, len(proc_unbundled.steps_ran))
yield proc_unbundled.step_until_terminated()
self.assertEqual([SavePauseProc.step2.__name__],
proc_unbundled.steps_ran)
def test_save_instance_state():
"""Test save instance's state."""
proc = test_processes.DummyProcess()
# Save the instance state
bundle = plumpy.Bundle(proc)
proc.close()
bundle.unbundle()
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(self):
process = DummyProcess()
saved_state = plumpy.Bundle(process)
process.close()
loaded_process = saved_state.unbundle()
run(loaded_process)
self.assertEqual(loaded_process.state, plumpy.ProcessState.FINISHED)
async def async_test():
await utils.run_until_paused(workchain)
self.assertTrue(workchain.ctx.s1)
self.assertFalse(workchain.ctx.s2)
# Now bundle the thing
bundle = plumpy.Bundle(workchain)
# Load from saved state
workchain2 = bundle.unbundle()
self.assertTrue(workchain2.ctx.s1)
self.assertFalse(workchain2.ctx.s2)
bundle2 = plumpy.Bundle(workchain2)
self.assertDictEqual(bundle, bundle2)
workchain.play()
await workchain.future()
self.assertTrue(workchain.ctx.s1)
self.assertTrue(workchain.ctx.s2)
def run_async(workchain):
yield run_until_paused(workchain)
self.assertTrue(workchain.ctx.s1)
self.assertFalse(workchain.ctx.s2)
# Now bundle the thing
bundle = plumpy.Bundle(workchain)
# Load from saved state
workchain2 = bundle.unbundle()
self.assertTrue(workchain2.ctx.s1)
self.assertFalse(workchain2.ctx.s2)
bundle2 = plumpy.Bundle(workchain2)
self.assertDictEqual(bundle, bundle2)
workchain.play()
yield workchain.future()
self.assertTrue(workchain.ctx.s1)
self.assertTrue(workchain.ctx.s2)
def run_async(workchain):
yield run_until_paused(workchain)
self.assertTrue(workchain.ctx.s1)
self.assertFalse(workchain.ctx.s2)
# Now bundle the thing
bundle = plumpy.Bundle(workchain)
# Need to close the process before recreating a new instance
workchain.close()
# Load from saved state
workchain2 = bundle.unbundle()
self.assertTrue(workchain2.ctx.s1)
self.assertFalse(workchain2.ctx.s2)
bundle2 = plumpy.Bundle(workchain2)
self.assertDictEqual(bundle, bundle2)
workchain.play()
yield workchain.future()
self.assertTrue(workchain.ctx.s1)
self.assertTrue(workchain.ctx.s2)
async def async_test():
await utils.run_until_waiting(proc)
saved_state = plumpy.Bundle(proc)
result = await proc.pause()
self.assertTrue(result)
self.assertTrue(proc.paused)
# Kill the process
proc.kill()
with self.assertRaises(plumpy.KilledError):
result = await proc.future()
async def async_test():
await utils.run_until_waiting(proc)
# Save the state of the process
saved_state = plumpy.Bundle(proc)
# Load a process from the saved state
loaded_proc = saved_state.unbundle()
self.assertEqual(loaded_proc.state, ProcessState.WAITING)
# Now resume it
loaded_proc.resume()
await loaded_proc.step_until_terminated()
self.assertEqual(loaded_proc.outputs, {'finished': True})
def test_restart_async(proc):
yield test_utils.run_until_waiting(proc)
# Save the state of the process
saved_state = plumpy.Bundle(proc)
# Load a process from the saved state
loaded_proc = saved_state.unbundle()
self.assertEqual(loaded_proc.state, ProcessState.WAITING)
# Now resume it
self.loop.add_callback(loaded_proc.step_until_terminated)
loaded_proc.resume()
yield loaded_proc.future()
self.assertEqual(loaded_proc.outputs, {'finished': True})
def run_async(proc):
yield test_utils.run_until_waiting(proc)
saved_state = plumpy.Bundle(proc)
result = yield proc.pause()
self.assertTrue(result)
self.assertTrue(proc.paused)
# Kill the process
proc.kill()
with self.assertRaises(plumpy.KilledError):
result = yield proc.future()
self.assertEqual(proc.state, ProcessState.KILLED)
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)
async def async_test():
await utils.run_until_paused(nsync_comeback)
# Create a checkpoint
bundle = plumpy.Bundle(nsync_comeback)
self.assertListEqual([SavePauseProc.run.__name__], nsync_comeback.steps_ran)
nsync_comeback.play()
await nsync_comeback.future()
self.assertListEqual([SavePauseProc.run.__name__, SavePauseProc.step2.__name__], nsync_comeback.steps_ran)
proc_unbundled = bundle.unbundle()
# At bundle time the Process was paused, the future of which will be persisted to the bundle.
# As a result the process, recreated from that bundle, will also be paused and will have to be played
proc_unbundled.play()
self.assertEqual(0, len(proc_unbundled.steps_ran))
await proc_unbundled.step_until_terminated()
self.assertEqual([SavePauseProc.step2.__name__], proc_unbundled.steps_ran)
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
def test_save_instance_state(self):
proc = test_processes.DummyProcess()
# Save the instance state
bundle = plumpy.Bundle(proc)
proc.close()
bundle.unbundle()
