# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This work is licensed under a Creative Commons
# Attribution-NonCommercial-ShareAlike 4.0 International License.
# You should have received a copy of the license along with this
# work. If not, see http://creativecommons.org/licenses/by-nc-sa/4.0/

import os
import torch
from datetime import timedelta
from accelerate import Accelerator
from accelerate import DistributedDataParallelKwargs
try:
    from accelerate import InitProcessGroupKwargs
except Exception:
    InitProcessGroupKwargs = None

#----------------------------------------------------------------------------

ACCELERATOR = None

def init():
    global ACCELERATOR
    timeout_minutes = int(os.environ.get("REOAC_DIST_TIMEOUT_MIN", "60"))
    if timeout_minutes < 1:
        timeout_minutes = 60
    ddp_kwargs = DistributedDataParallelKwargs(
        find_unused_parameters=False,
        broadcast_buffers=False,
    )
    kwargs_handlers = [ddp_kwargs]
    if InitProcessGroupKwargs is not None:
        kwargs_handlers.append(
            InitProcessGroupKwargs(timeout=timedelta(minutes=timeout_minutes))
        )
    ACCELERATOR = Accelerator(kwargs_handlers=kwargs_handlers)

#----------------------------------------------------------------------------

def get_accelerator():
    return ACCELERATOR

def get_rank():
    return ACCELERATOR.process_index

#----------------------------------------------------------------------------

def get_local_rank():
    return ACCELERATOR.local_process_index

#----------------------------------------------------------------------------

def get_world_size():
    return ACCELERATOR.num_processes

#----------------------------------------------------------------------------

def update_progress(cur, total):
    _ = cur, total

#----------------------------------------------------------------------------

def print0(*args, **kwargs):
    if get_rank() == 0:
        print(*args, **kwargs)

#----------------------------------------------------------------------------

if __name__ == "__main__":
    init()
    a = torch.zeros(3, device="cuda") + get_rank()
    aa = ACCELERATOR.gather(a)
    print(aa)