-
Notifications
You must be signed in to change notification settings - Fork 5.6k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[data] Disable block slicing for shuffle ops #40538
Changes from 1 commit
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -3,10 +3,13 @@ | |
|
||
import numpy as np | ||
|
||
from ray.data._internal.dataset_logger import DatasetLogger | ||
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder | ||
from ray.data._internal.planner.exchange.interfaces import ExchangeTaskSpec | ||
from ray.data.block import Block, BlockAccessor, BlockExecStats, BlockMetadata | ||
|
||
logger = DatasetLogger(__name__) | ||
|
||
|
||
class ShuffleTaskSpec(ExchangeTaskSpec): | ||
""" | ||
|
@@ -19,12 +22,18 @@ class ShuffleTaskSpec(ExchangeTaskSpec): | |
|
||
def __init__( | ||
self, | ||
target_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=[upstream_map_fn, random_shuffle, random_seed], | ||
map_args=[ | ||
target_max_block_size, | ||
upstream_map_fn, | ||
random_shuffle, | ||
random_seed, | ||
], | ||
reduce_args=[random_shuffle, random_seed], | ||
) | ||
|
||
|
@@ -33,22 +42,37 @@ def map( | |
idx: int, | ||
block: Block, | ||
output_num_blocks: int, | ||
target_max_block_size: int, | ||
upstream_map_fn: Optional[Callable[[Iterable[Block]], Iterable[Block]]], | ||
random_shuffle: bool, | ||
random_seed: Optional[int], | ||
) -> List[Union[BlockMetadata, Block]]: | ||
# TODO: Support fusion with other upstream operators. | ||
stats = BlockExecStats.builder() | ||
if upstream_map_fn: | ||
mapped_blocks = list(upstream_map_fn([block])) | ||
if len(mapped_blocks) > 1: | ||
builder = BlockAccessor.for_block(mapped_blocks[0]).builder() | ||
for b in mapped_blocks: | ||
builder.add_block(b) | ||
block = builder.build() | ||
else: | ||
block = mapped_blocks[0] | ||
# 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() > target_max_block_size: | ||
logger.get_logger().warn( | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I am worried about if this warning would confuse users. Can we log it as There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think print a warning makes sense. But we may want to make it less verbose by, for example, checking size_bytes > 2 * target_max_block_size. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'll change to 2 * target_max_block_size. But materialize() is the best solution here (see other comment). Hard to recommend an exact fix via increasing parallelism. |
||
"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_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." | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Setting different resources can also prevent fusion, without having to materialize all data. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think materializing is fine here because you have to materialize all data for an all-to-all op anyway. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. makes sense. |
||
) | ||
|
||
# Randomize the distribution of records to blocks. | ||
if random_shuffle: | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -36,12 +36,20 @@ def shuffle_repartition_fn( | |
map_transformer: Optional["MapTransformer"] = ctx.upstream_map_transformer | ||
upstream_map_fn = None | ||
if map_transformer: | ||
map_transformer.set_target_max_block_size(ctx.target_max_block_size) | ||
# NOTE(swang): We override the target block size with infinity, to | ||
# prevent the upstream map from slicing its output into smaller | ||
# blocks. Since the shuffle task will just fuse these back | ||
# together, the extra slicing and re-fusing can add high memory | ||
# overhead. This can be removed once dynamic block splitting is | ||
# supported for all-to-all ops. | ||
# See https://github.com/ray-project/ray/issues/40518. | ||
map_transformer.set_target_max_block_size(float("inf")) | ||
Comment on lines
+39
to
+46
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. nit: better to put this into a separate method, to avoid repeated comments. |
||
|
||
def upstream_map_fn(blocks): | ||
return map_transformer.apply_transform(blocks, ctx) | ||
|
||
shuffle_spec = ShuffleTaskSpec( | ||
ctx.target_max_block_size, | ||
random_shuffle=False, | ||
upstream_map_fn=upstream_map_fn, | ||
) | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.