Add the ability to register and unregister reinitialization hooks

python/rmm/__init__.py

@@ -22,8 +22,10 @@
     RMMNumbaManager,
     _numba_memory_manager,
     is_initialized,
+    register_reinitialize_hook,
     reinitialize,
     rmm_cupy_allocator,
+    unregister_reinitialize_hook,
 )
 
 __version__ = get_versions()["version"]

python/rmm/rmm.py

@@ -29,6 +29,9 @@ def __init__(self, errcode, msg):
         super(RMMError, self).__init__(msg)
 
 
+_reinitialize_hooks = {}
+
+
 def reinitialize(
     pool_allocator=False,
     managed_memory=False,
@@ -82,6 +85,10 @@ def reinitialize(
     Use `rmm.get_log_filenames()` to get the log file names
     corresponding to each device.
     """
+    [
+        hook(*args, **kwargs)
+        for (hook, (args, kwargs)) in reversed(_reinitialize_hooks.items())
+    ]
     rmm.mr._initialize(
         pool_allocator=pool_allocator,
         managed_memory=managed_memory,
@@ -231,3 +238,21 @@ def rmm_cupy_allocator(nbytes):
     ptr = cupy.cuda.memory.MemoryPointer(mem, 0)
 
     return ptr
+
+
+def register_reinitialize_hook(func, *args, **kwargs):
+    """
+    Register a hook to be called by `rmm.reinitialize()`.
+
+    Hooks are called in the *reverse* order they are registered.
+    """
+    _reinitialize_hooks[func] = (args, kwargs)
+    return func
+
+
+def unregister_reinitialize_hook(func):
+    """
+    Remove func from the list of hooks to be called by `rmm.reinitialize()`.
+    """
+    _reinitialize_hooks.pop(func, None)
+    return func

python/rmm/tests/test_rmm.py

@@ -765,3 +765,38 @@ def deallocate_func(ptr, size):
 
     captured = capsys.readouterr()
     assert captured.out == "Allocating 256 bytes\nDeallocating 256 bytes\n"
+
+
+def test_reinitialize_hooks(capsys):
+    def one():
+        print("one")
+
+    def two():
+        print("two")
+
+    def three(s):
+        print(s)
+
+    rmm.register_reinitialize_hook(one)
+    rmm.register_reinitialize_hook(two)
+    rmm.register_reinitialize_hook(three, "four")
+
+    rmm.reinitialize()
+    captured = capsys.readouterr()
+    assert captured.out == "four\ntwo\none\n"
+
+    rmm.unregister_reinitialize_hook(one)
+    rmm.unregister_reinitialize_hook(two)
+    rmm.unregister_reinitialize_hook(three)
+
+    rmm.register_reinitialize_hook(one)
+    rmm.register_reinitialize_hook(two)
+    rmm.unregister_reinitialize_hook(one)
+
+    rmm.reinitialize()
+    captured = capsys.readouterr()
+    assert captured.out == "two\n"
+
+    rmm.unregister_reinitialize_hook(one)
+    rmm.unregister_reinitialize_hook(two)
+    rmm.unregister_reinitialize_hook(three)