wehub-resource-sync
146 lines
4.3 KiB
Python
146 lines
4.3 KiB
Python
from dataclasses import dataclass
|
|
from typing import Iterable, List
|
|
|
|
import numpy
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.execution.interfaces import TaskContext
|
|
from ray.data.block import Block
|
|
from ray.data.datasource import Datasink
|
|
from ray.data.datasource.datasink import DummyOutputDatasink, WriteResult
|
|
|
|
|
|
def test_write_datasink(ray_start_regular_shared):
|
|
output = DummyOutputDatasink()
|
|
ds = ray.data.range(10, override_num_blocks=2)
|
|
ds.write_datasink(output)
|
|
assert output.num_ok == 1
|
|
assert output.num_failed == 0
|
|
assert ray.get(output.data_sink.get_rows_written.remote()) == 10
|
|
|
|
output.enabled = False
|
|
ds = ray.data.range(10, override_num_blocks=2)
|
|
with pytest.raises(ValueError):
|
|
ds.write_datasink(output, ray_remote_args={"max_retries": 0})
|
|
assert output.num_ok == 1
|
|
assert output.num_failed == 1
|
|
assert ray.get(output.data_sink.get_rows_written.remote()) == 10
|
|
|
|
|
|
@pytest.mark.parametrize("min_rows_per_write", [25, 50])
|
|
def test_min_rows_per_write(tmp_path, ray_start_regular_shared, min_rows_per_write):
|
|
class MockDatasink(Datasink[None]):
|
|
def __init__(self, min_rows_per_write):
|
|
self._min_rows_per_write = min_rows_per_write
|
|
|
|
def write(self, blocks: Iterable[Block], ctx: TaskContext) -> None:
|
|
assert sum(len(block) for block in blocks) == self._min_rows_per_write
|
|
|
|
@property
|
|
def min_rows_per_write(self):
|
|
return self._min_rows_per_write
|
|
|
|
ray.data.range(100, override_num_blocks=4).write_datasink(
|
|
MockDatasink(min_rows_per_write)
|
|
)
|
|
|
|
|
|
def test_write_result(ray_start_regular_shared):
|
|
"""Test the write_result argument in `on_write_complete`."""
|
|
|
|
@dataclass
|
|
class CustomWriteResult:
|
|
|
|
ids: List[int]
|
|
|
|
class CustomDatasink(Datasink[CustomWriteResult]):
|
|
def __init__(self) -> None:
|
|
self.ids = []
|
|
self.num_rows = 0
|
|
self.size_bytes = 0
|
|
|
|
def write(self, blocks: Iterable[Block], ctx: TaskContext):
|
|
ids = []
|
|
for b in blocks:
|
|
ids.extend(b["id"].to_pylist())
|
|
return CustomWriteResult(ids=ids)
|
|
|
|
def on_write_complete(self, write_result: WriteResult[CustomWriteResult]):
|
|
ids = []
|
|
for result in write_result.write_returns:
|
|
ids.extend(result.ids)
|
|
self.ids = sorted(ids)
|
|
self.num_rows = write_result.num_rows
|
|
self.size_bytes = write_result.size_bytes
|
|
|
|
num_items = 10
|
|
size_bytes_per_row = 500
|
|
|
|
def map_fn(row):
|
|
row["data"] = numpy.zeros(size_bytes_per_row, dtype=numpy.int8)
|
|
return row
|
|
|
|
ds = ray.data.range(num_items).map(map_fn)
|
|
|
|
datasink = CustomDatasink()
|
|
ds.write_datasink(datasink)
|
|
|
|
assert datasink.ids == list(range(num_items))
|
|
assert datasink.num_rows == num_items
|
|
assert datasink.size_bytes == pytest.approx(num_items * size_bytes_per_row, rel=0.1)
|
|
|
|
|
|
class NodeLoggerOutputDatasink(Datasink[None]):
|
|
"""A writable datasource that logs node IDs of write tasks, for testing."""
|
|
|
|
def __init__(self, node_id: str):
|
|
|
|
self.num_ok = 0
|
|
self.num_failed = 0
|
|
self.node_id = node_id
|
|
self.num_rows_written = 0
|
|
|
|
def write(
|
|
self,
|
|
blocks: Iterable[Block],
|
|
ctx: TaskContext,
|
|
) -> None:
|
|
|
|
node_id = ray.get_runtime_context().get_node_id()
|
|
assert node_id == self.node_id
|
|
|
|
def on_write_complete(self, write_result: WriteResult[None]):
|
|
self.num_ok += 1
|
|
self.num_rows_written += write_result.num_rows
|
|
|
|
def on_write_failed(self, error: Exception) -> None:
|
|
self.num_failed += 1
|
|
|
|
|
|
def test_write_datasink_ray_remote_args(ray_start_cluster):
|
|
ray.shutdown()
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(
|
|
resources={"foo": 100},
|
|
num_cpus=1,
|
|
)
|
|
bar_worker = cluster.add_node(resources={"bar": 100}, num_cpus=1)
|
|
bar_node_id = bar_worker.node_id
|
|
|
|
ray.init(cluster.address)
|
|
|
|
output = NodeLoggerOutputDatasink(bar_node_id)
|
|
ds = ray.data.range(100, override_num_blocks=10)
|
|
# Pin write tasks to node with "bar" resource.
|
|
ds.write_datasink(output, ray_remote_args={"resources": {"bar": 1}})
|
|
assert output.num_ok == 1
|
|
assert output.num_failed == 0
|
|
assert output.num_rows_written == 100
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|