def _process_persistent_checkpoint(self, checkpoint: _TrackedCheckpoint):
self.add_tune_checkpoint_id(checkpoint.dir_or_data)
# If inside a Tune Trainable, then checkpoint with Tune.
with tune.checkpoint_dir(
step=self._latest_checkpoint_id) as checkpoint_dir:
path = Path(checkpoint_dir)
checkpoint.commit(path)
return super()._process_persistent_checkpoint(checkpoint)
def handle_checkpoint(self, checkpoint: _TrackedCheckpoint):
# Set checkpoint ID
checkpoint.id = checkpoint.id or self._latest_checkpoint_id
self._latest_checkpoint_id += 1
if checkpoint.storage_mode == CheckpointStorage.MEMORY:
self._replace_latest_memory_checkpoint(checkpoint)
else:
assert checkpoint.storage_mode == CheckpointStorage.PERSISTENT
assert (self._checkpoint_strategy.num_to_keep is None
or self._checkpoint_strategy.num_to_keep > 0)
self._process_persistent_checkpoint(checkpoint)
def _process_persistent_checkpoint(self, checkpoint: _TrackedCheckpoint):
checkpoint.dir_or_data[PREPROCESSOR_KEY] = self.preprocessor
super(_DataParallelCheckpointManager, self)._process_persistent_checkpoint(
checkpoint=checkpoint
)