From 15b1ae113b5d04f4e8724c41397393dd5dd0871a Mon Sep 17 00:00:00 2001
From: Matthew Murray <matthewmurray711@gmail.com>
Date: Mon, 16 Sep 2024 11:07:37 -0700
Subject: [PATCH] deprecate rmm._lib

---
 python/rmm/rmm/__init__.py       | 15 +++++++++++++++
 python/rmm/rmm/_lib/__init__.py  | 22 ++++++++++++++++++++++
 python/rmm/rmm/tests/test_rmm.py |  5 +++++
 3 files changed, 42 insertions(+)
 create mode 100644 python/rmm/rmm/_lib/__init__.py

diff --git a/python/rmm/rmm/__init__.py b/python/rmm/rmm/__init__.py
index 61c5e4561..3c55a5cbd 100644
--- a/python/rmm/rmm/__init__.py
+++ b/python/rmm/rmm/__init__.py
@@ -52,3 +52,18 @@
     "should_log",
     "unregister_reinitialize_hook",
 ]
+
+
+def __getattr__(name):
+    if name == "_lib":
+        import importlib
+        import warnings
+
+        warnings.warn(
+            "The `rmm._lib` module is deprecated in will be removed in a future release. Use `rmm.python` instead.",
+            FutureWarning,
+        )
+        module = importlib.import_module("rmm.python")
+        return module
+    else:
+        raise AttributeError(f"Module '{__name__}' has no attribute '{name}'")
diff --git a/python/rmm/rmm/_lib/__init__.py b/python/rmm/rmm/_lib/__init__.py
new file mode 100644
index 000000000..9b6d1daf4
--- /dev/null
+++ b/python/rmm/rmm/_lib/__init__.py
@@ -0,0 +1,22 @@
+# Copyright (c) 2018-2024, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import warnings
+
+warnings.warn(
+    "The `rmm._lib` module is deprecated in will be removed in a future release. Use `rmm.python` instead.",
+    FutureWarning,
+)
+
+from rmm.python import *
diff --git a/python/rmm/rmm/tests/test_rmm.py b/python/rmm/rmm/tests/test_rmm.py
index 4d7b6c646..a1a973950 100644
--- a/python/rmm/rmm/tests/test_rmm.py
+++ b/python/rmm/rmm/tests/test_rmm.py
@@ -1076,3 +1076,8 @@ def test_available_device_memory():
     assert initial_memory[1] == final_memory[1]
     assert initial_memory[0] > 0
     assert final_memory[0] > 0
+
+
+def test_deprecate_rmm_lib():
+    with pytest.warns(FutureWarning):
+        rmm._lib.device_buffer.DeviceBuffer(size=100)