chore: import upstream snapshot with attribution

python/ray/data/tests/test_dynamic_block_split.py

@@ -0,0 +1,564 @@
+import os
+import sys
+import time
+from dataclasses import astuple, dataclass
+from typing import TYPE_CHECKING, List, Optional
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+import ray
+from ray.data._internal.arrow_block import ArrowBlockBuilder
+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
+from ray.data.datasource.datasource import ReadTask
+from ray.data.tests.conftest import (
+    CoreExecutionMetrics,
+    assert_blocks_expected_in_plasma,
+    assert_core_execution_metrics_equals,
+    get_initial_core_execution_metrics_snapshot,
+)
+from ray.tests.conftest import *  # noqa
+
+if TYPE_CHECKING:
+    from ray.data.context import DataContext
+
+
+# Data source generates random bytes data
+class RandomBytesDatasource(Datasource):
+    def __init__(
+        self,
+        num_tasks: int,
+        num_batches_per_task: int,
+        row_size: int,
+        num_rows_per_batch=None,
+        use_bytes=True,
+        use_arrow=False,
+    ):
+        self.num_tasks = num_tasks
+        self.num_batches_per_task = num_batches_per_task
+        self.row_size = row_size
+        if num_rows_per_batch is None:
+            num_rows_per_batch = 1
+        self.num_rows_per_batch = num_rows_per_batch
+        self.use_bytes = use_bytes
+        self.use_arrow = use_arrow
+
+    def estimate_inmemory_data_size(self):
+        return None
+
+    def get_read_tasks(
+        self,
+        parallelism: int,
+        per_task_row_limit: Optional[int] = None,
+        data_context: Optional["DataContext"] = None,
+    ) -> List[ReadTask]:
+        def _blocks_generator():
+            for _ in range(self.num_batches_per_task):
+                if self.use_bytes:
+                    # NOTE(swang): Each np object has some metadata bytes, so
+                    # actual size can be much more than num_rows_per_batch * row_size
+                    # if row_size is small.
+                    yield pd.DataFrame(
+                        {
+                            "one": [
+                                np.random.bytes(self.row_size)
+                                for _ in range(self.num_rows_per_batch)
+                            ]
+                        }
+                    )
+                elif self.use_arrow:
+                    batch = {
+                        "one": np.ones(
+                            (self.num_rows_per_batch, self.row_size), dtype=np.uint8
+                        )
+                    }
+                    block = ArrowBlockBuilder._table_from_pydict(batch)
+                    yield block
+                else:
+                    yield pd.DataFrame(
+                        {
+                            "one": [
+                                np.array2string(np.ones(self.row_size, dtype=int))
+                                for _ in range(self.num_rows_per_batch)
+                            ]
+                        }
+                    )
+
+        return self.num_tasks * [
+            ReadTask(
+                lambda: _blocks_generator(),
+                BlockMetadata(
+                    num_rows=self.num_batches_per_task * self.num_rows_per_batch,
+                    size_bytes=self.num_batches_per_task
+                    * self.num_rows_per_batch
+                    * self.row_size,
+                    input_files=None,
+                    exec_stats=None,
+                ),
+                per_task_row_limit=per_task_row_limit,
+            )
+        ]
+
+    def num_rows(self) -> int:
+        return self.num_tasks * self.num_batches_per_task * self.num_rows_per_batch
+
+
+class SlowCSVDatasource(CSVDatasource):
+    def _read_stream(self, f: "pa.NativeFile", path: str):
+        for block in super()._read_stream(f, path):
+            time.sleep(3)
+            yield block
+
+
+# Tests that we don't block on exponential rampup when doing bulk reads.
+# https://github.com/ray-project/ray/issues/20625
+@pytest.mark.parametrize("block_split", [False, True])
+def test_bulk_lazy_eval_split_mode(shutdown_only, block_split, tmp_path):
+    # Defensively shutdown Ray for the first test here to make sure there
+    # is no existing Ray cluster.
+    ray.shutdown()
+
+    ray.init(num_cpus=8)
+    ctx = ray.data.context.DataContext.get_current()
+
+    ray.data.range(8, override_num_blocks=8).write_csv(str(tmp_path))
+    if not block_split:
+        # Setting a huge block size effectively disables block splitting.
+        ctx.target_max_block_size = 2**64
+    ds = ray.data.read_datasource(
+        SlowCSVDatasource(str(tmp_path)), override_num_blocks=8
+    )
+
+    start = time.time()
+    ds.map(lambda x: x)
+    delta = time.time() - start
+
+    print("full read time", delta)
+    # Should run in ~3 seconds. It takes >9 seconds if bulk read is broken.
+    assert delta < 8, delta
+
+
+@pytest.mark.parametrize(
+    "compute",
+    [
+        "tasks",
+        "actors",
+    ],
+)
+def test_dataset(
+    shutdown_only,
+    restore_data_context,
+    compute,
+):
+    def identity_fn(x):
+        return x
+
+    def empty_fn(x):
+        return {}
+
+    class IdentityClass:
+        def __call__(self, x):
+            return x
+
+    class EmptyClass:
+        def __call__(self, x):
+            return {}
+
+    ctx = ray.data.DataContext.get_current()
+    # 1MiB.
+    ctx.target_max_block_size = 1024 * 1024
+
+    if compute == "tasks":
+        compute = ray.data._internal.compute.TaskPoolStrategy()
+        identity_func = identity_fn
+        empty_func = empty_fn
+        func_name = "identity_fn"
+        task_name = f"ReadRandomBytes->MapBatches({func_name})"
+    else:
+        compute = ray.data.ActorPoolStrategy()
+        identity_func = IdentityClass
+        empty_func = EmptyClass
+        func_name = "IdentityClass"
+        task_name = f"MapWorker(ReadRandomBytes->MapBatches({func_name})).submit"
+
+    ray.shutdown()
+    # We need at least 2 CPUs to run a actorpool streaming
+    ray.init(num_cpus=2, object_store_memory=1e9)
+    # Test 10 tasks, each task returning 10 blocks, each block has 1 row and each
+    # row has 1024 bytes.
+    num_blocks_per_task = 10
+    num_tasks = 10
+
+    @ray.remote
+    def warmup():
+        return np.zeros(ctx.target_max_block_size, dtype=np.uint8)
+
+    last_snapshot = get_initial_core_execution_metrics_snapshot()
+    ds = ray.data.read_datasource(
+        RandomBytesDatasource(
+            num_tasks=num_tasks,
+            num_batches_per_task=num_blocks_per_task,
+            row_size=ctx.target_max_block_size,
+        ),
+        override_num_blocks=num_tasks,
+    )
+    # Note the following calls to ds will not fully execute it.
+    assert ds.schema() is not None
+    assert ds.count() == num_blocks_per_task * num_tasks
+    assert ds._logical_plan.initial_num_blocks() == num_tasks
+    last_snapshot = assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(
+            task_count={
+                "ReadRandomBytes": lambda count: count <= num_tasks,
+            },
+            object_store_stats={
+                "cumulative_created_plasma_bytes": lambda count: True,
+                "cumulative_created_plasma_objects": lambda count: True,
+            },
+        ),
+        last_snapshot,
+    )
+
+    # Too-large blocks will get split to respect target max block size.
+    map_ds = ds.map_batches(identity_func, compute=compute)
+    map_ds = map_ds.materialize()
+    num_blocks_expected = num_tasks * num_blocks_per_task
+    assert map_ds._logical_plan.initial_num_blocks() == num_blocks_expected
+    expected_actor_name = f"MapWorker(ReadRandomBytes->MapBatches({func_name}))"
+    assert_core_execution_metrics_equals(
+        CoreExecutionMetrics(
+            task_count={
+                f"{expected_actor_name}.__init__": lambda count: True,
+                f"{expected_actor_name}.get_location": lambda count: True,
+                task_name: num_tasks,
+            },
+        ),
+        last_snapshot,
+    )
+    assert_blocks_expected_in_plasma(
+        last_snapshot,
+        num_blocks_expected,
+        block_size_expected=ctx.target_max_block_size,
+    )
+
+    # Blocks smaller than requested batch size will get coalesced.
+    map_ds = ds.map_batches(
+        empty_func,
+        batch_size=num_blocks_per_task * num_tasks,
+        compute=compute,
+    )
+    map_ds = map_ds.materialize()
+    assert map_ds._logical_plan.initial_num_blocks() == 1
+    map_ds = ds.map(identity_func, compute=compute)
+    map_ds = map_ds.materialize()
+    assert map_ds._logical_plan.initial_num_blocks() == num_blocks_per_task * num_tasks
+
+    ds_list = ds.split(5)
+    assert len(ds_list) == 5
+    for new_ds in ds_list:
+        assert (
+            new_ds._logical_plan.initial_num_blocks()
+            == num_blocks_per_task * num_tasks / 5
+        )
+
+    train, test = ds.train_test_split(test_size=0.25)
+    assert (
+        train._logical_plan.initial_num_blocks()
+        == num_blocks_per_task * num_tasks * 0.75
+    )
+    assert (
+        test._logical_plan.initial_num_blocks()
+        == num_blocks_per_task * num_tasks * 0.25
+    )
+
+    new_ds = ds.union(ds, ds)
+    assert new_ds._logical_plan.initial_num_blocks() == num_tasks * 3
+    new_ds = new_ds.materialize()
+    assert (
+        new_ds._logical_plan.initial_num_blocks() == num_blocks_per_task * num_tasks * 3
+    )
+
+    new_ds = ds.random_shuffle()
+    assert new_ds._logical_plan.initial_num_blocks() == num_tasks
+    new_ds = ds.randomize_block_order()
+    assert new_ds._logical_plan.initial_num_blocks() == num_tasks
+    assert ds.groupby("one").count().count() == num_blocks_per_task * num_tasks
+
+    new_ds = ds.zip(ds)
+    new_ds = new_ds.materialize()
+    assert new_ds._logical_plan.initial_num_blocks() == num_blocks_per_task * num_tasks
+
+    assert len(ds.take(5)) == 5
+    assert len(ds.take_all()) == num_blocks_per_task * num_tasks
+    for batch in ds.iter_batches(batch_size=10):
+        assert len(batch["one"]) == 10
+
+
+def test_filter(ray_start_regular_shared, target_max_block_size):
+    # Test 10 tasks, each task returning 10 blocks, each block has 1 row and each
+    # row has 1024 bytes.
+    num_blocks_per_task = 10
+    block_size = 1024
+
+    ds = ray.data.read_datasource(
+        RandomBytesDatasource(
+            num_tasks=1,
+            num_batches_per_task=num_blocks_per_task,
+            row_size=block_size,
+        ),
+        override_num_blocks=1,
+    )
+
+    ds = ds.filter(lambda _: True)
+    ds = ds.materialize()
+    assert ds.count() == num_blocks_per_task
+    assert ds._logical_plan.initial_num_blocks() == num_blocks_per_task
+
+    ds = ds.filter(lambda _: False)
+    ds = ds.materialize()
+    assert ds.count() == 0
+    assert ds._logical_plan.initial_num_blocks() == num_blocks_per_task
+
+
+@pytest.mark.skip("Needs zero-copy optimization for read->map_batches.")
+def test_read_large_data(ray_start_cluster):
+    # Test 20G input with single task
+    num_blocks_per_task = 20
+    block_size = 1024 * 1024 * 1024
+
+    cluster = ray_start_cluster
+    cluster.add_node(num_cpus=1)
+
+    ray.init(cluster.address)
+
+    def foo(batch):
+        return pd.DataFrame({"one": [1]})
+
+    ds = ray.data.read_datasource(
+        RandomBytesDatasource(
+            num_tasks=1,
+            num_batches_per_task=num_blocks_per_task,
+            row_size=block_size,
+        ),
+        override_num_blocks=1,
+    )
+
+    ds = ds.map_batches(foo, num_rows_per_batch=None)
+    assert ds.count() == num_blocks_per_task
+
+
+def _test_write_large_data(
+    tmp_path, ext, write_fn, read_fn, use_bytes, write_kwargs=None
+):
+    # Test 2G input with single task
+    num_blocks_per_task = 200
+    block_size = 10 * 1024 * 1024
+
+    ds = ray.data.read_datasource(
+        RandomBytesDatasource(
+            num_tasks=1,
+            num_batches_per_task=num_blocks_per_task,
+            row_size=block_size,
+            use_bytes=use_bytes,
+        ),
+        override_num_blocks=1,
+    )
+
+    # This should succeed without OOM.
+    # https://github.com/ray-project/ray/pull/37966.
+    out_dir = os.path.join(tmp_path, ext)
+    write_kwargs = {} if write_kwargs is None else write_kwargs
+    write_fn(ds, out_dir, **write_kwargs)
+
+    # Make sure we can read out a record.
+    if read_fn is not None:
+        assert read_fn(out_dir).count() == num_blocks_per_task
+
+
+def test_write_large_data_parquet(shutdown_only, tmp_path):
+    _test_write_large_data(
+        tmp_path,
+        "parquet",
+        Dataset.write_parquet,
+        ray.data.read_parquet,
+        use_bytes=True,
+    )
+
+
+def test_write_large_data_json(shutdown_only, tmp_path):
+    _test_write_large_data(
+        tmp_path, "json", Dataset.write_json, ray.data.read_json, use_bytes=False
+    )
+
+
+def test_write_large_data_numpy(shutdown_only, tmp_path):
+    _test_write_large_data(
+        tmp_path,
+        "numpy",
+        Dataset.write_numpy,
+        ray.data.read_numpy,
+        use_bytes=False,
+        write_kwargs={"column": "one"},
+    )
+
+
+def test_write_large_data_csv(shutdown_only, tmp_path):
+    _test_write_large_data(
+        tmp_path, "csv", Dataset.write_csv, ray.data.read_csv, use_bytes=False
+    )
+
+
+@pytest.mark.skipif(
+    sys.version_info >= (3, 12),
+    reason="Skip due to incompatibility tensorflow with Python 3.12+",
+)
+def test_write_large_data_tfrecords(shutdown_only, tmp_path):
+    _test_write_large_data(
+        tmp_path,
+        "tfrecords",
+        Dataset.write_tfrecords,
+        ray.data.read_tfrecords,
+        use_bytes=True,
+    )
+
+
+def test_write_large_data_webdataset(shutdown_only, tmp_path):
+    _test_write_large_data(
+        tmp_path,
+        "webdataset",
+        Dataset.write_webdataset,
+        ray.data.read_webdataset,
+        use_bytes=True,
+    )
+
+
+@dataclass
+class TestCase:
+    target_max_block_size: int
+    batch_size: int
+    num_batches: int
+    expected_num_blocks: int
+
+
+TEST_CASES = [
+    # Don't create blocks smaller than 50%.
+    TestCase(
+        target_max_block_size=1024,
+        batch_size=int(1024 * 1.125),
+        num_batches=1,
+        expected_num_blocks=1,
+    ),
+    # Split blocks larger than 150% the target block size.
+    TestCase(
+        target_max_block_size=1024,
+        batch_size=int(1024 * 1.8),
+        num_batches=1,
+        expected_num_blocks=2,
+    ),
+    # Huge batch will get split into multiple blocks.
+    TestCase(
+        target_max_block_size=1024,
+        batch_size=int(1024 * 10.125),
+        num_batches=1,
+        expected_num_blocks=10,
+    ),
+    # Different batch sizes but same total size should produce a similar number
+    # of blocks.
+    TestCase(
+        target_max_block_size=1024,
+        batch_size=int(1024 * 1.5),
+        num_batches=4,
+        expected_num_blocks=6,
+    ),
+    TestCase(
+        target_max_block_size=1024,
+        batch_size=int(1024 * 0.75),
+        num_batches=8,
+        expected_num_blocks=6,
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "target_max_block_size,batch_size,num_batches,expected_num_blocks",
+    [astuple(test) for test in TEST_CASES],
+)
+def test_block_slicing(
+    ray_start_regular_shared,
+    restore_data_context,
+    target_max_block_size,
+    batch_size,
+    num_batches,
+    expected_num_blocks,
+):
+    ctx = ray.data.context.DataContext.get_current()
+    ctx.target_max_block_size = target_max_block_size
+
+    # Row sizes smaller than this seem to add significant amounts of per-row
+    # metadata overhead.
+    row_size = 128
+    num_rows_per_batch = int(batch_size / row_size)
+    num_tasks = 1
+
+    ds = ray.data.read_datasource(
+        RandomBytesDatasource(
+            num_tasks=num_tasks,
+            num_batches_per_task=num_batches,
+            num_rows_per_batch=num_rows_per_batch,
+            row_size=row_size,
+            use_bytes=False,
+            use_arrow=True,
+        ),
+        override_num_blocks=num_tasks,
+    ).materialize()
+    assert ds._logical_plan.initial_num_blocks() == expected_num_blocks
+
+    block_sizes = []
+    num_rows = 0
+    for batch in ds.iter_batches(batch_size=None, batch_format="numpy"):
+        block_sizes.append(batch["one"].size)
+        num_rows += len(batch["one"])
+    assert num_rows == num_rows_per_batch * num_batches
+    for size in block_sizes:
+        # Blocks are not too big.
+        assert (
+            size <= target_max_block_size * ray.data.context.MAX_SAFE_BLOCK_SIZE_FACTOR
+        )
+        # Blocks are not too small.
+        assert size >= target_max_block_size / 2
+
+
+@pytest.mark.parametrize(
+    "target_max_block_size",
+    [128, 256, 512],
+)
+def test_dynamic_block_split_deterministic(
+    ray_start_regular_shared, target_max_block_size
+):
+    # Tests the determinism of block splitting.
+    TEST_ITERATIONS = 10
+    ctx = ray.data.DataContext.get_current()
+    ctx.target_max_block_size = target_max_block_size
+
+    # ~800 bytes per block
+    ds = ray.data.range(1000, override_num_blocks=10).map_batches(lambda x: x)
+    data = [ray.get(block) for block in ds.materialize()._cache._bundle.block_refs]
+    # Maps: first item of block -> block
+    block_map = {block["id"][0]: block for block in data}
+    # Iterate over multiple executions of the dataset,
+    # and check that blocks were split in the same way
+    for _ in range(TEST_ITERATIONS):
+        data = [ray.get(block) for block in ds.materialize()._cache._bundle.block_refs]
+        for block in data:
+            assert block_map[block["id"][0]] == block
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))