diff --git a/cpp/benchmarks/bench_comm.cpp b/cpp/benchmarks/bench_comm.cpp index 6df6747bb..43a872f09 100644 --- a/cpp/benchmarks/bench_comm.cpp +++ b/cpp/benchmarks/bench_comm.cpp @@ -326,7 +326,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..32b992b21 100644 --- a/cpp/benchmarks/bench_memory_resources.cpp +++ b/cpp/benchmarks/bench_memory_resources.cpp @@ -5,14 +5,18 @@ #include #include +#include #include +#include #include #include #include +#include #include +#include #include #include #include @@ -96,6 +100,19 @@ class NewDelete { friend void get_property(NewDelete const&, cuda::mr::host_accessible) noexcept {} }; +// Build a buffer resource with the given (optional) pinned pool properties. +std::shared_ptr make_pinned_buffer_resource( + std::optional 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,16 +121,20 @@ cuda::mr::any_resource create_host_memory_resource( case ResourceType::NEW_DELETE: return NewDelete{}; case ResourceType::HOST_MEMORY_RESOURCE: - return rapidsmpf::HostMemoryResource{}; + { + 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: { - 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", std::runtime_error ); - return *mr; + return *mr; // br is kept alive by the back-reference } default: RAPIDSMPF_FAIL("Unknown memory resource type"); @@ -442,9 +463,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 +519,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..2b48cede8 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -91,9 +91,12 @@ 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`. 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` @@ -102,10 +105,12 @@ 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, - 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 +222,11 @@ 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 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_resource` to extend the + * `BufferResource` lifetime. */ [[nodiscard]] rmm::host_async_resource_ref host_mr() noexcept; @@ -225,16 +235,22 @@ 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 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 The pinned host memory resource as an `any_resource`, or `std::nullopt` if - * pinned host memory is not available. + * @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 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..9038531ba 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. * @@ -29,9 +32,8 @@ namespace rapidsmpf { * buffers. The hint is applied via `madvise(MADV_HUGEPAGE)` and may be ignored * by the kernel depending on system configuration or resource availability. */ -class HostMemoryResource { +class HostMemoryResource : public BackRefMixin { public: - HostMemoryResource() = default; ~HostMemoryResource() = default; HostMemoryResource(HostMemoryResource const&) = default; ///< Copyable. @@ -138,6 +140,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..90a9c9b36 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. * @@ -75,8 +79,40 @@ 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 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. + * + * 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. * @@ -91,47 +127,12 @@ struct PinnedPoolProperties { */ 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. * @@ -217,16 +218,16 @@ 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: use `BufferResource` to construct instances. * - * @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/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/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index 5eb9e91f7..968181fe3 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -50,12 +50,26 @@ 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 ) { + std::optional pinned_mr; + 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}; + } + std::shared_ptr br{new BufferResource{ std::move(device_mr), std::move(pinned_mr), @@ -65,12 +79,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 +99,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 +120,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 +146,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 +153,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 +166,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..4474c23e8 100644 --- a/cpp/src/memory/pinned_memory_resource.cpp +++ b/cpp/src/memory/pinned_memory_resource.cpp @@ -41,57 +41,48 @@ 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 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< 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..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(), - PinnedMemoryResource::make_if_available(), + 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 d3b0fa1c0..ef3091cb0 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, + pinned_available ? PinnedPoolProperties{} : PinnedMemoryDisabled, {}, 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 79a6f7809..8a04726ae 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..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 @@ -11,7 +12,9 @@ #include #include +#include +#include #include #include #include @@ -19,20 +22,32 @@ 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(), std::move(pinned_pool_properties) + ); 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 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(), std::move(pinned_pool_properties) + ); 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 574185c94..3d95bb8d0 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 ); @@ -507,7 +507,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..7bf01eba7 100644 --- a/cpp/tests/test_statistics.cpp +++ b/cpp/tests/test_statistics.cpp @@ -180,10 +180,14 @@ TEST_F(StatisticsTest, ReportSorting) { } TEST_F(StatisticsTest, MemoryProfiler) { - auto br = - rapidsmpf::BufferResource::create(rmm::mr::get_current_device_resource_ref()); + 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(), pinned_pool_properties + ); 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 +207,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 +273,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/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: diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd index b735bdcbe..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 @@ -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 b95ad8fa3..19b13d0ed 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 @@ -44,8 +44,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 @@ -125,11 +127,16 @@ 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, 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. @@ -163,7 +170,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, @@ -201,14 +208,26 @@ cdef class BufferResource: # TODO: drop these 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. 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: + _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( any_resource[device_accessible](device_mr.get_mr()), - cpp_pinned_mr, + cpp_pinned_pool, move(_mem_limits), period, stream_pool.c_obj, @@ -246,10 +265,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), @@ -329,12 +358,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_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 2a957886d..83c67406c 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pxd @@ -1,11 +1,13 @@ -# 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 from libcpp cimport bool as bool_t from libcpp.optional cimport optional from rmm.librmm.cuda_stream_view cimport cuda_stream_view from rapidsmpf._detail.exception_handling cimport ex_handler +from rapidsmpf.config cimport cpp_Options cdef extern from "" nogil: @@ -13,7 +15,24 @@ cdef extern from "" nogil: void* allocate(cuda_stream_view, size_t) except +ex_handler void deallocate(cuda_stream_view, void*, size_t) + 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_handle( + const optional[cpp_PinnedMemoryResource]& handle + ) diff --git a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi index 0b059a730..b3d8f71e3 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyi @@ -1,23 +1,19 @@ -# 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 dataclasses import dataclass from rmm.pylibrmm.stream import Stream -from rapidsmpf.config import Options - 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: ... def allocate(self, nbytes: int, stream: Stream) -> int: ... def deallocate(self, ptr: int, nbytes: int, stream: Stream) -> 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 618a2c4ef..63b9604b1 100644 --- a/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/pinned_memory_resource.pyx @@ -1,28 +1,40 @@ -# 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 + from libcpp.optional cimport optional from rmm.pylibrmm.stream cimport Stream 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 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(std::optional const& src) { + return src; + } + } // namespace + """ + optional[cpp_PinnedMemoryResource] cpp_copy_pinned_mr( + const optional[cpp_PinnedMemoryResource]& + ) except +ex_handler nogil cpdef bool_t is_pinned_memory_resources_supported(): @@ -37,42 +49,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_properties=PinnedPoolProperties())`) and " + "obtain it via BufferResource.pinned_mr" + ) def __dealloc__(self): with nogil: @@ -81,30 +124,10 @@ 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): - """ - Create a pinned memory resource if the system supports pinned memory. - - 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 - def allocate(self, size_t nbytes, Stream stream not None) -> int: """ Allocate pinned host memory associated with a CUDA stream. @@ -141,26 +164,35 @@ cdef class PinnedMemoryResource: with nogil: self._handle.value().deallocate(stream.view(), ptr, nbytes) - @classmethod - def from_options(cls, Options options not None): + @staticmethod + cdef PinnedMemoryResource from_handle( + const optional[cpp_PinnedMemoryResource]& handle + ): """ - Construct from configuration options. + Create a Python ``PinnedMemoryResource`` by copying a back-ref'd C++ handle. + + 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 ---------- - options - Configuration options. + 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 ------- - The constructed PinnedMemoryResource instance if pinned memory is enabled - and supported by the system, otherwise ``None``. + A new Python ``PinnedMemoryResource`` wrapping the copied C++ handle. """ - cdef optional[cpp_PinnedMemoryResource] opt_handle + cdef PinnedMemoryResource ret = PinnedMemoryResource.__new__( + PinnedMemoryResource + ) + # 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: - 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 + ret._handle = cpp_copy_pinned_mr(handle) return ret diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyx b/python/rapidsmpf/rapidsmpf/statistics.pyx index 9106aa7cf..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 @@ -255,7 +255,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..b1888a81a 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 @@ -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,23 @@ 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) + # 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(): - assert pmr is not None + if expect_enabled_if_supported: + 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: