after_save_checkpoint hook for generic Callback

[favrei]

🚀 Feature

Create the after_save_checkpint hook for Callback in the same fashion as the hook in the logger.

Motivation

In our group, we will create a symlink whenever a better checkpoint is saved, so the automated scripts in the background can find the right checkpoint to start their job. Our in-house framework will handle the symlinks, but I find no entry point after we moved to Pytorch Lightning to implement the same feature.

Pitch

A callback hook e.g., Callback.after_save_checkpoint so users can define their custom tasks whenever a better checkpoint is saved. For example:

class UpdateBestCheckpointCB(Callback):
    def after_save_checkpint(self, trainer, pl_module, **others):
        update_symlink_to_best(trainer)

Alternatives

For now, we can always regard the latest checkpoint as the best, but it will be better if I can be sure about which is the best one.

[rohitgr7]

you can actually do

class CustomModelCheckpoint(ModelCheckpoint):
    def save_checkpoint(self, trainer):
        super().save_checkpoint(trainer)
        update_symlink_to_best(trainer)

although this seems a useful hook to have. Mind send a PR??

[favrei]

@rohitgr7 oh, this idea is great. Thanks for the advice.

[jinyoung-lim]

This issue seems to be open and unassigned. I'd love to take a stab at it as my first contribution!

[ananthsub]

I like @rohitgr7 's earlier suggestion and would prefer making it easier to extend the ModelCheckpoint callback to enable this. I personally don't think the logger hook for after save checkpoint is good precedent for this either.

• What would call after_save_checkpoint? Currently all callback hooks are run by the trainer.
• When does after_save_checkpoint actually run? After control returns from the model checkpoint callback entirely? At the end of the model checkpoint's save_checkpoint ? Whenever the trainer calls save_checkpoint ? After the IO plugin calls save_checkpoint ? after each of these? https://github.com/PyTorchLightning/pytorch-lightning/blob/2e17b471f865ed76c940b027f19a50176a69af4b/pytorch_lightning/callbacks/model_checkpoint.py#L382-L389
• What do people want to use after_save_checkpoint for? It looks like there's demand for passing state beyond the trainer and lightning module to this hook. What is the contract for **others ? since others can be highly customized per use case and per checkpoint callback instance, what actually populates the others argument? The Trainer? The chekpoint callback?
  -- Assuming there is an **others argument, passing it as kwargs doesn't give enough visibility to the callback hook implementor around what it contains. This means every user needs to write code to validate that the kwargs contains what they expect.
• Someone new could easily be confused as to what's the difference between on_save_checkpoint and after_save_checkpoint .

[jinyoung-lim]

Good points @ananthsub! I also do like the suggestion by @rohitgr7 to extend the ModelCheckpoint callback. Below are my reply to your points with the limited knowledge that I have from trying to understand this feature request:

• What would call after_save_checkpoint? Currently all callback hooks are run by the trainer.
  : The hook is called by the trainer
• When does after_save_checkpoint actually run? After control returns from the model checkpoint callback entirely? At the end of the model checkpoint's save_checkpoint ? Whenever the trainer calls save_checkpoint? After the IO plugin calls save_checkpoint ? after each of these?
  : This was where I had the most trouble figuring out when trying to implement this feature (so I put out a draft for the PR in the hopes of discussing with other folks). My initial choice was to go for the mixture of at the end of the model checkpoint's save_checkpoint and whenever the trainer calls save_checkpoint, which, after reading your points and reading the code more I do not believe is not the expected behavior. It seems like what @favrei was pitching for was to call this hook after a better checkpoint is saved (_save_top_k_checkpoint where k=1?). Which I do not think is generic enough unless we want to go for specific hooks such as after_save_top_k_checkpoint, which could seem too specific and convoluted. I think after_save_checkpoint should either be called every time save_checkpoint is called. This might be either by ModelCheckpoint or by Trainer.
• What do people want to use after_save_checkpoint for? It looks like there's demand for passing state beyond the trainer and lightning module to this hook. What is the contract for **others ? since others can be highly customized per use case and per checkpoint callback instance, what actually populates the others argument? The Trainer? The chekpoint callback? -- Assuming there is an **others argument, passing it as kwargs doesn't give enough visibility to the callback hook implementor around what it contains. This means every user needs to write code to validate that the kwargs contains what they expect.
  : I do not think after_save_checkpoint should take any **others. Rather I imagine it being overwritten/customized by users. However, this circles back to your point of, "Why can't we just on_save_checkpoint like @rohitgr7 suggested."
• Someone new could easily be confused as to what's the difference between on_save_checkpoint and after_save_checkpoint .
  : I think naming the hook as on_after_save_checkpoint or on_save_checkpoint_end should do the trick to lesson any confusion. With on_save_checkpoint, there is not a guarantee or context of when during the save_checkpoint call the on_save_checkpoint is called. Is it in the middle of saving the checkpoint? Before? After? Whereas if we have specific hooks such as on_after_save_checkpoint or on_before_save_checkpoint, when those are called should be clearer. Again, users can do this by:

  class CustomModelCheckpoint(ModelCheckpoint):
      def save_checkpoint(self, trainer):
          before_save_checkpoint(trainer)
          super().save_checkpoint(trainer)
          after_save_checkpoint(trainer)
  

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!