wehub-resource-sync
284 lines
9.8 KiB
Python
284 lines
9.8 KiB
Python
import collections
|
|
import itertools
|
|
import random
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.execution.interfaces import ExecutionOptions
|
|
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
|
|
from ray.data._internal.execution.operators.output_splitter import OutputSplitter
|
|
from ray.data._internal.execution.util import make_ref_bundles
|
|
from ray.data.context import DataContext
|
|
from ray.data.tests.conftest import noop_counter
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
@pytest.mark.parametrize("equal", [False, True])
|
|
@pytest.mark.parametrize("chunk_size", [1, 10])
|
|
def test_split_operator(ray_start_regular_shared, equal, chunk_size):
|
|
num_input_blocks = 100
|
|
num_splits = 3
|
|
# Add this many input blocks each time.
|
|
# Make sure it is greater than num_splits * 2,
|
|
# so we can test the output order of `OutputSplitter.get_next`.
|
|
num_add_input_blocks = 10
|
|
input_op = InputDataBuffer(
|
|
DataContext.get_current(),
|
|
make_ref_bundles([[i] * chunk_size for i in range(num_input_blocks)]),
|
|
)
|
|
op = OutputSplitter(
|
|
input_op,
|
|
num_splits,
|
|
equal=equal,
|
|
data_context=DataContext.get_current(),
|
|
)
|
|
|
|
# Feed data and implement streaming exec.
|
|
output_splits = [[] for _ in range(num_splits)]
|
|
op.start(ExecutionOptions(), noop_counter())
|
|
while input_op.has_next():
|
|
for _ in range(num_add_input_blocks):
|
|
if not input_op.has_next():
|
|
break
|
|
op.add_input(input_op.get_next(), 0)
|
|
while op.has_next():
|
|
ref = op.get_next()
|
|
assert ref.owns_blocks, ref
|
|
for block_ref in ref.block_refs:
|
|
assert ref.output_split_idx is not None
|
|
output_splits[ref.output_split_idx].extend(
|
|
list(ray.get(block_ref)["id"])
|
|
)
|
|
op.all_inputs_done()
|
|
|
|
expected_splits = [[] for _ in range(num_splits)]
|
|
for i in range(num_splits):
|
|
for j in range(i, num_input_blocks, num_splits):
|
|
expected_splits[i].extend([j] * chunk_size)
|
|
if equal:
|
|
min_len = min(len(expected_splits[i]) for i in range(num_splits))
|
|
for i in range(num_splits):
|
|
expected_splits[i] = expected_splits[i][:min_len]
|
|
for i in range(num_splits):
|
|
assert output_splits[i] == expected_splits[i], (
|
|
output_splits[i],
|
|
expected_splits[i],
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("equal", [False, True])
|
|
@pytest.mark.parametrize("random_seed", list(range(10)))
|
|
def test_split_operator_random(ray_start_regular_shared, equal, random_seed):
|
|
random.seed(random_seed)
|
|
inputs = make_ref_bundles([[i] * random.randint(0, 10) for i in range(100)])
|
|
num_inputs = sum(x.num_rows() for x in inputs)
|
|
input_op = InputDataBuffer(DataContext.get_current(), inputs)
|
|
op = OutputSplitter(
|
|
input_op, 3, equal=equal, data_context=DataContext.get_current()
|
|
)
|
|
|
|
# Feed data and implement streaming exec.
|
|
output_splits = collections.defaultdict(list)
|
|
op.start(ExecutionOptions(), noop_counter())
|
|
while input_op.has_next():
|
|
op.add_input(input_op.get_next(), 0)
|
|
op.all_inputs_done()
|
|
while op.has_next():
|
|
ref = op.get_next()
|
|
assert ref.owns_blocks, ref
|
|
for block_ref in ref.block_refs:
|
|
output_splits[ref.output_split_idx].extend(list(ray.get(block_ref)["id"]))
|
|
if equal:
|
|
actual = [len(output_splits[i]) for i in range(3)]
|
|
expected = [num_inputs // 3] * 3
|
|
assert actual == expected
|
|
else:
|
|
assert sum(len(output_splits[i]) for i in range(3)) == num_inputs, output_splits
|
|
|
|
|
|
def test_split_operator_locality_hints(ray_start_regular_shared):
|
|
input_op = InputDataBuffer(
|
|
DataContext.get_current(), make_ref_bundles([[i] for i in range(10)])
|
|
)
|
|
op = OutputSplitter(
|
|
input_op,
|
|
2,
|
|
equal=False,
|
|
data_context=DataContext.get_current(),
|
|
locality_hints=["node1", "node2"],
|
|
)
|
|
|
|
def get_fake_loc(item):
|
|
assert isinstance(item, int), item
|
|
if item in [0, 1, 4, 5, 8]:
|
|
return "node1"
|
|
else:
|
|
return "node2"
|
|
|
|
def get_bundle_loc(bundle):
|
|
block = ray.get(bundle.blocks[0].ref)
|
|
fval = list(block["id"])[0]
|
|
return [get_fake_loc(fval)]
|
|
|
|
op._get_locations = get_bundle_loc
|
|
|
|
# Feed data and implement streaming exec.
|
|
output_splits = collections.defaultdict(list)
|
|
op.start(ExecutionOptions(actor_locality_enabled=True), noop_counter())
|
|
while input_op.has_next():
|
|
op.add_input(input_op.get_next(), 0)
|
|
op.all_inputs_done()
|
|
while op.has_next():
|
|
ref = op.get_next()
|
|
assert ref.owns_blocks, ref
|
|
for block_ref in ref.block_refs:
|
|
output_splits[ref.output_split_idx].extend(list(ray.get(block_ref)["id"]))
|
|
|
|
total = 0
|
|
for i in range(2):
|
|
if i == 0:
|
|
node = "node1"
|
|
else:
|
|
node = "node2"
|
|
split = output_splits[i]
|
|
for item in split:
|
|
assert get_fake_loc(item) == node
|
|
total += 1
|
|
|
|
assert total == 10, total
|
|
assert "all objects local" in op.progress_str()
|
|
|
|
|
|
@pytest.mark.parametrize("equal", [False, True])
|
|
@pytest.mark.parametrize("random_seed", list(range(10)))
|
|
def test_split_operator_with_locality(ray_start_regular_shared, equal, random_seed):
|
|
"""Test locality-based dispatching with equal=True and equal=False modes.
|
|
|
|
This test verifies that the OutputSplitter:
|
|
1. Correctly buffers data to ensure equal distribution when equal=True
|
|
2. Respects locality hints in both modes
|
|
3. Yields blocks incrementally when locality is matched (streaming behavior)
|
|
4. The fix ensures that _can_safely_dispatch correctly calculates remaining
|
|
buffer requirements.
|
|
"""
|
|
|
|
random.seed(random_seed)
|
|
|
|
# Create bundles with varying sizes to test buffer management
|
|
input_bundles = make_ref_bundles([[i] * random.randint(1, 10) for i in range(100)])
|
|
num_inputs = sum(x.num_rows() for x in input_bundles)
|
|
|
|
input_op = InputDataBuffer(DataContext.get_current(), input_bundles)
|
|
op = OutputSplitter(
|
|
input_op,
|
|
3,
|
|
equal=equal,
|
|
data_context=DataContext.get_current(),
|
|
locality_hints=["node0", "node1", "node2"],
|
|
)
|
|
|
|
# Mock locality function: distribute items across 3 nodes
|
|
def _map_row_to_node(first_row_id_val) -> str:
|
|
return f"node{first_row_id_val % 3}"
|
|
|
|
def _get_fake_bundle_loc(bundle):
|
|
block = ray.get(bundle.block_refs[0])
|
|
first_row_id_val = block["id"][0]
|
|
return [_map_row_to_node(first_row_id_val)]
|
|
|
|
op._get_locations = _get_fake_bundle_loc
|
|
|
|
# Feed data and implement streaming exec
|
|
output_splits = [[] for _ in range(3)]
|
|
yielded_incrementally = 0
|
|
|
|
op.start(ExecutionOptions(actor_locality_enabled=True), noop_counter())
|
|
while input_op.has_next():
|
|
op.add_input(input_op.get_next(), 0)
|
|
|
|
# Drain some outputs to simulate streaming consumption
|
|
while op.has_next():
|
|
yielded_incrementally += 1
|
|
|
|
ref = op.get_next()
|
|
|
|
assert ref.owns_blocks, ref
|
|
|
|
for block_ref in ref.block_refs:
|
|
output_splits[ref.output_split_idx].extend(
|
|
list(ray.get(block_ref)["id"])
|
|
)
|
|
|
|
op.all_inputs_done()
|
|
|
|
# Collect remaining outputs
|
|
while op.has_next():
|
|
ref = op.get_next()
|
|
|
|
assert ref.owns_blocks, ref
|
|
|
|
for block_ref in ref.block_refs:
|
|
output_splits[ref.output_split_idx].extend(list(ray.get(block_ref)["id"]))
|
|
|
|
# Verify streaming behavior: outputs should be yielded before all inputs are done
|
|
# With locality hints, we should see outputs during input phase
|
|
assert yielded_incrementally > 0, (
|
|
f"Expected incremental output with locality hints, but got 0 outputs during "
|
|
f"{len(input_bundles)} input blocks. This suggests buffering all data instead of streaming."
|
|
)
|
|
|
|
# Verify equal distribution when equal=True
|
|
if equal:
|
|
actual = [len(output_splits[i]) for i in range(3)]
|
|
expected = [num_inputs // 3] * 3
|
|
assert (
|
|
actual == expected
|
|
), f"Expected equal distribution {expected}, got {actual}"
|
|
else:
|
|
# In non-equal mode, verify all data is output with correct row IDs
|
|
all_output_row_ids = set(itertools.chain.from_iterable(output_splits))
|
|
|
|
# Reconstruct expected row IDs from the input bundles
|
|
expected_row_ids = set()
|
|
for b in input_bundles:
|
|
id_col = ray.get(b.block_refs[0])["id"]
|
|
expected_row_ids.update(list(id_col))
|
|
|
|
assert all_output_row_ids == expected_row_ids
|
|
|
|
# Verify locality was respected (most items should be on their preferred node)
|
|
locality_hits = 0
|
|
total = 0
|
|
|
|
for split_idx in range(3):
|
|
actual_node = f"node{split_idx}"
|
|
|
|
for row_id in output_splits[split_idx]:
|
|
total += 1
|
|
expected_node = _map_row_to_node(row_id)
|
|
|
|
assert expected_node in ["node0", "node1", "node2"], expected_node
|
|
|
|
if expected_node == actual_node:
|
|
locality_hits += 1
|
|
|
|
# Should have excellent locality since bundles are dispatched based on locality hints.
|
|
# With perfect locality we'd get 100%, but buffering for equal distribution and
|
|
# occasional forced dispatches when buffer is full may cause some misses.
|
|
# We expect at least 85% locality hit rate, which validates the feature is working.
|
|
locality_ratio = locality_hits / total if total > 0 else 0
|
|
|
|
# NOTE: 90% is an observed locality ratio that should be fixed for this test
|
|
assert locality_ratio >= 0.85, (
|
|
f"Locality ratio {locality_ratio:.2f} too low. "
|
|
f"Expected >=85% with locality-aware dispatching. "
|
|
f"Hits: {locality_hits}/{total}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|