wehub-resource-sync
230 lines
7.0 KiB
Python
230 lines
7.0 KiB
Python
import gc
|
|
from typing import List
|
|
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.execution.interfaces import (
|
|
BlockEntry,
|
|
ExecutionOptions,
|
|
RefBundle,
|
|
)
|
|
from ray.data._internal.execution.operators.base_physical_operator import (
|
|
AllToAllOperator,
|
|
)
|
|
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
|
|
from ray.data._internal.execution.operators.map_operator import MapOperator
|
|
from ray.data._internal.execution.util import make_ref_bundles
|
|
from ray.data._internal.progress.base_progress import NoopSubProgressBar
|
|
from ray.data.block import BlockAccessor
|
|
from ray.data.context import DataContext
|
|
from ray.data.tests.conftest import noop_counter
|
|
from ray.data.tests.util import (
|
|
_get_blocks,
|
|
_mul2_transform,
|
|
_take_outputs,
|
|
create_map_transformer_from_block_fn,
|
|
run_one_op_task,
|
|
run_op_tasks_sync,
|
|
)
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
_mul2_map_data_prcessor = create_map_transformer_from_block_fn(_mul2_transform)
|
|
|
|
|
|
def test_name_and_repr(ray_start_regular_shared):
|
|
inputs = make_ref_bundles([[1, 2], [3], [4, 5]])
|
|
input_op = InputDataBuffer(DataContext.get_current(), inputs)
|
|
map_op1 = MapOperator.create(
|
|
_mul2_map_data_prcessor,
|
|
input_op,
|
|
DataContext.get_current(),
|
|
name="map1",
|
|
)
|
|
|
|
assert map_op1.name == "map1"
|
|
assert map_op1.dag_str == "InputDataBuffer[Input] -> TaskPoolMapOperator[map1]"
|
|
assert str(map_op1) == "TaskPoolMapOperator[map1]"
|
|
|
|
map_op2 = MapOperator.create(
|
|
_mul2_map_data_prcessor,
|
|
map_op1,
|
|
DataContext.get_current(),
|
|
name="map2",
|
|
)
|
|
assert map_op2.name == "map2"
|
|
assert (
|
|
map_op2.dag_str
|
|
== "InputDataBuffer[Input] -> TaskPoolMapOperator[map1] -> TaskPoolMapOperator[map2]"
|
|
)
|
|
assert str(map_op2) == "TaskPoolMapOperator[map2]"
|
|
|
|
|
|
def test_input_data_buffer(ray_start_regular_shared):
|
|
# Create with bundles.
|
|
inputs = make_ref_bundles([[1, 2], [3], [4, 5]])
|
|
op = InputDataBuffer(DataContext.get_current(), inputs)
|
|
|
|
# Check we return all bundles in order.
|
|
assert not op.has_completed()
|
|
assert _take_outputs(op) == [[1, 2], [3], [4, 5]]
|
|
assert op.has_completed()
|
|
|
|
|
|
def test_all_to_all_operator():
|
|
def dummy_all_transform(bundles: List[RefBundle], ctx):
|
|
assert len(ctx.sub_progress_bar_dict) == 2
|
|
assert list(ctx.sub_progress_bar_dict.keys()) == ["Test1", "Test2"]
|
|
return make_ref_bundles([[1, 2], [3, 4]]), {"FooStats": []}
|
|
|
|
input_op = InputDataBuffer(
|
|
DataContext.get_current(), make_ref_bundles([[i] for i in range(100)])
|
|
)
|
|
op = AllToAllOperator(
|
|
dummy_all_transform,
|
|
input_op,
|
|
DataContext.get_current(),
|
|
target_max_block_size_override=DataContext.get_current().target_max_block_size,
|
|
num_outputs=2,
|
|
sub_progress_bar_names=["Test1", "Test2"],
|
|
name="TestAll",
|
|
)
|
|
|
|
# Initialize progress bar.
|
|
for name in op.get_sub_progress_bar_names():
|
|
pg = NoopSubProgressBar(
|
|
name=name,
|
|
max_name_length=100,
|
|
)
|
|
op.set_sub_progress_bar(name, pg)
|
|
|
|
# Feed data.
|
|
op.start(ExecutionOptions(), noop_counter())
|
|
while input_op.has_next():
|
|
op.add_input(input_op.get_next(), 0)
|
|
op.all_inputs_done()
|
|
|
|
# Check we return transformed bundles.
|
|
assert not op.has_completed()
|
|
outputs = _take_outputs(op)
|
|
expected = [[1, 2], [3, 4]]
|
|
assert sorted(outputs) == expected, f"Expected {expected}, got {outputs}"
|
|
stats = op.get_stats()
|
|
assert "FooStats" in stats
|
|
assert op.has_completed()
|
|
|
|
|
|
def test_num_outputs_total():
|
|
# The number of outputs is always known for InputDataBuffer.
|
|
input_op = InputDataBuffer(
|
|
DataContext.get_current(), make_ref_bundles([[i] for i in range(100)])
|
|
)
|
|
assert input_op.num_outputs_total() == 100
|
|
|
|
# Prior to execution, the number of outputs is unknown
|
|
# for Map/AllToAllOperator operators.
|
|
op1 = MapOperator.create(
|
|
_mul2_map_data_prcessor,
|
|
input_op,
|
|
DataContext.get_current(),
|
|
name="TestMapper",
|
|
)
|
|
assert op1.num_outputs_total() is None
|
|
|
|
def dummy_all_transform(bundles: List[RefBundle]):
|
|
return make_ref_bundles([[1, 2], [3, 4]]), {"FooStats": []}
|
|
|
|
op2 = AllToAllOperator(
|
|
dummy_all_transform,
|
|
input_op=op1,
|
|
data_context=DataContext.get_current(),
|
|
target_max_block_size_override=DataContext.get_current().target_max_block_size,
|
|
name="TestAll",
|
|
)
|
|
assert op2.num_outputs_total() is None
|
|
|
|
# Feed data and implement streaming exec.
|
|
output = []
|
|
op1.start(ExecutionOptions(actor_locality_enabled=True), noop_counter())
|
|
while input_op.has_next():
|
|
op1.add_input(input_op.get_next(), 0)
|
|
while not op1.has_next():
|
|
run_one_op_task(op1)
|
|
while op1.has_next():
|
|
ref = op1.get_next()
|
|
assert ref.owns_blocks, ref
|
|
_get_blocks(ref, output)
|
|
# After op finishes, num_outputs_total is known.
|
|
assert op1.num_outputs_total() == 100
|
|
|
|
|
|
def test_all_to_all_estimated_num_output_bundles():
|
|
# Test all to all operator
|
|
input_op = InputDataBuffer(
|
|
DataContext.get_current(), make_ref_bundles([[i] for i in range(100)])
|
|
)
|
|
|
|
def all_transform(bundles: List[RefBundle], ctx):
|
|
return bundles, {}
|
|
|
|
estimated_output_blocks = 500
|
|
op1 = AllToAllOperator(
|
|
all_transform,
|
|
input_op,
|
|
DataContext.get_current(),
|
|
DataContext.get_current().target_max_block_size,
|
|
estimated_output_blocks,
|
|
)
|
|
op2 = AllToAllOperator(
|
|
all_transform,
|
|
op1,
|
|
DataContext.get_current(),
|
|
DataContext.get_current().target_max_block_size,
|
|
)
|
|
|
|
op1.start(ExecutionOptions(), noop_counter())
|
|
op2.start(ExecutionOptions(), noop_counter())
|
|
while input_op.has_next():
|
|
op1.add_input(input_op.get_next(), 0)
|
|
op1.all_inputs_done()
|
|
run_op_tasks_sync(op1)
|
|
|
|
while op1.has_next():
|
|
op2.add_input(op1.get_next(), 0)
|
|
op2.all_inputs_done()
|
|
run_op_tasks_sync(op2)
|
|
|
|
# estimated output blocks for op2 should fallback to op1
|
|
assert op2._estimated_num_output_bundles is None
|
|
assert op2.num_outputs_total() == estimated_output_blocks
|
|
|
|
|
|
def test_input_data_buffer_does_not_free_inputs():
|
|
# Tests https://github.com/ray-project/ray/issues/46282
|
|
block = pd.DataFrame({"id": [0]})
|
|
block_ref = ray.put(block)
|
|
metadata = BlockAccessor.for_block(block).get_metadata()
|
|
schema = BlockAccessor.for_block(block).schema()
|
|
op = InputDataBuffer(
|
|
DataContext.get_current(),
|
|
input_data=[
|
|
RefBundle(
|
|
[BlockEntry(block_ref, metadata)], owns_blocks=False, schema=schema
|
|
)
|
|
],
|
|
)
|
|
|
|
op.get_next()
|
|
gc.collect()
|
|
|
|
# `InputDataBuffer` should still hold a reference to the input block even after
|
|
# `get_next` is called.
|
|
assert len(gc.get_referrers(block_ref)) > 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|