chore: import upstream snapshot with attribution

python/ray/train/v2/tests/test_dataset_manager.py

@@ -0,0 +1,233 @@
+import asyncio
+from unittest.mock import MagicMock
+
+import pytest
+
+import ray
+import ray.data
+from ray.data import DataContext
+from ray.data._internal.iterator.stream_split_iterator import StreamSplitDataIterator
+from ray.train.v2._internal.data_integration.dataset_manager import (
+    DatasetManager,
+    DatasetShardMetadata,
+)
+
+
+async def get_dataset_shard_for_worker(
+    dataset_manager: DatasetManager,
+    dataset_name: str,
+    worker_rank: int,
+):
+    return await asyncio.create_task(
+        dataset_manager.get_dataset_shard(
+            DatasetShardMetadata(dataset_name=dataset_name, world_rank=worker_rank)
+        )
+    )
+
+
+async def get_dataset_shard_for_all_workers(
+    dataset_manager: DatasetManager,
+    dataset_name: str,
+    num_workers: int,
+):
+    return await asyncio.gather(
+        *[
+            get_dataset_shard_for_worker(dataset_manager, dataset_name, i)
+            for i in range(num_workers)
+        ]
+    )
+
+
+@pytest.mark.asyncio
+async def test_get_multiple_datasets_serially(ray_start_4_cpus):
+    """Tests DatasetManager.get_dataset_shard for multiple datasets,
+    called serially by each worker. This is the typical case.
+    Workers 0, 1:
+    ray.train.get_dataset_shard("sharded_1")
+    ray.train.get_dataset_shard("sharded_2")
+    ray.train.get_dataset_shard("unsharded")
+    """
+
+    NUM_ROWS = 100
+    NUM_TRAIN_WORKERS = 2
+
+    sharded_ds_1 = ray.data.range(NUM_ROWS)
+    sharded_ds_2 = ray.data.range(NUM_ROWS)
+    unsharded_ds = ray.data.range(NUM_ROWS)
+
+    dataset_manager = DatasetManager(
+        datasets={
+            "sharded_1": sharded_ds_1,
+            "sharded_2": sharded_ds_2,
+            "unsharded": unsharded_ds,
+        },
+        data_config=ray.train.DataConfig(datasets_to_split=["sharded_1", "sharded_2"]),
+        data_context=DataContext.get_current(),
+        world_size=NUM_TRAIN_WORKERS,
+        worker_node_ids=None,
+    )
+
+    shards = await get_dataset_shard_for_all_workers(
+        dataset_manager, "sharded_1", NUM_TRAIN_WORKERS
+    )
+    assert all(isinstance(shard, StreamSplitDataIterator) for shard in shards)
+    assert all(shard._get_dataset_tag().startswith("sharded_1_") for shard in shards)
+
+    shards = await get_dataset_shard_for_all_workers(
+        dataset_manager, "sharded_2", NUM_TRAIN_WORKERS
+    )
+    assert all(isinstance(shard, StreamSplitDataIterator) for shard in shards)
+    assert all(shard._get_dataset_tag().startswith("sharded_2_") for shard in shards)
+
+    shards = await get_dataset_shard_for_all_workers(
+        dataset_manager, "unsharded", NUM_TRAIN_WORKERS
+    )
+    assert not any(isinstance(shard, StreamSplitDataIterator) for shard in shards)
+    assert all(shard._get_dataset_tag().startswith("unsharded_") for shard in shards)
+
+
+@pytest.mark.asyncio
+async def test_get_multiple_datasets_interleaved(ray_start_4_cpus):
+    """Tests DatasetManager.get_dataset_shard for multiple datasets,
+    called in an interleaved order by workers.
+    Worker 0:
+    ray.train.get_dataset_shard("train")
+    ray.train.get_dataset_shard("valid")
+    Worker 1:
+    ray.train.get_dataset_shard("valid")
+    ray.train.get_dataset_shard("train")
+    """
+
+    NUM_ROWS = 100
+    NUM_TRAIN_WORKERS = 2
+
+    train_ds = ray.data.range(NUM_ROWS)
+    valid_ds = ray.data.range(NUM_ROWS)
+
+    dataset_manager = DatasetManager(
+        datasets={"train": train_ds, "valid": valid_ds},
+        data_config=ray.train.DataConfig(datasets_to_split="all"),
+        data_context=DataContext.get_current(),
+        world_size=NUM_TRAIN_WORKERS,
+        worker_node_ids=None,
+    )
+
+    tasks = [
+        get_dataset_shard_for_worker(dataset_manager, "train", 0),
+        get_dataset_shard_for_worker(dataset_manager, "valid", 1),
+        get_dataset_shard_for_worker(dataset_manager, "train", 1),
+        get_dataset_shard_for_worker(dataset_manager, "valid", 0),
+    ]
+    iterators = await asyncio.gather(*tasks)
+    assert all(isinstance(iterator, StreamSplitDataIterator) for iterator in iterators)
+    expected_names = ["train", "valid", "train", "valid"]
+    assert all(
+        it._get_dataset_tag().startswith(f"{name}_")
+        for it, name in zip(iterators, expected_names)
+    )
+
+
+@pytest.mark.asyncio
+async def test_get_multiple_datasets_rank_specific(ray_start_4_cpus):
+    """Tests rank-specific DatasetManager.get_dataset_shard calls.
+    # Epoch 1
+    ray.train.get_dataset_shard("train")
+    # Validation, which only happens on worker 0.
+    if world_rank == 0:
+        ray.train.get_dataset_shard("valid")
+    # Epoch 2
+    ray.train.get_dataset_shard("train")
+    """
+
+    NUM_ROWS = 100
+    NUM_TRAIN_WORKERS = 2
+
+    train_ds = ray.data.range(NUM_ROWS)
+    valid_ds = ray.data.range(NUM_ROWS)
+
+    dataset_manager = DatasetManager(
+        datasets={"train": train_ds, "valid": valid_ds},
+        data_config=ray.train.DataConfig(datasets_to_split=["train"]),
+        data_context=DataContext.get_current(),
+        world_size=NUM_TRAIN_WORKERS,
+        worker_node_ids=None,
+    )
+
+    # ray.train.get_dataset_shard("train")
+    iterators = await get_dataset_shard_for_all_workers(
+        dataset_manager, "train", NUM_TRAIN_WORKERS
+    )
+    assert all(isinstance(iterator, StreamSplitDataIterator) for iterator in iterators)
+    assert all(it._get_dataset_tag().startswith("train_") for it in iterators)
+
+    # if world_rank == 0:
+    #     ray.train.get_dataset_shard("valid")
+    iterator = await get_dataset_shard_for_worker(dataset_manager, "valid", 0)
+    assert not isinstance(iterator, StreamSplitDataIterator)
+    assert iterator._get_dataset_tag().startswith("valid_")
+
+    # ray.train.get_dataset_shard("train")
+    iterators = await get_dataset_shard_for_all_workers(
+        dataset_manager, "train", NUM_TRAIN_WORKERS
+    )
+    assert all(isinstance(iterator, StreamSplitDataIterator) for iterator in iterators)
+    assert all(it._get_dataset_tag().startswith("train_") for it in iterators)
+
+
+@pytest.mark.asyncio
+async def test_dataset_manager_shutdown_multiple_datasets(ray_start_4_cpus):
+    """The DatasetManager collects SplitCoordinator actors for sharded datasets
+    and triggers executor shutdown on them.
+    """
+
+    NUM_ROWS = 100
+    NUM_TRAIN_WORKERS = 2
+
+    sharded_ds_1 = ray.data.range(NUM_ROWS)
+    sharded_ds_2 = ray.data.range(NUM_ROWS)
+    unsharded_ds = ray.data.range(NUM_ROWS)
+
+    dataset_manager = DatasetManager(
+        datasets={
+            "sharded_1": sharded_ds_1,
+            "sharded_2": sharded_ds_2,
+            "unsharded": unsharded_ds,
+        },
+        data_config=ray.train.DataConfig(datasets_to_split=["sharded_1", "sharded_2"]),
+        data_context=DataContext.get_current(),
+        world_size=NUM_TRAIN_WORKERS,
+        worker_node_ids=None,
+    )
+
+    await get_dataset_shard_for_all_workers(
+        dataset_manager, "sharded_1", NUM_TRAIN_WORKERS
+    )
+    assert len(dataset_manager._coordinator_actors) == 1
+    assert isinstance(dataset_manager._coordinator_actors[0], ray.actor.ActorHandle)
+
+    await get_dataset_shard_for_all_workers(
+        dataset_manager, "sharded_2", NUM_TRAIN_WORKERS
+    )
+    assert len(dataset_manager._coordinator_actors) == 2
+    assert isinstance(dataset_manager._coordinator_actors[1], ray.actor.ActorHandle)
+
+    # Unsharded datasets are not tracked for shutdown.
+    await get_dataset_shard_for_all_workers(
+        dataset_manager, "unsharded", NUM_TRAIN_WORKERS
+    )
+    assert len(dataset_manager._coordinator_actors) == 2
+
+    mocks = [MagicMock() for _ in range(2)]
+    remote_mocks = [mock.shutdown_executor.remote for mock in mocks]
+    dataset_manager._coordinator_actors = mocks
+
+    dataset_manager.shutdown_data_executors()
+
+    for remote_mock in remote_mocks:
+        remote_mock.assert_called_once()
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", "-x", __file__]))