Function _distributed_worker

yolox/core/launch.py:168–219

(
    local_rank,
    main_func,
    world_size,
    num_gpus_per_machine,
    num_machines,
    machine_rank,
    backend,
    dist_url,
    args,
)

def _distributed_worker(
    local_rank,
    main_func,
    world_size,
    num_gpus_per_machine,
    num_machines,
    machine_rank,
    backend,
    dist_url,
    args,
):
    assert (
        torch.cuda.is_available()
    ), "cuda is not available. Please check your installation."
    configure_nccl()
    global_rank = machine_rank * num_gpus_per_machine + local_rank
    logger.info("Rank {} initialization finished.".format(global_rank))
    try:
        dist.init_process_group(
            backend=backend,
            init_method=dist_url,
            world_size=world_size,
            rank=global_rank,
        )
    except Exception:
        logger.error("Process group URL: {}".format(dist_url))
        raise
    # synchronize is needed here to prevent a possible timeout after calling init_process_group
    # See: https://github.com/facebookresearch/maskrcnn-benchmark/issues/172
    comm.synchronize()

    if global_rank == 0 and os.path.exists(
        "./" + args[1].experiment_name + "_ip_add.txt"
    ):
        os.remove("./" + args[1].experiment_name + "_ip_add.txt")

    assert num_gpus_per_machine <= torch.cuda.device_count()
    torch.cuda.set_device(local_rank)

    args[1].local_rank = local_rank
    args[1].num_machines = num_machines

    # Setup the local process group (which contains ranks within the same machine)
    # assert comm._LOCAL_PROCESS_GROUP is None
    # num_machines = world_size // num_gpus_per_machine
    # for i in range(num_machines):
    # ranks_on_i = list(range(i * num_gpus_per_machine, (i + 1) * num_gpus_per_machine))
    # pg = dist.new_group(ranks_on_i)
    # if i == machine_rank:
    # comm._LOCAL_PROCESS_GROUP = pg

    main_func(*args)