Skip to content

Memory logging resources#3051

Open
huuanhhuyn wants to merge 15 commits into
NVIDIA:mainfrom
huuanhhuyn:fix-alloc-mislabel
Open

Memory logging resources#3051
huuanhhuyn wants to merge 15 commits into
NVIDIA:mainfrom
huuanhhuyn:fix-alloc-mislabel

Conversation

@huuanhhuyn

@huuanhhuyn huuanhhuyn commented Jun 8, 2026

Copy link
Copy Markdown
Contributor

Add memory_logging_resources to attribute (de)allocations to the correct code block / nvtx range instead of estimating the attributions as with memory_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:

  • queue-based recording approach preserves all (de)allocation events in the queue and associate the full nvtx range to each event. Additionally, it labels a deallocation to the nvtx range where the corresponding allocation occurs. This should be used when labels are essential and certain overhead for debugging is accepted.
  • notification approach is less invasive with low overhead. This is used when label accuracy is not important and several dropped events are acceptable.

Unit test benchmark (H100, 64 threads, each thread 200x allocations, each allocation 256KiB):

  • recording approach: 34ms, all 25600 events recorded
  • sampling approach: 7ms, around 100 events recorded
image

@copy-pr-bot

copy-pr-bot Bot commented Jun 8, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@huuanhhuyn
huuanhhuyn force-pushed the fix-alloc-mislabel branch 3 times, most recently from ef5b83f to e852f30 Compare June 16, 2026 11:43
@huuanhhuyn huuanhhuyn changed the title [WIP] Reproduce allocation mislabelling issue [WIP] Extend memory tracking resources tool Jun 16, 2026
@huuanhhuyn
huuanhhuyn force-pushed the fix-alloc-mislabel branch from e852f30 to f4e4cee Compare June 30, 2026 09:16
@huuanhhuyn
huuanhhuyn requested review from a team as code owners June 30, 2026 09:16
@huuanhhuyn huuanhhuyn changed the title [WIP] Extend memory tracking resources tool [WIP] Extend memory tracking resources Jul 1, 2026
@huuanhhuyn
huuanhhuyn force-pushed the fix-alloc-mislabel branch 2 times, most recently from 32d3c6c to 455f959 Compare July 1, 2026 11:52
@huuanhhuyn
huuanhhuyn force-pushed the fix-alloc-mislabel branch from 455f959 to a315972 Compare July 1, 2026 12:43
@huuanhhuyn huuanhhuyn changed the title [WIP] Extend memory tracking resources Extend memory tracking resources Jul 1, 2026

@achirkin achirkin left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 tfeher added feature request New feature or request improvement Improvement / enhancement to an existing function and removed improvement Improvement / enhancement to an existing function labels Jul 2, 2026
Comment thread cpp/include/raft/mr/host_memory_resource.hpp
@jameslamb
jameslamb removed the request for review from a team July 10, 2026 16:14
@huuanhhuyn huuanhhuyn changed the title Extend memory tracking resources [WIP] Memory logging resources Jul 16, 2026
@coderabbitai

coderabbitai Bot commented Jul 16, 2026

Copy link
Copy Markdown

Review Change Stack

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review
📝 Walkthrough

Walkthrough

Adds 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.

Changes

Memory recording and monitoring

Layer / File(s) Summary
NVTX range path tracking
cpp/include/raft/core/detail/nvtx.hpp, cpp/include/raft/core/detail/nvtx_range_path_stack.hpp
Tracks nested NVTX range names with unique instance IDs, full paths, depth, and synchronized push/pop operations.
Allocation event recording pipeline
cpp/include/raft/mr/recording_monitor.hpp, cpp/include/raft/mr/recording_adaptor.hpp
Adds allocation events, a thread-safe queue, background CSV serialization, and sync/stream adaptors that record allocation deltas and NVTX metadata.
Resource logging integration
cpp/include/raft/core/memory_logging_resources.hpp
Adds a non-copyable resources wrapper that installs recording adaptors across host, device, pinned, managed, workspace, and large-workspace resources, then restores prior resources on destruction.
Recording and sampling coverage
cpp/tests/core/monitor_resources.cu
Adds single-thread and parallel allocation-recording tests and expands sampled memory-tracking coverage.

Estimated code review effort: 4 (Complex) | ~45 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 37.04% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Title check ✅ Passed The title is concise and clearly names the main change: adding memory logging resources.
Description check ✅ Passed The description is directly related to the change and explains the new logging approach and usage.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@huuanhhuyn
huuanhhuyn force-pushed the fix-alloc-mislabel branch from 3747994 to 9ac8dec Compare July 16, 2026 15:12

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 8

🧹 Nitpick comments (1)
cpp/tests/core/monitor_resources.cu (1)

147-183: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

SamplingParallelThreads correctness relies on scheduling/timing, risking CI flakiness.

The assertion 3 < num_lines && num_lines < max_num_rows only 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 before stop(), num_lines could 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.md referenced 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

📥 Commits

Reviewing files that changed from the base of the PR and between 07b5d39 and 3747994.

📒 Files selected for processing (9)
  • cpp/include/raft/core/detail/nvtx_range_stack.hpp
  • cpp/include/raft/core/memory_logging_resources.hpp
  • cpp/include/raft/mr/host_memory_resource.hpp
  • cpp/include/raft/mr/notifying_adaptor.hpp
  • cpp/include/raft/mr/recording_adaptor.hpp
  • cpp/include/raft/mr/recording_monitor.hpp
  • cpp/include/raft/mr/resource_monitor.hpp
  • cpp/tests/core/allocation_tracking.cpp
  • cpp/tests/core/monitor_resources.cu

Comment on lines +74 to +79
~memory_logging_resources() override
{
if (recorder_) recorder_->stop();
raft::mr::set_default_host_resource(old_host_);
rmm::mr::set_current_device_resource(old_device_);
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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

Comment thread cpp/include/raft/mr/recording_adaptor.hpp Outdated
Comment thread cpp/include/raft/mr/recording_adaptor.hpp
Comment on lines +123 to +125
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));

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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: call forget_allocation(ptr) before upstream_.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.

Comment thread cpp/include/raft/mr/recording_monitor.hpp
Comment thread cpp/include/raft/mr/recording_monitor.hpp
Comment on lines +141 to +146
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";

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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.

Comment thread cpp/include/raft/mr/recording_monitor.hpp Outdated

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

📥 Commits

Reviewing files that changed from the base of the PR and between 9ac8dec and 0f61edc.

📒 Files selected for processing (6)
  • cpp/include/raft/core/detail/nvtx.hpp
  • cpp/include/raft/core/detail/nvtx_range_path_stack.hpp
  • cpp/include/raft/core/memory_logging_resources.hpp
  • cpp/include/raft/mr/recording_adaptor.hpp
  • cpp/include/raft/mr/recording_monitor.hpp
  • cpp/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

Comment thread cpp/include/raft/core/detail/nvtx_range_path_stack.hpp
Comment thread cpp/include/raft/core/detail/nvtx.hpp Outdated
Comment on lines +154 to +155
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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 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/**' || true

Repository: 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.hpp

Repository: 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 || true

Repository: 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.hpp

Repository: 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.hpp

Repository: 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()
@huuanhhuyn
huuanhhuyn force-pushed the fix-alloc-mislabel branch from 0f61edc to 02d785c Compare July 17, 2026 11:49
logging resources is now splitted from memory tracking resources. Therefore, no need to be backward compatible with it in term of csv columns.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

♻️ Duplicate comments (1)
cpp/include/raft/mr/recording_adaptor.hpp (1)

98-105: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Reject a null event queue.

A null queue creates an adaptor that later dereferences queue_ 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

📥 Commits

Reviewing files that changed from the base of the PR and between 02d785c and 28368cb.

📒 Files selected for processing (3)
  • cpp/include/raft/mr/recording_adaptor.hpp
  • cpp/include/raft/mr/recording_monitor.hpp
  • cpp/tests/core/monitor_resources.cu
🚧 Files skipped from review as they are similar to previous changes (1)
  • cpp/tests/core/monitor_resources.cu

Comment on lines +83 to +94
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));

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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.

Comment on lines +120 to +121
upstream_.deallocate_sync(ptr, bytes, alignment);
emit(forget_allocation(ptr), -static_cast<std::int64_t>(bytes), timestamp);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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: call forget_allocation(ptr) before upstream_.deallocate_sync(), then emit the captured path.
  • cpp/include/raft/mr/recording_adaptor.hpp#L142-L143: apply the same ordering before upstream_.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.

Comment thread cpp/include/raft/mr/recording_monitor.hpp
@huuanhhuyn huuanhhuyn changed the title [WIP] Memory logging resources Memory logging resources Jul 17, 2026

@achirkin achirkin left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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();

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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};

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature request New feature or request

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

4 participants