diff --git a/python/cudf_polars/tests/conftest.py b/python/cudf_polars/tests/conftest.py index 47412f1b2966..9bdd755177fe 100644 --- a/python/cudf_polars/tests/conftest.py +++ b/python/cudf_polars/tests/conftest.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import pytest @@ -23,13 +23,38 @@ from cudf_polars.engine.core import StreamingEngine from cudf_polars.engine.dask import DaskEngine from cudf_polars.engine.options import StreamingOptions + from cudf_polars.engine.ray import RayEngine from cudf_polars.engine.spmd import SPMDEngine -# Number of ranks for multi-rank streaming engines that share one GPU -# (currently ``RayEngine``). Single-GPU dev hosts and CI runners require -# ``allow_gpu_sharing=True`` to oversubscribe one device across actors. -NUM_RANKS = 2 +@pytest.fixture(scope="session") +def ray_num_ranks() -> int: + """ + Number of ranks for multi-rank streaming engines that share one GPU + (currently ``RayEngine``). Single-GPU dev hosts and CI runners require + ``allow_gpu_sharing=True`` to oversubscribe one device across actors. + """ + return 2 + + +@pytest.fixture(scope="session") +def ray_init_options(ray_num_ranks: int) -> dict[str, Any]: + """ + Keyword arguments forwarded to ``ray.init`` for the test Ray cluster. + + When using this fixture, a ``RayEngine`` must be constructed with: + - ``num_ranks`` set + - ``engine_options={"allow_gpu_sharing": True}`` + + This is required because the cluster is configured with ``num_gpus=0``, + so Ray does not autodetect or track GPU resources. + """ + return { + "num_cpus": ray_num_ranks, + "num_gpus": 0, + "include_dashboard": False, + "object_store_memory": 256 * 1024 * 1024, # 256 MB + } @pytest.fixture(params=[False, True], ids=["no_nulls", "nulls"], scope="session") @@ -66,6 +91,8 @@ def _engine_param(request: pytest.FixtureRequest) -> EngineFixtureParam: @pytest.fixture(scope="session") def _unconfigured_engine( _engine_param: EngineFixtureParam, + ray_num_ranks: int, + ray_init_options: dict[str, Any], ) -> Generator[tuple[pl.GPUEngine, StreamingOptions | None], None, None]: """ Fixture generating an engine resource and options to apply before use. @@ -111,9 +138,9 @@ def _unconfigured_engine( # otherwise ``RayEngine`` defaults to # ``get_num_gpus_in_ray_cluster()`` engine = RayEngine( - num_ranks=NUM_RANKS, + num_ranks=ray_num_ranks, engine_options={"allow_gpu_sharing": True}, - ray_init_options={"include_dashboard": False}, + ray_init_options=ray_init_options, ) case _: # pragma: no cover raise ValueError( @@ -165,6 +192,15 @@ def factory(options: StreamingOptions) -> SPMDEngine: return factory +@pytest.fixture +def ray_engine( + _unconfigured_engine: tuple[RayEngine, StreamingOptions], +) -> RayEngine: + """Return the shared configured :class:`RayEngine`.""" + engine, options = _unconfigured_engine + return configure_streaming_engine(engine, options) + + @pytest.fixture def dask_engine( _unconfigured_engine: tuple[DaskEngine, StreamingOptions], @@ -304,6 +340,8 @@ def pytest_generate_tests(metafunc: pytest.Metafunc): if "spmd_engine" in fixtures or "spmd_engine_factory" in fixtures: engines = ["spmd"] + elif "ray_engine" in fixtures: + engines = ["ray"] elif "dask_engine" in fixtures: engines = ["dask"] elif "streaming_engine" in fixtures or "streaming_engine_factory" in fixtures: diff --git a/python/cudf_polars/tests/streaming/test_ray.py b/python/cudf_polars/tests/streaming/test_ray.py index 8427099fb021..a4e543d5c628 100644 --- a/python/cudf_polars/tests/streaming/test_ray.py +++ b/python/cudf_polars/tests/streaming/test_ray.py @@ -5,7 +5,7 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from unittest.mock import patch import pytest @@ -22,22 +22,6 @@ if TYPE_CHECKING: from collections.abc import Iterator -NUM_RANKS = 2 - - -@pytest.fixture(scope="module") -def engine() -> Iterator[RayEngine]: - """Create one Ray cluster + GPU actors shared across the test session.""" - with RayEngine( - # Use a small partition size so tests exercise the multi-partition - # code path deterministically, regardless of input size. - executor_options={"max_rows_per_partition": 10}, - engine_options={"allow_gpu_sharing": True}, - num_ranks=NUM_RANKS, - ray_init_options={"include_dashboard": False}, - ) as engine: - yield engine - pytestmark = [ pytest.mark.skipif( @@ -78,12 +62,12 @@ def test_raises_inside_rrun() -> None: RayEngine() -def test_num_ranks_requires_allow_gpu_sharing() -> None: +def test_num_ranks_requires_allow_gpu_sharing(ray_num_ranks: int) -> None: """num_ranks requires engine_options['allow_gpu_sharing']=True.""" with pytest.raises(ValueError, match="allow_gpu_sharing"): - RayEngine(num_ranks=NUM_RANKS) + RayEngine(num_ranks=ray_num_ranks) with pytest.raises(ValueError, match="allow_gpu_sharing"): - RayEngine(num_ranks=NUM_RANKS, engine_options={"allow_gpu_sharing": False}) + RayEngine(num_ranks=ray_num_ranks, engine_options={"allow_gpu_sharing": False}) def test_num_ranks_must_be_positive() -> None: @@ -93,138 +77,65 @@ def test_num_ranks_must_be_positive() -> None: # --------------------------------------------------------------------------- -# GPU tests — share a single Ray cluster + actor set for the whole session +# GPU tests — reuse the session-scoped Ray cluster from conftest # --------------------------------------------------------------------------- -def test_yields_engine( - engine: RayEngine, -) -> None: +def test_yields_engine(ray_engine: RayEngine) -> None: """RayEngine is a GPUEngine with at least one rank.""" - assert isinstance(engine, pl.GPUEngine) - assert engine.nranks >= 1 + assert isinstance(ray_engine, pl.GPUEngine) + assert ray_engine.nranks >= 1 -def test_executor_options_forwarded( - engine: RayEngine, -) -> None: +def test_executor_options_forwarded(ray_engine: RayEngine) -> None: """Reserved executor_options keys are injected into the engine config.""" - opts = engine.config["executor_options"] + opts = ray_engine.config["executor_options"] assert opts["cluster"] == "ray" assert isinstance(opts["ray_context"], RayContext) - assert engine.rank_actors == opts["ray_context"].rank_actors - assert len(engine.rank_actors) == engine.nranks + assert ray_engine.rank_actors == opts["ray_context"].rank_actors + assert len(ray_engine.rank_actors) == ray_engine.nranks -def test_gather_cluster_info(engine: RayEngine) -> None: +def test_gather_cluster_info(ray_engine: RayEngine) -> None: """gather_cluster_info returns one ClusterInfo per rank with expected fields.""" - infos = engine.gather_cluster_info() - assert len(infos) == engine.nranks + infos = ray_engine.gather_cluster_info() + assert len(infos) == ray_engine.nranks for info in infos: assert isinstance(info.hostname, str) assert isinstance(info.pid, int) # Each actor runs in its own process. - assert len({info.pid for info in infos}) == engine.nranks - - -def test_scan(engine: RayEngine) -> None: - """Input rows are partitioned across actors; total output equals input.""" - lf = pl.LazyFrame({"a": [1, 2, 3]}) - result = lf.collect(engine=engine) - assert result.shape == (3, 1) - assert sorted(result["a"].to_list()) == [1, 2, 3] - - -def test_filter(engine: RayEngine) -> None: - """Filter is applied correctly across all actors.""" - lf = pl.LazyFrame({"a": [1, 2, 3, 4, 5]}) - result = lf.filter(pl.col("a") > 3).collect(engine=engine) - assert result.shape == (2, 1) - assert sorted(result["a"].to_list()) == [4, 5] - - -def test_group_by(engine: RayEngine) -> None: - """Group-by produces the correct aggregation across all ranks.""" - # max_rows_per_partition=10 (set on the session fixture) gives each rank - # exactly 5 partitions, so the multi-partition path is always exercised. - n, n_keys = engine.nranks * 50, 5 - keys = [str(i % n_keys) for i in range(n)] - vals = list(range(n)) - lf = pl.LazyFrame({"key": keys, "val": vals}) - result = ( - lf.group_by("key").agg(pl.col("val").sum()).collect(engine=engine).sort("key") - ) - expected = ( - pl.LazyFrame({"key": keys, "val": vals}) - .group_by("key") - .agg(pl.col("val").sum()) - .collect() - .sort("key") - ) - assert result.shape == expected.shape - assert result["key"].to_list() == expected["key"].to_list() - assert result["val"].to_list() == expected["val"].to_list() - - -def test_join(engine: RayEngine) -> None: - """Hash join between two tables produces the correct result across all ranks.""" - # max_rows_per_partition=10 (set on the session fixture) gives each rank - # exactly 5 partitions, so the multi-partition path is always exercised. - n = engine.nranks * 50 - lf_left = pl.LazyFrame({"key": list(range(n)), "val_left": list(range(n))}) - lf_right = pl.LazyFrame( - {"key": list(range(n)), "val_right": [x * 2 for x in range(n)]} - ) - result = lf_left.join(lf_right, on="key").collect(engine=engine).sort("key") - assert result.shape == (n, 3) - assert result["val_left"].to_list() == list(range(n)) - assert result["val_right"].to_list() == [x * 2 for x in range(n)] - - -def test_empty_dataframe(engine: RayEngine) -> None: - """An empty LazyFrame produces an empty result with the correct schema.""" - lf = pl.LazyFrame( - {"a": pl.Series([], dtype=pl.Int32), "b": pl.Series([], dtype=pl.Float64)} - ) - result = lf.collect(engine=engine) - assert result.shape == (0, 2) - assert result.columns == ["a", "b"] - assert result.dtypes == [pl.Int32, pl.Float64] + assert len({info.pid for info in infos}) == ray_engine.nranks -def test_run(engine: RayEngine) -> None: - result = engine._run(os.getpid) - assert len(set(result)) == engine.nranks +def test_run(ray_engine: RayEngine) -> None: + result = ray_engine._run(os.getpid) + assert len(set(result)) == ray_engine.nranks -def test_num_ranks_oversubscribes() -> None: +def test_num_ranks_oversubscribes(ray_engine: RayEngine, ray_num_ranks: int) -> None: """num_ranks creates the requested number of actors sharing GPU 0.""" - n = 2 - with RayEngine( - executor_options={"max_rows_per_partition": 10}, - engine_options={"allow_gpu_sharing": True}, - num_ranks=n, - ray_init_options={"include_dashboard": False}, - ) as engine: - assert engine.nranks == n - assert len(engine.rank_actors) == n - result = pl.LazyFrame({"a": [1, 2, 3, 4]}).collect(engine=engine) - assert sorted(result["a"].to_list()) == [1, 2, 3, 4] + assert ray_engine.nranks == ray_num_ranks + assert len(ray_engine.rank_actors) == ray_num_ranks + result = pl.LazyFrame({"a": [1, 2, 3, 4]}).collect(engine=ray_engine) + assert sorted(result["a"].to_list()) == [1, 2, 3, 4] @pytest.fixture(scope="module") -def reset_engine() -> Iterator[RayEngine]: - """Module-scoped engine for reset tests — independent of ``engine``. +def reset_engine( + ray_num_ranks: int, + ray_init_options: dict[str, Any], +) -> Iterator[RayEngine]: + """Module-scoped engine for reset tests — independent of ``ray_engine``. These tests exercise :meth:`RayEngine._reset` (which mutates the - engine in-place) and the shutdown guard. A dedicated fixture keeps - those mutations from leaking into the other tests. + engine in-place). A dedicated fixture keeps those mutations from + leaking into the conftest-shared ``ray_engine``. """ with RayEngine( executor_options={"max_rows_per_partition": 10}, engine_options={"allow_gpu_sharing": True}, - num_ranks=NUM_RANKS, - ray_init_options={"include_dashboard": False}, + num_ranks=ray_num_ranks, + ray_init_options=ray_init_options, ) as e: yield e @@ -273,13 +184,16 @@ def test_reset_collects_after_options_change(reset_engine: RayEngine) -> None: assert sorted(result["a"].to_list()) == [1, 2, 3, 4, 5] -def test_reset_after_shutdown_raises() -> None: +def test_reset_after_shutdown_raises( + ray_num_ranks: int, + ray_init_options: dict[str, Any], +) -> None: """``shutdown`` is idempotent; ``_reset`` after shutdown raises every time.""" engine = RayEngine( executor_options={"max_rows_per_partition": 10}, engine_options={"allow_gpu_sharing": True}, - num_ranks=NUM_RANKS, - ray_init_options={"include_dashboard": False}, + num_ranks=ray_num_ranks, + ray_init_options=ray_init_options, ) engine.shutdown() engine.shutdown() # idempotent @@ -321,13 +235,16 @@ def test_reset_rejects_construction_time_engine_options( ) -def test_shutdown_skips_when_ray_not_initialized() -> None: +def test_shutdown_skips_when_ray_not_initialized( + ray_num_ranks: int, + ray_init_options: dict[str, Any], +) -> None: """``shutdown`` short-circuits if ``ray.is_initialized()`` is ``False``.""" engine = RayEngine( executor_options={"max_rows_per_partition": 10}, engine_options={"allow_gpu_sharing": True}, - num_ranks=NUM_RANKS, - ray_init_options={"include_dashboard": False}, + num_ranks=ray_num_ranks, + ray_init_options=ray_init_options, ) try: with patch("ray.is_initialized", return_value=False):