wehub-resource-sync
531 lines
18 KiB
Python
531 lines
18 KiB
Python
import sys
|
|
import threading
|
|
from typing import Dict
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
import torch
|
|
|
|
import ray
|
|
|
|
if sys.version_info <= (3, 12):
|
|
# Skip this test for Python 3.12+ due to to incompatibility tensorflow
|
|
import tensorflow as tf
|
|
|
|
|
|
def build_model():
|
|
import tensorflow as tf
|
|
|
|
model = tf.keras.Sequential(
|
|
[
|
|
tf.keras.layers.InputLayer(input_shape=()),
|
|
tf.keras.layers.Flatten(),
|
|
tf.keras.layers.Dense(1),
|
|
]
|
|
)
|
|
model.compile(optimizer="sgd", loss="mse")
|
|
return model
|
|
|
|
|
|
def test_basic_dataset(ray_start_regular_shared):
|
|
ds = ray.data.range(100)
|
|
it = ds.iterator()
|
|
for _ in range(2):
|
|
result = []
|
|
for batch in it.iter_batches():
|
|
batch = batch["id"]
|
|
result += batch.tolist()
|
|
assert result == list(range(100))
|
|
|
|
# TODO(swang): This check currently fails nondeterministically because
|
|
# stats are stored in an actor.
|
|
# https://github.com/ray-project/ray/issues/31571
|
|
# assert it.stats() == ds.stats()
|
|
|
|
|
|
def test_basic_dataset_multi_use_iterator(ray_start_regular_shared):
|
|
"""Tests that the iterable outputted by `iter_batches` can be used
|
|
multiple times."""
|
|
ds = ray.data.range(100)
|
|
it = ds.iterator().iter_batches()
|
|
for _ in range(2):
|
|
result = []
|
|
for batch in it:
|
|
batch = batch["id"]
|
|
result += batch.tolist()
|
|
assert result == list(range(100))
|
|
|
|
|
|
def test_basic_dataset_preemption(ray_start_regular_shared):
|
|
"""Tests that the iterable outputted by ``iter_batches``
|
|
can be used multiple times even if it is preempted during iteration."""
|
|
|
|
ds = ray.data.range(100)
|
|
it = ds.iterator().iter_batches(batch_size=50)
|
|
for _ in range(2):
|
|
result = []
|
|
for i, batch in enumerate(it):
|
|
if i > 0:
|
|
break
|
|
batch = batch["id"]
|
|
result += batch.tolist()
|
|
|
|
assert result == list(range(50))
|
|
|
|
|
|
def test_iter_batches_early_exit_shuts_down_executor(ray_start_regular_shared):
|
|
"""Tests that breaking out of ``iter_batches`` early shuts down the
|
|
streaming executor deterministically. Otherwise the executor's worker
|
|
thread keeps producing blocks that pile up in the object store and
|
|
holds resources that can starve other datasets (e.g., a validation
|
|
dataset waiting to run)."""
|
|
|
|
ds = ray.data.range(1000, override_num_blocks=20)
|
|
it = ds.iterator()
|
|
|
|
for i, _ in enumerate(it.iter_batches(batch_size=10)):
|
|
if i == 0:
|
|
break
|
|
|
|
executor = ds._current_executor
|
|
assert executor is not None
|
|
assert executor._shutdown is True
|
|
|
|
|
|
def test_iter_batches_early_break_flushes_metrics(ray_start_regular_shared):
|
|
"""Tests that an early ``break`` in the training loop still records
|
|
``iter_total_s`` and flushes metrics via ``update_iteration_metrics``."""
|
|
|
|
from ray.data._internal.stats import _StatsManager
|
|
|
|
ds = ray.data.range(100, override_num_blocks=5)
|
|
it = ds.iterator()
|
|
|
|
captured_stats = []
|
|
orig = _StatsManager.update_iteration_metrics
|
|
|
|
def spy(stats, dataset_tag):
|
|
captured_stats.append(stats)
|
|
return orig(stats, dataset_tag)
|
|
|
|
with patch.object(_StatsManager, "update_iteration_metrics", spy):
|
|
for i, _ in enumerate(it.iter_batches(batch_size=10)):
|
|
if i == 0:
|
|
break
|
|
|
|
# finally block should have called update_iteration_metrics
|
|
assert len(captured_stats) > 0
|
|
# iter_total_s was recorded even on early break
|
|
assert captured_stats[-1].iter_total_s.get() > 0
|
|
assert captured_stats[-1].iter_batches_total > 0
|
|
|
|
|
|
def test_iter_batches_full_iteration_shuts_down_executor(ray_start_regular_shared):
|
|
"""Tests that fully iterating ``iter_batches`` shuts down the
|
|
streaming executor (regression guard for the early-exit cleanup
|
|
path, which adds an idempotent shutdown to the iterator)."""
|
|
|
|
ds = ray.data.range(100, override_num_blocks=5)
|
|
it = ds.iterator()
|
|
|
|
for _ in it.iter_batches(batch_size=10):
|
|
pass
|
|
|
|
executor = ds._current_executor
|
|
assert executor is not None
|
|
assert executor._shutdown is True
|
|
|
|
|
|
def test_iter_batches_exception_shuts_down_executor(ray_start_regular_shared):
|
|
"""Tests that an exception raised inside the user's iteration loop
|
|
still triggers executor shutdown via the iterator's ``finally``."""
|
|
|
|
ds = ray.data.range(1000, override_num_blocks=20)
|
|
it = ds.iterator()
|
|
|
|
class _Sentinel(Exception):
|
|
pass
|
|
|
|
with pytest.raises(_Sentinel):
|
|
for _ in it.iter_batches(batch_size=10):
|
|
raise _Sentinel()
|
|
|
|
executor = ds._current_executor
|
|
assert executor is not None
|
|
assert executor._shutdown is True
|
|
|
|
|
|
def test_iter_batches_close_on_held_iterator_shuts_down_executor(
|
|
ray_start_regular_shared,
|
|
):
|
|
"""Tests that ``it.close()`` shuts down the executor when the caller
|
|
holds an explicit reference to the iterator. Without ``close()``,
|
|
Python defers cleanup until the reference is dropped — ``break``
|
|
inside the for-loop wouldn't fire ``GeneratorExit`` on a held
|
|
reference. Some libraries (e.g. PyTorch Lightning's batch fetchers)
|
|
keep an ``iter()`` reference internally; this is the documented
|
|
eager-cleanup escape hatch."""
|
|
|
|
ds = ray.data.range(1000, override_num_blocks=20)
|
|
batches = ds.iterator().iter_batches(batch_size=10)
|
|
|
|
it = iter(batches)
|
|
next(it)
|
|
it.close()
|
|
|
|
executor = ds._current_executor
|
|
assert executor is not None
|
|
assert executor._shutdown is True
|
|
|
|
|
|
def test_basic_dataset_iter_rows(ray_start_regular_shared):
|
|
ds = ray.data.range(100)
|
|
it = ds.iterator()
|
|
for _ in range(2):
|
|
result = []
|
|
for row in it.iter_rows():
|
|
row = row["id"]
|
|
result.append(row)
|
|
assert result == list(range(100))
|
|
|
|
# TODO(swang): This check currently fails nondeterministically because
|
|
# stats are stored in an actor.
|
|
# https://github.com/ray-project/ray/issues/31571
|
|
# assert it.stats() == ds.stats()
|
|
|
|
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_tf_conversion(ray_start_regular_shared, include_additional_columns):
|
|
ds = ray.data.range(5)
|
|
it = ds.iterator()
|
|
|
|
if include_additional_columns:
|
|
tf_dataset = it.to_tf("id", "id", additional_columns="id")
|
|
else:
|
|
tf_dataset = it.to_tf("id", "id")
|
|
|
|
for i, row in enumerate(tf_dataset):
|
|
assert all(row[0] == i)
|
|
assert all(row[1] == i)
|
|
assert isinstance(row[0], tf.Tensor)
|
|
assert isinstance(row[1], tf.Tensor)
|
|
if include_additional_columns:
|
|
assert all(row[2] == i)
|
|
assert isinstance(row[2], tf.Tensor)
|
|
|
|
|
|
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
|
def test_tf_e2e(ray_start_regular_shared, include_additional_columns):
|
|
ds = ray.data.range(5)
|
|
it = ds.iterator()
|
|
model = build_model()
|
|
if include_additional_columns:
|
|
model.fit(it.to_tf("id", "id", additional_columns="id"), epochs=3)
|
|
else:
|
|
model.fit(it.to_tf("id", "id"), epochs=3)
|
|
|
|
|
|
def test_torch_conversion(ray_start_regular_shared):
|
|
ds = ray.data.range(5)
|
|
it = ds.iterator()
|
|
it._iter_batches = MagicMock()
|
|
|
|
for batch in it.iter_torch_batches():
|
|
assert isinstance(batch["id"], torch.Tensor)
|
|
assert batch["id"].tolist() == list(range(5))
|
|
|
|
# When collate_fn is not specified, check that the default
|
|
# `_collate_fn` (handles formatting and Tensor creation)
|
|
# and `_finalize_fn` (handles host to device data transfer)
|
|
# are used in `DataIterator.iter_batches()`.
|
|
iter_batches_calls_kwargs = [a.kwargs for a in it._iter_batches.call_args_list]
|
|
assert all(
|
|
callable(kwargs["_collate_fn"]) and callable(kwargs["_finalize_fn"])
|
|
for kwargs in iter_batches_calls_kwargs
|
|
), iter_batches_calls_kwargs
|
|
|
|
|
|
def test_torch_multi_use_iterator(ray_start_regular_shared):
|
|
"""Tests that the iterator outputted by `iter_torch_batches` can be used
|
|
multiple times."""
|
|
ds = ray.data.range(5)
|
|
it = ds.iterator().iter_torch_batches()
|
|
|
|
for _ in range(2):
|
|
for batch in it:
|
|
assert isinstance(batch["id"], torch.Tensor)
|
|
assert batch["id"].tolist() == list(range(5))
|
|
|
|
|
|
def test_torch_conversion_collate_fn(ray_start_regular_shared):
|
|
def collate_fn(batch: Dict[str, np.ndarray]):
|
|
return torch.as_tensor(batch["id"] + 5)
|
|
|
|
ds = ray.data.range(5)
|
|
it = ds.iterator()
|
|
for batch in it.iter_torch_batches(collate_fn=collate_fn):
|
|
assert isinstance(batch, torch.Tensor)
|
|
assert batch.tolist() == list(range(5, 10))
|
|
|
|
# Should fail.
|
|
with pytest.raises(ValueError):
|
|
for batch in it.iter_torch_batches(collate_fn=collate_fn, dtypes=torch.float32):
|
|
assert isinstance(batch, torch.Tensor)
|
|
assert batch.tolist() == list(range(5, 10))
|
|
|
|
with pytest.raises(ValueError):
|
|
for batch in it.iter_torch_batches(collate_fn=collate_fn, device="cpu"):
|
|
assert isinstance(batch, torch.Tensor)
|
|
assert batch.tolist() == list(range(5, 10))
|
|
|
|
# Test that we don't automatically set device if collate_fn is specified.
|
|
with patch("ray.train.torch.get_device", lambda: torch.device("cuda")):
|
|
devices = ray.train.torch.get_device()
|
|
assert devices.type == "cuda"
|
|
|
|
it._iter_batches = MagicMock()
|
|
for batch in it.iter_torch_batches(collate_fn=collate_fn):
|
|
assert batch.device.type == "cpu"
|
|
assert isinstance(batch, torch.Tensor)
|
|
assert batch.tolist() == list(range(5, 10))
|
|
|
|
# Check that _finalize_fn is always used in `DataIterator.iter_batches()`.
|
|
iter_batches_calls_kwargs = [a.kwargs for a in it._iter_batches.call_args_list]
|
|
assert all(
|
|
kwargs["_finalize_fn"] is not None for kwargs in iter_batches_calls_kwargs
|
|
), iter_batches_calls_kwargs
|
|
|
|
|
|
@pytest.fixture(params=["regular", "chunked"])
|
|
def null_array_table(request):
|
|
"""Fixture that returns a PyArrow table with either a regular or chunked null array."""
|
|
if request.param == "regular":
|
|
# Regular array
|
|
return pa.table({"fruit_apple": pa.array([None, None, None], type=pa.null())})
|
|
else:
|
|
# Chunked array
|
|
return pa.table(
|
|
{
|
|
"fruit_apple": pa.chunked_array(
|
|
[
|
|
pa.array([None], type=pa.null()),
|
|
pa.array([None, None], type=pa.null()),
|
|
]
|
|
)
|
|
}
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def image_dataset(ray_start_regular_shared):
|
|
"""Fixture that returns a Ray dataset with image-like columns."""
|
|
num_rows = 10
|
|
num_images = 8
|
|
image_size_flat = 224 * 224 * 3 # Flattened image size
|
|
rows = []
|
|
for _ in range(num_rows):
|
|
row = {}
|
|
for i in range(num_images):
|
|
row[f"image_{i}"] = np.random.randint(
|
|
0, 255, size=image_size_flat, dtype=np.uint8
|
|
)
|
|
rows.append(row)
|
|
|
|
table = pa.Table.from_pylist(rows)
|
|
return ray.data.from_arrow([table])
|
|
|
|
|
|
def test_torch_conversion_null_type(ray_start_regular_shared, null_array_table):
|
|
"""Test iter_torch_batches with a PyArrow table containing null type arrays."""
|
|
ds = ray.data.from_arrow(null_array_table)
|
|
it = ds.iterator()
|
|
for batch in it.iter_torch_batches():
|
|
assert isinstance(batch, dict)
|
|
assert "fruit_apple" in batch
|
|
assert isinstance(batch["fruit_apple"], torch.Tensor)
|
|
assert torch.isnan(batch["fruit_apple"]).all()
|
|
assert batch["fruit_apple"].shape == (3,)
|
|
|
|
|
|
@pytest.fixture
|
|
def collate_fns_with_and_without_threading():
|
|
"""Fixture that provides DefaultCollateFn with and without threading."""
|
|
from ray.data.collate_fn import DefaultCollateFn
|
|
|
|
collate_fn_with_threading = DefaultCollateFn(num_workers=4)
|
|
collate_fn_no_threading = DefaultCollateFn(num_workers=0)
|
|
|
|
return {
|
|
"with_threading": collate_fn_with_threading,
|
|
"without_threading": collate_fn_no_threading,
|
|
}
|
|
|
|
|
|
def _collect_batches(ds, collate_fn, batch_size=5):
|
|
"""Helper function to collect batches from iter_torch_batches."""
|
|
batches = []
|
|
for batch in ds.iterator().iter_torch_batches(
|
|
collate_fn=collate_fn, batch_size=batch_size
|
|
):
|
|
batches.append(batch)
|
|
return batches
|
|
|
|
|
|
def test_torch_conversion_default_collate_fn_threading(
|
|
ray_start_regular_shared,
|
|
image_dataset,
|
|
collate_fns_with_and_without_threading,
|
|
):
|
|
"""Test DefaultCollateFn with/without threading produces same results."""
|
|
ds = image_dataset
|
|
collate_fns = collate_fns_with_and_without_threading
|
|
|
|
batches_with_threading = _collect_batches(ds, collate_fns["with_threading"])
|
|
batches_no_threading = _collect_batches(ds, collate_fns["without_threading"])
|
|
|
|
# Verify results are the same
|
|
assert len(batches_with_threading) == len(batches_no_threading)
|
|
for b1, b2 in zip(batches_with_threading, batches_no_threading):
|
|
assert set(b1.keys()) == set(b2.keys())
|
|
for col in b1.keys():
|
|
assert len(b1[col]) == len(b2[col])
|
|
for t1, t2 in zip(b1[col], b2[col]):
|
|
assert torch.equal(t1, t2)
|
|
|
|
|
|
class TestToTorch:
|
|
def test_features_only(self, ray_start_regular_shared):
|
|
ds = ray.data.from_items([{"a": 1, "b": 2}, {"a": 3, "b": 4}])
|
|
it = ds.iterator().to_torch(batch_size=2)
|
|
|
|
features, labels = next(iter(it))
|
|
|
|
assert labels is None
|
|
assert isinstance(features, torch.Tensor)
|
|
assert features.shape == (2, 2)
|
|
|
|
def test_with_label_column(self, ray_start_regular_shared):
|
|
ds = ray.data.from_items([{"x": 1, "y": 10}, {"x": 2, "y": 20}])
|
|
it = ds.iterator().to_torch(label_column="y", batch_size=2)
|
|
|
|
features, labels = next(iter(it))
|
|
|
|
assert features.shape == (2, 1)
|
|
assert labels.shape == (2, 1)
|
|
assert sorted(labels.flatten().tolist()) == [10, 20]
|
|
|
|
def test_squeezed_label(self, ray_start_regular_shared):
|
|
ds = ray.data.from_items([{"x": 1, "y": 10}, {"x": 2, "y": 20}])
|
|
it = ds.iterator().to_torch(
|
|
label_column="y", batch_size=2, unsqueeze_label_tensor=False
|
|
)
|
|
|
|
features, labels = next(iter(it))
|
|
|
|
assert labels.shape == (2,)
|
|
|
|
def test_feature_columns_as_list_of_lists(self, ray_start_regular_shared):
|
|
ds = ray.data.from_pandas(pd.DataFrame({"a": [[1, 2], [3, 4], [5, 6]]}))
|
|
it = ds.iterator().to_torch(feature_columns=["a"], batch_size=2)
|
|
|
|
features, labels = next(iter(it))
|
|
|
|
assert labels is None
|
|
assert features.shape == (2, 1, 2)
|
|
|
|
def test_feature_columns_subset(self, ray_start_regular_shared):
|
|
ds = ray.data.from_items([{"a": 1, "b": 2, "c": 3}])
|
|
it = ds.iterator().to_torch(feature_columns=["a", "c"], batch_size=1)
|
|
|
|
features, labels = next(iter(it))
|
|
|
|
assert features.shape == (1, 2)
|
|
|
|
def test_feature_columns_as_dict(self, ray_start_regular_shared):
|
|
ds = ray.data.from_items([{"a": 1, "b": 2, "c": 3}])
|
|
it = ds.iterator().to_torch(
|
|
feature_columns={"group1": ["a"], "group2": ["b", "c"]},
|
|
batch_size=1,
|
|
)
|
|
|
|
features, labels = next(iter(it))
|
|
|
|
assert isinstance(features, dict)
|
|
assert set(features.keys()) == {"group1", "group2"}
|
|
assert features["group1"].shape == (1, 1)
|
|
assert features["group2"].shape == (1, 2)
|
|
|
|
def test_explicit_dtype(self, ray_start_regular_shared):
|
|
ds = ray.data.from_items([{"a": 1, "b": 2}])
|
|
it = ds.iterator().to_torch(feature_column_dtypes=torch.float32, batch_size=1)
|
|
|
|
features, _ = next(iter(it))
|
|
|
|
assert features.dtype == torch.float32
|
|
|
|
def test_returns_iterable_dataset(self, ray_start_regular_shared):
|
|
"""to_torch should return a torch IterableDataset."""
|
|
ds = ray.data.range(1)
|
|
|
|
result = ds.iterator().to_torch(batch_size=1)
|
|
|
|
assert isinstance(result, torch.utils.data.IterableDataset)
|
|
|
|
|
|
def test_to_torch_emits_deprecation_warning(ray_start_regular_shared):
|
|
with pytest.warns(DeprecationWarning):
|
|
ray.data.range(1).iterator().to_torch()
|
|
|
|
|
|
@pytest.mark.parametrize("should_equalize", [True, False])
|
|
def test_iterator_to_materialized_dataset(ray_start_regular_shared, should_equalize):
|
|
"""Tests that `DataIterator.materialize` fully consumes the
|
|
iterator and returns a `MaterializedDataset` view of the data
|
|
that can be used to interact with the full dataset
|
|
(e.g. load it all into memory)."""
|
|
ds = ray.data.range(10)
|
|
num_splits = 2
|
|
iters = ds.streaming_split(num_splits, equal=should_equalize)
|
|
|
|
def consume_in_parallel(fn):
|
|
runners = [
|
|
threading.Thread(target=fn, args=(it, i)) for i, it in enumerate(iters)
|
|
]
|
|
[r.start() for r in runners]
|
|
[r.join() for r in runners]
|
|
|
|
materialized_ds = {}
|
|
shard_data = {}
|
|
|
|
def materialize(it, i):
|
|
materialized_ds[i] = it.materialize()
|
|
|
|
def iter_batches(it, i):
|
|
data = []
|
|
for batch in it.iter_batches():
|
|
data.extend(batch["id"].tolist())
|
|
shard_data[i] = data
|
|
|
|
consume_in_parallel(materialize)
|
|
consume_in_parallel(iter_batches)
|
|
|
|
# Check that the materialized datasets contain the same data as the
|
|
# original iterators.
|
|
for i in range(num_splits):
|
|
assert sorted(materialized_ds[i].to_pandas()["id"].tolist()) == sorted(
|
|
shard_data[i]
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
if sys.version_info >= (3, 12):
|
|
# Skip this test for Python 3.12+ due to to incompatibility tensorflow
|
|
sys.exit(0)
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|