From b0ddd707dbad3e7e34eed17f64103e8506814359 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Wed, 8 Jul 2026 21:29:51 +0200 Subject: [PATCH 1/3] Restore tracked device resources on their original GPU --- .../raft/core/memory_stats_resources.hpp | 5 +- .../raft/core/memory_tracking_resources.hpp | 5 +- cpp/tests/core/memory_stats_resources.cpp | 63 +++++++++++++++++++ cpp/tests/core/monitor_resources.cu | 62 ++++++++++++++++++ 4 files changed, 133 insertions(+), 2 deletions(-) diff --git a/cpp/include/raft/core/memory_stats_resources.hpp b/cpp/include/raft/core/memory_stats_resources.hpp index 136b48beb6..12e56aecb3 100644 --- a/cpp/include/raft/core/memory_stats_resources.hpp +++ b/cpp/include/raft/core/memory_stats_resources.hpp @@ -4,6 +4,7 @@ */ #pragma once +#include #include #include #include @@ -76,6 +77,7 @@ class memory_stats_resources : public resources { public: explicit memory_stats_resources(const resources& existing) : resources(existing), + device_id_(device_setter::get_current_device()), old_host_(mr::get_default_host_resource()), old_device_(rmm::mr::get_current_device_resource_ref()) { @@ -85,7 +87,7 @@ class memory_stats_resources : public resources { ~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{device_id_}, std::move(old_device_)); } memory_stats_resources(memory_stats_resources const&) = delete; @@ -143,6 +145,7 @@ class memory_stats_resources : public resources { std::vector> snapshot_; + int device_id_; raft::mr::host_resource old_host_; raft::mr::device_resource old_device_; diff --git a/cpp/include/raft/core/memory_tracking_resources.hpp b/cpp/include/raft/core/memory_tracking_resources.hpp index d06e14ef0e..3cbabc9334 100644 --- a/cpp/include/raft/core/memory_tracking_resources.hpp +++ b/cpp/include/raft/core/memory_tracking_resources.hpp @@ -5,6 +5,7 @@ #pragma once #include +#include #include #include #include @@ -108,7 +109,7 @@ 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{device_id_}, std::move(old_device_)); } memory_tracking_resources(memory_tracking_resources const&) = delete; @@ -127,6 +128,7 @@ class memory_tracking_resources : public resources { : resources(existing ? *existing : resources{}), owned_stream_(std::move(owned_stream)), report_(out_override ? *out_override : *owned_stream_, sample_interval), + device_id_(device_setter::get_current_device()), old_host_(raft::mr::get_default_host_resource()), old_device_(rmm::mr::get_current_device_resource_ref()) { @@ -141,6 +143,7 @@ class memory_tracking_resources : public resources { std::unique_ptr owned_stream_; raft::mr::resource_monitor report_; + int device_id_; raft::mr::host_resource old_host_; raft::mr::device_resource old_device_; diff --git a/cpp/tests/core/memory_stats_resources.cpp b/cpp/tests/core/memory_stats_resources.cpp index be615c9041..6b84a85cd5 100644 --- a/cpp/tests/core/memory_stats_resources.cpp +++ b/cpp/tests/core/memory_stats_resources.cpp @@ -4,19 +4,37 @@ */ #include +#include #include #include +#include #include +#include #include +#include #include #include #include +#include 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(res); + auto wrong_device = device_setter{device1}; + static_cast(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(¤t_mr), nullptr); + } + + { + auto scoped_device = device_setter{device1}; + auto current_mr = rmm::mr::get_current_device_resource_ref(); + EXPECT_NE(cuda::mr::resource_cast(¤t_mr), nullptr); + } +} + } // namespace raft diff --git a/cpp/tests/core/monitor_resources.cu b/cpp/tests/core/monitor_resources.cu index c7be4bd285..554f191385 100644 --- a/cpp/tests/core/monitor_resources.cu +++ b/cpp/tests/core/monitor_resources.cu @@ -4,19 +4,36 @@ */ #include +#include #include #include +#include +#include +#include + +#include #include #include #include +#include #include #include #include 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(oss); + auto wrong_device = raft::device_setter{device1}; + static_cast(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(¤t_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(¤t_mr), nullptr); + } +} + } // namespace From 7980f7c7cbfb6dd9713068ca0e2daebd473530d8 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Thu, 9 Jul 2026 19:07:11 +0200 Subject: [PATCH 2/3] Use the handle device-id resource for MR restore --- cpp/include/raft/core/memory_stats_resources.hpp | 10 +++++----- cpp/include/raft/core/memory_tracking_resources.hpp | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/cpp/include/raft/core/memory_stats_resources.hpp b/cpp/include/raft/core/memory_stats_resources.hpp index 12e56aecb3..1afdc82066 100644 --- a/cpp/include/raft/core/memory_stats_resources.hpp +++ b/cpp/include/raft/core/memory_stats_resources.hpp @@ -4,7 +4,7 @@ */ #pragma once -#include +#include #include #include #include @@ -77,9 +77,9 @@ class memory_stats_resources : public resources { public: explicit memory_stats_resources(const resources& existing) : resources(existing), - device_id_(device_setter::get_current_device()), 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(); } @@ -87,7 +87,8 @@ class memory_stats_resources : public resources { ~memory_stats_resources() override { mr::set_default_host_resource(old_host_); - rmm::mr::set_per_device_resource(rmm::cuda_device_id{device_id_}, 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; @@ -145,7 +146,6 @@ class memory_stats_resources : public resources { std::vector> snapshot_; - int device_id_; raft::mr::host_resource old_host_; raft::mr::device_resource old_device_; diff --git a/cpp/include/raft/core/memory_tracking_resources.hpp b/cpp/include/raft/core/memory_tracking_resources.hpp index 3cbabc9334..a358336862 100644 --- a/cpp/include/raft/core/memory_tracking_resources.hpp +++ b/cpp/include/raft/core/memory_tracking_resources.hpp @@ -5,7 +5,7 @@ #pragma once #include -#include +#include #include #include #include @@ -109,7 +109,8 @@ class memory_tracking_resources : public resources { { report_.stop(); raft::mr::set_default_host_resource(old_host_); - rmm::mr::set_per_device_resource(rmm::cuda_device_id{device_id_}, std::move(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,9 +129,9 @@ class memory_tracking_resources : public resources { : resources(existing ? *existing : resources{}), owned_stream_(std::move(owned_stream)), report_(out_override ? *out_override : *owned_stream_, sample_interval), - device_id_(device_setter::get_current_device()), 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(); } @@ -143,7 +144,6 @@ class memory_tracking_resources : public resources { std::unique_ptr owned_stream_; raft::mr::resource_monitor report_; - int device_id_; raft::mr::host_resource old_host_; raft::mr::device_resource old_device_; From 381f2d18cc86a188c96924bf4e7e5ae55fde75de Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Thu, 9 Jul 2026 19:22:32 +0200 Subject: [PATCH 3/3] Install tracked resources on the handle device --- .../raft/core/memory_stats_resources.hpp | 3 +- .../raft/core/memory_tracking_resources.hpp | 3 +- cpp/tests/core/memory_stats_resources.cpp | 68 ++++++++++++++++++ cpp/tests/core/monitor_resources.cu | 69 +++++++++++++++++++ 4 files changed, 141 insertions(+), 2 deletions(-) diff --git a/cpp/include/raft/core/memory_stats_resources.hpp b/cpp/include/raft/core/memory_stats_resources.hpp index 1afdc82066..38eaa12270 100644 --- a/cpp/include/raft/core/memory_stats_resources.hpp +++ b/cpp/include/raft/core/memory_stats_resources.hpp @@ -218,7 +218,8 @@ class memory_stats_resources : public resources { device_stats_adaptor_t sa{rmm::device_async_resource_ref{old_device_}}; device_stats_ = sa.get_stats(); device_adaptor_ = std::make_unique(std::move(sa)); - rmm::mr::set_current_device_resource(*device_adaptor_); + rmm::mr::set_per_device_resource(rmm::cuda_device_id{resource::get_device_id(*this)}, + *device_adaptor_); } // --- Workspace --- { diff --git a/cpp/include/raft/core/memory_tracking_resources.hpp b/cpp/include/raft/core/memory_tracking_resources.hpp index a358336862..d3a36e4b66 100644 --- a/cpp/include/raft/core/memory_tracking_resources.hpp +++ b/cpp/include/raft/core/memory_tracking_resources.hpp @@ -220,7 +220,8 @@ class memory_tracking_resources : public resources { device_stats_t sa{rmm::device_async_resource_ref{old_device_}}; report_.register_source("device", sa.get_stats()); device_adaptor_ = std::make_unique(std::move(sa), report_.get_notifier()); - rmm::mr::set_current_device_resource(*device_adaptor_); + rmm::mr::set_per_device_resource(rmm::cuda_device_id{resource::get_device_id(*this)}, + *device_adaptor_); } // --- Workspace (track upstream to preserve limiting_resource_adaptor) --- diff --git a/cpp/tests/core/memory_stats_resources.cpp b/cpp/tests/core/memory_stats_resources.cpp index 6b84a85cd5..0d16f93b64 100644 --- a/cpp/tests/core/memory_stats_resources.cpp +++ b/cpp/tests/core/memory_stats_resources.cpp @@ -34,6 +34,18 @@ struct device_resource_restore_guard { } }; +auto current_device_uses_pool_resource() -> bool +{ + auto current_mr = rmm::mr::get_current_device_resource_ref(); + return cuda::mr::resource_cast(¤t_mr) != nullptr; +} + +auto current_device_uses_default_cuda_resource() -> bool +{ + auto current_mr = rmm::mr::get_current_device_resource_ref(); + return cuda::mr::resource_cast(¤t_mr) != nullptr; +} + } // namespace TEST(MemoryStatsResources, IndependentCounting_DefaultWorkspace) @@ -156,4 +168,60 @@ TEST(MemoryStatsResources, RestoresDeviceResourceOnConstructionDevice) } } +TEST(MemoryStatsResources, InstallsTrackedResourceOnHandleDevice) +{ + 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; + static_cast(resource::get_device_id(res)); + + auto wrong_device = device_setter{device1}; + static_cast(wrong_device); + auto tracked = std::make_unique(res); + + { + auto verify_device0 = device_setter{device0}; + EXPECT_FALSE(current_device_uses_pool_resource()); + } + + { + auto verify_device1 = device_setter{device1}; + EXPECT_TRUE(current_device_uses_default_cuda_resource()); + } + + tracked.reset(); + } + + { + auto scoped_device = device_setter{device0}; + EXPECT_TRUE(current_device_uses_pool_resource()); + } + + { + auto scoped_device = device_setter{device1}; + EXPECT_TRUE(current_device_uses_default_cuda_resource()); + } +} + } // namespace raft diff --git a/cpp/tests/core/monitor_resources.cu b/cpp/tests/core/monitor_resources.cu index 554f191385..94f7ff9904 100644 --- a/cpp/tests/core/monitor_resources.cu +++ b/cpp/tests/core/monitor_resources.cu @@ -34,6 +34,18 @@ struct device_resource_restore_guard { } }; +auto current_device_uses_pool_resource() -> bool +{ + auto current_mr = rmm::mr::get_current_device_resource_ref(); + return cuda::mr::resource_cast(¤t_mr) != nullptr; +} + +auto current_device_uses_default_cuda_resource() -> bool +{ + auto current_mr = rmm::mr::get_current_device_resource_ref(); + return cuda::mr::resource_cast(¤t_mr) != nullptr; +} + TEST(MemoryTrackingResources, TracksDeviceAllocations) { using namespace std::chrono_literals; @@ -111,4 +123,61 @@ TEST(MemoryTrackingResources, RestoresDeviceResourceOnConstructionDevice) } } +TEST(MemoryTrackingResources, InstallsTrackedResourceOnHandleDevice) +{ + 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}; + raft::resources res; + static_cast(raft::resource::get_device_id(res)); + + auto wrong_device = raft::device_setter{device1}; + static_cast(wrong_device); + std::ostringstream oss; + auto tracked = std::make_unique(res, oss); + + { + auto verify_device0 = raft::device_setter{device0}; + EXPECT_FALSE(current_device_uses_pool_resource()); + } + + { + auto verify_device1 = raft::device_setter{device1}; + EXPECT_TRUE(current_device_uses_default_cuda_resource()); + } + + tracked.reset(); + } + + { + auto scoped_device = raft::device_setter{device0}; + EXPECT_TRUE(current_device_uses_pool_resource()); + } + + { + auto scoped_device = raft::device_setter{device1}; + EXPECT_TRUE(current_device_uses_default_cuda_resource()); + } +} + } // namespace