import os
import pickle
import time
import timeit

import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import numpy as np
import torch
import tempfile
import horovod.torch as hvd
from horovod.ray import RayExecutor

from ray_shuffling_data_loader.torch_dataset import (TorchShufflingDataset)
from ray_shuffling_data_loader.data_generation import (generate_data,
                                                       DATA_SPEC)

import argparse

DEFAULT_DATA_DIR = "s3://shuffling-data-loader-benchmarks/data/"

numpy_to_torch_dtype = {
    np.bool: torch.bool,
    np.uint8: torch.uint8,
    np.int8: torch.int8,
    np.int16: torch.int16,
    np.int32: torch.int32,
    np.int64: torch.int64,
    np.float16: torch.float16,
    np.float32: torch.float32,
    np.float64: torch.float64,
    np.complex64: torch.complex64,
    np.complex128: torch.complex128
}

# Training settings
parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
parser.add_argument(
    "--batch-size",
    type=int,
    default=250000,
    metavar="N",
    help="input batch size for training (default: 64)")
parser.add_argument(
    "--test-batch-size",
    type=int,
    default=250000,
    metavar="N",
    help="input batch size for testing (default: 1000)")
parser.add_argument(
    "--epochs",
    type=int,
    default=10,
    metavar="N",
    help="number of epochs to train (default: 10)")
parser.add_argument(
    "--lr",
    type=float,
    default=0.01,
    metavar="LR",
    help="learning rate (default: 0.01)")
parser.add_argument(
    "--momentum",
    type=float,
    default=0.5,
    metavar="M",
    help="SGD momentum (default: 0.5)")
parser.add_argument(
    "--no-cuda",
    action="store_true",
    default=False,
    help="disables CUDA training")
parser.add_argument(
    "--seed",
    type=int,
    default=42,
    metavar="S",
    help="random seed (default: 42)")
parser.add_argument(
    "--log-interval",
    type=int,
    default=10,
    metavar="N",
    help=("how many batches to wait before logging training "
          "status"))
parser.add_argument(
    "--fp16-allreduce",
    action="store_true",
    default=False,
    help="use fp16 compression during allreduce")
parser.add_argument(
    "--use-adasum",
    action="store_true",
    default=False,
    help="use adasum algorithm to do reduction")
parser.add_argument(
    "--gradient-predivide-factor",
    type=float,
    default=1.0,
    help=("apply gradient predivide factor in optimizer "
          "(default: 1.0)"))
parser.add_argument("--num-workers", type=int, default=None)
parser.add_argument("--num-hosts", type=int, default=None)
parser.add_argument("--num-workers-per-host", type=int, default=None)
parser.add_argument("--cpus-per-worker", type=int, default=1)
parser.add_argument("--mock-train-step-time", type=float, default=1.0)

# Synthetic training data generation settings.
parser.add_argument("--cache-files", action="store_true", default=False)
parser.add_argument("--num-rows", type=int, default=2 * (10**7))
parser.add_argument("--num-files", type=int, default=25)
parser.add_argument("--max-row-group-skew", type=float, default=0.0)
parser.add_argument("--num-row-groups-per-file", type=int, default=5)
parser.add_argument("--data-dir", type=str, default=DEFAULT_DATA_DIR)

# Shuffling data loader settings.
parser.add_argument("--num-reducers", type=int, default=32)
parser.add_argument("--max-concurrent-epochs", type=int, default=2)
parser.add_argument("--address", default="auto")


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)


def train_main(args, filenames):
    # Horovod: initialize library.
    hvd.init()
    torch.manual_seed(args.seed)

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

    # Horovod: limit # of CPU threads to be used per worker.
    torch.set_num_threads(1)
    rank = hvd.rank()
    train_dataset = create_dataset(
        filenames,
        batch_size=args.batch_size,
        rank=rank,
        num_epochs=args.epochs,
        world_size=hvd.size(),
        num_reducers=args.num_reducers,
        max_concurrent_epochs=args.max_concurrent_epochs)
    model = Net()
    # By default, Adasum doesn"t need scaling up learning rate.
    lr_scaler = hvd.size() if not args.use_adasum else 1

    if torch.cuda.is_available() and not args.no_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 = hvd.local_size()

    # Horovod: scale learning rate by lr_scaler.
    optimizer = optim.SGD(
        model.parameters(), lr=args.lr * lr_scaler, momentum=args.momentum)

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

    # 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,
        op=hvd.Adasum if args.use_adasum else hvd.Average,
        gradient_predivide_factor=args.gradient_predivide_factor)

    def _train(epoch):
        model.train()
        # Horovod: set epoch to sampler for shuffling.
        train_dataset.set_epoch(epoch)
        start_epoch = timeit.default_timer()
        last_batch_time = start_epoch
        batch_wait_times = []
        for batch_idx, (data, target) in enumerate(train_dataset):
            batch_wait_times.append(timeit.default_timer() - last_batch_time)
            if torch.cuda.is_available() and not args.no_cuda:
                if isinstance(data, list):
                    data = [t.cuda() for t in data]
                target = target.cuda()
            optimizer.zero_grad()
            # output = model(data)
            if batch_idx % args.log_interval == 0:
                print(
                    f"Processing batch {batch_idx} in epoch {epoch} on worker "
                    f"{rank}.")
            time.sleep(args.mock_train_step_time)
            # TODO(Clark): Add worker synchronization barrier here.
            # loss = F.nll_loss(output, target)
            # loss.backward()
            # optimizer.step()
            last_batch_time = timeit.default_timer()
        epoch_duration = timeit.default_timer() - start_epoch
        avg_batch_wait_time = np.mean(batch_wait_times)
        std_batch_wait_time = np.std(batch_wait_times)
        max_batch_wait_time = np.max(batch_wait_times)
        min_batch_wait_time = np.min(batch_wait_times)
        print(f"\nEpoch {epoch}, worker {rank} stats over "
              f"{len(batch_wait_times)} steps: {epoch_duration:.3f}")
        print(f"Mean batch wait time: {avg_batch_wait_time:.3f}s +- "
              f"{std_batch_wait_time}")
        print(f"Max batch wait time: {max_batch_wait_time:.3f}s")
        print(f"Min batch wait time: {min_batch_wait_time:.3f}s")
        return batch_wait_times

    print(f"Starting training on worker {rank}.")
    batch_wait_times = []
    for epoch in range(args.epochs):
        batch_wait_times.extend(_train(epoch))
    batch_wait_times.pop(0)
    print(f"Done training on worker {rank}.")
    avg_batch_wait_time = np.mean(batch_wait_times)
    std_batch_wait_time = np.std(batch_wait_times)
    max_batch_wait_time = np.max(batch_wait_times)
    min_batch_wait_time = np.min(batch_wait_times)
    print(f"\nWorker {rank} training stats over {args.epochs} epochs:")
    print(f"Mean batch wait time: {avg_batch_wait_time:.3f}s +- "
          f"{std_batch_wait_time}")
    print(f"Max batch wait time: {max_batch_wait_time:.3f}s")
    print(f"Min batch wait time: {min_batch_wait_time:.3f}s")
    # TODO(Clark): Add logic to the dataset abstraction so we don't have to do
    # this.
    if rank == 0:
        print("Waiting in rank 0 worker to let other workers consume queue...")
        time.sleep(10)
        print("Done waiting in rank 0 worker.")


def create_dataset(filenames, *, batch_size, rank, num_epochs, world_size,
                   num_reducers, max_concurrent_epochs):
    print(f"Creating Torch shuffling dataset for worker {rank} with "
          f"{batch_size} batch size, {num_epochs} epochs, {num_reducers} "
          f"reducers, and {world_size} trainers.")
    feature_columns = list(DATA_SPEC.keys())
    feature_types = [
        numpy_to_torch_dtype[dtype] for _, _, dtype in DATA_SPEC.values()
    ]
    label_column = feature_columns.pop()
    label_type = feature_types.pop()
    return TorchShufflingDataset(
        filenames,
        num_epochs,
        world_size,
        batch_size,
        rank,
        num_reducers=num_reducers,
        max_concurrent_epochs=max_concurrent_epochs,
        feature_columns=feature_columns,
        feature_types=feature_types,
        label_column=label_column,
        label_type=label_type)


if __name__ == "__main__":
    args = parser.parse_args()
    from ray_shuffling_data_loader.stats import human_readable_size
    import ray
    print("Connecting to Ray cluster...")
    ray.init(address=args.address)

    num_rows = args.num_rows
    num_files = args.num_files
    num_row_groups_per_file = args.num_row_groups_per_file
    max_row_group_skew = args.max_row_group_skew
    data_dir = args.data_dir

    cache_path = os.path.join(tempfile.gettempdir(), "data_cache")
    filenames = None
    if args.cache_files and os.path.exists(cache_path):
        try:
            with open(cache_path, "rb") as f:
                filenames, num_bytes = pickle.load(f)
        except Exception as exc:
            print(f"Cache load failed - {exc}")

    if not filenames:

        print(f"Generating {num_rows} rows over {num_files} files, with "
              f"{num_row_groups_per_file} row groups per file and at most "
              f"{100 * max_row_group_skew:.1f}% row group skew.")
        filenames, num_bytes = generate_data(num_rows, num_files,
                                             num_row_groups_per_file,
                                             max_row_group_skew, data_dir)
        if args.cache_files:
            with open(os.path.join(tempfile.gettempdir(), "data_cache"),
                      "wb") as f:
                pickle.dump((filenames, num_bytes), f)
    print(f"Generated {len(filenames)} files containing {num_rows} rows "
          f"with {num_row_groups_per_file} row groups per file, totalling "
          f"{human_readable_size(num_bytes)}.")

    print("Create Ray executor")
    worker_kwargs = {}
    num_workers = args.num_workers
    num_hosts = args.num_hosts
    num_workers_per_host = args.num_workers_per_host
    if num_workers is not None:
        if num_hosts is not None:
            raise ValueError(
                "Only one of --num-workers and --num-hosts should be used.")
        worker_kwargs["num_workers"] = num_workers
    elif num_hosts is not None:
        worker_kwargs["num_hosts"] = num_hosts
        if num_workers_per_host is None:
            raise ValueError("When giving --num-hosts, --num-workers-per-host "
                             "must also be given.")
        worker_kwargs["num_workers_per_host"] = num_workers_per_host
    cpus_per_worker = args.cpus_per_worker
    settings = RayExecutor.create_settings(timeout_s=30)
    executor = RayExecutor(
        settings,
        use_gpu=True,
        gpus_per_worker=1,
        cpus_per_worker=cpus_per_worker,
        **worker_kwargs)
    executor.start()
    executor.run(train_main, args=[args, filenames])
    executor.shutdown()

    print("Done consuming batches.")

[((922, 982), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""PyTorch MNIST Example"""'}), "(description='PyTorch MNIST Example')\n", (945, 982), False, 'import argparse\n'), ((4215, 4225), 'horovod.torch.init', 'hvd.init', ([], {}), '()\n', (4223, 4225), True, 'import horovod.torch as hvd\n'), ((4230, 4258), 'torch.manual_seed', 'torch.manual_seed', (['args.seed'], {}), '(args.seed)\n', (4247, 4258), False, 'import torch\n'), ((4513, 4537), 'torch.set_num_threads', 'torch.set_num_threads', (['(1)'], {}), '(1)\n', (4534, 4537), False, 'import torch\n'), ((4549, 4559), 'horovod.torch.rank', 'hvd.rank', ([], {}), '()\n', (4557, 4559), True, 'import horovod.torch as hvd\n'), ((5517, 5570), 'horovod.torch.broadcast_optimizer_state', 'hvd.broadcast_optimizer_state', (['optimizer'], {'root_rank': '(0)'}), '(optimizer, root_rank=0)\n', (5546, 5570), True, 'import horovod.torch as hvd\n'), ((8136, 8161), 'numpy.mean', 'np.mean', (['batch_wait_times'], {}), '(batch_wait_times)\n', (8143, 8161), True, 'import numpy as np\n'), ((8188, 8212), 'numpy.std', 'np.std', (['batch_wait_times'], {}), '(batch_wait_times)\n', (8194, 8212), True, 'import numpy as np\n'), ((8239, 8263), 'numpy.max', 'np.max', (['batch_wait_times'], {}), '(batch_wait_times)\n', (8245, 8263), True, 'import numpy as np\n'), ((8290, 8314), 'numpy.min', 'np.min', (['batch_wait_times'], {}), '(batch_wait_times)\n', (8296, 8314), True, 'import numpy as np\n'), ((9439, 9709), 'ray_shuffling_data_loader.torch_dataset.TorchShufflingDataset', 'TorchShufflingDataset', (['filenames', 'num_epochs', 'world_size', 'batch_size', 'rank'], {'num_reducers': 'num_reducers', 'max_concurrent_epochs': 'max_concurrent_epochs', 'feature_columns': 'feature_columns', 'feature_types': 'feature_types', 'label_column': 'label_column', 'label_type': 'label_type'}), '(filenames, num_epochs, world_size, batch_size, rank,\n    num_reducers=num_reducers, max_concurrent_epochs=max_concurrent_epochs,\n    feature_columns=feature_columns, feature_types=feature_types,\n    label_column=label_column, label_type=label_type)\n', (9460, 9709), False, 'from ray_shuffling_data_loader.torch_dataset import TorchShufflingDataset\n'), ((9976, 10006), 'ray.init', 'ray.init', ([], {'address': 'args.address'}), '(address=args.address)\n', (9984, 10006), False, 'import ray\n'), ((12163, 12204), 'horovod.ray.RayExecutor.create_settings', 'RayExecutor.create_settings', ([], {'timeout_s': '(30)'}), '(timeout_s=30)\n', (12190, 12204), False, 'from horovod.ray import RayExecutor\n'), ((12220, 12329), 'horovod.ray.RayExecutor', 'RayExecutor', (['settings'], {'use_gpu': '(True)', 'gpus_per_worker': '(1)', 'cpus_per_worker': 'cpus_per_worker'}), '(settings, use_gpu=True, gpus_per_worker=1, cpus_per_worker=\n    cpus_per_worker, **worker_kwargs)\n', (12231, 12329), False, 'from horovod.ray import RayExecutor\n'), ((3628, 3659), 'torch.nn.Conv2d', 'nn.Conv2d', (['(1)', '(10)'], {'kernel_size': '(5)'}), '(1, 10, kernel_size=5)\n', (3637, 3659), True, 'import torch.nn as nn\n'), ((3681, 3713), 'torch.nn.Conv2d', 'nn.Conv2d', (['(10)', '(20)'], {'kernel_size': '(5)'}), '(10, 20, kernel_size=5)\n', (3690, 3713), True, 'import torch.nn as nn\n'), ((3740, 3754), 'torch.nn.Dropout2d', 'nn.Dropout2d', ([], {}), '()\n', (3752, 3754), True, 'import torch.nn as nn\n'), ((3774, 3792), 'torch.nn.Linear', 'nn.Linear', (['(320)', '(50)'], {}), '(320, 50)\n', (3783, 3792), True, 'import torch.nn as nn\n'), ((3812, 3829), 'torch.nn.Linear', 'nn.Linear', (['(50)', '(10)'], {}), '(50, 10)\n', (3821, 3829), True, 'import torch.nn as nn\n'), ((4048, 4084), 'torch.nn.functional.dropout', 'F.dropout', (['x'], {'training': 'self.training'}), '(x, training=self.training)\n', (4057, 4084), True, 'import torch.nn.functional as F\n'), ((4124, 4140), 'torch.nn.functional.log_softmax', 'F.log_softmax', (['x'], {}), '(x)\n', (4137, 4140), True, 'import torch.nn.functional as F\n'), ((4267, 4292), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (4290, 4292), False, 'import torch\n'), ((4413, 4446), 'torch.cuda.manual_seed', 'torch.cuda.manual_seed', (['args.seed'], {}), '(args.seed)\n', (4435, 4446), False, 'import torch\n'), ((4929, 4939), 'horovod.torch.size', 'hvd.size', ([], {}), '()\n', (4937, 4939), True, 'import horovod.torch as hvd\n'), ((4978, 5003), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (5001, 5003), False, 'import torch\n'), ((6220, 6242), 'timeit.default_timer', 'timeit.default_timer', ([], {}), '()\n', (6240, 6242), False, 'import timeit\n'), ((7280, 7305), 'numpy.mean', 'np.mean', (['batch_wait_times'], {}), '(batch_wait_times)\n', (7287, 7305), True, 'import numpy as np\n'), ((7336, 7360), 'numpy.std', 'np.std', (['batch_wait_times'], {}), '(batch_wait_times)\n', (7342, 7360), True, 'import numpy as np\n'), ((7391, 7415), 'numpy.max', 'np.max', (['batch_wait_times'], {}), '(batch_wait_times)\n', (7397, 7415), True, 'import numpy as np\n'), ((7446, 7470), 'numpy.min', 'np.min', (['batch_wait_times'], {}), '(batch_wait_times)\n', (7452, 7470), True, 'import numpy as np\n'), ((8810, 8824), 'time.sleep', 'time.sleep', (['(10)'], {}), '(10)\n', (8820, 8824), False, 'import time\n'), ((9230, 9246), 'ray_shuffling_data_loader.data_generation.DATA_SPEC.keys', 'DATA_SPEC.keys', ([], {}), '()\n', (9244, 9246), False, 'from ray_shuffling_data_loader.data_generation import generate_data, DATA_SPEC\n'), ((10236, 10257), 'tempfile.gettempdir', 'tempfile.gettempdir', ([], {}), '()\n', (10255, 10257), False, 'import tempfile\n'), ((10322, 10348), 'os.path.exists', 'os.path.exists', (['cache_path'], {}), '(cache_path)\n', (10336, 10348), False, 'import os\n'), ((10815, 10908), 'ray_shuffling_data_loader.data_generation.generate_data', 'generate_data', (['num_rows', 'num_files', 'num_row_groups_per_file', 'max_row_group_skew', 'data_dir'], {}), '(num_rows, num_files, num_row_groups_per_file,\n    max_row_group_skew, data_dir)\n', (10828, 10908), False, 'from ray_shuffling_data_loader.data_generation import generate_data, DATA_SPEC\n'), ((4387, 4403), 'horovod.torch.local_rank', 'hvd.local_rank', ([], {}), '()\n', (4401, 4403), True, 'import horovod.torch as hvd\n'), ((4721, 4731), 'horovod.torch.size', 'hvd.size', ([], {}), '()\n', (4729, 4731), True, 'import horovod.torch as hvd\n'), ((5183, 5199), 'horovod.torch.nccl_built', 'hvd.nccl_built', ([], {}), '()\n', (5197, 5199), True, 'import horovod.torch as hvd\n'), ((5225, 5241), 'horovod.torch.local_size', 'hvd.local_size', ([], {}), '()\n', (5239, 5241), True, 'import horovod.torch as hvd\n'), ((6920, 6957), 'time.sleep', 'time.sleep', (['args.mock_train_step_time'], {}), '(args.mock_train_step_time)\n', (6930, 6957), False, 'import time\n'), ((7165, 7187), 'timeit.default_timer', 'timeit.default_timer', ([], {}), '()\n', (7185, 7187), False, 'import timeit\n'), ((7213, 7235), 'timeit.default_timer', 'timeit.default_timer', ([], {}), '()\n', (7233, 7235), False, 'import timeit\n'), ((9325, 9343), 'ray_shuffling_data_loader.data_generation.DATA_SPEC.values', 'DATA_SPEC.values', ([], {}), '()\n', (9341, 9343), False, 'from ray_shuffling_data_loader.data_generation import generate_data, DATA_SPEC\n'), ((6471, 6496), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (6494, 6496), False, 'import torch\n'), ((10448, 10462), 'pickle.load', 'pickle.load', (['f'], {}), '(f)\n', (10459, 10462), False, 'import pickle\n'), ((11147, 11185), 'pickle.dump', 'pickle.dump', (['(filenames, num_bytes)', 'f'], {}), '((filenames, num_bytes), f)\n', (11158, 11185), False, 'import pickle\n'), ((11197, 11227), 'ray_shuffling_data_loader.stats.human_readable_size', 'human_readable_size', (['num_bytes'], {}), '(num_bytes)\n', (11216, 11227), False, 'from ray_shuffling_data_loader.stats import human_readable_size\n'), ((6414, 6436), 'timeit.default_timer', 'timeit.default_timer', ([], {}), '()\n', (6434, 6436), False, 'import timeit\n'), ((11059, 11080), 'tempfile.gettempdir', 'tempfile.gettempdir', ([], {}), '()\n', (11078, 11080), False, 'import tempfile\n')]