From 51baa7dd11d100baee822743ba9a293717a75e7f Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Fri, 6 Feb 2026 02:25:52 -0800
Subject: [PATCH 1/9] Add rrun support for cudf-polars

---
 .../experimental/benchmarks/utils.py          | 108 ++++++++++++++----
 .../cudf_polars/experimental/parallel.py      |  10 +-
 .../experimental/rapidsmpf/bootstrap_ctx.py   |  98 ++++++++++++++++
 .../experimental/rapidsmpf/core.py            |  92 ++++++++++++++-
 .../cudf_polars/cudf_polars/utils/config.py   |  40 +++++--
 5 files changed, 309 insertions(+), 39 deletions(-)
 create mode 100644 python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py

diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
index 906413450668..6774a2f29bea 100644
--- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
+++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
@@ -547,6 +547,25 @@ def initialize_dask_cluster(run_config: RunConfig, args: argparse.Namespace):  #
     Client or None
         A Dask distributed Client, or None if not using distributed mode.
     """
+    # Check if running with rrun
+    try:
+        from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import (
+            get_nranks,
+            get_rank,
+            is_running_with_rrun,
+        )
+
+        if is_running_with_rrun():
+            rank = get_rank()
+            nranks = get_nranks()
+            if rank == 0:
+                print(
+                    f"[rrun] Detected rrun execution environment with {nranks} ranks"
+                )
+            return None  # No Dask client needed for rrun
+    except ImportError:
+        pass  # rapidsmpf not available, continue with normal path
+
     if run_config.cluster != "distributed":
         return None
 
@@ -1029,6 +1048,28 @@ def run_polars(
     validation_failures: list[int] = []
     query_failures: list[tuple[int, int]] = []
 
+    # Check if running with rrun
+    is_rrun = False
+    rank = 0
+    nranks = 1
+    try:
+        from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import (
+            get_nranks,
+            get_rank,
+            is_running_with_rrun,
+        )
+
+        if is_running_with_rrun():
+            is_rrun = True
+            rank = get_rank()
+            nranks = get_nranks()
+            # Update cluster and n_workers for rrun mode
+            run_config = dataclasses.replace(
+                run_config, cluster="rrun", n_workers=nranks
+            )
+    except ImportError:
+        pass  # rapidsmpf not available
+
     client = initialize_dask_cluster(run_config, args)
 
     # Update n_workers from the actual cluster when using scheduler file/address
@@ -1078,10 +1119,22 @@ def run_polars(
             try:
                 result = execute_query(q_id, i, q, run_config, args, engine)
             except Exception:
-                print(f"❌ query={q_id} iteration={i} failed!")
-                print(traceback.format_exc())
+                if not is_rrun or rank == 0:
+                    print(f"❌ query={q_id} iteration={i} failed!")
+                    print(traceback.format_exc())
                 query_failures.append((q_id, i))
                 continue
+
+            # In rrun mode, result is None for non-root ranks
+            if is_rrun and result is None:
+                # Non-root ranks: skip result processing but record timing
+                t1 = time.monotonic()
+                record = Record(
+                    query=q_id, iteration=i, duration=t1 - t0, shuffle_stats=None
+                )
+                records[q_id].append(record)
+                continue
+
             if run_config.shuffle == "rapidsmpf" and run_config.gather_shuffle_stats:
                 from rapidsmpf.integrations.dask.shuffler import (
                     clear_shuffle_statistics,
@@ -1101,28 +1154,31 @@ def run_polars(
                         executor=run_config.executor,
                         check_exact=False,
                     )
-                    print(f"✅ Query {q_id} passed validation!")
+                    if not is_rrun or rank == 0:
+                        print(f"✅ Query {q_id} passed validation!")
                 except AssertionError as e:
                     validation_failures.append(q_id)
-                    print(f"❌ Query {q_id} failed validation!\n{e}")
+                    if not is_rrun or rank == 0:
+                        print(f"❌ Query {q_id} failed validation!\n{e}")
 
             t1 = time.monotonic()
             record = Record(
                 query=q_id, iteration=i, duration=t1 - t0, shuffle_stats=shuffle_stats
             )
-            if args.print_results:
+            if args.print_results and (not is_rrun or rank == 0):
                 print(result)
 
-            if args.results_directory is not None and i == 0:
+            if args.results_directory is not None and i == 0 and (not is_rrun or rank == 0):
                 results_dir = Path(args.results_directory)
                 results_dir.mkdir(parents=True, exist_ok=True)
                 output_path = results_dir / f"q_{q_id:02d}.parquet"
                 result.write_parquet(output_path)
 
-            print(
-                f"Query {q_id} - Iteration {i} finished in {record.duration:0.4f}s",
-                flush=True,
-            )
+            if not is_rrun or rank == 0:
+                print(
+                    f"Query {q_id} - Iteration {i} finished in {record.duration:0.4f}s",
+                    flush=True,
+                )
             records[q_id].append(record)
 
     run_config = dataclasses.replace(run_config, records=dict(records))
@@ -1176,24 +1232,26 @@ def sort_key(x: dict) -> tuple[int, int]:
 
             run_config.records[query_id] = new_records
 
-    if args.summarize:
-        run_config.summarize()
+    # Only rank 0 should print summaries and write output
+    if not is_rrun or rank == 0:
+        if args.summarize:
+            run_config.summarize()
 
-    if client is not None:
-        client.close(timeout=60)
+        if args.validate and run_config.executor != "cpu":
+            print("\nValidation Summary")
+            print("==================")
+            if validation_failures:
+                print(
+                    f"{len(validation_failures)} queries failed validation: {sorted(set(validation_failures))}"
+                )
+            else:
+                print("All validated queries passed.")
 
-    if args.validate and run_config.executor != "cpu":
-        print("\nValidation Summary")
-        print("==================")
-        if validation_failures:
-            print(
-                f"{len(validation_failures)} queries failed validation: {sorted(set(validation_failures))}"
-            )
-        else:
-            print("All validated queries passed.")
+        args.output.write(json.dumps(run_config.serialize(engine=engine)))
+        args.output.write("\n")
 
-    args.output.write(json.dumps(run_config.serialize(engine=engine)))
-    args.output.write("\n")
+    if client is not None:
+        client.close(timeout=60)
 
     if query_failures or validation_failures:
         sys.exit(1)
diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/experimental/parallel.py
index 28533dd4ae3a..2d292651e5df 100644
--- a/python/cudf_polars/cudf_polars/experimental/parallel.py
+++ b/python/cudf_polars/cudf_polars/experimental/parallel.py
@@ -173,7 +173,15 @@ def get_scheduler(config_options: ConfigOptions) -> Any:
 
     cluster = config_options.executor.cluster
 
-    if (
+    # Check if running with rrun
+    from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import is_running_with_rrun
+
+    if is_running_with_rrun() or cluster == "rrun":
+        # rrun mode: use synchronous scheduler (already rank-aware via streaming)
+        from cudf_polars.experimental.scheduler import synchronous_scheduler
+
+        return synchronous_scheduler
+    elif (
         cluster == "distributed"
     ):  # pragma: no cover; block depends on executor type and Distributed cluster
         from distributed import get_client
diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
new file mode 100644
index 000000000000..e57793c03d99
--- /dev/null
+++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
@@ -0,0 +1,98 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+"""Bootstrap context management for rrun execution."""
+
+from __future__ import annotations
+
+import os
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from rapidsmpf.bootstrap.bootstrap import Context
+
+try:
+    import rapidsmpf.bootstrap as bootstrap
+
+    BOOTSTRAP_AVAILABLE = True
+except ImportError:
+    BOOTSTRAP_AVAILABLE = False
+
+_global_context: Context | None = None
+
+
+def is_running_with_rrun() -> bool:
+    """
+    Check if running under rrun.
+
+    Returns
+    -------
+    bool
+        True if the RAPIDSMPF_RANK environment variable is set,
+        indicating execution under rrun.
+    """
+    if not BOOTSTRAP_AVAILABLE:
+        return False
+    return bootstrap.is_running_with_rrun()
+
+
+def get_bootstrap_context() -> Context:
+    """
+    Get or initialize bootstrap context (singleton).
+
+    Returns
+    -------
+    Context
+        The RapidsMPF bootstrap context.
+
+    Raises
+    ------
+    RuntimeError
+        If rapidsmpf.bootstrap is not available or not running under rrun.
+    """
+    global _global_context
+    if _global_context is None:
+        if not BOOTSTRAP_AVAILABLE:
+            raise RuntimeError(
+                "rapidsmpf.bootstrap not available. "
+                "Please install rapidsmpf to use rrun execution."
+            )
+        if not is_running_with_rrun():
+            raise RuntimeError(
+                "Not running under rrun (RAPIDSMPF_RANK environment variable not set). "
+                "Use 'rrun -n <nranks> python ...' to launch with rrun."
+            )
+        # Initialize the bootstrap context
+        # The context is initialized based on environment variables set by rrun
+        _global_context = bootstrap.create_ucxx_comm(bootstrap.BackendType.FILE)
+    return _global_context
+
+
+def get_rank() -> int:
+    """
+    Get current rank.
+
+    Returns
+    -------
+    int
+        The rank of the current process (0 if not running under rrun).
+    """
+    if not is_running_with_rrun():
+        return 0
+    # Read directly from environment variable for efficiency
+    # This avoids initializing the full bootstrap context just to get rank
+    return int(os.environ.get("RAPIDSMPF_RANK", "0"))
+
+
+def get_nranks() -> int:
+    """
+    Get total number of ranks.
+
+    Returns
+    -------
+    int
+        The total number of ranks (1 if not running under rrun).
+    """
+    if not is_running_with_rrun():
+        return 1
+    # Read directly from environment variable for efficiency
+    return int(os.environ.get("RAPIDSMPF_NRANKS", "1"))
diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
index 376ce43c2a29..f9c669d734a2 100644
--- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
+++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
@@ -74,7 +74,7 @@ def evaluate_logical_plan(
     config_options: ConfigOptions,
     *,
     collect_metadata: bool = False,
-) -> tuple[pl.DataFrame, list[ChannelMetadata] | None]:
+) -> tuple[pl.DataFrame | None, list[ChannelMetadata] | None]:
     """
     Evaluate a logical plan with the RapidsMPF streaming runtime.
 
@@ -90,11 +90,20 @@ def evaluate_logical_plan(
     Returns
     -------
     The output DataFrame and metadata collector.
+    For rrun execution, non-root ranks return (None, None).
     """
     assert config_options.executor.name == "streaming", "Executor must be streaming"
     assert config_options.executor.runtime == "rapidsmpf", "Runtime must be rapidsmpf"
 
-    # Lower the IR graph on the client process (for now).
+    # Check if running with rrun
+    from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import (
+        get_rank,
+        is_running_with_rrun,
+    )
+
+    is_rrun = is_running_with_rrun()
+
+    # Lower the IR graph on all processes (for rrun) or client process (for dask).
     ir, partition_info, stats = lower_ir_graph(ir, config_options)
 
     # Log the query plan structure for tracing (no-op if tracing disabled)
@@ -104,8 +113,29 @@ def evaluate_logical_plan(
     with ReserveOpIDs(ir) as collective_id_map:
         # Build and execute the streaming pipeline.
         # This must be done on all worker processes
-        # for cluster == "distributed".
+        # for cluster == "distributed" or cluster == "rrun".
         if (
+            config_options.executor.cluster == "rrun" or is_rrun
+        ):  # pragma: no cover; block depends on executor type and rrun cluster
+            # SPMD execution: All ranks execute, only rank 0 returns result
+            result, metadata_collector = evaluate_pipeline(
+                ir,
+                partition_info,
+                config_options,
+                stats,
+                collective_id_map,
+                collect_metadata=collect_metadata,
+            )
+
+            # Only rank 0 returns result to caller
+            rank = get_rank()
+            if rank == 0:
+                return result, metadata_collector
+            else:
+                # Non-root ranks return None
+                return None, None
+
+        elif (
             config_options.executor.cluster == "distributed"
         ):  # pragma: no cover; block depends on executor type and Distributed cluster
             # Distributed execution: Use client.run
@@ -122,6 +152,7 @@ def evaluate_logical_plan(
                 collective_id_map,
                 collect_metadata=collect_metadata,
             )
+            return result, metadata_collector
         else:
             # Single-process execution: Run locally
             result, metadata_collector = evaluate_pipeline(
@@ -132,8 +163,7 @@ def evaluate_logical_plan(
                 collective_id_map,
                 collect_metadata=collect_metadata,
             )
-
-    return result, metadata_collector
+            return result, metadata_collector
 
 
 def evaluate_pipeline(
@@ -175,12 +205,62 @@ def evaluate_pipeline(
 
     _initial_mr: Any = None
     stream_pool: CudaStreamPool | bool = False
+
+    # Check if running with rrun
+    from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import (
+        get_bootstrap_context,
+        is_running_with_rrun,
+    )
+
+    is_rrun = is_running_with_rrun()
+
     if rmpf_context is not None:
-        # Using "distributed" mode.
+        # Using "distributed" mode (Dask).
         # Always use the RapidsMPF stream pool for now.
         br = rmpf_context.br()
         stream_pool = True
         rmpf_context_manager = contextlib.nullcontext(rmpf_context)
+    elif is_rrun and rmpf_context is None:
+        # Using "rrun" mode - initialize from bootstrap context
+        # Create a new distributed RapidsMPF context using the bootstrap communicator
+        _original_mr = rmm.mr.get_current_device_resource()
+        mr = RmmResourceAdaptor(_original_mr)
+        rmm.mr.set_current_device_resource(mr)
+
+        # Get the bootstrap-initialized communicator
+        bootstrap_ctx = get_bootstrap_context()
+
+        options = Options(
+            {
+                # By default, set the number of streaming threads to the max
+                # number of IO threads. The user may override this with an
+                # environment variable (i.e. RAPIDSMPF_NUM_STREAMING_THREADS)
+                "num_streaming_threads": str(
+                    max(config_options.executor.max_io_threads, 1)
+                )
+            }
+            | get_environment_variables()
+        )
+        pinned_mr = (
+            PinnedMemoryResource.make_if_available()
+            if config_options.executor.spill_to_pinned_memory
+            else None
+        )
+        if isinstance(config_options.cuda_stream_policy, CUDAStreamPoolConfig):
+            stream_pool = config_options.cuda_stream_policy.build()
+        else:
+            stream_pool = True  # Use stream pool for distributed execution
+
+        # Note: For rrun, we use the communicator from bootstrap_ctx
+        # The BufferResource is created but memory limits are not enforced
+        # in the same way as single-GPU mode
+        br = BufferResource(
+            mr,
+            pinned_mr=pinned_mr,
+            memory_available=None,  # No memory limits for distributed
+            stream_pool=stream_pool,
+        )
+        rmpf_context_manager = Context(bootstrap_ctx, br, options)
     else:
         # Using "single" mode.
         # Create a new local RapidsMPF context.
diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py
index 31bbec83fc2a..c245b205be54 100644
--- a/python/cudf_polars/cudf_polars/utils/config.py
+++ b/python/cudf_polars/cudf_polars/utils/config.py
@@ -167,10 +167,14 @@ class Cluster(str, enum.Enum):
     * ``Cluster.DISTRIBUTED`` : Multi-GPU distributed execution. Currently
       uses a Dask-based distributed scheduler and requires an
       active Dask cluster.
+    * ``Cluster.RRUN`` : Multi-GPU distributed execution using rrun-based
+      coordination with RapidsMPF bootstrap. Uses SPMD execution model
+      where all ranks execute the same code.
     """
 
     SINGLE = "single"
     DISTRIBUTED = "distributed"
+    RRUN = "rrun"
 
 
 class Scheduler(str, enum.Enum):
@@ -354,10 +358,10 @@ def default_blocksize(cluster: str) -> int:
         return 1_000_000_000
 
     if (
-        cluster == "distributed"
+        cluster in ("distributed", "rrun")
         or _env_get_int("POLARS_GPU_ENABLE_CUDA_MANAGED_MEMORY", default=1) == 0
     ):
-        # Distributed execution requires a conservative
+        # Distributed/rrun execution requires a conservative
         # blocksize for now. We are also more conservative
         # when UVM is disabled.
         blocksize = int(device_size * 0.025)
@@ -861,10 +865,28 @@ def __post_init__(self) -> None:  # noqa: D105
                 stacklevel=2,
             )
 
+        # rrun cluster requires rapidsmpf runtime
+        if self.cluster == "rrun":
+            if self.runtime != "rapidsmpf":
+                raise ValueError(
+                    "cluster='rrun' requires runtime='rapidsmpf'. "
+                    "Please set runtime='rapidsmpf' when using rrun."
+                )
+            # Force shuffle_method to rapidsmpf for rrun
+            if self.shuffle_method is not None and self.shuffle_method != "rapidsmpf":
+                warnings.warn(
+                    f"Ignoring shuffle_method='{self.shuffle_method}' for rrun cluster. "
+                    "Using shuffle_method='rapidsmpf' instead.",
+                    stacklevel=2,
+                )
+            object.__setattr__(self, "shuffle_method", "rapidsmpf")
+
         # Handle shuffle_method defaults for streaming executor
         if self.shuffle_method is None:
-            if self.cluster == "distributed" and rapidsmpf_distributed_available():
-                # For distributed cluster, prefer rapidsmpf if available
+            if (
+                self.cluster == "distributed" or self.cluster == "rrun"
+            ) and rapidsmpf_distributed_available():
+                # For distributed/rrun cluster, prefer rapidsmpf if available
                 object.__setattr__(self, "shuffle_method", "rapidsmpf")
             else:
                 # Otherwise, use task-based shuffle for now.
@@ -883,6 +905,10 @@ def __post_init__(self) -> None:  # noqa: D105
                 raise ValueError(
                     "rapidsmpf shuffle method requested, but rapidsmpf is not installed."
                 )
+            elif self.cluster == "rrun" and not rapidsmpf_single_available():
+                raise ValueError(
+                    "rrun cluster requires rapidsmpf to be installed."
+                )
             # Select "rapidsmpf-single" for single-GPU
             if self.cluster == "single":
                 object.__setattr__(self, "shuffle_method", "rapidsmpf-single")
@@ -902,7 +928,7 @@ def __post_init__(self) -> None:  # noqa: D105
                 self,
                 "broadcast_join_limit",
                 # Usually better to avoid shuffling for single gpu with UVM
-                2 if self.cluster == "distributed" else 32,
+                2 if self.cluster in ("distributed", "rrun") else 32,
             )
         object.__setattr__(self, "cluster", Cluster(self.cluster))
         object.__setattr__(self, "shuffle_method", ShuffleMethod(self.shuffle_method))
@@ -929,10 +955,10 @@ def __post_init__(self) -> None:  # noqa: D105
                 DynamicPlanningOptions(**self.dynamic_planning),
             )
 
-        if self.cluster == "distributed":
+        if self.cluster in ("distributed", "rrun"):
             if self.sink_to_directory is False:
                 raise ValueError(
-                    "The distributed cluster requires sink_to_directory=True"
+                    f"The {self.cluster} cluster requires sink_to_directory=True"
                 )
             object.__setattr__(self, "sink_to_directory", True)
         elif self.sink_to_directory is None:

From 163050549047fcdd175dce4cf36a1aef3a5f48cf Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Fri, 6 Feb 2026 03:42:03 -0800
Subject: [PATCH 2/9] Attempt Slurm fix

---
 .../experimental/rapidsmpf/bootstrap_ctx.py   | 132 +++++++++++++++++-
 1 file changed, 126 insertions(+), 6 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
index e57793c03d99..34dc1adfdb7e 100644
--- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
+++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
@@ -35,6 +35,67 @@ def is_running_with_rrun() -> bool:
     return bootstrap.is_running_with_rrun()
 
 
+def is_running_under_slurm() -> bool:
+    """
+    Check if running under Slurm.
+
+    Returns
+    -------
+    bool
+        True if Slurm environment variables are detected.
+    """
+    # Check for Slurm environment variables
+    return (
+        "SLURM_JOB_ID" in os.environ
+        or "SLURM_PROCID" in os.environ
+        or "PMIX_NAMESPACE" in os.environ
+    )
+
+
+def _detect_backend_type():
+    """
+    Detect the appropriate backend type based on environment.
+
+    Returns
+    -------
+    BackendType
+        The detected backend type.
+
+    Notes
+    -----
+    Detection logic:
+    1. If RAPIDSMPF_COORD_DIR is set -> use AUTO (will choose FILE)
+    2. If running under Slurm without COORD_DIR -> explicitly use SLURM if available
+    3. Otherwise use AUTO (default)
+    """
+    if not BOOTSTRAP_AVAILABLE:
+        raise RuntimeError("rapidsmpf.bootstrap not available")
+
+    # If COORD_DIR is set, FILE backend can work - use AUTO
+    if "RAPIDSMPF_COORD_DIR" in os.environ:
+        return bootstrap.BackendType.AUTO
+
+    # If running under Slurm without COORD_DIR, try to use SLURM backend explicitly
+    if is_running_under_slurm():
+        # Check if SLURM backend is available in the Python bindings
+        if hasattr(bootstrap.BackendType, "SLURM"):
+            print(
+                f"[Rank {get_rank()}] Detected Slurm environment, using SLURM backend",
+                flush=True,
+            )
+            return bootstrap.BackendType.SLURM
+        else:
+            # SLURM not in Python enum, but AUTO should still detect it in C++
+            print(
+                f"[Rank {get_rank()}] Detected Slurm environment, using AUTO backend (will select SLURM in C++)",
+                flush=True,
+            )
+            return bootstrap.BackendType.AUTO
+
+    # Default to AUTO
+    return bootstrap.BackendType.AUTO
+
+
 def get_bootstrap_context() -> Context:
     """
     Get or initialize bootstrap context (singleton).
@@ -61,9 +122,53 @@ def get_bootstrap_context() -> Context:
                 "Not running under rrun (RAPIDSMPF_RANK environment variable not set). "
                 "Use 'rrun -n <nranks> python ...' to launch with rrun."
             )
-        # Initialize the bootstrap context
-        # The context is initialized based on environment variables set by rrun
-        _global_context = bootstrap.create_ucxx_comm(bootstrap.BackendType.FILE)
+
+        # Detect and use appropriate backend
+        backend_type = _detect_backend_type()
+
+        # Debug: print environment info on rank 0
+        rank = get_rank()
+        if rank == 0:
+            print(f"[Bootstrap] Backend type: {backend_type}", flush=True)
+            print(
+                f"[Bootstrap] RAPIDSMPF_COORD_DIR: {os.environ.get('RAPIDSMPF_COORD_DIR', 'NOT SET')}",
+                flush=True,
+            )
+            print(
+                f"[Bootstrap] SLURM_JOB_ID: {os.environ.get('SLURM_JOB_ID', 'NOT SET')}",
+                flush=True,
+            )
+            print(
+                f"[Bootstrap] PMIX_NAMESPACE: {os.environ.get('PMIX_NAMESPACE', 'NOT SET')}",
+                flush=True,
+            )
+
+        try:
+            # Initialize the bootstrap context with detected backend
+            _global_context = bootstrap.create_ucxx_comm(backend_type)
+        except RuntimeError as e:
+            # Provide helpful error message
+            error_msg = f"Failed to initialize bootstrap context: {e}\n"
+            error_msg += "\nEnvironment variables:\n"
+            error_msg += f"  RAPIDSMPF_RANK: {os.environ.get('RAPIDSMPF_RANK', 'NOT SET')}\n"
+            error_msg += (
+                f"  RAPIDSMPF_NRANKS: {os.environ.get('RAPIDSMPF_NRANKS', 'NOT SET')}\n"
+            )
+            error_msg += f"  RAPIDSMPF_COORD_DIR: {os.environ.get('RAPIDSMPF_COORD_DIR', 'NOT SET')}\n"
+            error_msg += f"  SLURM_JOB_ID: {os.environ.get('SLURM_JOB_ID', 'NOT SET')}\n"
+            error_msg += (
+                f"  SLURM_PROCID: {os.environ.get('SLURM_PROCID', 'NOT SET')}\n"
+            )
+            error_msg += f"  SLURM_NPROCS: {os.environ.get('SLURM_NPROCS', 'NOT SET')}\n"
+            error_msg += (
+                f"  PMIX_NAMESPACE: {os.environ.get('PMIX_NAMESPACE', 'NOT SET')}\n"
+            )
+            error_msg += "\nFor Slurm, ensure you're using: srun --mpi=pmix ...\n"
+            error_msg += (
+                "For rrun, ensure RAPIDSMPF_COORD_DIR is set or rrun is launching.\n"
+            )
+            raise RuntimeError(error_msg) from e
+
     return _global_context
 
 
@@ -79,8 +184,15 @@ def get_rank() -> int:
     if not is_running_with_rrun():
         return 0
     # Read directly from environment variable for efficiency
-    # This avoids initializing the full bootstrap context just to get rank
-    return int(os.environ.get("RAPIDSMPF_RANK", "0"))
+    # Try RAPIDSMPF_RANK first, fall back to SLURM_PROCID for Slurm
+    rank = os.environ.get("RAPIDSMPF_RANK")
+    if rank is not None:
+        return int(rank)
+    # Fall back to Slurm env vars
+    rank = os.environ.get("SLURM_PROCID")
+    if rank is not None:
+        return int(rank)
+    return 0
 
 
 def get_nranks() -> int:
@@ -95,4 +207,12 @@ def get_nranks() -> int:
     if not is_running_with_rrun():
         return 1
     # Read directly from environment variable for efficiency
-    return int(os.environ.get("RAPIDSMPF_NRANKS", "1"))
+    # Try RAPIDSMPF_NRANKS first, fall back to SLURM_NPROCS/SLURM_NTASKS for Slurm
+    nranks = os.environ.get("RAPIDSMPF_NRANKS")
+    if nranks is not None:
+        return int(nranks)
+    # Fall back to Slurm env vars
+    nranks = os.environ.get("SLURM_NPROCS") or os.environ.get("SLURM_NTASKS")
+    if nranks is not None:
+        return int(nranks)
+    return 1

From 71a54290260204a9040d0562e92dc42dde4d6c62 Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Fri, 6 Feb 2026 05:02:07 -0800
Subject: [PATCH 3/9] Attempt Slurm fix part 2

---
 .../experimental/rapidsmpf/bootstrap_ctx.py   | 36 +++++++++++++++----
 1 file changed, 29 insertions(+), 7 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
index 34dc1adfdb7e..0b7f295deec1 100644
--- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
+++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
@@ -65,7 +65,7 @@ def _detect_backend_type():
     -----
     Detection logic:
     1. If RAPIDSMPF_COORD_DIR is set -> use AUTO (will choose FILE)
-    2. If running under Slurm without COORD_DIR -> explicitly use SLURM if available
+    2. If running under Slurm without COORD_DIR -> set up FILE backend with temp dir
     3. Otherwise use AUTO (default)
     """
     if not BOOTSTRAP_AVAILABLE:
@@ -75,7 +75,8 @@ def _detect_backend_type():
     if "RAPIDSMPF_COORD_DIR" in os.environ:
         return bootstrap.BackendType.AUTO
 
-    # If running under Slurm without COORD_DIR, try to use SLURM backend explicitly
+    # If running under Slurm without COORD_DIR, we need to set one up
+    # because the SLURM backend is not available in Python bindings yet
     if is_running_under_slurm():
         # Check if SLURM backend is available in the Python bindings
         if hasattr(bootstrap.BackendType, "SLURM"):
@@ -85,11 +86,32 @@ def _detect_backend_type():
             )
             return bootstrap.BackendType.SLURM
         else:
-            # SLURM not in Python enum, but AUTO should still detect it in C++
-            print(
-                f"[Rank {get_rank()}] Detected Slurm environment, using AUTO backend (will select SLURM in C++)",
-                flush=True,
-            )
+            # SLURM backend not available in Python bindings
+            # Create a temporary coordination directory as workaround
+
+            # Use a shared temp directory based on Slurm job ID
+            job_id = os.environ.get("SLURM_JOB_ID", "unknown")
+            coord_dir = f"/tmp/rapidsmpf-coord-{job_id}"
+
+            # All ranks create the directory (exist_ok=True avoids race)
+            rank = get_rank()
+            os.makedirs(coord_dir, exist_ok=True)
+
+            if rank == 0:
+                print(
+                    f"[Rank {rank}] SLURM backend not available in Python bindings, "
+                    f"using FILE backend with coordination directory: {coord_dir}",
+                    flush=True,
+                )
+                print(
+                    f"[Rank {rank}] NOTE: Clean up {coord_dir} after job completes "
+                    "(or it will be reused on next job with same ID)",
+                    flush=True,
+                )
+
+            # Set the environment variable for all ranks
+            os.environ["RAPIDSMPF_COORD_DIR"] = coord_dir
+
             return bootstrap.BackendType.AUTO
 
     # Default to AUTO

From 53c6498abd25e2ffa89958a8d169e50fcb7cc0a6 Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Fri, 6 Feb 2026 10:14:57 -0800
Subject: [PATCH 4/9] RMM resources

---
 .../experimental/benchmarks/utils.py          | 97 +++++++++++++++++++
 .../experimental/rapidsmpf/core.py            | 40 +++++++-
 2 files changed, 133 insertions(+), 4 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
index 6774a2f29bea..ef9d923de328 100644
--- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
+++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
@@ -464,6 +464,8 @@ def get_executor_options(
             executor_options["rapidsmpf_spill"] = run_config.rapidsmpf_spill
         if run_config.cluster == "distributed":
             executor_options["cluster"] = "distributed"
+        elif run_config.cluster == "rrun":
+            executor_options["cluster"] = "rrun"
         executor_options["stats_planning"] = {
             "use_reduction_planning": run_config.stats_planning,
             "use_sampling": (
@@ -526,6 +528,99 @@ def print_query_plan(
         )
 
 
+def _setup_rmm_for_rrun(args: argparse.Namespace, rank: int) -> None:  # type: ignore[no-untyped-def]
+    """
+    Set up RMM resources for rrun mode (similar to RMMPlugin for Dask workers).
+
+    This configures the RMM memory resource based on command-line arguments,
+    replicating what the Dask-CUDA RMMPlugin does for workers.
+
+    Parameters
+    ----------
+    args : argparse.Namespace
+        Parsed command line arguments containing RMM settings.
+    rank : int
+        The rank of the current process.
+    """
+    import rmm
+
+    if rank == 0:
+        print("[RMM Setup] Configuring RMM resources for rrun mode", flush=True)
+
+    # Parse pool size if specified
+    pool_size = None
+    if args.rmm_pool_size is not None:
+        if isinstance(args.rmm_pool_size, str):
+            from dask.utils import parse_bytes
+
+            pool_size = parse_bytes(args.rmm_pool_size)
+        else:
+            # Assume it's a fraction
+            total_memory = rmm.mr.available_device_memory()[1]
+            pool_size = int(total_memory * args.rmm_pool_size)
+
+    # Set up RMM resource based on configuration
+    if args.rmm_async:
+        # Use CudaAsyncMemoryResource (stream-ordered allocator)
+        if rank == 0:
+            print("[RMM Setup] Using CudaAsyncMemoryResource", flush=True)
+
+        # Create async resource with optional pool settings
+        if pool_size is not None:
+            # Align to 256 bytes as done in dask-cuda
+            initial_pool_size = (pool_size // 256) * 256
+            if rank == 0:
+                print(
+                    f"[RMM Setup] Initial pool size: {initial_pool_size / 1e9:.2f} GB",
+                    flush=True,
+                )
+            mr = rmm.mr.CudaAsyncMemoryResource(initial_pool_size=initial_pool_size)
+        else:
+            mr = rmm.mr.CudaAsyncMemoryResource()
+
+        # Set release threshold if specified
+        if hasattr(args, "rmm_release_threshold") and args.rmm_release_threshold:
+            if isinstance(args.rmm_release_threshold, str):
+                from dask.utils import parse_bytes
+
+                release_threshold = parse_bytes(args.rmm_release_threshold)
+            else:
+                release_threshold = args.rmm_release_threshold
+            mr.release_threshold = release_threshold
+            if rank == 0:
+                print(
+                    f"[RMM Setup] Release threshold: {release_threshold / 1e9:.2f} GB",
+                    flush=True,
+                )
+
+        rmm.mr.set_current_device_resource(mr)
+
+    elif pool_size is not None:
+        # Use PoolMemoryResource
+        if rank == 0:
+            print(
+                f"[RMM Setup] Using PoolMemoryResource with {pool_size / 1e9:.2f} GB",
+                flush=True,
+            )
+
+        # Use rmm.reinitialize similar to dask-cuda
+        rmm.reinitialize(
+            pool_allocator=True,
+            initial_pool_size=pool_size,
+            maximum_pool_size=None,  # No maximum limit
+        )
+    else:
+        # No specific RMM configuration, use default
+        if rank == 0:
+            print("[RMM Setup] Using default RMM configuration", flush=True)
+
+    # Enable RMM statistics if requested
+    if hasattr(args, "rmm_statistics") and args.rmm_statistics:
+        rmm.statistics.enable_statistics()
+        if rank == 0:
+            print("[RMM Setup] RMM statistics enabled", flush=True)
+
+
 def initialize_dask_cluster(run_config: RunConfig, args: argparse.Namespace):  # type: ignore[no-untyped-def]
     """
     Initialize a Dask distributed cluster.
@@ -1067,6 +1162,8 @@ def run_polars(
             run_config = dataclasses.replace(
                 run_config, cluster="rrun", n_workers=nranks
             )
+            # Set up RMM resources for this rank (similar to RMMPlugin for Dask)
+            _setup_rmm_for_rrun(args, rank)
     except ImportError:
         pass  # rapidsmpf not available
 
diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
index f9c669d734a2..227c0215d494 100644
--- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
+++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
@@ -230,6 +230,34 @@ def evaluate_pipeline(
         # Get the bootstrap-initialized communicator
         bootstrap_ctx = get_bootstrap_context()
 
+        # Configure memory limits for spilling (similar to Dask workers)
+        memory_available: MutableMapping[MemoryType, LimitAvailableMemory] | None = None
+        rrun_spill_device = config_options.executor.client_device_threshold
+        if rrun_spill_device > 0.0 and rrun_spill_device < 1.0:
+            total_memory = rmm.mr.available_device_memory()[1]
+            spill_threshold = int(total_memory * rrun_spill_device)
+            memory_available = {
+                MemoryType.DEVICE: LimitAvailableMemory(mr, limit=spill_threshold)
+            }
+
+            # Debug output on rank 0
+            from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import get_rank
+
+            rank = get_rank()
+            if rank == 0:
+                print(
+                    f"[RMM Config] Total device memory: {total_memory / 1e9:.2f} GB",
+                    flush=True,
+                )
+                print(
+                    f"[RMM Config] Spill threshold ({rrun_spill_device:.1%}): {spill_threshold / 1e9:.2f} GB",
+                    flush=True,
+                )
+                print(
+                    f"[RMM Config] Spill to pinned memory: {config_options.executor.spill_to_pinned_memory}",
+                    flush=True,
+                )
+
         options = Options(
             {
                 # By default, set the number of streaming threads to the max
@@ -251,16 +279,20 @@ def evaluate_pipeline(
         else:
             stream_pool = True  # Use stream pool for distributed execution
 
-        # Note: For rrun, we use the communicator from bootstrap_ctx
-        # The BufferResource is created but memory limits are not enforced
-        # in the same way as single-GPU mode
+        # Create BufferResource with memory limits for spilling
         br = BufferResource(
             mr,
             pinned_mr=pinned_mr,
-            memory_available=None,  # No memory limits for distributed
+            memory_available=memory_available,
             stream_pool=stream_pool,
         )
         rmpf_context_manager = Context(bootstrap_ctx, br, options)
+
+        # Enable RMM statistics for monitoring (similar to Dask)
+        try:
+            rmm.statistics.enable_statistics()
+        except Exception:
+            pass  # Statistics not available or already enabled
     else:
         # Using "single" mode.
         # Create a new local RapidsMPF context.

From 457cc4178e3e7b2cbb0f140a907edff3ae1cb7c2 Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Mon, 16 Feb 2026 13:40:08 -0800
Subject: [PATCH 5/9] Prevent resource reinitialization

---
 .../experimental/benchmarks/utils.py          |  21 +++
 .../experimental/rapidsmpf/bootstrap_ctx.py   |  98 ++++++++++++++
 .../experimental/rapidsmpf/core.py            | 125 +++++-------------
 3 files changed, 153 insertions(+), 91 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
index ef9d923de328..b8e9c3220958 100644
--- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
+++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
@@ -1164,6 +1164,27 @@ def run_polars(
             )
             # Set up RMM resources for this rank (similar to RMMPlugin for Dask)
             _setup_rmm_for_rrun(args, rank)
+
+            # Initialize the rrun worker context once (similar to
+            # bootstrap_dask_cluster -> dask_worker_setup for Dask).
+            # This creates BufferResource, ProgressThread, Statistics,
+            # and spill functions that are reused across all queries.
+            from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import (
+                setup_rrun_worker_context,
+            )
+
+            setup_rrun_worker_context(
+                spill_device=run_config.spill_device,
+                spill_to_pinned_memory=run_config.spill_to_pinned_memory,
+                oom_protection=run_config.rapidsmpf_oom_protection,
+                max_io_threads=run_config.max_io_threads,
+            )
+
+            # Enable RMM statistics (matches Dask's client.run(rmm.statistics.enable_statistics))
+            try:
+                rmm.statistics.enable_statistics()
+            except Exception:
+                pass
     except ImportError:
         pass  # rapidsmpf not available
 
diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
index 0b7f295deec1..fd6e458b726a 100644
--- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
+++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
@@ -9,6 +9,7 @@
 
 if TYPE_CHECKING:
     from rapidsmpf.bootstrap.bootstrap import Context
+    from rapidsmpf.integrations import WorkerContext
 
 try:
     import rapidsmpf.bootstrap as bootstrap
@@ -18,6 +19,17 @@
     BOOTSTRAP_AVAILABLE = False
 
 _global_context: Context | None = None
+_global_worker_context: WorkerContext | None = None
+
+
+class _RrunWorker:
+    """Sentinel object representing the rrun worker process for rmpf_worker_setup."""
+
+    def __init__(self, rank: int) -> None:
+        self._rank = rank
+
+    def __str__(self) -> str:
+        return f"rrun-rank-{self._rank}"
 
 
 def is_running_with_rrun() -> bool:
@@ -238,3 +250,89 @@ def get_nranks() -> int:
     if nranks is not None:
         return int(nranks)
     return 1
+
+
+def setup_rrun_worker_context(
+    *,
+    spill_device: float = 0.5,
+    spill_to_pinned_memory: bool = False,
+    oom_protection: bool = False,
+    max_io_threads: int = 2,
+) -> WorkerContext:
+    """
+    Initialize the rrun worker context once (singleton).
+
+    This calls ``rmpf_worker_setup()`` to create a ``WorkerContext`` with
+    properly configured ``BufferResource``, ``ProgressThread``, ``Statistics``,
+    and spill functions — matching the one-time setup that Dask performs via
+    ``bootstrap_dask_cluster()`` → ``dask_worker_setup()``.
+
+    Parameters
+    ----------
+    spill_device
+        Device memory threshold for spilling (fraction, 0.0-1.0).
+    spill_to_pinned_memory
+        Whether to spill to pinned host memory.
+    oom_protection
+        Whether to use managed memory fallback for OOM protection.
+    max_io_threads
+        Maximum number of IO threads.
+
+    Returns
+    -------
+    WorkerContext
+        The initialized worker context.
+    """
+    global _global_worker_context
+    if _global_worker_context is not None:
+        return _global_worker_context
+
+    from rapidsmpf.config import Options, get_environment_variables
+    from rapidsmpf.integrations.core import rmpf_worker_setup
+
+    comm = get_bootstrap_context()
+
+    options = Options(
+        {
+            "rrun_spill_device": str(spill_device),
+            "rrun_spill_to_pinned_memory": str(spill_to_pinned_memory),
+            "rrun_oom_protection": str(oom_protection),
+            "rrun_statistics": "False",
+            "rrun_print_statistics": "False",
+            "num_streaming_threads": str(max(max_io_threads, 1)),
+        }
+        | get_environment_variables()
+    )
+
+    rank = get_rank()
+    worker = _RrunWorker(rank)
+    _global_worker_context = rmpf_worker_setup(
+        worker, "rrun_", comm=comm, options=options
+    )
+
+    if rank == 0:
+        print("[rrun] Worker context initialized (one-time setup)", flush=True)
+
+    return _global_worker_context
+
+
+def get_rrun_worker_context() -> WorkerContext:
+    """
+    Get the initialized rrun worker context.
+
+    Returns
+    -------
+    WorkerContext
+        The rrun worker context.
+
+    Raises
+    ------
+    RuntimeError
+        If the worker context has not been initialized yet.
+    """
+    if _global_worker_context is None:
+        raise RuntimeError(
+            "rrun worker context not initialized. "
+            "Call setup_rrun_worker_context() first."
+        )
+    return _global_worker_context
diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
index 227c0215d494..481c9df623a2 100644
--- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
+++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
@@ -117,15 +117,38 @@ def evaluate_logical_plan(
         if (
             config_options.executor.cluster == "rrun" or is_rrun
         ):  # pragma: no cover; block depends on executor type and rrun cluster
-            # SPMD execution: All ranks execute, only rank 0 returns result
-            result, metadata_collector = evaluate_pipeline(
-                ir,
-                partition_info,
-                config_options,
-                stats,
-                collective_id_map,
-                collect_metadata=collect_metadata,
+            # SPMD execution: All ranks execute, only rank 0 returns result.
+            # Get the pre-initialized worker context (set up once in run_polars)
+            # and create a lightweight streaming Context from it — same pattern
+            # as Dask's _evaluate_pipeline_dask.
+            from rapidsmpf.config import Options as RmpfOptions
+            from rapidsmpf.config import get_environment_variables
+
+            from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import (
+                get_rrun_worker_context,
+            )
+
+            worker_ctx = get_rrun_worker_context()
+            options = RmpfOptions(
+                {
+                    "num_streaming_threads": str(
+                        max(config_options.executor.max_io_threads, 1)
+                    )
+                }
+                | get_environment_variables()
             )
+            with Context(
+                worker_ctx.comm, worker_ctx.br, options, worker_ctx.statistics
+            ) as rmpf_context:
+                result, metadata_collector = evaluate_pipeline(
+                    ir,
+                    partition_info,
+                    config_options,
+                    stats,
+                    collective_id_map,
+                    rmpf_context,
+                    collect_metadata=collect_metadata,
+                )
 
             # Only rank 0 returns result to caller
             rank = get_rank()
@@ -206,93 +229,13 @@ def evaluate_pipeline(
     _initial_mr: Any = None
     stream_pool: CudaStreamPool | bool = False
 
-    # Check if running with rrun
-    from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import (
-        get_bootstrap_context,
-        is_running_with_rrun,
-    )
-
-    is_rrun = is_running_with_rrun()
-
     if rmpf_context is not None:
-        # Using "distributed" mode (Dask).
-        # Always use the RapidsMPF stream pool for now.
+        # Using "distributed" or "rrun" mode.
+        # The caller has already set up the Context (from a pre-initialized
+        # WorkerContext). Always use the RapidsMPF stream pool for now.
         br = rmpf_context.br()
         stream_pool = True
         rmpf_context_manager = contextlib.nullcontext(rmpf_context)
-    elif is_rrun and rmpf_context is None:
-        # Using "rrun" mode - initialize from bootstrap context
-        # Create a new distributed RapidsMPF context using the bootstrap communicator
-        _original_mr = rmm.mr.get_current_device_resource()
-        mr = RmmResourceAdaptor(_original_mr)
-        rmm.mr.set_current_device_resource(mr)
-
-        # Get the bootstrap-initialized communicator
-        bootstrap_ctx = get_bootstrap_context()
-
-        # Configure memory limits for spilling (similar to Dask workers)
-        memory_available: MutableMapping[MemoryType, LimitAvailableMemory] | None = None
-        rrun_spill_device = config_options.executor.client_device_threshold
-        if rrun_spill_device > 0.0 and rrun_spill_device < 1.0:
-            total_memory = rmm.mr.available_device_memory()[1]
-            spill_threshold = int(total_memory * rrun_spill_device)
-            memory_available = {
-                MemoryType.DEVICE: LimitAvailableMemory(mr, limit=spill_threshold)
-            }
-
-            # Debug output on rank 0
-            from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import get_rank
-
-            rank = get_rank()
-            if rank == 0:
-                print(
-                    f"[RMM Config] Total device memory: {total_memory / 1e9:.2f} GB",
-                    flush=True,
-                )
-                print(
-                    f"[RMM Config] Spill threshold ({rrun_spill_device:.1%}): {spill_threshold / 1e9:.2f} GB",
-                    flush=True,
-                )
-                print(
-                    f"[RMM Config] Spill to pinned memory: {config_options.executor.spill_to_pinned_memory}",
-                    flush=True,
-                )
-
-        options = Options(
-            {
-                # By default, set the number of streaming threads to the max
-                # number of IO threads. The user may override this with an
-                # environment variable (i.e. RAPIDSMPF_NUM_STREAMING_THREADS)
-                "num_streaming_threads": str(
-                    max(config_options.executor.max_io_threads, 1)
-                )
-            }
-            | get_environment_variables()
-        )
-        pinned_mr = (
-            PinnedMemoryResource.make_if_available()
-            if config_options.executor.spill_to_pinned_memory
-            else None
-        )
-        if isinstance(config_options.cuda_stream_policy, CUDAStreamPoolConfig):
-            stream_pool = config_options.cuda_stream_policy.build()
-        else:
-            stream_pool = True  # Use stream pool for distributed execution
-
-        # Create BufferResource with memory limits for spilling
-        br = BufferResource(
-            mr,
-            pinned_mr=pinned_mr,
-            memory_available=memory_available,
-            stream_pool=stream_pool,
-        )
-        rmpf_context_manager = Context(bootstrap_ctx, br, options)
-
-        # Enable RMM statistics for monitoring (similar to Dask)
-        try:
-            rmm.statistics.enable_statistics()
-        except Exception:
-            pass  # Statistics not available or already enabled
     else:
         # Using "single" mode.
         # Create a new local RapidsMPF context.

From c192c4f72b3119bfa0cd0dfaa46d0f5cd051bf07 Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Tue, 17 Feb 2026 00:43:14 -0800
Subject: [PATCH 6/9] Fix RMM memory resource override

---
 python/cudf_polars/cudf_polars/callback.py                    | 4 ++--
 .../cudf_polars/cudf_polars/experimental/benchmarks/utils.py  | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/callback.py b/python/cudf_polars/cudf_polars/callback.py
index 17b8d691461a..23865e5b9d7a 100644
--- a/python/cudf_polars/cudf_polars/callback.py
+++ b/python/cudf_polars/cudf_polars/callback.py
@@ -331,8 +331,8 @@ def execute_with_cudf(
         if (
             memory_resource is None
             and translator.config_options.executor.name == "streaming"
-            and translator.config_options.executor.cluster == "distributed"
-        ):  # pragma: no cover; Requires distributed cluster
+            and translator.config_options.executor.cluster in ("distributed", "rrun")
+        ):  # pragma: no cover; Requires distributed cluster or rrun
             memory_resource = rmm.mr.get_current_device_resource()
         if len(ir_translation_errors):
             # TODO: Display these errors in user-friendly way.
diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
index b8e9c3220958..2f3380e87ca4 100644
--- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
+++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
@@ -1209,7 +1209,7 @@ def run_polars(
         engine = pl.GPUEngine(
             raise_on_fail=True,
             memory_resource=rmm.mr.CudaAsyncMemoryResource()
-            if run_config.rmm_async
+            if run_config.rmm_async and not is_rrun
             else None,
             cuda_stream_policy=run_config.stream_policy,
             executor=run_config.executor,

From 50f9be4a203bc39070d5b3f0682e53dc8ed53cd0 Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Wed, 18 Feb 2026 13:31:39 -0800
Subject: [PATCH 7/9] Cleanup

---
 .../experimental/benchmarks/utils.py          | 54 +------------------
 .../cudf_polars/experimental/parallel.py      |  1 -
 .../experimental/rapidsmpf/bootstrap_ctx.py   | 53 +++++++++---------
 .../experimental/rapidsmpf/core.py            |  6 +--
 4 files changed, 30 insertions(+), 84 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
index 2f3380e87ca4..24b9f10760f4 100644
--- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
+++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
@@ -544,10 +544,6 @@ def _setup_rmm_for_rrun(args: argparse.Namespace, rank: int) -> None:  # type: i
     """
     import rmm
 
-    if rank == 0:
-        print("[RMM Setup] Configuring RMM resources for rrun mode", flush=True)
-
-    # Parse pool size if specified
     pool_size = None
     if args.rmm_pool_size is not None:
         if isinstance(args.rmm_pool_size, str):
@@ -559,26 +555,13 @@ def _setup_rmm_for_rrun(args: argparse.Namespace, rank: int) -> None:  # type: i
             total_memory = rmm.mr.available_device_memory()[1]
             pool_size = int(total_memory * args.rmm_pool_size)
 
-    # Set up RMM resource based on configuration
     if args.rmm_async:
-        # Use CudaAsyncMemoryResource (stream-ordered allocator)
-        if rank == 0:
-            print("[RMM Setup] Using CudaAsyncMemoryResource", flush=True)
-
-        # Create async resource with optional pool settings
         if pool_size is not None:
-            # Align to 256 bytes as done in dask-cuda
             initial_pool_size = (pool_size // 256) * 256
-            if rank == 0:
-                print(
-                    f"[RMM Setup] Initial pool size: {initial_pool_size / 1e9:.2f} GB",
-                    flush=True,
-                )
             mr = rmm.mr.CudaAsyncMemoryResource(initial_pool_size=initial_pool_size)
         else:
             mr = rmm.mr.CudaAsyncMemoryResource()
 
-        # Set release threshold if specified
         if hasattr(args, "rmm_release_threshold") and args.rmm_release_threshold:
             if isinstance(args.rmm_release_threshold, str):
                 from dask.utils import parse_bytes
@@ -587,38 +570,17 @@ def _setup_rmm_for_rrun(args: argparse.Namespace, rank: int) -> None:  # type: i
             else:
                 release_threshold = args.rmm_release_threshold
             mr.release_threshold = release_threshold
-            if rank == 0:
-                print(
-                    f"[RMM Setup] Release threshold: {release_threshold / 1e9:.2f} GB",
-                    flush=True,
-                )
 
         rmm.mr.set_current_device_resource(mr)
-
     elif pool_size is not None:
-        # Use PoolMemoryResource
-        if rank == 0:
-            print(
-                f"[RMM Setup] Using PoolMemoryResource with {pool_size / 1e9:.2f} GB",
-                flush=True,
-            )
-
-        # Use rmm.reinitialize similar to dask-cuda
         rmm.reinitialize(
             pool_allocator=True,
             initial_pool_size=pool_size,
             maximum_pool_size=None,  # No maximum limit
         )
-    else:
-        # No specific RMM configuration, use default
-        if rank == 0:
-            print("[RMM Setup] Using default RMM configuration", flush=True)
 
-    # Enable RMM statistics if requested
     if hasattr(args, "rmm_statistics") and args.rmm_statistics:
         rmm.statistics.enable_statistics()
-        if rank == 0:
-            print("[RMM Setup] RMM statistics enabled", flush=True)
 
 
 def initialize_dask_cluster(run_config: RunConfig, args: argparse.Namespace):  # type: ignore[no-untyped-def]
@@ -653,13 +615,9 @@ def initialize_dask_cluster(run_config: RunConfig, args: argparse.Namespace):  #
         if is_running_with_rrun():
             rank = get_rank()
             nranks = get_nranks()
-            if rank == 0:
-                print(
-                    f"[rrun] Detected rrun execution environment with {nranks} ranks"
-                )
-            return None  # No Dask client needed for rrun
+            return None  # No Dask client is used with rrun
     except ImportError:
-        pass  # rapidsmpf not available, continue with normal path
+        pass
 
     if run_config.cluster != "distributed":
         return None
@@ -1162,13 +1120,8 @@ def run_polars(
             run_config = dataclasses.replace(
                 run_config, cluster="rrun", n_workers=nranks
             )
-            # Set up RMM resources for this rank (similar to RMMPlugin for Dask)
             _setup_rmm_for_rrun(args, rank)
 
-            # Initialize the rrun worker context once (similar to
-            # bootstrap_dask_cluster -> dask_worker_setup for Dask).
-            # This creates BufferResource, ProgressThread, Statistics,
-            # and spill functions that are reused across all queries.
             from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import (
                 setup_rrun_worker_context,
             )
@@ -1180,7 +1133,6 @@ def run_polars(
                 max_io_threads=run_config.max_io_threads,
             )
 
-            # Enable RMM statistics (matches Dask's client.run(rmm.statistics.enable_statistics))
             try:
                 rmm.statistics.enable_statistics()
             except Exception:
@@ -1245,7 +1197,6 @@ def run_polars(
 
             # In rrun mode, result is None for non-root ranks
             if is_rrun and result is None:
-                # Non-root ranks: skip result processing but record timing
                 t1 = time.monotonic()
                 record = Record(
                     query=q_id, iteration=i, duration=t1 - t0, shuffle_stats=None
@@ -1350,7 +1301,6 @@ def sort_key(x: dict) -> tuple[int, int]:
 
             run_config.records[query_id] = new_records
 
-    # Only rank 0 should print summaries and write output
     if not is_rrun or rank == 0:
         if args.summarize:
             run_config.summarize()
diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/experimental/parallel.py
index 2d292651e5df..ae760c5a7f75 100644
--- a/python/cudf_polars/cudf_polars/experimental/parallel.py
+++ b/python/cudf_polars/cudf_polars/experimental/parallel.py
@@ -173,7 +173,6 @@ def get_scheduler(config_options: ConfigOptions) -> Any:
 
     cluster = config_options.executor.cluster
 
-    # Check if running with rrun
     from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import is_running_with_rrun
 
     if is_running_with_rrun() or cluster == "rrun":
diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
index fd6e458b726a..c7389e1516ff 100644
--- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
+++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 """Bootstrap context management for rrun execution."""
 
@@ -34,7 +34,7 @@ def __str__(self) -> str:
 
 def is_running_with_rrun() -> bool:
     """
-    Check if running under rrun.
+    Check if running with rrun.
 
     Returns
     -------
@@ -47,9 +47,9 @@ def is_running_with_rrun() -> bool:
     return bootstrap.is_running_with_rrun()
 
 
-def is_running_under_slurm() -> bool:
+def is_running_with_slurm() -> bool:
     """
-    Check if running under Slurm.
+    Check if running with Slurm.
 
     Returns
     -------
@@ -75,6 +75,9 @@ def _detect_backend_type():
 
     Notes
     -----
+    The necessity of this function should be revisited. In an ideal case
+    BackendType.AUTO should suffice.
+
     Detection logic:
     1. If RAPIDSMPF_COORD_DIR is set -> use AUTO (will choose FILE)
     2. If running under Slurm without COORD_DIR -> set up FILE backend with temp dir
@@ -89,7 +92,7 @@ def _detect_backend_type():
 
     # If running under Slurm without COORD_DIR, we need to set one up
     # because the SLURM backend is not available in Python bindings yet
-    if is_running_under_slurm():
+    if is_running_with_slurm():
         # Check if SLURM backend is available in the Python bindings
         if hasattr(bootstrap.BackendType, "SLURM"):
             print(
@@ -132,7 +135,7 @@ def _detect_backend_type():
 
 def get_bootstrap_context() -> Context:
     """
-    Get or initialize bootstrap context (singleton).
+    Get or initialize bootstrap context.
 
     Returns
     -------
@@ -157,10 +160,8 @@ def get_bootstrap_context() -> Context:
                 "Use 'rrun -n <nranks> python ...' to launch with rrun."
             )
 
-        # Detect and use appropriate backend
         backend_type = _detect_backend_type()
 
-        # Debug: print environment info on rank 0
         rank = get_rank()
         if rank == 0:
             print(f"[Bootstrap] Backend type: {backend_type}", flush=True)
@@ -178,7 +179,6 @@ def get_bootstrap_context() -> Context:
             )
 
         try:
-            # Initialize the bootstrap context with detected backend
             _global_context = bootstrap.create_ucxx_comm(backend_type)
         except RuntimeError as e:
             # Provide helpful error message
@@ -214,9 +214,14 @@ def get_rank() -> int:
     -------
     int
         The rank of the current process (0 if not running under rrun).
+
+    Raises
+    ------
+    RuntimeError
+        If not running with rrun or rank could not be determined.
     """
     if not is_running_with_rrun():
-        return 0
+        raise RuntimeError("Not running with rrun.")
     # Read directly from environment variable for efficiency
     # Try RAPIDSMPF_RANK first, fall back to SLURM_PROCID for Slurm
     rank = os.environ.get("RAPIDSMPF_RANK")
@@ -226,7 +231,7 @@ def get_rank() -> int:
     rank = os.environ.get("SLURM_PROCID")
     if rank is not None:
         return int(rank)
-    return 0
+    raise RuntimeError("Could not determine rank.")
 
 
 def get_nranks() -> int:
@@ -237,19 +242,13 @@ def get_nranks() -> int:
     -------
     int
         The total number of ranks (1 if not running under rrun).
+
+    Raises
+    ------
+    RuntimeError
+        If not running with rrun or number of ranks could not be determined.
     """
-    if not is_running_with_rrun():
-        return 1
-    # Read directly from environment variable for efficiency
-    # Try RAPIDSMPF_NRANKS first, fall back to SLURM_NPROCS/SLURM_NTASKS for Slurm
-    nranks = os.environ.get("RAPIDSMPF_NRANKS")
-    if nranks is not None:
-        return int(nranks)
-    # Fall back to Slurm env vars
-    nranks = os.environ.get("SLURM_NPROCS") or os.environ.get("SLURM_NTASKS")
-    if nranks is not None:
-        return int(nranks)
-    return 1
+    return bootstrap.get_nranks()
 
 
 def setup_rrun_worker_context(
@@ -260,12 +259,12 @@ def setup_rrun_worker_context(
     max_io_threads: int = 2,
 ) -> WorkerContext:
     """
-    Initialize the rrun worker context once (singleton).
+    Initialize the rrun worker context once.
 
     This calls ``rmpf_worker_setup()`` to create a ``WorkerContext`` with
     properly configured ``BufferResource``, ``ProgressThread``, ``Statistics``,
-    and spill functions — matching the one-time setup that Dask performs via
-    ``bootstrap_dask_cluster()`` → ``dask_worker_setup()``.
+    and spill functions, matching the one-time setup that Dask performs via
+    ``bootstrap_dask_cluster()`` and ``dask_worker_setup()``.
 
     Parameters
     ----------
@@ -311,7 +310,7 @@ def setup_rrun_worker_context(
     )
 
     if rank == 0:
-        print("[rrun] Worker context initialized (one-time setup)", flush=True)
+        print("[rrun] Worker context initialized", flush=True)
 
     return _global_worker_context
 
diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
index 481c9df623a2..e487af8c5791 100644
--- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
+++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py
@@ -95,7 +95,6 @@ def evaluate_logical_plan(
     assert config_options.executor.name == "streaming", "Executor must be streaming"
     assert config_options.executor.runtime == "rapidsmpf", "Runtime must be rapidsmpf"
 
-    # Check if running with rrun
     from cudf_polars.experimental.rapidsmpf.bootstrap_ctx import (
         get_rank,
         is_running_with_rrun,
@@ -119,8 +118,8 @@ def evaluate_logical_plan(
         ):  # pragma: no cover; block depends on executor type and rrun cluster
             # SPMD execution: All ranks execute, only rank 0 returns result.
             # Get the pre-initialized worker context (set up once in run_polars)
-            # and create a lightweight streaming Context from it — same pattern
-            # as Dask's _evaluate_pipeline_dask.
+            # and create a lightweight streaming Context from it using the same
+            # pattern as Dask's _evaluate_pipeline_dask.
             from rapidsmpf.config import Options as RmpfOptions
             from rapidsmpf.config import get_environment_variables
 
@@ -155,7 +154,6 @@ def evaluate_logical_plan(
             if rank == 0:
                 return result, metadata_collector
             else:
-                # Non-root ranks return None
                 return None, None
 
         elif (

From 95b87a67e6f4b6a02cc9bbb86a2d21ab0daf835b Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Wed, 18 Feb 2026 14:02:25 -0800
Subject: [PATCH 8/9] Update to use new Python API

---
 .../experimental/rapidsmpf/bootstrap_ctx.py   | 42 +++++++------------
 1 file changed, 15 insertions(+), 27 deletions(-)

diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
index c7389e1516ff..89945ac0bc93 100644
--- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
+++ b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/bootstrap_ctx.py
@@ -54,14 +54,11 @@ def is_running_with_slurm() -> bool:
     Returns
     -------
     bool
-        True if Slurm environment variables are detected.
+        True if running under Slurm with PMIx.
     """
-    # Check for Slurm environment variables
-    return (
-        "SLURM_JOB_ID" in os.environ
-        or "SLURM_PROCID" in os.environ
-        or "PMIX_NAMESPACE" in os.environ
-    )
+    if not BOOTSTRAP_AVAILABLE:
+        return False
+    return bootstrap.is_running_with_slurm()
 
 
 def _detect_backend_type():
@@ -145,19 +142,20 @@ def get_bootstrap_context() -> Context:
     Raises
     ------
     RuntimeError
-        If rapidsmpf.bootstrap is not available or not running under rrun.
+        If rapidsmpf.bootstrap is not available or not running under a
+        bootstrap launcher (rrun).
     """
     global _global_context
     if _global_context is None:
         if not BOOTSTRAP_AVAILABLE:
             raise RuntimeError(
                 "rapidsmpf.bootstrap not available. "
-                "Please install rapidsmpf to use rrun execution."
+                "Please install rapidsmpf to use rrun or Slurm execution."
             )
-        if not is_running_with_rrun():
+        if not (is_running_with_rrun() or is_running_with_slurm()):
             raise RuntimeError(
-                "Not running under rrun (RAPIDSMPF_RANK environment variable not set). "
-                "Use 'rrun -n <nranks> python ...' to launch with rrun."
+                "Not running under a bootstrap launcher "
+                "(RAPIDSMPF_RANK / Slurm PMIx not detected)."
             )
 
         backend_type = _detect_backend_type()
@@ -208,30 +206,20 @@ def get_bootstrap_context() -> Context:
 
 def get_rank() -> int:
     """
-    Get current rank.
+    Get current bootstrap rank.
 
     Returns
     -------
     int
-        The rank of the current process (0 if not running under rrun).
+        The rank of the current process.
 
     Raises
     ------
     RuntimeError
-        If not running with rrun or rank could not be determined.
+        If not running with a bootstrap launcher (rrun) or rank
+        could not be determined.
     """
-    if not is_running_with_rrun():
-        raise RuntimeError("Not running with rrun.")
-    # Read directly from environment variable for efficiency
-    # Try RAPIDSMPF_RANK first, fall back to SLURM_PROCID for Slurm
-    rank = os.environ.get("RAPIDSMPF_RANK")
-    if rank is not None:
-        return int(rank)
-    # Fall back to Slurm env vars
-    rank = os.environ.get("SLURM_PROCID")
-    if rank is not None:
-        return int(rank)
-    raise RuntimeError("Could not determine rank.")
+    return bootstrap.get_rank()
 
 
 def get_nranks() -> int:

From fb4c3a401b55127be62a87e835fcc63f465f5ff8 Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Wed, 18 Feb 2026 14:12:43 -0800
Subject: [PATCH 9/9] Do not call initialize_dask_cluster when running with
 rrun

---
 .../cudf_polars/cudf_polars/experimental/benchmarks/utils.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
index 24b9f10760f4..c201047255a7 100644
--- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
+++ b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py
@@ -1140,7 +1140,10 @@ def run_polars(
     except ImportError:
         pass  # rapidsmpf not available
 
-    client = initialize_dask_cluster(run_config, args)
+    if is_rrun:
+        client = None
+    else:
+        client = initialize_dask_cluster(run_config, args)
 
     # Update n_workers from the actual cluster when using scheduler file/address
     if client is not None: