code: # Copyright (c) Facebook, Inc. and its affiliates.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import os
  import pickle
  import socket
  import struct
  import subprocess
  import warnings
  
  import torch
  import torch.distributed as dist
  
  
  
  def is_master(args):
      return args.distributed_rank == 0
  
  
  def infer_init_method(args):
      if args.distributed_init_method is not None:
          return
  
      # support torch.distributed.launch
      if all(key in os.environ for key in [
          'MASTER_ADDR', 'MASTER_PORT', 'WORLD_SIZE', 'RANK'
      ]):
          args.distributed_init_method = 'env://'
          args.distributed_world_size = int(os.environ['WORLD_SIZE'])
          args.distributed_rank = int(os.environ['RANK'])
  
      # we can determine the init method automatically for Slurm
      elif args.distributed_port > 0:
          node_list = os.environ.get('SLURM_STEP_NODELIST')
          if node_list is None:
              node_list = os.environ.get('SLURM_JOB_NODELIST')
          if node_list is not None:
              try:
                  hostnames = subprocess.check_output(['scontrol', 'show', 'hostnames', node_list])
                  args.distributed_init_method = 'tcp://{host}:{port}'.format(
                      host=hostnames.split()[0].decode('utf-8'),
                      port=args.distributed_port,
                  )
                  nnodes = int(os.environ.get('SLURM_NNODES'))
                  ntasks_per_node = os.environ.get('SLURM_NTASKS_PER_NODE')
                  if ntasks_per_node is not None:
                      ntasks_per_node = int(ntasks_per_node)
                  else:
                      ntasks = int(os.environ.get('SLURM_NTASKS'))
                      nnodes = int(os.environ.get('SLURM_NNODES'))
                      assert ntasks % nnodes == 0
                      ntasks_per_node = int(ntasks / nnodes)
                  if ntasks_per_node == 1:
                      assert args.distributed_world_size % nnodes == 0
                      gpus_per_node = args.distributed_world_size // nnodes
                      node_id = int(os.environ.get('SLURM_NODEID'))
                      args.distributed_rank = node_id * gpus_per_node
                  else:
                      assert ntasks_per_node == args.distributed_world_size // nnodes
                      args.distributed_no_spawn = True
                      args.distributed_rank = int(os.environ.get('SLURM_PROCID'))
                      args.device_id = int(os.environ.get('SLURM_LOCALID'))
              except subprocess.CalledProcessError as e:  # scontrol failed
                  raise e
              except FileNotFoundError:  # Slurm is not installed
                  pass
  
  
  def distributed_init(args):
      if args.distributed_world_size == 1:
          raise ValueError('Cannot initialize distributed with distributed_world_size=1')
  
      if torch.distributed.is_initialized():
          warnings.warn('Distributed is already initialized, cannot initialize twice!')
      else:
          print('| distributed init (rank {}): {}'.format(
              args.distributed_rank, args.distributed_init_method), flush=True)
          dist.init_process_group(
              backend=args.distributed_backend,
              init_method=args.distributed_init_method,
              world_size=args.distributed_world_size,
              rank=args.distributed_rank,
          )
          print('| initialized host {} as rank {}'.format(
              socket.gethostname(), args.distributed_rank), flush=True)
  
          # perform a dummy all-reduce to initialize the NCCL communicator
          if torch.cuda.is_available():
              dist.all_reduce(torch.zeros(1).cuda())
          else:
              dist.all_reduce(torch.zeros(1))
  
          suppress_output(is_master(args))
  
      args.distributed_rank = torch.distributed.get_rank()
      return args.distributed_rank
  
  
  def suppress_output(is_master):
      """Suppress printing on the current device. Force printing with `force=True`."""
      import builtins as __builtin__
      builtin_print = __builtin__.print
  
      def print(*args, **kwargs):
          force = kwargs.pop('force', False)
          if is_master or force:
              builtin_print(*args, **kwargs)
  
      __builtin__.print = print
  
  
  def get_rank():
      return dist.get_rank()
  
  
  def get_world_size():
      try:
          return dist.get_world_size()
      except AssertionError:
          return 1
  
  def get_default_group():
      return dist.group.WORLD
  
  
  def all_reduce(tensor, op='sum', group=None):
      if group is None:
          group = get_default_group()
      output = dist.all_reduce(tensor, group=group)
      if op == 'mean':
          return output / get_world_size()
      return output
  
  def all_gather_list(data, group=None, max_size=16384):
      """Gathers arbitrary data from all nodes into a list.
  
      Similar to :func:`~torch.distributed.all_gather` but for arbitrary Python
      data. Note that *data* must be picklable.
  
      Args:
          data (Any): data from the local worker to be gathered on other workers
          group (optional): group of the collective
          max_size (int, optional): maximum size of the data to be gathered
              across workers
      """
      rank = get_rank()
      world_size = get_world_size()
  
      buffer_size = max_size * world_size
      if not hasattr(all_gather_list, '_buffer') or \
              all_gather_list._buffer.numel() < buffer_size:
          all_gather_list._buffer = torch.cuda.ByteTensor(buffer_size)
          all_gather_list._cpu_buffer = torch.ByteTensor(max_size).pin_memory()
      buffer = all_gather_list._buffer
      buffer.zero_()
      cpu_buffer = all_gather_list._cpu_buffer
  
      enc = pickle.dumps(data)
      enc_size = len(enc)
      header_size = 4  # size of header that contains the length of the encoded data
      size = header_size + enc_size
      if size > max_size:
          raise ValueError('encoded data size ({}) exceeds max_size ({})'.format(size, max_size))
  
      header = struct.pack(">I", enc_size)
      cpu_buffer[:size] = torch.ByteTensor(list(header + enc))
      start = rank * max_size
      buffer[start:start + size].copy_(cpu_buffer[:size])
  
      all_reduce(buffer, group=group)
  
      try:
          result = []
          for i in range(world_size):
              out_buffer = buffer[i * max_size:(i + 1) * max_size]
              enc_size, = struct.unpack(">I", bytes(out_buffer[:header_size].tolist()))
              if enc_size > 0:
                  result.append(pickle.loads(bytes(out_buffer[header_size:header_size + enc_size].tolist())))
          return result
      except pickle.UnpicklingError:
          raise Exception(
              'Unable to unpickle data from other workers. all_gather_list requires all '
              'workers to enter the function together, so this error usually indicates '
              'that the workers have fallen out of sync somehow. Workers can fall out of '
              'sync if one of them runs out of memory, or if there are other conditions '
              'in your training script that can cause one worker to finish an epoch '
              'while other workers are still iterating over their portions of the data.'
          )
  
apis: torch.distributed.get_rank torch.distributed.get_world_size torch.distributed.init_process_group torch.cuda.ByteTensor torch.distributed.is_initialized torch.ByteTensor torch.cuda.is_available torch.distributed.all_reduce torch.zeros