code: """Utility code for running native pytorch distributed"""
  
  import os
  
  import torch.distributed as dist
  
  def init_workers_file():
      rank = int(os.environ['SLURM_PROCID'])
      n_ranks = int(os.environ['SLURM_NTASKS'])
      sync_file = 'file:///tmp/%s_%s_pytorch_sync' % (
          os.environ['USER'], os.environ['SLURM_JOB_ID'])
      dist.init_process_group(backend='nccl', world_size=n_ranks, rank=rank,
                              init_method=sync_file)
      return rank, n_ranks
  
  def init_workers_mpi():
      dist.init_process_group(backend='mpi')
      rank = dist.get_rank()
      n_ranks = dist.get_world_size()
      return rank, n_ranks
  
apis: torch.distributed.init_process_group torch.distributed.get_world_size torch.distributed.get_rank