Memory Profiling

[madsbk]

Use RMM's new memory profiler to profile all functions already decorated with _cudf_nvtx_annotate.

Example

import cudf
from cudf.utils.performance_tracking import print_memory_report

cudf.set_option("memory_profiling", True)

df1 = cudf.DataFrame({"a": [1, 2, 3]})
df2 = cudf.DataFrame({"a": [2, 2, 3]})
df3 = df1.merge(df2)

print_memory_report()

Output:

Memory Profiling
================

Ordered by: memory_peak

ncalls     memory_peak    memory_total  filename:lineno(function)
     1             272             688  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/dataframe.py:4072(DataFrame.merge)
     2              32              64  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/dataframe.py:1043(DataFrame._init_from_dict_like)
     2              32              64  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/dataframe.py:690(DataFrame.__init__)
     2               0               0  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/dataframe.py:1131(DataFrame._align_input_series_indices)
     7               0               0  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/index.py:214(RangeIndex.__init__)
     6               0               0  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/index.py:424(RangeIndex.__len__)
     4               0               0  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/frame.py:271(Frame.__len__)
     2               0               0  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/dataframe.py:3195(DataFrame._insert)
     2               0               0  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/index.py:270(RangeIndex.name)
     2               0               0  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/index.py:369(RangeIndex.copy)
     5               0               0  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/frame.py:134(Frame._from_data)
     2               0               0  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/frame.py:1039(Frame._copy_type_metadata)
     2               0               0  /home/mkristensen/apps/miniforge3/envs/rmm-cudf-0527/lib/python3.11/site-packages/cudf/core/indexed_frame.py:315(IndexedFrame._from_columns_like_self)

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• Rename _cudf_nvtx_annotate to _performance_tracking

[madsbk]

Should we rename _cudf_nvtx_annotate to something like _performance_tracking ?
And nvtx_annotation.py => performance_tracking.py

[mroeschke]

I think that would be a more generically clear name. +1

[madsbk]

I will do the rename when the PR is ready to be merged. It results in 591 function renames, which might be hard to keep merge-conflict free :)

[harrism]

I requested this in rapidsai/rmm#1563, but it seems not to have made it in. The current output shown above provides no way of knowing what the units are in the peak and total memory report. Suggest changing the headings to make this explicit, otherwise usability is reduced.

ncalls     memory_peak (MiB)    memory_total (MiB)  filename:lineno(function)
...

[madsbk]

@harrism, I added a Legends section, which include units

Memory Profiling
================

Legends:
  ncalls       - number of times the function or code block was called
  memory_peak  - peak memory allocated in function or code block (in bytes)
  memory_total - total memory allocated in function or code block (in bytes)

[bdice]

@madsbk Is there any measurable overhead to using the statistic profiler? The NVTX overhead is very small, which is why we have enabled it across cuDF. We may want to run benchmarks to verify this before merging. If it has a cost, perhaps use an environment variable such as CUDF_PROFILE_MEMORY_STATISTICS to control the behavior.

[madsbk]

Good point, it has a price of 1.7 microseconds.
The overhead goes from 0.18us to 1.9us. The extra overhead is dominated by two calls to rmm.mr.get_current_device_resource(), which takes 1.5us.

I think this is small enough to ignore?

[bdice]

The overhead goes from 0.18us to 1.9us.

@vyasr @galipremsagar Do you think this is small enough to enable memory statistics tracking on every NVTX-annotated function in cudf Python? I am pretty hesitant to enable this by default but don't want to be the only person weighing in.

[harrism]

Given the overhead is constant, when memory size is high, this will be a drop in the bucket. When it's low, it may be significant on some functions.

Can we just make it really easy to enable, rather than enabled by default?

My strong opinion, weakly held: nonzero cost profiling should not be enabled by default.

[madsbk]

Alright, I have added a memory_profiling option :)

It also makes it a bit more user-friendly. To enable memory profiling do:

import cudf
from cudf.utils.performance_tracking import print_memory_report

cudf.set_option("memory_profiling", True)

# my code

print_memory_report()

NB: before merging this PR, I will remove nvtx_annotation.py and use @_performance_tracking throughout.

[wence-]

I think we want to figure out a way to be able to disable the statistics MR, which is currently not possible.

[wence-]

issue: This doesn't behave like a normal context manager (leave the state unchanged after exit) because it doesn't unwind the change to the current memory resource that rmm.statistics.enable_statistics enacts. This means that the following code does an "unexpected" thing: side-effectfully changes the MR:

# statistics not enabled...
with cudf.options_context(("memory_profiling", True)):
    # code that profiles statistics

# expect stats memory resource to be popped, but it is not

It seems like the call to enable statistics should be paired with a matching disable statistics call. But this seems bad if it were done at every level of the call stack.

It feels like the right place for modification of the MR is in the set_option call. set_option("memory_profiling", True) would store the old MR, push the statistics MR, and then a later set_option("memory_profiling", False") would pop the statistics MR if it still matches.

This, however, would preclude use of the environment variable default (because that is set on import, and we don't want to have to call RMM on import).

Does this suggest that we need a specific enable_memory_profiling that can be used as a context manager that sets the option and sets up the MR?

[madsbk]

If this is an issue, RMM already have rmm.statistics.statistics() and rmm.statistics.enable_statistics().
Thus the user code would look something like:

import cudf
import rmm.statistics

# The must first enable the cudf option
cudf.set_option("memory_profiling", True)

# Then enable statistics within a context
with rmm.statistics.statistics():
    # by code

# Or globally
rmm.statistics.enable_statistics()

Of cause, we could wrap/alias the rmm functions in cudf native functions. However, I am not sure if the disadvantage of the "unexpected" side-effect outweigh the user-friendliness of just setting CUDF_MEMORY_PROFILING="1"?

[madsbk]

Alternatively, we could implement lazy environment variable in options.py. It might be worth streamlining the option module in any case.

[wence-]

In terms of solutions, it depends a little bit on how end-user-friendly we want this to be. If this is mostly for developer use, I am happy with the programmatic set_option + use RMM calls. If we want this as a debugging tool for end users "where is my code allocating lots of vram", then we likely need something a little more ergonomic.

As much as I am allergic to environment variables from a reproducibility point of view, I can see the attraction.

[madsbk]

I am not sure but I guess we could start with the set_option + use RMM calls and then later consider a more end-user-friendly solution?

[wence-]

I think that's a reasonable compromise.

[madsbk]

done

[harrism]

Doc comments.

[wence-]

Looks good, tiny typographic suggestions!

[vyasr]

One question, otherwise LGTM.

[wence-]

Approving dask-cudf changes

[madsbk]

Thanks all

[madsbk]

/merge