fix cython bindings

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Author: nirandaperera

python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py

@@ -175,7 +175,11 @@ def streaming_shuffle(
     consumer_thread.join(timeout=wait_timeout)
 
 
-def ucxx_mpi_setup(options: Options, progress_thread: ProgressThread) -> Communicator:
+def ucxx_mpi_setup(
+    options: Options,
+    progress_thread: ProgressThread,
+    statistics: Statistics | None = None,
+) -> Communicator:
     """
     Bootstrap UCXX cluster using MPI.
 
@@ -185,6 +189,9 @@ def ucxx_mpi_setup(options: Options, progress_thread: ProgressThread) -> Communi
         Configuration options.
     progress_thread
         Progress thread for the communicator.
+    statistics
+        Statistics object for recording communication metrics. If ``None``
+        (the default), statistics tracking is disabled.
 
     Returns
     -------
@@ -201,7 +208,7 @@ def ucxx_mpi_setup(options: Options, progress_thread: ProgressThread) -> Communi
 
     if MPI.COMM_WORLD.Get_rank() == 0:
         comm = new_communicator(
-            MPI.COMM_WORLD.size, None, None, options, progress_thread
+            MPI.COMM_WORLD.size, None, None, options, progress_thread, statistics
         )
         root_address_str = get_root_ucxx_address(comm)
     else:
@@ -212,7 +219,12 @@ def ucxx_mpi_setup(options: Options, progress_thread: ProgressThread) -> Communi
     if MPI.COMM_WORLD.Get_rank() != 0:
         root_address = ucx_api.UCXAddress.create_from_buffer(root_address_str)
         comm = new_communicator(
-            MPI.COMM_WORLD.size, None, root_address, options, progress_thread
+            MPI.COMM_WORLD.size,
+            None,
+            root_address,
+            options,
+            progress_thread,
+            statistics,
         )
 
     assert comm.nranks == MPI.COMM_WORLD.size
@@ -245,15 +257,15 @@ def setup_and_run(args: argparse.Namespace) -> None:
     progress_thread = ProgressThread(stats)
     if args.comm == "mpi":
         comm = rapidsmpf.communicator.mpi.new_communicator(
-            MPI.COMM_WORLD, options, progress_thread
+            MPI.COMM_WORLD, options, progress_thread, stats
         )
     elif args.comm == "ucxx":
         if rapidsmpf.bootstrap.is_running_with_rrun():
             raise ValueError(
                 "UCXX communicator is not supported with rrun yet, due to missing allreduce support"
             )
         else:
-            comm = ucxx_mpi_setup(options, progress_thread)
+            comm = ucxx_mpi_setup(options, progress_thread, stats)
 
     # Create a buffer resource that limits device memory if `--spill-device`
     # is not None.

python/rapidsmpf/rapidsmpf/coll/allgather.pxd

@@ -14,7 +14,6 @@ from rapidsmpf.memory.buffer_resource cimport (BufferResource,
                                                cpp_BufferResource)
 from rapidsmpf.memory.packed_data cimport cpp_PackedData
 from rapidsmpf.progress_thread cimport cpp_ProgressThread
-from rapidsmpf.statistics cimport cpp_Statistics
 
 
 cdef extern from "<rapidsmpf/coll/allgather.hpp>" namespace \
@@ -30,8 +29,7 @@ cdef extern from "<rapidsmpf/coll/allgather.hpp>" nogil:
         cpp_AllGather(
             shared_ptr[cpp_Communicator] comm,
             int32_t op_id,
-            cpp_BufferResource *br,
-            shared_ptr[cpp_Statistics] statistics
+            cpp_BufferResource *br
         ) except +ex_handler
         void insert(uint64_t sequence_number, cpp_PackedData packed_data) \
             except +ex_handler

python/rapidsmpf/rapidsmpf/coll/allgather.pyi

@@ -5,15 +5,13 @@ from __future__ import annotations
 from rapidsmpf.communicator.communicator import Communicator
 from rapidsmpf.memory.buffer_resource import BufferResource
 from rapidsmpf.memory.packed_data import PackedData
-from rapidsmpf.statistics import Statistics
 
 class AllGather:
     def __init__(
         self,
         comm: Communicator,
         op_id: int,
         br: BufferResource,
-        statistics: Statistics | None = None,
     ) -> None: ...
     @property
     def comm(self) -> Communicator: ...

python/rapidsmpf/rapidsmpf/coll/allgather.pyx

@@ -14,7 +14,6 @@ from rapidsmpf.memory.buffer_resource cimport (BufferResource,
                                                cpp_BufferResource)
 from rapidsmpf.memory.packed_data cimport (PackedData, cpp_PackedData,
                                            packed_data_vector_to_list)
-from rapidsmpf.statistics cimport Statistics
 
 
 cdef class AllGather:
@@ -37,8 +36,6 @@ cdef class AllGather:
         between 0 and 2^20 - 1.
     br
         Buffer resource for memory allocation.
-    statistics
-        Statistics collection instance. If None, statistics is disabled.
 
     Notes
     -----
@@ -52,19 +49,15 @@ cdef class AllGather:
         Communicator comm not None,
         int32_t op_id,
         BufferResource br not None,
-        Statistics statistics = None,
     ):
         self._br = br
         self._comm = comm
         cdef cpp_BufferResource* br_ = br.ptr()
-        if statistics is None:
-            statistics = Statistics(enable=False)  # Disables statistics.
         with nogil:
             self._handle = make_unique[cpp_AllGather](
                 comm._handle,
                 op_id,
                 br_,
-                statistics._handle,
             )
 
     def __dealloc__(self):

python/rapidsmpf/rapidsmpf/communicator/communicator.pxd

@@ -7,6 +7,7 @@ from libcpp.string cimport string
 
 from rapidsmpf._detail.exception_handling cimport ex_handler
 from rapidsmpf.progress_thread cimport cpp_ProgressThread
+from rapidsmpf.statistics cimport Statistics, cpp_Statistics
 
 
 cdef extern from "<rapidsmpf/communicator/communicator.hpp>" namespace \
@@ -37,6 +38,7 @@ cdef extern from "<rapidsmpf/communicator/communicator.hpp>" nogil:
         string str() except +ex_handler
         shared_ptr[cpp_ProgressThread] progress_thread() except +ex_handler
         shared_ptr[cpp_Logger] logger()
+        shared_ptr[cpp_Statistics] statistics()
 
 cdef class Communicator:
     cdef shared_ptr[cpp_Communicator] _handle

python/rapidsmpf/rapidsmpf/communicator/communicator.pyi

@@ -3,6 +3,7 @@
 from enum import IntEnum
 
 from rapidsmpf.progress_thread import ProgressThread
+from rapidsmpf.statistics import Statistics
 
 class LOG_LEVEL(IntEnum):
     NONE = ...
@@ -30,6 +31,8 @@ class Communicator:
     def logger(self) -> Logger: ...
     @property
     def progress_thread(self) -> ProgressThread: ...
+    @property
+    def statistics(self) -> Statistics: ...
     def get_str(self) -> str: ...
 
 def _available_communicators() -> tuple[str, ...]: ...

python/rapidsmpf/rapidsmpf/communicator/communicator.pyx

@@ -6,6 +6,7 @@ from libcpp.string cimport string
 from libcpp.utility cimport move
 
 from rapidsmpf.progress_thread cimport ProgressThread
+from rapidsmpf.statistics cimport Statistics
 
 
 cdef class Logger:
@@ -179,6 +180,21 @@ cdef class Communicator:
         pt._handle = deref(self._handle).progress_thread()
         return pt
 
+    @property
+    def statistics(self):
+        """
+        Get the statistics object associated with this communicator.
+
+        Returns
+        -------
+        Statistics
+            The statistics instance. If no statistics object was provided at
+            construction time, a disabled (no-op) instance is returned.
+        """
+        cdef Statistics stats = Statistics.__new__(Statistics)
+        stats._handle = deref(self._handle).statistics()
+        return stats
+
     def get_str(self):
         """
         Get a string representation of the communicator.

python/rapidsmpf/rapidsmpf/communicator/mpi.pyi

@@ -5,7 +5,11 @@ from mpi4py.MPI import Intracomm
 from rapidsmpf.communicator.communicator import Communicator
 from rapidsmpf.config import Options
 from rapidsmpf.progress_thread import ProgressThread
+from rapidsmpf.statistics import Statistics
 
 def new_communicator(
-    comm: Intracomm, options: Options, progress_thread: ProgressThread
+    comm: Intracomm,
+    options: Options,
+    progress_thread: ProgressThread,
+    statistics: Statistics | None = None,
 ) -> Communicator: ...

python/rapidsmpf/rapidsmpf/communicator/mpi.pyx

@@ -9,19 +9,24 @@ from rapidsmpf._detail.exception_handling cimport ex_handler
 from rapidsmpf.communicator.communicator cimport Communicator
 from rapidsmpf.config cimport Options, cpp_Options
 from rapidsmpf.progress_thread cimport ProgressThread, cpp_ProgressThread
+from rapidsmpf.statistics cimport Statistics, cpp_Statistics
 
 
 cdef extern from "<rapidsmpf/communicator/mpi.hpp>" nogil:
     cdef cppclass cpp_MPI_Communicator "rapidsmpf::MPI":
         cpp_MPI_Communicator(
-            libmpi.MPI_Comm comm, cpp_Options options, shared_ptr[cpp_ProgressThread]
+            libmpi.MPI_Comm comm,
+            cpp_Options options,
+            shared_ptr[cpp_ProgressThread],
+            shared_ptr[cpp_Statistics]
         ) except +ex_handler
 
 
 def new_communicator(
     Intracomm comm not None,
     Options options not None,
     ProgressThread progress_thread not None,
+    Statistics statistics=None,
 ):
     """
     Create a new RapidsMPF-MPI communicator based on an existing mpi4py communicator.
@@ -34,14 +39,19 @@ def new_communicator(
         Configuration options.
     progress_thread
         Progress thread for the communicator.
+    statistics
+        Statistics object for recording communication metrics. If ``None``
+        (the default), statistics tracking is disabled.
 
     Returns
     -------
     A new RapidsMPF-MPI communicator.
     """
+    if statistics is None:
+        statistics = Statistics(enable=False)
     cdef Communicator ret = Communicator.__new__(Communicator)
     with nogil:
         ret._handle = make_shared[cpp_MPI_Communicator](
-            comm.ob_mpi, options._handle, progress_thread._handle
+            comm.ob_mpi, options._handle, progress_thread._handle, statistics._handle
         )
     return ret

python/rapidsmpf/rapidsmpf/communicator/testing.py

@@ -15,6 +15,7 @@
     from rapidsmpf.communicator.communicator import Communicator
     from rapidsmpf.config import Options
     from rapidsmpf.progress_thread import ProgressThread
+    from rapidsmpf.statistics import Statistics
 
 
 def initialize_ucxx() -> ucx_api.UCXWorker:
@@ -36,7 +37,10 @@ def initialize_ucxx() -> ucx_api.UCXWorker:
 
 
 def ucxx_mpi_setup(
-    ucxx_worker: ucx_api.UCXWorker, options: Options, progress_thread: ProgressThread
+    ucxx_worker: ucx_api.UCXWorker,
+    options: Options,
+    progress_thread: ProgressThread,
+    statistics: Statistics | None = None,
 ) -> Communicator:
     """
     Bootstrap a UCXX communicator within an MPI rank.
@@ -49,6 +53,9 @@ def ucxx_mpi_setup(
         Configuration options.
     progress_thread
         Progress thread for the initialized communicator.
+    statistics
+        Statistics object for recording communication metrics. If ``None``
+        (the default), statistics tracking is disabled.
 
     Returns
     -------
@@ -63,7 +70,7 @@ def ucxx_mpi_setup(
 
     if MPI.COMM_WORLD.Get_rank() == 0:
         comm = new_communicator(
-            MPI.COMM_WORLD.size, ucxx_worker, None, options, progress_thread
+            MPI.COMM_WORLD.size, ucxx_worker, None, options, progress_thread, statistics
         )
         root_address_bytes = get_root_ucxx_address(comm)
     else:
@@ -74,7 +81,12 @@ def ucxx_mpi_setup(
     if MPI.COMM_WORLD.Get_rank() != 0:
         root_address = ucx_api.UCXAddress.create_from_buffer(root_address_bytes)
         comm = new_communicator(
-            MPI.COMM_WORLD.size, ucxx_worker, root_address, options, progress_thread
+            MPI.COMM_WORLD.size,
+            ucxx_worker,
+            root_address,
+            options,
+            progress_thread,
+            statistics,
         )
 
     assert comm.nranks == MPI.COMM_WORLD.size