From 9445b5056f7fe27a74dc207232d7d1f7fb382628 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 27 May 2026 13:31:40 -0700 Subject: [PATCH 01/13] making statistics mandatory arg --- cpp/benchmarks/bench_comm.cpp | 11 +--- cpp/benchmarks/bench_partition.cpp | 4 +- cpp/benchmarks/bench_shuffle.cpp | 8 +-- .../streaming/bench_streaming_shuffle.cpp | 20 ++----- cpp/benchmarks/streaming/ndsh/utils.cpp | 4 +- cpp/examples/example_shuffle.cpp | 2 +- cpp/include/rapidsmpf/bootstrap/ucxx.hpp | 4 +- .../rapidsmpf/communicator/communicator.hpp | 32 ++++++++++- cpp/include/rapidsmpf/communicator/mpi.hpp | 1 + cpp/include/rapidsmpf/communicator/single.hpp | 1 + cpp/include/rapidsmpf/communicator/ucxx.hpp | 2 + .../rapidsmpf/memory/buffer_resource.hpp | 25 +++++---- cpp/include/rapidsmpf/progress_thread.hpp | 5 +- .../rapidsmpf/streaming/core/context.hpp | 7 ++- cpp/src/coll/allreduce.cpp | 3 +- cpp/src/communicator/communicator.cpp | 14 +++++ cpp/src/communicator/mpi.cpp | 3 +- cpp/src/communicator/single.cpp | 3 +- cpp/src/communicator/ucxx.cpp | 3 +- cpp/src/memory/buffer_resource.cpp | 8 +-- cpp/src/streaming/core/context.cpp | 1 + cpp/tests/main/mpi.cpp | 8 ++- cpp/tests/main/single.cpp | 6 +- cpp/tests/main/ucxx.cpp | 7 ++- .../streaming/base_streaming_fixture.hpp | 5 +- cpp/tests/streaming/test_channel_metadata.cpp | 4 +- cpp/tests/streaming/test_fanout.cpp | 8 +-- cpp/tests/streaming/test_message.cpp | 5 +- cpp/tests/streaming/test_table_chunk.cpp | 18 +++--- cpp/tests/test_allgather.cpp | 9 ++- cpp/tests/test_allreduce.cu | 5 +- cpp/tests/test_buffer.cpp | 2 + cpp/tests/test_buffer_resource.cpp | 24 +++++--- cpp/tests/test_chunk.cpp | 5 +- cpp/tests/test_communicator.cpp | 5 +- cpp/tests/test_config.cpp | 40 ++++++-------- cpp/tests/test_metadata_payload_exchange.cpp | 4 +- cpp/tests/test_partition.cpp | 8 ++- cpp/tests/test_progress_thread.cpp | 2 +- cpp/tests/test_shuffler.cpp | 55 ++++++++++++++----- cpp/tests/test_shuffler_many_streams.cpp | 5 +- cpp/tests/test_sparse_alltoall.cpp | 5 +- cpp/tests/test_spill_manager.cpp | 1 + .../benchmarks/streaming_benchmark.py | 2 +- .../rapidsmpf/communicator/communicator.pxd | 2 + .../rapidsmpf/communicator/communicator.pyi | 3 + .../rapidsmpf/communicator/communicator.pyx | 14 +++++ .../rapidsmpf/examples/bulk_mpi_shuffle.py | 10 ++-- .../examples/ray/bulk_ray_shuffle.py | 5 +- .../examples/ray/ray_shuffle_example.py | 3 +- .../examples/streaming/basic_example.py | 6 +- .../rapidsmpf/memory/buffer_resource.pyi | 4 +- .../rapidsmpf/memory/buffer_resource.pyx | 26 ++++----- .../rapidsmpf/rapidsmpf/progress_thread.pyi | 2 +- .../rapidsmpf/rapidsmpf/progress_thread.pyx | 8 +-- python/rapidsmpf/rapidsmpf/statistics.pyx | 11 ++-- .../rapidsmpf/streaming/core/context.pyi | 2 +- .../rapidsmpf/streaming/core/context.pyx | 8 +-- python/rapidsmpf/rapidsmpf/tests/conftest.py | 13 ++++- .../rapidsmpf/tests/streaming/conftest.py | 3 +- .../streaming/test_memory_reserve_or_wait.py | 5 +- .../rapidsmpf/tests/test_allgather.py | 4 +- .../rapidsmpf/tests/test_buffer_resource.py | 30 +++++----- .../rapidsmpf/tests/test_communicator.py | 20 +++++-- .../rapidsmpf/rapidsmpf/tests/test_config.py | 45 ++++++++------- .../rapidsmpf/tests/test_examples.py | 3 +- .../tests/test_packed_data_host_bytes.py | 5 +- .../rapidsmpf/tests/test_partition.py | 9 +-- .../rapidsmpf/tests/test_shuffler.py | 5 +- .../rapidsmpf/tests/test_sparse_alltoall.py | 7 ++- .../rapidsmpf/tests/test_spill_manager.py | 15 +++-- 71 files changed, 415 insertions(+), 247 deletions(-) diff --git a/cpp/benchmarks/bench_comm.cpp b/cpp/benchmarks/bench_comm.cpp index 37a5039a1..7416d80e8 100644 --- a/cpp/benchmarks/bench_comm.cpp +++ b/cpp/benchmarks/bench_comm.cpp @@ -317,16 +317,7 @@ int main(int argc, char** argv) { set_current_rmm_resource(args.rmm_mr); rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); - BufferResource br{ - mr, - PinnedMemoryResource::Disabled, - {}, - std::chrono::milliseconds{1}, - std::make_shared( - 16, rmm::cuda_stream::flags::non_blocking - ), - stats - }; + BufferResource br{stats, mr}; // Print benchmark/hardware info. { diff --git a/cpp/benchmarks/bench_partition.cpp b/cpp/benchmarks/bench_partition.cpp index ca286885b..24f971a45 100644 --- a/cpp/benchmarks/bench_partition.cpp +++ b/cpp/benchmarks/bench_partition.cpp @@ -60,7 +60,7 @@ 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}; - rapidsmpf::BufferResource br{pool_mr}; + rapidsmpf::BufferResource br{rapidsmpf::Statistics::disabled(), pool_mr}; // Create input table auto table = create_int_table(num_rows, stream); @@ -111,7 +111,7 @@ 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}; - rapidsmpf::BufferResource br{pool_mr}; + rapidsmpf::BufferResource br{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 cd30993e3..c71603900 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(); rapidsmpf::BufferResource br{ + 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 5836b11ff..de1f6577b 100644 --- a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp +++ b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp @@ -324,7 +324,8 @@ 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::from_options(options); + auto progress_thread = std::make_shared(stats); std::shared_ptr comm; if (args.comm_type == "mpi") { @@ -364,20 +365,11 @@ 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 = std::make_shared( - stat_enabled_mr, - pinned_mr, - std::move(memory_limits), - std::nullopt, - std::make_shared( - 16, rmm::cuda_stream::flags::non_blocking - ), - stats + stats, stat_enabled_mr, pinned_mr, std::move(memory_limits) ); auto& log = *comm->logger(); @@ -411,7 +403,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; @@ -460,13 +452,13 @@ int main(int argc, char** argv) { 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 fb20dd8b1..20fe2a6ab 100644 --- a/cpp/benchmarks/streaming/ndsh/utils.cpp +++ b/cpp/benchmarks/streaming/ndsh/utils.cpp @@ -157,6 +157,7 @@ create_context( ); auto br = std::make_shared( + 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 7ba73f0cf..0fc8edf85 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(); - rapidsmpf::BufferResource br{mr}; + rapidsmpf::BufferResource br{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..5d96cec2c 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,21 @@ class Communicator { }; protected: - Communicator() = default; + /** + * @brief Construct the base Communicator with a statistics instance. + * + * Derived classes typically forward `progress_thread->statistics()` here so + * that the communicator and its progress thread share the same `Statistics` + * instance. + * + * @param statistics The statistics instance to associate with this + * communicator. Must not be null. + * + * @throws std::invalid_argument If @p statistics is null. + */ + explicit Communicator(std::shared_ptr statistics); + + std::shared_ptr statistics_; ///< Statistics instance. Never null. public: virtual ~Communicator() noexcept = default; @@ -619,6 +634,19 @@ class Communicator { [[nodiscard]] virtual std::shared_ptr const& progress_thread() const = 0; + /** + * @brief Retrieves the statistics instance associated with this communicator. + * + * Satisfies the `StatisticsProvider` concept. Returns the instance supplied + * to the protected `Communicator` constructor by the concrete implementation + * (typically the same instance owned by `progress_thread()`). + * + * @return Shared pointer to the statistics instance (never null). + */ + [[nodiscard]] std::shared_ptr statistics() const noexcept { + return statistics_; + } + /** * @brief Provides a string representation of the communicator. * @return A string describing the communicator. @@ -626,6 +654,8 @@ class Communicator { [[nodiscard]] virtual std::string str() const = 0; }; +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..07df73689 100644 --- a/cpp/include/rapidsmpf/communicator/mpi.hpp +++ b/cpp/include/rapidsmpf/communicator/mpi.hpp @@ -270,5 +270,6 @@ class MPI final : public Communicator { 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..b3b7369b9 100644 --- a/cpp/include/rapidsmpf/communicator/single.hpp +++ b/cpp/include/rapidsmpf/communicator/single.hpp @@ -209,5 +209,6 @@ class Single final : public Communicator { 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 fa607d96b..15fd71fd8 100644 --- a/cpp/include/rapidsmpf/communicator/ucxx.hpp +++ b/cpp/include/rapidsmpf/communicator/ucxx.hpp @@ -347,4 +347,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 5b01ef1fc..fdb636798 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -59,15 +59,15 @@ class BufferResource { /** * @brief Constructs a buffer resource. * - * Memory availability is computed per `MemoryType` as `limit - allocated`. + * The function should return the current available memory of a specific type and + * must be thread-safe iff used by multiple `BufferResource` instances concurrently. * - * Device and pinned-host allocations routed through this BufferResource are tracked - * automatically. Host memory allocations are not tracked, so the available memory - * always equals the configured limit. - * - * If pinned-host memory is disabled, available pinned-host memory is always reported - * as zero regardless of the configured limit. + * @warning Calling any `BufferResource` instance methods in the function might result + * in a deadlock. This is because the buffer resource is locked when the function is + * called. * + * @param statistics The statistics instance to use. Pass `Statistics::disabled()` + * to opt out of statistics collection. * @param device_mr The RMM device memory resource used for device allocations. * @param pinned_mr The pinned host memory resource used for `MemoryType::PINNED_HOST` * allocations. If disabled, pinned host allocations are unavailable regardless of @@ -80,16 +80,15 @@ class BufferResource { * periodic spill checking is performed. * @param stream_pool Pool of CUDA streams used throughout RapidsMPF for operations * that do not take an explicit CUDA stream. - * @param statistics The statistics instance to use (disabled by default). */ BufferResource( + 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) ); /** @@ -101,7 +100,9 @@ class BufferResource { * * @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. The caller is responsible for creating and + * owning this object. * * @return A shared pointer to a BufferResource instance configured according to the * options. @@ -109,7 +110,7 @@ class BufferResource { 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; diff --git a/cpp/include/rapidsmpf/progress_thread.hpp b/cpp/include/rapidsmpf/progress_thread.hpp index 44f3c539c..ceaa59b09 100644 --- a/cpp/include/rapidsmpf/progress_thread.hpp +++ b/cpp/include/rapidsmpf/progress_thread.hpp @@ -117,14 +117,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} ); diff --git a/cpp/include/rapidsmpf/streaming/core/context.hpp b/cpp/include/rapidsmpf/streaming/core/context.hpp index 559f58574..173d8aaf6 100644 --- a/cpp/include/rapidsmpf/streaming/core/context.hpp +++ b/cpp/include/rapidsmpf/streaming/core/context.hpp @@ -83,8 +83,9 @@ 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. The caller is responsible for creating and + * owning this object. * @return A fully initialized Context. * * @throws std::invalid_argument If an option value is invalid. @@ -105,7 +106,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 971337dba..43a4c7b0f 100644 --- a/cpp/src/communicator/communicator.cpp +++ b/cpp/src/communicator/communicator.cpp @@ -3,10 +3,24 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include +#include + #include +#include #include namespace rapidsmpf { + +Communicator::Communicator(std::shared_ptr statistics) + : statistics_{std::move(statistics)} { + RAPIDSMPF_EXPECTS( + statistics_ != nullptr, + "Communicator statistics cannot be null", + std::invalid_argument + ); +} + namespace { Communicator::Logger::LOG_LEVEL level_from_string(std::string const& str) { if (str.empty()) { diff --git a/cpp/src/communicator/mpi.cpp b/cpp/src/communicator/mpi.cpp index 6aa501662..3582efad2 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{progress_thread->statistics()}, + comm_{comm}, rank_{[comm]() { int r; RAPIDSMPF_MPI(MPI_Comm_rank(comm, &r)); diff --git a/cpp/src/communicator/single.cpp b/cpp/src/communicator/single.cpp index e8d594a69..6b80410bf 100644 --- a/cpp/src/communicator/single.cpp +++ b/cpp/src/communicator/single.cpp @@ -12,7 +12,8 @@ namespace rapidsmpf { Single::Single(config::Options options, std::shared_ptr progress_thread) - : logger_{std::make_shared(0, std::move(options))}, + : Communicator{progress_thread->statistics()}, + logger_{std::make_shared(0, std::move(options))}, progress_thread_{std::move(progress_thread)} {} std::unique_ptr Single::send( diff --git a/cpp/src/communicator/ucxx.cpp b/cpp/src/communicator/ucxx.cpp index 5b114c841..e1e31af60 100644 --- a/cpp/src/communicator/ucxx.cpp +++ b/cpp/src/communicator/ucxx.cpp @@ -1131,7 +1131,8 @@ UCXX::UCXX( config::Options options, std::shared_ptr progress_thread ) - : shared_resources_(ucxx_initialized_rank->shared_resources_), + : Communicator{progress_thread->statistics()}, + 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)} { diff --git a/cpp/src/memory/buffer_resource.cpp b/cpp/src/memory/buffer_resource.cpp index c2f24fd73..da7884bc5 100644 --- a/cpp/src/memory/buffer_resource.cpp +++ b/cpp/src/memory/buffer_resource.cpp @@ -21,12 +21,12 @@ namespace rapidsmpf { BufferResource::BufferResource( + 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 ) : device_adaptor_{std::move(device_mr)}, device_mr_{device_adaptor_}, // any_resource shares state via shared_resource @@ -59,12 +59,12 @@ std::shared_ptr BufferResource::from_options( }; return std::make_shared( + 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/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 d2dc7ea84..0f1dded2c 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 = std::make_shared( - 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 603ae2006..bfa958cbc 100644 --- a/cpp/tests/streaming/test_channel_metadata.cpp +++ b/cpp/tests/streaming/test_channel_metadata.cpp @@ -166,7 +166,9 @@ TEST_F(StreamingChannelMetadata, MessageRoundTrip) { class StreamingChannelMetadataGPU : public ::testing::Test { protected: rmm::cuda_stream_view stream{cudf::get_default_stream()}; - rapidsmpf::BufferResource br{cudf::get_current_device_resource_ref()}; + rapidsmpf::BufferResource br{ + 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 4d77dc966..ab3fccd2c 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 = std::make_shared( - mr_cuda, rapidsmpf::PinnedMemoryResource::Disabled, memory_limits + br = std::make_shared( + 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 165730e1d..092ba8464 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 = std::make_unique(cudf::get_current_device_resource_ref()); + br = std::make_unique( + 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 3b66e5e9e..da86c9e1b 100644 --- a/cpp/tests/streaming/test_table_chunk.cpp +++ b/cpp/tests/streaming/test_table_chunk.cpp @@ -29,12 +29,10 @@ 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( @@ -42,22 +40,20 @@ class StreamingTableChunk : public BaseStreamingFixture, ); stream = cudf::get_default_stream(); br = std::make_shared( + Statistics::disabled(), // statistics 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 + stream_pool // stream_pool ); + 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 667c8b4aa..f24d0fb23 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 = std::make_unique(rmm::mr::cuda_memory_resource{}); + br = std::make_unique( + 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 = std::make_unique( + rapidsmpf::Statistics::disabled(), DelayedMemoryResource{br->device_mr(), std::chrono::milliseconds(500)} ); allgather = @@ -363,7 +366,9 @@ TEST_F(BaseAllGatherTest, opid_reuse) { TEST(PostBox, spill_uses_remaining_amount) { auto stream = cudf::get_default_stream(); auto mr = std::make_unique(); - auto br = std::make_unique(*mr); + auto br = std::make_unique( + rapidsmpf::Statistics::disabled(), *mr + ); rapidsmpf::coll::detail::PostBox postbox; diff --git a/cpp/tests/test_allreduce.cu b/cpp/tests/test_allreduce.cu index 6c20b26c6..679c9a97b 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,9 @@ class BaseAllReduceTest : public ::testing::Test { protected: void SetUp() override { mr = std::make_unique(); - br = std::make_unique(*mr); + br = std::make_unique( + rapidsmpf::Statistics::disabled(), *mr + ); comm = GlobalEnvironment->comm_.get(); } diff --git a/cpp/tests/test_buffer.cpp b/cpp/tests/test_buffer.cpp index 770629aeb..c841b8824 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 = std::make_unique( + 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 cea208c88..772d6c45a 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -53,6 +53,7 @@ std::unique_ptr zeros( TEST(BufferResource, ReservationOverbooking) { // Create a buffer resource that always reports 10 KiB of available device memory. BufferResource br{ + Statistics::disabled(), cudf::get_current_device_resource_ref(), PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, 10_KiB}} @@ -120,6 +121,7 @@ TEST(BufferResource, ReservationReleasing) { // Create a buffer resource that always reports 10 KiB of available host and device // memory. BufferResource br{ + Statistics::disabled(), cudf::get_current_device_resource_ref(), PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, 10_KiB}, {MemoryType::HOST, 10_KiB}} @@ -173,7 +175,10 @@ TEST(BufferResource, MemoryLimit) { // Create a buffer resource that limits available device memory to 10 KiB. BufferResource br{ - 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); @@ -262,7 +267,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); } - BufferResource br{cuda_mr, pinned_mr, std::move(memory_limits)}; + BufferResource br{ + 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); @@ -290,12 +297,12 @@ TEST(BufferResource, AllocStatistics) { auto stats = Statistics::create(); auto pinned_mr = PinnedMemoryResource::make_if_available(); BufferResource br{ + 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(); @@ -355,6 +362,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 = std::make_unique( + Statistics::disabled(), mr_cuda, PinnedMemoryResource::Disabled, std::unordered_map{{MemoryType::DEVICE, 10_KiB}} @@ -417,7 +425,9 @@ TEST_F(BufferResourceReserveOrFailTest, MultipleTypes) { class BaseBufferResourceCopyTest : public ::testing::Test { protected: void SetUp() override { - br = std::make_unique(cudf::get_current_device_resource_ref()); + br = std::make_unique( + Statistics::disabled(), cudf::get_current_device_resource_ref() + ); stream = cudf::get_default_stream(); // initialize the host pattern @@ -639,11 +649,11 @@ class BufferResourceDifferentResourcesTest : public ::testing::Test { // Setup br1 with statistics for its device memory mr1 = std::make_unique(rmm::mr::cuda_memory_resource{}); - br1 = std::make_unique(*mr1); + br1 = std::make_unique(Statistics::disabled(), *mr1); // Setup br2 with statistics for its device memory mr2 = std::make_unique(rmm::mr::cuda_memory_resource{}); - br2 = std::make_unique(*mr2); + br2 = std::make_unique(Statistics::disabled(), *mr2); } std::unique_ptr create_source_buffer() { diff --git a/cpp/tests/test_chunk.cpp b/cpp/tests/test_chunk.cpp index 77af94a82..eb9bf7cfb 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 = std::make_unique(cudf::get_current_device_resource_ref()); + br = std::make_unique( + 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 dc7d88aa0..d4f05cd88 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,9 @@ class BaseCommunicatorTest : public ::testing::Test { void SetUp() override { comm = GlobalEnvironment->comm_.get(); mr = std::make_unique(); - br = std::make_unique(*mr); + br = std::make_unique( + 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 bab869315..4f708c6cc 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 = std::make_unique(*mr); - stream = rmm::cuda_stream_default; statistics = Statistics::create(); + br = std::make_unique(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 9bae901b4..d59215eda 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(); - rapidsmpf::BufferResource br{mr()}; + BufferResource br{Statistics::disabled(), mr()}; cudf::table expect = random_table_with_index(seed, static_cast(num_rows), 0, 10); @@ -72,7 +72,7 @@ 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(); - rapidsmpf::BufferResource br{cudf::get_current_device_resource_ref()}; + BufferResource br{Statistics::disabled(), cudf::get_current_device_resource_ref()}; cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); @@ -99,7 +99,9 @@ TEST_P(NumOfPartitions, split_and_pack) { class SpillingTest : public ::testing::Test { protected: void SetUp() override { - br = std::make_unique(cudf::get_current_device_resource_ref()); + br = std::make_unique( + 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..79bf96e15 100644 --- a/cpp/tests/test_progress_thread.cpp +++ b/cpp/tests/test_progress_thread.cpp @@ -89,7 +89,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([] { diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index df89b50d5..39ddc452b 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,9 @@ extern Environment* GlobalEnvironment; TEST(ReceivedChunks, spill_skips_control_messages) { auto mr = cudf::get_current_device_resource_ref(); - auto br = std::make_unique(mr); + auto br = std::make_unique( + rapidsmpf::Statistics::disabled(), mr + ); rapidsmpf::shuffler::detail::ReceivedChunks received; @@ -49,7 +52,9 @@ TEST(ReceivedChunks, spill_skips_control_messages) { TEST(ReceivedChunks, spill_respects_amount) { auto mr = cudf::get_current_device_resource_ref(); - auto br = std::make_unique(mr); + auto br = std::make_unique( + rapidsmpf::Statistics::disabled(), mr + ); auto stream = cudf::get_default_stream(); rapidsmpf::shuffler::detail::ReceivedChunks received; @@ -75,7 +80,9 @@ 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 = std::make_unique(mr); + auto br = std::make_unique( + rapidsmpf::Statistics::disabled(), mr + ); auto metadata = iota_vector(100); @@ -238,7 +245,10 @@ class MemoryLimits_NumPartition total_num_partitions = std::get<1>(GetParam()); total_num_rows = std::get<2>(GetParam()); br = std::make_unique( - mr(), rapidsmpf::PinnedMemoryResource::Disabled, memory_limits + rapidsmpf::Statistics::disabled(), + mr(), + rapidsmpf::PinnedMemoryResource::Disabled, + memory_limits ); shuffler = std::make_unique( @@ -309,7 +319,9 @@ class ConcurrentShuffleTest static_cast(std::get<1>(GetParam())); // these resources will be used by multiple threads to instantiate shufflers - br = std::make_unique(mr()); + br = std::make_unique( + rapidsmpf::Statistics::disabled(), mr() + ); stream = cudf::get_default_stream(); } @@ -415,6 +427,7 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { constexpr std::int64_t k_no_spill_limit = (1LL << 40); constexpr std::int64_t k_force_spill_limit = -(1LL << 40); rapidsmpf::BufferResource br{ + rapidsmpf::Statistics::disabled(), mr, rapidsmpf::PinnedMemoryResource::Disabled, {{rapidsmpf::MemoryType::DEVICE, k_no_spill_limit}}, @@ -676,7 +689,9 @@ TEST(Shuffler, ShutdownWhilePaused) { auto progress_thread = GlobalEnvironment->comm_->progress_thread(); auto mr = cudf::get_current_device_resource_ref(); - auto br = std::make_unique(mr); + auto br = std::make_unique( + rapidsmpf::Statistics::disabled(), mr + ); auto shuffler = rapidsmpf::shuffler::Shuffler(GlobalEnvironment->comm_, 0, 1, br.get()); @@ -711,7 +726,9 @@ class ExtractEmptyPartitionsTest : public cudf::test::BaseFixture { void SetUp() override { stream = cudf::get_default_stream(); - br = std::make_unique(mr()); + br = std::make_unique( + rapidsmpf::Statistics::disabled(), mr() + ); shuffler = std::make_unique( GlobalEnvironment->comm_, 0, nparts, br.get() @@ -810,7 +827,9 @@ TEST_F(ExtractEmptyPartitionsTest, SomeEmptyAndNonEmptyInsertions) { TEST(ShufflerTest, multiple_shutdowns) { auto& comm = GlobalEnvironment->comm_; - rapidsmpf::BufferResource br(cudf::get_current_device_resource_ref()); + rapidsmpf::BufferResource br( + rapidsmpf::Statistics::disabled(), cudf::get_current_device_resource_ref() + ); auto shuffler = std::make_unique(comm, 0, comm->nranks(), &br); @@ -834,7 +853,9 @@ TEST(ShufflerTest, multiple_shutdowns) { TEST(Shuffler, concurrent_wait) { auto const& comm = GlobalEnvironment->comm_; auto stream = cudf::get_default_stream(); - rapidsmpf::BufferResource br(cudf::get_current_device_resource_ref()); + rapidsmpf::BufferResource br( + 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. @@ -938,7 +959,9 @@ TEST(Shuffler, opid_reuse) { constexpr auto wait_timeout = std::chrono::seconds{30}; rmm::mr::cuda_memory_resource mr; - auto br = std::make_unique(mr); + auto br = std::make_unique( + rapidsmpf::Statistics::disabled(), mr + ); // On rank 0, wrap the device MR with a delayed version for the shuffler. std::unique_ptr delayed_mr; @@ -947,7 +970,9 @@ TEST(Shuffler, opid_reuse) { if (comm->rank() == 0) { delayed_mr = std::make_unique(mr, std::chrono::milliseconds(500)); - delayed_br = std::make_unique(*delayed_mr); + delayed_br = std::make_unique( + rapidsmpf::Statistics::disabled(), *delayed_mr + ); shuffler_br = delayed_br.get(); } @@ -1045,7 +1070,9 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) { constexpr auto wait_timeout = std::chrono::seconds{30}; rmm::mr::cuda_memory_resource mr; - auto br = std::make_unique(mr); + auto br = std::make_unique( + rapidsmpf::Statistics::disabled(), mr + ); // On rank 0, wrap the device MR with a delayed version for the shuffler. std::unique_ptr delayed_mr; @@ -1054,7 +1081,9 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) { if (comm->rank() == 0) { delayed_mr = std::make_unique(mr, std::chrono::milliseconds(500)); - delayed_br = std::make_unique(*delayed_mr); + delayed_br = std::make_unique( + 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 a42fea215..ea024b0ae 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 = std::make_unique(cudf::get_current_device_resource_ref()); + auto br = std::make_unique( + 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 69b4d9e80..d5c7b0e30 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 = std::make_unique(rmm::mr::cuda_memory_resource{}); + br = std::make_unique( + rapidsmpf::Statistics::disabled(), rmm::mr::cuda_memory_resource{} + ); } std::unique_ptr br; diff --git a/cpp/tests/test_spill_manager.cpp b/cpp/tests/test_spill_manager.cpp index f4fb6bf8c..6da926810 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; BufferResource br{ + Statistics::disabled(), cudf::get_current_device_resource_ref(), rapidsmpf::PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, mem_available}} 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.pyi b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyi index ce21b941f..ee94b49eb 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) -> DeviceMemoryResource: ... diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index 682d008e9..2146d82ea 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -136,18 +136,18 @@ cdef class BufferResource: 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. + The statistics instance to use. Required. Pass + ``Statistics.disabled()`` for a no-op recorder. """ 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: @@ -181,13 +181,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 # Stored for the Python device_mr/pinned_mr property accessors. # The C++ BufferResource owns the resource via any_resource. @@ -198,21 +194,21 @@ cdef class BufferResource: cpp_pinned_mr = self._pinned_mr._handle with nogil: self._handle = make_shared[cpp_BufferResource]( + 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. @@ -223,28 +219,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.pyi b/python/rapidsmpf/rapidsmpf/progress_thread.pyi index fba04eb35..3207f6ec9 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pyi +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pyi @@ -5,4 +5,4 @@ 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: ... diff --git a/python/rapidsmpf/rapidsmpf/progress_thread.pyx b/python/rapidsmpf/rapidsmpf/progress_thread.pyx index f4deee800..2de803241 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pyx +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pyx @@ -19,7 +19,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,11 +30,8 @@ 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) diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyx b/python/rapidsmpf/rapidsmpf/statistics.pyx index 0d64d0ad8..abd124ba2 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 *: """ @@ -186,14 +189,14 @@ cdef class Statistics: @classmethod def disabled(cls): """ - Returns a disabled (no-op) Statistics instance. + Get a shared, no-op Statistics instance. - Useful when you need to pass a Statistics argument but do not want to - collect any data. + Returns the same disabled instance on every call. Use this when an API + requires a Statistics object but no recording is desired. Returns ------- - A Statistics instance with tracking disabled. + A disabled Statistics instance. """ cdef Statistics ret = cls.__new__(cls) with nogil: 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 15801e171..54d9d3e4e 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_buffer_resource.py @@ -19,7 +19,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 @@ -30,7 +32,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 @@ -79,12 +81,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 @@ -92,20 +94,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) @@ -120,7 +122,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) @@ -134,7 +136,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) @@ -156,7 +158,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) @@ -178,7 +180,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) @@ -195,7 +197,9 @@ def test_opaque_memory_usage_partial_consumption(mem_type: MemoryType) -> 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 c7969a386..32235b5f3 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_config.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_config.py @@ -486,7 +486,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 @@ -497,7 +497,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 @@ -509,16 +510,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) @@ -530,7 +531,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 @@ -543,11 +545,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 @@ -557,9 +558,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 @@ -569,11 +571,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 @@ -581,9 +582,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 @@ -591,9 +593,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] From 2d93022d0a5b7535d032d34cbfd6e397e6a794b8 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 27 May 2026 14:25:05 -0700 Subject: [PATCH 02/13] missing changes Signed-off-by: niranda perera --- cpp/include/rapidsmpf/memory/buffer_resource.hpp | 15 ++++++++------- cpp/include/rapidsmpf/streaming/core/context.hpp | 3 +-- cpp/tests/streaming/test_table_chunk.cpp | 13 +++---------- 3 files changed, 12 insertions(+), 19 deletions(-) diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index fdb636798..bc0ad76e0 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -59,12 +59,14 @@ class BufferResource { /** * @brief Constructs a buffer resource. * - * The function should return the current available memory of a specific type and - * must be thread-safe iff used by multiple `BufferResource` instances concurrently. + * Memory availability is computed per `MemoryType` as `limit - allocated`. * - * @warning Calling any `BufferResource` instance methods in the function might result - * in a deadlock. This is because the buffer resource is locked when the function is - * called. + * Device and pinned-host allocations routed through this BufferResource are tracked + * automatically. Host memory allocations are not tracked, so the available memory + * always equals the configured limit. + * + * If pinned-host memory is disabled, available pinned-host memory is always reported + * as zero regardless of the configured limit. * * @param statistics The statistics instance to use. Pass `Statistics::disabled()` * to opt out of statistics collection. @@ -101,8 +103,7 @@ class BufferResource { * @param mr A device-accessible RMM memory resource. * @param options Configuration options. * @param statistics The statistics instance to use. Pass `Statistics::disabled()` - * to opt out of statistics collection. The caller is responsible for creating and - * owning this object. + * to opt out of statistics collection. * * @return A shared pointer to a BufferResource instance configured according to the * options. diff --git a/cpp/include/rapidsmpf/streaming/core/context.hpp b/cpp/include/rapidsmpf/streaming/core/context.hpp index 173d8aaf6..578f1e6c1 100644 --- a/cpp/include/rapidsmpf/streaming/core/context.hpp +++ b/cpp/include/rapidsmpf/streaming/core/context.hpp @@ -84,8 +84,7 @@ class Context { * @param options Configuration options used to initialize the Context and its * components. * @param statistics The statistics instance to use. Pass `Statistics::disabled()` - * to opt out of statistics collection. The caller is responsible for creating and - * owning this object. + * to opt out of statistics collection. * @return A fully initialized Context. * * @throws std::invalid_argument If an option value is invalid. diff --git a/cpp/tests/streaming/test_table_chunk.cpp b/cpp/tests/streaming/test_table_chunk.cpp index da86c9e1b..fc641bf95 100644 --- a/cpp/tests/streaming/test_table_chunk.cpp +++ b/cpp/tests/streaming/test_table_chunk.cpp @@ -34,18 +34,11 @@ class StreamingTableChunk : public BaseStreamingFixture, void SetUp() override { 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 = std::make_shared( - Statistics::disabled(), // statistics - 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(), + mr_cuda, + rapidsmpf::PinnedMemoryResource::make_if_available() ); ctx = std::make_shared(options, GlobalEnvironment->comm_->logger(), br); } From a6803e664d825f5cb1d94007121cc356ab489fa3 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 27 May 2026 16:47:18 -0700 Subject: [PATCH 03/13] set_stats Signed-off-by: niranda perera --- .../rapidsmpf/communicator/communicator.hpp | 18 ++++++++++ cpp/include/rapidsmpf/progress_thread.hpp | 25 +++++++++++++ cpp/include/rapidsmpf/statistics.hpp | 29 +++++++++++++++ cpp/src/communicator/communicator.cpp | 9 +++++ cpp/src/progress_thread.cpp | 16 +++++++++ .../rapidsmpf/communicator/communicator.pxd | 3 ++ .../rapidsmpf/communicator/communicator.pyi | 1 + .../rapidsmpf/communicator/communicator.pyx | 27 ++++++++++++++ .../rapidsmpf/rapidsmpf/progress_thread.pxd | 4 +++ .../rapidsmpf/rapidsmpf/progress_thread.pyi | 3 ++ .../rapidsmpf/rapidsmpf/progress_thread.pyx | 36 +++++++++++++++++++ 11 files changed, 171 insertions(+) diff --git a/cpp/include/rapidsmpf/communicator/communicator.hpp b/cpp/include/rapidsmpf/communicator/communicator.hpp index 5d96cec2c..4cef046b8 100644 --- a/cpp/include/rapidsmpf/communicator/communicator.hpp +++ b/cpp/include/rapidsmpf/communicator/communicator.hpp @@ -647,6 +647,23 @@ class Communicator { return statistics_; } + /** + * @brief Replace the statistics instance held by this communicator. + * + * Installs @p statistics in place of the current instance. Future + * `statistics()` calls return the new instance. + * + * @param statistics The new statistics instance. Must not be null. Pass + * `Statistics::disabled()` to opt out of statistics collection. + * + * @throws std::invalid_argument If @p statistics is null. + * + * @warning Concurrent calls to `set_statistics()` are not allowed; the + * caller must ensure this method is invoked from a single thread at a + * time. + */ + void set_statistics(std::shared_ptr statistics); + /** * @brief Provides a string representation of the communicator. * @return A string describing the communicator. @@ -655,6 +672,7 @@ class Communicator { }; static_assert(StatisticsProvider); +static_assert(MutableStatisticsProvider); /// @brief Whether RapidsMPF was built with the UCXX Communicator. #ifdef RAPIDSMPF_HAVE_UCXX diff --git a/cpp/include/rapidsmpf/progress_thread.hpp b/cpp/include/rapidsmpf/progress_thread.hpp index ceaa59b09..d45cd56b9 100644 --- a/cpp/include/rapidsmpf/progress_thread.hpp +++ b/cpp/include/rapidsmpf/progress_thread.hpp @@ -184,6 +184,30 @@ class ProgressThread { */ std::shared_ptr statistics() const noexcept; + /** + * @brief Replace the statistics instance held by this progress thread. + * + * This method pauses the progress loop (via `pause()`) for the duration + * of the swap so the loop cannot dereference `statistics_` while it is + * being replaced, then restores the prior running state (calls `resume()` + * only if the loop was active before the call). + * + * Because `pause()` synchronizes with the loop, this call may block + * briefly until the in-flight iteration completes. + * + * @param statistics The new statistics instance. Must not be null. Pass + * `Statistics::disabled()` to opt out of statistics collection. + * + * @throws std::invalid_argument If @p statistics is null. Validated + * before the pause so a null argument does not leave the progress loop + * paused. + * + * @warning Concurrent calls to `set_statistics()` are not allowed; the + * caller must ensure this method is invoked from a single thread at a + * time. + */ + void set_statistics(std::shared_ptr statistics); + private: /** * @brief The event loop progressing each of the functions. @@ -207,5 +231,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..3f4da0bee 100644 --- a/cpp/include/rapidsmpf/statistics.hpp +++ b/cpp/include/rapidsmpf/statistics.hpp @@ -656,6 +656,35 @@ 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. + * + * Only `ProgressThread` and `Communicator` — the *root* providers that drive the + * shared `Statistics` for an entire RapidsMPF runtime instance — satisfy this + * concept. Secondary providers such as `BufferResource` and `streaming::Context` + * deliberately do **not** offer mutation: their `Statistics` is fixed at + * construction. A typical caller-driven rebuild therefore swaps the instance on + * the progress thread and communicator and then constructs a fresh + * `streaming::Context` (and the `BufferResource` it owns) with the same new + * `Statistics`. + * + * @warning Concurrent calls to `set_statistics()` are not allowed; the + * caller must ensure this method is invoked from a single thread at a + * time. + */ +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/src/communicator/communicator.cpp b/cpp/src/communicator/communicator.cpp index 43a4c7b0f..b73947d6d 100644 --- a/cpp/src/communicator/communicator.cpp +++ b/cpp/src/communicator/communicator.cpp @@ -21,6 +21,15 @@ Communicator::Communicator(std::shared_ptr statistics) ); } +void Communicator::set_statistics(std::shared_ptr statistics) { + RAPIDSMPF_EXPECTS( + statistics != nullptr, + "Communicator statistics cannot be null", + std::invalid_argument + ); + statistics_ = std::move(statistics); +} + namespace { Communicator::Logger::LOG_LEVEL level_from_string(std::string const& str) { if (str.empty()) { diff --git a/cpp/src/progress_thread.cpp b/cpp/src/progress_thread.cpp index 42846b421..bc69f378e 100644 --- a/cpp/src/progress_thread.cpp +++ b/cpp/src/progress_thread.cpp @@ -119,6 +119,22 @@ std::shared_ptr ProgressThread::statistics() const noexcept { return statistics_; } +void ProgressThread::set_statistics(std::shared_ptr statistics) { + RAPIDSMPF_EXPECTS( + statistics != nullptr, + "the statistics pointer cannot be NULL", + std::invalid_argument + ); + // Pausing ensures the progress loop is not in the middle of dereferencing + // `statistics_` while we swap it. Restore the prior running state on exit. + bool const was_running = is_running(); + pause(); + statistics_ = std::move(statistics); + if (was_running) { + resume(); + } +} + void ProgressThread::event_loop() { auto const t0_event_loop = Clock::now(); { diff --git a/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd b/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd index 5b3079801..6f6e06d94 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd @@ -38,6 +38,9 @@ cdef extern from "" nogil: string str() except +ex_handler shared_ptr[cpp_ProgressThread] progress_thread() except +ex_handler shared_ptr[cpp_Statistics] statistics() except +ex_handler + void set_statistics( + 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 ba1be8f53..a1075b929 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi @@ -33,6 +33,7 @@ class Communicator: def progress_thread(self) -> ProgressThread: ... @property def statistics(self) -> Statistics: ... + def set_statistics(self, statistics: Statistics) -> None: ... 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 8995c4cbd..da3493005 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx @@ -193,6 +193,33 @@ cdef class Communicator: stats._handle = deref(self._handle).statistics() return stats + def set_statistics(self, Statistics statistics not None): + """ + Replace the statistics instance held by this communicator. + + Concrete communicator implementations typically share their statistics + instance with their progress thread; when swapping at runtime, the + caller is responsible for keeping the two in sync by also calling + :meth:`ProgressThread.set_statistics` with the same instance and + rebuilding any dependent secondary providers (e.g. + :class:`~rapidsmpf.memory.buffer_resource.BufferResource` / + :class:`~rapidsmpf.streaming.core.context.Context`) from the same + new instance. + + Parameters + ---------- + statistics + The new statistics instance. Must not be ``None``. Pass + ``Statistics.disabled()`` to opt out of statistics collection. + + Warnings + -------- + Concurrent calls to ``set_statistics`` are not allowed; the caller + must ensure this method is invoked from a single thread at a time. + """ + with nogil: + deref(self._handle).set_statistics(statistics._handle) + def get_str(self): """ Get a string representation of the communicator. 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 3207f6ec9..d49955a3c 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pyi +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pyi @@ -6,3 +6,6 @@ from rapidsmpf.statistics import Statistics class ProgressThread: 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 2de803241..13f5a05b8 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 @@ -35,6 +36,41 @@ cdef class ProgressThread: 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 underlying C++ implementation pauses the progress loop for the + duration of the swap so the loop cannot dereference the held instance + while it is being replaced, then restores the prior running state. + + Parameters + ---------- + statistics + The new statistics instance. Must not be ``None``. Pass + ``Statistics.disabled()`` to opt out of statistics collection. + + Warnings + -------- + Concurrent calls to ``set_statistics`` are not allowed; the caller + must ensure this method is invoked from a single thread at a time. + """ + with nogil: + deref(self._handle).set_statistics(statistics._handle) + def __dealloc__(self): with nogil: self._handle.reset() From 891af67605246562fed0aedfab99eab5f8a3cbf6 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 27 May 2026 22:32:33 -0700 Subject: [PATCH 04/13] Revert "set_stats" This reverts commit a6803e664d825f5cb1d94007121cc356ab489fa3. --- .../rapidsmpf/communicator/communicator.hpp | 18 ---------- cpp/include/rapidsmpf/progress_thread.hpp | 25 ------------- cpp/include/rapidsmpf/statistics.hpp | 29 --------------- cpp/src/communicator/communicator.cpp | 9 ----- cpp/src/progress_thread.cpp | 16 --------- .../rapidsmpf/communicator/communicator.pxd | 3 -- .../rapidsmpf/communicator/communicator.pyi | 1 - .../rapidsmpf/communicator/communicator.pyx | 27 -------------- .../rapidsmpf/rapidsmpf/progress_thread.pxd | 4 --- .../rapidsmpf/rapidsmpf/progress_thread.pyi | 3 -- .../rapidsmpf/rapidsmpf/progress_thread.pyx | 36 ------------------- 11 files changed, 171 deletions(-) diff --git a/cpp/include/rapidsmpf/communicator/communicator.hpp b/cpp/include/rapidsmpf/communicator/communicator.hpp index 4cef046b8..5d96cec2c 100644 --- a/cpp/include/rapidsmpf/communicator/communicator.hpp +++ b/cpp/include/rapidsmpf/communicator/communicator.hpp @@ -647,23 +647,6 @@ class Communicator { return statistics_; } - /** - * @brief Replace the statistics instance held by this communicator. - * - * Installs @p statistics in place of the current instance. Future - * `statistics()` calls return the new instance. - * - * @param statistics The new statistics instance. Must not be null. Pass - * `Statistics::disabled()` to opt out of statistics collection. - * - * @throws std::invalid_argument If @p statistics is null. - * - * @warning Concurrent calls to `set_statistics()` are not allowed; the - * caller must ensure this method is invoked from a single thread at a - * time. - */ - void set_statistics(std::shared_ptr statistics); - /** * @brief Provides a string representation of the communicator. * @return A string describing the communicator. @@ -672,7 +655,6 @@ class Communicator { }; static_assert(StatisticsProvider); -static_assert(MutableStatisticsProvider); /// @brief Whether RapidsMPF was built with the UCXX Communicator. #ifdef RAPIDSMPF_HAVE_UCXX diff --git a/cpp/include/rapidsmpf/progress_thread.hpp b/cpp/include/rapidsmpf/progress_thread.hpp index d45cd56b9..ceaa59b09 100644 --- a/cpp/include/rapidsmpf/progress_thread.hpp +++ b/cpp/include/rapidsmpf/progress_thread.hpp @@ -184,30 +184,6 @@ class ProgressThread { */ std::shared_ptr statistics() const noexcept; - /** - * @brief Replace the statistics instance held by this progress thread. - * - * This method pauses the progress loop (via `pause()`) for the duration - * of the swap so the loop cannot dereference `statistics_` while it is - * being replaced, then restores the prior running state (calls `resume()` - * only if the loop was active before the call). - * - * Because `pause()` synchronizes with the loop, this call may block - * briefly until the in-flight iteration completes. - * - * @param statistics The new statistics instance. Must not be null. Pass - * `Statistics::disabled()` to opt out of statistics collection. - * - * @throws std::invalid_argument If @p statistics is null. Validated - * before the pause so a null argument does not leave the progress loop - * paused. - * - * @warning Concurrent calls to `set_statistics()` are not allowed; the - * caller must ensure this method is invoked from a single thread at a - * time. - */ - void set_statistics(std::shared_ptr statistics); - private: /** * @brief The event loop progressing each of the functions. @@ -231,6 +207,5 @@ 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 3f4da0bee..adf47b4c7 100644 --- a/cpp/include/rapidsmpf/statistics.hpp +++ b/cpp/include/rapidsmpf/statistics.hpp @@ -656,35 +656,6 @@ 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. - * - * Only `ProgressThread` and `Communicator` — the *root* providers that drive the - * shared `Statistics` for an entire RapidsMPF runtime instance — satisfy this - * concept. Secondary providers such as `BufferResource` and `streaming::Context` - * deliberately do **not** offer mutation: their `Statistics` is fixed at - * construction. A typical caller-driven rebuild therefore swaps the instance on - * the progress thread and communicator and then constructs a fresh - * `streaming::Context` (and the `BufferResource` it owns) with the same new - * `Statistics`. - * - * @warning Concurrent calls to `set_statistics()` are not allowed; the - * caller must ensure this method is invoked from a single thread at a - * time. - */ -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/src/communicator/communicator.cpp b/cpp/src/communicator/communicator.cpp index b73947d6d..43a4c7b0f 100644 --- a/cpp/src/communicator/communicator.cpp +++ b/cpp/src/communicator/communicator.cpp @@ -21,15 +21,6 @@ Communicator::Communicator(std::shared_ptr statistics) ); } -void Communicator::set_statistics(std::shared_ptr statistics) { - RAPIDSMPF_EXPECTS( - statistics != nullptr, - "Communicator statistics cannot be null", - std::invalid_argument - ); - statistics_ = std::move(statistics); -} - namespace { Communicator::Logger::LOG_LEVEL level_from_string(std::string const& str) { if (str.empty()) { diff --git a/cpp/src/progress_thread.cpp b/cpp/src/progress_thread.cpp index bc69f378e..42846b421 100644 --- a/cpp/src/progress_thread.cpp +++ b/cpp/src/progress_thread.cpp @@ -119,22 +119,6 @@ std::shared_ptr ProgressThread::statistics() const noexcept { return statistics_; } -void ProgressThread::set_statistics(std::shared_ptr statistics) { - RAPIDSMPF_EXPECTS( - statistics != nullptr, - "the statistics pointer cannot be NULL", - std::invalid_argument - ); - // Pausing ensures the progress loop is not in the middle of dereferencing - // `statistics_` while we swap it. Restore the prior running state on exit. - bool const was_running = is_running(); - pause(); - statistics_ = std::move(statistics); - if (was_running) { - resume(); - } -} - void ProgressThread::event_loop() { auto const t0_event_loop = Clock::now(); { diff --git a/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd b/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd index 6f6e06d94..5b3079801 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd @@ -38,9 +38,6 @@ cdef extern from "" nogil: string str() except +ex_handler shared_ptr[cpp_ProgressThread] progress_thread() except +ex_handler shared_ptr[cpp_Statistics] statistics() except +ex_handler - void set_statistics( - 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 a1075b929..ba1be8f53 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi @@ -33,7 +33,6 @@ class Communicator: def progress_thread(self) -> ProgressThread: ... @property def statistics(self) -> Statistics: ... - def set_statistics(self, statistics: Statistics) -> None: ... 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 da3493005..8995c4cbd 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx @@ -193,33 +193,6 @@ cdef class Communicator: stats._handle = deref(self._handle).statistics() return stats - def set_statistics(self, Statistics statistics not None): - """ - Replace the statistics instance held by this communicator. - - Concrete communicator implementations typically share their statistics - instance with their progress thread; when swapping at runtime, the - caller is responsible for keeping the two in sync by also calling - :meth:`ProgressThread.set_statistics` with the same instance and - rebuilding any dependent secondary providers (e.g. - :class:`~rapidsmpf.memory.buffer_resource.BufferResource` / - :class:`~rapidsmpf.streaming.core.context.Context`) from the same - new instance. - - Parameters - ---------- - statistics - The new statistics instance. Must not be ``None``. Pass - ``Statistics.disabled()`` to opt out of statistics collection. - - Warnings - -------- - Concurrent calls to ``set_statistics`` are not allowed; the caller - must ensure this method is invoked from a single thread at a time. - """ - with nogil: - deref(self._handle).set_statistics(statistics._handle) - def get_str(self): """ Get a string representation of the communicator. diff --git a/python/rapidsmpf/rapidsmpf/progress_thread.pxd b/python/rapidsmpf/rapidsmpf/progress_thread.pxd index 344643063..7ee796b09 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pxd +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pxd @@ -17,10 +17,6 @@ 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 d49955a3c..3207f6ec9 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pyi +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pyi @@ -6,6 +6,3 @@ from rapidsmpf.statistics import Statistics class ProgressThread: 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 13f5a05b8..2de803241 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pyx +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pyx @@ -2,7 +2,6 @@ # 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 @@ -36,41 +35,6 @@ cdef class ProgressThread: 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 underlying C++ implementation pauses the progress loop for the - duration of the swap so the loop cannot dereference the held instance - while it is being replaced, then restores the prior running state. - - Parameters - ---------- - statistics - The new statistics instance. Must not be ``None``. Pass - ``Statistics.disabled()`` to opt out of statistics collection. - - Warnings - -------- - Concurrent calls to ``set_statistics`` are not allowed; the caller - must ensure this method is invoked from a single thread at a time. - """ - with nogil: - deref(self._handle).set_statistics(statistics._handle) - def __dealloc__(self): with nogil: self._handle.reset() From 2ce411f12725c3145c3af262e6ac70c1422e56ab Mon Sep 17 00:00:00 2001 From: niranda perera Date: Thu, 28 May 2026 00:27:15 -0700 Subject: [PATCH 05/13] reset from options Signed-off-by: niranda perera --- cpp/include/rapidsmpf/statistics.hpp | 18 ++++++++++++++ cpp/src/statistics.cpp | 29 ++++++++++++++++++++--- cpp/src/stream_ordered_timing.cpp | 2 +- python/rapidsmpf/rapidsmpf/statistics.pxd | 2 ++ python/rapidsmpf/rapidsmpf/statistics.pyi | 1 + python/rapidsmpf/rapidsmpf/statistics.pyx | 24 +++++++++++++++++++ 6 files changed, 72 insertions(+), 4 deletions(-) diff --git a/cpp/include/rapidsmpf/statistics.hpp b/cpp/include/rapidsmpf/statistics.hpp index adf47b4c7..e7e433384 100644 --- a/cpp/include/rapidsmpf/statistics.hpp +++ b/cpp/include/rapidsmpf/statistics.hpp @@ -144,6 +144,24 @@ class Statistics : public std::enable_shared_from_this { */ static std::shared_ptr from_options(config::Options options); + /** + * @brief Resets the current state of this instance based on @p new_options. + * + * @warning Live `MemoryRecorder` instances (constructed before this call + * but not yet destroyed) are *not* tracked. Their destructors will run + * after the reset and publish their scope into the freshly cleared + * `memory_records_` map, mixing pre-reset measurements into the new + * session. Likewise, a `StreamOrderedTiming` whose stop callback is + * already mid-execution at reset time may still record one stale stat + * (its global-map entry was extracted before `cancel_inflight_timings` + * could remove it). Callers that need clean boundaries must quiesce + * recorders and synchronise the relevant CUDA streams before calling. + * + * @param new_options Configuration options whose `"statistics"` field + * controls the new enabled state. + */ + void reset_from_options(config::Options new_options); + ~Statistics() noexcept; // `Statistics` is owned exclusively through `std::shared_ptr`. Use diff --git a/cpp/src/statistics.cpp b/cpp/src/statistics.cpp index 3bc7a6d1a..553ea0d7a 100644 --- a/cpp/src/statistics.cpp +++ b/cpp/src/statistics.cpp @@ -168,11 +168,34 @@ std::shared_ptr Statistics::create(Mode mode) { return std::shared_ptr(new Statistics(mode == Mode::Enabled)); } -std::shared_ptr Statistics::from_options(config::Options options) { - bool const statistics = options.get("statistics", [](auto const& s) { +namespace { +bool parse_statistics_option(config::Options& options) { + return options.get("statistics", [](auto const& s) { return s.empty() ? false : parse_string(s); }); - return create(statistics ? Mode::Enabled : Mode::Disabled); +} +} // namespace + +std::shared_ptr Statistics::from_options(config::Options options) { + return create(parse_statistics_option(options) ? Mode::Enabled : Mode::Disabled); +} + +void Statistics::reset_from_options(config::Options new_options) { + // Drop pending stream-ordered stop callbacks before the reset so a + // late-firing callback cannot write a pre-reset measurement into the + // post-reset `stats_`. + StreamOrderedTiming::cancel_inflight_timings(this); + { + std::lock_guard lock(mutex_); + stats_.clear(); + report_entries_.clear(); + memory_records_.clear(); + } + if (parse_statistics_option(new_options)) { + enable(); + } else { + disable(); + } } Statistics::Stat Statistics::get_stat(std::string const& name) const { diff --git a/cpp/src/stream_ordered_timing.cpp b/cpp/src/stream_ordered_timing.cpp index c25f68980..fc89fa8ae 100644 --- a/cpp/src/stream_ordered_timing.cpp +++ b/cpp/src/stream_ordered_timing.cpp @@ -79,7 +79,7 @@ void timing_stop_cb(void* data) { lock.unlock(); auto statistics = timing.statistics.lock(); - if (statistics == nullptr) { + if (statistics == nullptr || !statistics->enabled()) { // This can happen if the statistics object has gone out of scope before CUDA gets // to this stream operation. return; diff --git a/python/rapidsmpf/rapidsmpf/statistics.pxd b/python/rapidsmpf/rapidsmpf/statistics.pxd index 134620f2a..09b1d8e9d 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pxd +++ b/python/rapidsmpf/rapidsmpf/statistics.pxd @@ -11,6 +11,7 @@ from libcpp.unordered_map cimport unordered_map from libcpp.vector cimport vector from rapidsmpf._detail.exception_handling cimport ex_handler +from rapidsmpf.config cimport cpp_Options from rapidsmpf.memory.pinned_memory_resource cimport (PinnedMemoryResource, cpp_PinnedMemoryResource) from rapidsmpf.memory.scoped_memory_record cimport cpp_ScopedMemoryRecord @@ -30,6 +31,7 @@ cdef extern from "" nogil: cdef cppclass cpp_Statistics "rapidsmpf::Statistics": bool enabled() except +ex_handler + void reset_from_options(cpp_Options options) except +ex_handler void add_stat( string name, double value diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyi b/python/rapidsmpf/rapidsmpf/statistics.pyi index e569c5824..83f28cf96 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pyi +++ b/python/rapidsmpf/rapidsmpf/statistics.pyi @@ -28,6 +28,7 @@ class Statistics: def disabled(cls: type[Self]) -> Self: ... @property def enabled(self) -> bool: ... + def reset_from_options(self, new_options: Options) -> None: ... def report( self, *, diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyx b/python/rapidsmpf/rapidsmpf/statistics.pyx index abd124ba2..98151edb1 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pyx +++ b/python/rapidsmpf/rapidsmpf/statistics.pyx @@ -220,6 +220,30 @@ cdef class Statistics: """ return deref(self._handle).enabled() + def reset_from_options(self, Options new_options not None) -> None: + """ + Resets the current state of the instance based on ``new_options``. + + .. warning:: + Live :class:`MemoryRecorder` instances (constructed before this + call but not yet destroyed) are *not* tracked. Their destructors + will run after the reset and publish their scope into the + freshly cleared memory records, mixing pre-reset measurements + into the new session. Likewise, a stream-ordered timing whose + stop callback is already mid-execution at reset time may still + record one stale stat (its global-map entry was extracted + before the cancellation could remove it). Callers that need + clean boundaries must quiesce recorders and synchronise the + relevant CUDA streams before calling. + + Parameters + ---------- + new_options + New configuration options + """ + with nogil: + deref(self._handle).reset_from_options(new_options._handle) + def report( self, *, From 36cd37a94e0a210d67432fa3d8dc0fffc4d410fd Mon Sep 17 00:00:00 2001 From: niranda perera Date: Thu, 28 May 2026 10:57:35 -0700 Subject: [PATCH 06/13] revert 2ce411f12725c3145c3af262e6ac70c1422e56ab 891af67605246562fed0aedfab99eab5f8a3cbf6 Signed-off-by: niranda perera --- .../rapidsmpf/communicator/communicator.hpp | 18 +++++++ cpp/include/rapidsmpf/progress_thread.hpp | 25 ++++++++++ cpp/include/rapidsmpf/statistics.hpp | 47 ++++++++++++------- cpp/src/communicator/communicator.cpp | 9 ++++ cpp/src/progress_thread.cpp | 16 +++++++ cpp/src/statistics.cpp | 29 ++---------- cpp/src/stream_ordered_timing.cpp | 2 +- .../rapidsmpf/communicator/communicator.pxd | 3 ++ .../rapidsmpf/communicator/communicator.pyi | 1 + .../rapidsmpf/communicator/communicator.pyx | 27 +++++++++++ .../rapidsmpf/rapidsmpf/progress_thread.pxd | 4 ++ .../rapidsmpf/rapidsmpf/progress_thread.pyi | 3 ++ .../rapidsmpf/rapidsmpf/progress_thread.pyx | 36 ++++++++++++++ python/rapidsmpf/rapidsmpf/statistics.pxd | 2 - python/rapidsmpf/rapidsmpf/statistics.pyi | 1 - python/rapidsmpf/rapidsmpf/statistics.pyx | 24 ---------- 16 files changed, 175 insertions(+), 72 deletions(-) diff --git a/cpp/include/rapidsmpf/communicator/communicator.hpp b/cpp/include/rapidsmpf/communicator/communicator.hpp index 5d96cec2c..4cef046b8 100644 --- a/cpp/include/rapidsmpf/communicator/communicator.hpp +++ b/cpp/include/rapidsmpf/communicator/communicator.hpp @@ -647,6 +647,23 @@ class Communicator { return statistics_; } + /** + * @brief Replace the statistics instance held by this communicator. + * + * Installs @p statistics in place of the current instance. Future + * `statistics()` calls return the new instance. + * + * @param statistics The new statistics instance. Must not be null. Pass + * `Statistics::disabled()` to opt out of statistics collection. + * + * @throws std::invalid_argument If @p statistics is null. + * + * @warning Concurrent calls to `set_statistics()` are not allowed; the + * caller must ensure this method is invoked from a single thread at a + * time. + */ + void set_statistics(std::shared_ptr statistics); + /** * @brief Provides a string representation of the communicator. * @return A string describing the communicator. @@ -655,6 +672,7 @@ class Communicator { }; static_assert(StatisticsProvider); +static_assert(MutableStatisticsProvider); /// @brief Whether RapidsMPF was built with the UCXX Communicator. #ifdef RAPIDSMPF_HAVE_UCXX diff --git a/cpp/include/rapidsmpf/progress_thread.hpp b/cpp/include/rapidsmpf/progress_thread.hpp index ceaa59b09..d45cd56b9 100644 --- a/cpp/include/rapidsmpf/progress_thread.hpp +++ b/cpp/include/rapidsmpf/progress_thread.hpp @@ -184,6 +184,30 @@ class ProgressThread { */ std::shared_ptr statistics() const noexcept; + /** + * @brief Replace the statistics instance held by this progress thread. + * + * This method pauses the progress loop (via `pause()`) for the duration + * of the swap so the loop cannot dereference `statistics_` while it is + * being replaced, then restores the prior running state (calls `resume()` + * only if the loop was active before the call). + * + * Because `pause()` synchronizes with the loop, this call may block + * briefly until the in-flight iteration completes. + * + * @param statistics The new statistics instance. Must not be null. Pass + * `Statistics::disabled()` to opt out of statistics collection. + * + * @throws std::invalid_argument If @p statistics is null. Validated + * before the pause so a null argument does not leave the progress loop + * paused. + * + * @warning Concurrent calls to `set_statistics()` are not allowed; the + * caller must ensure this method is invoked from a single thread at a + * time. + */ + void set_statistics(std::shared_ptr statistics); + private: /** * @brief The event loop progressing each of the functions. @@ -207,5 +231,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 e7e433384..3f4da0bee 100644 --- a/cpp/include/rapidsmpf/statistics.hpp +++ b/cpp/include/rapidsmpf/statistics.hpp @@ -144,24 +144,6 @@ class Statistics : public std::enable_shared_from_this { */ static std::shared_ptr from_options(config::Options options); - /** - * @brief Resets the current state of this instance based on @p new_options. - * - * @warning Live `MemoryRecorder` instances (constructed before this call - * but not yet destroyed) are *not* tracked. Their destructors will run - * after the reset and publish their scope into the freshly cleared - * `memory_records_` map, mixing pre-reset measurements into the new - * session. Likewise, a `StreamOrderedTiming` whose stop callback is - * already mid-execution at reset time may still record one stale stat - * (its global-map entry was extracted before `cancel_inflight_timings` - * could remove it). Callers that need clean boundaries must quiesce - * recorders and synchronise the relevant CUDA streams before calling. - * - * @param new_options Configuration options whose `"statistics"` field - * controls the new enabled state. - */ - void reset_from_options(config::Options new_options); - ~Statistics() noexcept; // `Statistics` is owned exclusively through `std::shared_ptr`. Use @@ -674,6 +656,35 @@ 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. + * + * Only `ProgressThread` and `Communicator` — the *root* providers that drive the + * shared `Statistics` for an entire RapidsMPF runtime instance — satisfy this + * concept. Secondary providers such as `BufferResource` and `streaming::Context` + * deliberately do **not** offer mutation: their `Statistics` is fixed at + * construction. A typical caller-driven rebuild therefore swaps the instance on + * the progress thread and communicator and then constructs a fresh + * `streaming::Context` (and the `BufferResource` it owns) with the same new + * `Statistics`. + * + * @warning Concurrent calls to `set_statistics()` are not allowed; the + * caller must ensure this method is invoked from a single thread at a + * time. + */ +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/src/communicator/communicator.cpp b/cpp/src/communicator/communicator.cpp index 43a4c7b0f..b73947d6d 100644 --- a/cpp/src/communicator/communicator.cpp +++ b/cpp/src/communicator/communicator.cpp @@ -21,6 +21,15 @@ Communicator::Communicator(std::shared_ptr statistics) ); } +void Communicator::set_statistics(std::shared_ptr statistics) { + RAPIDSMPF_EXPECTS( + statistics != nullptr, + "Communicator statistics cannot be null", + std::invalid_argument + ); + statistics_ = std::move(statistics); +} + namespace { Communicator::Logger::LOG_LEVEL level_from_string(std::string const& str) { if (str.empty()) { diff --git a/cpp/src/progress_thread.cpp b/cpp/src/progress_thread.cpp index 42846b421..bc69f378e 100644 --- a/cpp/src/progress_thread.cpp +++ b/cpp/src/progress_thread.cpp @@ -119,6 +119,22 @@ std::shared_ptr ProgressThread::statistics() const noexcept { return statistics_; } +void ProgressThread::set_statistics(std::shared_ptr statistics) { + RAPIDSMPF_EXPECTS( + statistics != nullptr, + "the statistics pointer cannot be NULL", + std::invalid_argument + ); + // Pausing ensures the progress loop is not in the middle of dereferencing + // `statistics_` while we swap it. Restore the prior running state on exit. + bool const was_running = is_running(); + pause(); + statistics_ = std::move(statistics); + if (was_running) { + resume(); + } +} + void ProgressThread::event_loop() { auto const t0_event_loop = Clock::now(); { diff --git a/cpp/src/statistics.cpp b/cpp/src/statistics.cpp index 553ea0d7a..3bc7a6d1a 100644 --- a/cpp/src/statistics.cpp +++ b/cpp/src/statistics.cpp @@ -168,34 +168,11 @@ std::shared_ptr Statistics::create(Mode mode) { return std::shared_ptr(new Statistics(mode == Mode::Enabled)); } -namespace { -bool parse_statistics_option(config::Options& options) { - return options.get("statistics", [](auto const& s) { +std::shared_ptr Statistics::from_options(config::Options options) { + bool const statistics = options.get("statistics", [](auto const& s) { return s.empty() ? false : parse_string(s); }); -} -} // namespace - -std::shared_ptr Statistics::from_options(config::Options options) { - return create(parse_statistics_option(options) ? Mode::Enabled : Mode::Disabled); -} - -void Statistics::reset_from_options(config::Options new_options) { - // Drop pending stream-ordered stop callbacks before the reset so a - // late-firing callback cannot write a pre-reset measurement into the - // post-reset `stats_`. - StreamOrderedTiming::cancel_inflight_timings(this); - { - std::lock_guard lock(mutex_); - stats_.clear(); - report_entries_.clear(); - memory_records_.clear(); - } - if (parse_statistics_option(new_options)) { - enable(); - } else { - disable(); - } + return create(statistics ? Mode::Enabled : Mode::Disabled); } Statistics::Stat Statistics::get_stat(std::string const& name) const { diff --git a/cpp/src/stream_ordered_timing.cpp b/cpp/src/stream_ordered_timing.cpp index fc89fa8ae..c25f68980 100644 --- a/cpp/src/stream_ordered_timing.cpp +++ b/cpp/src/stream_ordered_timing.cpp @@ -79,7 +79,7 @@ void timing_stop_cb(void* data) { lock.unlock(); auto statistics = timing.statistics.lock(); - if (statistics == nullptr || !statistics->enabled()) { + if (statistics == nullptr) { // This can happen if the statistics object has gone out of scope before CUDA gets // to this stream operation. return; diff --git a/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd b/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd index 5b3079801..6f6e06d94 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd @@ -38,6 +38,9 @@ cdef extern from "" nogil: string str() except +ex_handler shared_ptr[cpp_ProgressThread] progress_thread() except +ex_handler shared_ptr[cpp_Statistics] statistics() except +ex_handler + void set_statistics( + 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 ba1be8f53..a1075b929 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi @@ -33,6 +33,7 @@ class Communicator: def progress_thread(self) -> ProgressThread: ... @property def statistics(self) -> Statistics: ... + def set_statistics(self, statistics: Statistics) -> None: ... 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 8995c4cbd..da3493005 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx @@ -193,6 +193,33 @@ cdef class Communicator: stats._handle = deref(self._handle).statistics() return stats + def set_statistics(self, Statistics statistics not None): + """ + Replace the statistics instance held by this communicator. + + Concrete communicator implementations typically share their statistics + instance with their progress thread; when swapping at runtime, the + caller is responsible for keeping the two in sync by also calling + :meth:`ProgressThread.set_statistics` with the same instance and + rebuilding any dependent secondary providers (e.g. + :class:`~rapidsmpf.memory.buffer_resource.BufferResource` / + :class:`~rapidsmpf.streaming.core.context.Context`) from the same + new instance. + + Parameters + ---------- + statistics + The new statistics instance. Must not be ``None``. Pass + ``Statistics.disabled()`` to opt out of statistics collection. + + Warnings + -------- + Concurrent calls to ``set_statistics`` are not allowed; the caller + must ensure this method is invoked from a single thread at a time. + """ + with nogil: + deref(self._handle).set_statistics(statistics._handle) + def get_str(self): """ Get a string representation of the communicator. 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 3207f6ec9..d49955a3c 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pyi +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pyi @@ -6,3 +6,6 @@ from rapidsmpf.statistics import Statistics class ProgressThread: 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 2de803241..13f5a05b8 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 @@ -35,6 +36,41 @@ cdef class ProgressThread: 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 underlying C++ implementation pauses the progress loop for the + duration of the swap so the loop cannot dereference the held instance + while it is being replaced, then restores the prior running state. + + Parameters + ---------- + statistics + The new statistics instance. Must not be ``None``. Pass + ``Statistics.disabled()`` to opt out of statistics collection. + + Warnings + -------- + Concurrent calls to ``set_statistics`` are not allowed; the caller + must ensure this method is invoked from a single thread at a time. + """ + with nogil: + deref(self._handle).set_statistics(statistics._handle) + def __dealloc__(self): with nogil: self._handle.reset() diff --git a/python/rapidsmpf/rapidsmpf/statistics.pxd b/python/rapidsmpf/rapidsmpf/statistics.pxd index 09b1d8e9d..134620f2a 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pxd +++ b/python/rapidsmpf/rapidsmpf/statistics.pxd @@ -11,7 +11,6 @@ from libcpp.unordered_map cimport unordered_map from libcpp.vector cimport vector from rapidsmpf._detail.exception_handling cimport ex_handler -from rapidsmpf.config cimport cpp_Options from rapidsmpf.memory.pinned_memory_resource cimport (PinnedMemoryResource, cpp_PinnedMemoryResource) from rapidsmpf.memory.scoped_memory_record cimport cpp_ScopedMemoryRecord @@ -31,7 +30,6 @@ cdef extern from "" nogil: cdef cppclass cpp_Statistics "rapidsmpf::Statistics": bool enabled() except +ex_handler - void reset_from_options(cpp_Options options) except +ex_handler void add_stat( string name, double value diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyi b/python/rapidsmpf/rapidsmpf/statistics.pyi index 83f28cf96..e569c5824 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pyi +++ b/python/rapidsmpf/rapidsmpf/statistics.pyi @@ -28,7 +28,6 @@ class Statistics: def disabled(cls: type[Self]) -> Self: ... @property def enabled(self) -> bool: ... - def reset_from_options(self, new_options: Options) -> None: ... def report( self, *, diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyx b/python/rapidsmpf/rapidsmpf/statistics.pyx index 98151edb1..abd124ba2 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pyx +++ b/python/rapidsmpf/rapidsmpf/statistics.pyx @@ -220,30 +220,6 @@ cdef class Statistics: """ return deref(self._handle).enabled() - def reset_from_options(self, Options new_options not None) -> None: - """ - Resets the current state of the instance based on ``new_options``. - - .. warning:: - Live :class:`MemoryRecorder` instances (constructed before this - call but not yet destroyed) are *not* tracked. Their destructors - will run after the reset and publish their scope into the - freshly cleared memory records, mixing pre-reset measurements - into the new session. Likewise, a stream-ordered timing whose - stop callback is already mid-execution at reset time may still - record one stale stat (its global-map entry was extracted - before the cancellation could remove it). Callers that need - clean boundaries must quiesce recorders and synchronise the - relevant CUDA streams before calling. - - Parameters - ---------- - new_options - New configuration options - """ - with nogil: - deref(self._handle).reset_from_options(new_options._handle) - def report( self, *, From 2eca536845d5942106226a4be1f6ca7b4055514c Mon Sep 17 00:00:00 2001 From: niranda perera Date: Thu, 28 May 2026 13:15:00 -0700 Subject: [PATCH 07/13] set stats and weak ptr Signed-off-by: niranda perera --- cpp/include/rapidsmpf/statistics.hpp | 41 +++++++++++-------- .../rapidsmpf/stream_ordered_timing.hpp | 22 ++++++++-- cpp/src/statistics.cpp | 14 ++++--- cpp/src/stream_ordered_timing.cpp | 24 +++++++---- cpp/tests/test_stream_ordered_timing.cpp | 7 +++- .../rapidsmpf/communicator/communicator.pyx | 16 ++------ .../rapidsmpf/rapidsmpf/progress_thread.pyx | 11 ++--- 7 files changed, 77 insertions(+), 58 deletions(-) diff --git a/cpp/include/rapidsmpf/statistics.hpp b/cpp/include/rapidsmpf/statistics.hpp index 3f4da0bee..7ea0b984a 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_{}; }; @@ -666,23 +680,14 @@ concept StatisticsProvider = requires(T const& t) { * `std::invalid_argument` (or similar) on a null argument and otherwise install * the new instance unconditionally. * - * Only `ProgressThread` and `Communicator` — the *root* providers that drive the - * shared `Statistics` for an entire RapidsMPF runtime instance — satisfy this - * concept. Secondary providers such as `BufferResource` and `streaming::Context` - * deliberately do **not** offer mutation: their `Statistics` is fixed at - * construction. A typical caller-driven rebuild therefore swaps the instance on - * the progress thread and communicator and then constructs a fresh - * `streaming::Context` (and the `BufferResource` it owns) with the same new - * `Statistics`. - * - * @warning Concurrent calls to `set_statistics()` are not allowed; the - * caller must ensure this method is invoked from a single thread at a - * time. + * @warning Concurrent calls to `set_statistics()` will result in undefined behavior. */ template concept MutableStatisticsProvider = StatisticsProvider && requires(T& t, std::shared_ptr stats) { - { t.set_statistics(std::move(stats)) } -> std::same_as; + { + t.set_statistics(std::move(stats)) + } -> std::same_as; }; /** 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/src/statistics.cpp b/cpp/src/statistics.cpp index 3bc7a6d1a..1ef0aee73 100644 --- a/cpp/src/statistics.cpp +++ b/cpp/src/statistics.cpp @@ -248,10 +248,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(); } @@ -260,13 +260,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/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/communicator/communicator.pyx b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx index da3493005..6fde62a92 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx @@ -197,25 +197,15 @@ cdef class Communicator: """ Replace the statistics instance held by this communicator. - Concrete communicator implementations typically share their statistics - instance with their progress thread; when swapping at runtime, the - caller is responsible for keeping the two in sync by also calling - :meth:`ProgressThread.set_statistics` with the same instance and - rebuilding any dependent secondary providers (e.g. - :class:`~rapidsmpf.memory.buffer_resource.BufferResource` / - :class:`~rapidsmpf.streaming.core.context.Context`) from the same - new instance. - Parameters ---------- statistics - The new statistics instance. Must not be ``None``. Pass - ``Statistics.disabled()`` to opt out of statistics collection. + The new statistics instance. Pass ``Statistics.disabled()`` to opt out of + statistics collection. Warnings -------- - Concurrent calls to ``set_statistics`` are not allowed; the caller - must ensure this method is invoked from a single thread at a time. + Concurrent calls to ``set_statistics`` will result in undefined behavior. """ with nogil: deref(self._handle).set_statistics(statistics._handle) diff --git a/python/rapidsmpf/rapidsmpf/progress_thread.pyx b/python/rapidsmpf/rapidsmpf/progress_thread.pyx index 13f5a05b8..822a82f78 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pyx +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pyx @@ -53,20 +53,15 @@ cdef class ProgressThread: """ Replace the statistics instance held by this progress thread. - The underlying C++ implementation pauses the progress loop for the - duration of the swap so the loop cannot dereference the held instance - while it is being replaced, then restores the prior running state. - Parameters ---------- statistics - The new statistics instance. Must not be ``None``. Pass - ``Statistics.disabled()`` to opt out of statistics collection. + The new statistics instance. Pass ``Statistics.disabled()`` to opt out of + statistics collection. Warnings -------- - Concurrent calls to ``set_statistics`` are not allowed; the caller - must ensure this method is invoked from a single thread at a time. + Concurrent calls to ``set_statistics`` will result in undefined behavior. """ with nogil: deref(self._handle).set_statistics(statistics._handle) From 9ae3d1cdfd1ca5dbc5a18a39c7785d6d7d62794f Mon Sep 17 00:00:00 2001 From: niranda perera Date: Thu, 28 May 2026 13:43:13 -0700 Subject: [PATCH 08/13] move progress thread to communicator Signed-off-by: niranda perera --- .../rapidsmpf/communicator/communicator.hpp | 59 ++++++++----------- cpp/include/rapidsmpf/communicator/mpi.hpp | 9 --- cpp/include/rapidsmpf/communicator/single.hpp | 9 --- cpp/include/rapidsmpf/communicator/ucxx.hpp | 9 --- cpp/src/communicator/communicator.cpp | 17 ++---- cpp/src/communicator/mpi.cpp | 5 +- cpp/src/communicator/single.cpp | 5 +- cpp/src/communicator/ucxx.cpp | 7 +-- .../rapidsmpf/communicator/communicator.pxd | 3 - .../rapidsmpf/communicator/communicator.pyi | 1 - .../rapidsmpf/communicator/communicator.pyx | 17 ------ 11 files changed, 35 insertions(+), 106 deletions(-) diff --git a/cpp/include/rapidsmpf/communicator/communicator.hpp b/cpp/include/rapidsmpf/communicator/communicator.hpp index 4cef046b8..714e07ed3 100644 --- a/cpp/include/rapidsmpf/communicator/communicator.hpp +++ b/cpp/include/rapidsmpf/communicator/communicator.hpp @@ -405,20 +405,18 @@ class Communicator { protected: /** - * @brief Construct the base Communicator with a statistics instance. + * @brief Construct the base Communicator with a progress thread. * - * Derived classes typically forward `progress_thread->statistics()` here so - * that the communicator and its progress thread share the same `Statistics` - * instance. + * 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 statistics The statistics instance to associate with this - * communicator. Must not be null. + * @param progress_thread The progress thread for this communicator. Must + * not be null. * - * @throws std::invalid_argument If @p statistics is null. + * @throws std::invalid_argument If @p progress_thread is null. */ - explicit Communicator(std::shared_ptr statistics); - - std::shared_ptr statistics_; ///< Statistics instance. Never null. + explicit Communicator(std::shared_ptr progress_thread); public: virtual ~Communicator() noexcept = default; @@ -629,50 +627,41 @@ 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. * - * Satisfies the `StatisticsProvider` concept. Returns the instance supplied - * to the protected `Communicator` constructor by the concrete implementation - * (typically the same instance owned by `progress_thread()`). + * 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 statistics_; + return progress_thread_->statistics(); } - /** - * @brief Replace the statistics instance held by this communicator. - * - * Installs @p statistics in place of the current instance. Future - * `statistics()` calls return the new instance. - * - * @param statistics The new statistics instance. Must not be null. Pass - * `Statistics::disabled()` to opt out of statistics collection. - * - * @throws std::invalid_argument If @p statistics is null. - * - * @warning Concurrent calls to `set_statistics()` are not allowed; the - * caller must ensure this method is invoked from a single thread at a - * time. - */ - void set_statistics(std::shared_ptr 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); -static_assert(MutableStatisticsProvider); /// @brief Whether RapidsMPF was built with the UCXX Communicator. #ifdef RAPIDSMPF_HAVE_UCXX diff --git a/cpp/include/rapidsmpf/communicator/mpi.hpp b/cpp/include/rapidsmpf/communicator/mpi.hpp index 07df73689..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,7 +259,6 @@ class MPI final : public Communicator { Rank rank_; Rank nranks_; std::shared_ptr logger_; - std::shared_ptr progress_thread_; }; static_assert(StatisticsProvider); diff --git a/cpp/include/rapidsmpf/communicator/single.hpp b/cpp/include/rapidsmpf/communicator/single.hpp index b3b7369b9..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,7 +198,6 @@ class Single final : public Communicator { private: std::shared_ptr logger_; - std::shared_ptr progress_thread_; }; static_assert(StatisticsProvider); diff --git a/cpp/include/rapidsmpf/communicator/ucxx.hpp b/cpp/include/rapidsmpf/communicator/ucxx.hpp index 15fd71fd8..f0050e933 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(); diff --git a/cpp/src/communicator/communicator.cpp b/cpp/src/communicator/communicator.cpp index b73947d6d..96667dcff 100644 --- a/cpp/src/communicator/communicator.cpp +++ b/cpp/src/communicator/communicator.cpp @@ -12,24 +12,15 @@ namespace rapidsmpf { -Communicator::Communicator(std::shared_ptr statistics) - : statistics_{std::move(statistics)} { +Communicator::Communicator(std::shared_ptr progress_thread) + : progress_thread_{std::move(progress_thread)} { RAPIDSMPF_EXPECTS( - statistics_ != nullptr, - "Communicator statistics cannot be null", + progress_thread_ != nullptr, + "Communicator progress thread cannot be null", std::invalid_argument ); } -void Communicator::set_statistics(std::shared_ptr statistics) { - RAPIDSMPF_EXPECTS( - statistics != nullptr, - "Communicator statistics cannot be null", - std::invalid_argument - ); - statistics_ = std::move(statistics); -} - namespace { Communicator::Logger::LOG_LEVEL level_from_string(std::string const& str) { if (str.empty()) { diff --git a/cpp/src/communicator/mpi.cpp b/cpp/src/communicator/mpi.cpp index 3582efad2..e435fe7fa 100644 --- a/cpp/src/communicator/mpi.cpp +++ b/cpp/src/communicator/mpi.cpp @@ -102,7 +102,7 @@ MPI::MPI( config::Options options, std::shared_ptr progress_thread ) - : Communicator{progress_thread->statistics()}, + : Communicator{std::move(progress_thread)}, comm_{comm}, rank_{[comm]() { int r; @@ -114,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 6b80410bf..3859bf0e6 100644 --- a/cpp/src/communicator/single.cpp +++ b/cpp/src/communicator/single.cpp @@ -12,9 +12,8 @@ namespace rapidsmpf { Single::Single(config::Options options, std::shared_ptr progress_thread) - : Communicator{progress_thread->statistics()}, - 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 e1e31af60..97c400eed 100644 --- a/cpp/src/communicator/ucxx.cpp +++ b/cpp/src/communicator/ucxx.cpp @@ -1131,11 +1131,10 @@ UCXX::UCXX( config::Options options, std::shared_ptr progress_thread ) - : Communicator{progress_thread->statistics()}, + : 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(); } @@ -1505,7 +1504,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/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd b/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd index 6f6e06d94..5b3079801 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pxd @@ -38,9 +38,6 @@ cdef extern from "" nogil: string str() except +ex_handler shared_ptr[cpp_ProgressThread] progress_thread() except +ex_handler shared_ptr[cpp_Statistics] statistics() except +ex_handler - void set_statistics( - 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 a1075b929..ba1be8f53 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyi @@ -33,7 +33,6 @@ class Communicator: def progress_thread(self) -> ProgressThread: ... @property def statistics(self) -> Statistics: ... - def set_statistics(self, statistics: Statistics) -> None: ... 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 6fde62a92..8995c4cbd 100644 --- a/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx +++ b/python/rapidsmpf/rapidsmpf/communicator/communicator.pyx @@ -193,23 +193,6 @@ cdef class Communicator: stats._handle = deref(self._handle).statistics() return stats - def set_statistics(self, Statistics statistics not None): - """ - Replace the statistics instance held by this communicator. - - Parameters - ---------- - statistics - The new statistics instance. Pass ``Statistics.disabled()`` to opt out of - statistics collection. - - Warnings - -------- - Concurrent calls to ``set_statistics`` will result in undefined behavior. - """ - with nogil: - deref(self._handle).set_statistics(statistics._handle) - def get_str(self): """ Get a string representation of the communicator. From 5659dc0d46ba091e3f1f281d07f0368482d48455 Mon Sep 17 00:00:00 2001 From: Niranda Perera Date: Thu, 28 May 2026 15:17:08 -0700 Subject: [PATCH 09/13] Apply suggestions from code review Co-authored-by: Peter Andreas Entschev --- cpp/benchmarks/streaming/bench_streaming_shuffle.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp index de1f6577b..bfa1450c8 100644 --- a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp +++ b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp @@ -450,7 +450,6 @@ int main(int argc, char** argv) { log.print(ss.str()); } - auto statistics = ctx->statistics(); if (args.enable_memory_profiler) { log.print(stats->report({ .mr = stat_enabled_mr, From 4e0d47ce343843a26f08481c13d17d317473244e Mon Sep 17 00:00:00 2001 From: niranda perera Date: Thu, 28 May 2026 15:19:21 -0700 Subject: [PATCH 10/13] addressing initial comments Signed-off-by: niranda perera --- cpp/include/rapidsmpf/progress_thread.hpp | 18 ++++++------------ .../rapidsmpf/memory/buffer_resource.pyx | 6 +++--- python/rapidsmpf/rapidsmpf/statistics.pyx | 8 ++++---- 3 files changed, 13 insertions(+), 19 deletions(-) diff --git a/cpp/include/rapidsmpf/progress_thread.hpp b/cpp/include/rapidsmpf/progress_thread.hpp index d45cd56b9..9cc83c1d9 100644 --- a/cpp/include/rapidsmpf/progress_thread.hpp +++ b/cpp/include/rapidsmpf/progress_thread.hpp @@ -187,24 +187,18 @@ class ProgressThread { /** * @brief Replace the statistics instance held by this progress thread. * - * This method pauses the progress loop (via `pause()`) for the duration - * of the swap so the loop cannot dereference `statistics_` while it is - * being replaced, then restores the prior running state (calls `resume()` - * only if the loop was active before the call). + * This method pauses the progress loop for the duration of the swap so the loop + * cannot dereference `statistics_` while it is being replaced, then restores the + * prior running state. * - * Because `pause()` synchronizes with the loop, this call may block - * briefly until the in-flight iteration completes. - * - * @param statistics The new statistics instance. Must not be null. Pass - * `Statistics::disabled()` to opt out of statistics collection. + * @param statistics The new statistics instance. Pass `Statistics::disabled()` to opt + * out of statistics collection. * * @throws std::invalid_argument If @p statistics is null. Validated * before the pause so a null argument does not leave the progress loop * paused. * - * @warning Concurrent calls to `set_statistics()` are not allowed; the - * caller must ensure this method is invoked from a single thread at a - * time. + * @warning Concurrent calls to `set_statistics()` will result in undefined behavior. */ void set_statistics(std::shared_ptr statistics); diff --git a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx index 4350265a1..ea961e892 100644 --- a/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/buffer_resource.pyx @@ -113,6 +113,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. The BufferResource transparently wraps this resource in an internal RMM @@ -135,9 +138,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. Required. Pass - ``Statistics.disabled()`` for a no-op recorder. """ def __cinit__( self, diff --git a/python/rapidsmpf/rapidsmpf/statistics.pyx b/python/rapidsmpf/rapidsmpf/statistics.pyx index abd124ba2..4a77625dc 100644 --- a/python/rapidsmpf/rapidsmpf/statistics.pyx +++ b/python/rapidsmpf/rapidsmpf/statistics.pyx @@ -189,14 +189,14 @@ cdef class Statistics: @classmethod def disabled(cls): """ - Get a shared, no-op Statistics instance. + Returns a disabled (no-op) Statistics instance. - Returns the same disabled instance on every call. Use this when an API - requires a Statistics object but no recording is desired. + Useful when you need to pass a Statistics argument but do not want to + collect any data. Returns ------- - A disabled Statistics instance. + A Statistics instance with tracking disabled. """ cdef Statistics ret = cls.__new__(cls) with nogil: From 37ff748ba116f5be45a74407cc0e0da67a7a7288 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Thu, 28 May 2026 15:47:30 -0700 Subject: [PATCH 11/13] addressing comments Signed-off-by: niranda perera --- cpp/benchmarks/streaming/bench_streaming_shuffle.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp index bfa1450c8..f0c0488cc 100644 --- a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp +++ b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp @@ -324,7 +324,8 @@ int main(int argc, char** argv) { // Initialize configuration options from environment variables. rapidsmpf::config::Options options{rapidsmpf::config::get_environment_variables()}; - auto stats = rapidsmpf::Statistics::from_options(options); + + auto stats = rapidsmpf::Statistics::create(); auto progress_thread = std::make_shared(stats); std::shared_ptr comm; From 78d8a04e9a906720832b17d43a8192ea0823c691 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Thu, 28 May 2026 17:25:51 -0700 Subject: [PATCH 12/13] add test Signed-off-by: niranda perera --- cpp/tests/test_progress_thread.cpp | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/cpp/tests/test_progress_thread.cpp b/cpp/tests/test_progress_thread.cpp index 79bf96e15..32a3291c2 100644 --- a/cpp/tests/test_progress_thread.cpp +++ b/cpp/tests/test_progress_thread.cpp @@ -10,6 +10,7 @@ #include +#include #include #include @@ -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()); +} From ab50f50859589f87bccc629274b7ff56d3fab143 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Fri, 29 May 2026 13:45:26 -0700 Subject: [PATCH 13/13] addressing comments Signed-off-by: niranda perera --- cpp/benchmarks/bench_comm.cpp | 9 ++++++++- .../streaming/bench_streaming_shuffle.cpp | 7 ++++++- cpp/include/rapidsmpf/progress_thread.hpp | 14 +++++--------- cpp/include/rapidsmpf/statistics.hpp | 4 +++- cpp/src/progress_thread.cpp | 19 ++++++++----------- .../rapidsmpf/rapidsmpf/progress_thread.pyx | 6 ++---- 6 files changed, 32 insertions(+), 27 deletions(-) diff --git a/cpp/benchmarks/bench_comm.cpp b/cpp/benchmarks/bench_comm.cpp index 7416d80e8..4727870cd 100644 --- a/cpp/benchmarks/bench_comm.cpp +++ b/cpp/benchmarks/bench_comm.cpp @@ -317,7 +317,14 @@ int main(int argc, char** argv) { set_current_rmm_resource(args.rmm_mr); rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); - BufferResource br{stats, mr}; + BufferResource br{ + stats, + mr, + PinnedMemoryResource::Disabled, + {}, + std::chrono::milliseconds{1}, + std::make_shared(16, rmm::cuda_stream::flags::non_blocking) + }; // Print benchmark/hardware info. { diff --git a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp index f0c0488cc..2f54a8ce2 100644 --- a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp +++ b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp @@ -370,7 +370,12 @@ int main(int argc, char** argv) { ? rapidsmpf::PinnedMemoryResource::Disabled : rapidsmpf::PinnedMemoryResource::make_if_available(); auto br = std::make_shared( - stats, stat_enabled_mr, pinned_mr, std::move(memory_limits) + stats, + stat_enabled_mr, + pinned_mr, + std::move(memory_limits), + std::nullopt, + std::make_shared(16, rmm::cuda_stream::flags::non_blocking) ); auto& log = *comm->logger(); diff --git a/cpp/include/rapidsmpf/progress_thread.hpp b/cpp/include/rapidsmpf/progress_thread.hpp index 9cc83c1d9..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 @@ -187,18 +188,13 @@ class ProgressThread { /** * @brief Replace the statistics instance held by this progress thread. * - * This method pauses the progress loop for the duration of the swap so the loop - * cannot dereference `statistics_` while it is being replaced, then restores the - * prior running state. + * 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. Validated - * before the pause so a null argument does not leave the progress loop - * paused. - * - * @warning Concurrent calls to `set_statistics()` will result in undefined behavior. + * @throws std::invalid_argument If @p statistics is null. */ void set_statistics(std::shared_ptr statistics); @@ -211,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_; diff --git a/cpp/include/rapidsmpf/statistics.hpp b/cpp/include/rapidsmpf/statistics.hpp index 7ea0b984a..92e777236 100644 --- a/cpp/include/rapidsmpf/statistics.hpp +++ b/cpp/include/rapidsmpf/statistics.hpp @@ -680,7 +680,9 @@ concept StatisticsProvider = requires(T const& t) { * `std::invalid_argument` (or similar) on a null argument and otherwise install * the new instance unconditionally. * - * @warning Concurrent calls to `set_statistics()` will result in undefined behavior. + * @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 = diff --git a/cpp/src/progress_thread.cpp b/cpp/src/progress_thread.cpp index bc69f378e..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,7 @@ 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) { @@ -125,14 +127,7 @@ void ProgressThread::set_statistics(std::shared_ptr statistics) { "the statistics pointer cannot be NULL", std::invalid_argument ); - // Pausing ensures the progress loop is not in the middle of dereferencing - // `statistics_` while we swap it. Restore the prior running state on exit. - bool const was_running = is_running(); - pause(); - statistics_ = std::move(statistics); - if (was_running) { - resume(); - } + statistics_.store(std::move(statistics)); } void ProgressThread::event_loop() { @@ -147,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/python/rapidsmpf/rapidsmpf/progress_thread.pyx b/python/rapidsmpf/rapidsmpf/progress_thread.pyx index 822a82f78..ebb0b5437 100644 --- a/python/rapidsmpf/rapidsmpf/progress_thread.pyx +++ b/python/rapidsmpf/rapidsmpf/progress_thread.pyx @@ -53,15 +53,13 @@ cdef class ProgressThread: """ 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. - - Warnings - -------- - Concurrent calls to ``set_statistics`` will result in undefined behavior. """ with nogil: deref(self._handle).set_statistics(statistics._handle)