chore: import upstream snapshot with attribution

python/ray/air/tests/test_new_dataset_config.py

@@ -0,0 +1,394 @@
+import random
+from typing import Optional
+from unittest.mock import MagicMock
+
+import pytest
+
+import ray
+from ray import train
+from ray.data import DataIterator
+from ray.data._internal.execution.interfaces.execution_options import (
+    ExecutionOptions,
+    ExecutionResources,
+)
+from ray.tests.conftest import *  # noqa
+from ray.train import DataConfig, ScalingConfig
+from ray.train.data_parallel_trainer import DataParallelTrainer
+
+
+@pytest.fixture
+def ray_start_4_cpus():
+    address_info = ray.init(num_cpus=4)
+    yield address_info
+    ray.shutdown()
+
+
+class TestBasic(DataParallelTrainer):
+    def __init__(
+        self, num_workers: int, expect_ds: bool, expect_sizes: Optional[dict], **kwargs
+    ):
+        def train_loop_per_worker():
+            data_shard = train.get_dataset_shard("train")
+            assert isinstance(data_shard, DataIterator), data_shard
+            for k, v in expect_sizes.items():
+                shard = train.get_dataset_shard(k)
+                if v == -1:
+                    assert shard is None, shard
+                else:
+                    count = 0
+                    for batch in shard.iter_batches():
+                        for arr in batch.values():
+                            count += arr.size
+                    assert count == v, shard
+
+        kwargs.pop("scaling_config", None)
+        super().__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            scaling_config=ScalingConfig(num_workers=num_workers),
+            **kwargs,
+        )
+
+
+def test_basic(ray_start_4_cpus):
+    ds = ray.data.range(10)
+
+    # Single worker basic case.
+    test = TestBasic(
+        1,
+        True,
+        {"train": 10, "test": 10},
+        datasets={"train": ds, "test": ds},
+    )
+    test.fit()
+
+    # Single worker, no test ds.
+    test = TestBasic(1, True, {"train": 10, "test": -1}, datasets={"train": ds})
+    test.fit()
+
+    # Two workers, train and test split.
+    test = TestBasic(
+        2, True, {"train": 5, "test": 5}, datasets={"train": ds, "test": ds}
+    )
+    test.fit()
+
+    # Two workers, both split.
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 5},
+        dataset_config=DataConfig(datasets_to_split=["train", "test"]),
+        datasets={"train": ds, "test": ds},
+    )
+    # Test get config.
+    assert isinstance(test.get_dataset_config(), DataConfig)
+    test.fit()
+
+
+def test_split(ray_start_4_cpus):
+    ds = ray.data.range(10)
+
+    # Split all by default
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 5, "val": 5},
+        datasets={"train": ds, "test": ds, "val": ds},
+    )
+    test.fit()
+
+    # Test flag "all"
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 5},
+        datasets={"train": ds, "test": ds},
+        dataset_config=DataConfig(datasets_to_split="all"),
+    )
+
+    # Test split train only.
+    test = TestBasic(
+        2,
+        True,
+        {"train": 5, "test": 10},
+        datasets={"train": ds, "test": ds},
+        dataset_config=DataConfig(datasets_to_split=["train"]),
+    )
+    test.fit()
+
+    # Test invalid arguments
+    for datasets_to_split in ["train", ("train"), {}]:
+        with pytest.raises(TypeError, match="`datasets_to_split` should be.*"):
+            test = TestBasic(
+                2,
+                True,
+                {"train": 5, "test": 10},
+                datasets={"train": ds, "test": ds},
+                dataset_config=DataConfig(datasets_to_split=datasets_to_split),
+            )
+
+    # Test empty `datasets_to_split` list
+    test = TestBasic(
+        2,
+        True,
+        {"train": 10, "test": 10},
+        datasets={"train": ds, "test": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+
+def test_configure_execution_options_carryover_context(ray_start_4_cpus):
+    """Tests that execution options in DataContext are carried over to DatConfig
+    automatically."""
+
+    ctx = ray.data.DataContext.get_current()
+    ctx.execution_options.preserve_order = True
+    ctx.execution_options.verbose_progress = True
+
+    data_config = DataConfig()
+
+    ingest_options = data_config.default_ingest_options()
+    assert ingest_options.preserve_order is True
+    assert ingest_options.verbose_progress is True
+
+
+@pytest.mark.parametrize("enable_locality", [True, False])
+def test_configure_locality(enable_locality):
+    data_config = DataConfig(enable_shard_locality=enable_locality)
+
+    mock_ds = MagicMock()
+    mock_ds.streaming_split = MagicMock()
+    mock_ds.copy = MagicMock(return_value=mock_ds)
+    world_size = 2
+    worker_handles = [MagicMock() for _ in range(world_size)]
+    worker_node_ids = ["node" + str(i) for i in range(world_size)]
+    data_config.configure(
+        datasets={"train": mock_ds},
+        world_size=world_size,
+        worker_handles=worker_handles,
+        worker_node_ids=worker_node_ids,
+    )
+    mock_ds.streaming_split.assert_called_once()
+    mock_ds.streaming_split.assert_called_with(
+        world_size,
+        equal=True,
+        locality_hints=worker_node_ids if enable_locality else None,
+    )
+
+
+class CustomConfig(DataConfig):
+    def __init__(self):
+        pass
+
+    def configure(self, *args, **kwargs):
+        ds = ray.data.range(10)
+        return [
+            {"train": ds.iterator()},
+            {"train": ds.iterator()},
+        ]
+
+
+def test_custom_config_subclass(ray_start_4_cpus):
+    test = TestBasic(
+        1,
+        True,
+        {"train": 10},
+        dataset_config=CustomConfig(),
+    )
+    test.fit()
+
+
+class TestRandom(DataParallelTrainer):
+    def __init__(self, num_workers: int, expect_random: bool, **kwargs):
+        def train_loop_per_worker():
+            data_shard = train.get_dataset_shard("train")
+            assert isinstance(data_shard, DataIterator), data_shard
+            epoch1 = list(data_shard.iter_rows())
+            epoch2 = list(data_shard.iter_rows())
+            print("Epochs", epoch1, "\n", epoch2)
+            if expect_random:
+                assert epoch1 != epoch2
+            else:
+                assert epoch1 == epoch2
+
+        kwargs.pop("scaling_config", None)
+        super().__init__(
+            train_loop_per_worker=train_loop_per_worker,
+            scaling_config=ScalingConfig(num_workers=num_workers),
+            **kwargs,
+        )
+
+
+def test_per_epoch_preprocessing(ray_start_4_cpus):
+    ds = ray.data.range(100, override_num_blocks=100).randomize_block_order()
+    test = TestRandom(2, True, datasets={"train": ds})
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).random_shuffle()
+    test = TestRandom(2, True, datasets={"train": ds})
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).map(
+        lambda x: {"id": x["id"] * random.random()}
+    )
+    test = TestRandom(2, True, datasets={"train": ds})
+    test.fit()
+
+
+def test_materialized_preprocessing(ray_start_4_cpus):
+    # TODO(ekl) we should test all these configs with splitting enabled, but this
+    # requires implementing deterministic streaming split.
+    ds = ray.data.range(100, override_num_blocks=100).randomize_block_order()
+    ds = ds.materialize()
+    test = TestRandom(
+        2,
+        False,
+        datasets={"train": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).random_shuffle()
+    ds = ds.materialize()
+    test = TestRandom(
+        2,
+        False,
+        datasets={"train": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+    ds = ray.data.range(100, override_num_blocks=100).map(
+        lambda x: {"id": x["id"] * random.random()}
+    )
+    ds = ds.materialize()
+    test = TestRandom(
+        2,
+        False,
+        datasets={"train": ds},
+        dataset_config=DataConfig(datasets_to_split=[]),
+    )
+    test.fit()
+
+
+def _run_data_config_resource_test(data_config):
+    cluster_cpus, cluster_gpus = 20, 10
+    num_workers = 2
+    # Resources used by training workers.
+    cpus_per_worker, gpus_per_worker = 2, 1
+
+    original_execution_options = data_config._get_execution_options("train")
+
+    ray.init(num_cpus=cluster_cpus, num_gpus=cluster_gpus)
+
+    class MyTrainer(DataParallelTrainer):
+        def __init__(self, **kwargs):
+            def train_loop_fn():
+                train_ds = train.get_dataset_shard("train")
+                new_execution_options = train_ds.get_context().execution_options
+                if original_execution_options.is_resource_limits_default():
+                    # If the original resource limits are default, the new resource
+                    # limits should be the default as well.
+                    assert new_execution_options.is_resource_limits_default()
+                    exclude_resources = new_execution_options.exclude_resources
+                    assert (
+                        exclude_resources.cpu
+                        == original_execution_options.exclude_resources.cpu
+                        + cpus_per_worker * num_workers
+                        + 1  # trainer coordinator
+                    )
+                    assert (
+                        exclude_resources.gpu
+                        == original_execution_options.exclude_resources.gpu
+                        + gpus_per_worker * num_workers
+                    )
+                else:
+                    # If the original resource limits are not default, the new resource
+                    # limits should be the same as the original ones.
+                    # And the new exclude_resources should be zero.
+                    assert (
+                        new_execution_options.resource_limits
+                        == original_execution_options.resource_limits
+                    )
+                    assert (
+                        new_execution_options.exclude_resources
+                        == ExecutionResources.zero()
+                    )
+
+            kwargs.pop("scaling_config", None)
+
+            super().__init__(
+                train_loop_per_worker=train_loop_fn,
+                scaling_config=ScalingConfig(
+                    num_workers=num_workers,
+                    use_gpu=True,
+                    resources_per_worker={
+                        "CPU": cpus_per_worker,
+                        "GPU": gpus_per_worker,
+                    },
+                ),
+                datasets={"train": ray.data.range(10)},
+                dataset_config=data_config,
+                **kwargs,
+            )
+
+    trainer = MyTrainer()
+    trainer.fit()
+
+
+def test_data_config_default_resource_limits(shutdown_only):
+    """Test that DataConfig preserves user-configured exclude_resources."""
+    execution_options = ExecutionOptions()
+    execution_options.exclude_resources = execution_options.exclude_resources.copy(
+        cpu=2, gpu=1
+    )
+    data_config = DataConfig(execution_options=execution_options)
+
+    _run_data_config_resource_test(data_config)
+
+
+def test_data_config_manual_resource_limits(shutdown_only):
+    """Test manually setting resource limits in DataConfig."""
+    execution_options = ExecutionOptions()
+    execution_options.resource_limits = execution_options.resource_limits.copy(
+        cpu=10, gpu=5
+    )
+    data_config = DataConfig(execution_options=execution_options)
+
+    _run_data_config_resource_test(data_config)
+
+
+def test_v1_train_with_v2_data_autoscaler_sets_exclude_resources(
+    shutdown_only, monkeypatch
+):
+    """Regression test for the Train V1 + V2 cluster autoscaler combination."""
+    monkeypatch.setenv("RAY_DATA_CLUSTER_AUTOSCALER", "V2")
+
+    ray.init(num_cpus=10, num_gpus=2)
+
+    num_train_cpus, num_train_gpus = 4.0, 2.0
+    data_config = DataConfig()
+    data_config.set_train_total_resources(
+        num_train_cpus=num_train_cpus, num_train_gpus=num_train_gpus
+    )
+
+    iterators = data_config.configure(
+        datasets={"train": ray.data.range(10)},
+        world_size=2,
+        worker_handles=None,
+        worker_node_ids=None,
+    )
+
+    exclude_resources = (
+        iterators[0]["train"].get_context().execution_options.exclude_resources
+    )
+    assert exclude_resources.cpu == num_train_cpus
+    assert exclude_resources.gpu == num_train_gpus
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))