clean up

rapidsai · May 17, 2024 · 0e2f19c · 0e2f19c
1 parent b9d57db
commit 0e2f19c
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 41 deletions.
diff --git a/include/rmm/mr/device/statistics_resource_adaptor.hpp b/include/rmm/mr/device/statistics_resource_adaptor.hpp
@@ -99,7 +99,6 @@ class statistics_resource_adaptor final : public device_memory_resource {
      */
     counter& operator+=(const counter& val)
     {
-      // We count the peak from value
       peak = std::max(value + val.peak, peak);
       value += val.value;
       total += val.total;

diff --git a/python/rmm/rmm/tests/test_statistics.py b/python/rmm/rmm/tests/test_statistics.py
@@ -57,46 +57,45 @@ def test_multiple_mr(stats_mr):
     # Push a new Tracking adaptor
     mr2 = rmm.mr.StatisticsResourceAdaptor(stats_mr)
     rmm.mr.set_current_device_resource(mr2)
-
-    for _ in range(2):
-        buffers.append(rmm.DeviceBuffer(size=1000))
-
-    assert mr2.allocation_counts == {
-        "current_bytes": 2016,
-        "current_count": 2,
-        "peak_bytes": 2016,
-        "peak_count": 2,
-        "total_bytes": 2016,
-        "total_count": 2,
-    }
-    assert stats_mr.allocation_counts == {
-        "current_bytes": 7056,
-        "current_count": 7,
-        "peak_bytes": 10080,
-        "peak_count": 10,
-        "total_bytes": 12096,
-        "total_count": 12,
-    }
-
-    del buffers
-
-    assert mr2.allocation_counts == {
-        "current_bytes": 0,
-        "current_count": 0,
-        "peak_bytes": 2016,
-        "peak_count": 2,
-        "total_bytes": 2016,
-        "total_count": 2,
-    }
-    assert stats_mr.allocation_counts == {
-        "current_bytes": 0,
-        "current_count": 0,
-        "peak_bytes": 10080,
-        "peak_count": 10,
-        "total_bytes": 12096,
-        "total_count": 12,
-    }
-    rmm.mr.set_current_device_resource(stats_mr)
+    try:
+        for _ in range(2):
+            buffers.append(rmm.DeviceBuffer(size=1000))
+
+        assert mr2.allocation_counts == {
+            "current_bytes": 2016,
+            "current_count": 2,
+            "peak_bytes": 2016,
+            "peak_count": 2,
+            "total_bytes": 2016,
+            "total_count": 2,
+        }
+        assert stats_mr.allocation_counts == {
+            "current_bytes": 7056,
+            "current_count": 7,
+            "peak_bytes": 10080,
+            "peak_count": 10,
+            "total_bytes": 12096,
+            "total_count": 12,
+        }
+        del buffers
+        assert mr2.allocation_counts == {
+            "current_bytes": 0,
+            "current_count": 0,
+            "peak_bytes": 2016,
+            "peak_count": 2,
+            "total_bytes": 2016,
+            "total_count": 2,
+        }
+        assert stats_mr.allocation_counts == {
+            "current_bytes": 0,
+            "current_count": 0,
+            "peak_bytes": 10080,
+            "peak_count": 10,
+            "total_bytes": 12096,
+            "total_count": 12,
+        }
+    finally:
+        rmm.mr.set_current_device_resource(stats_mr)
 
 
 def test_counter_stack(stats_mr):