# sudo docker run -it --name horovod-py3.6 --mount type=bind,source=/,target=/host_fs --network dtest --ip 172.28.0.4 horovod/horovod:0.19.1-tf2.1.0-torch1.4.0-mxnet1.6.0-py3.6-cpu
# sudo docker run -it --name horovod-py3.6-worker3 --cap-add=NET_ADMIN --mount type=bind,source=/,target=/host_fs --network dtest --ip 172.28.0.6 horovod/horovod:0.19.1-tf2.1.0-torch1.4.0-mxnet1.6.0-py3.6-cpu
# sudo docker run -it --name horovod-py3.6-worker3 --cap-add=NET_ADMIN --mount type=bind,source=/,target=/host_fs --mount type=bind,source=/mnt/raid10/hanjing/thesis/horovod_docker/.ssh, target=/root/.ssh --network dtest --ip 172.28.0.6 horovod/horovod:0.19.1-tf2.1.0-torch1.4.0-mxnet1.6.0-py3.6-cpu

# horovodrun -np 1 -H localhost:1 python /host_fs/mnt/raid10/hanjing/thesis/horovod/examples/pytorch_imagenet_resnet50.py
from __future__ import print_function

import torch
import argparse
import torch.backends.cudnn as cudnn
import torch.multiprocessing as mp
import torch.nn.functional as F
import torch.optim as optim
import torch.utils.data.distributed
from torchvision import datasets, transforms, models
import horovod.torch as hvd
import os
import math
from tqdm import tqdm
from distutils.version import LooseVersion

import time

import sys
print(f'sys.path: {sys.path}')
print(f'torch: {torch}')
print(f'hvd: {hvd}')
# Training settings
parser = argparse.ArgumentParser(description='PyTorch ImageNet Example',
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--train-dir', default=os.path.expanduser('~/imagenet/train'),
                    help='path to training data')
parser.add_argument('--val-dir', default=os.path.expanduser('~/imagenet/validation'),
                    help='path to validation data')
parser.add_argument('--log-dir', default='./logs',
                    help='tensorboard log directory')
parser.add_argument('--checkpoint-format', default='./checkpoint-{epoch}.pth.tar',
                    help='checkpoint file format')
parser.add_argument('--fp16-allreduce', action='store_true', default=False,
                    help='use fp16 compression during allreduce')
parser.add_argument('--batches-per-allreduce', type=int, default=1,
                    help='number of batches processed locally before '
                         'executing allreduce across workers; it multiplies '
                         'total batch size.')
parser.add_argument('--use-adasum', action='store_true', default=False,
                    help='use adasum algorithm to do reduction')

# Default settings from https://arxiv.org/abs/1706.02677.
parser.add_argument('--batch-size', type=int, default=32,
                    help='input batch size for training')
parser.add_argument('--val-batch-size', type=int, default=32,
                    help='input batch size for validation')
parser.add_argument('--epochs', type=int, default=90,
                    help='number of epochs to train')
parser.add_argument('--base-lr', type=float, default=0.0125,
                    help='learning rate for a single GPU')
parser.add_argument('--warmup-epochs', type=float, default=5,
                    help='number of warmup epochs')
parser.add_argument('--momentum', type=float, default=0.9,
                    help='SGD momentum')
parser.add_argument('--wd', type=float, default=0.00005,
                    help='weight decay')

parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='disables CUDA training')
parser.add_argument('--seed', type=int, default=42,
                    help='random seed')


def train(epoch, backward_hook, tracker):
    model.train()
    train_sampler.set_epoch(epoch)
    train_loss = Metric('train_loss')
    train_accuracy = Metric('train_accuracy')

    with tqdm(total=len(train_loader),
              desc='Train Epoch     #{}'.format(epoch + 1),
              disable=not verbose) as t:
        for batch_idx, (data, target) in enumerate(train_loader):
            adjust_learning_rate(epoch, batch_idx)

            if args.cuda:
                data, target = data.cuda(), target.cuda()
            optimizer.zero_grad()
            # Split data into sub-batches of size batch_size
            for i in range(0, len(data), args.batch_size):
                print(f"SUB BATCH {i}")
                data_batch = data[i:i + args.batch_size]
                target_batch = target[i:i + args.batch_size]
                output = model(data_batch)
                train_accuracy.update(accuracy(output, target_batch))
                loss = F.cross_entropy(output, target_batch)
                if batch_idx == 0:
                    backward_hook(model)
                train_loss.update(loss)
                # Average gradients among sub-batches
                loss.div_(math.ceil(float(len(data)) / args.batch_size))
                tracker.add_record("Start Backpropagation")
                loss.backward()
            # Gradient is applied across all ranks
            optimizer.step()
            t.set_postfix({'loss': train_loss.avg.item(),
                           'accuracy': 100. * train_accuracy.avg.item()})
            t.update(1)
            tracker.reset()

    if log_writer:
        log_writer.add_scalar('train/loss', train_loss.avg, epoch)
        log_writer.add_scalar('train/accuracy', train_accuracy.avg, epoch)


def validate(epoch):
    model.eval()
    val_loss = Metric('val_loss')
    val_accuracy = Metric('val_accuracy')

    with tqdm(total=len(val_loader),
              desc='Validate Epoch  #{}'.format(epoch + 1),
              disable=not verbose) as t:
        with torch.no_grad():
            for data, target in val_loader:
                if args.cuda:
                    data, target = data.cuda(), target.cuda()
                output = model(data)

                val_loss.update(F.cross_entropy(output, target))
                val_accuracy.update(accuracy(output, target))
                t.set_postfix({'loss': val_loss.avg.item(),
                               'accuracy': 100. * val_accuracy.avg.item()})
                t.update(1)

    if log_writer:
        log_writer.add_scalar('val/loss', val_loss.avg, epoch)
        log_writer.add_scalar('val/accuracy', val_accuracy.avg, epoch)


# Horovod: using `lr = base_lr * hvd.size()` from the very beginning leads to worse final
# accuracy. Scale the learning rate `lr = base_lr` ---> `lr = base_lr * hvd.size()` during
# the first five epochs. See https://arxiv.org/abs/1706.02677 for details.
# After the warmup reduce learning rate by 10 on the 30th, 60th and 80th epochs.
def adjust_learning_rate(epoch, batch_idx):
    if epoch < args.warmup_epochs:
        epoch += float(batch_idx + 1) / len(train_loader)
        lr_adj = 1. / hvd.size() * (epoch * (hvd.size() - 1) / args.warmup_epochs + 1)
    elif epoch < 30:
        lr_adj = 1.
    elif epoch < 60:
        lr_adj = 1e-1
    elif epoch < 80:
        lr_adj = 1e-2
    else:
        lr_adj = 1e-3
    for param_group in optimizer.param_groups:
        param_group['lr'] = args.base_lr * hvd.size() * args.batches_per_allreduce * lr_adj


def accuracy(output, target):
    # get the index of the max log-probability
    pred = output.max(1, keepdim=True)[1]
    return pred.eq(target.view_as(pred)).cpu().float().mean()


def save_checkpoint(epoch):
    if hvd.rank() == 0:
        filepath = args.checkpoint_format.format(epoch=epoch + 1)
        state = {
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
        }
        torch.save(state, filepath)


# Horovod: average metrics from distributed training.
class Metric(object):
    def __init__(self, name):
        self.name = name
        self.sum = torch.tensor(0.)
        self.n = torch.tensor(0.)

    def update(self, val):
        self.sum += hvd.allreduce(val.detach().cpu(), name=self.name)
        self.n += 1

    @property
    def avg(self):
        return self.sum / self.n


if __name__ == '__main__':
    args = parser.parse_args()
    args.cuda = not args.no_cuda and torch.cuda.is_available()

    allreduce_batch_size = args.batch_size * args.batches_per_allreduce

    hvd.init()
    torch.manual_seed(args.seed)

    if args.cuda:
        # Horovod: pin GPU to local rank.
        torch.cuda.set_device(hvd.local_rank())
        torch.cuda.manual_seed(args.seed)

    cudnn.benchmark = True

    # If set > 0, will resume training from a given checkpoint.
    resume_from_epoch = 0
    for try_epoch in range(args.epochs, 0, -1):
        if os.path.exists(args.checkpoint_format.format(epoch=try_epoch)):
            resume_from_epoch = try_epoch
            break

    # Horovod: broadcast resume_from_epoch from rank 0 (which will have
    # checkpoints) to other ranks.
    resume_from_epoch = hvd.broadcast(torch.tensor(resume_from_epoch), root_rank=0,
                                      name='resume_from_epoch').item()

    # Horovod: print logs on the first worker.
    verbose = 1 if hvd.rank() == 0 else 0

    # Horovod: write TensorBoard logs on first worker.
    try:
        if LooseVersion(torch.__version__) >= LooseVersion('1.2.0'):
            from torch.utils.tensorboard import SummaryWriter
        else:
            from tensorboardX import SummaryWriter
        log_writer = SummaryWriter(args.log_dir) if hvd.rank() == 0 else None
    except ImportError:
        log_writer = None

    # Horovod: limit # of CPU threads to be used per worker.
    torch.set_num_threads(4)

    kwargs = {'num_workers': 4, 'pin_memory': True} if args.cuda else {}
    # When supported, use 'forkserver' to spawn dataloader workers instead of 'fork' to prevent
    # issues with Infiniband implementations that are not fork-safe
    if (kwargs.get('num_workers', 0) > 0 and hasattr(mp, '_supports_context') and
            mp._supports_context and 'forkserver' in mp.get_all_start_methods()):
        kwargs['multiprocessing_context'] = 'forkserver'

    # train_dataset = \
    #     datasets.ImageFolder(args.train_dir,
    #                          transform=transforms.Compose([
    #                              transforms.RandomResizedCrop(224),
    #                              transforms.RandomHorizontalFlip(),
    #                              transforms.ToTensor(),
    #                              transforms.Normalize(mean=[0.485, 0.456, 0.406],
    #                                                   std=[0.229, 0.224, 0.225])
    #                          ]))
    train_dataset = datasets.CIFAR10(root='/host_fs/mnt/raid10/hanjing/thesis/horovod/examples/cifar10_data', train=True,
                                                download=True, transform=transforms.Compose([
                                 transforms.RandomResizedCrop(224),
                                 transforms.RandomHorizontalFlip(),
                                 transforms.ToTensor(),
                                 transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                                      std=[0.229, 0.224, 0.225])
                             ]))
    # Horovod: use DistributedSampler to partition data among workers. Manually specify
    # `num_replicas=hvd.size()` and `rank=hvd.rank()`.
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        train_dataset, num_replicas=hvd.size(), rank=hvd.rank())
    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=allreduce_batch_size,
        sampler=train_sampler, **kwargs)

    # val_dataset = \
    #     datasets.ImageFolder(args.val_dir,
    #                          transform=transforms.Compose([
    #                              transforms.Resize(256),
    #                              transforms.CenterCrop(224),
    #                              transforms.ToTensor(),
    #                              transforms.Normalize(mean=[0.485, 0.456, 0.406],
    #                                                   std=[0.229, 0.224, 0.225])
    #                          ]))
    val_dataset = datasets.CIFAR10(root='/host_fs/mnt/raid10/hanjing/thesis/horovod/examples/cifar10_data', train=False,
                                               download=True, transform=transforms.Compose([
                                 transforms.Resize(256),
                                 transforms.CenterCrop(224),
                                 transforms.ToTensor(),
                                 transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                                      std=[0.229, 0.224, 0.225])
                             ]))
    val_sampler = torch.utils.data.distributed.DistributedSampler(
        val_dataset, num_replicas=hvd.size(), rank=hvd.rank())
    val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=args.val_batch_size,
                                             sampler=val_sampler, **kwargs)


    # Set up standard ResNet-50 model.
    model = models.resnet50()

    # By default, Adasum doesn't need scaling up learning rate.
    # For sum/average with gradient Accumulation: scale learning rate by batches_per_allreduce
    lr_scaler = args.batches_per_allreduce * hvd.size() if not args.use_adasum else 1

    if args.cuda:
        # Move model to GPU.
        model.cuda()
        # If using GPU Adasum allreduce, scale learning rate by local_size.
        if args.use_adasum and hvd.nccl_built():
            lr_scaler = args.batches_per_allreduce * hvd.local_size()

    # Horovod: scale learning rate by the number of GPUs.
    optimizer = optim.SGD(model.parameters(),
                          lr=(args.base_lr *
                              lr_scaler),
                          momentum=args.momentum, weight_decay=args.wd)
    
    class Tracker:
        def __init__(self):
            self.layer = 0
            self.records = []
            self.m_layer = {} # m -> layer id
        
        def reset(self):
            if len(self.records) > 0:
                for record in self.records:
                    print(",".join(map(lambda x: str(x), record)))
            self.layer = 0
            self.records = []
            self.m_layer = {}

        # def track(self, attrs):
        #     attrs.append(self.layer)
        #     self.layer += 1
        #     self.records.append(attrs)
        def add_record(self, message):
            self.records.append([message, time.time()])
        
        def backward_track(self, module, parameter, parameter_name):
            fw_layer = self.m_layer[module]
            self.records.append(["BP", time.time(), fw_layer, module._get_name(), parameter_name, parameter.size()])
        
        def forward_track(self, module, parameter_weight_size):
            self.layer += 1 
            self.m_layer[module] = self.layer
            self.records.append(["FP",time.time(), self.layer, module._get_name(), parameter_weight_size])

    tracker = Tracker()
    def _make_hook(p, m, name):
        """Define hook for backward propagation. The hook wraps the allreduce OP as a ByteTask and
        posts it to Core.
        Arguments:
            p: the parameter.
        """
        # self._logger.debug("{} calls make_hook for {}".format(self._desc, self._parameter_names[p]))
        def hook(*args):
            # g = args[1][0].reshape([-1])[0] #first value of the gradient
            # print(f'parameter gradient ready: {time.time()} {name}: {p.size()} in layer {m._get_name()}, {g}')
            # print(f"backward m id {id(m)} = {m._get_name()}, p id {id(p)} name {name}")
            tracker.backward_track(m, p, name)
        return hook

    def _register_hooks(p, m, name):
        """Add a hook after the backward propagation of each layer to start allreduce"""
        if p.requires_grad:
            # print(f"{m._get_name()}:{name} grad_fn {p.grad_fn} next {p.grad_fn.next_functions}")
            p.grad = p.data.new(p.size()).zero_()
            # self._requires_update.add(p)
            p_tmp = p.expand_as(p)
            grad_acc = p_tmp.grad_fn.next_functions[0][0]
            grad_acc.register_hook(_make_hook(p, m, name))
            # self._grad_accs.append(grad_acc)
    
    
    
    def install_backward_hook(m):
        for name, p in m.named_parameters(recurse = False):
            _register_hooks(p, m, name)
    def install_backward_hooks(model):
        model.apply(install_backward_hook)
    
    @torch.no_grad()
    def install_time_hooks(m):
        # print(m)
        def print_layer(m, _):
            size = 0
            if hasattr(m, 'weight'):
                # size = math.prod(m.weight.size()) * m.weight.element_size()
                size = m.weight.element_size()
                for dim in m.weight.size():
                  size *= dim
            # tracker.track([time.time(), tracker.layer, m._get_name(), size])
            tracker.forward_track(m, size)
        m.register_forward_pre_hook(print_layer)
        
        #m.register_forward_hook()    

    model.apply(install_time_hooks)


    # Horovod: (optional) compression algorithm.
    compression = hvd.Compression.fp16 if args.fp16_allreduce else hvd.Compression.none

    # Horovod: wrap optimizer with DistributedOptimizer.
    optimizer = hvd.DistributedOptimizer(
        optimizer, named_parameters=model.named_parameters(),
        compression=compression,
        backward_passes_per_step=args.batches_per_allreduce,
        op=hvd.Adasum if args.use_adasum else hvd.Average)

    # Restore from a previous checkpoint, if initial_epoch is specified.
    # Horovod: restore on the first worker which will broadcast weights to other workers.
    if resume_from_epoch > 0 and hvd.rank() == 0:
        filepath = args.checkpoint_format.format(epoch=resume_from_epoch)
        checkpoint = torch.load(filepath)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])

    # Horovod: broadcast parameters & optimizer state.
    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
    hvd.broadcast_optimizer_state(optimizer, root_rank=0)

    for epoch in range(resume_from_epoch, args.epochs):
        train(epoch, install_backward_hooks, tracker)
        validate(epoch)
        save_checkpoint(epoch)
