class_docstr: Push-based shuffle merges intermediate map outputs on the reducer nodes
  while other map tasks are executing. The merged outputs are merged again
  during a final reduce stage. This works as follows:
  
  1. Submit rounds of concurrent map and merge tasks until all map inputs
  have been processed. In each round, we execute:
  
     M map tasks
       Each produces N outputs. Each output contains P blocks.
     N merge tasks
       Takes 1 output from each of M map tasks.
       Each produces P outputs.
     Where M and N are chosen to maximize parallelism across CPUs. Note that
     this assumes that all CPUs in the cluster will be dedicated to the
     shuffle job.
  
     Map and merge tasks are pipelined so that we always merge the previous
     round of map outputs while executing the next round of map tasks.
  
  2. In the final reduce stage:
     R reduce tasks
       Takes 1 output from one of the merge tasks from every round.
  
  Notes:
      N * P = R = total number of output blocks
      M / N = merge factor - the ratio of map : merge tasks is to improve
        pipelined parallelism. For example, if map takes twice as long to
        execute as merge, then we should set this to 2. If pipeline bubbles
        appear and the merge tasks are much longer than the map tasks, then
        the merge factor should be decreased, and vice versa.
      See paper at https://arxiv.org/abs/2203.05072 for more details.
class_docstr_tok: ["Push-based","shuffle","merges","intermediate","map","outputs","on","the","reducer","nodes","while","other","map","tasks","are","executing",".","The","merged","outputs","are","merged","again","during","a","final","reduce","stage",".","This","works","as","follows",":","1",".","Submit","rounds","of","concurrent","map","and","merge","tasks","until","all","map","inputs","have","been","processed",".","In","each","round",",","we","execute",":","M","map","tasks","Each","produces","N","outputs",".","Each","output","contains","P","blocks",".","N","merge","tasks","Takes","1","output","from","each","of","M","map","tasks",".","Each","produces","P","outputs",".","Where","M","and","N","are","chosen","to","maximize","parallelism","across","CPUs",".","Note","that","this","assumes","that","all","CPUs","in","the","cluster","will","be","dedicated","to","the","shuffle","job",".","Map","and","merge","tasks","are","pipelined","so","that","we","always","merge","the","previous","round","of","map","outputs","while","executing","the","next","round","of","map","tasks",".","2",".","In","the","final","reduce","stage",":","R","reduce","tasks","Takes","1","output","from","one","of","the","merge","tasks","from","every","round",".","Notes",":","N","*","P","=","R","=","total","number","of","output","blocks","M","\/","N","=","merge","factor","-","the","ratio","of","map",":","merge","tasks","is","to","improve","pipelined","parallelism",".","For","example",",","if","map","takes","twice","as","long","to","execute","as","merge",",","then","we","should","set","this","to","2",".","If","pipeline","bubbles","appear","and","the","merge","tasks","are","much","longer","than","the","map","tasks",",","then","the","merge","factor","should","be","decreased",",","and","vice","versa",".","See","paper","at","https",":","\/\/arxiv.org\/abs\/2203.05072","for","more","details","."]
code: def execute(
      self,
      refs: List[RefBundle],
      output_num_blocks: int,
      ctx: TaskContext,
      map_ray_remote_args: Optional[Dict[str, Any]] = None,
      reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
      merge_factor: float = 2,
      _debug_limit_execution_to_num_blocks: int = None,
  ) -> Tuple[List[RefBundle], StatsDict]:
      logger.get_logger().info("Using experimental push-based shuffle.")
      # TODO: Preemptively clear the blocks list since we will incrementally delete
      # the last remaining references as we submit the dependent map tasks during the
      # map-merge stage.
  
      # TODO(swang): For jobs whose reduce work is heavier than the map work,
      # we should support fractional merge factors.
      # TODO(swang): For large jobs, we should try to choose the merge factor
      # automatically, e.g., by running one test round of map and merge tasks
      # and comparing their run times.
      # TODO(swang): Add option to automatically reduce write amplification
      # during map-merge stage, by limiting how many partitions can be
      # processed concurrently.
      input_blocks_list = []
      for ref_bundle in refs:
          for block, _ in ref_bundle.blocks:
              input_blocks_list.append(block)
      input_owned = all(b.owns_blocks for b in refs)
  
      if map_ray_remote_args is None:
          map_ray_remote_args = {}
      if reduce_ray_remote_args is None:
          reduce_ray_remote_args = {}
      # The placement strategy for reduce tasks is overwritten to colocate
      # them with their inputs from the merge stage, so remove any
      # pre-specified scheduling strategy here.
      reduce_ray_remote_args = reduce_ray_remote_args.copy()
      reduce_ray_remote_args.pop("scheduling_strategy", None)
  
      # Compute all constants used for task scheduling.
      num_cpus_per_node_map = _get_num_cpus_per_node_map()
      stage = self._compute_shuffle_schedule(
          num_cpus_per_node_map,
          len(input_blocks_list),
          merge_factor,
          output_num_blocks,
      )
  
      # TODO(swang): Use INFO level. Currently there is no easy way to set
      # the logging level to DEBUG from a driver script, so just print
      # verbosely for now.
      # See https://github.com/ray-project/ray/issues/42002.
      logger.get_logger().info(f"Push-based shuffle schedule:\n{stage}")
  
      map_fn = self._map_partition
      merge_fn = self._merge
  
      def map_partition(*args, **kwargs):
          return map_fn(self._exchange_spec.map, *args, **kwargs)
  
      def merge(*args, **kwargs):
          return merge_fn(self._exchange_spec.reduce, *args, **kwargs)
  
      shuffle_map = cached_remote_fn(map_partition)
      shuffle_map = shuffle_map.options(
          **map_ray_remote_args,
          num_returns=1
          + stage.merge_schedule.num_merge_tasks_per_round,
      )
  
      if _debug_limit_execution_to_num_blocks is not None:
          input_blocks_list = input_blocks_list[
              :_debug_limit_execution_to_num_blocks
          ]
          logger.get_logger().info(
              f"Limiting execution to {len(input_blocks_list)} map tasks"
          )
      map_stage_iter = _MapStageIterator(
          input_blocks_list,
          shuffle_map,
          [
              output_num_blocks,
              stage.merge_schedule,
              *self._exchange_spec._map_args,
          ],
      )
  
      sub_progress_bar_dict = ctx.sub_progress_bar_dict
      bar_name = ExchangeTaskSpec.MAP_SUB_PROGRESS_BAR_NAME
      assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
      map_bar = sub_progress_bar_dict[bar_name]
      map_stage_executor = _PipelinedStageExecutor(
          map_stage_iter,
          stage.num_map_tasks_per_round,
          progress_bar=map_bar,
      )
  
      shuffle_merge = cached_remote_fn(merge)
      merge_stage_iter = _MergeStageIterator(
          map_stage_iter,
          shuffle_merge,
          stage,
          self._exchange_spec._reduce_args,
      )
      merge_stage_executor = _PipelinedStageExecutor(
          merge_stage_iter,
          stage.merge_schedule.num_merge_tasks_per_round,
          max_concurrent_rounds=2,
      )
      # Execute the map-merge stage. This submits tasks in rounds of M map
      # tasks and N merge tasks each. Task execution between map and merge is
      # pipelined, so that while executing merge for one round of inputs, we
      # also execute the map tasks for the following round.
      map_done = False
      merge_done = False
      map_stage_metadata = []
      merge_stage_metadata = []
      while not (map_done and merge_done):
          try:
              map_stage_metadata += next(map_stage_executor)
          except StopIteration:
              map_done = True
              break
  
          try:
              merge_stage_metadata += next(merge_stage_executor)
          except StopIteration:
              merge_done = True
              break
      all_merge_results = merge_stage_iter.pop_merge_results()
  
      if _debug_limit_execution_to_num_blocks is not None:
          for merge_idx in range(len(all_merge_results)):
              while (
                  len(all_merge_results[merge_idx]) < stage.num_rounds
              ):
                  # Repeat the first merge task's results.
                  all_merge_results[merge_idx].append(
                      all_merge_results[merge_idx][0][:]
                  )
  
      # Execute and wait for the reduce stage.
      bar_name = ExchangeTaskSpec.REDUCE_SUB_PROGRESS_BAR_NAME
      assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
      reduce_bar = sub_progress_bar_dict[bar_name]
  
      shuffle_reduce = cached_remote_fn(self._exchange_spec.reduce)
      reduce_stage_iter = _ReduceStageIterator(
          stage,
          shuffle_reduce,
          all_merge_results,
          reduce_ray_remote_args,
          self._exchange_spec._reduce_args,
          _debug_limit_execution_to_num_blocks,
      )
  
      max_reduce_tasks_in_flight = output_num_blocks
      ctx = DataContext.get_current()
      if ctx.pipeline_push_based_shuffle_reduce_tasks:
          # If pipelining is enabled, we should still try to utilize all
          # cores.
          max_reduce_tasks_in_flight = min(
              max_reduce_tasks_in_flight,
              sum(num_cpus_per_node_map.values()),
          )
  
      reduce_stage_executor = _PipelinedStageExecutor(
          reduce_stage_iter,
          max_reduce_tasks_in_flight,
          max_concurrent_rounds=2,
          progress_bar=reduce_bar,
      )
      reduce_stage_metadata = []
      while True:
          try:
              reduce_stage_metadata += next(reduce_stage_executor)
          except StopIteration:
              break
  
      new_blocks = reduce_stage_iter.pop_reduce_results()
      sorted_blocks = [
          (block[0], block[1], reduce_stage_metadata[i])
          for i, block in enumerate(new_blocks)
      ]
      sorted_blocks.sort(key=lambda x: x[0])
  
      new_blocks, reduce_stage_metadata = [], []
      if sorted_blocks:
          _, new_blocks, reduce_stage_metadata = zip(*sorted_blocks)
      del sorted_blocks
  
      if _debug_limit_execution_to_num_blocks is not None:
          output_num_blocks = min(
              _debug_limit_execution_to_num_blocks, output_num_blocks
          )
  
      assert (
          len(new_blocks) == output_num_blocks
      ), f"Expected {output_num_blocks} outputs, produced {len(new_blocks)}"
  
      output = []
      for block, meta in zip(new_blocks, reduce_stage_metadata):
          output.append(
              RefBundle(
                  [
                      (
                          block,
                          meta,
                      )
                  ],
                  owns_blocks=input_owned,
              )
          )
      stats = {
          "map": map_stage_metadata,
          "merge": merge_stage_metadata,
          "reduce": reduce_stage_metadata,
      }
  
      return (output, stats)
  
imports: import math
  from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
  import ray
  from ray.data._internal.dataset_logger import DatasetLogger
  from ray.data._internal.execution.interfaces import RefBundle, TaskContext
  from ray.data._internal.planner.exchange.interfaces import ExchangeTaskScheduler, ExchangeTaskSpec
  from ray.data._internal.progress_bar import ProgressBar
  from ray.data._internal.remote_fn import cached_remote_fn
  from ray.data._internal.stats import StatsDict
  from ray.data.block import Block, BlockAccessor, BlockExecStats, BlockMetadata
  from ray.data.context import DataContext
  from ray.types import ObjectRef
  from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
label_desc: Use image node_id 1 for calling the PushBasedShuffleTaskScheduler obj's underlying member method code with example usage: obj.execute(refs, output_num_blocks, ctx, map_ray_remote_args, reduce_ray_remote_args, merge_factor, _debug_limit_execution_to_num_blocks) and returns: output,  stats, map_fn, merge_fn