Callback memory resource #980

shwina · 2022-02-15T00:13:04Z

This PR adds a CallbackMemoryResource that accepts Python callback functions that are responsible for allocating and deallocating memory (warning: this should not be used in production for performance reasons).

import rmm

# using a CudaMemoryResource as the backing MR,
# define allocation and deallocation functions that
# print the amount of memory being (de)allocated

base_mr = rmm.mr.CudaMemoryResource()

def allocate(size):
    print(f"Allocating {size}")
    return base_mr.allocate(size)

def deallocate(ptr, size):
    print(f"Deallocating {size}")    
    return base_mr.deallocate(ptr, size)

# create a CallbackMemoryResource and set it to be
# the default memory resource used by RMM:

mr = rmm.mr.CallbackMemoryResource(allocate, deallocate)
rmm.mr.set_current_device_resource(mr)

# All allocations/deallocations go through the callback:
s = cudf.Series([0, 1, 2]) 
# prints "Allocating 24"

include/rmm/mr/device/callback_memory_resource.hpp

github-actions · 2022-03-17T03:01:50Z

This PR has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this PR if it is no longer required. Otherwise, please respond with a comment indicating any updates. This PR will be labeled inactive-90d if there is no activity in the next 60 days.

harrism · 2022-03-21T20:26:14Z

@shwina should we push this to next release?

… callback-memory-resource

Co-authored-by: Bradley Dice <bdice@bradleydice.com>

… callback-memory-resource

bdice

No other comments from me -- thanks again. Eager to make use of all the functionality this enables!

hyperbolic2346

The only thing left are nits and slight improvements to testing. I'm happy to see this merged.

leofang

LGTM. This is pretty close to what cuQuantum expects too (except for the missing stream argument which we can work around at the Python layer):
https://docs.nvidia.com/cuda/cuquantum/python/api/generated/cuquantum.cutensornet.set_device_mem_handler.html#cuquantum.cutensornet.set_device_mem_handler
so we probably could use RMM at least for testing purpose in the next release.

vyasr · 2022-04-13T20:03:52Z

python/rmm/_lib/memory_resource.pxd

+        void* allocate(size_t bytes) except +
+        void deallocate(void* ptr, size_t bytes) except +


AFAICT these methods are being added just for the purpose of the test. I'm not opposed to exposing these functions in the Python API, but that seems like it merits discussion beyond this largely unrelated PR.

I'm also fine if you want to push back here and just get this done, but I wanted to at least have that discussion recorded if so.

I think it's probably a good idea to expose these anyway for users that want to use RMM but not necessarily construct memory owning DeviceBuffers. Also see my comment below.

I'm fine with that. Could we add some explicit tests of these APIs for other memory resources then? It seems awkward that the only test of the allocate function of memory allocators (which sounds like a pretty core feature...) would only be tested in this one callback test where we don't even actually validate the allocation.

I added a test for allocate/deallocate

python/rmm/_lib/memory_resource.pyx

include/rmm/mr/device/callback_memory_resource.hpp

python/rmm/_lib/memory_resource.pyx

vyasr · 2022-04-13T20:31:46Z

python/rmm/_lib/memory_resource.pyx

+    ``CallbackMemoryResource`` should really only be used for
+    debugging memory issues, as there is a significant performance
+    penalty associated with using a Python function for each memory
+    allocation and deallocation.


Probably out of scope for this PR, but would it be possible to instead accept a cdef function as the allocator (as a void * pointer at that point) that wouldn't have these performance implications?

That would preclude passing Python functions as the callbacks, which is the primary motivation for the CallbackMemoryResource.

Correct. I'm not suggesting that we could do it with this same class or in this PR. I'm asking if this is a useful feature for future work and another class CythonCallbackMemoryResource (or if there's some way to make this signature polymorphic). Mostly asking if we should open a follow-up issue.

python/rmm/_lib/memory_resource.pyx

python/rmm/tests/test_rmm.py

vyasr · 2022-04-13T20:35:45Z

python/rmm/tests/test_rmm.py

+    dbuf = rmm.DeviceBuffer(size=256)
+    del dbuf


Could we actually check the total memory allocation here instead of just looking for printed output?

Not for any arbitrary base_mr.

My thought here is that this test doesn't need to check that base_mr is behaving correctly, or really test what happens inside the callbacks. This test should just ensure that the callbacks are indeed invoked as expected.

Maybe there's a better way to do that I'm missing. Modifying a global is one approach I guess?

I think that's a reasonable expectation for the test, but in that case why even have a base mr? The test could remove the actual allocation from the callback entirely. Is calling the allocate function really testing anything other than the fact that you can run arbitrary Python from the callback?

Setting a global would work, but I'm also OK with output capturing for this purpose.

Mainly because it serves as a useful example of how to use CallbackMemoryResource, although I realized that probably belongs in the docstring, so I added it there.

Co-authored-by: Vyas Ramasubramani <vyas.ramasubramani@gmail.com>

… callback-memory-resource

shwina · 2022-04-14T14:51:31Z

rerun tests

vyasr

I went through and resolved all the conversations that have been addressed. There are still a few outstanding items around adding tests. Also, many of these files need their copyrights updated. mock_resource.hpp needs the entire license header added, it's not currently present. I'll be out tomorrow and Monday and I don't want to be the only blocking review in case you get a chance to wrap things up, and these issues are pretty minor, so I'm going to approve assuming those remaining issues get addressed. Thanks! Super cool feature.

shwina · 2022-04-19T22:40:41Z

@gpucibot merge

In rapidsai#980, the DeviceMemoryResource class in Python gained allocation and deallocation routines. This was to facility writing Python allocate/deallocate callbacks for the CallbackMemoryResource. These routines should, to match the C++ API, accept a stream parameter such that one can use them for stream-ordered allocation. Although we recommend that users allocate on the Python side using the DeviceBuffer interface, exposing these routines implicitly makes them public. To fix this, add an optional stream argument defaulting to the default stream. - Closes rapidsai#1493

…1494) In #980, the DeviceMemoryResource class in Python gained allocation and deallocation routines. This was to facilitate writing Python allocate/deallocate callbacks for the CallbackMemoryResource. These routines should, to match the C++ API, accept a stream parameter such that one can use them for stream-ordered allocation. Although we recommend that users allocate on the Python side using the DeviceBuffer interface, exposing these routines implicitly makes them public. To fix this, add an optional stream argument defaulting to the default stream. - Closes #1493 Authors: - Lawrence Mitchell (https://github.com/wence-) - Bradley Dice (https://github.com/bdice) Approvers: - Mark Harris (https://github.com/harrism) URL: #1494

shwina added 3 commits November 25, 2021 10:14

Draft of callback_memory_resource

ea6f54c

building...

ba5a43d

Some fixes

46c6507

github-actions bot added cpp Pertains to C++ code Python Related to RMM Python API labels Feb 15, 2022

harrism reviewed Feb 15, 2022

View reviewed changes

include/rmm/mr/device/callback_memory_resource.hpp Outdated Show resolved Hide resolved

github-actions bot added the inactive-30d label Mar 17, 2022

github-actions bot removed the inactive-30d label Mar 21, 2022

shwina added 10 commits April 8, 2022 12:57

Draft of callback_memory_resource

eb8f9e1

building...

8f1a62d

Some fixes

e7c8c88

Expose allocate() and deallocate()

1e2a10c

Add first cpp test

d2fca0d

Add another test

63b9690

Use fmt instead

cdc53e7

C++ docs

e1014d7

Add python test

f53ccf1

Merge branch 'callback-memory-resource' of github.com:shwina/rmm into…

de221d6

… callback-memory-resource

github-actions bot added CMake conda gpuCI labels Apr 8, 2022

shwina changed the base branch from branch-22.04 to branch-22.06 April 8, 2022 17:50

shwina marked this pull request as ready for review April 8, 2022 17:51

shwina requested review from a team as code owners April 8, 2022 17:51

shwina and others added 4 commits April 13, 2022 14:45

Correctly pass base_mr using ctx

a776631

Update include/rmm/mr/device/callback_memory_resource.hpp

169fa1f

Co-authored-by: Bradley Dice <bdice@bradleydice.com>

Merge branch 'callback-memory-resource' of github.com:shwina/rmm into…

a647f99

… callback-memory-resource

Address various review comments

fa22ecb

bdice approved these changes Apr 13, 2022

View reviewed changes

hyperbolic2346 approved these changes Apr 13, 2022

View reviewed changes

leofang approved these changes Apr 13, 2022

View reviewed changes

vyasr requested changes Apr 13, 2022

View reviewed changes

shwina and others added 7 commits April 14, 2022 08:53

Update include/rmm/mr/device/callback_memory_resource.hpp

8d69937

Co-authored-by: Vyas Ramasubramani <vyas.ramasubramani@gmail.com>

Update include/rmm/mr/device/callback_memory_resource.hpp

ba3c1bd

Co-authored-by: Vyas Ramasubramani <vyas.ramasubramani@gmail.com>

Update python/rmm/_lib/memory_resource.pyx

c1ba187

Co-authored-by: Vyas Ramasubramani <vyas.ramasubramani@gmail.com>

Default to nullptr

9d94861

Merge branch 'callback-memory-resource' of github.com:shwina/rmm into…

1933e9a

… callback-memory-resource

Move docs

347dc9c

Use mock class in test

b32dfa2

vyasr approved these changes Apr 15, 2022

View reviewed changes

shwina added 5 commits April 19, 2022 11:59

Add a test for allocate/deallocate

02a05d6

Add header

f9e445b

Add example

d080157

Copyright

5529ca1

Copyright

3c73ca1

rapids-bot bot merged commit 17bdbcb into rapidsai:branch-22.06 Apr 19, 2022

wence- mentioned this pull request Mar 5, 2024

[BUG] Python memory_resource allocate/deallocate do not accept a stream argument #1493

Closed

wence- mentioned this pull request Mar 5, 2024

Accept stream argument in DeviceMemoryResource allocate/deallocate #1494

Merged

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Callback memory resource #980

Callback memory resource #980

shwina commented Feb 15, 2022 •

edited

Loading

github-actions bot commented Mar 17, 2022

harrism commented Mar 21, 2022

bdice left a comment

hyperbolic2346 left a comment

leofang left a comment

vyasr Apr 13, 2022

vyasr Apr 13, 2022

shwina Apr 14, 2022 •

edited

Loading

vyasr Apr 14, 2022

shwina Apr 19, 2022

vyasr Apr 13, 2022

shwina Apr 14, 2022

vyasr Apr 14, 2022

vyasr Apr 13, 2022

shwina Apr 14, 2022 •

edited

Loading

vyasr Apr 14, 2022

shwina Apr 19, 2022

shwina commented Apr 14, 2022

vyasr left a comment

shwina commented Apr 19, 2022

		void* allocate(size_t bytes) except +
		void deallocate(void* ptr, size_t bytes) except +

Callback memory resource #980

Callback memory resource #980

Conversation

shwina commented Feb 15, 2022 • edited Loading

github-actions bot commented Mar 17, 2022

harrism commented Mar 21, 2022

bdice left a comment

Choose a reason for hiding this comment

hyperbolic2346 left a comment

Choose a reason for hiding this comment

leofang left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shwina Apr 14, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shwina Apr 14, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shwina commented Apr 14, 2022

vyasr left a comment

Choose a reason for hiding this comment

shwina commented Apr 19, 2022

shwina commented Feb 15, 2022 •

edited

Loading

shwina Apr 14, 2022 •

edited

Loading

shwina Apr 14, 2022 •

edited

Loading