Add the ability to register and unregister reinitialization hooks #1072
Conversation
jakirkham
added
non-breaking
Co-authored-by: Lawrence Mitchell <wence@gmx.li>
Co-authored-by: Lawrence Mitchell <wence@gmx.li>
python/rmm/rmm.py
Outdated
| """ | ||
| global _reinitialize_hooks | ||
| _reinitialize_hooks = list( | ||
| filterfalse(lambda x: x[0] == func, _reinitialize_hooks) |
There was a problem hiding this comment.
Why filterfalse with == when filter with != would do the same thing? Then we don't need to import from itertools.
| filterfalse(lambda x: x[0] == func, _reinitialize_hooks) | |
| filter(lambda x: x[0] != func, _reinitialize_hooks) |
There was a problem hiding this comment.
Also we could use a comprehension, which tends to be a bit faster as well as more succinct.
_reinitialize_hooks = [x for x in _reinitialize_hooks if x[0] != func]There was a problem hiding this comment.
This (the list comprehension) seems less succinct than what is already here
There was a problem hiding this comment.
The comparison I was making was:
_reinitialize_hooks = list(
filterfalse(lambda x: x[0] == func, _reinitialize_hooks)
)(3 lines, 92 chars)
vs.
_reinitialize_hooks = [x for x in _reinitialize_hooks if x[0] != func](1 line, 70 chars)
The difference is small. I cared more about removing the from itertools import filterfalse when we just needed to swap filterfalse, == for filter, !=.
There was a problem hiding this comment.
Yeah the filter change seems fine (assuming that check behaves as expected). Tried to clarify that above
|
@gpucibot merge |
This PR introduces the ability to register one or more "hooks" that will be invoked before each call to
rmm.reinitialize().The motivation is to allow downstream libraries to register hooks that are responsible for deleting any internal objects that that keep references to an RMM memory resource. In particular, this is useful for rapidsai/cudf#11246.