Fix pool_memory_resource crash during process exit

cpp/CMakeLists.txt

@@ -120,7 +120,8 @@ add_library(
   src/mr/statistics_resource_adaptor.cpp
   src/mr/thread_safe_resource_adaptor.cpp
   src/mr/tracking_resource_adaptor.cpp
-  src/prefetch.cpp)
+  src/prefetch.cpp
+  src/runtime_shutdown.cpp)
 add_library(rmm::rmm ALIAS rmm)
 
 target_include_directories(

cpp/include/rmm/detail/runtime_shutdown.hpp

@@ -0,0 +1,22 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <rmm/detail/export.hpp>
+
+namespace RMM_NAMESPACE {
+namespace detail {
+
+/**
+ * @brief Register the atexit callback that flips the flag observed by `rmm::process_is_exiting()`.
+ *
+ * This registers the single process-exit hook used to make resources held in RMM's internal
+ * per-device resource map safe to destruct during process termination. It is not a general
+ * per-static-object registration facility.
+ */
+RMM_EXPORT void register_process_exit_hook() noexcept;
+
+}  // namespace detail
+}  // namespace RMM_NAMESPACE

cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp

@@ -10,6 +10,7 @@
 #include <rmm/detail/export.hpp>
 #include <rmm/detail/format.hpp>
 #include <rmm/logger.hpp>
+#include <rmm/process_is_exiting.hpp>
 
 #include <cuda/stream_ref>
 #include <cuda_runtime_api.h>
@@ -486,6 +487,8 @@ class stream_ordered_memory_resource : public crtp<PoolResource> {
   {
     lock_guard lock(mtx_);
 
+    if (rmm::process_is_exiting()) { return; }
+
     for (auto s_e : stream_events_) {
       RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaEventSynchronize(s_e.second.event));
       RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaEventDestroy(s_e.second.event));

cpp/include/rmm/mr/per_device_resource.hpp

@@ -7,6 +7,7 @@
 
 #include <rmm/cuda_device.hpp>
 #include <rmm/detail/export.hpp>
+#include <rmm/detail/runtime_shutdown.hpp>
 #include <rmm/mr/cuda_memory_resource.hpp>
 #include <rmm/resource_ref.hpp>
 
@@ -76,6 +77,8 @@ RMM_EXPORT inline auto& get_ref_map()
 {
   static std::map<cuda_device_id::value_type, cuda::mr::any_resource<cuda::mr::device_accessible>>
     device_id_to_resource;
+  // Register the process-exit hook immediately after constructing the map.
+  rmm::detail::register_process_exit_hook();
   return device_id_to_resource;
 }
 
@@ -136,6 +139,10 @@ inline device_async_resource_ref get_per_device_resource_ref(cuda_device_id devi
  * resource is undefined if used while the active CUDA device is a different device from the one
  * that was active when the memory resource was created.
  *
+ * @note The per-device resource map keeps the provided resource alive until process exit. Its
+ * destructor may therefore run during process termination. If the destructor may call CUDA APIs,
+ * it must consult `rmm::process_is_exiting()` and skip those calls when it returns `true`.
+ *
  * @param device_id The id of the target device
  * @param new_resource New resource to use for `device_id`
  * @return An owning `any_resource` holding the previous resource for `device_id`
@@ -161,9 +168,8 @@ inline cuda::mr::any_resource<cuda::mr::device_accessible> set_per_device_resour
  * `device_id.value()` must be in the range `[0, cudaGetDeviceCount())`, otherwise behavior is
  * undefined.
  *
- * The object referenced by `new_resource_ref` must outlive the last use of the resource, otherwise
- * behavior is undefined. It is the caller's responsibility to maintain the lifetime of the resource
- * object.
+ * The referenced resource is copied into an owning `any_resource` and moved into the per-device
+ * resource map.
  *
  * This function is thread-safe with respect to concurrent calls to `set_per_device_resource`,
  * `set_per_device_resource_ref`, `get_per_device_resource_ref`,
@@ -176,6 +182,10 @@ inline cuda::mr::any_resource<cuda::mr::device_accessible> set_per_device_resour
  * `device_async_resource_ref` is undefined if used while the active CUDA device is a different
  * device from the one that was active when the memory resource was created.
  *
+ * @note The per-device resource map keeps the underlying resource alive until process exit. Its
+ * destructor may therefore run during process termination. If it may call CUDA APIs, it must
+ * consult `rmm::process_is_exiting()` and skip those calls when it returns `true`.
+ *
  * @param device_id The id of the target device
  * @param new_resource_ref new `device_async_resource_ref` to use as new resource for `device_id`
  * @return An owning `any_resource` holding the previous resource for `device_id`
@@ -231,6 +241,10 @@ inline device_async_resource_ref get_current_device_resource_ref()
  * The behavior of a memory resource is undefined if used while the active CUDA device is a
  * different device from the one that was active when the memory resource was created.
  *
+ * @note The per-device resource map keeps the provided resource alive until process exit. Its
+ * destructor may therefore run during process termination. If the destructor may call CUDA APIs,
+ * it must consult `rmm::process_is_exiting()` and skip those calls when it returns `true`.
+ *
  * @param new_resource New resource to use for the current device
  * @return An owning `any_resource` holding the previous resource for the current device
  */
@@ -247,9 +261,8 @@ inline cuda::mr::any_resource<cuda::mr::device_accessible> set_current_device_re
  *
  * The "current device" is the device returned by `cudaGetDevice`.
  *
- * The object referenced by `new_resource_ref` must outlive the last use of the resource, otherwise
- * behavior is undefined. It is the caller's responsibility to maintain the lifetime of the resource
- * object.
+ * The referenced resource is copied into an owning `any_resource` and moved into the per-device
+ * resource map.
  *
  * This function is thread-safe with respect to concurrent calls to `set_per_device_resource`,
  * `set_per_device_resource_ref`, `get_per_device_resource_ref`,
@@ -261,6 +274,10 @@ inline cuda::mr::any_resource<cuda::mr::device_accessible> set_current_device_re
  * device. The behavior of a `device_async_resource_ref` is undefined if used while the active CUDA
  * device is a different device from the one that was active when the memory resource was created.
  *
+ * @note The per-device resource map keeps the underlying resource alive until process exit. Its
+ * destructor may therefore run during process termination. If it may call CUDA APIs, it must
+ * consult `rmm::process_is_exiting()` and skip those calls when it returns `true`.
+ *
  * @param new_resource_ref New `device_async_resource_ref` to use for the current device
  * @return An owning `any_resource` holding the previous resource for the current device
  */

cpp/include/rmm/process_is_exiting.hpp

@@ -0,0 +1,64 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <rmm/detail/export.hpp>
+
+namespace RMM_EXPORT rmm {
+
+/**
+ * @addtogroup memory_resources
+ * @{
+ * @file
+ */
+
+/**
+ * @brief Returns `true` if the process has entered `exit()` / atexit handler execution.
+ *
+ * Destructors of static objects, as well as atexit handlers registered by other DSOs, run
+ * during process termination after `main()` has returned. At that point calling into the CUDA
+ * runtime or driver is undefined behavior: the primary context may already be destroyed, and
+ * CUDA API calls may dereference released state and crash inside libcuda rather than returning
+ * an error.
+ *
+ * Use this function from a memory resource destructor (or a helper invoked by a destructor, such
+ * as a `release()` method) when the resource may be held in RMM's internal per-device resource
+ * map and destroyed during process termination. In that case the destructor may run after the
+ * CUDA primary context has been destroyed, and calling into the CUDA runtime is undefined
+ * behavior. Destructors can avoid that by:
+ *
+ * 1. Never calling CUDA APIs from the destructor at all, or
+ * 2. Consulting `rmm::process_is_exiting()` in the destructor (and in any helper invoked by
+ *    the destructor, such as a `release()` method) and skipping CUDA API calls when it
+ *    returns `true`. In that case, resources that would have been explicitly released should be
+ *    leaked; the OS reclaims them when the process exits.
+ *
+ * Storing RMM objects with static or thread-local scope is unsupported. Users should not create
+ * their own static containers of RMM objects and rely on `rmm::process_is_exiting()` to make
+ * those destructors safe.
+ *
+ * Calling `rmm::process_is_exiting()` from a resource destructor is always safe: it performs a
+ * single atomic load (acquire semantics) and never calls into CUDA.
+ *
+ * Example:
+ * @code{.cpp}
+ * class my_resource final : public ... {
+ *   ~my_resource() override
+ *   {
+ *     if (rmm::process_is_exiting()) {
+ *       return;
+ *     }
+ *     RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaFree(ptr_));
+ *   }
+ * };
+ * @endcode
+ *
+ * @return `true` if `exit()` has begun; `false` otherwise.
+ */
+bool process_is_exiting() noexcept;
+
+/** @} */  // end of group
+
+}  // namespace RMM_EXPORT rmm

cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp

@@ -9,6 +9,7 @@
 #include <rmm/detail/error.hpp>
 #include <rmm/detail/runtime_capabilities.hpp>
 #include <rmm/mr/detail/cuda_async_memory_resource_impl.hpp>
+#include <rmm/process_is_exiting.hpp>
 
 #include <cuda_runtime_api.h>
 
@@ -66,6 +67,8 @@ cuda_async_memory_resource_impl::cuda_async_memory_resource_impl(
 
 cuda_async_memory_resource_impl::~cuda_async_memory_resource_impl()
 {
+  if (rmm::process_is_exiting()) { return; }
+
   RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaMemPoolDestroy(pool_handle()));
 }
 

cpp/src/mr/detail/fixed_size_memory_resource_impl.cpp

@@ -7,6 +7,7 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/detail/logging_assert.hpp>
 #include <rmm/mr/detail/fixed_size_memory_resource_impl.hpp>
+#include <rmm/process_is_exiting.hpp>
 
 #include <cuda/iterator>
 
@@ -93,6 +94,8 @@ void fixed_size_memory_resource_impl::release()
 {
   lock_guard lock(this->get_mutex());
 
+  if (rmm::process_is_exiting()) { return; }
+
   for (auto block : upstream_blocks_) {
     upstream_mr_.deallocate_sync(
       block.pointer(), upstream_chunk_size_, rmm::CUDA_ALLOCATION_ALIGNMENT);

cpp/src/mr/detail/pool_memory_resource_impl.cpp

@@ -10,6 +10,7 @@
 #include <rmm/detail/logging_assert.hpp>
 #include <rmm/logger.hpp>
 #include <rmm/mr/detail/pool_memory_resource_impl.hpp>
+#include <rmm/process_is_exiting.hpp>
 
 #include <algorithm>
 #include <cstddef>
@@ -170,6 +171,8 @@ void pool_memory_resource_impl::release()
 {
   lock_guard lock(this->get_mutex());
 
+  if (rmm::process_is_exiting()) { return; }
+
   for (auto block : upstream_blocks_) {
     get_upstream_resource().deallocate_sync(
       block.pointer(), block.size(), rmm::CUDA_ALLOCATION_ALIGNMENT);

cpp/src/runtime_shutdown.cpp

@@ -0,0 +1,42 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <rmm/detail/error.hpp>
+#include <rmm/detail/runtime_shutdown.hpp>
+#include <rmm/process_is_exiting.hpp>
+
+#include <atomic>
+#include <cstdlib>
+#include <mutex>
+
+namespace RMM_NAMESPACE {
+
+namespace {
+
+std::atomic<bool> exiting{false};
+std::once_flag registered;
+
+}  // namespace
+
+bool process_is_exiting() noexcept { return exiting.load(std::memory_order_acquire); }
+
+namespace detail {
+
+void register_process_exit_hook() noexcept
+{
+  // The C++ standard guarantees that if the completion of the initialization of a static object A
+  // is sequenced-before a call to std::atexit(F), then F runs before A's destructor at
+  // termination. RMM exploits this by registering this flag-setter immediately after the internal
+  // per-device resource map is constructed, so the flag is observed as true during that map's
+  // destructor.
+  std::call_once(registered, []() {
+    auto const registration_status =
+      std::atexit([]() noexcept { exiting.store(true, std::memory_order_release); });
+    RMM_EXPECTS(registration_status == 0, "Unable to register process-exit hook");
+  });
+}
+
+}  // namespace detail
+}  // namespace RMM_NAMESPACE

cpp/tests/CMakeLists.txt

@@ -243,6 +243,29 @@ ConfigureTest(CONTAINER_MULTIDEVICE_TEST container_multidevice_tests.cu)
 # error macros tests
 ConfigureTest(ERROR_MACROS_TEST error_macros_tests.cpp)
 
+# process_is_exiting API tests (CPU-only, no CUDA calls).
+ConfigureTest(PROCESS_IS_EXITING_TEST process_is_exiting_tests.cpp)
+
+# Standalone shutdown test: a binary (no gtest) that verifies the flag is observed as true from a C
+# atexit handler registered before register_process_exit_hook(), exercising the C++ std::atexit
+# sequencing guarantee the fix relies on. CTest checks the exit code.
+add_executable(PROCESS_IS_EXITING_SHUTDOWN_TEST process_is_exiting_shutdown_test.cpp)
+target_include_directories(PROCESS_IS_EXITING_SHUTDOWN_TEST
+                           PRIVATE "$<BUILD_INTERFACE:${RMM_SOURCE_DIR}>")
+target_link_libraries(PROCESS_IS_EXITING_SHUTDOWN_TEST PRIVATE rmm
+                                                               $<TARGET_NAME_IF_EXISTS:conda_env>)
+set_target_properties(
+  PROCESS_IS_EXITING_SHUTDOWN_TEST
+  PROPERTIES POSITION_INDEPENDENT_CODE ON
+             RUNTIME_OUTPUT_DIRECTORY "$<BUILD_INTERFACE:${RMM_BINARY_DIR}/gtests>"
+             INSTALL_RPATH "\$ORIGIN/../../../lib"
+             CXX_STANDARD 20
+             CXX_STANDARD_REQUIRED ON)
+rapids_test_add(
+  NAME PROCESS_IS_EXITING_SHUTDOWN_TEST
+  COMMAND PROCESS_IS_EXITING_SHUTDOWN_TEST
+  INSTALL_COMPONENT_SET testing)
+
 # Resource ref construction and conversion tests
 ConfigureTest(RESOURCE_REF_CONVERSION_TEST mr/resource_ref_conversion_tests.cpp)
 

cpp/tests/process_is_exiting_shutdown_test.cpp

@@ -0,0 +1,51 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <rmm/mr/per_device_resource.hpp>
+#include <rmm/process_is_exiting.hpp>
+
+#include <cstdlib>
+
+/*
+ * Standalone test that verifies rmm::process_is_exiting() reports true while the C runtime is
+ * executing atexit handlers. Resource destructors running during exit() must observe the flag
+ * as set so that they can skip CUDA calls.
+ *
+ * This binary exercises the real integration point without any CUDA involvement:
+ *
+ *   1. main() registers `checker` via std::atexit FIRST.
+ *   2. main() then calls detail::get_ref_map(). This constructs the static per-device map and,
+ *      immediately afterward, calls register_process_exit_hook(), which registers the internal
+ *      flag-setter via std::atexit SECOND.
+ *   3. main() returns 0.
+ *   4. At termination, atexit handlers run in LIFO order: first the flag-setter (sets the
+ *      flag to true), then `checker` (reads the flag).
+ *   5. If `checker` observes the flag as true, the process exits with status 0; otherwise it
+ *      calls _Exit(1) to signal failure.
+ *
+ * CTest checks the exit status.
+ */
+
+namespace {
+
+void checker() noexcept
+{
+  if (!rmm::process_is_exiting()) {
+    // The process is already running atexit handlers. Use std::_Exit rather than std::exit to
+    // report failure without re-entering normal termination or running more cleanup.
+    std::_Exit(1);
+  }
+}
+
+}  // namespace
+
+int main()
+{
+  // Register the checker before get_ref_map() so that the checker runs AFTER the flag-setter at
+  // termination.
+  if (std::atexit(checker) != 0) { return 2; }
+  [[maybe_unused]] auto& map = rmm::mr::detail::get_ref_map();
+  return 0;
+}