Initial stats cleanup

cpp/benchmarks/bench_shuffle.cpp

@@ -543,7 +543,7 @@ int main(int argc, char** argv) {
         memory_limits[rapidsmpf::MemoryType::DEVICE] = args.device_mem_limit_mb << 20;
     }
 
-    auto stats = std::make_shared<rapidsmpf::Statistics>(/* enable = */ true);
+    auto stats = rapidsmpf::Statistics::create();
 
     // We're only going to measure the last run, so disable initially.
     stats->disable();

cpp/benchmarks/streaming/bench_streaming_shuffle.cpp

@@ -364,7 +364,7 @@ int main(int argc, char** argv) {
         memory_limits[rapidsmpf::MemoryType::DEVICE] = args.device_mem_limit_mb << 20;
     }
 
-    auto stats = std::make_shared<rapidsmpf::Statistics>(/* enable = */ true);
+    auto stats = rapidsmpf::Statistics::create();
 
     auto pinned_mr = args.pinned_mem_disable
                          ? rapidsmpf::PinnedMemoryResource::Disabled
@@ -458,14 +458,15 @@ int main(int argc, char** argv) {
         log.print(ss.str());
     }
 
+    auto statistics = ctx->statistics();
     if (args.enable_memory_profiler) {
-        log.print(ctx->statistics()->report({
+        log.print(statistics->report({
             .mr = stat_enabled_mr,
             .pinned_mr = pinned_mr,
             .header = "Statistics (of the last run):",
         }));
     } else {
-        log.print(ctx->statistics()->report({.header = "Statistics (of the last run):"}));
+        log.print(statistics->report({.header = "Statistics (of the last run):"}));
     }
 
     if (!use_bootstrap) {

cpp/benchmarks/streaming/ndsh/bench_read.cpp

@@ -411,10 +411,11 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report(
+        auto statistics = ctx->statistics();
+        comm->logger()->print(statistics->report(
             {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()}
         ));
-        ctx->statistics()->clear();
+        statistics->clear();
     }
 
     if (comm->rank() == 0) {

cpp/benchmarks/streaming/ndsh/q01.cpp

@@ -433,10 +433,11 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report(
+        auto statistics = ctx->statistics();
+        comm->logger()->print(statistics->report(
             {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()}
         ));
-        ctx->statistics()->clear();
+        statistics->clear();
     }
 
     if (comm->rank() == 0) {

cpp/benchmarks/streaming/ndsh/q03.cpp

@@ -695,10 +695,11 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report(
+        auto statistics = ctx->statistics();
+        comm->logger()->print(statistics->report(
             {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()}
         ));
-        ctx->statistics()->clear();
+        statistics->clear();
     }
     if (comm->rank() == 0) {
         for (int i = 0; i < arguments.num_iterations; i++) {

cpp/benchmarks/streaming/ndsh/q04.cpp

@@ -518,10 +518,11 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report(
+        auto statistics = ctx->statistics();
+        comm->logger()->print(statistics->report(
             {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()}
         ));
-        ctx->statistics()->clear();
+        statistics->clear();
     }
 
     if (comm->rank() == 0) {

cpp/benchmarks/streaming/ndsh/q09.cpp

@@ -653,10 +653,11 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report(
+        auto statistics = ctx->statistics();
+        comm->logger()->print(statistics->report(
             {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()}
         ));
-        ctx->statistics()->clear();
+        statistics->clear();
     }
     if (comm->rank() == 0) {
         for (std::size_t i = 0; i < safe_cast<std::size_t>(arguments.num_iterations); i++)

cpp/benchmarks/streaming/ndsh/q21.cpp

@@ -995,10 +995,11 @@ int main(int argc, char** argv) {
         std::chrono::duration<double> compute = end - start;
         timings.push_back(pipeline.count());
         timings.push_back(compute.count());
-        comm->logger()->print(ctx->statistics()->report(
+        auto statistics = ctx->statistics();
+        comm->logger()->print(statistics->report(
             {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()}
         ));
-        ctx->statistics()->clear();
+        statistics->clear();
     }
 
     if (comm->rank() == 0) {

cpp/benchmarks/streaming/ndsh/utils.cpp

@@ -144,7 +144,7 @@ create_context(
 
         memory_limits[MemoryType::DEVICE] = static_cast<std::int64_t>(limit_size);
     }
-    auto statistics = std::make_shared<Statistics>(/* enable = */ true);
+    auto statistics = Statistics::create();
 
     RAPIDSMPF_EXPECTS(
         arguments.no_pinned_host_memory || is_pinned_memory_resources_supported(),

cpp/examples/example_shuffle.cpp

@@ -27,7 +27,7 @@ int main(int argc, char** argv) {
     rapidsmpf::config::Options options{rapidsmpf::config::get_environment_variables()};
 
     // Create a statistics instance for the shuffler that tracks useful information.
-    auto stats = std::make_shared<rapidsmpf::Statistics>();
+    auto stats = rapidsmpf::Statistics::create();
 
     // The communicator has a progress thread where the shuffler event loop executes. A
     // single progress thread may be used by multiple shufflers simultaneously.

cpp/include/rapidsmpf/coll/allgather.hpp

@@ -23,7 +23,6 @@
 #include <rapidsmpf/memory/packed_data.hpp>
 #include <rapidsmpf/memory/spill_manager.hpp>
 #include <rapidsmpf/progress_thread.hpp>
-#include <rapidsmpf/statistics.hpp>
 
 /**
  * @namespace rapidsmpf::coll
@@ -101,8 +100,6 @@ class AllGather {
      * @param comm The communicator for communication.
      * @param op_id Unique operation identifier for this allgather.
      * @param br Buffer resource for memory allocation.
-     * @param statistics Statistics collection instance (disabled by
-     * default).
      * @param finished_callback Optional callback run when partitions are locally
      * finished. The callback is guaranteed to be called by the progress thread exactly
      * once when the allgather is locally ready.
@@ -118,7 +115,6 @@ class AllGather {
         std::shared_ptr<Communicator> comm,
         OpID op_id,
         BufferResource* br,
-        std::shared_ptr<Statistics> statistics = Statistics::disabled(),
         std::function<void(void)>&& finished_callback = nullptr
     );
 
@@ -195,7 +191,6 @@ class AllGather {
 
     std::shared_ptr<Communicator> comm_;  ///< Communicator
     BufferResource* br_;  ///< Buffer resource for memory allocation
-    std::shared_ptr<Statistics> statistics_;  ///< Statistics collection instance
     std::function<void(void)> finished_callback_{
         nullptr
     };  ///< Optional callback to run when allgather is finished and ready for extraction.

cpp/include/rapidsmpf/memory/buffer_resource.hpp

@@ -101,12 +101,15 @@ class BufferResource {
      *
      * @param mr A device-accessible RMM memory resource.
      * @param options Configuration options.
+     * @param statistics The statistics instance to use (disabled by default).
      *
      * @return A shared pointer to a BufferResource instance configured according to the
      * options.
      */
     static std::shared_ptr<BufferResource> from_options(
-        cuda::mr::any_resource<cuda::mr::device_accessible> mr, config::Options options
+        cuda::mr::any_resource<cuda::mr::device_accessible> mr,
+        config::Options options,
+        std::shared_ptr<Statistics> statistics = Statistics::disabled()
     );
 
     ~BufferResource() noexcept = default;
@@ -416,7 +419,7 @@ class BufferResource {
      *
      * @return Shared pointer the Statistics instance.
      */
-    std::shared_ptr<Statistics> statistics();
+    std::shared_ptr<Statistics> statistics() const noexcept;
 
   private:
     std::mutex mutex_;
@@ -435,6 +438,8 @@ class BufferResource {
     std::shared_ptr<Statistics> statistics_;
 };
 
+static_assert(StatisticsProvider<BufferResource>);
+
 /**
  * @brief Parse the `spill_device_limit` parameter from configuration options.
  *

cpp/include/rapidsmpf/progress_thread.hpp

@@ -205,4 +205,6 @@ class ProgressThread {
     detail::PausableThreadLoop thread_;
 };
 
+static_assert(StatisticsProvider<ProgressThread>);
+
 }  // namespace rapidsmpf