Memory logging resources#3051
Conversation
ef5b83f to
e852f30
Compare
e852f30 to
f4e4cee
Compare
32d3c6c to
455f959
Compare
455f959 to
a315972
Compare
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.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds NVTX full-path tracking, asynchronous CSV allocation recording, memory-resource adaptors, global resource integration, and single-thread and parallel tests for recording and sampling behavior. ChangesMemory recording and monitoring
Estimated code review effort: 4 (Complex) | ~45 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
3747994 to
9ac8dec
Compare
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
| 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.
| 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.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/detail/nvtx_range_path_stack.hpp`:
- Around line 47-62: Make the noexcept recording paths emit() and
record_allocation() resilient to allocation failures from
inner_range_and_depth() and thread_local_nvtx_full_path(). Catch std::bad_alloc
at the recording boundary and continue without the optional NVTX labels,
ensuring these paths never propagate allocation exceptions or terminate the
process.
In `@cpp/include/raft/core/detail/nvtx.hpp`:
- Around line 154-155: Update the shadow-stack updates surrounding
nvtxDomainRangePushEx so the NVTX range depth and both stacks remain consistent
if either push throws. Prefer completing the potentially throwing pushes before
committing the NVTX push, or catch failures and roll back any completed stack
update while restoring the prior NVTX state; ensure no partial update remains.
🪄 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: 415d7566-f070-4ec6-96e3-a272c05b07d0
📒 Files selected for processing (6)
cpp/include/raft/core/detail/nvtx.hppcpp/include/raft/core/detail/nvtx_range_path_stack.hppcpp/include/raft/core/memory_logging_resources.hppcpp/include/raft/mr/recording_adaptor.hppcpp/include/raft/mr/recording_monitor.hppcpp/tests/core/monitor_resources.cu
🚧 Files skipped from review as they are similar to previous changes (4)
- cpp/include/raft/mr/recording_adaptor.hpp
- cpp/include/raft/mr/recording_monitor.hpp
- cpp/include/raft/core/memory_logging_resources.hpp
- cpp/tests/core/monitor_resources.cu
| detail::range_name_stack_instance.push(name); // sampling tracker (name + depth, cross-thread) | ||
| detail::full_range_stack_instance.push(name); // recording tracker (full path, thread-local) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== file outline =="
ast-grep outline cpp/include/raft/core/detail/nvtx.hpp --view expanded || true
echo
echo "== relevant lines =="
nl -ba cpp/include/raft/core/detail/nvtx.hpp | sed -n '1,260p'
echo
echo "== search stack types/usages =="
rg -n "range_name_stack_instance|full_range_stack_instance|push\\(|pop\\(|class .*stack|struct .*stack|thread_local" cpp/include/raft/core/detail/nvtx.hpp cpp -g '!**/build/**' || trueRepository: NVIDIA/raft
Length of output: 925
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== nvtx.hpp excerpt =="
sed -n '130,230p' cpp/include/raft/core/detail/nvtx.hpp
echo
echo "== surrounding helper definitions =="
rg -n "struct .*stack|class .*stack|thread_local|push_range_name|push_range\\(|pop_range\\(" cpp/include/raft/core/detail -g '!**/build/**' || true
echo
echo "== full file if small enough =="
wc -l cpp/include/raft/core/detail/nvtx.hppRepository: NVIDIA/raft
Length of output: 4168
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate any stack implementation used by nvtx.hpp =="
rg -n "push_range_name|full_range_stack_instance|range_name_stack_instance|thread_local" cpp/include cpp/src -g '!**/build/**' || true
echo
echo "== inspect any matching helper headers =="
fd -a "nvtx" cpp/include cpp/src 2>/dev/null || trueRepository: NVIDIA/raft
Length of output: 3256
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== nvtx_range_stack.hpp =="
sed -n '1,220p' cpp/include/raft/core/detail/nvtx_range_stack.hpp
echo
echo "== nvtx_range_path_stack.hpp =="
sed -n '1,240p' cpp/include/raft/core/detail/nvtx_range_path_stack.hppRepository: NVIDIA/raft
Length of output: 5166
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== grep for exception specs and allocations in NVTX helpers =="
rg -n "noexcept|std::string|std::vector|push\\(|pop\\(|reserve|make_shared|allocate|new " cpp/include/raft/core/detail/nvtx_range_stack.hpp cpp/include/raft/core/detail/nvtx_range_path_stack.hpp cpp/include/raft/mr/recording_adaptor.hpp cpp/include/raft/mr/resource_monitor.hppRepository: NVIDIA/raft
Length of output: 6337
Make NVTX push and shadow-stack updates atomic nvtxDomainRangePushEx() happens before both shadow push() calls, and either push() can allocate and throw. If one fails, the NVTX depth and shadow stacks diverge, so later pops and allocation labels can use stale state. Roll back completed updates on failure, or make the shadow updates non-throwing before committing the NVTX push.
🤖 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/detail/nvtx.hpp` around lines 154 - 155, Update the
shadow-stack updates surrounding nvtxDomainRangePushEx so the NVTX range depth
and both stacks remain consistent if either push throws. Prefer completing the
potentially throwing pushes before committing the NVTX push, or catch failures
and roll back any completed stack update while restoring the prior NVTX state;
ensure no partial update remains.
Sources: Coding guidelines, Path instructions
Pros: clarity, mutex free cons: similar code duplication, double hooks from nvtx::push_range_name()
0f61edc to
02d785c
Compare
logging resources is now splitted from memory tracking resources. Therefore, no need to be backward compatible with it in term of csv columns.
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
cpp/include/raft/mr/recording_adaptor.hpp (1)
98-105: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winReject a null event queue.
A null
queuecreates an adaptor that later dereferencesqueue_at Line 94. Validate it in the constructor with RAFT’s established precondition mechanism.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 98 - 105, Update recording_adaptor’s constructor to validate the queue parameter with RAFT’s established precondition mechanism before storing or using queue_. Reject null shared pointers immediately, while preserving the existing initialization of upstream_, alloc_map_, source_id_, and current_bytes_.Source: Coding guidelines
🤖 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/mr/recording_adaptor.hpp`:
- Around line 120-121: In both deallocation paths of recording_adaptor.hpp
(120-121 and 142-143), capture the result of forget_allocation(ptr) before
calling the corresponding upstream_.deallocate_sync() or upstream_.deallocate(),
then emit the captured path after deallocation using the existing byte and
timestamp values.
- Around line 83-94: Serialize event admission in the recording adaptor by
adding or reusing a mutex shared across adaptor copies, then lock it around both
the current_bytes_ fetch_add in the event construction and queue_->push. Ensure
every copy of the adaptor shares this mutex so concurrent producers preserve
atomic-update and enqueue order, while leaving the event contents and queue
behavior unchanged.
In `@cpp/include/raft/mr/recording_monitor.hpp`:
- Around line 135-140: Update CSV generation in the recording monitor, including
the header emitted near the source loop and row output in register_source/NVTX
handling, to escape every string field using CSV quoting: wrap fields containing
special characters and double any embedded quotes. Apply the same helper or
logic consistently to source names and NVTX labels, preserving numeric fields
unchanged.
---
Duplicate comments:
In `@cpp/include/raft/mr/recording_adaptor.hpp`:
- Around line 98-105: Update recording_adaptor’s constructor to validate the
queue parameter with RAFT’s established precondition mechanism before storing or
using queue_. Reject null shared pointers immediately, while preserving the
existing initialization of upstream_, alloc_map_, source_id_, and
current_bytes_.
🪄 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: a7f2d990-478e-46b5-a293-d6cc3cedd69e
📒 Files selected for processing (3)
cpp/include/raft/mr/recording_adaptor.hppcpp/include/raft/mr/recording_monitor.hppcpp/tests/core/monitor_resources.cu
🚧 Files skipped from review as they are similar to previous changes (1)
- cpp/tests/core/monitor_resources.cu
| auto [name, depth] = raft::common::nvtx::thread_local_inner_range_and_depth(); | ||
| allocation_event event{ | ||
| .timestamp = timestamp, | ||
| .source_id = source_id_, | ||
| .current_bytes = | ||
| current_bytes_->fetch_add(signed_bytes, std::memory_order_relaxed) + signed_bytes, | ||
| .delta_bytes = signed_bytes, | ||
| .nvtx_depth = depth, | ||
| .nvtx_inner_range = std::move(name), | ||
| .nvtx_full_range = std::move(nvtx_full_range), | ||
| }; | ||
| queue_->push(std::move(event)); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Serialize byte-state updates with queue admission.
Line 88 updates current_bytes_ before acquiring the queue lock. Concurrent producers can enqueue snapshots out of atomic-update order, so recording_monitor::source_current_ may end with a stale value. Share an emission mutex across adaptor copies and hold it across fetch_add and queue_->push.
Proposed synchronization
+ std::shared_ptr<std::mutex> emission_mtx_;
std::shared_ptr<std::atomic<std::int64_t>> current_bytes_;
void emit(...) noexcept
{
auto [name, depth] = raft::common::nvtx::thread_local_inner_range_and_depth();
+ std::lock_guard<std::mutex> lock(*emission_mtx_);
allocation_event event{
.current_bytes =
current_bytes_->fetch_add(signed_bytes, std::memory_order_relaxed) + signed_bytes,🤖 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 83 - 94, Serialize
event admission in the recording adaptor by adding or reusing a mutex shared
across adaptor copies, then lock it around both the current_bytes_ fetch_add in
the event construction and queue_->push. Ensure every copy of the adaptor shares
this mutex so concurrent producers preserve atomic-update and enqueue order,
while leaving the event contents and queue behavior unchanged.
| upstream_.deallocate_sync(ptr, bytes, alignment); | ||
| emit(forget_allocation(ptr), -static_cast<std::int64_t>(bytes), timestamp); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Remove the allocation mapping before upstream deallocation.
Freeing upstream first permits another thread to reuse ptr and install its new NVTX path before the old deallocation erases that mapping.
cpp/include/raft/mr/recording_adaptor.hpp#L120-L121: callforget_allocation(ptr)beforeupstream_.deallocate_sync(), then emit the captured path.cpp/include/raft/mr/recording_adaptor.hpp#L142-L143: apply the same ordering beforeupstream_.deallocate().
📍 Affects 1 file
cpp/include/raft/mr/recording_adaptor.hpp#L120-L121(this comment)cpp/include/raft/mr/recording_adaptor.hpp#L142-L143
🤖 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 120 - 121, In both
deallocation paths of recording_adaptor.hpp (120-121 and 142-143), capture the
result of forget_allocation(ptr) before calling the corresponding
upstream_.deallocate_sync() or upstream_.deallocate(), then emit the captured
path after deallocation using the existing byte and timestamp values.
achirkin
left a comment
There was a problem hiding this comment.
Thank you for the updates!
Let's focus on performance of the new/changed NVTX code for now (independently of the memory tracking/logging resources). If you think we absolutely have to present the whole range sequence stack on every record, we may consider removing the old range_name_stack_instance to reduce the overheads on each NVTX push/pop. Let's do some benchmarks to make sure we don't slow down the downstream libraries with our new additions.
| inline void pop_range() | ||
| { | ||
| detail::range_name_stack_instance.pop(); | ||
| detail::full_range_stack_instance.pop(); |
There was a problem hiding this comment.
Please note, the push/pop functions here are called all the time, independently of whether the memory tracking/logging is enabled. We must ensure they are fast, as they may appear in performance-critical code sections (a user code may create a range in a tight loop that does no allocations at all).
Let's add a new bench module specifically for this and post benchmark results in the discussion thread (before/after the change).
| * same block of code. User can aggregate them into a distribution of allocations | ||
| * over multiple runs. | ||
| */ | ||
| RAFT_EXPORT inline std::atomic<std::uint64_t> range_instance_counter{0}; |
There was a problem hiding this comment.
On the performance note: let's test what is faster: having a single global atomic vs using a thread-local counter + thread id reported in the csv file.
Add
memory_logging_resourcesto attribute (de)allocations to the correct code block / nvtx range instead of estimating the attributions as withmemory_tracking_resources. In addition, all (de)allocation events are garanteed to be captured - no dropping allowed.Usage:
raft::memory_logging_resources tracked(res, oss);raft::memory_tracking_resources tracked(res, oss, 1ms);Compare two approaches:
Unit test benchmark (H100, 64 threads, each thread 200x allocations, each allocation 256KiB):