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 load_checkpoint(self, pid, tag=None):
"""Load a process from a persisted checkpoint by its process id.
:param pid: the process id of the :class:`plumpy.Process`
:param tag: optional checkpoint identifier to allow retrieving a specific sub checkpoint
:return: a bundle with the process state
:rtype: :class:`plumpy.Bundle`
:raises: :class:`plumpy.PersistenceError` Raised if there was a problem loading the checkpoint
"""
from aiida.common.exceptions import MultipleObjectsError, NotExistent
from aiida.orm import load_node
if tag is not None:
raise NotImplementedError('Checkpoint tags not supported yet')
try:
calculation = load_node(pid)
except (MultipleObjectsError, NotExistent):
raise plumpy.PersistenceError(
f'Failed to load the node for process<{pid}>: {traceback.format_exc()}'
)
checkpoint = calculation.checkpoint
if checkpoint is None:
raise plumpy.PersistenceError(
f'Calculation<{calculation.pk}> does not have a saved checkpoint'
)
try:
bundle = serialize.deserialize(checkpoint)
except Exception:
raise plumpy.PersistenceError(
f'Failed to load the checkpoint for process<{pid}>: {traceback.format_exc()}'
)
return bundle
def load_checkpoint(self, pid, tag=None):
"""
Load a process from a persisted checkpoint by its process id
:param pid: the process id of the :class:`plumpy.Process`
:param tag: optional checkpoint identifier to allow retrieving a specific sub checkpoint
:return: a bundle with the process state
:rtype: :class:`plumpy.Bundle`
:raises: :class:`plumpy.PersistenceError` Raised if there was a problem loading the checkpoint
"""
from aiida.orm import load_node
if tag is not None:
raise NotImplementedError('Checkpoint tags not supported yet')
calculation = load_node(pid)
checkpoint = calculation.checkpoint
if checkpoint is None:
raise plumpy.PersistenceError('Calculation<{}> does not have a saved checkpoint'.format(calculation.pk))
bundle = yaml.load(checkpoint)
return bundle
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