wehub-resource-sync
885 lines
29 KiB
Python
885 lines
29 KiB
Python
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.block_builder import BlockBuilder
|
|
from ray.data._internal.datasource.csv_datasink import CSVDatasink
|
|
from ray.data._internal.datasource.csv_datasource import CSVDatasource
|
|
from ray.data._internal.datasource.range_datasource import RangeDatasource
|
|
from ray.data._internal.execution.interfaces.ref_bundle import (
|
|
_ref_bundles_iterator_to_block_refs_list,
|
|
)
|
|
from ray.data.block import BlockAccessor
|
|
from ray.data.dataset import Dataset, MaterializedDataset
|
|
from ray.data.datasource.util import (
|
|
_validate_head_node_resources_for_local_scheduling,
|
|
)
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.data.tests.conftest import (
|
|
CoreExecutionMetrics,
|
|
assert_core_execution_metrics_equals,
|
|
get_initial_core_execution_metrics_snapshot,
|
|
)
|
|
from ray.data.tests.util import column_udf, extract_values
|
|
from ray.tests.conftest import * # noqa
|
|
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
|
|
|
|
|
def test_schema(ray_start_regular):
|
|
last_snapshot = get_initial_core_execution_metrics_snapshot()
|
|
|
|
ds2 = ray.data.range(10, override_num_blocks=10)
|
|
ds3 = ds2.repartition(5)
|
|
ds3 = ds3.materialize()
|
|
last_snapshot = assert_core_execution_metrics_equals(
|
|
CoreExecutionMetrics(
|
|
task_count={
|
|
"ReadRange": 10,
|
|
"reduce": 5,
|
|
}
|
|
),
|
|
last_snapshot,
|
|
)
|
|
|
|
ds4 = ds3.map(lambda x: {"a": "hi", "b": 1.0}).limit(5).repartition(1)
|
|
ds4 = ds4.materialize()
|
|
last_snapshot = assert_core_execution_metrics_equals(
|
|
CoreExecutionMetrics(
|
|
task_count={
|
|
"Map(<lambda>)": lambda count: count <= 5,
|
|
"slice_fn": 1,
|
|
"reduce": 1,
|
|
}
|
|
),
|
|
last_snapshot,
|
|
)
|
|
|
|
ds2_schema = ds2.schema(fetch_if_missing=False)
|
|
assert ds2_schema is not None
|
|
assert ds2_schema.names == ["id"]
|
|
assert not isinstance(ds2, MaterializedDataset)
|
|
last_snapshot = assert_core_execution_metrics_equals(
|
|
CoreExecutionMetrics(task_count={}), last_snapshot
|
|
)
|
|
|
|
ds3_schema = ds3.schema(fetch_if_missing=False)
|
|
assert ds3_schema is not None
|
|
assert ds3_schema.names == ["id"]
|
|
assert isinstance(ds3, MaterializedDataset)
|
|
last_snapshot = assert_core_execution_metrics_equals(
|
|
CoreExecutionMetrics(task_count={}), last_snapshot
|
|
)
|
|
ds4_schema = ds4.schema(fetch_if_missing=False)
|
|
assert ds4_schema is not None
|
|
assert ds4_schema.names == ["a", "b"]
|
|
assert isinstance(ds4, MaterializedDataset)
|
|
last_snapshot = assert_core_execution_metrics_equals(
|
|
CoreExecutionMetrics(task_count={}), last_snapshot
|
|
)
|
|
|
|
|
|
def test_schema_no_execution(ray_start_regular):
|
|
last_snapshot = get_initial_core_execution_metrics_snapshot()
|
|
ds = ray.data.range(100, override_num_blocks=10)
|
|
last_snapshot = assert_core_execution_metrics_equals(
|
|
CoreExecutionMetrics(task_count={}),
|
|
last_snapshot,
|
|
)
|
|
# We do not kick off the read task by default.
|
|
schema = ds.schema()
|
|
assert schema.names == ["id"]
|
|
|
|
# Fetching the schema does not trigger execution, since
|
|
# the schema is known beforehand for RangeDatasource.
|
|
last_snapshot = assert_core_execution_metrics_equals(
|
|
CoreExecutionMetrics(task_count={}), last_snapshot
|
|
)
|
|
# Fetching the schema should not trigger execution of extra read tasks.
|
|
|
|
|
|
def test_schema_cached(ray_start_regular):
|
|
def check_schema_cached(ds, expected_task_count, last_snapshot):
|
|
schema = ds.schema()
|
|
last_snapshot = assert_core_execution_metrics_equals(
|
|
CoreExecutionMetrics(expected_task_count), last_snapshot
|
|
)
|
|
assert schema.names == ["a"]
|
|
cached_schema = ds.schema(fetch_if_missing=False)
|
|
assert cached_schema is not None
|
|
assert schema == cached_schema
|
|
last_snapshot = assert_core_execution_metrics_equals(
|
|
CoreExecutionMetrics({}), last_snapshot
|
|
)
|
|
return last_snapshot
|
|
|
|
last_snapshot = get_initial_core_execution_metrics_snapshot()
|
|
ds = ray.data.from_items([{"a": i} for i in range(100)], override_num_blocks=10)
|
|
last_snapshot = check_schema_cached(ds, {}, last_snapshot)
|
|
|
|
# Add a map_batches operator so that we are forced to compute the schema.
|
|
ds = ds.map_batches(lambda x: x)
|
|
last_snapshot = check_schema_cached(
|
|
ds,
|
|
{
|
|
"MapBatches(<lambda>)": lambda count: count <= 5,
|
|
"slice_fn": 1,
|
|
},
|
|
last_snapshot,
|
|
)
|
|
|
|
|
|
def test_avoid_placement_group_capture(shutdown_only):
|
|
ray.init(num_cpus=2)
|
|
|
|
@ray.remote
|
|
def run():
|
|
ds = ray.data.range(5)
|
|
assert sorted(
|
|
extract_values("id", ds.map(column_udf("id", lambda x: x + 1)).take())
|
|
) == [1, 2, 3, 4, 5]
|
|
assert ds.count() == 5
|
|
assert sorted(extract_values("id", ds.iter_rows())) == [0, 1, 2, 3, 4]
|
|
|
|
pg = ray.util.placement_group([{"CPU": 1}])
|
|
ray.get(
|
|
run.options(
|
|
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
|
placement_group=pg, placement_group_capture_child_tasks=True
|
|
)
|
|
).remote()
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def remove_named_placement_groups():
|
|
yield
|
|
for info in ray.util.placement_group_table().values():
|
|
if info["name"]:
|
|
pg = ray.util.get_placement_group(info["name"])
|
|
ray.util.remove_placement_group(pg)
|
|
|
|
|
|
def test_ray_remote_args_fn(shutdown_only, remove_named_placement_groups):
|
|
ray.init()
|
|
|
|
pg = ray.util.placement_group([{"CPU": 1}], name="test_pg")
|
|
|
|
def ray_remote_args_fn():
|
|
scheduling_strategy = PlacementGroupSchedulingStrategy(placement_group=pg)
|
|
return {"scheduling_strategy": scheduling_strategy}
|
|
|
|
class ActorClass:
|
|
def __call__(self, batch):
|
|
assert ray.util.get_current_placement_group() == pg
|
|
return batch
|
|
|
|
ray.data.range(1).map_batches(
|
|
ActorClass, concurrency=1, ray_remote_args_fn=ray_remote_args_fn
|
|
).take_all()
|
|
|
|
|
|
def test_dataset_lineage_serialization(shutdown_only):
|
|
ray.init()
|
|
ds = ray.data.range(10)
|
|
ds = ds.map(column_udf("id", lambda x: x + 1))
|
|
ds = ds.map(column_udf("id", lambda x: x + 1))
|
|
ds = ds.random_shuffle()
|
|
uuid = ds._get_uuid()
|
|
plan_uuid = ds._uuid
|
|
|
|
serialized_ds = ds.serialize_lineage()
|
|
|
|
ray.shutdown()
|
|
ray.init()
|
|
|
|
ds = Dataset.deserialize_lineage(serialized_ds)
|
|
# Check Dataset state.
|
|
assert ds._get_uuid() == uuid
|
|
assert ds._uuid == plan_uuid
|
|
# Check Dataset content.
|
|
assert ds.count() == 10
|
|
assert sorted(extract_values("id", ds.take())) == list(range(2, 12))
|
|
|
|
|
|
def test_dataset_lineage_serialization_unsupported(shutdown_only):
|
|
ray.init()
|
|
# In-memory data sources not supported.
|
|
ds = ray.data.from_items(list(range(10)))
|
|
ds = ds.map(column_udf("item", lambda x: x + 1))
|
|
ds = ds.map(column_udf("item", lambda x: x + 1))
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.serialize_lineage()
|
|
|
|
# In-memory data source unions not supported.
|
|
ds = ray.data.from_items(list(range(10)))
|
|
ds1 = ray.data.from_items(list(range(10, 20)))
|
|
ds2 = ds.union(ds1)
|
|
|
|
with pytest.raises(ValueError):
|
|
ds2.serialize_lineage()
|
|
|
|
# Lazy read unions supported.
|
|
ds = ray.data.range(10)
|
|
ds1 = ray.data.range(20)
|
|
ds2 = ds.union(ds1)
|
|
|
|
serialized_ds = ds2.serialize_lineage()
|
|
ds3 = Dataset.deserialize_lineage(serialized_ds)
|
|
assert set(extract_values("id", ds3.take(30))) == set(
|
|
list(range(10)) + list(range(20))
|
|
)
|
|
|
|
# Zips not supported.
|
|
ds = ray.data.from_items(list(range(10)))
|
|
ds1 = ray.data.from_items(list(range(10, 20)))
|
|
ds2 = ds.zip(ds1)
|
|
|
|
with pytest.raises(ValueError):
|
|
ds2.serialize_lineage()
|
|
|
|
|
|
def test_basic(ray_start_regular_shared):
|
|
ds = ray.data.range(5)
|
|
assert sorted(
|
|
extract_values("id", ds.map(column_udf("id", lambda x: x + 1)).take())
|
|
) == [1, 2, 3, 4, 5]
|
|
assert ds.count() == 5
|
|
assert sorted(extract_values("id", ds.iter_rows())) == [0, 1, 2, 3, 4]
|
|
|
|
|
|
def test_range(ray_start_regular_shared):
|
|
ds = ray.data.range(10, override_num_blocks=10)
|
|
assert ds._logical_plan.initial_num_blocks() == 10
|
|
assert ds.count() == 10
|
|
assert ds.take() == [{"id": i} for i in range(10)]
|
|
|
|
ds = ray.data.range(10, override_num_blocks=2)
|
|
assert ds._logical_plan.initial_num_blocks() == 2
|
|
assert ds.count() == 10
|
|
assert ds.take() == [{"id": i} for i in range(10)]
|
|
|
|
|
|
def test_empty_dataset(ray_start_regular_shared):
|
|
ds = ray.data.range(0)
|
|
assert ds.count() == 0
|
|
assert ds.size_bytes() == 0
|
|
assert ds.schema() is None
|
|
|
|
ds = ray.data.range(1)
|
|
ds = ds.filter(lambda x: x["id"] > 1)
|
|
ds = ds.materialize()
|
|
assert (
|
|
str(ds)
|
|
== "MaterializedDataset(num_blocks=1, num_rows=0, schema=Unknown schema)"
|
|
)
|
|
|
|
# Test map on empty dataset.
|
|
ds = ray.data.from_items([])
|
|
ds = ds.map(lambda x: x)
|
|
ds = ds.materialize()
|
|
assert ds.count() == 0
|
|
|
|
# Test filter on empty dataset.
|
|
ds = ray.data.from_items([])
|
|
ds = ds.filter(lambda: True)
|
|
ds = ds.materialize()
|
|
assert ds.count() == 0
|
|
|
|
|
|
@ray.remote
|
|
class Counter:
|
|
def __init__(self):
|
|
self.value = 0
|
|
|
|
def increment(self):
|
|
self.value += 1
|
|
return self.value
|
|
|
|
|
|
def test_cache_dataset(ray_start_regular_shared):
|
|
|
|
c = Counter.remote()
|
|
|
|
def inc(x):
|
|
ray.get(c.increment.remote())
|
|
return x
|
|
|
|
ds = ray.data.range(1)
|
|
ds = ds.map(inc)
|
|
assert not isinstance(ds, MaterializedDataset)
|
|
ds2 = ds.materialize()
|
|
assert isinstance(ds2, MaterializedDataset)
|
|
assert not isinstance(ds, MaterializedDataset)
|
|
|
|
# Tests standard iteration uses the materialized blocks.
|
|
for _ in range(10):
|
|
ds2.take_all()
|
|
|
|
assert ray.get(c.increment.remote()) == 2
|
|
|
|
# Tests streaming iteration uses the materialized blocks.
|
|
for _ in range(10):
|
|
list(ds2.streaming_split(1)[0].iter_batches())
|
|
|
|
assert ray.get(c.increment.remote()) == 3
|
|
|
|
|
|
def test_columns(ray_start_regular_shared):
|
|
ds = ray.data.range(1)
|
|
assert ds.columns() == ds.schema().names
|
|
assert ds.columns() == ["id"]
|
|
|
|
ds = ds.map(lambda x: x)
|
|
assert ds.columns(fetch_if_missing=False) is None
|
|
|
|
|
|
def test_schema_repr(ray_start_regular_shared):
|
|
ds = ray.data.from_items([{"text": "spam", "number": 0}])
|
|
# fmt: off
|
|
expected_repr = (
|
|
"Column Type\n"
|
|
"------ ----\n"
|
|
"text string\n"
|
|
"number int64"
|
|
)
|
|
# fmt:on
|
|
assert repr(ds.schema()) == expected_repr
|
|
|
|
ds = ray.data.from_items([{"long_column_name": "spam"}])
|
|
# fmt: off
|
|
expected_repr = (
|
|
"Column Type\n"
|
|
"------ ----\n"
|
|
"long_column_name string"
|
|
)
|
|
# fmt: on
|
|
assert repr(ds.schema()) == expected_repr
|
|
|
|
|
|
def _check_none_computed(ds):
|
|
# In streaming executor, ds.take() will not invoke partial execution
|
|
# in LazyBlocklist.
|
|
assert not ds._has_computed_output()
|
|
|
|
|
|
def test_lazy_loading_exponential_rampup(ray_start_regular_shared):
|
|
ds = ray.data.range(100, override_num_blocks=20)
|
|
_check_none_computed(ds)
|
|
assert extract_values("id", ds.take(10)) == list(range(10))
|
|
_check_none_computed(ds)
|
|
assert extract_values("id", ds.take(20)) == list(range(20))
|
|
_check_none_computed(ds)
|
|
assert extract_values("id", ds.take(30)) == list(range(30))
|
|
_check_none_computed(ds)
|
|
assert extract_values("id", ds.take(50)) == list(range(50))
|
|
_check_none_computed(ds)
|
|
assert extract_values("id", ds.take(100)) == list(range(100))
|
|
_check_none_computed(ds)
|
|
|
|
|
|
def test_dataset_repr_not_materialized(ray_start_regular_shared, restore_data_context):
|
|
ds = ray.data.range(5)
|
|
assert repr(ds) == (
|
|
"shape: (5, 1)\n"
|
|
"╭───────╮\n"
|
|
"│ id │\n"
|
|
"│ --- │\n"
|
|
"│ int64 │\n"
|
|
"╰───────╯\n"
|
|
"(Dataset isn't materialized)"
|
|
)
|
|
|
|
|
|
def test_dataset_repr_materialized(ray_start_regular_shared, restore_data_context):
|
|
materialized = ray.data.range(5).materialize()
|
|
assert repr(materialized) == (
|
|
"shape: (5, 1)\n"
|
|
"╭───────╮\n"
|
|
"│ id │\n"
|
|
"│ --- │\n"
|
|
"│ int64 │\n"
|
|
"╞═══════╡\n"
|
|
"│ 0 │\n"
|
|
"│ 1 │\n"
|
|
"│ 2 │\n"
|
|
"│ 3 │\n"
|
|
"│ 4 │\n"
|
|
"╰───────╯\n"
|
|
"(Showing 5 of 5 rows)"
|
|
)
|
|
|
|
|
|
def test_dataset_repr_gap(ray_start_regular_shared, restore_data_context):
|
|
ds_with_gap = ray.data.range(20).materialize()
|
|
assert repr(ds_with_gap) == (
|
|
"shape: (20, 1)\n"
|
|
"╭───────╮\n"
|
|
"│ id │\n"
|
|
"│ --- │\n"
|
|
"│ int64 │\n"
|
|
"╞═══════╡\n"
|
|
"│ 0 │\n"
|
|
"│ 1 │\n"
|
|
"│ 2 │\n"
|
|
"│ 3 │\n"
|
|
"│ 4 │\n"
|
|
"│ … │\n"
|
|
"│ 15 │\n"
|
|
"│ 16 │\n"
|
|
"│ 17 │\n"
|
|
"│ 18 │\n"
|
|
"│ 19 │\n"
|
|
"╰───────╯\n"
|
|
"(Showing 10 of 20 rows)"
|
|
)
|
|
|
|
|
|
def test_dataset_explain(ray_start_regular_shared, capsys):
|
|
ds = ray.data.range(10, override_num_blocks=10)
|
|
ds = ds.map(lambda x: x)
|
|
|
|
ds.explain()
|
|
captured = capsys.readouterr()
|
|
assert captured.out.strip() == (
|
|
"-------- Logical Plan --------\n"
|
|
"MapRows[Map(<lambda>)]\n"
|
|
"+- Read[ReadRange]\n"
|
|
"\n-------- Logical Plan (Optimized) --------\n"
|
|
"MapRows[Map(<lambda>)]\n"
|
|
"+- Read[ReadRange]\n"
|
|
"\n-------- Physical Plan --------\n"
|
|
"TaskPoolMapOperator[Map(<lambda>)]\n"
|
|
"+- TaskPoolMapOperator[ReadRange]\n"
|
|
" +- InputDataBuffer[Input]\n"
|
|
"\n-------- Physical Plan (Optimized) --------\n"
|
|
"TaskPoolMapOperator[ReadRange->Map(<lambda>)]\n"
|
|
"+- InputDataBuffer[Input]"
|
|
)
|
|
|
|
ds = ds.filter(lambda x: x["id"] > 0)
|
|
ds.explain()
|
|
captured = capsys.readouterr()
|
|
assert captured.out.strip() == (
|
|
"-------- Logical Plan --------\n"
|
|
"Filter[Filter(<lambda>)]\n"
|
|
"+- MapRows[Map(<lambda>)]\n"
|
|
" +- Read[ReadRange]\n"
|
|
"\n-------- Logical Plan (Optimized) --------\n"
|
|
"Filter[Filter(<lambda>)]\n"
|
|
"+- MapRows[Map(<lambda>)]\n"
|
|
" +- Read[ReadRange]\n"
|
|
"\n-------- Physical Plan --------\n"
|
|
"TaskPoolMapOperator[Filter(<lambda>)]\n"
|
|
"+- TaskPoolMapOperator[Map(<lambda>)]\n"
|
|
" +- TaskPoolMapOperator[ReadRange]\n"
|
|
" +- InputDataBuffer[Input]\n"
|
|
"\n-------- Physical Plan (Optimized) --------\n"
|
|
"TaskPoolMapOperator[ReadRange->Map(<lambda>)->Filter(<lambda>)]\n"
|
|
"+- InputDataBuffer[Input]"
|
|
)
|
|
ds = ds.random_shuffle().map(lambda x: x)
|
|
ds.explain()
|
|
captured = capsys.readouterr()
|
|
assert captured.out.strip() == (
|
|
"-------- Logical Plan --------\n"
|
|
"MapRows[Map(<lambda>)]\n"
|
|
"+- RandomShuffle[RandomShuffle]\n"
|
|
" +- Filter[Filter(<lambda>)]\n"
|
|
" +- MapRows[Map(<lambda>)]\n"
|
|
" +- Read[ReadRange]\n"
|
|
"\n-------- Logical Plan (Optimized) --------\n"
|
|
"MapRows[Map(<lambda>)]\n"
|
|
"+- RandomShuffle[RandomShuffle]\n"
|
|
" +- Filter[Filter(<lambda>)]\n"
|
|
" +- MapRows[Map(<lambda>)]\n"
|
|
" +- Read[ReadRange]\n"
|
|
"\n-------- Physical Plan --------\n"
|
|
"TaskPoolMapOperator[Map(<lambda>)]\n"
|
|
"+- AllToAllOperator[RandomShuffle]\n"
|
|
" +- TaskPoolMapOperator[Filter(<lambda>)]\n"
|
|
" +- TaskPoolMapOperator[Map(<lambda>)]\n"
|
|
" +- TaskPoolMapOperator[ReadRange]\n"
|
|
" +- InputDataBuffer[Input]\n"
|
|
"\n-------- Physical Plan (Optimized) --------\n"
|
|
"TaskPoolMapOperator[Map(<lambda>)]\n"
|
|
"+- AllToAllOperator[ReadRange->Map(<lambda>)->Filter(<lambda>)->RandomShuffle]\n"
|
|
" +- InputDataBuffer[Input]"
|
|
)
|
|
|
|
|
|
def test_convert_types(ray_start_regular_shared):
|
|
plain_ds = ray.data.range(1)
|
|
arrow_ds = plain_ds.map(lambda x: {"a": x["id"]})
|
|
assert arrow_ds.take() == [{"a": 0}]
|
|
assert "dict" in str(arrow_ds.map(lambda x: {"out": str(type(x))}).take()[0])
|
|
|
|
arrow_ds = ray.data.range(1)
|
|
assert arrow_ds.map(lambda x: {"out": "plain_{}".format(x["id"])}).take() == [
|
|
{"out": "plain_0"}
|
|
]
|
|
assert arrow_ds.map(lambda x: {"a": (x["id"],)}).take() == [{"a": [0]}]
|
|
|
|
|
|
def test_take_batch(ray_start_regular_shared):
|
|
ds = ray.data.range(10, override_num_blocks=2)
|
|
assert ds.take_batch(3)["id"].tolist() == [0, 1, 2]
|
|
assert ds.take_batch(6)["id"].tolist() == [0, 1, 2, 3, 4, 5]
|
|
assert ds.take_batch(100)["id"].tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
|
|
assert isinstance(ds.take_batch(3, batch_format="pandas"), pd.DataFrame)
|
|
assert isinstance(ds.take_batch(3, batch_format="numpy"), dict)
|
|
|
|
ds = ray.data.range_tensor(10, override_num_blocks=2)
|
|
assert np.all(ds.take_batch(3)["data"] == np.array([[0], [1], [2]]))
|
|
assert isinstance(ds.take_batch(3, batch_format="pandas"), pd.DataFrame)
|
|
assert isinstance(ds.take_batch(3, batch_format="numpy"), dict)
|
|
|
|
with pytest.raises(ValueError):
|
|
ray.data.range(0).take_batch()
|
|
|
|
|
|
def test_take_all(ray_start_regular_shared):
|
|
assert extract_values("id", ray.data.range(5).take_all()) == [0, 1, 2, 3, 4]
|
|
|
|
with pytest.raises(ValueError):
|
|
assert ray.data.range(5).take_all(4)
|
|
|
|
|
|
def test_union(ray_start_regular_shared, restore_data_context):
|
|
# Set aggregator CPU to 0 to avoid deadlock in resource-constrained test env.
|
|
# Without this, the shuffle task (1 CPU) + aggregator actor (~0.25 CPU) would
|
|
# exceed the 1 CPU available in the test cluster, causing scheduler deadlock.
|
|
restore_data_context.hash_aggregate_operator_actor_num_cpus_override = 0
|
|
|
|
ds = ray.data.range(20, override_num_blocks=10).materialize()
|
|
|
|
# Test lazy union.
|
|
ds = ds.union(ds, ds, ds, ds)
|
|
assert ds._logical_plan.initial_num_blocks() == 50
|
|
assert ds.count() == 100
|
|
assert ds.sum() == 950
|
|
|
|
ds = ds.union(ds)
|
|
assert ds.count() == 200
|
|
assert ds.sum() == (950 * 2)
|
|
|
|
# Test materialized union.
|
|
ds2 = ray.data.from_items([1, 2, 3, 4, 5])
|
|
assert ds2.count() == 5
|
|
ds2 = ds2.union(ds2)
|
|
assert ds2.count() == 10
|
|
ds2 = ds2.union(ds)
|
|
assert ds2.count() == 210
|
|
|
|
|
|
def test_block_builder_for_block(ray_start_regular_shared):
|
|
# pandas dataframe
|
|
builder = BlockBuilder.for_block(pd.DataFrame())
|
|
b1 = pd.DataFrame({"A": [1], "B": ["a"]})
|
|
builder.add_block(b1)
|
|
assert builder.build().equals(b1)
|
|
b2 = pd.DataFrame({"A": [2, 3], "B": ["c", "d"]})
|
|
builder.add_block(b2)
|
|
expected = pd.DataFrame({"A": [1, 2, 3], "B": ["a", "c", "d"]})
|
|
assert builder.build().equals(expected)
|
|
|
|
# pyarrow table
|
|
builder = BlockBuilder.for_block(pa.Table.from_arrays(list()))
|
|
b1 = pa.Table.from_pydict({"A": [1], "B": ["a"]})
|
|
builder.add_block(b1)
|
|
builder.build().equals(b1)
|
|
b2 = pa.Table.from_pydict({"A": [2, 3], "B": ["c", "d"]})
|
|
builder.add_block(b2)
|
|
expected = pa.Table.from_pydict({"A": [1, 2, 3], "B": ["a", "c", "d"]})
|
|
builder.build().equals(expected)
|
|
|
|
# wrong type
|
|
with pytest.raises(TypeError):
|
|
BlockBuilder.for_block(str())
|
|
|
|
|
|
def test_len(ray_start_regular_shared):
|
|
ds = ray.data.range(1)
|
|
with pytest.raises(AttributeError):
|
|
len(ds)
|
|
|
|
|
|
def test_pandas_block_select():
|
|
df = pd.DataFrame({"one": [10, 11, 12], "two": [11, 12, 13], "three": [14, 15, 16]})
|
|
block_accessor = BlockAccessor.for_block(df)
|
|
|
|
block = block_accessor.select(["two"])
|
|
assert block.equals(df[["two"]])
|
|
|
|
block = block_accessor.select(["two", "one"])
|
|
assert block.equals(df[["two", "one"]])
|
|
|
|
with pytest.raises(ValueError):
|
|
block = block_accessor.select([lambda x: x % 3, "two"])
|
|
|
|
|
|
# NOTE: All tests above share a Ray cluster, while the tests below do not. These
|
|
# tests should only be carefully reordered to retain this invariant!
|
|
|
|
|
|
def test_read_write_local_node_ray_client(ray_start_cluster_enabled):
|
|
cluster = ray_start_cluster_enabled
|
|
cluster.add_node(num_cpus=4)
|
|
cluster.head_node._ray_params.ray_client_server_port = "10004"
|
|
cluster.head_node.start_ray_client_server()
|
|
address = "ray://localhost:10004"
|
|
|
|
import tempfile
|
|
|
|
data_path = tempfile.mkdtemp()
|
|
df = pd.DataFrame({"one": list(range(0, 10)), "two": list(range(10, 20))})
|
|
path = os.path.join(data_path, "test.parquet")
|
|
df.to_parquet(path)
|
|
|
|
# Read/write from Ray Client will result in error.
|
|
ray.init(address)
|
|
with pytest.raises(ValueError):
|
|
ds = ray.data.read_parquet("local://" + path).materialize()
|
|
ds = ray.data.from_pandas(df)
|
|
with pytest.raises(ValueError):
|
|
ds.write_parquet("local://" + data_path).materialize()
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info >= (3, 12), reason="No tensorflow for Python 3.12+"
|
|
)
|
|
def test_read_warning_large_parallelism(
|
|
ray_start_regular_shared, propagate_logs, caplog
|
|
):
|
|
with caplog.at_level(logging.WARNING, logger="ray.data.read_api"):
|
|
ray.data.range(5000, override_num_blocks=5000).materialize()
|
|
assert (
|
|
"The requested number of read blocks of 5000 is "
|
|
"more than 4x the number of available CPU slots in the cluster" in caplog.text
|
|
), caplog.text
|
|
|
|
|
|
def test_read_write_local_node(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(
|
|
resources={"bar:1": 100},
|
|
num_cpus=10,
|
|
_system_config={"max_direct_call_object_size": 0},
|
|
)
|
|
cluster.add_node(resources={"bar:2": 100}, num_cpus=10)
|
|
cluster.add_node(resources={"bar:3": 100}, num_cpus=10)
|
|
|
|
ray.shutdown()
|
|
ray.init(cluster.address)
|
|
|
|
import os
|
|
import tempfile
|
|
|
|
data_path = tempfile.mkdtemp()
|
|
num_files = 5
|
|
for idx in range(num_files):
|
|
df = pd.DataFrame(
|
|
{"one": list(range(idx, idx + 10)), "two": list(range(idx + 10, idx + 20))}
|
|
)
|
|
path = os.path.join(data_path, f"test{idx}.parquet")
|
|
df.to_parquet(path)
|
|
|
|
ctx = ray.data.context.DataContext.get_current()
|
|
ctx.read_write_local_node = True
|
|
|
|
def check_dataset_is_local(ds):
|
|
bundles = ds.iter_internal_ref_bundles()
|
|
block_refs = _ref_bundles_iterator_to_block_refs_list(bundles)
|
|
ray.wait(block_refs, num_returns=len(block_refs), fetch_local=False)
|
|
location_data = ray.experimental.get_object_locations(block_refs)
|
|
locations = []
|
|
for block in block_refs:
|
|
locations.extend(location_data[block]["node_ids"])
|
|
assert set(locations) == {ray.get_runtime_context().get_node_id()}
|
|
|
|
local_path = "local://" + data_path
|
|
# Plain read.
|
|
ds = ray.data.read_parquet(local_path).materialize()
|
|
check_dataset_is_local(ds)
|
|
|
|
# SPREAD scheduling got overridden when read local scheme.
|
|
ds = ray.data.read_parquet(
|
|
local_path, ray_remote_args={"scheduling_strategy": "SPREAD"}
|
|
).materialize()
|
|
check_dataset_is_local(ds)
|
|
|
|
# With fusion.
|
|
ds = (
|
|
ray.data.read_parquet(local_path, override_num_blocks=1)
|
|
.map(lambda x: x)
|
|
.materialize()
|
|
)
|
|
check_dataset_is_local(ds)
|
|
|
|
# Write back to local scheme.
|
|
output = os.path.join(local_path, "test_read_write_local_node")
|
|
ds.write_parquet(output)
|
|
assert "1 nodes used" in ds.stats(), ds.stats()
|
|
ray.data.read_parquet(output).take_all() == ds.take_all()
|
|
|
|
# Mixing paths of local and non-local scheme is invalid.
|
|
with pytest.raises(ValueError):
|
|
ds = ray.data.read_parquet(
|
|
[local_path + "/test1.parquet", data_path + "/test2.parquet"]
|
|
).materialize()
|
|
with pytest.raises(ValueError):
|
|
ds = ray.data.read_parquet(
|
|
[local_path + "/test1.parquet", "example://iris.parquet"]
|
|
).materialize()
|
|
with pytest.raises(ValueError):
|
|
ds = ray.data.read_parquet(
|
|
["example://iris.parquet", local_path + "/test1.parquet"]
|
|
).materialize()
|
|
|
|
|
|
def test_validate_head_node_resources_zero_head_cpu(ray_start_cluster):
|
|
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=0)
|
|
cluster.wait_for_nodes()
|
|
|
|
ray.shutdown()
|
|
ray.init(address=cluster.address)
|
|
|
|
with pytest.raises(ValueError, match=r"head node doesn't have enough resources"):
|
|
_validate_head_node_resources_for_local_scheduling(
|
|
{}, op_description="read local"
|
|
)
|
|
|
|
|
|
class FlakyCSVDatasource(CSVDatasource):
|
|
def __init__(self, paths, **csv_datasource_kwargs):
|
|
super().__init__(paths, **csv_datasource_kwargs)
|
|
|
|
self.counter = Counter.remote()
|
|
|
|
def _read_stream(self, f: "pa.NativeFile", path: str):
|
|
count = self.counter.increment.remote()
|
|
if ray.get(count) == 1:
|
|
raise ValueError("oops")
|
|
else:
|
|
for block in CSVDatasource._read_stream(self, f, path):
|
|
yield block
|
|
|
|
|
|
class FlakyCSVDatasink(CSVDatasink):
|
|
def __init__(self, path, **csv_datasink_kwargs):
|
|
super().__init__(path, **csv_datasink_kwargs)
|
|
|
|
self.counter = Counter.remote()
|
|
|
|
def write_block_to_file(self, block: BlockAccessor, file):
|
|
count = self.counter.increment.remote()
|
|
if ray.get(count) == 1:
|
|
raise ValueError("oops")
|
|
else:
|
|
super().write_block_to_file(block, file)
|
|
|
|
|
|
def test_datasource(ray_start_regular):
|
|
source = ray.data.datasource.RandomIntRowDatasource(n=10, num_columns=2)
|
|
assert len(ray.data.read_datasource(source).take()) == 10
|
|
source = RangeDatasource(n=10)
|
|
assert extract_values(
|
|
"value",
|
|
ray.data.read_datasource(source).take(),
|
|
) == list(range(10))
|
|
|
|
|
|
@pytest.mark.skip(reason="")
|
|
def test_polars_lazy_import(shutdown_only):
|
|
import sys
|
|
|
|
ctx = ray.data.context.DataContext.get_current()
|
|
|
|
try:
|
|
original_use_polars = ctx.use_polars
|
|
ctx.use_polars = True
|
|
|
|
num_items = 100
|
|
parallelism = 4
|
|
ray.init(num_cpus=4)
|
|
|
|
@ray.remote
|
|
def f(should_import_polars):
|
|
# Sleep to spread the tasks.
|
|
time.sleep(1)
|
|
polars_imported = "polars" in sys.modules.keys()
|
|
return polars_imported == should_import_polars
|
|
|
|
# We should not use polars for non-Arrow sort.
|
|
_ = ray.data.range(num_items, override_num_blocks=parallelism).sort()
|
|
assert all(ray.get([f.remote(False) for _ in range(parallelism)]))
|
|
|
|
a = range(100)
|
|
dfs = []
|
|
partition_size = num_items // parallelism
|
|
for i in range(parallelism):
|
|
dfs.append(
|
|
pd.DataFrame({"a": a[i * partition_size : (i + 1) * partition_size]})
|
|
)
|
|
# At least one worker should have imported polars.
|
|
_ = (
|
|
ray.data.from_pandas(dfs)
|
|
.map_batches(lambda t: t, batch_format="pyarrow", batch_size=None)
|
|
.sort(key="a")
|
|
.materialize()
|
|
)
|
|
assert any(ray.get([f.remote(True) for _ in range(parallelism)]))
|
|
|
|
finally:
|
|
ctx.use_polars = original_use_polars
|
|
|
|
|
|
def test_batch_formats(shutdown_only):
|
|
ds = ray.data.range(100)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format=None))), pa.Table)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="default"))), dict)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="pandas"))), pd.DataFrame)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="pyarrow"))), pa.Table)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="numpy"))), dict)
|
|
|
|
ds = ray.data.range_tensor(100)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format=None))), pa.Table)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="default"))), dict)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="pandas"))), pd.DataFrame)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="pyarrow"))), pa.Table)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="numpy"))), dict)
|
|
|
|
df = pd.DataFrame({"foo": ["a", "b"], "bar": [0, 1]})
|
|
ds = ray.data.from_pandas(df)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format=None))), pd.DataFrame)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="default"))), dict)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="pandas"))), pd.DataFrame)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="pyarrow"))), pa.Table)
|
|
assert isinstance(next(iter(ds.iter_batches(batch_format="numpy"))), dict)
|
|
|
|
|
|
def test_dataset_schema_after_read_stats(ray_start_cluster):
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=1)
|
|
ray.init(cluster.address)
|
|
cluster.add_node(num_cpus=1, resources={"foo": 1})
|
|
ds = ray.data.read_csv(
|
|
"example://iris.csv", ray_remote_args={"resources": {"foo": 1}}
|
|
)
|
|
schema = ds.schema()
|
|
ds.stats()
|
|
assert schema == ds.schema()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|