[air] pyarrow.fs persistence: Remove dependence on rank 0 worker checkpoint reporting

[justinvyu]

Why are these changes needed?

This PR removes the need for the rank 0 worker to report a checkpoint in order for a checkpoint to be tracked by Train/Tune.

Before (ray <= 2.6):

def train_fn_per_worker(config):
    ...
    tmpdir = tempfile.mkdtemp()
    if session.get_world_rank() == 0:
        # write to tmpdir
        # global rank 0 MUST report. otherwise it's as if you didn't checkpoint
        checkpoint = Checkpoint.from_directory(...)
    else:
        # create an "empty" checkpoint...
        # otherwise, if you just reported None, we throw an error
        # even worse, if you report a dict checkpoint here... unknown territory
        checkpoint = Checkpoint.from_directory(...)
    session.report(..., checkpoint)

After:

def train_fn_per_worker(config):
    ...
    # ANY combination of workers can report a checkpoint
    if train.get_context().get_world_rank() in [2, 4, 6]:
        with tempfile.TemporaryDirectory() as tempdir:
            # write to tmpdir
            train.report(metrics, Checkpoint.from_directory(tempdir))
    else:
        train.report(metrics)

Note: the reported metrics are still pulled from the global rank 0 worker (same behavior as before). This PR does not remove that restriction.

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

[ericl]

Is results synchronously gathered from all workers (i.e., a barrier?)

[justinvyu]

Yep. Results are returned by the training iterator via a ray.get([worker.get_next_result.remote() for worker in workers])

[ericl]

This change makes sense to me, but could you elaborate on the motivation/use cases for this? Are there new edge cases this would expose, or is it a pretty straightforward win-win?

[ericl]

Should we validate/assert all the checkpoints are pointing to the same path?

[justinvyu]

I think it's purely win-win here. The motivating use case is with deepspeed checkpoints:

• Lightning saves multiple ranks' deepspeed checkpoint shards to a single directory and it's hard to de-aggregate, so our integration should just have each local rank 0 worker upload the directory instead of every worker doing it. The global rank 0 worker may not be the local rank 0 worker, and without this PR, that would mean the checkpoint doesn't get tracked by Tune internally if global rank 0 doesn't report.

I think this is better for UX as well -- we don't enforce which worker can report a checkpoint.

@woshiyyya can elaborate more here.

[justinvyu]

One consequence of this: if users end up using our new framework checkpoints, it'll lose the class here and always end up as a generic checkpoint.

[woshiyyya]

Thanks @justinvyu! This change is mostly for distributed checkpointing. In deepspeed, there should be one worker on each node to report checkpoints, not necessarily be global rank 0.

This change gives us more general and flexible interface, as long as one of the workers report checkpoint, we can always correctly track it.