Allow users to provide custom exception handling

[daniellepintz]

🚀 Feature

Allow users to provide custom exception handling via a new callback hook, similar to on_keyboard_interrupt.

Motivation

Users should be able to implement their own error handling if they want.

Pitch

Create a new callback hook here:
https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/trainer.py#L507-L515

Alternatives

Additional context

---

If you enjoy Lightning, check out our other projects! ⚡

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.

• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, finetuning and solving problems with deep learning

• Bolts: Pretrained SOTA Deep Learning models, callbacks and more for research and production with PyTorch Lightning and PyTorch

• Lightning Transformers: Flexible interface for high performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

[yopknopixx]

Hello! I would like to work on this. Please assign this to me.

[daniellepintz]

Thanks @yopknopixx! I'd like to first collect some feedback from the community before we start working on it. Sorry, didn't mean to add the "help wanted" tag just yet, I think it is automatic.

[awaelchli]

@daniellepintz would the idea be to pass the exception in as an argument to the hook? would it replace the on_keyboard_interrupt?

[daniellepintz]

@awaelchli yes the idea would be to pass the exception in as an argument. I hadn't thought of it but actually it could replace on_keyboard_interrupt, and then we could just call on_exception (new callback hook) with "KeyboardInterrupt" as the exception.

[awaelchli]

that sounds good to me!

[hugoperrin]

I was in the middle of creating my own ticket for that when I saw yours, so I will add the complementary info I would have had added on that subject.

Proposed refactoring or deprecation

feat: add a callback hook for whenever a crash happens
This could be implemented in the trainer like for the keyboard interrupt callback.

It does not seem that the teardown callback is called in the case of a crash, but only in the case of a completed training. It feels more appropriate to have a separate callback rather than forcing the use of teardown.
It seems the teardown call does not seem to happen since the error is re-raised, thus halting the script before getting to teardown.

Motivation

The motivation behind this is to allow running code on failure of training.
Use cases:

• Add Telegram/Slack notifications on training failure
• Print useful debugging information through a debugging callback
• Specific MLFlow/WANDB/etc. logging added on crash

[aurelien-clu]

https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/trainer.py#L1386

Example of a trainer with this behavior:

from contextlib import suppress

from pytorch_lightning import Trainer


class TrainerExtended(Trainer):
    # FIXME: fix typo in pytorch-lightning _on_expection -> _on_exception
    def _on_expection(self):
        super()._on_expection()
        for callback in self.callbacks:
            with suppress(Exception):
                callback.on_exception(self, self.model)

So to update the pytorch_lightning.Trainer one could (may not be the standard way) do the following:

        for callback in self.callbacks:
            with suppress(Exception):
                callback.on_exception()

I used suppress as not all Callbacks have the on_exception() method but also because we don't want to fail because of a callback imho.

It may be replace with a try except to log an error if need be.

As for callbacks:

    def on_exception(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        """Called when an exception is raised"""
        ...

seems sufficient to me.

(I can do a PR if need be - I did not try with passing the training and module but self, self.model would seem to be ok)

[awaelchli]

@aurelien-clu I'm not sure I understand. I believe we should never suppress any exceptions. We will raise the Exception no matter what, but the point of the on_exception method is to run before we re-raise it. And if something inside the on_exception raises an other exception, well that's the problem of the implementation there and should be addressed by the author.

[aurelien-clu]

Sure, I did a quick implementation without checking whether each callback has the function because it is not yet implemented in PL 😉

Let's scratch:

It may be replaced with a try except to log an error if need be.

by

It should be replaced

EDIT:

My "quick implementation" is only subclassing a Trainer and having one callback with on_exception I have not forked PL thus I don't have on_exception for the parent Callback class
EDIT 2:
I can easily do an implementation but someone seems to have called dibs on it 😇