Auto resumption

[dakinggg]

What does this PR do?

Add support for autoresume with distributed training.

What issue(s) does this change relate to?

Fixes CO-1238
While fixing this issue, we discovered a race condition that is related to this, that is CO-1270

Manual testing

Manually tested a script that makes a trainer, trains it, makes another trainer with autoresume, and verifies that the params and run name are the same. This was tested on multiples gpus on one node.

Multi node does not currently work, but I also can't really test/debug it because of limited capacity on our cluster.

Before submitting

• Have you read the contributor guidelines?
• Is this change a documentation change or typo fix? If so, skip the rest of this checklist.
• Was this change discussed/approved in a GitHub issue first? It is much more likely to be merged if so.
• Did you update any related docs and document your change?
• Did you update any related tests and add any new tests related to your change? (see testing)
• Did you run the tests locally to make sure they pass?
• Did you run pre-commit on your change? (see the pre-commit section of prerequisites)

[mvpatel2000]

Thanks for fixing this!! Approving to unblock, left some really minor nits for conventions we've used in repo. Actually, I'm not sure we're consistent with these... but it'd be nice to do so 🤷

Discussed offline -- lets actually hold this PR until we can test multinode instead of doing it in two parts

[mvpatel2000]

I'll let Hanlin approve since it'd be nice to get two pairs of eyes on this. LGTM tho. Amazing job!

[hanlint]

LGTM, one suggestion to clean up the code a bit

[dakinggg]

@hanlint I changed the num_concurrent_uploads default to 1. I experimented a bit, and at least on the demo example, if you checkpoint enough to want multiple upload workers, it seems that you end up bottlenecked on network communication with s3 (streaming data in and sending checkpoints out). With 1 worker, you should only end up with an invalid symlink (pointing to one checkpoint before the actual latest) if the job dies between finishing uploading the checkpoint and uploading the symlink (which is tiny). This would be quite unlikely, and even this issue could be resolved by having save_overwrite=True when using autoresume.

[eracah]

Approving to unblock!