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