[tune] Remove temporary checkpoint directories after restore

[krfricke]

Why are these changes needed?

FunctionTrainable.restore_from_object creates a temporary checkpoint directory.

This directory is kept around as we don't control how the user interacts with the checkpoint - they might load it several times, or no time at all.

Once a new checkpoint is tracked in the status reporter, there is no need to keep the temporary object around anymore.

In this PR, we add functionality to remove these temporary directories. Additionally we adjust the number of checkpoints to keep in pytorch_pbt_failure to 10 to reduce disk pressure in the release test. It looks like this lead to recent failures of the test. By removing the total number of checkpoints and fixing the issue with temporary directories we should see much less disk usage.

Related issue number

(Hopefully) closes #36561

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 :(

[krfricke]

In the running release test it looks like the temp checkpoints are removed correctly. However, the num_to_keep param doesn't seem to be working well:

~/ray_results/TorchTrainer_2023-07-06_17-33-05/TorchTrainer_d6a87_00000_0_lr=0.0010_2023-07-06_17-33-07$ ls -alp | grep checkpoint_0 | wc -l
206

This looks like a separate problem though. I'll address this in a follow-up

[justinvyu]

Thanks, I think this makes sense. I have one minor suggestion and some clarifications:

1. Just to confirm: We don't ever track the tmp directory anywhere (ex: in a checkpoint manager), so we don't run the risk of trying to restore from this temp checkpoint ever again, after the initial restore_from_object. Otherwise, it may be deleted already. Answer: Yes.
2. Possible to include some disk usage before/after plots in the PR description? Answer: Grafana is broken :(

[justinvyu]

LGTM. Waiting for tests to pass.