_cudf_nvtx_annotate: add RMM's memory profiler

See <rapidsai/rmm#1563>
rapidsai · May 27, 2024 · 61093a4 · 61093a4
1 parent 8458306
commit 61093a4
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/python/cudf/cudf/utils/nvtx_annotation.py b/python/cudf/cudf/utils/nvtx_annotation.py
@@ -1,10 +1,12 @@
-# Copyright (c) 2023, NVIDIA CORPORATION.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION.
 
 import hashlib
 from functools import partial
 
 from nvtx import annotate
 
+import rmm.statistics
+
 _NVTX_COLORS = ["green", "blue", "purple", "rapids"]
 
 
@@ -18,11 +20,12 @@ def _get_color_for_nvtx(name):
 
 def _cudf_nvtx_annotate(func, domain="cudf_python"):
     """Decorator for applying nvtx annotations to methods in cudf."""
-    return annotate(
+    nvtx_wrapped = annotate(
         message=func.__qualname__,
         color=_get_color_for_nvtx(func.__qualname__),
         domain=domain,
     )(func)
+    return rmm.statistics.profiler()(nvtx_wrapped)
 
 
 _dask_cudf_nvtx_annotate = partial(