text: # Owner(s): ["oncall: distributed"]
  
  import copy
  import os
  import sys
  import tempfile
  import threading
  import time
  from datetime import timedelta
  from itertools import product
  from sys import platform
  
  import torch
  import torch.distributed as dist
  
  if not dist.is_available():
      print("distributed package not available, skipping tests", file=sys.stderr)
      sys.exit(0)
  
  import torch.distributed.distributed_c10d as c10d
  import torch.distributed.algorithms.ddp_comm_hooks.powerSGD_hook as powerSGD
  import torch.nn.functional as F
  import torch.testing._internal.common_utils as common
  from torch import nn
  from torch.nn.parallel import DistributedDataParallel
  from torch.testing._internal.common_distributed import (
      MultiProcessTestCase,
  )
  from torch.testing._internal.common_utils import (
      TestCase,
      load_tests,
      run_tests,
      TEST_WITH_DEV_DBG_ASAN,
  )
  
  if TEST_WITH_DEV_DBG_ASAN:
      print("Multiprocessing spawn is not compatible with dev/dbg asan", file=sys.stderr)
      sys.exit(0)
  
  # load_tests from common_utils is used to automatically filter tests for
  # sharding on sandcastle. This line silences flake warnings
  load_tests = load_tests
  
  if platform == "darwin":
      LOOPBACK = "lo0"
  else:
      LOOPBACK = "lo"
  
  torch.backends.cuda.matmul.allow_tf32 = False
  
  
  def gpus_for_rank(world_size):
      """Multigpu tests are designed to simulate the multi nodes with multi
      GPUs on each node. Nccl backend requires equal #GPUs in each process.
      On a single node, all visible GPUs are evenly
      divided to subsets, each process only uses a subset.
      """
      visible_devices = list(range(torch.cuda.device_count()))
      gpus_per_process = torch.cuda.device_count() // world_size
      gpus_for_rank = []
      for rank in range(world_size):
          gpus_for_rank.append(
              visible_devices[rank * gpus_per_process : (rank + 1) * gpus_per_process]
          )
      return gpus_for_rank
  
  
  class AbstractTimeoutTest(object):
      def _test_store_timeout(self, backend, init_method, c2p):
          try:
              dist.init_process_group(
                  backend=backend,
                  init_method=init_method,
                  world_size=1,
                  rank=0,
                  timeout=timedelta(seconds=1),
              )
              default_store = c10d._get_default_store()
              tik = time.time()
              with self.assertRaisesRegex(RuntimeError, "Timeout"):
                  default_store.get("nonexistent key")
              tok = time.time()
              dist.destroy_process_group()
              c2p.append(float(tok - tik))
          except RuntimeError as e:
              # catch "Address already in use" error and report it to the main
              # thread
              c2p.append(e)
  
      def _init_methods(self):
          f = tempfile.NamedTemporaryFile(delete=False)
          if sys.platform == "win32":
              yield "file:///%s" % f.name.replace("\\", "/")
              f.close()
          else:
              yield "file://%s" % f.name
              f.close()
              yield "tcp://127.0.0.1:%d" % common.find_free_port()
  
      def _test_default_store_timeout(self, backend):
          for init_method in self._init_methods():
              c2p = []
              t = threading.Thread(
                  target=self._test_store_timeout, args=(backend, init_method, c2p)
              )
              t.daemon = True
              t.start()
              t.join(5)
  
              self.assertEqual(1, len(c2p))
              if isinstance(c2p[0], float):
                  # waiting time should be 1s, use 3s to rule out false alarm
                  self.assertGreater(3, c2p[0])
              elif isinstance(c2p[0], RuntimeError):
                  # let @retry_on_connect_failures handle the error
                  raise c2p[0]
              else:
                  raise RuntimeError("Unexpected type {}".format(type(c2p[0])))
  
  
  class Net(nn.Module):
      def __init__(self):
          super(Net, self).__init__()
          self.fc1 = nn.Linear(2, 10, bias=False)
          self.fc2 = nn.Linear(10, 50, bias=False)
          self.fc3 = nn.Linear(50, 4, bias=False)
          self.relu = nn.ReLU()
  
      def forward(self, x):
          x = self.relu(self.fc1(x))
          x = self.relu(self.fc2(x))
          x = self.fc3(x)
          return F.softmax(x, dim=1)
  
  
  class DoubleGpuNet(nn.Module):
      def __init__(self, gpus):
          super(DoubleGpuNet, self).__init__()
          self.fc1 = nn.Linear(2, 10, bias=False).to(gpus[0])
          self.fc2 = nn.Linear(10, 50, bias=False).to(gpus[1])
          self.fc3 = nn.Linear(50, 4, bias=False).to(gpus[1])
          self.relu = nn.ReLU()
          self.no_grad_param = nn.Parameter(
              torch.tensor([2, 2]).long(), requires_grad=False
          ).to(gpus[0])
  
      def forward(self, x):
          dev0 = self.fc1.weight.device
          dev1 = self.fc2.weight.device
          x = self.relu(self.fc1(x.to(dev0)))
          x = self.relu(self.fc2(x.to(dev1)))
          x = self.fc3(x)
          return F.softmax(x, dim=1).to(dev0)
  
  
  class QuadraGpuNet(nn.Module):
      def __init__(self, gpus):
          super(QuadraGpuNet, self).__init__()
          self.fc1 = nn.Linear(2, 10, bias=False).to(gpus[0])
          self.fc2 = nn.Linear(10, 50, bias=False).to(gpus[1])
          self.fc3 = nn.Linear(50, 4, bias=False).to(gpus[2])
          self.fc4 = nn.Linear(4, 4, bias=False).to(gpus[3])
          self.relu = nn.ReLU()
          self.no_grad_param = nn.Parameter(
              torch.tensor([2, 2]).long(), requires_grad=False
          ).to(gpus[0])
  
      def forward(self, x):
          dev0 = self.fc1.weight.device
          dev1 = self.fc2.weight.device
          dev2 = self.fc3.weight.device
          dev3 = self.fc4.weight.device
          x = self.relu(self.fc1(x.to(dev0)))
          x = self.relu(self.fc2(x.to(dev1)))
          x = self.relu(self.fc3(x.to(dev2)))
          x = self.fc4(x.to(dev3))
          return F.softmax(x, dim=1).to(dev0)
  
  
  class ConvNet(nn.Module):
      def __init__(self, gpus, layouts, dtypes):
          super(ConvNet, self).__init__()
          self.dtypes = dtypes
          if isinstance(gpus, list):
              self.layer_gpus = gpus
          else:
              gpus = [gpus] * 4
          self.conv0 = torch.nn.Conv2d(8, 16, (2, 2)).to(
              device=gpus[0], memory_format=layouts[0], dtype=dtypes[0]
          )
          self.conv1 = torch.nn.Conv2d(16, 32, (2, 2)).to(
              device=gpus[1], memory_format=layouts[1], dtype=dtypes[1]
          )
          self.conv2 = torch.nn.Conv2d(32, 16, (2, 2)).to(
              device=gpus[2], memory_format=layouts[2], dtype=dtypes[2]
          )
          self.conv3 = torch.nn.Conv2d(16, 8, (2, 2)).to(
              device=gpus[3], memory_format=layouts[3], dtype=dtypes[3]
          )
  
      def forward(self, x):
          x = x.to(self.dtypes[0])
          # Could say
          # x = self.conv0(x).to(device=self.conv1.weight.device, dtype=self.dtypes[1])
          # etc.  But I don't want to appeal to the weights' devices directly, because part of this test's purpose
          # is to verify weights are where expected if the model gets replicated.
          gpus = self.layer_gpus if hasattr(self, "layer_gpus") else [x.device] * 4
          x = self.conv0(x).to(device=gpus[1], dtype=self.dtypes[1])
          x = self.conv1(x).to(device=gpus[2], dtype=self.dtypes[2])
          x = self.conv2(x).to(device=gpus[3], dtype=self.dtypes[3])
          return self.conv3(x)
  
  
  class Task(nn.Module):
      def __init__(self):
          super().__init__()
          self.p = nn.Parameter(torch.ones(2, 2))
  
      def forward(self, x):
          return self.p + x
  
  
  class ModuleForDdpCommHook(nn.Module):
      def __init__(self):
          super().__init__()
          self.t0 = Task()
  
      def forward(self, x, rank):
          return self.t0(x + rank)
  
  
  class SparseGradientModule(nn.Module):
      def __init__(self):
          super(SparseGradientModule, self).__init__()
          self.embedding = nn.EmbeddingBag(10, 10, sparse=True)
  
      def forward(self, x):
          return F.softmax(self.embedding(x), dim=1)
  
  
  class AbstractDistributedDataParallelTest(object):
      def tearDown(self):
          # DistributedDataParallel test doesn't seem to call FileStore destructor
          # TODO: investigate this test and the test is known to have issues
          # Use this hack to remove files for that test
          try:
              os.remove(self.file_name)
          except OSError:
              pass
  
      @property
      def world_size(self):
          return 2
  
      def _prepare_single_device_module(
          self,
          process_group,
          devices,
          device_ids,
          global_batch_size,
          gradient_as_bucket_view=False,
      ):
          model = Net()
          device = devices[0] if devices else torch.device("cuda:%d" % self.rank)
          ddp_model = DistributedDataParallel(
              copy.deepcopy(model).to(device),
              device_ids=device_ids,
              process_group=process_group,
              bucket_cap_mb=0.001,
              gradient_as_bucket_view=gradient_as_bucket_view,
          )
  
          model.to(device)
  
          input = torch.randn(global_batch_size, 2).to(device)
          target = torch.randn(global_batch_size, 4).to(device)
  
          return model, ddp_model, input, target
  
      def _prepare_multi_device_module(
          self,
          process_group,
          devices,
          device_ids,
          global_batch_size,
          gradient_as_bucket_view=False,
      ):
          self.assertTrue(
              len(devices) == 2 or len(devices) == 4,
              "unexpected devices for ddp tests {}".format(devices),
          )
          if len(devices) == 2:
              model = DoubleGpuNet(devices)
          elif len(devices) == 4:
              model = QuadraGpuNet(devices)
  
          ddp_model = DistributedDataParallel(
              copy.deepcopy(model),
              device_ids=device_ids,
              process_group=process_group,
              bucket_cap_mb=0.001,
              gradient_as_bucket_view=gradient_as_bucket_view,
          )
  
          input = torch.randn(global_batch_size, 2).cuda(devices[0])
          target = torch.randn(global_batch_size, 4)
  
          return model, ddp_model, input, target
  
      def _test_ddp_with_process_group(
          self,
          process_group,
          devices,
          device_ids,
          multi_device=False,
          gradient_as_bucket_view=False,
      ):
          """
          Note: we pass down `device_ids` all the way to DistributedDataParallel
          as part of the test. Below you find tests that either use a list of
          integers, a list of `torch.Device` instances, or an empty list.
          The `devices` argument is used to control placement of the model and
          must always be specified as list of `torch.Device` instances.
          """
          local_batch_size = 1 if devices is None else len(devices)
          global_batch_size = self.world_size * local_batch_size
  
          if multi_device:
              model, ddp_model, input, target = self._prepare_multi_device_module(
                  process_group,
                  devices,
                  device_ids,
                  global_batch_size,
                  gradient_as_bucket_view,
              )
              ddp_logging_data = ddp_model._get_ddp_logging_data()
              self.assertTrue(ddp_logging_data.get("is_multi_device_module"))
          else:
              model, ddp_model, input, target = self._prepare_single_device_module(
                  process_group,
                  devices,
                  device_ids,
                  global_batch_size,
                  gradient_as_bucket_view,
              )
              ddp_logging_data = ddp_model._get_ddp_logging_data()
              self.assertFalse(ddp_logging_data.get("is_multi_device_module"))
  
          def step_model(model, input, target):
              model.train()
              output = model(input)
              loss = F.mse_loss(output, target.to(output.device))
              loss.backward()
  
          def update_parameters(model):
              for param in model.parameters():
                  with torch.no_grad():
                      param -= param.grad
                  param.grad = None
  
          # check two model parameters over 2 iterations
          for iteration in range(2):
              # single cpu/gpu training
              step_model(model, input, target)
  
              # DDP training, DDP scatters subsets of input_cpu to nodes/GPUs
              step_model(
                  ddp_model,
                  input[
                      self.rank * local_batch_size : (self.rank + 1) * local_batch_size
                  ],
                  target[
                      self.rank * local_batch_size : (self.rank + 1) * local_batch_size
                  ],
              )
  
              # Update weights and run a second iteration to shake out errors
              update_parameters(model)
              update_parameters(ddp_model)
              self.assertEqual(
                  len(list(model.parameters())), len(list(ddp_model.parameters()))
              )
              for i, j in zip(model.parameters(), ddp_model.parameters()):
                  self.assertEqual(i, j, rtol=1.3e-06, atol=5e-5)
  
              # Shuffle the input so that DDP input is different
              torch.manual_seed(1337 + iteration)
              input = input[torch.randperm(global_batch_size)]
  
      def _gpu_model_with_ddp_comm_hook(
          self, process_group, hook=None, gradient_as_bucket_view=False, state=None
      ):
          device_id = gpus_for_rank(self.world_size)[self.rank][0]
          gpu_model = DistributedDataParallel(
              ModuleForDdpCommHook().to(device_id),
              device_ids=[device_id],
              process_group=process_group,
              gradient_as_bucket_view=gradient_as_bucket_view,
          )
  
          # Register a DDP communication hook if any.
          if hook is not None:
              gpu_model.register_comm_hook(state, hook)
  
          return gpu_model
  
      def _gpu_model_with_builtin_ddp_comm_hook(
          self, process_group, hook=None, gradient_as_bucket_view=False
      ):
          device_id = gpus_for_rank(self.world_size)[self.rank][0]
          gpu_model = DistributedDataParallel(
              ModuleForDdpCommHook().to(device_id),
              device_ids=[device_id],
              process_group=process_group,
              gradient_as_bucket_view=gradient_as_bucket_view,
          )
  
          # Register a built-in DDP communication hook if defined
          if hook is not None:
              gpu_model._register_builtin_comm_hook(hook)
  
          return gpu_model
  
      def _run_and_verify_hook(self, model, input, expected_grad):
          # Run forward
          output = model(input, self.rank)
  
          # Run backward
          output.mean().backward()
  
          [self.assertEqual(p.grad, expected_grad) for p in model.parameters()]
  
      def _simple_hook(
          self, state: object, bucket: dist.GradBucket
      ) -> torch.futures.Future[torch.Tensor]:
          fut = torch.futures.Future()
          fut.set_result(torch.ones_like(bucket.buffer()))
  
          def fut_then(fut):
              # Add ones to fut's result.
              t = fut.value()
              return t + torch.ones_like(t)
  
          return fut.then(fut_then)
  
  
  class DistributedDataParallelTest(
      AbstractDistributedDataParallelTest, MultiProcessTestCase
  ):
      def setUp(self):
          super(DistributedDataParallelTest, self).setUp()
          self._spawn_processes()
  
      def test_invalid_powerSGD_state(self):
          for start_powerSGD_iter, use_error_feedback, warm_start in product(
              [0, 1], [True, False], [True, False]
          ):
              if not use_error_feedback and not warm_start:
                  continue
              with self.assertRaisesRegex(
                  ValueError,
                  "Expect `start_powerSGD_iter` > 1 if `use_error_feedback` or `warm_start` is enabled, "
                  "because PowerSGD can only be applied after the first two iterations in DDP.",
              ):
                  state = powerSGD.PowerSGDState(
                      process_group=None,
                      matrix_approximation_rank=1,
                      start_powerSGD_iter=start_powerSGD_iter,
                      use_error_feedback=use_error_feedback,
                      warm_start=warm_start,
                  )
  
  
  class ComputeBucketAssignmentTest(TestCase):
      def test_single_limit_single_dtype(self):
          tensors = [
              torch.empty([100], dtype=torch.float),
              torch.empty([200], dtype=torch.float),
              torch.empty([100], dtype=torch.float),
              torch.empty([50], dtype=torch.float),
          ]
          result, per_bucket_size_limits = dist._compute_bucket_assignment_by_size(
              tensors, [400]
          )
          self.assertTrue(all(size_lim == 400 for size_lim in per_bucket_size_limits))
          self.assertEqual([[0], [1], [2], [3]], result)
  
      def test_single_limit_multi_dtype(self):
          tensors = [
              torch.empty([50], dtype=torch.float),
              torch.empty([25], dtype=torch.double),
              torch.empty([50], dtype=torch.float),
              torch.empty([25], dtype=torch.double),
              torch.empty([50], dtype=torch.float),
              torch.empty([25], dtype=torch.double),
          ]
          result, per_bucket_size_limits = dist._compute_bucket_assignment_by_size(
              tensors, [400]
          )
          self.assertTrue(all(size_lim == 400 for size_lim in per_bucket_size_limits))
          self.assertEqual([[0, 2], [1, 3], [4], [5]], result)
  
      def test_multi_limit_single_dtype(self):
          tensors = [
              torch.empty([10], dtype=torch.float),
              torch.empty([10], dtype=torch.float),
              torch.empty([10], dtype=torch.float),
              torch.empty([10], dtype=torch.float),
          ]
          result, per_bucket_size_limits = dist._compute_bucket_assignment_by_size(
              tensors, [40, 80]
          )
          self.assertEqual(per_bucket_size_limits, [40, 80, 80])
          self.assertEqual([[0], [1, 2], [3]], result)
  
      def test_multi_limit_multi_dtype(self):
          tensors = [
              torch.empty([50], dtype=torch.float),
              torch.empty([25], dtype=torch.double),
              torch.empty([50], dtype=torch.float),
              torch.empty([25], dtype=torch.double),
              torch.empty([50], dtype=torch.float),
              torch.empty([25], dtype=torch.double),
          ]
          result, per_bucket_size_limits = dist._compute_bucket_assignment_by_size(
              tensors, [200, 400]
          )
          self.assertEqual([[0], [1], [2, 4], [3, 5]], result)
          self.assertEqual(per_bucket_size_limits, [200, 200, 400, 400])
  
  
  class AbstractCommTest(object):
      @property
      def op_timeout_sec(self):
          return 1
  
      @property
      def world_size(self):
          return 2
  
      def _verify_sequence_number_across_pg(self, pg, verify_pg):
  
          seq_num = pg._get_sequence_number_for_group()
          obj_list = [None for _ in range(dist.get_world_size(verify_pg))]
          # We use a separate pg to verify the sequence numbers, otherwise these
          # collectives will themselves increment the sequence number.
          dist.all_gather_object(obj_list, seq_num, group=verify_pg)
          self.assertEqual(len(set(obj_list)), 1)
          return obj_list[0]
  
      def _test_sequence_num_incremented(self, process_group, ranks):
          # verify initial sequence numbers. Use a distinct process group for
          # verification to keep counts as expected with respect to process_group.
          verify_pg = dist.new_group(
              ranks=ranks,
              backend="gloo",
          )
          assert dist.get_world_size(process_group) == dist.get_world_size(verify_pg)
  
          initial_num = (
              self._verify_sequence_number_across_pg(
                  pg=process_group, verify_pg=verify_pg
              )
              if not c10d._rank_not_in_group(process_group)
              else -1
          )
  
          # Verify sequence numbers are appropriately incremented
          for i in range(10):
              t = torch.ones(1, device=torch.cuda.current_device())
              dist.all_reduce(t, group=process_group)
              if not c10d._rank_not_in_group(process_group):
                  seq_num = self._verify_sequence_number_across_pg(
                      pg=process_group,
                      verify_pg=verify_pg,
                  )
                  self.assertEqual(initial_num + i + 1, seq_num)
  
          if dist.get_world_size(process_group) > 2:
              # Test when certain ranks don't call collectives
              if dist.get_rank(process_group) not in [0, 2]:
                  dist.all_reduce(t, group=process_group, async_op=True)
              # Now ranks 0 and 2 should be lagging by 1.
              if not c10d._rank_not_in_group(process_group):
                  seq_num = process_group._get_sequence_number_for_group()
                  rank = dist.get_rank(process_group)
                  obj_list = [None for _ in range(dist.get_world_size(verify_pg))]
                  dist.all_gather_object(obj_list, (rank, seq_num), group=verify_pg)
                  rank_to_seq_num = {rank: num for (rank, num) in obj_list}
                  self.assertEqual(len(set(rank_to_seq_num.values())), 2)
                  self.assertEqual(rank_to_seq_num[0], rank_to_seq_num[2])
                  expected_same = {
                      rank_to_seq_num[i]
                      for i in rank_to_seq_num.keys()
                      if i not in [0, 2]
                  }
                  self.assertEqual(len(expected_same), 1)
                  self.assertEqual(rank_to_seq_num[0] + 1, rank_to_seq_num[1])
  
      def _test_sequence_num_incremented_default_group(self, backend_name):
          torch.cuda.set_device(self.rank)
          store = dist.FileStore(self.file_name, self.world_size)
          dist.init_process_group(
              backend_name,
              world_size=self.world_size,
              rank=self.rank,
              store=store,
          )
          self._test_sequence_num_incremented(
              c10d._get_default_group(),
              ranks=list(i for i in range(dist.get_world_size())),
          )
  
      def _test_sequence_num_incremented_subgroup(self, backend_name):
          torch.cuda.set_device(self.rank)
          store = dist.FileStore(self.file_name, self.world_size)
          dist.init_process_group(
              backend_name,
              world_size=self.world_size,
              rank=self.rank,
              store=store,
          )
          subgroup_ranks = [0, 1, 2]
          subgroup = dist.new_group(subgroup_ranks)
          self._test_sequence_num_incremented(subgroup, subgroup_ranks)
  
      def _test_sequence_num_set_default_pg(self, backend):
          store = dist.FileStore(self.file_name, self.world_size)
          dist.init_process_group(
              backend,
              world_size=self.world_size,
              rank=self.rank,
              store=store,
          )
  
          default_pg = c10d._get_default_group()
          seq_num = default_pg._get_sequence_number_for_group()
          obj_list = [None for _ in range(dist.get_world_size())]
          dist.all_gather_object(obj_list, seq_num)
          self.assertEqual(len(set(obj_list)), 1)
  
      def _test_sequence_num_set_new_group(self, backend):
          store = dist.FileStore(self.file_name, self.world_size)
          dist.init_process_group(
              backend,
              world_size=self.world_size,
              rank=self.rank,
              store=store,
          )
  
          subgroup = dist.new_group([0, 1])
  
          if not c10d._rank_not_in_group(subgroup):
              subgroup_seq = subgroup._get_sequence_number_for_group()
              obj_list = [None for _ in range(dist.get_world_size(subgroup))]
              dist.all_gather_object(obj_list, subgroup_seq, group=subgroup)
              self.assertEqual(len(set(obj_list)), 1)
  
      def _test_warn_not_in_group(self, backend):
          store = dist.FileStore(self.file_name, self.world_size)
          dist.init_process_group(
              backend,
              world_size=self.world_size,
              rank=self.rank,
              store=store,
          )
          in_group_ranks = list(filter(lambda x: x % 2 == 0, range(self.world_size)))
          group = dist.new_group(in_group_ranks)
  
          x = torch.zeros(2, 2).cuda(self.rank)
          xs = [torch.zeros(2, 2).cuda(self.rank) for _ in range(len(in_group_ranks))]
          if self.rank not in in_group_ranks:
              msg = ".*{}.*does not belong to.*"
              with self.assertWarnsOnceRegex(UserWarning, msg.format("all_gather")):
                  dist.all_gather(xs, x, group=group)
              with self.assertWarnsOnceRegex(UserWarning, msg.format("all_reduce")):
                  dist.all_reduce(x, group=group)
              with self.assertWarnsOnceRegex(UserWarning, msg.format("barrier")):
                  dist.barrier(group=group)
              with self.assertWarnsOnceRegex(UserWarning, msg.format("broadcast")):
                  dist.broadcast(x, src=0, group=group)
          else:
              dist.all_gather(xs, x, group=group)
              dist.all_reduce(x, group=group)
              dist.barrier(group=group)
              dist.broadcast(x, src=0, group=group)
  
  
  class CommTest(AbstractCommTest, MultiProcessTestCase):
      def setUp(self):
          super(CommTest, self).setUp()
          self._spawn_processes()
  
      def tearDown(self):
          super(CommTest, self).tearDown()
          try:
              os.remove(self.file_name)
          except OSError:
              pass
  
      def test_debug_level(self):
          try:
              del os.environ["TORCH_DISTRIBUTED_DEBUG"]
          except KeyError:
              pass
  
          dist.set_debug_level_from_env()
          # Default should be off
          default_debug_mode = dist.get_debug_level()
          self.assertEqual(default_debug_mode, dist.DebugLevel.OFF)
          mapping = {
              "OFF": dist.DebugLevel.OFF,
              "off": dist.DebugLevel.OFF,
              "oFf": dist.DebugLevel.OFF,
              "INFO": dist.DebugLevel.INFO,
              "info": dist.DebugLevel.INFO,
              "INfO": dist.DebugLevel.INFO,
              "DETAIL": dist.DebugLevel.DETAIL,
              "detail": dist.DebugLevel.DETAIL,
              "DeTaIl": dist.DebugLevel.DETAIL,
          }
          invalid_debug_modes = ["foo", 0, 1, -1]
  
          for mode in mapping.keys():
              os.environ["TORCH_DISTRIBUTED_DEBUG"] = str(mode)
              dist.set_debug_level_from_env()
              set_debug_mode = dist.get_debug_level()
              self.assertEqual(
                  set_debug_mode,
                  mapping[mode],
                  f"Expected {mode} to map to {mapping[mode]} but got {set_debug_mode}",
              )
  
          for mode in invalid_debug_modes:
              os.environ["TORCH_DISTRIBUTED_DEBUG"] = str(mode)
              with self.assertRaisesRegex(RuntimeError, "The value of TORCH_DISTRIBUTED_DEBUG must"):
                  dist.set_debug_level_from_env()
  
  
  class DummyWork(dist._Work):
      def wait(self, timeout=5.0):
          if torch.cuda.is_available():
              torch.cuda.current_stream().synchronize()
          return True
  
  
  class DummyProcessGroup(dist.ProcessGroup):
      def getBackendName(self):
          return "Dummy"
  
      def allgather(self, output_tensor_lists, input_tensor_list, opts=None):
          for output_tensor_list, input_tensor in zip(output_tensor_lists, input_tensor_list):
              for output_tensor in output_tensor_list:
                  output_tensor.copy_(input_tensor)
  
          return DummyWork()
  
      def allreduce(self, tensor_list, opts=None):
          for tensor in tensor_list:
              tensor.add_(2)
  
          return DummyWork()
  
      def barrier(self, opts=None):
          store = c10d._get_default_store()
          key = "TEST:DummyProcessGroup:barrier"
          if self.rank() == 0:
              worker_count = 0
              # By default, TCPServer lives on rank 0. So rank 0 needs to make
              # sure that it does not exit too early before other ranks finish
              # using the store.
              # Note that, _store_based_barrier does not solve this problem, as
              # all ranks need to run at least one store.add(key, 0) before
              # exiting, but there is no guarantee that rank 0 is still alive at
              # that point.
              while worker_count < self.size() - 1:
                  worker_count = store.add(key, 0)
          else:
              store.add(key, 1)
  
          return DummyWork()
  
      def broadcast(self, tensor_list, opts=None):
          for tensor in tensor_list:
              tensor.add_(1)
  
          return DummyWork()
  
      def reduce_scatter(self, output_tensor_list, input_tensor_lists, opts=None):
          for output_tensor, input_tensor_list in zip(output_tensor_list, input_tensor_lists):
              output_tensor.copy_(input_tensor_list[self.rank()])
  
          return DummyWork()
  
      def send(self, tensor_list, dst, tag=0):
          for tensor in tensor_list:
              tensor.add_(1)
  
          return DummyWork()
  
      def recv(self, tensor_list, src, tag=0):
          for tensor in tensor_list:
              tensor.add_(2)
  
          return DummyWork()
  
  
  class PythonProcessGroupExtensionTest(MultiProcessTestCase):
      def setUp(self):
          super(PythonProcessGroupExtensionTest, self).setUp()
          self._spawn_processes()
  
      def tearDown(self):
          super(PythonProcessGroupExtensionTest, self).tearDown()
          try:
              os.remove(self.file_name)
          except OSError:
              pass
  
      def test_get_backend_name(self):
          dpg = DummyProcessGroup(0, 1)
          self.assertEqual("Dummy", dpg.name())
  
      def test_backend_class_attr(self):
          dist.Backend.register_backend(
              "dummy",
              PythonProcessGroupExtensionTest.create_dummy
          )
          self.assertEqual(dist.Backend.DUMMY, "DUMMY")
          self.assertEqual(
              dist.Backend._plugins["DUMMY"],
              PythonProcessGroupExtensionTest.create_dummy
          )
  
      @staticmethod
      def create_dummy(store, rank, size, timeout):
          return DummyProcessGroup(rank, size)
  
      def test_collectives(self):
          dist.Backend.register_backend("dummy", PythonProcessGroupExtensionTest.create_dummy)
  
          os.environ['MASTER_ADDR'] = 'localhost'
          os.environ['MASTER_PORT'] = '6789'
          dist.init_process_group("dummy", rank=self.rank, world_size=self.world_size)
  
          # test all_gather
          input_tensor = torch.ones(2, 2) * 7
          output_tensor_list = [torch.zeros(2, 2) for _ in range(self.world_size)]
          dist.all_gather(output_tensor_list, input_tensor)
  
          for tensor in output_tensor_list:
              self.assertEqual(tensor, input_tensor)
  
          # test all_reduce
          input_tensor = torch.ones(2, 2) * 7
          dist.all_reduce(input_tensor)
          self.assertEqual(input_tensor, torch.ones(2, 2) * 7 + 2)
  
          # test broadcast
          input_tensor = torch.zeros(2, 2)
          dist.broadcast(input_tensor, 0, async_op=True).wait()
          self.assertEqual(torch.ones(2, 2), input_tensor)
  
          # test reduce_scatter
          output_tensor = torch.zeros(2, 2)
          input_tensor_list = [torch.ones(2, 2) for _ in range(self.world_size)]
          dist.reduce_scatter(output_tensor, input_tensor_list)
          self.assertEqual(output_tensor, torch.zeros(2, 2) + 1)
  
          dist.barrier()
          dist.destroy_process_group()
  
      def test_send_recv(self):
          dist.Backend.register_backend("dummy", PythonProcessGroupExtensionTest.create_dummy)
  
          os.environ['MASTER_ADDR'] = 'localhost'
          os.environ['MASTER_PORT'] = '6789'
          dist.init_process_group("dummy", rank=self.rank, world_size=self.world_size)
  
          # test send
          input_tensor = torch.zeros(2, 2)
          dist.send(input_tensor, (self.rank + 1) % self.world_size)
          self.assertEqual(input_tensor, torch.zeros(2, 2) + 1)
  
          # test recv
          input_tensor = torch.zeros(2, 2)
          dist.recv(input_tensor, (self.rank + 1) % self.world_size)
          self.assertEqual(input_tensor, torch.zeros(2, 2) + 2)
  
          dist.barrier()
          # intentionally not calling into `destroy_process_group` as not all
          # user applications would explicitly that.
  
  
  
  if __name__ == "__main__":
      assert (
          not torch.cuda._initialized
      ), "test_distributed must not have initialized CUDA context on main process"
  
      run_tests()
  