Repository: Lightning-AI/pytorch-lightning

Issue:
Avoid running `on_training_end` after Keyboard Interrupt
## 🚀 Feature
Right now, due to #795 , `on_training_end` runs after a `KeyboardInterrupt`. This ends up running code that's meant to be run only after successful training completion. This feature should either be reverted, or an alternative should be provided, so as to run some code only after successful training.

### Motivation

I am training a model on Sagemaker and have added a notebook shutdown code within the `on_training_end` method. There were times where I had to manually cancel my model training because some parameters were incorrect. However, If I do that, the notebook shuts down immediately. This is because the `on_training_end` method runs even after a `KeyboardInterrupt`.  I don't want my notebook shutting down after a keyboard interrupt, only after successful training completion.

### Pitch
Maybe add an `on_training_completed` method for code that's meant to be run after successful training.


Discussion:
Hi! thanks for your contribution!, great first issue!
Good point, it makes sense to skip eval of the unfinished train also when a user is it he doesn't want to wait for another hour lol Mind send a PR? :robot: 
Thanks for bringing this use-case to our attention @lezwon , definitely something we'd want to consider.

What if the `Trainer` object had a `status` property? Similar to how compute jobs will have a status `{PENDING, RUNNING, COMPLETE, FAILED}` we could apply something similar here for the training job. In this case we could differentiate the two cases with a status of INTERRUPTED vs COMPLETED and your callback logic can check for the proper status before closing (eg. shutdown notebook on FAILED/COMPLETED but don't shutdown for INTERRUPTED). 

What do you think?
@jeremyjordan that sounds great. We could definitely do that.
Not sure if we really need to add complexity to the existing callbacks...
Also `on_training_completed` is standard behavior so rather some cleanup for interrupted?
@jeremyjordan status sounds cool, by my opinion it is much better signal the returned 0/1 from e.g. `.fit` =) 
@Borda are you suggesting that we add a new callback (`on_training_completed` that @lezwon mentioned) which runs conditionally according to the value returned from `fit()`? My only worry is that it might be confusing to know the difference between `on_training_end` (existing) and `on_training_complete` (proposed).
in this moment I was just thinking about adding Trainer status... 
not sure if we want to add complexity with new callback...
Ok yes I agree, I think the Trainer status would be a simple solution that may also be useful in other situations as well :)
Cool, does anyone send a PR with Trainer status? I may suggest to implement it as enum/numbering similar like `logging.INFO/DEBUG/...`
Will this status also account for validation/test completed/interrupted?
yes we could do something like:

```
>>> from enum import Enum
>>> class TrainerStatus(Enum):
...     PENDING = "Initializing Trainer object"
...     TRAINING = "Optimizing model via train loop"
...     VALIDATING = "Running model on validation set"
...     TESTING = "Running model on test set"
...     FAILED = "Trainer failed to complete a successful run"
...     INTERRUPTED = "Training was interrupted by the user"
...     COMPLETED = "Training completed successfully"
```

I can work on getting this into a PR
Hey @jeremyjordan, what if we would like to execute different actions if the trainer failed during a validation task? Can we find out which task it failed at from the status? Just wondering if this would be flexible in such a scenario.
@lezwon do you have an example in mind of where you'd need that? 

i'd prefer to keep the implementation simple and generic enough such that we don't keep adding new status types.
Don't really have an example yet. Was just considering a scenario like that though. 

I guess we could go ahead with the current proposal you mentioned. That should solve my issue for sure 😊
probably want to consider building upon the Trainer state defined in #770 

fyi @xingzhaolee - what do you think? (we don't need to include it in #770 just thinking about building off of that)
@jeremyjordan sounds like a good idea! also I was thinking that since we have TrainerStatus, can it be the main class while TrainerMode servers as a nested enum instead of putting them all under a single enum? something like:
```python
class TrainerMode(enum.Enum):
    TRAINING = enum.auto()
    VALIDATING = enum.auto()
    TESTING = enum.auto()

class TrainerStatus(enum.Enum):
    PENDING = enum.auto()
    FAILED = enum.auto()
    INTERRUPTED = enum.auto()
    ...
    MODE = TrainerMode
```

then user can access the current status and mode through:
```python
# Interrupted when validation is running
if ... is TrainerStatus.INTERRUPTED and ... is TrainerStatus.MODE.VALIDATING:
```
sure! that would be flexible enough to address @lezwon 's earlier comment. i had started a branch to work on TrainerStatus, once #770 is merged i can pick up that work with your new suggestion here :)
if I understand it correctly the `.auto` means that it can be every run different `enum` value, right? Even it seems to be starting from 1 and incrementally increase...
if so I would recommend using some exact numbers in case of exporting/importing :]
not too sure about that, but probably the order matters too when new status or mode is added. so I guess we should use exact numbers like you suggested in case any import or export is involved! 🙂

Resolving patch:
diff --git a/pytorch_lightning/trainer/__init__.py b/pytorch_lightning/trainer/__init__.py
--- a/pytorch_lightning/trainer/__init__.py
+++ b/pytorch_lightning/trainer/__init__.py
@@ -30,6 +30,7 @@
     trainer = Trainer()
     trainer.fit(model)
 
+
 --------
 
 Best Practices
@@ -59,6 +60,14 @@ def main(hparams):
 
     $ python main.py --gpus 2
 
+
+.. note::
+    If you want to stop a training run early, you can press "Ctrl + C" on your keyboard.
+    The trainer will catch the `KeyboardInterrupt` and attempt a graceful shutdown, including
+    running callbacks such as `on_train_end`. The trainer object will also set an attribute
+    `interrupted` to `True` in such cases. If you have a callback which shuts down compute
+    resources, for example, you can conditionally run the shutdown logic for only uninterrupted runs.
+
 ------------
 
 Testing
diff --git a/pytorch_lightning/trainer/trainer.py b/pytorch_lightning/trainer/trainer.py
--- a/pytorch_lightning/trainer/trainer.py
+++ b/pytorch_lightning/trainer/trainer.py
@@ -364,6 +364,7 @@ def __init__(
         self.global_step = 0
         self.current_epoch = 0
         self.total_batches = 0
+        self.interrupted = False
 
         # configure logger
         self.configure_logger(logger)
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
@@ -203,6 +203,7 @@ class TrainerTrainLoopMixin(ABC):
     use_amp: bool
     track_grad_norm: ...
     model: LightningModule
+    interrupted: bool
     running_loss: ...
     training_tqdm_dict: ...
     reduce_lr_on_plateau_scheduler: ...
@@ -387,6 +388,7 @@ def train(self):
 
         except KeyboardInterrupt:
             log.info('Detected KeyboardInterrupt, attempting graceful shutdown...')
+            self.interrupted = True
             self.run_training_teardown()
 
     def run_training_epoch(self):


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
@@ -12,6 +12,7 @@
     EarlyStopping,
     ModelCheckpoint,
 )
+from pytorch_lightning import Callback
 from pytorch_lightning.core.lightning import load_hparams_from_tags_csv
 from pytorch_lightning.trainer.logging import TrainerLoggingMixin
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
@@ -630,3 +631,33 @@ def on_after_backward(self):
     # after aborting the training loop, model still has nan-valued params
     params = torch.cat([param.view(-1) for param in model.parameters()])
     assert not torch.isfinite(params).all()
+
+
+def test_trainer_interrupted_flag(tmpdir):
+    """Test the flag denoting that a user interrupted training."""
+
+    model = DictHparamsModel({'in_features': 28 * 28, 'out_features': 10})
+
+    class InterruptCallback(Callback):
+        def __init__(self):
+            super().__init__()
+
+        def on_batch_start(self, trainer, pl_module):
+            raise KeyboardInterrupt
+
+    interrupt_callback = InterruptCallback()
+
+    trainer_options = {
+        'callbacks': [interrupt_callback],
+        'max_epochs': 1,
+        'val_percent_check': 0.1,
+        'train_percent_check': 0.2,
+        'progress_bar_refresh_rate': 0,
+        'logger': False,
+        'default_save_path': tmpdir,
+    }
+
+    trainer = Trainer(**trainer_options)
+    assert not trainer.interrupted
+    trainer.fit(model)
+    assert trainer.interrupted