Repository: Lightning-AI/pytorch-lightning
File: src/lightning/pytorch/callbacks/model_checkpoint.py (Python)

Code before:
            monitor_candidates = self._monitor_candidates(trainer)
            if self._every_n_epochs >= 1 and (trainer.current_epoch + 1) % self._every_n_epochs == 0:
                self._save_topk_checkpoint(trainer, monitor_candidates)
            self._save_last_checkpoint(trainer, monitor_candidates)

    @override
    def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        """Save a checkpoint at the end of the validation stage."""
        if not self._should_skip_saving_checkpoint(trainer) and not self._should_save_on_train_epoch_end(trainer):
            monitor_candidates = self._monitor_candidates(trainer)
            if self._every_n_epochs >= 1 and (trainer.current_epoch + 1) % self._every_n_epochs == 0:
                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 {
            "monitor": self.monitor,
            "best_model_score": self.best_model_score,
            "best_model_path": self.best_model_path,
            "current_score": self.current_score,
            "dirpath": self.dirpath,
            "best_k_models": self.best_k_models,
            "kth_best_model_path": self.kth_best_model_path,
            "kth_value": self.kth_value,
            "last_model_path": self.last_model_path,
        }

    @override
    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
        dirpath_from_ckpt = state_dict.get("dirpath", self.dirpath)

        if self.dirpath == dirpath_from_ckpt:
            self.best_model_score = state_dict["best_model_score"]
            self.kth_best_model_path = state_dict.get("kth_best_model_path", self.kth_best_model_path)
            self.kth_value = state_dict.get("kth_value", self.kth_value)
            self.best_k_models = state_dict.get("best_k_models", self.best_k_models)
            self.last_model_path = state_dict.get("last_model_path", self.last_model_path)

Review comment (suggestion):
```suggestion
        if not self._should_save_on_exception(trainer):
            return
        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}"
        )
```

Code after:

    @override
    def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        """Save a checkpoint at the end of the training epoch."""
        if not self._should_skip_saving_checkpoint(trainer) and self._should_save_on_train_epoch_end(trainer):
            monitor_candidates = self._monitor_candidates(trainer)
            if self._every_n_epochs >= 1 and (trainer.current_epoch + 1) % self._every_n_epochs == 0:
                self._save_topk_checkpoint(trainer, monitor_candidates)
            self._save_last_checkpoint(trainer, monitor_candidates)

    @override
    def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        """Save a checkpoint at the end of the validation stage."""
        if not self._should_skip_saving_checkpoint(trainer) and not self._should_save_on_train_epoch_end(trainer):
            monitor_candidates = self._monitor_candidates(trainer)
            if self._every_n_epochs >= 1 and (trainer.current_epoch + 1) % self._every_n_epochs == 0:
                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: BaseException) -> None:
        """Save a checkpoint when an exception is raised."""
        if not self._should_save_on_exception(trainer):
            return
        monitor_candidates = self._monitor_candidates(trainer)
        filepath = self.format_checkpoint_name(metrics=monitor_candidates)
        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}"
        )

    def on_train_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        """Ensure save_last=True is applied when training ends."""
        if self.save_last and not self._last_checkpoint_saved:
            monitor_candidates = self._monitor_candidates(trainer)
            self._save_last_checkpoint(trainer, monitor_candidates)

    @override
    def state_dict(self) -> dict[str, Any]:
        return {
            "monitor": self.monitor,
            "best_model_score": self.best_model_score,
            "best_model_path": self.best_model_path,
            "current_score": self.current_score,
            "dirpath": self.dirpath,
            "best_k_models": self.best_k_models,
            "kth_best_model_path": self.kth_best_model_path,
            "kth_value": self.kth_value,
            "last_model_path": self.last_model_path,