-
Notifications
You must be signed in to change notification settings - Fork 193
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add the ability to register and unregister reinitialization hooks #1072
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks Ashwin! 🙏
Generally this seems good. Had a couple comments below
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Mostly docstring-related quibbles.
Co-authored-by: Lawrence Mitchell <wence@gmx.li>
Co-authored-by: Lawrence Mitchell <wence@gmx.li>
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks @shwina for the work here, and patience for the discussions!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
A couple minor nits -- I'm really happy with how this progressed during review from others' inputs.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This (the list
comprehension) seems less succinct than what is already here
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yeah the filter
change seems fine (assuming that check behaves as expected). Tried to clarify that above
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thanks for the changes and doc clarifications, looks good!
@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.