From a29fe96f5d44d849545b6f33e49741a87dfa6269 Mon Sep 17 00:00:00 2001
From: Richard Liaw <rliaw@berkeley.edu>
Date: Wed, 4 Aug 2021 17:59:03 -0700
Subject: [PATCH 1/3] checkin

Signed-off-by: Richard Liaw <rliaw@berkeley.edu>
---
 examples/horovod/dataset_shuffle.py     | 352 ++++++++++++++++++++++++
 examples/horovod/embedding_model.py     | 200 ++++++++++++++
 examples/horovod/generate_data_utils.py |  96 +++++++
 examples/horovod/ray_torch_shuffle.py   |  42 ++-
 4 files changed, 681 insertions(+), 9 deletions(-)
 create mode 100644 examples/horovod/dataset_shuffle.py
 create mode 100644 examples/horovod/embedding_model.py
 create mode 100644 examples/horovod/generate_data_utils.py

diff --git a/examples/horovod/dataset_shuffle.py b/examples/horovod/dataset_shuffle.py
new file mode 100644
index 0000000..b479253
--- /dev/null
+++ b/examples/horovod/dataset_shuffle.py
@@ -0,0 +1,352 @@
+import argparse
+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
+import ray
+
+import os
+
+import pandas as pd
+import numpy as np
+
+import ray
+from generate_data_utils import DATA_SPEC, generate_data
+from embedding_model import MyModel, annotation, huber_loss
+
+
+# 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(
+    "--debug",
+    action="store_true",
+    default=False,
+    help="disables hvd")
+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=4)
+parser.add_argument("--cpus-per-worker", type=int, default=2)
+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")
+
+def construct_optimizers(model):
+    sparse_params = []
+    dense_params = []
+    for k,v in model.named_parameters():
+        if "input.embeddings.embeddings" in k:
+            sparse_params.append((k,v))
+        else:
+            dense_params.append((k,v))
+
+    optimizers = []
+    if len(dense_params) > 0:
+        opt = optim.Adam([v for _,v in dense_params], lr=0.001)
+        opt = hvd.DistributedOptimizer(opt, dense_params)
+        optimizers.append(opt)
+    if len(sparse_params) > 0:
+        opt = optim.SparseAdam([v for _,v in sparse_params], lr=0.001)
+        opt = hvd.DistributedOptimizer(opt, sparse_params)
+        optimizers.append(opt)
+
+    if hvd.rank() == 0:
+        print(optimizers)
+
+    return optimizers
+
+
+def train_main(args, splits):
+    # 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()
+    my_split = splits[rank]
+    train_dataset = create_torch_iterator(my_split, args.batch_size)
+
+    model = MyModel(annotation)
+    # By default, Adasum doesn"t need scaling up learning rate.
+    if torch.cuda.is_available() and not args.no_cuda:
+        # Move model to GPU.
+        model.cuda()
+
+    optimizers = construct_optimizers(model)
+    loss_function = huber_loss
+    # Horovod: broadcast parameters & optimizer state.
+    hvd.broadcast_parameters(model.state_dict(), root_rank=0)
+    for opt in optimizers:
+        hvd.broadcast_optimizer_state(opt, root_rank=0)
+
+    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:
+                data = [t.cuda() for t in data]
+                target = target.cuda()
+            optimizer.zero_grad()
+            batch = OrderedDict()
+            for name, tensor in zip(annotation["embeddings"], data):
+                batch[name] = tensor
+
+            batch_pred = model(batch)
+
+            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 = loss_function(batch_pred, target, delta=60)
+            loss.mean().backward()
+            for opt in optimizers:
+                opt.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):
+        new_batch_times = _train(epoch)
+        new_batch_times.pop(0)
+        batch_wait_times.extend(new_batch_times)
+
+    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")
+
+    with open(f"/tmp/dataset_shuffle_worker_{rank}.csv", "wt") as fp:
+        fp.writelines([f"{f:.6f}\n" for f in batch_wait_times])
+
+    # 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.")
+######################################################
+
+
+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
+}
+
+
+def create_torch_iterator(split, batch_size, rank=None):
+    print(f"Creating Torch shuffling dataset for worker {rank} with "
+          f"{batch_size} batch size.")
+    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()
+
+    torch_iterator = split.to_torch(
+         label_column=label_column,
+         feature_columns=feature_columns,
+         label_column_dtype=label_type,
+         feature_column_dtypes=feature_types,
+         batch_size=batch_size,
+         # prefetch_blocks: int = 0,
+         # drop_last: bool = False
+    )
+    return torch_iterator
+
+
+@ray.remote
+def consume(split, rank=None, batch_size=None):
+    torch_iterator = create_torch_iterator(split, batch_size=batch_size, rank=rank)
+    for i, (x, y) in enumerate(torch_iterator):
+        if i % 10 == 0:
+            print(i)
+    return
+
+def create_dataset(filenames):
+    pipeline = ray.data.read_parquet(list(filenames))\
+        .repeat().random_shuffle()
+    return pipeline
+
+
+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="auto")
+
+    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:
+        filenames = [
+            f's3://shuffling-data-loader-benchmarks/data/input_data_{i}.parquet.snappy' for i in range(0, 25)
+        ]
+    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(filenames)
+    pipeline = create_dataset(filenames)
+    splits = pipeline.split(args.num_workers)
+
+    if args.debug:
+        tasks = [
+            consume.remote(split, rank=idx, batch_size=args.batch_size)
+            for idx, split in enumerate(splits)
+        ]
+        ray.get(tasks)
+    else:
+
+        print("Create Ray executor")
+        settings = RayExecutor.create_settings(timeout_s=30)
+        executor = RayExecutor(
+            settings,
+            num_workers=args.num_workers,
+            use_gpu=not args.no_cuda)
+        executor.start()
+        executor.run(train_main, args=[args, splits])
+        executor.shutdown()
diff --git a/examples/horovod/embedding_model.py b/examples/horovod/embedding_model.py
new file mode 100644
index 0000000..2f340e3
--- /dev/null
+++ b/examples/horovod/embedding_model.py
@@ -0,0 +1,200 @@
+
+class EnumType:
+    def __init__(self, name, maxCategories, embed_dim=0):
+        self.name = name
+        self.maxCategories = maxCategories
+        self.embed_dim = embed_dim
+
+
+class EmbeddingsCollection(nn.Module):
+    def __init__(self, annotation, concat=True, use_batch_norm=True):
+        super().__init__()
+
+        self.output_size = 0
+        self.embeddings = {}
+        for k, embed in annotation.items():
+            self.embeddings[k] = torch.nn.Embedding(
+                embed.maxCategories,
+                embed.embed_dim,
+                sparse=True,
+            )
+            self.output_size += embed.embed_dim
+        self.embeddings = torch.nn.ModuleDict(self.embeddings)
+        self.concat = concat
+
+        if use_batch_norm and concat:
+            self.input_normalizer = torch.nn.BatchNorm1d(self.output_size)
+        else:
+            self.input_normalizer = None
+
+    def forward(self, input: OrderedDict):
+        embeds = []
+        k: str
+        v: nn.modules.sparse.Embedding
+        for k, v in self.embeddings.items():
+            embeds.append(v(input[k]))
+        if self.concat:
+            embeds = torch.cat(embeds, -1)
+        if self.input_normalizer:
+            return self.input_normalizer(embeds)
+        return embeds
+
+
+class OneHotEncoding(nn.Module):
+    def __init__(self, maxSize):
+        super().__init__()
+        assert isinstance(maxSize, int)
+        self.maxSize = maxSize
+
+    def forward(self, x):
+        result = torch.zeros(*x.shape, self.maxSize, device=x.device, dtype=torch.float32)
+        result.scatter_(-1, x.unsqueeze(-1), 1.0)
+        return result
+
+
+class OneHotEncodingCollection(nn.Module):
+    def __init__(self, annotation, use_batch_norm=True):
+        super().__init__()
+
+        self.output_size = 0
+        self.one_hot = {}
+        for k, embed in annotation.items():
+            self.one_hot[k] = OneHotEncoding(embed.maxCategories)
+            self.output_size += embed.maxCategories
+
+        self.one_hot = torch.nn.ModuleDict(self.one_hot)
+
+    def forward(self, input: OrderedDict):
+        embeds = []
+        k: str
+        v: OneHotEncoding
+        for k, v in self.one_hot.items():
+            embeds.append(v(input[k]))
+        return embeds
+
+
+class MyInput(nn.Module):
+    def __init__(self, annotation, use_bn=False):
+        super().__init__()
+
+        self.output_size = 0
+        if 'embeddings' in annotation:
+            self.embeddings = EmbeddingsCollection(annotation["embeddings"])
+            self.output_size += self.embeddings.output_size
+
+        if 'one_hot' in annotation:
+            self.one_hot = OneHotEncodingCollection(annotation["one_hot"])
+            self.output_size += self.one_hot.output_size
+
+        if self.output_size == 0:
+            raise ValueError("MyInput was not able to process " + str(annotation.keys()))
+
+    def forward(self, buffer: OrderedDict):
+        features: List[torch.Tensor] = []
+
+        if hasattr(self, "embeddings"):
+            features.append(self.embeddings(buffer["embeddings"]))
+
+        if hasattr(self, "one_hot"):
+            features += self.one_hot(buffer['one_hot'])
+
+        return torch.cat(features, -1)
+
+
+class MyLinear(nn.Linear):
+    def forward(self, x):
+        if len(self.weight.shape) == 2:
+            return super().forward(x)
+        else:
+            return torch.baddbmm(self.bias.unsqueeze(-2), x, self.weight.transpose(-1, -2))
+
+
+class MySequential(nn.Module):
+    def __init__(self, input_size, layers=None,
+                 use_dropout=False, dropout_rate=0.5, use_bn=False,
+                 activation=nn.ReLU, activate_final=False):
+        super().__init__()
+        self.layers = []
+
+        if isinstance(input_size, (tuple, list)):
+            input_size, = input_size
+
+        if not isinstance(activation, (tuple, list)):
+            activation = [activation] * len(layers)
+            if not activate_final:
+                activation[-1] = None
+
+        hidden_dim = input_size
+        for l_dim, act in zip(layers or [], activation):
+            self.layers.append(MyLinear(hidden_dim, l_dim))
+            if act is not None:
+                if use_bn:
+                    self.layers.append(nn.BatchNorm1d(l_dim))
+                self.layers.append(act())
+                if use_dropout:
+                    self.layers.append(nn.Dropout(dropout_rate))
+            hidden_dim = l_dim
+        self.layers = nn.Sequential(*self.layers)
+
+    def forward(self, inputs):
+        return self.layers(inputs)
+
+
+class MyModel(nn.Module):
+    def __init__(self, annotation: OrderedDict, layers=(512, 1), layer_size=512, num_heads=4,
+                 t_size=128, use_dropout=False, use_bn=True, num_segments=48):
+
+        super().__init__()
+        annotation = annotation
+        if not layers:
+            layers = [layer_size, layer_size, 1]
+        self.num_segments = num_segments
+        self.input = MyInput(annotation)
+
+        self.sequential = MySequential(self.input.output_size, layers=layers, use_dropout=use_dropout, use_bn=use_bn)
+        self.sequential2 = MySequential(layers[-1]+self.num_segments+2, layers=[1], use_dropout=use_dropout, use_bn=use_bn,)
+
+    def forward(self, buffer):
+        hot1 = buffer["one_hot"]["hot1"]
+        hot1_ref = arange(self.num_segments).reshape(1, self.num_segments).float()
+        hot0 = buffer["one_hot"]["hot0"]
+        hot0_ref = arange(2).reshape(1, 2).float()
+        if hot1.is_cuda:
+            hot1_ref = hot1_ref.cuda()
+            hot0_ref = hot0_ref.cuda()
+        hot1_onehot = (hot1.float().unsqueeze(-1) == hot1_ref).float()
+        hot0_onehot = (hot0.float().unsqueeze(-1) == hot0_ref).float()
+        network = self.sequential(self.input(buffer))
+        link = self.sequential2(cat([network, hot1_onehot, hot0_onehot], dim=1))
+        return link.squeeze(-1)
+
+
+def huber_loss(a, b, delta=20):
+    err = (a - b).abs()
+    mask = err < delta
+    return (0.5 * mask * (err ** 2)) + ~mask * (err * delta - 0.5 * (delta ** 2))
+
+
+annotation = OrderedDict()
+annotation["embeddings"] = OrderedDict()
+annotation["embeddings"]["name0"] = EnumType("name0", 2385, 12)
+annotation["embeddings"]["name1"] = EnumType("name1", 201, 8)
+annotation["embeddings"]["name2"] = EnumType("name2", 201, 8)
+annotation["embeddings"]["name3"] = EnumType("name3", 6, 3)
+annotation["embeddings"]["name4"] = EnumType("name4", 19, 5)
+annotation["embeddings"]["name5"] = EnumType("name5", 1441, 11)
+annotation["embeddings"]["name6"] = EnumType("name6", 201, 8)
+annotation["embeddings"]["name7"] = EnumType("name7", 22, 5)
+annotation["embeddings"]["name8"] = EnumType("name8", 156, 8)
+annotation["embeddings"]["name9"] = EnumType("name9", 1216, 11)
+annotation["embeddings"]["name10"] = EnumType("name10", 9216, 14)
+annotation["embeddings"]["name11"] = EnumType("name11", 88999, 17)
+annotation["embeddings"]["name12"] = EnumType("name12", 941792, 20)
+annotation["embeddings"]["name13"] = EnumType("name13", 9405, 14)
+annotation["embeddings"]["name14"] = EnumType("name14", 83332, 17)
+annotation["embeddings"]["name15"] = EnumType("name15", 828767, 20)
+annotation["embeddings"]["name16"] = EnumType("name16", 945195, 20)
+
+annotation["one_hot"] = OrderedDict()
+annotation["one_hot"]["hot0"] = EnumType("hot0", 3) # one_hot doesn't use dimension
+annotation["one_hot"]["hot1"] = EnumType("hot1", 50) # one_hot doesn't use dimension
diff --git a/examples/horovod/generate_data_utils.py b/examples/horovod/generate_data_utils.py
new file mode 100644
index 0000000..ffaa4a3
--- /dev/null
+++ b/examples/horovod/generate_data_utils.py
@@ -0,0 +1,96 @@
+import ray
+
+import os
+
+import pandas as pd
+import numpy as np
+
+#######################################################################################
+#
+# Data generation utilities for the shuffling data loader.
+#
+
+def generate_data(num_rows, num_files, num_row_groups_per_file,
+                  max_row_group_skew, data_dir):
+    assert max_row_group_skew == 0.0
+    # TODO(Clark): Generate skewed row groups according to max_row_group_skew.
+    results = []
+    for file_index, global_row_index in enumerate(
+            range(0, num_rows, num_rows // num_files)):
+        num_rows_in_file = min(num_rows // num_files,
+                               num_rows - global_row_index)
+        results.append(
+            generate_file.remote(file_index, global_row_index,
+                                 num_rows_in_file, num_row_groups_per_file,
+                                 data_dir))
+    filenames, data_sizes = zip(*ray.get(results))
+    return filenames, sum(data_sizes)
+
+
+@ray.remote
+def generate_file(file_index, global_row_index, num_rows_in_file,
+                  num_row_groups_per_file, data_dir):
+    # TODO(Clark): Generate skewed row groups according to max_row_group_skew.
+    # TODO(Clark): Optimize this data generation to reduce copies and
+    # progressively write smaller buffers to the Parquet file.
+    buffs = []
+    for group_index, group_global_row_index in enumerate(
+            range(0, num_rows_in_file,
+                  num_rows_in_file // num_row_groups_per_file)):
+        num_rows_in_group = min(num_rows_in_file // num_row_groups_per_file,
+                                num_rows_in_file - group_global_row_index)
+        buffs.append(
+            generate_row_group(group_index, group_global_row_index,
+                               num_rows_in_group))
+    df = pd.concat(buffs)
+    data_size = df.memory_usage(deep=True).sum()
+    filename = os.path.join(data_dir,
+                            f"input_data_{file_index}.parquet.snappy")
+    df.to_parquet(
+        filename,
+        engine="pyarrow",
+        compression="snappy",
+        row_group_size=num_rows_in_file // num_row_groups_per_file)
+    return filename, data_size
+
+
+DATA_SPEC = {
+    "embeddings_name0": (0, 2385, np.int64),
+    "embeddings_name1": (0, 201, np.int64),
+    "embeddings_name2": (0, 201, np.int64),
+    "embeddings_name3": (0, 6, np.int64),
+    "embeddings_name4": (0, 19, np.int64),
+    "embeddings_name5": (0, 1441, np.int64),
+    "embeddings_name6": (0, 201, np.int64),
+    "embeddings_name7": (0, 22, np.int64),
+    "embeddings_name8": (0, 156, np.int64),
+    "embeddings_name9": (0, 1216, np.int64),
+    "embeddings_name10": (0, 9216, np.int64),
+    "embeddings_name11": (0, 88999, np.int64),
+    "embeddings_name12": (0, 941792, np.int64),
+    "embeddings_name13": (0, 9405, np.int64),
+    "embeddings_name14": (0, 83332, np.int64),
+    "embeddings_name15": (0, 828767, np.int64),
+    "embeddings_name16": (0, 945195, np.int64),
+    "one_hot0": (0, 3, np.int64),
+    "one_hot1": (0, 50, np.int64),
+    "labels": (0, 1, np.float64),
+}
+
+
+def generate_row_group(group_index, global_row_index, num_rows_in_group):
+    buffer = {
+        "key": np.array(
+            range(global_row_index, global_row_index + num_rows_in_group)),
+    }
+    for col, (low, high, dtype) in DATA_SPEC.items():
+        if dtype in (np.int16, np.int32, np.int64):
+            buffer[col] = np.random.randint(
+                low, high, num_rows_in_group, dtype=dtype)
+        elif dtype in (np.float32, np.float64):
+            buffer[col] = (
+                (high - low) * np.random.rand(num_rows_in_group) + low)
+
+    return pd.DataFrame(buffer)
+
+################################################################################
\ No newline at end of file
diff --git a/examples/horovod/ray_torch_shuffle.py b/examples/horovod/ray_torch_shuffle.py
index 1d4a7bf..eecd05e 100644
--- a/examples/horovod/ray_torch_shuffle.py
+++ b/examples/horovod/ray_torch_shuffle.py
@@ -101,8 +101,10 @@
     default=1.0,
     help=("apply gradient predivide factor in optimizer "
           "(default: 1.0)"))
-parser.add_argument("--num-workers", type=int, default=4)
-parser.add_argument("--cpus-per-worker", type=int, default=2)
+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.
@@ -199,9 +201,12 @@ def _train(epoch):
         batch_wait_times = []
         for batch_idx, (data, target) in enumerate(train_dataset):
             batch_wait_times.append(timeit.default_timer() - last_batch_time)
+            print(type(data), type(target), args.no_cuda)
             if torch.cuda.is_available() and not args.no_cuda:
+                print("checking isinstance")
                 if isinstance(data, list):
-                    data = [t.cuda() for t in data]
+                    data = torch.tensor(data).cuda()
+
                 target = target.cuda()
             optimizer.zero_grad()
             # output = model(data)
@@ -231,9 +236,9 @@ def _train(epoch):
     print(f"Starting training on worker {rank}.")
     batch_wait_times = []
     for epoch in range(args.epochs):
-        # TODO(Clark): Don't include stats from first epoch since we already
-        # expect that epoch to be cold?
-        batch_wait_times.extend(_train(epoch))
+        new_batch_times = _train(epoch)
+        new_batch_times.pop(0)
+        batch_wait_times.extend(new_batch_times)
     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)
@@ -244,6 +249,10 @@ def _train(epoch):
           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")
+
+    with open(f"/tmp/ray_torch_shuffle_worker_{rank}.csv", "wt") as fp:
+        fp.writelines([f"{f:.6f}\n" for f in batch_wait_times])
+
     # TODO(Clark): Add logic to the dataset abstraction so we don't have to do
     # this.
     if rank == 0:
@@ -316,14 +325,29 @@ def create_dataset(filenames, *, batch_size, rank, num_epochs, world_size,
           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,
-        num_workers=num_workers,
-        use_gpu=True,
-        cpus_per_worker=cpus_per_worker)
+        use_gpu=not args.no_cuda,
+        gpus_per_worker=int(not args.no_cuda),
+        cpus_per_worker=cpus_per_worker,
+        **worker_kwargs)
     executor.start()
     executor.run(train_main, args=[args, filenames])
     executor.shutdown()

From 232f708dacf190c228c2b10448885c47d2c7ee23 Mon Sep 17 00:00:00 2001
From: Richard Liaw <rliaw@berkeley.edu>
Date: Thu, 5 Aug 2021 01:53:05 -0700
Subject: [PATCH 2/3] simples

Signed-off-by: Richard Liaw <rliaw@berkeley.edu>
---
 examples/horovod/dataset_shuffle.py           |  49 ++++----
 examples/horovod/generate_data_utils.py       |  96 ---------------
 examples/horovod/mini_shuffle_s3.py           | 113 ++++++++++++++++++
 .../embedding_model.py                        |  14 ++-
 4 files changed, 148 insertions(+), 124 deletions(-)
 delete mode 100644 examples/horovod/generate_data_utils.py
 create mode 100644 examples/horovod/mini_shuffle_s3.py
 rename {examples/horovod => ray_shuffling_data_loader}/embedding_model.py (96%)

diff --git a/examples/horovod/dataset_shuffle.py b/examples/horovod/dataset_shuffle.py
index b479253..7f11849 100644
--- a/examples/horovod/dataset_shuffle.py
+++ b/examples/horovod/dataset_shuffle.py
@@ -1,3 +1,4 @@
+from collections import OrderedDict
 import argparse
 import os
 import pickle
@@ -21,9 +22,9 @@
 import numpy as np
 
 import ray
-from generate_data_utils import DATA_SPEC, generate_data
-from embedding_model import MyModel, annotation, huber_loss
-
+from ray_shuffling_data_loader.data_generation import DATA_SPEC, generate_data
+from ray_shuffling_data_loader.embedding_model import MyModel, annotation, huber_loss
+DEFAULT_DATA_DIR = "s3://shuffling-data-loader-benchmarks/data/"
 
 # Training settings
 parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
@@ -101,9 +102,9 @@
 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("--read-cache", action="store_true", default=False)
+parser.add_argument("--num-rows", type=int, default=2 * (10**8))
+parser.add_argument("--num-files", type=int, default=50)
 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)
@@ -154,7 +155,7 @@ def train_main(args, splits):
     my_split = splits[rank]
     train_dataset = create_torch_iterator(my_split, args.batch_size)
 
-    model = MyModel(annotation)
+    model = MyModel(annotation, use_bn=False)
     # By default, Adasum doesn"t need scaling up learning rate.
     if torch.cuda.is_available() and not args.no_cuda:
         # Move model to GPU.
@@ -179,10 +180,16 @@ def _train(epoch):
             if torch.cuda.is_available() and not args.no_cuda:
                 data = [t.cuda() for t in data]
                 target = target.cuda()
-            optimizer.zero_grad()
+            for opt in optimizers:
+                opt.zero_grad()
             batch = OrderedDict()
+            batch["embeddings"] = OrderedDict()
+            batch["one_hot"] = OrderedDict()
             for name, tensor in zip(annotation["embeddings"], data):
-                batch[name] = tensor
+                batch["embeddings"][name] = tensor
+            hot0, hot1 = data[-2:]
+            batch["one_hot"]["hot0"] = hot0
+            batch["one_hot"]["hot1"] = hot1
 
             batch_pred = model(batch)
 
@@ -241,8 +248,6 @@ def _train(epoch):
 ######################################################
 
 
-DEFAULT_DATA_DIR = "s3://shuffling-data-loader-benchmarks/data/"
-
 numpy_to_torch_dtype = {
     np.bool: torch.bool,
     np.uint8: torch.uint8,
@@ -288,8 +293,8 @@ def consume(split, rank=None, batch_size=None):
             print(i)
     return
 
-def create_dataset(filenames):
-    pipeline = ray.data.read_parquet(list(filenames))\
+def create_dataset(data_dir):
+    pipeline = ray.data.read_parquet(data_dir)\
         .repeat().random_shuffle()
     return pipeline
 
@@ -305,32 +310,26 @@ def create_dataset(filenames):
     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:
-        filenames = [
-            f's3://shuffling-data-loader-benchmarks/data/input_data_{i}.parquet.snappy' for i in range(0, 25)
-        ]
-    if not filenames:
+    data_dir = os.path.join(args.data_dir, f"r{num_rows:_}-f{num_files}/")
+
+    if not args.read_cache:
         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(f"Saved to: {data_dir}")
 
-    print(filenames)
-    pipeline = create_dataset(filenames)
+    pipeline = create_dataset(data_dir)
     splits = pipeline.split(args.num_workers)
 
     if args.debug:
diff --git a/examples/horovod/generate_data_utils.py b/examples/horovod/generate_data_utils.py
deleted file mode 100644
index ffaa4a3..0000000
--- a/examples/horovod/generate_data_utils.py
+++ /dev/null
@@ -1,96 +0,0 @@
-import ray
-
-import os
-
-import pandas as pd
-import numpy as np
-
-#######################################################################################
-#
-# Data generation utilities for the shuffling data loader.
-#
-
-def generate_data(num_rows, num_files, num_row_groups_per_file,
-                  max_row_group_skew, data_dir):
-    assert max_row_group_skew == 0.0
-    # TODO(Clark): Generate skewed row groups according to max_row_group_skew.
-    results = []
-    for file_index, global_row_index in enumerate(
-            range(0, num_rows, num_rows // num_files)):
-        num_rows_in_file = min(num_rows // num_files,
-                               num_rows - global_row_index)
-        results.append(
-            generate_file.remote(file_index, global_row_index,
-                                 num_rows_in_file, num_row_groups_per_file,
-                                 data_dir))
-    filenames, data_sizes = zip(*ray.get(results))
-    return filenames, sum(data_sizes)
-
-
-@ray.remote
-def generate_file(file_index, global_row_index, num_rows_in_file,
-                  num_row_groups_per_file, data_dir):
-    # TODO(Clark): Generate skewed row groups according to max_row_group_skew.
-    # TODO(Clark): Optimize this data generation to reduce copies and
-    # progressively write smaller buffers to the Parquet file.
-    buffs = []
-    for group_index, group_global_row_index in enumerate(
-            range(0, num_rows_in_file,
-                  num_rows_in_file // num_row_groups_per_file)):
-        num_rows_in_group = min(num_rows_in_file // num_row_groups_per_file,
-                                num_rows_in_file - group_global_row_index)
-        buffs.append(
-            generate_row_group(group_index, group_global_row_index,
-                               num_rows_in_group))
-    df = pd.concat(buffs)
-    data_size = df.memory_usage(deep=True).sum()
-    filename = os.path.join(data_dir,
-                            f"input_data_{file_index}.parquet.snappy")
-    df.to_parquet(
-        filename,
-        engine="pyarrow",
-        compression="snappy",
-        row_group_size=num_rows_in_file // num_row_groups_per_file)
-    return filename, data_size
-
-
-DATA_SPEC = {
-    "embeddings_name0": (0, 2385, np.int64),
-    "embeddings_name1": (0, 201, np.int64),
-    "embeddings_name2": (0, 201, np.int64),
-    "embeddings_name3": (0, 6, np.int64),
-    "embeddings_name4": (0, 19, np.int64),
-    "embeddings_name5": (0, 1441, np.int64),
-    "embeddings_name6": (0, 201, np.int64),
-    "embeddings_name7": (0, 22, np.int64),
-    "embeddings_name8": (0, 156, np.int64),
-    "embeddings_name9": (0, 1216, np.int64),
-    "embeddings_name10": (0, 9216, np.int64),
-    "embeddings_name11": (0, 88999, np.int64),
-    "embeddings_name12": (0, 941792, np.int64),
-    "embeddings_name13": (0, 9405, np.int64),
-    "embeddings_name14": (0, 83332, np.int64),
-    "embeddings_name15": (0, 828767, np.int64),
-    "embeddings_name16": (0, 945195, np.int64),
-    "one_hot0": (0, 3, np.int64),
-    "one_hot1": (0, 50, np.int64),
-    "labels": (0, 1, np.float64),
-}
-
-
-def generate_row_group(group_index, global_row_index, num_rows_in_group):
-    buffer = {
-        "key": np.array(
-            range(global_row_index, global_row_index + num_rows_in_group)),
-    }
-    for col, (low, high, dtype) in DATA_SPEC.items():
-        if dtype in (np.int16, np.int32, np.int64):
-            buffer[col] = np.random.randint(
-                low, high, num_rows_in_group, dtype=dtype)
-        elif dtype in (np.float32, np.float64):
-            buffer[col] = (
-                (high - low) * np.random.rand(num_rows_in_group) + low)
-
-    return pd.DataFrame(buffer)
-
-################################################################################
\ No newline at end of file
diff --git a/examples/horovod/mini_shuffle_s3.py b/examples/horovod/mini_shuffle_s3.py
new file mode 100644
index 0000000..7b44340
--- /dev/null
+++ b/examples/horovod/mini_shuffle_s3.py
@@ -0,0 +1,113 @@
+import time
+import numpy as np
+import torch
+import ray
+import argparse
+import os
+
+DEFAULT_DATA_DIR = "s3://shuffling-data-loader-benchmarks/data/"
+
+
+def create_parser():
+    parser = argparse.ArgumentParser(description="Eric Example")
+    parser.add_argument("--address")
+    parser.add_argument("--num-rows", type=int, default=2 * (10**8))
+    parser.add_argument("--num-files", type=int, default=50)
+    parser.add_argument("--read-cache", action="store_true", default=False)
+    parser.add_argument("--data-dir", type=str, default=DEFAULT_DATA_DIR)
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=250000,
+        metavar="N",
+        help="input batch size for training (default: 64)")
+    parser.add_argument("--num-workers", type=int, default=4)
+    return parser
+
+
+def create_torch_iterator(split, batch_size, rank=None):
+    print(f"Creating Torch shuffling dataset for worker {rank} with "
+          f"{batch_size} batch size.")
+    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
+    }
+    DATA_SPEC = {
+        "embeddings_name0": (0, 2385, np.int64),
+        "embeddings_name1": (0, 201, np.int64),
+        "embeddings_name2": (0, 201, np.int64),
+        "embeddings_name3": (0, 6, np.int64),
+        "embeddings_name4": (0, 19, np.int64),
+        "embeddings_name5": (0, 1441, np.int64),
+        "embeddings_name6": (0, 201, np.int64),
+        "embeddings_name7": (0, 22, np.int64),
+        "embeddings_name8": (0, 156, np.int64),
+        "embeddings_name9": (0, 1216, np.int64),
+        "embeddings_name10": (0, 9216, np.int64),
+        "embeddings_name11": (0, 88999, np.int64),
+        "embeddings_name12": (0, 941792, np.int64),
+        "embeddings_name13": (0, 9405, np.int64),
+        "embeddings_name14": (0, 83332, np.int64),
+        "embeddings_name15": (0, 828767, np.int64),
+        "embeddings_name16": (0, 945195, np.int64),
+        "one_hot0": (0, 3, np.int64),
+        "one_hot1": (0, 50, np.int64),
+        "labels": (0, 1, np.float64),
+    }
+    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()
+
+    torch_iterator = split.to_torch(
+         label_column=label_column,
+         feature_columns=feature_columns,
+         label_column_dtype=label_type,
+         feature_column_dtypes=feature_types,
+         batch_size=batch_size,
+         # prefetch_blocks: int = 0,
+         # drop_last: bool = False
+    )
+    return torch_iterator
+
+
+def create_dataset(data_dir):
+    pipeline = ray.data.read_parquet(data_dir)\
+        .repeat().random_shuffle()
+    return pipeline
+
+if __name__ == '__main__':
+    parser = create_parser()
+    args = parser.parse_args()
+    print("Connecting to Ray cluster...")
+    ray.init(address=args.address)
+
+    data_dir = os.path.join(args.data_dir, f"r{args.num_rows:_}-f{args.num_files}/")
+    pipeline = create_dataset(data_dir)
+    splits = pipeline.split(args.num_workers, equal=True)
+
+    @ray.remote
+    def consume(split, rank=None, batch_size=None):
+        torch_iterator = create_torch_iterator(split, batch_size=batch_size, rank=rank)
+        for i, (x, y) in enumerate(torch_iterator):
+            time.sleep(1)
+            if i % 10 == 0:
+                print(i)
+        return
+
+    tasks = [
+        consume.remote(split, rank=idx, batch_size=args.batch_size)
+        for idx, split in enumerate(splits)
+    ]
+    ray.get(tasks)
diff --git a/examples/horovod/embedding_model.py b/ray_shuffling_data_loader/embedding_model.py
similarity index 96%
rename from examples/horovod/embedding_model.py
rename to ray_shuffling_data_loader/embedding_model.py
index 2f340e3..c76acc7 100644
--- a/examples/horovod/embedding_model.py
+++ b/ray_shuffling_data_loader/embedding_model.py
@@ -1,3 +1,11 @@
+from typing import List
+import argparse
+from collections import OrderedDict
+
+import torch
+from torch import nn, arange, cat, optim
+import torch.nn.functional as F
+import horovod.torch as hvd
 
 class EnumType:
     def __init__(self, name, maxCategories, embed_dim=0):
@@ -79,11 +87,11 @@ def __init__(self, annotation, use_bn=False):
 
         self.output_size = 0
         if 'embeddings' in annotation:
-            self.embeddings = EmbeddingsCollection(annotation["embeddings"])
+            self.embeddings = EmbeddingsCollection(annotation["embeddings"], use_batch_norm=use_bn)
             self.output_size += self.embeddings.output_size
 
         if 'one_hot' in annotation:
-            self.one_hot = OneHotEncodingCollection(annotation["one_hot"])
+            self.one_hot = OneHotEncodingCollection(annotation["one_hot"], use_batch_norm=use_bn)
             self.output_size += self.one_hot.output_size
 
         if self.output_size == 0:
@@ -165,7 +173,7 @@ def forward(self, buffer):
         hot1_onehot = (hot1.float().unsqueeze(-1) == hot1_ref).float()
         hot0_onehot = (hot0.float().unsqueeze(-1) == hot0_ref).float()
         network = self.sequential(self.input(buffer))
-        link = self.sequential2(cat([network, hot1_onehot, hot0_onehot], dim=1))
+        link = self.sequential2(cat([network, hot1_onehot, hot0_onehot], dim=2))
         return link.squeeze(-1)
 
 

From a3632664224b122cda07d87cc235042bd15a4885 Mon Sep 17 00:00:00 2001
From: Richard Liaw <rliaw@berkeley.edu>
Date: Fri, 13 Aug 2021 16:28:29 -0700
Subject: [PATCH 3/3] update

Signed-off-by: Richard Liaw <rliaw@berkeley.edu>
---
 examples/horovod/dataset_shuffle.py          | 4 ++--
 examples/horovod/mini_shuffle_s3.py          | 6 ++++--
 ray_shuffling_data_loader/data_generation.py | 2 +-
 3 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/examples/horovod/dataset_shuffle.py b/examples/horovod/dataset_shuffle.py
index 7f11849..b13a079 100644
--- a/examples/horovod/dataset_shuffle.py
+++ b/examples/horovod/dataset_shuffle.py
@@ -103,8 +103,8 @@
 
 # Synthetic training data generation settings.
 parser.add_argument("--read-cache", action="store_true", default=False)
-parser.add_argument("--num-rows", type=int, default=2 * (10**8))
-parser.add_argument("--num-files", type=int, default=50)
+parser.add_argument("--num-rows", type=int, default=2 * (10**9))
+parser.add_argument("--num-files", type=int, default=200)
 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)
diff --git a/examples/horovod/mini_shuffle_s3.py b/examples/horovod/mini_shuffle_s3.py
index 7b44340..f04c19b 100644
--- a/examples/horovod/mini_shuffle_s3.py
+++ b/examples/horovod/mini_shuffle_s3.py
@@ -83,8 +83,9 @@ def create_torch_iterator(split, batch_size, rank=None):
 
 
 def create_dataset(data_dir):
-    pipeline = ray.data.read_parquet(data_dir)\
-        .repeat().random_shuffle()
+    # pipeline = ray.data.read_parquet(data_dir)\
+    #     .repeat().random_shuffle()
+    print(ray.data.read_parquet(data_dir).take(1))
     return pipeline
 
 if __name__ == '__main__':
@@ -94,6 +95,7 @@ def create_dataset(data_dir):
     ray.init(address=args.address)
 
     data_dir = os.path.join(args.data_dir, f"r{args.num_rows:_}-f{args.num_files}/")
+    print(f"creating dataset from: {data_dir}")
     pipeline = create_dataset(data_dir)
     splits = pipeline.split(args.num_workers, equal=True)
 
diff --git a/ray_shuffling_data_loader/data_generation.py b/ray_shuffling_data_loader/data_generation.py
index c4725c4..d2a5e84 100644
--- a/ray_shuffling_data_loader/data_generation.py
+++ b/ray_shuffling_data_loader/data_generation.py
@@ -27,7 +27,7 @@ def generate_data(num_rows, num_files, num_row_groups_per_file,
     return filenames, sum(data_sizes)
 
 
-@ray.remote
+@ray.remote(num_cpus=4)
 def generate_file(file_index, global_row_index, num_rows_in_file,
                   num_row_groups_per_file, data_dir):
     # TODO(Clark): Generate skewed row groups according to max_row_group_skew.