Restore tracked device resources on their original GPU#3077
Restore tracked device resources on their original GPU#3077fallintoplace wants to merge 4 commits into
Conversation
📝 WalkthroughSummary by CodeRabbit
Walkthrough
ChangesPer-device resource restoration
Estimated code review effort: 3 (Moderate) | ~25 minutes 🚥 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? |
2 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