Add save_on_exception option to ModelCheckpoint

src/lightning/pytorch/callbacks/model_checkpoint.py

@@ -97,6 +97,7 @@ class ModelCheckpoint(Checkpoint):
             collisions unless ``enable_version_counter`` is set to False. The version counter is unrelated to the top-k
             ranking of the checkpoint, and we recommend formatting the filename to include the monitored metric to avoid
             collisions.
+        save_on_exception: Whether to save a checkpoint when an exception is raised. Default: ``False``.
         mode: one of {min, max}.
             If ``save_top_k != 0``, the decision to overwrite the current save file is made
             based on either the maximization or the minimization of the monitored quantity.
@@ -213,6 +214,7 @@ class ModelCheckpoint(Checkpoint):
     CHECKPOINT_JOIN_CHAR = "-"
     CHECKPOINT_EQUALS_CHAR = "="
     CHECKPOINT_NAME_LAST = "last"
+    CHECKPOINT_EXCEPTION_PREFIX = "exception"
     FILE_EXTENSION = ".ckpt"
     STARTING_VERSION = 1
 
@@ -224,6 +226,7 @@ def __init__(
         verbose: bool = False,
         save_last: Optional[Union[bool, Literal["link"]]] = None,
         save_top_k: int = 1,
+        save_on_exception: bool = False,
         save_weights_only: bool = False,
         mode: str = "min",
         auto_insert_metric_name: bool = True,
@@ -238,6 +241,7 @@ def __init__(
         self.verbose = verbose
         self.save_last = save_last
         self.save_top_k = save_top_k
+        self.save_on_exception = save_on_exception
         self.save_weights_only = save_weights_only
         self.auto_insert_metric_name = auto_insert_metric_name
         self._save_on_train_epoch_end = save_on_train_epoch_end
@@ -338,6 +342,19 @@ def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModul
                 self._save_topk_checkpoint(trainer, monitor_candidates)
             self._save_last_checkpoint(trainer, monitor_candidates)
 
+    @override
+    def on_exception(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", exception: Exception) -> None:
+        """Save a checkpoint when an exception is raised."""
+        if self._should_save_on_exception(trainer):
+            monitor_candidates = self._monitor_candidates(trainer)
+            filepath = self.format_checkpoint_name(metrics=monitor_candidates, prefix=self.CHECKPOINT_EXCEPTION_PREFIX)
+            self._save_checkpoint(trainer, filepath)
+            self._save_last_checkpoint(trainer, monitor_candidates)
+            rank_zero_info(
+                f"An {type(exception).__name__} was raised with message: \
+                {str(exception)}, saved checkpoint to {filepath}"
+            )
+
     @override
     def state_dict(self) -> dict[str, Any]:
         return {
@@ -426,6 +443,14 @@ def _should_skip_saving_checkpoint(self, trainer: "pl.Trainer") -> bool:
             or self._last_global_step_saved == trainer.global_step  # already saved at the last step
         )
 
+    def _should_save_on_exception(self, trainer: "pl.Trainer") -> bool:
+        return (
+            self.save_on_exception
+            and not bool(trainer.fast_dev_run)  # disable checkpointing with fast_dev_run
+            and not trainer.sanity_checking  # don't save anything during sanity check
+            and self._last_global_step_saved != trainer.global_step  # already saved at the last step
+        )
+
     def _should_save_on_train_epoch_end(self, trainer: "pl.Trainer") -> bool:
         if self._save_on_train_epoch_end is not None:
             return self._save_on_train_epoch_end
@@ -571,7 +596,11 @@ def _format_checkpoint_name(
         return filename
 
     def format_checkpoint_name(
-        self, metrics: dict[str, Tensor], filename: Optional[str] = None, ver: Optional[int] = None
+        self,
+        metrics: dict[str, Tensor],
+        filename: Optional[str] = None,
+        prefix: Optional[str] = None,
+        ver: Optional[int] = None,
     ) -> str:
         """Generate a filename according to the defined template.
 
@@ -603,7 +632,9 @@ def format_checkpoint_name(
 
         """
         filename = filename or self.filename
-        filename = self._format_checkpoint_name(filename, metrics, auto_insert_metric_name=self.auto_insert_metric_name)
+        filename = self._format_checkpoint_name(
+            filename, metrics, prefix, auto_insert_metric_name=self.auto_insert_metric_name
+        )
 
         if ver is not None:
             filename = self.CHECKPOINT_JOIN_CHAR.join((filename, f"v{ver}"))