Restore tracked device resources on their original GPU - #3077
Restore tracked device resources on their original GPU#3077fallintoplace wants to merge 5 commits into
Conversation
|
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:
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 (5)
📝 WalkthroughSummary by CodeRabbit
Walkthrough
ChangesPer-device resource restoration
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
achirkin
left a comment
achirkin
left a comment
There was a problem hiding this comment.
Thank you for uncovering this! The memory tracking/statistics APIs are rather new and haven't been properly tested in multi-gpu environment yet.
In the context of raft::resources handle we already have access to a "device id" resource (core/resource/device_id.hpp) stored per resources handle. Therefore, there's no need for an extra member variable duplicating this information (see below). Please adjust both files to use this resource instead.
| { | ||
| mr::set_default_host_resource(old_host_); | ||
| rmm::mr::set_current_device_resource(std::move(old_device_)); | ||
| rmm::mr::set_per_device_resource(rmm::cuda_device_id{device_id_}, std::move(old_device_)); |
There was a problem hiding this comment.
core/resource/device_id.hpp is the source of truth for the GPU device associated with the handle:
| rmm::mr::set_per_device_resource(rmm::cuda_device_id{device_id_}, std::move(old_device_)); | |
| rmm::mr::set_per_device_resource(resource::get_device_id(*this), std::move(old_device_)); |
| : resources(existing), | ||
| device_id_(device_setter::get_current_device()), | ||
| old_host_(mr::get_default_host_resource()), | ||
| old_device_(rmm::mr::get_current_device_resource_ref()) |
There was a problem hiding this comment.
Calling resource::get_device_id() on the current handle ensures the device id resource is instantiated and won't mutate until the resource is destroyed.
| old_device_(rmm::mr::get_current_device_resource_ref()) | |
| old_device_(rmm::mr::get_per_device_resource_ref(resource::get_device_id(existing))) |
|
/ok to test b0ddd70 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/include/raft/core/memory_tracking_resources.hpp (1)
108-114: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winInstall the tracking MR on the handle’s device id
init()still usesrmm::mr::set_current_device_resource(*device_adaptor_), so a copiedresourceshandle can install the adaptor on the wrong GPU while capture/restore is keyed offresource::get_device_id(*this). Useset_per_device_resource(rmm::cuda_device_id{resource::get_device_id(*this)}, *device_adaptor_)here too.🤖 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_tracking_resources.hpp` around lines 108 - 114, The destructor in memory_tracking_resources currently restores the per-device resource using the tracked device id, but the corresponding installation path in init() still uses the current device instead of the handle’s device id. Update the resource installation logic in memory_tracking_resources::init() to use rmm::mr::set_per_device_resource with rmm::cuda_device_id{resource::get_device_id(*this)} and the device adaptor, matching the handle’s device binding used elsewhere.Sources: Coding guidelines, 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_stats_resources.hpp`:
- Around line 81-91: The device-specific resource management in
memory_stats_resources is using resource::get_device_id(*this) for restore while
set_current_device_resource() installs on the active CUDA device, so they can
target different GPUs. Update memory_stats_resources so the same device id is
used consistently for both installation and restoration, and keep the
save/restore logic aligned with the device captured in the constructor and
destructor.
---
Outside diff comments:
In `@cpp/include/raft/core/memory_tracking_resources.hpp`:
- Around line 108-114: The destructor in memory_tracking_resources currently
restores the per-device resource using the tracked device id, but the
corresponding installation path in init() still uses the current device instead
of the handle’s device id. Update the resource installation logic in
memory_tracking_resources::init() to use rmm::mr::set_per_device_resource with
rmm::cuda_device_id{resource::get_device_id(*this)} and the device adaptor,
matching the handle’s device binding used elsewhere.
🪄 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: 113e8e7c-4498-49d5-bd87-7c71dd889a4c
📒 Files selected for processing (2)
cpp/include/raft/core/memory_stats_resources.hppcpp/include/raft/core/memory_tracking_resources.hpp
There was a problem hiding this comment.
🧹 Nitpick comments (1)
cpp/tests/core/memory_stats_resources.cpp (1)
37-48: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated test helper code across two files.
current_device_uses_pool_resource()/current_device_uses_default_cuda_resource()and thedevice_resource_restore_guardstruct (lines 27-48) are duplicated verbatim incpp/tests/core/monitor_resources.cu. Consider extracting these into a shared test utility header (e.g.,cpp/tests/core/device_resource_test_utils.hpp) to avoid future divergence between the two test files.🤖 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/memory_stats_resources.cpp` around lines 37 - 48, The helper functions current_device_uses_pool_resource(), current_device_uses_default_cuda_resource(), and the device_resource_restore_guard in this test file are duplicated in monitor_resources.cu, so extract them into a shared test utility header (for example, device_resource_test_utils.hpp) and update both test files to include and use the shared helpers instead of maintaining separate copies. Keep the existing behavior intact by moving the duplicated device resource checks and guard logic into the common utility.
🤖 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.
Nitpick comments:
In `@cpp/tests/core/memory_stats_resources.cpp`:
- Around line 37-48: The helper functions current_device_uses_pool_resource(),
current_device_uses_default_cuda_resource(), and the
device_resource_restore_guard in this test file are duplicated in
monitor_resources.cu, so extract them into a shared test utility header (for
example, device_resource_test_utils.hpp) and update both test files to include
and use the shared helpers instead of maintaining separate copies. Keep the
existing behavior intact by moving the duplicated device resource checks and
guard logic into the common utility.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0f5967f3-e6ac-4fd2-bde8-da3e92e49405
📒 Files selected for processing (4)
cpp/include/raft/core/memory_stats_resources.hppcpp/include/raft/core/memory_tracking_resources.hppcpp/tests/core/memory_stats_resources.cppcpp/tests/core/monitor_resources.cu
🚧 Files skipped from review as they are similar to previous changes (2)
- cpp/include/raft/core/memory_tracking_resources.hpp
- cpp/include/raft/core/memory_stats_resources.hpp
|
/ok to test |
@achirkin, there was an error processing your request: See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/1/ |
|
If something comes up, let me know. |
|
/ok to test 381f2d1 |
achirkin
left a comment
achirkin
left a comment
There was a problem hiding this comment.
Thanks for the updates! LGTM
|
/ok to test 3242e55 |
|
@fallintoplace could you please apply the style fixes? |
9c78a19 to
7c41eaa
Compare
|
@achirkin Sorry, can you please check again? |
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
cpp/tests/core/memory_stats_resources.cpp (1)
27-47: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated test helpers across two translation units.
device_resource_restore_guard,current_device_uses_pool_resource, andcurrent_device_uses_default_cuda_resourceare defined identically in both files.
cpp/tests/core/memory_stats_resources.cpp#L27-L47: move these into a shared test-utility header included by both files.cpp/tests/core/monitor_resources.cu#L28-L48: remove the duplicate definitions and include the shared header instead.🤖 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/memory_stats_resources.cpp` around lines 27 - 47, Move device_resource_restore_guard, current_device_uses_pool_resource, and current_device_uses_default_cuda_resource from cpp/tests/core/memory_stats_resources.cpp:27-47 into a shared test-utility header, then include that header there. Remove the identical definitions from cpp/tests/core/monitor_resources.cu:28-48 and include the shared header instead, preserving the existing behavior and symbols.
🤖 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_stats_resources.hpp`:
- Around line 87-92: Guard the rmm::mr::set_per_device_resource restore call in
both memory_stats_resources::~memory_stats_resources and
memory_tracking_resources::~memory_tracking_resources using RAFT’s no-throw
handling, ensuring exceptions cannot escape either destructor during teardown.
Apply the change in cpp/include/raft/core/memory_stats_resources.hpp lines 87-92
and cpp/include/raft/core/memory_tracking_resources.hpp lines 108-114.
In `@cpp/tests/core/memory_stats_resources.cpp`:
- Around line 176-183: Update the device0 setup around the device0_guard lambda
so the pool_memory_resource and raft::mr::device_resource objects outlive
device0_guard and remain valid while the test uses the restored resource. Store
these resource objects in appropriately scoped variables before calling
set_current_device_resource, matching the lifetime fix used for the analogous
setup near the earlier test block.
In `@cpp/tests/core/monitor_resources.cu`:
- Around line 136-143: Fix the dangling temporary in the device0_guard
initializer by ensuring the pool_memory_resource and raft::mr::device_resource
objects outlive the device_resource_restore_guard, matching the safe lifetime
pattern used around the earlier guard setup. Update the surrounding setup rather
than returning guards that reference temporaries.
---
Nitpick comments:
In `@cpp/tests/core/memory_stats_resources.cpp`:
- Around line 27-47: Move device_resource_restore_guard,
current_device_uses_pool_resource, and current_device_uses_default_cuda_resource
from cpp/tests/core/memory_stats_resources.cpp:27-47 into a shared test-utility
header, then include that header there. Remove the identical definitions from
cpp/tests/core/monitor_resources.cu:28-48 and include the shared header instead,
preserving the existing behavior and symbols.
🪄 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: 8cc561f3-0568-4ed5-8878-7059177b802a
📒 Files selected for processing (4)
cpp/include/raft/core/memory_stats_resources.hppcpp/include/raft/core/memory_tracking_resources.hppcpp/tests/core/memory_stats_resources.cppcpp/tests/core/monitor_resources.cu
|
/ok to test 38ca3a2 |
3 participants
Summary
memory_stats_resourcesandmemory_tracking_resourcesdevice memory resources to the CUDA device they wrapped at construction timermm::mr::set_per_device_resource(...)using the captured device idWhy
Both wrappers used
rmm::mr::set_current_device_resource(...)during teardown. In multi-GPU flows that writes the saved resource into whichever CUDA device is current at destruction time, which can leave the construction device pointing at the tracking adaptor and install the wrong resource on another GPU.Impact
This keeps per-device RMM state stable across wrapper teardown and closes a nasty multi-GPU correctness hole that could otherwise lead to wrong-device allocators being reused.
Validation
git diff --checkPARALLEL_LEVEL=8 ./build.sh tests -n --limit-tests=CORE_TESTnvcc/CUDAToolkit_ROOT, so the new tests were not run here