From 8d49fffdb93b55ce70c72981d2e1d5511692eaa2 Mon Sep 17 00:00:00 2001
From: Matthew Murray <41342305+Matt711@users.noreply.github.com>
Date: Thu, 31 Oct 2024 19:44:38 -0400
Subject: [PATCH] Deprecate `rmm._lib` (#1713)

Follows up #1676 to add deprecation warnings to the `rmm._lib` sub package.

Authors:
  - Matthew Murray (https://github.com/Matt711)

Approvers:
  - Vyas Ramasubramani (https://github.com/vyasr)

URL: https://github.com/rapidsai/rmm/pull/1713
---
 python/rmm/rmm/__init__.py       | 8 ++++++++
 python/rmm/rmm/_lib/__init__.py  | 8 ++++++++
 python/rmm/rmm/tests/test_rmm.py | 6 ++++++
 3 files changed, 22 insertions(+)

diff --git a/python/rmm/rmm/__init__.py b/python/rmm/rmm/__init__.py
index b23ad68f9..832fec095 100644
--- a/python/rmm/rmm/__init__.py
+++ b/python/rmm/rmm/__init__.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
+
 from rmm import mr
 from rmm._version import __git_commit__, __version__
 from rmm.mr import disable_logging, enable_logging, get_log_filenames
@@ -58,6 +60,12 @@ def __getattr__(name):
     if name == "_lib":
         import importlib
 
+        warnings.warn(
+            "The `rmm._lib` module is deprecated in will be removed in a future release. Use `rmm.pylibrmm` instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+
         module = importlib.import_module("rmm.pylibrmm")
         return module
     else:
diff --git a/python/rmm/rmm/_lib/__init__.py b/python/rmm/rmm/_lib/__init__.py
index 7cfddab60..7e01bda77 100644
--- a/python/rmm/rmm/_lib/__init__.py
+++ b/python/rmm/rmm/_lib/__init__.py
@@ -12,4 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
+
 from rmm.pylibrmm import *
+
+warnings.warn(
+    "The `rmm._lib` module is deprecated in will be removed in a future release. Use `rmm.pylibrmm` instead.",
+    FutureWarning,
+    stacklevel=2,
+)
diff --git a/python/rmm/rmm/tests/test_rmm.py b/python/rmm/rmm/tests/test_rmm.py
index c03b9e501..9872ba89d 100644
--- a/python/rmm/rmm/tests/test_rmm.py
+++ b/python/rmm/rmm/tests/test_rmm.py
@@ -1076,3 +1076,9 @@ def test_available_device_memory():
     assert initial_memory[1] == final_memory[1]
     assert initial_memory[0] > 0
     assert final_memory[0] > 0
+
+
+# TODO: Remove test when rmm._lib is removed in 25.02
+def test_deprecate_rmm_lib():
+    with pytest.warns(FutureWarning):
+        rmm._lib.device_buffer.DeviceBuffer(size=100)