From 8660942ac8003579261fb682a50fbd87b18612c7 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Wed, 22 Apr 2026 21:32:57 -0500 Subject: [PATCH 01/10] Fix pool_memory_resource crash during process exit pool_memory_resource's destructor unconditionally calls cudaEventSynchronize and cudaEventDestroy from stream_ordered_memory_resource::release(). When an owning pool_memory_resource is held in the static per-device resource map returned by detail::get_ref_map(), the destructor runs at exit() time after the CUDA primary context has been destroyed. At that point CUDA calls dereference released context state inside libcuda and crash, producing intermittent SIGSEGV at process exit. Introduce rmm::process_is_exiting(), a public API that resource destructors can consult to detect this condition. It is backed by an atexit callback registered the first time the per-device resource map is constructed; atexit/__cxa_atexit LIFO ordering ensures the callback runs before the map's destructor during teardown. stream_ordered_memory_resource::release() and pool_memory_resource_impl::release() both skip their CUDA calls when process_is_exiting() returns true, allowing associated state to leak; the OS reclaims it when the process exits. Document the memory-resource author contract in the new public header and in docs/cpp/memory_resources/memory_resources.md: all memory resources must be safe to destroy at shutdown, either by not calling CUDA APIs from destructors or by consulting rmm::process_is_exiting() and skipping CUDA calls when it returns true. --- cpp/CMakeLists.txt | 3 +- cpp/include/rmm/detail/runtime_shutdown.hpp | 24 ++++++++ .../detail/stream_ordered_memory_resource.hpp | 14 ++++- cpp/include/rmm/mr/per_device_resource.hpp | 6 ++ cpp/include/rmm/process_is_exiting.hpp | 57 +++++++++++++++++++ cpp/src/cuda_stream.cpp | 8 ++- .../cuda_async_memory_resource_impl.cpp | 8 ++- .../fixed_size_memory_resource_impl.cpp | 12 +++- .../mr/detail/pool_memory_resource_impl.cpp | 12 +++- cpp/src/runtime_shutdown.cpp | 43 ++++++++++++++ docs/cpp/memory_resources/memory_resources.md | 45 +++++++++++++++ 11 files changed, 220 insertions(+), 12 deletions(-) create mode 100644 cpp/include/rmm/detail/runtime_shutdown.hpp create mode 100644 cpp/include/rmm/process_is_exiting.hpp create mode 100644 cpp/src/runtime_shutdown.cpp diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 3cd8b2303..a032ba515 100644 --- a/cpp/CMakeLists.txt +++ b/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( diff --git a/cpp/include/rmm/detail/runtime_shutdown.hpp b/cpp/include/rmm/detail/runtime_shutdown.hpp new file mode 100644 index 000000000..6a438a19c --- /dev/null +++ b/cpp/include/rmm/detail/runtime_shutdown.hpp @@ -0,0 +1,24 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +namespace RMM_NAMESPACE { +namespace detail { + +/** + * @brief Register the atexit callback that flips the flag observed by `rmm::process_is_exiting()`. + * + * Idempotent: safe to call from multiple places; only the first call registers the callback. + * Callers must invoke this after constructing any static object whose destructor may run during + * process exit and needs to consult `rmm::process_is_exiting()`. The C++ standard guarantees + * that if a static object's construction is sequenced-before a `std::atexit` call, the atexit + * callback runs before that object's destructor at termination. + */ +RMM_EXPORT void register_process_exit_hook() noexcept; + +} // namespace detail +} // namespace RMM_NAMESPACE diff --git a/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp b/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp index e8de8b65e..72c2ca005 100644 --- a/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp +++ b/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -486,9 +487,16 @@ class stream_ordered_memory_resource : public crtp { { lock_guard lock(mtx_); - 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)); + // This destructor can run from a static destructor after the CUDA primary context has been + // destroyed (e.g., when an owning MR is held in the static per-device resource map). + // Calling into the CUDA runtime after exit() has begun is undefined behavior: the driver + // may dereference destroyed context state and crash inside libcuda rather than returning + // an error. In that case, leak the events; the OS reclaims them when the process exits. + if (!rmm::process_is_exiting()) { + 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)); + } } stream_events_.clear(); diff --git a/cpp/include/rmm/mr/per_device_resource.hpp b/cpp/include/rmm/mr/per_device_resource.hpp index df4451be1..0dfda6d1c 100644 --- a/cpp/include/rmm/mr/per_device_resource.hpp +++ b/cpp/include/rmm/mr/per_device_resource.hpp @@ -7,6 +7,7 @@ #include #include +#include #include #include @@ -76,6 +77,11 @@ RMM_EXPORT inline auto& get_ref_map() { static std::map> device_id_to_resource; + // The map's destructor may run owning resources' destructors during process exit, when calling + // into the CUDA runtime is undefined behavior. Register the process-exit hook here, right after + // the map is constructed: the C++ standard guarantees that an atexit callback registered after + // a static object's construction runs before that object's destructor at termination. + rmm::detail::register_process_exit_hook(); return device_id_to_resource; } diff --git a/cpp/include/rmm/process_is_exiting.hpp b/cpp/include/rmm/process_is_exiting.hpp new file mode 100644 index 000000000..c84d0cd0f --- /dev/null +++ b/cpp/include/rmm/process_is_exiting.hpp @@ -0,0 +1,57 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +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, can run + * 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. + * + * @par Memory resource author contract + * All RMM memory resources must be safe to destroy at process shutdown. An MR destructor may + * run during normal program flow (when calling CUDA APIs is safe) or after `exit()` has been + * called (when it is not). Authors must satisfy this by either: + * + * 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 + * simply be leaked; the OS reclaims them when the process exits. + * + * Calling `rmm::process_is_exiting()` from a resource destructor is always safe: it performs a + * single relaxed atomic load and never calls into CUDA. + * + * @par Example + * @code + * class my_resource final : public ... { + * ~my_resource() override + * { + * if (!rmm::process_is_exiting()) { + * 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 diff --git a/cpp/src/cuda_stream.cpp b/cpp/src/cuda_stream.cpp index a0ff49872..f2122b4cf 100644 --- a/cpp/src/cuda_stream.cpp +++ b/cpp/src/cuda_stream.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include @@ -22,7 +23,12 @@ cuda_stream::cuda_stream(cuda_stream::flags flags) return stream; }(), [](cudaStream_t* stream) { - RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamDestroy(*stream)); + // During process exit, calling into CUDA is undefined behavior once the primary + // context has been destroyed. Leak the stream in that case; the OS reclaims it + // when the process exits. See rmm/process_is_exiting.hpp. + if (!rmm::process_is_exiting()) { + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamDestroy(*stream)); + } delete stream; // NOLINT(cppcoreguidelines-owning-memory) }} { diff --git a/cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp b/cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp index 48f5dad03..5f3d1b89f 100644 --- a/cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include @@ -66,7 +67,12 @@ cuda_async_memory_resource_impl::cuda_async_memory_resource_impl( cuda_async_memory_resource_impl::~cuda_async_memory_resource_impl() { - RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaMemPoolDestroy(pool_handle())); + // During process exit, calling into CUDA is undefined behavior once the primary context has + // been destroyed. Leak the memory pool in that case; the OS reclaims it when the process + // exits. See rmm/process_is_exiting.hpp. + if (!rmm::process_is_exiting()) { + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaMemPoolDestroy(pool_handle())); + } } cudaMemPool_t cuda_async_memory_resource_impl::pool_handle() const noexcept diff --git a/cpp/src/mr/detail/fixed_size_memory_resource_impl.cpp b/cpp/src/mr/detail/fixed_size_memory_resource_impl.cpp index e17f6ecdd..56b116831 100644 --- a/cpp/src/mr/detail/fixed_size_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/fixed_size_memory_resource_impl.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include @@ -93,9 +94,14 @@ void fixed_size_memory_resource_impl::release() { lock_guard lock(this->get_mutex()); - for (auto block : upstream_blocks_) { - upstream_mr_.deallocate_sync( - block.pointer(), upstream_chunk_size_, rmm::CUDA_ALLOCATION_ALIGNMENT); + // During process exit, calling into CUDA (even via the upstream resource) is undefined + // behavior once the primary context has been destroyed. Leak upstream blocks in that case; + // the OS reclaims them when the process exits. See rmm/process_is_exiting.hpp. + if (!rmm::process_is_exiting()) { + for (auto block : upstream_blocks_) { + upstream_mr_.deallocate_sync( + block.pointer(), upstream_chunk_size_, rmm::CUDA_ALLOCATION_ALIGNMENT); + } } upstream_blocks_.clear(); } diff --git a/cpp/src/mr/detail/pool_memory_resource_impl.cpp b/cpp/src/mr/detail/pool_memory_resource_impl.cpp index f00c7fcb8..c2212d51f 100644 --- a/cpp/src/mr/detail/pool_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/pool_memory_resource_impl.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -170,9 +171,14 @@ void pool_memory_resource_impl::release() { lock_guard lock(this->get_mutex()); - for (auto block : upstream_blocks_) { - get_upstream_resource().deallocate_sync( - block.pointer(), block.size(), rmm::CUDA_ALLOCATION_ALIGNMENT); + // During process exit, calling into CUDA (even via the upstream resource) is undefined + // behavior once the primary context has been destroyed. Leak upstream blocks in that case; + // the OS reclaims them when the process exits. See rmm/process_is_exiting.hpp. + if (!rmm::process_is_exiting()) { + for (auto block : upstream_blocks_) { + get_upstream_resource().deallocate_sync( + block.pointer(), block.size(), rmm::CUDA_ALLOCATION_ALIGNMENT); + } } upstream_blocks_.clear(); #ifdef RMM_POOL_TRACK_ALLOCATIONS diff --git a/cpp/src/runtime_shutdown.cpp b/cpp/src/runtime_shutdown.cpp new file mode 100644 index 000000000..6fc7361dd --- /dev/null +++ b/cpp/src/runtime_shutdown.cpp @@ -0,0 +1,43 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include +#include + +namespace RMM_NAMESPACE { + +namespace { + +std::atomic& exiting_flag() noexcept +{ + static std::atomic flag{false}; + return flag; +} + +} // namespace + +bool process_is_exiting() noexcept { return exiting_flag().load(std::memory_order_acquire); } + +namespace detail { + +void register_process_exit_hook() noexcept +{ + // The C++ standard guarantees that if a static object's construction is sequenced-before a + // call to std::atexit, the atexit callback runs before that object's destructor at + // termination (see https://en.cppreference.com/cpp/utility/program/exit). Callers must + // therefore invoke this function after constructing the static object whose destructor needs + // to observe the exit flag. + static std::atomic registered{false}; + bool expected = false; + if (registered.compare_exchange_strong(expected, true, std::memory_order_acq_rel)) { + std::atexit([]() noexcept { exiting_flag().store(true, std::memory_order_release); }); + } +} + +} // namespace detail +} // namespace RMM_NAMESPACE diff --git a/docs/cpp/memory_resources/memory_resources.md b/docs/cpp/memory_resources/memory_resources.md index 73f7197ca..4e4f7086f 100644 --- a/docs/cpp/memory_resources/memory_resources.md +++ b/docs/cpp/memory_resources/memory_resources.md @@ -1,5 +1,50 @@ # Memory Resources +## Authoring custom memory resources + +All memory resources — both those provided by RMM and those written by library authors against +RMM's resource concepts — must be safe to destroy at process shutdown. + +A memory resource destructor may run in two very different contexts: + +1. **During normal program execution**, where calling into the CUDA runtime is safe. +2. **After `main()` has returned**, as part of `exit()` / atexit handler processing. This + happens whenever a memory resource is owned by a static or `thread_local` object, including + the per-device resource map that RMM uses internally and any static container maintained by + a consuming library. In this context, calling into the CUDA runtime or driver is + **undefined behavior**: the primary context may already have been destroyed, and CUDA API + calls may dereference released state and crash inside `libcuda` rather than returning an + error. + +To satisfy this contract, authors of memory resources (and of any destructor reachable from an +MR destructor, such as a `release()` helper) must follow one of these rules: + +- **Do not call CUDA APIs from destructors at all.** This is the simplest option and should be + preferred when feasible. + +- **Consult {cpp:func}`rmm::process_is_exiting()`** in any destructor that would otherwise call + a CUDA API. When `rmm::process_is_exiting()` returns `true`, skip the CUDA calls and allow + the associated resources to leak; the operating system reclaims process memory, file + descriptors, and driver-owned state when the process exits. + +Calling `rmm::process_is_exiting()` is always safe from a destructor: it performs a single +atomic load and never calls into CUDA. + +### Example + +```cpp +class my_resource final : public rmm::mr::device_memory_resource { + public: + ~my_resource() override + { + if (!rmm::process_is_exiting()) { + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaFree(ptr_)); + } + } + ... +}; +``` + ```{doxygengroup} memory_resources :members: :undoc-members: From 03d48050e667e59bb4a6fbee965b26170e0b15b9 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Wed, 22 Apr 2026 22:37:47 -0500 Subject: [PATCH 02/10] Move shutdown-safety contract onto per-device resource setters In response to PR #2367 review: the requirement that resources be safe to destroy at process shutdown only applies to resources that are going to be held in the per-device resource map, so the contract belongs on the setter functions. Add a @note to each of set_per_device_resource, set_per_device_resource_ref, set_current_device_resource, and set_current_device_resource_ref stating the requirement and pointing at rmm::process_is_exiting(). Remove the duplicate narrative from docs/cpp/memory_resources/memory_resources.md; the page now only renders the doxygen memory_resources group, which includes the updated setters. Also align the rmm::process_is_exiting() docstring wording with the implementation's acquire memory ordering (per coderabbit review comment). --- cpp/include/rmm/mr/per_device_resource.hpp | 26 +++++++++++ cpp/include/rmm/process_is_exiting.hpp | 2 +- docs/cpp/memory_resources/memory_resources.md | 45 ------------------- 3 files changed, 27 insertions(+), 46 deletions(-) diff --git a/cpp/include/rmm/mr/per_device_resource.hpp b/cpp/include/rmm/mr/per_device_resource.hpp index 0dfda6d1c..fffe91002 100644 --- a/cpp/include/rmm/mr/per_device_resource.hpp +++ b/cpp/include/rmm/mr/per_device_resource.hpp @@ -142,6 +142,12 @@ 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 resource passed in `new_resource` must be safe to destroy at process shutdown. The + * per-device resource map keeps the resource alive until process exit, so its destructor may run + * after the CUDA primary context has been destroyed -- at that point calling into the CUDA + * runtime is undefined behavior. Resource destructors must either not call CUDA APIs at all, or + * consult `rmm::process_is_exiting()` and skip CUDA 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` @@ -182,6 +188,13 @@ inline cuda::mr::any_resource 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 underlying resource referenced by `new_resource_ref` must be safe to destroy at + * process shutdown. The per-device resource map keeps a reference alive until process exit, so + * the underlying resource's destructor may run after the CUDA primary context has been + * destroyed -- at that point calling into the CUDA runtime is undefined behavior. Resource + * destructors must either not call CUDA APIs at all, or consult `rmm::process_is_exiting()` and + * skip CUDA 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` @@ -237,6 +250,12 @@ 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 resource passed in `new_resource` must be safe to destroy at process shutdown. The + * per-device resource map keeps the resource alive until process exit, so its destructor may run + * after the CUDA primary context has been destroyed -- at that point calling into the CUDA + * runtime is undefined behavior. Resource destructors must either not call CUDA APIs at all, or + * consult `rmm::process_is_exiting()` and skip CUDA 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 */ @@ -267,6 +286,13 @@ inline cuda::mr::any_resource 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 underlying resource referenced by `new_resource_ref` must be safe to destroy at + * process shutdown. The per-device resource map keeps a reference alive until process exit, so + * the underlying resource's destructor may run after the CUDA primary context has been + * destroyed -- at that point calling into the CUDA runtime is undefined behavior. Resource + * destructors must either not call CUDA APIs at all, or consult `rmm::process_is_exiting()` and + * skip CUDA 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 */ diff --git a/cpp/include/rmm/process_is_exiting.hpp b/cpp/include/rmm/process_is_exiting.hpp index c84d0cd0f..53f1f4af6 100644 --- a/cpp/include/rmm/process_is_exiting.hpp +++ b/cpp/include/rmm/process_is_exiting.hpp @@ -34,7 +34,7 @@ namespace RMM_EXPORT rmm { * simply be leaked; the OS reclaims them when the process exits. * * Calling `rmm::process_is_exiting()` from a resource destructor is always safe: it performs a - * single relaxed atomic load and never calls into CUDA. + * single atomic load (acquire semantics) and never calls into CUDA. * * @par Example * @code diff --git a/docs/cpp/memory_resources/memory_resources.md b/docs/cpp/memory_resources/memory_resources.md index 4e4f7086f..73f7197ca 100644 --- a/docs/cpp/memory_resources/memory_resources.md +++ b/docs/cpp/memory_resources/memory_resources.md @@ -1,50 +1,5 @@ # Memory Resources -## Authoring custom memory resources - -All memory resources — both those provided by RMM and those written by library authors against -RMM's resource concepts — must be safe to destroy at process shutdown. - -A memory resource destructor may run in two very different contexts: - -1. **During normal program execution**, where calling into the CUDA runtime is safe. -2. **After `main()` has returned**, as part of `exit()` / atexit handler processing. This - happens whenever a memory resource is owned by a static or `thread_local` object, including - the per-device resource map that RMM uses internally and any static container maintained by - a consuming library. In this context, calling into the CUDA runtime or driver is - **undefined behavior**: the primary context may already have been destroyed, and CUDA API - calls may dereference released state and crash inside `libcuda` rather than returning an - error. - -To satisfy this contract, authors of memory resources (and of any destructor reachable from an -MR destructor, such as a `release()` helper) must follow one of these rules: - -- **Do not call CUDA APIs from destructors at all.** This is the simplest option and should be - preferred when feasible. - -- **Consult {cpp:func}`rmm::process_is_exiting()`** in any destructor that would otherwise call - a CUDA API. When `rmm::process_is_exiting()` returns `true`, skip the CUDA calls and allow - the associated resources to leak; the operating system reclaims process memory, file - descriptors, and driver-owned state when the process exits. - -Calling `rmm::process_is_exiting()` is always safe from a destructor: it performs a single -atomic load and never calls into CUDA. - -### Example - -```cpp -class my_resource final : public rmm::mr::device_memory_resource { - public: - ~my_resource() override - { - if (!rmm::process_is_exiting()) { - RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaFree(ptr_)); - } - } - ... -}; -``` - ```{doxygengroup} memory_resources :members: :undoc-members: From 68ebf5490bda3d3c0cf0033517772c0f6a6303ac Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Wed, 22 Apr 2026 23:13:23 -0500 Subject: [PATCH 03/10] Match rmm doxygen style --- cpp/include/rmm/process_is_exiting.hpp | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/cpp/include/rmm/process_is_exiting.hpp b/cpp/include/rmm/process_is_exiting.hpp index 53f1f4af6..5d6f640ff 100644 --- a/cpp/include/rmm/process_is_exiting.hpp +++ b/cpp/include/rmm/process_is_exiting.hpp @@ -22,7 +22,6 @@ namespace RMM_EXPORT rmm { * 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. * - * @par Memory resource author contract * All RMM memory resources must be safe to destroy at process shutdown. An MR destructor may * run during normal program flow (when calling CUDA APIs is safe) or after `exit()` has been * called (when it is not). Authors must satisfy this by either: @@ -36,8 +35,8 @@ namespace RMM_EXPORT rmm { * 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. * - * @par Example - * @code + * Example: + * @code{.cpp} * class my_resource final : public ... { * ~my_resource() override * { From 8bd42cc5578c718db34d50ede7bcd4138280c6d1 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Wed, 22 Apr 2026 23:27:26 -0500 Subject: [PATCH 04/10] Add tests for process_is_exiting() Two binaries: - PROCESS_IS_EXITING_TEST (gtest): verifies process_is_exiting() returns false during normal execution, is idempotent, and that register_process_exit_hook() can be called repeatedly without effect. - PROCESS_IS_EXITING_SHUTDOWN_TEST (standalone, no gtest): exercises the C++ std::atexit sequencing guarantee the fix depends on. main() registers a C atexit checker first and then calls register_process_exit_hook(), so at termination the flag-setter runs before the checker. The checker calls _Exit(1) if the flag is not set; CTest verifies exit status 0. Neither test touches CUDA. --- cpp/tests/CMakeLists.txt | 23 ++++++++ .../process_is_exiting_shutdown_test.cpp | 53 +++++++++++++++++++ cpp/tests/process_is_exiting_tests.cpp | 40 ++++++++++++++ 3 files changed, 116 insertions(+) create mode 100644 cpp/tests/process_is_exiting_shutdown_test.cpp create mode 100644 cpp/tests/process_is_exiting_tests.cpp diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index d56a5a258..baf0c3873 100644 --- a/cpp/tests/CMakeLists.txt +++ b/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 "$") +target_link_libraries(PROCESS_IS_EXITING_SHUTDOWN_TEST PRIVATE rmm + $) +set_target_properties( + PROCESS_IS_EXITING_SHUTDOWN_TEST + PROPERTIES POSITION_INDEPENDENT_CODE ON + RUNTIME_OUTPUT_DIRECTORY "$" + 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) diff --git a/cpp/tests/process_is_exiting_shutdown_test.cpp b/cpp/tests/process_is_exiting_shutdown_test.cpp new file mode 100644 index 000000000..b05475f85 --- /dev/null +++ b/cpp/tests/process_is_exiting_shutdown_test.cpp @@ -0,0 +1,53 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +/* + * Standalone test that verifies rmm::process_is_exiting() reports true while the C runtime is + * executing atexit handlers. The fix in rapidsai/rmm#2367 depends on this behavior: resource + * destructors running during exit() must observe the flag as set so that they can skip CUDA + * calls. + * + * 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 having register_process_exit_hook() register a flag-setter + * via std::atexit on first use; it is called from get_ref_map() immediately after the static + * map is constructed, so the flag is observed as true during that map's destructor. + * + * This binary exercises the same mechanism without any per-device map or CUDA involvement: + * + * 1. main() registers `checker` via std::atexit FIRST. + * 2. main() then calls register_process_exit_hook(), which registers its 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()) { std::_Exit(1); } +} + +} // namespace + +int main() +{ + // Register the checker before rmm's flag-setter so that the checker runs AFTER the + // flag-setter at termination. + if (std::atexit(checker) != 0) { return 2; } + rmm::detail::register_process_exit_hook(); + return 0; +} diff --git a/cpp/tests/process_is_exiting_tests.cpp b/cpp/tests/process_is_exiting_tests.cpp new file mode 100644 index 000000000..afe501f20 --- /dev/null +++ b/cpp/tests/process_is_exiting_tests.cpp @@ -0,0 +1,40 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include + +/* + * Tests for rmm::process_is_exiting() and the internal register_process_exit_hook(). + * + * These tests cover the observable behavior of the API during normal program execution. The + * complementary property -- that the flag flips to true before the destructor of a static + * object whose construction was sequenced-before register_process_exit_hook() -- is exercised + * by the standalone PROCESS_IS_EXITING_SHUTDOWN_TEST binary, whose exit code is checked by + * CTest. + */ + +// Flag should be false while the test process is running normally. +TEST(ProcessIsExitingTest, FalseDuringNormalExecution) { EXPECT_FALSE(rmm::process_is_exiting()); } + +// Calling the query is safe to call repeatedly and does not modify state. +TEST(ProcessIsExitingTest, QueryIsIdempotent) +{ + EXPECT_FALSE(rmm::process_is_exiting()); + EXPECT_FALSE(rmm::process_is_exiting()); + EXPECT_FALSE(rmm::process_is_exiting()); +} + +// register_process_exit_hook is idempotent: repeated calls are safe and do not flip the flag +// before the actual exit. +TEST(ProcessIsExitingTest, RegisterHookIsIdempotent) +{ + rmm::detail::register_process_exit_hook(); + rmm::detail::register_process_exit_hook(); + rmm::detail::register_process_exit_hook(); + EXPECT_FALSE(rmm::process_is_exiting()); +} From 561ce688256e8b5fc9e7a3f546d470d30ed6b21c Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Fri, 24 Apr 2026 11:15:21 -0500 Subject: [PATCH 05/10] Narrow shutdown guards and simplify cleanup paths --- .../mr/detail/stream_ordered_memory_resource.hpp | 15 +++++---------- cpp/src/cuda_stream.cpp | 8 +------- .../mr/detail/cuda_async_memory_resource_impl.cpp | 9 +++------ .../mr/detail/fixed_size_memory_resource_impl.cpp | 13 +++++-------- cpp/src/mr/detail/pool_memory_resource_impl.cpp | 13 +++++-------- 5 files changed, 19 insertions(+), 39 deletions(-) diff --git a/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp b/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp index 72c2ca005..17c51ab6b 100644 --- a/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp +++ b/cpp/include/rmm/mr/detail/stream_ordered_memory_resource.hpp @@ -487,16 +487,11 @@ class stream_ordered_memory_resource : public crtp { { lock_guard lock(mtx_); - // This destructor can run from a static destructor after the CUDA primary context has been - // destroyed (e.g., when an owning MR is held in the static per-device resource map). - // Calling into the CUDA runtime after exit() has begun is undefined behavior: the driver - // may dereference destroyed context state and crash inside libcuda rather than returning - // an error. In that case, leak the events; the OS reclaims them when the process exits. - if (!rmm::process_is_exiting()) { - 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)); - } + 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)); } stream_events_.clear(); diff --git a/cpp/src/cuda_stream.cpp b/cpp/src/cuda_stream.cpp index f2122b4cf..a0ff49872 100644 --- a/cpp/src/cuda_stream.cpp +++ b/cpp/src/cuda_stream.cpp @@ -6,7 +6,6 @@ #include #include #include -#include #include @@ -23,12 +22,7 @@ cuda_stream::cuda_stream(cuda_stream::flags flags) return stream; }(), [](cudaStream_t* stream) { - // During process exit, calling into CUDA is undefined behavior once the primary - // context has been destroyed. Leak the stream in that case; the OS reclaims it - // when the process exits. See rmm/process_is_exiting.hpp. - if (!rmm::process_is_exiting()) { - RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamDestroy(*stream)); - } + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaStreamDestroy(*stream)); delete stream; // NOLINT(cppcoreguidelines-owning-memory) }} { diff --git a/cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp b/cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp index 5f3d1b89f..c1ad542a5 100644 --- a/cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/cuda_async_memory_resource_impl.cpp @@ -67,12 +67,9 @@ cuda_async_memory_resource_impl::cuda_async_memory_resource_impl( cuda_async_memory_resource_impl::~cuda_async_memory_resource_impl() { - // During process exit, calling into CUDA is undefined behavior once the primary context has - // been destroyed. Leak the memory pool in that case; the OS reclaims it when the process - // exits. See rmm/process_is_exiting.hpp. - if (!rmm::process_is_exiting()) { - RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaMemPoolDestroy(pool_handle())); - } + if (rmm::process_is_exiting()) { return; } + + RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaMemPoolDestroy(pool_handle())); } cudaMemPool_t cuda_async_memory_resource_impl::pool_handle() const noexcept diff --git a/cpp/src/mr/detail/fixed_size_memory_resource_impl.cpp b/cpp/src/mr/detail/fixed_size_memory_resource_impl.cpp index 56b116831..4d432c8d5 100644 --- a/cpp/src/mr/detail/fixed_size_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/fixed_size_memory_resource_impl.cpp @@ -94,14 +94,11 @@ void fixed_size_memory_resource_impl::release() { lock_guard lock(this->get_mutex()); - // During process exit, calling into CUDA (even via the upstream resource) is undefined - // behavior once the primary context has been destroyed. Leak upstream blocks in that case; - // the OS reclaims them when the process exits. See rmm/process_is_exiting.hpp. - if (!rmm::process_is_exiting()) { - for (auto block : upstream_blocks_) { - upstream_mr_.deallocate_sync( - block.pointer(), upstream_chunk_size_, rmm::CUDA_ALLOCATION_ALIGNMENT); - } + if (rmm::process_is_exiting()) { return; } + + for (auto block : upstream_blocks_) { + upstream_mr_.deallocate_sync( + block.pointer(), upstream_chunk_size_, rmm::CUDA_ALLOCATION_ALIGNMENT); } upstream_blocks_.clear(); } diff --git a/cpp/src/mr/detail/pool_memory_resource_impl.cpp b/cpp/src/mr/detail/pool_memory_resource_impl.cpp index c2212d51f..4297ee056 100644 --- a/cpp/src/mr/detail/pool_memory_resource_impl.cpp +++ b/cpp/src/mr/detail/pool_memory_resource_impl.cpp @@ -171,14 +171,11 @@ void pool_memory_resource_impl::release() { lock_guard lock(this->get_mutex()); - // During process exit, calling into CUDA (even via the upstream resource) is undefined - // behavior once the primary context has been destroyed. Leak upstream blocks in that case; - // the OS reclaims them when the process exits. See rmm/process_is_exiting.hpp. - if (!rmm::process_is_exiting()) { - for (auto block : upstream_blocks_) { - get_upstream_resource().deallocate_sync( - block.pointer(), block.size(), rmm::CUDA_ALLOCATION_ALIGNMENT); - } + if (rmm::process_is_exiting()) { return; } + + for (auto block : upstream_blocks_) { + get_upstream_resource().deallocate_sync( + block.pointer(), block.size(), rmm::CUDA_ALLOCATION_ALIGNMENT); } upstream_blocks_.clear(); #ifdef RMM_POOL_TRACK_ALLOCATIONS From 573dd782864f87671088b25435aac130cb97f669 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Fri, 24 Apr 2026 12:30:27 -0500 Subject: [PATCH 06/10] Simplify runtime shutdown registration --- cpp/include/rmm/detail/runtime_shutdown.hpp | 4 ---- cpp/include/rmm/mr/per_device_resource.hpp | 7 +++---- cpp/src/runtime_shutdown.cpp | 23 +++++++-------------- 3 files changed, 10 insertions(+), 24 deletions(-) diff --git a/cpp/include/rmm/detail/runtime_shutdown.hpp b/cpp/include/rmm/detail/runtime_shutdown.hpp index 6a438a19c..97c28fa58 100644 --- a/cpp/include/rmm/detail/runtime_shutdown.hpp +++ b/cpp/include/rmm/detail/runtime_shutdown.hpp @@ -13,10 +13,6 @@ namespace detail { * @brief Register the atexit callback that flips the flag observed by `rmm::process_is_exiting()`. * * Idempotent: safe to call from multiple places; only the first call registers the callback. - * Callers must invoke this after constructing any static object whose destructor may run during - * process exit and needs to consult `rmm::process_is_exiting()`. The C++ standard guarantees - * that if a static object's construction is sequenced-before a `std::atexit` call, the atexit - * callback runs before that object's destructor at termination. */ RMM_EXPORT void register_process_exit_hook() noexcept; diff --git a/cpp/include/rmm/mr/per_device_resource.hpp b/cpp/include/rmm/mr/per_device_resource.hpp index fffe91002..67c8bd701 100644 --- a/cpp/include/rmm/mr/per_device_resource.hpp +++ b/cpp/include/rmm/mr/per_device_resource.hpp @@ -77,10 +77,9 @@ RMM_EXPORT inline auto& get_ref_map() { static std::map> device_id_to_resource; - // The map's destructor may run owning resources' destructors during process exit, when calling - // into the CUDA runtime is undefined behavior. Register the process-exit hook here, right after - // the map is constructed: the C++ standard guarantees that an atexit callback registered after - // a static object's construction runs before that object's destructor at termination. + // Register the process-exit hook immediately after constructing the map. The C++ standard + // guarantees that an atexit callback registered after a static object's construction runs + // before that object's destructor at termination. rmm::detail::register_process_exit_hook(); return device_id_to_resource; } diff --git a/cpp/src/runtime_shutdown.cpp b/cpp/src/runtime_shutdown.cpp index 6fc7361dd..bf4abb4df 100644 --- a/cpp/src/runtime_shutdown.cpp +++ b/cpp/src/runtime_shutdown.cpp @@ -8,35 +8,26 @@ #include #include +#include namespace RMM_NAMESPACE { namespace { -std::atomic& exiting_flag() noexcept -{ - static std::atomic flag{false}; - return flag; -} +std::atomic exiting{false}; +std::once_flag registered; } // namespace -bool process_is_exiting() noexcept { return exiting_flag().load(std::memory_order_acquire); } +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 a static object's construction is sequenced-before a - // call to std::atexit, the atexit callback runs before that object's destructor at - // termination (see https://en.cppreference.com/cpp/utility/program/exit). Callers must - // therefore invoke this function after constructing the static object whose destructor needs - // to observe the exit flag. - static std::atomic registered{false}; - bool expected = false; - if (registered.compare_exchange_strong(expected, true, std::memory_order_acq_rel)) { - std::atexit([]() noexcept { exiting_flag().store(true, std::memory_order_release); }); - } + std::call_once(registered, []() { + std::atexit([]() noexcept { exiting.store(true, std::memory_order_release); }); + }); } } // namespace detail From 6ecc14be918ce1570f2e9ac0e30beabebe0f8f81 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Fri, 24 Apr 2026 12:45:14 -0500 Subject: [PATCH 07/10] Narrow shutdown docs and strengthen tests --- cpp/include/rmm/mr/per_device_resource.hpp | 36 ++++++++----------- cpp/include/rmm/process_is_exiting.hpp | 11 +++--- .../process_is_exiting_shutdown_test.cpp | 13 +++---- cpp/tests/process_is_exiting_tests.cpp | 20 +++-------- 4 files changed, 31 insertions(+), 49 deletions(-) diff --git a/cpp/include/rmm/mr/per_device_resource.hpp b/cpp/include/rmm/mr/per_device_resource.hpp index 67c8bd701..72595a766 100644 --- a/cpp/include/rmm/mr/per_device_resource.hpp +++ b/cpp/include/rmm/mr/per_device_resource.hpp @@ -141,11 +141,9 @@ 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 resource passed in `new_resource` must be safe to destroy at process shutdown. The - * per-device resource map keeps the resource alive until process exit, so its destructor may run - * after the CUDA primary context has been destroyed -- at that point calling into the CUDA - * runtime is undefined behavior. Resource destructors must either not call CUDA APIs at all, or - * consult `rmm::process_is_exiting()` and skip CUDA calls when it returns `true`. + * @note The per-device resource map keeps `new_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` @@ -187,12 +185,10 @@ inline cuda::mr::any_resource 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 underlying resource referenced by `new_resource_ref` must be safe to destroy at - * process shutdown. The per-device resource map keeps a reference alive until process exit, so - * the underlying resource's destructor may run after the CUDA primary context has been - * destroyed -- at that point calling into the CUDA runtime is undefined behavior. Resource - * destructors must either not call CUDA APIs at all, or consult `rmm::process_is_exiting()` and - * skip CUDA calls when it returns `true`. + * @note The per-device resource map keeps a reference to `new_resource_ref` alive until process + * exit. The underlying resource's 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` @@ -249,11 +245,9 @@ 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 resource passed in `new_resource` must be safe to destroy at process shutdown. The - * per-device resource map keeps the resource alive until process exit, so its destructor may run - * after the CUDA primary context has been destroyed -- at that point calling into the CUDA - * runtime is undefined behavior. Resource destructors must either not call CUDA APIs at all, or - * consult `rmm::process_is_exiting()` and skip CUDA calls when it returns `true`. + * @note The per-device resource map keeps `new_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 @@ -285,12 +279,10 @@ inline cuda::mr::any_resource 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 underlying resource referenced by `new_resource_ref` must be safe to destroy at - * process shutdown. The per-device resource map keeps a reference alive until process exit, so - * the underlying resource's destructor may run after the CUDA primary context has been - * destroyed -- at that point calling into the CUDA runtime is undefined behavior. Resource - * destructors must either not call CUDA APIs at all, or consult `rmm::process_is_exiting()` and - * skip CUDA calls when it returns `true`. + * @note The per-device resource map keeps a reference to `new_resource_ref` alive until process + * exit. The underlying resource's 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 diff --git a/cpp/include/rmm/process_is_exiting.hpp b/cpp/include/rmm/process_is_exiting.hpp index 5d6f640ff..e235b978d 100644 --- a/cpp/include/rmm/process_is_exiting.hpp +++ b/cpp/include/rmm/process_is_exiting.hpp @@ -22,15 +22,16 @@ namespace RMM_EXPORT rmm { * 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. * - * All RMM memory resources must be safe to destroy at process shutdown. An MR destructor may - * run during normal program flow (when calling CUDA APIs is safe) or after `exit()` has been - * called (when it is not). Authors must satisfy this by either: + * Use this function from a destructor (or a helper invoked by a destructor, such as a + * `release()` method) when the object may be 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 - * simply be leaked; the OS reclaims them when the process exits. + * returns `true`. In that case, resources that would have been explicitly released should be + * leaked; the OS reclaims them when the process exits. * * 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. diff --git a/cpp/tests/process_is_exiting_shutdown_test.cpp b/cpp/tests/process_is_exiting_shutdown_test.cpp index b05475f85..e3c9561b4 100644 --- a/cpp/tests/process_is_exiting_shutdown_test.cpp +++ b/cpp/tests/process_is_exiting_shutdown_test.cpp @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include +#include #include #include @@ -20,10 +20,11 @@ * via std::atexit on first use; it is called from get_ref_map() immediately after the static * map is constructed, so the flag is observed as true during that map's destructor. * - * This binary exercises the same mechanism without any per-device map or CUDA involvement: + * This binary exercises the real integration point without any CUDA involvement: * * 1. main() registers `checker` via std::atexit FIRST. - * 2. main() then calls register_process_exit_hook(), which registers its internal + * 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 @@ -45,9 +46,9 @@ void checker() noexcept int main() { - // Register the checker before rmm's flag-setter so that the checker runs AFTER the - // flag-setter at termination. + // 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; } - rmm::detail::register_process_exit_hook(); + [[maybe_unused]] auto& map = rmm::mr::detail::get_ref_map(); return 0; } diff --git a/cpp/tests/process_is_exiting_tests.cpp b/cpp/tests/process_is_exiting_tests.cpp index afe501f20..03fcb1fd0 100644 --- a/cpp/tests/process_is_exiting_tests.cpp +++ b/cpp/tests/process_is_exiting_tests.cpp @@ -3,19 +3,17 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include #include #include /* - * Tests for rmm::process_is_exiting() and the internal register_process_exit_hook(). + * Tests for rmm::process_is_exiting(). * * These tests cover the observable behavior of the API during normal program execution. The - * complementary property -- that the flag flips to true before the destructor of a static - * object whose construction was sequenced-before register_process_exit_hook() -- is exercised - * by the standalone PROCESS_IS_EXITING_SHUTDOWN_TEST binary, whose exit code is checked by - * CTest. + * complementary property -- that the flag flips to true before the destructor of the + * per-device resource map returned by detail::get_ref_map() -- is exercised by the standalone + * PROCESS_IS_EXITING_SHUTDOWN_TEST binary, whose exit code is checked by CTest. */ // Flag should be false while the test process is running normally. @@ -28,13 +26,3 @@ TEST(ProcessIsExitingTest, QueryIsIdempotent) EXPECT_FALSE(rmm::process_is_exiting()); EXPECT_FALSE(rmm::process_is_exiting()); } - -// register_process_exit_hook is idempotent: repeated calls are safe and do not flip the flag -// before the actual exit. -TEST(ProcessIsExitingTest, RegisterHookIsIdempotent) -{ - rmm::detail::register_process_exit_hook(); - rmm::detail::register_process_exit_hook(); - rmm::detail::register_process_exit_hook(); - EXPECT_FALSE(rmm::process_is_exiting()); -} From 4a814f36ed4a838ffea7a6893da7a1ddf5490429 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Mon, 27 Apr 2026 14:50:33 -0500 Subject: [PATCH 08/10] Document shutdown support is limited to the per-device map --- cpp/include/rmm/process_is_exiting.hpp | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/cpp/include/rmm/process_is_exiting.hpp b/cpp/include/rmm/process_is_exiting.hpp index e235b978d..7c2b891ce 100644 --- a/cpp/include/rmm/process_is_exiting.hpp +++ b/cpp/include/rmm/process_is_exiting.hpp @@ -17,10 +17,11 @@ namespace RMM_EXPORT rmm { /** * @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, can run - * 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. + * 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 destructor (or a helper invoked by a destructor, such as a * `release()` method) when the object may be destroyed during process termination. In that @@ -33,6 +34,11 @@ namespace RMM_EXPORT rmm { * 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 otherwise unsupported. The intended + * use of this facility is for resources owned by RMM's internal per-device resource map; 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. * From 520d6f9f152b83a860fc79bddd40d817a1c6395a Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Mon, 27 Apr 2026 15:34:41 -0500 Subject: [PATCH 09/10] Abort if atexit registration fails --- cpp/src/runtime_shutdown.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/cpp/src/runtime_shutdown.cpp b/cpp/src/runtime_shutdown.cpp index bf4abb4df..de0fe8fa8 100644 --- a/cpp/src/runtime_shutdown.cpp +++ b/cpp/src/runtime_shutdown.cpp @@ -3,6 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include #include #include @@ -26,7 +27,9 @@ namespace detail { void register_process_exit_hook() noexcept { std::call_once(registered, []() { - std::atexit([]() noexcept { exiting.store(true, std::memory_order_release); }); + 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"); }); } From 963e54ae00492bdb72071f2cdb6e39fbc0bb90ce Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Thu, 30 Apr 2026 13:37:01 +0000 Subject: [PATCH 10/10] Address reviews --- cpp/include/rmm/detail/runtime_shutdown.hpp | 4 +- cpp/include/rmm/mr/per_device_resource.hpp | 40 ++++++++----------- cpp/include/rmm/process_is_exiting.hpp | 21 +++++----- cpp/src/runtime_shutdown.cpp | 5 +++ .../process_is_exiting_shutdown_test.cpp | 17 ++++---- 5 files changed, 43 insertions(+), 44 deletions(-) diff --git a/cpp/include/rmm/detail/runtime_shutdown.hpp b/cpp/include/rmm/detail/runtime_shutdown.hpp index 97c28fa58..2f6ae1403 100644 --- a/cpp/include/rmm/detail/runtime_shutdown.hpp +++ b/cpp/include/rmm/detail/runtime_shutdown.hpp @@ -12,7 +12,9 @@ namespace detail { /** * @brief Register the atexit callback that flips the flag observed by `rmm::process_is_exiting()`. * - * Idempotent: safe to call from multiple places; only the first call registers the callback. + * 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; diff --git a/cpp/include/rmm/mr/per_device_resource.hpp b/cpp/include/rmm/mr/per_device_resource.hpp index 72595a766..996b7eb13 100644 --- a/cpp/include/rmm/mr/per_device_resource.hpp +++ b/cpp/include/rmm/mr/per_device_resource.hpp @@ -77,9 +77,7 @@ RMM_EXPORT inline auto& get_ref_map() { static std::map> device_id_to_resource; - // Register the process-exit hook immediately after constructing the map. The C++ standard - // guarantees that an atexit callback registered after a static object's construction runs - // before that object's destructor at termination. + // Register the process-exit hook immediately after constructing the map. rmm::detail::register_process_exit_hook(); return device_id_to_resource; } @@ -141,9 +139,9 @@ 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 `new_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`. + * @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` @@ -170,9 +168,8 @@ inline cuda::mr::any_resource 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`, @@ -185,10 +182,9 @@ inline cuda::mr::any_resource 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 a reference to `new_resource_ref` alive until process - * exit. The underlying resource's 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`. + * @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` @@ -245,9 +241,9 @@ 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 `new_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`. + * @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 @@ -265,9 +261,8 @@ inline cuda::mr::any_resource 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`, @@ -279,10 +274,9 @@ inline cuda::mr::any_resource 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 a reference to `new_resource_ref` alive until process - * exit. The underlying resource's 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`. + * @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 diff --git a/cpp/include/rmm/process_is_exiting.hpp b/cpp/include/rmm/process_is_exiting.hpp index 7c2b891ce..37acf1c1c 100644 --- a/cpp/include/rmm/process_is_exiting.hpp +++ b/cpp/include/rmm/process_is_exiting.hpp @@ -23,10 +23,11 @@ namespace RMM_EXPORT rmm { * CUDA API calls may dereference released state and crash inside libcuda rather than returning * an error. * - * Use this function from a destructor (or a helper invoked by a destructor, such as a - * `release()` method) when the object may be 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: + * 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 @@ -34,10 +35,9 @@ namespace RMM_EXPORT rmm { * 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 otherwise unsupported. The intended - * use of this facility is for resources owned by RMM's internal per-device resource map; users - * should not create their own static containers of RMM objects and rely on - * `rmm::process_is_exiting()` to make those destructors safe. + * 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. @@ -47,9 +47,10 @@ namespace RMM_EXPORT rmm { * class my_resource final : public ... { * ~my_resource() override * { - * if (!rmm::process_is_exiting()) { - * RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaFree(ptr_)); + * if (rmm::process_is_exiting()) { + * return; * } + * RMM_ASSERT_CUDA_SUCCESS_SAFE_SHUTDOWN(cudaFree(ptr_)); * } * }; * @endcode diff --git a/cpp/src/runtime_shutdown.cpp b/cpp/src/runtime_shutdown.cpp index de0fe8fa8..0e76be4aa 100644 --- a/cpp/src/runtime_shutdown.cpp +++ b/cpp/src/runtime_shutdown.cpp @@ -26,6 +26,11 @@ 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); }); diff --git a/cpp/tests/process_is_exiting_shutdown_test.cpp b/cpp/tests/process_is_exiting_shutdown_test.cpp index e3c9561b4..fd625235b 100644 --- a/cpp/tests/process_is_exiting_shutdown_test.cpp +++ b/cpp/tests/process_is_exiting_shutdown_test.cpp @@ -10,15 +10,8 @@ /* * Standalone test that verifies rmm::process_is_exiting() reports true while the C runtime is - * executing atexit handlers. The fix in rapidsai/rmm#2367 depends on this behavior: resource - * destructors running during exit() must observe the flag as set so that they can skip CUDA - * calls. - * - * 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 having register_process_exit_hook() register a flag-setter - * via std::atexit on first use; it is called from get_ref_map() immediately after the static - * map is constructed, so the flag is observed as true during that map's destructor. + * 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: * @@ -39,7 +32,11 @@ namespace { void checker() noexcept { - if (!rmm::process_is_exiting()) { std::_Exit(1); } + 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