Repository: Lightning-AI/pytorch-lightning
File: tests/tests_pytorch/checkpointing/test_model_checkpoint.py (Python)

Code before:
        logger=False,
    )

    with pytest.raises(RuntimeError, match="Trouble!"):
        trainer.fit(model)
    assert not os.path.isfile(tmp_path / "exception-fast_dev_run.ckpt")

    # Don't save checkpoint if already saved a checkpoint
    class TroubledModelAlreadySavedCheckpoint(BoringModel):
        def on_train_batch_start(self, batch, batch_idx) -> None:
            if self.trainer.global_step == 1:
                raise RuntimeError("Trouble!")

    model = TroubledModelAlreadySavedCheckpoint()
    checkpoint_callback = ModelCheckpoint(
        dirpath=tmp_path, filename="already_saved", save_on_exception=True, every_n_train_steps=1
    )
    trainer = Trainer(
        default_root_dir=tmp_path, limit_train_batches=2, callbacks=[checkpoint_callback], max_epochs=2, logger=False
    )

    with pytest.raises(RuntimeError, match="Trouble!"):
        trainer.fit(model)

    assert not os.path.isfile(tmp_path / "exception-already_saved.ckpt")
    assert os.path.isfile(tmp_path / "already_saved.ckpt")


class TroubledModelInTrainingStep(BoringModel):
    def training_step(self, batch, batch_idx):
        if batch_idx == 1:
            raise RuntimeError("Trouble!")


class TroubledModelInValidationStep(BoringModel):
    def validation_step(self, batch, batch_idx):
        if not self.trainer.sanity_checking and batch_idx == 1:
            raise RuntimeError("Trouble!")


class TroubledModelBackward(BoringModel):
    def backward(self, loss):
        if self.current_epoch == 1:
            raise RuntimeError("Trouble!")


class TroubledModelOnBeforeBackward(BoringModel):
    def on_before_backward(self, loss):
        if self.current_epoch == 1:
            raise RuntimeError("Trouble!")


Review comment (suggestion):
```suggestion
```
This seems to be the same as the parametrization bellow

Code after:
    trainer = Trainer(
        default_root_dir=tmp_path,
        fast_dev_run=2,
        limit_train_batches=2,
        callbacks=[checkpoint_callback],
        max_epochs=2,
        logger=False,
    )

    with pytest.raises(RuntimeError, match="Trouble!"):
        trainer.fit(model)
    assert not os.path.isfile(tmp_path / "exception-fast_dev_run.ckpt")


def test_model_checkpoint_on_exception_run_condition_on_train_batch_start(tmp_path):
    """Test that no checkpoint is saved when an exception is raised during a sanity check or a fast dev run, or when a
    checkpoint has already been saved at the current training step."""

    # Don't save checkpoint if already saved a checkpoint
    class TroubledModelAlreadySavedCheckpoint(BoringModel):
        def on_train_batch_start(self, batch, batch_idx) -> None:
            if self.trainer.global_step == 1:
                raise RuntimeError("Trouble!")

    model = TroubledModelAlreadySavedCheckpoint()
    checkpoint_callback = ModelCheckpoint(
        dirpath=tmp_path, filename="already_saved", save_on_exception=True, every_n_train_steps=1
    )
    trainer = Trainer(
        default_root_dir=tmp_path, limit_train_batches=2, callbacks=[checkpoint_callback], max_epochs=2, logger=False
    )

    with pytest.raises(RuntimeError, match="Trouble!"):
        trainer.fit(model)

    assert not os.path.isfile(tmp_path / "exception-already_saved.ckpt")
    assert os.path.isfile(tmp_path / "already_saved.ckpt")


class TroubledModelInTrainingStep(BoringModel):
    def training_step(self, batch, batch_idx):
        if batch_idx == 1:
            raise RuntimeError("Trouble!")


class TroubledModelInValidationStep(BoringModel):
    def validation_step(self, batch, batch_idx):
        if not self.trainer.sanity_checking and batch_idx == 1:
            raise RuntimeError("Trouble!")