diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index ccb119f5e..dba727891 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -206,7 +206,6 @@ add_library( src/memory/spill_manager.cpp src/pausable_thread_loop.cpp src/progress_thread.cpp - src/rmm_resource_adaptor.cpp src/rrun/rrun.cpp src/shuffler/chunk.cpp src/shuffler/finish_counter.cpp diff --git a/cpp/benchmarks/bench_shuffle.cpp b/cpp/benchmarks/bench_shuffle.cpp index cd30993e3..e2e5e9202 100644 --- a/cpp/benchmarks/bench_shuffle.cpp +++ b/cpp/benchmarks/bench_shuffle.cpp @@ -536,7 +536,7 @@ int main(int argc, char** argv) { rapidsmpf::config::Options options{rapidsmpf::config::get_environment_variables()}; set_current_rmm_resource(args.rmm_mr); - rapidsmpf::RmmResourceAdaptor stat_enabled_mr = set_device_mem_resource_with_stats(); + auto stat_enabled_mr = set_device_mem_resource_with_stats(); std::unordered_map memory_limits{}; if (args.device_mem_limit_mb >= 0) { @@ -548,7 +548,7 @@ int main(int argc, char** argv) { // We're only going to measure the last run, so disable initially. stats->disable(); rapidsmpf::BufferResource br{ - stat_enabled_mr, + std::move(stat_enabled_mr), args.pinned_mem_disable ? rapidsmpf::PinnedMemoryResource::Disabled : rapidsmpf::PinnedMemoryResource::make_if_available(), std::move(memory_limits), @@ -663,7 +663,7 @@ int main(int argc, char** argv) { << " | out_parts: " << args.num_output_partitions << " | nranks: " << comm->nranks(); if (args.enable_memory_profiler) { - auto record = stat_enabled_mr.get_main_record(); + auto record = br.get_main_record(); ss << " | device memory peak: " << rapidsmpf::format_nbytes(record.peak()) << " | device memory total: " << rapidsmpf::format_nbytes( @@ -675,9 +675,7 @@ int main(int argc, char** argv) { } if (args.enable_memory_profiler) { - log->print(stats->report( - {.mr = stat_enabled_mr, .header = "Statistics (of the last run):"} - )); + log->print(stats->report({.mr = br, .header = "Statistics (of the last run):"})); } else { log->print(stats->report({.header = "Statistics (of the last run):"})); } diff --git a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp index 5836b11ff..847f96107 100644 --- a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp +++ b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp @@ -370,7 +370,7 @@ int main(int argc, char** argv) { ? rapidsmpf::PinnedMemoryResource::Disabled : rapidsmpf::PinnedMemoryResource::make_if_available(); auto br = std::make_shared( - stat_enabled_mr, + std::move(stat_enabled_mr), pinned_mr, std::move(memory_limits), std::nullopt, @@ -447,7 +447,7 @@ int main(int argc, char** argv) { << " | out_parts: " << args.num_output_partitions << " | nranks: " << comm->nranks(); if (args.enable_memory_profiler) { - auto record = stat_enabled_mr.get_main_record(); + auto record = br->get_main_record(); ss << " | device memory peak: " << rapidsmpf::format_nbytes(record.peak()) << " | device memory total: " << rapidsmpf::format_nbytes( @@ -461,7 +461,7 @@ int main(int argc, char** argv) { auto statistics = ctx->statistics(); if (args.enable_memory_profiler) { log.print(statistics->report({ - .mr = stat_enabled_mr, + .mr = *br, .pinned_mr = pinned_mr, .header = "Statistics (of the last run):", })); diff --git a/cpp/benchmarks/streaming/ndsh/q21.cpp b/cpp/benchmarks/streaming/ndsh/q21.cpp index 3b7e6a052..fbb1aff48 100644 --- a/cpp/benchmarks/streaming/ndsh/q21.cpp +++ b/cpp/benchmarks/streaming/ndsh/q21.cpp @@ -30,7 +30,6 @@ #include #include #include -#include #include #include #include diff --git a/cpp/benchmarks/streaming/ndsh/utils.cpp b/cpp/benchmarks/streaming/ndsh/utils.cpp index fb20dd8b1..644af393c 100644 --- a/cpp/benchmarks/streaming/ndsh/utils.cpp +++ b/cpp/benchmarks/streaming/ndsh/utils.cpp @@ -33,7 +33,6 @@ #include #include #include -#include #include #include diff --git a/cpp/benchmarks/utils/rmm_utils.hpp b/cpp/benchmarks/utils/rmm_utils.hpp index 754463aad..28926b0d4 100644 --- a/cpp/benchmarks/utils/rmm_utils.hpp +++ b/cpp/benchmarks/utils/rmm_utils.hpp @@ -6,6 +6,8 @@ #include +#include + #include #include #include @@ -13,7 +15,6 @@ #include #include -#include /** * @brief Create and set a RMM memory resource as the current device resource. @@ -45,13 +46,18 @@ inline void set_current_rmm_resource(std::string const& name) { } /** - * @brief Create a statistics-enabled device memory resource wrapping the current - * device resource, and set it as the current device resource. + * @brief Return the current device resource as a CCCL `any_resource`. + * + * Compatibility shim for benchmarks that previously wrapped the current device + * resource in `RmmResourceAdaptor` to gain statistics. Tracking is now part of + * `BufferResource` itself, so callers pass the returned `any_resource` directly + * to a `BufferResource` constructor. * - * @return A RmmResourceAdaptor (shared ownership) for accessing statistics. + * @return The current device resource as a type-erased CCCL resource. */ -[[nodiscard]] inline rapidsmpf::RmmResourceAdaptor set_device_mem_resource_with_stats() { - rapidsmpf::RmmResourceAdaptor adaptor{rmm::mr::get_current_device_resource_ref()}; - rmm::mr::set_current_device_resource(adaptor); - return adaptor; +[[nodiscard]] inline cuda::mr::any_resource +set_device_mem_resource_with_stats() { + return cuda::mr::any_resource{ + rmm::mr::get_current_device_resource_ref() + }; } diff --git a/cpp/include/rapidsmpf/detail/rmm_resource_adaptor_impl.hpp b/cpp/include/rapidsmpf/detail/rmm_resource_adaptor_impl.hpp index 2307a632e..73fd6ec18 100644 --- a/cpp/include/rapidsmpf/detail/rmm_resource_adaptor_impl.hpp +++ b/cpp/include/rapidsmpf/detail/rmm_resource_adaptor_impl.hpp @@ -29,11 +29,12 @@ namespace rapidsmpf::detail { /** - * @brief Implementation class for RmmResourceAdaptor. + * @brief Implementation class for instrumented RMM memory resources. * * Holds all mutable state for memory tracking. This class satisfies the CCCL - * `cuda::mr::resource` concept and is held by `RmmResourceAdaptor` via - * `cuda::mr::shared_resource` for reference-counted ownership. + * `cuda::mr::resource` concept and is the building block used internally by + * both `BufferResource` (for device memory tracking) and `PinnedMemoryResource` + * (for pinned-host tracking with in-place storage of `cuda::pinned_memory_pool`). * * @tparam PrimaryMR The type of the primary memory resource. Use a concrete * resource type (e.g. `cuda::pinned_memory_pool`) to store the resource @@ -96,25 +97,37 @@ class RmmResourceAdaptorImpl { return primary_mr_; } - /// @copydoc RmmResourceAdaptor::get_main_record + /** + * @brief Returns a copy of the main memory record (lifetime-of-resource stats). + * + * @return A copy of the main `ScopedMemoryRecord`. + */ [[nodiscard]] ScopedMemoryRecord get_main_record() const { std::lock_guard lock(mutex_); return main_record_; } - /// @copydoc RmmResourceAdaptor::current_allocated + /** + * @brief Total number of currently allocated bytes. + * + * @return Currently outstanding allocated bytes. + */ [[nodiscard]] std::int64_t current_allocated() const noexcept { std::lock_guard lock(mutex_); return main_record_.current(); } - /// @copydoc RmmResourceAdaptor::begin_scoped_memory_record + /// @brief Push a new scoped memory record onto the current thread's stack. void begin_scoped_memory_record() { std::lock_guard lock(mutex_); record_stacks_[std::this_thread::get_id()].emplace(); } - /// @copydoc RmmResourceAdaptor::end_scoped_memory_record + /** + * @brief Pop and return the topmost scoped memory record on the current thread. + * + * @return The popped `ScopedMemoryRecord`. + */ ScopedMemoryRecord end_scoped_memory_record() { std::lock_guard lock(mutex_); auto& stack = record_stacks_.at(std::this_thread::get_id()); diff --git a/cpp/include/rapidsmpf/memory/buffer.hpp b/cpp/include/rapidsmpf/memory/buffer.hpp index 44215e1d3..e29ffb394 100644 --- a/cpp/include/rapidsmpf/memory/buffer.hpp +++ b/cpp/include/rapidsmpf/memory/buffer.hpp @@ -24,6 +24,10 @@ namespace rapidsmpf { +namespace detail { +class BufferResourceImpl; +} // namespace detail + /** * @brief Buffer representing device or host memory. * @@ -46,6 +50,7 @@ namespace rapidsmpf { */ class Buffer { friend class BufferResource; + friend class detail::BufferResourceImpl; public: /// @brief Storage type for a device buffer. diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index 5b01ef1fc..7c7576a6d 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -5,11 +5,9 @@ #pragma once -#include -#include +#include #include #include -#include #include #include #include @@ -21,53 +19,55 @@ #include #include -#include +#include #include #include #include +#include #include -#include #include #include namespace rapidsmpf { /** - * @brief Policy controlling whether a memory reservation is allowed to overbook. + * @brief CCCL-compatible memory resource managing all memory operations in RapidsMPF. * - * This enum is used throughout RapidsMPF to specify the overbooking behavior of - * a memory reservation request. The exact semantics depend on the specific API - * and execution context in which it is used. - */ -enum class AllowOverbooking : bool { - NO, ///< Overbooking is not allowed. - YES, ///< Overbooking is allowed. -}; - -/** - * @brief Class managing buffer resources. + * `BufferResource` handles allocations and transfers between different memory + * types (device, host, pinned host). All memory operations in RapidsMPF — for + * example those performed by the Shuffler — flow through a `BufferResource`. + * + * `BufferResource` itself satisfies the CCCL + * `cuda::mr::resource_with` concept, so an instance can be + * passed directly anywhere an RMM-compatible device memory resource is + * expected (e.g. as the `mr` argument to `rmm::device_buffer`). Buffers + * allocated through it hold an owning ref to the resource, which transitively + * keeps the underlying stream pool alive. * - * This class handles memory allocation and transfers between different memory types - * (e.g., host and device). All memory operations in rapidsmpf, such as those performed - * by the Shuffler, rely on a buffer resource for memory management. + * The class is held by reference-counted shared ownership through + * `cuda::mr::shared_resource`; copies of a `BufferResource` are cheap and + * refer to the same underlying state. * - * @note Similar to RMM's memory resource, the `BufferResource` instance must outlive all - * allocated buffers and memory reservations. + * Memory availability is computed per `MemoryType` as `limit - allocated`. + * Device and pinned-host allocations routed through this `BufferResource` are + * tracked automatically. Host memory allocations are not tracked, so the + * available memory always equals the configured limit. If pinned-host memory + * is disabled, available pinned-host memory is always reported as zero + * regardless of the configured limit. */ -class BufferResource { +class BufferResource : public cuda::mr::shared_resource { + using shared_base = cuda::mr::shared_resource; + using any_device_resource = cuda::mr::any_resource; + public: + /// @brief Tag this resource as device-accessible for the CCCL concept. + friend void get_property( + BufferResource const&, cuda::mr::device_accessible + ) noexcept {} + /** * @brief Constructs a buffer resource. * - * Memory availability is computed per `MemoryType` as `limit - allocated`. - * - * Device and pinned-host allocations routed through this BufferResource are tracked - * automatically. Host memory allocations are not tracked, so the available memory - * always equals the configured limit. - * - * If pinned-host memory is disabled, available pinned-host memory is always reported - * as zero regardless of the configured limit. - * * @param device_mr The RMM device memory resource used for device allocations. * @param pinned_mr The pinned host memory resource used for `MemoryType::PINNED_HOST` * allocations. If disabled, pinned host allocations are unavailable regardless of @@ -83,7 +83,7 @@ class BufferResource { * @param statistics The statistics instance to use (disabled by default). */ BufferResource( - cuda::mr::any_resource device_mr, + any_device_resource device_mr, std::optional pinned_mr = PinnedMemoryResource::Disabled, std::unordered_map memory_limits = {}, std::optional periodic_spill_check = std::chrono::milliseconds{1}, @@ -96,8 +96,7 @@ class BufferResource { * @brief Construct a BufferResource from configuration options. * * This factory method creates a BufferResource using configuration options to - * initialize all components. The supplied device memory resource is wrapped in - * an internal `RmmResourceAdaptor` for allocation tracking. + * initialize all components. * * @param mr A device-accessible RMM memory resource. * @param options Configuration options. @@ -107,26 +106,120 @@ class BufferResource { * options. */ static std::shared_ptr from_options( - cuda::mr::any_resource mr, + any_device_resource mr, config::Options options, std::shared_ptr statistics = Statistics::disabled() ); + /// @brief Default destructor. ~BufferResource() noexcept = default; + /// @brief Default copy constructor (refcounted shared ownership). + BufferResource(BufferResource const&) noexcept = default; + /// @brief Default move constructor (refcounted shared ownership). + BufferResource(BufferResource&&) noexcept = default; + /** + * @brief Default copy assignment (refcounted shared ownership). + * @return Reference to this. + */ + BufferResource& operator=(BufferResource const&) noexcept = default; + /** + * @brief Default move assignment (refcounted shared ownership). + * @return Reference to this. + */ + BufferResource& operator=(BufferResource&&) noexcept = default; + + /** + * @brief Equality by identity. + * + * Two `BufferResource` handles are equal iff they share the same underlying + * impl instance. + * + * @param other The other `BufferResource` handle. + * @return True iff both handles refer to the same impl. + */ + [[nodiscard]] bool operator==(BufferResource const& other) const noexcept { + return std::addressof(get()) == std::addressof(other.get()); + } + + // --- Per-allocation tracking (was RmmResourceAdaptor) -------------------- + + /** + * @brief Returns a copy of the main memory record. + * + * Lifetime-of-resource allocation statistics, covering all device + * allocations made through this `BufferResource` since its construction. + * + * @return A copy of the main `ScopedMemoryRecord`. + */ + [[nodiscard]] ScopedMemoryRecord get_main_record() const { + return get().get_main_record(); + } + + /** + * @brief Total number of device bytes currently allocated through this resource. + * + * @return Currently outstanding allocated bytes. + */ + [[nodiscard]] std::int64_t current_allocated() const noexcept { + return get().current_allocated(); + } + + /** + * @brief Begin a new scoped memory record on the current thread. + * + * Pushes a fresh `ScopedMemoryRecord` onto the per-thread record stack. + * Subsequent allocations and deallocations on this thread are accumulated + * into the new record (in addition to the main record) until a matching + * `end_scoped_memory_record()` pops it. + * + * @see end_scoped_memory_record() + */ + void begin_scoped_memory_record() { + get().begin_scoped_memory_record(); + } + + /** + * @brief End the topmost scoped memory record on the current thread. + * + * Pops the top of the per-thread record stack and returns it. If another + * scoped record is still active on this thread, the popped record is added + * to it as a sub-scope. + * + * @return The popped `ScopedMemoryRecord`. + * + * @throws std::out_of_range if the current thread's record stack is empty. + * + * @see begin_scoped_memory_record() + */ + ScopedMemoryRecord end_scoped_memory_record() { + return get().end_scoped_memory_record(); + } + + // --- Memory-resource accessors ------------------------------------------- + /** * @brief Get the RMM device memory resource. * * @return Reference to the RMM resource used for device allocations. */ - [[nodiscard]] rmm::device_async_resource_ref device_mr() const noexcept; + [[nodiscard]] rmm::device_async_resource_ref device_mr() const noexcept { + // Wrap *this — BufferResource is itself a CCCL-compatible resource. + // Allocations through the returned ref flow through + // `shared_resource::allocate` → impl tracker, so they are counted by + // `current_allocated()`. The const_cast is safe: allocations are + // logically non-const operations on the underlying state. + return rmm::device_async_resource_ref{const_cast(*this)}; + } /** * @brief Get the RMM host memory resource. * * @return Reference to the RMM resource used for host allocations. */ - [[nodiscard]] rmm::host_async_resource_ref host_mr() noexcept; + [[nodiscard]] rmm::host_async_resource_ref host_mr() noexcept { + return get().host_mr(); + } /** * @brief Get the RMM pinned host memory resource. @@ -134,28 +227,34 @@ class BufferResource { * @throws std::invalid_argument if no pinned memory resource is available. * @return Reference to the RMM resource used for pinned host allocations. */ - [[nodiscard]] rmm::host_device_async_resource_ref pinned_mr(); + [[nodiscard]] rmm::host_device_async_resource_ref pinned_mr() { + return get().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 pinned host memory resource as an `any_resource`, or + * `std::nullopt` if pinned host memory is not available. */ - [[nodiscard]] std::optional try_pinned_mr() const noexcept; + [[nodiscard]] std::optional try_pinned_mr() const noexcept { + return get().try_pinned_mr(); + } + + // --- Memory availability ------------------------------------------------- /** * @brief Returns the currently available memory for a given memory type, in bytes. * - * Computed as `limit - allocated`, where `allocated` is reported by the - * memory type's allocation counter (see the constructor documentation for - * how each memory type is tracked). The value may be negative when + * Computed as `limit - allocated`. The value may be negative when * allocations exceed the configured limit. * * @param mem_type The memory type to query. * @return The available memory in bytes. */ - [[nodiscard]] std::int64_t memory_available(MemoryType mem_type) const noexcept; + [[nodiscard]] std::int64_t memory_available(MemoryType mem_type) const noexcept { + return get().memory_available(mem_type); + } /** * @brief Updates the memory limit for a given memory type at runtime. @@ -171,16 +270,18 @@ class BufferResource { * `memory_available(mem_type)` always negative and so trigger continuous * spilling. */ - void set_memory_limit(MemoryType mem_type, std::int64_t limit) noexcept; + void set_memory_limit(MemoryType mem_type, std::int64_t limit) noexcept { + get().set_memory_limit(mem_type, limit); + } /** - * @brief Get the current reserved memory of the specified memory type. + * @brief Currently reserved bytes for the given memory type. * - * @param mem_type The target memory type. - * @return The memory reserved. + * @param mem_type The memory type to query. + * @return Bytes reserved (but not necessarily allocated) for @p mem_type. */ [[nodiscard]] std::size_t memory_reserved(MemoryType mem_type) const { - return memory_reserved_[static_cast(mem_type)]; + return get().memory_reserved(mem_type); } /** @@ -207,7 +308,9 @@ class BufferResource { */ std::pair reserve( MemoryType mem_type, std::size_t size, AllowOverbooking allow_overbooking - ); + ) { + return get().reserve(this, mem_type, size, allow_overbooking); + } /** * @brief Reserve device memory and spill if necessary. @@ -228,7 +331,9 @@ class BufferResource { */ MemoryReservation reserve_device_memory_and_spill( std::size_t size, AllowOverbooking allow_overbooking - ); + ) { + return get().reserve_device_memory_and_spill(this, size, allow_overbooking); + } /** * @brief Make a memory reservation or fail based on the given order of memory types. @@ -246,10 +351,8 @@ class BufferResource { template requires std::convertible_to, MemoryType> [[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 && !get().try_pinned_mr().has_value()) { // Pinned host memory is only available if the memory resource is // available. @@ -267,8 +370,8 @@ class BufferResource { * @brief Make a memory reservation or fail. * * @param size The size of the buffer to allocate. - * @param mem_type The memory type to try to reserve memory from. - * @return A memory reservation. + * @param mem_type The single memory type to attempt. + * @return A memory reservation in @p mem_type. * * @throws std::runtime_error if no memory reservation was made. */ @@ -290,7 +393,11 @@ class BufferResource { * @throws rapidsmpf::reservation_error if the released size exceeds the size of the * reservation. */ - std::size_t release(MemoryReservation& reservation, std::size_t size); + std::size_t release(MemoryReservation& reservation, std::size_t size) { + return get().release(reservation, size); + } + + // --- Buffer allocation / movement ---------------------------------------- /** * @brief Allocate a buffer of the specified memory type by the reservation. @@ -305,7 +412,9 @@ class BufferResource { */ std::unique_ptr make_buffer( std::size_t size, rmm::cuda_stream_view stream, MemoryReservation& reservation - ); + ) { + return get().make_buffer(this, size, stream, reservation); + } /** * @brief Allocate a buffer consuming the entire reservation. @@ -319,7 +428,9 @@ class BufferResource { */ std::unique_ptr make_buffer( rmm::cuda_stream_view stream, MemoryReservation&& reservation - ); + ) { + return get().make_buffer(this, stream, std::move(reservation)); + } /** * @brief Move device or pinned host buffer data into a Buffer. @@ -336,19 +447,22 @@ class BufferResource { * - the device buffer's current stream is updated to @p stream. * * @param data Unique pointer to the device or pinned host buffer. - * @param stream CUDA stream associated with the new Buffer. Use or synchronize with - * this stream when operating on the Buffer. + * @param stream CUDA stream associated with the new Buffer. Use or + * synchronize with this stream when operating on the Buffer. * @return Unique pointer to the resulting Buffer. */ std::unique_ptr move( std::unique_ptr data, rmm::cuda_stream_view stream - ); + ) { + return get().move(std::move(data), stream); + } /** * @brief Move a Buffer to the memory type specified by the reservation. * - * If the Buffer already resides in the target memory type, a cheap move is performed. - * Otherwise, the Buffer is copied to the target memory using its own CUDA stream. + * If the Buffer already resides in the target memory type, a cheap move + * is performed. Otherwise, the Buffer is copied to the target memory using + * its own CUDA stream. * * @param buffer Buffer to move. * @param reservation Memory reservation used if a copy is required. @@ -359,7 +473,9 @@ class BufferResource { */ std::unique_ptr move( std::unique_ptr buffer, MemoryReservation& reservation - ); + ) { + return get().move(this, std::move(buffer), reservation); + } /** * @brief Move a Buffer to a device buffer. @@ -371,13 +487,16 @@ class BufferResource { * @param reservation Memory reservation used if a copy is required. * @return A unique pointer to the resulting device buffer. * - * @throws std::invalid_argument If the reservation's memory type isn't device memory. - * @throws rapidsmpf::reservation_error if the memory requirement exceeds the - * reservation. + * @throws std::invalid_argument If the reservation's memory type isn't + * device memory. + * @throws rapidsmpf::reservation_error if the memory requirement exceeds + * the reservation. */ std::unique_ptr move_to_device_buffer( std::unique_ptr buffer, MemoryReservation& reservation - ); + ) { + return get().move_to_device_buffer(this, std::move(buffer), reservation); + } /** * @brief Move a Buffer into a host buffer. @@ -389,13 +508,18 @@ class BufferResource { * @param reservation Memory reservation used if a copy is required. * @return Unique pointer to the resulting host buffer. * - * @throws std::invalid_argument If the reservation's memory type isn't host memory. + * @throws std::invalid_argument If the reservation's memory type isn't + * host memory. * @throws rapidsmpf::reservation_error If the allocation size exceeds the * reservation. */ std::unique_ptr move_to_host_buffer( std::unique_ptr buffer, MemoryReservation& reservation - ); + ) { + return get().move_to_host_buffer(this, std::move(buffer), reservation); + } + + // --- Stream pool / spill manager / statistics ---------------------------- /** * @brief Returns the CUDA stream pool used by this buffer resource. @@ -404,14 +528,18 @@ class BufferResource { * * @return Reference to the underlying CUDA stream pool. */ - rmm::cuda_stream_pool const& stream_pool() const; + [[nodiscard]] rmm::cuda_stream_pool const& stream_pool() const { + return get().stream_pool(); + } /** * @brief Gets a reference to the spill manager used. * * @return Reference to the SpillManager instance. */ - SpillManager& spill_manager(); + SpillManager& spill_manager() { + return get().spill_manager(); + } /** * @brief Gets a shared pointer to the statistics associated with this buffer @@ -419,25 +547,12 @@ class BufferResource { * * @return Shared pointer the Statistics instance. */ - std::shared_ptr statistics() const noexcept; - - private: - std::mutex mutex_; - // The internal RmmResourceAdaptor wraps the user's device MR so that - // allocations are tracked for the DEVICE memory_available calculation. - // Declared before device_mr_ because device_mr_ is initialized from it. - RmmResourceAdaptor device_adaptor_; - cuda::mr::any_resource device_mr_; - std::optional pinned_mr_; - HostMemoryResource host_mr_; - std::array, MEMORY_TYPES.size()> memory_limits_; - // Zero initialized reserved counters. - std::array memory_reserved_ = {}; - std::shared_ptr stream_pool_; - SpillManager spill_manager_; - std::shared_ptr statistics_; + [[nodiscard]] std::shared_ptr statistics() const noexcept { + return get().statistics(); + } }; +static_assert(cuda::mr::resource_with); static_assert(StatisticsProvider); /** diff --git a/cpp/include/rapidsmpf/memory/detail/buffer_resource_impl.hpp b/cpp/include/rapidsmpf/memory/detail/buffer_resource_impl.hpp new file mode 100644 index 000000000..378f7bb92 --- /dev/null +++ b/cpp/include/rapidsmpf/memory/detail/buffer_resource_impl.hpp @@ -0,0 +1,416 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace rapidsmpf { + +class BufferResource; // defined in buffer_resource.hpp. + +/** + * @brief Policy controlling whether a memory reservation is allowed to overbook. + * + * This enum is used throughout RapidsMPF to specify the overbooking behavior of + * a memory reservation request. The exact semantics depend on the specific API + * and execution context in which it is used. + */ +enum class AllowOverbooking : bool { + NO, ///< Overbooking is not allowed. + YES, ///< Overbooking is allowed. +}; + +namespace detail { + +/** + * @brief Implementation class for `BufferResource`. + * + * Holds all of `BufferResource`'s state: the user-provided device MR plus + * per-allocation tracking (lifetime stats + scoped records), pinned/host + * MRs, reservation bookkeeping, stream pool, spill manager, and statistics. + * + * This class satisfies the CCCL `cuda::mr::resource` concept and is held by + * `BufferResource` via `cuda::mr::shared_resource` for reference-counted + * ownership. + */ +class BufferResourceImpl { + public: + /// @brief Type-erased device-accessible memory resource. + using any_device_resource = cuda::mr::any_resource; + + /** + * @brief Construct the impl. + * + * @param device_mr Primary device memory resource. + * @param pinned_mr Optional pinned host memory resource. If + * `PinnedMemoryResource::Disabled`, pinned host allocations fail regardless + * of `memory_limits`. + * @param memory_limits Maximum bytes per memory type. Missing entries + * default to unlimited. + * @param periodic_spill_check Pause between periodic spill checks + * (`std::nullopt` disables the dedicated spill thread). + * @param stream_pool CUDA stream pool used for operations without an + * explicit stream. + * @param statistics Statistics instance. + */ + BufferResourceImpl( + any_device_resource device_mr, + std::optional pinned_mr, + std::unordered_map memory_limits, + std::optional periodic_spill_check, + std::shared_ptr stream_pool, + std::shared_ptr statistics + ); + + ~BufferResourceImpl() = default; + + BufferResourceImpl(BufferResourceImpl const&) = delete; + BufferResourceImpl(BufferResourceImpl&&) = delete; + BufferResourceImpl& operator=(BufferResourceImpl const&) = delete; + BufferResourceImpl& operator=(BufferResourceImpl&&) = delete; + + /** + * @brief CCCL concept: async allocation. Records the alloc and forwards + * to the user-provided device MR. + * + * @param stream The CUDA stream for the allocation. + * @param bytes Number of bytes to allocate. + * @param alignment Alignment requirement. + * @return Pointer to the allocated memory. + */ + void* allocate( + cuda::stream_ref stream, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT + ) { + void* ret = device_mr_.allocate(stream, bytes, alignment); + std::lock_guard lock(mutex_); + main_record_.record_allocation(safe_cast(bytes)); + if (!record_stacks_.empty()) { + auto const thread_id = std::this_thread::get_id(); + auto& record = record_stacks_[thread_id]; + if (!record.empty()) { + record.top().record_allocation(safe_cast(bytes)); + RAPIDSMPF_EXPECTS( + allocating_threads_.insert({ret, thread_id}).second, + "duplicate memory pointer" + ); + } + } + return ret; + } + + /** + * @brief CCCL concept: async deallocation. Records the dealloc and + * forwards to the user-provided device MR. + * + * @param stream The CUDA stream for the deallocation. + * @param ptr Pointer to the memory to deallocate. + * @param bytes Number of bytes to deallocate. + * @param alignment Alignment of the original allocation. + */ + void deallocate( + cuda::stream_ref stream, + void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT + ) noexcept { + { + std::lock_guard lock(mutex_); + main_record_.record_deallocation(safe_cast(bytes)); + if (!allocating_threads_.empty()) { + auto const node = allocating_threads_.extract(ptr); + if (node) { + auto thread_id = node.mapped(); + auto& record = record_stacks_[thread_id]; + if (!record.empty()) { + record.top().record_deallocation(safe_cast(bytes)); + } + } + } + } + device_mr_.deallocate(stream, ptr, bytes, alignment); + } + + /** + * @brief CCCL concept: sync allocation. Allocates on the internal sync + * stream and synchronizes before returning. + * + * @param bytes Number of bytes to allocate. + * @param alignment Alignment requirement. + * @return Pointer to the allocated memory. + */ + void* allocate_sync( + std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT + ) { + auto* ptr = allocate(sync_stream_, bytes, alignment); + sync_stream_.synchronize(); + return ptr; + } + + /** + * @brief CCCL concept: sync deallocation. + * + * @param ptr Pointer to the memory to deallocate. + * @param bytes Number of bytes to deallocate. + * @param alignment Alignment of the original allocation. + */ + void deallocate_sync( + void* ptr, + std::size_t bytes, + std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT + ) noexcept { + deallocate(sync_stream_, ptr, bytes, alignment); + } + + /** + * @brief Equality by identity (two impls are equal iff they are the same instance). + * + * @param other The other impl to compare. + * @return True iff @p other is this same instance. + */ + [[nodiscard]] bool operator==(BufferResourceImpl const& other) const noexcept { + return this == std::addressof(other); + } + + /// @brief Tag this resource as device-accessible for the CCCL concept. + friend void get_property( + BufferResourceImpl const&, cuda::mr::device_accessible + ) noexcept {} + + // --- Per-allocation tracking ------------------------------------------- + + /// @copydoc rapidsmpf::BufferResource::get_main_record + [[nodiscard]] ScopedMemoryRecord get_main_record() const { + std::lock_guard lock(mutex_); + return main_record_; + } + + /// @copydoc rapidsmpf::BufferResource::current_allocated + [[nodiscard]] std::int64_t current_allocated() const noexcept { + std::lock_guard lock(mutex_); + return main_record_.current(); + } + + /// @copydoc rapidsmpf::BufferResource::begin_scoped_memory_record + void begin_scoped_memory_record() { + std::lock_guard lock(mutex_); + record_stacks_[std::this_thread::get_id()].emplace(); + } + + /// @copydoc rapidsmpf::BufferResource::end_scoped_memory_record + ScopedMemoryRecord end_scoped_memory_record() { + std::lock_guard lock(mutex_); + auto& stack = record_stacks_.at(std::this_thread::get_id()); + RAPIDSMPF_EXPECTS( + !stack.empty(), + "calling end_scoped_memory_record() on an empty stack.", + std::out_of_range + ); + auto ret = stack.top(); + stack.pop(); + if (!stack.empty()) { + stack.top().add_subscope(ret); + } + return ret; + } + + // --- BufferResource public API (rich operations) ----------------------- + + /// @copydoc rapidsmpf::BufferResource::host_mr + [[nodiscard]] rmm::host_async_resource_ref host_mr() noexcept; + + /// @copydoc rapidsmpf::BufferResource::pinned_mr + [[nodiscard]] rmm::host_device_async_resource_ref pinned_mr(); + + /// @copydoc rapidsmpf::BufferResource::try_pinned_mr + [[nodiscard]] std::optional try_pinned_mr() const noexcept; + + /// @copydoc rapidsmpf::BufferResource::memory_available + [[nodiscard]] std::int64_t memory_available(MemoryType mem_type) const noexcept; + + /// @copydoc rapidsmpf::BufferResource::set_memory_limit + void set_memory_limit(MemoryType mem_type, std::int64_t limit) noexcept; + + /// @copydoc rapidsmpf::BufferResource::memory_reserved + [[nodiscard]] std::size_t memory_reserved(MemoryType mem_type) const { + return memory_reserved_[static_cast(mem_type)]; + } + + /** + * @copydoc rapidsmpf::BufferResource::reserve + * + * @param outer_br Outer `BufferResource` handle stored in the returned + * `MemoryReservation` so that `MemoryReservation::br()` keeps working + * with the public `BufferResource` API. The caller must ensure the outer + * handle outlives the returned reservation. + */ + std::pair reserve( + BufferResource* outer_br, + MemoryType mem_type, + std::size_t size, + AllowOverbooking allow_overbooking + ); + + /** + * @copydoc rapidsmpf::BufferResource::reserve_device_memory_and_spill + * + * @param outer_br Outer `BufferResource` handle stored in the returned + * `MemoryReservation`. The caller must ensure the outer handle outlives + * the returned reservation. + */ + MemoryReservation reserve_device_memory_and_spill( + BufferResource* outer_br, std::size_t size, AllowOverbooking allow_overbooking + ); + + /// @copydoc rapidsmpf::BufferResource::release + std::size_t release(MemoryReservation& reservation, std::size_t size); + + // The buffer-creating methods need the outer `BufferResource` handle to + // construct `rmm::device_buffer` instances against an + // `rmm::device_async_resource_ref` (which requires a copyable resource — + // the impl itself is non-copyable; the outer handle's shared-ownership + // state satisfies the requirement). + + // clang-format off + /** + * @copydoc rapidsmpf::BufferResource::make_buffer(std::size_t,rmm::cuda_stream_view,MemoryReservation&) + * + * @param outer_br Outer `BufferResource` handle used to construct an + * `rmm::device_async_resource_ref` for device allocations (the impl + * itself is non-copyable; the outer handle's shared-ownership state + * satisfies the ref's copyable requirement). + */ + std::unique_ptr make_buffer( + BufferResource* outer_br, + std::size_t size, + rmm::cuda_stream_view stream, + MemoryReservation& reservation + ); + + /** + * @copydoc rapidsmpf::BufferResource::make_buffer(rmm::cuda_stream_view,MemoryReservation&&) + * + * @param outer_br Outer `BufferResource` handle. See the other + * `make_buffer` overload for the rationale. + */ + std::unique_ptr make_buffer( + BufferResource* outer_br, + rmm::cuda_stream_view stream, + MemoryReservation&& reservation + ); + + /** @copydoc rapidsmpf::BufferResource::move(std::unique_ptr,rmm::cuda_stream_view) */ + std::unique_ptr move( + std::unique_ptr data, rmm::cuda_stream_view stream + ); + + /** + * @copydoc rapidsmpf::BufferResource::move(std::unique_ptr,MemoryReservation&) + * + * @param outer_br Outer `BufferResource` handle. Forwarded to + * `make_buffer` when a copy across memory types is needed. + */ + std::unique_ptr move( + BufferResource* outer_br, + std::unique_ptr buffer, + MemoryReservation& reservation + ); + // clang-format on + + /** + * @copydoc rapidsmpf::BufferResource::move_to_device_buffer + * + * @param outer_br Outer `BufferResource` handle. Forwarded to + * `make_buffer` when a copy to device memory is needed. + */ + std::unique_ptr move_to_device_buffer( + BufferResource* outer_br, + std::unique_ptr buffer, + MemoryReservation& reservation + ); + + /** + * @copydoc rapidsmpf::BufferResource::move_to_host_buffer + * + * @param outer_br Outer `BufferResource` handle. Forwarded to + * `make_buffer` when a copy to host memory is needed. + */ + std::unique_ptr move_to_host_buffer( + BufferResource* outer_br, + std::unique_ptr buffer, + MemoryReservation& reservation + ); + + /// @copydoc rapidsmpf::BufferResource::stream_pool + [[nodiscard]] rmm::cuda_stream_pool const& stream_pool() const { + return *stream_pool_; + } + + /// @copydoc rapidsmpf::BufferResource::spill_manager + SpillManager& spill_manager() { + return spill_manager_; + } + + /// @copydoc rapidsmpf::BufferResource::statistics + [[nodiscard]] std::shared_ptr statistics() const noexcept { + return statistics_; + } + + private: + /// @brief Protects all tracking + reservation state. + mutable std::mutex mutex_; + + /// @brief User-provided primary device MR. + any_device_resource device_mr_; + /// @brief Lifetime-of-resource allocation stats. + ScopedMemoryRecord main_record_; + /// @brief Per-thread scoped record stacks. + std::unordered_map> record_stacks_; + /// @brief Map from allocation ptr → originating thread, used to credit + /// deallocations back to the right per-thread scoped record. + std::unordered_map allocating_threads_; + /// @brief Stream used for synchronous allocations/deallocations. + rmm::cuda_stream sync_stream_{rmm::cuda_stream::flags::non_blocking}; + + std::optional pinned_mr_; + HostMemoryResource host_mr_; + std::array, MEMORY_TYPES.size()> memory_limits_; + /// @brief Zero-initialized reserved counters per memory type. + std::array memory_reserved_ = {}; + std::shared_ptr stream_pool_; + SpillManager spill_manager_; + std::shared_ptr statistics_; +}; + +} // namespace detail + +} // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/memory/memory_reservation.hpp b/cpp/include/rapidsmpf/memory/memory_reservation.hpp index 97bdd67d4..7932297ba 100644 --- a/cpp/include/rapidsmpf/memory/memory_reservation.hpp +++ b/cpp/include/rapidsmpf/memory/memory_reservation.hpp @@ -1,5 +1,5 @@ /** - * SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -14,6 +14,10 @@ namespace rapidsmpf { class BufferResource; +namespace detail { +class BufferResourceImpl; +} // namespace detail + /** * @brief Represents a reservation for future memory allocation. * @@ -22,6 +26,7 @@ class BufferResource; */ class MemoryReservation { friend class BufferResource; + friend class detail::BufferResourceImpl; public: /** diff --git a/cpp/include/rapidsmpf/memory/spill_manager.hpp b/cpp/include/rapidsmpf/memory/spill_manager.hpp index 43892fb7a..cf75fd147 100644 --- a/cpp/include/rapidsmpf/memory/spill_manager.hpp +++ b/cpp/include/rapidsmpf/memory/spill_manager.hpp @@ -14,7 +14,9 @@ namespace rapidsmpf { -class BufferResource; +namespace detail { +class BufferResourceImpl; +} // namespace detail /** * @brief Manages memory spilling to free up device memory when needed. @@ -29,6 +31,12 @@ class SpillManager { * * A spill function takes a requested spill amount as input and returns the actual * amount of memory (in bytes) that was spilled. + * + * @warning Spill functions must NOT capture an owning reference (shared_ptr, a + * `BufferResource` by value, etc.) to the `BufferResource` that owns this + * `SpillManager`. Doing so closes a reference cycle: the BR holds the SpillManager + * which holds the function which owns the BR. Capture raw pointers/references only, + * and unregister via `remove_spill_function` from the owner's destructor. */ using SpillFunction = std::function; @@ -40,14 +48,15 @@ class SpillManager { /** * @brief Constructs a SpillManager instance. * - * @param br Buffer resource used to retrieve current available memory. + * @param br_impl Buffer-resource impl used to retrieve current available memory. * @param periodic_spill_check Enable periodic spill checks. A dedicated thread * continuously checks and perform spilling based on the current available memory as * reported by the buffer resource. The value of `periodic_spill_check` is used as the * pause between checks. If `std::nullopt`, no periodic spill check is performed. */ SpillManager( - BufferResource* br, std::optional periodic_spill_check = std::nullopt + detail::BufferResourceImpl* br_impl, + std::optional periodic_spill_check = std::nullopt ); /** @@ -114,11 +123,23 @@ class SpillManager { private: mutable std::mutex mutex_; - BufferResource* br_; std::size_t spill_function_id_counter_{0}; std::map spill_functions_; std::multimap> spill_function_priorities_; std::optional periodic_spill_thread_; + + /// @brief Non-owning back-pointer to the owning impl. + /// + /// `BufferResource` is a thin `shared_resource` handle: copies share + /// the same impl but each carries a different `this`, so there is no + /// canonical outer pointer to point at. The impl IS the resource identity. + /// + /// We cannot hold a `BufferResource` by value either that would close a + /// refcount cycle `BR → Impl → SpillManager → BR`. + /// + /// A raw pointer is safe because `SpillManager` is a member of the impl, + /// so the pointer is stable for `SpillManager`'s entire lifetime. + detail::BufferResourceImpl* br_impl_; }; diff --git a/cpp/include/rapidsmpf/rmm_resource_adaptor.hpp b/cpp/include/rapidsmpf/rmm_resource_adaptor.hpp deleted file mode 100644 index 408820492..000000000 --- a/cpp/include/rapidsmpf/rmm_resource_adaptor.hpp +++ /dev/null @@ -1,127 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include -#include - -#include - -#include - -#include -#include - -namespace rapidsmpf { - -/** - * @brief A RMM memory resource adaptor tailored to RapidsMPF. - * - * This adaptor wraps a primary device memory resource and adds memory usage - * tracking (lifetime stats plus per-thread scoped records). - * - * This class is copyable and shares ownership of its internal state via - * `cuda::mr::shared_resource`. - */ -class RmmResourceAdaptor - : public cuda::mr::shared_resource>> { - using any_device_resource = cuda::mr::any_resource; - using shared_base = - cuda::mr::shared_resource>; - - public: - /// @brief Tag this resource as device-accessible for the CCCL concept. - friend void get_property( - RmmResourceAdaptor const&, cuda::mr::device_accessible - ) noexcept {} - - /** - * @brief Construct with the specified primary memory resource. - * - * @param primary_mr The primary memory resource. - */ - explicit RmmResourceAdaptor( - cuda::mr::any_resource primary_mr - ); - - ~RmmResourceAdaptor() = default; - - /** - * @brief Equality comparison. - * - * Two adaptors are equal if and only if they share the same underlying shared state. - * - * @param other The other adaptor to compare. - * @return True if both adaptors refer to the same shared resource instance. - */ - [[nodiscard]] bool operator==(RmmResourceAdaptor const& other) const noexcept { - return get() == other.get(); - } - - /** - * @brief Get a reference to the primary upstream resource. - * - * @return Reference to the RMM memory resource. - */ - [[nodiscard]] rmm::device_async_resource_ref get_upstream_resource() const noexcept; - - /** - * @brief Returns a copy of the main memory record. - * - * The main record tracks memory statistics for the lifetime of the resource. - * - * @return A copy of the current main memory record. - */ - [[nodiscard]] ScopedMemoryRecord get_main_record() const; - - /** - * @brief Get the total current allocated memory through this resource. - * - * @return Total number of currently allocated bytes. - */ - [[nodiscard]] std::int64_t current_allocated() const noexcept; - - /** - * @brief Begin recording a new scoped memory usage record for the current thread. - * - * This method pushes a new empty `ScopedMemoryRecord` onto the thread-local - * record stack, allowing for nested memory tracking scopes. - * - * Must be paired with a matching call to `end_scoped_memory_record()`. - * - * @see end_scoped_memory_record() - */ - void begin_scoped_memory_record(); - - /** - * @brief End the current scoped memory record and return it. - * - * Pops the top `ScopedMemoryRecord` from the thread-local stack and returns it. - * If this scope was nested within another (i.e. if `begin_scoped_memory_record()` was - * called multiple times in a row), the returned scope is automatically added as a - * subscope to the next scope remaining on the stack. - * - * This allows nesting of scoped memory tracking, where each scope can contain one or - * more subscopes. When analyzing or reporting memory statistics, the memory usage - * of each scope can be calculated **inclusive of its subscopes**. This behavior - * mimics standard hierarchical memory profilers, where the total memory attributed to - * a scope includes all allocations made within it, plus those made in its nested - * regions. - * - * @return The scope that was just ended. - * - * @throws std::out_of_range if called without a matching - * `begin_scoped_memory_record()`. - * - * @see begin_scoped_memory_record() - */ - ScopedMemoryRecord end_scoped_memory_record(); -}; - -static_assert(cuda::mr::resource_with); - -} // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/statistics.hpp b/cpp/include/rapidsmpf/statistics.hpp index adf47b4c7..fe0fecd91 100644 --- a/cpp/include/rapidsmpf/statistics.hpp +++ b/cpp/include/rapidsmpf/statistics.hpp @@ -25,11 +25,12 @@ #include #include #include -#include +#include #include namespace rapidsmpf { +class BufferResource; class StreamOrderedTiming; /** @@ -551,7 +552,7 @@ class Statistics : public std::enable_shared_from_this { */ void clear(); - // TODO: move MemoryRecord and MemoryRecorder to RmmResourceAdaptor? + // TODO: move MemoryRecord and MemoryRecorder to BufferResource? /** * @brief Holds memory profiling information for a named scope. @@ -578,11 +579,13 @@ class Statistics : public std::enable_shared_from_this { * enabled) publishes it under @p name. * * @param stats Owning Statistics. Must not be null. - * @param mr RMM resource adaptor providing scoped memory statistics. + * @param br The `BufferResource` providing scoped memory statistics. * @param name Name of the scope. */ MemoryRecorder( - std::shared_ptr stats, RmmResourceAdaptor mr, std::string name + std::shared_ptr stats, + std::shared_ptr br, + std::string name ); ~MemoryRecorder(); @@ -593,9 +596,9 @@ class Statistics : public std::enable_shared_from_this { MemoryRecorder& operator=(MemoryRecorder&&) = delete; private: - /// No-op recorder iff `mr_` is `std::nullopt`. - std::optional mr_{std::nullopt}; - /// stats_ != nullptr iff `mr_.has_value()`. + /// No-op recorder iff `br_` is null. + std::shared_ptr br_; + /// stats_ is non-null iff `br_` is non-null. std::shared_ptr stats_{nullptr}; std::string name_{}; }; @@ -604,10 +607,10 @@ class Statistics : public std::enable_shared_from_this { * @brief Creates a scoped memory recorder for the given name. * * @param mr Type-erased device memory resource. Recording is only active - * when the underlying resource is an `RmmResourceAdaptor`. + * when the underlying resource is a `BufferResource`. * @param name Name of the scope. - * @return A MemoryRecorder instance. If `!enabled()` or @p mr is not backed by an - * `RmmResourceAdaptor`, returns a no-op recorder. + * @return A MemoryRecorder instance. If `!enabled()` or @p mr is not backed by a + * `BufferResource`, returns a no-op recorder. */ MemoryRecorder create_memory_recorder(any_device_resource mr, std::string name); @@ -669,9 +672,9 @@ concept StatisticsProvider = requires(T const& t) { * * Example usage: * @code - * void foo(std::shared_ptr stats, RmmResourceAdaptor& mr) { - * RAPIDSMPF_MEMORY_PROFILE(stats, mr); - * RAPIDSMPF_MEMORY_PROFILE(stats, mr, "custom_name"); + * void foo(std::shared_ptr stats, BufferResource& br) { + * RAPIDSMPF_MEMORY_PROFILE(stats, br); + * RAPIDSMPF_MEMORY_PROFILE(stats, br, "custom_name"); * } * @endcode * @@ -680,7 +683,7 @@ concept StatisticsProvider = requires(T const& t) { * (`create_memory_recorder` returns a no-op recorder when the statistics * instance is disabled). * The second argument is the device memory resource. Recording is only active - * when the underlying resource is an `RmmResourceAdaptor`; other device + * when the underlying resource is a `BufferResource`; other device * resources yield a no-op recorder. * The third argument (optional) is a custom function name string to use instead of * __func__. diff --git a/cpp/include/rapidsmpf/streaming/core/context.hpp b/cpp/include/rapidsmpf/streaming/core/context.hpp index 559f58574..faf83d426 100644 --- a/cpp/include/rapidsmpf/streaming/core/context.hpp +++ b/cpp/include/rapidsmpf/streaming/core/context.hpp @@ -8,6 +8,8 @@ #include #include +#include + #include #include @@ -79,7 +81,7 @@ class Context { * @note The current CUDA device must be set prior to calling this function. * Options that depend on device memory availability query the current device. * - * @param mr Device memory resource adaptor used by RapidsMPF. + * @param device_mr Primary device memory resource used by RapidsMPF. * @param logger The logger to use. * @param options Configuration options used to initialize the Context and its * components. @@ -102,7 +104,7 @@ class Context { * thread. */ static std::shared_ptr from_options( - RmmResourceAdaptor mr, + cuda::mr::any_resource device_mr, std::shared_ptr logger, config::Options options, std::shared_ptr statistics = Statistics::disabled() diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index c2f24fd73..b2171b29c 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -20,16 +20,17 @@ namespace rapidsmpf { -BufferResource::BufferResource( - cuda::mr::any_resource device_mr, +namespace detail { + +BufferResourceImpl::BufferResourceImpl( + any_device_resource device_mr, std::optional pinned_mr, std::unordered_map memory_limits, std::optional periodic_spill_check, std::shared_ptr stream_pool, std::shared_ptr statistics ) - : device_adaptor_{std::move(device_mr)}, - device_mr_{device_adaptor_}, // any_resource shares state via shared_resource + : device_mr_{std::move(device_mr)}, pinned_mr_{std::move(pinned_mr)}, host_mr_{}, stream_pool_{std::move(stream_pool)}, @@ -48,33 +49,13 @@ BufferResource::BufferResource( RAPIDSMPF_EXPECTS(statistics_ != nullptr, "the statistics pointer cannot be NULL"); } -std::shared_ptr BufferResource::from_options( - cuda::mr::any_resource mr, - 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 std::make_shared( - std::move(mr), - std::move(pinned_mr), - std::move(memory_limits), - periodic_spill_check_from_options(options), - stream_pool_from_options(options), - std::move(statistics) - ); -} - -std::int64_t BufferResource::memory_available(MemoryType mem_type) const noexcept { +std::int64_t BufferResourceImpl::memory_available(MemoryType mem_type) const noexcept { std::int64_t const limit = memory_limits_[static_cast(mem_type)].load( std::memory_order_acquire ); switch (mem_type) { case MemoryType::DEVICE: - return limit - device_adaptor_.current_allocated(); + return limit - current_allocated(); case MemoryType::PINNED_HOST: if (pinned_mr_ == PinnedMemoryResource::Disabled) { return 0; @@ -87,37 +68,37 @@ std::int64_t BufferResource::memory_available(MemoryType mem_type) const noexcep return std::numeric_limits::max(); } -void BufferResource::set_memory_limit(MemoryType mem_type, std::int64_t limit) noexcept { +void BufferResourceImpl::set_memory_limit( + MemoryType mem_type, std::int64_t limit +) noexcept { memory_limits_[static_cast(mem_type)].store( limit, std::memory_order_release ); } -rmm::device_async_resource_ref BufferResource::device_mr() const noexcept { - return rmm::device_async_resource_ref{ - const_cast&>(device_mr_) - }; -} - -rmm::host_async_resource_ref BufferResource::host_mr() noexcept { +rmm::host_async_resource_ref BufferResourceImpl::host_mr() noexcept { return host_mr_; } -rmm::host_device_async_resource_ref BufferResource::pinned_mr() { +rmm::host_device_async_resource_ref BufferResourceImpl::pinned_mr() { RAPIDSMPF_EXPECTS( pinned_mr_, "no pinned memory resource is available", std::invalid_argument ); return *pinned_mr_; } -std::optional BufferResource::try_pinned_mr() const noexcept { +std::optional +BufferResourceImpl::try_pinned_mr() const noexcept { // since any_host_device_resource is constructible from // host_device_async_resource_ref, optional can be returned as-is. return pinned_mr_; } -std::pair BufferResource::reserve( - MemoryType mem_type, std::size_t size, AllowOverbooking allow_overbooking +std::pair BufferResourceImpl::reserve( + BufferResource* outer_br, + MemoryType mem_type, + std::size_t size, + AllowOverbooking allow_overbooking ) { RAPIDSMPF_EXPECTS( mem_type != MemoryType::PINNED_HOST @@ -138,18 +119,19 @@ std::pair BufferResource::reserve( headroom < 0 ? safe_cast(std::abs(headroom)) : 0; if (overbooking > 0 && allow_overbooking == AllowOverbooking::NO) { // Cancel the reservation, overbooking isn't allowed. - return {MemoryReservation(mem_type, this, 0), overbooking}; + return {MemoryReservation(mem_type, outer_br, 0), overbooking}; } // Make the reservation. reserved += size; - return {MemoryReservation(mem_type, this, size), overbooking}; + return {MemoryReservation(mem_type, outer_br, size), overbooking}; } -MemoryReservation BufferResource::reserve_device_memory_and_spill( - std::size_t size, AllowOverbooking allow_overbooking +MemoryReservation BufferResourceImpl::reserve_device_memory_and_spill( + BufferResource* outer_br, std::size_t size, AllowOverbooking allow_overbooking ) { // reserve device memory with overbooking - auto [reservation, ob] = reserve(MemoryType::DEVICE, size, AllowOverbooking::YES); + auto [reservation, ob] = + reserve(outer_br, MemoryType::DEVICE, size, AllowOverbooking::YES); // ask the spill manager to make room for overbooking if (ob > 0) { @@ -166,7 +148,9 @@ MemoryReservation BufferResource::reserve_device_memory_and_spill( return std::move(reservation); } -std::size_t BufferResource::release(MemoryReservation& reservation, std::size_t size) { +std::size_t BufferResourceImpl::release( + MemoryReservation& reservation, std::size_t size +) { std::lock_guard const lock(mutex_); RAPIDSMPF_EXPECTS( size <= reservation.size_, @@ -181,8 +165,11 @@ std::size_t BufferResource::release(MemoryReservation& reservation, std::size_t return reservation.size_ -= size; } -std::unique_ptr BufferResource::make_buffer( - std::size_t size, rmm::cuda_stream_view stream, MemoryReservation& reservation +std::unique_ptr BufferResourceImpl::make_buffer( + BufferResource* outer_br, + std::size_t size, + rmm::cuda_stream_view stream, + MemoryReservation& reservation ) { auto const mem_type = reservation.mem_type_; StreamOrderedTiming timing{stream, statistics_}; @@ -204,7 +191,9 @@ std::unique_ptr BufferResource::make_buffer( break; case MemoryType::DEVICE: ret = std::unique_ptr(new Buffer( - std::make_unique(size, stream, device_mr()), + std::make_unique( + size, stream, rmm::device_async_resource_ref{*outer_br} + ), MemoryType::DEVICE )); break; @@ -216,13 +205,15 @@ std::unique_ptr BufferResource::make_buffer( return ret; } -std::unique_ptr BufferResource::make_buffer( - rmm::cuda_stream_view stream, MemoryReservation&& reservation +std::unique_ptr BufferResourceImpl::make_buffer( + BufferResource* outer_br, + rmm::cuda_stream_view stream, + MemoryReservation&& reservation ) { - return make_buffer(reservation.size(), stream, reservation); + return make_buffer(outer_br, reservation.size(), stream, reservation); } -std::unique_ptr BufferResource::move( +std::unique_ptr BufferResourceImpl::move( std::unique_ptr data, rmm::cuda_stream_view stream ) { auto upstream = data->stream(); @@ -242,20 +233,24 @@ std::unique_ptr BufferResource::move( return std::unique_ptr(new Buffer(std::move(data), MemoryType::DEVICE)); } -std::unique_ptr BufferResource::move( - std::unique_ptr buffer, MemoryReservation& reservation +std::unique_ptr BufferResourceImpl::move( + BufferResource* outer_br, + std::unique_ptr buffer, + MemoryReservation& reservation ) { if (reservation.mem_type_ != buffer->mem_type()) { auto const nbytes = buffer->size; - auto ret = make_buffer(nbytes, buffer->stream(), reservation); + auto ret = make_buffer(outer_br, nbytes, buffer->stream(), reservation); buffer_copy(statistics_, *ret, *buffer, nbytes); return ret; } return buffer; } -std::unique_ptr BufferResource::move_to_device_buffer( - std::unique_ptr buffer, MemoryReservation& reservation +std::unique_ptr BufferResourceImpl::move_to_device_buffer( + BufferResource* outer_br, + std::unique_ptr buffer, + MemoryReservation& reservation ) { RAPIDSMPF_EXPECTS( reservation.mem_type_ == MemoryType::DEVICE, @@ -263,7 +258,7 @@ std::unique_ptr BufferResource::move_to_device_buffer( std::invalid_argument ); auto stream = buffer->stream(); - auto ret = move(std::move(buffer), reservation)->release_device_buffer(); + auto ret = move(outer_br, std::move(buffer), reservation)->release_device_buffer(); RAPIDSMPF_EXPECTS( ret->stream().value() == stream.value(), "something went wrong, the Buffer's stream and the device_buffer's stream " @@ -272,27 +267,55 @@ std::unique_ptr BufferResource::move_to_device_buffer( return ret; } -std::unique_ptr BufferResource::move_to_host_buffer( - std::unique_ptr buffer, MemoryReservation& reservation +std::unique_ptr BufferResourceImpl::move_to_host_buffer( + BufferResource* outer_br, + std::unique_ptr buffer, + MemoryReservation& reservation ) { RAPIDSMPF_EXPECTS( reservation.mem_type_ == MemoryType::HOST, "the memory type of MemoryReservation doesn't match", std::invalid_argument ); - return move(std::move(buffer), reservation)->release_host_buffer(); + return move(outer_br, std::move(buffer), reservation)->release_host_buffer(); } -rmm::cuda_stream_pool const& BufferResource::stream_pool() const { - return *stream_pool_; -} +} // namespace detail -SpillManager& BufferResource::spill_manager() { - return spill_manager_; -} +BufferResource::BufferResource( + any_device_resource device_mr, + std::optional pinned_mr, + std::unordered_map memory_limits, + std::optional periodic_spill_check, + std::shared_ptr stream_pool, + std::shared_ptr statistics +) + : shared_base{cuda::mr::make_shared_resource( + std::move(device_mr), + std::move(pinned_mr), + std::move(memory_limits), + periodic_spill_check, + std::move(stream_pool), + std::move(statistics) + )} {} -std::shared_ptr BufferResource::statistics() const noexcept { - return statistics_; +std::shared_ptr BufferResource::from_options( + any_device_resource mr, + 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 std::make_shared( + std::move(mr), + std::move(pinned_mr), + std::move(memory_limits), + periodic_spill_check_from_options(options), + stream_pool_from_options(options), + std::move(statistics) + ); } std::int64_t device_limit_from_options(config::Options options) { diff --git a/cpp/src/memory/spill_manager.cpp b/cpp/src/memory/spill_manager.cpp index b9a1af6e0..31ba0d890 100644 --- a/cpp/src/memory/spill_manager.cpp +++ b/cpp/src/memory/spill_manager.cpp @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -14,9 +15,9 @@ namespace rapidsmpf { SpillManager::SpillManager( - BufferResource* br, std::optional periodic_spill_check + detail::BufferResourceImpl* br_impl, std::optional periodic_spill_check ) - : br_{br} { + : br_impl_{br_impl} { if (periodic_spill_check.has_value()) { periodic_spill_thread_.emplace( [this]() { spill_to_make_headroom(0); }, *periodic_spill_check @@ -80,7 +81,7 @@ std::size_t SpillManager::spill(std::size_t amount) { std::size_t SpillManager::spill_to_make_headroom(std::int64_t headroom) { // TODO: check other memory types. - std::int64_t available = br_->memory_available(MemoryType::DEVICE); + std::int64_t available = br_impl_->memory_available(MemoryType::DEVICE); if (headroom <= available) { return 0; } diff --git a/cpp/src/rmm_resource_adaptor.cpp b/cpp/src/rmm_resource_adaptor.cpp deleted file mode 100644 index f7a223e8c..000000000 --- a/cpp/src/rmm_resource_adaptor.cpp +++ /dev/null @@ -1,43 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include - -namespace rapidsmpf { - -RmmResourceAdaptor::RmmResourceAdaptor( - cuda::mr::any_resource primary_mr -) - : shared_base( - cuda::mr::make_shared_resource< - detail::RmmResourceAdaptorImpl>(std::move(primary_mr)) - ) {} - -rmm::device_async_resource_ref -RmmResourceAdaptor::get_upstream_resource() const noexcept { - return rmm::device_async_resource_ref{ - const_cast(get().get_upstream_resource()) - }; -} - -ScopedMemoryRecord RmmResourceAdaptor::get_main_record() const { - return get().get_main_record(); -} - -std::int64_t RmmResourceAdaptor::current_allocated() const noexcept { - return get().current_allocated(); -} - -void RmmResourceAdaptor::begin_scoped_memory_record() { - get().begin_scoped_memory_record(); -} - -ScopedMemoryRecord RmmResourceAdaptor::end_scoped_memory_record() { - return get().end_scoped_memory_record(); -} - -} // namespace rapidsmpf diff --git a/cpp/src/statistics.cpp b/cpp/src/statistics.cpp index 3bc7a6d1a..f721b83bb 100644 --- a/cpp/src/statistics.cpp +++ b/cpp/src/statistics.cpp @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -248,20 +249,23 @@ void Statistics::clear() { } Statistics::MemoryRecorder::MemoryRecorder( - std::shared_ptr stats, RmmResourceAdaptor mr, std::string name + std::shared_ptr stats, + std::shared_ptr br, + std::string name ) - : mr_{std::move(mr)}, stats_{std::move(stats)}, name_{std::move(name)} { + : br_{std::move(br)}, stats_{std::move(stats)}, name_{std::move(name)} { RAPIDSMPF_EXPECTS(stats_ != nullptr, "the statistics cannot be null"); - mr_->begin_scoped_memory_record(); + RAPIDSMPF_EXPECTS(br_ != nullptr, "the buffer resource cannot be null"); + br_->begin_scoped_memory_record(); } Statistics::MemoryRecorder::~MemoryRecorder() { - if (!mr_.has_value()) { + if (br_ == nullptr) { return; // no-op recorder; nothing was pushed. } - // Always pop to keep the RMM adaptor's per-thread stack balanced, even if + // Always pop to keep the per-thread scoped-record stack balanced, even if // statistics were disabled after construction (in which case skip publish). - auto const scope = mr_->end_scoped_memory_record(); + auto const scope = br_->end_scoped_memory_record(); if (!stats_->enabled()) { return; } @@ -275,11 +279,13 @@ Statistics::MemoryRecorder::~MemoryRecorder() { Statistics::MemoryRecorder Statistics::create_memory_recorder( any_device_resource mr, std::string name ) { - auto* rma = cuda::mr::resource_cast(&mr); - if (!enabled() || !rma) { + auto* br = cuda::mr::resource_cast(&mr); + if (!enabled() || !br) { return MemoryRecorder{}; } - return MemoryRecorder{shared_from_this(), *rma, std::move(name)}; + return MemoryRecorder{ + shared_from_this(), std::make_shared(*br), std::move(name) + }; } std::unordered_map const& @@ -359,7 +365,7 @@ std::string Statistics::report(ReportArgs report_args) const { // Print memory profiling. ss << "Memory Profiling\n"; ss << "----------------\n"; - auto* dev_adaptor = get_optional_resource_as(report_args.mr); + auto* dev_adaptor = get_optional_resource_as(report_args.mr); auto* pinned_adaptor = get_optional_resource_as(report_args.pinned_mr); if (!dev_adaptor) { @@ -375,7 +381,7 @@ std::string Statistics::report(ReportArgs report_args) const { // Insert the "main" record, which is the overall statistics from `mr`. auto const main_record = dev_adaptor->get_main_record(); sorted_records.emplace_back( - "main (all allocations using RmmResourceAdaptor)", + "main (all allocations using BufferResource)", MemoryRecord{ .scoped = main_record, .global_peak = main_record.peak(), .num_calls = 1 } diff --git a/cpp/src/streaming/core/context.cpp b/cpp/src/streaming/core/context.cpp index e2e00d5b8..136595bcf 100644 --- a/cpp/src/streaming/core/context.cpp +++ b/cpp/src/streaming/core/context.cpp @@ -104,7 +104,7 @@ Context::Context( ) {} std::shared_ptr Context::from_options( - RmmResourceAdaptor mr, + cuda::mr::any_resource device_mr, std::shared_ptr logger, config::Options options, std::shared_ptr statistics @@ -112,7 +112,7 @@ std::shared_ptr Context::from_options( return std::make_shared( options, std::move(logger), - BufferResource::from_options(std::move(mr), options, std::move(statistics)) + BufferResource::from_options(std::move(device_mr), options, std::move(statistics)) ); } diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 2f5a52979..cd12c4dcf 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -88,7 +88,7 @@ target_sources( test_partition.cpp test_pausable_thread_loop.cpp test_progress_thread.cpp - test_rmm_resource_adaptor.cpp + test_buffer_resource_tracking.cpp test_sparse_alltoall.cpp test_shuffler_many_streams.cpp test_shuffler.cpp diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index cea208c88..570e4fa1e 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -638,12 +638,18 @@ class BufferResourceDifferentResourcesTest : public ::testing::Test { } // Setup br1 with statistics for its device memory - mr1 = std::make_unique(rmm::mr::cuda_memory_resource{}); - br1 = std::make_unique(*mr1); + br1 = std::make_unique( + cuda::mr::any_resource{ + rmm::mr::cuda_memory_resource{} + } + ); // Setup br2 with statistics for its device memory - mr2 = std::make_unique(rmm::mr::cuda_memory_resource{}); - br2 = std::make_unique(*mr2); + br2 = std::make_unique( + cuda::mr::any_resource{ + rmm::mr::cuda_memory_resource{} + } + ); } std::unique_ptr create_source_buffer() { @@ -660,23 +666,21 @@ class BufferResourceDifferentResourcesTest : public ::testing::Test { ); }); buf1->stream().synchronize(); - EXPECT_EQ(mr1->get_main_record().total(), buffer_size); + EXPECT_EQ(br1->get_main_record().total(), buffer_size); return buf1; } void verify_memory_allocation( std::size_t expected_br1_total, std::size_t expected_br2_total ) { - EXPECT_EQ(mr1->get_main_record().total(), expected_br1_total); - EXPECT_EQ(mr2->get_main_record().total(), expected_br2_total); + EXPECT_EQ(br1->get_main_record().total(), expected_br1_total); + EXPECT_EQ(br2->get_main_record().total(), expected_br2_total); } std::size_t buffer_size; rmm::cuda_stream_view stream; std::vector host_pattern; - std::unique_ptr mr1; - std::unique_ptr mr2; std::unique_ptr br1; std::unique_ptr br2; }; diff --git a/cpp/tests/test_rmm_resource_adaptor.cpp b/cpp/tests/test_buffer_resource_tracking.cpp similarity index 83% rename from cpp/tests/test_rmm_resource_adaptor.cpp rename to cpp/tests/test_buffer_resource_tracking.cpp index 56bccef4b..c100f971b 100644 --- a/cpp/tests/test_rmm_resource_adaptor.cpp +++ b/cpp/tests/test_buffer_resource_tracking.cpp @@ -19,7 +19,7 @@ #include #include -#include +#include #include "utils.hpp" @@ -109,9 +109,9 @@ struct throw_at_limit_resource } }; -TEST(RmmResourceAdaptor, TracksAllocations) { +TEST(BufferResourceTracking, TracksAllocations) { throw_at_limit_resource primary_mr{4_MiB}; - RmmResourceAdaptor mr{primary_mr}; + BufferResource mr{cuda::mr::any_resource{primary_mr}}; EXPECT_EQ(mr.current_allocated(), 0); @@ -124,23 +124,23 @@ TEST(RmmResourceAdaptor, TracksAllocations) { EXPECT_EQ(mr.current_allocated(), 0); } -TEST(RmmResourceAdaptor, OOMPropagates) { +TEST(BufferResourceTracking, OOMPropagates) { throw_at_limit_resource primary_mr{1_MiB}; - RmmResourceAdaptor mr{primary_mr}; + BufferResource mr{cuda::mr::any_resource{primary_mr}}; EXPECT_THROW((void)mr.allocate_sync(8_MiB), rmm::out_of_memory); } -TEST(RmmResourceAdaptor, PropagatesNonOutOfMemoryExceptions) { +TEST(BufferResourceTracking, PropagatesNonOutOfMemoryExceptions) { throw_at_limit_resource primary_mr{1_MiB}; - RmmResourceAdaptor mr{primary_mr}; + BufferResource mr{cuda::mr::any_resource{primary_mr}}; EXPECT_THROW(std::ignore = mr.allocate_sync(2_MiB), std::logic_error); } -TEST(RmmResourceAdaptor, RecordReflectsCorrectStatistics) { +TEST(BufferResourceTracking, RecordReflectsCorrectStatistics) { throw_at_limit_resource primary_mr{4_MiB}; - RmmResourceAdaptor mr{primary_mr}; + BufferResource mr{cuda::mr::any_resource{primary_mr}}; auto main_record_before = mr.get_main_record(); EXPECT_EQ(main_record_before.num_total_allocs(), 0); @@ -219,8 +219,10 @@ TEST(ScopedMemoryRecord, AddScopeMergesSiblingScopesCorrectly) { EXPECT_EQ(scope1.num_total_allocs(), 2); } -TEST(RmmResourceAdaptor, EmptyScopedMemoryRecord) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; +TEST(BufferResourceTracking, EmptyScopedMemoryRecord) { + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; mr.begin_scoped_memory_record(); auto scope = mr.end_scoped_memory_record(); @@ -230,8 +232,10 @@ TEST(RmmResourceAdaptor, EmptyScopedMemoryRecord) { EXPECT_EQ(scope.num_total_allocs(), 0); } -TEST(RmmResourceAdaptorScopedMemory, SingleScopedAllocationTracksCorrectly) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; +TEST(BufferResourceTrackingScopedMemory, SingleScopedAllocationTracksCorrectly) { + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; mr.begin_scoped_memory_record(); void* p = mr.allocate_sync(1_MiB); @@ -245,8 +249,10 @@ TEST(RmmResourceAdaptorScopedMemory, SingleScopedAllocationTracksCorrectly) { mr.deallocate_sync(p, 1_MiB); } -TEST(RmmResourceAdaptorScopedMemory, NestedScopedAllocationsMerged) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; +TEST(BufferResourceTrackingScopedMemory, NestedScopedAllocationsMerged) { + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; mr.begin_scoped_memory_record(); // Outer @@ -271,8 +277,10 @@ TEST(RmmResourceAdaptorScopedMemory, NestedScopedAllocationsMerged) { mr.deallocate_sync(p1, 1_MiB); } -TEST(RmmResourceAdaptorScopedMemory, NestedScopedTracksAllocsAndDeallocs) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; +TEST(BufferResourceTrackingScopedMemory, NestedScopedTracksAllocsAndDeallocs) { + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; mr.begin_scoped_memory_record(); // Outer @@ -298,8 +306,10 @@ TEST(RmmResourceAdaptorScopedMemory, NestedScopedTracksAllocsAndDeallocs) { mr.deallocate_sync(p1, 1_MiB); } -TEST(RmmResourceAdaptorScopedMemory, NestedDeallocationYieldsNegativeStats) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; +TEST(BufferResourceTrackingScopedMemory, NestedDeallocationYieldsNegativeStats) { + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; // Allocate in outer scope mr.begin_scoped_memory_record(); // Outer @@ -322,12 +332,14 @@ TEST(RmmResourceAdaptorScopedMemory, NestedDeallocationYieldsNegativeStats) { EXPECT_EQ(outer.current(), 0); // Net usage is zero } -TEST(RmmResourceAdaptorScopedMemory, MultiThreadedScopedAllocations) { +TEST(BufferResourceTrackingScopedMemory, MultiThreadedScopedAllocations) { constexpr int num_threads = 8; constexpr int num_allocs_per_thread = 8; constexpr std::size_t alloc_size = 1_MiB; - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; std::vector threads; std::vector> allocations(num_threads); std::vector records(num_threads); @@ -378,25 +390,31 @@ TEST(RmmResourceAdaptorScopedMemory, MultiThreadedScopedAllocations) { EXPECT_EQ(mr.current_allocated(), 0); // All allocations have been released } -TEST(RmmResourceAdaptor, EqualityWithCudaMemoryResource) { +TEST(BufferResourceTracking, EqualityWithCudaMemoryResource) { rmm::mr::cuda_memory_resource cuda_mr{}; - RmmResourceAdaptor adaptor_a{cuda_mr}; - RmmResourceAdaptor adaptor_b{cuda_mr}; + BufferResource adaptor_a{ + cuda::mr::any_resource{cuda_mr} + }; + BufferResource adaptor_b{ + cuda::mr::any_resource{cuda_mr} + }; // Both wrap same resouce but have difference shared states EXPECT_NE(adaptor_a, adaptor_b); // A copy shares the same control block -> equal. - RmmResourceAdaptor adaptor_a_copy = adaptor_a; + BufferResource adaptor_a_copy = adaptor_a; EXPECT_EQ(adaptor_a, adaptor_a_copy); } -TEST(RmmResourceAdaptorScopedMemory, CrossThreadNestedScopesNotMerged) { +TEST(BufferResourceTrackingScopedMemory, CrossThreadNestedScopesNotMerged) { constexpr std::size_t outer_alloc_size = 1_MiB; constexpr std::size_t inner_alloc_size = 2_MiB; - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; void* outer_alloc = nullptr; void* inner_alloc = nullptr; rapidsmpf::ScopedMemoryRecord inner_record; diff --git a/cpp/tests/test_config.cpp b/cpp/tests/test_config.cpp index c78d840cf..8561e9f22 100644 --- a/cpp/tests/test_config.cpp +++ b/cpp/tests/test_config.cpp @@ -15,7 +15,6 @@ #include #include #include -#include #include #include #include @@ -622,7 +621,7 @@ TEST(OptionsTest, BufferResourceFromOptionsCreatesInstanceWithExplicitOptions) { config::Options opts(strings); rmm::mr::cuda_memory_resource cuda_mr; - RmmResourceAdaptor mr{cuda_mr}; + cuda::mr::any_resource mr{cuda_mr}; auto br = BufferResource::from_options(mr, opts, Statistics::from_options(opts)); EXPECT_TRUE(br->statistics()->enabled()); @@ -634,7 +633,7 @@ TEST(OptionsTest, BufferResourceFromOptionsUsesDefaultWhenOptionsEmpty) { config::Options opts; // Empty options rmm::mr::cuda_memory_resource cuda_mr; - RmmResourceAdaptor mr{cuda_mr}; + cuda::mr::any_resource mr{cuda_mr}; auto br = BufferResource::from_options(mr, opts); EXPECT_FALSE(br->statistics()->enabled()); EXPECT_EQ(br->stream_pool().get_pool_size(), 16); @@ -648,7 +647,7 @@ TEST(OptionsTest, BufferResourceFromOptionsEnablesStatisticsWhenRequested) { config::Options opts(strings); rmm::mr::cuda_memory_resource cuda_mr; - RmmResourceAdaptor mr{cuda_mr}; + cuda::mr::any_resource mr{cuda_mr}; auto br = BufferResource::from_options(mr, opts, Statistics::from_options(opts)); EXPECT_TRUE(br->statistics()->enabled()); @@ -661,7 +660,7 @@ TEST(OptionsTest, BufferResourceFromOptionsAcceptsPercentageForDeviceLimit) { config::Options opts(strings); rmm::mr::cuda_memory_resource cuda_mr; - RmmResourceAdaptor mr{cuda_mr}; + cuda::mr::any_resource mr{cuda_mr}; auto br = BufferResource::from_options(mr, opts); // Verify device memory limit is 50% of total @@ -679,7 +678,7 @@ TEST(OptionsTest, BufferResourceFromOptionsEnablesPinnedMemoryWhenSupported) { config::Options opts(strings); rmm::mr::cuda_memory_resource cuda_mr; - RmmResourceAdaptor mr{cuda_mr}; + cuda::mr::any_resource mr{cuda_mr}; auto br = BufferResource::from_options(mr, opts); // Should not throw when accessing pinned_mr @@ -695,7 +694,7 @@ TEST(OptionsTest, ContextFromOptionsCreatesInstanceWithExplicitOptions) { config::Options opts(strings); rmm::mr::cuda_memory_resource cuda_mr; - RmmResourceAdaptor mr{cuda_mr}; + cuda::mr::any_resource mr{cuda_mr}; auto comm = std::make_shared(opts, std::make_shared()); auto ctx = streaming::Context::from_options( @@ -711,7 +710,7 @@ TEST(OptionsTest, ContextFromOptionsUsesDefaultWhenOptionsEmpty) { config::Options opts; rmm::mr::cuda_memory_resource cuda_mr; - RmmResourceAdaptor mr{cuda_mr}; + cuda::mr::any_resource mr{cuda_mr}; auto comm = std::make_shared(opts, std::make_shared()); auto ctx = streaming::Context::from_options(mr, comm->logger(), opts); @@ -726,7 +725,7 @@ TEST(OptionsTest, ContextFromOptionsEnablesStatisticsWhenRequested) { config::Options opts(strings); rmm::mr::cuda_memory_resource cuda_mr; - RmmResourceAdaptor mr{cuda_mr}; + cuda::mr::any_resource mr{cuda_mr}; auto comm = std::make_shared(opts, std::make_shared()); auto ctx = streaming::Context::from_options( @@ -741,7 +740,7 @@ TEST(OptionsTest, ContextFromOptionsCreatesProgressThread) { config::Options opts; rmm::mr::cuda_memory_resource cuda_mr; - RmmResourceAdaptor mr{cuda_mr}; + cuda::mr::any_resource mr{cuda_mr}; auto comm = std::make_shared(opts, std::make_shared()); auto ctx = streaming::Context::from_options(mr, comm->logger(), opts); @@ -753,7 +752,7 @@ TEST(OptionsTest, ContextFromOptionsCreatesExecutor) { config::Options opts; rmm::mr::cuda_memory_resource cuda_mr; - RmmResourceAdaptor mr{cuda_mr}; + cuda::mr::any_resource mr{cuda_mr}; auto comm = std::make_shared(opts, std::make_shared()); auto ctx = streaming::Context::from_options(mr, comm->logger(), opts); diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index df89b50d5..e3de230fb 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -403,19 +403,17 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { cudf::hash_id const hash_fn = cudf::hash_id::HASH_MURMUR3; auto stream = cudf::get_default_stream(); - // Use RapidsMPF's memory resource adaptor so the test can observe per-rank - // allocation counts via `get_main_record().num_current_allocs()`. - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; - // Control spilling by adjusting the DEVICE memory limit at runtime. // `memory_available(DEVICE)` is computed as `limit - current_allocated()`, so a // sufficiently large positive limit reliably keeps available memory > 0 (no spill), // while a sufficiently large negative limit reliably keeps available memory < 0 - // (force spill), regardless of how many bytes are currently allocated from `mr`. + // (force spill), regardless of how many bytes are currently allocated. constexpr std::int64_t k_no_spill_limit = (1LL << 40); constexpr std::int64_t k_force_spill_limit = -(1LL << 40); rapidsmpf::BufferResource br{ - mr, + cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }, rapidsmpf::PinnedMemoryResource::Disabled, {{rapidsmpf::MemoryType::DEVICE, k_no_spill_limit}}, std::nullopt // disable periodic spill check @@ -446,10 +444,10 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { // Insert spills does nothing when device memory is available, we start // with 2 device allocations. - EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2); + EXPECT_EQ(br.get_main_record().num_current_allocs(), 2); shuffler.insert(std::move(input_chunks)); // And we end with two 2 device allocations. - EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2); + EXPECT_EQ(br.get_main_record().num_current_allocs(), 2); // Let's force spilling. br.set_memory_limit(rapidsmpf::MemoryType::DEVICE, k_force_spill_limit); @@ -459,24 +457,24 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { std::vector output_chunks = rapidsmpf::unspill_partitions( shuffler.extract(0), &br, rapidsmpf::AllowOverbooking::YES ); - EXPECT_EQ(mr.get_main_record().num_current_allocs(), 1); + EXPECT_EQ(br.get_main_record().num_current_allocs(), 1); // And insert also triggers spilling. We end up with zero device allocations. std::unordered_map chunk; chunk.emplace(0, std::move(output_chunks.at(0))); shuffler.insert(std::move(chunk)); - EXPECT_EQ(mr.get_main_record().num_current_allocs(), 0); + EXPECT_EQ(br.get_main_record().num_current_allocs(), 0); } // Extract and unspill both partitions. std::vector out0 = rapidsmpf::unspill_partitions( shuffler.extract(0), &br, rapidsmpf::AllowOverbooking::YES ); - EXPECT_EQ(mr.get_main_record().num_current_allocs(), 1); + EXPECT_EQ(br.get_main_record().num_current_allocs(), 1); std::vector out1 = rapidsmpf::unspill_partitions( shuffler.extract(1), &br, rapidsmpf::AllowOverbooking::YES ); - EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2); + EXPECT_EQ(br.get_main_record().num_current_allocs(), 2); // Disable spilling and insert the first partition. br.set_memory_limit(rapidsmpf::MemoryType::DEVICE, k_no_spill_limit); @@ -485,7 +483,7 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { chunk.emplace(0, std::move(out0.at(0))); shuffler.insert(std::move(chunk)); } - EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2); + EXPECT_EQ(br.get_main_record().num_current_allocs(), 2); // Enable spilling and insert the second partition, which should trigger spilling // of both the first partition already in the shuffler and the second partition @@ -496,7 +494,7 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { chunk.emplace(1, std::move(out1.at(0))); shuffler.insert(std::move(chunk)); } - EXPECT_EQ(mr.get_main_record().num_current_allocs(), 0); + EXPECT_EQ(br.get_main_record().num_current_allocs(), 0); shuffler.insert_finished(); } diff --git a/cpp/tests/test_statistics.cpp b/cpp/tests/test_statistics.cpp index 5b2ecc2f1..3420f6e0c 100644 --- a/cpp/tests/test_statistics.cpp +++ b/cpp/tests/test_statistics.cpp @@ -15,7 +15,7 @@ #include #include -#include +#include #include #include @@ -178,7 +178,9 @@ TEST_F(StatisticsTest, ReportSorting) { } TEST_F(StatisticsTest, MemoryProfiler) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; auto pinned_mr = rapidsmpf::PinnedMemoryResource::make_if_available(); auto stats = rapidsmpf::Statistics::create(); auto stream = cudf::get_default_stream(); @@ -241,7 +243,7 @@ TEST_F(StatisticsTest, MemoryProfiler) { std::istringstream ss(report); std::string line; while (std::getline(ss, line) && (main_line.empty() || pinned_line.empty())) { - if (line.find("main (all allocations using RmmResourceAdaptor)") + if (line.find("main (all allocations using BufferResource)") != std::string::npos) { main_line = line; @@ -262,7 +264,7 @@ TEST_F(StatisticsTest, MemoryProfiler) { // For the main record: num_calls=1, peak=2 MiB, g-peak=2 MiB, accum=4 MiB. static constexpr std::string_view kExpectedMainLine = " 1 2 MiB 2 MiB 4 MiB 1 MiB" - " main (all allocations using RmmResourceAdaptor)"; + " main (all allocations using BufferResource)"; EXPECT_EQ(main_line, kExpectedMainLine); static const std::string_view kExpectedPinnedLine = pinned_mr == PinnedMemoryResource::Disabled @@ -273,14 +275,16 @@ TEST_F(StatisticsTest, MemoryProfiler) { } TEST_F(StatisticsTest, MemoryProfilerDisabled) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; auto stats = rapidsmpf::Statistics::disabled(); { auto const& records = stats->get_memory_records(); EXPECT_TRUE(records.empty()); } - // Outer scope — disabled stats make the recorder a no-op even when an - // `RmmResourceAdaptor` is provided. + // Outer scope — disabled stats make the recorder a no-op even when a + // `BufferResource` is provided. { auto outer = stats->create_memory_recorder(mr, "outer"); void* ptr1 = mr.allocate_sync(1_MiB); // +1 MiB @@ -304,12 +308,14 @@ TEST_F(StatisticsTest, MemoryProfilerDisabled) { // // Invariants checked: // 1. The toggled-off recorder publishes no entry. -// 2. The recorder still pops its scope so the `RmmResourceAdaptor`'s +// 2. The recorder still pops its scope so the `BufferResource`'s // per-thread record stack is balanced after the scope exits. Pre-fix, // the dtor early-returned and the frame stayed on the stack. // 3. A follow-up recorder works correctly against the balanced stack. TEST_F(StatisticsTest, MemoryProfilerToggledMidScope) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; auto stats = rapidsmpf::Statistics::create(); { @@ -335,7 +341,9 @@ TEST_F(StatisticsTest, MemoryProfilerToggledMidScope) { } TEST_F(StatisticsTest, MemoryProfilerMacro) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; auto stats = rapidsmpf::Statistics::create(); { RAPIDSMPF_MEMORY_PROFILE(stats, mr); @@ -350,7 +358,9 @@ TEST_F(StatisticsTest, MemoryProfilerMacro) { } TEST_F(StatisticsTest, MemoryProfilerMacroDisabled) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; auto stats = rapidsmpf::Statistics::disabled(); { RAPIDSMPF_MEMORY_PROFILE(stats, mr); @@ -390,7 +400,9 @@ TEST_F(StatisticsTest, InvalidStatNames) { } TEST_F(StatisticsTest, InvalidMemoryRecordNames) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; auto stats = rapidsmpf::Statistics::create(); std::ignore = stats->create_memory_recorder(mr, "bad\"name"); std::ostringstream ss; @@ -398,7 +410,9 @@ TEST_F(StatisticsTest, InvalidMemoryRecordNames) { } TEST_F(StatisticsTest, JsonMemoryRecords) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::BufferResource mr{cuda::mr::any_resource{ + cudf::get_current_device_resource_ref() + }}; auto stats = rapidsmpf::Statistics::create(); { auto rec = stats->create_memory_recorder(mr, "alloc"); diff --git a/python/rapidsmpf/rapidsmpf/CMakeLists.txt b/python/rapidsmpf/rapidsmpf/CMakeLists.txt index 841570bad..b92bba8bf 100644 --- a/python/rapidsmpf/rapidsmpf/CMakeLists.txt +++ b/python/rapidsmpf/rapidsmpf/CMakeLists.txt @@ -5,8 +5,8 @@ # cmake-format: on # ================================================================================= -set(cython_modules config.pyx cuda_stream.pyx error.pyx progress_thread.pyx - rmm_resource_adaptor.pyx shuffler.pyx statistics.pyx +set(cython_modules config.pyx cuda_stream.pyx error.pyx progress_thread.pyx shuffler.pyx + statistics.pyx ) # Add cupti module conditionally if CUPTI support is enabled diff --git a/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py b/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py index c57af8cce..501ef8ba2 100644 --- a/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py +++ b/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py @@ -24,7 +24,6 @@ from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.progress_thread import ProgressThread -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.shuffler import Shuffler from rapidsmpf.statistics import Statistics from rapidsmpf.utils.string import format_bytes, parse_bytes @@ -230,13 +229,11 @@ def setup_and_run(args: argparse.Namespace) -> None: """ options = Options(get_environment_variables()) - # Create a RMM stack with both a device pool and statistics. - mr = RmmResourceAdaptor( - rmm.mr.PoolMemoryResource( - rmm.mr.CudaMemoryResource(), - initial_pool_size=args.rmm_pool_size, - maximum_pool_size=args.rmm_pool_size, - ) + # Create a RMM device pool (BufferResource provides tracking on top). + mr = rmm.mr.PoolMemoryResource( + rmm.mr.CudaMemoryResource(), + initial_pool_size=args.rmm_pool_size, + maximum_pool_size=args.rmm_pool_size, ) rmm.mr.set_current_device_resource(mr) @@ -287,7 +284,7 @@ def setup_and_run(args: argparse.Namespace) -> None: ) if args.statistics: - comm.logger.print(stats.report(mr=mr)) + comm.logger.print(stats.report(br=br)) def parse_args( diff --git a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py index 149217c05..48f636eca 100644 --- a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py +++ b/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py @@ -27,7 +27,6 @@ from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.progress_thread import ProgressThread -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.shuffler import Shuffler from rapidsmpf.statistics import Statistics from rapidsmpf.utils.string import format_bytes, parse_bytes @@ -301,13 +300,11 @@ def setup_and_run(args: argparse.Namespace) -> None: """ options = Options(get_environment_variables()) - # Create a RMM stack with both a device pool and statistics. - mr = RmmResourceAdaptor( - rmm.mr.PoolMemoryResource( - rmm.mr.CudaMemoryResource(), - initial_pool_size=args.rmm_pool_size, - maximum_pool_size=args.rmm_pool_size, - ) + # Create a RMM device pool (BufferResource provides tracking on top). + mr = rmm.mr.PoolMemoryResource( + rmm.mr.CudaMemoryResource(), + initial_pool_size=args.rmm_pool_size, + maximum_pool_size=args.rmm_pool_size, ) rmm.mr.set_current_device_resource(mr) @@ -411,7 +408,7 @@ def setup_and_run(args: argparse.Namespace) -> None: f"elapsed: {elapsed_time:.2f} sec | rmm device memory peak: {mem_peak}" ) if stats.enabled: - comm.logger.print(stats.report(mr=mr)) + comm.logger.print(stats.report(br=br)) def dir_path(path: str) -> Path: diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py index dcdf5f3fd..56cf12113 100644 --- a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py +++ b/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py @@ -24,7 +24,6 @@ from rapidsmpf.integrations.ray import RapidsMPFActor, setup_ray_ucxx_cluster from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.shuffler import Shuffler from rapidsmpf.statistics import Statistics from rapidsmpf.utils.cudf import pylibcudf_to_cudf_dataframe @@ -79,12 +78,10 @@ def __init__( self.spill_device = spill_device # Initialize actor-local resources (statistics, memory resource) - self.mr = RmmResourceAdaptor( - rmm.mr.PoolMemoryResource( - rmm.mr.CudaMemoryResource(), - initial_pool_size=self.rmm_pool_size, - maximum_pool_size=self.rmm_pool_size, - ) + self.mr = rmm.mr.PoolMemoryResource( + rmm.mr.CudaMemoryResource(), + initial_pool_size=self.rmm_pool_size, + maximum_pool_size=self.rmm_pool_size, ) rmm.mr.set_current_device_resource(self.mr) # Create a buffer resource that limits device memory if `--spill-device` diff --git a/python/rapidsmpf/rapidsmpf/integrations/core.py b/python/rapidsmpf/rapidsmpf/integrations/core.py new file mode 100644 index 000000000..f809326e3 --- /dev/null +++ b/python/rapidsmpf/rapidsmpf/integrations/core.py @@ -0,0 +1,791 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Shuffler integration with external libraries.""" + +from __future__ import annotations + +import threading +import weakref +from contextlib import suppress +from dataclasses import dataclass, field +from functools import cached_property, partial +from typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, Protocol, TypeVar + +import rmm.mr +from rmm.pylibrmm.stream import DEFAULT_STREAM + +from rapidsmpf.config import ( + OptionalBytes, + Options, +) +from rapidsmpf.memory.buffer_resource import ( + BufferResource, +) +from rapidsmpf.memory.spill_collection import SpillCollection +from rapidsmpf.shuffler import Shuffler + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.statistics import Statistics + + +DataFrameT = TypeVar("DataFrameT") + + +# Set of available shuffle IDs +_shuffle_id_vacancy: set[int] = set(range(Shuffler.max_concurrent_shuffles)) +_shuffle_id_vacancy_lock: threading.Lock = threading.Lock() + + +def get_new_shuffle_id(get_occupied_ids: Callable[[], Sequence[set[int]]]) -> int: + """ + Get a new available shuffle ID. + + Since RapidsMPF only supports a limited number of shuffler instances at + any given time, this function maintains a shared pool of shuffle IDs. + + If no IDs are available locally, it calls get_occupied_ids to query all + workers for IDs in use, updates the vacancy set accordingly, and retries. + If all IDs are in use across all workers, an error is raised. + + Parameters + ---------- + get_occupied_ids + Callable function that returns the occupied shuffle IDs. + + Returns + ------- + A unique shuffle ID not currently in use. + + Raises + ------ + ValueError + If all shuffle IDs are currently in use. + """ + global _shuffle_id_vacancy # noqa: PLW0603 + + with _shuffle_id_vacancy_lock: + if not _shuffle_id_vacancy: + # We start with setting all IDs as vacant and then subtract all + # IDs occupied on any one worker. + _shuffle_id_vacancy = set(range(Shuffler.max_concurrent_shuffles)) + _shuffle_id_vacancy.difference_update(*get_occupied_ids()) + if not _shuffle_id_vacancy: + raise ValueError( + f"Cannot shuffle more than {Shuffler.max_concurrent_shuffles} " + "times in a single query." + ) + + return _shuffle_id_vacancy.pop() + + +@dataclass +class WorkerContext: + """ + RapidsMPF specific attributes for a worker. + + Attributes + ---------- + lock + The global worker lock. Must be acquired before accessing attributes + that might be modified while the worker is running such as the shufflers. + br + The buffer resource used by the worker exclusively. + statistics + The statistics used by the worker. If None, statistics is disabled. + comm + The communicator connected to all other workers. + spill_collection + A collection of Python objects that can be spilled to free up device memory. + shufflers + A mapping from shuffler IDs to active shuffler instances. + options + Configuration options. + python_object_spill_function_id + ID from ``SpillManager.add_spill_function`` for ``spill_func``; cleared by + ``unregister_python_spill_callback``. + """ + + lock: ClassVar[threading.RLock] = threading.RLock() + br: BufferResource + statistics: Statistics + comm: Communicator | None = None + spill_collection: SpillCollection = field(default_factory=SpillCollection) + shufflers: dict[int, Shuffler] = field(default_factory=dict) + options: Options = field(default_factory=Options) + python_object_spill_function_id: int | None = field(default=None, init=False) + + def unregister_python_spill_callback(self) -> None: + """ + Remove the Python-object spill callback from the buffer resource. + + Safe to call more than once. Call this from integration teardown so + the C++ periodic spill thread cannot invoke ``spill_func`` during + interpreter shutdown, when attribute access on this object may be + unreliable. + """ + fid = self.python_object_spill_function_id + if fid is None: + return + with suppress(Exception): + self.br.spill_manager.remove_spill_function(fid) + self.python_object_spill_function_id = None + + def get_statistics(self) -> dict[str, dict[str, int | float]]: + """ + Get the statistics from the worker context. + + Returns + ------- + statistics + A dictionary of statistics. The keys are the names of the statistics. + The values are dictionaries with two keys: + + - "count" is the number of times the statistic was recorded. + - "value" is the value of the statistic. + + Notes + ----- + Statistics are global across all shuffles. To measure statistics for any + given shuffle, gather statistics before and after the shuffle and compute + the difference. + """ + return { + stat: self.statistics.get_stat(stat) + for stat in self.statistics.list_stat_names() + } + + +class ShufflerIntegration(Protocol[DataFrameT]): + """Shuffle-integration protocol.""" + + @staticmethod + def insert_partition( + df: DataFrameT, + partition_id: int, + partition_count: int, + shuffler: Shuffler, + options: Any, + *other: Any, + ) -> None: + """ + Add a partition to a RapidsMPF Shuffler. + + Parameters + ---------- + df + DataFrame partition to add to a RapidsMPF shuffler. + partition_id + The input partition id of ``df``. + partition_count + Number of output partitions for the current shuffle. + shuffler + The RapidsMPF Shuffler object to extract from. + options + Additional options. + *other + Other data needed for partitioning. For example, + this may be boundary values needed for sorting. + """ + ... + + @staticmethod + def extract_partition( + partition_id: int, + shuffler: Shuffler, + options: Any, + ) -> DataFrameT: + """ + Extract a DataFrame partition from a RapidsMPF Shuffler. + + Parameters + ---------- + partition_id + Partition id to extract. + shuffler + The RapidsMPF Shuffler object to extract from. + options + Additional options. + + Returns + ------- + A shuffled DataFrame partition. + """ + ... + + +def get_shuffler( + ctx: WorkerContext, + shuffle_id: int, + *, + partition_count: int | None = None, + worker: Any = None, +) -> Shuffler: + """ + Return the appropriate :class:`Shuffler` object. + + Parameters + ---------- + ctx + The worker context. + shuffle_id + Unique ID for the shuffle operation. + partition_count + Output partition count for the shuffle operation. + worker + The current worker. + + Returns + ------- + The active RapidsMPF :class:`Shuffler` object associated with + the specified ``shuffle_id``, ``partition_count`` and + ``worker``. + + Notes + ----- + Whenever a new :class:`Shuffler` object is created, it is + saved as ``WorkerContext.shufflers[shuffle_id]``. + """ + with ctx.lock: + if shuffle_id not in ctx.shufflers: + if partition_count is None: + raise ValueError( + "Need partition_count to create new shuffler." + f" shuffle_id: {shuffle_id}\n" + f" Shufflers: {ctx.shufflers}" + ) + assert ctx.br is not None + assert ctx.comm is not None + ctx.shufflers[shuffle_id] = Shuffler( + ctx.comm, + op_id=shuffle_id, + total_num_partitions=partition_count, + br=ctx.br, + ) + return ctx.shufflers[shuffle_id] + + +def insert_partition( + get_context: Callable[..., WorkerContext], + callback: Callable[ + [ + DataFrameT, + int, + int, + Shuffler, + Any, + *tuple[str | tuple[str, int], ...], + ], + None, + ], + df: DataFrameT, + partition_id: int, + partition_count: int, + shuffle_id: int, + options: Any, + *other_keys: str | tuple[str, int], +) -> None: + """ + Add a partition to a RapidsMPF Shuffler. + + Parameters + ---------- + get_context + Callable function to fetch the worker context. + callback + Insertion callback function. This function must be the + `insert_partition` attribute of a `ShufflerIntegration` + protocol. + df + DataFrame partition to add to a RapidsMPF shuffler. + partition_id + The input partition id of ``df``. + partition_count + Number of output partitions for the current shuffle. + shuffle_id + The RapidsMPF shuffle id. + options + Optional key-word arguments. + *other_keys + Other keys needed by ``callback``. + """ + callback( + df, + partition_id, + partition_count, + get_shuffler(get_context(), shuffle_id), + options, + *other_keys, + ) + + +def extract_partition( + get_context: Callable[..., WorkerContext], + callback: Callable[ + [int, Shuffler, Any], + DataFrameT, + ], + shuffle_id: int, + partition_id: int, + worker_barrier: tuple[int, ...], + options: Any, +) -> DataFrameT: + """ + Extract a partition from a RapidsMPF Shuffler. + + Parameters + ---------- + get_context + Callable function to fetch the worker context. + callback + Insertion callback function. This function must be the + `extract_partition` attribute of a `ShufflerIntegration` + protocol. + shuffle_id + The RapidsMPF shuffle id. + partition_id + Partition id to extract. + worker_barrier + Worker-barrier task dependency. This value should + not be used for compute logic. + options + Additional options. + + Returns + ------- + Extracted DataFrame partition. + """ + shuffler = get_shuffler(get_context(), shuffle_id) + try: + return callback( + partition_id, + shuffler, + options, + ) + finally: + if shuffler.finished(): + ctx = get_context() + with ctx.lock: + if shuffle_id in ctx.shufflers: + del ctx.shufflers[shuffle_id] + + +@dataclass +class BCastJoinInfo: # pragma: no cover; TODO: Cover in follow-up + """ + Broadcast join information. + + Parameters + ---------- + bcast_side + The side of the join being broadcasted. + bcast_count + The number of broadcasted partitions. + need_local_repartition + Whether to locally repartition on the broadcasted table. + This is not necessary for inner joins. + """ + + bcast_side: Literal["left", "right"] + bcast_count: int = 1 + need_local_repartition: bool = False + + +class JoinIntegration(Protocol[DataFrameT]): + """Join-integration protocol.""" + + @staticmethod + def get_shuffler_integration() -> ShufflerIntegration[DataFrameT]: + """Return the shuffler integration.""" + ... + + @staticmethod + def join_partition( + left_input: Callable[[int], DataFrameT], + right_input: Callable[[int], DataFrameT], + bcast_info: BCastJoinInfo | None, + options: Any, + ) -> DataFrameT: + """ + Produce a joined table partition. + + Parameters + ---------- + left_input + A callable that produces chunks of the left partition. + The ``bcast_info.bcast_count`` parameter corresponds + to the number of chunks the callable can produce. + right_input + A callable that produces chunks of the right partition. + The ``bcast_info.bcast_count`` parameter corresponds + to the number of chunks the callable can produce. + bcast_info + The broadcast join information. This should be None + for a regular hash join. + options + Additional join options. + + Returns + ------- + A joined DataFrame partition. + """ + ... + + +class FetchJoinChunk(Generic[DataFrameT]): + """ + Fetch the data for one side of a join operation. + + Parameters + ---------- + side + The side of the join being fetched. + output_partition_id + The output partition id for the join operation. + get_context + Callable function to fetch the worker context. + integration + The JoinIntegration protocol to use. + op_id + The operation id. + barrier + The barrier to fetch the partition from. + bcast_info + The broadcast join information. + n_worker_tasks + The number of join_partition tasks to be called on this worker. + options + Additional options. + + Notes + ----- + A ``FetchJoinChunk`` object only fetches data needed for a single + output partition. For in-memory or shuffled data, there will only be + one chunk to return. For broadcast joins, there may be multiple chunks. + """ + + def __init__( + self, + side: Literal["left", "right"], + output_partition_id: int, + get_worker_context: Callable[..., WorkerContext], + integration: JoinIntegration[DataFrameT], + op_id: int | None, + barrier: DataFrameT | tuple[int, ...], + bcast_info: BCastJoinInfo | None, + n_worker_tasks: int, + options: Any, + ): + if bcast_info is not None: # pragma: no cover + raise NotImplementedError("Broadcast join not yet supported.") + + self.side = side + self.output_partition_id = output_partition_id + self.get_worker_context = get_worker_context + self.integration = integration + self.op_id = op_id + self.barrier = barrier + self.bcast_info = bcast_info + self.n_worker_tasks = n_worker_tasks + self.options = options + + @cached_property + def _data(self) -> dict[int, DataFrameT]: + """Return a dictionary of DataFrame chunks.""" + op_id = self.op_id + data: DataFrameT + if op_id is None: + assert not isinstance(self.barrier, tuple), "Barrier must be a DataFrame." + data = self.barrier + else: + ctx = self.get_worker_context() + shuffler = get_shuffler(ctx, op_id) + try: + data = self.integration.get_shuffler_integration().extract_partition( + self.output_partition_id, + shuffler, + self.options, + ) + finally: + if shuffler.finished(): + with ctx.lock: + if op_id in ctx.shufflers: + del ctx.shufflers[op_id] + return {0: data} + + def __call__(self, chunk_id: int) -> Any: + """ + Return the DataFrame associated with the given chunk id. + + Parameters + ---------- + chunk_id + The id of the local chunk to fetch for a join operation. + There will only be one chunk to return for a hash join. + There may be multiple chunks to return for a broadcast join. + + Returns + ------- + A DataFrame chunk to be used in a join operation. + """ + if self.bcast_info is None: + # Fetch a chunk of a non-broadcasted partition. + # The partition_id is ignored, because we only have a single chunk. + return self._data[0] + else: # pragma: no cover + # Fetch a chunk of the broadcasted partition. + raise NotImplementedError("Broadcast join not implemented.") + + +def join_partition( + get_context: Callable[..., WorkerContext], + integration: JoinIntegration[DataFrameT], + bcast_info: BCastJoinInfo | None, + left_op_id: int | None, + right_op_id: int | None, + left_dependency: DataFrameT | tuple[int, ...], + right_dependency: DataFrameT | tuple[int, ...], + part_id: int, + n_worker_tasks: int, + left_options: Any, + right_options: Any, + join_options: Any, +) -> DataFrameT: + """ + Produce a joined table partition. + + Parameters + ---------- + get_context + Callable function to fetch the worker context. + integration + The JoinIntegration protocol to use. + bcast_info + The broadcast join information. + This should be None for a regular hash join. + left_op_id + The left-table operation id. The operation may correspond + to an allgather or a shuffle operation. If None, the + left_dependency argument must be the left partition. + right_op_id + The right-table operation id. The operation may correspond + to an allgather or a shuffle operation. If None, the + right_dependency argument must be the right partition. + left_dependency + Task dependency for the left table. If left_op_id is None, + this will correspond to the left partition. Otherwise, this + argument is only used to enforce task ordering. The left_op_id + argument should be used to fetch the real left partition. + right_dependency + Task dependency for the right table. If right_op_id is None, + this will correspond to the right partition. Otherwise, this + argument is only used to enforce task ordering. The right_op_id + argument should be used to fetch the real right partition. + part_id + The output partition id. + This information is needed to extract shuffled partitions. + n_worker_tasks + The number of join_partition tasks to be called on this worker. + This information may be used for cleanup. + left_options + Additional options for extracting the left table. + right_options + Additional options for extracting the right table. + join_options + Additional options for the join. + + Returns + ------- + A joined DataFrame partition. + """ + + def _get_input(side: Literal["left", "right"]) -> FetchJoinChunk: + """Return the input for one side of the join.""" + if side == "left": + op_id = left_op_id + barrier = left_dependency + options = left_options + elif side == "right": + op_id = right_op_id + barrier = right_dependency + options = right_options + else: + raise ValueError(f"Invalid side: {side}") + + return FetchJoinChunk( + side, + part_id, + get_context, + integration, + op_id, + barrier, + bcast_info, + n_worker_tasks, + options, + ) + + return integration.join_partition( + _get_input("left"), + _get_input("right"), + bcast_info, + join_options, + ) + + +# Create a spill function that spills the python objects in the spill- +# collection. This way, we have a central place (the worker) to track +# and trigger spilling of python objects. +def spill_func( + amount: int, + *, + staging_buffer: rmm.DeviceBuffer | None, + lock: threading.Lock, + mr: rmm.mr.DeviceMemoryResource, + ctx: WorkerContext, +) -> int: + """ + Spill a specified amount of data from the Python object spill collection. + + This function attempts to use a preallocated staging device buffer to + spill Python objects from the spill collection. If the staging buffer + is currently in use, it will fall back to spilling without it. + + Parameters + ---------- + amount + The amount of data to spill, in bytes. + staging_buffer + Optional buffer to stage data through. + lock + Lock to protect access to the staging buffer. + mr + Memory resource for device allocations. + ctx + The worker context to spill from. + + Returns + ------- + The actual amount of data spilled, in bytes. + """ + spill_collection: SpillCollection | None = getattr(ctx, "spill_collection", None) + if spill_collection is None: + return 0 + if staging_buffer is not None and lock.acquire(blocking=False): + try: + return spill_collection.spill( + amount, + stream=DEFAULT_STREAM, + device_mr=mr, + staging_device_buffer=staging_buffer, + ) + finally: + lock.release() + return spill_collection.spill(amount, stream=DEFAULT_STREAM, device_mr=mr) + + +def rmpf_worker_local_setup( + worker: Any, + option_prefix: str, + *, + options: Options, +) -> WorkerContext: + """ + Create per-worker local RapidsMPF attributes on a remote worker. + + After creating the local context, a communicator must be bootstrapped. + + Parameters + ---------- + worker + The current worker process. + option_prefix + Prefix for config-option names. + options + Configuration options. + + Returns + ------- + WorkerContext + New local worker context + """ + # Use the current device resource directly; BufferResource provides + # tracking on top. + mr = rmm.mr.get_current_device_resource() + + options_map = options.get_strings() + # Map prefixed integration keys to internal RapidsMPF option names. + for suffix, rmpf_key in ( + ("statistics", "statistics"), + ("spill_to_pinned_memory", "pinned_memory"), + ("periodic_spill_check", "periodic_spill_check"), + ): + custom_key = f"{option_prefix}{suffix}" + if custom_key in options_map: + options_map[rmpf_key] = options_map.pop(custom_key) + + # Convert spill_device (legacy float fraction, e.g. "0.5") to the + # spill_device_limit if spill_device is set. + spill_device_key = f"{option_prefix}spill_device" + if spill_device_key in options_map: + val = float(options_map.pop(spill_device_key)) + total_memory = rmm.mr.available_device_memory()[1] + options_map["spill_device_limit"] = f"{int(total_memory * val)}" + + # overwrite the options with the new options map + options = Options(options_map) + + # use options to create the buffer resource + br = BufferResource.from_options(mr, options) + statistics = br.statistics + + # If enabled, create a staging device buffer for the spilling to reduce + # device memory pressure. + # TODO: maybe have a pool of staging buffers? + spill_staging_buffer_size = options.get_or_default( + f"{option_prefix}staging_spill_buffer", + default_value=OptionalBytes("128 MiB"), + ).value + spill_staging_buffer = ( + None + if spill_staging_buffer_size is None + else rmm.DeviceBuffer( + size=spill_staging_buffer_size, stream=DEFAULT_STREAM, mr=mr + ) + ) + + ctx = WorkerContext( + br=br, + statistics=statistics, + options=options, + ) + + if ( + options.get_or_default(f"{option_prefix}print_statistics", default_value=True) + and statistics.enabled + ): + ref, name = ( + (ctx, "rapidsmpf_worker_ctx") if worker is None else (worker, str(worker)) + ) + weakref.finalize( + ref, + lambda name, stats: print(name, stats.report()), + name=name, + stats=statistics, + ) + + # Add the spill function using a negative priority (-10) such that spilling + # of internal shuffle buffers (non-python objects) have higher priority than + # spilling of the Python objects in the collection. + ctx.python_object_spill_function_id = br.spill_manager.add_spill_function( + func=partial( + spill_func, + staging_buffer=spill_staging_buffer, + lock=threading.Lock(), + mr=mr, + ctx=ctx, + ), + priority=-10, + ) + return ctx diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd index 29fefa2a6..2c6b7690b 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd @@ -33,8 +33,9 @@ cdef extern from "" nogil: const cuda_stream_pool &stream_pool() except +ex_handler size_t release(cpp_MemoryReservation&, size_t) except +ex_handler shared_ptr[cpp_Statistics] statistics() except +ex_handler + int64_t current_allocated() except +ex_handler -cdef class BufferResource: +cdef class BufferResource(DeviceMemoryResource): cdef object __weakref__ cdef shared_ptr[cpp_BufferResource] _handle cdef readonly SpillManager spill_manager diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi index ce21b941f..d8e6762c6 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi @@ -14,7 +14,7 @@ from rapidsmpf.memory.pinned_memory_resource import PinnedMemoryResource from rapidsmpf.memory.spill_manager import SpillManager from rapidsmpf.statistics import Statistics -class BufferResource: +class BufferResource(DeviceMemoryResource): def __init__( self, device_mr: DeviceMemoryResource, @@ -33,9 +33,13 @@ class BufferResource: statistics: Statistics | None = None, ) -> Self: ... @property - def device_mr(self) -> DeviceMemoryResource: ... + def device_mr(self) -> BufferResource: ... + @property + def primary_mr(self) -> DeviceMemoryResource: ... @property def pinned_mr(self) -> PinnedMemoryResource | None: ... + @property + def current_allocated(self) -> int: ... def memory_available(self, mem_type: MemoryType) -> int: ... def memory_reserved(self, mem_type: MemoryType) -> int: ... def set_memory_limit(self, mem_type: MemoryType, limit: int) -> None: ... diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index 682d008e9..33b5a7bc2 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -15,7 +15,8 @@ from rmm.librmm.cuda_stream_pool cimport cuda_stream_pool from rmm.pylibrmm import CudaStreamFlags -from rmm.librmm.memory_resource cimport make_any_device_resource +from rmm.librmm.memory_resource cimport (device_async_resource_ref, + make_any_device_resource) from rmm.pylibrmm.cuda_stream_pool cimport CudaStreamPool from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource @@ -43,6 +44,25 @@ cdef extern from *: ) except +ex_handler +# Expose `make_device_async_resource_ref(BufferResource&)` to Cython so the +# Python wrapper can set `self.c_ref` to the C++ `BufferResource`. The rmm +# template is C++-generic; here we just instantiate it for `cpp_BufferResource`. +cdef extern from * nogil: + """ + #include + #include + #include + inline std::optional + cpp_make_device_async_resource_ref_for_br(rapidsmpf::BufferResource& br) { + return std::optional( + rmm::device_async_resource_ref(br)); + } + """ + optional[device_async_resource_ref] cpp_make_device_async_resource_ref_for_br( + cpp_BufferResource& + ) except +ex_handler + + cdef extern from * nogil: """ namespace { @@ -103,7 +123,7 @@ cdef extern from * nogil: @no_gc_clear -cdef class BufferResource: +cdef class BufferResource(DeviceMemoryResource): """ Class managing buffer resources. @@ -111,6 +131,12 @@ cdef class BufferResource: (e.g., host and device). All memory operations in RapidsMPF, such as those performed by the Shuffler, rely on a buffer resource for memory management. + BufferResource subclasses :class:`rmm.pylibrmm.DeviceMemoryResource`, so an + instance can be passed directly anywhere an RMM device memory resource is + expected (e.g. as the ``mr`` argument to ``rmm.DeviceBuffer``). Buffers + allocated through it hold an owning ref to the resource, which transitively + keeps the underlying stream pool alive. + Parameters ---------- device_mr @@ -205,6 +231,11 @@ cdef class BufferResource: cpp_stream_pool, stats_handle, ) + # Expose this BufferResource as a CCCL-conformant resource through the + # inherited rmm `DeviceMemoryResource` slot. Means `rmm.DeviceBuffer(mr=br)` + # works and the DeviceBuffer holds a ref to this Python BufferResource, + # transitively keeping the C++ state alive. + self.c_ref = cpp_make_device_async_resource_ref_for_br(deref(self._handle)) self.spill_manager = SpillManager._create(self) @classmethod @@ -274,7 +305,20 @@ cdef class BufferResource: @property def device_mr(self): """ - The memory resource used for device memory allocations. + Back-compat accessor; returns ``self``. + + The Python `BufferResource` *is* a CCCL-conformant RMM + `DeviceMemoryResource`, so the natural way to pass it where an RMM + MR is expected is the object itself (``rmm.DeviceBuffer(mr=br)``). + This property returns ``self`` so that code that still does + ``rmm.DeviceBuffer(mr=br.device_mr)`` keeps working. + """ + return self + + @property + def primary_mr(self): + """ + The primary memory resource passed to the constructor. Returns ------- @@ -294,6 +338,21 @@ cdef class BufferResource: """ return self._pinned_mr + @property + def current_allocated(self): + """ + Total number of device bytes currently allocated through this BufferResource. + + Returns + ------- + int + Currently outstanding allocated bytes. + """ + cdef int64_t ret + with nogil: + ret = deref(self._handle).current_allocated() + return ret + def memory_reserved(self, MemoryType mem_type): """ Get the current reserved memory of the specified memory type. diff --git a/python/rapidsmpf/rapidsmpf/rmm_resource_adaptor.pxd b/python/rapidsmpf/rapidsmpf/rmm_resource_adaptor.pxd deleted file mode 100644 index 5990a225a..000000000 --- a/python/rapidsmpf/rapidsmpf/rmm_resource_adaptor.pxd +++ /dev/null @@ -1,35 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from libc.stdint cimport uint64_t -from libcpp.memory cimport unique_ptr -from libcpp.optional cimport optional -from rmm.librmm.memory_resource cimport (any_resource, device_accessible, - device_async_resource_ref, - make_device_async_resource_ref) -from rmm.pylibrmm.memory_resource cimport UpstreamResourceAdaptor - -from rapidsmpf._detail.exception_handling cimport ex_handler -from rapidsmpf.memory.scoped_memory_record cimport cpp_ScopedMemoryRecord - - -cdef extern from "" nogil: - cdef cppclass cpp_RmmResourceAdaptor"rapidsmpf::RmmResourceAdaptor": - cpp_RmmResourceAdaptor( - any_resource[device_accessible] primary_mr, - ) except +ex_handler - - cpp_ScopedMemoryRecord get_main_record() except +ex_handler - uint64_t current_allocated() noexcept - - -# The make_device_async_resource_ref C++ template (declared in RMM's -# memory_resource.pxd) also covers RmmResourceAdaptor; declare the overload. -cdef extern from *: - optional[device_async_resource_ref] make_device_async_resource_ref( - cpp_RmmResourceAdaptor&) except + - - -cdef class RmmResourceAdaptor(UpstreamResourceAdaptor): - cdef unique_ptr[cpp_RmmResourceAdaptor] c_obj - cdef cpp_RmmResourceAdaptor* get_handle(self) diff --git a/python/rapidsmpf/rapidsmpf/rmm_resource_adaptor.pyi b/python/rapidsmpf/rapidsmpf/rmm_resource_adaptor.pyi deleted file mode 100644 index 06553445c..000000000 --- a/python/rapidsmpf/rapidsmpf/rmm_resource_adaptor.pyi +++ /dev/null @@ -1,20 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from rmm.pylibrmm.memory_resource import DeviceMemoryResource -from rmm.pylibrmm.stream import Stream - -from rapidsmpf.memory.scoped_memory_record import ScopedMemoryRecord - -class RmmResourceAdaptor(DeviceMemoryResource): - def __init__( - self, - upstream_mr: DeviceMemoryResource, - ): ... - @property - def get_upstream(self) -> DeviceMemoryResource: ... - def allocate(self, nbytes: int, stream: Stream = ...) -> int: ... - def deallocate(self, ptr: int, nbytes: int, stream: Stream = ...) -> None: ... - def get_main_record(self) -> ScopedMemoryRecord: ... - @property - def current_allocated(self) -> int: ... diff --git a/python/rapidsmpf/rapidsmpf/rmm_resource_adaptor.pyx b/python/rapidsmpf/rapidsmpf/rmm_resource_adaptor.pyx deleted file mode 100644 index 69dbc801a..000000000 --- a/python/rapidsmpf/rapidsmpf/rmm_resource_adaptor.pyx +++ /dev/null @@ -1,73 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from cython.operator cimport dereference as deref -from libc.stdint cimport uint64_t -from rmm.librmm.memory_resource cimport make_any_device_resource -from rmm.pylibrmm.memory_resource cimport (DeviceMemoryResource, - UpstreamResourceAdaptor) - -from rapidsmpf.memory.scoped_memory_record cimport ScopedMemoryRecord - - -cdef class RmmResourceAdaptor(UpstreamResourceAdaptor): - """A RMM memory resource adaptor tailored to RapidsMPF.""" - def __cinit__( - self, - DeviceMemoryResource upstream_mr, - ): - """ - A RMM memory resource adaptor tailored to RapidsMPF. - - Wraps a primary device memory resource and adds memory usage tracking - (lifetime stats plus per-thread scoped records). - - Parameters - ---------- - upstream_mr - The primary device memory resource used for allocations and deallocations. - """ - self.c_obj.reset( - new cpp_RmmResourceAdaptor( - make_any_device_resource(upstream_mr.get_mr()) - ) - ) - self.c_ref = make_device_async_resource_ref(deref(self.c_obj)) - - def __dealloc__(self): - with nogil: - self.c_obj.reset() - - cdef cpp_RmmResourceAdaptor* get_handle(self): - return self.c_obj.get() - - def get_main_record(self): - """Returns a copy of the main memory record. - - The main record tracks memory statistics for the lifetime of the resource. - - Returns - ------- - A copy of the current main memory record. - """ - cdef cpp_RmmResourceAdaptor* mr = self.get_handle() - cdef cpp_ScopedMemoryRecord ret - with nogil: - ret = deref(mr).get_main_record() - return ScopedMemoryRecord.from_handle(ret) - - @property - def current_allocated(self) -> int: - """Get the total number of currently allocated bytes. - - This includes both allocations on the primary and fallback memory resources. - - Returns - ------- - Total number of currently allocated bytes. - """ - cdef cpp_RmmResourceAdaptor* mr = self.get_handle() - cdef uint64_t ret - with nogil: - ret = deref(mr).current_allocated() - return ret diff --git a/python/rapidsmpf/rapidsmpf/statistics.pxd b/python/rapidsmpf/rapidsmpf/statistics.pxd index 134620f2a..6a03704c7 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pxd +++ b/python/rapidsmpf/rapidsmpf/statistics.pxd @@ -11,11 +11,11 @@ from libcpp.unordered_map cimport unordered_map from libcpp.vector cimport vector from rapidsmpf._detail.exception_handling cimport ex_handler +from rapidsmpf.memory.buffer_resource cimport (BufferResource, + cpp_BufferResource) from rapidsmpf.memory.pinned_memory_resource cimport (PinnedMemoryResource, cpp_PinnedMemoryResource) from rapidsmpf.memory.scoped_memory_record cimport cpp_ScopedMemoryRecord -from rapidsmpf.rmm_resource_adaptor cimport (RmmResourceAdaptor, - cpp_RmmResourceAdaptor) cdef extern from "" nogil: @@ -52,7 +52,7 @@ cdef extern from "" nogil: cdef cppclass cpp_MemoryRecorder "rapidsmpf::Statistics::MemoryRecorder": cpp_MemoryRecorder( shared_ptr[cpp_Statistics] stats, - cpp_RmmResourceAdaptor mr, + shared_ptr[cpp_BufferResource] br, string name ) except +ex_handler @@ -63,5 +63,5 @@ cdef class Statistics: cdef class MemoryRecorder: cdef unique_ptr[cpp_MemoryRecorder] _handle cdef Statistics _stats - cdef RmmResourceAdaptor _mr + cdef BufferResource _br cdef string _name diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyi b/python/rapidsmpf/rapidsmpf/statistics.pyi index e569c5824..e2f562f2c 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pyi +++ b/python/rapidsmpf/rapidsmpf/statistics.pyi @@ -9,9 +9,9 @@ from os import PathLike from typing import Any, Self from rapidsmpf.config import Options +from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.memory.pinned_memory_resource import PinnedMemoryResource from rapidsmpf.memory.scoped_memory_record import ScopedMemoryRecord -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor class Formatter(IntEnum): Default = ... @@ -31,7 +31,7 @@ class Statistics: def report( self, *, - mr: RmmResourceAdaptor | None = None, + br: BufferResource | None = None, pinned_mr: PinnedMemoryResource | None = None, header: str | None = None, ) -> str: ... @@ -47,7 +47,7 @@ class Statistics: ) -> None: ... def get_memory_records(self) -> dict[str, MemoryRecord]: ... def memory_profiling( - self, mr: RmmResourceAdaptor | None, name: str + self, br: BufferResource | None, name: str ) -> MemoryRecorder: ... def copy(self) -> Statistics: ... @staticmethod diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyx b/python/rapidsmpf/rapidsmpf/statistics.pyx index 0d64d0ad8..a03e06368 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pyx +++ b/python/rapidsmpf/rapidsmpf/statistics.pyx @@ -17,11 +17,11 @@ from dataclasses import dataclass from rapidsmpf._detail.exception_handling cimport ex_handler from rapidsmpf.config cimport Options, cpp_Options +from rapidsmpf.memory.buffer_resource cimport (BufferResource, + cpp_BufferResource) from rapidsmpf.memory.pinned_memory_resource cimport (PinnedMemoryResource, cpp_PinnedMemoryResource) from rapidsmpf.memory.scoped_memory_record cimport ScopedMemoryRecord -from rapidsmpf.rmm_resource_adaptor cimport (RmmResourceAdaptor, - cpp_RmmResourceAdaptor) import os @@ -52,21 +52,21 @@ cdef extern from *: } std::string cpp_report( rapidsmpf::Statistics const& stats, - rapidsmpf::RmmResourceAdaptor* mr_ptr, + rapidsmpf::BufferResource* br_ptr, std::optional const& pinned_mr ) { - std::optional mr = - mr_ptr ? std::make_optional(*mr_ptr) : std::nullopt; + std::optional mr = + br_ptr ? std::make_optional(*br_ptr) : std::nullopt; return stats.report({.mr = std::move(mr), .pinned_mr = pinned_mr}); } std::string cpp_report( rapidsmpf::Statistics const& stats, - rapidsmpf::RmmResourceAdaptor* mr_ptr, + rapidsmpf::BufferResource* br_ptr, std::optional const& pinned_mr, std::string const& header ) { - std::optional mr = - mr_ptr ? std::make_optional(*mr_ptr) : std::nullopt; + std::optional mr = + br_ptr ? std::make_optional(*br_ptr) : std::nullopt; return stats.report({.mr = std::move(mr), .pinned_mr = pinned_mr, .header = header}); } @@ -124,12 +124,12 @@ cdef extern from *: shared_ptr[cpp_Statistics] cpp_disabled() except +ex_handler nogil string cpp_report( cpp_Statistics stats, - cpp_RmmResourceAdaptor* mr_ptr, + cpp_BufferResource* br_ptr, optional[cpp_PinnedMemoryResource] pinned_mr, ) except +ex_handler nogil string cpp_report( cpp_Statistics stats, - cpp_RmmResourceAdaptor* mr_ptr, + cpp_BufferResource* br_ptr, optional[cpp_PinnedMemoryResource] pinned_mr, string header, ) except +ex_handler nogil @@ -220,7 +220,7 @@ cdef class Statistics: def report( self, *, - RmmResourceAdaptor mr = None, + BufferResource br = None, PinnedMemoryResource pinned_mr = None, header = None, ): @@ -231,7 +231,7 @@ cdef class Statistics: Parameters ---------- - mr + br When provided, a memory profiling section is included in the report. When ``None``, the memory profiling section shows "Disabled". @@ -247,20 +247,20 @@ cdef class Statistics: A string representing the formatted statistics report. """ cdef string ret - cdef cpp_RmmResourceAdaptor* mr_ptr = NULL + cdef cpp_BufferResource* br_ptr = NULL cdef optional[cpp_PinnedMemoryResource] cpp_pinned cdef string cpp_header - if mr is not None: - mr_ptr = mr.get_handle() + if br is not None: + br_ptr = br.ptr() if pinned_mr is not None: cpp_pinned = pinned_mr._handle if header is None: with nogil: - ret = cpp_report(deref(self._handle), mr_ptr, cpp_pinned) + ret = cpp_report(deref(self._handle), br_ptr, cpp_pinned) else: cpp_header = header.encode() with nogil: - ret = cpp_report(deref(self._handle), mr_ptr, cpp_pinned, cpp_header) + ret = cpp_report(deref(self._handle), br_ptr, cpp_pinned, cpp_header) return ret.decode('UTF-8') def get_stat(self, name): @@ -439,12 +439,12 @@ cdef class Statistics: preincrement(it) return ret - def memory_profiling(self, RmmResourceAdaptor mr, name): + def memory_profiling(self, BufferResource br, name): """ Create a scoped memory profiling context for a named code region. Returns a context manager that tracks memory allocations and - deallocations made through the associated memory resource while + deallocations made through the associated `BufferResource` while the context is active. The profiling data is aggregated under the provided ``name`` and made available via :meth:`Statistics.get_memory_records()`. @@ -454,12 +454,12 @@ cdef class Statistics: - Global peak memory usage during the scope (``global_peak``) - Number of times the named scope was entered (``num_calls``) - Pass ``mr=None`` to get a no-op recorder. + Pass ``br=None`` to get a no-op recorder. Parameters ---------- - mr - The memory resource through which allocations are tracked. + br + The `BufferResource` through which allocations are tracked. Pass ``None`` to get a no-op recorder. name A unique identifier for the profiling scope. Used as a key @@ -472,12 +472,12 @@ cdef class Statistics: Examples -------- >>> import rmm - >>> mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) + >>> br = BufferResource(rmm.mr.CudaMemoryResource()) >>> stats = Statistics(enable=True) - >>> with stats.memory_profiling(mr, "outer"): - ... b1 = rmm.DeviceBuffer(size=1024, mr=mr) - ... with stats.memory_profiling(mr, "inner"): - ... b2 = rmm.DeviceBuffer(size=1024, mr=mr) + >>> with stats.memory_profiling(br, "outer"): + ... b1 = rmm.DeviceBuffer(size=1024, mr=br) + ... with stats.memory_profiling(br, "inner"): + ... b2 = rmm.DeviceBuffer(size=1024, mr=br) >>> inner = stats.get_memory_records()["inner"] >>> print(inner.scoped.peak()) 1024 @@ -485,7 +485,7 @@ cdef class Statistics: >>> print(outer.scoped.peak()) 2048 """ - return MemoryRecorder(self, mr, name) + return MemoryRecorder(self, br, name) def clear(self) -> None: """ @@ -685,30 +685,30 @@ cdef class MemoryRecorder: ---------- stats The statistics object responsible for aggregating memory profiling data. - mr - The memory resource through which allocations are tracked. + br + The `BufferResource` through which allocations are tracked. name The name of the profiling scope. Used as a key in the statistics record. """ def __cinit__( - self, Statistics stats not None, RmmResourceAdaptor mr not None, name + self, Statistics stats not None, BufferResource br not None, name ): self._stats = stats - self._mr = mr + self._br = br self._name = str.encode(name) def __enter__(self): - if self._mr is None: + if self._br is None: return - cdef cpp_RmmResourceAdaptor* mr = self._mr.get_handle() + cdef shared_ptr[cpp_BufferResource] br = self._br._handle with nogil: self._handle = make_unique[cpp_MemoryRecorder]( - self._stats._handle, deref(mr), self._name + self._stats._handle, br, self._name ) def __exit__(self, exc_type, exc_value, traceback): - if self._mr is not None: + if self._br is not None: with nogil: self._handle.reset() return False # do not suppress exceptions diff --git a/python/rapidsmpf/rapidsmpf/streaming/core/context.pyi b/python/rapidsmpf/rapidsmpf/streaming/core/context.pyi index 4ed3e4a65..27dfb310e 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/core/context.pyi +++ b/python/rapidsmpf/rapidsmpf/streaming/core/context.pyi @@ -5,13 +5,13 @@ from __future__ import annotations from typing import Any, Self +from rmm.pylibrmm.memory_resource import DeviceMemoryResource from rmm.pylibrmm.stream import Stream from rapidsmpf.communicator.communicator import Logger from rapidsmpf.config import Options from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.memory_reserve_or_wait import MemoryReserveOrWait @@ -29,7 +29,7 @@ class Context: def from_options( cls: type[Self], logger: Logger, - mr: RmmResourceAdaptor, + device_mr: DeviceMemoryResource, options: Options, statistics: Statistics | None = None, ) -> Self: ... diff --git a/python/rapidsmpf/rapidsmpf/streaming/core/context.pyx b/python/rapidsmpf/rapidsmpf/streaming/core/context.pyx index 512bd6f22..014e1c0e2 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/core/context.pyx +++ b/python/rapidsmpf/rapidsmpf/streaming/core/context.pyx @@ -12,9 +12,9 @@ from rapidsmpf.memory.buffer_resource cimport BufferResource from rapidsmpf.config import get_environment_variables from libcpp.memory cimport make_shared +from rmm.pylibrmm.memory_resource cimport DeviceMemoryResource from rmm.pylibrmm.stream cimport Stream -from rapidsmpf.rmm_resource_adaptor cimport RmmResourceAdaptor from rapidsmpf.streaming.core.channel cimport Channel, cpp_Channel from rapidsmpf.streaming.core.memory_reserve_or_wait cimport \ MemoryReserveOrWait @@ -97,7 +97,7 @@ cdef class Context: def from_options( cls, Logger logger not None, - RmmResourceAdaptor mr not None, + DeviceMemoryResource device_mr not None, Options options not None, statistics=None, ): @@ -105,7 +105,7 @@ cdef class Context: statistics = Statistics.disabled() return cls( logger=logger, - br=BufferResource.from_options(mr, options, statistics), + br=BufferResource.from_options(device_mr, options, statistics), options=options, ) diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/conftest.py b/python/rapidsmpf/rapidsmpf/tests/streaming/conftest.py index 879b73ca2..adba4063d 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/conftest.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/conftest.py @@ -10,7 +10,6 @@ from rapidsmpf.config import Options, get_environment_variables from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.streaming.core.context import Context if TYPE_CHECKING: @@ -25,8 +24,7 @@ def context(comm: Communicator) -> Generator[Context, None, None]: Fixture to get a streaming context. """ options = Options(get_environment_variables()) - mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) - br = BufferResource(mr) + br = BufferResource(rmm.mr.CudaMemoryResource()) with Context(comm.logger, br, options) as ctx: yield ctx diff --git a/python/rapidsmpf/rapidsmpf/tests/test_rmm_resource_adaptor.py b/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource_tracking.py similarity index 75% rename from python/rapidsmpf/rapidsmpf/tests/test_rmm_resource_adaptor.py rename to python/rapidsmpf/rapidsmpf/tests/test_buffer_resource_tracking.py index f241836bb..e9d33e677 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_rmm_resource_adaptor.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource_tracking.py @@ -9,8 +9,8 @@ import rmm import rmm.mr +from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.memory.scoped_memory_record import ScopedMemoryRecord -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor if TYPE_CHECKING: from rmm.pylibrmm.stream import Stream @@ -35,18 +35,18 @@ def dealloc_cb(ptr: int, size: int, stream: Stream) -> None: track.append(ptr) upstream_mr = rmm.mr.CallbackMemoryResource(alloc_cb, dealloc_cb) - mr_adaptor = RmmResourceAdaptor(upstream_mr=upstream_mr) + br = BufferResource(upstream_mr) - # Delete upstream to check that adaptor keeps it alive. + # Delete upstream to check that BR keeps it alive. del upstream_mr - buf = rmm.DeviceBuffer(size=100 * KIB, mr=mr_adaptor) - assert mr_adaptor.current_allocated == 100 * KIB + buf = rmm.DeviceBuffer(size=100 * KIB, mr=br) + assert br.current_allocated == 100 * KIB assert state["bytes"] == 100 * KIB assert len(track) == 1 del buf - assert mr_adaptor.current_allocated == 0 + assert br.current_allocated == 0 assert state["bytes"] == 0 assert len(track) == 2 # alloc + dealloc @@ -58,12 +58,24 @@ def alloc_cb(size: int, stream: Stream) -> int: def dealloc_cb(ptr: int, size: int, stream: Stream) -> None: return None - mr = RmmResourceAdaptor( - upstream_mr=rmm.mr.CallbackMemoryResource(alloc_cb, dealloc_cb), - ) + br = BufferResource(rmm.mr.CallbackMemoryResource(alloc_cb, dealloc_cb)) with pytest.raises(RuntimeError, match="not a MemoryError"): - mr.allocate(1024) + br.allocate(1024) + + +def test_lifetime_no_dangling_stream_pool() -> None: + # rmm.DeviceBuffer keeps the BufferResource alive via the inherited + # DeviceMemoryResource owning ref. Dropping the local handle to BR + # must not invalidate the buffer's underlying stream pool. + def make() -> rmm.DeviceBuffer: + br = BufferResource(rmm.mr.CudaMemoryResource()) + return rmm.DeviceBuffer(size=1024, mr=br) + + buf = make() + # Before the merge, this would dangle once the local `br` went out of + # scope. Now the buffer keeps the BR (and its stream pool) alive. + buf.copy_from_host(b"x" * 1024) def test_initial_state() -> None: diff --git a/python/rapidsmpf/rapidsmpf/tests/test_config.py b/python/rapidsmpf/rapidsmpf/tests/test_config.py index c7969a386..63b38ca45 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_config.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_config.py @@ -25,7 +25,6 @@ is_pinned_memory_resources_supported, ) from rapidsmpf.progress_thread import ProgressThread -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.core.context import Context @@ -542,7 +541,7 @@ def test_context_from_options_creates_instance_with_explicit_options() -> None: "spill_device_limit": "1GiB", } ) - mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) + mr = rmm.mr.CudaMemoryResource() comm = single_comm.new_communicator(opts, ProgressThread()) with Context.from_options( @@ -556,7 +555,7 @@ def test_context_from_options_creates_instance_with_explicit_options() -> None: def test_context_from_options_uses_default_when_options_empty() -> None: opts = Options() - mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) + mr = rmm.mr.CudaMemoryResource() comm = single_comm.new_communicator(opts, ProgressThread()) with Context.from_options(comm.logger, mr, opts) as ctx: @@ -568,7 +567,7 @@ def test_context_from_options_uses_default_when_options_empty() -> None: def test_context_from_options_enables_statistics_when_requested() -> None: opts = Options({"statistics": "on"}) - mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) + mr = rmm.mr.CudaMemoryResource() comm = single_comm.new_communicator(opts, ProgressThread()) with Context.from_options( @@ -580,7 +579,7 @@ def test_context_from_options_enables_statistics_when_requested() -> None: def test_context_from_options_creates_buffer_resource() -> None: opts = Options() - mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) + mr = rmm.mr.CudaMemoryResource() comm = single_comm.new_communicator(opts, ProgressThread()) with Context.from_options(comm.logger, mr, opts) as ctx: @@ -590,7 +589,7 @@ def test_context_from_options_creates_buffer_resource() -> None: def test_context_from_options_can_create_channel() -> None: opts = Options() - mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) + mr = rmm.mr.CudaMemoryResource() comm = single_comm.new_communicator(opts, ProgressThread()) with Context.from_options(comm.logger, mr, opts) as ctx: diff --git a/python/rapidsmpf/rapidsmpf/tests/test_statistics.py b/python/rapidsmpf/rapidsmpf/tests/test_statistics.py index 362c79c7c..ff68188f7 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_statistics.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_statistics.py @@ -9,7 +9,7 @@ import pytest -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor +from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.statistics import Formatter, Statistics if TYPE_CHECKING: @@ -46,15 +46,15 @@ def test_get_empty_memory_records() -> None: def test_memory_profiling(device_mr: rmm.mr.CudaMemoryResource) -> None: - mr = RmmResourceAdaptor(device_mr) + br = BufferResource(device_mr) stats = Statistics(enable=True) - with stats.memory_profiling(mr, "outer"): - b1 = mr.allocate(1024) - with stats.memory_profiling(mr, "inner"): - mr.deallocate(mr.allocate(512), 512) - mr.deallocate(mr.allocate(512), 512) - mr.deallocate(b1, 1024) - mr.deallocate(mr.allocate(1024), 1024) + with stats.memory_profiling(br, "outer"): + b1 = br.allocate(1024) + with stats.memory_profiling(br, "inner"): + br.deallocate(br.allocate(512), 512) + br.deallocate(br.allocate(512), 512) + br.deallocate(b1, 1024) + br.deallocate(br.allocate(1024), 1024) inner = stats.get_memory_records()["inner"] assert inner.scoped.num_total_allocs() == 2 @@ -149,10 +149,10 @@ def test_write_json_string_matches_file(tmp_path: pathlib.Path) -> None: def test_write_json_memory_records(device_mr: rmm.mr.CudaMemoryResource) -> None: - mr = RmmResourceAdaptor(device_mr) + br = BufferResource(device_mr) stats = Statistics(enable=True) - with stats.memory_profiling(mr, "alloc"): - mr.deallocate(mr.allocate(1024), 1024) + with stats.memory_profiling(br, "alloc"): + br.deallocate(br.allocate(1024), 1024) data = json.loads(stats.write_json_string()) assert "memory_records" in data @@ -164,9 +164,9 @@ def test_write_json_memory_records(device_mr: rmm.mr.CudaMemoryResource) -> None def test_invalid_memory_record_names(device_mr: rmm.mr.CudaMemoryResource) -> None: - mr = RmmResourceAdaptor(device_mr) + br = BufferResource(device_mr) stats = Statistics(enable=True) - with stats.memory_profiling(mr, 'bad"name'): + with stats.memory_profiling(br, 'bad"name'): pass with pytest.raises(ValueError): stats.write_json_string()