[Multi-GPU Polars] Use current rmm resource in SPMD mode (rapidsai#21842)

Fix `spmd_execution()` MR lifecycle: wrap, set as current, restore

`spmd_execution()` wrapped the current device resource in `RmmResourceAdaptor`, but never installed it as the _current_ resource. As a result, libcudf temporary allocations bypassed the adaptor.

* Adds a `set_memory_resource(mr)` context manager in `rapidsmpf/utils.py` that saves, sets, and restores the current device resource.
* Uses it in `spmd_execution()` so both the RapidsMPF `Context` and libcudf share the same `RmmResourceAdaptor`.

Authors:
  - Mads R. B. Kristensen (https://github.com/madsbk)

Approvers:
  - Tom Augspurger (https://github.com/TomAugspurger)

URL: rapidsai#21842

Author: madsbk

python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py

@@ -1864,19 +1864,23 @@ def run_polars_spmd(
     """Run benchmark queries using SPMD execution via the ``rrun`` launcher."""
     if run_config.collect_traces:
         raise NotImplementedError(
-            "--collect-traces is not yet supported with --cluster spmd."
+            "--collect-traces is not yet supported with --cluster spmd"
         )
+    if run_config.rmm_async:
+        raise NotImplementedError("--rmm-async is not supported with --cluster spmd")
     executor_options = get_executor_options(run_config, benchmark=benchmark)
     # "runtime" and "cluster" are reserved — spmd_execution sets them
     executor_options.pop("runtime", None)
     executor_options.pop("cluster", None)
+    rmm.mr.set_current_device_resource(
+        rmm.mr.CudaAsyncMemoryResource(release_threshold=args.rmm_release_threshold)
+    )
     with spmd_execution(
-        mr=rmm.mr.CudaAsyncMemoryResource(release_threshold=args.rmm_release_threshold)
-        if run_config.rmm_async
-        else None,
         executor_options=executor_options,
-        parquet_options=parquet_options,
-        cuda_stream_policy=run_config.stream_policy,
+        engine_options={
+            "parquet_options": parquet_options,
+            "cuda_stream_policy": run_config.stream_policy,
+        },
     ) as (comm, ctx, engine):
         from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id
         from cudf_polars.experimental.rapidsmpf.frontend.spmd import (

python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py

@@ -28,7 +28,10 @@
 from cudf_polars.containers import DataFrame
 from cudf_polars.dsl.ir import IRExecutionContext
 from cudf_polars.experimental.rapidsmpf.core import generate_network
-from cudf_polars.experimental.rapidsmpf.utils import empty_table_chunk
+from cudf_polars.experimental.rapidsmpf.utils import (
+    empty_table_chunk,
+    set_memory_resource,
+)
 from cudf_polars.experimental.utils import _concat
 from cudf_polars.utils.config import SPMDContext
 
@@ -87,11 +90,11 @@ def evaluate_pipeline_spmd_mode(
     """
     if config_options.executor.runtime != "rapidsmpf":
         raise RuntimeError("Runtime must be rapidsmpf")
-    if config_options.executor.spmd is None:
-        raise RuntimeError("spmd must be set for SPMD mode")
-    comm = config_options.executor.spmd.comm
-    context = config_options.executor.spmd.context
-    py_executor = config_options.executor.spmd.py_executor
+    if config_options.executor.spmd_context is None:
+        raise RuntimeError("spmd_context must be set for SPMD mode")
+    comm = config_options.executor.spmd_context.comm
+    context = config_options.executor.spmd_context.context
+    py_executor = config_options.executor.spmd_context.py_executor
 
     ir_context = IRExecutionContext(get_cuda_stream=context.get_stream_from_pool)
 
@@ -210,10 +213,9 @@ def allgather_polars_dataframe(
 @contextmanager
 def spmd_execution(
     *,
-    mr: rmm.mr.DeviceMemoryResource | None = None,
     rapidsmpf_options: Options | None = None,
     executor_options: dict[str, object] | None = None,
-    **engine_kwargs: Any,
+    engine_options: dict[str, Any] | None = None,
 ) -> Iterator[tuple[Communicator, Context, pl.GPUEngine]]:
     """
     Context manager that bootstraps a RapidsMPF SPMD context and a matching GPUEngine.
@@ -239,6 +241,27 @@ def spmd_execution(
 
     All resources (communicator, stream pool, thread-pool) are released on exit.
 
+    **Memory resource**
+
+    ``spmd_execution`` captures ``rmm.mr.get_current_device_resource()`` at entry,
+    wraps it in ``RmmResourceAdaptor`` (so libcudf temporary allocations and the
+    RapidsMPF ``Context`` share the same resource), sets the wrapped resource as
+    current, and restores the original on exit.
+
+    To use a custom allocator, call ``rmm.mr.set_current_device_resource(your_mr)``
+    before entering ``spmd_execution()``. Do not pre-wrap it in ``RmmResourceAdaptor``.
+
+    .. code-block:: python
+
+        import rmm
+
+        # Optional: install a pool allocator before entering spmd_execution.
+        # rmm.mr.set_current_device_resource(
+        #     rmm.mr.PoolMemoryResource(rmm.mr.CudaMemoryResource())
+        # )
+        with spmd_execution(...) as (comm, ctx, engine):
+            ...
+
     **DataFrame and LazyFrame semantics**
 
     Because every rank runs an independent Python process, a :class:`~polars.DataFrame`
@@ -283,9 +306,6 @@ def spmd_execution(
 
     Parameters
     ----------
-    mr
-        RMM device memory resource to use. Defaults to
-        ``rmm.mr.CudaAsyncMemoryResource()`` when ``None``.
     rapidsmpf_options
         RapidsMPF options. Defaults to ``Options(get_environment_variables())``
         when ``None``.
@@ -294,20 +314,20 @@ def spmd_execution(
         :class:`~polars.lazyframe.engine_config.GPUEngine`. The keys
         ``"runtime"``, ``"cluster"``, and ``"spmd"`` are reserved and may not
         be overridden.
-    **engine_kwargs
+    engine_options
         Extra keyword arguments forwarded directly to
         :class:`~polars.lazyframe.engine_config.GPUEngine`.  For example,
-        pass ``parquet_options={"use_rapidsmpf_native": True}`` to enable
-        native Parquet reads. The keys ``"memory_resource"`` and
+        pass ``engine_options={"parquet_options": {"use_rapidsmpf_native": True}}``
+        to enable native Parquet reads. The keys ``"memory_resource"`` and
         ``"executor"`` are reserved and may not be overridden.
 
     Yields
     ------
-    comm : Communicator
+    comm
         The active RapidsMPF communicator.
-    ctx : Context
+    ctx
         The active RapidsMPF context.
-    engine : pl.GPUEngine
+    engine
         A Polars GPU engine wired to ``comm`` and ``ctx``. Pass it to
         ``LazyFrame.collect(engine=engine)`` on each rank.
 
@@ -320,8 +340,8 @@ def spmd_execution(
         If ``executor_options`` contains any of the reserved keys
         ``"runtime"``, ``"cluster"``, or ``"spmd"``.
     TypeError
-        If ``engine_kwargs`` contains any of the reserved keys
-        ``"raise_on_fail"``, ``"memory_resource"``, or ``"executor"``.
+        If ``engine_options`` contains any of the reserved keys
+        ``"memory_resource"`` or ``"executor"``.
 
     Examples
     --------
@@ -341,20 +361,20 @@ def spmd_execution(
         )
 
     executor_options = executor_options or {}
-    engine_kwargs = engine_kwargs or {}
+    engine_options = engine_options or {}
 
     # Check for reserved keys.
-    if bad := {"runtime", "cluster", "spmd"} & executor_options.keys():
+    if bad := {"runtime", "cluster", "spmd_context"} & executor_options.keys():
         raise TypeError(f"executor_options may not contain reserved keys: {bad}")
-    if bad := {"memory_resource", "executor"} & engine_kwargs.keys():
-        raise TypeError(f"engine_kwargs may not contain reserved keys: {bad}")
+    if bad := {"memory_resource", "executor"} & engine_options.keys():
+        raise TypeError(f"engine_options may not contain reserved keys: {bad}")
 
     rapidsmpf_options = (
         rapidsmpf_options
         if rapidsmpf_options is not None
         else Options(get_environment_variables())
     )
-    mr = RmmResourceAdaptor(mr if mr is not None else rmm.mr.CudaAsyncMemoryResource())
+    mr = RmmResourceAdaptor(rmm.mr.get_current_device_resource())
     comm = bootstrap.create_ucxx_comm(
         progress_thread=ProgressThread(),
         type=bootstrap.BackendType.AUTO,
@@ -367,19 +387,22 @@ def spmd_execution(
         thread_name_prefix="spmd-executor",
     )
     try:
-        with Context.from_options(comm.logger, mr, rapidsmpf_options) as ctx:
+        with (
+            set_memory_resource(mr),
+            Context.from_options(comm.logger, mr, rapidsmpf_options) as ctx,
+        ):
             engine = pl.GPUEngine(
                 memory_resource=ctx.br().device_mr,
                 executor="streaming",
                 executor_options={
                     **executor_options,
                     "runtime": "rapidsmpf",
                     "cluster": "spmd",
-                    "spmd": SPMDContext(
+                    "spmd_context": SPMDContext(
                         comm=comm, context=ctx, py_executor=py_executor
                     ),
                 },
-                **engine_kwargs,
+                **engine_options,
             )
             yield comm, ctx, engine
     finally:

python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py

@@ -5,13 +5,16 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import operator
 import struct
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from functools import reduce
 from typing import TYPE_CHECKING, Any
 
+import rmm.mr
+
 from cudf_polars.dsl.tracing import LOG_TRACES, Scope
 
 try:
@@ -42,7 +45,7 @@
 from cudf_polars.experimental.utils import _concat
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncIterator, Callable
+    from collections.abc import AsyncIterator, Callable, Iterator
 
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.streaming.core.channel import Channel
@@ -57,6 +60,27 @@
     from cudf_polars.typing import DataType, Schema
 
 
+@contextlib.contextmanager
+def set_memory_resource(mr: rmm.mr.DeviceMemoryResource) -> Iterator[None]:
+    """
+    Context manager that temporarily sets ``mr`` as the current device resource.
+
+    On entry, ``mr`` is installed via ``rmm.mr.set_current_device_resource(mr)``.
+    On exit, the previously active resource is restored unconditionally.
+
+    Parameters
+    ----------
+    mr
+        The memory resource to activate for the duration of the block.
+    """
+    old = rmm.mr.get_current_device_resource()
+    rmm.mr.set_current_device_resource(mr)
+    try:
+        yield
+    finally:
+        rmm.mr.set_current_device_resource(old)
+
+
 @asynccontextmanager
 async def shutdown_on_error(
     context: Context,

python/cudf_polars/cudf_polars/utils/config.py

@@ -872,7 +872,7 @@ class StreamingExecutor:
             f"{_env_prefix}__RAPIDSMPF_PY_EXECUTOR_MAX_WORKERS", int, default=None
         )
     )
-    spmd: SPMDContext | None = None
+    spmd_context: SPMDContext | None = None
     ray_context: RayContext | None = None
 
     def __post_init__(self) -> None:  # noqa: D105

python/cudf_polars/docs/cudf-polars-mp.md

@@ -356,35 +356,41 @@ The result is guaranteed to be a `pl.DataFrame` containing rows from all ranks i
 
 ### Passing options
 
-`mr`, `rapidsmpf_options`, `executor_options`, and `engine_kwargs` accept pass-through
+`rapidsmpf_options`, `executor_options`, and `engine_options` accept pass-through
 arguments:
 
 ```python
 import rmm
 from rapidsmpf.integrations.cudf_polars import Options
 
 with spmd_execution(
-    mr=rmm.mr.PoolMemoryResource(rmm.mr.CudaMemoryResource()),
     rapidsmpf_options=Options(num_streaming_threads=8),
     executor_options={
         "max_rows_per_partition": 500_000,
         "rapidsmpf_spill": True,
         "rapidsmpf_py_executor_max_workers": 2,
     },
-    parquet_options={"use_rapidsmpf_native": True},
+    engine_options={"parquet_options": {"use_rapidsmpf_native": True}},
 ) as (comm, ctx, engine):
     ...
 ```
 
-`mr` is an `rmm.mr.DeviceMemoryResource` used as the GPU memory resource for the
-RapidsMPF `Context`. Defaults to `None` (uses the current device resource).
+**Memory resource:** `spmd_execution` captures `rmm.mr.get_current_device_resource()`
+at entry, wraps it in `RmmResourceAdaptor` (so libcudf temporary allocations and the
+RapidsMPF `Context` share the same resource), sets the wrapped resource as current, and
+restores the original resource on exit. To use a custom allocator, call
+`rmm.mr.set_current_device_resource(your_mr)` **before** entering `spmd_execution()`.
+Do not pre-wrap it in `RmmResourceAdaptor`.
 
 `rapidsmpf_options` is an `Options` object passed to the RapidsMPF `Context`. Defaults
 to `None` (uses RapidsMPF defaults).
 
 `executor_options` is forwarded directly to `pl.GPUEngine` as its `executor_options`
 argument; user-supplied keys are merged with reserved entries set by `spmd_execution()`.
-Any additional keyword arguments to `spmd_execution()` are also forwarded to `pl.GPUEngine`.
+
+`engine_options` is forwarded as keyword arguments to `pl.GPUEngine`. For example,
+pass `engine_options={"parquet_options": {"use_rapidsmpf_native": True}}` to enable
+native Parquet reads.
 
 Notable `executor_options` keys:
 

python/cudf_polars/tests/experimental/rapidsmpf/test_spmd.py

@@ -4,10 +4,9 @@
 
 from __future__ import annotations
 
-from typing import Any
-
 import pytest
 from rapidsmpf.bootstrap import is_running_with_rrun
+from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor
 
 import polars as pl
 
@@ -35,39 +34,30 @@ def test_yields_context_and_engine() -> None:
 
 def test_reserved_keys() -> None:
     """executor_options rejects reserved keys."""
-    for key in ("runtime", "cluster", "spmd"):
+    for key in ("runtime", "cluster", "spmd_context"):
         with (
             pytest.raises(TypeError, match="reserved"),
             spmd_execution(executor_options={key: "anything"}),
         ):
             pass
 
 
-def test_engine_kwargs_reserved_keys() -> None:
-    """engine_kwargs rejects keys that are set explicitly by spmd_execution."""
+def test_engine_options_reserved_keys() -> None:
+    """engine_options rejects keys that are set explicitly by spmd_execution."""
     for key in ("memory_resource", "executor"):
-        kwargs: dict[str, Any] = {key: "anything"}
         with (
             pytest.raises(TypeError, match="reserved"),
-            spmd_execution(**kwargs),
+            spmd_execution(engine_options={key: "anything"}),
         ):
             pass
 
 
-def test_engine_kwargs_parquet_options() -> None:
-    """engine_kwargs forwards parquet_options to GPUEngine without error."""
-    with spmd_execution(parquet_options={}) as (comm, ctx, engine):
+def test_engine_options_parquet_options() -> None:
+    """engine_options forwards parquet_options to GPUEngine without error."""
+    with spmd_execution(engine_options={"parquet_options": {}}) as (comm, ctx, engine):
         assert isinstance(engine, pl.GPUEngine)
 
 
-def test_custom_mr() -> None:
-    """spmd_execution accepts a custom memory resource."""
-    mr = rmm.mr.CudaMemoryResource()
-    with spmd_execution(mr=mr) as (comm, ctx, engine):
-        result = pl.LazyFrame({"a": [1, 2, 3]}).collect(engine=engine)
-    assert result.shape == (3, 1)
-
-
 def test_scan() -> None:
     """Each rank scans its own single-row LazyFrame and gets that row back."""
     with spmd_execution() as (comm, ctx, engine):
@@ -158,6 +148,20 @@ def test_allgather_polars_dataframe_empty() -> None:
     assert result.dtypes == [pl.Int32, pl.Float64]
 
 
+def test_mr_wrapped_as_current_inside_context() -> None:
+    """Inside spmd_execution the current device resource is RmmResourceAdaptor."""
+    with spmd_execution() as (comm, ctx, engine):
+        assert isinstance(rmm.mr.get_current_device_resource(), RmmResourceAdaptor)
+
+
+def test_mr_restored_after_context() -> None:
+    """After spmd_execution exits the original device resource is restored."""
+    original = rmm.mr.get_current_device_resource()
+    with spmd_execution() as (comm, ctx, engine):
+        pass
+    assert rmm.mr.get_current_device_resource() is original
+
+
 def test_allgather_polars_dataframe_multi_column() -> None:
     """allgather preserves column names, count, and dtypes for multi-column DataFrames."""
     with spmd_execution() as (comm, ctx, _):