wehub-resource-sync
94 lines
3.0 KiB
Python
94 lines
3.0 KiB
Python
import pytest
|
|
import torch
|
|
|
|
import ray
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.data.tests.util import extract_values
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
@pytest.mark.parametrize("local_read", [True, False])
|
|
def test_from_torch_map_style_dataset(ray_start_10_cpus_shared, local_read):
|
|
class StubDataset(torch.utils.data.Dataset):
|
|
def __len__(self):
|
|
return 1
|
|
|
|
def __getitem__(self, index):
|
|
return index
|
|
|
|
torch_dataset = StubDataset()
|
|
|
|
ray_dataset = ray.data.from_torch(torch_dataset, local_read=local_read)
|
|
|
|
actual_data = ray_dataset.take_all()
|
|
assert actual_data == [{"item": 0}]
|
|
|
|
|
|
def test_from_torch_iterable_style_dataset(ray_start_10_cpus_shared):
|
|
class StubIterableDataset(torch.utils.data.IterableDataset):
|
|
def __len__(self):
|
|
return 1
|
|
|
|
def __iter__(self):
|
|
return iter([0])
|
|
|
|
iter_torch_dataset = StubIterableDataset()
|
|
|
|
ray_dataset = ray.data.from_torch(iter_torch_dataset)
|
|
|
|
actual_data = ray_dataset.take_all()
|
|
assert actual_data == [{"item": 0}]
|
|
|
|
|
|
@pytest.mark.parametrize("local_read", [True, False])
|
|
def test_from_torch_boundary_conditions(ray_start_10_cpus_shared, local_read):
|
|
"""
|
|
Tests that from_torch respects __len__ for map-style datasets
|
|
"""
|
|
from torch.utils.data import Dataset
|
|
|
|
class BoundaryTestMapDataset(Dataset):
|
|
"""A map-style dataset where __len__ is less than the underlying data size."""
|
|
|
|
def __init__(self, data, length):
|
|
super().__init__()
|
|
self._data = data
|
|
self._length = length
|
|
assert self._length <= len(
|
|
self._data
|
|
), "Length must be <= data size to properly test boundary conditions"
|
|
|
|
def __len__(self):
|
|
return self._length
|
|
|
|
def __getitem__(self, index):
|
|
if not (0 <= index < self._length):
|
|
# Note: don't use IndexError because we want to fail clearly if
|
|
# Ray Data tries to access beyond __len__ - 1
|
|
raise RuntimeError(
|
|
f"Index {index} out of bounds for dataset with length {self._length}"
|
|
)
|
|
return self._data[index]
|
|
|
|
source_data = list(range(10))
|
|
dataset_len = 8 # Intentionally less than len(source_data)
|
|
|
|
# --- Test MapDataset ---
|
|
map_ds = BoundaryTestMapDataset(source_data, dataset_len)
|
|
# Expected data only includes elements up to dataset_len - 1
|
|
expected_items = source_data[:dataset_len]
|
|
|
|
ray_ds_map = ray.data.from_torch(map_ds, local_read=local_read)
|
|
actual_items_map = extract_values("item", list(ray_ds_map.take_all()))
|
|
|
|
# This assertion verifies that ray_ds_map didn't try to access index 8 or 9,
|
|
# which would have raised an IndexError in BoundaryTestMapDataset.__getitem__
|
|
assert actual_items_map == expected_items
|
|
assert len(actual_items_map) == dataset_len
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|