Pass all running stages to DataModule.setup
🚀 Feature
Currently. DataModule.setup is only called with stages fit or test. But we have several more:
Stages:

  
    
      pytorch-lightning/pytorch_lightning/trainer/states.py
    
    
        Lines 39 to 49
      in
      5f33728
    
  
  
    

        
          
           class RunningStage(LightningEnum): 
        

        
          
               """Type of train phase. 
        

        
          
            
        

        
          
               >>> # you can match the Enum with string 
        

        
          
               >>> RunningStage.TRAINING == 'train' 
        

        
          
               True 
        

        
          
               """ 
        

        
          
               TRAINING = 'train' 
        

        
          
               EVALUATING = 'eval' 
        

        
          
               TESTING = 'test' 
        

        
          
               TUNING = 'tune' 
        
    
  


Note that it's a bit tricky because fit is not a RunningStage. It indicates train or eval
Motivation
Allows having custom logic for each stage
Pitch
def setup(stage: Optional[str] = None):
    assert stage in list(RunningStage)
    ...
Additional context
We are passing 'test' when predicting as seen in #5579

  
    
      pytorch-lightning/pytorch_lightning/trainer/trainer.py
    
    
         Line 909
      in
      9137b16
    
  
  
    

        
          
           self.data_connector.attach_datamodule(model or self.get_model(), datamodule, 'test')