[air] pyarrow.fs persistence (4/n): Introduce a simplified checkpoint manager

[justinvyu]

Why are these changes needed?

This PR condenses the functionality of ray.air._internal.checkpoint_manager into a few simplified classes:

1. Introduces a new _TrainingResult in place of the old _TrackedCheckpoint, which holds onto a (checkpoint, metrics) checkpoint result reported by the user.
2. The logic of _CheckpointManager is stripped down to only handle the top K heap of checkpoints. This includes deleting the bottom checkpoints.
   • Previously, the old _CheckpointManager also handled "committing" checkpoints to disk. This was mostly for in-memory checkpoints, which don't apply to the new persistence mode. This behavior has been removed.
   • The implementation of delete_fn is now standardized using the pyarrow fs.

This is a prerequisite cleanup for #37888, which shouldn't deal with the old _TrackedCheckpoint interface anymore. The old _TrackedCheckpoint holds onto the node that a checkpoint lives on + implements "commit" / "delete" and has a bunch of other functionality that doesn't apply anymore.

Future cleanup / considerations

• We should remove ray.air._internal.checkpoint_manager, ray.train._internal.checkpoint, and ray.tune.execution.checkpoint_manager.
• This checkpoint manager can also be moved from the driver to the trial process. (That is the direction we're trying to head with the experiment layering effort.)

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• 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 added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• 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 :(

[justinvyu]

This was adapted from ray/air/test_checkpoint_manager.py

[justinvyu]

This was adapted from ray/air/_internal/checkpoint_manager.py

[ericl]

Why not cleanup checkpoints right away instead of deferring?

[justinvyu]

This is behavior from before: we always keep the latest checkpoint around, even if it's not a top k checkpoint. Restoration from fault tolerance restores from the latest checkpoint at the moment.

[justinvyu]

@ericl Are you ok with keeping the existing functionality for now? This test walks through an example.

[ericl]

Hmm I'm more talking about the implementation. I think now's a good time to clean up some of this code. I believe we can replace this class with the following, which is much simpler and unlikely to be much slower:

class _CheckpointManager:

    def __init__(
        self,
        checkpoint_config: Optional[CheckpointConfig],
        latest_checkpoint_index: int = 0,
    ):
        self._checkpoint_config = checkpoint_config or CheckpointConfig()
        self._checkpoints: List[_TrackedCheckpoint)] = []
        self._latest_checkpoint_index = latest_checkpoint_index

    def register_checkpoint(self, tracked_checkpoint: _TrackedCheckpoint):
        self._checkpoints.append(tracked_checkpoint)

        if self._checkpoint_config.num_to_keep:
            # Could cache score as optimization if we wanted
            candidates = sorted(self._checkpoints[:-1], key=self._get_checkpoint_score)
            while candidates and len(self._checkpoints) > self._checkpoint_config.num_to_keep:
                checkpoint = candidates.pop(0).checkpoint
                _delete_fs_path(fs=checkpoint.filesystem, fs_path=checkpoint.path)

    @property
    def best_checkpoint(self) -> Optional[_TrackedCheckpoint]:
        if not self._checkpoints:
            return None
        return sorted(self._checkpoints, key=self._get_checkpoint_score)[-1]

    @property
    def latest_checkpoint(self) -> Optional[_TrackedCheckpoint]:
        if not self._checkpoints:
            return None
        return self._checkpoints[-1]

[justinvyu]

Ok, simplified it in a slightly different way that doesn't need to sort the list. PTAL

There was also a slight behavioral change from above: the K+1th checkpoint wouldn't be deleted if the best checkpoint is always excluded from the candidates.

[justinvyu]

@ericl Any thoughts on a new name for _TrackedCheckpoint? I'm using it in the next PR and aliasing it to _NewTrackedCheckpoint, but I think this name is confusing, and it's clearer if we just choose a new name for it.

I'm thinking of changing it to TrainingResult, since it's an object that holds the (metrics, checkpoint) reported from the user:

• We can use this in place of ray.train._internal.session.TrainingResult (which is pretty much the same thing).
• This is in line with unifying the 2 session reports.
• This is also in line with combining the train and save steps in the Tune control loop.

[ericl]

+1, was thinking the name is a bit confusing too

…

On Tue, Aug 1, 2023, 3:14 PM Justin Yu ***@***.***> wrote: @ericl <https://github.com/ericl> Any thoughts on a new name for _TrackedCheckpoint? I'm using it in the next PR and aliasing it to _NewTrackedCheckpoint, but I think this name is confusing, and it's clearer if we just choose a new name for it. I'm thinking of changing it to TrainingResult, since it's an object that holds the (metrics, checkpoint) reported from the user: - We can use this in place of ray.train._internal.session.TrainingResult (which is pretty much the same thing). - This is in line with unifying the 2 session reports. - This is also in line with combining the train and save steps in the Tune control loop. — Reply to this email directly, view it on GitHub <#37962 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAADUSUPROFZ4UJMRCMO7JLXTF5VVANCNFSM6AAAAAA27BA3ZU> . You are receiving this because you were mentioned.Message ID: ***@***.***>