wehub-resource-sync
1040 lines
37 KiB
Python
1040 lines
37 KiB
Python
import itertools
|
|
import math
|
|
import random
|
|
import threading
|
|
import time
|
|
from typing import Any, List, Tuple
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.equalize import _equalize
|
|
from ray.data._internal.execution.interfaces import BlockEntry, RefBundle
|
|
from ray.data._internal.execution.interfaces.ref_bundle import (
|
|
_ref_bundles_iterator_to_block_refs_list,
|
|
)
|
|
from ray.data._internal.logical.interfaces import LogicalPlan
|
|
from ray.data._internal.logical.operators import InputData
|
|
from ray.data._internal.split import (
|
|
_drop_empty_block_split,
|
|
_generate_global_split_results,
|
|
_generate_per_block_split_indices,
|
|
_generate_valid_indices,
|
|
_split_at_indices,
|
|
_split_single_block,
|
|
)
|
|
from ray.data._internal.stats import DatasetStats
|
|
from ray.data.block import Block, BlockAccessor, BlockMetadata
|
|
from ray.data.context import DataContext
|
|
from ray.data.dataset import Dataset
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.data.tests.util import extract_values
|
|
from ray.tests.conftest import * # noqa
|
|
from ray.types import ObjectRef
|
|
|
|
|
|
@ray.remote
|
|
class Counter:
|
|
def __init__(self):
|
|
self.value = 0
|
|
|
|
def increment(self):
|
|
self.value += 1
|
|
return self.value
|
|
|
|
|
|
def test_equal_split(shutdown_only):
|
|
ray.init(num_cpus=2)
|
|
|
|
def range2x(n):
|
|
return ray.data.range(2 * n)
|
|
|
|
def counts(shards):
|
|
@ray.remote(num_cpus=0)
|
|
def count(s):
|
|
return s.count()
|
|
|
|
return ray.get([count.remote(s) for s in shards])
|
|
|
|
r1 = counts(range2x(10).split(3, equal=True))
|
|
assert all(c == 6 for c in r1), r1
|
|
|
|
# The following test is failing and may be a regression.
|
|
# Splits appear to be based on existing block boundaries ([10, 5, 5], [8, 8, 4]).
|
|
|
|
# r2 = counts(range2x(10).split(3, equal=False))
|
|
# assert all(c >= 6 for c in r2), r2
|
|
# assert not all(c == 6 for c in r2), r2
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"block_sizes,num_splits",
|
|
[
|
|
([3, 6, 3], 3), # Test baseline.
|
|
([3, 3, 3], 3), # Already balanced.
|
|
([3, 6, 4], 3), # Row truncation.
|
|
([3, 6, 2, 3], 3), # Row truncation, smaller number of blocks.
|
|
([5, 6, 2, 5], 5), # Row truncation, larger number of blocks.
|
|
([1, 1, 1, 1, 6], 5), # All smaller but one.
|
|
([4, 4, 4, 4, 1], 5), # All larger but one.
|
|
([2], 2), # Single block.
|
|
([2, 5], 1), # Single split.
|
|
],
|
|
)
|
|
def test_equal_split_balanced(ray_start_regular_shared_2_cpus, block_sizes, num_splits):
|
|
_test_equal_split_balanced(block_sizes, num_splits)
|
|
|
|
|
|
def _test_equal_split_balanced(block_sizes, num_splits):
|
|
ctx = DataContext.get_current()
|
|
|
|
blocks = []
|
|
metadata = []
|
|
ref_bundles = []
|
|
total_rows = 0
|
|
for block_size in block_sizes:
|
|
block = pd.DataFrame({"id": list(range(total_rows, total_rows + block_size))})
|
|
blocks.append(ray.put(block))
|
|
metadata.append(BlockAccessor.for_block(block).get_metadata())
|
|
schema = BlockAccessor.for_block(block).schema()
|
|
blk = BlockEntry(blocks[-1], metadata[-1])
|
|
ref_bundles.append(RefBundle((blk,), owns_blocks=True, schema=schema))
|
|
total_rows += block_size
|
|
|
|
logical_plan = LogicalPlan(InputData(input_data=ref_bundles), ctx)
|
|
stats = DatasetStats(metadata={"TODO": []}, parent=None)
|
|
ds = Dataset(logical_plan, ctx, stats)
|
|
|
|
splits = ds.split(num_splits, equal=True)
|
|
split_counts = [split.count() for split in splits]
|
|
assert len(split_counts) == num_splits
|
|
expected_block_size = total_rows // num_splits
|
|
# Check that all splits are the expected size.
|
|
assert all([count == expected_block_size for count in split_counts])
|
|
expected_total_rows = sum(split_counts)
|
|
# Check that the expected number of rows were dropped.
|
|
assert total_rows - expected_total_rows == total_rows % num_splits
|
|
# Check that all rows are unique (content check).
|
|
split_rows = [row for split in splits for row in split.take(total_rows)]
|
|
assert len(set(extract_values("id", split_rows))) == len(split_rows)
|
|
|
|
|
|
def test_equal_split_balanced_grid(ray_start_regular_shared_2_cpus):
|
|
# Tests balanced equal splitting over a grid of configurations.
|
|
# Grid: num_blocks x num_splits x num_rows_block_1 x ... x num_rows_block_n
|
|
seed = int(time.time())
|
|
print(f"Seeding RNG for test_equal_split_balanced_grid with: {seed}")
|
|
random.seed(seed)
|
|
max_num_splits = 15
|
|
num_splits_samples = 3
|
|
max_num_blocks = 50
|
|
max_num_rows_per_block = 100
|
|
num_blocks_samples = 3
|
|
block_sizes_samples = 3
|
|
for num_splits in np.random.randint(2, max_num_splits + 1, size=num_splits_samples):
|
|
for num_blocks in np.random.randint(
|
|
1, max_num_blocks + 1, size=num_blocks_samples
|
|
):
|
|
block_sizes_list = [
|
|
np.random.randint(1, max_num_rows_per_block + 1, size=num_blocks)
|
|
for _ in range(block_sizes_samples)
|
|
]
|
|
for block_sizes in block_sizes_list:
|
|
if sum(block_sizes) < num_splits:
|
|
min_ = math.ceil(num_splits / num_blocks)
|
|
block_sizes = np.random.randint(
|
|
min_, max_num_rows_per_block + 1, size=num_blocks
|
|
)
|
|
_test_equal_split_balanced(block_sizes, num_splits)
|
|
|
|
|
|
def test_split_small(ray_start_regular_shared_2_cpus):
|
|
x = [Counter.remote() for _ in range(4)]
|
|
data = ["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"]
|
|
fail = []
|
|
|
|
@ray.remote(num_cpus=0)
|
|
def take(s):
|
|
return extract_values("item", s.take())
|
|
|
|
for m in [1, 3]:
|
|
for n in [1, 3]:
|
|
for locality_hints in [None, x[:n]]:
|
|
for equal in [True, False]:
|
|
print("Testing", m, n, equal, locality_hints)
|
|
ds = ray.data.from_items(data, override_num_blocks=m)
|
|
splits = ds.split(n, equal=equal, locality_hints=locality_hints)
|
|
assert len(splits) == n
|
|
outs = ray.get([take.remote(s) for s in splits])
|
|
out = []
|
|
for r in outs:
|
|
out.extend(r)
|
|
if equal:
|
|
lens = set([len(s) for s in outs]) # noqa
|
|
limit = len(data) - (len(data) % n)
|
|
allowed = [limit]
|
|
# Allow for some pipelining artifacts.
|
|
print(len(out), len(set(out)), allowed)
|
|
if (
|
|
len(out) not in allowed
|
|
or len(set(out)) != len(out)
|
|
# TODO(ekl) we should be able to enable this check, but
|
|
# there are some edge condition bugs in split.
|
|
# or len(lens) != 1
|
|
):
|
|
print("FAIL", m, n, equal, locality_hints)
|
|
fail.append((m, n, equal, locality_hints))
|
|
else:
|
|
if sorted(out) != data:
|
|
print("FAIL", m, n, equal, locality_hints)
|
|
fail.append((m, n, equal, locality_hints))
|
|
|
|
assert not fail, fail
|
|
|
|
|
|
def test_split_at_indices_simple(ray_start_regular_shared_2_cpus, restore_data_context):
|
|
# NOTE: It's critical to preserve ordering for assertions in this test to work
|
|
DataContext.get_current().execution_options.preserve_order = True
|
|
|
|
ds = ray.data.range(10, override_num_blocks=3)
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.split_at_indices([])
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.split_at_indices([-1])
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.split_at_indices([3, 1])
|
|
|
|
splits = ds.split_at_indices([5])
|
|
r = [extract_values("id", s.take()) for s in splits]
|
|
assert r == [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]
|
|
|
|
splits = ds.split_at_indices([2, 5])
|
|
r = [extract_values("id", s.take()) for s in splits]
|
|
assert r == [[0, 1], [2, 3, 4], [5, 6, 7, 8, 9]]
|
|
|
|
splits = ds.split_at_indices([2, 5, 5, 100])
|
|
r = [extract_values("id", s.take()) for s in splits]
|
|
assert r == [[0, 1], [2, 3, 4], [], [5, 6, 7, 8, 9], []]
|
|
|
|
splits = ds.split_at_indices([100])
|
|
r = [extract_values("id", s.take()) for s in splits]
|
|
assert r == [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], []]
|
|
|
|
splits = ds.split_at_indices([0])
|
|
r = [extract_values("id", s.take()) for s in splits]
|
|
assert r == [[], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]]
|
|
|
|
|
|
@pytest.mark.parametrize("num_blocks", list(range(1, 20)) + [25, 40])
|
|
@pytest.mark.parametrize(
|
|
"indices",
|
|
[
|
|
# Two-splits.
|
|
[5],
|
|
[10],
|
|
[15],
|
|
# Three-splits.
|
|
[5, 12],
|
|
[1, 18],
|
|
[9, 10],
|
|
# Misc.
|
|
[3, 10, 17],
|
|
[2, 4, 11, 12, 19],
|
|
list(range(20)),
|
|
list(range(0, 20, 2)),
|
|
# Empty splits.
|
|
[10, 10],
|
|
[5, 10, 10, 15],
|
|
# Out-of-bounds.
|
|
[25],
|
|
[7, 11, 23, 33],
|
|
],
|
|
)
|
|
def test_split_at_indices_coverage(
|
|
ray_start_regular_shared_2_cpus, num_blocks, indices, restore_data_context
|
|
):
|
|
# Test that split_at_indices() creates the expected splits on a set of partition and
|
|
# indices configurations.
|
|
|
|
# NOTE: It's critical to preserve ordering for assertions in this test to work
|
|
DataContext.get_current().execution_options.preserve_order = True
|
|
|
|
ds = ray.data.range(20, override_num_blocks=num_blocks)
|
|
splits = ds.split_at_indices(indices)
|
|
r = [extract_values("id", s.sort("id").take_all()) for s in splits]
|
|
# Use np.array_split() semantics as our correctness ground-truth.
|
|
assert r == [arr.tolist() for arr in np.array_split(list(range(20)), indices)]
|
|
|
|
|
|
@pytest.mark.parametrize("num_blocks", [1, 3, 5, 10])
|
|
@pytest.mark.parametrize(
|
|
"indices",
|
|
[
|
|
[2], # Single split
|
|
[1, 3], # Two splits
|
|
[0, 2, 4], # Three splits
|
|
[1, 2, 3, 4], # Four splits
|
|
[1, 2, 3, 4, 7], # Five splits
|
|
[1, 2, 3, 4, 6, 9], # Six splits
|
|
]
|
|
+ [
|
|
list(x) for x in itertools.combinations_with_replacement([1, 3, 4], 2)
|
|
] # Selected two-split cases
|
|
+ [
|
|
list(x) for x in itertools.combinations_with_replacement([0, 2, 4], 3)
|
|
], # Selected three-split cases
|
|
)
|
|
def test_split_at_indices_coverage_complete(
|
|
ray_start_regular_shared_2_cpus, num_blocks, indices, restore_data_context
|
|
):
|
|
# NOTE: It's critical to preserve ordering for assertions in this test to work
|
|
DataContext.get_current().execution_options.preserve_order = True
|
|
|
|
# Test that split_at_indices() creates the expected splits on a set of partition and
|
|
# indices configurations.
|
|
ds = ray.data.range(10, override_num_blocks=num_blocks)
|
|
splits = ds.split_at_indices(indices)
|
|
r = [extract_values("id", s.take_all()) for s in splits]
|
|
# Use np.array_split() semantics as our correctness ground-truth.
|
|
assert r == [arr.tolist() for arr in np.array_split(list(range(10)), indices)]
|
|
|
|
|
|
def test_split_proportionately(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.range(10, override_num_blocks=3)
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.split_proportionately([])
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.split_proportionately([-1])
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.split_proportionately([0])
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.split_proportionately([1])
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.split_proportionately([0.5, 0.5])
|
|
|
|
splits = ds.split_proportionately([0.5])
|
|
r = [extract_values("id", s.take()) for s in splits]
|
|
assert r == [[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]
|
|
|
|
splits = ds.split_proportionately([0.2, 0.3])
|
|
r = [extract_values("id", s.take()) for s in splits]
|
|
assert r == [[0, 1], [2, 3, 4], [5, 6, 7, 8, 9]]
|
|
|
|
splits = ds.split_proportionately([0.2, 0.3, 0.3])
|
|
r = [extract_values("id", s.take()) for s in splits]
|
|
assert r == [[0, 1], [2, 3, 4], [5, 6, 7], [8, 9]]
|
|
|
|
splits = ds.split_proportionately([0.98, 0.01])
|
|
r = [extract_values("id", s.take()) for s in splits]
|
|
assert r == [[0, 1, 2, 3, 4, 5, 6, 7], [8], [9]]
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.split_proportionately([0.90] + ([0.001] * 90))
|
|
|
|
|
|
def test_split(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.range(20, override_num_blocks=10)
|
|
assert ds._logical_plan.initial_num_blocks() == 10
|
|
assert ds.sum() == 190
|
|
assert ds._block_num_rows() == [2] * 10
|
|
|
|
datasets = ds.split(5)
|
|
assert [2] * 5 == [len(dataset._execute().blocks) for dataset in datasets]
|
|
assert 190 == sum([dataset.sum("id") for dataset in datasets])
|
|
|
|
datasets = ds.split(3)
|
|
assert [4, 3, 3] == [len(dataset._execute().blocks) for dataset in datasets]
|
|
assert 190 == sum([dataset.sum("id") for dataset in datasets])
|
|
|
|
datasets = ds.split(1)
|
|
assert [10] == [len(dataset._execute().blocks) for dataset in datasets]
|
|
assert 190 == sum([dataset.sum("id") for dataset in datasets])
|
|
|
|
datasets = ds.split(10)
|
|
assert [1] * 10 == [len(dataset._execute().blocks) for dataset in datasets]
|
|
assert 190 == sum([dataset.sum("id") for dataset in datasets])
|
|
|
|
datasets = ds.split(11)
|
|
assert [1] * 10 + [0] == [len(dataset._execute().blocks) for dataset in datasets]
|
|
assert 190 == sum([dataset.sum("id") or 0 for dataset in datasets])
|
|
|
|
|
|
def test_split_hints(ray_start_regular_shared_2_cpus):
|
|
@ray.remote
|
|
class Actor(object):
|
|
def __init__(self):
|
|
pass
|
|
|
|
def assert_split_assignment(
|
|
block_node_ids: list, actor_node_ids: list, expected_split_result: list
|
|
):
|
|
"""Helper function to setup split hints test.
|
|
|
|
Args:
|
|
block_node_ids: a list of blocks with their locations. For
|
|
example ["node1", "node2"] represents two blocks with
|
|
"node1", "node2" as their location respectively.
|
|
actor_node_ids: a list of actors with their locations. For
|
|
example ["node1", "node2"] represents two actors with
|
|
"node1", "node2" as their location respectively.
|
|
expected_split_result: a list of allocation result, each entry
|
|
in the list stores the block_index in the split dataset.
|
|
For example, [[0, 1], [2]] represents the split result has
|
|
two datasets, datasets[0] contains block 0 and 1; and
|
|
datasets[1] contains block 2.
|
|
"""
|
|
num_blocks = len(block_node_ids)
|
|
ds = ray.data.range(num_blocks, override_num_blocks=num_blocks).materialize()
|
|
bundles = ds.iter_internal_ref_bundles()
|
|
blocks = _ref_bundles_iterator_to_block_refs_list(bundles)
|
|
assert len(block_node_ids) == len(blocks)
|
|
actors = [Actor.remote() for i in range(len(actor_node_ids))]
|
|
with patch("ray.experimental.get_object_locations") as location_mock:
|
|
with patch("ray._private.state.actors") as state_mock:
|
|
block_locations = {}
|
|
for i, node_id in enumerate(block_node_ids):
|
|
if node_id:
|
|
block_locations[blocks[i]] = {"node_ids": [node_id]}
|
|
location_mock.return_value = block_locations
|
|
|
|
actor_state = {}
|
|
for i, node_id in enumerate(actor_node_ids):
|
|
actor_state[actors[i]._actor_id.hex()] = {
|
|
"Address": {"NodeID": node_id}
|
|
}
|
|
|
|
state_mock.return_value = actor_state
|
|
|
|
datasets = ds.split(len(actors), locality_hints=actors)
|
|
assert len(datasets) == len(actors)
|
|
for i in range(len(actors)):
|
|
assert {blocks[j] for j in expected_split_result[i]} == set(
|
|
_ref_bundles_iterator_to_block_refs_list(
|
|
datasets[i].iter_internal_ref_bundles()
|
|
)
|
|
)
|
|
|
|
assert_split_assignment(
|
|
["node2", "node1", "node1"], ["node1", "node2"], [[1, 2], [0]]
|
|
)
|
|
assert_split_assignment(
|
|
["node1", "node1", "node1"], ["node1", "node2"], [[2, 1], [0]]
|
|
)
|
|
assert_split_assignment(["node2", "node2", None], ["node1", "node2"], [[0, 2], [1]])
|
|
assert_split_assignment(["node2", "node2", None], [None, None], [[2, 1], [0]])
|
|
assert_split_assignment(
|
|
["n1", "n2", "n3", "n1", "n2"], ["n1", "n2"], [[0, 2, 3], [1, 4]]
|
|
)
|
|
|
|
assert_split_assignment(["n1", "n2"], ["n1", "n2", "n3"], [[0], [1], []])
|
|
|
|
# perfect split:
|
|
#
|
|
# split 300 blocks
|
|
# with node_ids interleaving between "n0", "n1", "n2"
|
|
#
|
|
# to 3 actors
|
|
# with has node_id "n1", "n2", "n0"
|
|
#
|
|
# expect that block 1, 4, 7... are assigned to actor with node_id n1
|
|
# block 2, 5, 8... are assigned to actor with node_id n2
|
|
# block 0, 3, 6... are assigned to actor with node_id n0
|
|
assert_split_assignment(
|
|
["n0", "n1", "n2"] * 100,
|
|
["n1", "n2", "n0"],
|
|
[range(1, 300, 3), range(2, 300, 3), range(0, 300, 3)],
|
|
)
|
|
|
|
# even split regardless of locality:
|
|
#
|
|
# split 301 blocks
|
|
# with block 0 to block 50 on "n0",
|
|
# block 51 to block 300 on "n1"
|
|
#
|
|
# to 3 actors
|
|
# with node_ids "n1", "n2", "n0"
|
|
#
|
|
# expect that block 200 to block 300 are assigned to actor with node_id n1
|
|
# block 100 to block 199 are assigned to actor with node_id n2
|
|
# block 0 to block 99 are assigned to actor with node_id n0
|
|
assert_split_assignment(
|
|
["n0"] * 50 + ["n1"] * 251,
|
|
["n1", "n2", "n0"],
|
|
[range(200, 301), range(100, 200), list(range(0, 50)) + list(range(50, 100))],
|
|
)
|
|
|
|
|
|
def test_generate_valid_indices():
|
|
assert [1, 2, 3] == _generate_valid_indices([10], [1, 2, 3])
|
|
assert [1, 2, 2] == _generate_valid_indices([1, 1], [1, 2, 3])
|
|
|
|
|
|
def test_generate_per_block_split_indices():
|
|
assert [[1], [1, 2], [], []] == _generate_per_block_split_indices(
|
|
[3, 3, 3, 1], [1, 4, 5]
|
|
)
|
|
assert [[3], [], [], [1, 1]] == _generate_per_block_split_indices(
|
|
[3, 3, 3, 1], [3, 10, 10]
|
|
)
|
|
assert [[], [], [], []] == _generate_per_block_split_indices([3, 3, 3, 1], [])
|
|
|
|
|
|
def _create_meta(num_rows):
|
|
return BlockMetadata(
|
|
num_rows=num_rows,
|
|
size_bytes=None,
|
|
input_files=None,
|
|
exec_stats=None,
|
|
)
|
|
|
|
|
|
def _create_block_and_metadata(data: Any) -> Tuple[ObjectRef[Block], BlockMetadata]:
|
|
block = pd.DataFrame({"id": data})
|
|
metadata = BlockAccessor.for_block(block).get_metadata()
|
|
return (ray.put(block), metadata)
|
|
|
|
|
|
def _create_bundle(blocks: List[List[Any]]) -> RefBundle:
|
|
schema = BlockAccessor.for_block(pd.DataFrame({"id": []})).schema()
|
|
return RefBundle(
|
|
[BlockEntry(*_create_block_and_metadata(block)) for block in blocks],
|
|
owns_blocks=True,
|
|
schema=schema,
|
|
)
|
|
|
|
|
|
def _create_blocks_with_metadata(blocks):
|
|
# Returns the legacy 2-tuple shape (BlockPartition) consumed by the
|
|
# split helpers in ``ray.data._internal.split``.
|
|
bundle = _create_bundle(blocks)
|
|
return [(entry.ref, entry.metadata) for entry in bundle.blocks]
|
|
|
|
|
|
def test_split_single_block(ray_start_regular_shared_2_cpus):
|
|
block = pd.DataFrame({"id": [1, 2, 3]})
|
|
metadata = _create_meta(3)
|
|
|
|
results = ray.get(
|
|
ray.remote(_split_single_block)
|
|
.options(num_returns=2)
|
|
.remote(234, block, metadata, [])
|
|
)
|
|
block_id, meta = results[0]
|
|
blocks = results[1:]
|
|
assert 234 == block_id
|
|
assert len(blocks) == 1
|
|
assert list(blocks[0]["id"]) == [1, 2, 3]
|
|
assert meta[0].num_rows == 3
|
|
|
|
results = ray.get(
|
|
ray.remote(_split_single_block)
|
|
.options(num_returns=3)
|
|
.remote(234, block, metadata, [1])
|
|
)
|
|
block_id, meta = results[0]
|
|
blocks = results[1:]
|
|
assert 234 == block_id
|
|
assert len(blocks) == 2
|
|
assert list(blocks[0]["id"]) == [1]
|
|
assert meta[0].num_rows == 1
|
|
assert list(blocks[1]["id"]) == [2, 3]
|
|
assert meta[1].num_rows == 2
|
|
|
|
results = ray.get(
|
|
ray.remote(_split_single_block)
|
|
.options(num_returns=6)
|
|
.remote(234, block, metadata, [0, 1, 1, 3])
|
|
)
|
|
block_id, meta = results[0]
|
|
blocks = results[1:]
|
|
assert 234 == block_id
|
|
assert len(blocks) == 5
|
|
assert list(blocks[0]["id"]) == []
|
|
assert list(blocks[1]["id"]) == [1]
|
|
assert list(blocks[2]["id"]) == []
|
|
assert list(blocks[3]["id"]) == [2, 3]
|
|
assert list(blocks[4]["id"]) == []
|
|
|
|
block = pd.DataFrame({"id": []})
|
|
metadata = _create_meta(0)
|
|
|
|
results = ray.get(
|
|
ray.remote(_split_single_block)
|
|
.options(num_returns=3)
|
|
.remote(234, block, metadata, [0])
|
|
)
|
|
block_id, meta = results[0]
|
|
blocks = results[1:]
|
|
assert 234 == block_id
|
|
assert len(blocks) == 2
|
|
assert list(blocks[0]["id"]) == []
|
|
assert list(blocks[1]["id"]) == []
|
|
|
|
|
|
def test_drop_empty_block_split():
|
|
assert [1, 2] == _drop_empty_block_split([0, 1, 2, 3], 3)
|
|
assert [1, 2] == _drop_empty_block_split([1, 1, 2, 2], 3)
|
|
assert [] == _drop_empty_block_split([0], 0)
|
|
|
|
|
|
def verify_splits(splits, blocks_by_split):
|
|
assert len(splits) == len(blocks_by_split)
|
|
for blocks, (block_refs, metas) in zip(blocks_by_split, splits):
|
|
assert len(blocks) == len(block_refs)
|
|
assert len(blocks) == len(metas)
|
|
for block, block_ref, meta in zip(blocks, block_refs, metas):
|
|
assert list(ray.get(block_ref)["id"]) == block
|
|
assert meta.num_rows == len(block)
|
|
|
|
|
|
def test_generate_global_split_results(ray_start_regular_shared_2_cpus):
|
|
inputs = [
|
|
_create_block_and_metadata([1]),
|
|
_create_block_and_metadata([2, 3]),
|
|
_create_block_and_metadata([4]),
|
|
]
|
|
|
|
splits = list(zip(*_generate_global_split_results(iter(inputs), [1, 2, 1])))
|
|
verify_splits(splits, [[[1]], [[2, 3]], [[4]]])
|
|
|
|
splits = list(zip(*_generate_global_split_results(iter(inputs), [3, 1])))
|
|
verify_splits(splits, [[[1], [2, 3]], [[4]]])
|
|
|
|
splits = list(zip(*_generate_global_split_results(iter(inputs), [3, 0, 1])))
|
|
verify_splits(splits, [[[1], [2, 3]], [], [[4]]])
|
|
|
|
inputs = []
|
|
splits = list(zip(*_generate_global_split_results(iter(inputs), [0, 0])))
|
|
verify_splits(splits, [[], []])
|
|
|
|
|
|
def test_private_split_at_indices(ray_start_regular_shared_2_cpus):
|
|
inputs = _create_blocks_with_metadata([])
|
|
splits = list(zip(*_split_at_indices(inputs, [0], True)))
|
|
verify_splits(splits, [[], []])
|
|
|
|
splits = list(zip(*_split_at_indices(inputs, [], True)))
|
|
verify_splits(splits, [[]])
|
|
|
|
inputs = _create_blocks_with_metadata([[1], [2, 3], [4]])
|
|
|
|
splits = list(zip(*_split_at_indices(inputs, [1], True)))
|
|
verify_splits(splits, [[[1]], [[2, 3], [4]]])
|
|
|
|
inputs = _create_blocks_with_metadata([[1], [2, 3], [4]])
|
|
splits = list(zip(*_split_at_indices(inputs, [2], True)))
|
|
verify_splits(splits, [[[1], [2]], [[3], [4]]])
|
|
|
|
inputs = _create_blocks_with_metadata([[1], [2, 3], [4]])
|
|
splits = list(zip(*_split_at_indices(inputs, [1], True)))
|
|
verify_splits(splits, [[[1]], [[2, 3], [4]]])
|
|
|
|
inputs = _create_blocks_with_metadata([[1], [2, 3], [4]])
|
|
splits = list(zip(*_split_at_indices(inputs, [2, 2], True)))
|
|
verify_splits(splits, [[[1], [2]], [], [[3], [4]]])
|
|
|
|
inputs = _create_blocks_with_metadata([[1], [2, 3], [4]])
|
|
splits = list(zip(*_split_at_indices(inputs, [], True)))
|
|
verify_splits(splits, [[[1], [2, 3], [4]]])
|
|
|
|
inputs = _create_blocks_with_metadata([[1], [2, 3], [4]])
|
|
splits = list(zip(*_split_at_indices(inputs, [0, 4], True)))
|
|
verify_splits(splits, [[], [[1], [2, 3], [4]], []])
|
|
|
|
|
|
def equalize_helper(input_block_lists: List[List[List[Any]]]):
|
|
result = _equalize(
|
|
[_create_bundle(block_list) for block_list in input_block_lists],
|
|
owned_by_consumer=True,
|
|
)
|
|
result_block_lists = []
|
|
for bundle in result:
|
|
block_list = []
|
|
for block_ref in bundle.block_refs:
|
|
block = ray.get(block_ref)
|
|
block_accessor = BlockAccessor.for_block(block)
|
|
block_list.append(list(block_accessor.to_default()["id"]))
|
|
result_block_lists.append(block_list)
|
|
return result_block_lists
|
|
|
|
|
|
def verify_equalize_result(input_block_lists, expected_block_lists):
|
|
result_block_lists = equalize_helper(input_block_lists)
|
|
assert result_block_lists == expected_block_lists
|
|
|
|
|
|
def test_equalize(ray_start_regular_shared_2_cpus):
|
|
verify_equalize_result([], [])
|
|
verify_equalize_result([[]], [[]])
|
|
verify_equalize_result([[[1]], []], [[], []])
|
|
verify_equalize_result([[[1], [2, 3]], [[4]]], [[[1], [2]], [[4], [3]]])
|
|
verify_equalize_result([[[1], [2, 3]], []], [[[1]], [[2]]])
|
|
verify_equalize_result(
|
|
[[[1], [2, 3], [4, 5]], [[6]], []], [[[1], [2]], [[6], [3]], [[4, 5]]]
|
|
)
|
|
verify_equalize_result(
|
|
[[[1, 2, 3], [4, 5]], [[6]], []], [[[4, 5]], [[6], [1]], [[2, 3]]]
|
|
)
|
|
|
|
|
|
def test_equalize_randomized(ray_start_regular_shared_2_cpus):
|
|
# verify the entries in the splits are in the range of 0 .. num_rows,
|
|
# unique, and the total number matches num_rows if exact_num == True.
|
|
def assert_unique_and_inrange(splits, num_rows, exact_num=False):
|
|
unique_set = set()
|
|
for split in splits:
|
|
for block in split:
|
|
for entry in block:
|
|
assert entry not in unique_set
|
|
assert entry >= 0 and entry < num_rows
|
|
unique_set.add(entry)
|
|
if exact_num:
|
|
assert len(unique_set) == num_rows
|
|
|
|
# verify that splits are equalized.
|
|
def assert_equal_split(splits, num_rows, num_split):
|
|
split_size = num_rows // num_split
|
|
for split in splits:
|
|
assert len((list(itertools.chain.from_iterable(split)))) == split_size
|
|
|
|
# create randomized splits contains entries from 0 ... num_rows.
|
|
def random_split(num_rows, num_split):
|
|
split_point = [int(random.random() * num_rows) for _ in range(num_split - 1)]
|
|
split_index_helper = [0] + sorted(split_point) + [num_rows]
|
|
splits = []
|
|
for i in range(1, len(split_index_helper)):
|
|
split_start = split_index_helper[i - 1]
|
|
split_end = split_index_helper[i]
|
|
num_entries = split_end - split_start
|
|
split = []
|
|
num_block_split = int(random.random() * num_entries)
|
|
block_split_point = [
|
|
split_start + int(random.random() * num_entries)
|
|
for _ in range(num_block_split)
|
|
]
|
|
block_index_helper = [split_start] + sorted(block_split_point) + [split_end]
|
|
for j in range(1, len(block_index_helper)):
|
|
split.append(
|
|
list(range(block_index_helper[j - 1], block_index_helper[j]))
|
|
)
|
|
splits.append(split)
|
|
assert_unique_and_inrange(splits, num_rows, exact_num=True)
|
|
return splits
|
|
|
|
for i in range(100):
|
|
num_rows = int(random.random() * 100)
|
|
num_split = int(random.random() * 10) + 1
|
|
input_splits = random_split(num_rows, num_split)
|
|
print(input_splits)
|
|
equalized_splits = equalize_helper(input_splits)
|
|
assert_unique_and_inrange(equalized_splits, num_rows)
|
|
assert_equal_split(equalized_splits, num_rows, num_split)
|
|
|
|
|
|
def test_train_test_split(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.range(8)
|
|
|
|
# float
|
|
train, test = ds.train_test_split(test_size=0.25)
|
|
assert extract_values("id", train.take()) == [0, 1, 2, 3, 4, 5]
|
|
assert extract_values("id", test.take()) == [6, 7]
|
|
|
|
# int
|
|
train, test = ds.train_test_split(test_size=2)
|
|
assert extract_values("id", train.take()) == [0, 1, 2, 3, 4, 5]
|
|
assert extract_values("id", test.take()) == [6, 7]
|
|
|
|
# shuffle
|
|
train, test = ds.train_test_split(test_size=0.25, shuffle=True, seed=1)
|
|
assert extract_values("id", train.take()) == [7, 4, 6, 0, 5, 2]
|
|
assert extract_values("id", test.take()) == [1, 3]
|
|
|
|
# error handling
|
|
with pytest.raises(TypeError):
|
|
ds.train_test_split(test_size=[1])
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.train_test_split(test_size=-1)
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.train_test_split(test_size=0)
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.train_test_split(test_size=1.1)
|
|
|
|
with pytest.raises(ValueError):
|
|
ds.train_test_split(test_size=9)
|
|
|
|
|
|
def test_train_test_split_stratified(ray_start_regular_shared_2_cpus):
|
|
# Test basic stratification with simple dataset
|
|
data = [
|
|
{"id": 0, "label": "A"},
|
|
{"id": 1, "label": "A"},
|
|
{"id": 2, "label": "B"},
|
|
{"id": 3, "label": "B"},
|
|
{"id": 4, "label": "C"},
|
|
{"id": 5, "label": "C"},
|
|
]
|
|
ds = ray.data.from_items(data)
|
|
|
|
# Test stratified split
|
|
train, test = ds.train_test_split(test_size=0.5, stratify="label")
|
|
|
|
# Check that we have the right number of samples
|
|
assert train.count() == 3
|
|
assert test.count() == 3
|
|
|
|
# Check that class proportions are preserved
|
|
train_labels = [row["label"] for row in train.take()]
|
|
test_labels = [row["label"] for row in test.take()]
|
|
|
|
train_label_counts = {label: train_labels.count(label) for label in ["A", "B", "C"]}
|
|
test_label_counts = {label: test_labels.count(label) for label in ["A", "B", "C"]}
|
|
|
|
# Each class should have exactly 1 sample in each split
|
|
assert train_label_counts == {"A": 1, "B": 1, "C": 1}
|
|
assert test_label_counts == {"A": 1, "B": 1, "C": 1}
|
|
|
|
|
|
def test_train_test_split_shuffle_stratify_error(ray_start_regular_shared_2_cpus):
|
|
# Test that shuffle=True and stratify cannot be used together
|
|
data = [
|
|
{"id": 0, "label": "A"},
|
|
{"id": 1, "label": "A"},
|
|
{"id": 2, "label": "B"},
|
|
{"id": 3, "label": "B"},
|
|
]
|
|
ds = ray.data.from_items(data)
|
|
|
|
# Test that combining shuffle=True and stratify raises ValueError
|
|
with pytest.raises(
|
|
ValueError, match="Cannot specify both 'shuffle=True' and 'stratify'"
|
|
):
|
|
ds.train_test_split(test_size=0.5, shuffle=True, stratify="label")
|
|
|
|
|
|
def test_train_test_split_stratified_imbalanced(ray_start_regular_shared_2_cpus):
|
|
# Test stratified split with imbalanced class distribution
|
|
data = [
|
|
{"id": 0, "label": "A"},
|
|
{"id": 1, "label": "A"},
|
|
{"id": 2, "label": "A"},
|
|
{"id": 3, "label": "A"},
|
|
{"id": 4, "label": "A"},
|
|
{"id": 5, "label": "A"}, # 6 samples of class A
|
|
{"id": 6, "label": "B"},
|
|
{"id": 7, "label": "B"}, # 2 samples of class B
|
|
{"id": 8, "label": "C"}, # 1 sample of class C
|
|
]
|
|
ds = ray.data.from_items(data)
|
|
|
|
# Test with 0.3 test size
|
|
train, test = ds.train_test_split(test_size=0.3, stratify="label")
|
|
|
|
train_labels = [row["label"] for row in train.take()]
|
|
test_labels = [row["label"] for row in test.take()]
|
|
|
|
train_label_counts = {label: train_labels.count(label) for label in ["A", "B", "C"]}
|
|
test_label_counts = {label: test_labels.count(label) for label in ["A", "B", "C"]}
|
|
|
|
# Check proportions are maintained as closely as possible
|
|
# Class A: 6 samples -> test_count = int(6 * 0.3) = 1 -> train: 5, test: 1
|
|
# Class B: 2 samples -> test_count = int(2 * 0.3) = 0 -> train: 2, test: 0
|
|
# Class C: 1 sample -> test_count = int(1 * 0.3) = 0 -> train: 1, test: 0
|
|
assert train_label_counts["A"] == 5
|
|
assert test_label_counts["A"] == 1
|
|
assert train_label_counts["B"] == 2
|
|
assert test_label_counts["B"] == 0
|
|
assert train_label_counts["C"] == 1
|
|
assert test_label_counts["C"] == 0
|
|
|
|
|
|
def test_split_is_not_disruptive(ray_start_cluster):
|
|
ray.shutdown()
|
|
ds = ray.data.range(100, override_num_blocks=10).map_batches(lambda x: x)
|
|
|
|
def verify_integrity(splits):
|
|
for dss in splits:
|
|
for batch in dss.iter_batches():
|
|
pass
|
|
for batch in ds.iter_batches():
|
|
pass
|
|
|
|
# No block splitting invovled: split 10 even blocks into 2 groups.
|
|
verify_integrity(ds.split(2, equal=True))
|
|
# Block splitting invovled: split 10 even blocks into 3 groups.
|
|
verify_integrity(ds.split(3, equal=True))
|
|
|
|
# Same as above but having tranforms post converting to lazy.
|
|
verify_integrity(ds.map_batches(lambda x: x).split(2, equal=True))
|
|
verify_integrity(ds.map_batches(lambda x: x).split(3, equal=True))
|
|
|
|
# Same as above but having in-place tranforms post converting to lazy.
|
|
verify_integrity(ds.randomize_block_order().split(2, equal=True))
|
|
verify_integrity(ds.randomize_block_order().split(3, equal=True))
|
|
|
|
|
|
def test_streaming_train_test_split_hash(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.range(10000000, override_num_blocks=10)
|
|
|
|
ds_train, ds_test = ds.streaming_train_test_split(
|
|
test_size=0.2, split_type="hash", hash_column="id"
|
|
)
|
|
|
|
np.testing.assert_almost_equal(float(ds_train.count()) / 10000000.0, 0.8, decimal=3)
|
|
np.testing.assert_almost_equal(float(ds_test.count()) / 10000000.0, 0.2, decimal=3)
|
|
|
|
# Check if train and test are disjoint
|
|
assert (
|
|
ds_train.join(ds_test, join_type="inner", on=("id",), num_partitions=1).count()
|
|
== 0
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("seed", [None, 42])
|
|
def test_streaming_train_test_split_random(ray_start_regular_shared_2_cpus, seed):
|
|
ds = ray.data.range(10000000, override_num_blocks=10)
|
|
|
|
ds_train, ds_test = ds.streaming_train_test_split(
|
|
test_size=0.2, split_type="random", seed=seed
|
|
)
|
|
|
|
np.testing.assert_almost_equal(float(ds_train.count()) / 10000000.0, 0.8, decimal=3)
|
|
np.testing.assert_almost_equal(float(ds_test.count()) / 10000000.0, 0.2, decimal=3)
|
|
|
|
# Check if train and test are disjoint
|
|
assert (
|
|
ds_train.join(ds_test, join_type="inner", on=("id",), num_partitions=1).count()
|
|
== 0
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"test_size,split_type,hash_column,seed,error_msg",
|
|
[
|
|
(0.2, "hash", None, None, "hash_column is required for hash split"),
|
|
(0.2, "hash", "id", 42, "seed is not supported for hash split"),
|
|
(0, "hash", "id", None, "test_size must be between 0 and 1"),
|
|
(1, "hash", "id", None, "test_size must be between 0 and 1"),
|
|
(0.2, "random", "id", None, "hash_column is not supported for random split"),
|
|
(0, "random", None, None, "test_size must be between 0 and 1"),
|
|
(1, "random", None, None, "test_size must be between 0 and 1"),
|
|
(0.2, "unknown", "id", None, "Invalid split type: unknown"),
|
|
],
|
|
)
|
|
def test_streaming_train_test_split_wrong_params(
|
|
ray_start_regular_shared_2_cpus, test_size, split_type, hash_column, seed, error_msg
|
|
):
|
|
ds = ray.data.range(10)
|
|
|
|
with pytest.raises(ValueError, match=error_msg):
|
|
ds.streaming_train_test_split(
|
|
test_size=test_size,
|
|
split_type=split_type,
|
|
hash_column=hash_column,
|
|
seed=seed,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("prefetch_batches", [0, 2])
|
|
def test_streaming_split_reports_and_clears_prefetched_bytes(
|
|
ray_start_regular_shared_2_cpus, prefetch_batches
|
|
):
|
|
"""Test streaming_split reports and clears prefetched bytes.
|
|
|
|
Verifies that:
|
|
1. Prefetched bytes are tracked during iteration
|
|
2. When StopIteration is raised in SplitCoordinator.get(), the
|
|
prefetched bytes are cleared to avoid stale backpressure data
|
|
"""
|
|
ds = ray.data.range(20, override_num_blocks=4)
|
|
splits = ds.streaming_split(2, equal=True)
|
|
|
|
# Get the coordinator actor to check prefetched bytes
|
|
coord = splits[0]._coord_actor
|
|
|
|
results = []
|
|
|
|
def consume(split, idx):
|
|
count = 0
|
|
for batch in split.iter_batches(
|
|
batch_size=5, prefetch_batches=prefetch_batches
|
|
):
|
|
count += len(batch["id"])
|
|
results.append((idx, count))
|
|
|
|
threads = [
|
|
threading.Thread(target=consume, args=(splits[0], 0)),
|
|
threading.Thread(target=consume, args=(splits[1], 1)),
|
|
]
|
|
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# Verify both splits consumed all data (epoch completed via StopIteration)
|
|
total_rows = sum(r[1] for r in results)
|
|
assert total_rows == 20, f"Expected 20 total rows, got {total_rows}"
|
|
|
|
# Verify client prefetched bytes are cleared after epoch end
|
|
# (cleared when StopIteration is raised in SplitCoordinator.get())
|
|
client_bytes = ray.get(coord.get_client_prefetched_bytes.remote())
|
|
for split_idx, bytes_val in client_bytes.items():
|
|
assert (
|
|
bytes_val == 0
|
|
), f"Split {split_idx} has stale prefetched bytes: {bytes_val}"
|
|
|
|
# Run a second epoch to verify cleared bytes don't cause issues
|
|
results.clear()
|
|
|
|
threads = [
|
|
threading.Thread(target=consume, args=(splits[0], 0)),
|
|
threading.Thread(target=consume, args=(splits[1], 1)),
|
|
]
|
|
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
# Second epoch should also consume all data
|
|
total_rows = sum(r[1] for r in results)
|
|
assert total_rows == 20, f"Second epoch: expected 20 rows, got {total_rows}"
|
|
|
|
# Verify prefetched bytes cleared again after second epoch
|
|
client_bytes = ray.get(coord.get_client_prefetched_bytes.remote())
|
|
for split_idx, bytes_val in client_bytes.items():
|
|
assert (
|
|
bytes_val == 0
|
|
), f"Split {split_idx} stale bytes after 2nd epoch: {bytes_val}"
|
|
|
|
|
|
def test_streaming_splits_schema_access(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.range(20, override_num_blocks=4)
|
|
|
|
iter_1, iter_2 = ds.streaming_split(2)
|
|
|
|
expected_schema = pa.schema([pa.field("id", pa.int64())])
|
|
|
|
assert expected_schema.equals(iter_1.schema().base_schema)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|