chore: import upstream snapshot with attribution

python/ray/data/tests/test_push_based_shuffle.py

@@ -0,0 +1,318 @@
+from collections import defaultdict
+
+import pytest
+
+import ray
+from ray._raylet import NodeID
+from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import (
+    PushBasedShuffleTaskScheduler,
+)
+
+
+def test_push_based_shuffle_schedule():
+    def _test(num_input_blocks, merge_factor, num_cpus_per_node_map):
+        num_cpus = sum(v for v in num_cpus_per_node_map.values())
+        op_cls = PushBasedShuffleTaskScheduler
+        schedule = op_cls._compute_shuffle_schedule(
+            num_cpus_per_node_map, num_input_blocks, merge_factor, num_input_blocks
+        )
+        # All input blocks will be processed.
+        assert (
+            schedule.num_rounds * schedule.num_map_tasks_per_round >= num_input_blocks
+        )
+        # Each round of tasks does not over-subscribe CPUs.
+        assert (
+            schedule.num_map_tasks_per_round
+            + schedule.merge_schedule.num_merge_tasks_per_round
+            <= max(num_cpus, 2)
+        )
+        print(
+            "map",
+            schedule.num_map_tasks_per_round,
+            "merge",
+            schedule.merge_schedule.num_merge_tasks_per_round,
+            "num_cpus",
+            num_cpus,
+            "merge_factor",
+            merge_factor,
+        )
+        # Merge factor between map : merge tasks is approximately correct.
+        if schedule.num_map_tasks_per_round > merge_factor:
+            actual_merge_factor = (
+                schedule.num_map_tasks_per_round
+                / schedule.merge_schedule.num_merge_tasks_per_round
+            )
+            next_highest_merge_factor = schedule.num_map_tasks_per_round / (
+                schedule.merge_schedule.num_merge_tasks_per_round + 1
+            )
+            assert actual_merge_factor - 1 <= merge_factor <= actual_merge_factor + 1, (
+                next_highest_merge_factor,
+                merge_factor,
+                actual_merge_factor,
+            )
+        else:
+            assert schedule.merge_schedule.num_merge_tasks_per_round == 1, (
+                schedule.num_map_tasks_per_round,
+                merge_factor,
+            )
+
+        # Tasks are evenly distributed.
+        tasks_per_node = defaultdict(int)
+        for i in range(schedule.merge_schedule.num_merge_tasks_per_round):
+            task_options = schedule.get_merge_task_options(i)
+            node_id = task_options["scheduling_strategy"].node_id
+            tasks_per_node[node_id] += 1
+        low = min(tasks_per_node.values())
+        high = low + 1
+        assert low <= max(tasks_per_node.values()) <= high
+
+        # Reducers are evenly distributed across mergers.
+        num_reducers_per_merge_idx = [
+            schedule.merge_schedule.get_num_reducers_per_merge_idx(i)
+            for i in range(schedule.merge_schedule.num_merge_tasks_per_round)
+        ]
+        high = max(num_reducers_per_merge_idx)
+        for num_reducers in num_reducers_per_merge_idx:
+            assert num_reducers == high or num_reducers == high - 1
+
+        for merge_idx in range(schedule.merge_schedule.num_merge_tasks_per_round):
+            assert isinstance(
+                schedule.merge_schedule.get_num_reducers_per_merge_idx(merge_idx), int
+            )
+            assert schedule.merge_schedule.get_num_reducers_per_merge_idx(merge_idx) > 0
+
+        reduce_idxs = list(range(schedule.merge_schedule.output_num_blocks))
+        actual_num_reducers_per_merge_idx = [
+            0 for _ in range(schedule.merge_schedule.num_merge_tasks_per_round)
+        ]
+        for reduce_idx in schedule.merge_schedule.round_robin_reduce_idx_iterator():
+            reduce_idxs.pop(reduce_idxs.index(reduce_idx))
+            actual_num_reducers_per_merge_idx[
+                schedule.merge_schedule.get_merge_idx_for_reducer_idx(reduce_idx)
+            ] += 1
+        # Check that each reduce task is submitted exactly once.
+        assert len(reduce_idxs) == 0
+        # Check that each merge and reduce task are correctly paired.
+        for i, num_reducers in enumerate(actual_num_reducers_per_merge_idx):
+            assert (
+                num_reducers == num_reducers_per_merge_idx[i]
+            ), f"""Merge task [{i}] has {num_reducers} downstream reduce tasks,
+            expected {num_reducers_per_merge_idx[i]}."""
+            assert num_reducers > 0
+
+    node_id_1 = NodeID.from_random().hex()
+    node_id_2 = NodeID.from_random().hex()
+    node_id_3 = NodeID.from_random().hex()
+    for num_cpus in range(1, 20):
+        _test(20, 3, {node_id_1: num_cpus})
+    _test(20, 3, {node_id_1: 100})
+    _test(100, 3, {node_id_1: 10, node_id_2: 10, node_id_3: 10})
+    _test(100, 10, {node_id_1: 10, node_id_2: 10, node_id_3: 10})
+    # Regression test for https://github.com/ray-project/ray/issues/25863.
+    _test(1000, 2, {NodeID.from_random().hex(): 16 for i in range(20)})
+    # Regression test for https://github.com/ray-project/ray/issues/37754.
+    _test(260, 2, {node_id_1: 128})
+    _test(1, 2, {node_id_1: 128})
+
+    # Test float merge_factor.
+    for cluster_config in [
+        {node_id_1: 10},
+        {node_id_1: 10, node_id_2: 10},
+    ]:
+        _test(100, 1, cluster_config)
+        _test(100, 1.3, cluster_config)
+        _test(100, 1.6, cluster_config)
+        _test(100, 1.75, cluster_config)
+        _test(100, 2, cluster_config)
+
+        _test(1, 1.2, cluster_config)
+        _test(2, 1.2, cluster_config)
+
+
+def test_push_based_shuffle_stats(ray_start_cluster):
+    ctx = ray.data.context.DataContext.get_current()
+    try:
+        original = ctx.use_push_based_shuffle
+        ctx.use_push_based_shuffle = True
+
+        cluster = ray_start_cluster
+        cluster.add_node(
+            resources={"bar:1": 100},
+            num_cpus=10,
+            _system_config={"max_direct_call_object_size": 0},
+        )
+        cluster.add_node(resources={"bar:2": 100}, num_cpus=10)
+        cluster.add_node(resources={"bar:3": 100}, num_cpus=0)
+
+        ray.init(cluster.address)
+
+        parallelism = 100
+        ds = ray.data.range(1000, override_num_blocks=parallelism).random_shuffle()
+        ds = ds.materialize()
+        assert "RandomShuffleMerge" in ds.stats()
+        # Check all nodes used.
+        assert "2 nodes used" in ds.stats()
+        assert "1 nodes used" not in ds.stats()
+
+        # Check all merge tasks are included in stats.
+        internal_stats = ds._raw_stats()
+        num_merge_tasks = len(internal_stats.metadata["RandomShuffleMerge"])
+        # Merge factor is 2 for random_shuffle ops.
+        merge_factor = 2
+        assert (
+            parallelism // (merge_factor + 1)
+            <= num_merge_tasks
+            <= parallelism // merge_factor
+        )
+
+    finally:
+        ctx.use_push_based_shuffle = original
+
+
+def test_sort_multinode(ray_start_cluster, configure_shuffle_method):
+    cluster = ray_start_cluster
+    cluster.add_node(
+        resources={"bar:1": 100},
+        num_cpus=10,
+        _system_config={"max_direct_call_object_size": 0},
+    )
+    cluster.add_node(resources={"bar:2": 100}, num_cpus=10)
+    cluster.add_node(resources={"bar:3": 100}, num_cpus=0)
+
+    ray.init(cluster.address)
+
+    parallelism = 100
+    ds = (
+        ray.data.range(1000, override_num_blocks=parallelism)
+        .random_shuffle()
+        .sort("id")
+    )
+    for i, row in enumerate(ds.iter_rows()):
+        assert row["id"] == i
+
+
+def patch_ray_remote(condition, callback):
+    original_ray_remote = ray.remote
+
+    def ray_remote_override(*args, **kwargs):
+        def wrapper(fn):
+            remote_fn = original_ray_remote(*args, **kwargs)(fn)
+            if condition(fn):
+                original_remote_options = remote_fn.options
+
+                def options(**task_options):
+                    callback(task_options)
+                    original_options = original_remote_options(**task_options)
+                    return original_options
+
+                remote_fn.options = options
+            return remote_fn
+
+        return wrapper
+
+    ray.remote = ray_remote_override
+    return original_ray_remote
+
+
+def patch_ray_get(callback):
+    original_ray_get = ray.get
+
+    def ray_get_override(object_refs, *args, **kwargs):
+        callback(object_refs)
+        return original_ray_get(object_refs, *args, **kwargs)
+
+    ray.get = ray_get_override
+    return original_ray_get
+
+
+@pytest.mark.parametrize("pipeline", [False, True])
+def test_push_based_shuffle_reduce_stage_scheduling(ray_start_cluster, pipeline):
+    ctx = ray.data.context.DataContext.get_current()
+    try:
+        original = ctx.use_push_based_shuffle
+        ctx.use_push_based_shuffle = True
+        ctx.pipeline_push_based_shuffle_reduce_tasks = pipeline
+
+        num_cpus_per_node = 8
+        num_nodes = 3
+        num_output_blocks = 100
+
+        task_context = {
+            "reduce_options_submitted": [],
+            # The total number of CPUs available.
+            "pipelined_parallelism": num_cpus_per_node * num_nodes,
+            # The total number of reduce tasks.
+            "total_parallelism": num_output_blocks,
+            "num_instances_below_parallelism": 0,
+        }
+
+        def reduce_options_patch(task_options):
+            task_context["reduce_options_submitted"].append(task_options)
+
+        def check_pipelined(refs):
+            if task_context["reduce_options_submitted"]:
+                # Check that we have the correct number of tasks in flight.
+                if pipeline:
+                    # When pipelining, we should limit the number of reduce
+                    # tasks in flight based on how many CPUs are in the
+                    # cluster.
+                    if not (
+                        task_context["pipelined_parallelism"]
+                        <= len(task_context["reduce_options_submitted"])
+                        <= 2 * task_context["pipelined_parallelism"]
+                    ):
+                        task_context["num_instances_below_parallelism"] += 1
+                else:
+                    # When not pipelining, we should submit all reduce tasks at
+                    # once.
+                    assert (
+                        len(task_context["reduce_options_submitted"])
+                        == task_context["total_parallelism"]
+                    )
+
+                # Check that tasks are close to evenly spread across the nodes.
+                nodes = defaultdict(int)
+                for options in task_context["reduce_options_submitted"]:
+                    nodes[options["scheduling_strategy"].node_id] += 1
+                assert len(nodes) > 1
+                assert min(nodes.values()) >= max(nodes.values()) // 2
+
+            task_context["reduce_options_submitted"].clear()
+
+        ray_remote = patch_ray_remote(
+            lambda fn: "reduce" in fn.__name__, reduce_options_patch
+        )
+        ray_get = patch_ray_get(check_pipelined)
+
+        cluster = ray_start_cluster
+        for _ in range(num_nodes):
+            cluster.add_node(
+                num_cpus=num_cpus_per_node,
+            )
+
+        ray.init(cluster.address)
+
+        ds = ray.data.range(
+            1000, override_num_blocks=num_output_blocks
+        ).random_shuffle()
+        # Only the last round should have fewer tasks in flight.
+        assert task_context["num_instances_below_parallelism"] <= 1
+        task_context["num_instances_below_parallelism"] = 0
+
+        ds = ds.sort("id")
+        # Only the last round should have fewer tasks in flight.
+        assert task_context["num_instances_below_parallelism"] <= 1
+        task_context["num_instances_below_parallelism"] = 0
+        for i, row in enumerate(ds.iter_rows()):
+            assert row["id"] == i
+
+    finally:
+        ctx.use_push_based_shuffle = original
+        ray.remote = ray_remote
+        ray.get = ray_get
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-sv", __file__]))