text: import unittest
  import torch
  import os
  from bagua.torch_api.communication import (
      init_bagua_communicator,
      allreduce,
      send,
      recv,
      allgather,
  )
  from tests.internal.common_utils import find_free_port
  import multiprocessing
  import bagua.torch_api as bagua
  import threading
  import time
  from tests import skip_if_cuda_not_available
  
  
  class Result(object):
      def __init__(self):
          self.ret = torch.Tensor([False]).bool()
          self.data = torch.Tensor([0.0])
  
  
  def init_env(rank, env):
      os.environ["WORLD_SIZE"] = env["WORLD_SIZE"]
      os.environ["LOCAL_WORLD_SIZE"] = env["LOCAL_WORLD_SIZE"]
      os.environ["MASTER_ADDR"] = env["MASTER_ADDR"]
      os.environ["MASTER_PORT"] = env["MASTER_PORT"]
      os.environ["BAGUA_SERVICE_PORT"] = env["BAGUA_SERVICE_PORT"]
  
      os.environ["RANK"] = str(rank)
      os.environ["LOCAL_RANK"] = str(rank)
  
      # init bagua distributed process group
      torch.cuda.set_device(rank)
      bagua.init_process_group()
  
  
  def run_abort(rank, nprocs, results, env):
      init_env(rank, env)
  
      os.environ["NCCL_PROTO"] = "^LL128"
  
      comm_stream = torch.cuda.Stream()
      comm = init_bagua_communicator(model_name="test_comm", stream=comm_stream)
  
      def abort():
          time.sleep(10)
          comm.abort()
  
      threading.Thread(target=abort).start()
  
      with torch.cuda.stream(comm_stream):
          data = torch.rand(10).cuda()
  
      for _ in range(rank + 1):
          comm.allreduce_inplace(data.to_bagua_tensor().bagua_backend_tensor(), 10)
  
      comm_stream.synchronize()
  
  
  def run_allreduce(rank, nprocs, results, env):
      init_env(rank, env)
  
      send_tensor = torch.rand(100).cuda()
      recv_tensor = torch.zeros_like(send_tensor)
  
      tensor = send_tensor.clone()
  
      allreduce(send_tensor, recv_tensor)
  
      torch.distributed.all_reduce(tensor)
  
      results[rank].ret[0] = torch.equal(recv_tensor, tensor)
  
  
  def run_p2p(rank, nprocs, results, env):
      init_env(rank, env)
  
      send_tensor = torch.rand(100).cuda()
      recv_tensor = torch.zeros_like(send_tensor)
  
      if rank % 2 == 0:
          send(send_tensor, dst=(rank + 1) % nprocs)
          results[rank].data.copy_(torch.norm(send_tensor))
      else:
          recv(recv_tensor, src=(rank - 1 + nprocs) % nprocs)
          results[rank].data.copy_(torch.norm(recv_tensor))
  
  
  def run_allgather(rank, nprocs, results, env):
      init_env(rank, env)
  
      send_tensor = torch.rand(100).cuda()
      recv_tensor = torch.zeros(
          [nprocs, 100], device=send_tensor.device, dtype=send_tensor.dtype
      )
  
      tensor = send_tensor.clone()
      tensor_list = [torch.zeros_like(tensor) for _ in range(nprocs)]
  
      allgather(send_tensor, recv_tensor)
  
      torch.distributed.all_gather(tensor_list, tensor)
  
      ret = True
      for i in range(nprocs):
          ret = ret and torch.equal(recv_tensor[i], tensor_list[i])
  
      results[rank].ret[0] = ret
  
  
  class TestCommunication(unittest.TestCase):
      def run_test_locally(self, fn):
          nprocs = torch.cuda.device_count()
          env = {
              "WORLD_SIZE": str(nprocs),
              "LOCAL_WORLD_SIZE": str(nprocs),
              "MASTER_ADDR": "127.0.0.1",
              "MASTER_PORT": str(find_free_port(8000, 8100)),
              "BAGUA_SERVICE_PORT": str(find_free_port(9000, 9100)),
          }
  
          mp = multiprocessing.get_context("spawn")
          results = [Result() for _ in range(nprocs)]
          processes = []
          for i in range(nprocs):
              p = mp.Process(
                  target=fn,
                  args=(i, nprocs, results, env),
              )
              p.start()
              processes.append(p)
  
          for p in processes:
              p.join(timeout=60)
              self.assertTrue(p.exitcode == 0)
  
          return results
  
      @skip_if_cuda_not_available()
      def test_abort(self):
          self.run_test_locally(run_abort)
  
      @skip_if_cuda_not_available()
      def test_allreduce(self):
          results = self.run_test_locally(run_allreduce)
          for ret in results:
              self.assertTrue(ret.ret.item())
  
      @skip_if_cuda_not_available()
      def test_p2p(self):
          results = self.run_test_locally(run_p2p)
  
          i = 1
          while i < len(results):
              self.assertTrue(torch.equal(results[i].data, results[i - 1].data))
              i += 2
  
      @skip_if_cuda_not_available()
      def test_allgather(self):
          results = self.run_test_locally(run_allgather)
          for ret in results:
              self.assertTrue(ret.ret.item())
  
  
  if __name__ == "__main__":
      unittest.main()
  