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