From cf15a02a52e5245bc72e522bc45b78acf8138576 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 27 May 2024 17:05:23 +0200 Subject: [PATCH] _cudf_nvtx_annotate: add RMM's memory profiler See --- python/cudf/cudf/utils/nvtx_annotation.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/python/cudf/cudf/utils/nvtx_annotation.py b/python/cudf/cudf/utils/nvtx_annotation.py index a4404e51232..d30774f6ef5 100644 --- 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"] @@ -17,12 +19,13 @@ 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( + """Decorator for applying nvtx annotations and memory statistics.""" + 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(