Restore tracked device resources on their original GPU

cpp/include/raft/core/memory_stats_resources.hpp

@@ -4,6 +4,7 @@
  */
 #pragma once
 
+#include <raft/core/resource/device_id.hpp>
 #include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/core/resource/managed_memory_resource.hpp>
 #include <raft/core/resource/pinned_memory_resource.hpp>
@@ -77,15 +78,17 @@ class memory_stats_resources : public resources {
   explicit memory_stats_resources(const resources& existing)
     : resources(existing),
       old_host_(mr::get_default_host_resource()),
-      old_device_(rmm::mr::get_current_device_resource_ref())
+      old_device_(rmm::mr::get_per_device_resource_ref(
+        rmm::cuda_device_id{resource::get_device_id(*this)}))
   {
     init();
   }
 
   ~memory_stats_resources() override
   {
     mr::set_default_host_resource(old_host_);
-    rmm::mr::set_current_device_resource(std::move(old_device_));
+    rmm::mr::set_per_device_resource(rmm::cuda_device_id{resource::get_device_id(*this)},
+                                     std::move(old_device_));
   }
 
   memory_stats_resources(memory_stats_resources const&)            = delete;

cpp/include/raft/core/memory_tracking_resources.hpp

@@ -5,6 +5,7 @@
 #pragma once
 
 #include <raft/core/detail/macros.hpp>
+#include <raft/core/resource/device_id.hpp>
 #include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/core/resource/managed_memory_resource.hpp>
 #include <raft/core/resource/pinned_memory_resource.hpp>
@@ -108,7 +109,8 @@ class memory_tracking_resources : public resources {
   {
     report_.stop();
     raft::mr::set_default_host_resource(old_host_);
-    rmm::mr::set_current_device_resource(old_device_);
+    rmm::mr::set_per_device_resource(rmm::cuda_device_id{resource::get_device_id(*this)},
+                                     std::move(old_device_));
   }
 
   memory_tracking_resources(memory_tracking_resources const&)            = delete;
@@ -128,7 +130,8 @@ class memory_tracking_resources : public resources {
       owned_stream_(std::move(owned_stream)),
       report_(out_override ? *out_override : *owned_stream_, sample_interval),
       old_host_(raft::mr::get_default_host_resource()),
-      old_device_(rmm::mr::get_current_device_resource_ref())
+      old_device_(rmm::mr::get_per_device_resource_ref(
+        rmm::cuda_device_id{resource::get_device_id(*this)}))
   {
     init();
   }

cpp/tests/core/memory_stats_resources.cpp

@@ -4,19 +4,37 @@
  */
 
 #include <raft/core/memory_stats_resources.hpp>
+#include <raft/core/device_setter.hpp>
 #include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/core/resources.hpp>
 
+#include <rmm/mr/cuda_memory_resource.hpp>
 #include <rmm/mr/per_device_resource.hpp>
+#include <rmm/mr/pool_memory_resource.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cuda/memory_resource>
 #include <cuda/stream_ref>
 
 #include <gtest/gtest.h>
 
 #include <cstddef>
+#include <memory>
 
 namespace raft {
+namespace {
+
+struct device_resource_restore_guard {
+  int device_id;
+  raft::mr::device_resource resource;
+
+  ~device_resource_restore_guard()
+  {
+    rmm::mr::set_per_device_resource(rmm::cuda_device_id{device_id}, std::move(resource));
+  }
+};
+
+}  // namespace
 
 TEST(MemoryStatsResources, IndependentCounting_DefaultWorkspace)
 {
@@ -93,4 +111,49 @@ TEST(MemoryStatsResources, IndependentCounting_PoolWorkspace)
   dev_mr.deallocate(cuda::stream_ref{cudaStreamLegacy}, dev_ptr, kGlobalSize);
 }
 
+TEST(MemoryStatsResources, RestoresDeviceResourceOnConstructionDevice)
+{
+  if (device_setter::get_device_count() < 2) {
+    GTEST_SKIP() << "Requires at least 2 CUDA devices";
+  }
+
+  auto device0 = 0;
+  auto device1 = 1;
+
+  auto device0_guard = [&]() {
+    auto scoped_device = device_setter{device0};
+    auto upstream      = rmm::mr::get_current_device_resource_ref();
+    return device_resource_restore_guard{
+      device0,
+      rmm::mr::set_current_device_resource(
+        raft::mr::device_resource{rmm::mr::pool_memory_resource(upstream, 1 << 20, 2 << 20)})};
+  }();
+
+  auto device1_guard = [&]() {
+    auto scoped_device = device_setter{device1};
+    return device_resource_restore_guard{device1, rmm::mr::reset_current_device_resource()};
+  }();
+
+  {
+    auto scoped_device = device_setter{device0};
+    raft::resources res;
+    auto tracked = std::make_unique<memory_stats_resources>(res);
+    auto wrong_device = device_setter{device1};
+    static_cast<void>(wrong_device);
+    tracked.reset();
+  }
+
+  {
+    auto scoped_device = device_setter{device0};
+    auto current_mr    = rmm::mr::get_current_device_resource_ref();
+    EXPECT_NE(cuda::mr::resource_cast<rmm::mr::pool_memory_resource>(&current_mr), nullptr);
+  }
+
+  {
+    auto scoped_device = device_setter{device1};
+    auto current_mr    = rmm::mr::get_current_device_resource_ref();
+    EXPECT_NE(cuda::mr::resource_cast<rmm::mr::cuda_memory_resource>(&current_mr), nullptr);
+  }
+}
+
 }  // namespace raft

cpp/tests/core/monitor_resources.cu

@@ -4,19 +4,36 @@
  */
 
 #include <raft/core/device_mdarray.hpp>
+#include <raft/core/device_setter.hpp>
 #include <raft/core/memory_tracking_resources.hpp>
 #include <raft/core/resources.hpp>
 
+#include <rmm/mr/cuda_memory_resource.hpp>
+#include <rmm/mr/per_device_resource.hpp>
+#include <rmm/mr/pool_memory_resource.hpp>
+
+#include <cuda/memory_resource>
 #include <gtest/gtest.h>
 
 #include <algorithm>
 #include <chrono>
+#include <memory>
 #include <sstream>
 #include <string>
 #include <thread>
 
 namespace {
 
+struct device_resource_restore_guard {
+  int device_id;
+  raft::mr::device_resource resource;
+
+  ~device_resource_restore_guard()
+  {
+    rmm::mr::set_per_device_resource(rmm::cuda_device_id{device_id}, std::move(resource));
+  }
+};
+
 TEST(MemoryTrackingResources, TracksDeviceAllocations)
 {
   using namespace std::chrono_literals;
@@ -49,4 +66,49 @@ TEST(MemoryTrackingResources, TracksDeviceAllocations)
                           << output;
 }
 
+TEST(MemoryTrackingResources, RestoresDeviceResourceOnConstructionDevice)
+{
+  if (raft::device_setter::get_device_count() < 2) {
+    GTEST_SKIP() << "Requires at least 2 CUDA devices";
+  }
+
+  auto device0 = 0;
+  auto device1 = 1;
+
+  auto device0_guard = [&]() {
+    auto scoped_device = raft::device_setter{device0};
+    auto upstream      = rmm::mr::get_current_device_resource_ref();
+    return device_resource_restore_guard{
+      device0,
+      rmm::mr::set_current_device_resource(
+        raft::mr::device_resource{rmm::mr::pool_memory_resource(upstream, 1 << 20, 2 << 20)})};
+  }();
+
+  auto device1_guard = [&]() {
+    auto scoped_device = raft::device_setter{device1};
+    return device_resource_restore_guard{device1, rmm::mr::reset_current_device_resource()};
+  }();
+
+  {
+    auto scoped_device = raft::device_setter{device0};
+    std::ostringstream oss;
+    auto tracked = std::make_unique<raft::memory_tracking_resources>(oss);
+    auto wrong_device = raft::device_setter{device1};
+    static_cast<void>(wrong_device);
+    tracked.reset();
+  }
+
+  {
+    auto scoped_device = raft::device_setter{device0};
+    auto current_mr    = rmm::mr::get_current_device_resource_ref();
+    EXPECT_NE(cuda::mr::resource_cast<rmm::mr::pool_memory_resource>(&current_mr), nullptr);
+  }
+
+  {
+    auto scoped_device = raft::device_setter{device1};
+    auto current_mr    = rmm::mr::get_current_device_resource_ref();
+    EXPECT_NE(cuda::mr::resource_cast<rmm::mr::cuda_memory_resource>(&current_mr), nullptr);
+  }
+}
+
 }  // namespace