diff --git a/cpp/benchmarks/bench_comm.cpp b/cpp/benchmarks/bench_comm.cpp index afc70faf7..419f4d04b 100644 --- a/cpp/benchmarks/bench_comm.cpp +++ b/cpp/benchmarks/bench_comm.cpp @@ -318,14 +318,12 @@ int main(int argc, char** argv) { rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); auto br = BufferResource::create( + stats, mr, PinnedMemoryResource::Disabled, {}, std::chrono::milliseconds{1}, - std::make_shared( - 16, rmm::cuda_stream::flags::non_blocking - ), - stats + std::make_shared(16, rmm::cuda_stream::flags::non_blocking) ); // Print benchmark/hardware info. diff --git a/cpp/benchmarks/bench_partition.cpp b/cpp/benchmarks/bench_partition.cpp index 909744f24..4aa7440de 100644 --- a/cpp/benchmarks/bench_partition.cpp +++ b/cpp/benchmarks/bench_partition.cpp @@ -60,7 +60,8 @@ static void BM_PartitionAndPack(benchmark::State& state) { // Create a pool memory resource with 50% of GPU memory rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_memory_resource{}, pool_size}; - auto br = rapidsmpf::BufferResource::create(pool_mr); + auto br = + rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), pool_mr); // Create input table auto table = create_int_table(num_rows, stream); @@ -111,7 +112,8 @@ static void BM_PartitionAndPackCurrentImpl(benchmark::State& state) { // Create a pool memory resource with 50% of GPU memory rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_memory_resource{}, pool_size}; - auto br = rapidsmpf::BufferResource::create(pool_mr); + auto br = + rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), pool_mr); // Create input table auto table = create_int_table(num_rows, stream); diff --git a/cpp/benchmarks/bench_shuffle.cpp b/cpp/benchmarks/bench_shuffle.cpp index 36710c4fb..be49848ee 100644 --- a/cpp/benchmarks/bench_shuffle.cpp +++ b/cpp/benchmarks/bench_shuffle.cpp @@ -548,15 +548,11 @@ int main(int argc, char** argv) { // We're only going to measure the last run, so disable initially. stats->disable(); auto br = rapidsmpf::BufferResource::create( + stats, stat_enabled_mr, args.pinned_mem_disable ? rapidsmpf::PinnedMemoryResource::Disabled : rapidsmpf::PinnedMemoryResource::make_if_available(), - std::move(memory_limits), - std::chrono::milliseconds{1}, - std::make_shared( - 16, rmm::cuda_stream::flags::non_blocking - ), - stats + std::move(memory_limits) ); std::shared_ptr comm; diff --git a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp index 8b3b20cd8..331c58a72 100644 --- a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp +++ b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp @@ -324,7 +324,9 @@ int main(int argc, char** argv) { // Initialize configuration options from environment variables. rapidsmpf::config::Options options{rapidsmpf::config::get_environment_variables()}; - auto progress_thread = std::make_shared(); + + auto stats = rapidsmpf::Statistics::create(); + auto progress_thread = std::make_shared(stats); std::shared_ptr comm; if (args.comm_type == "mpi") { @@ -364,20 +366,16 @@ int main(int argc, char** argv) { memory_limits[rapidsmpf::MemoryType::DEVICE] = args.device_mem_limit_mb << 20; } - auto stats = rapidsmpf::Statistics::create(); - auto pinned_mr = args.pinned_mem_disable ? rapidsmpf::PinnedMemoryResource::Disabled : rapidsmpf::PinnedMemoryResource::make_if_available(); auto br = rapidsmpf::BufferResource::create( + stats, stat_enabled_mr, pinned_mr, std::move(memory_limits), std::nullopt, - std::make_shared( - 16, rmm::cuda_stream::flags::non_blocking - ), - stats + std::make_shared(16, rmm::cuda_stream::flags::non_blocking) ); auto& log = *comm->logger(); @@ -411,7 +409,7 @@ int main(int argc, char** argv) { for (std::uint64_t i = 0; i < total_num_runs; ++i) { // Clear statistics before the last run so only the final run is reported. if (i == total_num_runs - 1) { - ctx->statistics()->clear(); + stats->clear(); } double const elapsed = run(ctx, comm, args, stream).count(); std::stringstream ss; @@ -458,15 +456,14 @@ int main(int argc, char** argv) { log.print(ss.str()); } - auto statistics = ctx->statistics(); if (args.enable_memory_profiler) { - log.print(statistics->report({ + log.print(stats->report({ .mr = stat_enabled_mr, .pinned_mr = pinned_mr, .header = "Statistics (of the last run):", })); } else { - log.print(statistics->report({.header = "Statistics (of the last run):"})); + log.print(stats->report({.header = "Statistics (of the last run):"})); } if (!use_bootstrap) { diff --git a/cpp/benchmarks/streaming/ndsh/utils.cpp b/cpp/benchmarks/streaming/ndsh/utils.cpp index 8716da3be..210251473 100644 --- a/cpp/benchmarks/streaming/ndsh/utils.cpp +++ b/cpp/benchmarks/streaming/ndsh/utils.cpp @@ -157,6 +157,7 @@ create_context( ); auto br = BufferResource::create( + statistics, std::move(mr), arguments.no_pinned_host_memory ? PinnedMemoryResource::Disabled : PinnedMemoryResource::make_if_available(), @@ -164,8 +165,7 @@ create_context( arguments.periodic_spill, std::make_shared( arguments.num_streams, rmm::cuda_stream::flags::non_blocking - ), - statistics + ) ); auto environment = config::get_environment_variables(); environment["NUM_STREAMING_THREADS"] = diff --git a/cpp/examples/example_shuffle.cpp b/cpp/examples/example_shuffle.cpp index 144c3bd59..ae51c9870 100644 --- a/cpp/examples/example_shuffle.cpp +++ b/cpp/examples/example_shuffle.cpp @@ -46,7 +46,7 @@ int main(int argc, char** argv) { // We will use the same stream, memory, and buffer resource throughout the example. rmm::cuda_stream_view stream = cudf::get_default_stream(); rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); - auto br = rapidsmpf::BufferResource::create(mr); + auto br = rapidsmpf::BufferResource::create(stats, mr); // As input data, we use a helper function from the benchmark suite. It creates a // random cudf table with 2 columns and 100 rows. In this example, each MPI rank diff --git a/cpp/include/rapidsmpf/bootstrap/ucxx.hpp b/cpp/include/rapidsmpf/bootstrap/ucxx.hpp index 52d7d7623..35718e1ab 100644 --- a/cpp/include/rapidsmpf/bootstrap/ucxx.hpp +++ b/cpp/include/rapidsmpf/bootstrap/ucxx.hpp @@ -41,7 +41,9 @@ namespace bootstrap { * @throws std::runtime_error if initialization fails. * * @code - * auto progress = std::make_shared(); + * auto progress = std::make_shared( + * rapidsmpf::Statistics::disabled() + * ); * auto comm = rapidsmpf::bootstrap::create_ucxx_comm(progress); * comm->logger().print("Hello from rank " + std::to_string(comm->rank())); * @endcode diff --git a/cpp/include/rapidsmpf/communicator/communicator.hpp b/cpp/include/rapidsmpf/communicator/communicator.hpp index 2b236b8cd..714e07ed3 100644 --- a/cpp/include/rapidsmpf/communicator/communicator.hpp +++ b/cpp/include/rapidsmpf/communicator/communicator.hpp @@ -17,6 +17,7 @@ #include #include #include +#include /** * @namespace rapidsmpf @@ -403,7 +404,19 @@ class Communicator { }; protected: - Communicator() = default; + /** + * @brief Construct the base Communicator with a progress thread. + * + * The communicator delegates `statistics()` to @p progress_thread so the + * two always agree on the current `Statistics` instance, including after + * `ProgressThread::set_statistics()` swaps it. + * + * @param progress_thread The progress thread for this communicator. Must + * not be null. + * + * @throws std::invalid_argument If @p progress_thread is null. + */ + explicit Communicator(std::shared_ptr progress_thread); public: virtual ~Communicator() noexcept = default; @@ -614,18 +627,42 @@ class Communicator { /** * @brief Retrieves the progress thread associated with this communicator. - * @return Shared pointer to the progress thread. + * @return Shared pointer to the progress thread (never null). */ - [[nodiscard]] virtual std::shared_ptr const& - progress_thread() const = 0; + [[nodiscard]] std::shared_ptr const& + progress_thread() const noexcept { + return progress_thread_; + } + + /** + * @brief Retrieves the statistics instance associated with this communicator. + * + * Forwards to `progress_thread()->statistics()`, so the communicator and + * its progress thread always agree on the current instance — including + * after `ProgressThread::set_statistics()` swaps it. Satisfies the + * `StatisticsProvider` concept. + * + * @return Shared pointer to the statistics instance (never null). + */ + [[nodiscard]] std::shared_ptr statistics() const noexcept { + return progress_thread_->statistics(); + } /** * @brief Provides a string representation of the communicator. * @return A string describing the communicator. */ [[nodiscard]] virtual std::string str() const = 0; + + private: + /// Progress thread owning this communicator's `Statistics` instance. + /// Never null after construction. Accessed only through `progress_thread()` + /// and `statistics()` so derived classes cannot bypass that contract. + std::shared_ptr progress_thread_; }; +static_assert(StatisticsProvider); + /// @brief Whether RapidsMPF was built with the UCXX Communicator. #ifdef RAPIDSMPF_HAVE_UCXX constexpr bool COMM_HAVE_UCXX = true; diff --git a/cpp/include/rapidsmpf/communicator/mpi.hpp b/cpp/include/rapidsmpf/communicator/mpi.hpp index cc57a2812..63b0560e8 100644 --- a/cpp/include/rapidsmpf/communicator/mpi.hpp +++ b/cpp/include/rapidsmpf/communicator/mpi.hpp @@ -249,14 +249,6 @@ class MPI final : public Communicator { return logger_; } - /** - * @copydoc Communicator::progress_thread - */ - [[nodiscard]] std::shared_ptr const& - progress_thread() const override { - return progress_thread_; - } - /** * @copydoc Communicator::str */ @@ -267,8 +259,8 @@ class MPI final : public Communicator { Rank rank_; Rank nranks_; std::shared_ptr logger_; - std::shared_ptr progress_thread_; }; +static_assert(StatisticsProvider); } // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/communicator/single.hpp b/cpp/include/rapidsmpf/communicator/single.hpp index 9580bd6d9..390db753d 100644 --- a/cpp/include/rapidsmpf/communicator/single.hpp +++ b/cpp/include/rapidsmpf/communicator/single.hpp @@ -191,14 +191,6 @@ class Single final : public Communicator { return logger_; } - /** - * @copydoc Communicator::progress_thread - */ - [[nodiscard]] std::shared_ptr const& - progress_thread() const override { - return progress_thread_; - } - /** * @copydoc Communicator::str */ @@ -206,8 +198,8 @@ class Single final : public Communicator { private: std::shared_ptr logger_; - std::shared_ptr progress_thread_; }; +static_assert(StatisticsProvider); } // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/communicator/ucxx.hpp b/cpp/include/rapidsmpf/communicator/ucxx.hpp index 6b551670c..91980ccec 100644 --- a/cpp/include/rapidsmpf/communicator/ucxx.hpp +++ b/cpp/include/rapidsmpf/communicator/ucxx.hpp @@ -293,14 +293,6 @@ class UCXX final : public Communicator { return logger_; } - /** - * @copydoc Communicator::progress_thread - */ - [[nodiscard]] std::shared_ptr const& - progress_thread() const override { - return progress_thread_; - } - /** * @copydoc Communicator::str */ @@ -339,7 +331,6 @@ class UCXX final : public Communicator { std::shared_ptr shared_resources_; config::Options options_; std::shared_ptr logger_; - std::shared_ptr progress_thread_; std::shared_ptr<::ucxx::Endpoint> get_endpoint(Rank rank); void progress_worker(); @@ -347,4 +338,6 @@ class UCXX final : public Communicator { } // namespace ucxx +static_assert(StatisticsProvider); + } // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index d477f903a..3ad7bf632 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -88,6 +88,8 @@ class BufferResource : public std::enable_shared_from_this { * If pinned-host memory is disabled, available pinned-host memory is always reported * as zero regardless of the configured limit. * + * @param statistics The statistics instance to use. Pass `Statistics::disabled()` + * to opt out of statistics collection. * @param device_mr Device memory resource used for device allocations. To ensure * allocations are tracked for memory-limit accounting and statistics, use * `BufferResource::device_mr()` instead of the original memory resource after @@ -101,17 +103,16 @@ class BufferResource : public std::enable_shared_from_this { * disables the dedicated spill-check thread. * @param stream_pool CUDA stream pool used for operations that do not take an * explicit CUDA stream. - * @param statistics Statistics instance used for runtime metrics. * @return A newly constructed `BufferResource` owned by `std::shared_ptr`. */ [[nodiscard]] static std::shared_ptr create( + std::shared_ptr statistics, cuda::mr::any_resource device_mr, std::optional pinned_mr = PinnedMemoryResource::Disabled, std::unordered_map memory_limits = {}, std::optional periodic_spill_check = std::chrono::milliseconds{1}, std::shared_ptr stream_pool = std::make_shared< - rmm::cuda_stream_pool>(16, rmm::cuda_stream::flags::non_blocking), - std::shared_ptr statistics = Statistics::disabled() + rmm::cuda_stream_pool>(16, rmm::cuda_stream::flags::non_blocking) ); /** @@ -123,14 +124,15 @@ class BufferResource : public std::enable_shared_from_this { * * @param mr A device-accessible RMM memory resource. * @param options Configuration options. - * @param statistics The statistics instance to use (disabled by default). + * @param statistics The statistics instance to use. Pass `Statistics::disabled()` + * to opt out of statistics collection. * @return A shared pointer to a BufferResource instance configured according to the * options. */ static std::shared_ptr from_options( cuda::mr::any_resource mr, config::Options options, - std::shared_ptr statistics = Statistics::disabled() + std::shared_ptr statistics ); ~BufferResource() noexcept = default; @@ -508,11 +510,11 @@ class BufferResource : public std::enable_shared_from_this { private: /** @brief Private constructor, use `create()` or `from_options()`. */ BufferResource( + std::shared_ptr statistics, std::optional pinned_mr, std::unordered_map memory_limits, std::optional periodic_spill_check, - std::shared_ptr stream_pool, - std::shared_ptr statistics + std::shared_ptr stream_pool ); std::mutex mutex_; diff --git a/cpp/include/rapidsmpf/progress_thread.hpp b/cpp/include/rapidsmpf/progress_thread.hpp index 44f3c539c..9a4efdb8c 100644 --- a/cpp/include/rapidsmpf/progress_thread.hpp +++ b/cpp/include/rapidsmpf/progress_thread.hpp @@ -4,6 +4,7 @@ */ #pragma once +#include #include #include #include @@ -117,14 +118,15 @@ class ProgressThread { /** * @brief Construct a new progress thread that can handle multiple functions. * - * @param statistics The statistics instance to use (disabled by default). + * @param statistics The statistics instance to use. Pass `Statistics::disabled()` + * to opt out of statistics collection. * @param sleep The duration to sleep between each progress loop iteration. * If 0, the thread yields execution instead of sleeping. Anecdotally, a 1 us * sleep time (the default) is sufficient to avoid starvation and get smooth * progress. */ ProgressThread( - std::shared_ptr statistics = Statistics::disabled(), + std::shared_ptr statistics, Duration sleep = std::chrono::microseconds{1} ); @@ -183,6 +185,19 @@ class ProgressThread { */ std::shared_ptr statistics() const noexcept; + /** + * @brief Replace the statistics instance held by this progress thread. + * + * The swap is performed atomically so the progress loop can safely + * dereference `statistics_` concurrently with this call. + * + * @param statistics The new statistics instance. Pass `Statistics::disabled()` to opt + * out of statistics collection. + * + * @throws std::invalid_argument If @p statistics is null. + */ + void set_statistics(std::shared_ptr statistics); + private: /** * @brief The event loop progressing each of the functions. @@ -192,7 +207,7 @@ class ProgressThread { */ void event_loop(); - std::shared_ptr statistics_; + std::atomic> statistics_; bool is_thread_initialized_{false}; bool active_{false}; mutable std::mutex mutex_; @@ -206,5 +221,6 @@ class ProgressThread { }; static_assert(StatisticsProvider); +static_assert(MutableStatisticsProvider); } // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/statistics.hpp b/cpp/include/rapidsmpf/statistics.hpp index adf47b4c7..92e777236 100644 --- a/cpp/include/rapidsmpf/statistics.hpp +++ b/cpp/include/rapidsmpf/statistics.hpp @@ -575,14 +575,23 @@ class Statistics : public std::enable_shared_from_this { /** * @brief Constructs an active MemoryRecorder. Pushes a scoped record at * construction; the destructor pops it and (if @p stats is still - * enabled) publishes it under @p name. + * alive and enabled) publishes it under @p name. * - * @param stats Owning Statistics. Must not be null. + * The recorder holds @p stats weakly so it does not extend the + * lifetime of the Statistics instance. If the instance has been + * destroyed by the time the recorder is destructed (for example + * because a `MutableStatisticsProvider` swapped in a new one and + * dropped the last strong reference), the destructor still pops the + * scoped memory record but does not publish. A `std::shared_ptr` + * may be passed directly thanks to the implicit conversion to + * `std::weak_ptr`. + * + * @param stats Owning Statistics. Must not be expired/empty. * @param mr RMM resource adaptor providing scoped memory statistics. * @param name Name of the scope. */ MemoryRecorder( - std::shared_ptr stats, RmmResourceAdaptor mr, std::string name + std::weak_ptr stats, RmmResourceAdaptor mr, std::string name ); ~MemoryRecorder(); @@ -595,8 +604,13 @@ class Statistics : public std::enable_shared_from_this { private: /// No-op recorder iff `mr_` is `std::nullopt`. std::optional mr_{std::nullopt}; - /// stats_ != nullptr iff `mr_.has_value()`. - std::shared_ptr stats_{nullptr}; + /// Weak reference to the owning Statistics. Held weakly so that a + /// `MutableStatisticsProvider` swap (or any other release of the last + /// strong reference) can destroy the original instance without being + /// blocked by outstanding recorders. The destructor locks this and + /// skips publishing if the instance is no longer alive. Has been + /// assigned (possibly to a now-expired pointer) iff `mr_.has_value()`. + std::weak_ptr stats_; std::string name_{}; }; @@ -656,6 +670,28 @@ concept StatisticsProvider = requires(T const& t) { } noexcept -> std::same_as>; }; +/** + * @brief Refinement of `StatisticsProvider` for providers that allow swapping + * their `Statistics` instance at runtime. + * + * Extends `StatisticsProvider` with a `set_statistics(std::shared_ptr)` + * method that replaces the provider's currently held instance with the given one. + * The new instance must not be null; implementations are expected to throw + * `std::invalid_argument` (or similar) on a null argument and otherwise install + * the new instance unconditionally. + * + * @warning Concurrent calls to `set_statistics()` are implementation-defined and + * likely result in undefined behavior unless the implementation explicitly + * documents otherwise. Callers should assume external synchronization is required. + */ +template +concept MutableStatisticsProvider = + StatisticsProvider && requires(T& t, std::shared_ptr stats) { + { + t.set_statistics(std::move(stats)) + } -> std::same_as; + }; + /** * @brief Macro for automatic memory profiling of a code scope. * diff --git a/cpp/include/rapidsmpf/stream_ordered_timing.hpp b/cpp/include/rapidsmpf/stream_ordered_timing.hpp index 3ebdf3666..a92990d26 100644 --- a/cpp/include/rapidsmpf/stream_ordered_timing.hpp +++ b/cpp/include/rapidsmpf/stream_ordered_timing.hpp @@ -37,15 +37,21 @@ class StreamOrderedTiming { /** * @brief Constructs a StreamOrderedTiming and marks the start position in the stream. * - * If `statistics` is disabled (e.g. constructed with + * If @p statistics is disabled (e.g. constructed with * `Statistics::Mode::Disabled`), this is a no-op and subsequent calls to * `stop_and_record` will also be no-ops. * + * @p statistics is held weakly: the timing does not extend the lifetime + * of the `Statistics` instance. If the instance is destroyed before the + * stream actually executes the stop callback, the recording is skipped + * (see `stop_and_record` and `cancel_inflight_timings`). + * * @param stream The CUDA stream to time. - * @param statistics The Statistics object that will receive the duration entry. + * @param statistics The Statistics object that will receive the duration + * entry. Must not be expired/empty. */ StreamOrderedTiming( - rmm::cuda_stream_view stream, std::shared_ptr statistics + rmm::cuda_stream_view stream, std::weak_ptr statistics ); /** @@ -94,7 +100,15 @@ class StreamOrderedTiming { private: std::uintptr_t uid_{0}; rmm::cuda_stream_view stream_; - std::shared_ptr statistics_; + /// Weak reference to the owning Statistics. Held weakly so that this + /// timing does not extend the lifetime of the instance; the global + /// timings map and `cancel_inflight_timings` handle the case where the + /// instance is destroyed before the stop callback fires. + /// + /// Doubles as the active/no-op sentinel: empty (default-constructed) + /// for no-op instances (constructor was given a disabled or already- + /// destroyed Statistics), non-empty for active instances. + std::weak_ptr statistics_; }; } // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/streaming/core/context.hpp b/cpp/include/rapidsmpf/streaming/core/context.hpp index 559f58574..578f1e6c1 100644 --- a/cpp/include/rapidsmpf/streaming/core/context.hpp +++ b/cpp/include/rapidsmpf/streaming/core/context.hpp @@ -83,8 +83,8 @@ class Context { * @param logger The logger to use. * @param options Configuration options used to initialize the Context and its * components. - * @param statistics The statistics instance to use (disabled by default). - * + * @param statistics The statistics instance to use. Pass `Statistics::disabled()` + * to opt out of statistics collection. * @return A fully initialized Context. * * @throws std::invalid_argument If an option value is invalid. @@ -105,7 +105,7 @@ class Context { RmmResourceAdaptor mr, std::shared_ptr logger, config::Options options, - std::shared_ptr statistics = Statistics::disabled() + std::shared_ptr statistics ); // No copy constructor and assignment operator. diff --git a/cpp/src/coll/allreduce.cpp b/cpp/src/coll/allreduce.cpp index 6e075f767..bfbb5d551 100644 --- a/cpp/src/coll/allreduce.cpp +++ b/cpp/src/coll/allreduce.cpp @@ -59,8 +59,7 @@ AllReduce::AllReduce( out_buffer_->rebind_stream(in_buffer_->stream()); // Note: after this copy, we must check out_buffer's write event before receiving into // in_buffer. See StartPreRemainder in the event loop. - // TODO: make Communicator statistics-aware, and pass its statistics instance here. - buffer_copy(Statistics::disabled(), *out_buffer_, *in_buffer_, in_buffer_->size); + buffer_copy(comm_->statistics(), *out_buffer_, *in_buffer_, in_buffer_->size); auto const rank = comm_->rank(); if (rank < 2 * non_pow2_remainder_) { diff --git a/cpp/src/communicator/communicator.cpp b/cpp/src/communicator/communicator.cpp index ede8dff7c..3b2de55d0 100644 --- a/cpp/src/communicator/communicator.cpp +++ b/cpp/src/communicator/communicator.cpp @@ -3,11 +3,25 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include +#include + #include #include +#include #include namespace rapidsmpf { + +Communicator::Communicator(std::shared_ptr progress_thread) + : progress_thread_{std::move(progress_thread)} { + RAPIDSMPF_EXPECTS( + progress_thread_ != nullptr, + "Communicator progress thread cannot be null", + std::invalid_argument + ); +} + namespace { Communicator::Logger::LOG_LEVEL level_from_string(std::string const& str) { auto const trimmed = to_upper(trim(str)); diff --git a/cpp/src/communicator/mpi.cpp b/cpp/src/communicator/mpi.cpp index 6aa501662..e435fe7fa 100644 --- a/cpp/src/communicator/mpi.cpp +++ b/cpp/src/communicator/mpi.cpp @@ -102,7 +102,8 @@ MPI::MPI( config::Options options, std::shared_ptr progress_thread ) - : comm_{comm}, + : Communicator{std::move(progress_thread)}, + comm_{comm}, rank_{[comm]() { int r; RAPIDSMPF_MPI(MPI_Comm_rank(comm, &r)); @@ -113,8 +114,7 @@ MPI::MPI( RAPIDSMPF_MPI(MPI_Comm_size(comm, &n)); return Rank(n); }()}, - logger_{std::make_shared(rank_, std::move(options))}, - progress_thread_{std::move(progress_thread)} { + logger_{std::make_shared(rank_, std::move(options))} { check_mpi_thread_support(); } diff --git a/cpp/src/communicator/single.cpp b/cpp/src/communicator/single.cpp index e8d594a69..3859bf0e6 100644 --- a/cpp/src/communicator/single.cpp +++ b/cpp/src/communicator/single.cpp @@ -12,8 +12,8 @@ namespace rapidsmpf { Single::Single(config::Options options, std::shared_ptr progress_thread) - : logger_{std::make_shared(0, std::move(options))}, - progress_thread_{std::move(progress_thread)} {} + : Communicator{std::move(progress_thread)}, + logger_{std::make_shared(0, std::move(options))} {} std::unique_ptr Single::send( std::unique_ptr>, Rank, Tag diff --git a/cpp/src/communicator/ucxx.cpp b/cpp/src/communicator/ucxx.cpp index 294178e30..d7b29909f 100644 --- a/cpp/src/communicator/ucxx.cpp +++ b/cpp/src/communicator/ucxx.cpp @@ -1130,10 +1130,10 @@ UCXX::UCXX( config::Options options, std::shared_ptr progress_thread ) - : shared_resources_(ucxx_initialized_rank->shared_resources_), + : Communicator{std::move(progress_thread)}, + shared_resources_(ucxx_initialized_rank->shared_resources_), options_{std::move(options)}, - logger_{std::make_shared(shared_resources_->rank(), options_)}, - progress_thread_{std::move(progress_thread)} { + logger_{std::make_shared(shared_resources_->rank(), options_)} { shared_resources_->logger = logger_.get(); } @@ -1503,7 +1503,7 @@ std::shared_ptr UCXX::split() { // Create the new UCXX instance auto initialized_rank = std::make_unique(shared_resources); return std::make_shared( - std::move(initialized_rank), options_, progress_thread_ + std::move(initialized_rank), options_, progress_thread() ); } diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index 0fb385a5f..97abcd44c 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -23,11 +23,11 @@ namespace rapidsmpf { BufferResource::BufferResource( + std::shared_ptr statistics, std::optional pinned_mr, std::unordered_map memory_limits, std::optional periodic_spill_check, - std::shared_ptr stream_pool, - std::shared_ptr statistics + std::shared_ptr stream_pool ) : pinned_mr_{std::move(pinned_mr)}, host_mr_{}, @@ -48,19 +48,19 @@ BufferResource::BufferResource( } std::shared_ptr BufferResource::create( + std::shared_ptr statistics, cuda::mr::any_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 + std::shared_ptr stream_pool ) { std::shared_ptr br{new BufferResource{ + std::move(statistics), std::move(pinned_mr), std::move(memory_limits), periodic_spill_check, - std::move(stream_pool), - std::move(statistics) + std::move(stream_pool) }}; // Install the owning adaptor *after* construction so that `weak_from_this()` is @@ -88,12 +88,12 @@ std::shared_ptr BufferResource::from_options( {MemoryType::DEVICE, device_limit_from_options(options)} }; return create( + std::move(statistics), 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) + stream_pool_from_options(options) ); } diff --git a/cpp/src/progress_thread.cpp b/cpp/src/progress_thread.cpp index 42846b421..6d114e915 100644 --- a/cpp/src/progress_thread.cpp +++ b/cpp/src/progress_thread.cpp @@ -37,7 +37,9 @@ ProgressThread::ProgressThread(std::shared_ptr statistics, Duration }, sleep ) { - RAPIDSMPF_EXPECTS(statistics_ != nullptr, "the statistics pointer cannot be NULL"); + RAPIDSMPF_EXPECTS( + statistics_.load() != nullptr, "the statistics pointer cannot be NULL" + ); } ProgressThread::~ProgressThread() { @@ -116,7 +118,16 @@ bool ProgressThread::is_running() const { } std::shared_ptr ProgressThread::statistics() const noexcept { - return statistics_; + return statistics_.load(); +} + +void ProgressThread::set_statistics(std::shared_ptr statistics) { + RAPIDSMPF_EXPECTS( + statistics != nullptr, + "the statistics pointer cannot be NULL", + std::invalid_argument + ); + statistics_.store(std::move(statistics)); } void ProgressThread::event_loop() { @@ -131,7 +142,9 @@ void ProgressThread::event_loop() { // Notify all waiting functions that we've completed an iteration cv_.notify_all(); - statistics_->add_duration_stat("event-loop-total", Clock::now() - t0_event_loop); + statistics_.load()->add_duration_stat( + "event-loop-total", Clock::now() - t0_event_loop + ); } } // namespace rapidsmpf diff --git a/cpp/src/statistics.cpp b/cpp/src/statistics.cpp index f9d3b96d3..676cba635 100644 --- a/cpp/src/statistics.cpp +++ b/cpp/src/statistics.cpp @@ -247,10 +247,10 @@ void Statistics::clear() { } Statistics::MemoryRecorder::MemoryRecorder( - std::shared_ptr stats, RmmResourceAdaptor mr, std::string name + std::weak_ptr stats, RmmResourceAdaptor mr, std::string name ) : mr_{std::move(mr)}, stats_{std::move(stats)}, name_{std::move(name)} { - RAPIDSMPF_EXPECTS(stats_ != nullptr, "the statistics cannot be null"); + RAPIDSMPF_EXPECTS(!stats_.expired(), "the statistics cannot be null"); mr_->begin_scoped_memory_record(); } @@ -259,13 +259,15 @@ Statistics::MemoryRecorder::~MemoryRecorder() { return; // no-op recorder; nothing was pushed. } // Always pop to keep the RMM adaptor's per-thread stack balanced, even if - // statistics were disabled after construction (in which case skip publish). + // statistics were disabled after construction or the owning Statistics + // was destroyed/replaced in the meantime (in which case skip publish). auto const scope = mr_->end_scoped_memory_record(); - if (!stats_->enabled()) { + auto const stats = stats_.lock(); + if (!stats || !stats->enabled()) { return; } - std::lock_guard lock(stats_->mutex_); - auto& record = stats_->memory_records_[name_]; + std::lock_guard lock(stats->mutex_); + auto& record = stats->memory_records_[name_]; ++record.num_calls; record.scoped.add_scope(scope); record.global_peak = std::max(record.global_peak, scope.peak()); diff --git a/cpp/src/stream_ordered_timing.cpp b/cpp/src/stream_ordered_timing.cpp index c25f68980..dbf3dcdf7 100644 --- a/cpp/src/stream_ordered_timing.cpp +++ b/cpp/src/stream_ordered_timing.cpp @@ -98,14 +98,15 @@ void timing_stop_cb(void* data) { } // namespace StreamOrderedTiming::StreamOrderedTiming( - rmm::cuda_stream_view stream, std::shared_ptr statistics + rmm::cuda_stream_view stream, std::weak_ptr statistics ) : stream_{stream}, statistics_{std::move(statistics)} { - RAPIDSMPF_EXPECTS(statistics_ != nullptr, "the statistics pointer cannot be NULL"); - if (!statistics_->enabled()) { - // Reset to nullptr to make the disabled state permanent for this instance. - // Otherwise, statistics could be re-enabled between now and stop_and_record(). - statistics_ = nullptr; + auto const stats = statistics_.lock(); + RAPIDSMPF_EXPECTS(stats != nullptr, "the statistics pointer cannot be NULL"); + if (!stats->enabled()) { + // Reset to an empty weak_ptr so this instance stays a no-op for its + // entire lifetime. + statistics_.reset(); return; } TimePoint const now = Clock::now(); @@ -122,7 +123,16 @@ StreamOrderedTiming::StreamOrderedTiming( void StreamOrderedTiming::stop_and_record( std::string const& name, std::optional stream_delay_name ) { - if (statistics_ == nullptr) { + // Locking the weak_ptr here serves two purposes: (1) it gates the + // disabled / no-op case (empty weak_ptr) and the case where the + // Statistics has already been destroyed, and (2) it pins the Statistics + // alive for the body below, so `~Statistics` (which calls + // `cancel_inflight_timings`) cannot run concurrently with our access of + // `global_timings_.at(uid_)`. + auto const stats = statistics_.lock(); + if (stats == nullptr) { + // Note: intentionally not checking enabled() here so that the timing will be + // cleared from the global timings map even if the Statistics is disabled. return; } { diff --git a/cpp/src/streaming/core/context.cpp b/cpp/src/streaming/core/context.cpp index e2e00d5b8..5dc875fc4 100644 --- a/cpp/src/streaming/core/context.cpp +++ b/cpp/src/streaming/core/context.cpp @@ -109,6 +109,7 @@ std::shared_ptr Context::from_options( config::Options options, std::shared_ptr statistics ) { + RAPIDSMPF_EXPECTS(statistics != nullptr, "the statistics pointer cannot be NULL"); return std::make_shared( options, std::move(logger), diff --git a/cpp/tests/main/mpi.cpp b/cpp/tests/main/mpi.cpp index a1fd1f2bf..4e23bde5e 100644 --- a/cpp/tests/main/mpi.cpp +++ b/cpp/tests/main/mpi.cpp @@ -7,6 +7,8 @@ #include #include +#include +#include #include "../environment.hpp" @@ -26,7 +28,11 @@ void Environment::SetUp() { options_ = rapidsmpf::config::Options(rapidsmpf::config::get_environment_variables()); comm_ = std::make_shared( - mpi_comm_, options_, std::make_shared() + mpi_comm_, + options_, + std::make_shared( + rapidsmpf::Statistics::from_options(options_) + ) ); } diff --git a/cpp/tests/main/single.cpp b/cpp/tests/main/single.cpp index a8c81ac2e..ddb672b4b 100644 --- a/cpp/tests/main/single.cpp +++ b/cpp/tests/main/single.cpp @@ -8,6 +8,7 @@ #include #include +#include #include "../environment.hpp" @@ -22,7 +23,10 @@ TestEnvironmentType Environment::type() const { void Environment::SetUp() { options_ = rapidsmpf::config::Options(rapidsmpf::config::get_environment_variables()); comm_ = std::make_shared( - options_, std::make_shared() + options_, + std::make_shared( + rapidsmpf::Statistics::from_options(options_) + ) ); split_comm_ = comm_; } diff --git a/cpp/tests/main/ucxx.cpp b/cpp/tests/main/ucxx.cpp index 82e891a58..9d3a5589a 100644 --- a/cpp/tests/main/ucxx.cpp +++ b/cpp/tests/main/ucxx.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include "../environment.hpp" @@ -36,7 +37,11 @@ void Environment::SetUp() { options_ = rapidsmpf::config::Options(rapidsmpf::config::get_environment_variables()); comm_ = rapidsmpf::ucxx::init_using_mpi( - MPI_COMM_WORLD, options_, std::make_shared() + MPI_COMM_WORLD, + options_, + std::make_shared( + rapidsmpf::Statistics::from_options(options_) + ) ); } diff --git a/cpp/tests/streaming/base_streaming_fixture.hpp b/cpp/tests/streaming/base_streaming_fixture.hpp index 23cbf2b6e..2e4fae5ec 100644 --- a/cpp/tests/streaming/base_streaming_fixture.hpp +++ b/cpp/tests/streaming/base_streaming_fixture.hpp @@ -44,7 +44,10 @@ class BaseStreamingFixture : public ::testing::Test { stream = cudf::get_default_stream(); br = rapidsmpf::BufferResource::create( - mr_cuda, rapidsmpf::PinnedMemoryResource::Disabled, std::move(memory_limits) + rapidsmpf::Statistics::disabled(), + mr_cuda, + rapidsmpf::PinnedMemoryResource::Disabled, + std::move(memory_limits) ); ctx = std::make_shared( std::move(options), GlobalEnvironment->comm_->logger(), br diff --git a/cpp/tests/streaming/test_channel_metadata.cpp b/cpp/tests/streaming/test_channel_metadata.cpp index 6c97ce3c5..12792d3dc 100644 --- a/cpp/tests/streaming/test_channel_metadata.cpp +++ b/cpp/tests/streaming/test_channel_metadata.cpp @@ -166,8 +166,9 @@ TEST_F(StreamingChannelMetadata, MessageRoundTrip) { class StreamingChannelMetadataGPU : public ::testing::Test { protected: rmm::cuda_stream_view stream{cudf::get_default_stream()}; - std::shared_ptr br = - rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + std::shared_ptr br = rapidsmpf::BufferResource::create( + rapidsmpf::Statistics::disabled(), cudf::get_current_device_resource_ref() + ); std::shared_ptr make_chunk(std::vector vals) { rmm::device_buffer buf(vals.data(), vals.size() * sizeof(int32_t), stream); diff --git a/cpp/tests/streaming/test_fanout.cpp b/cpp/tests/streaming/test_fanout.cpp index 3653f643e..33a1d1fd6 100644 --- a/cpp/tests/streaming/test_fanout.cpp +++ b/cpp/tests/streaming/test_fanout.cpp @@ -547,13 +547,11 @@ class SpillingStreamingFanout : public BaseStreamingFixture { std::unordered_map memory_limits = { {MemoryType::DEVICE, 0}, }; - br = rapidsmpf::BufferResource::create( - mr_cuda, rapidsmpf::PinnedMemoryResource::Disabled, memory_limits + br = BufferResource::create( + Statistics::disabled(), mr_cuda, PinnedMemoryResource::Disabled, memory_limits ); auto options = ctx->options(); - ctx = std::make_shared( - options, GlobalEnvironment->comm_->logger(), br - ); + ctx = std::make_shared(options, GlobalEnvironment->comm_->logger(), br); } }; diff --git a/cpp/tests/streaming/test_message.cpp b/cpp/tests/streaming/test_message.cpp index ba2a9454f..e755f0750 100644 --- a/cpp/tests/streaming/test_message.cpp +++ b/cpp/tests/streaming/test_message.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include @@ -18,7 +19,9 @@ using namespace rapidsmpf::streaming; class StreamingMessage : public ::testing::Test { protected: void SetUp() override { - br = BufferResource::create(cudf::get_current_device_resource_ref()); + br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref() + ); stream = cudf::get_default_stream(); } diff --git a/cpp/tests/streaming/test_table_chunk.cpp b/cpp/tests/streaming/test_table_chunk.cpp index 5d42285bd..26c9db612 100644 --- a/cpp/tests/streaming/test_table_chunk.cpp +++ b/cpp/tests/streaming/test_table_chunk.cpp @@ -29,35 +29,24 @@ using namespace rapidsmpf; using namespace rapidsmpf::streaming; class StreamingTableChunk : public BaseStreamingFixture, - public ::testing::WithParamInterface { + public ::testing::WithParamInterface { protected: void SetUp() override { - rapidsmpf::config::Options options( - rapidsmpf::config::get_environment_variables() - ); + config::Options options(config::get_environment_variables()); - std::unordered_map memory_limits{}; - auto stream_pool = std::make_shared( - 16, rmm::cuda_stream::flags::non_blocking - ); stream = cudf::get_default_stream(); br = rapidsmpf::BufferResource::create( - mr_cuda, // device_mr - rapidsmpf::PinnedMemoryResource::make_if_available(), // pinned_mr - memory_limits, // memory_limits - std::chrono::milliseconds{1}, // periodic_spill_check - stream_pool, // stream_pool - Statistics::disabled() // statistics - ); - ctx = std::make_shared( - options, GlobalEnvironment->comm_->logger(), br + Statistics::disabled(), + mr_cuda, + rapidsmpf::PinnedMemoryResource::make_if_available() ); + ctx = std::make_shared(options, GlobalEnvironment->comm_->logger(), br); } rmm::cuda_stream_view stream; rmm::mr::cuda_memory_resource mr_cuda; - std::shared_ptr br; - std::shared_ptr ctx; + std::shared_ptr br; + std::shared_ptr ctx; }; TEST_F(StreamingTableChunk, FromTable) { diff --git a/cpp/tests/test_allgather.cpp b/cpp/tests/test_allgather.cpp index 99472b499..daebbe14f 100644 --- a/cpp/tests/test_allgather.cpp +++ b/cpp/tests/test_allgather.cpp @@ -34,7 +34,9 @@ class BaseAllGatherTest : public ::testing::Test { protected: void SetUp() override { stream = cudf::get_default_stream(); - br = rapidsmpf::BufferResource::create(rmm::mr::cuda_memory_resource{}); + br = rapidsmpf::BufferResource::create( + rapidsmpf::Statistics::disabled(), rmm::mr::cuda_memory_resource{} + ); } void TearDown() override { @@ -295,6 +297,7 @@ TEST_F(BaseAllGatherTest, opid_reuse) { if (this_rank == 0) { // Recreate the buffer resource and allgather with the delayed MR. delay_br = rapidsmpf::BufferResource::create( + rapidsmpf::Statistics::disabled(), DelayedMemoryResource{br->device_mr(), std::chrono::milliseconds(500)} ); allgather = @@ -363,7 +366,7 @@ TEST_F(BaseAllGatherTest, opid_reuse) { TEST(PostBox, spill_uses_remaining_amount) { auto stream = cudf::get_default_stream(); auto mr = std::make_unique(); - auto br = rapidsmpf::BufferResource::create(*mr); + auto br = rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), *mr); rapidsmpf::coll::detail::PostBox postbox; diff --git a/cpp/tests/test_allreduce.cu b/cpp/tests/test_allreduce.cu index 5d4b9bd93..e0ce5d79c 100644 --- a/cpp/tests/test_allreduce.cu +++ b/cpp/tests/test_allreduce.cu @@ -28,6 +28,7 @@ #include #include #include +#include #include "environment.hpp" @@ -207,7 +208,7 @@ class BaseAllReduceTest : public ::testing::Test { protected: void SetUp() override { mr = std::make_unique(); - br = rapidsmpf::BufferResource::create(*mr); + br = rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), *mr); comm = GlobalEnvironment->comm_.get(); } diff --git a/cpp/tests/test_buffer.cpp b/cpp/tests/test_buffer.cpp index 742339c6f..2e0b7ab32 100644 --- a/cpp/tests/test_buffer.cpp +++ b/cpp/tests/test_buffer.cpp @@ -21,6 +21,7 @@ #include #include #include +#include #include "utils.hpp" @@ -54,6 +55,7 @@ class BufferRebindStreamTest : public ::testing::TestWithParam { } br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref(), PinnedMemoryResource::make_if_available(), std::unordered_map{}, diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index 95cb74edc..0ed337aa1 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -55,6 +55,7 @@ std::unique_ptr zeros( TEST(BufferResource, ReservationOverbooking) { // Create a buffer resource that always reports 10 KiB of available device memory. auto br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref(), PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, 10_KiB}} @@ -122,6 +123,7 @@ TEST(BufferResource, ReservationReleasing) { // Create a buffer resource that always reports 10 KiB of available host and device // memory. auto br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref(), PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, 10_KiB}, {MemoryType::HOST, 10_KiB}} @@ -175,7 +177,10 @@ TEST(BufferResource, MemoryLimit) { // Create a buffer resource that limits available device memory to 10 KiB. auto br = BufferResource::create( - mr_cuda, PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, 10_KiB}} + Statistics::disabled(), + mr_cuda, + PinnedMemoryResource::Disabled, + {{MemoryType::DEVICE, 10_KiB}} ); EXPECT_EQ(br->memory_available(MemoryType::DEVICE), 10_KiB); EXPECT_EQ(br->memory_reserved(MemoryType::DEVICE), 0); @@ -264,7 +269,9 @@ TEST_P(PinnedMaxPoolSizeReservationLimitTest, TwoReservations) { if (max_pool_size.has_value() && *max_pool_size > 0) { memory_limits[MemoryType::PINNED_HOST] = safe_cast(*max_pool_size); } - auto br = BufferResource::create(cuda_mr, pinned_mr, std::move(memory_limits)); + auto br = BufferResource::create( + Statistics::disabled(), cuda_mr, pinned_mr, std::move(memory_limits) + ); // First 1 KiB reservation always succeeds. auto [r1, ob1] = br->reserve(MemoryType::PINNED_HOST, 1_KiB, AllowOverbooking::NO); @@ -292,12 +299,12 @@ TEST(BufferResource, AllocStatistics) { auto stats = Statistics::create(); auto pinned_mr = PinnedMemoryResource::make_if_available(); auto br = BufferResource::create( + stats, mr_cuda, pinned_mr, {}, std::nullopt, - std::make_shared(1, rmm::cuda_stream::flags::non_blocking), - stats + std::make_shared(1, rmm::cuda_stream::flags::non_blocking) ); auto stream = cudf::get_default_stream(); @@ -357,6 +364,7 @@ class BufferResourceReserveOrFailTest : public ::testing::Test { // Create a buffer resource with limited device memory (10 KiB) and unlimited // host memory. BufferResource auto-wraps mr_cuda for allocation tracking. br = BufferResource::create( + Statistics::disabled(), mr_cuda, PinnedMemoryResource::Disabled, std::unordered_map{{MemoryType::DEVICE, 10_KiB}} @@ -419,7 +427,9 @@ TEST_F(BufferResourceReserveOrFailTest, MultipleTypes) { class BaseBufferResourceCopyTest : public ::testing::Test { protected: void SetUp() override { - br = BufferResource::create(cudf::get_current_device_resource_ref()); + br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref() + ); stream = cudf::get_default_stream(); // initialize the host pattern @@ -641,11 +651,11 @@ class BufferResourceDifferentResourcesTest : public ::testing::Test { // Setup br1 with statistics for its device memory mr1 = std::make_unique(rmm::mr::cuda_memory_resource{}); - br1 = BufferResource::create(*mr1); + br1 = BufferResource::create(Statistics::disabled(), *mr1); // Setup br2 with statistics for its device memory mr2 = std::make_unique(rmm::mr::cuda_memory_resource{}); - br2 = BufferResource::create(*mr2); + br2 = BufferResource::create(Statistics::disabled(), *mr2); } std::unique_ptr create_source_buffer() { @@ -792,7 +802,9 @@ TEST_F(BufferCopyEdgeCases, SameBufferIsDisallowed) { TEST(BufferResource, DeviceMrKeepsBufferResourceAlive) { constexpr std::size_t N = 1024; - auto br = BufferResource::create(cudf::get_current_device_resource_ref()); + auto br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref() + ); std::weak_ptr weak_br = br; auto stream = cudf::get_default_stream(); @@ -823,7 +835,9 @@ TEST(OwningResourceAdaptor, CopyThrowsWhenBackRefExpired) { std::weak_ptr weak_br; { - auto br = BufferResource::create(cudf::get_current_device_resource_ref()); + auto br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref() + ); weak_br = br; } ASSERT_TRUE(weak_br.expired()); diff --git a/cpp/tests/test_chunk.cpp b/cpp/tests/test_chunk.cpp index 79ea8ab9d..3434a5d9c 100644 --- a/cpp/tests/test_chunk.cpp +++ b/cpp/tests/test_chunk.cpp @@ -16,6 +16,7 @@ #include #include #include +#include using namespace rapidsmpf; using namespace rapidsmpf::shuffler; @@ -24,7 +25,9 @@ using namespace rapidsmpf::shuffler::detail; class ChunkTest : public ::testing::Test { protected: void SetUp() override { - br = BufferResource::create(cudf::get_current_device_resource_ref()); + br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref() + ); stream = cudf::get_default_stream(); } diff --git a/cpp/tests/test_communicator.cpp b/cpp/tests/test_communicator.cpp index d1c2eea54..376488702 100644 --- a/cpp/tests/test_communicator.cpp +++ b/cpp/tests/test_communicator.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include "environment.hpp" #include "utils.hpp" @@ -24,7 +25,7 @@ class BaseCommunicatorTest : public ::testing::Test { void SetUp() override { comm = GlobalEnvironment->comm_.get(); mr = std::make_unique(); - br = rapidsmpf::BufferResource::create(*mr); + br = rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), *mr); stream = rmm::cuda_stream_default; } diff --git a/cpp/tests/test_config.cpp b/cpp/tests/test_config.cpp index c78d840cf..91f1c5667 100644 --- a/cpp/tests/test_config.cpp +++ b/cpp/tests/test_config.cpp @@ -635,7 +635,7 @@ TEST(OptionsTest, BufferResourceFromOptionsUsesDefaultWhenOptionsEmpty) { rmm::mr::cuda_memory_resource cuda_mr; RmmResourceAdaptor mr{cuda_mr}; - auto br = BufferResource::from_options(mr, opts); + auto br = BufferResource::from_options(mr, opts, Statistics::from_options(opts)); EXPECT_FALSE(br->statistics()->enabled()); EXPECT_EQ(br->stream_pool().get_pool_size(), 16); auto [_, total_mem] = rmm::available_device_memory(); @@ -662,7 +662,7 @@ TEST(OptionsTest, BufferResourceFromOptionsAcceptsPercentageForDeviceLimit) { rmm::mr::cuda_memory_resource cuda_mr; RmmResourceAdaptor mr{cuda_mr}; - auto br = BufferResource::from_options(mr, opts); + auto br = BufferResource::from_options(mr, opts, Statistics::from_options(opts)); // Verify device memory limit is 50% of total auto [_, total_mem] = rmm::available_device_memory(); @@ -680,7 +680,7 @@ TEST(OptionsTest, BufferResourceFromOptionsEnablesPinnedMemoryWhenSupported) { rmm::mr::cuda_memory_resource cuda_mr; RmmResourceAdaptor mr{cuda_mr}; - auto br = BufferResource::from_options(mr, opts); + auto br = BufferResource::from_options(mr, opts, Statistics::from_options(opts)); // Should not throw when accessing pinned_mr EXPECT_NO_THROW(std::ignore = br->pinned_mr()); @@ -696,11 +696,9 @@ TEST(OptionsTest, ContextFromOptionsCreatesInstanceWithExplicitOptions) { rmm::mr::cuda_memory_resource cuda_mr; RmmResourceAdaptor mr{cuda_mr}; - auto comm = - std::make_shared(opts, std::make_shared()); - auto ctx = streaming::Context::from_options( - mr, comm->logger(), opts, Statistics::from_options(opts) - ); + auto stats = Statistics::from_options(opts); + auto comm = std::make_shared(opts, std::make_shared(stats)); + auto ctx = streaming::Context::from_options(mr, comm->logger(), opts, stats); ASSERT_NE(ctx, nullptr); EXPECT_TRUE(ctx->statistics()->enabled()); @@ -712,9 +710,9 @@ TEST(OptionsTest, ContextFromOptionsUsesDefaultWhenOptionsEmpty) { rmm::mr::cuda_memory_resource cuda_mr; RmmResourceAdaptor mr{cuda_mr}; - auto comm = - std::make_shared(opts, std::make_shared()); - auto ctx = streaming::Context::from_options(mr, comm->logger(), opts); + auto stats = Statistics::from_options(opts); + auto comm = std::make_shared(opts, std::make_shared(stats)); + auto ctx = streaming::Context::from_options(mr, comm->logger(), opts, stats); ASSERT_NE(ctx, nullptr); EXPECT_FALSE(ctx->statistics()->enabled()); @@ -727,11 +725,9 @@ TEST(OptionsTest, ContextFromOptionsEnablesStatisticsWhenRequested) { rmm::mr::cuda_memory_resource cuda_mr; RmmResourceAdaptor mr{cuda_mr}; - auto comm = - std::make_shared(opts, std::make_shared()); - auto ctx = streaming::Context::from_options( - mr, comm->logger(), opts, Statistics::from_options(opts) - ); + auto stats = Statistics::from_options(opts); + auto comm = std::make_shared(opts, std::make_shared(stats)); + auto ctx = streaming::Context::from_options(mr, comm->logger(), opts, stats); ASSERT_NE(ctx, nullptr); EXPECT_TRUE(ctx->statistics()->enabled()); @@ -742,9 +738,9 @@ TEST(OptionsTest, ContextFromOptionsCreatesProgressThread) { rmm::mr::cuda_memory_resource cuda_mr; RmmResourceAdaptor mr{cuda_mr}; - auto comm = - std::make_shared(opts, std::make_shared()); - auto ctx = streaming::Context::from_options(mr, comm->logger(), opts); + auto stats = Statistics::from_options(opts); + auto comm = std::make_shared(opts, std::make_shared(stats)); + auto ctx = streaming::Context::from_options(mr, comm->logger(), opts, stats); ASSERT_NE(ctx, nullptr); } @@ -754,9 +750,9 @@ TEST(OptionsTest, ContextFromOptionsCreatesExecutor) { rmm::mr::cuda_memory_resource cuda_mr; RmmResourceAdaptor mr{cuda_mr}; - auto comm = - std::make_shared(opts, std::make_shared()); - auto ctx = streaming::Context::from_options(mr, comm->logger(), opts); + auto stats = Statistics::from_options(opts); + auto comm = std::make_shared(opts, std::make_shared(stats)); + auto ctx = streaming::Context::from_options(mr, comm->logger(), opts, stats); ASSERT_NE(ctx, nullptr); EXPECT_NE(ctx->executor(), nullptr); diff --git a/cpp/tests/test_metadata_payload_exchange.cpp b/cpp/tests/test_metadata_payload_exchange.cpp index fe6d647e4..1739d0999 100644 --- a/cpp/tests/test_metadata_payload_exchange.cpp +++ b/cpp/tests/test_metadata_payload_exchange.cpp @@ -27,9 +27,9 @@ class MetadataPayloadExchangeTest : public ::testing::Test { void SetUp() override { comm = GlobalEnvironment->comm_.get(); mr = std::make_unique(); - br = BufferResource::create(*mr); - stream = rmm::cuda_stream_default; statistics = Statistics::create(); + br = BufferResource::create(statistics, *mr); + stream = rmm::cuda_stream_default; auto allocate_fn = [this](std::size_t size) { return allocate_receive_buffer(size); diff --git a/cpp/tests/test_partition.cpp b/cpp/tests/test_partition.cpp index b4e68a247..0611547ef 100644 --- a/cpp/tests/test_partition.cpp +++ b/cpp/tests/test_partition.cpp @@ -44,7 +44,7 @@ TEST_P(NumOfPartitions, partition_and_pack) { std::int64_t const seed = 42; cudf::hash_id const hash_fn = cudf::hash_id::HASH_MURMUR3; auto stream = cudf::get_default_stream(); - auto br = rapidsmpf::BufferResource::create(mr()); + auto br = BufferResource::create(Statistics::disabled(), mr()); cudf::table expect = random_table_with_index(seed, static_cast(num_rows), 0, 10); @@ -73,7 +73,9 @@ TEST_P(NumOfPartitions, split_and_pack) { int const num_rows = std::get<1>(GetParam()); std::int64_t const seed = 42; auto stream = cudf::get_default_stream(); - auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + auto br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref() + ); cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); @@ -101,7 +103,9 @@ TEST_P(NumOfPartitions, split_and_pack) { class SpillingTest : public ::testing::Test { protected: void SetUp() override { - br = BufferResource::create(cudf::get_current_device_resource_ref()); + br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref() + ); stream = cudf::get_default_stream(); } diff --git a/cpp/tests/test_progress_thread.cpp b/cpp/tests/test_progress_thread.cpp index 1435fe2ec..32a3291c2 100644 --- a/cpp/tests/test_progress_thread.cpp +++ b/cpp/tests/test_progress_thread.cpp @@ -10,6 +10,7 @@ #include +#include #include #include @@ -89,7 +90,7 @@ TEST_P(ProgressThreadEvents, events) { } TEST(ProgressThreadTests, RemoveFunctionWithDelayedPause) { - ProgressThread progress_thread{}; + ProgressThread progress_thread{rapidsmpf::Statistics::disabled()}; // add a function to the progress thread that never completes auto id = progress_thread.add_function([] { @@ -109,3 +110,26 @@ TEST(ProgressThreadTests, RemoveFunctionWithDelayedPause) { future.get(); } + +TEST(ProgressThreadTests, SetStatisticsSingleComm) { + if (GlobalEnvironment->type() != TestEnvironmentType::SINGLE) { + GTEST_SKIP() << "This test is only valid for single communicator"; + } + + rapidsmpf::config::Options options{}; + auto stats = rapidsmpf::Statistics::create(); + auto progress = std::make_shared(stats); + auto comm = std::make_shared(options, progress); + + auto check_stats = [&](const auto* stats_ptr) { + EXPECT_EQ(stats_ptr, progress->statistics().get()); + EXPECT_EQ(stats_ptr, comm->statistics().get()); + }; + + check_stats(stats.get()); + + auto new_stats = rapidsmpf::Statistics::create(); + progress->set_statistics(new_stats); + + check_stats(new_stats.get()); +} diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index 024820a90..c397b30b7 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include "environment.hpp" @@ -32,7 +33,7 @@ extern Environment* GlobalEnvironment; TEST(ReceivedChunks, spill_skips_control_messages) { auto mr = cudf::get_current_device_resource_ref(); - auto br = rapidsmpf::BufferResource::create(mr); + auto br = rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), mr); rapidsmpf::shuffler::detail::ReceivedChunks received; @@ -49,7 +50,7 @@ TEST(ReceivedChunks, spill_skips_control_messages) { TEST(ReceivedChunks, spill_respects_amount) { auto mr = cudf::get_current_device_resource_ref(); - auto br = rapidsmpf::BufferResource::create(mr); + auto br = rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), mr); auto stream = cudf::get_default_stream(); rapidsmpf::shuffler::detail::ReceivedChunks received; @@ -75,7 +76,7 @@ TEST(ReceivedChunks, spill_respects_amount) { TEST(MetadataMessage, round_trip) { auto stream = cudf::get_default_stream(); auto mr = cudf::get_current_device_resource_ref(); - auto br = rapidsmpf::BufferResource::create(mr); + auto br = rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), mr); auto metadata = iota_vector(100); @@ -238,7 +239,10 @@ class MemoryLimits_NumPartition total_num_partitions = std::get<1>(GetParam()); total_num_rows = std::get<2>(GetParam()); br = rapidsmpf::BufferResource::create( - mr(), rapidsmpf::PinnedMemoryResource::Disabled, memory_limits + rapidsmpf::Statistics::disabled(), + mr(), + rapidsmpf::PinnedMemoryResource::Disabled, + memory_limits ); shuffler = std::make_unique( @@ -309,7 +313,7 @@ class ConcurrentShuffleTest static_cast(std::get<1>(GetParam())); // these resources will be used by multiple threads to instantiate shufflers - br = rapidsmpf::BufferResource::create(mr()); + br = rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), mr()); stream = cudf::get_default_stream(); } @@ -415,6 +419,7 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { constexpr std::int64_t k_no_spill_limit = (1LL << 40); constexpr std::int64_t k_force_spill_limit = -(1LL << 40); auto br = rapidsmpf::BufferResource::create( + rapidsmpf::Statistics::disabled(), mr, rapidsmpf::PinnedMemoryResource::Disabled, {{rapidsmpf::MemoryType::DEVICE, k_no_spill_limit}}, @@ -676,7 +681,7 @@ TEST(Shuffler, ShutdownWhilePaused) { auto progress_thread = GlobalEnvironment->comm_->progress_thread(); auto mr = cudf::get_current_device_resource_ref(); - auto br = rapidsmpf::BufferResource::create(mr); + auto br = rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), mr); auto shuffler = rapidsmpf::shuffler::Shuffler(GlobalEnvironment->comm_, 0, 1, br.get()); @@ -711,7 +716,7 @@ class ExtractEmptyPartitionsTest : public cudf::test::BaseFixture { void SetUp() override { stream = cudf::get_default_stream(); - br = rapidsmpf::BufferResource::create(mr()); + br = rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), mr()); shuffler = std::make_unique( GlobalEnvironment->comm_, 0, nparts, br.get() @@ -810,7 +815,9 @@ TEST_F(ExtractEmptyPartitionsTest, SomeEmptyAndNonEmptyInsertions) { TEST(ShufflerTest, multiple_shutdowns) { auto& comm = GlobalEnvironment->comm_; - auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + auto br = rapidsmpf::BufferResource::create( + rapidsmpf::Statistics::disabled(), cudf::get_current_device_resource_ref() + ); auto shuffler = std::make_unique( comm, 0, comm->nranks(), br.get() ); @@ -835,7 +842,9 @@ TEST(ShufflerTest, multiple_shutdowns) { TEST(Shuffler, concurrent_wait) { auto const& comm = GlobalEnvironment->comm_; auto stream = cudf::get_default_stream(); - auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + auto br = rapidsmpf::BufferResource::create( + rapidsmpf::Statistics::disabled(), cudf::get_current_device_resource_ref() + ); // Use more partitions than ranks so each rank owns multiple partitions, ensuring // multiple threads call wait() concurrently on the same shuffler. @@ -939,7 +948,7 @@ TEST(Shuffler, opid_reuse) { constexpr auto wait_timeout = std::chrono::seconds{30}; rmm::mr::cuda_memory_resource mr; - auto br = rapidsmpf::BufferResource::create(mr); + auto br = rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), mr); // On rank 0, wrap the device MR with a delayed version for the shuffler. std::unique_ptr delayed_mr; @@ -948,7 +957,9 @@ TEST(Shuffler, opid_reuse) { if (comm->rank() == 0) { delayed_mr = std::make_unique(mr, std::chrono::milliseconds(500)); - delayed_br = rapidsmpf::BufferResource::create(*delayed_mr); + delayed_br = rapidsmpf::BufferResource::create( + rapidsmpf::Statistics::disabled(), *delayed_mr + ); shuffler_br = delayed_br.get(); } @@ -1046,7 +1057,7 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) { constexpr auto wait_timeout = std::chrono::seconds{30}; rmm::mr::cuda_memory_resource mr; - auto br = rapidsmpf::BufferResource::create(mr); + auto br = rapidsmpf::BufferResource::create(rapidsmpf::Statistics::disabled(), mr); // On rank 0, wrap the device MR with a delayed version for the shuffler. std::unique_ptr delayed_mr; @@ -1055,7 +1066,9 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) { if (comm->rank() == 0) { delayed_mr = std::make_unique(mr, std::chrono::milliseconds(500)); - delayed_br = rapidsmpf::BufferResource::create(*delayed_mr); + delayed_br = rapidsmpf::BufferResource::create( + rapidsmpf::Statistics::disabled(), *delayed_mr + ); shuffler_br = delayed_br.get(); } diff --git a/cpp/tests/test_shuffler_many_streams.cpp b/cpp/tests/test_shuffler_many_streams.cpp index 43c3f6402..41a432412 100644 --- a/cpp/tests/test_shuffler_many_streams.cpp +++ b/cpp/tests/test_shuffler_many_streams.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include "environment.hpp" #include "utils.hpp" @@ -50,7 +51,9 @@ TEST(ShufflerManyStreams, Test) { std::mt19937 random_generator{42}; constexpr std::size_t chunksize = 1 << 20; constexpr int num_partitions = 100; - auto br = BufferResource::create(cudf::get_current_device_resource_ref()); + auto br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref() + ); // Create a CUDA stream for each partition. // To stress-test stream handling, assign random priorities so streams are more diff --git a/cpp/tests/test_sparse_alltoall.cpp b/cpp/tests/test_sparse_alltoall.cpp index 325060937..c307e2ee2 100644 --- a/cpp/tests/test_sparse_alltoall.cpp +++ b/cpp/tests/test_sparse_alltoall.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include "environment.hpp" @@ -105,7 +106,9 @@ int decode_payload(rapidsmpf::PackedData const& packed_data) { class SparseAlltoallTest : public ::testing::Test { protected: void SetUp() override { - br = rapidsmpf::BufferResource::create(rmm::mr::cuda_memory_resource{}); + br = rapidsmpf::BufferResource::create( + rapidsmpf::Statistics::disabled(), rmm::mr::cuda_memory_resource{} + ); } std::shared_ptr br; diff --git a/cpp/tests/test_spill_manager.cpp b/cpp/tests/test_spill_manager.cpp index 401b26f2f..3d2f7745f 100644 --- a/cpp/tests/test_spill_manager.cpp +++ b/cpp/tests/test_spill_manager.cpp @@ -29,6 +29,7 @@ TEST(SpillManager, SpillFunction) { // currently configured limit. std::int64_t mem_available = 10_KiB; auto br = BufferResource::create( + Statistics::disabled(), cudf::get_current_device_resource_ref(), rapidsmpf::PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, mem_available}} diff --git a/cpp/tests/test_stream_ordered_timing.cpp b/cpp/tests/test_stream_ordered_timing.cpp index 7661b166d..1345733cc 100644 --- a/cpp/tests/test_stream_ordered_timing.cpp +++ b/cpp/tests/test_stream_ordered_timing.cpp @@ -134,13 +134,16 @@ TEST(StreamOrderedTiming, StreamDelay) { TEST(StreamOrderedTiming, StatisticsDestroyedBeforeStreamSync) { // If the Statistics object is destroyed before the stream callbacks execute, // the callbacks detect the expired weak_ptr and return without crashing. + // `StreamOrderedTiming` only holds a weak reference, so destroying the last + // shared_ptr here destroys `Statistics` immediately and `~Statistics` runs + // `cancel_inflight_timings(this)` to drop the in-flight entry. rmm::cuda_stream stream; { auto stats = Statistics::create(); StreamOrderedTiming timing{stream.view(), stats}; timing.stop_and_record("my-timing"); - // Both `timing` and `stats` go out of scope here. The shared_ptr count - // drops to zero, destroying Statistics before the stream callbacks run. + // `timing` and `stats` go out of scope here; `~Statistics` then + // cancels the still-pending stream callback. } EXPECT_NO_THROW(stream.synchronize()); } diff --git a/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py b/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py index c57af8cce..4a1f1f763 100644 --- a/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py +++ b/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py @@ -259,7 +259,7 @@ def setup_and_run(args: argparse.Namespace) -> None: memory_limits = ( None if args.spill_device is None else {MemoryType.DEVICE: args.spill_device} ) - br = BufferResource(mr, memory_limits=memory_limits, statistics=stats) + br = BufferResource(stats, mr, memory_limits=memory_limits) args.out_nparts = args.out_nparts if args.out_nparts is not None else comm.nranks args.part_size = args.part_size if args.part_size is not None else args.local_size diff --git a/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd b/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd index 86d7ee1aa..5b3079801 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd @@ -7,6 +7,7 @@ from libcpp.string cimport string from rapidsmpf._detail.exception_handling cimport ex_handler from rapidsmpf.progress_thread cimport cpp_ProgressThread +from rapidsmpf.statistics cimport cpp_Statistics cdef extern from "" namespace \ @@ -36,6 +37,7 @@ cdef extern from "" nogil: Rank nranks() except +ex_handler string str() except +ex_handler shared_ptr[cpp_ProgressThread] progress_thread() except +ex_handler + shared_ptr[cpp_Statistics] statistics() except +ex_handler shared_ptr[cpp_Logger] logger() cdef class Communicator: diff --git a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi index 66d0e258c..ba1be8f53 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi @@ -3,6 +3,7 @@ from enum import IntEnum from rapidsmpf.progress_thread import ProgressThread +from rapidsmpf.statistics import Statistics class LOG_LEVEL(IntEnum): NONE = ... @@ -30,6 +31,8 @@ class Communicator: def logger(self) -> Logger: ... @property def progress_thread(self) -> ProgressThread: ... + @property + def statistics(self) -> Statistics: ... def get_str(self) -> str: ... def _available_communicators() -> tuple[str, ...]: ... diff --git a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx index 884245fb5..8995c4cbd 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx @@ -6,6 +6,7 @@ from libcpp.string cimport string from libcpp.utility cimport move from rapidsmpf.progress_thread cimport ProgressThread +from rapidsmpf.statistics cimport Statistics cdef class Logger: @@ -179,6 +180,19 @@ cdef class Communicator: pt._handle = deref(self._handle).progress_thread() return pt + @property + def statistics(self): + """ + Get the statistics instance associated with this communicator. + + Returns + ------- + The Statistics instance. + """ + cdef Statistics stats = Statistics.__new__(Statistics) + stats._handle = deref(self._handle).statistics() + return stats + def get_str(self): """ Get a string representation of the communicator. diff --git a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py index 149217c05..945d95ddf 100644 --- a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py +++ b/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py @@ -162,7 +162,7 @@ def bulk_mpi_shuffle( baseline Whether to skip the shuffle and run a simple IO baseline. statistics - The statistics instance to use. If None, statistics is disabled. + The statistics instance to use. If None, ``Statistics.disabled()`` is used. Notes ----- @@ -196,7 +196,10 @@ def bulk_mpi_shuffle( columns, ) else: - br = BufferResource(rmm.mr.get_current_device_resource()) + br = BufferResource( + statistics or Statistics.disabled(), + rmm.mr.get_current_device_resource(), + ) shuffler = Shuffler( comm, op_id=0, @@ -316,9 +319,8 @@ def setup_and_run(args: argparse.Namespace) -> None: memory_limits = ( None if args.spill_device is None else {MemoryType.DEVICE: args.spill_device} ) - br = BufferResource(mr, memory_limits=memory_limits) - stats = Statistics(enable=args.statistics) + br = BufferResource(stats, mr, memory_limits=memory_limits) progress_thread = ProgressThread(stats) if args.cluster_type == "mpi": diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py index dcdf5f3fd..af137210a 100644 --- a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py +++ b/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py @@ -93,9 +93,10 @@ def __init__( if self.spill_device is None else {MemoryType.DEVICE: self.spill_device} ) - br = BufferResource(self.mr, memory_limits=memory_limits) + stats = Statistics(enable=enable_statistics) + br = BufferResource(stats, self.mr, memory_limits=memory_limits) self.br = br - super().__init__(nranks, Statistics(enable=enable_statistics)) + super().__init__(nranks, stats) def setup_worker(self, root_address_bytes: bytes) -> None: """ diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py b/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py index 526b1c2f6..01eb0c2d8 100644 --- a/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py +++ b/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py @@ -21,6 +21,7 @@ from rapidsmpf.integrations.ray import RapidsMPFActor, setup_ray_ucxx_cluster from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.shuffler import Shuffler +from rapidsmpf.statistics import Statistics from rapidsmpf.testing import assert_eq @@ -89,7 +90,7 @@ def run(self) -> None: columns_to_hash = (1,) mr = rmm.mr.get_current_device_resource() - br = BufferResource(mr) + br = BufferResource(Statistics.disabled(), mr) stream = DEFAULT_STREAM # use the default stream # Calculate the expected output partitions on all ranks diff --git a/python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py b/python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py index 0f10b9735..477b1b4c4 100644 --- a/python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py +++ b/python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py @@ -16,6 +16,7 @@ from rapidsmpf.config import Options, get_environment_variables from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.progress_thread import ProgressThread +from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.core.actor import ( define_actor, run_actor_network, @@ -38,10 +39,11 @@ def main() -> int: options = Options(get_environment_variables()) # Create a communicator and context that will be used by all streaming actors. - comm = single_process_comm(options, ProgressThread()) + stats = Statistics.from_options(options) + comm = single_process_comm(options, ProgressThread(stats)) ctx = Context( logger=comm.logger, - br=BufferResource(rmm.mr.get_current_device_resource()), + br=BufferResource(stats, rmm.mr.get_current_device_resource()), options=options, ) diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd index 0dcb45c4a..a188315ec 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pxd @@ -33,12 +33,12 @@ cdef extern from "" nogil: cdef cppclass cpp_BufferResource "rapidsmpf::BufferResource": @staticmethod shared_ptr[cpp_BufferResource] create( + shared_ptr[cpp_Statistics], any_resource[device_accessible], optional[cpp_PinnedMemoryResource], unordered_map[MemoryType, int64_t], optional[cpp_Duration], - shared_ptr[cuda_stream_pool], - shared_ptr[cpp_Statistics], + shared_ptr[cuda_stream_pool] ) except +ex_handler size_t memory_reserved(MemoryType mem_type) except +ex_handler int64_t memory_available(MemoryType mem_type) except +ex_handler diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi index 5f7ded0fa..429c668f0 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi @@ -17,20 +17,20 @@ from rapidsmpf.statistics import Statistics class BufferResource: def __init__( self, + statistics: Statistics, device_mr: DeviceMemoryResource, *, pinned_mr: PinnedMemoryResource | None = None, memory_limits: Mapping[MemoryType, int] | None = None, periodic_spill_check: float | None = 1e-3, stream_pool: CudaStreamPool | None = None, - statistics: Statistics | None = None, ) -> None: ... @classmethod def from_options( cls: type[Self], + statistics: Statistics, mr: DeviceMemoryResource, options: Options, - statistics: Statistics | None = None, ) -> Self: ... @property def device_mr(self) -> OwningDeviceMemoryResource: ... diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index cdaf9f8cb..589bfa212 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -136,6 +136,9 @@ cdef class BufferResource: Parameters ---------- + statistics + The statistics instance to use. Required. Pass ``Statistics.disabled()`` for a + no-op Statistics instance. device_mr The RMM device memory resource used for device allocations. To ensure allocations are tracked for memory-limit accounting and statistics, use @@ -158,9 +161,6 @@ cdef class BufferResource: Optional CUDA stream pool to use. If None, a new pool with 16 streams will be created. Must be an instance of ``rmm.pylibrmm.cuda_stream_pool.CudaStreamPool``. - statistics - The statistics instance to use. If None, a disabled statistics instance - will be created. Notes ----- @@ -177,13 +177,13 @@ cdef class BufferResource: """ def __cinit__( self, + Statistics statistics not None, DeviceMemoryResource device_mr not None, *, PinnedMemoryResource pinned_mr = None, memory_limits = None, periodic_spill_check = 1e-3, stream_pool = None, - statistics = None, ): cdef unordered_map[MemoryType, int64_t] _mem_limits if memory_limits is not None: @@ -217,13 +217,9 @@ cdef class BufferResource: ) ) - if statistics is None: - statistics = Statistics(enable=False) - # Keep statistics alive self._statistics = statistics - # checked cast requires the GIL - stats_handle = (statistics)._handle + stats_handle = statistics._handle # Keep the original Python memory resources alive while the C++ # BufferResource holds resources derived from them. These anchors are @@ -240,21 +236,21 @@ cdef class BufferResource: cpp_pinned_mr = self._pinned_mr._handle with nogil: self._handle = cpp_BufferResource.create( + stats_handle, make_any_device_resource(device_mr.get_mr()), cpp_pinned_mr, move(_mem_limits), period, cpp_stream_pool, - stats_handle, ) self.spill_manager = SpillManager._create(self) @classmethod def from_options( cls, + Statistics statistics not None, DeviceMemoryResource mr not None, Options options not None, - statistics=None, ): """ Construct a BufferResource from configuration options. @@ -265,28 +261,26 @@ cdef class BufferResource: Parameters ---------- + statistics + The statistics instance to use. The caller is responsible for creating and + owning this object. Pass ``Statistics.disabled()`` for a no-op recorder. mr A device-accessible RMM memory resource. options Configuration options. - statistics - The statistics instance to use. The caller is responsible for creating and - owning this object. Defaults to ``Statistics.disabled()``. Returns ------- A BufferResource instance configured according to the options. """ - if statistics is None: - statistics = Statistics.disabled() cdef PinnedMemoryResource pinned_mr = PinnedMemoryResource.from_options(options) return cls( + statistics=statistics, device_mr=mr, pinned_mr=pinned_mr, memory_limits={MemoryType.DEVICE: device_limit_from_options(options)}, periodic_spill_check=periodic_spill_check_from_options(options), stream_pool=stream_pool_from_options(options), - statistics=statistics, ) def __dealloc__(self): diff --git a/python/rapidsmpf/rapidsmpf/progress_thread.pxd b/python/rapidsmpf/rapidsmpf/progress_thread.pxd index 7ee796b09..344643063 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pxd +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pxd @@ -17,6 +17,10 @@ cdef extern from "" nogil: cpp_ProgressThread( shared_ptr[cpp_Statistics] statistics, ) except +ex_handler + shared_ptr[cpp_Statistics] statistics() except +ex_handler + void set_statistics( + shared_ptr[cpp_Statistics] statistics + ) except +ex_handler cdef class ProgressThread: diff --git a/python/rapidsmpf/rapidsmpf/progress_thread.pyi b/python/rapidsmpf/rapidsmpf/progress_thread.pyi index fba04eb35..d49955a3c 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pyi +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pyi @@ -5,4 +5,7 @@ from __future__ import annotations from rapidsmpf.statistics import Statistics class ProgressThread: - def __init__(self, statistics: Statistics | None = None) -> None: ... + def __init__(self, statistics: Statistics) -> None: ... + @property + def statistics(self) -> Statistics: ... + def set_statistics(self, statistics: Statistics) -> None: ... diff --git a/python/rapidsmpf/rapidsmpf/progress_thread.pyx b/python/rapidsmpf/rapidsmpf/progress_thread.pyx index f4deee800..ebb0b5437 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pyx +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pyx @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """The ProgressThread interface for RapidsMPF.""" +from cython.operator cimport dereference as deref from libcpp.memory cimport make_shared from rapidsmpf.statistics cimport Statistics @@ -19,7 +20,8 @@ cdef class ProgressThread: Parameters ---------- statistics - The statistics instance to use. If None, statistics is disabled. + The statistics instance to use. Required. Pass + ``Statistics.disabled()`` for a no-op recorder. Notes ----- @@ -29,14 +31,39 @@ cdef class ProgressThread: """ def __init__( self, - Statistics statistics = None, + Statistics statistics not None, ): - if statistics is None: - statistics = Statistics(enable=False) # Disables statistics. - with nogil: self._handle = make_shared[cpp_ProgressThread](statistics._handle) + @property + def statistics(self): + """ + Get the statistics instance associated with this progress thread. + + Returns + ------- + The Statistics instance. + """ + cdef Statistics stats = Statistics.__new__(Statistics) + stats._handle = deref(self._handle).statistics() + return stats + + def set_statistics(self, Statistics statistics not None): + """ + Replace the statistics instance held by this progress thread. + + The swap is performed atomically. + + Parameters + ---------- + statistics + The new statistics instance. Pass ``Statistics.disabled()`` to opt out of + statistics collection. + """ + with nogil: + deref(self._handle).set_statistics(statistics._handle) + def __dealloc__(self): with nogil: self._handle.reset() diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyx b/python/rapidsmpf/rapidsmpf/statistics.pyx index 0d64d0ad8..4a77625dc 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pyx +++ b/python/rapidsmpf/rapidsmpf/statistics.pyx @@ -32,6 +32,9 @@ cdef extern from "" nogil: cpp_Options options, ) except +ex_handler + cdef shared_ptr[cpp_Statistics] cpp_disabled \ + "rapidsmpf::Statistics::disabled"() except +ex_handler + cdef extern from *: """ diff --git a/python/rapidsmpf/rapidsmpf/streaming/core/context.pyi b/python/rapidsmpf/rapidsmpf/streaming/core/context.pyi index 4ed3e4a65..8408924d3 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/core/context.pyi +++ b/python/rapidsmpf/rapidsmpf/streaming/core/context.pyi @@ -31,7 +31,7 @@ class Context: logger: Logger, mr: RmmResourceAdaptor, options: Options, - statistics: Statistics | None = None, + statistics: Statistics, ) -> Self: ... def __enter__(self) -> Context: ... def __exit__( diff --git a/python/rapidsmpf/rapidsmpf/streaming/core/context.pyx b/python/rapidsmpf/rapidsmpf/streaming/core/context.pyx index 512bd6f22..5920e0a8d 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/core/context.pyx +++ b/python/rapidsmpf/rapidsmpf/streaming/core/context.pyx @@ -15,12 +15,12 @@ from libcpp.memory cimport make_shared from rmm.pylibrmm.stream cimport Stream from rapidsmpf.rmm_resource_adaptor cimport RmmResourceAdaptor +from rapidsmpf.statistics cimport Statistics from rapidsmpf.streaming.core.channel cimport Channel, cpp_Channel from rapidsmpf.streaming.core.memory_reserve_or_wait cimport \ MemoryReserveOrWait from rapidsmpf.memory.buffer import MemoryType as py_MemoryType -from rapidsmpf.statistics import Statistics @no_gc_clear @@ -99,13 +99,11 @@ cdef class Context: Logger logger not None, RmmResourceAdaptor mr not None, Options options not None, - statistics=None, + Statistics statistics not None, ): - if statistics is None: - statistics = Statistics.disabled() return cls( logger=logger, - br=BufferResource.from_options(mr, options, statistics), + br=BufferResource.from_options(statistics, mr, options), options=options, ) diff --git a/python/rapidsmpf/rapidsmpf/tests/conftest.py b/python/rapidsmpf/rapidsmpf/tests/conftest.py index db2618660..566d59c3f 100644 --- a/python/rapidsmpf/rapidsmpf/tests/conftest.py +++ b/python/rapidsmpf/rapidsmpf/tests/conftest.py @@ -12,6 +12,7 @@ from rapidsmpf.communicator import COMMUNICATORS from rapidsmpf.config import Options, get_environment_variables from rapidsmpf.progress_thread import ProgressThread +from rapidsmpf.statistics import Statistics if TYPE_CHECKING: from collections.abc import Generator @@ -86,8 +87,11 @@ def _mpi_comm(*, _mpi_disabled: bool) -> Communicator: from rapidsmpf.communicator.mpi import new_communicator + options = Options(get_environment_variables()) return new_communicator( - MPI.COMM_WORLD, Options(get_environment_variables()), ProgressThread() + MPI.COMM_WORLD, + options, + ProgressThread(Statistics.from_options(options)), ) @@ -103,7 +107,12 @@ def _ucxx_comm() -> Communicator: """ from rapidsmpf.communicator.testing import ucxx_mpi_setup - return ucxx_mpi_setup(None, Options(get_environment_variables()), ProgressThread()) + options = Options(get_environment_variables()) + return ucxx_mpi_setup( + None, + options, + ProgressThread(Statistics.from_options(options)), + ) @pytest.fixture( diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/conftest.py b/python/rapidsmpf/rapidsmpf/tests/streaming/conftest.py index 879b73ca2..7724c9055 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/conftest.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/conftest.py @@ -11,6 +11,7 @@ from rapidsmpf.config import Options, get_environment_variables from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor +from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.core.context import Context if TYPE_CHECKING: @@ -26,7 +27,7 @@ def context(comm: Communicator) -> Generator[Context, None, None]: """ options = Options(get_environment_variables()) mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) - br = BufferResource(mr) + br = BufferResource(Statistics.from_options(options), mr) with Context(comm.logger, br, options) as ctx: yield ctx diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_memory_reserve_or_wait.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_memory_reserve_or_wait.py index a08e033f6..c523b65ca 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_memory_reserve_or_wait.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_memory_reserve_or_wait.py @@ -17,6 +17,7 @@ from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.progress_thread import ProgressThread +from rapidsmpf.statistics import Statistics from rapidsmpf.streaming.core.actor import define_actor, run_actor_network from rapidsmpf.streaming.core.context import Context from rapidsmpf.streaming.core.memory_reserve_or_wait import ( @@ -32,8 +33,10 @@ def make_context( if overwrite_options is not None: env.update(overwrite_options) options = Options(env) - comm = single_process_comm(options, ProgressThread()) + stats = Statistics.from_options(options) + comm = single_process_comm(options, ProgressThread(stats)) br = BufferResource( + stats, rmm.mr.CudaMemoryResource(), memory_limits={MemoryType.DEVICE: dev_limit}, ) diff --git a/python/rapidsmpf/rapidsmpf/tests/test_allgather.py b/python/rapidsmpf/rapidsmpf/tests/test_allgather.py index 93ea25f34..dcef9450d 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_allgather.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_allgather.py @@ -135,7 +135,7 @@ def test_basic_allgather( Each rank inserts n_inserts pieces of data, then all ranks should receive all data from all ranks. """ - br = BufferResource(device_mr) + br = BufferResource(comm.statistics, device_mr) # Create AllGather instance allgather = AllGather( @@ -197,7 +197,7 @@ def test_insert_finished_raises_in_context( device_mr: rmm.mr.CudaMemoryResource, ) -> None: """Test that insert_finished raises when called inside a context manager.""" - br = BufferResource(device_mr) + br = BufferResource(comm.statistics, device_mr) ag = AllGather(comm=comm, op_id=0, br=br) with ( ag, diff --git a/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource.py b/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource.py index 95cabae14..a6cdf87ed 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource.py @@ -21,7 +21,9 @@ def KiB(x: int) -> int: def test_buffer_resource() -> None: mr = rmm.mr.CudaMemoryResource() - br = BufferResource(mr, memory_limits={MemoryType.DEVICE: KiB(100)}) + br = BufferResource( + Statistics.disabled(), mr, memory_limits={MemoryType.DEVICE: KiB(100)} + ) assert br.memory_reserved(MemoryType.DEVICE) == 0 assert br.memory_reserved(MemoryType.HOST) == 0 @@ -32,7 +34,7 @@ def test_buffer_resource() -> None: @pytest.mark.parametrize("mem_type", [MemoryType.DEVICE, MemoryType.HOST]) def test_memory_reservation(mem_type: MemoryType) -> None: mr = rmm.mr.CudaMemoryResource() - br = BufferResource(mr, memory_limits={mem_type: KiB(100)}) + br = BufferResource(Statistics.disabled(), mr, memory_limits={mem_type: KiB(100)}) res1, ob = br.reserve(mem_type, KiB(100), allow_overbooking=False) assert res1.br is br assert res1.mem_type == mem_type @@ -81,12 +83,12 @@ def test_stream_pool() -> None: mr = rmm.mr.CudaMemoryResource() # Test with default stream pool size (16) - br_default = BufferResource(mr) + br_default = BufferResource(Statistics.disabled(), mr) assert br_default.stream_pool_size() == 16 # Test with custom stream pool custom_pool = rmm.pylibrmm.cuda_stream_pool.CudaStreamPool(pool_size=32) - br_custom = BufferResource(mr, stream_pool=custom_pool) + br_custom = BufferResource(Statistics.disabled(), mr, stream_pool=custom_pool) assert br_custom.stream_pool_size() == 32 @@ -94,20 +96,20 @@ def test_statistics() -> None: """Test that statistics parameter can be configured.""" mr = rmm.mr.CudaMemoryResource() - # Disabled by default - br_default = BufferResource(mr) + # Disabled + br_default = BufferResource(Statistics.disabled(), mr) assert not br_default.statistics.enabled # Test with enabled statistics (with memory profiling) stats = Statistics(enable=True) - br_with_mr = BufferResource(mr, statistics=stats) + br_with_mr = BufferResource(stats, mr) assert br_with_mr.statistics is stats @pytest.mark.parametrize("mem_type", [MemoryType.DEVICE, MemoryType.HOST]) def test_opaque_memory_usage_basic(mem_type: MemoryType) -> None: mr = rmm.mr.CudaMemoryResource() - br = BufferResource(mr, memory_limits={mem_type: KiB(100)}) + br = BufferResource(Statistics.disabled(), mr, memory_limits={mem_type: KiB(100)}) res, _ = br.reserve(mem_type, KiB(50), allow_overbooking=False) assert res.size == KiB(50) @@ -122,7 +124,7 @@ def test_opaque_memory_usage_basic(mem_type: MemoryType) -> None: @pytest.mark.parametrize("mem_type", [MemoryType.DEVICE, MemoryType.HOST]) def test_opaque_memory_usage_clears_on_exception(mem_type: MemoryType) -> None: mr = rmm.mr.CudaMemoryResource() - br = BufferResource(mr, memory_limits={mem_type: KiB(100)}) + br = BufferResource(Statistics.disabled(), mr, memory_limits={mem_type: KiB(100)}) res, _ = br.reserve(mem_type, KiB(60), allow_overbooking=False) assert res.size == KiB(60) @@ -136,7 +138,7 @@ def test_opaque_memory_usage_clears_on_exception(mem_type: MemoryType) -> None: @pytest.mark.parametrize("mem_type", [MemoryType.DEVICE, MemoryType.HOST]) def test_opaque_memory_usage_multiple_reservations(mem_type: MemoryType) -> None: mr = rmm.mr.CudaMemoryResource() - br = BufferResource(mr, memory_limits={mem_type: KiB(200)}) + br = BufferResource(Statistics.disabled(), mr, memory_limits={mem_type: KiB(200)}) res1, _ = br.reserve(mem_type, KiB(50), allow_overbooking=False) assert res1.size == KiB(50) @@ -158,7 +160,7 @@ def test_opaque_memory_usage_multiple_reservations(mem_type: MemoryType) -> None @pytest.mark.parametrize("mem_type", [MemoryType.DEVICE, MemoryType.HOST]) def test_opaque_memory_usage_nested(mem_type: MemoryType) -> None: mr = rmm.mr.CudaMemoryResource() - br = BufferResource(mr, memory_limits={mem_type: KiB(200)}) + br = BufferResource(Statistics.disabled(), mr, memory_limits={mem_type: KiB(200)}) res1, _ = br.reserve(mem_type, KiB(40), allow_overbooking=False) res2, _ = br.reserve(mem_type, KiB(60), allow_overbooking=False) @@ -180,7 +182,7 @@ def test_opaque_memory_usage_nested(mem_type: MemoryType) -> None: @pytest.mark.parametrize("mem_type", [MemoryType.DEVICE, MemoryType.HOST]) def test_opaque_memory_usage_partial_consumption(mem_type: MemoryType) -> None: mr = rmm.mr.CudaMemoryResource() - br = BufferResource(mr, memory_limits={mem_type: KiB(100)}) + br = BufferResource(Statistics.disabled(), mr, memory_limits={mem_type: KiB(100)}) res, _ = br.reserve(mem_type, KiB(80), allow_overbooking=False) assert res.size == KiB(80) @@ -197,7 +199,9 @@ def test_opaque_memory_usage_partial_consumption(mem_type: MemoryType) -> None: def test_device_mr_property_tracks_allocations() -> None: mr = rmm.mr.CudaMemoryResource() - br = BufferResource(mr, memory_limits={MemoryType.DEVICE: KiB(1024)}) + br = BufferResource( + Statistics.disabled(), mr, memory_limits={MemoryType.DEVICE: KiB(1024)} + ) assert isinstance(br.device_mr, OwningDeviceMemoryResource) avail_before = br.memory_available(MemoryType.DEVICE) @@ -220,7 +224,7 @@ def test_device_mr_outlives_buffer_resource() -> None: # Write data and delete the buffer resource. mr = rmm.mr.CudaMemoryResource() - br = BufferResource(mr) + br = BufferResource(Statistics.disabled(), mr) buf = rmm.DeviceBuffer(size=expected.nbytes, mr=br.device_mr) buf.copy_from_host(expected.tobytes()) del br @@ -233,7 +237,9 @@ def test_device_mr_outlives_buffer_resource() -> None: def test_reserve_or_fail() -> None: mr = rmm.mr.CudaMemoryResource() - br = BufferResource(mr, memory_limits={MemoryType.DEVICE: KiB(100)}) + br = BufferResource( + Statistics.disabled(), mr, memory_limits={MemoryType.DEVICE: KiB(100)} + ) # Succeeds on the first available memory type. res = br.reserve_or_fail(KiB(50), [MemoryType.DEVICE, MemoryType.HOST]) diff --git a/python/rapidsmpf/rapidsmpf/tests/test_communicator.py b/python/rapidsmpf/rapidsmpf/tests/test_communicator.py index 62de32f30..6a5d95664 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_communicator.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_communicator.py @@ -5,6 +5,7 @@ import pytest from rapidsmpf.progress_thread import ProgressThread +from rapidsmpf.statistics import Statistics MPI = pytest.importorskip("mpi4py.MPI") from typing import TYPE_CHECKING # noqa: E402 @@ -28,8 +29,9 @@ def test_log_level(capfd: pytest.CaptureFixture[str], level: LOG_LEVEL) -> None: with pytest.MonkeyPatch.context() as monkeypatch: monkeypatch.setenv("RAPIDSMPF_LOG", level.name) + options = Options(get_environment_variables()) comm = new_communicator( - MPI.COMM_WORLD, Options(get_environment_variables()), ProgressThread() + MPI.COMM_WORLD, options, ProgressThread(Statistics.from_options(options)) ) assert comm.logger.verbosity_level is level comm.logger.print("PRINT") @@ -50,8 +52,9 @@ def test_log_level(capfd: pytest.CaptureFixture[str], level: LOG_LEVEL) -> None: def test_mpi() -> None: + options = Options(get_environment_variables()) comm = new_communicator( - MPI.COMM_WORLD, Options(get_environment_variables()), ProgressThread() + MPI.COMM_WORLD, options, ProgressThread(Statistics.from_options(options)) ) assert comm.nranks == MPI.COMM_WORLD.size assert comm.rank == MPI.COMM_WORLD.rank @@ -59,8 +62,9 @@ def test_mpi() -> None: def test_ucxx() -> None: ucxx_worker = initialize_ucxx() + options = Options(get_environment_variables()) comm = ucxx_mpi_setup( - ucxx_worker, Options(get_environment_variables()), ProgressThread() + ucxx_worker, options, ProgressThread(Statistics.from_options(options)) ) assert comm.nranks == MPI.COMM_WORLD.size @@ -68,14 +72,20 @@ def test_ucxx() -> None: def test_single_process() -> None: - comm = single_process_comm(Options(get_environment_variables()), ProgressThread()) + options = Options(get_environment_variables()) + comm = single_process_comm( + options, ProgressThread(Statistics.from_options(options)) + ) assert comm.nranks == 1 assert comm.rank == 0 def test_logger_survives_communicator(capfd: pytest.CaptureFixture[str]) -> None: def get_logger() -> Logger: - comm = single_process_comm(Options(), ProgressThread()) + options = Options() + comm = single_process_comm( + options, ProgressThread(Statistics.from_options(options)) + ) return comm.logger get_logger().print("Logger should survive") diff --git a/python/rapidsmpf/rapidsmpf/tests/test_config.py b/python/rapidsmpf/rapidsmpf/tests/test_config.py index d60099978..f4b1a18c5 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_config.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_config.py @@ -492,7 +492,7 @@ def test_buffer_resource_from_options_creates_instance_with_explicit_options() - } ) br = BufferResource.from_options( - rmm.mr.CudaMemoryResource(), opts, Statistics.from_options(opts) + Statistics.from_options(opts), rmm.mr.CudaMemoryResource(), opts ) assert br.statistics.enabled @@ -503,7 +503,8 @@ def test_buffer_resource_from_options_creates_instance_with_explicit_options() - def test_buffer_resource_from_options_uses_default_when_options_empty() -> None: opts = Options() - br = BufferResource.from_options(rmm.mr.CudaMemoryResource(), opts) + mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) + br = BufferResource.from_options(Statistics.from_options(opts), mr, opts) assert not br.statistics.enabled assert br.stream_pool_size() == 16 @@ -515,16 +516,16 @@ def test_buffer_resource_from_options_uses_default_when_options_empty() -> None: def test_buffer_resource_from_options_enables_statistics_when_requested() -> None: opts = Options({"statistics": "ON"}) - br = BufferResource.from_options( - rmm.mr.CudaMemoryResource(), opts, Statistics.from_options(opts) - ) + mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) + br = BufferResource.from_options(Statistics.from_options(opts), mr, opts) assert br.statistics.enabled def test_buffer_resource_from_options_accepts_percentage_for_device_limit() -> None: opts = Options({"spill_device_limit": "50%"}) - br = BufferResource.from_options(rmm.mr.CudaMemoryResource(), opts) + mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) + br = BufferResource.from_options(Statistics.from_options(opts), mr, opts) # Verify device memory limit is set (exact value depends on system memory) mem_avail = br.memory_available(MemoryType.DEVICE) @@ -536,7 +537,8 @@ def test_buffer_resource_from_options_enables_pinned_memory_when_supported() -> pytest.skip("Pinned memory not supported on this system") opts = Options({"pinned_memory": "True"}) - br = BufferResource.from_options(rmm.mr.CudaMemoryResource(), opts) + mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) + br = BufferResource.from_options(Statistics.from_options(opts), mr, opts) assert br.pinned_mr is not None @@ -549,11 +551,10 @@ def test_context_from_options_creates_instance_with_explicit_options() -> None: } ) mr = RmmResourceAdaptor(rmm.mr.CudaMemoryResource()) - comm = single_comm.new_communicator(opts, ProgressThread()) + stats = Statistics.from_options(opts) + comm = single_comm.new_communicator(opts, ProgressThread(stats)) - with Context.from_options( - comm.logger, mr, opts, Statistics.from_options(opts) - ) as ctx: + with Context.from_options(comm.logger, mr, opts, stats) as ctx: assert ctx is not None assert ctx.statistics().enabled assert ctx.stream_pool_size() == 8 @@ -563,9 +564,10 @@ 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()) - comm = single_comm.new_communicator(opts, ProgressThread()) + stats = Statistics.from_options(opts) + comm = single_comm.new_communicator(opts, ProgressThread(stats)) - with Context.from_options(comm.logger, mr, opts) as ctx: + with Context.from_options(comm.logger, mr, opts, stats) as ctx: assert ctx is not None assert not ctx.statistics().enabled assert ctx.stream_pool_size() == 16 # Default @@ -575,11 +577,10 @@ 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()) - comm = single_comm.new_communicator(opts, ProgressThread()) + stats = Statistics.from_options(opts) + comm = single_comm.new_communicator(opts, ProgressThread(stats)) - with Context.from_options( - comm.logger, mr, opts, Statistics.from_options(opts) - ) as ctx: + with Context.from_options(comm.logger, mr, opts, stats) as ctx: assert ctx is not None assert ctx.statistics().enabled @@ -587,9 +588,10 @@ 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()) - comm = single_comm.new_communicator(opts, ProgressThread()) + stats = Statistics.from_options(opts) + comm = single_comm.new_communicator(opts, ProgressThread(stats)) - with Context.from_options(comm.logger, mr, opts) as ctx: + with Context.from_options(comm.logger, mr, opts, stats) as ctx: assert ctx is not None assert ctx.br() is not None @@ -597,9 +599,10 @@ 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()) - comm = single_comm.new_communicator(opts, ProgressThread()) + stats = Statistics.from_options(opts) + comm = single_comm.new_communicator(opts, ProgressThread(stats)) - with Context.from_options(comm.logger, mr, opts) as ctx: + with Context.from_options(comm.logger, mr, opts, stats) as ctx: assert ctx is not None channel: Channel[Any] = ctx.create_channel() assert channel is not None diff --git a/python/rapidsmpf/rapidsmpf/tests/test_examples.py b/python/rapidsmpf/rapidsmpf/tests/test_examples.py index 128bb66e3..578cf1bce 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_examples.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_examples.py @@ -11,6 +11,7 @@ import pylibcudf as plc from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.statistics import Statistics from rapidsmpf.testing import assert_eq MPI = pytest.importorskip("mpi4py.MPI") @@ -93,7 +94,7 @@ def test_bulk_shuffle( output_dir = str(mpi_tmpdir.join("output")) # Use a default buffer resource. - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) # Perform a the shuffle bulk_mpi_shuffle( diff --git a/python/rapidsmpf/rapidsmpf/tests/test_packed_data_host_bytes.py b/python/rapidsmpf/rapidsmpf/tests/test_packed_data_host_bytes.py index 374b5f727..bb8ba1537 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_packed_data_host_bytes.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_packed_data_host_bytes.py @@ -12,13 +12,14 @@ from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.memory.packed_data import PackedData +from rapidsmpf.statistics import Statistics def test_packed_data_host_bytes_roundtrip( device_mr: rmm.mr.CudaMemoryResource, ) -> None: """Test creating PackedData from host bytes and extracting them back.""" - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) # Test with raw bytes original = b"hello world" @@ -38,7 +39,7 @@ def test_packed_data_empty_bytes( device_mr: rmm.mr.CudaMemoryResource, ) -> None: """Test creating PackedData from empty bytes.""" - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) original = b"" packed = PackedData.from_host_bytes(original, br) result = packed.to_host_bytes() diff --git a/python/rapidsmpf/rapidsmpf/tests/test_partition.py b/python/rapidsmpf/rapidsmpf/tests/test_partition.py index 9caa7f46c..815e29451 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_partition.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_partition.py @@ -18,6 +18,7 @@ unspill_partitions, ) from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.statistics import Statistics from rapidsmpf.testing import assert_eq if TYPE_CHECKING: @@ -39,7 +40,7 @@ def _make_table(cols: list[list[int]]) -> plc.Table: def test_partition_and_pack_unpack( device_mr: rmm.mr.CudaMemoryResource, cols: list[list[int]], num_partitions: int ) -> None: - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) expect = _make_table(cols) partitions = partition_and_pack( expect, @@ -69,7 +70,7 @@ def test_partition_and_pack_unpack( def test_split_and_pack_unpack( device_mr: rmm.mr.CudaMemoryResource, cols: list[list[int]], num_partitions: int ) -> None: - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) expect = _make_table(cols) splits = np.linspace(0, expect.num_rows(), num_partitions, endpoint=False)[ 1: @@ -94,7 +95,7 @@ def test_split_and_pack_unpack( def test_split_and_pack_unpack_out_of_range( device_mr: rmm.mr.CudaMemoryResource, cols: list[list[int]], num_partitions: int ) -> None: - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) expect = _make_table(cols) with pytest.raises(IndexError): split_and_pack( @@ -110,7 +111,7 @@ def test_split_and_pack_unpack_out_of_range( def test_spill_unspill_roundtrip( device_mr: rmm.mr.CudaMemoryResource, cols: list[list[int]], num_partitions: int ) -> None: - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) expect = _make_table(cols) partitions = partition_and_pack( expect, diff --git a/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py b/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py index 6fcbd89c1..2d86b378b 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py @@ -19,6 +19,7 @@ from rapidsmpf.shuffler import ( Shuffler, ) +from rapidsmpf.statistics import Statistics from rapidsmpf.testing import assert_eq if TYPE_CHECKING: @@ -35,7 +36,7 @@ def test_shuffler_single_nonempty_partition( stream: Stream, total_num_partitions: int, ) -> None: - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) shuffler = Shuffler( comm, @@ -99,7 +100,7 @@ def test_shuffler_uniform( batch_size: int | None, total_num_partitions: int, ) -> None: - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) # Every rank creates the full input dataframe and all the expected partitions # (also partitions this rank might not get after the shuffle). diff --git a/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py b/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py index 3990751df..3a55a9719 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py @@ -15,6 +15,7 @@ from rapidsmpf.integrations.cudf.partition import unpack_and_concat from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.memory.packed_data import PackedData +from rapidsmpf.statistics import Statistics from rapidsmpf.testing import assert_eq if TYPE_CHECKING: @@ -71,7 +72,7 @@ def test_basic( Each rank sends `n_inserts` messages to each adjacent rank and then checks that extraction by source returns the expected number of payloads. """ - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) srcs, dsts = expected_peers(comm) sparse_alltoall = SparseAlltoall( comm=comm, @@ -123,7 +124,7 @@ def test_non_participating_ranks( if comm.nranks < 2: pytest.skip("Need at least two ranks") - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) if comm.rank == 0: srcs = [] @@ -173,7 +174,7 @@ def test_invalid_peers_raise( ) -> None: rank = comm.rank size = comm.nranks - br = BufferResource(device_mr) + br = BufferResource(Statistics.disabled(), device_mr) for src, dst in [([], [rank]), ([rank], []), ([], [size]), ([size], [])]: with pytest.raises( IndexError, match=r"SparseAlltoall invalid (source|destination) rank" diff --git a/python/rapidsmpf/rapidsmpf/tests/test_spill_manager.py b/python/rapidsmpf/rapidsmpf/tests/test_spill_manager.py index eb023a2e6..3a983aa5d 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_spill_manager.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_spill_manager.py @@ -10,6 +10,7 @@ from rapidsmpf.error import BadAlloc, OutOfMemory, ReservationError from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.statistics import Statistics if TYPE_CHECKING: import rmm.mr @@ -37,7 +38,7 @@ def test_error_handling( def spill(amount: int) -> int: raise error - br = BufferResource(device_mr, periodic_spill_check=None) + br = BufferResource(Statistics.disabled(), device_mr, periodic_spill_check=None) br.spill_manager.add_spill_function(spill, 0) with pytest.raises(error): br.spill_manager.spill(10) @@ -46,7 +47,7 @@ def spill(amount: int) -> int: def test_spill_function( device_mr: rmm.mr.CudaMemoryResource, ) -> None: - br = BufferResource(device_mr, periodic_spill_check=None) + br = BufferResource(Statistics.disabled(), device_mr, periodic_spill_check=None) track_spilled = [0] def spill_unlimited(amount: int) -> int: @@ -90,7 +91,9 @@ def spill_limited(amount: int) -> int: def test_spill_function_outlive_buffer_resource( device_mr: rmm.mr.CudaMemoryResource, ) -> None: - spill_manager = BufferResource(device_mr, periodic_spill_check=None).spill_manager + spill_manager = BufferResource( + Statistics.disabled(), device_mr, periodic_spill_check=None + ).spill_manager with pytest.raises(ValueError): spill_manager.add_spill_function(lambda x: x, 0) with pytest.raises(ValueError): @@ -106,6 +109,7 @@ def test_periodic_spill_check( # a periodic spill check enabled. (memory_available = limit - allocated; with # limit = -100 and no allocations, the available value is always negative.) br = BufferResource( + Statistics.disabled(), device_mr, memory_limits={MemoryType.DEVICE: -100}, periodic_spill_check=1e-3, @@ -128,6 +132,7 @@ def test_spill_to_make_headroom( ) -> None: # Create a buffer resource with a fixed limit of 100 bytes. br = BufferResource( + Statistics.disabled(), device_mr, memory_limits={MemoryType.DEVICE: 100}, periodic_spill_check=None, @@ -152,9 +157,7 @@ def test_reserve_device_memory_and_spill( ) -> None: # Create a buffer resource with a fixed limit of 100 bytes. br = BufferResource( - device_mr, - memory_limits={MemoryType.DEVICE: 100}, - periodic_spill_check=None, + Statistics.disabled(), device_mr, memory_limits={MemoryType.DEVICE: 100} ) track_spilled = [0]