[RFC] Support checkpointing multiple callbacks of the same type
🚀 Feature
Currently when dumping the checkpoint dict, we overwrite callback states if there are multiple callbacks of the same type: 
  
    
      pytorch-lightning/pytorch_lightning/trainer/callback_hook.py
    
    
        Lines 209 to 224
      in
      1c013b4
    
  
  
    

        
          
           def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> Dict[Type, dict]: 
        

        
          
               """Called when saving a model checkpoint.""" 
        

        
          
               callback_states = {} 
        

        
          
               for callback in self.callbacks: 
        

        
          
                   if self.__is_old_signature(callback.on_save_checkpoint): 
        

        
          
                       rank_zero_warn( 
        

        
          
                           "`Callback.on_save_checkpoint` signature has changed in v1.3." 
        

        
          
                           " A `checkpoint` parameter has been added." 
        

        
          
                           " Support for the old signature will be removed in v1.5", DeprecationWarning 
        

        
          
                       ) 
        

        
          
                       state = callback.on_save_checkpoint(self, self.lightning_module)  # noqa: parameter-unfilled 
        

        
          
                   else: 
        

        
          
                       state = callback.on_save_checkpoint(self, self.lightning_module, checkpoint) 
        

        
          
                   if state: 
        

        
          
                       callback_states[type(callback)] = state 
        

        
          
               return callback_states 
        
    
  


Motivation
This blocks #2908 and is an existing bug or at least unexpected behavior with checkpointing callback states today.
Pitch
[WIP]

We add a state_identifier property to the base Callback class

@property
def state_identifier(self) -> str:
    return ""


TODO: should this raise NotImplementedError? Reasoning: as it stands, we'd call this only when saving/loading callback states. Could we be even more selective and call this only when we have multiple callback instances of the same type? Is that complexity warranted?


Callback implementations define this property based on the business logic of the class in order to determine uniqueness, while also preserving flexibility for development.

For instance
class ModelCheckpoint(Callback):

...

@property
def state_identifier(self):
    return f"monitor={self.monitor}"  # handle special case where this is not available, include other params from here

TODO: provide better default implementation based on args here

At save time, we still partition the callback state dicts by the type, but include the state_identifier to further disambiguate.

# we include the checkpoint dict in order to look at the lightning version or other metadata that can be used to better preserve backwards compatibility
def _compute_callback_state_key(self, checkpoint: Dict[str, Any], callback: Callback) -> str:
    return type(callback) + callback.state_identifier()

def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> Dict[Type, dict]:
        """Called when saving a model checkpoint."""
        callback_states = {}
        for callback in self.callbacks:
            if self.__is_old_signature(callback.on_save_checkpoint):
                rank_zero_warn(
                    "`Callback.on_save_checkpoint` signature has changed in v1.3."
                    " A `checkpoint` parameter has been added."
                    " Support for the old signature will be removed in v1.5", DeprecationWarning
                )
                state = callback.on_save_checkpoint(self, self.lightning_module)  # noqa: parameter-unfilled
            else:
                state = callback.on_save_checkpoint(self, self.lightning_module, checkpoint)
            if state:
                key = self._compute_callback_state_key(checkpoint, callback)
                callback_states[key] = state
        return callback_states


At load time, we do something similar:

def on_load_checkpoint(self, checkpoint):
        """Called when loading a model checkpoint."""
        callback_states = checkpoint.get('callbacks')
        # Todo: the `callback_states` are dropped with TPUSpawn as they
        # can't be saved using `xm.save`
        # https://github.com/pytorch/xla/issues/2773
        if callback_states is not None:
            for callback in self.callbacks:
                key = self._compute_callback_state_key(checkpoint, callback)
                state = callback_states.get(key)
                if state:
                    state = deepcopy(state)
                    callback.on_load_checkpoint(state)

Alternatives

Why don't Callback implementations simply hash the value of the constructor args?

from @awaelchli : The problematic part is around Trainer(resume_from_checkpoint=..., callbacks=[...]), which would restore the trainer state plus callback state. If the source code changes, what will happen? Well in the normal case you wouldn't modify the code to perfectly resume training, right? But I think it is still worth discussing what kind of flexibility we could allow, if any? For instance, if we hashed all the constructor args, a callback that changes a verbose flag from False to True would no longer be able to access its previously checkpointed state.
Additional context