code: 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
  
      for num_cpus in range(1, 20):
          _test(20, 3, {"node1": num_cpus})
      _test(20, 3, {"node1": 100})
      _test(100, 3, {"node1": 10, "node2": 10, "node3": 10})
      _test(100, 10, {"node1": 10, "node2": 10, "node3": 10})
      # Regression test for https://github.com/ray-project/ray/issues/25863.
      _test(1000, 2, {f"node{i}": 16 for i in range(20)})
      # Regression test for https://github.com/ray-project/ray/issues/37754.
      _test(260, 2, {"node1": 128})
      _test(1, 2, {"node1": 128})
  
      # Test float merge_factor.
      for cluster_config in [
          {"node1": 10},
          {"node1": 10, "node2": 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)
  
imports: import random
  from collections import defaultdict
  import numpy
  import pandas
  import pyarrow
  import pytest
  import ray
  from ray.data import Dataset
  from ray.data._internal.planner.exchange.push_based_shuffle_task_scheduler import PushBasedShuffleTaskScheduler
  from ray.data._internal.sort import SortKey
  from ray.data.block import BlockAccessor
  from ray.data.context import DataContext
  from ray.data.tests.conftest import *
  from ray.data.tests.util import extract_values
  from ray.tests.conftest import *
label_desc: Use image node_id 10 for calling a global function with example usage: test_push_based_shuffle_schedule() without return types