[Train] implement CheckpointStrategy

[matthewdeng]

Implements CheckpointStrategy to delete checkpoints from disk.

Why are these changes needed?

For a training function with many epochs, persisted checkpoints may lead to a high disk usage. The user should be able to define in their training run to clean up checkpoints. Specifically, they may want to keep N checkpoints, removing checkpoints based on:

1. Timestamp - remove old and outdated checkpoints.
2. A user-defined metric - remove "bad" or low-scoring checkpoints.

Example

This example is also included in the docs.

from ray import train
from ray.train import CheckpointStrategy, Trainer


def train_func():
    # first checkpoint
    train.save_checkpoint(loss=2)
    # second checkpoint
    train.save_checkpoint(loss=4)
    # third checkpoint
    train.save_checkpoint(loss=1)
    # fourth checkpoint
    train.save_checkpoint(loss=3)

# Keep the 2 checkpoints with the smallest "loss" value.
checkpoint_strategy = CheckpointStrategy(num_to_keep=2,
                                         checkpoint_score_attribute="loss",
                                         checkpoint_score_order="min")

trainer = Trainer(backend="torch", num_workers=2)
trainer.start()
trainer.run(train_func, checkpoint_strategy=checkpoint_strategy)
print(trainer.best_checkpoint_path)
# /home/ray_results/train_2021-09-01_12-00-00/run_001/checkpoints/checkpoint_000003
print(trainer.latest_checkpoint_dir)
# /home/ray_results/train_2021-09-01_12-00-00/run_001/checkpoints
print([checkpoint_path for checkpoint_path in trainer.latest_checkpoint_dir.iterdir()])
# [PosixPath('/home/ray_results/train_2021-09-01_12-00-00/run_001/checkpoints/checkpoint_000003'),
# PosixPath('/home/ray_results/train_2021-09-01_12-00-00/run_001/checkpoints/checkpoint_000001')]
trainer.shutdown()

Implementation

The implementation loosely follows that of Tune's checkpoint_manager.

1. Add a _top_persisted_checkpoints priority queue in CheckpointManager to keep track of the "best" checkpoints. Checkpoint paths are wrapped in a PersistedCheckpoint object.
2. The user can customize checkpoint retention logic in CheckpointStrategy. This will determine when and which checkpoints to delete. Note: If two checkpoints have the same priority, the original one will be kept and the latest one will never be persisted.
3. Store a single _best_persisted_checkpoint, which is exposed in Trainer.best_checkpoint_path. This replaces Trainer.last_checkpoint_path, as the most chronologically recent persistent checkpoint path may be already deleted based on the CheckpointStrategy. The default strategy uses timestamp, which would allow this field to effectively remain the same.
4. A default _timestamp attribute is added to checkpoints for the purpose of supporting timestamp based retention.

Related issue number

Checks

• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[amogkam]

Nice! Is there docs for this?

[amogkam]

Can we consolidate the implementation for latest_checkpoint_path and best_checkpoint_path? Currently best_checkpoint_path uses the PersistedCheckpoint object to get the path, but latest_checkpoint_path does path manipulation to get the path given the directory and file name. I think we can have both of these use the path stored in the PersistedCheckpoint object.

[amogkam]

Actually do we even need latest_checkpoint_path anymore?

[matthewdeng]

Great point, latest_checkpoint_path is actually just used for writing while best_checkpoint_path is used for reads. I renamed latest_checkpoint_path to next_checkpoint_path (with a minor change to increment the ID).

[amogkam]

Do we also want to expose a best_checkpoint object the same way we have for latest_checkpoint?

[matthewdeng]

Hmmm, I didn't think this was necessary right now and thought it might convolute the API. Is there a use-case where we'd need to expose this?

[amogkam]

LGTM! Just left some minor documentation points