-
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] Report first block size per Operator
#39656
Changes from all commits
86ca5d9
9dd7d19
36b0bc8
65450df
fe35502
fa7c34e
b95bcf3
5cc6f37
632917a
3353451
e933602
e7db8e4
955819c
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 |
---|---|---|
|
@@ -9,15 +9,14 @@ | |
import ray | ||
from ray.data._internal.block_batching.interfaces import Batch | ||
from ray.data._internal.block_batching.util import ( | ||
Queue, | ||
_calculate_ref_hits, | ||
blocks_to_batches, | ||
collate, | ||
finalize_batches, | ||
format_batches, | ||
resolve_block_refs, | ||
) | ||
from ray.data._internal.util import make_async_gen | ||
from ray.data._internal.util import Queue, make_async_gen | ||
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. fixes breaking test import (unrelated to changes in this PR) |
||
|
||
|
||
def block_generator(num_rows: int, num_blocks: int): | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,11 +1,13 @@ | ||
import collections | ||
import time | ||
from unittest.mock import MagicMock | ||
from unittest.mock import MagicMock, patch | ||
|
||
import numpy as np | ||
import pytest | ||
|
||
import ray | ||
from ray._private.test_utils import wait_for_condition | ||
from ray.data._internal.dataset_logger import DatasetLogger | ||
from ray.data._internal.execution.interfaces import ( | ||
ExecutionOptions, | ||
ExecutionResources, | ||
|
@@ -33,6 +35,7 @@ | |
update_operator_states, | ||
) | ||
from ray.data._internal.execution.util import make_ref_bundles | ||
from ray.data.context import DataContext | ||
from ray.data.tests.conftest import * # noqa | ||
from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy | ||
|
||
|
@@ -684,6 +687,56 @@ def test_execution_allowed_nothrottle(): | |
) | ||
|
||
|
||
def test__check_first_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. is it a typo to have 2 underscores? 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. the function name is |
||
# Check that block size of the first block is recorded and reported. | ||
logger = DatasetLogger( | ||
"ray.data._internal.execution.streaming_executor_state" | ||
).get_logger() | ||
with patch.object(logger, "info") as mock_logger: | ||
# Generate 1MB and 2MB blocks. | ||
inputs = make_ref_bundles( | ||
[ | ||
np.zeros(1024 * 1024, dtype=np.uint8), | ||
np.zeros(1024 * 1024 * 2, dtype=np.uint8), | ||
] | ||
) | ||
o1 = InputDataBuffer(inputs) | ||
o2 = MapOperator.create(make_map_transformer(lambda block: block), o1) | ||
|
||
topo, _ = build_streaming_topology(o2, ExecutionOptions()) | ||
process_completed_tasks(topo) | ||
update_operator_states(topo) | ||
|
||
logger_args, logger_kwargs = mock_logger.call_args | ||
assert logger_args[0] == ("Input in-memory block size: 1.00 MB") | ||
|
||
# Test case where the block size is much larger than the target max block size. | ||
ctx = DataContext.get_current() | ||
ctx.target_max_block_size = 1024 * 1024 | ||
|
||
with patch.object(logger, "warning") as mock_logger: | ||
# Generate a 4 MB block, more than 2 times larger than | ||
# the configured max block size of 1 MB. | ||
inputs = make_ref_bundles( | ||
[ | ||
np.zeros(1024 * 1024 * 4, dtype=np.uint8), | ||
] | ||
) | ||
o1 = InputDataBuffer(inputs) | ||
o2 = MapOperator.create(make_map_transformer(lambda block: block), o1) | ||
topo, _ = build_streaming_topology(o2, ExecutionOptions()) | ||
process_completed_tasks(topo) | ||
update_operator_states(topo) | ||
|
||
logger_args, logger_kwargs = mock_logger.call_args | ||
assert logger_args[0].startswith( | ||
( | ||
"Input in-memory block size of 4.00 MB is significantly larger " | ||
"than the maximium target block size of 1.00 MB." | ||
) | ||
) | ||
|
||
|
||
if __name__ == "__main__": | ||
import sys | ||
|
||
|
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.
I'm still unclear how useful this would be. 1) in which case the block size would be larger than
target_max_block_size
? It shouldn't happen unless there is a bug? 2) printing this as an info may be too verbose, and may cause confusions to the users.@ericl @scottjlee could one of you elaborate on what exact issue this is trying to solve?
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.
@ericl what was the original scenario where this issue appeared? The block splitting should be ensuring we don't get block sizes larger than ``target_max_block_size`, are there operators where this is not being handled?