[air] Add Checkpoint.as_directory() for efficient checkpoint fs processing

[krfricke]

Why are these changes needed?

This PR adds a Checkpoint_as_directory() context manager that either returns the local path (if checkpoint is already a directory) or a temporary directory path containing the checkpoint data, which is cleaned up after use. The path should be considered as a read-only source for loading data from the checkpoint.

A common use case for processing checkpoint data is to convert it into a directory with Checkpoint.to_directory() and then do some read-only processing (e.g. restoring a ML model).

This process has two flaws: First, to_directory() creates a temporary directory that has to be explicitly cleaned up by the user after use. Secondly, if the checkpoint is already a directory checkpoint, it is copied over, which is inefficient for large checkpoints (e.g. huggingface models) and then even more prone to unwanted side effects if not cleaned up properly.

With this context manager that effectively returns a directory that is to be used as a read-only data source, we can avoid manual cleaning up and unnecessary data copies (or avoid internal inspection as e.g. in https://github.com/ray-project/ray/pull/23876/files#diff-47db2f054ca359879f77306e7b054dd8b780aab994961e3b4911330ae15eeae3R57-R60)

See also discussion in https://github.com/ray-project/ray/pull/23850/files#r850036905

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]

Should we change the current usages of to_directory with a tmpfile to use this context manager instead?