Refactor memory profiler plugin for pytest

Updated the memory profiler plugin to utilize GPU memory statistics and issue warnings for high memory usage during test execution. Removed the previous memory tracking methods and streamlined the reporting mechanism to focus on peak GPU memory usage.

Author: csadorf

python/cuml/cuml/testing/plugins/__init__.py

@@ -13,42 +13,34 @@
 # limitations under the License.
 #
 
-import os
-import psutil
+import warnings
+from rmm.statistics import statistics, get_statistics
+import pytest
 
-# Memory threshold in MB for reporting memory usage
-MEMORY_REPORT_THRESHOLD_MB = 1024
 
+class HighMemoryUsageWarning(UserWarning):
+    """Warning emitted when a test exceeds the memory usage threshold."""
 
-def get_process_memory():
-    """Get the current process memory usage in MB."""
-    process = psutil.Process(os.getpid())
-    return process.memory_info().rss / 1024 / 1024  # Convert to MB
+    pass
 
 
-class MemoryProfiler:
-    def __init__(self):
-        self.start_memory = None
-        self.max_memory = 0
+# Memory threshold in MB for reporting memory usage
+MEMORY_REPORT_THRESHOLD_MB = 1024
 
-    def pytest_runtest_setup(self, item):
-        """Record memory usage at test setup."""
-        self.start_memory = get_process_memory()
 
-    def pytest_runtest_teardown(self, item):
-        """Record memory usage at test teardown and report if significant."""
-        end_memory = get_process_memory()
-        if self.start_memory is not None:
-            memory_used = end_memory - self.start_memory
-            self.max_memory = max(self.max_memory, end_memory)
-            if memory_used > MEMORY_REPORT_THRESHOLD_MB:
-                print(f"\nMemory usage for {item.nodeid}:")
-                print(f"  Start: {self.start_memory:.2f} MB")
-                print(f"  End: {end_memory:.2f} MB")
-                print(f"  Delta: {memory_used:.2f} MB")
-                print(f"  Max: {self.max_memory:.2f} MB")
+@pytest.hookimpl(hookwrapper=True)
+def pytest_runtest_call(item):
+    """Wrap test execution with GPU memory profiler."""
+    with statistics():
+        yield
 
+        # Check memory usage after test completion
+        stats = get_statistics()
+        peak_memory_mb = stats.peak_bytes / (1024 * 1024)
 
-def pytest_configure(config):
-    """Register the memory profiler plugin."""
-    config.pluginmanager.register(MemoryProfiler())
+        if peak_memory_mb > MEMORY_REPORT_THRESHOLD_MB:
+            msg = (
+                f"Test {item.nodeid} used {peak_memory_mb:.2f} MB of GPU memory, "
+                f"exceeding threshold of {MEMORY_REPORT_THRESHOLD_MB} MB"
+            )
+            warnings.warn(msg, HighMemoryUsageWarning)