Fix RF multi-stream fits on non-default device

[csadorf]

Ensures random forest fitting sets the RAFT handle's CUDA device in both the main thread and OpenMP workers, avoiding invalid multi-stream CUDA calls when training data is created under cp.cuda.Device(1). Adds regression coverage for the reported crash and removes the n_streams=1 workaround from the FIL device-selection test.

Closes #5983

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

[csadorf]

/ok to test 3c2352f

[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: d4e677f7-b305-4b50-bca0-5f5f245d9a22

📥 Commits

Reviewing files that changed from the base of the PR and between 3c2352f and 7969078.

📒 Files selected for processing (3)

• cpp/src/randomforest/randomforest.cuh
• python/cuml/tests/test_fil.py
• python/cuml/tests/test_random_forest.py

---

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes

  • Improved GPU device handling during Random Forest training to ensure correct device context when running on multi-GPU systems.

• Tests

  • Added multi-GPU regression tests verifying device selection and that the CUDA device context is restored after training.
  • Added test coverage using CUDA arrays on non-default devices to validate multi-device workflows.

Walkthrough

RandomForest::fit now sets the CUDA device from the provided raft::handle_t using raft::device_setter for initial pointer validation and installs per-thread device_setter guards inside the OpenMP tree-building loop. Python tests were updated: a prior n_streams=1 workaround was removed and a multi-device CuPy test was added.

Changes

Device Context Management in RandomForest Training

Layer / File(s)	Summary
Device context guard in RandomForest::fit cpp/src/randomforest/randomforest.cuh	Adds raft/core/device_setter.hpp; RandomForest::fit extracts the device from raft::handle_t, constructs a raft::device_setter for initial GPU pointer validation, and installs per-thread raft::device_setter(handle_device) inside the OpenMP parallel tree-building loop before selecting per-thread streams.
Test updates and multi-device regression python/cuml/tests/test_fil.py, python/cuml/tests/test_random_forest.py	Removes the n_streams=1 argument in test_device_selection, adds import cupy as cp, and introduces test_rf_fit_with_cupy_nondefault_device which fits RandomForest on CuPy arrays allocated on device 1 and verifies the original CUDA device context is restored.

---

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% 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
Title check	✅ Passed	The title clearly and concisely summarizes the main change: fixing random forest multi-stream fits when training data is created under a non-default CUDA device.
Description check	✅ Passed	The description is directly related to the changeset, explaining the fix for CUDA device handling in random forest fitting and mentioning test additions and workaround removal.
Linked Issues check	✅ Passed	The PR addresses all objectives from issue #5983: setting RAFT handle's CUDA device in both main and worker threads, adding regression test coverage, and removing the n_streams=1 workaround.
Out of Scope Changes check	✅ Passed	All changes are directly related to fixing the multi-stream CUDA device issue: the randomforest.cuh changes implement the fix, test_fil.py removes the workaround, and test_random_forest.py adds regression coverage.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[jcrist]

One concern on correctness.

Beyond that, there are a few other places in the C++ code where we use stream pools across threads. I assume those also need this change applied?

Also, in the original issue (and last time we discussed this) we decided we didn't necessarily want to support users changing the device at runtime this way (and instead wanted users to use CUDA_VISIBLE_DEVICES). If this change doesn't result in any downsides (perf, maintainability, ...) it's probably fine, but otherwise we may want to fix this with docs (and maybe a different way of preventing runtime errors) rather than supporting this.

[jcrist]

Is this value cached by the handle? If it is, then it won't get the correct device id on repeated calls if cp.cuda.Device is used to change the device.

[csadorf]

Good catch! The device id is usually cached because the handle itself is cached thread-local. However, since by default RandomForestClassifier set the n_streams value to nonzero, it happens to not be cached.

With n_streams > 0, the handle we will be created with a new CUDA stream pool and use the per thread default stream for each thread. The device id will be set on the first call to get_device() based on whatever is the currently active device. That is all to say that this will indeed not work well if we create arrays with one device active, and then run fit with a different device active, but it will generally work for switching devices, but only incidentally because n_streams > 0.

The changes on this PR still have merit since we should generally use the device set by the handle, but we probably have to apply this in a few other cases as well.

Given that – as you correctly assert - we currently don't generally support a workflow for setting devices via cp.cuda.Device() I think we should not merge this PR as-is and instead perform a broader evaluation of this failure class.

[csadorf]

We will document the expected behavior in #8134 , so I'll just retarget this PR to 26.08 and pick it back up then.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.