code: def test_async_checkpoint_plugin(tmpdir):
      """Ensure that the custom checkpoint IO plugin and torch checkpoint IO plugin is called when async saving and
      loading."""
  
      checkpoint_plugin = AsyncCheckpointIO()
  
      checkpoint_plugin.save_checkpoint = Mock(
          wraps=checkpoint_plugin.save_checkpoint
      )
      checkpoint_plugin.remove_checkpoint = Mock(
          wraps=checkpoint_plugin.remove_checkpoint
      )
  
      class CustomBoringModel(BoringModel):
          def on_fit_start(self):
              base_ckpt_io = (
                  self.trainer.strategy.checkpoint_io.checkpoint_io
              )
              base_ckpt_io.save_checkpoint = Mock(
                  wraps=base_ckpt_io.save_checkpoint
              )
              base_ckpt_io.remove_checkpoint = Mock(
                  wraps=base_ckpt_io.remove_checkpoint
              )
  
      ck = ModelCheckpoint(
          dirpath=tmpdir, save_top_k=2, monitor="step", mode="max"
      )
  
      model = CustomBoringModel()
      trainer = Trainer(
          default_root_dir=tmpdir,
          plugins=[checkpoint_plugin],
          callbacks=ck,
          max_epochs=3,
          limit_train_batches=1,
          limit_val_batches=0,
          enable_progress_bar=False,
          enable_model_summary=False,
      )
      trainer.fit(model)
  
      assert checkpoint_plugin.save_checkpoint.call_count == 3
      assert checkpoint_plugin.remove_checkpoint.call_count == 1
  
      base_ckpt_io = trainer.strategy.checkpoint_io.checkpoint_io
      assert base_ckpt_io.save_checkpoint.call_count == 3
      assert base_ckpt_io.remove_checkpoint.call_count == 1
  
file_path: pytorch-lightning/tests/tests_pytorch/plugins/test_checkpoint_io_plugin.py
imports: import os
  from pathlib import Path
  from typing import Any, Dict, Optional
  from unittest.mock import MagicMock, Mock
  import torch
  from lightning.fabric.plugins import CheckpointIO, TorchCheckpointIO
  from lightning.fabric.utilities.types import _PATH
  from lightning.pytorch import Trainer
  from lightning.pytorch.callbacks import ModelCheckpoint
  from lightning.pytorch.demos.boring_classes import BoringModel
  from lightning.pytorch.plugins.io.async_plugin import AsyncCheckpointIO
  from lightning.pytorch.strategies import SingleDeviceStrategy
label_desc: Use image node_id 2 for calling a global function with example usage: test_async_checkpoint_plugin(tmpdir) without return types