[WIP] Rapidsmpf Stats API changes

python/cudf_polars/cudf_polars/engine/dask.py

@@ -109,6 +109,17 @@ def dask_setup(nanny: distributed.Nanny) -> None:
 class _WorkerContext:
     """Per-worker GPU resources stored on each Dask worker."""
 
+    # Statistics instance currently shared by the worker's progress thread
+    # and its streaming Context. Built in _setup_root / _setup_worker and
+    # replaced in-place in _reset_worker by installing a fresh Statistics
+    # on ``progress`` (the communicator forwards ``statistics()`` to its
+    # progress thread, so the swap automatically updates the comm's view).
+    stats: Statistics | None
+    # Owning reference to the worker's ProgressThread. Kept here so
+    # _reset_worker can call ``progress.set_statistics(...)`` without
+    # routing through ``comm.progress_thread`` (which would also work but
+    # adds an extra hop).
+    progress: ProgressThread | None
     comm: Communicator | None
     ctx: Context | None
     py_executor: ThreadPoolExecutor | None
@@ -158,17 +169,28 @@ def _setup_root(
     memory_resource_config = memory_resource_config or MemoryResourceConfig.default()
     base_mr = memory_resource_config.create_memory_resource()
     mr = RmmResourceAdaptor(base_mr)
+    # Build the single Statistics instance shared by the progress thread and
+    # the streaming Context this worker will create in :func:`_setup_worker`.
+    stats = Statistics.from_options(options)
+    progress = ProgressThread(stats)
     comm = new_communicator(
         nranks=nranks,
         ucx_worker=None,
         root_ucxx_address=None,
         options=options,
-        progress_thread=ProgressThread(),
+        progress_thread=progress,
     )
     setattr(
         dask_worker,
         f"_cudf_polars_mp_context_{uid}",
-        _WorkerContext(comm=comm, ctx=None, py_executor=None, mr=mr),
+        _WorkerContext(
+            stats=stats,
+            progress=progress,
+            comm=comm,
+            ctx=None,
+            py_executor=None,
+            mr=mr,
+        ),
     )
     return get_root_ucxx_address(comm)
 
@@ -225,23 +247,32 @@ def _setup_worker(
         )
         base_mr = memory_resource_config.create_memory_resource()
         mr = RmmResourceAdaptor(base_mr)
+        # Build the single Statistics instance shared by the progress thread
+        # and the streaming Context built below.
+        stats = Statistics.from_options(options)
+        progress = ProgressThread(stats)
         root_addr = ucx_api.UCXAddress.create_from_buffer(root_ucxx_address_as_bytes)
         comm = new_communicator(
             nranks=nranks,
             ucx_worker=None,
             root_ucxx_address=root_addr,
             options=options,
-            progress_thread=ProgressThread(),
+            progress_thread=progress,
         )
     else:
-        # Root worker: comm and mr were created in _setup_root.
-        mr = mp_ctx.mr
-        assert mp_ctx.comm is not None
-        comm = mp_ctx.comm
+        # Root worker: comm, mr, stats, and progress were created in _setup_root.
+        stats, progress, comm, mr = (
+            mp_ctx.stats,
+            mp_ctx.progress,
+            mp_ctx.comm,
+            mp_ctx.mr,
+        )
+        assert stats is not None
+        assert progress is not None
+        assert comm is not None
 
     barrier(comm)
-    statistics = Statistics.from_options(options)
-    ctx = Context.from_options(comm.logger, mr, options, statistics)
+    ctx = Context.from_options(comm.logger, mr, options, stats)
     # Set the current RMM device resource so all temporary allocations
     # in libcudf also use the same memory resource.
     rmm.mr.set_current_device_resource(ctx.br().device_mr)
@@ -255,7 +286,14 @@ def _setup_worker(
     setattr(
         dask_worker,
         attr,
-        _WorkerContext(comm=comm, ctx=ctx, py_executor=py_executor, mr=mr),
+        _WorkerContext(
+            stats=stats,
+            progress=progress,
+            comm=comm,
+            ctx=ctx,
+            py_executor=py_executor,
+            mr=mr,
+        ),
     )
 
 
@@ -287,8 +325,10 @@ def _teardown_worker(
             if mp_ctx.ctx is not None:
                 mp_ctx.ctx.shutdown()
         finally:
-            mp_ctx.ctx = None
+            mp_ctx.stats = None
+            mp_ctx.progress = None
             mp_ctx.comm = None
+            mp_ctx.ctx = None
             mp_ctx.mr = None
             delattr(dask_worker, attr)
 
@@ -319,6 +359,7 @@ def _reset_worker(
     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.progress is not None
     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.
@@ -330,9 +371,12 @@ def _reset_worker(
     mp_ctx.ctx.shutdown()
     mp_ctx.ctx = None
     options = Options.deserialize(rapidsmpf_options_as_bytes)
-    statistics = Statistics.from_options(options)
+    # Build a fresh Statistics from the new options and set it on the
+    # progress thread. Communicator will infer from progress thread.
+    mp_ctx.stats = Statistics.from_options(options)
+    mp_ctx.progress.set_statistics(mp_ctx.stats)
     mp_ctx.ctx = Context.from_options(
-        mp_ctx.comm.logger, mp_ctx.mr, options, statistics
+        mp_ctx.comm.logger, mp_ctx.mr, options, mp_ctx.stats
     )
     rmm.mr.set_current_device_resource(mp_ctx.ctx.br().device_mr)
 

python/cudf_polars/cudf_polars/engine/default_singleton_engine.py

@@ -15,6 +15,7 @@
     new_communicator as single_communicator,
 )
 from rapidsmpf.progress_thread import ProgressThread
+from rapidsmpf.statistics import Statistics
 
 from cudf_polars.engine.core import (
     resolve_rapidsmpf_options,
@@ -141,9 +142,10 @@ def _build_engine() -> DefaultSingletonEngine:
     """
     with _state.lock:
         try:
+            options = resolve_rapidsmpf_options(None)
             comm = single_communicator(
-                progress_thread=ProgressThread(),
-                options=resolve_rapidsmpf_options(None),
+                progress_thread=ProgressThread(Statistics.from_options(options)),
+                options=options,
             )
             instance = DefaultSingletonEngine(comm=comm)
             assert instance.nranks == 1

python/cudf_polars/cudf_polars/engine/ray.py

@@ -192,6 +192,9 @@ def __init__(
             max_workers=num_py_executors,
             thread_name_prefix="ray-executor",
         )
+        # Same instance shared by all resources.
+        self._stats: Statistics | None = None
+        self._progress: ProgressThread | None = None
         self._comm: Communicator | None = None
         self._ctx: Context | None = None
 
@@ -207,12 +210,14 @@ def setup_root(self) -> bytes:
         -------
         Serialized UCXX root address for communicator bootstrap.
         """
+        self._stats = Statistics.from_options(self._rapidsmpf_options)
+        self._progress = ProgressThread(self._stats)
         self._comm = new_communicator(
             nranks=self._nranks,
             ucx_worker=None,
             root_ucxx_address=None,
             options=self._rapidsmpf_options,
-            progress_thread=ProgressThread(),
+            progress_thread=self._progress,
         )
         return get_root_ucxx_address(self._comm)
 
@@ -231,6 +236,10 @@ def setup_worker(self, root_ucxx_address_as_bytes: bytes) -> None:
             Serialized UCXX root address returned by :meth:`setup_root`.
         """
         if self._comm is None:
+            # Non-root actor: ``setup_root`` did not run here, so build the
+            # shared Statistics instance now before creating the communicator.
+            self._stats = Statistics.from_options(self._rapidsmpf_options)
+            self._progress = ProgressThread(self._stats)
             root_ucxx_address = ucx_api.UCXAddress.create_from_buffer(
                 root_ucxx_address_as_bytes
             )
@@ -239,14 +248,13 @@ def setup_worker(self, root_ucxx_address_as_bytes: bytes) -> None:
                 ucx_worker=None,
                 root_ucxx_address=root_ucxx_address,
                 options=self._rapidsmpf_options,
-                progress_thread=ProgressThread(),
+                progress_thread=self._progress,
             )
+        assert self._stats is not None
+        assert self._progress is not None
         barrier(self._comm)
         self._ctx = Context.from_options(
-            self._comm.logger,
-            self._mr,
-            self._rapidsmpf_options,
-            self._rapidsmpf_statistics,
+            self._comm.logger, self._mr, self._rapidsmpf_options, self._stats
         )
         # Set the current RMM device resource so all temporary allocations
         # in libcudf also use the same memory resource.
@@ -267,18 +275,19 @@ def reset(self, *, rapidsmpf_options_as_bytes: bytes) -> None:
         if self._ctx is None:
             raise RuntimeError("reset() requires setup_worker() to have run")
         assert self._comm is not None
+        assert self._progress is not None
         # Collective: all ranks idle before any rank tears down its Context.
         if self._comm.nranks > 1:
             barrier(self._comm)
         self._ctx.shutdown()
         self._ctx = None
         self._rapidsmpf_options = Options.deserialize(rapidsmpf_options_as_bytes)
-        self._rapidsmpf_statistics = Statistics.from_options(self._rapidsmpf_options)
+        # Build a fresh Statistics from the new options and set it on the
+        # progress thread. Communicator will infer from progress thread.
+        self._stats = Statistics.from_options(self._rapidsmpf_options)
+        self._progress.set_statistics(self._stats)
         self._ctx = Context.from_options(
-            self._comm.logger,
-            self._mr,
-            self._rapidsmpf_options,
-            self._rapidsmpf_statistics,
+            self._comm.logger, self._mr, self._rapidsmpf_options, self._stats
         )
 
     def shutdown(self) -> None:
@@ -298,6 +307,8 @@ def shutdown(self) -> None:
             self._ctx = None
             self._comm = None
             self._mr = None
+            self._progress = None
+            self._stats = None
             ray.actor.exit_actor()
 
     def get_info(self) -> ClusterInfo:

python/cudf_polars/cudf_polars/engine/spmd.py

@@ -353,21 +353,26 @@ def __init__(
         base_mr = mr_config.create_memory_resource()
         mr = RmmResourceAdaptor(base_mr)
         if comm is None:
+            self._stats: Statistics = Statistics.from_options(self.rapidsmpf_options)
+            progress = ProgressThread(self._stats)
             if bootstrap.is_running_with_rrun():
                 comm = bootstrap.create_ucxx_comm(
-                    progress_thread=ProgressThread(),
+                    progress_thread=progress,
                     type=bootstrap.BackendType.AUTO,
                     options=self.rapidsmpf_options,
                 )
             else:
                 comm = single_communicator(
-                    progress_thread=ProgressThread(),
+                    progress_thread=progress,
                     options=self.rapidsmpf_options,
                 )
-        # else: caller-provided comm; the caller retains ownership
+        else:  # Caller-provided comm
+            progress = comm.progress_thread
+            self._stats = progress.statistics
 
         self._mr: RmmResourceAdaptor = mr
         self._comm: Communicator | None = comm
+        self._progress: ProgressThread | None = progress
         self._ctx: Context | None = None
         self._py_executor: ThreadPoolExecutor | None = None
         exit_stack = contextlib.ExitStack()
@@ -376,13 +381,8 @@ def __init__(
 
             # Register `_cleanup_ctx`, which shuts down whatever `self._ctx` points
             # to at engine shutdown time, i.e. the `Context` from the latest reset.
-            if self.rapidsmpf_options is not None:
-                statistics = Statistics.from_options(self.rapidsmpf_options)
-            else:
-                statistics = None
-
             self._ctx = Context.from_options(
-                comm.logger, mr, self.rapidsmpf_options, statistics
+                comm.logger, mr, self.rapidsmpf_options, self._stats
             )
             exit_stack.callback(self._cleanup_ctx)
 
@@ -477,6 +477,7 @@ def _reset(
         if self._ctx is None:
             raise RuntimeError("Cannot reset a shut-down engine")
         assert self._comm is not None
+        assert self._progress is not None
         super()._reset(
             rapidsmpf_options=rapidsmpf_options,
             executor_options=executor_options,
@@ -493,14 +494,15 @@ def _reset(
         # Context (the test driver's main thread). The per-engine RMM
         # resource is kept alive across resets, see :meth:`_cleanup_ctx`.
         self._ctx.shutdown()
-
-        if rapidsmpf_options is not None:
-            statistics = Statistics.from_options(rapidsmpf_options)
-        else:
-            statistics = None
-
+        # Build a fresh Statistics from the new options and install it on the
+        # progress thread. The communicator's ``statistics()`` forwards to its
+        # progress thread (see ``Communicator::statistics()`` in C++), so this
+        # single swap atomically updates the comm's view too. We then hand the
+        # new instance to the rebuilt Context.
+        self._stats = Statistics.from_options(rapidsmpf_options)
+        self._progress.set_statistics(self._stats)
         self._ctx = Context.from_options(
-            self._comm.logger, self._mr, rapidsmpf_options, statistics
+            self._comm.logger, self._mr, rapidsmpf_options, self._stats
         )
 
         # Re-run ``StreamingEngine.__init__`` on the existing instance to
@@ -652,6 +654,7 @@ def shutdown(self) -> None:
         # Clear the references only after shutdown completes.
         super().shutdown()
         self._comm = None
+        self._progress = None
         self._ctx = None
         self._py_executor = None
 

python/cudf_polars/tests/streaming/test_spmd.py

@@ -329,6 +329,23 @@ def test_reset_keeps_comm_alive(comm: Communicator) -> None:
         assert sorted(result["a"].to_list()) == [1, 2, 3]
 
 
+def test_reset_propagates_new_stats_through_comm(comm: Communicator) -> None:
+    """After ``_reset``, ``Communicator.statistics`` forwards to the swapped instance."""
+    from rapidsmpf.config import Options
+
+    with SPMDEngine(comm=comm) as engine:
+        # Force the post-reset Statistics to be enabled so writes are
+        # observable through ``add_stat`` / ``list_stat_names``.
+        engine._reset(rapidsmpf_options=Options({"statistics": "True"}))
+        assert engine._stats.enabled
+
+        # A stat added directly to ``engine._stats`` is visible through
+        # ``engine.comm.statistics``
+        marker = "foo_marker"
+        engine._stats.add_stat(marker, 1.0)
+        assert marker in engine.comm.statistics.list_stat_names()
+
+
 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