Skip to content

Fix RF multi-stream fits on non-default device #8131

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
csadorf wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Fix RF multi-stream fits on non-default device #8131

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions cpp/src/randomforest/randomforest.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@

#include <cuml/ensemble/randomforest.hpp>

#include <raft/core/device_setter.hpp>
#include <raft/core/handle.hpp>
#include <raft/core/nvtx.hpp>
#include <raft/random/permute.cuh>
Expand Down Expand Up @@ -120,9 +121,11 @@ class RandomForest {
bool* bootstrap_masks = nullptr)
{
raft::common::nvtx::range fun_scope("RandomForest::fit @randomforest.cuh");
this->error_checking(input, labels, n_rows, n_cols, false);
const raft::handle_t& handle = user_handle;
int n_sampled_rows = 0;
int handle_device = handle.get_device();

@jcrist jcrist May 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

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 csadorf May 19, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

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 csadorf May 19, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

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

auto device_guard = raft::device_setter(handle_device);
this->error_checking(input, labels, n_rows, n_cols, false);
int n_sampled_rows = 0;
if (this->rf_params.bootstrap) {
n_sampled_rows = std::round(this->rf_params.max_samples * n_rows);
} else {
Expand Down Expand Up @@ -162,8 +165,9 @@ class RandomForest {

#pragma omp parallel for num_threads(n_streams)
for (int i = 0; i < this->rf_params.n_trees; i++) {
int stream_id = omp_get_thread_num();
auto s = handle.get_stream_from_stream_pool(stream_id);
auto thread_device_guard = raft::device_setter(handle_device);
int stream_id = omp_get_thread_num();
auto s = handle.get_stream_from_stream_pool(stream_id);

this->get_row_sample(i, n_rows, &selected_rows[stream_id], s);

Expand Down
3 changes: 0 additions & 3 deletions python/cuml/tests/test_fil.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -935,13 +935,10 @@ def test_device_selection(device_id, model_kind, tmp_path):
)

with device_context:
# TODO(hcho3): Remove n_streams=1 argument once the bug
# https://github.com/rapidsai/cuml/issues/5983 is resolved
cuml_model = cumlRandomForestClassifier(
max_depth=3,
random_state=0,
n_estimators=n_estimators,
n_streams=1,
)
cuml_model.fit(cp.array(X), cp.array(y))
fm = cuml_model.as_fil()
Expand Down
31 changes: 31 additions & 0 deletions python/cuml/tests/test_random_forest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
import warnings

import cudf
import cupy as cp
import numpy as np
import pytest
import treelite
Expand Down Expand Up @@ -1046,6 +1047,36 @@ def test_rf_predict_returns_int():
assert pred.dtype == np.int64


@pytest.mark.skipif(
cp.cuda.runtime.getDeviceCount() < 2,
reason="test requires at least two visible CUDA devices",
)
def test_rf_fit_with_cupy_nondefault_device():
# Verify that cuml operations respect the current device context and
# restore it upon completion. This test checks that cuML correctly saves
# and restores the CUDA device ID, ensuring it does not inadvertently
# modify the device set by external code.

# IMPORTANT: This test does **not** document or endorse ``cp.cuda.Device``
# (or ``cp.cuda.runtime.setDevice``) as a supported API for switching
# devices in cuML. Switching devices via CuPy is generally **not
# supported**.
current_device = cp.cuda.runtime.getDevice()

with cp.cuda.Device(1):
X = cp.random.normal(size=(10, 4)).astype(cp.float32)
y = cp.asarray([0, 1] * 5, dtype=cp.int32)

clf = cuml.ensemble.RandomForestClassifier(
max_features=1.0,
n_bins=8,
n_estimators=2,
)
clf.fit(X, y)

assert cp.cuda.runtime.getDevice() == current_device


def test_ensemble_estimator_length():
X, y = make_classification()
clf = cuml.ensemble.RandomForestClassifier(n_estimators=3)
Expand Down
Loading