code: def test_loop_state_on_exception(
      accumulate_grad_batches, stop_epoch, stop_batch, tmpdir
  ):
      n_epochs = 3
      n_batches = 3
  
      class TestModel(BoringModel):
          def training_step(self, batch, batch_idx):
              if (
                  self.trainer.current_epoch == stop_epoch
                  and batch_idx == stop_batch
              ):
                  raise CustomException
              return super().training_step(batch, batch_idx)
  
      model = TestModel()
  
      trainer = Trainer(
          default_root_dir=tmpdir,
          max_epochs=n_epochs,
          limit_train_batches=n_batches,
          limit_val_batches=0,
          accumulate_grad_batches=accumulate_grad_batches,
          enable_progress_bar=False,
          logger=False,
          callbacks=OnExceptionCheckpoint(tmpdir),
      )
  
      # simulate a failure
      with pytest.raises(CustomException):
          trainer.fit(model)
  
      ckpt_path = str(tmpdir / "on_exception.ckpt")
      assert os.path.exists(ckpt_path)
      checkpoint = torch.load(ckpt_path)
  
      optim_progress = (
          trainer.fit_loop.epoch_loop.automatic_optimization.optim_progress
      )
      sch_progress = trainer.fit_loop.epoch_loop.scheduler_progress
  
      # `nbe_`: non-breaking epoch, as in, no exception will be raised. `be_`: breaking epoch
      nbe_batches_completed = stop_epoch * n_batches
      be_batches_completed = stop_batch
      # lightning applies leftover accumulated gradients when the epoch ends
      has_leftover_accumulation_batches = (
          n_batches % accumulate_grad_batches != 0
      )
      # number of batches that will call `optimizer.step()` during non-breaking and breaking epochs
      nbe_stepping_batches = (
          nbe_batches_completed // accumulate_grad_batches
      )
      be_stepping_batches = (
          be_batches_completed // accumulate_grad_batches
      )
  
      nbe_total_opt_steps = (
          nbe_stepping_batches + has_leftover_accumulation_batches
      )
      be_total_opt_steps = be_stepping_batches
      assert (
          optim_progress.optimizer_steps
          == nbe_total_opt_steps + be_total_opt_steps
      )
      assert (
          optim_progress.optimizer.step.current.completed
          == be_total_opt_steps
      )
      has_opt_stepped_in_be = stop_batch + 1 >= accumulate_grad_batches
  
      nbe_total_zero_grad = (
          nbe_stepping_batches + has_leftover_accumulation_batches
      )
      # `max` because the first batch always zero-grads
      be_total_zero_grad = max(1, be_stepping_batches)
      assert (
          optim_progress.optimizer.zero_grad.total.completed
          == nbe_total_zero_grad + be_total_zero_grad
      )
      assert (
          optim_progress.optimizer.zero_grad.current.completed
          == be_total_zero_grad
      )
  
      nbe_sch_steps = stop_epoch
      be_sch_steps = 0  # the current epoch did not complete
      assert (
          sch_progress.total.completed == nbe_sch_steps + be_sch_steps
      )
      assert sch_progress.current.completed == be_sch_steps
  
      expected = {
          "state_dict": ANY,
          "epoch_progress": {
              "total": {
                  "ready": stop_epoch + 1,
                  "started": stop_epoch + 1,
                  "processed": stop_epoch,
                  "completed": stop_epoch,
              },
              "current": {
                  "ready": stop_epoch + 1,
                  "started": stop_epoch + 1,
                  "processed": stop_epoch,
                  "completed": stop_epoch,
              },
          },
          "epoch_loop.state_dict": ANY,
          "epoch_loop.batch_progress": {
              "total": {
                  "ready": nbe_batches_completed
                  + be_batches_completed
                  + 1,
                  "started": nbe_batches_completed
                  + be_batches_completed
                  + 1,
                  "processed": nbe_batches_completed
                  + be_batches_completed,
                  "completed": nbe_batches_completed
                  + be_batches_completed,
              },
              "current": {
                  "ready": stop_batch + 1,
                  "started": stop_batch + 1,
                  "processed": stop_batch,
                  "completed": stop_batch,
              },
              "is_last_batch": (stop_batch + 1) == n_batches,
          },
          "epoch_loop.scheduler_progress": {
              "total": {
                  "ready": nbe_sch_steps + be_sch_steps,
                  "completed": nbe_sch_steps + be_sch_steps,
              },
              "current": {
                  "ready": be_sch_steps,
                  "completed": be_sch_steps,
              },
          },
          "epoch_loop.manual_optimization.state_dict": ANY,
          "epoch_loop.manual_optimization.optim_step_progress": {
              "total": {"ready": 0, "completed": 0},
              "current": {"ready": 0, "completed": 0},
          },
          "epoch_loop.automatic_optimization.state_dict": {},
          "epoch_loop.automatic_optimization.optim_progress": {
              "optimizer": {
                  "step": {
                      "total": {
                          "ready": nbe_total_opt_steps
                          + be_total_opt_steps
                          + has_opt_stepped_in_be,
                          "completed": nbe_total_opt_steps
                          + be_total_opt_steps,
                      },
                      "current": {
                          "ready": be_total_opt_steps
                          + has_opt_stepped_in_be,
                          "completed": be_total_opt_steps,
                      },
                  },
                  "zero_grad": {
                      "total": {
                          "ready": nbe_total_zero_grad
                          + be_total_zero_grad,
                          "started": nbe_total_zero_grad
                          + be_total_zero_grad,
                          "completed": nbe_total_zero_grad
                          + be_total_zero_grad,
                      },
                      "current": {
                          "ready": be_total_zero_grad,
                          "started": be_total_zero_grad,
                          "completed": be_total_zero_grad,
                      },
                  },
              },
          },
          "epoch_loop.val_loop.state_dict": ANY,
          "epoch_loop.val_loop.batch_progress": ANY,
      }
      assert checkpoint["loops"]["fit_loop"] == expected
  
      trainer.fit_loop.load_state_dict(checkpoint["loops"]["fit_loop"])
      state_dict = trainer.fit_loop.state_dict()
  
      # need to remove these elements for comparison; comparing with `fit_loop.state_dict()` would require the
      # fit loop to have an iterator, which is only available during training
      state_dict["epoch_loop.state_dict"]["dataloader_state_dict"] = ANY
      checkpoint["loops"]["fit_loop"]["epoch_loop.state_dict"][
          "dataloader_state_dict"
      ] = ANY
      assert state_dict == checkpoint["loops"]["fit_loop"]
  
      trainer.fit_loop.load_state_dict(checkpoint["loops"]["fit_loop"])
      # test resetting manually, we expect all `ready` counters to be reset to `completed`
      trainer.fit_loop.reset()
      trainer.fit_loop.epoch_loop.reset()
  
      epoch_progress = trainer.fit_loop.epoch_progress
      assert epoch_progress.current.ready == stop_epoch
      assert epoch_progress.current.completed == stop_epoch
  
      batch_progress = trainer.fit_loop.epoch_loop.batch_progress
      assert batch_progress.current.ready == be_batches_completed
      assert batch_progress.current.completed == be_batches_completed
  
      optim_progress = (
          trainer.fit_loop.epoch_loop.automatic_optimization.optim_progress
      )
      assert (
          optim_progress.optimizer.step.current.ready
          == be_total_opt_steps
      )
      assert (
          optim_progress.optimizer.step.current.completed
          == be_total_opt_steps
      )
      assert (
          optim_progress.optimizer.zero_grad.current.ready
          == be_total_zero_grad
      )
      assert (
          optim_progress.optimizer.zero_grad.current.completed
          == be_total_zero_grad
      )
  
      state_dict = trainer.fit_loop.state_dict()
      assert state_dict != checkpoint["loops"]["fit_loop"]
      assert (
          state_dict["epoch_progress"]["total"]["started"]
          == stop_epoch + 1
      )
      assert (
          state_dict["epoch_progress"]["current"]["started"]
          == stop_epoch
      )
  
imports: import os
  from copy import deepcopy
  from dataclasses import dataclass
  from typing import Dict, Iterator
  from unittest.mock import ANY, Mock
  import pytest
  import torch
  from lightning.pytorch import LightningModule, Trainer
  from lightning.pytorch.callbacks import Callback, ModelCheckpoint, OnExceptionCheckpoint
  from lightning.pytorch.demos.boring_classes import BoringModel, RandomDataset
  from lightning.pytorch.loops import _Loop
  from lightning.pytorch.loops.progress import _BaseProgress
  from torch.utils.data.dataloader import DataLoader, _MultiProcessingDataLoaderIter
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 5 for calling a global function with example usage: test_loop_state_on_exception(accumulate_grad_batches, stop_epoch, stop_batch, tmpdir) and returns: super