code: def test_fault_tolerance(fault0: Fault, fault1: Fault):
      def _make_tensors():
          return [
              torch.rand(16, 1024),
              -torch.rand(3, 8192),
              2 * torch.randn(4, 4, 4),
              torch.randn(1024, 1024),
          ]
  
      dht = hivemind.DHT(start=True)
  
      averagers = []
      for i in range(5):
          averager = FaultyAverager(
              _make_tensors(),
              hivemind.DHT(
                  initial_peers=dht.get_visible_maddrs(), start=True
              ),
              prefix="test",
              request_timeout=0.3,
              min_matchmaking_time=1.0,
              next_chunk_timeout=0.5,
              allreduce_timeout=5,
              part_size_bytes=2**16,
              client_mode=(i == 1),
              start=True,
              fault=fault0
              if i == 0
              else fault1
              if i == 1
              else Fault.NONE,
          )
          averagers.append(averager)
  
      ref_numerators = [0, 0, 0, 0]
      ref_denominator = 0
  
      for averager in averagers:
          if averager.fault not in (Fault.FAIL_BEFORE, Fault.CANCEL):
              with averager.get_tensors() as tensors:
                  for i, tensor in enumerate(tensors):
                      ref_numerators[i] = (
                          ref_numerators[i] + tensor.clone()
                      )
                  ref_denominator += 1
  
      ref_tensors = [
          ref_numerator / ref_denominator
          for ref_numerator in ref_numerators
      ]
      flat_ref = torch.cat(list(map(torch.flatten, ref_tensors)))
  
      flat_local_tensors = []
      for averager in averagers:
          with averager.get_tensors() as tensors:
              flat_local_tensors.append(
                  torch.cat(list(map(torch.flatten, tensors)))
              )
  
      futures = [
          averager.step(timeout=5, wait=False, allow_retries=False)
          for averager in averagers
      ]
      for i, averager in enumerate(averagers):
          if averager.fault == Fault.CANCEL:
              futures[i].cancel()
  
      for future in futures[2:]:
          assert future.result()
  
      for averager, prev_local_tensors in zip(
          averagers[2:], flat_local_tensors[2:]
      ):
          with averager.get_tensors() as tensors:
              flat_tensors = torch.cat(
                  list(map(torch.flatten, tensors))
              )
  
          diff_with_reference = abs(flat_ref - flat_tensors)
  
          if all(
              fault == (Fault.FAIL_SENDING, Fault.SLOW_SENDING)
              for fault in (fault0, fault1)
          ):
              assert (
                  fault0 != Fault.FAIL_REDUCING
                  and fault1 != Fault.FAIL_REDUCING
              )
              assert (
                  diff_with_reference[
                      : len(diff_with_reference) // 2
                  ].max()
                  < 1e-5
              )
          elif all(
              fault in (Fault.FAIL_REDUCING, Fault.SLOW_REDUCING)
              for fault in (fault0, fault1)
          ):
              diff_to_reference = abs(flat_ref - flat_tensors)
              diff_to_local = abs(prev_local_tensors - flat_tensors)
              assert (diff_with_reference < 1e-5).numpy().mean() > 0.5
              assert torch.all(
                  torch.minimum(diff_to_reference, diff_to_local) < 1e-5
              ).item()
          elif any(fault == Fault.CANCEL for fault in (fault0, fault1)):
              pass  # late cancel may result in an arbitrary mix of averaging results with and without the cancelled peer
          elif (
              fault0 == Fault.NONE
          ):  # only peer1 in client mode may have failed
              assert diff_with_reference.max() < 1e-5
          else:
              assert (diff_with_reference < 1e-5).numpy().mean() > 0.5
  
      for averager in averagers:
          averager.shutdown()
  
imports: from __future__ import annotations
  from enum import Enum, auto
  import pytest
  import hivemind
  from hivemind.averaging.averager import *
  from hivemind.averaging.group_info import GroupInfo
  from hivemind.averaging.load_balancing import load_balance_peers
  from hivemind.averaging.matchmaking import MatchmakingException
  from hivemind.proto import averaging_pb2
  from hivemind.utils.asyncio import aenumerate, as_aiter, azip, enter_asynchronously
  from hivemind.utils.logging import get_logger
label_desc: Use image node_id 1 for calling a global function with example usage: test_fault_tolerance(fault0, fault1) and returns: list