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