Add CUPTI Python API

conda/recipes/rapidsmpf/recipe.yaml

@@ -58,10 +58,12 @@ requirements:
     - ${{ compiler("c") }}
     - ${{ compiler("cxx") }}
     - ${{ compiler("cuda") }}
+    - cuda-cupti-dev
     - cuda-version =${{ cuda_version }}
     - ${{ stdlib("c") }}
   host:
     - cuda-cudart-dev
+    - cuda-cupti-dev
     - cuda-version =${{ cuda_version }}
     - cython >=3.0.0
     - librapidsmpf =${{ version }}
@@ -74,10 +76,11 @@ requirements:
     - rmm =${{ minor_version }}
     - scikit-build-core >=0.10.0
   run:
-    - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
+    - cuda-cupti
     - if: cuda_major == "12"
       then: cuda-python >=12.9.2,<13.0a0
       else: cuda-python >=13.0.1,<14.0a0
+    - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }}
     - cudf =${{ minor_version }}
     - cupy >=13.6.0
     - librapidsmpf =${{ version }}
@@ -98,6 +101,7 @@ requirements:
       - ${{ stdlib("c") }}
     by_name:
       - cuda-cudart
+      - cuda-cupti
       - cuda-version
       - openmpi
       - mpi4py

cpp/examples/example_cupti_monitor.cpp

@@ -58,7 +58,7 @@ int main() {
                 // Allocate device memory using rmm::device_buffer
                 rmm::device_buffer buf(allocation_size, rmm::cuda_stream_default);
                 device_buffers.push_back(std::move(buf));
-            } catch (rmm::bad_alloc& e) {
+            } catch (rmm::bad_alloc const& e) {
                 std::cerr << "rmm::device_buffer allocation failed: " << e.what()
                           << std::endl;
                 break;

python/rapidsmpf/rapidsmpf/CMakeLists.txt

@@ -7,6 +7,11 @@ set(cython_modules config.pyx cuda_stream.pyx progress_thread.pyx rmm_resource_a
                    shuffler.pyx statistics.pyx
 )
 
+# Add cupti module conditionally if CUPTI support is enabled
+if(RAPIDSMPF_HAVE_CUPTI)
+  list(APPEND cython_modules cupti.pyx)
+endif()
+
 rapids_cython_create_modules(
   CXX
   SOURCE_FILES "${cython_modules}"

python/rapidsmpf/rapidsmpf/__init__.py

@@ -4,6 +4,8 @@
 
 from __future__ import annotations
 
+import contextlib
+
 from rapidsmpf._version import __git_commit__, __version__  # noqa: F401
 
 # If librapidsmpf was installed as a wheel, we must request it to load the
@@ -16,3 +18,7 @@
 else:
     librapidsmpf.load_library()
     del librapidsmpf
+
+# Import CUPTI module if available (only built when CUPTI is found)
+with contextlib.suppress(ImportError):
+    from rapidsmpf import cupti  # noqa: F401

python/rapidsmpf/rapidsmpf/cupti.pxd

@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+from libc.stddef cimport size_t
+from libcpp cimport bool as cpp_bool
+from libcpp.memory cimport unique_ptr
+from libcpp.string cimport string
+from libcpp.unordered_map cimport unordered_map
+from libcpp.vector cimport vector
+
+
+cdef extern from "<chrono>" namespace "std::chrono" nogil:
+    cdef cppclass milliseconds:
+        milliseconds(long long) except +
+
+cdef extern from "<cupti.h>" nogil:
+    ctypedef enum CUpti_CallbackId:
+        pass
+
+
+cdef extern from "<rapidsmpf/cupti.hpp>" nogil:
+    cdef struct cpp_MemoryDataPoint "rapidsmpf::MemoryDataPoint":
+        double timestamp
+        size_t free_memory
+        size_t total_memory
+        size_t used_memory
+
+    cdef cppclass cpp_CuptiMonitor "rapidsmpf::CuptiMonitor":
+        cpp_CuptiMonitor(
+            cpp_bool enable_periodic_sampling,
+            milliseconds sampling_interval_ms
+        ) except +
+        void start_monitoring() except +
+        void stop_monitoring() except +
+        cpp_bool is_monitoring() except +
+        void capture_memory_sample() except +
+        const vector[cpp_MemoryDataPoint]& get_memory_samples() except +
+        void clear_samples() except +
+        size_t get_sample_count() except +
+        void write_csv(const string& filename) except +
+        void set_debug_output(cpp_bool enabled, size_t threshold_mb) except +
+        unordered_map[CUpti_CallbackId, size_t] get_callback_counters() except +
+        void clear_callback_counters() except +
+        size_t get_total_callback_count() except +
+        string get_callback_summary() except +
+
+
+cdef class MemoryDataPoint:
+    cdef cpp_MemoryDataPoint _data
+
+    @staticmethod
+    cdef MemoryDataPoint from_cpp(cpp_MemoryDataPoint data)
+
+
+cdef class CuptiMonitor:
+    cdef unique_ptr[cpp_CuptiMonitor] _handle

python/rapidsmpf/rapidsmpf/cupti.pyi

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+class MemoryDataPoint:
+    @property
+    def timestamp(self) -> float: ...
+    @property
+    def free_memory(self) -> int: ...
+    @property
+    def total_memory(self) -> int: ...
+    @property
+    def used_memory(self) -> int: ...
+    def __repr__(self) -> str: ...
+
+class CuptiMonitor:
+    def __init__(
+        self, enable_periodic_sampling: bool = False, sampling_interval_ms: int = 100
+    ) -> None: ...
+    def start_monitoring(self) -> None: ...
+    def stop_monitoring(self) -> None: ...
+    def is_monitoring(self) -> bool: ...
+    def capture_memory_sample(self) -> None: ...
+    def get_memory_samples(self) -> list[MemoryDataPoint]: ...
+    def clear_samples(self) -> None: ...
+    def get_sample_count(self) -> int: ...
+    def write_csv(self, filename: str) -> None: ...
+    def set_debug_output(self, enabled: bool, threshold_mb: int = 10) -> None: ...
+    def get_callback_counters(self) -> dict[int, int]: ...
+    def clear_callback_counters(self) -> None: ...
+    def get_total_callback_count(self) -> int: ...
+    def get_callback_summary(self) -> str: ...