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..2f6ae1403 --- /dev/null +++ b/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 + +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 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..17c51ab6b 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,6 +487,8 @@ class stream_ordered_memory_resource : public crtp { { 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)); diff --git a/cpp/include/rmm/mr/per_device_resource.hpp b/cpp/include/rmm/mr/per_device_resource.hpp index df4451be1..996b7eb13 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,8 @@ RMM_EXPORT inline auto& get_ref_map() { static std::map> 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 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 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 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 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 */ diff --git a/cpp/include/rmm/process_is_exiting.hpp b/cpp/include/rmm/process_is_exiting.hpp new file mode 100644 index 000000000..37acf1c1c --- /dev/null +++ b/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 + +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 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..c1ad542a5 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,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())); } 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..4d432c8d5 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,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); diff --git a/cpp/src/mr/detail/pool_memory_resource_impl.cpp b/cpp/src/mr/detail/pool_memory_resource_impl.cpp index f00c7fcb8..4297ee056 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,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); diff --git a/cpp/src/runtime_shutdown.cpp b/cpp/src/runtime_shutdown.cpp new file mode 100644 index 000000000..0e76be4aa --- /dev/null +++ b/cpp/src/runtime_shutdown.cpp @@ -0,0 +1,42 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include + +#include +#include +#include + +namespace RMM_NAMESPACE { + +namespace { + +std::atomic 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 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..fd625235b --- /dev/null +++ b/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 +#include + +#include + +/* + * 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; +} diff --git a/cpp/tests/process_is_exiting_tests.cpp b/cpp/tests/process_is_exiting_tests.cpp new file mode 100644 index 000000000..03fcb1fd0 --- /dev/null +++ b/cpp/tests/process_is_exiting_tests.cpp @@ -0,0 +1,28 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +/* + * 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 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. +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()); +}