[REVIEW] Add cudf.set_allocator function for easier RMM config

[shwina]

Without memory pool:

In [1]: import cudf                                                                                                                                                                           

In [2]: import numpy as np                                                                                                                                                                    

In [3]: a = cudf.Series(np.arange(100_000_000))                                                                                                                                               

In [4]: pos = cudf.Series(np.arange(50_000_000))                                                                                                                                              

In [5]: %timeit a[pos]                                                                                                                                                                        
32.6 ms ± 18.6 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

With memory pool (default allocator), and initial pool size of 1/2 GPU memory:

In [1]: import cudf                                                                                                                                                                           

In [2]: import numpy as np                                                                                                                                                                    

In [3]: cudf.set_allocator("default", pool=True)                                                                                                                                              

In [4]: a = cudf.Series(np.arange(100_000_000))                                                                                                                                               

In [5]: pos = cudf.Series(np.arange(50_000_000))                                                                                                                                              

In [6]: %timeit a[pos]                                                                                                                                                                        
9.14 ms ± 324 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

With memory pool (default allocator), and initial pool size of 0 bytes (i.e., caching allocator):

In [1]: import cudf                                                                                                                                                                           

In [2]: import numpy as np                                                                                                                                                                    

In [3]: cudf.set_allocator(pool=True, initial_pool_size=0)                                                                                                                                    

In [4]: a = cudf.Series(np.arange(100_000_000))                                                                                                                                               

In [5]: pos = cudf.Series(np.arange(50_000_000))                                                                                                                                              

In [6]: %timeit a[pos] 
10.5 ms ± 221 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

cc: @jrhemstad @harrism @jakirkham

[jrhemstad]

As @harrism mentioned in #2676 (comment), when you specify a pool size of 0, that actually defaults to half the GPU memory. Are you translating 0 from Python into a small, non-zero value passed to RMM?

[shwina]

@jrhemstad - yes, 0 bytes is translated to 1 byte.

[jrhemstad]

So a word of caution. Calling set_allocator multiple times throughout a workflow can have potentially surprising and disastrous results.

For example:

cudf.set_allocator("default", pool=True)
# do stuff, allocate memory, etc

#any additional calls to `set_allocator` will cause all previously allocated memory to be destroyed
cudf.set_allocator("default", pool=True)

Any time memory was allocated with pool=True, then the next invocation of set_allocator will cause all previously allocated memory to be freed out from underneath you.

As such, would it be possible to make it such that set_allocator can only be called once in a workflow?

[shwina]

Yup, that's the reason I have the _initfunc decorator, which ensures that the function is called only once and ignores any subsequent calls to it.

https://github.com/rapidsai/cudf/pull/2682/files#diff-cd29b5f971f5dbde2615245fd859be27R74

import cudf
cudf.set_allocator(...)
a = cudf.Series(...)
cudf.set_allocator(...)  # ignored

That being said, this doesn't protect the user from doing something like this:

import cudf
a = cudf.Series(...)
cudf.set_allocator(...)  # frees memory associated with `a`

[jrhemstad]

Yup, that's the reason I have the _initfunc decorator, which ensures that the function is called only once and ignores any subsequent calls to it.

Nifty, glad that exists in Python.

That being said, this doesn't protect the user from doing something like this:

import cudf
a = cudf.Series(...)
cudf.set_allocator(...)  # frees memory associated with `a`

This is problematic only if default allocator is pool=True. Otherwise, a would have been allocated with a non-pool/non-caching allocator and as such would not be freed by set_allocator.

.set_allocator calls rmmFinalize, which will only free memory if RMM was previously initialized with pool=True.

However, your example does illustrate another potential problem. If the user allocates a bunch of memory w/ non-pool mode, and then calls set_allocator(pool=True), this can fail if there isn't sufficient free GPU memory available to allocate the pool. I think this falls into the user error bucket though.

[shwina]

Nifty, glad that exists in Python.

It doesn't :) but it's trivial to write:

cudf/python/cudf/cudf/utils/utils.py

Line 182 in 05e297e

def initfunc(f):

[codecov]

Codecov Report

Merging #2682 into branch-0.10 will increase coverage by 0.07%.
The diff coverage is 32.14%.

@@               Coverage Diff               @@
##           branch-0.10    #2682      +/-   ##
===============================================
+ Coverage        86.36%   86.43%   +0.07%     
===============================================
  Files               48       48              
  Lines             8808     8854      +46     
===============================================
+ Hits              7607     7653      +46     
  Misses            1201     1201

Impacted Files	Coverage Δ
python/cudf/cudf/utils/utils.py	76% <32.14%> (-12.78%)	⬇️
python/cudf/cudf/utils/ioutils.py	84.28% <0%> (-8.94%)	⬇️
python/cudf/cudf/core/dataframe.py	93.83% <0%> (-0.14%)	⬇️
python/cudf/cudf/io/avro.py	80% <0%> (ø)	⬆️
python/cudf/cudf/utils/applyutils.py	100% <0%> (ø)	⬆️
python/cudf/cudf/core/series.py	92.99% <0%> (+0.02%)	⬆️
python/dask_cudf/dask_cudf/io/tests/test_s3.py	94.28% <0%> (+80%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update e764d36...67fb100. Read the comment docs.

[shwina]

Needs a complete docstring

[thomcom]

Nice

[harrism]

Would be nice to expose logging.

[kkraus14]

python/cudf/cudf/utils/utils.py:278:1: W293 blank line contains whitespace

[shwina]

Thanks @kkraus14 :)

[jakirkham]

Thanks all! Will this then be included in the nightlies?

[harrism]

Yep.