chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,123 @@
|
||||
from typing import List, Tuple, Union
|
||||
|
||||
from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec
|
||||
from ray.data._internal.planner.exchange.sort_task_spec import SortKey
|
||||
from ray.data._internal.table_block import TableBlockAccessor
|
||||
from ray.data.aggregate import AggregateFn, AggregateFnV2, Count
|
||||
from ray.data.block import (
|
||||
Block,
|
||||
BlockAccessor,
|
||||
BlockExecStats,
|
||||
BlockMetadataWithSchema,
|
||||
KeyType,
|
||||
)
|
||||
|
||||
|
||||
class SortAggregateTaskSpec(ExchangeTaskSpec):
|
||||
"""
|
||||
The implementation for sort-based aggregate tasks.
|
||||
|
||||
Aggregate is done in 2 steps: partial aggregate of individual blocks, and
|
||||
final aggregate of sorted blocks.
|
||||
|
||||
Partial aggregate (`map`): each block is sorted locally, then partitioned into
|
||||
smaller blocks according to the boundaries. Each partitioned block is aggregated
|
||||
separately, then passed to a final aggregate task.
|
||||
|
||||
Final aggregate (`reduce`): each task would receive a block from every worker that
|
||||
consists of items in a certain range. It then merges the sorted blocks and
|
||||
aggregates on-the-fly.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
boundaries: List[KeyType],
|
||||
key: SortKey,
|
||||
aggs: List[AggregateFn],
|
||||
):
|
||||
super().__init__(
|
||||
map_args=[boundaries, key, aggs],
|
||||
reduce_args=[key, aggs],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def map(
|
||||
idx: int,
|
||||
block: Block,
|
||||
output_num_blocks: int,
|
||||
boundaries: List[KeyType],
|
||||
sort_key: SortKey,
|
||||
aggs: List[AggregateFn],
|
||||
) -> List[Union[Block, "BlockMetadataWithSchema"]]:
|
||||
stats = BlockExecStats.builder()
|
||||
|
||||
block = SortAggregateTaskSpec._prune_unused_columns(block, sort_key, aggs)
|
||||
if sort_key.get_columns():
|
||||
partitions = BlockAccessor.for_block(block).sort_and_partition(
|
||||
boundaries,
|
||||
sort_key,
|
||||
)
|
||||
else:
|
||||
partitions = [block]
|
||||
parts = [
|
||||
BlockAccessor.for_block(p)._aggregate(sort_key, aggs) for p in partitions
|
||||
]
|
||||
from ray.data.block import BlockMetadataWithSchema
|
||||
|
||||
meta_with_schema = BlockMetadataWithSchema.from_block(
|
||||
block, block_exec_stats=stats.build()
|
||||
)
|
||||
return parts + [meta_with_schema]
|
||||
|
||||
@staticmethod
|
||||
def reduce(
|
||||
key: SortKey,
|
||||
aggs: List[AggregateFn],
|
||||
*mapper_outputs: List[Block],
|
||||
partial_reduce: bool = False,
|
||||
) -> Tuple[Block, "BlockMetadataWithSchema"]:
|
||||
normalized_blocks = TableBlockAccessor.normalize_block_types(
|
||||
mapper_outputs,
|
||||
target_block_type=None,
|
||||
)
|
||||
blocks, meta_with_schema = BlockAccessor.for_block(
|
||||
normalized_blocks[0]
|
||||
)._combine_aggregated_blocks(
|
||||
list(normalized_blocks), key, aggs, finalize=not partial_reduce
|
||||
)
|
||||
return blocks, meta_with_schema
|
||||
|
||||
@staticmethod
|
||||
def _prune_unused_columns(
|
||||
block: Block,
|
||||
sort_key: SortKey,
|
||||
aggs: Tuple[AggregateFn],
|
||||
) -> Block:
|
||||
"""Prune unused columns from block before aggregate."""
|
||||
prune_columns = True
|
||||
columns = set()
|
||||
key = sort_key.get_columns()
|
||||
|
||||
if isinstance(key, str):
|
||||
columns.add(key)
|
||||
elif isinstance(key, list):
|
||||
columns.update(key)
|
||||
elif callable(key):
|
||||
prune_columns = False
|
||||
|
||||
for agg in aggs:
|
||||
if isinstance(agg, AggregateFnV2) and agg.get_target_column():
|
||||
columns.add(agg.get_target_column())
|
||||
elif not isinstance(agg, Count):
|
||||
# Don't prune columns if any aggregate key is not string.
|
||||
prune_columns = False
|
||||
|
||||
block_accessor = BlockAccessor.for_block(block)
|
||||
if (
|
||||
prune_columns
|
||||
and isinstance(block_accessor, TableBlockAccessor)
|
||||
and block_accessor.num_rows() > 0
|
||||
):
|
||||
return block_accessor.select(list(columns))
|
||||
else:
|
||||
return block
|
||||
@@ -0,0 +1,135 @@
|
||||
import logging
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import ray._private.worker
|
||||
from ray.data._internal.execution.interfaces import RefBundle
|
||||
from ray.data._internal.stats import StatsDict
|
||||
from ray.data._internal.util import convert_bytes_to_human_readable_str
|
||||
from ray.data.block import Block
|
||||
from ray.data.context import DataContext
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
from ray.data.block import BlockMetadataWithSchema
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ExchangeTaskSpec:
|
||||
"""
|
||||
An interface to specify the exchange map and reduce tasks.
|
||||
|
||||
Subclasses should implement the `map` and `reduce` static methods.
|
||||
`map` method is to transform one input block into multiple output blocks.
|
||||
`reduce` is to combine multiple map output blocks. Both methods are
|
||||
single-task operations. See `ExchangeScheduler` for how to distribute
|
||||
the operations across multiple tasks.
|
||||
|
||||
Any custom arguments for `map` and `reduce` methods should be specified by
|
||||
setting `map_args` and `reduce_args`.
|
||||
|
||||
The concept here is similar to the exchange operator described in
|
||||
"Volcano - An Extensible and Parallel Query Evaluation System"
|
||||
(https://dl.acm.org/doi/10.1109/69.273032).
|
||||
"""
|
||||
|
||||
MAP_SUB_PROGRESS_BAR_NAME = "Shuffle Map"
|
||||
REDUCE_SUB_PROGRESS_BAR_NAME = "Shuffle Reduce"
|
||||
|
||||
def __init__(self, map_args: List[Any] = None, reduce_args: List[Any] = None):
|
||||
self._map_args = map_args or []
|
||||
self._reduce_args = reduce_args or []
|
||||
assert isinstance(self._map_args, list)
|
||||
assert isinstance(self._reduce_args, list)
|
||||
|
||||
@staticmethod
|
||||
def map(
|
||||
idx: int,
|
||||
block: Block,
|
||||
output_num_blocks: int,
|
||||
) -> List[Union[Block, "BlockMetadataWithSchema"]]:
|
||||
"""
|
||||
Map function to be run on each input block.
|
||||
|
||||
Returns list of [BlockMetadata, Block1, Block2, ..., BlockN].
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
@staticmethod
|
||||
def reduce(
|
||||
*mapper_outputs: List[Block],
|
||||
partial_reduce: bool = False,
|
||||
) -> Tuple[Block, "BlockMetadataWithSchema"]:
|
||||
"""
|
||||
Reduce function to be run for each output block.
|
||||
|
||||
Args:
|
||||
*mapper_outputs: List of map output blocks to reduce.
|
||||
partial_reduce: Whether should partially or fully reduce.
|
||||
|
||||
Returns:
|
||||
The reduced block and its metadata.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class ExchangeTaskScheduler:
|
||||
"""
|
||||
An interface to schedule exchange tasks (`exchange_spec`) for multi-nodes
|
||||
execution.
|
||||
"""
|
||||
|
||||
def __init__(self, exchange_spec: ExchangeTaskSpec):
|
||||
"""Initialize the scheduler.
|
||||
|
||||
Args:
|
||||
exchange_spec: The implementation of exchange tasks to execute.
|
||||
"""
|
||||
self._exchange_spec = exchange_spec
|
||||
# If driver memory exceeds this threshold, warn the user. For now, this
|
||||
# only applies to shuffle ops because most other ops are unlikely to use as
|
||||
# much driver memory.
|
||||
self.warn_on_driver_memory_usage_bytes: Optional[
|
||||
int
|
||||
] = DataContext.get_current().warn_on_driver_memory_usage_bytes
|
||||
|
||||
def execute(
|
||||
self,
|
||||
refs: List[RefBundle],
|
||||
output_num_blocks: int,
|
||||
map_ray_remote_args: Optional[Dict[str, Any]] = None,
|
||||
reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
|
||||
warn_on_driver_memory_usage: Optional[int] = None,
|
||||
) -> Tuple[List[RefBundle], StatsDict]:
|
||||
"""
|
||||
Execute the exchange tasks on input `refs`.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def warn_on_high_local_memory_store_usage(self) -> None:
|
||||
ray_core_worker = ray._private.worker.global_worker.core_worker
|
||||
local_memory_store_bytes_used = (
|
||||
ray_core_worker.get_local_memory_store_bytes_used()
|
||||
)
|
||||
self.warn_on_driver_memory_usage(
|
||||
local_memory_store_bytes_used,
|
||||
"More than "
|
||||
f"{convert_bytes_to_human_readable_str(local_memory_store_bytes_used)} "
|
||||
"of driver memory used to store Ray Data block data and metadata. "
|
||||
"This job may exit if driver memory is insufficient.\n\n"
|
||||
"This can happen when many tiny blocks are created. "
|
||||
"Check the block size using Dataset.stats() and see "
|
||||
"https://docs.ray.io/en/latest/data/performance-tips.html"
|
||||
" for mitigation.",
|
||||
)
|
||||
|
||||
def warn_on_driver_memory_usage(
|
||||
self, memory_usage_bytes: int, log_str: str
|
||||
) -> None:
|
||||
if self.warn_on_driver_memory_usage_bytes is None:
|
||||
return
|
||||
|
||||
if memory_usage_bytes > self.warn_on_driver_memory_usage_bytes:
|
||||
logger.warning(log_str)
|
||||
# Double the threshold to avoid verbose warnings.
|
||||
self.warn_on_driver_memory_usage_bytes = memory_usage_bytes * 2
|
||||
@@ -0,0 +1,155 @@
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from ray._private.ray_constants import CALLER_MEMORY_USAGE_PER_OBJECT_REF
|
||||
from ray.data._internal.execution.interfaces import BlockEntry, RefBundle, TaskContext
|
||||
from ray.data._internal.planner.exchange.interfaces import (
|
||||
ExchangeTaskScheduler,
|
||||
ExchangeTaskSpec,
|
||||
)
|
||||
from ray.data._internal.remote_fn import cached_remote_fn
|
||||
from ray.data._internal.stats import StatsDict
|
||||
from ray.data._internal.util import (
|
||||
convert_bytes_to_human_readable_str,
|
||||
unzip,
|
||||
)
|
||||
from ray.data.block import BlockMetadataWithSchema, to_stats
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PullBasedShuffleTaskScheduler(ExchangeTaskScheduler):
|
||||
"""
|
||||
The pull-based map-reduce shuffle scheduler.
|
||||
|
||||
Map tasks are first scheduled to generate map output blocks. After all map output
|
||||
are generated, then reduce tasks are scheduled to combine map output blocks
|
||||
together.
|
||||
|
||||
The concept here is similar to
|
||||
"MapReduce: Simplified Data Processing on Large Clusters"
|
||||
(https://dl.acm.org/doi/10.1145/1327452.1327492).
|
||||
"""
|
||||
|
||||
def execute(
|
||||
self,
|
||||
refs: List[RefBundle],
|
||||
output_num_blocks: int,
|
||||
task_ctx: TaskContext,
|
||||
map_ray_remote_args: Optional[Dict[str, Any]] = None,
|
||||
reduce_ray_remote_args: Optional[Dict[str, Any]] = None,
|
||||
_debug_limit_execution_to_num_blocks: Optional[int] = None,
|
||||
) -> Tuple[List[RefBundle], StatsDict]:
|
||||
|
||||
# TODO: eagerly delete the input and map output block references in order to
|
||||
# eagerly release the blocks' memory.
|
||||
input_blocks_list = []
|
||||
for ref_bundle in refs:
|
||||
input_blocks_list.extend(ref_bundle.block_refs)
|
||||
input_num_blocks = len(input_blocks_list)
|
||||
input_owned = all(b.owns_blocks for b in refs)
|
||||
|
||||
caller_memory_usage = (
|
||||
input_num_blocks * output_num_blocks * CALLER_MEMORY_USAGE_PER_OBJECT_REF
|
||||
)
|
||||
self.warn_on_driver_memory_usage(
|
||||
caller_memory_usage,
|
||||
"Execution is estimated to use at least "
|
||||
f"{convert_bytes_to_human_readable_str(caller_memory_usage)} "
|
||||
"of driver memory. Ensure that the driver machine has at least "
|
||||
"this much memory to ensure job completion.\n\n"
|
||||
"To reduce the "
|
||||
"amount of driver memory needed, enable push-based shuffle using "
|
||||
"RAY_DATA_PUSH_BASED_SHUFFLE=1 "
|
||||
"(https://docs.ray.io/en/latest/data/performance-tips.html"
|
||||
").",
|
||||
)
|
||||
|
||||
if map_ray_remote_args is None:
|
||||
map_ray_remote_args = {}
|
||||
if reduce_ray_remote_args is None:
|
||||
reduce_ray_remote_args = {}
|
||||
if "scheduling_strategy" not in reduce_ray_remote_args:
|
||||
reduce_ray_remote_args = reduce_ray_remote_args.copy()
|
||||
reduce_ray_remote_args["scheduling_strategy"] = "SPREAD"
|
||||
|
||||
shuffle_map = cached_remote_fn(self._exchange_spec.map)
|
||||
shuffle_reduce = cached_remote_fn(self._exchange_spec.reduce)
|
||||
|
||||
sub_progress_bar_dict = task_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]
|
||||
|
||||
if _debug_limit_execution_to_num_blocks is not None:
|
||||
input_blocks_list = input_blocks_list[:_debug_limit_execution_to_num_blocks]
|
||||
logger.debug(f"Limiting execution to {len(input_blocks_list)} map tasks")
|
||||
|
||||
shuffle_map_out = [
|
||||
shuffle_map.options(
|
||||
**map_ray_remote_args,
|
||||
num_returns=1 + output_num_blocks,
|
||||
).remote(i, block, output_num_blocks, *self._exchange_spec._map_args)
|
||||
for i, block in enumerate(input_blocks_list)
|
||||
]
|
||||
|
||||
# The first item returned is the BlockMetadata.
|
||||
shuffle_map_metadata_schema = []
|
||||
for i, refs in enumerate(shuffle_map_out):
|
||||
shuffle_map_metadata_schema.append(refs[-1])
|
||||
shuffle_map_out[i] = refs[:-1]
|
||||
|
||||
if _debug_limit_execution_to_num_blocks is not None:
|
||||
while len(shuffle_map_out) < output_num_blocks:
|
||||
# Repeat the first map task's results.
|
||||
shuffle_map_out.append(shuffle_map_out[0][:])
|
||||
|
||||
shuffle_map_metadata_schema = map_bar.fetch_until_complete(
|
||||
shuffle_map_metadata_schema
|
||||
)
|
||||
|
||||
self.warn_on_high_local_memory_store_usage()
|
||||
|
||||
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]
|
||||
|
||||
if _debug_limit_execution_to_num_blocks is not None:
|
||||
output_num_blocks = _debug_limit_execution_to_num_blocks
|
||||
logger.debug(f"Limiting execution to {output_num_blocks} reduce tasks")
|
||||
shuffle_reduce_out = [
|
||||
shuffle_reduce.options(**reduce_ray_remote_args, num_returns=2).remote(
|
||||
*self._exchange_spec._reduce_args,
|
||||
*[shuffle_map_out[i][j] for i in range(input_num_blocks)],
|
||||
)
|
||||
for j in range(output_num_blocks)
|
||||
]
|
||||
|
||||
# Release map task outputs from the Ray object store.
|
||||
del shuffle_map_out
|
||||
|
||||
new_blocks, new_metadata_schema = [], []
|
||||
if shuffle_reduce_out:
|
||||
new_blocks, new_metadata_schema = unzip(shuffle_reduce_out)
|
||||
new_metadata_schema: List[
|
||||
"BlockMetadataWithSchema"
|
||||
] = reduce_bar.fetch_until_complete(list(new_metadata_schema))
|
||||
|
||||
self.warn_on_high_local_memory_store_usage()
|
||||
|
||||
output = []
|
||||
for block, meta_with_schema in zip(new_blocks, new_metadata_schema):
|
||||
output.append(
|
||||
RefBundle(
|
||||
[BlockEntry(block, meta_with_schema.metadata)],
|
||||
owns_blocks=input_owned,
|
||||
schema=meta_with_schema.schema,
|
||||
)
|
||||
)
|
||||
|
||||
stats = {
|
||||
"map": to_stats(shuffle_map_metadata_schema),
|
||||
"reduce": to_stats(new_metadata_schema),
|
||||
}
|
||||
|
||||
return (output, stats)
|
||||
@@ -0,0 +1,847 @@
|
||||
import logging
|
||||
import math
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, TypeVar, Union
|
||||
|
||||
import ray
|
||||
from ray._private.ray_constants import CALLER_MEMORY_USAGE_PER_OBJECT_REF
|
||||
from ray.data._internal.execution.interfaces import BlockEntry, RefBundle, TaskContext
|
||||
from ray.data._internal.planner.exchange.interfaces import (
|
||||
ExchangeTaskScheduler,
|
||||
ExchangeTaskSpec,
|
||||
)
|
||||
from ray.data._internal.remote_fn import cached_remote_fn
|
||||
from ray.data._internal.stats import StatsDict
|
||||
from ray.data._internal.util import (
|
||||
convert_bytes_to_human_readable_str,
|
||||
unzip,
|
||||
)
|
||||
from ray.data.block import (
|
||||
Block,
|
||||
BlockAccessor,
|
||||
BlockExecStats,
|
||||
BlockMetadata,
|
||||
BlockMetadataWithSchema,
|
||||
_take_first_non_empty_schema,
|
||||
to_stats,
|
||||
)
|
||||
from ray.data.context import DataContext
|
||||
from ray.types import ObjectRef
|
||||
from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data._internal.progress.base_progress import BaseProgressBar
|
||||
from ray.data.block import BlockMetadataWithSchema
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
U = TypeVar("U")
|
||||
|
||||
|
||||
class _MergeTaskSchedule:
|
||||
def __init__(self, output_num_blocks: int, num_merge_tasks_per_round: int):
|
||||
self.output_num_blocks = output_num_blocks
|
||||
self.num_merge_tasks_per_round = num_merge_tasks_per_round
|
||||
self.num_reducers_per_merger = output_num_blocks // num_merge_tasks_per_round
|
||||
self._num_mergers_with_extra_reducer = (
|
||||
output_num_blocks % num_merge_tasks_per_round
|
||||
)
|
||||
|
||||
if self.num_reducers_per_merger == 0:
|
||||
self.num_merge_tasks_per_round = self._num_mergers_with_extra_reducer
|
||||
self.num_reducers_per_merger = 1
|
||||
self._num_mergers_with_extra_reducer = 0
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
f" num merge tasks per round: {self.num_merge_tasks_per_round}\n"
|
||||
f" num reduce tasks per merge task: {self.num_reducers_per_merger}\n"
|
||||
" num merge tasks with extra reduce task: "
|
||||
f"{self._num_mergers_with_extra_reducer}"
|
||||
)
|
||||
|
||||
def get_num_reducers_per_merge_idx(self, merge_idx: int) -> int:
|
||||
"""
|
||||
Each intermediate merge task will produce outputs for a partition of P
|
||||
final reduce tasks. This helper function returns P based on the merge
|
||||
task index.
|
||||
"""
|
||||
assert merge_idx < self.num_merge_tasks_per_round
|
||||
num_reducers_for_cur_merger = self.num_reducers_per_merger
|
||||
if merge_idx < self._num_mergers_with_extra_reducer:
|
||||
num_reducers_for_cur_merger += 1
|
||||
return num_reducers_for_cur_merger
|
||||
|
||||
def get_merge_idx_for_reducer_idx(self, reducer_idx: int) -> int:
|
||||
if (
|
||||
reducer_idx
|
||||
< (self.num_reducers_per_merger + 1) * self._num_mergers_with_extra_reducer
|
||||
):
|
||||
merge_idx = reducer_idx // (self.num_reducers_per_merger + 1)
|
||||
else:
|
||||
reducer_idx -= (
|
||||
self.num_reducers_per_merger + 1
|
||||
) * self._num_mergers_with_extra_reducer
|
||||
merge_idx = (
|
||||
self._num_mergers_with_extra_reducer
|
||||
+ reducer_idx // self.num_reducers_per_merger
|
||||
)
|
||||
assert merge_idx < self.num_merge_tasks_per_round
|
||||
return merge_idx
|
||||
|
||||
def round_robin_reduce_idx_iterator(self):
|
||||
"""
|
||||
When there are multiple nodes, merge tasks are spread throughout the
|
||||
cluster to improve load-balancing. Each merge task produces outputs for
|
||||
a contiguous partition of reduce tasks. This method creates an iterator
|
||||
that returns reduce task indices round-robin across the merge tasks.
|
||||
This can be used to submit reduce tasks in a way that spreads the load
|
||||
evenly across the cluster.
|
||||
"""
|
||||
idx = 0
|
||||
round_idx = 0
|
||||
while idx < self.output_num_blocks:
|
||||
for merge_idx in range(self.num_merge_tasks_per_round):
|
||||
if merge_idx < self._num_mergers_with_extra_reducer:
|
||||
reduce_idx = merge_idx * (self.num_reducers_per_merger + 1)
|
||||
num_reducers_for_cur_merger = self.num_reducers_per_merger + 1
|
||||
else:
|
||||
reduce_idx = self._num_mergers_with_extra_reducer * (
|
||||
self.num_reducers_per_merger + 1
|
||||
)
|
||||
merge_idx -= self._num_mergers_with_extra_reducer
|
||||
reduce_idx += merge_idx * self.num_reducers_per_merger
|
||||
num_reducers_for_cur_merger = self.num_reducers_per_merger
|
||||
|
||||
if round_idx >= num_reducers_for_cur_merger:
|
||||
continue
|
||||
|
||||
reduce_idx += round_idx
|
||||
yield reduce_idx
|
||||
idx += 1
|
||||
round_idx += 1
|
||||
|
||||
|
||||
class _PushBasedShuffleStage:
|
||||
def __init__(
|
||||
self,
|
||||
output_num_blocks: int,
|
||||
num_rounds: int,
|
||||
num_map_tasks_per_round: int,
|
||||
merge_task_placement: List[str],
|
||||
):
|
||||
# The number of rounds of map-merge tasks. Reducer tasks are given the
|
||||
# outputs of the merge tasks as inputs. Reducer tasks receive one input
|
||||
# per round.
|
||||
self.num_rounds = num_rounds
|
||||
# The number of map tasks per round of map-merge tasks. The map task
|
||||
# produces one output per merge task in the same round.
|
||||
self.num_map_tasks_per_round = num_map_tasks_per_round
|
||||
|
||||
node_strategies = {
|
||||
node_id: {
|
||||
"scheduling_strategy": NodeAffinitySchedulingStrategy(
|
||||
node_id, soft=True
|
||||
)
|
||||
}
|
||||
for node_id in set(merge_task_placement)
|
||||
}
|
||||
self._merge_task_options = [
|
||||
node_strategies[node_id] for node_id in merge_task_placement
|
||||
]
|
||||
|
||||
self.merge_schedule = _MergeTaskSchedule(
|
||||
output_num_blocks, len(merge_task_placement)
|
||||
)
|
||||
|
||||
def get_estimated_num_refs(self) -> int:
|
||||
# Number of intermediate blocks = Number of rounds x (map tasks per
|
||||
# round * merge tasks per round).
|
||||
num_intermediate_refs = self.num_rounds * (
|
||||
self.num_map_tasks_per_round * self.merge_schedule.num_merge_tasks_per_round
|
||||
)
|
||||
# Number of input blocks + intermediate blocks + output blocks.
|
||||
num_refs_total = (
|
||||
(self.num_rounds * self.num_map_tasks_per_round)
|
||||
+ num_intermediate_refs
|
||||
+ self.merge_schedule.output_num_blocks
|
||||
)
|
||||
return num_refs_total
|
||||
|
||||
def get_merge_task_options(self, merge_idx):
|
||||
return self._merge_task_options[merge_idx]
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
"num map tasks per round (num args per merge task): "
|
||||
f"{self.num_map_tasks_per_round}\n"
|
||||
f"num rounds (num args per reduce task): {self.num_rounds}\n"
|
||||
"merge task placement: \n"
|
||||
f"{self.merge_schedule}"
|
||||
)
|
||||
|
||||
|
||||
class _PipelinedStageExecutor:
|
||||
def __init__(
|
||||
self,
|
||||
stage_iter,
|
||||
num_tasks_per_round: int,
|
||||
max_concurrent_rounds: int = 1,
|
||||
progress_bar: Optional["BaseProgressBar"] = None,
|
||||
):
|
||||
self._stage_iter = stage_iter
|
||||
self._num_tasks_per_round = num_tasks_per_round
|
||||
self._max_concurrent_rounds = max_concurrent_rounds
|
||||
self._progress_bar = progress_bar
|
||||
|
||||
self._rounds: List[List[ObjectRef]] = []
|
||||
self._task_idx = 0
|
||||
|
||||
self._submit_round()
|
||||
|
||||
self._num_block_bytes_stored_at_driver = 0
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self) -> List["BlockMetadataWithSchema"]:
|
||||
"""
|
||||
Submit one round of tasks. If we already have the max concurrent rounds
|
||||
in flight, first wait for the oldest round of tasks to finish.
|
||||
"""
|
||||
prev_metadata_and_schema = []
|
||||
if all(len(r) == 0 for r in self._rounds):
|
||||
raise StopIteration
|
||||
|
||||
if len(self._rounds) >= self._max_concurrent_rounds:
|
||||
prev_metadata_schema_refs = self._rounds.pop(0)
|
||||
if prev_metadata_schema_refs:
|
||||
if self._progress_bar is not None:
|
||||
prev_metadata_and_schema = self._progress_bar.fetch_until_complete(
|
||||
prev_metadata_schema_refs
|
||||
)
|
||||
# TODO(swang): Eagerly free the previous round's args.
|
||||
# See https://github.com/ray-project/ray/issues/42145.
|
||||
else:
|
||||
prev_metadata_and_schema = ray.get(prev_metadata_schema_refs)
|
||||
|
||||
self._submit_round()
|
||||
|
||||
return prev_metadata_and_schema
|
||||
|
||||
def _submit_round(self):
|
||||
assert len(self._rounds) < self._max_concurrent_rounds
|
||||
task_round = []
|
||||
for _ in range(self._num_tasks_per_round):
|
||||
try:
|
||||
task_round.append(next(self._stage_iter))
|
||||
except StopIteration:
|
||||
break
|
||||
self._rounds.append(task_round)
|
||||
|
||||
|
||||
class _MapStageIterator:
|
||||
def __init__(self, input_blocks_list, shuffle_map, map_args):
|
||||
self._input_blocks_list = input_blocks_list
|
||||
self._shuffle_map = shuffle_map
|
||||
self._map_args = map_args
|
||||
|
||||
self._mapper_idx = 0
|
||||
self._map_results = []
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if not self._input_blocks_list:
|
||||
raise StopIteration
|
||||
|
||||
block = self._input_blocks_list.pop(0)
|
||||
# NOTE(swang): Results are shuffled between map and merge tasks, so
|
||||
# there is no advantage to colocating specific map and merge tasks.
|
||||
# Therefore, we do not specify a node affinity policy for map tasks
|
||||
# in case the caller or Ray has a better scheduling strategy, e.g.,
|
||||
# based on data locality.
|
||||
map_result = self._shuffle_map.remote(
|
||||
self._mapper_idx,
|
||||
block,
|
||||
*self._map_args,
|
||||
)
|
||||
metadata_schema_ref = map_result.pop(-1)
|
||||
self._map_results.append(map_result)
|
||||
self._mapper_idx += 1
|
||||
return metadata_schema_ref
|
||||
|
||||
def pop_map_results(self) -> List[List[ObjectRef]]:
|
||||
map_results = self._map_results
|
||||
self._map_results = []
|
||||
return map_results
|
||||
|
||||
|
||||
class _MergeStageIterator:
|
||||
def __init__(
|
||||
self,
|
||||
map_stage_iter: _MapStageIterator,
|
||||
shuffle_merge,
|
||||
stage: _PushBasedShuffleStage,
|
||||
reduce_args,
|
||||
):
|
||||
self._map_stage_iter = map_stage_iter
|
||||
self._shuffle_merge = shuffle_merge
|
||||
self._stage = stage
|
||||
self._reduce_args = reduce_args
|
||||
|
||||
self._merge_idx = 0
|
||||
self._map_result_buffer = None
|
||||
# Final outputs from the map-merge stage.
|
||||
# This is a map from merge task index to a nested list of merge results
|
||||
# (ObjectRefs). Each merge task index corresponds to a partition of P
|
||||
# final reduce tasks.
|
||||
self._all_merge_results = [
|
||||
[] for _ in range(self._stage.merge_schedule.num_merge_tasks_per_round)
|
||||
]
|
||||
|
||||
def __next__(self):
|
||||
if not self._map_result_buffer or not self._map_result_buffer[0]:
|
||||
assert self._merge_idx == 0
|
||||
self._map_result_buffer = self._map_stage_iter.pop_map_results()
|
||||
|
||||
if not self._map_result_buffer:
|
||||
raise StopIteration
|
||||
|
||||
# Shuffle the map results for the merge tasks.
|
||||
merge_args = [map_result.pop(0) for map_result in self._map_result_buffer]
|
||||
num_merge_returns = self._stage.merge_schedule.get_num_reducers_per_merge_idx(
|
||||
self._merge_idx
|
||||
)
|
||||
merge_result = self._shuffle_merge.options(
|
||||
num_returns=1 + num_merge_returns,
|
||||
**self._stage.get_merge_task_options(self._merge_idx),
|
||||
).remote(
|
||||
*merge_args,
|
||||
reduce_args=self._reduce_args,
|
||||
)
|
||||
metadata_schema_ref = merge_result.pop(-1)
|
||||
self._all_merge_results[self._merge_idx].append(merge_result)
|
||||
del merge_result
|
||||
|
||||
self._merge_idx += 1
|
||||
self._merge_idx %= self._stage.merge_schedule.num_merge_tasks_per_round
|
||||
return metadata_schema_ref
|
||||
|
||||
def pop_merge_results(self) -> List[List[ObjectRef]]:
|
||||
"""Return a nested list of merge task results. The list at index i
|
||||
stores the outputs of the i-th merge task submitted during each
|
||||
map-merge round. Each merge task returns a list of outputs because it
|
||||
may produce outputs for multiple downstream reduce tasks.
|
||||
"""
|
||||
all_merge_results = self._all_merge_results
|
||||
self._all_merge_results = []
|
||||
return all_merge_results
|
||||
|
||||
|
||||
class _ReduceStageIterator:
|
||||
def __init__(
|
||||
self,
|
||||
stage: _PushBasedShuffleStage,
|
||||
shuffle_reduce,
|
||||
all_merge_results: List[List[List[ObjectRef]]],
|
||||
ray_remote_args,
|
||||
reduce_args: List[Any],
|
||||
_debug_limit_execution_to_num_blocks: Optional[int],
|
||||
):
|
||||
self._shuffle_reduce = shuffle_reduce
|
||||
self._stage = stage
|
||||
self._reduce_arg_blocks: List[Tuple[int, List[ObjectRef]]] = []
|
||||
self._ray_remote_args = ray_remote_args
|
||||
self._reduce_args = reduce_args
|
||||
|
||||
for reduce_idx in self._stage.merge_schedule.round_robin_reduce_idx_iterator():
|
||||
merge_idx = self._stage.merge_schedule.get_merge_idx_for_reducer_idx(
|
||||
reduce_idx
|
||||
)
|
||||
reduce_arg_blocks = [
|
||||
merge_results.pop(0) for merge_results in all_merge_results[merge_idx]
|
||||
]
|
||||
self._reduce_arg_blocks.append((reduce_idx, reduce_arg_blocks))
|
||||
|
||||
assert len(self._reduce_arg_blocks) == stage.merge_schedule.output_num_blocks
|
||||
|
||||
if _debug_limit_execution_to_num_blocks is not None:
|
||||
self._reduce_arg_blocks = self._reduce_arg_blocks[
|
||||
:_debug_limit_execution_to_num_blocks
|
||||
]
|
||||
logger.debug(
|
||||
f"Limiting execution to {len(self._reduce_arg_blocks)} reduce tasks"
|
||||
)
|
||||
|
||||
for merge_idx, merge_results in enumerate(all_merge_results):
|
||||
assert all(len(merge_result) == 0 for merge_result in merge_results), (
|
||||
"Reduce stage did not process outputs from merge tasks at index: "
|
||||
f"{merge_idx}"
|
||||
)
|
||||
|
||||
self._reduce_results: List[Tuple[int, ObjectRef]] = []
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if not self._reduce_arg_blocks:
|
||||
raise StopIteration
|
||||
|
||||
reduce_idx, reduce_arg_blocks = self._reduce_arg_blocks.pop(0)
|
||||
merge_idx = self._stage.merge_schedule.get_merge_idx_for_reducer_idx(reduce_idx)
|
||||
# Submit one partition of reduce tasks, one for each of the P
|
||||
# outputs produced by the corresponding merge task.
|
||||
# We also add the merge task arguments so that the reduce task
|
||||
# is colocated with its inputs.
|
||||
block, meta_with_schema = self._shuffle_reduce.options(
|
||||
**self._ray_remote_args,
|
||||
**self._stage.get_merge_task_options(merge_idx),
|
||||
num_returns=2,
|
||||
).remote(*self._reduce_args, *reduce_arg_blocks, partial_reduce=False)
|
||||
self._reduce_results.append((reduce_idx, block))
|
||||
return meta_with_schema
|
||||
|
||||
def pop_reduce_results(self):
|
||||
reduce_results = self._reduce_results
|
||||
self._reduce_results = []
|
||||
return reduce_results
|
||||
|
||||
|
||||
class PushBasedShuffleTaskScheduler(ExchangeTaskScheduler):
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
|
||||
def execute(
|
||||
self,
|
||||
refs: List[RefBundle],
|
||||
output_num_blocks: int,
|
||||
task_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.debug("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:
|
||||
input_blocks_list.extend(ref_bundle.block_refs)
|
||||
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,
|
||||
)
|
||||
|
||||
caller_memory_usage = (
|
||||
stage.get_estimated_num_refs() * CALLER_MEMORY_USAGE_PER_OBJECT_REF
|
||||
)
|
||||
self.warn_on_driver_memory_usage(
|
||||
caller_memory_usage,
|
||||
"Execution is estimated to use at least "
|
||||
f"{convert_bytes_to_human_readable_str(caller_memory_usage)}"
|
||||
" of driver memory. Ensure that the driver machine has at least "
|
||||
"this much memory to ensure job completion.",
|
||||
)
|
||||
|
||||
# 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.debug(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.debug(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 = task_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_schema = []
|
||||
merge_stage_metadata_schema = []
|
||||
while not (map_done and merge_done):
|
||||
try:
|
||||
map_stage_metadata_schema += next(map_stage_executor)
|
||||
except StopIteration:
|
||||
map_done = True
|
||||
break
|
||||
|
||||
try:
|
||||
merge_stage_metadata_schema += next(merge_stage_executor)
|
||||
except StopIteration:
|
||||
merge_done = True
|
||||
break
|
||||
|
||||
self.warn_on_high_local_memory_store_usage()
|
||||
|
||||
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_schema = []
|
||||
while True:
|
||||
try:
|
||||
reduce_stage_metadata_schema += next(reduce_stage_executor)
|
||||
except StopIteration:
|
||||
break
|
||||
|
||||
self.warn_on_high_local_memory_store_usage()
|
||||
|
||||
new_blocks = reduce_stage_iter.pop_reduce_results()
|
||||
sorted_blocks = [
|
||||
(block[0], block[1], reduce_stage_metadata_schema[i])
|
||||
for i, block in enumerate(new_blocks)
|
||||
]
|
||||
sorted_blocks.sort(key=lambda x: x[0])
|
||||
|
||||
new_blocks, reduce_stage_metadata_schema = [], []
|
||||
if sorted_blocks:
|
||||
res: Tuple[
|
||||
List[Any], List[ObjectRef[Block]], List[BlockMetadataWithSchema]
|
||||
] = unzip(sorted_blocks)
|
||||
_, new_blocks, reduce_stage_metadata_schema = res
|
||||
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_with_schema in zip(new_blocks, reduce_stage_metadata_schema):
|
||||
output.append(
|
||||
RefBundle(
|
||||
[BlockEntry(block, meta_with_schema.metadata)],
|
||||
owns_blocks=input_owned,
|
||||
schema=meta_with_schema.schema,
|
||||
)
|
||||
)
|
||||
|
||||
stats = {
|
||||
"map": to_stats(map_stage_metadata_schema),
|
||||
"merge": to_stats(merge_stage_metadata_schema),
|
||||
"reduce": to_stats(reduce_stage_metadata_schema),
|
||||
}
|
||||
|
||||
return (output, stats)
|
||||
|
||||
@staticmethod
|
||||
def _map_partition(
|
||||
map_fn,
|
||||
idx: int,
|
||||
block: Block,
|
||||
output_num_blocks: int,
|
||||
schedule: _MergeTaskSchedule,
|
||||
*map_args: List[Any],
|
||||
) -> List[Union[Block, "BlockMetadataWithSchema"]]:
|
||||
mapper_outputs = map_fn(idx, block, output_num_blocks, *map_args)
|
||||
|
||||
# A merge task may produce results for multiple downstream reducer
|
||||
# tasks. Therefore, each map task should give each merge task a
|
||||
# partition of its outputs, where the length of the partition is equal
|
||||
# to the number of reducers downstream to the merge task.
|
||||
partition = []
|
||||
merge_idx = 0
|
||||
while merge_idx < schedule.num_merge_tasks_per_round and mapper_outputs:
|
||||
output = mapper_outputs.pop(0)
|
||||
partition.append(output)
|
||||
|
||||
if len(partition) == schedule.get_num_reducers_per_merge_idx(merge_idx):
|
||||
yield partition
|
||||
|
||||
partition = []
|
||||
merge_idx += 1
|
||||
|
||||
assert not partition
|
||||
assert len(mapper_outputs) == 1, (
|
||||
mapper_outputs,
|
||||
"The last output should be a BlockMetadataWithSchema",
|
||||
)
|
||||
assert isinstance(mapper_outputs[0], BlockMetadataWithSchema)
|
||||
yield mapper_outputs[0]
|
||||
|
||||
assert merge_idx == schedule.num_merge_tasks_per_round, (
|
||||
merge_idx,
|
||||
schedule.num_merge_tasks_per_round,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _merge(
|
||||
reduce_fn,
|
||||
*all_mapper_outputs: List[List[Block]],
|
||||
reduce_args: Optional[List[Any]] = None,
|
||||
) -> List[Union["BlockMetadataWithSchema", Block]]:
|
||||
"""
|
||||
Returns list of [BlockMetadata, O1, O2, O3, ...output_num_blocks].
|
||||
"""
|
||||
assert (
|
||||
len({len(mapper_outputs) for mapper_outputs in all_mapper_outputs}) == 1
|
||||
), "Received different number of map inputs"
|
||||
stats = BlockExecStats.builder()
|
||||
if not reduce_args:
|
||||
reduce_args = []
|
||||
|
||||
num_rows = 0
|
||||
size_bytes = 0
|
||||
schemas = []
|
||||
for i, mapper_outputs in enumerate(zip(*all_mapper_outputs)):
|
||||
block_meta_with_schema: Tuple[Block, "BlockMetadataWithSchema"] = reduce_fn(
|
||||
*reduce_args, *mapper_outputs, partial_reduce=True
|
||||
)
|
||||
block, meta_with_schema = block_meta_with_schema
|
||||
yield block
|
||||
|
||||
block = BlockAccessor.for_block(block)
|
||||
num_rows += block.num_rows()
|
||||
size_bytes += block.size_bytes()
|
||||
del block
|
||||
schemas.append(meta_with_schema.schema)
|
||||
|
||||
schema = _take_first_non_empty_schema(iter(schemas))
|
||||
|
||||
meta = BlockMetadata(
|
||||
num_rows=num_rows,
|
||||
size_bytes=size_bytes,
|
||||
input_files=None,
|
||||
exec_stats=stats.build(),
|
||||
)
|
||||
meta_with_schema = BlockMetadataWithSchema.from_metadata(meta, schema=schema)
|
||||
yield meta_with_schema
|
||||
|
||||
@staticmethod
|
||||
def _compute_shuffle_schedule(
|
||||
num_cpus_per_node_map: Dict[str, int],
|
||||
num_input_blocks: int,
|
||||
merge_factor: float,
|
||||
num_output_blocks: int,
|
||||
) -> _PushBasedShuffleStage:
|
||||
num_cpus_total = sum(v for v in num_cpus_per_node_map.values())
|
||||
logger.debug(
|
||||
f"Found {num_cpus_total} CPUs available CPUs for push-based shuffle."
|
||||
)
|
||||
num_tasks_per_map_merge_group = merge_factor + 1
|
||||
num_total_merge_tasks = math.ceil(num_input_blocks / merge_factor)
|
||||
|
||||
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.
|
||||
num_input_blocks_remaining = num_input_blocks
|
||||
for node, num_cpus in num_cpus_per_node_map.items():
|
||||
# First find how many merge tasks we should run on this node.
|
||||
# We take the min of the number of CPUs on this node and the number
|
||||
# of input blocks that we haven't scheduled yet, in case there are
|
||||
# fewer input blocks than CPU slots on this node.
|
||||
num_cpu_slots = min(num_cpus, num_input_blocks_remaining)
|
||||
num_merge_tasks_on_cur_node = round(
|
||||
num_cpu_slots / num_tasks_per_map_merge_group
|
||||
)
|
||||
# For small datasets, the number of tasks to run may be less than
|
||||
# the total CPU slots available.
|
||||
num_merge_tasks_on_cur_node = min(
|
||||
num_merge_tasks_on_cur_node, num_total_merge_tasks
|
||||
)
|
||||
for i in range(num_merge_tasks_on_cur_node):
|
||||
merge_task_placement.append(node)
|
||||
# We schedule `merge_factor` many map tasks for every merge
|
||||
# task. Subtract from the number of input blocks remaining to
|
||||
# account for cases where the number of map tasks is smaller
|
||||
# than the available CPU slots.
|
||||
num_input_blocks_remaining -= merge_factor
|
||||
num_cpus -= num_tasks_per_map_merge_group
|
||||
num_merge_tasks_per_round += num_merge_tasks_on_cur_node
|
||||
|
||||
# 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 += num_cpus
|
||||
if (
|
||||
leftover_cpus >= num_tasks_per_map_merge_group
|
||||
and num_merge_tasks_per_round < num_total_merge_tasks
|
||||
):
|
||||
merge_task_placement.append(node)
|
||||
num_merge_tasks_per_round += 1
|
||||
leftover_cpus -= num_tasks_per_map_merge_group
|
||||
num_input_blocks_remaining -= merge_factor
|
||||
|
||||
num_input_blocks_remaining = max(0, num_input_blocks_remaining)
|
||||
|
||||
if num_merge_tasks_per_round == 0:
|
||||
# For small datasets, make sure we have at least one merge task.
|
||||
for node, num_cpus in num_cpus_per_node_map.items():
|
||||
if num_cpus >= 1:
|
||||
merge_task_placement.append(node)
|
||||
num_merge_tasks_per_round = 1
|
||||
break
|
||||
|
||||
assert num_merge_tasks_per_round == len(merge_task_placement)
|
||||
assert num_merge_tasks_per_round > 0, num_merge_tasks_per_round
|
||||
# Use the remaining CPUs to execute map tasks.
|
||||
num_map_tasks_per_round = num_cpus_total - num_merge_tasks_per_round
|
||||
num_map_tasks_per_round = min(num_map_tasks_per_round, num_input_blocks)
|
||||
# Make sure there is at least one map task in each round.
|
||||
num_map_tasks_per_round = max(num_map_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,
|
||||
)
|
||||
|
||||
|
||||
def _get_num_cpus_per_node_map() -> Dict[str, int]:
|
||||
total_resources_by_node = ray.state.total_resources_per_node()
|
||||
# Map from per-node resource name to number of CPUs available on that
|
||||
# node.
|
||||
num_cpus_per_node_map = {}
|
||||
for node_id, resources in total_resources_by_node.items():
|
||||
num_cpus = int(resources.get("CPU", 0))
|
||||
if num_cpus == 0:
|
||||
continue
|
||||
num_cpus_per_node_map[node_id] = num_cpus
|
||||
return num_cpus_per_node_map
|
||||
@@ -0,0 +1,152 @@
|
||||
import logging
|
||||
import math
|
||||
from typing import Callable, Iterable, List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
|
||||
from ray.data._internal.execution.interfaces.task_context import TaskContext
|
||||
from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec
|
||||
from ray.data.block import (
|
||||
Block,
|
||||
BlockAccessor,
|
||||
BlockExecStats,
|
||||
BlockMetadata,
|
||||
BlockMetadataWithSchema,
|
||||
)
|
||||
from ray.data.context import MAX_SAFE_BLOCK_SIZE_FACTOR
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ShuffleTaskSpec(ExchangeTaskSpec):
|
||||
"""
|
||||
The implementation for shuffle tasks.
|
||||
|
||||
This is used by random_shuffle() and repartition().
|
||||
"""
|
||||
|
||||
SPLIT_REPARTITION_SUB_PROGRESS_BAR_NAME = "Split Repartition"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
target_shuffle_max_block_size: int,
|
||||
random_shuffle: bool = False,
|
||||
random_seed: Optional[int] = None,
|
||||
upstream_map_fn: Optional[Callable[[Iterable[Block]], Iterable[Block]]] = None,
|
||||
):
|
||||
super().__init__(
|
||||
map_args=[
|
||||
target_shuffle_max_block_size,
|
||||
upstream_map_fn,
|
||||
random_shuffle,
|
||||
random_seed,
|
||||
],
|
||||
reduce_args=[random_shuffle, random_seed],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def map(
|
||||
idx: int,
|
||||
block: Block,
|
||||
output_num_blocks: int,
|
||||
target_shuffle_max_block_size: int,
|
||||
upstream_map_fn: Optional[Callable[[Iterable[Block]], Iterable[Block]]],
|
||||
random_shuffle: bool,
|
||||
random_seed: Optional[int],
|
||||
) -> List[Union[Block, "BlockMetadataWithSchema"]]:
|
||||
stats = BlockExecStats.builder()
|
||||
if upstream_map_fn:
|
||||
# Create a local TaskContext for the upstream map function.
|
||||
# May be used by expressions that depend on task-level state.
|
||||
local_ctx = TaskContext(task_idx=idx, op_name="shuffle_map")
|
||||
with TaskContext.current(local_ctx):
|
||||
# TODO: Support dynamic block splitting in
|
||||
# all-to-all ops, to avoid having to re-fuse
|
||||
# upstream blocks together.
|
||||
upstream_map_iter = upstream_map_fn([block])
|
||||
mapped_block = next(upstream_map_iter)
|
||||
builder = BlockAccessor.for_block(mapped_block).builder()
|
||||
builder.add_block(mapped_block)
|
||||
for mapped_block in upstream_map_iter:
|
||||
builder.add_block(mapped_block)
|
||||
# Drop the upstream inputs to reduce memory usage.
|
||||
del mapped_block
|
||||
block = builder.build()
|
||||
|
||||
block = BlockAccessor.for_block(block)
|
||||
if (
|
||||
block.size_bytes()
|
||||
> MAX_SAFE_BLOCK_SIZE_FACTOR * target_shuffle_max_block_size
|
||||
):
|
||||
logger.warning(
|
||||
"Input block to map task has size "
|
||||
f"{block.size_bytes() // (1024 * 1024)}MiB, which exceeds "
|
||||
"DataContext.get_current().target_shuffle_max_block_size="
|
||||
f"{target_shuffle_max_block_size // (1024 * 1024)}MiB. "
|
||||
"This can lead to out-of-memory errors and can happen "
|
||||
"when map tasks are fused to the shuffle operation. "
|
||||
"To prevent fusion, call Dataset.materialize() on the "
|
||||
"dataset before shuffling."
|
||||
)
|
||||
|
||||
# Randomize the distribution of records to blocks.
|
||||
if random_shuffle:
|
||||
seed_i = random_seed + idx if random_seed is not None else None
|
||||
block = block.random_shuffle(seed_i)
|
||||
block = BlockAccessor.for_block(block)
|
||||
|
||||
# Build a list of slices to return. It's okay to put the results in a
|
||||
# list instead of yielding them as a generator because slicing the
|
||||
# ArrowBlock is zero-copy.
|
||||
slice_sz = max(1, math.ceil(block.num_rows() / output_num_blocks))
|
||||
slices = []
|
||||
for i in range(output_num_blocks):
|
||||
slices.append(block.slice(i * slice_sz, (i + 1) * slice_sz))
|
||||
|
||||
# Randomize the distribution order of the blocks (this prevents empty
|
||||
# outputs when input blocks are very small).
|
||||
if random_shuffle:
|
||||
random = np.random.RandomState(seed_i)
|
||||
random.shuffle(slices)
|
||||
|
||||
num_rows = sum(BlockAccessor.for_block(s).num_rows() for s in slices)
|
||||
assert num_rows == block.num_rows(), (num_rows, block.num_rows())
|
||||
from ray.data.block import BlockMetadataWithSchema
|
||||
|
||||
meta = block.get_metadata(block_exec_stats=stats.build())
|
||||
schema = block.schema()
|
||||
meta_with_schema = BlockMetadataWithSchema.from_metadata(meta, schema=schema)
|
||||
return slices + [meta_with_schema]
|
||||
|
||||
@staticmethod
|
||||
def reduce(
|
||||
random_shuffle: bool,
|
||||
random_seed: Optional[int],
|
||||
*mapper_outputs: List[Block],
|
||||
partial_reduce: bool = False,
|
||||
) -> Tuple[Block, "BlockMetadataWithSchema"]:
|
||||
# TODO: Support fusion with other downstream operators.
|
||||
stats = BlockExecStats.builder()
|
||||
builder = DelegatingBlockBuilder()
|
||||
for block in mapper_outputs:
|
||||
builder.add_block(block)
|
||||
new_block = builder.build()
|
||||
accessor = BlockAccessor.for_block(new_block)
|
||||
if random_shuffle:
|
||||
new_block = accessor.random_shuffle(
|
||||
random_seed if random_seed is not None else None
|
||||
)
|
||||
accessor = BlockAccessor.for_block(new_block)
|
||||
new_metadata = BlockMetadata(
|
||||
num_rows=accessor.num_rows(),
|
||||
size_bytes=accessor.size_bytes(),
|
||||
input_files=None,
|
||||
exec_stats=stats.build(),
|
||||
)
|
||||
from ray.data.block import BlockMetadataWithSchema
|
||||
|
||||
meta_with_schema = BlockMetadataWithSchema.from_metadata(
|
||||
new_metadata, schema=accessor.schema()
|
||||
)
|
||||
return new_block, meta_with_schema
|
||||
@@ -0,0 +1,240 @@
|
||||
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, TypeVar, Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
|
||||
from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec
|
||||
from ray.data._internal.progress.progress_bar import ProgressBar
|
||||
from ray.data._internal.remote_fn import cached_remote_fn
|
||||
from ray.data._internal.table_block import TableBlockAccessor
|
||||
from ray.data._internal.util import NULL_SENTINEL
|
||||
from ray.data.block import Block, BlockAccessor, BlockExecStats
|
||||
from ray.types import ObjectRef
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import pyarrow
|
||||
|
||||
from ray.data.block import BlockMetadataWithSchema
|
||||
|
||||
|
||||
class SortKey:
|
||||
"""SortKey class to convert between different sort args formats."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
key: Optional[Union[str, List[str]]] = None,
|
||||
descending: Union[bool, List[bool]] = False,
|
||||
boundaries: Optional[List[T]] = None,
|
||||
):
|
||||
if key is None:
|
||||
key = []
|
||||
if isinstance(key, str):
|
||||
key = [key]
|
||||
if not (isinstance(key, list) and all(isinstance(k, str) for k in key)):
|
||||
raise ValueError(
|
||||
f"Key must be a string or a list of strings, but got {key}."
|
||||
)
|
||||
if isinstance(descending, bool):
|
||||
descending = [descending for _ in key]
|
||||
elif isinstance(descending, list):
|
||||
if len(descending) != len(key):
|
||||
raise ValueError(
|
||||
"Length of `descending` does not match the length of the key."
|
||||
)
|
||||
self._columns = key
|
||||
self._descending = descending
|
||||
if boundaries:
|
||||
for item in boundaries:
|
||||
if not isinstance(item, (int, float)):
|
||||
raise ValueError(
|
||||
"The type of items in boundaries must be int or float."
|
||||
)
|
||||
boundaries = list(set(boundaries))
|
||||
boundaries.sort()
|
||||
self._boundaries = boundaries
|
||||
|
||||
def get_columns(self) -> List[str]:
|
||||
return self._columns
|
||||
|
||||
def get_descending(self) -> List[bool]:
|
||||
return self._descending
|
||||
|
||||
def to_arrow_sort_args(self) -> List[Tuple[str, str]]:
|
||||
return [
|
||||
(key, "descending" if desc else "ascending")
|
||||
for key, desc in zip(self._columns, self._descending)
|
||||
]
|
||||
|
||||
def to_pandas_sort_args(self) -> Tuple[List[str], List[bool]]:
|
||||
return self._columns, [not desc for desc in self._descending]
|
||||
|
||||
def validate_schema(self, schema: Optional[Union[type, "pyarrow.lib.Schema"]]):
|
||||
"""Check the key function is valid on the given schema."""
|
||||
if schema is None:
|
||||
# Dataset is empty/cleared, validation not possible.
|
||||
return
|
||||
|
||||
if self._columns and len(schema.names) > 0:
|
||||
schema_names_set = set(schema.names)
|
||||
for column in self._columns:
|
||||
if column not in schema_names_set:
|
||||
raise ValueError(
|
||||
f"You specified the column '{column}', but there's no such "
|
||||
"column in the dataset. The dataset has columns: "
|
||||
f"{schema.names}"
|
||||
)
|
||||
|
||||
@property
|
||||
def boundaries(self):
|
||||
return self._boundaries
|
||||
|
||||
|
||||
class SortTaskSpec(ExchangeTaskSpec):
|
||||
"""
|
||||
The implementation for distributed sort tasks.
|
||||
|
||||
The algorithm is similar to [External Merge Sort]
|
||||
(https://en.wikipedia.org/wiki/External_sorting).
|
||||
Sorting is done in 3 steps: sampling, sorting individual blocks, and
|
||||
merging sorted blocks.
|
||||
|
||||
Sampling (`sample_boundaries`): we get a number of sample items from each block,
|
||||
sort them, and use them to compute boundaries that would partition all items into
|
||||
approximately equal ranges.
|
||||
|
||||
Sorting (`map`): each block is sorted locally, then partitioned into smaller
|
||||
blocks according to the boundaries. Each partitioned block is passed to a merge
|
||||
task.
|
||||
|
||||
Merging (`reduce`): a merge task would receive a block from every worker that
|
||||
consists of items in a certain range. It then merges the sorted blocks into one
|
||||
sorted block and becomes part of the new, sorted block.
|
||||
"""
|
||||
|
||||
SORT_SAMPLE_SUB_PROGRESS_BAR_NAME = "Sort Sample"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
boundaries: List[T],
|
||||
sort_key: SortKey,
|
||||
):
|
||||
super().__init__(
|
||||
map_args=[boundaries, sort_key],
|
||||
reduce_args=[sort_key],
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def map(
|
||||
idx: int,
|
||||
block: Block,
|
||||
output_num_blocks: int,
|
||||
boundaries: List[T],
|
||||
sort_key: SortKey,
|
||||
) -> List[Union[Block, "BlockMetadataWithSchema"]]:
|
||||
stats = BlockExecStats.builder()
|
||||
accessor = BlockAccessor.for_block(block)
|
||||
out = accessor.sort_and_partition(boundaries, sort_key)
|
||||
from ray.data.block import BlockMetadataWithSchema
|
||||
|
||||
meta_with_schema = BlockMetadataWithSchema.from_block(
|
||||
block, block_exec_stats=stats.build()
|
||||
)
|
||||
return out + [meta_with_schema]
|
||||
|
||||
@staticmethod
|
||||
def reduce(
|
||||
sort_key: SortKey,
|
||||
*mapper_outputs: List[Block],
|
||||
partial_reduce: bool = False,
|
||||
) -> Tuple[Block, "BlockMetadataWithSchema"]:
|
||||
normalized_blocks = TableBlockAccessor.normalize_block_types(
|
||||
mapper_outputs,
|
||||
target_block_type=None,
|
||||
)
|
||||
blocks, meta_with_schema = BlockAccessor.for_block(
|
||||
normalized_blocks[0]
|
||||
).merge_sorted_blocks(normalized_blocks, sort_key)
|
||||
return blocks, meta_with_schema
|
||||
|
||||
@staticmethod
|
||||
def sample_boundaries(
|
||||
blocks: List[ObjectRef[Block]],
|
||||
sort_key: SortKey,
|
||||
num_reducers: int,
|
||||
sample_bar: Optional[ProgressBar] = None,
|
||||
label_selector: Optional[Dict[str, str]] = None,
|
||||
) -> List[T]:
|
||||
"""
|
||||
Return (num_reducers - 1) items in ascending order from the blocks that
|
||||
partition the domain into ranges with approximately equally many elements.
|
||||
Each boundary item is a tuple of a form (col1_value, col2_value, ...).
|
||||
"""
|
||||
columns = sort_key.get_columns()
|
||||
n_samples = int(num_reducers * 10 / len(blocks))
|
||||
|
||||
sample_block = cached_remote_fn(_sample_block)
|
||||
if label_selector:
|
||||
sample_block = sample_block.options(label_selector=label_selector)
|
||||
|
||||
sample_results = [
|
||||
sample_block.remote(block, n_samples, sort_key) for block in blocks
|
||||
]
|
||||
if sample_bar is None:
|
||||
sample_bar = ProgressBar(
|
||||
SortTaskSpec.SORT_SAMPLE_SUB_PROGRESS_BAR_NAME,
|
||||
len(blocks) * n_samples,
|
||||
unit="rows",
|
||||
)
|
||||
# TODO(zhilong): Update sort sample bar before finished.
|
||||
samples = sample_bar.fetch_until_complete(sample_results)
|
||||
del sample_results
|
||||
samples: List[Block] = [s for s in samples if len(s) > 0]
|
||||
# The dataset is empty
|
||||
if len(samples) == 0:
|
||||
return [None] * (num_reducers - 1)
|
||||
|
||||
# Convert samples to a sorted list[tuple[...]] where each tuple represents a
|
||||
# sample.
|
||||
# TODO: Once we deprecate pandas blocks, we can avoid this conversion and
|
||||
# directly sort the samples.
|
||||
builder = DelegatingBlockBuilder()
|
||||
for sample in samples:
|
||||
builder.add_block(sample)
|
||||
samples_table = builder.build()
|
||||
samples_dict = BlockAccessor.for_block(samples_table).to_numpy(columns=columns)
|
||||
# This zip does the transposition from list of column values to list of tuples.
|
||||
samples_list = list(zip(*samples_dict.values()))
|
||||
|
||||
def is_na(x):
|
||||
# Check if x is None or NaN. Type casting to np.array first to avoid
|
||||
# isnan failing on strings and other types.
|
||||
if x is None:
|
||||
return True
|
||||
x = np.asarray(x)
|
||||
if np.issubdtype(x.dtype, np.number):
|
||||
return np.isnan(x)
|
||||
return False
|
||||
|
||||
# To allow multi-directional sort, we utilize Python's stable sort: we
|
||||
# sort several times with different directions. We do this in reverse, so
|
||||
# that the last key we sort by is the primary sort key passed by the user.
|
||||
for i, desc in list(enumerate(sort_key.get_descending()))[::-1]:
|
||||
# Sort the list, but Nones should be NULL_SENTINEL to ensure safe sorting.
|
||||
samples_list.sort(
|
||||
key=lambda sample: NULL_SENTINEL if is_na(sample[i]) else sample[i],
|
||||
reverse=desc,
|
||||
)
|
||||
|
||||
# Each boundary corresponds to a quantile of the data.
|
||||
quantile_indices = [
|
||||
int(q * (len(samples_list) - 1))
|
||||
for q in np.linspace(0, 1, num_reducers + 1)
|
||||
]
|
||||
# Exclude the first and last quantiles because they're 0 and 1.
|
||||
return [samples_list[i] for i in quantile_indices[1:-1]]
|
||||
|
||||
|
||||
def _sample_block(block: Block, n_samples: int, sort_key: SortKey) -> Block:
|
||||
return BlockAccessor.for_block(block).sample(n_samples, sort_key)
|
||||
@@ -0,0 +1,171 @@
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import ray
|
||||
from ray.data._internal.execution.interfaces import BlockEntry, RefBundle, TaskContext
|
||||
from ray.data._internal.execution.interfaces.transform_fn import (
|
||||
AllToAllTransformFnResult,
|
||||
)
|
||||
from ray.data._internal.planner.exchange.interfaces import (
|
||||
ExchangeTaskScheduler,
|
||||
)
|
||||
from ray.data._internal.planner.exchange.shuffle_task_spec import ShuffleTaskSpec
|
||||
from ray.data._internal.remote_fn import cached_remote_fn
|
||||
from ray.data._internal.split import _split_at_indices
|
||||
from ray.data._internal.util import unzip
|
||||
from ray.data.block import (
|
||||
Block,
|
||||
BlockMetadata,
|
||||
BlockMetadataWithSchema,
|
||||
)
|
||||
from ray.types import ObjectRef
|
||||
|
||||
|
||||
class SplitRepartitionTaskScheduler(ExchangeTaskScheduler):
|
||||
"""
|
||||
The split (non-shuffle) repartition scheduler.
|
||||
|
||||
First, we calculate global splits needed to produce `output_num_blocks` blocks.
|
||||
After the split blocks are generated accordingly, reduce tasks are scheduled
|
||||
to combine split blocks together.
|
||||
"""
|
||||
|
||||
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,
|
||||
) -> AllToAllTransformFnResult:
|
||||
input_num_rows = 0
|
||||
input_owned_by_consumer = True
|
||||
for ref_bundle in refs:
|
||||
block_num_rows = ref_bundle.num_rows()
|
||||
if block_num_rows is None:
|
||||
raise ValueError(
|
||||
"Cannot split partition on blocks with unknown number of rows."
|
||||
)
|
||||
input_num_rows += block_num_rows
|
||||
if not ref_bundle.owns_blocks:
|
||||
input_owned_by_consumer = False
|
||||
|
||||
# Compute the (output_num_blocks) indices needed for an equal split of the
|
||||
# input blocks. When output_num_blocks=1, the total number of
|
||||
# input rows is used as the end index during the split calculation,
|
||||
# so that we can combine all input blocks into a single output block.
|
||||
indices = []
|
||||
if output_num_blocks == 1:
|
||||
indices = [input_num_rows]
|
||||
else:
|
||||
cur_idx = 0
|
||||
for _ in range(output_num_blocks - 1):
|
||||
cur_idx += input_num_rows / output_num_blocks
|
||||
indices.append(int(cur_idx))
|
||||
assert len(indices) <= output_num_blocks, (indices, output_num_blocks)
|
||||
|
||||
if map_ray_remote_args is None:
|
||||
map_ray_remote_args = {}
|
||||
if reduce_ray_remote_args is None:
|
||||
reduce_ray_remote_args = {}
|
||||
if "scheduling_strategy" not in reduce_ray_remote_args:
|
||||
reduce_ray_remote_args = reduce_ray_remote_args.copy()
|
||||
reduce_ray_remote_args["scheduling_strategy"] = "SPREAD"
|
||||
|
||||
blocks_with_metadata: List[Tuple[ObjectRef[Block], BlockMetadata]] = []
|
||||
for ref_bundle in refs:
|
||||
blocks_with_metadata.extend(
|
||||
(entry.ref, entry.metadata) for entry in ref_bundle.blocks
|
||||
)
|
||||
split_return = _split_at_indices(
|
||||
blocks_with_metadata,
|
||||
indices,
|
||||
input_owned_by_consumer,
|
||||
label_selector=map_ray_remote_args.get("label_selector"),
|
||||
)
|
||||
split_block_refs, split_metadata = [], []
|
||||
for b, m in zip(*split_return):
|
||||
split_block_refs.append(b)
|
||||
split_metadata.extend(m)
|
||||
|
||||
sub_progress_bar_dict = ctx.sub_progress_bar_dict
|
||||
bar_name = ShuffleTaskSpec.SPLIT_REPARTITION_SUB_PROGRESS_BAR_NAME
|
||||
assert bar_name in sub_progress_bar_dict, sub_progress_bar_dict
|
||||
reduce_bar = sub_progress_bar_dict[bar_name]
|
||||
|
||||
reduce_task = cached_remote_fn(self._exchange_spec.reduce)
|
||||
reduce_return = [
|
||||
reduce_task.options(**reduce_ray_remote_args, num_returns=2).remote(
|
||||
*self._exchange_spec._reduce_args,
|
||||
*split_block_refs[j],
|
||||
)
|
||||
for j in range(output_num_blocks)
|
||||
# Only process splits which contain blocks.
|
||||
if len(split_block_refs[j]) > 0
|
||||
]
|
||||
|
||||
reduce_block_refs, reduce_metadata_schema = [], []
|
||||
if reduce_return:
|
||||
reduce_block_refs, reduce_metadata_schema = unzip(reduce_return)
|
||||
reduce_metadata_schema: List[
|
||||
"BlockMetadataWithSchema"
|
||||
] = reduce_bar.fetch_until_complete(list(reduce_metadata_schema))
|
||||
reduce_block_refs = list(reduce_block_refs)
|
||||
|
||||
# Handle empty blocks.
|
||||
if len(reduce_block_refs) < output_num_blocks:
|
||||
import pyarrow as pa
|
||||
|
||||
from ray.data._internal.arrow_block import ArrowBlockBuilder
|
||||
from ray.data._internal.pandas_block import (
|
||||
PandasBlockBuilder,
|
||||
PandasBlockSchema,
|
||||
)
|
||||
|
||||
num_empty_blocks = output_num_blocks - len(reduce_block_refs)
|
||||
if len(reduce_metadata_schema) > 0:
|
||||
first_block_schema = reduce_metadata_schema[0].schema
|
||||
if isinstance(first_block_schema, pa.Schema):
|
||||
builder = ArrowBlockBuilder()
|
||||
elif isinstance(first_block_schema, PandasBlockSchema):
|
||||
builder = PandasBlockBuilder()
|
||||
else:
|
||||
raise ValueError(
|
||||
"Cannot split partition on blocks with unknown block schema:"
|
||||
f" {first_block_schema}."
|
||||
)
|
||||
else:
|
||||
# If the result is empty, default to Arrow format for the empty blocks.
|
||||
builder = ArrowBlockBuilder()
|
||||
|
||||
empty_block = builder.build()
|
||||
empty_meta_with_schema = BlockMetadataWithSchema.from_block(
|
||||
empty_block
|
||||
) # No stats for empty block.
|
||||
empty_block_refs, empty_metadata = zip(
|
||||
*[
|
||||
(ray.put(empty_block), empty_meta_with_schema)
|
||||
for _ in range(num_empty_blocks)
|
||||
]
|
||||
)
|
||||
reduce_block_refs.extend(empty_block_refs)
|
||||
reduce_metadata_schema.extend(empty_metadata)
|
||||
|
||||
output = []
|
||||
assert len(reduce_block_refs) == len(reduce_metadata_schema), (
|
||||
len(reduce_block_refs),
|
||||
len(reduce_metadata_schema),
|
||||
)
|
||||
for block, meta_with_schema in zip(reduce_block_refs, reduce_metadata_schema):
|
||||
output.append(
|
||||
RefBundle(
|
||||
[BlockEntry(block, meta_with_schema.metadata)],
|
||||
owns_blocks=input_owned_by_consumer,
|
||||
schema=meta_with_schema.schema,
|
||||
)
|
||||
)
|
||||
stats = {
|
||||
"split": split_metadata,
|
||||
"reduce": reduce_metadata_schema,
|
||||
}
|
||||
|
||||
return (output, stats)
|
||||
Reference in New Issue
Block a user