From ec86fe5e8eacaf4a71583602ef63874f033de21c Mon Sep 17 00:00:00 2001 From: niranda perera Date: Mon, 22 Jun 2026 16:18:16 -0700 Subject: [PATCH 01/13] adding back ref to host and pinned MRs Signed-off-by: niranda perera --- cpp/benchmarks/bench_comm.cpp | 2 +- cpp/benchmarks/bench_memory_resources.cpp | 44 +++-- .../rapidsmpf/memory/buffer_resource.hpp | 34 +++- .../rapidsmpf/memory/host_memory_resource.hpp | 21 ++- .../memory/pinned_memory_resource.hpp | 100 +++++----- cpp/src/memory/buffer_resource.cpp | 44 +++-- cpp/src/memory/pinned_memory_resource.cpp | 56 +++--- .../streaming/base_streaming_fixture.hpp | 2 +- cpp/tests/streaming/test_fanout.cpp | 2 +- cpp/tests/test_buffer.cpp | 2 +- cpp/tests/test_buffer_resource.cpp | 79 ++++++-- cpp/tests/test_config.cpp | 33 ++-- cpp/tests/test_host_buffer.cpp | 90 +++++---- cpp/tests/test_memory_resources.cpp | 22 ++- cpp/tests/test_shuffler.cpp | 4 +- cpp/tests/test_spill_manager.cpp | 2 +- cpp/tests/test_statistics.cpp | 11 +- .../rapidsmpf/memory/buffer_resource.pxd | 7 +- .../rapidsmpf/memory/buffer_resource.pyi | 7 +- .../rapidsmpf/memory/buffer_resource.pyx | 80 +++++--- .../memory/pinned_memory_resource.pxd | 17 ++ .../memory/pinned_memory_resource.pyi | 18 +- .../memory/pinned_memory_resource.pyx | 172 ++++++++++-------- python/rapidsmpf/rapidsmpf/statistics.pyx | 3 +- .../rapidsmpf/rapidsmpf/tests/test_config.py | 12 +- 25 files changed, 544 insertions(+), 320 deletions(-) diff --git a/cpp/benchmarks/bench_comm.cpp b/cpp/benchmarks/bench_comm.cpp index 7e1a4c60a..d188e61f7 100644 --- a/cpp/benchmarks/bench_comm.cpp +++ b/cpp/benchmarks/bench_comm.cpp @@ -324,7 +324,7 @@ int main(int argc, char** argv) { rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource_ref(); auto br = BufferResource::create( mr, - PinnedMemoryResource::Disabled, + PinnedMemoryDisabled, {}, std::chrono::milliseconds{1}, std::make_shared( diff --git a/cpp/benchmarks/bench_memory_resources.cpp b/cpp/benchmarks/bench_memory_resources.cpp index 8178dca24..8ffe745a1 100644 --- a/cpp/benchmarks/bench_memory_resources.cpp +++ b/cpp/benchmarks/bench_memory_resources.cpp @@ -8,11 +8,14 @@ #include +#include #include #include #include +#include #include +#include #include #include #include @@ -96,6 +99,22 @@ class NewDelete { friend void get_property(NewDelete const&, cuda::mr::host_accessible) noexcept {} }; +// Build a lightweight `BufferResource` that owns a pinned pool with the given +// properties. Host and pinned memory resources are constructible only by a +// `BufferResource`. The dedicated spill-check thread is disabled and a +// single-stream pool is used to keep overhead minimal for micro-benchmarks. +std::shared_ptr make_pinned_buffer_resource( + rapidsmpf::PinnedPoolProperties props +) { + return rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), + std::move(props), + {}, + std::nullopt, // disable the periodic spill-check thread + std::make_shared(1, rmm::cuda_stream::flags::non_blocking) + ); +} + // Helper function to create a type-erased host memory resource. cuda::mr::any_resource create_host_memory_resource( ResourceType const& resource_type @@ -104,10 +123,16 @@ cuda::mr::any_resource create_host_memory_resource( case ResourceType::NEW_DELETE: return NewDelete{}; case ResourceType::HOST_MEMORY_RESOURCE: - return rapidsmpf::HostMemoryResource{}; + { + // The returned `any_resource` carries a back-reference that keeps the + // (local) `BufferResource` alive, so no external anchor is needed. + auto br = make_pinned_buffer_resource(rapidsmpf::PinnedPoolProperties{}); + return br->host_mr(); + } case ResourceType::PINNED_MEMORY_RESOURCE: { - auto mr = rapidsmpf::PinnedMemoryResource::make_if_available(); + auto br = make_pinned_buffer_resource(rapidsmpf::PinnedPoolProperties{}); + auto mr = br->try_pinned_mr(); RAPIDSMPF_EXPECTS( mr.has_value(), "pinned memory is not supported on this system", @@ -442,9 +467,8 @@ void BM_PinnedFirstAlloc_InitialPoolSize(benchmark::State& state) { for (auto _ : state) { state.PauseTiming(); - auto mr = rapidsmpf::PinnedMemoryResource::make_if_available( - rapidsmpf::get_current_numa_node(), props - ); + auto br = make_pinned_buffer_resource(props); + auto mr = br->try_pinned_mr(); state.ResumeTiming(); void* ptr = mr->allocate(stream, allocation_size); stream.synchronize(); @@ -499,13 +523,11 @@ void BM_PinnedPoolInit_InitialPoolSize(benchmark::State& state) { }; for (auto _ : state) { - auto mr = rapidsmpf::PinnedMemoryResource::make_if_available( - rapidsmpf::get_current_numa_node(), props - ); - benchmark::DoNotOptimize(mr); - // Destroy mr at end of iteration (pool teardown excluded from timing). + auto br = make_pinned_buffer_resource(props); + benchmark::DoNotOptimize(br); + // Destroy br at end of iteration (pool teardown excluded from timing). state.PauseTiming(); - mr.reset(); + br.reset(); state.ResumeTiming(); } diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 9190a87f7..6d106d1f4 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -91,9 +91,11 @@ class BufferResource : public std::enable_shared_from_this { * allocations are tracked for memory-limit accounting and statistics, use * `BufferResource::device_mr()` instead of the original memory resource after * construction. - * @param pinned_mr Pinned host memory resource used for - * `MemoryType::PINNED_HOST` allocations, or `PinnedMemoryResource::Disabled` to - * disable pinned allocations. + * @param pinned_pool_properties Configuration for the pinned host memory pool + * used for `MemoryType::PINNED_HOST` allocations, or `PinnedMemoryDisabled` to + * disable pinned allocations. The pinned resource is constructed internally and + * owned by the `BufferResource`; it is only created when pinned host memory is + * supported on the system (see `is_pinned_memory_resources_supported()`). * @param memory_limits Maximum allocation limits in bytes per `MemoryType`. Missing * entries are treated as unlimited. * @param periodic_spill_check Interval between periodic spill checks. `std::nullopt` @@ -105,7 +107,7 @@ class BufferResource : public std::enable_shared_from_this { */ [[nodiscard]] static std::shared_ptr create( cuda::mr::any_resource device_mr, - std::optional pinned_mr = PinnedMemoryResource::Disabled, + std::optional pinned_pool_properties = PinnedMemoryDisabled, std::unordered_map memory_limits = {}, std::optional periodic_spill_check = std::chrono::milliseconds{1}, std::shared_ptr stream_pool = std::make_shared< @@ -217,6 +219,12 @@ class BufferResource : public std::enable_shared_from_this { * @brief Get the RMM host memory resource. * * @return Reference to the RMM resource used for host allocations. + * + * @note The returned non-owning `resource_ref` does not by itself keep this + * `BufferResource` alive, but the underlying `HostMemoryResource` carries a + * back-reference: promoting the ref to an owning `cuda::mr::any_resource` (as + * RMM/CCCL containers do internally) keeps this `BufferResource` alive for the + * lifetime of that owning copy. */ [[nodiscard]] rmm::host_async_resource_ref host_mr() noexcept; @@ -225,16 +233,24 @@ class BufferResource : public std::enable_shared_from_this { * * @throws std::invalid_argument if no pinned memory resource is available. * @return Reference to the RMM resource used for pinned host allocations. + * + * @note The returned non-owning `resource_ref` does not by itself keep this + * `BufferResource` alive, but the underlying `PinnedMemoryResource` carries a + * back-reference: promoting the ref to an owning `cuda::mr::any_resource` (as + * RMM/CCCL containers do internally) keeps this `BufferResource` alive for the + * lifetime of that owning copy. */ [[nodiscard]] rmm::host_device_async_resource_ref pinned_mr(); /** * @brief Get the pinned host memory resource if available. * - * @return The pinned host memory resource as an `any_resource`, or `std::nullopt` if - * pinned host memory is not available. + * @return A back-referenced copy of the pinned host memory resource, or + * `std::nullopt` if pinned host memory is not available. The returned handle + * keeps this `BufferResource` alive for as long as the handle (or any copy of + * it) lives. */ - [[nodiscard]] std::optional try_pinned_mr() const noexcept; + [[nodiscard]] std::optional try_pinned_mr() const; /** * @brief Returns the currently available memory for a given memory type, in bytes. @@ -340,9 +356,7 @@ class BufferResource : public std::enable_shared_from_this { [[nodiscard]] MemoryReservation reserve_or_fail(std::size_t size, Range mem_types) { // try to reserve memory from the given order for (auto const& mem_type : mem_types) { - if (mem_type == MemoryType::PINNED_HOST - && pinned_mr_ == PinnedMemoryResource::Disabled) - { + if (mem_type == MemoryType::PINNED_HOST && !pinned_mr_.has_value()) { // Pinned host memory is only available if the memory resource is // available. continue; diff --git a/cpp/include/rapidsmpf/memory/host_memory_resource.hpp b/cpp/include/rapidsmpf/memory/host_memory_resource.hpp index 3fa9d1975..80506e6f9 100644 --- a/cpp/include/rapidsmpf/memory/host_memory_resource.hpp +++ b/cpp/include/rapidsmpf/memory/host_memory_resource.hpp @@ -13,9 +13,12 @@ #include #include +#include namespace rapidsmpf { +class BufferResource; + /** * @brief Host memory resource using standard CPU allocation. * @@ -28,10 +31,14 @@ namespace rapidsmpf { * region. THP can improve device-host memory transfer performance for large * buffers. The hint is applied via `madvise(MADV_HUGEPAGE)` and may be ignored * by the kernel depending on system configuration or resource availability. + * + * @note Instances are constructed only by `BufferResource`, which installs a + * back-reference (via `BackRefMixin`) so that any copy of the + * resource keeps its owning `BufferResource` alive. There is no public + * constructor; obtain the resource through a `BufferResource`. */ -class HostMemoryResource { +class HostMemoryResource : public BackRefMixin { public: - HostMemoryResource() = default; ~HostMemoryResource() = default; HostMemoryResource(HostMemoryResource const&) = default; ///< Copyable. @@ -112,7 +119,9 @@ class HostMemoryResource { * @brief Compares this resource to another resource. * * All instances are stateless and interchangeable, so this always returns - * true. + * true. Instances can only be created by `BufferResource` (which installs + * the back-reference exactly once), so the back-reference does not affect + * equality. * * @param other The resource to compare with. * @return true @@ -138,6 +147,12 @@ class HostMemoryResource { friend void get_property( HostMemoryResource const&, cuda::mr::host_accessible ) noexcept {} + + private: + /// @brief Default construct. Private: only `BufferResource` creates instances. + HostMemoryResource() = default; + + friend class BufferResource; }; static_assert(cuda::mr::resource); diff --git a/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp b/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp index 653c49f20..63bd6e36d 100644 --- a/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp +++ b/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp @@ -6,6 +6,7 @@ #include #include +#include #include #include @@ -20,6 +21,7 @@ #include #include #include +#include #include #include @@ -32,6 +34,8 @@ namespace rapidsmpf { +class BufferResource; + /** * @brief Checks if the PinnedMemoryResource is supported for the current CUDA version. * @@ -64,6 +68,16 @@ inline bool is_pinned_memory_resources_supported() { return supported; } +/** + * @brief Sentinel used to disable pinned host memory. + * + * Pass this in place of a `PinnedPoolProperties` (e.g. to `BufferResource::create()`) + * to disable pinned host memory allocations. It is simply `std::nullopt` typed as + * `std::nullopt_t`, so it implicitly converts to an empty + * `std::optional`. + */ +constexpr std::nullopt_t PinnedMemoryDisabled = std::nullopt; + /** * @brief Properties for configuring a pinned memory pool. */ @@ -75,8 +89,32 @@ struct PinnedPoolProperties { /// @brief Maximum size of the pool. `std::nullopt` means no limit. std::optional max_pool_size = std::nullopt; + + /// @brief NUMA node from which pinned memory should be allocated. Defaults to + /// the NUMA node of the calling thread. + int numa_id = get_current_numa_node(); }; +/** + * @brief Parse pinned memory pool properties from configuration options. + * + * Recognized options: + * - "pinned_memory": enable pinned memory. + * - "pinned_initial_pool_size" (bytes or percentage): initial pool size. + * - Byte values (e.g. "1 MiB") are applied literally. + * - Percentages (e.g. "10%") are relative to `get_host_memory_per_gpu()`. + * - "pinned_max_pool_size" (bytes, percentage, or disabled): maximum pool size. + * - Byte and percentages uses the same parsing rules as "pinned_initial_pool_size". + * - A disabled value (e.g. "off") leaves the pool unbounded. + * + * @param options Configuration options. + * @return The parsed `PinnedPoolProperties` when "pinned_memory" is enabled, + * otherwise `std::nullopt` (pinned host memory disabled). + */ +std::optional pinned_pool_properties_from_options( + config::Options options +); + /** * @brief Memory resource that provides pinned (page-locked) host memory using a pool. * @@ -88,50 +126,21 @@ struct PinnedPoolProperties { * This resource allocates and deallocates pinned host memory asynchronously through * CUDA streams. It offers higher bandwidth and lower latency for device transfers * compared to regular pageable host memory. + * + * @note Instances are constructed only by `BufferResource`, which installs a + * back-reference (via `BackRefMixin`) so that any copy of the + * resource keeps its owning `BufferResource` alive. There is no public + * constructor; obtain the resource through a `BufferResource`. Use the free + * function `is_pinned_memory_resources_supported()` to query availability. */ class PinnedMemoryResource final : public cuda::mr::shared_resource< - detail::RmmResourceAdaptorImpl> { + detail::RmmResourceAdaptorImpl>, + public BackRefMixin { using shared_base = cuda::mr::shared_resource< detail::RmmResourceAdaptorImpl>; public: - /// @brief Sentinel value indicating that pinned host memory is disabled. - static constexpr std::nullopt_t Disabled = std::nullopt; - - /** - * @brief Create a pinned memory resource if the system supports pinned memory. - * - * @param numa_id The NUMA node to associate with the resource. Defaults to the - * current NUMA node. - * @param pool_properties Properties for configuring the pinned memory pool. - * - * @return A `PinnedMemoryResource` when supported, otherwise `std::nullopt`. - * - * @see PinnedMemoryResource::PinnedMemoryResource - */ - static std::optional make_if_available( - int numa_id = get_current_numa_node(), PinnedPoolProperties pool_properties = {} - ); - - /** - * @brief Construct from configuration options. - * - * Recognized options: - * - "pinned_memory": enable pinned memory. - * - "pinned_initial_pool_size" (bytes or percentage): initial pool size. - * - Byte values (e.g. "1 MiB") are applied literally. - * - Percentages (e.g. "10%") are relative to `get_host_memory_per_gpu()`. - * - "pinned_max_pool_size" (bytes, percentage, or disabled): maximum pool size. - * - Byte and percentages uses the same parsing rules as "pinned_initial_pool_size". - * - A disabled value (e.g. "off") leaves the pool unbounded. - * - * @param options Configuration options. - * @return A `PinnedMemoryResource` if pinned memory is enabled and supported; - * otherwise `std::nullopt`. - */ - static std::optional from_options(config::Options options); - /** * @brief Allocates pinned host memory associated with a CUDA stream. * @@ -172,6 +181,10 @@ class PinnedMemoryResource final /** * @brief Equality comparison. * + * Instances can only be created by `BufferResource` (which installs the + * back-reference exactly once), so equality is determined solely by the + * underlying shared state. + * * @param other The other resource to compare. * @return True if the two resources share the same underlying shared state. */ @@ -217,16 +230,17 @@ class PinnedMemoryResource final /** * @brief Construct a pinned (page-locked) host memory resource. * - * Private — use `make_if_available` or `from_options` to obtain an instance. + * Private — only `BufferResource` constructs instances (it installs the + * back-reference immediately after construction). * - * @param numa_id NUMA node from which memory should be allocated. - * @param pool_properties Properties for configuring the pinned memory pool. + * @param pool_properties Properties for configuring the pinned memory pool, + * including the NUMA node from which memory should be allocated. * * @throws std::invalid_argument If pinned host memory pools are not supported. */ - PinnedMemoryResource( - int numa_id = get_current_numa_node(), PinnedPoolProperties pool_properties = {} - ); + explicit PinnedMemoryResource(PinnedPoolProperties pool_properties); + + friend class BufferResource; PinnedPoolProperties pool_properties_; ///< properties used to configure the pool }; diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index 5eb9e91f7..5500f07e0 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -50,12 +50,21 @@ BufferResource::BufferResource( std::shared_ptr BufferResource::create( cuda::mr::any_resource device_mr, - std::optional pinned_mr, + std::optional pinned_pool_properties, std::unordered_map memory_limits, std::optional periodic_spill_check, std::shared_ptr stream_pool, std::shared_ptr statistics ) { + // Construct the pinned resource in place (`BufferResource` is a friend of + // `PinnedMemoryResource`), so no observable, back-reference-less instance is + // ever minted. Only create it when pinned memory is both requested and + // supported on this system. + std::optional pinned_mr; + if (pinned_pool_properties.has_value() && is_pinned_memory_resources_supported()) { + pinned_mr = PinnedMemoryResource{*pinned_pool_properties}; + } + std::shared_ptr br{new BufferResource{ std::move(device_mr), std::move(pinned_mr), @@ -65,12 +74,18 @@ std::shared_ptr BufferResource::create( std::move(statistics) }}; - // Install the back-reference on the device adaptor *after* construction so - // that `weak_from_this()` is valid. The adaptor holds only a `weak_ptr`, + // Install the back-reference on the owned resources *after* construction so + // that `weak_from_this()` is valid. Each resource holds only a `weak_ptr`, // avoiding a reference cycle. When downstream code promotes a non-owning - // `resource_ref` returned by `device_mr()` to an owning - // `cuda::mr::any_resource`. - br->owning_mr_.set_backref(br->weak_from_this()); + // `resource_ref` returned by `device_mr()`/`host_mr()`/`pinned_mr()` to an + // owning `cuda::mr::any_resource`, the copy promotes the weak reference to a + // strong one and keeps this `BufferResource` alive. + auto const weak = br->weak_from_this(); + br->owning_mr_.set_backref(weak); + br->host_mr_.set_backref(weak); + if (br->pinned_mr_.has_value()) { + br->pinned_mr_->set_backref(weak); + } return br; } @@ -79,13 +94,12 @@ std::shared_ptr BufferResource::from_options( config::Options options, std::shared_ptr statistics ) { - auto pinned_mr = PinnedMemoryResource::from_options(options); std::unordered_map memory_limits{ {MemoryType::DEVICE, device_limit_from_options(options)} }; return create( std::move(mr), - std::move(pinned_mr), + pinned_pool_properties_from_options(options), std::move(memory_limits), periodic_spill_check_from_options(options), stream_pool_from_options(options), @@ -101,7 +115,7 @@ std::int64_t BufferResource::memory_available(MemoryType mem_type) const noexcep case MemoryType::DEVICE: return limit - owning_mr_.current_allocated(); case MemoryType::PINNED_HOST: - if (pinned_mr_ == PinnedMemoryResource::Disabled) { + if (!pinned_mr_.has_value()) { return 0; } else { return limit - pinned_mr_->current_allocated(); @@ -127,7 +141,6 @@ RmmResourceAdaptor& BufferResource::device_mr_adaptor() noexcept { } rmm::host_async_resource_ref BufferResource::host_mr() noexcept { - // TODO: returned ref will not keep the BufferResource alive return host_mr_; } @@ -135,14 +148,12 @@ rmm::host_device_async_resource_ref BufferResource::pinned_mr() { RAPIDSMPF_EXPECTS( pinned_mr_, "no pinned memory resource is available", std::invalid_argument ); - // TODO: returned ref will not keep the BufferResource alive return *pinned_mr_; } -std::optional BufferResource::try_pinned_mr() const noexcept { - // since any_host_device_resource is constructible from - // host_device_async_resource_ref, optional can be returned as-is. - // TODO: returned ref will not keep the BufferResource alive +std::optional BufferResource::try_pinned_mr() const { + // Returning by value copies the back-referenced `PinnedMemoryResource`, so the + // returned handle (and any copy of it) keeps this `BufferResource` alive. return pinned_mr_; } @@ -150,8 +161,7 @@ std::pair BufferResource::reserve( MemoryType mem_type, std::size_t size, AllowOverbooking allow_overbooking ) { RAPIDSMPF_EXPECTS( - mem_type != MemoryType::PINNED_HOST - || pinned_mr_ != PinnedMemoryResource::Disabled, + mem_type != MemoryType::PINNED_HOST || pinned_mr_.has_value(), "pinned memory resource is not available", std::invalid_argument ); diff --git a/cpp/src/memory/pinned_memory_resource.cpp b/cpp/src/memory/pinned_memory_resource.cpp index 4141b38b7..45d254ef8 100644 --- a/cpp/src/memory/pinned_memory_resource.cpp +++ b/cpp/src/memory/pinned_memory_resource.cpp @@ -41,57 +41,47 @@ cuda::memory_pool_properties get_memory_pool_properties( } // namespace -PinnedMemoryResource::PinnedMemoryResource( - int numa_id, PinnedPoolProperties pool_properties -) +PinnedMemoryResource::PinnedMemoryResource(PinnedPoolProperties pool_properties) : shared_base([&] { RAPIDSMPF_EXPECTS( is_pinned_memory_resources_supported(), "Pinned host memory is not supported on this system. " "CUDA " RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION_STR " is one of the requirements, but additional platform or driver " - "constraints may apply. If needed, use `PinnedMemoryResource::Disabled` " - "to disable pinned host memory, noting that this may significantly " - "degrade spilling performance.", + "constraints may apply. If needed, disable pinned host memory (pass " + "`std::nullopt` for the pinned pool to `BufferResource`), noting that " + "this may significantly degrade spilling performance.", std::invalid_argument ); return cuda::mr::make_shared_resource< detail::RmmResourceAdaptorImpl>( - std::in_place, numa_id, get_memory_pool_properties(pool_properties) + std::in_place, + pool_properties.numa_id, + get_memory_pool_properties(pool_properties) ); }()), pool_properties_{std::move(pool_properties)} {} -std::optional PinnedMemoryResource::make_if_available( - int numa_id, PinnedPoolProperties pool_properties -) { - if (is_pinned_memory_resources_supported()) { - return PinnedMemoryResource{numa_id, std::move(pool_properties)}; - } - return PinnedMemoryResource::Disabled; -} - -std::optional PinnedMemoryResource::from_options( +std::optional pinned_pool_properties_from_options( config::Options options ) { bool const pinned_memory = options.get("pinned_memory", parse_string); - - if (pinned_memory && is_pinned_memory_resources_supported()) { - auto const host_memory_per_gpu = get_host_memory_per_gpu(); - auto const total = safe_cast(host_memory_per_gpu); - PinnedPoolProperties pool_properties{ - .initial_pool_size = options.get( - "pinned_initial_pool_size", - [total](auto const& s) { return parse_nbytes_or_percent(s, total); } - ), - .max_pool_size = options.get>( - "pinned_max_pool_size", - [total](auto const& s) { return parse_nbytes_or_percent(s, total); } - ) - }; - return PinnedMemoryResource{get_current_numa_node(), std::move(pool_properties)}; + if (!pinned_memory) { + return PinnedMemoryDisabled; } - return PinnedMemoryResource::Disabled; + + auto const host_memory_per_gpu = get_host_memory_per_gpu(); + auto const total = safe_cast(host_memory_per_gpu); + return PinnedPoolProperties{ + .initial_pool_size = options.get( + "pinned_initial_pool_size", + [total](auto const& s) { return parse_nbytes_or_percent(s, total); } + ), + .max_pool_size = options.get>( + "pinned_max_pool_size", + [total](auto const& s) { return parse_nbytes_or_percent(s, total); } + ) + }; } } // namespace rapidsmpf diff --git a/cpp/tests/streaming/base_streaming_fixture.hpp b/cpp/tests/streaming/base_streaming_fixture.hpp index 311423aa6..7d67f9dc2 100644 --- a/cpp/tests/streaming/base_streaming_fixture.hpp +++ b/cpp/tests/streaming/base_streaming_fixture.hpp @@ -44,7 +44,7 @@ class BaseStreamingFixture : public ::testing::Test { stream = rmm::cuda_stream_view{}; br = rapidsmpf::BufferResource::create( - mr_cuda, rapidsmpf::PinnedMemoryResource::Disabled, std::move(memory_limits) + mr_cuda, rapidsmpf::PinnedMemoryDisabled, std::move(memory_limits) ); ctx = std::make_shared( std::move(options), GlobalEnvironment->comm_->logger(), br diff --git a/cpp/tests/streaming/test_fanout.cpp b/cpp/tests/streaming/test_fanout.cpp index 3f5e3898c..ec5c212bb 100644 --- a/cpp/tests/streaming/test_fanout.cpp +++ b/cpp/tests/streaming/test_fanout.cpp @@ -546,7 +546,7 @@ class SpillingStreamingFanout : public BaseStreamingFixture { {MemoryType::DEVICE, 0}, }; br = rapidsmpf::BufferResource::create( - mr_cuda, rapidsmpf::PinnedMemoryResource::Disabled, memory_limits + mr_cuda, rapidsmpf::PinnedMemoryDisabled, memory_limits ); auto options = ctx->options(); ctx = std::make_shared( diff --git a/cpp/tests/test_buffer.cpp b/cpp/tests/test_buffer.cpp index a0814c497..dbe749ae3 100644 --- a/cpp/tests/test_buffer.cpp +++ b/cpp/tests/test_buffer.cpp @@ -55,7 +55,7 @@ class BufferRebindStreamTest : public ::testing::TestWithParam { br = BufferResource::create( rmm::mr::get_current_device_resource_ref(), - PinnedMemoryResource::make_if_available(), + PinnedPoolProperties{}, std::unordered_map{}, std::nullopt, stream_pool diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index d3b0fa1c0..6b4ad8bba 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -55,7 +55,7 @@ TEST(BufferResource, ReservationOverbooking) { // Create a buffer resource that always reports 10 KiB of available device memory. auto br = BufferResource::create( rmm::mr::get_current_device_resource_ref(), - PinnedMemoryResource::Disabled, + PinnedMemoryDisabled, {{MemoryType::DEVICE, 10_KiB}} ); EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0); @@ -122,7 +122,7 @@ TEST(BufferResource, ReservationReleasing) { // memory. auto br = BufferResource::create( rmm::mr::get_current_device_resource_ref(), - PinnedMemoryResource::Disabled, + PinnedMemoryDisabled, {{MemoryType::DEVICE, 10_KiB}, {MemoryType::HOST, 10_KiB}} ); EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0); @@ -174,7 +174,7 @@ TEST(BufferResource, MemoryLimit) { // Create a buffer resource that limits available device memory to 10 KiB. auto br = BufferResource::create( - mr_cuda, PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, 10_KiB}} + mr_cuda, PinnedMemoryDisabled, {{MemoryType::DEVICE, 10_KiB}} ); EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 10_KiB); EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0); @@ -252,18 +252,18 @@ TEST_P(PinnedMaxPoolSizeReservationLimitTest, TwoReservations) { rmm::mr::cuda_memory_resource cuda_mr; - auto pinned_mr = PinnedMemoryResource::make_if_available( - get_current_numa_node(), PinnedPoolProperties{.max_pool_size = max_pool_size} - ); - ASSERT_NE(pinned_mr, PinnedMemoryResource::Disabled); - // Wire the PINNED_HOST limit to the pool's max_pool_size (or unlimited if the // pool is unbounded) so reservations respect the same ceiling as allocations. std::unordered_map memory_limits; if (max_pool_size.has_value() && *max_pool_size > 0) { memory_limits[MemoryType::PINNED_HOST] = safe_cast(*max_pool_size); } - auto br = BufferResource::create(cuda_mr, pinned_mr, std::move(memory_limits)); + auto br = BufferResource::create( + cuda_mr, + PinnedPoolProperties{.max_pool_size = max_pool_size}, + std::move(memory_limits) + ); + ASSERT_TRUE(br->try_pinned_mr().has_value()); // First 1 KiB reservation always succeeds. auto [r1, ob1] = br->reserve(MemoryType::PINNED_HOST, 1_KiB, AllowOverbooking::NO); @@ -289,10 +289,10 @@ INSTANTIATE_TEST_SUITE_P( TEST(BufferResource, AllocStatistics) { rmm::mr::cuda_memory_resource mr_cuda; auto stats = Statistics::create(); - auto pinned_mr = PinnedMemoryResource::make_if_available(); + bool const pinned_available = is_pinned_memory_resources_supported(); auto br = BufferResource::create( mr_cuda, - pinned_mr, + PinnedPoolProperties{}, {}, std::nullopt, std::make_shared(1, rmm::cuda_stream::flags::non_blocking), @@ -314,7 +314,7 @@ TEST(BufferResource, AllocStatistics) { br->make_buffer(device_size, stream, r); } // Allocate pinned_host memory once (if available). - if (pinned_mr != PinnedMemoryResource::Disabled) { + if (pinned_available) { auto [r, _] = br->reserve(MemoryType::PINNED_HOST, pinned_size, AllowOverbooking::YES); br->make_buffer(pinned_size, stream, r); @@ -333,7 +333,7 @@ TEST(BufferResource, AllocStatistics) { EXPECT_EQ(dev_bytes.value(), static_cast(2 * device_size)); // pinned_host: 1 allocation of pinned_size (if available). - if (pinned_mr != PinnedMemoryResource::Disabled) { + if (pinned_available) { auto const pinned_bytes = stats->get_stat("alloc-pinned_host-bytes"); EXPECT_EQ(pinned_bytes.count(), 1u); EXPECT_EQ(pinned_bytes.value(), static_cast(pinned_size)); @@ -357,7 +357,7 @@ class BufferResourceReserveOrFailTest : public ::testing::Test { // host memory. BufferResource auto-wraps mr_cuda for allocation tracking. br = BufferResource::create( mr_cuda, - PinnedMemoryResource::Disabled, + PinnedMemoryDisabled, std::unordered_map{{MemoryType::DEVICE, 10_KiB}} ); } @@ -816,6 +816,57 @@ TEST(BufferResource, DeviceMrKeepsBufferResourceAlive) { EXPECT_TRUE(weak_br.expired()) << "BR not destructed, refcount cycle?"; } +TEST(BufferResource, HostMrKeepsBufferResourceAlive) { + constexpr std::size_t N = 1024; + + auto br = BufferResource::create(rmm::mr::get_current_device_resource_ref()); + std::weak_ptr weak_br = br; + auto stream = rmm::cuda_stream_view{}; + + // Allocate a HOST buffer. The underlying `HostBuffer` stores the host memory + // resource as an owning `any_resource`, which copies the `HostMemoryResource`. + // Its `BackRefMixin` base promotes the installed weak ref to a + // `shared_ptr` during the copy. + auto buf = br->make_buffer(stream, br->reserve_or_fail(N, MemoryType::HOST)); + + // Drop the original shared_ptr. The buffer should keep the BR alive. + br.reset(); + EXPECT_FALSE(weak_br.expired()) << "BR freed while host buffer still holds it"; + + EXPECT_NO_THROW(buf.reset()); + + // After the buffer is destroyed, no shared ownership should remain. + EXPECT_TRUE(weak_br.expired()) << "BR not destructed, refcount cycle?"; +} + +TEST(BufferResource, PinnedMrKeepsBufferResourceAlive) { + if (!is_pinned_memory_resources_supported()) { + GTEST_SKIP() << "Pinned memory not supported on this system"; + } + constexpr std::size_t N = 1024; + + auto br = BufferResource::create( + rmm::mr::get_current_device_resource_ref(), PinnedPoolProperties{} + ); + std::weak_ptr weak_br = br; + auto stream = rmm::cuda_stream_view{}; + + // Allocate a PINNED_HOST buffer. The underlying `HostBuffer` stores the pinned + // memory resource as an owning `any_resource`, which copies the + // `PinnedMemoryResource`. Its `BackRefMixin` base promotes the + // installed weak ref to a `shared_ptr` during the copy. + auto buf = br->make_buffer(stream, br->reserve_or_fail(N, MemoryType::PINNED_HOST)); + + // Drop the original shared_ptr. The buffer should keep the BR alive. + br.reset(); + EXPECT_FALSE(weak_br.expired()) << "BR freed while pinned buffer still holds it"; + + EXPECT_NO_THROW(buf.reset()); + + // After the buffer is destroyed, no shared ownership should remain. + EXPECT_TRUE(weak_br.expired()) << "BR not destructed, refcount cycle?"; +} + TEST(RmmResourceAdaptor, EqualityAcrossCopiesAndAccessPaths) { auto br = BufferResource::create(rmm::mr::get_current_device_resource_ref()); any_device_resource copy1{br->device_mr()}; diff --git a/cpp/tests/test_config.cpp b/cpp/tests/test_config.cpp index 4cc8e6000..750234be5 100644 --- a/cpp/tests/test_config.cpp +++ b/cpp/tests/test_config.cpp @@ -476,39 +476,30 @@ TEST(OptionsTest, StatisticsFromOptionsDisabledByDefault) { EXPECT_FALSE(stats->enabled()); } -TEST(OptionsTest, PinnedMemoryResourceFromOptionsEnabledWhenSetToTrue) { +TEST(OptionsTest, PinnedPoolPropertiesFromOptionsEnabledWhenSetToTrue) { std::unordered_map strings = {{"pinned_memory", "True"}}; Options opts(strings); - auto pmr = PinnedMemoryResource::from_options(opts); - - // Should be enabled if system supports it, or Disabled (nullopt) if not - if (is_pinned_memory_resources_supported()) { - EXPECT_NE(pmr, PinnedMemoryResource::Disabled); - EXPECT_TRUE(pmr.has_value()); - } else { - EXPECT_EQ(pmr, PinnedMemoryResource::Disabled); - EXPECT_FALSE(pmr.has_value()); - } + // The parsed properties carry the pinned-memory request regardless of whether + // the system supports pinned memory; `BufferResource` decides whether to + // actually construct the resource. + auto props = pinned_pool_properties_from_options(opts); + EXPECT_TRUE(props.has_value()); } -TEST(OptionsTest, PinnedMemoryResourceFromOptionsDisabledWhenSetToFalse) { +TEST(OptionsTest, PinnedPoolPropertiesFromOptionsDisabledWhenSetToFalse) { std::unordered_map strings = {{"pinned_memory", "False"}}; Options opts(strings); - auto pmr = PinnedMemoryResource::from_options(opts); - - EXPECT_EQ(pmr, PinnedMemoryResource::Disabled); - EXPECT_FALSE(pmr.has_value()); + auto props = pinned_pool_properties_from_options(opts); + EXPECT_FALSE(props.has_value()); } -TEST(OptionsTest, PinnedMemoryResourceFromOptionsDisabledByDefault) { +TEST(OptionsTest, PinnedPoolPropertiesFromOptionsDisabledByDefault) { Options opts; // Empty options - auto pmr = PinnedMemoryResource::from_options(opts); - - EXPECT_EQ(pmr, PinnedMemoryResource::Disabled); - EXPECT_FALSE(pmr.has_value()); + auto props = pinned_pool_properties_from_options(opts); + EXPECT_FALSE(props.has_value()); } TEST(OptionsTest, DeviceLimitFromOptionsReturnsConfiguredLimit) { diff --git a/cpp/tests/test_host_buffer.cpp b/cpp/tests/test_host_buffer.cpp index ccd9e7eae..48691825c 100644 --- a/cpp/tests/test_host_buffer.cpp +++ b/cpp/tests/test_host_buffer.cpp @@ -15,10 +15,12 @@ #include #include +#include #include #include #include +#include #include #include #include @@ -83,10 +85,13 @@ class HostMemoryResource : public ::testing::TestWithParam { GTEST_SKIP() << "HostBuffer is not supported for CUDA versions " "below " RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION_STR; } + // `HostMemoryResource` is constructible only via a `BufferResource`. + br = + rapidsmpf::BufferResource::create(rmm::mr::get_current_device_resource_ref()); } rmm::cuda_stream_view stream{}; - rapidsmpf::HostMemoryResource mr; + std::shared_ptr br; }; // Test with various buffer sizes @@ -125,7 +130,8 @@ TEST_P(HostMemoryResource, from_uint8_vector) { auto source_data = random_vector(0, buffer_size); // Create a host buffer by copying the vector - auto buffer = rapidsmpf::HostBuffer::from_uint8_vector(source_data, stream, mr); + auto buffer = + rapidsmpf::HostBuffer::from_uint8_vector(source_data, stream, br->host_mr()); EXPECT_NO_THROW(test_buffer(std::move(buffer), source_data)); } @@ -133,14 +139,17 @@ TEST_P(HostMemoryResource, from_uint8_vector) { class PinnedResource : public ::testing::TestWithParam { protected: void SetUp() override { - mr = rapidsmpf::PinnedMemoryResource::make_if_available(); - if (mr == rapidsmpf::PinnedMemoryResource::Disabled) { + if (!rapidsmpf::is_pinned_memory_resources_supported()) { GTEST_SKIP() << "PinnedMemoryResource is not supported"; } - } - - void TearDown() override { - mr.reset(); + // `PinnedMemoryResource` is constructible only via a `BufferResource`. The + // handle carries a back-reference that keeps the (local) `BufferResource` + // alive, so there is no need to store the `BufferResource` separately. + mr = rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), + rapidsmpf::PinnedPoolProperties{} + ) + ->try_pinned_mr(); } rmm::cuda_stream_view stream{}; @@ -208,21 +217,35 @@ TEST_P(PinnedResource, from_rmm_device_buffer) { } TEST(PinnedResource, equality) { - auto mr1 = rapidsmpf::PinnedMemoryResource::make_if_available(); - if (mr1 == rapidsmpf::PinnedMemoryResource::Disabled) { + if (!rapidsmpf::is_pinned_memory_resources_supported()) { GTEST_SKIP() << "PinnedMemoryResource is not supported"; } + // Two handles obtained from the same BufferResource share the same pool and + // owner, so they compare equal. + auto br1 = rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{} + ); + auto mr1 = br1->try_pinned_mr(); rapidsmpf::PinnedMemoryResource mr2 = *mr1; EXPECT_EQ(*mr1, mr2); - auto mr3 = rapidsmpf::PinnedMemoryResource::make_if_available(); - EXPECT_NE(mr1, mr3); + + // A handle from a different BufferResource references a different pool and + // owner, so it compares unequal. + auto br2 = rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{} + ); + auto mr3 = br2->try_pinned_mr(); + EXPECT_NE(*mr1, *mr3); } TEST(PinnedResource, transient_mr) { - auto mr = rapidsmpf::PinnedMemoryResource::make_if_available(); - if (mr == rapidsmpf::PinnedMemoryResource::Disabled) { + if (!rapidsmpf::is_pinned_memory_resources_supported()) { GTEST_SKIP() << "PinnedMemoryResource is not supported"; } + auto br = rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{} + ); + auto mr = br->try_pinned_mr(); rmm::cuda_stream_view stream{}; auto source_data = random_vector(0, 1024); @@ -232,8 +255,11 @@ TEST(PinnedResource, transient_mr) { source_data.data(), source_data.size(), stream, *mr ); - // now reset mr, but pinned_host_buffer should keep the shared mr alive + // Now drop both the handle and the owning BufferResource; the device buffer + // holds an owning copy of the pinned resource (which keeps the pool and, via + // the back-reference, the BufferResource alive). mr.reset(); + br.reset(); auto buffer = rapidsmpf::HostBuffer::from_rmm_device_buffer( std::move(pinned_host_buffer), stream @@ -250,10 +276,11 @@ namespace { std::size_t discover_pinned_pool_actual_size( rmm::cuda_stream_view stream, std::size_t requested_max_pool_size = 1_MiB ) { - auto pinned_mr = rapidsmpf::PinnedMemoryResource::make_if_available( - rapidsmpf::get_current_numa_node(), + auto br = rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{.max_pool_size = requested_max_pool_size} ); + auto pinned_mr = br->try_pinned_mr(); auto can_allocate = [&](size_t size) -> bool { try { @@ -299,13 +326,14 @@ TEST(PinnedResource, max_pool_size_limit) { auto stream = rmm::cuda_stream_view{}; // Create a PinnedMemoryResource with max pool size of 1 MiB; driver may round up. - auto pinned_mr = rapidsmpf::PinnedMemoryResource::make_if_available( - rapidsmpf::get_current_numa_node(), - rapidsmpf::PinnedPoolProperties{.initial_pool_size = 0, .max_pool_size = 1_MiB} - ); - if (pinned_mr == rapidsmpf::PinnedMemoryResource::Disabled) { + if (!rapidsmpf::is_pinned_memory_resources_supported()) { GTEST_SKIP() << "PinnedMemoryResource is not supported"; } + auto br = rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), + rapidsmpf::PinnedPoolProperties{.initial_pool_size = 0, .max_pool_size = 1_MiB} + ); + auto pinned_mr = br->try_pinned_mr(); auto alloc_and_dealloc = [&](std::size_t size) { void* ptr = pinned_mr->allocate(stream, size); @@ -324,28 +352,26 @@ TEST(PinnedResource, max_pool_size_limit) { TEST(PinnedResource, from_default_options) { { // disabled by default - auto mr = - rapidsmpf::PinnedMemoryResource::from_options(rapidsmpf::config::Options{}); - EXPECT_EQ(mr, rapidsmpf::PinnedMemoryResource::Disabled); + auto props = + rapidsmpf::pinned_pool_properties_from_options(rapidsmpf::config::Options{}); + EXPECT_FALSE(props.has_value()); } - // check default pool values, if enabled + // check default pool values, when enabled std::unordered_map strings = {{"pinned_memory", "True"}}; - auto mr = rapidsmpf::PinnedMemoryResource::from_options( + auto props = rapidsmpf::pinned_pool_properties_from_options( rapidsmpf::config::Options(strings) ); - if (mr == rapidsmpf::PinnedMemoryResource::Disabled) { - GTEST_SKIP() << "PinnedMemoryResource is not supported"; - } + ASSERT_TRUE(props.has_value()); EXPECT_EQ( - mr->properties().initial_pool_size, + props->initial_pool_size, rapidsmpf::parse_nbytes_or_percent( rapidsmpf::config::DEFAULTS.at("pinned_initial_pool_size"), static_cast(rapidsmpf::get_host_memory_per_gpu()) ) ); EXPECT_EQ( - mr->properties().max_pool_size.value(), + props->max_pool_size.value(), rapidsmpf::parse_nbytes_or_percent( rapidsmpf::config::DEFAULTS.at("pinned_max_pool_size"), static_cast(rapidsmpf::get_host_memory_per_gpu()) diff --git a/cpp/tests/test_memory_resources.cpp b/cpp/tests/test_memory_resources.cpp index 5a72bb89f..0cf371038 100644 --- a/cpp/tests/test_memory_resources.cpp +++ b/cpp/tests/test_memory_resources.cpp @@ -11,28 +11,40 @@ #include #include +#include +#include #include #include #include namespace { +// `HostMemoryResource` and `PinnedMemoryResource` are constructible only by a +// `BufferResource`. The `any_resource` copies returned below each carry a +// back-reference that keeps the (otherwise local) `BufferResource` alive, so no +// external anchor is needed. std::vector> make_host_resources() { + auto br = rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{} + ); std::vector> resources; - resources.emplace_back(rapidsmpf::HostMemoryResource{}); - if (rapidsmpf::is_pinned_memory_resources_supported()) { - resources.emplace_back(*rapidsmpf::PinnedMemoryResource::make_if_available()); + resources.emplace_back(br->host_mr()); + if (auto pinned = br->try_pinned_mr(); pinned.has_value()) { + resources.emplace_back(*pinned); } return resources; } std::vector> make_device_resources() { + auto br = rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{} + ); std::vector> resources; resources.emplace_back(rmm::mr::cuda_memory_resource{}); resources.emplace_back(rmm::mr::cuda_async_memory_resource{}); - if (rapidsmpf::is_pinned_memory_resources_supported()) { - resources.emplace_back(*rapidsmpf::PinnedMemoryResource::make_if_available()); + if (auto pinned = br->try_pinned_mr(); pinned.has_value()) { + resources.emplace_back(*pinned); } return resources; } diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index fd36a06fb..c2a24f929 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -314,7 +314,7 @@ class MemoryLimits_NumPartition std::tie(memory_limits, total_num_partitions, total_num_rows) = GetParam(); br = rapidsmpf::BufferResource::create( rmm::mr::get_current_device_resource_ref(), - rapidsmpf::PinnedMemoryResource::Disabled, + rapidsmpf::PinnedMemoryDisabled, memory_limits ); @@ -460,7 +460,7 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { // allocation counts via `get_main_record().num_current_allocs()`. auto br = rapidsmpf::BufferResource::create( rmm::mr::get_current_device_resource_ref(), - rapidsmpf::PinnedMemoryResource::Disabled, + rapidsmpf::PinnedMemoryDisabled, {{rapidsmpf::MemoryType::DEVICE, k_no_spill_limit}}, std::nullopt // disable periodic spill check ); diff --git a/cpp/tests/test_spill_manager.cpp b/cpp/tests/test_spill_manager.cpp index b3abdab17..6f825d8d1 100644 --- a/cpp/tests/test_spill_manager.cpp +++ b/cpp/tests/test_spill_manager.cpp @@ -27,7 +27,7 @@ TEST(SpillManager, SpillFunction) { std::int64_t mem_available = 10_KiB; auto br = BufferResource::create( rmm::mr::get_current_device_resource_ref(), - rapidsmpf::PinnedMemoryResource::Disabled, + PinnedMemoryDisabled, {{MemoryType::DEVICE, mem_available}} ); EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 10_KiB); diff --git a/cpp/tests/test_statistics.cpp b/cpp/tests/test_statistics.cpp index 065db3be6..06cd1d3ab 100644 --- a/cpp/tests/test_statistics.cpp +++ b/cpp/tests/test_statistics.cpp @@ -180,10 +180,11 @@ TEST_F(StatisticsTest, ReportSorting) { } TEST_F(StatisticsTest, MemoryProfiler) { - auto br = - rapidsmpf::BufferResource::create(rmm::mr::get_current_device_resource_ref()); + auto br = rapidsmpf::BufferResource::create( + rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{} + ); auto mr = br->device_mr_adaptor(); - auto pinned_mr = rapidsmpf::PinnedMemoryResource::make_if_available(); + auto pinned_mr = br->try_pinned_mr(); auto stats = rapidsmpf::Statistics::create(); auto stream = rmm::cuda_stream_view{}; @@ -203,7 +204,7 @@ TEST_F(StatisticsTest, MemoryProfiler) { } // pinned host memory - if (pinned_mr != PinnedMemoryResource::Disabled) { + if (pinned_mr.has_value()) { void* ptr3 = pinned_mr->allocate(stream, 1_MiB); // +1 MiB void* ptr4 = pinned_mr->allocate(stream, 2_MiB); // +2 MiB pinned_mr->deallocate(stream, ptr3, 1_MiB); // -1 MiB @@ -269,7 +270,7 @@ TEST_F(StatisticsTest, MemoryProfiler) { " main (all allocations using RmmResourceAdaptor)"; EXPECT_EQ(main_line, kExpectedMainLine); static const std::string_view kExpectedPinnedLine = - pinned_mr == PinnedMemoryResource::Disabled + !pinned_mr.has_value() ? "" : " 1 3 MiB 3 MiB 4 MiB 2 MiB" " main (all allocations using PinnedMemoryResource)"; diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd index b735bdcbe..755045fdb 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd @@ -18,7 +18,8 @@ from rapidsmpf.config cimport Options, cpp_Options from rapidsmpf.memory.buffer cimport MemoryType from rapidsmpf.memory.memory_reservation cimport cpp_MemoryReservation from rapidsmpf.memory.pinned_memory_resource cimport (PinnedMemoryResource, - cpp_PinnedMemoryResource) + cpp_PinnedMemoryResource, + cpp_PinnedPoolProperties) from rapidsmpf.memory.spill_manager cimport SpillManager, cpp_SpillManager from rapidsmpf.rmm_resource_adaptor cimport (RmmResourceAdaptor, cpp_RmmResourceAdaptor) @@ -36,7 +37,7 @@ cdef extern from "" nogil: @staticmethod shared_ptr[cpp_BufferResource] create( any_resource[device_accessible], - optional[cpp_PinnedMemoryResource], + optional[cpp_PinnedPoolProperties], unordered_map[MemoryType, int64_t], optional[cpp_Duration], shared_ptr[cuda_stream_pool], @@ -51,6 +52,7 @@ cdef extern from "" nogil: shared_ptr[cpp_Statistics] statistics() except +ex_handler device_async_resource_ref device_mr() noexcept cpp_RmmResourceAdaptor& device_mr_adaptor() noexcept + optional[cpp_PinnedMemoryResource] try_pinned_mr() except +ex_handler cdef class BufferResource: cdef object __weakref__ @@ -58,7 +60,6 @@ cdef class BufferResource: cdef readonly SpillManager spill_manager cdef cpp_BufferResource* ptr(self) cdef DeviceMemoryResource _device_mr - cdef PinnedMemoryResource _pinned_mr cdef CudaStreamPool _stream_pool cdef Statistics _statistics cpdef RmmResourceAdaptor device_mr_adaptor(self) diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi index 4f99d5fdb..12a489e5c 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi @@ -10,7 +10,10 @@ from rmm.pylibrmm.memory_resource import DeviceMemoryResource from rapidsmpf.config import Options from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.memory.memory_reservation import MemoryReservation -from rapidsmpf.memory.pinned_memory_resource import PinnedMemoryResource +from rapidsmpf.memory.pinned_memory_resource import ( + PinnedMemoryResource, + PinnedPoolProperties, +) from rapidsmpf.memory.spill_manager import SpillManager from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.statistics import Statistics @@ -20,7 +23,7 @@ class BufferResource: self, device_mr: DeviceMemoryResource, *, - pinned_mr: PinnedMemoryResource | None = None, + pinned_pool_properties: PinnedPoolProperties | None = None, memory_limits: Mapping[MemoryType, int] | None = None, periodic_spill_check: float | None = 1e-3, stream_pool: CudaStreamPool | None = None, diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index fca241b3b..d5a6e4833 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -43,8 +43,10 @@ cdef extern from *: from rapidsmpf._detail.exception_handling cimport ex_handler from rapidsmpf.memory.memory_reservation cimport MemoryReservation -from rapidsmpf.memory.pinned_memory_resource cimport (PinnedMemoryResource, - cpp_PinnedMemoryResource) +from rapidsmpf.memory.pinned_memory_resource cimport ( + PinnedMemoryResource, cpp_PinnedMemoryResource, cpp_PinnedPoolProperties, + create_pinned_pool_properties_from_cpp, + pinned_pool_properties_from_options) from rapidsmpf.rmm_resource_adaptor cimport RmmResourceAdaptor from rapidsmpf.statistics cimport Statistics @@ -124,11 +126,15 @@ cdef class BufferResource: allocations are tracked for memory-limit accounting and statistics, use ``BufferResource.device_mr`` instead of the original ``device_mr`` after construction. - pinned_mr - The pinned host memory resource used for :attr:`~.MemoryType.PINNED_HOST` - allocations. If None, pinned host allocations are disabled. In that case, - any attempt to allocate pinned memory will fail regardless of any - ``memory_limits`` entry for ``PINNED_HOST``. + pinned_pool_properties + Configuration for the pinned host memory pool used for + :attr:`~.MemoryType.PINNED_HOST` allocations, as a + :class:`~rapidsmpf.memory.pinned_memory_resource.PinnedPoolProperties`. + When ``None`` (the default), pinned host allocations are disabled and any + attempt to allocate pinned memory will fail regardless of any + ``memory_limits`` entry for ``PINNED_HOST``. When provided, the pool is + created only if pinned host memory is supported on this system (see + :func:`~rapidsmpf.memory.pinned_memory_resource.is_pinned_memory_resources_supported`). memory_limits Optional mapping from :class:`~.MemoryType` to an integer byte limit. Memory types not present in the mapping are treated as unlimited. @@ -162,7 +168,7 @@ cdef class BufferResource: self, DeviceMemoryResource device_mr not None, *, - PinnedMemoryResource pinned_mr = None, + pinned_pool_properties = None, memory_limits = None, periodic_spill_check = 1e-3, CudaStreamPool stream_pool = None, @@ -191,23 +197,33 @@ cdef class BufferResource: # checked cast requires the GIL stats_handle = (statistics)._handle - # Keep the original Python memory resources alive while the C++ - # BufferResource holds resources derived from them. These anchors are - # likely redundant: `make_any_device_resource(...)` deep-copies the - # device MR into a self-sufficient owning any_resource, and - # `cpp_PinnedMemoryResource` is copied by value into the C++ BR. Kept - # defensively for now. - # TODO: drop these once verified against pool/upstream-adaptor MRs. + # Keep the original Python device memory resource alive while the C++ + # BufferResource holds resources derived from it. This anchor is likely + # redundant: `make_any_device_resource(...)` deep-copies the device MR + # into a self-sufficient owning any_resource. Kept defensively for now. + # TODO: drop this once verified against pool/upstream-adaptor MRs. # https://github.com/rapidsai/rapidsmpf/issues/1074 self._device_mr = device_mr - self._pinned_mr = pinned_mr - cdef optional[cpp_PinnedMemoryResource] cpp_pinned_mr - if self._pinned_mr is not None: - cpp_pinned_mr = self._pinned_mr._handle + + # The pinned resource is constructed internally by the C++ + # `BufferResource` from these properties (and only when pinned host + # memory is supported on this system). A missing `pinned_pool_properties` + # (None) leaves the optional empty, disabling pinned host memory. A default + # constructed `cpp_PinnedPoolProperties` already carries the C++ default + # NUMA node, so a `numa_id` of None keeps that default. + cdef cpp_PinnedPoolProperties _props + cdef optional[cpp_PinnedPoolProperties] cpp_pinned_pool + if pinned_pool_properties is not None: + _props.initial_pool_size = pinned_pool_properties.initial_pool_size + if pinned_pool_properties.max_pool_size is not None: + _props.max_pool_size = pinned_pool_properties.max_pool_size + if pinned_pool_properties.numa_id is not None: + _props.numa_id = pinned_pool_properties.numa_id + cpp_pinned_pool = _props with nogil: self._handle = cpp_BufferResource.create( make_any_device_resource(device_mr.get_mr()), - cpp_pinned_mr, + cpp_pinned_pool, move(_mem_limits), period, stream_pool.c_obj, @@ -245,10 +261,20 @@ cdef class BufferResource: """ if statistics is None: statistics = Statistics.disabled() - cdef PinnedMemoryResource pinned_mr = PinnedMemoryResource.from_options(options) + + # Derive the pinned pool configuration from the options; an empty optional + # means pinned host memory is disabled. + cdef optional[cpp_PinnedPoolProperties] props = \ + pinned_pool_properties_from_options(options._handle) + pinned_pool_properties = None + if props.has_value(): + pinned_pool_properties = create_pinned_pool_properties_from_cpp( + props.value() + ) + return cls( device_mr=mr, - pinned_mr=pinned_mr, + pinned_pool_properties=pinned_pool_properties, memory_limits={MemoryType.DEVICE: device_limit_from_options(options)}, periodic_spill_check=periodic_spill_check_from_options(options), stream_pool=stream_pool_from_options(options), @@ -328,12 +354,20 @@ cdef class BufferResource: """ The memory resource used for pinned host memory allocations. + The returned handle holds shared ownership of this ``BufferResource``, + keeping it alive for as long as the handle (or any copy of it) lives. + Returns ------- The pinned host memory resource, or None if pinned host allocations are disabled. """ - return self._pinned_mr + cdef optional[cpp_PinnedMemoryResource] opt + with nogil: + opt = deref(self._handle).try_pinned_mr() + if not opt.has_value(): + return None + return PinnedMemoryResource._from_cpp(opt.value()) def memory_reserved(self, MemoryType mem_type): """ diff --git a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd index 08fa00591..6a281e758 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd @@ -1,17 +1,34 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 +from libc.stddef cimport size_t from libcpp cimport bool as bool_t from libcpp.optional cimport optional from rapidsmpf._detail.exception_handling cimport ex_handler +from rapidsmpf.config cimport cpp_Options cdef extern from "" nogil: cdef cppclass cpp_PinnedMemoryResource"rapidsmpf::PinnedMemoryResource": pass + cdef cppclass cpp_PinnedPoolProperties"rapidsmpf::PinnedPoolProperties": + size_t initial_pool_size + optional[size_t] max_pool_size + int numa_id + + optional[cpp_PinnedPoolProperties] pinned_pool_properties_from_options \ + "rapidsmpf::pinned_pool_properties_from_options"( + cpp_Options options + ) except +ex_handler + cpdef bool_t is_pinned_memory_resources_supported() +cdef object create_pinned_pool_properties_from_cpp(cpp_PinnedPoolProperties props) + cdef class PinnedMemoryResource: cdef optional[cpp_PinnedMemoryResource] _handle + + @staticmethod + cdef PinnedMemoryResource _from_cpp(const cpp_PinnedMemoryResource& src) diff --git a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi index 97d4690e3..c35cd187a 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi @@ -1,19 +1,15 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 -from typing import Self - -from rapidsmpf.config import Options +from dataclasses import dataclass def is_pinned_memory_resources_supported() -> bool: ... +@dataclass +class PinnedPoolProperties: + initial_pool_size: int = 0 + max_pool_size: int | None = None + numa_id: int | None = None class PinnedMemoryResource: - def __init__(self, numa_id: int | None = None): ... @property def enabled(self) -> bool: ... - @staticmethod - def make_if_available( - numa_id: int | None = None, - ) -> PinnedMemoryResource | None: ... - @classmethod - def from_options(cls: type[Self], options: Options) -> Self | None: ... diff --git a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx index 3386ff2f6..cdebab5cd 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx @@ -1,27 +1,39 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass + from libcpp.optional cimport optional from rapidsmpf._detail.exception_handling cimport ex_handler -from rapidsmpf.config cimport Options, cpp_Options - -from rapidsmpf.utils.system_info import get_current_numa_node cdef extern from "" nogil: cdef bool_t cpp_is_pinned_memory_resources_supported \ "rapidsmpf::is_pinned_memory_resources_supported"(...) except +ex_handler - cdef optional[cpp_PinnedMemoryResource] cpp_make_if_available \ - "rapidsmpf::PinnedMemoryResource::make_if_available"( - int numa_id - ) except +ex_handler - cdef optional[cpp_PinnedMemoryResource] cpp_from_options \ - "rapidsmpf::PinnedMemoryResource::from_options"( - cpp_Options options - ) except +ex_handler +cdef extern from *: + """ + #include + + #include + + namespace { + // Copy a back-referenced `PinnedMemoryResource` into an optional. The copy + // promotes the back-reference, so the result keeps the owning + // `BufferResource` alive. Throws `std::bad_weak_ptr` if `src` carries no + // back-reference (which never happens for instances obtained from a + // `BufferResource`). + std::optional + cpp_copy_pinned_mr(rapidsmpf::PinnedMemoryResource const& src) { + return std::optional(src); + } + } // namespace + """ + optional[cpp_PinnedMemoryResource] cpp_copy_pinned_mr( + const cpp_PinnedMemoryResource& + ) except +ex_handler cpdef bool_t is_pinned_memory_resources_supported(): @@ -36,42 +48,73 @@ cpdef bool_t is_pinned_memory_resources_supported(): return ret -cdef class PinnedMemoryResource: +@dataclass +class PinnedPoolProperties: """ - Memory resource that provides pinned (page-locked) host memory using a pool. - - The resource allocates and deallocates pinned host memory asynchronously - through CUDA streams. Pinned memory enables higher bandwidth and lower - latency for device transfers compared to regular pageable host memory. + Configuration for a pinned (page-locked) host memory pool. - The pool has no maximum size. To limit its growth, pass an explicit - ``PINNED_HOST`` entry in :class:`BufferResource`'s ``memory_limits``. + Pass an instance to + :class:`~rapidsmpf.memory.buffer_resource.BufferResource` to enable pinned + host memory; passing ``None`` instead disables it. The pool is only created + when pinned host memory is supported on this system (see + :func:`is_pinned_memory_resources_supported`). - Parameters + Attributes ---------- + initial_pool_size + Initial size of the pinned host memory pool in bytes. The initial size + is important for pinned-memory performance, especially for the first + allocation. Defaults to ``0``. + max_pool_size + Maximum size of the pinned host memory pool in bytes, or ``None`` for no + limit. Defaults to ``None``. numa_id - NUMA node from which memory should be allocated. By default, the - resource uses the NUMA node of the calling thread. - - Raises - ------ - RuntimeError - If pinned host memory pools are not supported by the current CUDA - version. + NUMA node from which pinned host memory should be allocated, or ``None`` + to use the NUMA node of the calling thread. Defaults to ``None``. """ - def __init__(self, numa_id = None): - cdef optional[cpp_PinnedMemoryResource] opt - cdef int c_numa_id = get_current_numa_node() if numa_id is None \ - else numa_id - with nogil: - opt = cpp_make_if_available(c_numa_id) - if not opt.has_value(): - raise RuntimeError( - "Pinned host memory is not supported on this system. " - "CUDA v12.6 is one of the requirements, but additional platform " - "or driver constraints may apply." - ) - self._handle = opt + initial_pool_size: int = 0 + max_pool_size: object = None + numa_id: object = None + + +cdef object create_pinned_pool_properties_from_cpp(cpp_PinnedPoolProperties props): + """Build a Python ``PinnedPoolProperties`` from a C++ ``PinnedPoolProperties``.""" + cdef object max_pool_size = None + if props.max_pool_size.has_value(): + max_pool_size = props.max_pool_size.value() + return PinnedPoolProperties( + initial_pool_size=props.initial_pool_size, + max_pool_size=max_pool_size, + numa_id=props.numa_id, + ) + + +cdef class PinnedMemoryResource: + """ + Opaque handle to a pinned (page-locked) host memory resource. + + The resource provides pinned host memory using a pool, enabling higher + bandwidth and lower latency for device transfers compared to regular + pageable host memory. + + .. rubric:: Construction + + This class cannot be constructed directly. A pinned memory resource is owned + by a :class:`~rapidsmpf.memory.buffer_resource.BufferResource` (which installs + the back-reference that makes the handle copyable). Configure pinned memory on + a ``BufferResource`` and obtain the handle via + :attr:`~rapidsmpf.memory.buffer_resource.BufferResource.pinned_mr`. + + The returned handle holds shared ownership of its owning ``BufferResource``, + so it (and any copy of it) keeps the ``BufferResource`` alive. + """ + def __init__(self, *args, **kwargs): + raise TypeError( + "PinnedMemoryResource cannot be constructed directly; configure pinned " + "memory on a BufferResource (e.g. " + "`BufferResource(mr, pinned_pool=PinnedPoolProperties())`) and obtain it " + "via BufferResource.pinned_mr" + ) def __dealloc__(self): with nogil: @@ -80,50 +123,31 @@ cdef class PinnedMemoryResource: @property def enabled(self) -> bool: """ - Check if pinned memory resource is enabled. ie. if pinned memory is supported - by the system and a valid instance is created. + Whether this handle wraps a valid pinned memory resource. """ return self._handle.has_value() @staticmethod - def make_if_available(numa_id = None): + cdef PinnedMemoryResource _from_cpp(const cpp_PinnedMemoryResource& src): """ - Create a pinned memory resource if the system supports pinned memory. + Create a Python ``PinnedMemoryResource`` by copying a back-ref'd C++ resource. - Parameters - ---------- - numa_id - NUMA node to associate with the resource. Defaults to the current - NUMA node. - - Returns - ------- - A pinned memory resource when supported, otherwise None. - """ - if is_pinned_memory_resources_supported(): - return PinnedMemoryResource(numa_id) - return None - - @classmethod - def from_options(cls, Options options not None): - """ - Construct from configuration options. + The copy acquires shared ownership of the owning ``BufferResource``, + keeping it alive for the lifetime of the returned Python object. Parameters ---------- - options - Configuration options. + src + The C++ ``PinnedMemoryResource`` to copy from. Must have a + back-reference installed (i.e. it must have been obtained from a + ``BufferResource``); otherwise a ``std::bad_weak_ptr`` is raised. Returns ------- - The constructed PinnedMemoryResource instance if pinned memory is enabled - and supported by the system, otherwise ``None``. + A new Python ``PinnedMemoryResource`` wrapping the copied C++ resource. """ - cdef optional[cpp_PinnedMemoryResource] opt_handle - with nogil: - opt_handle = cpp_from_options(options._handle) - if not opt_handle.has_value(): - return None - cdef PinnedMemoryResource ret = cls.__new__(cls) - ret._handle = opt_handle + cdef PinnedMemoryResource ret = PinnedMemoryResource.__new__( + PinnedMemoryResource + ) + ret._handle = cpp_copy_pinned_mr(src) return ret diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyx b/python/rapidsmpf/rapidsmpf/statistics.pyx index 144040d37..d990aaef0 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pyx +++ b/python/rapidsmpf/rapidsmpf/statistics.pyx @@ -237,7 +237,8 @@ cdef class Statistics: "Disabled". pinned_mr When provided, a pinned memory section is included in the - report. + report. Obtain the handle from + :attr:`rapidsmpf.memory.buffer_resource.BufferResource.pinned_mr`. header Header line prepended to the report. When ``None``, the C++ default is used. diff --git a/python/rapidsmpf/rapidsmpf/tests/test_config.py b/python/rapidsmpf/rapidsmpf/tests/test_config.py index d224a3ea3..4a0ef4f5c 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_config.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_config.py @@ -21,7 +21,6 @@ stream_pool_from_options, ) from rapidsmpf.memory.pinned_memory_resource import ( - PinnedMemoryResource, is_pinned_memory_resources_supported, ) from rapidsmpf.progress_thread import ProgressThread @@ -418,15 +417,18 @@ def test_statistics_from_options(*, opts: Options, expected_enabled: bool) -> No (Options(), False), # Default case (disabled by default) ], ) -def test_pinned_memory_resource_from_options( +def test_pinned_memory_from_options( *, opts: Options, expect_enabled_if_supported: bool ) -> None: - pmr = PinnedMemoryResource.from_options(opts) + # Pinned memory is now configured through the BufferResource; the resource is + # only available via `BufferResource.pinned_mr` (not constructible directly). + br = BufferResource.from_options(rmm.mr.CudaMemoryResource(), opts) if expect_enabled_if_supported and is_pinned_memory_resources_supported(): - assert pmr is not None + assert br.pinned_mr is not None + assert br.pinned_mr.enabled else: - assert pmr is None + assert br.pinned_mr is None def test_device_limit_from_options_returns_configured_limit() -> None: From d65147ce9e91cecf987805c1767069b0c47ce744 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 1 Jul 2026 16:32:23 -0700 Subject: [PATCH 02/13] fixing docs Signed-off-by: niranda perera --- docs/source/python/api.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/docs/source/python/api.md b/docs/source/python/api.md index 488a1c7e8..2f8db7800 100644 --- a/docs/source/python/api.md +++ b/docs/source/python/api.md @@ -67,6 +67,9 @@ libraries. .. automodule:: rapidsmpf.memory.buffer_resource :members: +.. automodule:: rapidsmpf.memory.pinned_memory_resource + :members: + .. automodule:: rapidsmpf.memory.packed_data :members: From db1e73cac49491cb43402e3c56acca3f88a933d0 Mon Sep 17 00:00:00 2001 From: Niranda Perera Date: Mon, 13 Jul 2026 08:58:28 -0700 Subject: [PATCH 03/13] Apply suggestions from code review Co-authored-by: Mads R. B. Kristensen Co-authored-by: Lawrence Mitchell --- .../rapidsmpf/memory/buffer_resource.hpp | 7 ++++++- .../rapidsmpf/memory/host_memory_resource.hpp | 9 +-------- .../rapidsmpf/memory/pinned_memory_resource.hpp | 17 ++--------------- cpp/src/memory/buffer_resource.cpp | 4 ---- .../rapidsmpf/memory/pinned_memory_resource.pyx | 2 +- 5 files changed, 10 insertions(+), 29 deletions(-) diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 6d106d1f4..b74c96581 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -224,7 +224,12 @@ class BufferResource : public std::enable_shared_from_this { * `BufferResource` alive, but the underlying `HostMemoryResource` carries a * back-reference: promoting the ref to an owning `cuda::mr::any_resource` (as * RMM/CCCL containers do internally) keeps this `BufferResource` alive for the - * lifetime of that owning copy. + * + * @note Lifetime semantics are identical to `device_mr()`. See its + * `@par CCCL lifetime semantics` section for details. In brief, the returned + * `resource_ref` is non-owning. Promote it to a + * `cuda::mr::any_resource` to extend the + * `BufferResource` lifetime. */ [[nodiscard]] rmm::host_async_resource_ref host_mr() noexcept; diff --git a/cpp/include/rapidsmpf/memory/host_memory_resource.hpp b/cpp/include/rapidsmpf/memory/host_memory_resource.hpp index 80506e6f9..9038531ba 100644 --- a/cpp/include/rapidsmpf/memory/host_memory_resource.hpp +++ b/cpp/include/rapidsmpf/memory/host_memory_resource.hpp @@ -31,11 +31,6 @@ class BufferResource; * region. THP can improve device-host memory transfer performance for large * buffers. The hint is applied via `madvise(MADV_HUGEPAGE)` and may be ignored * by the kernel depending on system configuration or resource availability. - * - * @note Instances are constructed only by `BufferResource`, which installs a - * back-reference (via `BackRefMixin`) so that any copy of the - * resource keeps its owning `BufferResource` alive. There is no public - * constructor; obtain the resource through a `BufferResource`. */ class HostMemoryResource : public BackRefMixin { public: @@ -119,9 +114,7 @@ class HostMemoryResource : public BackRefMixin { * @brief Compares this resource to another resource. * * All instances are stateless and interchangeable, so this always returns - * true. Instances can only be created by `BufferResource` (which installs - * the back-reference exactly once), so the back-reference does not affect - * equality. + * true. * * @param other The resource to compare with. * @return true diff --git a/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp b/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp index 63bd6e36d..20d282bbd 100644 --- a/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp +++ b/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp @@ -72,9 +72,7 @@ inline bool is_pinned_memory_resources_supported() { * @brief Sentinel used to disable pinned host memory. * * Pass this in place of a `PinnedPoolProperties` (e.g. to `BufferResource::create()`) - * to disable pinned host memory allocations. It is simply `std::nullopt` typed as - * `std::nullopt_t`, so it implicitly converts to an empty - * `std::optional`. + * to disable pinned host memory allocations. */ constexpr std::nullopt_t PinnedMemoryDisabled = std::nullopt; @@ -126,12 +124,6 @@ std::optional pinned_pool_properties_from_options( * This resource allocates and deallocates pinned host memory asynchronously through * CUDA streams. It offers higher bandwidth and lower latency for device transfers * compared to regular pageable host memory. - * - * @note Instances are constructed only by `BufferResource`, which installs a - * back-reference (via `BackRefMixin`) so that any copy of the - * resource keeps its owning `BufferResource` alive. There is no public - * constructor; obtain the resource through a `BufferResource`. Use the free - * function `is_pinned_memory_resources_supported()` to query availability. */ class PinnedMemoryResource final : public cuda::mr::shared_resource< @@ -181,10 +173,6 @@ class PinnedMemoryResource final /** * @brief Equality comparison. * - * Instances can only be created by `BufferResource` (which installs the - * back-reference exactly once), so equality is determined solely by the - * underlying shared state. - * * @param other The other resource to compare. * @return True if the two resources share the same underlying shared state. */ @@ -230,8 +218,7 @@ class PinnedMemoryResource final /** * @brief Construct a pinned (page-locked) host memory resource. * - * Private — only `BufferResource` constructs instances (it installs the - * back-reference immediately after construction). + * Private: use `BufferResource` to construct instances. * * @param pool_properties Properties for configuring the pinned memory pool, * including the NUMA node from which memory should be allocated. diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index 5500f07e0..0e4faabac 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -56,10 +56,6 @@ std::shared_ptr BufferResource::create( std::shared_ptr stream_pool, std::shared_ptr statistics ) { - // Construct the pinned resource in place (`BufferResource` is a friend of - // `PinnedMemoryResource`), so no observable, back-reference-less instance is - // ever minted. Only create it when pinned memory is both requested and - // supported on this system. std::optional pinned_mr; if (pinned_pool_properties.has_value() && is_pinned_memory_resources_supported()) { pinned_mr = PinnedMemoryResource{*pinned_pool_properties}; diff --git a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx index cdebab5cd..1c0750756 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx @@ -112,7 +112,7 @@ cdef class PinnedMemoryResource: raise TypeError( "PinnedMemoryResource cannot be constructed directly; configure pinned " "memory on a BufferResource (e.g. " - "`BufferResource(mr, pinned_pool=PinnedPoolProperties())`) and obtain it " + "`BufferResource(mr, pinned_pool_properties=PinnedPoolProperties())`) and obtain it " "via BufferResource.pinned_mr" ) From d6e690d66cacc0bdbfdaffda7b99f1a38987122e Mon Sep 17 00:00:00 2001 From: niranda perera Date: Mon, 13 Jul 2026 11:38:56 -0700 Subject: [PATCH 04/13] addressing comments Signed-off-by: niranda perera --- cpp/benchmarks/bench_memory_resources.cpp | 5 +- .../rapidsmpf/memory/buffer_resource.hpp | 24 +++------ .../rapidsmpf/memory/resource_types.hpp | 3 ++ cpp/src/memory/pinned_memory_resource.cpp | 7 +-- cpp/tests/test_memory_resources.cpp | 4 -- .../rapidsmpf/memory/buffer_resource.pyx | 2 +- .../memory/pinned_memory_resource.pxd | 4 +- .../memory/pinned_memory_resource.pyx | 50 +++++++++++-------- 8 files changed, 49 insertions(+), 50 deletions(-) diff --git a/cpp/benchmarks/bench_memory_resources.cpp b/cpp/benchmarks/bench_memory_resources.cpp index 8ffe745a1..24ceef3ec 100644 --- a/cpp/benchmarks/bench_memory_resources.cpp +++ b/cpp/benchmarks/bench_memory_resources.cpp @@ -99,10 +99,7 @@ class NewDelete { friend void get_property(NewDelete const&, cuda::mr::host_accessible) noexcept {} }; -// Build a lightweight `BufferResource` that owns a pinned pool with the given -// properties. Host and pinned memory resources are constructible only by a -// `BufferResource`. The dedicated spill-check thread is disabled and a -// single-stream pool is used to keep overhead minimal for micro-benchmarks. +// Build a buffer resource with a pinned pool with the given properties. std::shared_ptr make_pinned_buffer_resource( rapidsmpf::PinnedPoolProperties props ) { diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index b74c96581..1852df7a0 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -220,15 +220,9 @@ class BufferResource : public std::enable_shared_from_this { * * @return Reference to the RMM resource used for host allocations. * - * @note The returned non-owning `resource_ref` does not by itself keep this - * `BufferResource` alive, but the underlying `HostMemoryResource` carries a - * back-reference: promoting the ref to an owning `cuda::mr::any_resource` (as - * RMM/CCCL containers do internally) keeps this `BufferResource` alive for the - * * @note Lifetime semantics are identical to `device_mr()`. See its * `@par CCCL lifetime semantics` section for details. In brief, the returned - * `resource_ref` is non-owning. Promote it to a - * `cuda::mr::any_resource` to extend the + * `resource_ref` is non-owning. Promote it to a `any_host_resource` to extend the * `BufferResource` lifetime. */ [[nodiscard]] rmm::host_async_resource_ref host_mr() noexcept; @@ -239,21 +233,19 @@ class BufferResource : public std::enable_shared_from_this { * @throws std::invalid_argument if no pinned memory resource is available. * @return Reference to the RMM resource used for pinned host allocations. * - * @note The returned non-owning `resource_ref` does not by itself keep this - * `BufferResource` alive, but the underlying `PinnedMemoryResource` carries a - * back-reference: promoting the ref to an owning `cuda::mr::any_resource` (as - * RMM/CCCL containers do internally) keeps this `BufferResource` alive for the - * lifetime of that owning copy. + * @note Lifetime semantics are identical to `device_mr()`. See its + * `@par CCCL lifetime semantics` section for details. In brief, the returned + * `resource_ref` is non-owning. Promote it to a `any_host_device_resource` to extend + * the `BufferResource` lifetime. */ [[nodiscard]] rmm::host_device_async_resource_ref pinned_mr(); /** * @brief Get the pinned host memory resource if available. * - * @return A back-referenced copy of the pinned host memory resource, or - * `std::nullopt` if pinned host memory is not available. The returned handle - * keeps this `BufferResource` alive for as long as the handle (or any copy of - * it) lives. + * @return The `PinnedMemoryResource` is available, or `std::nullopt` if pinned host + * memory is not available. The returned handle keeps this `BufferResource` alive as + * long as the handle (or any copy) exists. */ [[nodiscard]] std::optional try_pinned_mr() const; diff --git a/cpp/include/rapidsmpf/memory/resource_types.hpp b/cpp/include/rapidsmpf/memory/resource_types.hpp index 85c4a7911..634dcdfd8 100644 --- a/cpp/include/rapidsmpf/memory/resource_types.hpp +++ b/cpp/include/rapidsmpf/memory/resource_types.hpp @@ -16,6 +16,9 @@ using any_device_resource = cuda::mr::any_resource; using any_host_device_resource = cuda::mr::any_resource; +/// @brief Owning type-erased host memory resource. +using any_host_resource = cuda::mr::any_resource; + /** * @brief Check whether a type-erased memory resource is host-accessible. * diff --git a/cpp/src/memory/pinned_memory_resource.cpp b/cpp/src/memory/pinned_memory_resource.cpp index 45d254ef8..4474c23e8 100644 --- a/cpp/src/memory/pinned_memory_resource.cpp +++ b/cpp/src/memory/pinned_memory_resource.cpp @@ -48,9 +48,10 @@ PinnedMemoryResource::PinnedMemoryResource(PinnedPoolProperties pool_properties) "Pinned host memory is not supported on this system. " "CUDA " RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION_STR " is one of the requirements, but additional platform or driver " - "constraints may apply. If needed, disable pinned host memory (pass " - "`std::nullopt` for the pinned pool to `BufferResource`), noting that " - "this may significantly degrade spilling performance.", + "constraints may apply. If needed, disable pinned host memory by passing " + "`PinnedMemoryDisabled/ std::nullopt` for the `BufferResource` " + "`pinned_pool_properties`, noting that this may significantly degrade " + "spilling performance.", std::invalid_argument ); return cuda::mr::make_shared_resource< diff --git a/cpp/tests/test_memory_resources.cpp b/cpp/tests/test_memory_resources.cpp index 0cf371038..31a354e87 100644 --- a/cpp/tests/test_memory_resources.cpp +++ b/cpp/tests/test_memory_resources.cpp @@ -20,10 +20,6 @@ namespace { -// `HostMemoryResource` and `PinnedMemoryResource` are constructible only by a -// `BufferResource`. The `any_resource` copies returned below each carry a -// back-reference that keeps the (otherwise local) `BufferResource` alive, so no -// external anchor is needed. std::vector> make_host_resources() { auto br = rapidsmpf::BufferResource::create( rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{} diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index eb0f4f393..ccaf49b89 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -370,7 +370,7 @@ cdef class BufferResource: opt = deref(self._handle).try_pinned_mr() if not opt.has_value(): return None - return PinnedMemoryResource._from_cpp(opt.value()) + return PinnedMemoryResource.from_handle(opt) def memory_reserved(self, MemoryType mem_type): """ diff --git a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd index 6a281e758..5af8be5e0 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd @@ -31,4 +31,6 @@ cdef class PinnedMemoryResource: cdef optional[cpp_PinnedMemoryResource] _handle @staticmethod - cdef PinnedMemoryResource _from_cpp(const cpp_PinnedMemoryResource& src) + cdef PinnedMemoryResource from_handle( + const optional[cpp_PinnedMemoryResource]& handle + ) diff --git a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx index 1c0750756..307b43b6a 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx @@ -20,20 +20,20 @@ cdef extern from *: #include namespace { - // Copy a back-referenced `PinnedMemoryResource` into an optional. The copy - // promotes the back-reference, so the result keeps the owning - // `BufferResource` alive. Throws `std::bad_weak_ptr` if `src` carries no - // back-reference (which never happens for instances obtained from a - // `BufferResource`). + // Copy an optional back-referenced `PinnedMemoryResource`. When the source + // holds a value, the copy promotes its back-reference, so the result keeps + // the owning `BufferResource` alive; an empty source yields an empty + // result. Throws `std::bad_weak_ptr` if the contained resource carries no + // back-reference. std::optional - cpp_copy_pinned_mr(rapidsmpf::PinnedMemoryResource const& src) { - return std::optional(src); + cpp_copy_pinned_mr(std::optional const& src) { + return src; } } // namespace """ optional[cpp_PinnedMemoryResource] cpp_copy_pinned_mr( - const cpp_PinnedMemoryResource& - ) except +ex_handler + const optional[cpp_PinnedMemoryResource]& + ) except +ex_handler nogil cpdef bool_t is_pinned_memory_resources_supported(): @@ -112,8 +112,8 @@ cdef class PinnedMemoryResource: raise TypeError( "PinnedMemoryResource cannot be constructed directly; configure pinned " "memory on a BufferResource (e.g. " - "`BufferResource(mr, pinned_pool_properties=PinnedPoolProperties())`) and obtain it " - "via BufferResource.pinned_mr" + "`BufferResource(mr, pinned_pool_properties=PinnedPoolProperties())`) and " + "obtain it via BufferResource.pinned_mr" ) def __dealloc__(self): @@ -128,26 +128,34 @@ cdef class PinnedMemoryResource: return self._handle.has_value() @staticmethod - cdef PinnedMemoryResource _from_cpp(const cpp_PinnedMemoryResource& src): + cdef PinnedMemoryResource from_handle( + const optional[cpp_PinnedMemoryResource]& handle + ): """ - Create a Python ``PinnedMemoryResource`` by copying a back-ref'd C++ resource. + Create a Python ``PinnedMemoryResource`` by copying a back-ref'd C++ handle. - The copy acquires shared ownership of the owning ``BufferResource``, - keeping it alive for the lifetime of the returned Python object. + When ``handle`` holds a value, the copy acquires shared ownership of the + owning ``BufferResource``, keeping it alive for the lifetime of the + returned Python object. An empty ``handle`` produces a disabled resource. Parameters ---------- - src - The C++ ``PinnedMemoryResource`` to copy from. Must have a - back-reference installed (i.e. it must have been obtained from a - ``BufferResource``); otherwise a ``std::bad_weak_ptr`` is raised. + handle + The optional C++ ``PinnedMemoryResource`` to copy from. When it holds + a value, that value must have a back-reference installed (i.e. it must + have been obtained from a ``BufferResource``); otherwise a + ``std::bad_weak_ptr`` is raised. Returns ------- - A new Python ``PinnedMemoryResource`` wrapping the copied C++ resource. + A new Python ``PinnedMemoryResource`` wrapping the copied C++ handle. """ cdef PinnedMemoryResource ret = PinnedMemoryResource.__new__( PinnedMemoryResource ) - ret._handle = cpp_copy_pinned_mr(src) + # The copy promotes the contained resource's back-reference, keeping the + # owning BufferResource alive. Done via an extern helper so the + # std::bad_weak_ptr is translated into a Python exception. + with nogil: + ret._handle = cpp_copy_pinned_mr(handle) return ret From 08711309af5a6b84e00ab1c5ce7cfa89532b539b Mon Sep 17 00:00:00 2001 From: niranda perera Date: Mon, 13 Jul 2026 13:07:22 -0700 Subject: [PATCH 05/13] remove nullopt Signed-off-by: niranda perera --- .../rapidsmpf/memory/pinned_memory_resource.hpp | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp b/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp index 20d282bbd..90a9c9b36 100644 --- a/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp +++ b/cpp/include/rapidsmpf/memory/pinned_memory_resource.hpp @@ -68,14 +68,6 @@ inline bool is_pinned_memory_resources_supported() { return supported; } -/** - * @brief Sentinel used to disable pinned host memory. - * - * Pass this in place of a `PinnedPoolProperties` (e.g. to `BufferResource::create()`) - * to disable pinned host memory allocations. - */ -constexpr std::nullopt_t PinnedMemoryDisabled = std::nullopt; - /** * @brief Properties for configuring a pinned memory pool. */ @@ -93,6 +85,14 @@ struct PinnedPoolProperties { int numa_id = get_current_numa_node(); }; +/** + * @brief Sentinel used to disable pinned host memory. + * + * Pass this in place of a `PinnedPoolProperties` (e.g. to `BufferResource::create()`) + * to disable pinned host memory allocations. + */ +inline constexpr std::optional PinnedMemoryDisabled{}; + /** * @brief Parse pinned memory pool properties from configuration options. * From 41b902225b6007c7afa7fcd02855c3c2f8b575ca Mon Sep 17 00:00:00 2001 From: niranda perera Date: Mon, 13 Jul 2026 13:36:31 -0700 Subject: [PATCH 06/13] precommit Signed-off-by: niranda perera --- python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd | 2 +- python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx | 2 +- python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd | 2 +- python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx | 2 +- python/rapidsmpf/rapidsmpf/statistics.pyx | 2 +- python/rapidsmpf/rapidsmpf/tests/test_config.py | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd index 755045fdb..a9d54969a 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libc.stddef cimport size_t diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index ccaf49b89..4b8d8f71e 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from cython cimport no_gc_clear diff --git a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd index 5af8be5e0..b9087007f 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from libc.stddef cimport size_t diff --git a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx index 307b43b6a..e814a70e9 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyx b/python/rapidsmpf/rapidsmpf/statistics.pyx index 2eefd6bba..675cdf277 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pyx +++ b/python/rapidsmpf/rapidsmpf/statistics.pyx @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from cpython.bytes cimport PyBytes_FromStringAndSize diff --git a/python/rapidsmpf/rapidsmpf/tests/test_config.py b/python/rapidsmpf/rapidsmpf/tests/test_config.py index 4a0ef4f5c..c97aa3081 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_config.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_config.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations From aced22fe6095b08c4ee4067d9421c643b957a228 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Mon, 13 Jul 2026 14:21:08 -0700 Subject: [PATCH 07/13] increasing cython lint line len Signed-off-by: niranda perera --- pyproject.toml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 8d22591b5..9b17ca2ea 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # See file LICENSE for terms. @@ -157,3 +157,6 @@ rapids = ["rmm"] [tool.ruff.format] docstring-code-format = true + +[tool.cython-lint] +max-line-length = 120 From 2546b0576dc574710809278196074f5c577b821b Mon Sep 17 00:00:00 2001 From: niranda perera Date: Tue, 14 Jul 2026 10:04:41 -0700 Subject: [PATCH 08/13] make br create throw if pinned mr not supported Signed-off-by: niranda perera --- cpp/benchmarks/bench_memory_resources.cpp | 13 ++++++------- .../rapidsmpf/memory/buffer_resource.hpp | 7 +++++-- cpp/src/memory/buffer_resource.cpp | 11 ++++++++++- cpp/tests/test_memory_resources.cpp | 11 +++++++++-- .../rapidsmpf/memory/buffer_resource.pyx | 17 +++++++++-------- 5 files changed, 39 insertions(+), 20 deletions(-) diff --git a/cpp/benchmarks/bench_memory_resources.cpp b/cpp/benchmarks/bench_memory_resources.cpp index 24ceef3ec..32b992b21 100644 --- a/cpp/benchmarks/bench_memory_resources.cpp +++ b/cpp/benchmarks/bench_memory_resources.cpp @@ -5,6 +5,7 @@ #include #include +#include #include @@ -99,9 +100,9 @@ class NewDelete { friend void get_property(NewDelete const&, cuda::mr::host_accessible) noexcept {} }; -// Build a buffer resource with a pinned pool with the given properties. +// Build a buffer resource with the given (optional) pinned pool properties. std::shared_ptr make_pinned_buffer_resource( - rapidsmpf::PinnedPoolProperties props + std::optional props ) { return rapidsmpf::BufferResource::create( rmm::mr::get_current_device_resource_ref(), @@ -121,10 +122,8 @@ cuda::mr::any_resource create_host_memory_resource( return NewDelete{}; case ResourceType::HOST_MEMORY_RESOURCE: { - // The returned `any_resource` carries a back-reference that keeps the - // (local) `BufferResource` alive, so no external anchor is needed. - auto br = make_pinned_buffer_resource(rapidsmpf::PinnedPoolProperties{}); - return br->host_mr(); + auto br = make_pinned_buffer_resource(rapidsmpf::PinnedMemoryDisabled); + return br->host_mr(); // br is kept alive by the back-reference } case ResourceType::PINNED_MEMORY_RESOURCE: { @@ -135,7 +134,7 @@ cuda::mr::any_resource create_host_memory_resource( "pinned memory is not supported on this system", std::runtime_error ); - return *mr; + return *mr; // br is kept alive by the back-reference } default: RAPIDSMPF_FAIL("Unknown memory resource type"); diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 1852df7a0..2b48cede8 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -94,8 +94,9 @@ class BufferResource : public std::enable_shared_from_this { * @param pinned_pool_properties Configuration for the pinned host memory pool * used for `MemoryType::PINNED_HOST` allocations, or `PinnedMemoryDisabled` to * disable pinned allocations. The pinned resource is constructed internally and - * owned by the `BufferResource`; it is only created when pinned host memory is - * supported on the system (see `is_pinned_memory_resources_supported()`). + * owned by the `BufferResource`. When a value is provided, pinned host memory + * must be supported on the system (see `is_pinned_memory_resources_supported()`); + * otherwise a `std::runtime_error` is thrown. * @param memory_limits Maximum allocation limits in bytes per `MemoryType`. Missing * entries are treated as unlimited. * @param periodic_spill_check Interval between periodic spill checks. `std::nullopt` @@ -104,6 +105,8 @@ class BufferResource : public std::enable_shared_from_this { * explicit CUDA stream. * @param statistics Statistics instance used for runtime metrics. * @return A newly constructed `BufferResource` owned by `std::shared_ptr`. + * @throws std::runtime_error if `pinned_pool_properties` has a value but pinned + * host memory is not supported on this system. */ [[nodiscard]] static std::shared_ptr create( cuda::mr::any_resource device_mr, diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index 0e4faabac..968181fe3 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -57,7 +57,16 @@ std::shared_ptr BufferResource::create( std::shared_ptr statistics ) { std::optional pinned_mr; - if (pinned_pool_properties.has_value() && is_pinned_memory_resources_supported()) { + if (pinned_pool_properties.has_value()) { + RAPIDSMPF_EXPECTS( + is_pinned_memory_resources_supported(), + "pinned host memory was requested (via `PinnedPoolProperties`) but is not " + "supported on this system. " + "CUDA " RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION_STR + " is one of the requirements, but additional platform or driver constraints " + "may apply. Pass `PinnedMemoryDisabled` to disable pinned host memory.", + std::runtime_error + ); pinned_mr = PinnedMemoryResource{*pinned_pool_properties}; } diff --git a/cpp/tests/test_memory_resources.cpp b/cpp/tests/test_memory_resources.cpp index 31a354e87..0e8d08131 100644 --- a/cpp/tests/test_memory_resources.cpp +++ b/cpp/tests/test_memory_resources.cpp @@ -3,6 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include #include #include @@ -21,8 +22,11 @@ namespace { std::vector> make_host_resources() { + auto pinned_pool_properties = rapidsmpf::is_pinned_memory_resources_supported() + ? rapidsmpf::PinnedPoolProperties{} + : rapidsmpf::PinnedMemoryDisabled; auto br = rapidsmpf::BufferResource::create( - rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{} + rmm::mr::get_current_device_resource_ref(), std::move(pinned_pool_properties) ); std::vector> resources; resources.emplace_back(br->host_mr()); @@ -33,8 +37,11 @@ std::vector> make_host_resourc } std::vector> make_device_resources() { + auto pinned_pool_properties = rapidsmpf::is_pinned_memory_resources_supported() + ? rapidsmpf::PinnedPoolProperties{} + : rapidsmpf::PinnedMemoryDisabled; auto br = rapidsmpf::BufferResource::create( - rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{} + rmm::mr::get_current_device_resource_ref(), std::move(pinned_pool_properties) ); std::vector> resources; resources.emplace_back(rmm::mr::cuda_memory_resource{}); diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index 4b8d8f71e..19b13d0ed 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -133,9 +133,10 @@ cdef class BufferResource: :class:`~rapidsmpf.memory.pinned_memory_resource.PinnedPoolProperties`. When ``None`` (the default), pinned host allocations are disabled and any attempt to allocate pinned memory will fail regardless of any - ``memory_limits`` entry for ``PINNED_HOST``. When provided, the pool is - created only if pinned host memory is supported on this system (see - :func:`~rapidsmpf.memory.pinned_memory_resource.is_pinned_memory_resources_supported`). + ``memory_limits`` entry for ``PINNED_HOST``. When provided, pinned host + memory must be supported on this system (see + :func:`~rapidsmpf.memory.pinned_memory_resource.is_pinned_memory_resources_supported`); + otherwise a ``RuntimeError`` is raised. memory_limits Optional mapping from :class:`~.MemoryType` to an integer byte limit. Memory types not present in the mapping are treated as unlimited. @@ -209,11 +210,11 @@ cdef class BufferResource: self._device_mr = device_mr # The pinned resource is constructed internally by the C++ - # `BufferResource` from these properties (and only when pinned host - # memory is supported on this system). A missing `pinned_pool_properties` - # (None) leaves the optional empty, disabling pinned host memory. A default - # constructed `cpp_PinnedPoolProperties` already carries the C++ default - # NUMA node, so a `numa_id` of None keeps that default. + # `BufferResource` from these properties. A None `pinned_pool_properties` + # leaves the optional empty, disabling pinned host memory. Providing + # properties on a system without pinned-host-memory support raises a + # RuntimeError. A default constructed `cpp_PinnedPoolProperties` already + # carries the C++ default NUMA node, so a `numa_id` of None keeps that default. cdef cpp_PinnedPoolProperties _props cdef optional[cpp_PinnedPoolProperties] cpp_pinned_pool if pinned_pool_properties is not None: From b21966abe9056c53f1a58f5da69e3bfa951ff9a6 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Tue, 14 Jul 2026 10:59:33 -0700 Subject: [PATCH 09/13] test failures Signed-off-by: niranda perera --- cpp/tests/test_buffer.cpp | 6 +++++- cpp/tests/test_buffer_resource.cpp | 2 +- cpp/tests/test_statistics.cpp | 5 ++++- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/cpp/tests/test_buffer.cpp b/cpp/tests/test_buffer.cpp index dbe749ae3..2a8ea3f8e 100644 --- a/cpp/tests/test_buffer.cpp +++ b/cpp/tests/test_buffer.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -53,9 +54,12 @@ class BufferRebindStreamTest : public ::testing::TestWithParam { GTEST_SKIP() << "Pinned memory resources are not supported on this system"; } + auto pinned_pool_properties = is_pinned_memory_resources_supported() + ? PinnedPoolProperties{} + : PinnedMemoryDisabled; br = BufferResource::create( rmm::mr::get_current_device_resource_ref(), - PinnedPoolProperties{}, + std::move(pinned_pool_properties), std::unordered_map{}, std::nullopt, stream_pool diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index 6b4ad8bba..ef3091cb0 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -292,7 +292,7 @@ TEST(BufferResource, AllocStatistics) { bool const pinned_available = is_pinned_memory_resources_supported(); auto br = BufferResource::create( mr_cuda, - PinnedPoolProperties{}, + pinned_available ? PinnedPoolProperties{} : PinnedMemoryDisabled, {}, std::nullopt, std::make_shared(1, rmm::cuda_stream::flags::non_blocking), diff --git a/cpp/tests/test_statistics.cpp b/cpp/tests/test_statistics.cpp index 06cd1d3ab..7bf01eba7 100644 --- a/cpp/tests/test_statistics.cpp +++ b/cpp/tests/test_statistics.cpp @@ -180,8 +180,11 @@ TEST_F(StatisticsTest, ReportSorting) { } TEST_F(StatisticsTest, MemoryProfiler) { + auto pinned_pool_properties = rapidsmpf::is_pinned_memory_resources_supported() + ? rapidsmpf::PinnedPoolProperties{} + : rapidsmpf::PinnedMemoryDisabled; auto br = rapidsmpf::BufferResource::create( - rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedPoolProperties{} + rmm::mr::get_current_device_resource_ref(), pinned_pool_properties ); auto mr = br->device_mr_adaptor(); auto pinned_mr = br->try_pinned_mr(); From 8e42edcb39d908d6e4623dbf5dceed2c288f7a13 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Tue, 14 Jul 2026 11:40:22 -0700 Subject: [PATCH 10/13] fix test Signed-off-by: niranda perera --- python/rapidsmpf/rapidsmpf/tests/test_config.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/rapidsmpf/rapidsmpf/tests/test_config.py b/python/rapidsmpf/rapidsmpf/tests/test_config.py index c97aa3081..b1888a81a 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_config.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_config.py @@ -420,11 +420,16 @@ def test_statistics_from_options(*, opts: Options, expected_enabled: bool) -> No def test_pinned_memory_from_options( *, opts: Options, expect_enabled_if_supported: bool ) -> None: + # Requesting pinned memory on a system that doesn't support it now raises, so + # skip the case that would enable it. + if expect_enabled_if_supported and not is_pinned_memory_resources_supported(): + pytest.skip("Pinned memory not supported on this system") + # Pinned memory is now configured through the BufferResource; the resource is # only available via `BufferResource.pinned_mr` (not constructible directly). br = BufferResource.from_options(rmm.mr.CudaMemoryResource(), opts) - if expect_enabled_if_supported and is_pinned_memory_resources_supported(): + if expect_enabled_if_supported: assert br.pinned_mr is not None assert br.pinned_mr.enabled else: From 200975006f5bd256fe8f3b1ad295c1cd4ecfcc10 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Tue, 14 Jul 2026 11:52:55 -0700 Subject: [PATCH 11/13] trigger build From 0802d24398013352e20151675e7afcb66a8873ea Mon Sep 17 00:00:00 2001 From: niranda perera Date: Tue, 14 Jul 2026 13:33:00 -0700 Subject: [PATCH 12/13] trigger build From 316087afb59540771906de20d4471b1382e31738 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Sat, 18 Jul 2026 09:46:15 -0700 Subject: [PATCH 13/13] precommit Signed-off-by: niranda perera --- pyproject.toml | 3 --- python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi | 4 ---- 2 files changed, 7 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7c6613790..9cfc8ae6c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -160,6 +160,3 @@ rapids = ["rmm"] [tool.ruff.format] docstring-code-format = true - -[tool.cython-lint] -max-line-length = 120 diff --git a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi index 4b3b23d9f..b3d8f71e3 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi @@ -3,12 +3,8 @@ from dataclasses import dataclass -from typing import Self - from rmm.pylibrmm.stream import Stream -from rapidsmpf.config import Options - def is_pinned_memory_resources_supported() -> bool: ... @dataclass class PinnedPoolProperties: