diff --git a/cpp/src/randomforest/randomforest.cuh b/cpp/src/randomforest/randomforest.cuh index 356bddef4f..ba152dea27 100644 --- a/cpp/src/randomforest/randomforest.cuh +++ b/cpp/src/randomforest/randomforest.cuh @@ -7,6 +7,7 @@ #include +#include #include #include #include @@ -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(); + 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 { @@ -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); diff --git a/python/cuml/tests/test_fil.py b/python/cuml/tests/test_fil.py index 0997074f48..429fde355b 100644 --- a/python/cuml/tests/test_fil.py +++ b/python/cuml/tests/test_fil.py @@ -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() diff --git a/python/cuml/tests/test_random_forest.py b/python/cuml/tests/test_random_forest.py index 73e15403c0..1501ce3c98 100644 --- a/python/cuml/tests/test_random_forest.py +++ b/python/cuml/tests/test_random_forest.py @@ -9,6 +9,7 @@ import warnings import cudf +import cupy as cp import numpy as np import pytest import treelite @@ -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)