wehub-resource-sync
303 lines
9.4 KiB
Python
303 lines
9.4 KiB
Python
import os
|
|
from typing import List
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pytest
|
|
|
|
import ray
|
|
from ray._private.internal_api import memory_summary
|
|
from ray.data._internal.datasource.csv_datasource import CSVDatasource
|
|
from ray.data.block import BlockMetadata
|
|
from ray.data.dataset import Dataset
|
|
from ray.data.datasource import Datasource, ReadTask
|
|
from ray.data.tests.util import column_udf, extract_values
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
@ray.remote
|
|
class Counter:
|
|
def __init__(self):
|
|
self.value = 0
|
|
|
|
def increment(self):
|
|
self.value += 1
|
|
return self.value
|
|
|
|
def get(self):
|
|
return self.value
|
|
|
|
def reset(self):
|
|
self.value = 0
|
|
|
|
|
|
class MySource(CSVDatasource):
|
|
def __init__(self, paths, counter):
|
|
super().__init__(paths)
|
|
|
|
self.counter = counter
|
|
|
|
def _read_stream(self, f, path: str):
|
|
count = self.counter.increment.remote()
|
|
ray.get(count)
|
|
for block in super()._read_stream(f, path):
|
|
yield block
|
|
|
|
|
|
def expect_stages(pipe, num_stages_expected, stage_names):
|
|
stats = pipe.stats()
|
|
for name in stage_names:
|
|
name = " " + name + ":"
|
|
assert name in stats, (name, stats)
|
|
if isinstance(pipe, Dataset):
|
|
pass
|
|
else:
|
|
assert (
|
|
len(pipe._optimized_stages) == num_stages_expected
|
|
), pipe._optimized_stages
|
|
|
|
|
|
def dummy_map(x):
|
|
"""Dummy function used in calls to map_batches in these tests."""
|
|
return x
|
|
|
|
|
|
def test_memory_sanity(shutdown_only):
|
|
info = ray.init(num_cpus=1, object_store_memory=500e6)
|
|
ds = ray.data.range(10)
|
|
ds = ds.map(lambda x: {"data": np.ones(100 * 1024 * 1024, dtype=np.uint8)})
|
|
ds.materialize()
|
|
meminfo = memory_summary(info.address_info["address"], stats_only=True)
|
|
|
|
# Sanity check spilling is happening as expected.
|
|
assert "Spilled" in meminfo, meminfo
|
|
|
|
|
|
class OnesSource(Datasource):
|
|
def prepare_read(
|
|
self,
|
|
parallelism: int,
|
|
n_per_block: int,
|
|
) -> List[ReadTask]:
|
|
read_tasks: List[ReadTask] = []
|
|
meta = BlockMetadata(
|
|
num_rows=1,
|
|
size_bytes=n_per_block,
|
|
input_files=None,
|
|
exec_stats=None,
|
|
)
|
|
|
|
for _ in range(parallelism):
|
|
read_tasks.append(
|
|
ReadTask(lambda: [[np.ones(n_per_block, dtype=np.uint8)]], meta)
|
|
)
|
|
return read_tasks
|
|
|
|
|
|
def test_memory_release(shutdown_only):
|
|
info = ray.init(num_cpus=1, object_store_memory=1500e6)
|
|
ds = ray.data.range(10)
|
|
|
|
# Should get fused into single operator.
|
|
ds = ds.map(lambda x: {"data": np.ones(100 * 1024 * 1024, dtype=np.uint8)})
|
|
ds = ds.map(lambda x: {"data": np.ones(100 * 1024 * 1024, dtype=np.uint8)})
|
|
ds = ds.map(lambda x: {"data": np.ones(100 * 1024 * 1024, dtype=np.uint8)})
|
|
ds.materialize()
|
|
meminfo = memory_summary(info.address_info["address"], stats_only=True)
|
|
assert "Spilled" not in meminfo, meminfo
|
|
|
|
|
|
@pytest.mark.skip(reason="Flaky, see https://github.com/ray-project/ray/issues/24757")
|
|
def test_memory_release_shuffle(shutdown_only):
|
|
# TODO(ekl) why is this flaky? Due to eviction delay?
|
|
error = None
|
|
for trial in range(3):
|
|
print("Try", trial)
|
|
try:
|
|
info = ray.init(num_cpus=1, object_store_memory=1800e6)
|
|
ds = ray.data.range(10)
|
|
|
|
# Should get fused into single stage.
|
|
ds = ds.map(lambda x: {"data": np.ones(100 * 1024 * 1024, dtype=np.uint8)})
|
|
ds.random_shuffle().materialize()
|
|
meminfo = memory_summary(info.address_info["address"], stats_only=True)
|
|
assert "Spilled" not in meminfo, meminfo
|
|
return
|
|
except Exception as e:
|
|
error = e
|
|
print("Failed", e)
|
|
finally:
|
|
ray.shutdown()
|
|
raise error
|
|
|
|
|
|
def test_lazy_fanout(shutdown_only, local_path):
|
|
ray.init(num_cpus=1)
|
|
map_counter = Counter.remote()
|
|
|
|
def inc(row):
|
|
map_counter.increment.remote()
|
|
row["one"] += 1
|
|
return row
|
|
|
|
df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
|
path = os.path.join(local_path, "test.csv")
|
|
df.to_csv(path, index=False, storage_options={})
|
|
read_counter = Counter.remote()
|
|
source = MySource(path, read_counter)
|
|
|
|
# Test that fan-out of a lazy dataset results in re-execution up to the datasource,
|
|
# due to block move semantics.
|
|
ds = ray.data.read_datasource(source, override_num_blocks=1)
|
|
ds1 = ds.map(inc)
|
|
ds2 = ds1.map(inc)
|
|
ds3 = ds1.map(inc)
|
|
# Test content.
|
|
assert ds2.materialize().take() == [
|
|
{"one": 3, "two": "a"},
|
|
{"one": 4, "two": "b"},
|
|
{"one": 5, "two": "c"},
|
|
]
|
|
assert ds3.materialize().take() == [
|
|
{"one": 3, "two": "a"},
|
|
{"one": 4, "two": "b"},
|
|
{"one": 5, "two": "c"},
|
|
]
|
|
# Test that data is read twice.
|
|
assert ray.get(read_counter.get.remote()) == 2
|
|
# Test that first map is executed twice.
|
|
assert ray.get(map_counter.get.remote()) == 2 * 3 + 3 + 3
|
|
|
|
# Test that fan-out of a lazy dataset with a non-lazy datasource results in
|
|
# re-execution up to the datasource, due to block move semantics.
|
|
ray.get(map_counter.reset.remote())
|
|
|
|
def inc(x):
|
|
map_counter.increment.remote()
|
|
return {"item": x["item"] + 1}
|
|
|
|
# The source data shouldn't be cleared since it's non-lazy.
|
|
ds = ray.data.from_items(list(range(10)))
|
|
ds1 = ds.map(inc)
|
|
ds2 = ds1.map(inc)
|
|
ds3 = ds1.map(inc)
|
|
# Test content.
|
|
assert extract_values("item", ds2.materialize().take()) == list(range(2, 12))
|
|
assert extract_values("item", ds3.materialize().take()) == list(range(2, 12))
|
|
# Test that first map is executed twice.
|
|
assert ray.get(map_counter.get.remote()) == 2 * 10 + 10 + 10
|
|
|
|
ray.get(map_counter.reset.remote())
|
|
# The source data shouldn't be cleared since it's non-lazy.
|
|
ds = ray.data.from_items(list(range(10)))
|
|
# Add extra transformation after being lazy.
|
|
ds = ds.map(inc)
|
|
ds1 = ds.map(inc)
|
|
ds2 = ds.map(inc)
|
|
# Test content.
|
|
assert extract_values("item", ds1.materialize().take()) == list(range(2, 12))
|
|
assert extract_values("item", ds2.materialize().take()) == list(range(2, 12))
|
|
# Test that first map is executed twice, because ds1.materialize()
|
|
# clears up the previous snapshot blocks, and ds2.materialize()
|
|
# has to re-execute ds.map(inc) again.
|
|
assert ray.get(map_counter.get.remote()) == 2 * 10 + 10 + 10
|
|
|
|
|
|
def test_spread_hint_inherit(ray_start_regular_shared):
|
|
ds = ray.data.range(10)
|
|
ds = ds.map(column_udf("id", lambda x: x + 1))
|
|
ds = ds.random_shuffle()
|
|
|
|
shuffle_op = ds._logical_plan.dag
|
|
read_op = shuffle_op.input_dependencies[0].input_dependencies[0]
|
|
assert read_op.ray_remote_args == {"scheduling_strategy": "SPREAD"}
|
|
|
|
|
|
def test_optimize_randomize_block_order(ray_start_regular_shared):
|
|
"""Test that randomize_block_order is not fused with other operators."""
|
|
ds = (
|
|
ray.data.range(10)
|
|
.map_batches(dummy_map)
|
|
.randomize_block_order()
|
|
.map_batches(dummy_map)
|
|
.materialize()
|
|
)
|
|
expect_stages(
|
|
ds,
|
|
2,
|
|
[
|
|
"ReadRange->MapBatches(dummy_map)",
|
|
"RandomizeBlockOrder",
|
|
"MapBatches(dummy_map)",
|
|
],
|
|
)
|
|
|
|
ds2 = (
|
|
ray.data.range(10)
|
|
.randomize_block_order()
|
|
.repartition(10)
|
|
.map_batches(dummy_map)
|
|
.materialize()
|
|
)
|
|
expect_stages(
|
|
ds2,
|
|
3,
|
|
["ReadRange", "RandomizeBlockOrder", "Repartition", "MapBatches(dummy_map)"],
|
|
)
|
|
|
|
|
|
def test_write_fusion(ray_start_regular_shared, tmp_path):
|
|
path = os.path.join(tmp_path, "out")
|
|
ds = ray.data.range(100).map_batches(lambda x: x)
|
|
ds.write_csv(path)
|
|
stats = ds._write_ds.stats()
|
|
assert "ReadRange->MapBatches(<lambda>)->Write" in stats, stats
|
|
|
|
ds = (
|
|
ray.data.range(100)
|
|
.map_batches(lambda x: x)
|
|
.random_shuffle()
|
|
.map_batches(lambda x: x)
|
|
)
|
|
ds.write_csv(path)
|
|
stats = ds._write_ds.stats()
|
|
assert "ReadRange->MapBatches(<lambda>)" in stats, stats
|
|
assert "RandomShuffle" in stats, stats
|
|
assert "MapBatches(<lambda>)->Write" in stats, stats
|
|
|
|
|
|
@pytest.mark.skip(reason="reusing base data not enabled")
|
|
@pytest.mark.parametrize("with_shuffle", [True, False])
|
|
def test_optimize_lazy_reuse_base_data(
|
|
ray_start_regular_shared, local_path, enable_dynamic_splitting, with_shuffle
|
|
):
|
|
num_blocks = 4
|
|
dfs = [pd.DataFrame({"one": list(range(i, i + 4))}) for i in range(num_blocks)]
|
|
paths = [os.path.join(local_path, f"test{i}.csv") for i in range(num_blocks)]
|
|
for df, path in zip(dfs, paths):
|
|
df.to_csv(path, index=False)
|
|
counter = Counter.remote()
|
|
source = MySource(paths, counter)
|
|
ds = ray.data.read_datasource(source, override_num_blocks=4)
|
|
num_reads = ray.get(counter.get.remote())
|
|
assert num_reads == 1, num_reads
|
|
ds = ds.map(column_udf("id", lambda x: x))
|
|
if with_shuffle:
|
|
ds = ds.random_shuffle()
|
|
ds.take()
|
|
num_reads = ray.get(counter.get.remote())
|
|
assert num_reads == num_blocks, num_reads
|
|
|
|
|
|
def test_require_preserve_order(ray_start_regular_shared):
|
|
ds1 = ray.data.range(100).map_batches(lambda x: x).zip(ray.data.range(100))
|
|
assert ds1._logical_plan.require_preserve_order()
|
|
ds2 = ray.data.range(100).map_batches(lambda x: x).repartition(10)
|
|
assert not ds2._logical_plan.require_preserve_order()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|