Repository: Lightning-AI/pytorch-lightning

Issue:
Handle KeyboardInterrupt during training
## 🚀 Feature

It should be possible to examine the stack  if the training is interrupted manually.

### Motivation

In my case, the program hangs at some point and I have to cancel it manually. Because Lightning catches KeyboardInterrupt, I don't get to know where the program was hanging.

### Alternatives

1. **Re-raise KeyboardInterrupt.** This would break the current behavior. This could be made configurable.
2. **Save result of `sys.exc_info()` for later examination.** This could introduce memory issues like memory leaks or circular references.
3. **Call a handler.** The users could then call `sys.exc_info()` in the handler. This would make memory issues less likely.

I will prepare a pull request for the third option.


Discussion:
The least invasive change would be to save the result of sys.exc_info() for later examination.

Resolving patch:
diff --git a/pytorch_lightning/callbacks/base.py b/pytorch_lightning/callbacks/base.py
--- a/pytorch_lightning/callbacks/base.py
+++ b/pytorch_lightning/callbacks/base.py
@@ -85,3 +85,6 @@ def on_test_start(self, trainer, pl_module):
     def on_test_end(self, trainer, pl_module):
         """Called when the test ends."""
         pass
+
+    def on_keyboard_interrupt(self, trainer, pl_module):
+        """Called when the training is interrupted by KeyboardInterrupt."""
diff --git a/pytorch_lightning/trainer/callback_hook.py b/pytorch_lightning/trainer/callback_hook.py
--- a/pytorch_lightning/trainer/callback_hook.py
+++ b/pytorch_lightning/trainer/callback_hook.py
@@ -100,3 +100,8 @@ def on_test_end(self):
         """Called when the test ends."""
         for callback in self.callbacks:
             callback.on_test_end(self, self.get_model())
+
+    def on_keyboard_interrupt(self):
+        """Called when the training is interrupted by KeyboardInterrupt."""
+        for callback in self.callbacks:
+            callback.on_keyboard_interrupt(self, self.get_model())
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -237,6 +237,7 @@ class TrainerTrainLoopMixin(ABC):
     checkpoint_callback: ...
     terminate_on_nan: bool
     tpu_id: int
+    interactive_ddp_procs: ...
 
     # Callback system
     callbacks: List[Callback]
@@ -247,6 +248,7 @@ class TrainerTrainLoopMixin(ABC):
     on_epoch_start: Callable
     on_epoch_end: Callable
     on_validation_end: Callable
+    on_keyboard_interrupt: Callable
 
     @abstractmethod
     def get_model(self) -> LightningModule:
@@ -395,6 +397,7 @@ def train(self):
             # user could press ctrl+c many times... only shutdown once
             if not self.interrupted:
                 self.interrupted = True
+                self.on_keyboard_interrupt()
 
                 for proc in self.interactive_ddp_procs:
                     subprocess.Popen.kill(proc)


Test patch:
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
--- a/tests/trainer/test_trainer.py
+++ b/tests/trainer/test_trainer.py
@@ -3,6 +3,7 @@
 import os
 import pickle
 import types
+import sys
 from argparse import Namespace
 
 import cloudpickle
@@ -10,10 +11,10 @@
 import torch
 
 import tests.base.utils as tutils
-from pytorch_lightning import Callback, LightningModule
-from pytorch_lightning import Trainer
+from pytorch_lightning import Callback, LightningModule, Trainer
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
-from pytorch_lightning.core.saving import load_hparams_from_tags_csv, load_hparams_from_yaml, save_hparams_to_tags_csv
+from pytorch_lightning.core.saving import (
+    load_hparams_from_tags_csv, load_hparams_from_yaml, save_hparams_to_tags_csv)
 from pytorch_lightning.loggers import TensorBoardLogger
 from pytorch_lightning.trainer.logging import TrainerLoggingMixin
 from pytorch_lightning.utilities.io import load as pl_load
@@ -660,10 +661,19 @@ def __init__(self):
         def on_batch_start(self, trainer, pl_module):
             raise KeyboardInterrupt
 
+    class HandleInterruptCallback(Callback):
+        def __init__(self):
+            super().__init__()
+            self.exc_info = None
+
+        def on_keyboard_interrupt(self, trainer, pl_module):
+            self.exc_info = sys.exc_info()
+
     interrupt_callback = InterruptCallback()
+    handle_interrupt_callback = HandleInterruptCallback()
 
     trainer = Trainer(
-        callbacks=[interrupt_callback],
+        callbacks=[interrupt_callback, handle_interrupt_callback],
         max_epochs=1,
         val_percent_check=0.1,
         train_percent_check=0.2,
@@ -672,8 +682,10 @@ def on_batch_start(self, trainer, pl_module):
         default_root_dir=tmpdir,
     )
     assert not trainer.interrupted
+    assert handle_interrupt_callback.exc_info is None
     trainer.fit(model)
     assert trainer.interrupted
+    assert isinstance(handle_interrupt_callback.exc_info[1], KeyboardInterrupt)
 
 
 def test_gradient_clipping(tmpdir):