Method run_process

tests/test_utils.py:534–575

(self, func, local_rank, args, kwargs, results)

        self.verbose = verbose

    def run_process(self, func, local_rank, args, kwargs, results):
        _env = os.environ.copy()  # keep the original system env
        try:
            os.environ["MASTER_ADDR"] = self.master_addr
            os.environ["MASTER_PORT"] = str(self.master_port)
            os.environ["LOCAL_RANK"] = str(local_rank)
            if self.verbose:
                os.environ["NCCL_DEBUG"] = "INFO"
                os.environ["NCCL_DEBUG_SUBSYS"] = "ALL"
            os.environ["TORCH_NCCL_BLOCKING_WAIT"] = str(1)
            os.environ["OMP_NUM_THREADS"] = str(1)
            os.environ["WORLD_SIZE"] = str(self.nproc_per_node * self.nnodes)
            os.environ["RANK"] = str(self.nproc_per_node * self.node_rank + local_rank)

            if torch.cuda.is_available():
                torch.cuda.set_device(int(local_rank))  # using device ids from CUDA_VISIBILE_DEVICES

            if self.init_method != "no_init":
                dist.init_process_group(
                    backend=self.backend,
                    init_method=self.init_method,
                    timeout=self.timeout,
                    world_size=int(os.environ["WORLD_SIZE"]),
                    rank=int(os.environ["RANK"]),
                )
            func(*args, **kwargs)
            # the primary node lives longer to
            # avoid _store_based_barrier, RuntimeError: Broken pipe
            # as the TCP store daemon is on the rank 0
            if int(os.environ["RANK"]) == 0:
                time.sleep(0.1)
            results.put(True)
        except Exception as e:
            results.put(str(e))
            raise e
        finally:
            os.environ.clear()
            os.environ.update(_env)
            try:
                dist.destroy_process_group()
            except RuntimeError as e:
                warnings.warn(f"While closing process group: {e}.")

    def __call__(self, obj):
        if not torch.distributed.is_available():