From 833103717ed6de883ed65032d4c5606693dea0f8 Mon Sep 17 00:00:00 2001 From: Kai Fricke Date: Wed, 14 Jul 2021 13:41:38 +0100 Subject: [PATCH 1/3] test shuffle_map --- .../tests/test_shuffle.py | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 ray_shuffling_data_loader/tests/test_shuffle.py diff --git a/ray_shuffling_data_loader/tests/test_shuffle.py b/ray_shuffling_data_loader/tests/test_shuffle.py new file mode 100644 index 0000000..c8bd992 --- /dev/null +++ b/ray_shuffling_data_loader/tests/test_shuffle.py @@ -0,0 +1,71 @@ +import math +import tempfile +import unittest +import pytest + +import ray + +from ray_shuffling_data_loader.data_generation import generate_data +from ray_shuffling_data_loader.shuffle import shuffle_map + + +class DataLoaderShuffleTest(unittest.TestCase): + """This test suite validates core RayDMatrix functionality.""" + + def setUp(self): + self.num_rows = 10**4 + self.num_files = 1 + self.num_row_groups_per_file = 1 + self.max_row_group_skew = 0.0 + self.data_dir = tempfile.mkdtemp() + + self.filenames, self.num_bytes = generate_data( + self.num_rows, self.num_files, self.num_row_groups_per_file, + self.max_row_group_skew, self.data_dir) + + @classmethod + def setUpClass(cls): + ray.init(num_cpus=2) + + @classmethod + def tearDownClass(cls): + ray.shutdown() + + def testShuffleMap(self): + num_reducers = 4 + + # Calculate mean and SD of rows assigned to each reducer + p = 1 / num_reducers + mean = self.num_rows * p + sd = math.sqrt(self.num_rows * p * (1 - p)) + + reducer_parts = shuffle_map.remote( + filename=self.filenames[0], + num_reducers=4, + stats_collector=None, + epoch=0) + + fetched_parts = ray.get(reducer_parts) + + all_keys = [] + for i, part in enumerate(fetched_parts): + part_keys = part["key"].to_numpy() + + # 3sd = 99.7% chance of passing + assert mean - 3 * sd < len(part_keys) < mean + 3 * sd, \ + f"Not enough rows in partition {i}" + + assert len(set(part_keys)) == len(part_keys), \ + f"Keys in partition {i} are not distinct" + + all_keys.extend(part_keys) + + assert len(all_keys) == self.num_rows, "Not all rows were returned." + + assert len(set(all_keys)) == len(all_keys), \ + "Keys in full dataset are not distinct." + + +if __name__ == "__main__": + import sys + sys.exit(pytest.main(["-v", __file__])) From 49e88c1026f6447b0e7e7c02ee43f09df8aa511d Mon Sep 17 00:00:00 2001 From: Kai Fricke Date: Wed, 14 Jul 2021 14:18:16 +0100 Subject: [PATCH 2/3] test shuffle_reduce --- .../tests/test_shuffle.py | 45 ++++++++++++++++++- 1 file changed, 43 insertions(+), 2 deletions(-) diff --git a/ray_shuffling_data_loader/tests/test_shuffle.py b/ray_shuffling_data_loader/tests/test_shuffle.py index c8bd992..4cd49d5 100644 --- a/ray_shuffling_data_loader/tests/test_shuffle.py +++ b/ray_shuffling_data_loader/tests/test_shuffle.py @@ -3,10 +3,12 @@ import unittest import pytest +import pandas as pd + import ray from ray_shuffling_data_loader.data_generation import generate_data -from ray_shuffling_data_loader.shuffle import shuffle_map +from ray_shuffling_data_loader.shuffle import shuffle_map, shuffle_reduce class DataLoaderShuffleTest(unittest.TestCase): @@ -41,7 +43,7 @@ def testShuffleMap(self): reducer_parts = shuffle_map.remote( filename=self.filenames[0], - num_reducers=4, + num_reducers=num_reducers, stats_collector=None, epoch=0) @@ -65,6 +67,45 @@ def testShuffleMap(self): assert len(set(all_keys)) == len(all_keys), \ "Keys in full dataset are not distinct." + def testShuffleReduce(self): + num_reducers = 4 + num_shufflers = 2 + + reducer_parts = shuffle_map.remote( + filename=self.filenames[0], + num_reducers=num_reducers, + stats_collector=None, + epoch=0) + + fetched_parts = ray.get(reducer_parts) + + # We cannot get the original references here, so we just push + # to the object store again as a workaround + fetched_refs = [ray.put(part) for part in fetched_parts] + + parts_per_shuffler = num_reducers // num_shufflers + for i in range(num_shufflers): + unshuffled_refs = fetched_refs[(i * parts_per_shuffler):( + i + 1 * parts_per_shuffler)] + unshuffled_parts = fetched_parts[(i * parts_per_shuffler):( + i + 1 * parts_per_shuffler)] + + shuffled = ray.get( + shuffle_reduce.remote( + 0, + None, + 0, + *unshuffled_refs, + )) + + unshuffled = pd.concat(unshuffled_parts, copy=False) + + assert len(unshuffled) == len(shuffled), \ + "Length mismatch between unshuffled and shuffled parts" + + assert set(unshuffled) == set(shuffled), \ + "Key mismatch between unshuffled and shuffled parts" + if __name__ == "__main__": import sys From cee22d228b4eafc9aff0ebdee61a3bc7f8ede698 Mon Sep 17 00:00:00 2001 From: Kai Fricke Date: Thu, 15 Jul 2021 15:05:39 +0100 Subject: [PATCH 3/3] Shuffle end to end test --- .../tests/test_shuffle.py | 65 ++++++++++++++++++- run_ci_tests.sh | 1 + 2 files changed, 65 insertions(+), 1 deletion(-) diff --git a/ray_shuffling_data_loader/tests/test_shuffle.py b/ray_shuffling_data_loader/tests/test_shuffle.py index 4cd49d5..bce2ab2 100644 --- a/ray_shuffling_data_loader/tests/test_shuffle.py +++ b/ray_shuffling_data_loader/tests/test_shuffle.py @@ -1,6 +1,8 @@ import math import tempfile import unittest +from collections import defaultdict + import pytest import pandas as pd @@ -8,7 +10,8 @@ import ray from ray_shuffling_data_loader.data_generation import generate_data -from ray_shuffling_data_loader.shuffle import shuffle_map, shuffle_reduce +from ray_shuffling_data_loader.shuffle import shuffle_map, shuffle_reduce, \ + BatchConsumer, shuffle class DataLoaderShuffleTest(unittest.TestCase): @@ -106,6 +109,66 @@ def testShuffleReduce(self): assert set(unshuffled) == set(shuffled), \ "Key mismatch between unshuffled and shuffled parts" + def testShuffleEndToEnd(self): + class EndToEndConsumer(BatchConsumer): + def __init__(self): + self.rank_epoch_batches = defaultdict(dict) + + def consume(self, rank, epoch, batches): + self.rank_epoch_batches[rank][epoch] = ray.get(batches) + + def producer_done(self, rank, epoch): + pass + + def wait_until_ready(self, epoch): + return True + + def wait_until_all_epochs_done(self): + return True + + consumer = EndToEndConsumer() + num_epochs = 2 + num_reducers = 8 + num_trainers = 4 + + shuffle( + self.filenames, + batch_consumer=consumer, + num_epochs=num_epochs, + num_reducers=num_reducers, + num_trainers=num_trainers) + + assert len(consumer.rank_epoch_batches) == num_trainers, \ + "Trainer count mismatch" + + assert all(len(consumer.rank_epoch_batches[t]) == num_epochs + for t in consumer.rank_epoch_batches), \ + "Epoch count mismatch" + + for tid, epoch_batches in consumer.rank_epoch_batches.items(): + for i in range(len(epoch_batches) - 1): + assert len(epoch_batches[i]) == len( + epoch_batches[+1]) == num_epochs, \ + "Length mismatch in epoch batches" + + df1 = pd.concat(epoch_batches[i], copy=False) + df2 = pd.concat(epoch_batches[i], copy=False) + + keys1 = df1["key"].to_numpy() + keys2 = df2["key"].to_numpy() + + set1 = set(keys1) + set2 = set(keys2) + + assert len(set1) == len(keys1), \ + "Keys in dataset are not distinct." + + assert len(set2) == len(keys2), \ + "Keys in dataset are not distinct." + + assert set1 == set2, \ + "Shuffled key sets are not equal." + if __name__ == "__main__": import sys diff --git a/run_ci_tests.sh b/run_ci_tests.sh index b0573a1..a0cdc19 100755 --- a/run_ci_tests.sh +++ b/run_ci_tests.sh @@ -7,6 +7,7 @@ END_STATUS=0 # ADD TEST SCRIPTS HERE ### if ! python -m pytest -v --durations=0 -x "test_batch_queue.py" ; then END_STATUS=1; fi +if ! python -m pytest -v --durations=0 -x "test_shuffle.py" ; then END_STATUS=1; fi ### # END ###