[WIP] Memory logging resources#3051
Conversation
ef5b83f to
e852f30
Compare
huuanhhuyn
changed the title
e852f30 to
f4e4cee
Compare
huuanhhuyn
changed the title
32d3c6c to
455f959
Compare
455f959 to
a315972
Compare
huuanhhuyn
changed the title
There was a problem hiding this comment.
From the description it follows that the new queue approach attaches NVTX annotations to the allocation events one-to-one. If so, it seems logical that NVTX annotations should be pulled from the allocating thread during the allocation rather than from the main thread. Then you'd also not need mutex locking of nvtx records since they are not accessed across threads.
Please refactor this as a separate resource type rather than changing the behavior of the existing resource, because the difference is significant.
tfeher
added
feature request
huuanhhuyn
changed the title
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe change adds NVTX path tracking, asynchronous allocation-event recording, and ChangesMemory recording and monitoring
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
3747994 to
9ac8dec
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🧹 Nitpick comments (1)
cpp/tests/core/monitor_resources.cu (1)
147-183: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
SamplingParallelThreadscorrectness relies on scheduling/timing, risking CI flakiness.The assertion
3 < num_lines && num_lines < max_num_rowsonly proves "some drop happened," but the lower bound of 3 is tight: under heavy CI load, if the sampling thread is starved and the 1us interval barely fires beforestop(),num_linescould plausibly sit at or near the header + final flush row, making the test brittle. Also the upper bound is trivially satisfied by any dropped events, so the test provides weak signal for the "sampling drops many rows" claim being validated.As per path instructions (
docs/source/developer_guide.mdreferenced guidance): "ensure tests validate functional requirements (including sampling drop behavior)... keep tests deterministic/robust." Consider asserting a stronger, still-lenient bound (e.g.,num_lines < max_num_rows / 2) to better demonstrate drops, or documenting the acceptable flake tolerance.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tests/core/monitor_resources.cu` around lines 147 - 183, Strengthen the assertions in SamplingParallelThreads so the test deterministically verifies substantial sampling, not merely that a few rows were emitted. Keep a lenient positive lower bound for sampled output, and require num_lines to remain below roughly half of max_num_rows (or use an equivalently documented tolerance) while preserving the existing diagnostic output.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@cpp/include/raft/core/memory_logging_resources.hpp`:
- Around line 74-79: Update memory_logging_resources construction and
destruction to serialize or correctly nest process-wide host/device resource
swaps, using a process-wide guard around the lifetime of each active instance.
Ensure teardown restores resources only in strict LIFO order, preventing
concurrent or out-of-order instances from restoring stale resources while
another logger remains active.
In `@cpp/include/raft/mr/recording_adaptor.hpp`:
- Around line 96-103: Update the public recording_adaptor constructor to
validate the queue argument before storing or using it, rejecting a null
shared_ptr immediately with the project’s established precondition/error
mechanism. Preserve normal construction for non-null queues and ensure queue_
cannot be dereferenced by an invalid instance.
- Around line 52-92: Make record_allocation, forget_allocation, and emit robust
against allocation failures: contain exceptions from string operations,
unordered_map updates, event construction, and queue insertion so these noexcept
recording paths never terminate. Treat each failed recording operation as
best-effort by dropping only that record, while preserving successful
bookkeeping and event delivery.
- Around line 123-125: In both synchronous deallocation sites at
cpp/include/raft/mr/recording_adaptor.hpp:123-125 and stream deallocation sites
at cpp/include/raft/mr/recording_adaptor.hpp:145-147, capture and remove the
allocation mapping via forget_allocation(ptr) before calling the corresponding
upstream deallocation, then emit the captured path afterward while preserving
the existing byte accounting.
In `@cpp/include/raft/mr/recording_monitor.hpp`:
- Around line 141-146: Update the CSV output logic in the recording monitor,
including the header generation around source_names_ and the row-writing logic
for nvtx_range and alloc_range, to escape string fields according to CSV rules.
Double every embedded quote before writing registered source names, NVTX range
names, and allocation paths, while preserving the existing field delimiters and
output format.
- Around line 118-130: Document all newly exposed public APIs with Doxygen: in
cpp/include/raft/mr/recording_monitor.hpp lines 118-130, document start() and
stop(), including lifecycle behavior and the requirement to shut down producers
before stopping; in cpp/include/raft/mr/recording_adaptor.hpp lines 96-103,
document the constructor’s queue ownership and non-null requirement; and in
lines 111-157, document allocation/deallocation recording semantics and accessor
behavior.
- Around line 172-174: Update the serialization loop in the recording monitor to
output a true per-source peak value for the <source>_peak column instead of
repeating v.current. Track and retain each source’s maximum observed allocation
across frees, then write that retained value while preserving the existing
current, total_alloc, and total_freed columns.
- Around line 72-79: Update RecordingMonitor::stop and the producer path in push
so shutdown coordinates with active producers: prevent new events from being
accepted once shutdown begins, wait for in-progress producers to finish
enqueueing, and only complete after the worker drains all accepted events.
Preserve the existing mutex and condition-variable synchronization while
ensuring no event can remain queued after stop returns.
---
Nitpick comments:
In `@cpp/tests/core/monitor_resources.cu`:
- Around line 147-183: Strengthen the assertions in SamplingParallelThreads so
the test deterministically verifies substantial sampling, not merely that a few
rows were emitted. Keep a lenient positive lower bound for sampled output, and
require num_lines to remain below roughly half of max_num_rows (or use an
equivalently documented tolerance) while preserving the existing diagnostic
output.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: e3ca39c2-0b14-4544-9cd4-f20115772ac5
📒 Files selected for processing (9)
cpp/include/raft/core/detail/nvtx_range_stack.hppcpp/include/raft/core/memory_logging_resources.hppcpp/include/raft/mr/host_memory_resource.hppcpp/include/raft/mr/notifying_adaptor.hppcpp/include/raft/mr/recording_adaptor.hppcpp/include/raft/mr/recording_monitor.hppcpp/include/raft/mr/resource_monitor.hppcpp/tests/core/allocation_tracking.cppcpp/tests/core/monitor_resources.cu
| ~memory_logging_resources() override | ||
| { | ||
| if (recorder_) recorder_->stop(); | ||
| raft::mr::set_default_host_resource(old_host_); | ||
| rmm::mr::set_current_device_resource(old_device_); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== file outline ==\n'
ast-grep outline cpp/include/raft/core/memory_logging_resources.hpp --view expanded || true
printf '\n== relevant lines ==\n'
nl -ba cpp/include/raft/core/memory_logging_resources.hpp | sed -n '1,240p'
printf '\n== search for memory_logging_resources usages/constructors ==\n'
rg -n "memory_logging_resources|set_default_host_resource|set_current_device_resource|recording_adaptor" cpp -g '!**/build/**' -g '!**/dist/**'
printf '\n== search for mutex/guard/nesting hints in file ==\n'
rg -n "mutex|lock|atomic|nested|thread|global|default_host_resource|current_device_resource" cpp/include/raft/core/memory_logging_resources.hpp cpp -g '!**/build/**' -g '!**/dist/**'Repository: NVIDIA/raft
Length of output: 712
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path("cpp/include/raft/core/memory_logging_resources.hpp")
for start,end in [(1,220)]:
print(f"\n== {path} lines {start}-{end} ==\n")
with path.open() as f:
for i, line in enumerate(f, start=1):
if start <= i <= end:
print(f"{i:4d}: {line.rstrip()}")
PY
printf '\n== related declarations ==\n'
rg -n "class memory_logging_resources|snapshot_|old_host_|old_device_|recorder_|host_record_adaptor_|device_record_adaptor_|set_default_host_resource|set_current_device_resource" cpp/include/raft/core/memory_logging_resources.hpp cpp/src cpp -g '!**/build/**' -g '!**/dist/**'Repository: NVIDIA/raft
Length of output: 20939
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
paths = [
Path("cpp/include/raft/mr/host_memory_resource.hpp"),
Path("cpp/include/raft/core/memory_stats_resources.hpp"),
Path("cpp/include/raft/core/memory_tracking_resources.hpp"),
]
for path in paths:
print(f"\n== {path} ==")
if not path.exists():
print("missing")
continue
with path.open() as f:
for i, line in enumerate(f, start=1):
if i <= 260:
print(f"{i:4d}: {line.rstrip()}")
PY
printf '\n== search for swap guards / mutexes around set_current_device_resource ==\n'
rg -n "mutex|lock_guard|scoped_lock|set_current_device_resource|set_default_host_resource|global.*resource|single instance|one instance" cpp/include/raft cpp/tests -g '!**/build/**' -g '!**/dist/**'Repository: NVIDIA/raft
Length of output: 33904
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '== declarations/uses of set_current_device_resource ==\n'
rg -n "set_current_device_resource\(" -g '!**/build/**' -g '!**/dist/**' .
printf '\n== host_resource and device_resource type aliases ==\n'
rg -n "using .*device_resource|using .*host_resource|get_current_device_resource_ref|set_current_device_resource" cpp/include/raft cpp/tests -g '!**/build/**' -g '!**/dist/**'Repository: NVIDIA/raft
Length of output: 7484
Avoid overlapping memory_logging_resources lifetimes.
This wrapper swaps process-wide host/device resources and restores the saved values on teardown. That is only safe for strictly nested LIFO use; concurrent instances or out-of-order destruction can restore a stale adaptor while another logger is still active, dropping tracked allocations or leaving the global device resource wrong. Add a process-wide guard or make the swap nest-aware.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/include/raft/core/memory_logging_resources.hpp` around lines 74 - 79,
Update memory_logging_resources construction and destruction to serialize or
correctly nest process-wide host/device resource swaps, using a process-wide
guard around the lifetime of each active instance. Ensure teardown restores
resources only in strict LIFO order, preventing concurrent or out-of-order
instances from restoring stale resources while another logger remains active.
Source: Coding guidelines
| auto record_allocation(void* ptr) noexcept -> std::string | ||
| { | ||
| std::string path = ""; | ||
| if (ptr != nullptr) { | ||
| path = raft::common::nvtx::thread_local_current_path(); | ||
| if (!path.empty()) { | ||
| std::lock_guard<std::mutex> lock(alloc_map_->mtx); | ||
| alloc_map_->paths[ptr] = path; | ||
| } | ||
| } | ||
| return path; | ||
| } | ||
|
|
||
| // Returns the NVTX path recorded at alloc time for this pointer, then removes it. | ||
| auto forget_allocation(void* ptr) noexcept -> std::string | ||
| { | ||
| std::string path = ""; | ||
| std::lock_guard<std::mutex> lock(alloc_map_->mtx); | ||
| auto it = alloc_map_->paths.find(ptr); | ||
| if (it != alloc_map_->paths.end()) { | ||
| path = std::move(it->second); | ||
| alloc_map_->paths.erase(it); | ||
| } | ||
| return path; | ||
| } | ||
|
|
||
| // Enqueue an event. Called on the allocating/deallocating thread — mutex-free NVTX read. | ||
| void emit(std::string alloc_range, std::int64_t signed_bytes) noexcept | ||
| { | ||
| auto [name, depth] = raft::common::nvtx::thread_local_current_name_and_depth(); | ||
| allocation_event event; | ||
| event.source_id = source_id_; | ||
| event.current = stats_->bytes_current.load(std::memory_order_relaxed); | ||
| event.total_alloc = stats_->bytes_total_allocated.load(std::memory_order_relaxed); | ||
| event.total_freed = stats_->bytes_total_deallocated.load(std::memory_order_relaxed); | ||
| event.timestamp = std::chrono::steady_clock::now(); | ||
| event.event_bytes = signed_bytes; | ||
| event.nvtx_range = std::move(name); | ||
| event.nvtx_depth = depth; | ||
| event.alloc_range = std::move(alloc_range); | ||
| queue_->push(std::move(event)); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Do not terminate while recording an event.
These noexcept helpers copy strings, insert into unordered_map, and grow the queue vector. Any allocation failure calls std::terminate; this is especially dangerous from the noexcept deallocation paths after the upstream resource has already freed memory. Make recording best-effort in no-throw paths by containing logging failures and dropping only the failed record.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/include/raft/mr/recording_adaptor.hpp` around lines 52 - 92, Make
record_allocation, forget_allocation, and emit robust against allocation
failures: contain exceptions from string operations, unordered_map updates,
event construction, and queue insertion so these noexcept recording paths never
terminate. Treat each failed recording operation as best-effort by dropping only
that record, while preserving successful bookkeeping and event delivery.
| recording_adaptor(Upstream upstream, std::shared_ptr<allocation_event_queue> queue, int source_id) | ||
| : upstream_(std::move(upstream)), | ||
| stats_(std::make_shared<resource_stats>()), | ||
| queue_(std::move(queue)), | ||
| alloc_map_(std::make_shared<address_range_map>()), | ||
| source_id_(source_id) | ||
| { | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Reject a null event queue at construction.
queue_ is dereferenced on every allocation/deallocation, but this public constructor accepts a null shared_ptr; that creates an object which crashes later instead of failing at its boundary.
As per coding guidelines, “Add input validation for invalid dimensions, null pointers, and other obvious precondition failures where they can cause incorrect behavior.”
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/include/raft/mr/recording_adaptor.hpp` around lines 96 - 103, Update the
public recording_adaptor constructor to validate the queue argument before
storing or using it, rejecting a null shared_ptr immediately with the project’s
established precondition/error mechanism. Preserve normal construction for
non-null queues and ensure queue_ cannot be dereferenced by an invalid instance.
Source: Coding guidelines
| upstream_.deallocate_sync(ptr, bytes, alignment); | ||
| stats_->record_deallocate(static_cast<std::int64_t>(bytes)); | ||
| emit(forget_allocation(ptr), -static_cast<std::int64_t>(bytes)); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Remove the allocation mapping before releasing the pointer upstream.
A concurrent allocation can reuse ptr after upstream deallocation but before forget_allocation(ptr). The old deallocation then removes the new allocation’s NVTX path, mislabeling both later records.
cpp/include/raft/mr/recording_adaptor.hpp#L123-L125: callforget_allocation(ptr)beforeupstream_.deallocate_sync(...), then emit the captured path afterward.cpp/include/raft/mr/recording_adaptor.hpp#L145-L147: apply the same ordering to stream deallocation.
📍 Affects 1 file
cpp/include/raft/mr/recording_adaptor.hpp#L123-L125(this comment)cpp/include/raft/mr/recording_adaptor.hpp#L145-L147
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/include/raft/mr/recording_adaptor.hpp` around lines 123 - 125, In both
synchronous deallocation sites at
cpp/include/raft/mr/recording_adaptor.hpp:123-125 and stream deallocation sites
at cpp/include/raft/mr/recording_adaptor.hpp:145-147, capture and remove the
allocation mapping via forget_allocation(ptr) before calling the corresponding
upstream deallocation, then emit the captured path afterward while preserving
the existing byte accounting.
| void stop() | ||
| { | ||
| { | ||
| std::lock_guard<std::mutex> lock(mtx_); | ||
| stopped_ = true; | ||
| } | ||
| cv_.notify_all(); | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Prevent events from being accepted after shutdown begins.
A producer can call push() after stopped_ is set and after the worker drains its final batch, leaving events permanently queued and unrecorded. Coordinate producer teardown before stopping, or track active producers so shutdown only completes after all accepted events drain.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/include/raft/mr/recording_monitor.hpp` around lines 72 - 79, Update
RecordingMonitor::stop and the producer path in push so shutdown coordinates
with active producers: prevent new events from being accepted once shutdown
begins, wait for in-progress producers to finish enqueueing, and only complete
after the worker drains all accepted events. Preserve the existing mutex and
condition-variable synchronization while ensuring no event can remain queued
after stop returns.
| void start() | ||
| { | ||
| if (worker_.joinable()) { return; } | ||
| write_header(); | ||
| worker_ = std::thread([this] { run(); }); | ||
| } | ||
|
|
||
| void stop() | ||
| { | ||
| if (!worker_.joinable()) { return; } | ||
| queue_->stop(); // drains the queue and causes the worker to exit its loop | ||
| worker_.join(); | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Document the newly exposed public APIs.
The new public lifecycle and adaptor APIs lack Doxygen documentation.
cpp/include/raft/mr/recording_monitor.hpp#L118-L130: add Doxygen forstart()andstop(), including lifecycle and producer-shutdown requirements.cpp/include/raft/mr/recording_adaptor.hpp#L96-L103: document constructor ownership and non-null queue requirements.cpp/include/raft/mr/recording_adaptor.hpp#L111-L157: document allocation/deallocation recording semantics and accessor behavior.
As per path instructions, “Public APIs always require documentation” for newly added exported public functions/classes.
📍 Affects 2 files
cpp/include/raft/mr/recording_monitor.hpp#L118-L130(this comment)cpp/include/raft/mr/recording_adaptor.hpp#L96-L103cpp/include/raft/mr/recording_adaptor.hpp#L111-L157
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/include/raft/mr/recording_monitor.hpp` around lines 118 - 130, Document
all newly exposed public APIs with Doxygen: in
cpp/include/raft/mr/recording_monitor.hpp lines 118-130, document start() and
stop(), including lifecycle behavior and the requirement to shut down producers
before stopping; in cpp/include/raft/mr/recording_adaptor.hpp lines 96-103,
document the constructor’s queue ownership and non-null requirement; and in
lines 111-157, document allocation/deallocation recording semantics and accessor
behavior.
Source: Path instructions
| out_ << "timestamp_us"; | ||
| for (auto const& name : source_names_) { | ||
| out_ << ',' << name << "_current," << name << "_peak," << name << "_total_alloc," << name | ||
| << "_total_freed"; | ||
| } | ||
| out_ << ",nvtx_depth,nvtx_range,event_source,event_bytes,alloc_range\n"; |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Escape CSV string fields before writing them.
NVTX range names and allocation paths can contain ", which terminates the quoted CSV field and corrupts the row. Escape embedded quotes by doubling them for nvtx_range, alloc_range, and registered source names.
Also applies to: 175-181
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/include/raft/mr/recording_monitor.hpp` around lines 141 - 146, Update the
CSV output logic in the recording monitor, including the header generation
around source_names_ and the row-writing logic for nvtx_range and alloc_range,
to escape string fields according to CSV rules. Double every embedded quote
before writing registered source names, NVTX range names, and allocation paths,
while preserving the existing field delimiters and output format.
| for (auto const& v : view_) { | ||
| out_ << ',' << v.current << ',' << v.current << ',' << v.total_alloc << ',' << v.total_freed; | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Write an actual peak value.
The <source>_peak column always repeats v.current; after a source frees memory, rows report a lower value as its peak. Capture and retain a peak value per source, or remove the misleading column.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@cpp/include/raft/mr/recording_monitor.hpp` around lines 172 - 174, Update the
serialization loop in the recording monitor to output a true per-source peak
value for the <source>_peak column instead of repeating v.current. Track and
retain each source’s maximum observed allocation across frees, then write that
retained value while preserving the existing current, total_alloc, and
total_freed columns.
4 participants
Extend memory tracking resources with another queue-based approach (as opposed to the current sampling-rate notification approach).
Usage:
If the sampling interval is given, the notification approach is selected:
raft::memory_tracking_resources tracked(res, oss, 1ms);If it is NOT given, the queue-based approach is selected:
raft::memory_tracking_resources tracked(res, oss);Compare two approaches:
Unit test benchmark (H100, 64 threads, each thread 200x allocations, each allocation 256KiB):