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.
      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",".","See","paper","at","https",":","\/\/arxiv.org\/abs\/2203.05072","for","more","details","."]
code: def _compute_shuffle_schedule(
      num_cpus_per_node_map: Dict[str, int],
      num_input_blocks: int,
      merge_factor: int,
      num_output_blocks: int,
  ) -> _PushBasedShuffleStage:
      num_cpus_total = sum(v for v in num_cpus_per_node_map.values())
      task_parallelism = min(num_cpus_total, num_input_blocks)
  
      num_tasks_per_map_merge_group = merge_factor + 1
      num_merge_tasks_per_round = 0
      merge_task_placement = []
      leftover_cpus = 0
      # Compute the total number of merge tasks and their node placement.
      # Each merge task should be grouped with `merge_factor` map tasks for
      # pipelining. These groups should then be spread across nodes according
      # to CPU availability for load-balancing.
      for node, num_cpus in num_cpus_per_node_map.items():
          node_parallelism = min(
              num_cpus, num_input_blocks // len(num_cpus_per_node_map)
          )
          num_merge_tasks = (
              node_parallelism // num_tasks_per_map_merge_group
          )
          for i in range(num_merge_tasks):
              merge_task_placement.append(node)
          num_merge_tasks_per_round += num_merge_tasks
  
          # Handle the case where a single node cannot fit a group of map and
          # merge tasks, but we can spread the group across multiple distinct
          # nodes.
          leftover_cpus += (
              node_parallelism % num_tasks_per_map_merge_group
          )
          if (
              num_merge_tasks == 0
              and leftover_cpus > num_tasks_per_map_merge_group
          ):
              merge_task_placement.append(node)
              num_merge_tasks_per_round += 1
              leftover_cpus -= num_tasks_per_map_merge_group
      if num_merge_tasks_per_round == 0:
          merge_task_placement.append(list(num_cpus_per_node_map)[0])
          num_merge_tasks_per_round = 1
  
      assert num_merge_tasks_per_round == len(merge_task_placement)
      num_map_tasks_per_round = max(
          task_parallelism - num_merge_tasks_per_round, 1
      )
  
      num_rounds = math.ceil(num_input_blocks / num_map_tasks_per_round)
      return _PushBasedShuffleStage(
          num_output_blocks,
          num_rounds,
          num_map_tasks_per_round,
          merge_task_placement,
      )
  
imports: import logging
  import math
  from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
  import ray
  from ray.data._internal.block_list import BlockList
  from ray.data._internal.execution.interfaces import TaskContext
  from ray.data._internal.progress_bar import ProgressBar
  from ray.data._internal.remote_fn import cached_remote_fn
  from ray.data._internal.shuffle import ShuffleOp
  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 4 for calling the PushBasedShufflePlan obj's underlying member method code with example usage: obj._compute_shuffle_schedule(num_cpus_per_node_map, num_input_blocks, merge_factor, num_output_blocks) and returns: _PushBasedShuffleStage