code: def test_workers_are_shutdown(
      tmpdir, should_fail, persistent_workers
  ):
      # `num_workers == 1` uses `_MultiProcessingDataLoaderIter`
      # `persistent_workers` makes sure `self._iterator` gets set on the `DataLoader` instance
  
      class TestCallback(Callback):
          def on_train_epoch_end(self, trainer, *_):
              if trainer.current_epoch == 1:
                  raise CustomException
  
      max_epochs = 3
  
      model = BoringModel()
      trainer = Trainer(
          default_root_dir=tmpdir,
          limit_train_batches=2,
          limit_val_batches=2,
          max_epochs=max_epochs,
          callbacks=TestCallback() if should_fail else None,
          enable_checkpointing=False,
          enable_model_summary=False,
          enable_progress_bar=False,
          logger=False,
      )
  
      class _TestMultiProcessingDataLoaderIter(
          _MultiProcessingDataLoaderIter
      ):
          def __init__(self, *args, dataloader, **kwargs):
              super().__init__(*args, **kwargs)
              self.dataloader = dataloader
  
          def _shutdown_workers(self):
              self.dataloader.shutdown_workers_epochs.append(
                  trainer.current_epoch
              )
              super()._shutdown_workers()
  
      class TestDataLoader(DataLoader):
          def __init__(self, *args, **kwargs):
              super().__init__(*args, **kwargs)
              self.shutdown_workers_epochs = []
  
          def _get_iterator(self):
              if self.num_workers == 0:
                  return super()._get_iterator()
  
              self.check_worker_number_rationality()
              return _TestMultiProcessingDataLoaderIter(
                  self, dataloader=self
              )
  
      train_dataloader = TestDataLoader(
          RandomDataset(32, 64),
          num_workers=1,
          persistent_workers=persistent_workers,
      )
      val_dataloader = TestDataLoader(
          RandomDataset(32, 64),
          num_workers=1,
          persistent_workers=persistent_workers,
      )
  
      if should_fail:
          with pytest.raises(CustomException):
              trainer.fit(model, train_dataloader, val_dataloader)
      else:
          trainer.fit(model, train_dataloader, val_dataloader)
  
      if persistent_workers:
          # workers get created and persist until the teardown in the final epoch
          expected = [
              trainer.current_epoch,
              trainer.current_epoch,
          ]  # once epoch end, once on teardown
      elif should_fail:
          expected = [
              # <-- iter() on epoch 0, workers get created
              1,  # iter() on epoch 1, workers from epoch 0 get destroyed
              1,  # teardown on failed epoch 1, workers from epoch 1 get destroyed
          ]
      else:
          expected = [
              # <-- iter() on epoch 0, workers get created
              1,  # iter() on epoch 1, workers from epoch 0 get destroyed
              2,  # iter() on epoch 2, workers from epoch 1 get destroyed
              3,  # teardown on epoch 2, workers from epoch 2 get destroyed
          ]
      assert train_dataloader.shutdown_workers_epochs == expected
  
      if persistent_workers:
          # workers get created and persist until the teardown in the final epoch
          expected = [
              trainer.current_epoch,
              trainer.current_epoch,
          ]  # once epoch end, once on teardown
      elif should_fail:
          expected = [
              # <-- iter() on sanity check, workers get created
              0,  # iter() on epoch 0, workers from sanity check get destroyed
              1,  # iter() on epoch 1, workers from epoch 0 get destroyed
              1,  # teardown on failed epoch 1, workers from epoch 1 get destroyed
          ]
      else:
          expected = [
              # <-- iter() on sanity check, workers get created
              0,  # iter() on epoch 0, workers from sanity check get destroyed
              1,  # iter() on epoch 1, workers from epoch 0 get destroyed
              2,  # iter() on epoch 2, workers from epoch 1 get destroyed
              3,  # teardown on epoch 2, workers from epoch 2 get destroyed
          ]
      assert val_dataloader.shutdown_workers_epochs == expected
  
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 9 for calling a global function with example usage: test_workers_are_shutdown(tmpdir, should_fail, persistent_workers) and returns: _TestMultiProcessingDataLoaderIter, super