Support both approaches

Author: huuanhhuyn

cpp/include/raft/core/memory_tracking_resources.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -69,7 +69,8 @@ RAFT_EXPORT inline auto get_default_host_resource() -> raft::mr::host_resource_r
  * @param res The resource to install.
  * @return The previous default host resource.
  */
-RAFT_EXPORT inline auto set_default_host_resource(raft::mr::host_resource res) -> raft::mr::host_resource
+RAFT_EXPORT inline auto set_default_host_resource(raft::mr::host_resource res)
+  -> raft::mr::host_resource
 {
   return detail::default_host_resource_holder_.set(res);
 }

cpp/include/raft/mr/host_memory_resource.hpp

@@ -67,7 +67,7 @@ class notifier {
  *        or deallocation.
  *
  * Forwards all calls to the upstream resource, then calls notifier::notify().
- * A separate consumer (e.g. resource_monitor) can call notifier::wait() to
+ * A separate consumer (e.g. sampling_monitor) can call notifier::wait() to
  * block until activity occurs.
  *
  * @tparam Upstream  Stored by value.  Use a concrete resource type for owning

cpp/include/raft/mr/notifying_adaptor.hpp

@@ -6,7 +6,7 @@
 
 #include <raft/core/detail/macros.hpp>
 #include <raft/core/detail/nvtx_range_stack.hpp>  // thread_local_current_range
-#include <raft/mr/allocation_event_monitor.hpp>   // allocation_event, allocation_event_queue
+#include <raft/mr/recording_monitor.hpp>   // allocation_event, allocation_event_queue
 #include <raft/mr/statistics_adaptor.hpp>         // resource_stats (atomic counters, reused)
 
 #include <cuda/memory_resource>

cpp/include/raft/mr/recording_adaptor.hpp

@@ -89,14 +89,14 @@ class allocation_event_queue {
  * @brief Consumes allocation_events from a queue and writes one CSV row per
  *        event from a background thread.
  */
-class allocation_event_monitor {
+class recording_monitor {
  public:
-  explicit allocation_event_monitor(std::ostream& out) : out_(out) {}
+  explicit recording_monitor(std::ostream& out) : out_(out) {}
 
-  ~allocation_event_monitor() { stop(); }
+  ~recording_monitor() { stop(); }
 
-  allocation_event_monitor(allocation_event_monitor const&)            = delete;
-  allocation_event_monitor& operator=(allocation_event_monitor const&) = delete;
+  recording_monitor(recording_monitor const&)            = delete;
+  recording_monitor& operator=(recording_monitor const&) = delete;
 
   [[nodiscard]] auto get_queue() const noexcept -> std::shared_ptr<allocation_event_queue>
   {

…ude/raft/mr/allocation_event_monitor.hpp cpp/include/raft/mr/recording_monitor.hppcpp/include/raft/mr/allocation_event_monitor.hpp renamed to cpp/include/raft/mr/recording_monitor.hpp cpp/include/raft/mr/allocation_event_monitor.hpp renamed to cpp/include/raft/mr/recording_monitor.hpp

@@ -39,7 +39,7 @@ namespace mr {
  *
  * start() and stop() are idempotent.
  */
-class resource_monitor {
+class sampling_monitor {
   std::ostream& out_;
   std::chrono::steady_clock::duration sample_interval_;
   std::shared_ptr<notifier> notifier_;
@@ -55,18 +55,18 @@ class resource_monitor {
    * @param out             Output stream for CSV rows.
    * @param sample_interval Minimum time between successive samples.
    */
-  explicit resource_monitor(std::ostream& out, std::chrono::steady_clock::duration sample_interval)
+  explicit sampling_monitor(std::ostream& out, std::chrono::steady_clock::duration sample_interval)
     : out_(out),
       sample_interval_(sample_interval),
       notifier_(std::make_shared<notifier>()),
       nvtx_range_(raft::common::nvtx::thread_local_current_range())
   {
   }
 
-  ~resource_monitor() { stop(); }
+  ~sampling_monitor() { stop(); }
 
-  resource_monitor(resource_monitor const&)            = delete;
-  resource_monitor& operator=(resource_monitor const&) = delete;
+  sampling_monitor(sampling_monitor const&)            = delete;
+  sampling_monitor& operator=(sampling_monitor const&) = delete;
 
   /**
    * @brief Shared notifier for notifying_adaptor instances.

cpp/include/raft/mr/resource_monitor.hpp cpp/include/raft/mr/sampling_monitor.hppcpp/include/raft/mr/resource_monitor.hpp renamed to cpp/include/raft/mr/sampling_monitor.hpp

@@ -5,7 +5,7 @@
 
 #include <raft/mr/host_memory_resource.hpp>
 #include <raft/mr/notifying_adaptor.hpp>
-#include <raft/mr/resource_monitor.hpp>
+#include <raft/mr/sampling_monitor.hpp>
 #include <raft/mr/statistics_adaptor.hpp>
 
 #include <gtest/gtest.h>
@@ -69,7 +69,7 @@ TEST(AllocationReport, WritesCSVOnDirty)
   using namespace std::chrono_literals;
 
   std::ostringstream oss;
-  raft::mr::resource_monitor report(oss, 1ms);
+  raft::mr::sampling_monitor report(oss, 1ms);
 
   auto host_stats   = std::make_shared<raft::mr::resource_stats>();
   auto pinned_stats = std::make_shared<raft::mr::resource_stats>();
@@ -98,7 +98,7 @@ TEST(AllocationReport, StartStopIdempotent)
   using namespace std::chrono_literals;
 
   std::ostringstream oss;
-  raft::mr::resource_monitor report(oss, 1ms);
+  raft::mr::sampling_monitor report(oss, 1ms);
 
   auto stats = std::make_shared<raft::mr::resource_stats>();
   report.register_source("test", stats);
@@ -124,7 +124,7 @@ TEST(AllocationReport, DestructorCallsStop)
   std::ostringstream oss;
   {
     auto stats = std::make_shared<raft::mr::resource_stats>();
-    raft::mr::resource_monitor report(oss, 1ms);
+    raft::mr::sampling_monitor report(oss, 1ms);
     report.register_source("test", stats);
 
     stats->record_allocate(256);

cpp/tests/core/allocation_tracking.cpp

@@ -28,7 +28,9 @@ namespace nvtx = raft::common::nvtx;
 using namespace std::chrono_literals;
 constexpr std::size_t MiB = std::size_t{1024} * 1024;
 
-TEST(MemoryTrackingResources, TracksDeviceAllocations)
+// TODO improve tests (coverage + multiple allocating threads)
+
+TEST(MemoryTrackingResources, Sampling)
 {
   std::ostringstream oss;
   {
@@ -58,34 +60,40 @@ TEST(MemoryTrackingResources, TracksDeviceAllocations)
                           << output;
 }
 
-TEST(MemoryTrackingResources, MismatchedRangeLabeling)
+TEST(MemoryTrackingResources, Recording)
 {
-  const std::string csv_path = "mismatch_range_label.csv";
-
+  std::ostringstream oss;
   {
     raft::resources res;
-
-    raft::memory_tracking_resources tracked(res, csv_path, 1ms);
+    raft::memory_tracking_resources tracked(res, oss);
     {
       nvtx::range r{"1. expect 10 KB"};
       auto matrix = raft::make_host_vector<uint8_t>(tracked, 10 * 1024);
     }
     {
-      // Deliberately huge & slow: allocating/freeing 10 GiB of host memory takes
-      // several ms, which makes the background sampler lag past this range's end.
-      // As a result this allocation's peak is mis-attributed to the NEXT range in
-      // the CSV (the range-labeling race discussed in the file header). Source
-      // attribution (host) stays correct; only the nvtx_range label is wrong.
-      nvtx::range r{"2. expect 10 GiB"};
-      auto vector = raft::make_host_vector<uint8_t>(tracked, 10 * 1024 * MiB);
+      // Deliberately large allocation to test that the memory tracking
+      // resources can handle a large allocation and labels it correctly
+      nvtx::range r{"2. expect 100 MiB"};
+      auto vector = raft::make_host_vector<uint8_t>(tracked, 100 * MiB);
     }
     {
       nvtx::range r{"3. expect 4 MiB"};
       auto matrix = raft::make_host_vector<uint8_t>(tracked, 4 * MiB);
     }
   }  // tracked destroyed here: stops the sampler and flushes the file
 
-  std::cout << "Wrote allocation statistics to " << csv_path << "\n";
+  auto output                         = oss.str();
+  auto num_lines                      = std::count(output.begin(), output.end(), '\n');
+  constexpr size_t NUM_ALLOCS         = 3;
+  constexpr size_t NUM_DEALLOCS       = NUM_ALLOCS;
+  constexpr size_t NUM_HEADER_LINES   = 1;
+  constexpr size_t NUM_LINES_EXPECTED = NUM_ALLOCS + NUM_DEALLOCS + NUM_HEADER_LINES;
+  EXPECT_GE(num_lines, NUM_LINES_EXPECTED)
+    << "Expected at least " << NUM_LINES_EXPECTED
+    << " data records (allocation + deallocation + header); got " << num_lines << " lines"
+    << std::endl
+    << "content: " << std::endl
+    << output;
 }
 
 }  // namespace