wehub-resource-sync
1224 lines
45 KiB
Python
1224 lines
45 KiB
Python
from dataclasses import dataclass
|
|
from typing import Any, Dict, Optional
|
|
from unittest.mock import MagicMock, call, patch
|
|
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data import DataContext, ExecutionResources
|
|
from ray.data._internal.execution.interfaces import (
|
|
BlockEntry,
|
|
PhysicalOperator,
|
|
RefBundle,
|
|
)
|
|
from ray.data._internal.execution.operators.hash_aggregate import HashAggregateOperator
|
|
from ray.data._internal.execution.operators.hash_shuffle import HashShuffleOperator
|
|
from ray.data._internal.execution.operators.join import JoinOperator
|
|
from ray.data._internal.logical.interfaces import LogicalOperator
|
|
from ray.data._internal.logical.operators import JoinType
|
|
from ray.data._internal.planner.exchange.sort_task_spec import SortKey
|
|
from ray.data._internal.util import GiB, MiB
|
|
from ray.data.aggregate import AggregateFnV2, Count, Sum
|
|
from ray.data.block import BlockAccessor, BlockMetadata
|
|
|
|
|
|
def _create_aggregator_pool_for_test(op, estimated_dataset_bytes: Optional[int]):
|
|
pool = op._create_aggregator_pool(
|
|
estimated_dataset_bytes=estimated_dataset_bytes,
|
|
)
|
|
op._aggregator_pool = pool
|
|
return pool
|
|
|
|
|
|
def _make_ref_bundle(size_bytes: int) -> RefBundle:
|
|
block = pa.table({"id": [1]})
|
|
return RefBundle(
|
|
[
|
|
BlockEntry(
|
|
ray.put(block),
|
|
BlockMetadata(
|
|
num_rows=1,
|
|
size_bytes=size_bytes,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
),
|
|
)
|
|
],
|
|
schema=block.schema,
|
|
owns_blocks=False,
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class JoinTestCase:
|
|
# Expected outputs
|
|
expected_ray_remote_args: Dict[str, Any]
|
|
expected_num_partitions: int
|
|
expected_num_aggregators: int
|
|
|
|
# Input dataset configurations
|
|
left_size_bytes: Optional[int]
|
|
right_size_bytes: Optional[int]
|
|
left_num_blocks: Optional[int]
|
|
right_num_blocks: Optional[int]
|
|
|
|
# Join configuration
|
|
target_num_partitions: Optional[int]
|
|
|
|
# Cluster resources (for testing different resource scenarios)
|
|
total_cpu: float = 4.0
|
|
total_memory: int = 32 * GiB
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"tc",
|
|
[
|
|
# Case 1: Auto-derived partitions with limited CPUs
|
|
JoinTestCase(
|
|
left_size_bytes=1 * GiB,
|
|
right_size_bytes=2 * GiB,
|
|
left_num_blocks=10,
|
|
right_num_blocks=5,
|
|
target_num_partitions=None, # Auto-derive
|
|
total_cpu=4.0,
|
|
expected_num_partitions=10, # max(10, 5)
|
|
expected_num_aggregators=4, # min(10 partitions, 4 CPUs) = 4
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 3, # ceil(10 partitions / 4 aggregators)
|
|
"num_cpus": 0.25, # 4 CPUs * 25% / 4 aggregators
|
|
"memory": 1771674012,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 2: Single partition (much higher memory overhead)
|
|
JoinTestCase(
|
|
left_size_bytes=1 * GiB,
|
|
right_size_bytes=1 * GiB,
|
|
left_num_blocks=10,
|
|
right_num_blocks=10,
|
|
target_num_partitions=1,
|
|
total_cpu=4.0,
|
|
expected_num_partitions=1,
|
|
expected_num_aggregators=1, # min(1 partition, 4 CPUs) = 1
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 1,
|
|
"num_cpus": 1.0, # 4 CPUs * 25% / 1 aggregator
|
|
"memory": 8589934592,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 3: Limited CPU resources affecting num_cpus calculation
|
|
JoinTestCase(
|
|
left_size_bytes=2 * GiB,
|
|
right_size_bytes=2 * GiB,
|
|
left_num_blocks=20,
|
|
right_num_blocks=20,
|
|
target_num_partitions=40,
|
|
total_cpu=2.0, # Only 2 CPUs available
|
|
expected_num_partitions=40,
|
|
expected_num_aggregators=2, # min(40 partitions, 2 CPUs) = 2
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 8, # min(ceil(40/2), 8) = 8
|
|
"num_cpus": 0.25, # 2 CPUs * 25% / 2 aggregators
|
|
"memory": 2469606197,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 4: Testing with many CPUs and partitions
|
|
JoinTestCase(
|
|
left_size_bytes=10 * GiB,
|
|
right_size_bytes=10 * GiB,
|
|
left_num_blocks=100,
|
|
right_num_blocks=100,
|
|
target_num_partitions=100,
|
|
total_cpu=32.0,
|
|
expected_num_partitions=100,
|
|
expected_num_aggregators=32, # min(100 partitions, 32 CPUs)
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 4, # ceil(100 / 32)
|
|
"num_cpus": 0.25, # 32 CPUs * 25% / 32 aggregators
|
|
"memory": 1315333735,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 5: Testing max aggregators cap (128 default)
|
|
JoinTestCase(
|
|
left_size_bytes=50 * GiB,
|
|
right_size_bytes=50 * GiB,
|
|
left_num_blocks=200,
|
|
right_num_blocks=200,
|
|
target_num_partitions=200,
|
|
total_cpu=256.0, # Many CPUs
|
|
expected_num_partitions=200,
|
|
expected_num_aggregators=128, # min(200, min(256, 128 (default max))
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 2, # ceil(200 / 128)
|
|
"num_cpus": 0.5, # 256 CPUs * 25% / 128 aggregators
|
|
"memory": 2449473536,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 6: Testing num_cpus derived from memory allocation
|
|
JoinTestCase(
|
|
left_size_bytes=50 * GiB,
|
|
right_size_bytes=50 * GiB,
|
|
left_num_blocks=200,
|
|
right_num_blocks=200,
|
|
target_num_partitions=None,
|
|
total_cpu=1024, # Many CPUs
|
|
expected_num_partitions=200,
|
|
expected_num_aggregators=128, # min(200, min(1000, 128 (default max))
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 2, # ceil(200 / 128)
|
|
"num_cpus": 0.57, # ~2.5Gb / 4Gb = ~0.57
|
|
"memory": 2449473536,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 7: No dataset size estimates available (fallback to default memory request)
|
|
# Memory calculation (fallback):
|
|
# max_mem_per_agg = 32 GiB / 32 = 1 GiB
|
|
# modest_mem = 1 GiB / 2 = 512 MiB
|
|
# memory = min(512 MiB, DEFAULT_1GiB) = 512 MiB = 536870912
|
|
# CPU calculation:
|
|
# cap = min(4.0, 32.0 * 0.25 / 32) = 0.25
|
|
# target = min(0.25, 536870912 / 4 GiB) = 0.12
|
|
JoinTestCase(
|
|
left_size_bytes=None,
|
|
right_size_bytes=None,
|
|
left_num_blocks=None,
|
|
right_num_blocks=None,
|
|
target_num_partitions=None,
|
|
total_cpu=32,
|
|
expected_num_partitions=200, # default parallelism
|
|
expected_num_aggregators=32, # min(200, min(1000, 128 (default max))
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 7, # ceil(200 / 32)
|
|
"num_cpus": 0.12,
|
|
"memory": 536870912,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
],
|
|
)
|
|
def test_join_aggregator_remote_args(
|
|
ray_start_regular,
|
|
tc,
|
|
):
|
|
"""Test that join operator correctly estimates memory, CPU, and other resources
|
|
for Aggregator actors based on dataset size estimates as well as cluster resources.
|
|
"""
|
|
|
|
left_logical_op_mock = MagicMock(LogicalOperator)
|
|
left_logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=tc.left_size_bytes,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
left_logical_op_mock.estimated_num_outputs.return_value = tc.left_num_blocks
|
|
|
|
left_op_mock = MagicMock(PhysicalOperator)
|
|
left_op_mock._output_dependencies = []
|
|
left_op_mock._logical_operators = [left_logical_op_mock]
|
|
left_op_mock.num_output_splits.return_value = 1
|
|
|
|
right_logical_op_mock = MagicMock(LogicalOperator)
|
|
right_logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=tc.right_size_bytes,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
right_logical_op_mock.estimated_num_outputs.return_value = tc.right_num_blocks
|
|
|
|
right_op_mock = MagicMock(PhysicalOperator)
|
|
right_op_mock._output_dependencies = []
|
|
right_op_mock._logical_operators = [right_logical_op_mock]
|
|
right_op_mock.num_output_splits.return_value = 1
|
|
|
|
# Patch the total cluster resources
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle.ray.cluster_resources",
|
|
return_value={"CPU": tc.total_cpu, "memory": tc.total_memory},
|
|
):
|
|
# Create the join operator
|
|
op = JoinOperator(
|
|
left_input_op=left_op_mock,
|
|
right_input_op=right_op_mock,
|
|
data_context=DataContext.get_current(),
|
|
left_key_columns=("id",),
|
|
right_key_columns=("id",),
|
|
join_type=JoinType.INNER,
|
|
num_partitions=tc.target_num_partitions,
|
|
)
|
|
|
|
# Validate the estimations
|
|
assert op._num_partitions == tc.expected_num_partitions
|
|
assert op._aggregator_pool is None
|
|
|
|
estimated_dataset_bytes = (
|
|
tc.left_size_bytes + tc.right_size_bytes
|
|
if tc.left_size_bytes is not None and tc.right_size_bytes is not None
|
|
else None
|
|
)
|
|
pool = _create_aggregator_pool_for_test(op, estimated_dataset_bytes)
|
|
|
|
assert pool.num_aggregators == tc.expected_num_aggregators
|
|
assert pool._aggregator_ray_remote_args == tc.expected_ray_remote_args
|
|
|
|
|
|
@dataclass
|
|
class HashOperatorTestCase:
|
|
# Expected outputs
|
|
expected_ray_remote_args: Dict[str, Any]
|
|
expected_num_partitions: int
|
|
expected_num_aggregators: int
|
|
# Input dataset configuration
|
|
input_size_bytes: Optional[int]
|
|
input_num_blocks: Optional[int]
|
|
# Operator configuration
|
|
target_num_partitions: Optional[int]
|
|
# Cluster resources (for testing different resource scenarios)
|
|
total_cpu: float = 4.0
|
|
total_memory: int = 32 * GiB
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"tc",
|
|
[
|
|
# Case 1: Auto-derived partitions with limited CPUs
|
|
HashOperatorTestCase(
|
|
input_size_bytes=2 * GiB,
|
|
input_num_blocks=16,
|
|
target_num_partitions=None,
|
|
total_cpu=4.0,
|
|
expected_num_partitions=16,
|
|
expected_num_aggregators=4,
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 4,
|
|
# estimate-derived 640MiB is floored up to the modest default:
|
|
# min(32GiB / 4 / 2, DEFAULT_ALLOCATION=1GiB) = 1GiB
|
|
"num_cpus": 0.25, # min(min(4, 4*0.25/4), 1GiB / 4GiB) = 0.25
|
|
"memory": 1073741824,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 2: Single partition produced
|
|
HashOperatorTestCase(
|
|
input_size_bytes=512 * MiB,
|
|
input_num_blocks=8,
|
|
target_num_partitions=1,
|
|
total_cpu=8.0,
|
|
expected_num_partitions=1,
|
|
expected_num_aggregators=1,
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 1,
|
|
"num_cpus": 0.25,
|
|
"memory": 1073741824,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 3: Many CPUs
|
|
HashOperatorTestCase(
|
|
input_size_bytes=16 * GiB,
|
|
input_num_blocks=128,
|
|
target_num_partitions=32,
|
|
total_cpu=256.0,
|
|
expected_num_partitions=32,
|
|
expected_num_aggregators=32,
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 1,
|
|
"num_cpus": 0.25,
|
|
"memory": 1073741824,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 4: Testing num_cpus derived from memory allocation
|
|
HashOperatorTestCase(
|
|
input_size_bytes=50 * GiB,
|
|
input_num_blocks=200,
|
|
target_num_partitions=None,
|
|
total_cpu=1024, # Many CPUs
|
|
expected_num_partitions=200,
|
|
expected_num_aggregators=128, # min(200, min(1000, 128 (default max))
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 2, # ceil(200 / 128)
|
|
"num_cpus": 0.16, # ~0.6Gb / 4Gb = ~0.16
|
|
"memory": 687865856,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 6: No dataset size estimate inferred (fallback to default memory request)
|
|
# Memory calculation (fallback):
|
|
# max_mem_per_agg = 32 GiB / 32 = 1 GiB
|
|
# modest_mem = 1 GiB / 2 = 512 MiB
|
|
# memory = min(512 MiB, DEFAULT_1GiB) = 512 MiB = 536870912
|
|
# CPU calculation:
|
|
# cap = min(4.0, 32.0 * 0.25 / 32) = 0.25
|
|
# target = min(0.25, 536870912 / 4 GiB) = 0.12
|
|
HashOperatorTestCase(
|
|
input_size_bytes=None,
|
|
input_num_blocks=None,
|
|
target_num_partitions=None,
|
|
total_cpu=32.0,
|
|
expected_num_partitions=200,
|
|
expected_num_aggregators=32,
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 7,
|
|
"num_cpus": 0.12,
|
|
"memory": 536870912,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
],
|
|
)
|
|
def test_hash_aggregate_operator_remote_args(
|
|
ray_start_regular,
|
|
tc,
|
|
):
|
|
"""Test that HashAggregateOperator correctly estimates memory, CPU, and other resources
|
|
for aggregator actors based on dataset size estimates as well as cluster resources.
|
|
"""
|
|
logical_op_mock = MagicMock(LogicalOperator)
|
|
logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=tc.input_size_bytes,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
logical_op_mock.estimated_num_outputs.return_value = tc.input_num_blocks
|
|
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_op_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
|
|
# Create some test aggregation functions
|
|
agg_fns = [Sum("value"), Count()]
|
|
|
|
# Patch the total cluster resources
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle.ray.cluster_resources",
|
|
return_value={"CPU": tc.total_cpu, "memory": tc.total_memory},
|
|
):
|
|
# Create the hash aggregate operator
|
|
op = HashAggregateOperator(
|
|
input_op=op_mock,
|
|
data_context=DataContext.get_current(),
|
|
aggregation_fns=agg_fns,
|
|
key_columns=("id",),
|
|
num_partitions=tc.target_num_partitions,
|
|
)
|
|
|
|
# Validate the estimations
|
|
assert op._num_partitions == tc.expected_num_partitions
|
|
assert op._aggregator_pool is None
|
|
|
|
pool = _create_aggregator_pool_for_test(op, tc.input_size_bytes)
|
|
|
|
assert pool.num_aggregators == tc.expected_num_aggregators
|
|
assert pool._aggregator_ray_remote_args == tc.expected_ray_remote_args
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"tc",
|
|
[
|
|
# Case 1: Auto-derived partitions with limited CPUs
|
|
# Memory calculation:
|
|
# max_partitions_per_agg = ceil(16 / 4) = 4
|
|
# partition_size = ceil(2 GiB / 16) = 128 MiB
|
|
# shuffle + output = 2 * (128 MiB * 4) = 1024 MiB
|
|
# with 1.3x skew factor: ceil(1024 MiB * 1.3) = 1395864372
|
|
# CPU calculation:
|
|
# cap = min(4.0, 4.0 * 0.25 / 4) = 0.25
|
|
# target = min(0.25, 1395864372 / 4 GiB) = 0.25
|
|
HashOperatorTestCase(
|
|
input_size_bytes=2 * GiB,
|
|
input_num_blocks=16,
|
|
target_num_partitions=None,
|
|
total_cpu=4.0,
|
|
expected_num_partitions=16,
|
|
expected_num_aggregators=4,
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 4,
|
|
"num_cpus": 0.25,
|
|
"memory": 1395864372,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 2: Single partition produced
|
|
# Memory calculation:
|
|
# max_partitions_per_agg = ceil(1 / 1) = 1
|
|
# partition_size = ceil(512 MiB / 1) = 512 MiB
|
|
# shuffle + output = 2 * (512 MiB * 1) = 1024 MiB
|
|
# with 1.3x skew factor: ceil(1024 MiB * 1.3) = 1395864372
|
|
# CPU calculation:
|
|
# cap = min(4.0, 8.0 * 0.25 / 1) = 2.0
|
|
# target = min(2.0, 1395864372 / 4 GiB) = 0.33
|
|
HashOperatorTestCase(
|
|
input_size_bytes=512 * MiB,
|
|
input_num_blocks=8,
|
|
target_num_partitions=1,
|
|
total_cpu=8.0,
|
|
expected_num_partitions=1,
|
|
expected_num_aggregators=1,
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 1,
|
|
"num_cpus": 0.33,
|
|
"memory": 1395864372,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 3: Many CPUs
|
|
# Memory calculation:
|
|
# max_partitions_per_agg = ceil(32 / 32) = 1
|
|
# partition_size = ceil(16 GiB / 32) = 512 MiB
|
|
# shuffle + output = 2 * (512 MiB * 1) = 1024 MiB
|
|
# with 1.3x skew factor: ceil(1024 MiB * 1.3) = 1395864372
|
|
# CPU calculation:
|
|
# cap = min(4.0, 256.0 * 0.25 / 32) = 2.0
|
|
# target = min(2.0, 1395864372 / 4 GiB) = 0.33
|
|
HashOperatorTestCase(
|
|
input_size_bytes=16 * GiB,
|
|
input_num_blocks=128,
|
|
target_num_partitions=32,
|
|
total_cpu=256.0,
|
|
expected_num_partitions=32,
|
|
expected_num_aggregators=32,
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 1,
|
|
"num_cpus": 0.33,
|
|
"memory": 1395864372,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 4: Testing num_cpus derived from memory allocation
|
|
# Memory calculation:
|
|
# max_partitions_per_agg = ceil(200 / 128) = 2
|
|
# partition_size = ceil(50 GiB / 200) = 256 MiB
|
|
# shuffle + output = 2 * (256 MiB * 2) = 1024 MiB
|
|
# with 1.3x skew factor: ceil(1024 MiB * 1.3) = 1395864372
|
|
# CPU calculation:
|
|
# cap = min(4.0, 1024 * 0.25 / 128) = 2.0
|
|
# target = min(2.0, 1395864372 / 4 GiB) = 0.33
|
|
HashOperatorTestCase(
|
|
input_size_bytes=50 * GiB,
|
|
input_num_blocks=200,
|
|
target_num_partitions=None,
|
|
total_cpu=1024, # Many CPUs
|
|
expected_num_partitions=200,
|
|
expected_num_aggregators=128, # min(200, min(1000, 128 (default max))
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 2, # ceil(200 / 128)
|
|
"num_cpus": 0.33,
|
|
"memory": 1395864372,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
# Case 5: No dataset size estimate inferred (fallback to default memory request)
|
|
# Memory calculation (fallback):
|
|
# max_mem_per_agg = 32 GiB / 32 = 1 GiB
|
|
# modest_mem = 1 GiB / 2 = 512 MiB
|
|
# memory = min(512 MiB, DEFAULT_1GiB) = 512 MiB = 536870912
|
|
# CPU calculation:
|
|
# cap = min(4.0, 32.0 * 0.25 / 32) = 0.25
|
|
# target = min(0.25, 536870912 / 4 GiB) = 0.12
|
|
HashOperatorTestCase(
|
|
input_size_bytes=None,
|
|
input_num_blocks=None,
|
|
target_num_partitions=None,
|
|
total_cpu=32.0,
|
|
expected_num_partitions=200,
|
|
expected_num_aggregators=32,
|
|
expected_ray_remote_args={
|
|
"max_concurrency": 7,
|
|
"num_cpus": 0.12,
|
|
"memory": 536870912,
|
|
"scheduling_strategy": "SPREAD",
|
|
"allow_out_of_order_execution": True,
|
|
},
|
|
),
|
|
],
|
|
)
|
|
def test_hash_shuffle_operator_remote_args(
|
|
ray_start_regular,
|
|
tc,
|
|
):
|
|
"""Test that HashShuffleOperator correctly estimates memory, CPU, and other resources
|
|
for aggregator actors based on dataset size estimates as well as cluster resources.
|
|
"""
|
|
logical_op_mock = MagicMock(LogicalOperator)
|
|
logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=tc.input_size_bytes,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
logical_op_mock.estimated_num_outputs.return_value = tc.input_num_blocks
|
|
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_op_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
|
|
# Patch the total cluster resources
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle.ray.cluster_resources",
|
|
return_value={"CPU": tc.total_cpu, "memory": tc.total_memory},
|
|
):
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle._get_total_cluster_resources"
|
|
) as mock_resources:
|
|
mock_resources.return_value = ExecutionResources(
|
|
cpu=tc.total_cpu, memory=tc.total_memory
|
|
)
|
|
|
|
# Create the hash shuffle operator
|
|
op = HashShuffleOperator(
|
|
input_op=op_mock,
|
|
data_context=DataContext.get_current(),
|
|
key_columns=("id",),
|
|
num_partitions=tc.target_num_partitions,
|
|
)
|
|
|
|
# Validate the estimations
|
|
assert op._num_partitions == tc.expected_num_partitions
|
|
assert op._aggregator_pool is None
|
|
|
|
pool = _create_aggregator_pool_for_test(op, tc.input_size_bytes)
|
|
|
|
assert pool.num_aggregators == tc.expected_num_aggregators
|
|
assert pool._aggregator_ray_remote_args == tc.expected_ray_remote_args
|
|
|
|
|
|
def test_aggregator_ray_remote_args_includes_context_label_selector(
|
|
ray_start_regular, restore_data_context
|
|
):
|
|
"""ExecutionOptions.label_selector should appear on aggregator actor args."""
|
|
DataContext.get_current().execution_options.label_selector = {"subcluster": "train"}
|
|
|
|
logical_op_mock = MagicMock(LogicalOperator)
|
|
logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=2 * GiB,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
logical_op_mock.estimated_num_outputs.return_value = 16
|
|
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_op_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle.ray.cluster_resources",
|
|
return_value={"CPU": 4.0, "memory": 32 * GiB},
|
|
):
|
|
op = HashAggregateOperator(
|
|
input_op=op_mock,
|
|
data_context=DataContext.get_current(),
|
|
aggregation_fns=[Count()],
|
|
key_columns=("id",),
|
|
)
|
|
|
|
pool = _create_aggregator_pool_for_test(op, 2 * GiB)
|
|
|
|
assert pool._aggregator_ray_remote_args["label_selector"] == {"subcluster": "train"}
|
|
|
|
|
|
def test_aggregator_ray_remote_args_partial_override(ray_start_regular):
|
|
"""Test that partial override of aggregator_ray_remote_args retains default values.
|
|
|
|
This tests the behavior where a user provides only some values (e.g., num_cpus)
|
|
in aggregator_ray_remote_args_override, and the system should retain the default
|
|
values for other parameters (e.g., scheduling_strategy, allow_out_of_order_execution).
|
|
"""
|
|
logical_op_mock = MagicMock(LogicalOperator)
|
|
logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=2 * GiB,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
logical_op_mock.estimated_num_outputs.return_value = 16
|
|
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_op_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
|
|
# Patch the total cluster resources
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle.ray.cluster_resources",
|
|
return_value={"CPU": 4.0, "memory": 32 * GiB},
|
|
):
|
|
# Create operator with partial override (only num_cpus)
|
|
op = HashAggregateOperator(
|
|
input_op=op_mock,
|
|
data_context=DataContext.get_current(),
|
|
aggregation_fns=[Count()],
|
|
key_columns=("id",),
|
|
aggregator_ray_remote_args_override={
|
|
"num_cpus": 0.5
|
|
}, # Only override num_cpus
|
|
)
|
|
|
|
pool = _create_aggregator_pool_for_test(op, 2 * GiB)
|
|
|
|
# Verify that num_cpus was overridden
|
|
assert pool._aggregator_ray_remote_args["num_cpus"] == 0.5
|
|
|
|
# Verify that default values are retained
|
|
assert pool._aggregator_ray_remote_args["scheduling_strategy"] == "SPREAD"
|
|
assert pool._aggregator_ray_remote_args["allow_out_of_order_execution"] is True
|
|
|
|
# Verify that max_concurrency is still present
|
|
assert "max_concurrency" in pool._aggregator_ray_remote_args
|
|
|
|
# Verify that memory is still present
|
|
assert "memory" in pool._aggregator_ray_remote_args
|
|
|
|
|
|
def test_hash_shuffle_does_not_infer_metadata_for_memory_during_construction(
|
|
ray_start_regular,
|
|
):
|
|
logical_op_mock = MagicMock(LogicalOperator)
|
|
logical_op_mock.infer_metadata.side_effect = AssertionError(
|
|
"logical metadata should not be used for hash-shuffle memory sizing"
|
|
)
|
|
logical_op_mock.estimated_num_outputs.return_value = 16
|
|
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_op_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle"
|
|
"._get_total_cluster_resources",
|
|
return_value=ExecutionResources(cpu=4.0, memory=32 * GiB),
|
|
):
|
|
op = HashShuffleOperator(
|
|
input_op=op_mock,
|
|
data_context=DataContext.get_current(),
|
|
key_columns=("id",),
|
|
)
|
|
|
|
assert op._aggregator_pool is None
|
|
logical_op_mock.infer_metadata.assert_not_called()
|
|
|
|
|
|
def test_hash_shuffle_buffers_inputs_to_estimate_memory_at_execution(
|
|
ray_start_regular,
|
|
):
|
|
logical_op_mock = MagicMock(LogicalOperator)
|
|
logical_op_mock.estimated_num_outputs.return_value = 16
|
|
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_op_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
op_mock.num_outputs_total.return_value = 4
|
|
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle"
|
|
"._get_total_cluster_resources",
|
|
return_value=ExecutionResources(cpu=4.0, memory=32 * GiB),
|
|
):
|
|
op = HashShuffleOperator(
|
|
input_op=op_mock,
|
|
data_context=DataContext.get_current(),
|
|
key_columns=("id",),
|
|
)
|
|
|
|
first_bundle = _make_ref_bundle(100)
|
|
second_bundle = _make_ref_bundle(200)
|
|
|
|
mock_pool = MagicMock()
|
|
with patch.object(
|
|
op, "_create_aggregator_pool", return_value=mock_pool
|
|
) as create_pool, patch.object(op, "_do_add_input_inner") as replay:
|
|
op._add_input_inner(first_bundle, input_index=0)
|
|
op._add_input_inner(second_bundle, input_index=0)
|
|
|
|
assert not op._shuffle_started
|
|
create_pool.assert_not_called()
|
|
replay.assert_not_called()
|
|
|
|
op.all_inputs_done()
|
|
|
|
create_pool.assert_called_once_with(estimated_dataset_bytes=300)
|
|
mock_pool.start.assert_called_once()
|
|
assert replay.call_args_list == [
|
|
call(first_bundle, 0),
|
|
call(second_bundle, 0),
|
|
]
|
|
|
|
|
|
def test_hash_shuffle_starts_from_empty_buffer_when_inputs_done(
|
|
ray_start_regular,
|
|
):
|
|
logical_op_mock = MagicMock(LogicalOperator)
|
|
logical_op_mock.estimated_num_outputs.return_value = 16
|
|
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_op_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle"
|
|
"._get_total_cluster_resources",
|
|
return_value=ExecutionResources(cpu=4.0, memory=32 * GiB),
|
|
):
|
|
op = HashShuffleOperator(
|
|
input_op=op_mock,
|
|
data_context=DataContext.get_current(),
|
|
key_columns=("id",),
|
|
)
|
|
|
|
mock_pool = MagicMock()
|
|
with patch.object(
|
|
op, "_create_aggregator_pool", return_value=mock_pool
|
|
) as create_pool:
|
|
op.all_inputs_done()
|
|
|
|
create_pool.assert_called_once_with(estimated_dataset_bytes=0)
|
|
mock_pool.start.assert_called_once()
|
|
|
|
|
|
def _make_shuffle_op(upstream_num_outputs):
|
|
"""Build a HashShuffleOperator whose single upstream reports
|
|
``upstream_num_outputs`` bundles (None/0 => unknown)."""
|
|
logical_op_mock = MagicMock(LogicalOperator)
|
|
logical_op_mock.estimated_num_outputs.return_value = 16
|
|
# No logical size estimate by default (mirrors DataSource V2); tests that
|
|
# exercise the logical-estimate path override this.
|
|
logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=None,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_op_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
op_mock.num_outputs_total.return_value = upstream_num_outputs
|
|
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle"
|
|
"._get_total_cluster_resources",
|
|
return_value=ExecutionResources(cpu=4.0, memory=32 * GiB),
|
|
):
|
|
op = HashShuffleOperator(
|
|
input_op=op_mock,
|
|
data_context=DataContext.get_current(),
|
|
key_columns=("id",),
|
|
)
|
|
return op, logical_op_mock
|
|
|
|
|
|
def test_hash_shuffle_sample_window_trips_on_max_bundles_when_count_unknown(
|
|
ray_start_regular,
|
|
):
|
|
# Upstream output count never materializes -> defer to MAX_BUNDLES.
|
|
op, _ = _make_shuffle_op(upstream_num_outputs=None)
|
|
|
|
mock_pool = MagicMock()
|
|
with patch.object(
|
|
op, "_create_aggregator_pool", return_value=mock_pool
|
|
) as create_pool, patch.object(op, "_do_add_input_inner") as replay:
|
|
# First MAX_BUNDLES - 1 bundles only buffer; none trip the window.
|
|
for _ in range(op._MEMORY_ESTIMATION_SAMPLE_MAX_BUNDLES - 1):
|
|
op._add_input_inner(_make_ref_bundle(10), input_index=0)
|
|
create_pool.assert_not_called()
|
|
assert not op._shuffle_started
|
|
|
|
# The MAX_BUNDLES-th bundle trips the window mid-stream.
|
|
op._add_input_inner(_make_ref_bundle(10), input_index=0)
|
|
assert op._shuffle_started
|
|
create_pool.assert_called_once()
|
|
# All buffered bundles are replayed exactly once.
|
|
assert replay.call_count == op._MEMORY_ESTIMATION_SAMPLE_MAX_BUNDLES
|
|
|
|
# Subsequent bundles stream straight through (not buffered).
|
|
op._add_input_inner(_make_ref_bundle(10), input_index=0)
|
|
assert replay.call_count == op._MEMORY_ESTIMATION_SAMPLE_MAX_BUNDLES + 1
|
|
assert not op._buffered_input_bundles[0]
|
|
|
|
|
|
def test_hash_shuffle_sample_window_trips_on_byte_ceiling(ray_start_regular):
|
|
op, _ = _make_shuffle_op(upstream_num_outputs=10_000)
|
|
|
|
mock_pool = MagicMock()
|
|
with patch.object(
|
|
op, "_create_aggregator_pool", return_value=mock_pool
|
|
) as create_pool, patch.object(op, "_do_add_input_inner"):
|
|
# A single bundle at/above the byte ceiling trips the window
|
|
# immediately, regardless of bundle count.
|
|
op._add_input_inner(_make_ref_bundle(op._sample_byte_limit), input_index=0)
|
|
|
|
assert op._shuffle_started
|
|
create_pool.assert_called_once()
|
|
|
|
|
|
def test_hash_shuffle_extrapolates_dataset_bytes_from_sample(ray_start_regular):
|
|
op, _ = _make_shuffle_op(upstream_num_outputs=100)
|
|
# 8 bundles averaging 200 bytes, inputs still streaming.
|
|
op._sample_bytes = 1600
|
|
op._sample_bundles = 8
|
|
assert not op._inputs_complete
|
|
|
|
# avg_bytes_per_bundle (200) * max(upstream=100, sample=8) = 20_000
|
|
assert op._extrapolate_dataset_bytes() == 20_000
|
|
|
|
|
|
def test_hash_shuffle_extrapolation_never_below_observed(ray_start_regular):
|
|
# Upstream reports fewer bundles than we've already sampled -> use sample.
|
|
op, _ = _make_shuffle_op(upstream_num_outputs=2)
|
|
op._sample_bytes = 1000
|
|
op._sample_bundles = 10
|
|
assert not op._inputs_complete
|
|
|
|
# avg (100) * max(upstream=2, sample=10) = 1000 (not 200)
|
|
assert op._extrapolate_dataset_bytes() == 1000
|
|
|
|
|
|
def test_hash_shuffle_extrapolation_falls_back_to_logical_estimate(
|
|
ray_start_regular,
|
|
):
|
|
# Upstream count unknown (DataSource-V2-like) -> logical fallback.
|
|
op, logical_op_mock = _make_shuffle_op(upstream_num_outputs=0)
|
|
logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=4242,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
op._sample_bytes = 500
|
|
op._sample_bundles = 4
|
|
assert not op._inputs_complete
|
|
|
|
assert op._extrapolate_dataset_bytes() == 4242
|
|
|
|
|
|
def test_hash_shuffle_extrapolation_modest_default_when_no_signal(
|
|
ray_start_regular,
|
|
):
|
|
# Upstream count unknown AND logical estimate unavailable, and nothing has
|
|
# been buffered yet -> None, so the remote-args builder falls back to its
|
|
# modest default.
|
|
op, logical_op_mock = _make_shuffle_op(upstream_num_outputs=0)
|
|
logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=None,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
op._sample_bytes = 0
|
|
op._sample_bundles = 0
|
|
assert not op._inputs_complete
|
|
|
|
assert op._extrapolate_dataset_bytes() is None
|
|
|
|
|
|
def test_hash_shuffle_extrapolation_floors_at_observed_bytes_without_count(
|
|
ray_start_regular,
|
|
):
|
|
# Regression: the bounded sample window can trip before any upstream output
|
|
# count materializes (e.g. DataSource V2), so neither the sampled-ratio nor
|
|
# the logical estimate is available. We must still size off the bytes already
|
|
# buffered rather than collapsing to the modest default.
|
|
op, logical_op_mock = _make_shuffle_op(upstream_num_outputs=0)
|
|
logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=None,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
op._sample_bytes = 500
|
|
op._sample_bundles = 4
|
|
assert not op._inputs_complete
|
|
|
|
assert op._extrapolate_dataset_bytes() == 500
|
|
|
|
|
|
def test_hash_shuffle_prefers_logical_estimate_when_sample_underestimates(
|
|
ray_start_regular,
|
|
):
|
|
# Regression: a shuffle fed by a read whose runtime output count is still a
|
|
# large under-estimate when sampled. The accurate logical estimate must win
|
|
# so aggregators aren't under-sized (cf. the aggregate_groups timeout).
|
|
op, logical_op_mock = _make_shuffle_op(upstream_num_outputs=10)
|
|
logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=50 * GiB,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
# Small early sample -> sampled estimate ~= 100 bytes * max(10, 4) = 1000.
|
|
op._sample_bytes = 400
|
|
op._sample_bundles = 4
|
|
assert not op._inputs_complete
|
|
|
|
# max(sampled=1000, logical=50 GiB) -> logical.
|
|
assert op._extrapolate_dataset_bytes() == 50 * GiB
|
|
|
|
|
|
def test_hash_shuffle_floors_aggregator_memory_at_modest_default(
|
|
ray_start_regular,
|
|
):
|
|
# Regression: a tiny dataset estimate (e.g. a chained join whose upstream is
|
|
# back-loaded and sampled at ~0 bytes) must not size aggregators below the
|
|
# modest default (cf. the TPC-H Q3 0.0GiB -> 11.5MiB OOM).
|
|
op, _ = _make_shuffle_op(upstream_num_outputs=16)
|
|
|
|
cluster = ExecutionResources(cpu=4.0, memory=32 * GiB)
|
|
args = op._get_default_aggregator_ray_remote_args(
|
|
num_partitions=16,
|
|
num_aggregators=16,
|
|
total_available_cluster_resources=cluster,
|
|
estimated_dataset_bytes=1 * MiB, # absurdly small
|
|
)
|
|
|
|
# modest default = min(32GiB / 16 / 2, DEFAULT_ALLOCATION=1GiB) = 1 GiB
|
|
assert args["memory"] == 1 * GiB
|
|
|
|
|
|
def test_hash_shuffle_partition_size_hint_skips_sampling(ray_start_regular):
|
|
op, _ = _make_shuffle_op(upstream_num_outputs=10_000)
|
|
# Simulate a partition-size hint (not exposed on HashShuffleOperator ctor).
|
|
op._partition_size_hint = 1234
|
|
|
|
mock_pool = MagicMock()
|
|
with patch.object(
|
|
op, "_create_aggregator_pool", return_value=mock_pool
|
|
) as create_pool, patch.object(op, "_do_add_input_inner"):
|
|
# The very first bundle trips the window since the estimate is exact.
|
|
op._add_input_inner(_make_ref_bundle(10), input_index=0)
|
|
|
|
assert op._shuffle_started
|
|
create_pool.assert_called_once_with(
|
|
estimated_dataset_bytes=1234 * op._num_partitions
|
|
)
|
|
|
|
|
|
def test_hash_shuffle_does_not_double_start(ray_start_regular):
|
|
# Window trips mid-stream, then all_inputs_done fires: pool created once.
|
|
op, _ = _make_shuffle_op(upstream_num_outputs=None)
|
|
|
|
mock_pool = MagicMock()
|
|
with patch.object(
|
|
op, "_create_aggregator_pool", return_value=mock_pool
|
|
) as create_pool, patch.object(op, "_do_add_input_inner"):
|
|
for _ in range(op._MEMORY_ESTIMATION_SAMPLE_MAX_BUNDLES):
|
|
op._add_input_inner(_make_ref_bundle(10), input_index=0)
|
|
assert op._shuffle_started
|
|
create_pool.assert_called_once()
|
|
|
|
op.all_inputs_done()
|
|
|
|
create_pool.assert_called_once()
|
|
mock_pool.start.assert_called_once()
|
|
|
|
|
|
def test_hash_shuffle_aggregate_sampling_across_input_sequences(ray_start_regular):
|
|
# Sampling is aggregated across input sequences (e.g. joins): the byte
|
|
# ceiling accounts for bundles arriving on multiple input indices.
|
|
def _join_input_mock():
|
|
logical_op_mock = MagicMock(LogicalOperator)
|
|
logical_op_mock.estimated_num_outputs.return_value = 16
|
|
logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=None,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_op_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
op_mock.num_outputs_total.return_value = 10_000
|
|
return op_mock
|
|
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle.ray.cluster_resources",
|
|
return_value={"CPU": 4.0, "memory": 32 * GiB},
|
|
):
|
|
op = JoinOperator(
|
|
left_input_op=_join_input_mock(),
|
|
right_input_op=_join_input_mock(),
|
|
data_context=DataContext.get_current(),
|
|
left_key_columns=("id",),
|
|
right_key_columns=("id",),
|
|
join_type=JoinType.INNER,
|
|
num_partitions=16,
|
|
)
|
|
|
|
half = op._sample_byte_limit // 2
|
|
mock_pool = MagicMock()
|
|
with patch.object(
|
|
op, "_create_aggregator_pool", return_value=mock_pool
|
|
) as create_pool, patch.object(op, "_do_add_input_inner"):
|
|
op._add_input_inner(_make_ref_bundle(half), input_index=0)
|
|
create_pool.assert_not_called()
|
|
# A bundle on the *other* input sequence pushes the aggregate sample
|
|
# over the ceiling, proving the sample is shared across sequences.
|
|
op._add_input_inner(_make_ref_bundle(half + 1), input_index=1)
|
|
|
|
assert op._shuffle_started
|
|
create_pool.assert_called_once()
|
|
|
|
|
|
def _make_join_op(left_num_outputs, right_num_outputs):
|
|
"""Build a JoinOperator whose two upstreams report the given output-bundle
|
|
counts (None/0 => unknown)."""
|
|
|
|
def _join_input_mock(num_outputs):
|
|
logical_op_mock = MagicMock(LogicalOperator)
|
|
logical_op_mock.estimated_num_outputs.return_value = 16
|
|
# No logical size estimate (mirrors DataSource V2), so the sampled
|
|
# estimate is what `_extrapolate_dataset_bytes` returns.
|
|
logical_op_mock.infer_metadata.return_value = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=None,
|
|
exec_stats=None,
|
|
input_files=None,
|
|
)
|
|
op_mock = MagicMock(PhysicalOperator)
|
|
op_mock._output_dependencies = []
|
|
op_mock._logical_operators = [logical_op_mock]
|
|
op_mock.num_output_splits.return_value = 1
|
|
op_mock.num_outputs_total.return_value = num_outputs
|
|
return op_mock
|
|
|
|
with patch(
|
|
"ray.data._internal.execution.operators.hash_shuffle.ray.cluster_resources",
|
|
return_value={"CPU": 4.0, "memory": 32 * GiB},
|
|
):
|
|
op = JoinOperator(
|
|
left_input_op=_join_input_mock(left_num_outputs),
|
|
right_input_op=_join_input_mock(right_num_outputs),
|
|
data_context=DataContext.get_current(),
|
|
left_key_columns=("id",),
|
|
right_key_columns=("id",),
|
|
join_type=JoinType.INNER,
|
|
num_partitions=16,
|
|
)
|
|
return op
|
|
|
|
|
|
def test_hash_shuffle_extrapolates_per_input_sequence(ray_start_regular):
|
|
# Each input is extrapolated against its OWN output count using its OWN
|
|
# sampled per-bundle average, so a sample dominated by one (light) input
|
|
# doesn't drag down the estimate for the other (heavy) input.
|
|
op = _make_join_op(left_num_outputs=100, right_num_outputs=100)
|
|
# input 0: 8 bundles @ 200 bytes/bundle. input 1: 2 bundles @ 2000 bytes.
|
|
op._sample_bytes_by_input[0] = 1600
|
|
op._sample_bundles_by_input[0] = 8
|
|
op._sample_bytes_by_input[1] = 4000
|
|
op._sample_bundles_by_input[1] = 2
|
|
op._sample_bytes = 5600
|
|
op._sample_bundles = 10
|
|
assert not op._inputs_complete
|
|
|
|
# Per-input: 200 * 100 + 2000 * 100 = 220_000.
|
|
# (A global average of 560 bytes/bundle would give only 560 * 200 = 112_000,
|
|
# badly under-sizing the pool for the heavy right input.)
|
|
assert op._extrapolate_dataset_bytes() == 220_000
|
|
|
|
|
|
def test_hash_shuffle_unsampled_input_falls_back_to_global_average(
|
|
ray_start_regular,
|
|
):
|
|
# An input with a known output count but no sampled bundles yet uses the
|
|
# global average bundle size rather than being dropped from the estimate.
|
|
op = _make_join_op(left_num_outputs=100, right_num_outputs=50)
|
|
# Only input 0 has been sampled: 8 bundles @ 200 bytes.
|
|
op._sample_bytes_by_input[0] = 1600
|
|
op._sample_bundles_by_input[0] = 8
|
|
op._sample_bytes = 1600
|
|
op._sample_bundles = 8
|
|
assert not op._inputs_complete
|
|
|
|
# input 0: 200 * 100 = 20_000; input 1 (unsampled): global avg 200 * 50 =
|
|
# 10_000. Total = 30_000.
|
|
assert op._extrapolate_dataset_bytes() == 30_000
|
|
|
|
|
|
def test_partial_aggregate_preserves_sort_after_builder_compaction(
|
|
ray_start_regular,
|
|
monkeypatch,
|
|
):
|
|
"""Regression test for HashAggregate producing duplicate group rows when
|
|
`TableBlockBuilder.build()` reorders rows across an internal compaction.
|
|
|
|
For an `AggregateFnV2` whose accumulator can vary in size between groups
|
|
(here, an empty list vs. a non-empty list), the per-block partial-aggregate
|
|
output is built by `_aggregate` row-by-row via the table builder. If
|
|
`_compact_if_needed` triggered mid-loop, the legacy `build()` placed the
|
|
still-uncompacted (newest) rows in front of the compacted (older) tables,
|
|
breaking the "blocks reaching `_combine_aggregated_blocks` are sorted by
|
|
key" precondition that `heapq.merge` relies on. That precondition violation
|
|
surfaced as duplicate group rows whose count varied with the parallelism
|
|
arg (since parallelism changes per-block row count, and therefore whether
|
|
compaction triggers inside the partial-aggregate loop).
|
|
|
|
We force compaction on every row via `MAX_UNCOMPACTED_SIZE_BYTES=1` and
|
|
assert that the partial-aggregate output is still sorted by the group key.
|
|
"""
|
|
import ray.data._internal.table_block as table_block
|
|
|
|
class EmptyAccumulatorForOddKeys(AggregateFnV2):
|
|
def __init__(self):
|
|
super().__init__(
|
|
name="items",
|
|
on=None,
|
|
ignore_nulls=False,
|
|
zero_factory=lambda: [],
|
|
)
|
|
|
|
def aggregate_block(self, block):
|
|
table = BlockAccessor.for_block(block).to_arrow()
|
|
group_key = table.column("A")[0].as_py()
|
|
return [] if group_key % 2 else ["value"]
|
|
|
|
def combine(self, current, new):
|
|
return current + new
|
|
|
|
monkeypatch.setattr(table_block, "MAX_UNCOMPACTED_SIZE_BYTES", 1)
|
|
|
|
source = pa.table({"A": [1, 2, 3, 4], "B": [0, 0, 0, 0]})
|
|
partial = BlockAccessor.for_block(source)._aggregate(
|
|
SortKey("A"), (EmptyAccumulatorForOddKeys(),)
|
|
)
|
|
|
|
assert partial.column("A").to_pylist() == [1, 2, 3, 4]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|