wehub-resource-sync
112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
import logging
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data.context import DataContext, ShuffleStrategy
|
|
from ray.data.dataset import Dataset
|
|
|
|
SHUFFLE_ALL_TO_ALL_OPS = [
|
|
Dataset.random_shuffle,
|
|
lambda ds: ds.sort(key="id"),
|
|
lambda ds: ds.groupby("id").map_groups(lambda group: group),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"shuffle_op",
|
|
SHUFFLE_ALL_TO_ALL_OPS,
|
|
)
|
|
def test_debug_limit_shuffle_execution_to_num_blocks(
|
|
ray_start_regular, restore_data_context, configure_shuffle_method, shuffle_op
|
|
):
|
|
if configure_shuffle_method == ShuffleStrategy.HASH_SHUFFLE:
|
|
pytest.skip("Not supported by hash-shuffle")
|
|
|
|
shuffle_fn = shuffle_op
|
|
|
|
parallelism = 100
|
|
ds = ray.data.range(1000, override_num_blocks=parallelism)
|
|
shuffled_ds = shuffle_fn(ds).materialize()
|
|
shuffled_ds = shuffled_ds.materialize()
|
|
assert shuffled_ds._logical_plan.initial_num_blocks() == parallelism
|
|
|
|
ds.context.set_config("debug_limit_shuffle_execution_to_num_blocks", 1)
|
|
shuffled_ds = shuffle_fn(ds).materialize()
|
|
shuffled_ds = shuffled_ds.materialize()
|
|
assert shuffled_ds._logical_plan.initial_num_blocks() == 1
|
|
|
|
|
|
@pytest.mark.parametrize("under_threshold", [False, True])
|
|
def test_sort_object_ref_warnings(
|
|
ray_start_regular,
|
|
restore_data_context,
|
|
configure_shuffle_method,
|
|
under_threshold,
|
|
propagate_logs,
|
|
caplog,
|
|
):
|
|
# Test that we warn iff expected driver memory usage from
|
|
# storing ObjectRefs is higher than the configured
|
|
# threshold.
|
|
warning_str = "Execution is estimated to use"
|
|
warning_str_with_bytes = (
|
|
"Execution is estimated to use at least "
|
|
f"{90 if configure_shuffle_method == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED else 300}KB"
|
|
)
|
|
|
|
if not under_threshold:
|
|
DataContext.get_current().warn_on_driver_memory_usage_bytes = 10_000
|
|
|
|
ds = ray.data.range(int(1e8), override_num_blocks=10)
|
|
with caplog.at_level(logging.WARNING, logger="ray.data.dataset"):
|
|
ds = ds.random_shuffle().materialize()
|
|
|
|
if under_threshold:
|
|
assert warning_str not in caplog.text
|
|
assert warning_str_with_bytes not in caplog.text
|
|
else:
|
|
assert warning_str in caplog.text
|
|
assert warning_str_with_bytes in caplog.text
|
|
|
|
|
|
@pytest.mark.parametrize("under_threshold", [False, True])
|
|
def test_sort_inlined_objects_warnings(
|
|
ray_start_regular,
|
|
restore_data_context,
|
|
configure_shuffle_method,
|
|
under_threshold,
|
|
propagate_logs,
|
|
caplog,
|
|
):
|
|
# Test that we warn iff expected driver memory usage from
|
|
# storing tiny Ray objects on driver heap is higher than
|
|
# the configured threshold.
|
|
if configure_shuffle_method == ShuffleStrategy.SORT_SHUFFLE_PUSH_BASED:
|
|
warning_strs = [
|
|
"More than 3MB of driver memory used",
|
|
"More than 7MB of driver memory used",
|
|
]
|
|
else:
|
|
warning_strs = [
|
|
"More than 8MB of driver memory used",
|
|
]
|
|
|
|
if not under_threshold:
|
|
DataContext.get_current().warn_on_driver_memory_usage_bytes = 3_000_000
|
|
|
|
ds = ray.data.range(int(1e6), override_num_blocks=10)
|
|
with caplog.at_level(logging.WARNING, logger="ray.data.dataset"):
|
|
ds = ds.random_shuffle().materialize()
|
|
|
|
if under_threshold:
|
|
assert all(warning_str not in caplog.text for warning_str in warning_strs)
|
|
else:
|
|
assert all(warning_str in caplog.text for warning_str in warning_strs)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|