Restore tracked device resources on their original GPU

[fallintoplace]

Summary

• restore memory_stats_resources and memory_tracking_resources device memory resources to the CUDA device they wrapped at construction time
• switch the destructor restore path to rmm::mr::set_per_device_resource(...) using the captured device id
• add guarded 2-GPU regressions that destroy each wrapper while a different CUDA device is current

Why

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 --check
• attempted PARALLEL_LEVEL=8 ./build.sh tests -n --limit-tests=CORE_TEST
• local build is currently blocked by missing nvcc / CUDAToolkit_ROOT, so the new tests were not run here

[copy-pr-bot]

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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 229a428d-d6bd-4907-a813-4e828edd5944

📥 Commits

Reviewing files that changed from the base of the PR and between c515501 and b0ddd70.

📒 Files selected for processing (4)

• cpp/include/raft/core/memory_stats_resources.hpp
• cpp/include/raft/core/memory_tracking_resources.hpp
• cpp/tests/core/memory_stats_resources.cpp
• cpp/tests/core/monitor_resources.cu

---

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes
  • Improved GPU device-specific memory resource handling so memory state is restored correctly after use.
  • Added multi-GPU coverage to verify memory resources remain isolated per device after switching contexts.

Walkthrough

Both memory_stats_resources and memory_tracking_resources now capture the active CUDA device ID at construction and restore the previous memory resource for that specific device via rmm::mr::set_per_device_resource on destruction, replacing the prior set_current_device_resource approach. New multi-GPU tests validate correct per-device restoration.

Changes

Per-device resource restoration

Layer / File(s)	Summary
memory_stats_resources per-device restore cpp/include/raft/core/memory_stats_resources.hpp	Includes device_setter.hpp, stores device_id_ at construction from device_setter::get_current_device(), and destructor now calls rmm::mr::set_per_device_resource with the saved device ID instead of set_current_device_resource.
memory_tracking_resources per-device restore cpp/include/raft/core/memory_tracking_resources.hpp	Same pattern applied: includes device_setter.hpp, adds device_id_ member set at construction, and destructor restores the resource via set_per_device_resource using the saved device ID.
Multi-device restoration tests cpp/tests/core/memory_stats_resources.cpp, cpp/tests/core/monitor_resources.cu	Adds a device_resource_restore_guard RAII helper and new GoogleTest cases (RestoresDeviceResourceOnConstructionDevice) that skip on <2 GPUs, construct/destroy the respective resource classes under a device scope, and assert correct per-device memory resource types after destruction.

Estimated code review effort: 3 (Moderate) | ~25 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main multi-GPU resource restoration change.
Description check	✅ Passed	The description matches the code changes and explains the motivation, impact, and validation.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[achirkin]

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.

[achirkin]

core/resource/device_id.hpp is the source of truth for the GPU device associated with the handle:

Suggested change

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

[achirkin]

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.

Suggested change

	old_device_(rmm::mr::get_current_device_resource_ref())
	old_device_(rmm::mr::get_per_device_resource_ref(resource::get_device_id(existing)))

[achirkin]

/ok to test b0ddd70