def main():
with kiwipy.connect(
'amqp://127.0.0.1') as communicator, tempfile.TemporaryDirectory(
) as tmpdir:
loop = asyncio.get_event_loop()
persister = plumpy.PicklePersister(tmpdir)
task_receiver = plumpy.ProcessLauncher(loop=loop, persister=persister)
def callback(*args, **kwargs):
fut = plumpy.create_task(functools.partial(task_receiver, *args,
**kwargs),
loop=loop)
return fut
communicator.add_task_subscriber(callback)
process_controller = plumpy.RemoteProcessThreadController(communicator)
future = process_controller.launch_process(DummyProcessWithOutput)
async def task():
result = await asyncio.wrap_future(future.result())
print(result)
loop.run_until_complete(task())
def persister():
_tmppath = tempfile.mkdtemp()
persister = plumpy.PicklePersister(_tmppath)
yield persister
shutil.rmtree(_tmppath)
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')
with tempfile.TemporaryDirectory() as directory:
persister = plumpy.PicklePersister(directory)
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_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()
with tempfile.TemporaryDirectory() as directory:
persister = plumpy.PicklePersister(directory)
persister.save_checkpoint(process)
bundle = persister.load_checkpoint(process.pid)
load_context = plumpy.LoadSaveContext(loop=self.loop)
recreated = bundle.unbundle(load_context)
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]
with tempfile.TemporaryDirectory() as directory:
persister = plumpy.PicklePersister(directory)
persister.save_checkpoint(process_a)
persister.save_checkpoint(process_b)
retrieved_checkpoints = persister.get_checkpoints()
self.assertSetEqual(set(retrieved_checkpoints), set(checkpoints))