diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py index 69e6a36dca29..7bc8dabddecb 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py @@ -13,6 +13,7 @@ import cuda.core from rapidsmpf.coll import AllGather +from rapidsmpf.config import Options, get_environment_variables from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.core.actor import run_actor_network @@ -50,6 +51,39 @@ T = TypeVar("T") +def resolve_rapidsmpf_options(rapidsmpf_options: Options | None) -> Options: + """ + Resolve ``rapidsmpf_options`` and apply cross-frontend defaults. + + If ``None`` is passed, constructs an ``Options`` instance from + environment variables. Then applies defaults that should be consistent + across SPMD, Ray, and Dask. Defaults are set via + ``Options.insert_if_absent``, so explicit values or environment + variables always take precedence. + + Defaults applied: + + - ``num_streaming_threads=4``: moderate worker count for the rapidsmpf + streaming runtime, shared across frontends. + + Parameters + ---------- + rapidsmpf_options + Existing options to resolve, or ``None`` to construct from environment + variables. + + Returns + ------- + Options + Resolved options with cross-frontend defaults applied. + """ + if rapidsmpf_options is None: + rapidsmpf_options = Options(get_environment_variables()) + + rapidsmpf_options.insert_if_absent({"num_streaming_threads": "4"}) + return rapidsmpf_options + + @dataclasses.dataclass(frozen=True) class ClusterInfo: """ @@ -201,6 +235,66 @@ def global_statistics(self, *, clear: bool = False) -> Statistics: """ return Statistics.merge(self.gather_statistics(clear=clear)) + def _reset( + self, + *, + rapidsmpf_options: Options | None = None, + executor_options: dict[str, Any] | None = None, + engine_options: dict[str, Any] | None = None, + ) -> None: + """ + Reset the engine with new options, keeping cluster resources alive. + + The following inputs are fixed at construction time and cannot change: + - ``num_ranks`` + - ``num_py_executors`` (in ``executor_options``) + - ``hardware_binding`` (in ``engine_options``) + - ``memory_resource_config`` (in ``engine_options``) + + Subclasses must override this method. The override should: + 1. Raise :class:`RuntimeError` if the engine is already shut down. + 2. Call ``super()._reset(...)`` to apply the universal option validation below. + 3. Perform the backend-specific rebuild. + + Parameters + ---------- + rapidsmpf_options + New :class:`Options` for each rank's :class:`Context`. + ``None`` is treated as an empty dict. + executor_options + New executor options for the polars ``GPUEngine`` layer. + ``None`` is treated as an empty dict. + engine_options + New engine options for the polars ``GPUEngine`` layer. + ``None`` is treated as an empty dict. + + Raises + ------ + ValueError + If ``executor_options`` or ``engine_options`` contains a + construction-time-only key (see list above), or if a + reserved key is set (via :func:`check_reserved_keys`). + """ + executor_options = executor_options or {} + engine_options = engine_options or {} + check_reserved_keys(executor_options, engine_options) + + _disallowed_exec = {"num_py_executors"} & executor_options.keys() + if _disallowed_exec: + raise ValueError( + f"executor_options keys {sorted(_disallowed_exec)} cannot be " + "changed via _reset(). Construct a fresh engine instead." + ) + _disallowed_engine = { + "hardware_binding", + "memory_resource_config", + } & engine_options.keys() + if _disallowed_engine: + raise ValueError( + f"engine_options keys {sorted(_disallowed_engine)} cannot be " + "changed via _reset(). Construct a fresh engine instead." + ) + def shutdown(self) -> None: """ Shut down engine and release all owned resources. diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py index eb32abcf375c..49810e998fd2 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py @@ -19,10 +19,7 @@ import ucxx._lib.libucxx as ucx_api from rapidsmpf import bootstrap from rapidsmpf.communicator.ucxx import barrier, get_root_ucxx_address, new_communicator -from rapidsmpf.config import ( - Options, - get_environment_variables, -) +from rapidsmpf.config import Options from rapidsmpf.progress_thread import ProgressThread from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.streaming.core.context import Context @@ -36,6 +33,7 @@ StreamingEngine, check_reserved_keys, evaluate_on_rank, + resolve_rapidsmpf_options, ) from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( HardwareBindingPolicy, @@ -294,6 +292,47 @@ def _teardown_worker( delattr(dask_worker, attr) +def _reset_worker( + rapidsmpf_options_as_bytes: bytes, + *, + uid: str, + dask_worker: distributed.Worker | None = None, +) -> None: + """ + Rebuild the streaming Context with new options. + + Must be called collectively on all workers. A barrier ensures no + worker tears down its Context while peers may still be using it. + + Parameters + ---------- + rapidsmpf_options_as_bytes + Serialized :class:`Options` to install. + uid + Cluster instance identifier used to look up the per-worker context. + dask_worker + Injected by ``distributed`` when called via :meth:`distributed.Client.run`. + """ + assert dask_worker is not None + attr = f"_cudf_polars_mp_context_{uid}" + mp_ctx: _WorkerContext | None = getattr(dask_worker, attr, None) + if mp_ctx is None: + raise RuntimeError(f"_reset_worker called before _setup_worker for uid={uid}") + assert mp_ctx.comm is not None + assert mp_ctx.ctx is not None + # Collective: all ranks idle before any rank tears down its Context. + if mp_ctx.comm.nranks > 1: + barrier(mp_ctx.comm) + # Explicit shutdown is thread-affine. ``distributed.worker.run`` + # dispatches sync work onto the worker's event-loop thread, which is + # the same thread that built the Context in ``_setup_worker``. + mp_ctx.ctx.shutdown() + mp_ctx.ctx = None + options = Options.deserialize(rapidsmpf_options_as_bytes) + mp_ctx.ctx = Context.from_options(mp_ctx.comm.logger, mp_ctx.mr, options) + rmm.mr.set_current_device_resource(mp_ctx.ctx.br().device_mr) + + def _get_statistics( *, clear: bool, uid: str, dask_worker: distributed.Worker | None = None ) -> tuple[int, Statistics]: @@ -563,13 +602,9 @@ def __init__( "memory_resource_config", None ) - rapidsmpf_options = ( + rapidsmpf_options_as_bytes = resolve_rapidsmpf_options( rapidsmpf_options - if rapidsmpf_options is not None - else Options(get_environment_variables()) - ) - rapidsmpf_options.insert_if_absent({"num_streaming_threads": "4"}) - rapidsmpf_options_as_bytes = rapidsmpf_options.serialize() + ).serialize() # Unique identifier for this cluster instance; namespaces the per-worker # attribute so multiple DaskEngine contexts can coexist on the same workers. @@ -660,6 +695,55 @@ def __init__( engine_options={**engine_options, "memory_resource": None}, ) + def _reset( + self, + *, + rapidsmpf_options: Options | None = None, + executor_options: dict[str, Any] | None = None, + engine_options: dict[str, Any] | None = None, + ) -> None: + """Reset the engine; see :meth:`StreamingEngine._reset` for the contract.""" + if self._dask_context is None: + raise RuntimeError("Cannot reset a shut-down engine") + super()._reset( + rapidsmpf_options=rapidsmpf_options, + executor_options=executor_options, + engine_options=engine_options, + ) + executor_options = executor_options or {} + engine_options = engine_options or {} + + rapidsmpf_options_as_bytes = resolve_rapidsmpf_options( + rapidsmpf_options + ).serialize() + + ctx = self._dask_context + # Reset all worker Contexts collectively. ``client.run`` blocks + # until every worker's reset returns; the per-worker barrier + # inside :func:`_reset_worker` synchronizes the teardown across + # workers. + ctx.client.run( + functools.partial(_reset_worker, uid=ctx.rapidsmpf_id), + rapidsmpf_options_as_bytes, + ) + + # Re-run ``StreamingEngine.__init__`` on the existing instance to + # reconfigure the polars ``GPUEngine`` layer (``self.config``, + # ``self.device``, etc.) with the new options. Pass the existing + # ``self._exit_stack`` so any registered callbacks survive. + StreamingEngine.__init__( + self, + nranks=self._nranks, + executor_options={ + **executor_options, + "runtime": "rapidsmpf", + "cluster": "dask", + "dask_context": ctx, + }, + engine_options={**engine_options, "memory_resource": None}, + exit_stack=self._exit_stack, + ) + @classmethod def from_options( cls, diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py index 47c882491239..1ba92de3e493 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py @@ -14,10 +14,7 @@ import ucxx._lib.libucxx as ucx_api from rapidsmpf import bootstrap from rapidsmpf.communicator.ucxx import barrier, get_root_ucxx_address, new_communicator -from rapidsmpf.config import ( - Options, - get_environment_variables, -) +from rapidsmpf.config import Options from rapidsmpf.progress_thread import ProgressThread from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.streaming.core.context import Context @@ -31,6 +28,7 @@ StreamingEngine, check_reserved_keys, evaluate_on_rank, + resolve_rapidsmpf_options, ) from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( HardwareBindingPolicy, @@ -256,18 +254,6 @@ def reset(self, *, rapidsmpf_options_as_bytes: bytes) -> None: """ Rebuild the streaming Context with new options. - Keeps the UCXX communicator, the :class:`RmmResourceAdaptor`, - and the Python thread-pool executor alive — only the rapidsmpf - :class:`Context` is replaced. Used by :meth:`RayEngine._reset` - to amortize actor startup and UCX bootstrap costs across engines - that differ only in streaming options. - - The RMM resource is *not* rebuilt: UCX maps CUDA IPC buffers - against it (notably for pool memory resources) and never - releases those mappings during the application lifetime, so a - rebuilt MR would silently leak pool memory. Construct a fresh - :class:`RayEngine` if you need to swap the memory resource. - Must be called collectively on all actors. A barrier ensures no rank tears down its Context while peers may still be using it. @@ -280,7 +266,8 @@ def reset(self, *, rapidsmpf_options_as_bytes: bytes) -> None: raise RuntimeError("reset() requires setup_worker() to have run") assert self._comm is not None # Collective: all ranks idle before any rank tears down its Context. - barrier(self._comm) + if self._comm.nranks > 1: + barrier(self._comm) self._ctx.shutdown() self._ctx = None self._rapidsmpf_options = Options.deserialize(rapidsmpf_options_as_bytes) @@ -544,13 +531,9 @@ def __init__( "memory_resource_config", None ) - rapidsmpf_options = ( + rapidsmpf_options_as_bytes = resolve_rapidsmpf_options( rapidsmpf_options - if rapidsmpf_options is not None - else Options(get_environment_variables()) - ) - rapidsmpf_options.insert_if_absent({"num_streaming_threads": "4"}) - rapidsmpf_options_as_bytes = rapidsmpf_options.serialize() + ).serialize() exit_stack = contextlib.ExitStack() if not ray.is_initialized(): @@ -621,73 +604,23 @@ def _reset( executor_options: dict[str, Any] | None = None, engine_options: dict[str, Any] | None = None, ) -> None: - """ - Reset the engine with new options. - - Fast path for consecutive ``RayEngine`` uses that differ only in - streaming options. Avoids Ray actor startup and UCX bootstrap. - - Replaces engine state in full, similar to :meth:`__init__`. - ``StreamingEngine`` revalidates invariants on each reset, so callers - must pass required options (for example, ``allow_gpu_sharing=True`` - when ``num_ranks > 1``). - - The following inputs are fixed at construction time and cannot change: - - ``num_ranks`` - - ``num_py_executors`` (in ``executor_options``) - - ``hardware_binding`` (in ``engine_options``) - - ``memory_resource_config`` (in ``engine_options``) - - ``ray_init_options`` - - Parameters - ---------- - rapidsmpf_options - New :class:`Options` for each actor's ``Context``. Defaults to - ``Options(get_environment_variables())`` if ``None``. - executor_options - Polars ``GPUEngine`` executor options. ``None`` is treated as - an empty dict. - engine_options - Polars ``GPUEngine`` options. ``None`` is treated as an empty - dict. - """ + """Reset the engine; see :meth:`StreamingEngine._reset` for the contract.""" if self._rank_actors is None: raise RuntimeError("Cannot reset a shut-down engine") - + super()._reset( + rapidsmpf_options=rapidsmpf_options, + executor_options=executor_options, + engine_options=engine_options, + ) executor_options = executor_options or {} engine_options = engine_options or {} - check_reserved_keys(executor_options, engine_options) - - # Reject keys that cannot be changed. - _disallowed_exec = {"num_py_executors"} & executor_options.keys() - if _disallowed_exec: - raise ValueError( - f"executor_options keys {sorted(_disallowed_exec)} cannot be " - "changed via _reset(). Construct a fresh RayEngine instead." - ) - _disallowed_engine = { - "hardware_binding", - "memory_resource_config", - } & engine_options.keys() - if _disallowed_engine: - raise ValueError( - f"engine_options keys {sorted(_disallowed_engine)} cannot be " - "changed via _reset(). Construct a fresh RayEngine instead." - ) - - rapidsmpf_options = ( + rapidsmpf_options_as_bytes = resolve_rapidsmpf_options( rapidsmpf_options - if rapidsmpf_options is not None - else Options(get_environment_variables()) - ) - rapidsmpf_options.insert_if_absent({"num_streaming_threads": "4"}) - rapidsmpf_options_as_bytes = rapidsmpf_options.serialize() + ).serialize() # Reset all actor Contexts collectively. ``ray.get`` blocks until # every actor's reset returns; the per-actor barrier inside # :meth:`RankActor.reset` synchronizes the teardown across ranks. - # The per-actor RMM resource is kept alive across resets — see - # :meth:`RankActor.reset`. ray.get( [ rank.reset.remote( diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py index 0f52f83c1a1b..65e3eb8b1e7f 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py +++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py @@ -15,7 +15,7 @@ from rapidsmpf.communicator.single import ( new_communicator as single_communicator, ) -from rapidsmpf.config import Options, get_environment_variables +from rapidsmpf.communicator.ucxx import barrier from rapidsmpf.integrations.cudf.partition import unpack_and_concat from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.progress_thread import ProgressThread @@ -36,6 +36,7 @@ all_gather_host_data, check_reserved_keys, evaluate_on_rank, + resolve_rapidsmpf_options, ) from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( HardwareBindingPolicy, @@ -49,6 +50,7 @@ from collections.abc import Callable from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.config import Options from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata from cudf_polars.dsl.ir import IR @@ -341,11 +343,7 @@ def __init__( ) bind_to_gpu(hw_binding) - rapidsmpf_options = ( - rapidsmpf_options - if rapidsmpf_options is not None - else Options(get_environment_variables()) - ) + rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) mr_config: MemoryResourceConfig | None = engine_options.get( "memory_resource_config", None ) @@ -369,17 +367,22 @@ def __init__( ) # else: caller-provided comm; the caller retains ownership - py_executor = ThreadPoolExecutor( + self._py_executor: ThreadPoolExecutor = ThreadPoolExecutor( max_workers=cast(int, executor_options.get("num_py_executors", 8)), thread_name_prefix="spmd-executor", ) + self._mr: RmmResourceAdaptor = mr exit_stack = contextlib.ExitStack() try: - exit_stack.callback(py_executor.shutdown, wait=False) + exit_stack.callback(self._py_executor.shutdown, wait=False) exit_stack.enter_context(set_memory_resource(mr)) - ctx = exit_stack.enter_context( - Context.from_options(comm.logger, mr, rapidsmpf_options) - ) + # ``Context`` is *not* registered as a context manager so that + # :meth:`_reset` can swap it mid-life without leaving the + # exit-stack holding a stale reference. ``_cleanup_ctx`` is + # registered instead — it shuts down whatever ``self._ctx`` is + # at engine-shutdown time (i.e. the latest reset's Context). + ctx = Context.from_options(comm.logger, mr, rapidsmpf_options) + exit_stack.callback(self._cleanup_ctx) self._comm: Communicator | None = comm self._ctx: Context | None = ctx super().__init__( @@ -389,7 +392,7 @@ def __init__( "runtime": "rapidsmpf", "cluster": "spmd", "spmd_context": SPMDContext( - comm=comm, context=ctx, py_executor=py_executor + comm=comm, context=ctx, py_executor=self._py_executor ), }, engine_options={ @@ -402,6 +405,17 @@ def __init__( exit_stack.close() raise + def _cleanup_ctx(self) -> None: + """ + Shut down the current ``self._ctx`` if any; called from exit-stack. + + ``Context.shutdown()`` is idempotent on the rapidsmpf C++ side, so this is + safe even if a prior ``_reset`` already shut down a now-replaced Context. + """ + if self._ctx is not None: + self._ctx.shutdown() + self._ctx = None + @classmethod def from_options(cls, options: StreamingOptions) -> SPMDEngine: """ @@ -436,6 +450,65 @@ def from_options(cls, options: StreamingOptions) -> SPMDEngine: engine_options=options.to_engine_options(), ) + def _reset( + self, + *, + rapidsmpf_options: Options | None = None, + executor_options: dict[str, Any] | None = None, + engine_options: dict[str, Any] | None = None, + ) -> None: + """ + Reset the engine; see :meth:`StreamingEngine._reset` for the contract. + + Must be called collectively on all ranks. A barrier ensures no + rank tears down its Context while peers may still be using it. + """ + if self._ctx is None: + raise RuntimeError("Cannot reset a shut-down engine") + assert self._comm is not None + super()._reset( + rapidsmpf_options=rapidsmpf_options, + executor_options=executor_options, + engine_options=engine_options, + ) + executor_options = executor_options or {} + engine_options = engine_options or {} + rapidsmpf_options = resolve_rapidsmpf_options(rapidsmpf_options) + + # Collective: synchronize all ranks before tearing down the Context. + if self._comm.nranks > 1: + barrier(self._comm) + # Same-thread shutdown, _reset runs on the thread that built the + # Context (the test driver's main thread). The per-engine RMM + # resource is kept alive across resets, see :meth:`_cleanup_ctx`. + self._ctx.shutdown() + self._ctx = Context.from_options(self._comm.logger, self._mr, rapidsmpf_options) + + # Re-run ``StreamingEngine.__init__`` on the existing instance to + # reconfigure the polars ``GPUEngine`` layer (``self.config``, + # ``self.device``, etc.) with the new options. Pass the existing + # ``self._exit_stack`` so any registered callbacks (notably + # ``_cleanup_ctx`` and ``set_memory_resource``) survive. + StreamingEngine.__init__( + self, + nranks=self._comm.nranks, + executor_options={ + **executor_options, + "runtime": "rapidsmpf", + "cluster": "spmd", + "spmd_context": SPMDContext( + comm=self._comm, + context=self._ctx, + py_executor=self._py_executor, + ), + }, + engine_options={ + **engine_options, + "memory_resource": self._ctx.br().device_mr, + }, + exit_stack=self._exit_stack, + ) + @property def rank(self) -> int: """ @@ -536,9 +609,14 @@ def shutdown(self) -> None: """ if self._ctx is None: return # already shut down + + # Order matters: ``super().shutdown()`` closes ``self._exit_stack``, + # which invokes ``self._cleanup_ctx``. That requires ``self._ctx`` to + # still be set so the rapidsmpf Context can be shut down correctly. + # Clear the references only after shutdown completes. + super().shutdown() self._comm = None self._ctx = None - super().shutdown() def _run(self, func: Callable[..., T], *args: Any, **kwargs: Any) -> list[T]: data = json.dumps(func(*args, **kwargs)).encode() diff --git a/python/cudf_polars/cudf_polars/testing/engine_utils.py b/python/cudf_polars/cudf_polars/testing/engine_utils.py index ec216dc6d88f..c36bcf2ed27a 100644 --- a/python/cudf_polars/cudf_polars/testing/engine_utils.py +++ b/python/cudf_polars/cudf_polars/testing/engine_utils.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, Any, Literal if TYPE_CHECKING: - from rapidsmpf.communicator.communicator import Communicator + from collections.abc import Mapping import polars as pl @@ -112,39 +112,49 @@ def create_streaming_options( def build_streaming_engine( param: EngineFixtureParam, - spmd_comm: Communicator, + engines: Mapping[str, StreamingEngine], options: StreamingOptions | None = None, ) -> StreamingEngine: """ - Build a :class:`StreamingEngine` from an engine fixture parameter. + Return ``engines``'s entry for ``param``, ``_reset``-ed. + + ``engines`` must already contain a slot for ``param.engine_name`` — + seeded by the ``streaming_engines`` session-scoped fixture. The + fixture owns mutation; this function only reads and ``_reset``-s. Parameters ---------- param Decoded engine fixture parameter describing the backend and block size mode. - spmd_comm - Communicator used when constructing an :class:`SPMDEngine`. + engines + Streaming-engine collection keyed by backend name. Provided by + the ``streaming_engines`` test fixture. options Optional streaming options to merge on top of the baseline selected by ``param.blocksize_mode``. Returns ------- - A streaming engine matching ``param``. - """ - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + The shared :class:`StreamingEngine`, ``_reset`` to the requested options. + Raises + ------ + RuntimeError + If ``engines`` has no slot for ``param.engine_name``. + """ streaming_options = create_streaming_options(param.blocksize_mode, options) - match param.engine_name: - case "spmd": - return SPMDEngine( - comm=spmd_comm, - rapidsmpf_options=streaming_options.to_rapidsmpf_options(), - executor_options=streaming_options.to_executor_options(), - engine_options=streaming_options.to_engine_options(), - ) - case _: # pragma: no cover - raise AssertionError(f"Unknown streaming backend: {param.engine_name!r}") + engine = engines.get(param.engine_name) + if engine is None: # pragma: no cover + raise RuntimeError( + f"No streaming engine for {param.engine_name!r}. The corresponding " + "session-scoped fixture must populate the collection before tests run." + ) + engine._reset( + rapidsmpf_options=streaming_options.to_rapidsmpf_options(), + executor_options=streaming_options.to_executor_options(), + engine_options=streaming_options.to_engine_options(), + ) + return engine def get_blocksize_mode(obj: pl.GPUEngine) -> Literal["medium", "small"]: diff --git a/python/cudf_polars/tests/conftest.py b/python/cudf_polars/tests/conftest.py index 7ad45c066057..7f00684638f0 100644 --- a/python/cudf_polars/tests/conftest.py +++ b/python/cudf_polars/tests/conftest.py @@ -18,12 +18,18 @@ ) if TYPE_CHECKING: - from collections.abc import Callable, Generator - - from rapidsmpf.communicator.communicator import Communicator + from collections.abc import Callable, Generator, Mapping + from typing import TypeAlias from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions + from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + + # Read-only view over the per-backend streaming engines owned by the + # ``streaming_engines`` session fixture. Only that fixture mutates the + # underlying dict; consumers (``spmd_engine``, ``streaming_engine_factory``, + # ``engine``) only look up by backend name. + StreamingEngines: TypeAlias = Mapping[str, StreamingEngine] @pytest.fixture(params=[False, True], ids=["no_nulls", "nulls"], scope="session") @@ -66,12 +72,12 @@ def _skip_unless_spmd(request: pytest.FixtureRequest) -> None: @pytest.fixture(scope="session") -def spmd_comm() -> Communicator: - """Session-scoped communicator — bootstrapped once and shared across all tests. +def streaming_engines() -> Generator[StreamingEngines, None, None]: + """Return a session-scoped mapping of engine name to engine instance. - Sharing a single communicator avoids the file-based bootstrap race that can - cause hangs when ``create_ucxx_comm()`` is called repeatedly in the same - ``rrun`` session (stale barrier files / stale ``ucxx_root_address`` KV entry). + The returned :class:`StreamingEngines` is a dict that maps each engine + name to a single shared engine instance, which is reused across the entire + test session. """ pytest.importorskip("rapidsmpf") from rapidsmpf import bootstrap @@ -79,12 +85,36 @@ def spmd_comm() -> Communicator: from rapidsmpf.config import Options, get_environment_variables from rapidsmpf.progress_thread import ProgressThread + from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + if bootstrap.is_running_with_rrun(): - return bootstrap.create_ucxx_comm( + comm = bootstrap.create_ucxx_comm( progress_thread=ProgressThread(), type=bootstrap.BackendType.AUTO, ) - return single_communicator(Options(get_environment_variables()), ProgressThread()) + else: + comm = single_communicator( + Options(get_environment_variables()), ProgressThread() + ) + + engines: dict[str, StreamingEngine] = {"spmd": SPMDEngine(comm=comm)} + try: + yield engines + finally: + while engines: + _, engine = engines.popitem() + engine.shutdown() + + +@pytest.fixture +def spmd_engine(streaming_engines: StreamingEngines) -> SPMDEngine: + """Return the shared :class:`SPMDEngine` reset to default options.""" + from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + + engine = streaming_engines["spmd"] + assert isinstance(engine, SPMDEngine) + engine._reset() + return engine @pytest.fixture(params=STREAMING_ENGINE_FIXTURE_PARAMS) @@ -102,38 +132,29 @@ def _all_engine_param(request: pytest.FixtureRequest) -> EngineFixtureParam: @pytest.fixture def streaming_engine_factory( _streaming_engine_param: EngineFixtureParam, - spmd_comm: Communicator, -) -> Generator[Callable[..., StreamingEngine], None, None]: + streaming_engines: StreamingEngines, +) -> Callable[..., StreamingEngine]: """ - Yield a factory that constructs :class:`StreamingEngine` instances for tests. - - The fixture is parametrized over :data:`STREAMING_ENGINE_FIXTURE_PARAMS`. - Created engines are tracked and automatically shut down after the test. + Return a factory that yields a shared :class:`StreamingEngine`. Parameters ---------- _streaming_engine_param Parametrized engine descriptor controlling backend and block size mode. - spmd_comm - Communicator used when constructing SPMD-based engines. - - Yields - ------ - Factory function that creates :class:`StreamingEngine` instances. The - factory accepts optional :class:`StreamingOptions`, which are merged on - top of the parametrized blocksize baseline. + streaming_engines + Session-scoped engine collection to look up the shared engine in. + + Returns + ------- + Factory function that returns the shared :class:`StreamingEngine`. """ - engines: list[StreamingEngine] = [] def factory(options: StreamingOptions | None = None) -> StreamingEngine: - engine = build_streaming_engine(_streaming_engine_param, spmd_comm, options) - engines.append(engine) - return engine - - yield factory + return build_streaming_engine( + _streaming_engine_param, streaming_engines, options + ) - for engine in reversed(engines): - engine.shutdown() + return factory @pytest.fixture @@ -164,9 +185,9 @@ def streaming_engine( def engine( request: pytest.FixtureRequest, _all_engine_param: EngineFixtureParam, -) -> Generator[pl.GPUEngine, None, None]: +) -> pl.GPUEngine: """ - Yield a :class:`polars.GPUEngine` for each engine variant under test. + Return a :class:`polars.GPUEngine` for each engine variant under test. Parameters ---------- @@ -176,8 +197,8 @@ def engine( Parametrized engine descriptor covering both in-memory and streaming variants. - Yields - ------ + Returns + ------- Engine instance matching the parametrized variant. Notes @@ -186,15 +207,10 @@ def engine( :func:`streaming_engine` fixture instead. """ if _all_engine_param.engine_name == "in-memory": - yield pl.GPUEngine(executor="in-memory", raise_on_fail=True) - return + return pl.GPUEngine(executor="in-memory", raise_on_fail=True) - spmd_comm: Communicator = request.getfixturevalue("spmd_comm") - engine = build_streaming_engine(_all_engine_param, spmd_comm) - try: - yield engine - finally: - engine.shutdown() + engines: StreamingEngines = request.getfixturevalue("streaming_engines") + return build_streaming_engine(_all_engine_param, engines) @pytest.fixture diff --git a/python/cudf_polars/tests/experimental/test_all_gather_host_data.py b/python/cudf_polars/tests/experimental/test_all_gather_host_data.py index aad7b3416769..8f09a82c4bd2 100644 --- a/python/cudf_polars/tests/experimental/test_all_gather_host_data.py +++ b/python/cudf_polars/tests/experimental/test_all_gather_host_data.py @@ -14,7 +14,6 @@ all_gather_host_data, ) from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine pytestmark = pytest.mark.spmd @@ -36,15 +35,14 @@ def _struct(rank: int) -> bytes: @pytest.mark.parametrize("make_data", [_empty, _text, _bytearray, _struct]) -def test_all_gather_host_data(spmd_comm, make_data) -> None: +def test_all_gather_host_data(spmd_engine, make_data) -> None: """Each rank sends rank-specific data; results are correct and ordered.""" - with SPMDEngine(comm=spmd_comm) as spmd_engine: - comm = spmd_engine.comm - br = spmd_engine.context.br() - result = all_gather_host_data(comm, br, op_id=0, data=make_data(comm.rank)) - assert len(result) == comm.nranks - for i, item in enumerate(result): - assert item == bytes(make_data(i)) + comm = spmd_engine.comm + br = spmd_engine.context.br() + result = all_gather_host_data(comm, br, op_id=0, data=make_data(comm.rank)) + assert len(result) == comm.nranks + for i, item in enumerate(result): + assert item == bytes(make_data(i)) def test_gather_cluster_info(streaming_engine) -> None: diff --git a/python/cudf_polars/tests/experimental/test_allgather.py b/python/cudf_polars/tests/experimental/test_allgather.py index 514276c66472..52c353044eb5 100644 --- a/python/cudf_polars/tests/experimental/test_allgather.py +++ b/python/cudf_polars/tests/experimental/test_allgather.py @@ -13,7 +13,6 @@ import pylibcudf as plc from cudf_polars.experimental.rapidsmpf.collectives.allgather import AllGatherManager -from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine from cudf_polars.experimental.rapidsmpf.utils import allgather_reduce @@ -53,9 +52,8 @@ async def _test_allgather(engine) -> None: assert col.type().id().value == plc.types.TypeId.INT32.value -def test_allgather(spmd_comm) -> None: - with SPMDEngine(comm=spmd_comm) as engine: - asyncio.run(_test_allgather(engine)) +def test_allgather(spmd_engine) -> None: + asyncio.run(_test_allgather(spmd_engine)) async def _test_allgather_reduce(engine) -> None: @@ -72,6 +70,5 @@ async def _test_allgather_reduce(engine) -> None: assert results == (10, 20, 30) # Single rank, so sums are just the local values -def test_allgather_reduce(spmd_comm) -> None: - with SPMDEngine(comm=spmd_comm) as engine: - asyncio.run(_test_allgather_reduce(engine)) +def test_allgather_reduce(spmd_engine) -> None: + asyncio.run(_test_allgather_reduce(spmd_engine)) diff --git a/python/cudf_polars/tests/experimental/test_dask.py b/python/cudf_polars/tests/experimental/test_dask.py index d923edd37cf9..5ccdde864ef6 100644 --- a/python/cudf_polars/tests/experimental/test_dask.py +++ b/python/cudf_polars/tests/experimental/test_dask.py @@ -153,3 +153,97 @@ def test_empty_dataframe(engine: DaskEngine) -> None: def test_run(engine: DaskEngine) -> None: result = engine._run(os.getpid) assert len(set(result)) == engine.nranks + + +@pytest.fixture(scope="module") +def reset_engine() -> Iterator[DaskEngine]: + """Module-scoped engine for reset tests — independent of ``engine``. + + These tests exercise :meth:`DaskEngine._reset` (which mutates the + engine in-place). A dedicated fixture keeps those mutations from + leaking into the other tests. + """ + with DaskEngine( + executor_options={"max_rows_per_partition": 10}, + ) as e: + yield e + + +def test_reset_keeps_workers_alive(reset_engine: DaskEngine) -> None: + """``_reset`` must not respawn dask workers.""" + workers_before = sorted( + reset_engine._dask_ctx.client.scheduler_info(n_workers=-1)["workers"] + ) + pids_before = sorted(reset_engine._run(os.getpid)) + + reset_engine._reset(executor_options={"max_rows_per_partition": 7}) + + workers_after = sorted( + reset_engine._dask_ctx.client.scheduler_info(n_workers=-1)["workers"] + ) + pids_after = sorted(reset_engine._run(os.getpid)) + + # Same worker addresses … + assert workers_before == workers_after + # … and the workers are running in the same OS processes. + assert pids_before == pids_after + + +def test_reset_updates_executor_options(reset_engine: DaskEngine) -> None: + """``_reset`` updates the polars-layer config to the new options.""" + reset_engine._reset(executor_options={"max_rows_per_partition": 42}) + + opts = reset_engine.config["executor_options"] + assert opts["max_rows_per_partition"] == 42 + # Reserved keys are still injected by ``_reset``. + assert opts["runtime"] == "rapidsmpf" + assert opts["cluster"] == "dask" + assert isinstance(opts["dask_context"], DaskContext) + + +def test_reset_collects_after_options_change(reset_engine: DaskEngine) -> None: + """The engine still drives a real query after ``_reset``.""" + reset_engine._reset(executor_options={"max_rows_per_partition": 3}) + assert_gpu_result_equal( + pl.LazyFrame({"a": [1, 2, 3, 4, 5]}), + engine=reset_engine, + check_row_order=False, + ) + + +def test_reset_after_shutdown_raises() -> None: + """``shutdown`` is idempotent; ``_reset`` after shutdown raises every time.""" + engine = DaskEngine(executor_options={"max_rows_per_partition": 10}) + engine.shutdown() + engine.shutdown() # idempotent + with pytest.raises(RuntimeError, match="shut-down"): + engine._reset() + with pytest.raises(RuntimeError, match="shut-down"): + engine._reset() # still raises on a second attempt + engine.shutdown() # still safe after a failed _reset + + +def test_reset_rejects_construction_time_executor_options( + reset_engine: DaskEngine, +) -> None: + """``_reset`` rejects ``executor_options`` keys read at worker setup.""" + with pytest.raises(ValueError, match="num_py_executors"): + reset_engine._reset(executor_options={"num_py_executors": 4}) + + +def test_reset_rejects_construction_time_engine_options( + reset_engine: DaskEngine, +) -> None: + """``_reset`` rejects ``engine_options`` keys read at worker setup.""" + from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( + HardwareBindingPolicy, + ) + + with pytest.raises(ValueError, match="hardware_binding"): + reset_engine._reset( + engine_options={ + "hardware_binding": HardwareBindingPolicy(enabled=False), + }, + ) + with pytest.raises(ValueError, match="memory_resource_config"): + reset_engine._reset(engine_options={"memory_resource_config": None}) diff --git a/python/cudf_polars/tests/experimental/test_io_multirank.py b/python/cudf_polars/tests/experimental/test_io_multirank.py index 631f12fd85c6..2208cc673169 100644 --- a/python/cudf_polars/tests/experimental/test_io_multirank.py +++ b/python/cudf_polars/tests/experimental/test_io_multirank.py @@ -19,8 +19,6 @@ from collections.abc import Iterator from pathlib import Path - from rapidsmpf.communicator.communicator import Communicator - from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine # Runs the spmd variant even under rrun with nranks > 1. The ray/dask @@ -44,7 +42,7 @@ def df() -> pl.LazyFrame: @pytest.fixture(params=["spmd", "ray", "dask"]) def engine( request: pytest.FixtureRequest, - spmd_comm: Communicator, + spmd_engine: SPMDEngine, ) -> Iterator[StreamingEngine]: """Yield each supported streaming engine.""" backend = request.param @@ -52,7 +50,7 @@ def engine( if backend == "spmd": with SPMDEngine( - comm=spmd_comm, + comm=spmd_engine.comm, executor_options=executor_options, ) as eng: yield eng diff --git a/python/cudf_polars/tests/experimental/test_ray.py b/python/cudf_polars/tests/experimental/test_ray.py index 7365be733b3c..ded4903c5940 100644 --- a/python/cudf_polars/tests/experimental/test_ray.py +++ b/python/cudf_polars/tests/experimental/test_ray.py @@ -275,7 +275,7 @@ def test_reset_collects_after_options_change(reset_engine: RayEngine) -> None: def test_reset_after_shutdown_raises() -> None: - """``_reset`` after ``shutdown`` raises ``RuntimeError``.""" + """``shutdown`` is idempotent; ``_reset`` after shutdown raises every time.""" engine = RayEngine( executor_options={"max_rows_per_partition": 10}, engine_options={"allow_gpu_sharing": True}, @@ -283,8 +283,12 @@ def test_reset_after_shutdown_raises() -> None: ray_init_options={"include_dashboard": False}, ) engine.shutdown() + engine.shutdown() # idempotent with pytest.raises(RuntimeError, match="shut-down"): engine._reset() + with pytest.raises(RuntimeError, match="shut-down"): + engine._reset() # still raises on a second attempt + engine.shutdown() # still safe after a failed _reset def test_reset_rejects_construction_time_executor_options( diff --git a/python/cudf_polars/tests/experimental/test_sink.py b/python/cudf_polars/tests/experimental/test_sink.py index 9b0573d2cb44..df68b7c199ae 100644 --- a/python/cudf_polars/tests/experimental/test_sink.py +++ b/python/cudf_polars/tests/experimental/test_sink.py @@ -92,7 +92,7 @@ def test_sink_parquet_directory( assert len(list(check_path.iterdir())) == expected_file_count -def test_sink_parquet_raises_spmd(spmd_comm): +def test_sink_parquet_raises_spmd(spmd_engine): from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine with ( @@ -100,7 +100,7 @@ def test_sink_parquet_raises_spmd(spmd_comm): ValueError, match="The spmd cluster requires sink_to_directory=True" ), SPMDEngine( - comm=spmd_comm, executor_options={"sink_to_directory": False} + comm=spmd_engine.comm, executor_options={"sink_to_directory": False} ) as engine, ): ConfigOptions.from_polars_engine(engine) diff --git a/python/cudf_polars/tests/experimental/test_spilling.py b/python/cudf_polars/tests/experimental/test_spilling.py index 799d19402e6a..6aa118011320 100644 --- a/python/cudf_polars/tests/experimental/test_spilling.py +++ b/python/cudf_polars/tests/experimental/test_spilling.py @@ -9,7 +9,6 @@ import numpy as np import pytest -from rapidsmpf.config import Options from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.memory.pinned_memory_resource import is_pinned_memory_resources_supported from rapidsmpf.streaming.core.message import Message @@ -18,7 +17,7 @@ import pylibcudf as plc -from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine +from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions from cudf_polars.experimental.rapidsmpf.utils import ( make_spill_function, ) @@ -51,109 +50,104 @@ def create_test_table(nbytes: int, stream: Stream) -> plc.Table: ], ) def test_make_spill_function( - spmd_comm, + streaming_engine_factory, *, pinned_memory: bool, spilled_host_mem_type: MemoryType, ) -> None: """Test that spilling prioritizes longest queues and newest messages.""" - with SPMDEngine( - comm=spmd_comm, - rapidsmpf_options=Options({"pinned_memory": str(pinned_memory).lower()}), - ) as spmd_engine: - context = spmd_engine.context - - if spilled_host_mem_type == MemoryType.PINNED_HOST: - assert spmd_engine.context.br().pinned_mr is not None - other_host_mem_type = MemoryType.HOST - else: - assert spmd_engine.context.br().pinned_mr is None - other_host_mem_type = MemoryType.PINNED_HOST - - # Create 3 spillable message containers simulating fanout buffers - # Buffer 0: Fast consumer (2 messages) - # Buffer 1: Slow consumer (5 messages) <- should spill from here first - # Buffer 2: Medium consumer (3 messages) - buffers = [SpillableMessages(context.br()) for _ in range(3)] - messages_per_buffer = [2, 5, 3] - - # Track message IDs for each buffer - message_ids: dict[int, list[int]] = {} - - # Populate buffers with messages - stream = context.get_stream_from_pool() - for buffer_idx, (sm, count) in enumerate( - zip(buffers, messages_per_buffer, strict=False) - ): - message_ids[buffer_idx] = [] - for msg_idx in range(count): - # Create 1MB messages - table = create_test_table(1024 * 1024, stream) - chunk = TableChunk.from_pylibcudf_table( - table, stream, exclusive_view=True, br=context.br() - ) - msg = Message(msg_idx, chunk) - mid = sm.insert(msg) - message_ids[buffer_idx].append(mid) - - # Register spill function - spill_func = make_spill_function(buffers, context) - func_id = context.br().spill_manager.add_spill_function(spill_func, priority=0) - - try: - # Manually trigger spilling of 3MB - # Expected: Buffer 1 (longest) should spill newest messages first - amount_to_spill = 3 * 1024 * 1024 - actual_spilled = context.br().spill_manager.spill(amount_to_spill) - - # Allow some tolerance - assert actual_spilled >= amount_to_spill * 0.95 - - # Verify Buffer 1 (longest queue): newest 3 messages should be spilled - buffer_1_descs = buffers[1].get_content_descriptions() - for i in range(3, 5): # Messages 3, 4 (newest) - mid = message_ids[1][i] - desc = buffer_1_descs[mid] - # Should be in HOST memory (spilled) - assert desc.content_sizes[spilled_host_mem_type] > 0 - assert desc.content_sizes[other_host_mem_type] == 0 - assert desc.content_sizes[MemoryType.DEVICE] == 0 - - # Buffer 1: oldest messages should still be in device - for i in range(2): # Messages 0, 1 (oldest) - mid = message_ids[1][i] - desc = buffer_1_descs[mid] - # Should still be in DEVICE memory - assert desc.content_sizes[MemoryType.DEVICE] > 0 - assert desc.content_sizes[spilled_host_mem_type] == 0 - assert desc.content_sizes[other_host_mem_type] == 0 - - # Buffer 0 (shortest queue): all messages should still be on device - buffer_0_descs = buffers[0].get_content_descriptions() - for mid in message_ids[0]: - desc = buffer_0_descs[mid] - assert desc.content_sizes[MemoryType.DEVICE] > 0 - assert desc.content_sizes[spilled_host_mem_type] == 0 - assert desc.content_sizes[other_host_mem_type] == 0 - - # Verify we can extract and make available a spilled message - spilled_mid = message_ids[1][4] # Newest message from longest queue - spilled_msg = buffers[1].extract(mid=spilled_mid) - - chunk = TableChunk.from_message(spilled_msg, br=context.br()) - assert not chunk.is_available() # Should be on host - - # Make it available should bring it back to device - cost = chunk.make_available_cost() - assert cost > 0 - res, _ = context.br().reserve( - MemoryType.DEVICE, cost, allow_overbooking=True + engine = streaming_engine_factory(StreamingOptions(pinned_memory=pinned_memory)) + context = engine.context + + if spilled_host_mem_type == MemoryType.PINNED_HOST: + assert engine.context.br().pinned_mr is not None + other_host_mem_type = MemoryType.HOST + else: + assert engine.context.br().pinned_mr is None + other_host_mem_type = MemoryType.PINNED_HOST + + # Create 3 spillable message containers simulating fanout buffers + # Buffer 0: Fast consumer (2 messages) + # Buffer 1: Slow consumer (5 messages) <- should spill from here first + # Buffer 2: Medium consumer (3 messages) + buffers = [SpillableMessages(context.br()) for _ in range(3)] + messages_per_buffer = [2, 5, 3] + + # Track message IDs for each buffer + message_ids: dict[int, list[int]] = {} + + # Populate buffers with messages + stream = context.get_stream_from_pool() + for buffer_idx, (sm, count) in enumerate( + zip(buffers, messages_per_buffer, strict=False) + ): + message_ids[buffer_idx] = [] + for msg_idx in range(count): + # Create 1MB messages + table = create_test_table(1024 * 1024, stream) + chunk = TableChunk.from_pylibcudf_table( + table, stream, exclusive_view=True, br=context.br() ) - chunk_available = chunk.make_available(res) - - assert chunk_available.is_available() - # Verify we got a valid table back - assert chunk_available.table_view().num_rows() > 0 - - finally: - context.br().spill_manager.remove_spill_function(func_id) + msg = Message(msg_idx, chunk) + mid = sm.insert(msg) + message_ids[buffer_idx].append(mid) + + # Register spill function + spill_func = make_spill_function(buffers, context) + func_id = context.br().spill_manager.add_spill_function(spill_func, priority=0) + + try: + # Manually trigger spilling of 3MB + # Expected: Buffer 1 (longest) should spill newest messages first + amount_to_spill = 3 * 1024 * 1024 + actual_spilled = context.br().spill_manager.spill(amount_to_spill) + + # Allow some tolerance + assert actual_spilled >= amount_to_spill * 0.95 + + # Verify Buffer 1 (longest queue): newest 3 messages should be spilled + buffer_1_descs = buffers[1].get_content_descriptions() + for i in range(3, 5): # Messages 3, 4 (newest) + mid = message_ids[1][i] + desc = buffer_1_descs[mid] + # Should be in HOST memory (spilled) + assert desc.content_sizes[spilled_host_mem_type] > 0 + assert desc.content_sizes[other_host_mem_type] == 0 + assert desc.content_sizes[MemoryType.DEVICE] == 0 + + # Buffer 1: oldest messages should still be in device + for i in range(2): # Messages 0, 1 (oldest) + mid = message_ids[1][i] + desc = buffer_1_descs[mid] + # Should still be in DEVICE memory + assert desc.content_sizes[MemoryType.DEVICE] > 0 + assert desc.content_sizes[spilled_host_mem_type] == 0 + assert desc.content_sizes[other_host_mem_type] == 0 + + # Buffer 0 (shortest queue): all messages should still be on device + buffer_0_descs = buffers[0].get_content_descriptions() + for mid in message_ids[0]: + desc = buffer_0_descs[mid] + assert desc.content_sizes[MemoryType.DEVICE] > 0 + assert desc.content_sizes[spilled_host_mem_type] == 0 + assert desc.content_sizes[other_host_mem_type] == 0 + + # Verify we can extract and make available a spilled message + spilled_mid = message_ids[1][4] # Newest message from longest queue + spilled_msg = buffers[1].extract(mid=spilled_mid) + + chunk = TableChunk.from_message(spilled_msg, br=context.br()) + assert not chunk.is_available() # Should be on host + + # Make it available should bring it back to device + cost = chunk.make_available_cost() + assert cost > 0 + res, _ = context.br().reserve(MemoryType.DEVICE, cost, allow_overbooking=True) + chunk_available = chunk.make_available(res) + + assert chunk_available.is_available() + # Verify we got a valid table back + assert chunk_available.table_view().num_rows() > 0 + + finally: + context.br().spill_manager.remove_spill_function(func_id) diff --git a/python/cudf_polars/tests/experimental/test_spmd.py b/python/cudf_polars/tests/experimental/test_spmd.py index a1970c8e92f6..9fef0e003504 100644 --- a/python/cudf_polars/tests/experimental/test_spmd.py +++ b/python/cudf_polars/tests/experimental/test_spmd.py @@ -30,12 +30,22 @@ pytestmark = pytest.mark.spmd -def test_yields_context_and_engine(spmd_comm: Communicator) -> None: +@pytest.fixture +def comm(spmd_engine: SPMDEngine) -> Communicator: + """Communicator from the shared :class:`SPMDEngine` for local construction. + + Most tests in this module need to construct their own + :class:`SPMDEngine` to exercise lifecycle, construction-time + options, MR-state semantics, or :meth:`SPMDEngine._reset`. + """ + return spmd_engine.comm + + +def test_yields_context_and_engine(spmd_engine: SPMDEngine) -> None: """SPMDEngine has comm and context properties.""" - with SPMDEngine(comm=spmd_comm) as engine: - assert engine.comm is not None - assert engine.context is not None - assert isinstance(engine, pl.GPUEngine) + assert spmd_engine.comm is not None + assert spmd_engine.context is not None + assert isinstance(spmd_engine, pl.GPUEngine) def test_from_options() -> None: @@ -74,31 +84,29 @@ def test_engine_options_reserved_keys() -> None: pass -def test_engine_options_parquet_options(spmd_comm: Communicator) -> None: +def test_engine_options_parquet_options(comm: Communicator) -> None: """engine_options forwards parquet_options to GPUEngine without error.""" - with SPMDEngine(comm=spmd_comm, engine_options={"parquet_options": {}}) as engine: + with SPMDEngine(comm=comm, engine_options={"parquet_options": {}}) as engine: assert isinstance(engine, pl.GPUEngine) -def test_scan(spmd_comm: Communicator) -> None: +def test_scan(spmd_engine: SPMDEngine) -> None: """Each rank scans its own single-row LazyFrame and gets that row back.""" - with SPMDEngine(comm=spmd_comm) as engine: - lf = pl.LazyFrame({"a": [engine.rank], "b": [engine.rank * 10]}) - result = lf.collect(engine=engine) - assert result.shape == (1, 2) - assert result["a"].to_list() == [engine.rank] - assert result["b"].to_list() == [engine.rank * 10] + lf = pl.LazyFrame({"a": [spmd_engine.rank], "b": [spmd_engine.rank * 10]}) + result = lf.collect(engine=spmd_engine) + assert result.shape == (1, 2) + assert result["a"].to_list() == [spmd_engine.rank] + assert result["b"].to_list() == [spmd_engine.rank * 10] -def test_basic_query(spmd_comm: Communicator) -> None: +def test_basic_query(spmd_engine: SPMDEngine) -> None: """A simple in-memory LazyFrame can be collected.""" - with SPMDEngine(comm=spmd_comm) as engine: - result = pl.LazyFrame({"a": [1, 2, 3], "b": [4, 5, 6]}).collect(engine=engine) + result = pl.LazyFrame({"a": [1, 2, 3], "b": [4, 5, 6]}).collect(engine=spmd_engine) assert result.shape == (3, 2) assert result["a"].to_list() == [1, 2, 3] -def test_collect_then_lazy_equivalent(spmd_comm: Communicator) -> None: +def test_collect_then_lazy_equivalent(spmd_engine: SPMDEngine) -> None: """collect().lazy() preserves SPMD semantics: an intermediate materialize is a no-op. In SPMD mode a DataFrame is always rank-local. When it is wrapped back @@ -106,111 +114,105 @@ def test_collect_then_lazy_equivalent(spmd_comm: Communicator) -> None: re-slicing it across ranks. So ``lf.collect().lazy().op.collect()`` must produce the same result as ``lf.op.collect()``. """ - with SPMDEngine(comm=spmd_comm) as engine: - lf = pl.LazyFrame( - {"a": [engine.rank, engine.rank + 1, engine.rank + 2], "b": [0, 1, 2]} - ) + rank = spmd_engine.rank + lf = pl.LazyFrame({"a": [rank, rank + 1, rank + 2], "b": [0, 1, 2]}) - # One-step - one_step = lf.filter(pl.col("b") >= 1).collect(engine=engine) + # One-step + one_step = lf.filter(pl.col("b") >= 1).collect(engine=spmd_engine) - # Two-step: materialize then re-wrap - intermediate = lf.collect(engine=engine) - two_step = intermediate.lazy().filter(pl.col("b") >= 1).collect(engine=engine) + # Two-step: materialize then re-wrap + intermediate = lf.collect(engine=spmd_engine) + two_step = intermediate.lazy().filter(pl.col("b") >= 1).collect(engine=spmd_engine) assert one_step.sort("a").equals(two_step.sort("a")) -def test_group_by(spmd_comm: Communicator) -> None: +def test_group_by(spmd_engine: SPMDEngine) -> None: """Group-by on rank-local data, then allgather to verify the global result.""" - with SPMDEngine(comm=spmd_comm) as engine: - lf = pl.LazyFrame({"a": [engine.rank], "b": [engine.rank * 10]}) - local_result = lf.group_by("a").agg(pl.col("b").sum()).collect(engine=engine) - with reserve_op_id() as op_id: - global_result = allgather_polars_dataframe( - engine=engine, local_df=local_result, op_id=op_id - ) - assert global_result.shape == (engine.nranks, 2) - assert global_result.sort("a")["a"].to_list() == list(range(engine.nranks)) - assert global_result.sort("a")["b"].to_list() == [ - r * 10 for r in range(engine.nranks) - ] + lf = pl.LazyFrame({"a": [spmd_engine.rank], "b": [spmd_engine.rank * 10]}) + local_result = lf.group_by("a").agg(pl.col("b").sum()).collect(engine=spmd_engine) + with reserve_op_id() as op_id: + global_result = allgather_polars_dataframe( + engine=spmd_engine, local_df=local_result, op_id=op_id + ) + assert global_result.shape == (spmd_engine.nranks, 2) + assert global_result.sort("a")["a"].to_list() == list(range(spmd_engine.nranks)) + assert global_result.sort("a")["b"].to_list() == [ + r * 10 for r in range(spmd_engine.nranks) + ] -def test_allgather_polars_dataframe(spmd_comm: Communicator) -> None: +def test_allgather_polars_dataframe(spmd_engine: SPMDEngine) -> None: """allgather_polars_dataframe collects every rank's contribution in rank order.""" - with SPMDEngine(comm=spmd_comm) as engine: - local = pl.DataFrame({"rank": [engine.rank], "val": [engine.rank * 2]}) - with reserve_op_id() as op_id: - result = allgather_polars_dataframe( - engine=engine, local_df=local, op_id=op_id - ) - assert result.shape == (engine.nranks, 2) - assert result["rank"].to_list() == list(range(engine.nranks)) - assert result["val"].to_list() == [r * 2 for r in range(engine.nranks)] + local = pl.DataFrame({"rank": [spmd_engine.rank], "val": [spmd_engine.rank * 2]}) + with reserve_op_id() as op_id: + result = allgather_polars_dataframe( + engine=spmd_engine, local_df=local, op_id=op_id + ) + assert result.shape == (spmd_engine.nranks, 2) + assert result["rank"].to_list() == list(range(spmd_engine.nranks)) + assert result["val"].to_list() == [r * 2 for r in range(spmd_engine.nranks)] -def test_num_py_executors(spmd_comm: Communicator) -> None: +def test_num_py_executors(comm: Communicator) -> None: """executor_options forwards num_py_executors to the thread pool.""" with SPMDEngine( - comm=spmd_comm, + comm=comm, executor_options={"num_py_executors": 2}, ) as engine: result = pl.LazyFrame({"a": [1, 2, 3]}).collect(engine=engine) assert result.shape == (3, 1) -def test_allgather_polars_dataframe_empty(spmd_comm: Communicator) -> None: +def test_allgather_polars_dataframe_empty(spmd_engine: SPMDEngine) -> None: """allgather handles an empty (zero-row) local DataFrame on every rank.""" - with SPMDEngine(comm=spmd_comm) as engine: - local = pl.DataFrame( - {"a": pl.Series([], dtype=pl.Int32), "b": pl.Series([], dtype=pl.Float64)} + local = pl.DataFrame( + {"a": pl.Series([], dtype=pl.Int32), "b": pl.Series([], dtype=pl.Float64)} + ) + with reserve_op_id() as op_id: + result = allgather_polars_dataframe( + engine=spmd_engine, local_df=local, op_id=op_id ) - with reserve_op_id() as op_id: - result = allgather_polars_dataframe( - engine=engine, local_df=local, op_id=op_id - ) assert result.shape == (0, 2) assert result.columns == ["a", "b"] assert result.dtypes == [pl.Int32, pl.Float64] -def test_mr_wrapped_as_current_inside_context(spmd_comm: Communicator) -> None: +def test_mr_wrapped_as_current_inside_context(comm: Communicator) -> None: """Inside SPMDEngine the current device resource is RmmResourceAdaptor.""" - with SPMDEngine(comm=spmd_comm): + with SPMDEngine(comm=comm): assert isinstance(rmm.mr.get_current_device_resource(), RmmResourceAdaptor) -def test_mr_restored_after_context(spmd_comm: Communicator) -> None: +def test_mr_restored_after_context(comm: Communicator) -> None: """After SPMDEngine exits the original device resource is restored.""" original = rmm.mr.get_current_device_resource() - with SPMDEngine(comm=spmd_comm): + with SPMDEngine(comm=comm): pass assert rmm.mr.get_current_device_resource() is original -def test_allgather_polars_dataframe_multi_column(spmd_comm: Communicator) -> None: +def test_allgather_polars_dataframe_multi_column(spmd_engine: SPMDEngine) -> None: """allgather preserves column names, count, and dtypes for multi-column DataFrames.""" - with SPMDEngine(comm=spmd_comm) as engine: - local = pl.DataFrame( - { - "rank": [engine.rank], - "x": [float(engine.rank)], - "label": [f"r{engine.rank}"], - } + local = pl.DataFrame( + { + "rank": [spmd_engine.rank], + "x": [float(spmd_engine.rank)], + "label": [f"r{spmd_engine.rank}"], + } + ) + with reserve_op_id() as op_id: + result = allgather_polars_dataframe( + engine=spmd_engine, local_df=local, op_id=op_id ) - with reserve_op_id() as op_id: - result = allgather_polars_dataframe( - engine=engine, local_df=local, op_id=op_id - ) - assert result.shape == (engine.nranks, 3) - assert result.columns == ["rank", "x", "label"] - sorted_result = result.sort("rank") - assert sorted_result["rank"].to_list() == list(range(engine.nranks)) - assert sorted_result["x"].to_list() == [float(r) for r in range(engine.nranks)] - assert sorted_result["label"].to_list() == [ - f"r{r}" for r in range(engine.nranks) - ] + assert result.shape == (spmd_engine.nranks, 3) + assert result.columns == ["rank", "x", "label"] + sorted_result = result.sort("rank") + assert sorted_result["rank"].to_list() == list(range(spmd_engine.nranks)) + assert sorted_result["x"].to_list() == [float(r) for r in range(spmd_engine.nranks)] + assert sorted_result["label"].to_list() == [ + f"r{r}" for r in range(spmd_engine.nranks) + ] # --------------------------------------------------------------------------- @@ -218,44 +220,44 @@ def test_allgather_polars_dataframe_multi_column(spmd_comm: Communicator) -> Non # --------------------------------------------------------------------------- -def test_comm_argument_reuses_communicator(spmd_comm: Communicator) -> None: +def test_comm_argument_reuses_communicator(comm: Communicator) -> None: """Passing comm= reuses the communicator across two engine lifetimes.""" - with SPMDEngine(comm=spmd_comm) as engine1: + with SPMDEngine(comm=comm) as engine1: nranks = engine1.nranks rank = engine1.rank - # engine1 is shut down; spmd_comm is still alive - with SPMDEngine(comm=spmd_comm) as engine2: + # engine1 is shut down; the shared comm is still alive + with SPMDEngine(comm=comm) as engine2: assert engine2.nranks == nranks assert engine2.rank == rank -def test_comm_not_closed_after_engine_shutdown(spmd_comm: Communicator) -> None: +def test_comm_not_closed_after_engine_shutdown(comm: Communicator) -> None: """The caller-provided comm survives engine.shutdown().""" - with SPMDEngine(comm=spmd_comm): + with SPMDEngine(comm=comm): pass # engine.shutdown() is called on __exit__ - # spmd_comm must still be accessible — not destroyed by engine teardown - assert spmd_comm.rank >= 0 + # comm must still be accessible — not destroyed by engine teardown + assert comm.rank >= 0 -def test_comm_argument_mr_still_wrapped(spmd_comm: Communicator) -> None: +def test_comm_argument_mr_still_wrapped(comm: Communicator) -> None: """MR wrapping still happens even when comm is provided externally.""" - with SPMDEngine(comm=spmd_comm): + with SPMDEngine(comm=comm): assert isinstance(rmm.mr.get_current_device_resource(), RmmResourceAdaptor) -def test_comm_sequential_queries(spmd_comm: Communicator) -> None: +def test_comm_sequential_queries(comm: Communicator) -> None: """Two engines sharing a comm can each execute a query without interference.""" - with SPMDEngine(comm=spmd_comm) as engine: + with SPMDEngine(comm=comm) as engine: r1 = pl.LazyFrame({"a": [1, 2]}).collect(engine=engine) - with SPMDEngine(comm=spmd_comm) as engine: + with SPMDEngine(comm=comm) as engine: r2 = pl.LazyFrame({"a": [3, 4]}).collect(engine=engine) assert r1["a"].to_list() == [1, 2] assert r2["a"].to_list() == [3, 4] -def test_shutdown_idempotent(spmd_comm: Communicator) -> None: +def test_shutdown_idempotent(comm: Communicator) -> None: """Calling shutdown() twice does not raise.""" - engine = SPMDEngine(comm=spmd_comm) + engine = SPMDEngine(comm=comm) engine.shutdown() engine.shutdown() @@ -277,9 +279,9 @@ def test_memory_resource_config() -> None: mock_create.assert_called_once() -def test_comm_and_context_unavailable_after_shutdown(spmd_comm: Communicator) -> None: +def test_comm_and_context_unavailable_after_shutdown(comm: Communicator) -> None: """Accessing comm or context after shutdown raises RuntimeError.""" - engine = SPMDEngine(comm=spmd_comm) + engine = SPMDEngine(comm=comm) engine.shutdown() with pytest.raises(RuntimeError, match="shutdown"): _ = engine.comm @@ -287,8 +289,89 @@ def test_comm_and_context_unavailable_after_shutdown(spmd_comm: Communicator) -> _ = engine.context -def test_run(spmd_comm): - with SPMDEngine(comm=spmd_comm) as engine: - result = engine._run(os.getpid) - +def test_run(spmd_engine: SPMDEngine) -> None: + result = spmd_engine._run(os.getpid) assert result == [os.getpid()] + + +def test_reset_keeps_comm_alive(comm: Communicator) -> None: + """``_reset`` must not rebuild the communicator.""" + with SPMDEngine( + comm=comm, executor_options={"max_rows_per_partition": 10} + ) as engine: + comm_before = engine.comm + engine._reset(executor_options={"max_rows_per_partition": 7}) + # Same Communicator instance — caller-provided comm is preserved. + assert engine.comm is comm_before + # Engine still drives a real query. + result = pl.LazyFrame({"a": [1, 2, 3]}).collect(engine=engine) + assert sorted(result["a"].to_list()) == [1, 2, 3] + + +def test_reset_updates_executor_options(comm: Communicator) -> None: + """``_reset`` updates the polars-layer config to the new options.""" + from cudf_polars.utils.config import SPMDContext + + with SPMDEngine( + comm=comm, executor_options={"max_rows_per_partition": 10} + ) as engine: + engine._reset(executor_options={"max_rows_per_partition": 42}) + + opts = engine.config["executor_options"] + assert opts["max_rows_per_partition"] == 42 + # Reserved keys are still injected by ``_reset``. + assert opts["runtime"] == "rapidsmpf" + assert opts["cluster"] == "spmd" + assert isinstance(opts["spmd_context"], SPMDContext) + + +def test_reset_collects_after_options_change(comm: Communicator) -> None: + """The engine still drives a real query after ``_reset``.""" + with SPMDEngine( + comm=comm, executor_options={"max_rows_per_partition": 10} + ) as engine: + engine._reset(executor_options={"max_rows_per_partition": 3}) + result = pl.LazyFrame({"a": [1, 2, 3, 4, 5]}).collect(engine=engine) + assert sorted(result["a"].to_list()) == [1, 2, 3, 4, 5] + + +def test_reset_after_shutdown_raises(comm: Communicator) -> None: + """``shutdown`` is idempotent; ``_reset`` after shutdown raises every time.""" + engine = SPMDEngine(comm=comm) + engine.shutdown() + engine.shutdown() # idempotent + with pytest.raises(RuntimeError, match="shut-down"): + engine._reset() + with pytest.raises(RuntimeError, match="shut-down"): + engine._reset() # still raises on a second attempt + engine.shutdown() # still safe after a failed _reset + + +def test_reset_rejects_construction_time_executor_options( + comm: Communicator, +) -> None: + """``_reset`` rejects ``executor_options`` keys read at engine construction.""" + with ( + SPMDEngine(comm=comm) as engine, + pytest.raises(ValueError, match="num_py_executors"), + ): + engine._reset(executor_options={"num_py_executors": 4}) + + +def test_reset_rejects_construction_time_engine_options( + comm: Communicator, +) -> None: + """``_reset`` rejects ``engine_options`` keys read at engine construction.""" + from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( + HardwareBindingPolicy, + ) + + with SPMDEngine(comm=comm) as engine: + with pytest.raises(ValueError, match="hardware_binding"): + engine._reset( + engine_options={ + "hardware_binding": HardwareBindingPolicy(enabled=False), + }, + ) + with pytest.raises(ValueError, match="memory_resource_config"): + engine._reset(engine_options={"memory_resource_config": None}) diff --git a/python/cudf_polars/tests/experimental/test_statistics.py b/python/cudf_polars/tests/experimental/test_statistics.py index 965449b80f04..82c121d5830e 100644 --- a/python/cudf_polars/tests/experimental/test_statistics.py +++ b/python/cudf_polars/tests/experimental/test_statistics.py @@ -16,8 +16,6 @@ if TYPE_CHECKING: from collections.abc import Iterator - from rapidsmpf.communicator.communicator import Communicator - from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine # Runs the spmd variant even under rrun with nranks > 1. The ray/dask @@ -30,7 +28,7 @@ @pytest.fixture(params=["spmd", "ray", "dask"]) def engine( request: pytest.FixtureRequest, - spmd_comm: Communicator, + spmd_engine: SPMDEngine, ) -> Iterator[StreamingEngine]: """Yield each supported streaming engine with statistics enabled.""" backend = request.param @@ -39,7 +37,7 @@ def engine( if backend == "spmd": with SPMDEngine( - comm=spmd_comm, + comm=spmd_engine.comm, rapidsmpf_options=rapidsmpf_options, executor_options=executor_options, ) as engine: