Predictable raft::resources

cpp/.clang-format

@@ -56,7 +56,7 @@ BreakConstructorInitializers: BeforeColon
 BreakInheritanceList: BeforeColon
 BreakStringLiterals: true
 ColumnLimit: 100
-CommentPragmas: '^ IWYU pragma:'
+CommentPragmas: '^ ?[*]? ?(IWYU pragma:|SPDX-)'
 CompactNamespaces: false
 ConstructorInitializerAllOnOneLineOrOnePerLine: true
 # Kept the below 2 to be the same as `IndentWidth` to keep everything uniform

cpp/CMakeLists.txt

@@ -1,6 +1,6 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 cmake_minimum_required(VERSION 4.0 FATAL_ERROR)
@@ -191,7 +191,7 @@ target_include_directories(
 # Keep RAFT as lightweight as possible. Only CUDA libs and rmm should be used in global target.
 target_link_libraries(raft INTERFACE rapids_logger::rapids_logger rmm::rmm CCCL::CCCL)
 
-target_compile_features(raft INTERFACE cxx_std_17 $<BUILD_INTERFACE:cuda_std_17>)
+target_compile_features(raft INTERFACE cxx_std_20 $<BUILD_INTERFACE:cuda_std_20>)
 target_compile_options(
   raft INTERFACE $<$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>:--expt-extended-lambda
                  --expt-relaxed-constexpr>

cpp/include/raft/core/device_resources.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -37,10 +37,8 @@
 #include <cusparse.h>
 
 #include <memory>
-#include <mutex>
 #include <optional>
 #include <string>
-#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -60,9 +58,10 @@ class device_resources : public resources {
     resource::set_workspace_resource(*this, std::move(workspace_resource), allocation_limit);
   }
 
-  device_resources(const device_resources& handle) : resources{handle} {}
-  device_resources(device_resources&&)            = delete;
-  device_resources& operator=(device_resources&&) = delete;
+  device_resources(const device_resources&)            = default;
+  device_resources(device_resources&&)                 = default;
+  device_resources& operator=(const device_resources&) = default;
+  device_resources& operator=(device_resources&&)      = default;
 
   /**
    * @brief Construct a resources instance with a stream view and stream pool

cpp/include/raft/core/handle.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -27,10 +27,10 @@ class handle_t : public raft::device_resources {
   {
   }
 
-  handle_t(const handle_t& handle) : device_resources{handle} {}
-
-  handle_t(handle_t&&)            = delete;
-  handle_t& operator=(handle_t&&) = delete;
+  handle_t(const handle_t&)            = default;
+  handle_t(handle_t&&)                 = default;
+  handle_t& operator=(const handle_t&) = default;
+  handle_t& operator=(handle_t&&)      = default;
 
   /**
    * @brief Construct a resources instance with a stream view and stream pool

cpp/include/raft/core/memory_stats_resources.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -141,7 +141,7 @@ class memory_stats_resources : public resources {
     };
   }
 
-  std::vector<pair_resource> snapshot_;
+  std::vector<std::shared_ptr<resource::resource_cell>> snapshot_;
 
   raft::mr::host_resource old_host_;
   raft::mr::device_resource old_device_;
@@ -182,7 +182,7 @@ class memory_stats_resources : public resources {
     auto pinned_ref  = resource::get_pinned_memory_resource_ref(*this);
     auto managed_ref = resource::get_managed_memory_resource_ref(*this);
 
-    snapshot_ = resources_;
+    snapshot_ = cells_;
 
     // --- Host (global) ---
     {
@@ -207,11 +207,10 @@ class memory_stats_resources : public resources {
 
     // --- Device (global) ---
     // Invalidate the cached thrust policy (the resource_ref it captured
-    // will be stale once we replace the global device resource).
-    factories_.at(resource::resource_type::THRUST_POLICY) = std::make_pair(
-      resource::resource_type::LAST_KEY, std::make_shared<resource::empty_resource_factory>());
-    resources_.at(resource::resource_type::THRUST_POLICY) = std::make_pair(
-      resource::resource_type::LAST_KEY, std::make_shared<resource::empty_resource>());
+    // will be stale once we replace the global device resource).  Swapping in a
+    // fresh cell drops the cached factory/resource locally while snapshot_ keeps
+    // the originals alive, so it gets lazily rebuilt against the new device MR.
+    cells_[resource::resource_type::THRUST_POLICY] = std::make_shared<resource::resource_cell>();
     {
       device_stats_adaptor_t sa{rmm::device_async_resource_ref{old_device_}};
       device_stats_   = sa.get_stats();

cpp/include/raft/core/memory_tracking_resources.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -137,7 +137,7 @@ class memory_tracking_resources : public resources {
   // snapshot_ is destroyed last  (keeps original resource shared_ptrs alive).
   // owned_stream_ outlives report_ (report_ writes to it).
   // report_ is destroyed first of the three (stops background thread).
-  std::vector<pair_resource> snapshot_;
+  std::vector<std::shared_ptr<resource::resource_cell>> snapshot_;
   std::unique_ptr<std::ofstream> owned_stream_;
   raft::mr::resource_monitor report_;
 
@@ -177,7 +177,7 @@ class memory_tracking_resources : public resources {
     auto managed_ref  = raft::resource::get_managed_memory_resource_ref(*this);
 
     // Keeps original resource objects alive while tracking refs point into them.
-    snapshot_ = resources_;
+    snapshot_ = cells_;
 
     // --- Host (global) ---
     {
@@ -209,11 +209,10 @@ class memory_tracking_resources : public resources {
 
     // --- Device (global) ---
     // Invalidate the cached thrust policy (the resource_ref it captured
-    // will be stale once we replace the global device resource).
-    factories_.at(resource::resource_type::THRUST_POLICY) = std::make_pair(
-      resource::resource_type::LAST_KEY, std::make_shared<resource::empty_resource_factory>());
-    resources_.at(resource::resource_type::THRUST_POLICY) = std::make_pair(
-      resource::resource_type::LAST_KEY, std::make_shared<resource::empty_resource>());
+    // will be stale once we replace the global device resource).  Swapping in a
+    // fresh cell drops the cached factory/resource locally while snapshot_ keeps
+    // the originals alive, so it gets lazily rebuilt against the new device MR.
+    cells_[resource::resource_type::THRUST_POLICY] = std::make_shared<resource::resource_cell>();
     {
       device_stats_t sa{rmm::device_async_resource_ref{old_device_}};
       report_.register_source("device", sa.get_stats());

cpp/include/raft/core/resource/comms.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -58,7 +58,7 @@ inline comms::comms_t const& get_comms(resources const& res)
   return *(*res.get_resource<std::shared_ptr<comms::comms_t>>(resource_type::COMMUNICATOR));
 }
 
-inline void set_comms(resources const& res, std::shared_ptr<comms::comms_t> communicator)
+inline void set_comms(resources& res, std::shared_ptr<comms::comms_t> communicator)
 {
   res.add_resource_factory(std::make_shared<comms_resource_factory>(communicator));
 }

cpp/include/raft/core/resource/cublas_handle.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -60,8 +60,7 @@ class cublas_resource_factory : public resource_factory {
 inline cublasHandle_t get_cublas_handle(resources const& res)
 {
   if (!res.has_resource_factory(resource_type::CUBLAS_HANDLE)) {
-    cudaStream_t stream = get_cuda_stream(res);
-    res.add_resource_factory(std::make_shared<cublas_resource_factory>(stream));
+    res.ensure_default_factory(std::make_shared<cublas_resource_factory>(get_cuda_stream(res)));
   }
   auto ret = *res.get_resource<cublasHandle_t>(resource_type::CUBLAS_HANDLE);
   RAFT_CUBLAS_TRY(cublasSetStream(ret, get_cuda_stream(res)));

cpp/include/raft/core/resource/cublaslt_handle.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -46,11 +46,8 @@ class cublaslt_resource_factory : public resource_factory {
  */
 inline auto get_cublaslt_handle(resources const& res) -> cublasLtHandle_t
 {
-  if (!res.has_resource_factory(resource_type::CUBLASLT_HANDLE)) {
-    res.add_resource_factory(std::make_shared<cublaslt_resource_factory>());
-  }
-  auto ret = *res.get_resource<cublasLtHandle_t>(resource_type::CUBLASLT_HANDLE);
-  return ret;
+  res.ensure_default_factory(std::make_shared<cublaslt_resource_factory>());
+  return *res.get_resource<cublasLtHandle_t>(resource_type::CUBLASLT_HANDLE);
 };
 
 /**

cpp/include/raft/core/resource/cuda_stream.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -59,9 +59,7 @@ class cuda_stream_resource_factory : public resource_factory {
  */
 inline rmm::cuda_stream_view get_cuda_stream(resources const& res)
 {
-  if (!res.has_resource_factory(resource_type::CUDA_STREAM_VIEW)) {
-    res.add_resource_factory(std::make_shared<cuda_stream_resource_factory>());
-  }
+  res.ensure_default_factory(std::make_shared<cuda_stream_resource_factory>());
   return *res.get_resource<rmm::cuda_stream_view>(resource_type::CUDA_STREAM_VIEW);
 };
 
@@ -71,7 +69,7 @@ inline rmm::cuda_stream_view get_cuda_stream(resources const& res)
  * @param[in] res raft resources object for managing resources
  * @param[in] stream_view cuda stream view
  */
-inline void set_cuda_stream(resources const& res, rmm::cuda_stream_view stream_view)
+inline void set_cuda_stream(resources& res, rmm::cuda_stream_view stream_view)
 {
   res.add_resource_factory(std::make_shared<cuda_stream_resource_factory>(stream_view));
 };

cpp/include/raft/core/resource/cuda_stream_pool.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -69,9 +69,7 @@ inline bool is_stream_pool_initialized(const resources& res)
  */
 inline const rmm::cuda_stream_pool& get_cuda_stream_pool(const resources& res)
 {
-  if (!res.has_resource_factory(resource_type::CUDA_STREAM_POOL)) {
-    res.add_resource_factory(std::make_shared<cuda_stream_pool_resource_factory>());
-  }
+  res.ensure_default_factory(std::make_shared<cuda_stream_pool_resource_factory>());
   return *(
     *res.get_resource<std::shared_ptr<rmm::cuda_stream_pool>>(resource_type::CUDA_STREAM_POOL));
 };
@@ -82,8 +80,7 @@ inline const rmm::cuda_stream_pool& get_cuda_stream_pool(const resources& res)
  * @param res
  * @param stream_pool
  */
-inline void set_cuda_stream_pool(const resources& res,
-                                 std::shared_ptr<rmm::cuda_stream_pool> stream_pool)
+inline void set_cuda_stream_pool(resources& res, std::shared_ptr<rmm::cuda_stream_pool> stream_pool)
 {
   res.add_resource_factory(std::make_shared<cuda_stream_pool_resource_factory>(stream_pool));
 };
@@ -165,9 +162,7 @@ inline void sync_stream_pool(const resources& res, const std::vector<std::size_t
  */
 inline void wait_stream_pool_on_stream(const resources& res)
 {
-  if (!res.has_resource_factory(resource_type::CUDA_STREAM_POOL)) {
-    res.add_resource_factory(std::make_shared<cuda_stream_pool_resource_factory>());
-  }
+  res.ensure_default_factory(std::make_shared<cuda_stream_pool_resource_factory>());
 
   cudaEvent_t event = detail::get_cuda_stream_sync_event(res);
   RAFT_CUDA_TRY(cudaEventRecord(event, get_cuda_stream(res)));

cpp/include/raft/core/resource/cusolver_dn_handle.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -66,8 +66,8 @@ class cusolver_dn_resource_factory : public resource_factory {
 inline cusolverDnHandle_t get_cusolver_dn_handle(resources const& res)
 {
   if (!res.has_resource_factory(resource_type::CUSOLVER_DN_HANDLE)) {
-    cudaStream_t stream = get_cuda_stream(res);
-    res.add_resource_factory(std::make_shared<cusolver_dn_resource_factory>(stream));
+    res.ensure_default_factory(
+      std::make_shared<cusolver_dn_resource_factory>(get_cuda_stream(res)));
   }
   return *res.get_resource<cusolverDnHandle_t>(resource_type::CUSOLVER_DN_HANDLE);
 };

cpp/include/raft/core/resource/cusolver_sp_handle.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -63,8 +63,8 @@ class cusolver_sp_resource_factory : public resource_factory {
 inline cusolverSpHandle_t get_cusolver_sp_handle(resources const& res)
 {
   if (!res.has_resource_factory(resource_type::CUSOLVER_SP_HANDLE)) {
-    cudaStream_t stream = get_cuda_stream(res);
-    res.add_resource_factory(std::make_shared<cusolver_sp_resource_factory>(stream));
+    res.ensure_default_factory(
+      std::make_shared<cusolver_sp_resource_factory>(get_cuda_stream(res)));
   }
   return *res.get_resource<cusolverSpHandle_t>(resource_type::CUSOLVER_SP_HANDLE);
 };

cpp/include/raft/core/resource/cusparse_handle.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -58,8 +58,7 @@ class cusparse_resource_factory : public resource_factory {
 inline cusparseHandle_t get_cusparse_handle(resources const& res)
 {
   if (!res.has_resource_factory(resource_type::CUSPARSE_HANDLE)) {
-    rmm::cuda_stream_view stream = get_cuda_stream(res);
-    res.add_resource_factory(std::make_shared<cusparse_resource_factory>(stream));
+    res.ensure_default_factory(std::make_shared<cusparse_resource_factory>(get_cuda_stream(res)));
   }
   return *res.get_resource<cusparseHandle_t>(resource_type::CUSPARSE_HANDLE);
 };

cpp/include/raft/core/resource/custom_resource.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -71,9 +71,7 @@ template <typename ResourceT>
 auto get_custom_resource(resources const& res) -> ResourceT*
 {
   static_assert(std::is_default_constructible_v<ResourceT>);
-  if (!res.has_resource_factory(resource_type::CUSTOM)) {
-    res.add_resource_factory(std::make_shared<custom_resource_factory>());
-  }
+  res.ensure_default_factory(std::make_shared<custom_resource_factory>());
   return res.get_resource<custom_resource>(resource_type::CUSTOM)->load<ResourceT>();
 };
 

cpp/include/raft/core/resource/detail/stream_sync_event.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -33,9 +33,7 @@ class cuda_stream_sync_event_resource_factory : public resource_factory {
  */
 inline cudaEvent_t& get_cuda_stream_sync_event(resources const& res)
 {
-  if (!res.has_resource_factory(resource_type::CUDA_STREAM_SYNC_EVENT)) {
-    res.add_resource_factory(std::make_shared<cuda_stream_sync_event_resource_factory>());
-  }
+  res.ensure_default_factory(std::make_shared<cuda_stream_sync_event_resource_factory>());
   return *res.get_resource<cudaEvent_t>(resource_type::CUDA_STREAM_SYNC_EVENT);
 };
 

cpp/include/raft/core/resource/device_id.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -55,9 +55,7 @@ class device_id_resource_factory : public resource_factory {
  */
 inline int get_device_id(resources const& res)
 {
-  if (!res.has_resource_factory(resource_type::DEVICE_ID)) {
-    res.add_resource_factory(std::make_shared<device_id_resource_factory>());
-  }
+  res.ensure_default_factory(std::make_shared<device_id_resource_factory>());
   return *res.get_resource<int>(resource_type::DEVICE_ID);
 };