code: def _test_dist_failure(self, callback, kwargs):
      bad_ranks = next(iter(kwargs.values())) if len(kwargs) > 0 else []
  
      # Empty bad_ranks means it must work
      if len(bad_ranks) == 0:
          callback()
      else:
          with self.assertRaises(CheckpointException) as cm:
              callback()
          e = cast(CheckpointException, cm.exception)
          for rank, wrapped_ex in e.failures.items():
              ex = wrapped_ex[0]
              self.assertTrue(
                  rank in bad_ranks, msg=f"{rank} did not fail"
              )
              if not kwargs.get("ignore_exception_type", False):
                  self.assertEqual(ValueError, type(ex), str(ex))
  
          failed_ranks = e.failures.keys()
          for rank in bad_ranks:
              self.assertTrue(
                  rank in failed_ranks,
                  msg=f"{rank} was supposed to fail was fine",
              )
  
imports: import sys
  from typing import cast, List, Optional
  import torch
  import torch.distributed
  import torch.futures
  import torch.nn
  from torch.distributed._shard import sharded_tensor
  from torch.distributed._shard.sharded_tensor import ShardedTensor, state_dict_hook
  from torch.distributed._shard.sharding_spec import ChunkShardingSpec
  from torch.distributed.checkpoint import CheckpointException, load_state_dict, save_state_dict, StorageReader, StorageWriter
  from torch.distributed.checkpoint.default_planner import _create_default_local_metadata
  from torch.distributed.checkpoint.metadata import BytesStorageMetadata, Metadata, TensorStorageMetadata
  from torch.distributed.checkpoint.planner import LoadPlan, LoadPlanner, SavePlan, SavePlanner
  from torch.distributed.checkpoint.storage import WriteResult
  from torch.futures import Future
  from torch.testing._internal.common_distributed import requires_nccl, skip_if_lt_x_gpu
  from torch.testing._internal.common_utils import run_tests, TEST_WITH_DEV_DBG_ASAN
  from torch.testing._internal.distributed._shard.sharded_tensor import ShardedTensorTestBase, with_comms
label_desc: Use image node_id 4 for calling the TestDistributedFailure obj's underlying member method code with example usage: obj._test_dist_failure(callback, kwargs) without return types