[train][checkpoint] Add checkpoint_upload_function to ray.train.report

[TimothySeah]

Summary

After #55637, ray.train.report will allow users to upload checkpoints from disk to remote storage asynchronously.

If they want to use framework-specific async checkpointing like torch.async_save, they can manage torch.async_save themselves and then call ray.train.report(…checkpoint_upload_mode=CheckpointUploadMode.NO_UPLOAD).

However, it would also be nice to allow ray.train.report to handle rate limiting and report ordering for framework-specific async checkpointing as well. This PR achieves this by exposing a checkpoint_upload_function argument that can replace the persist_current_checkpoint call.

Tests

Tested with torch.distributed's async_save in a workspace. The code looks like

storage_context = ray.train.get_context().get_storage()
checkpoint_path = f"s3://{storage_context.build_checkpoint_path_from_name(str(epoch))}"
storage_writer = S3StorageWriter(region='us-west-2', path=checkpoint_path)
model_dict, opt_dict = get_state_dict(model=model, optimizers=optimizer)
ckpt_ref = async_save(
    {"model": model_dict, "opt": opt_dict},
    storage_writer=storage_writer,
)

def wait_async_save(checkpoint, checkpoint_dir):
    t2 = time.time()
    ckpt_ref.result()
    return checkpoint

checkpoint = Checkpoint(checkpoint_path)
ray.train.report(
    metrics={"loss": 10, "accuracy": 0.5},
    checkpoint=checkpoint,
    checkpoint_upload_mode=CheckpointUploadMode.ASYNC,
    checkpoint_upload_function=wait_async_save,
)

and correctly returned a Result object with 3 checkpoints and their associated metrics. Note that the checkpoint path must be inside the "experiment path" which is configured in the RunConfig and accessed through build_checkpoint_path_from_name.

---

Note

Adds an optional checkpoint_upload_fn to ray.train.report (and internal plumbing) to override checkpoint upload/persistence, with corresponding tests.

• API:
  • ray.train.report(...) adds checkpoint_upload_fn to allow custom checkpoint upload logic.
• Internal Execution:
  • Thread checkpoint_upload_fn through TrainContext.report and _upload_checkpoint; if provided, use it instead of StorageContext.persist_current_checkpoint.
  • Propagate parameter via TrainFnUtils (abstract, distributed, local) and public API wrappers.
• Tests:
  • Add test_report_checkpoint_upload_fn validating custom upload function integrates with storage and result checkpoint loading.
  • Minor test adjustments (utilities imports, shutil usage).

^{Written by Cursor Bugbot for commit e21ce19. This will update automatically on new commits. Configure here.}

[justinvyu]

API looks good to me now!

[justinvyu]

📈