[Train] Support Out of Band Checkpoints#63161
Conversation
Signed-off-by: Mark Towers <mark@anyscale.com>
Signed-off-by: Mark Towers <mark@anyscale.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for out-of-band checkpoints in Ray Train v2, allowing checkpoints to be stored outside the primary experiment directory. Key changes include updating the _TrainingResultState schema, implementing logic to identify and reconstruct filesystems for external checkpoints, and modifying validation routines to handle out-of-band data gracefully. Review feedback highlights the need for safer deletion policies for out-of-band checkpoints to avoid unintended data loss and suggests more robust error handling when encountering unsupported filesystem types during state restoration.
I am having trouble creating individual review comments. Click here to see my feedback.
python/ray/train/v2/_internal/execution/checkpoint/checkpoint_manager.py (94-102)
As noted in the PR description, there's a question of how to handle deletion of out-of-band checkpoints when num_to_keep is used. Automatically deleting data from user-specified, out-of-band locations could be risky and lead to unintended data loss.
A safer approach would be to avoid deleting out-of-band checkpoints automatically. Instead, Ray Train could log a warning indicating that a checkpoint has been marked for deletion but requires manual cleanup because it's out-of-band. This would give users full control over their data in external storage.
This function _is_in_band could be used within _save_state_and_delete_old_checkpoints to implement this safer behavior.
python/ray/train/v2/_internal/execution/checkpoint/checkpoint_manager.py (142-146)
The code currently accesses _TYPE_NAME_FS_CLS with a key directly. If state.checkpoint_filesystem_type contains a value that is not in _TYPE_NAME_FS_CLS (e.g., for a less common or custom filesystem), this will raise a KeyError and cause the training resumption to fail. It would be more robust to handle this case gracefully by checking if the key exists and raising a more informative error message if it doesn't.
fs_factory = _TYPE_NAME_FS_CLS.get(state.checkpoint_filesystem_type)
if not fs_factory:
raise CheckpointManagerInitializationError(
f"Unsupported filesystem type for out-of-band checkpoint: "
f"'{state.checkpoint_filesystem_type}'. Supported types are: "
f"{list(_TYPE_NAME_FS_CLS.keys())}"
)
return _TrainingResult(
checkpoint=Checkpoint(state.checkpoint_dir_name, filesystem=fs_factory()),
metrics=state.metrics,
)
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Signed-off-by: Mark Towers <mark@anyscale.com>
Signed-off-by: Mark Towers <mark@anyscale.com>
Signed-off-by: Mark Towers <mark@anyscale.com>
| "s3": lambda: pyarrow.fs.S3FileSystem(), | ||
| "gcs": lambda: pyarrow.fs.GcsFileSystem(), | ||
| "hdfs": lambda: pyarrow.fs.HadoopFileSystem("missing-host"), | ||
| "abfs": lambda: pyarrow.fs.AzureFileSystem("missing-host"), |
There was a problem hiding this comment.
HDFS/ABFS filesystem reconstruction crashes on construction
Medium Severity
The _TYPE_NAME_FS_CLS entries for "hdfs" and "abfs" use pyarrow.fs.HadoopFileSystem("missing-host") and pyarrow.fs.AzureFileSystem("missing-host"). HadoopFileSystem attempts to connect to the HDFS namenode at construction time, so providing an invalid host will raise an immediate connection error. When _get_training_result_from_state calls fs() at line 135, this crashes the entire state restoration in _load_state, preventing resumption of any run that ever had an HDFS or Azure out-of-band checkpoint—even if other checkpoints in the same run are healthy.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit 49d3860. Configure here.
There was a problem hiding this comment.
This is intentional
TimothySeah
left a comment
TimothySeah
left a comment
There was a problem hiding this comment.
I reviewed everything except the unit test, which I will review on my next pass.
Can you clarify what Side note: Using relative paths allows user copied a training result's data include the checkpoint data to a new directory and restore it later. For in-band checkpoints, we keep this feature while for out of band checkpoints, we assume that the checkpoints are in a long-term storage that won't be moved. means?
| # Out-of-band: checkpoint_dir_name holds the full path; | ||
| # reconstruct the filesystem from type_name. | ||
| if state.checkpoint_filesystem_type != "local": | ||
| logger.warning( |
There was a problem hiding this comment.
Will users see this warning even in the "happy" out of band path i.e. user saves checkpoint with recoverable filesystem? I wonder if there's an easy way to only warn when the user actually encounters a problem e.g. they fail to download the checkpoint.
| continue | ||
|
|
||
| if is_out_of_band: | ||
| logger.warning( |
There was a problem hiding this comment.
Similar comment here and everywhere else we are warning - I only want users to see warnings when they actually encounter something that fails. Loading an out of band checkpoint into self._checkpoint_results is fine, but actually downloading that with to_directory should surface the real error.
There was a problem hiding this comment.
Agreed, I was a tad worried about the console spam that this causes.
One option is to add a keyword only argument to Checkpoint that specifies it was created with an out of band filesystem. Then apply a try / except to to_directory that can add this debug information.
This could reduce the number of warnings that we're creating.
Thoughts?
There was a problem hiding this comment.
Do we need to add a new argument to Checkpoint? Ideally we should avoid adding new arguments - even if they are backwards compatible by default - and whether a checkpoint is in band or out of band should be invisible to users. I think if to_directory fails the error message can just mention the "out of band different filesystem" case as a potential cause.
One of the features of the checkpoint manager that I didn't understand intuitively is that it saves the checkpoint dir names as the relative path to the storage path rather than the absolute path. This was an unexpected design decision that I post-hoc rationalised as allowing users to copy the experiment directory to a new directory, users can resume or load the results with the new storage path. If we used absolute paths then this wouldn't be possible. If this is an incorrect understand of why this decision was taken then I would be fascinated to understand its true intentions |
Signed-off-by: Mark Towers <mark@anyscale.com>
There was a problem hiding this comment.
One of the features of the checkpoint manager that I didn't understand intuitively is that it saves the checkpoint dir names as the relative path to the storage path rather than the absolute path. This was an unexpected design decision that I post-hoc rationalised as allowing users to copy the experiment directory to a new directory, users can resume or load the results with the new storage path. If we used absolute paths then this wouldn't be possible.
If this is an incorrect understand of why this decision was taken then I would be fascinated to understand its true intentions
Ah gotcha thanks for clarifying. Tbh I'm not sure what the original motivation was, but my mental model is that in the past the CheckpointManager only managed in band (i.e. all with the same filesystem and directory prefix) checkpoints so storing any information that included the filesystem and directory prefix was unnecessary.
| continue | ||
|
|
||
| if is_out_of_band: | ||
| logger.warning( |
There was a problem hiding this comment.
Do we need to add a new argument to Checkpoint? Ideally we should avoid adding new arguments - even if they are backwards compatible by default - and whether a checkpoint is in band or out of band should be invisible to users. I think if to_directory fails the error message can just mention the "out of band different filesystem" case as a potential cause.
Possibly, the place it would make most sense to insert the error as an extra exception case in ray/python/ray/train/_checkpoint.py Line 212 in 85c03be This is just a preference of more logging or better error message. |
I think minimizing API changes and proactive logging in favor of raising informative errors when something actually goes wrong would be the least intrusive/confusing to users |
Signed-off-by: Mark Towers <mark@anyscale.com>
I think we are agreed on the current implementation then? |
Signed-off-by: Mark Towers <mark@anyscale.com>
Signed-off-by: Mark Towers <mark@anyscale.com>
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
There are 2 total unresolved issues (including 1 from previous review).
Reviewed by Cursor Bugbot for commit d3d633e. Configure here.
| ckpt_path = Path(checkpoint_path) | ||
|
|
||
| if ckpt_path.is_relative_to(exp_path): | ||
| return str(ckpt_path.relative_to(exp_path)) |
There was a problem hiding this comment.
Path uses OS-specific separators instead of POSIX format
Medium Severity
extract_checkpoint_dir_name_from_path returns str(ckpt_path.relative_to(exp_path)) which uses OS-specific path separators. On Windows this produces backslashes (e.g., "trial\\checkpoint"), but pyarrow filesystem paths always use forward slashes. The old code used string slicing that preserved the original forward-slash format. The result is stored in JSON state and later joined with Path(...).as_posix() on restore — if the state file is created on Windows, the backslash becomes part of the filename literal on Linux, corrupting the path. Using .as_posix() instead of str() would maintain consistency with all other path construction in this file.
Reviewed by Cursor Bugbot for commit d3d633e. Configure here.
TimothySeah
left a comment
TimothySeah
left a comment
There was a problem hiding this comment.
I think minimizing API changes and proactive logging in favor of raising informative errors when something actually goes wrong would be the least intrusive/confusing to users
I think we are agreed on the current implementation then?
Can you do a train run with an out of band checkpoint and show me the logs? Basically I think it should be something like:
- report it with NO_UPLOAD: no logs from uploading (no upload performed) or serializing into checkpoint manager state
- recover it from new run +
get_checkpoint. No logs from deserializing from checkpointmanager. - try to download it with
to_directoryorwithcontextmanager. Doesn't work and only raiseValueErrorthen.
Iiuc, right now this might still produce logs for the first 2 steps.
I also looked at the unit tests and gave some feedback - I think this should be good to go after 1 final pass.
| CreateBucketConfiguration={"LocationConstraint": region}, | ||
| ) | ||
| # S3 with default upload | ||
| s3.put_object(Bucket=bucket_name, Key="epoch-6/results.txt", Body="6") |
There was a problem hiding this comment.
Wait what does this test? Iiuc, report will upload a checkpoint to the local directory, and this s3.put_object isn't tracked at all?
There was a problem hiding this comment.
The checkpoint is originally stored on S3 but copied into the local experiment directory.
I don't test that all checkpoint files exist, just the checkpoint exists. I've added this testing that the local checkpoint contents exists.
| assert isinstance(checkpoint.filesystem, pyarrow.fs.S3FileSystem) | ||
|
|
||
| reported_checkpoints = ray.train.get_all_reported_checkpoints() | ||
| assert len(reported_checkpoints) == 8 |
There was a problem hiding this comment.
Should we also test that to_directory or with works from some checkpoints but raises a helpful error from others?
There was a problem hiding this comment.
Yes, added to the tests now
| # Verify the snapshot left on S3 is loadable and points to the | ||
| # correct (in-band/out-of-band) checkpoint locations. | ||
| restored_result = Result.from_path(str(in_band_uri / "test-s3-in-out-of-band")) | ||
| assert len(restored_result.best_checkpoints) == 7 |
There was a problem hiding this comment.
same here: consider calling to_directory/with from these checkpoints and verifying that some work and some error as expected
There was a problem hiding this comment.
Oh and we can test that the pattern we are recommending to users in the error message (making a new checkpoint like Checkpoint(path=<old checkpoint's path>, filesystem=<new filesystem that they just made>) works.
There was a problem hiding this comment.
Yes, added to the tests to confirm the checkpoint's contents
Signed-off-by: Mark Towers <mark@anyscale.com>
Signed-off-by: Mark Towers <mark@anyscale.com>
|
@TimothySeah Here the full logs from You can see at the end, the checkpoint resumption that for the two out of band s3 checkpoints we fail to confirm that it exists. Note: The RuntimeError is shown twice for some reason which isn't expected. Not sure why |
There was a problem hiding this comment.
Can we just force the user to report a "pointer checkpoint" instead? Then, we still keep track of the "in-band" checkpoint and the user just needs to manage loading the path to the out of band checkpoint themselves.
# Instead of "NO_UPLOAD", force the user to report an in-band pointer checkpoint to the out of band checkpoint
torch.distributed.checkpoint.save(..., "s3://a/b/c")
with TemporaryDirectory() as temp_ckpt_dir:
pointer_dict = {"uri": "s3://a/b/c"}
json.dump(pointer_dict, temp_ckpt_dir)
ray.train.report(..., checkpoint=pointer_dict, checkpoint_upload_mode=SYNC)
# Force custom_upload_fn to return in-band checkpoint.
# The contract of custom_upload_fn is that it must return an in-band checkpoint.
def custom_upload_fn(checkpoint_source, checkpoint_destination):
do_out_of_band_checkpoint()
pointer_dict = {"uri": "s3://a/b/c"}
write_pointer_checkpoint_to_destination(pointer_dict)If users need us to clean up checkpoints ("top k checkpoints"), we can add a callback hook that fires when a checkpoint gets evicted from the top k checkpoints and they can add the out of band cleanup themselves.
I'm thinking we can tighten the API contract a bit to avoid adding complexity handling out of band vs in band, different filesystem vs same filesystem. Also would like to avoid the reverse lookup of filesystem string name -> filesystem.
I agree we have a choice with this PR, either open the API to fully support in and out of band checkpoints, or restrict the API to only support in-band checkpoints with a documentation like your solution for users that need out of band checkpoint. Out of band checkpointing does feel like an edge by users to support, however callback being necessary for in-band only support is off putting to me. To me, I think we should support out of band checkpoint as the API needs to expand minimally to support them with the biggest difficulty coming from different out of band filesystems which we could disallow. This would restrict the API to only support in and out of band checkpoints on the same filesystem. We have three options
@justinvyu Any other considerations? If we want to support out of band checkpointing, this is the best solution IMO, if we want to disallow out of band checkpoint all together, then we can do this in a separate PR. |
I was proposing callback only being needed for out of band support. We would trigger the hook right before deleting the "in-band pointer checkpoint" and then the user can clean up the out of band checkpoint themselves. I like option 1 the most. I think detecting the same filesystem is a bit tricky due to the wide variety of s3 clients people could be using. Then, our "same filesystem out of band" download with |
|
Closing, in favor of an implementation that restricts the API to only support in-band checkpoints and raise an error when out of band checkpoints are used |
…nd checkpoints (#63645) ## Description Currently `ray.train.report` expects that all checkpoints are saved within the trainer's `storage_path`. However, with `CheckpointUploadMode.NO_UPLOAD` or a custom `checkpoint_upload_fn` allows users to specify paths outside the storage path. This causes a problem on resumption for two reasons, we only store the checkpoints path relative to the storage path (which wouldn't be valid for out of band checkpoints) and we resume from the storage path's filesystem (which wouldn't be valid for checkpoints stored on a different filesystem). Therefore, we have three options to either restrict the API to 1. Only save in-band checkpoints and raise errors on all others 2. Support in and out of band checkpoints on the same filesystems but not support checkpoints on a different filesystem 3. Support all in and out of band checkpoints on any filesystem. (This was implemented in #63161) The challenge is that for non-local file systems (gcp, s3, etc) often passwords, env variables, etc are necessary to config a filesystem to read the data. As pyarrow doesn't have `pyarrow.fs.FileSystem.to_uri()` we can't serialise the filesystem as a string. It is possible to serialising the filesystem with pickle but there are obvious issue with this. As a result, this PR restricts the API to only support in-band checkpoint to avoid several problems with out of band checkpoints (option 1). To access the storage context, its possible for users to do `ray.train.get_context().get_storage().experiment_fs_path` and `ray.train.get_context().get_storage().storage_filesystem` --------- Signed-off-by: Mark Towers <mark@anyscale.com> Co-authored-by: Mark Towers <mark@anyscale.com>
3 participants
Description
Currently
ray.train.reportexpects that all checkpoints are saved within the trainer'sstorage_path. However, withCheckpointUploadMode.NO_UPLOADor a customcheckpoint_upload_fnallows users to specify paths outside the storage path.This causes a problem on resumption for two reasons, we only store the checkpoints path relative to the storage path (which wouldn't be valid for out of band checkpoints) and we resume from the storage path's filesystem (which wouldn't be valid for checkpoints stored on a different filesystem).
Side note: Using relative paths allows user copied a training result's data include the checkpoint data to a new directory and restore it later. For in-band checkpoints, we keep this feature while for out of band checkpoints, we assume that the checkpoints are in a long-term storage that won't be moved.
Checkpoint options
The challenge is that for non-local file systems (gcp, s3, etc) then often passwords, env variables, etc are necessary to config a filesystem to read the data. Annoyingly,
pyarrow.fs.FileSystem.to_uri()doesn't exist so we can't serialise the filesystem as a string. It is possible to serialising the filesystem with pickle but there are issue with this that I comment on at the end.Implementation options
We have three options
ray.train.report(checkpoint=...)Implementing all in and out of band checking
If we want to support all checkpoint options, we need to handle them differently.
Checkpoint.to_directory(), therefore, we warn users about this on resumption if when we try to confirm the checkpoint still exist, this fails.For
CheckpointConfig(num_to_keep=1), then we will try to delete in and out of band checkpoints, however if this fails for any reason, a warning is now produced.Alternative implementations
It is feasible to pickle the filesystem and save this information alongside the path allowing the easiest resumptions with
Checkpoint(path, pickle.dump(filesystem)).However, first this is put a security risk as
pickle.dumpwill run whatever it passed allowing RCEs and second that saved some passwords, env vars, might become invalid.