# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import shutil

import torch
from omegaconf import OmegaConf
from pytorch_lightning import Trainer
from pytorch_lightning.utilities.distributed import rank_zero_only

from nemo.core import ModelPT
from nemo.utils import logging
from nemo.utils.exp_manager import ExpManagerConfig, exp_manager


class OnesDataset(torch.utils.data.Dataset):
    def __init__(self, dataset_len):
        super().__init__()
        self.__dataset_len = dataset_len

    def __getitem__(self, *args):
        return torch.ones(2)

    def __len__(self):
        return self.__dataset_len


class ExampleModel(ModelPT):
    def __init__(self, *args, **kwargs):
        cfg = OmegaConf.structured({})
        super().__init__(cfg, trainer=kwargs.get('trainer', None))
        # dummy parameter in order to allow DDP to execute
        self.l1 = torch.nn.modules.Linear(in_features=2, out_features=1)

    def train_dataloader(self):
        return None

    def val_dataloader(self):
        return None

    def predict_dataloader(self):
        dataset = OnesDataset(2)
        return torch.utils.data.DataLoader(dataset, batch_size=2)

    def forward(self, batch):
        return batch.mean()

    def validation_step(self, batch, batch_idx):
        return self(batch)

    def training_step(self, batch, batch_idx):
        return self(batch)

    def list_available_models(self):
        pass

    def setup_training_data(self):
        pass

    def setup_validation_data(self):
        pass

    def validation_epoch_end(self, loss):
        self.log("val_loss", torch.stack(loss).mean())


def instantiate_multinode_ddp_if_possible():
    num_gpus = torch.cuda.device_count()
    trainer = Trainer(gpus=num_gpus, accelerator='ddp', logger=None, checkpoint_callback=None)

    exp_manager_cfg = ExpManagerConfig(exp_dir='./ddp_check/', use_datetime_version=False, version="")
    exp_manager(trainer, cfg=OmegaConf.structured(exp_manager_cfg))
    return trainer


def setup_model(trainer: Trainer):
    model = ExampleModel(trainer=trainer)

    logging.info(f"M.Global Rank:{model.global_rank}")
    logging.info(f"M.Local Rank:{model.local_rank}")
    logging.info(f"M.World Size:{model.trainer.world_size}")

    trainer.predict(model)
    return model


def get_rank_info(texts: list, rank_key: str) -> int:
    for line in texts:
        if rank_key in line:
            rank_value = line.split(":")[-1]
            rank_value = int(rank_value)
            return rank_value

    print("Could not find the correct rank key !")
    exit(1)


@rank_zero_only
def check_model_ranks(model: ExampleModel):
    basedir = os.path.join('./ddp_check/', 'default', 'version_0')
    file_template = "nemo_log_globalrank-{rank}_localrank-{rank}.txt"

    world_size = torch.cuda.device_count()
    for rank in range(world_size):
        filename = file_template.format(rank=rank)
        filepath = os.path.join(basedir, filename)

        with open(filepath, 'r') as f:
            texts = f.readlines()
            texts = [t.replace("\n", "") for t in texts]

        log_global_rank = get_rank_info(texts, rank_key='M.Global Rank')
        log_world_size = get_rank_info(texts, rank_key='M.World Size')

        if log_global_rank != rank:
            print("Logged global rank is not equal to trainer.global_rank !")
            exit(1)

        if log_world_size != world_size:
            print("Logged world size if not equal to trainer.world_size !")
            exit(1)


@rank_zero_only
def cleanup():
    if os.path.exists('./ddp_check'):
        shutil.rmtree('./ddp_check', ignore_errors=True)


def run_checks():
    cleanup()

    trainer = instantiate_multinode_ddp_if_possible()
    model = setup_model(trainer)
    check_model_ranks(model)

    print("DDP checks passed !")

    cleanup()


if __name__ == '__main__':
    run_checks()

[((2257, 2282), 'torch.cuda.device_count', 'torch.cuda.device_count', ([], {}), '()\n', (2280, 2282), False, 'import torch\n'), ((2297, 2382), 'pytorch_lightning.Trainer', 'Trainer', ([], {'gpus': 'num_gpus', 'accelerator': '"""ddp"""', 'logger': 'None', 'checkpoint_callback': 'None'}), "(gpus=num_gpus, accelerator='ddp', logger=None, checkpoint_callback=None\n    )\n", (2304, 2382), False, 'from pytorch_lightning import Trainer\n'), ((2401, 2486), 'nemo.utils.exp_manager.ExpManagerConfig', 'ExpManagerConfig', ([], {'exp_dir': '"""./ddp_check/"""', 'use_datetime_version': '(False)', 'version': '""""""'}), "(exp_dir='./ddp_check/', use_datetime_version=False, version=''\n    )\n", (2417, 2486), False, 'from nemo.utils.exp_manager import ExpManagerConfig, exp_manager\n'), ((2653, 2703), 'nemo.utils.logging.info', 'logging.info', (['f"""M.Global Rank:{model.global_rank}"""'], {}), "(f'M.Global Rank:{model.global_rank}')\n", (2665, 2703), False, 'from nemo.utils import logging\n'), ((2708, 2756), 'nemo.utils.logging.info', 'logging.info', (['f"""M.Local Rank:{model.local_rank}"""'], {}), "(f'M.Local Rank:{model.local_rank}')\n", (2720, 2756), False, 'from nemo.utils import logging\n'), ((2761, 2817), 'nemo.utils.logging.info', 'logging.info', (['f"""M.World Size:{model.trainer.world_size}"""'], {}), "(f'M.World Size:{model.trainer.world_size}')\n", (2773, 2817), False, 'from nemo.utils import logging\n'), ((3227, 3279), 'os.path.join', 'os.path.join', (['"""./ddp_check/"""', '"""default"""', '"""version_0"""'], {}), "('./ddp_check/', 'default', 'version_0')\n", (3239, 3279), False, 'import os\n'), ((3368, 3393), 'torch.cuda.device_count', 'torch.cuda.device_count', ([], {}), '()\n', (3391, 3393), False, 'import torch\n'), ((4120, 4149), 'os.path.exists', 'os.path.exists', (['"""./ddp_check"""'], {}), "('./ddp_check')\n", (4134, 4149), False, 'import os\n'), ((1115, 1128), 'torch.ones', 'torch.ones', (['(2)'], {}), '(2)\n', (1125, 1128), False, 'import torch\n'), ((1273, 1297), 'omegaconf.OmegaConf.structured', 'OmegaConf.structured', (['{}'], {}), '({})\n', (1293, 1297), False, 'from omegaconf import OmegaConf\n'), ((1442, 1496), 'torch.nn.modules.Linear', 'torch.nn.modules.Linear', ([], {'in_features': '(2)', 'out_features': '(1)'}), '(in_features=2, out_features=1)\n', (1465, 1496), False, 'import torch\n'), ((1684, 1734), 'torch.utils.data.DataLoader', 'torch.utils.data.DataLoader', (['dataset'], {'batch_size': '(2)'}), '(dataset, batch_size=2)\n', (1711, 1734), False, 'import torch\n'), ((3499, 3530), 'os.path.join', 'os.path.join', (['basedir', 'filename'], {}), '(basedir, filename)\n', (3511, 3530), False, 'import os\n'), ((4159, 4207), 'shutil.rmtree', 'shutil.rmtree', (['"""./ddp_check"""'], {'ignore_errors': '(True)'}), "('./ddp_check', ignore_errors=True)\n", (4172, 4207), False, 'import shutil\n'), ((2511, 2548), 'omegaconf.OmegaConf.structured', 'OmegaConf.structured', (['exp_manager_cfg'], {}), '(exp_manager_cfg)\n', (2531, 2548), False, 'from omegaconf import OmegaConf\n'), ((2169, 2186), 'torch.stack', 'torch.stack', (['loss'], {}), '(loss)\n', (2180, 2186), False, 'import torch\n')]