Runtime- and Assertionerror handling in trainer.run_train
🐛 Bug
Runtime- (e.g. Out of Memory Errors) /Assertionerror are ignored when running trainer.fit(...). The error will not be raised correctly and the script will continue.
This can waste a lot of time in some cases:
# prepare code
...

# train (model raises an error, e.g. Out of Memory which is not raised by trainer)
trainer.fit(...)

# will continue here
# time intensive computation / evaluation / prediction
...
While the error is printed correctly, it is not raised and thus the script will continue.

  
    
      pytorch-lightning/pytorch_lightning/trainer/trainer.py
    
    
        Lines 621 to 634
      in
      bb9ace4
    
  
  
    

        
          
           except KeyboardInterrupt: 
        

        
          
               rank_zero_warn('Detected KeyboardInterrupt, attempting graceful shutdown...') 
        

        
          
               # user could press Ctrl+c many times... only shutdown once 
        

        
          
               if not self.interrupted: 
        

        
          
                   self.state = TrainerState.INTERRUPTED 
        

        
          
                   self.on_keyboard_interrupt() 
        

        
          
           except (RuntimeError, AssertionError): 
        

        
          
               # if an exception is raised, the finally block is executed and can hide the actual exception 
        

        
          
               # that was initially raised if `on_train_end` also raises an exception. we want to avoid that 
        

        
          
               # for assertions and other runtime errors so we aren't misled while debugging 
        

        
          
               print_exc() 
        

        
          
           finally: 
        

        
          
               # hook 
        

        
          
               self.train_loop.on_train_end() 
        
    
  



Please reproduce using the BoringModel

To Reproduce
Use following BoringModel and post here

Expected behavior
The script should stop after the trainer cleaned up when an Assertion or Runtime error occurs during training.

Environment
Note: Bugs with code are solved faster ! Colab Notebook should be made public !


IDE: Please, use our python bug_report_model.py template.


Colab Notebook: Please copy and paste the output from our environment collection script (or fill out the checklist below manually).


You can get the script and run it with:
wget https://raw.githubusercontent.com/PyTorchLightning/pytorch-lightning/master/tests/collect_env_details.py
# For security purposes, please check the contents of collect_env_details.py before running it.
python collect_env_details.py


PyTorch Version (e.g., 1.0):
OS (e.g., Linux):
How you installed PyTorch (conda, pip, source):
Build command you used (if compiling from source):
Python version:
CUDA/cuDNN version:
GPU models and configuration:
Any other relevant information:

Additional context
Simply saving the exception and raising it after the trainer called the final hook should be sufficient.