ray-project--ray/python/ray/data/tests/test_push_based_shuffle.py

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__]))