chore: import upstream snapshot with attribution

python/ray/data/tests/test_auto_parallelism.py

@@ -0,0 +1,184 @@
+from dataclasses import astuple, dataclass
+
+import pytest
+
+import ray
+from ray.data._internal.util import _autodetect_parallelism
+from ray.data.context import DataContext
+from ray.tests.conftest import *  # noqa
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+
+@dataclass
+class TestCase:
+    avail_cpus: int
+    target_max_block_size: int
+    data_size: int
+    expected_parallelism: int
+
+
+MiB = 1024 * 1024
+GiB = 1024 * MiB
+
+TEST_CASES = [
+    TestCase(
+        avail_cpus=4,
+        target_max_block_size=DataContext.get_current().target_max_block_size,
+        data_size=1024,
+        expected_parallelism=8,  # avail_cpus has precedence
+    ),
+    TestCase(
+        avail_cpus=4,
+        target_max_block_size=DataContext.get_current().target_max_block_size,
+        data_size=10 * MiB,
+        expected_parallelism=10,  # MIN_BLOCK_SIZE has precedence
+    ),
+    TestCase(
+        avail_cpus=4,
+        target_max_block_size=DataContext.get_current().target_max_block_size,
+        data_size=20 * MiB,
+        expected_parallelism=20,  # MIN_BLOCK_SIZE has precedence
+    ),
+    TestCase(
+        avail_cpus=4,
+        target_max_block_size=DataContext.get_current().target_max_block_size,
+        data_size=100 * MiB,
+        expected_parallelism=100,  # MIN_BLOCK_SIZE has precedence
+    ),
+    TestCase(
+        avail_cpus=4,
+        target_max_block_size=DataContext.get_current().target_max_block_size,
+        data_size=1 * GiB,
+        expected_parallelism=200,  # MIN_PARALLELISM has precedence
+    ),
+    TestCase(
+        avail_cpus=4,
+        target_max_block_size=DataContext.get_current().target_max_block_size,
+        data_size=10 * GiB,
+        expected_parallelism=200,  # MIN_PARALLELISM has precedence
+    ),
+    TestCase(
+        avail_cpus=150,
+        target_max_block_size=DataContext.get_current().target_max_block_size,
+        data_size=10 * GiB,
+        expected_parallelism=300,  # avail_cpus has precedence
+    ),
+    TestCase(
+        avail_cpus=400,
+        target_max_block_size=DataContext.get_current().target_max_block_size,
+        data_size=10 * GiB,
+        expected_parallelism=800,  # avail_cpus has precedence
+    ),
+    TestCase(
+        avail_cpus=400,
+        target_max_block_size=DataContext.get_current().target_max_block_size,
+        data_size=1 * MiB,
+        expected_parallelism=800,  # avail_cpus has precedence
+    ),
+    TestCase(
+        avail_cpus=4,
+        target_max_block_size=DataContext.get_current().target_max_block_size,
+        data_size=1000 * GiB,
+        expected_parallelism=8000,  # MAX_BLOCK_SIZE has precedence
+    ),
+    TestCase(
+        avail_cpus=4,
+        target_max_block_size=DataContext.get_current().target_max_block_size,
+        data_size=10000 * GiB,
+        expected_parallelism=80000,  # MAX_BLOCK_SIZE has precedence
+    ),
+    TestCase(
+        avail_cpus=4,
+        target_max_block_size=512 * MiB,
+        data_size=1000 * GiB,
+        expected_parallelism=2000,  # passed max_block_size has precedence
+    ),
+    TestCase(
+        avail_cpus=4,
+        target_max_block_size=512 * MiB,
+        data_size=10000 * GiB,
+        expected_parallelism=20000,  # passed max_block_size has precedence
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "avail_cpus,target_max_block_size,data_size,expected",
+    [astuple(test) for test in TEST_CASES],
+)
+def test_autodetect_parallelism(
+    shutdown_only, avail_cpus, target_max_block_size, data_size, expected
+):
+    class MockReader:
+        def estimate_inmemory_data_size(self):
+            return data_size
+
+    result, _, _ = _autodetect_parallelism(
+        parallelism=-1,
+        target_max_block_size=target_max_block_size,
+        ctx=DataContext.get_current(),
+        datasource_or_legacy_reader=MockReader(),
+        avail_cpus=avail_cpus,
+    )
+    assert result == expected, (result, expected)
+
+
+def test_auto_parallelism_basic(shutdown_only):
+    ray.init(num_cpus=8)
+    context = DataContext.get_current()
+    context.read_op_min_num_blocks = 1
+    # Datasource bound.
+    ds = ray.data.range_tensor(5, shape=(100,), override_num_blocks=-1)
+    assert ds._logical_plan.initial_num_blocks() == 5, ds
+    # CPU bound. TODO(ekl) we should fix range datasource to respect parallelism more
+    # properly, currently it can go a little over.
+    ds = ray.data.range_tensor(10000, shape=(100,), override_num_blocks=-1)
+    assert ds._logical_plan.initial_num_blocks() == 16, ds
+    # Block size bound.
+    ds = ray.data.range_tensor(100000000, shape=(100,), override_num_blocks=-1)
+    assert ds._logical_plan.initial_num_blocks() >= 590, ds
+    assert ds._logical_plan.initial_num_blocks() <= 600, ds
+
+
+def test_auto_parallelism_placement_group(shutdown_only):
+    ray.init(num_cpus=16, num_gpus=8)
+
+    @ray.remote
+    def run():
+        context = DataContext.get_current()
+        context.min_parallelism = 1
+        ds = ray.data.range_tensor(2000, shape=(100,), override_num_blocks=-1)
+        return ds._logical_plan.initial_num_blocks()
+
+    # 1/16 * 4 * 16 = 4
+    pg = ray.util.placement_group([{"CPU": 1}])
+    num_blocks = ray.get(
+        run.options(
+            scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
+        ).remote()
+    )
+    assert num_blocks == 4, num_blocks
+
+    # 2/16 * 4 * 16 = 8
+    pg = ray.util.placement_group([{"CPU": 2}])
+    num_blocks = ray.get(
+        run.options(
+            scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
+        ).remote()
+    )
+    assert num_blocks == 8, num_blocks
+
+    # 1/8 * 4 * 16 = 8
+    pg = ray.util.placement_group([{"CPU": 1, "GPU": 1}])
+    num_blocks = ray.get(
+        run.options(
+            scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
+        ).remote()
+    )
+    assert num_blocks == 8, num_blocks
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))