def test_delete_process_checkpoints(self): """ """ process_a = ProcessWithCheckpoint() process_b = ProcessWithCheckpoint() checkpoint_a1 = plumpy.PersistedCheckpoint(process_a.pid, '1') checkpoint_a2 = plumpy.PersistedCheckpoint(process_a.pid, '2') checkpoint_b1 = plumpy.PersistedCheckpoint(process_b.pid, '1') checkpoint_b2 = plumpy.PersistedCheckpoint(process_b.pid, '2') persister = plumpy.InMemoryPersister() persister.save_checkpoint(process_a, tag='1') persister.save_checkpoint(process_a, tag='2') persister.save_checkpoint(process_b, tag='1') persister.save_checkpoint(process_b, tag='2') checkpoints = [checkpoint_a1, checkpoint_a2] retrieved_checkpoints = persister.get_process_checkpoints( process_a.pid) self.assertSetEqual(set(retrieved_checkpoints), set(checkpoints)) persister.delete_process_checkpoints(process_a.pid) checkpoints = [] retrieved_checkpoints = persister.get_process_checkpoints( process_a.pid) self.assertSetEqual(set(retrieved_checkpoints), set(checkpoints))
def test_loader_is_used(self): """ Make sure that the provided class loader is used by the process launcher """ loader = CustomObjectLoader() proc = Process() persister = plumpy.InMemoryPersister(loader=loader) persister.save_checkpoint(proc) launcher = plumpy.ProcessLauncher(persister=persister, loader=loader) continue_task = plumpy.create_continue_body(proc.pid) yield launcher._continue(None, **continue_task[process_comms.TASK_ARGS])
async def test_loader_is_used(): """ Make sure that the provided class loader is used by the process launcher """ loader = CustomObjectLoader() proc = Process() persister = plumpy.InMemoryPersister(loader=loader) persister.save_checkpoint(proc) launcher = plumpy.ProcessLauncher(persister=persister, loader=loader) continue_task = plumpy.create_continue_body(proc.pid) result = await launcher._continue(None, **continue_task[process_comms.TASK_ARGS]) assert result == utils.DummyProcess.EXPECTED_OUTPUTS
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_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
def test_get_checkpoints_without_tags(self): """ """ process_a = ProcessWithCheckpoint() process_b = ProcessWithCheckpoint() checkpoint_a = plumpy.PersistedCheckpoint(process_a.pid, None) checkpoint_b = plumpy.PersistedCheckpoint(process_b.pid, None) checkpoints = [checkpoint_a, checkpoint_b] persister = plumpy.InMemoryPersister() persister.save_checkpoint(process_a) persister.save_checkpoint(process_b) retrieved_checkpoints = persister.get_checkpoints() self.assertSetEqual(set(retrieved_checkpoints), set(checkpoints))
def test_get_checkpoints_with_tags(self): """ """ process_a = ProcessWithCheckpoint() process_b = ProcessWithCheckpoint() tag_a = 'tag_a' tag_b = 'tag_b' checkpoint_a = plumpy.PersistedCheckpoint(process_a.pid, tag_a) checkpoint_b = plumpy.PersistedCheckpoint(process_b.pid, tag_b) # import pdb; pdb.set_trace() checkpoints = [checkpoint_a, checkpoint_b] persister = plumpy.InMemoryPersister() persister.save_checkpoint(process_a, tag=tag_a) persister.save_checkpoint(process_b, tag=tag_b) retrieved_checkpoints = persister.get_checkpoints() self.assertSetEqual(set(retrieved_checkpoints), set(checkpoints))