def setup_mpi_gpus()

in coinrun/main_utils.py [0:0]


def setup_mpi_gpus():
    if 'RCALL_NUM_GPU' not in os.environ:
        return
    num_gpus = int(os.environ['RCALL_NUM_GPU'])
    node_id = platform.node()
    nodes = MPI.COMM_WORLD.allgather(node_id)
    local_rank = len([n for n in nodes[:MPI.COMM_WORLD.Get_rank()] if n == node_id])
    os.environ['CUDA_VISIBLE_DEVICES'] = str(local_rank % num_gpus)