Lightning-AI / Lightning-AI/pytorch-lightning
A new after on_save_checkpoint callback for AsyncCheckpointIO
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 31.4k
- Forks
- 3.8k
- Avg merge
- 6d 7h
- Merged PRs (30d)
- 6
Description
### Description & Motivation
We want to log_artifact, publish tags (e.g. with `MLFlowLogger`) only _after_ saving the checkpoint is complete (successfully). With `AsyncCheckpointIO` with does not seem to have a way to know when checkpoint saving is actually complete.
### Pitch
It would help to either add a new `on_after_save_checkpoint` to `CheckpointHooks` or to a new `AsyncCheckpointHooks`.
Then `AsyncCheckpointIO` can provide a way to register callbacks for the `Future`'s `add_done_callback`. When initializing the `Trainer` can iterate `Callback`s implementing the new method/hooks and register them with `Strategy` or `AsyncCheckpointIO`.
Regular `CheckpointIO` can also allow registration but will call the callbacks as soon as `save_checkpoint` is done.
### Alternatives
Currently the users are forced to copy-paste `AsyncCheckpointIO` and alter it to provide such hook.
### Additional context
I'm willing to contribute if there are no concerns with the proposed solution.
cc @lantiga @borda
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by tracing how AsyncCheckpointIO saves checkpoints and how Trainer, Strategy, Callback, and CheckpointHooks interact. Review the Future.add_done_callback path and compare it with regular CheckpointIO completion. Done means callbacks can observe successful asynchronous completion while regular saves invoke them after saving finishes.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100