From 87ed689940cd7ea9254b4c245e36384f7ecc92b6 Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Wed, 21 Oct 2026 22:00:55 -0700 Subject: [PATCH 1/3] Support pickling for nvForest objects --- python/nvforest/nvforest/_factory.py | 2 +- python/nvforest/nvforest/_forest_inference.py | 44 ------- .../nvforest/nvforest/detail/device_infer.py | 49 ++++++++ .../nvforest/detail/forest_inference.pyx | 70 +++++++++++ python/nvforest/tests/test_pickle.py | 118 ++++++++++++++++++ 5 files changed, 238 insertions(+), 45 deletions(-) create mode 100644 python/nvforest/nvforest/detail/device_infer.py create mode 100644 python/nvforest/tests/test_pickle.py diff --git a/python/nvforest/nvforest/_factory.py b/python/nvforest/nvforest/_factory.py index 46ae1ee..beb7965 100644 --- a/python/nvforest/nvforest/_factory.py +++ b/python/nvforest/nvforest/_factory.py @@ -14,11 +14,11 @@ CPUForestInferenceRegressor, GPUForestInferenceClassifier, GPUForestInferenceRegressor, - infer_device, infer_is_classifier, ) from nvforest._handle import Handle from nvforest._typing import StreamLike +from nvforest.detail.device_infer import infer_device from nvforest.detail.raft_stream import _handle_deprecated_handle_arg diff --git a/python/nvforest/nvforest/_forest_inference.py b/python/nvforest/nvforest/_forest_inference.py index 8885275..017cef8 100644 --- a/python/nvforest/nvforest/_forest_inference.py +++ b/python/nvforest/nvforest/_forest_inference.py @@ -18,7 +18,6 @@ import numpy as np import treelite -from cuda.bindings import runtime from nvforest._base import ForestInferenceClassifier, ForestInferenceRegressor from nvforest._typing import DataType, StreamLike @@ -42,49 +41,6 @@ def infer_is_classifier(treelite_model: treelite.Model) -> bool: ) -def detect_current_device(require: bool) -> Optional[int]: - """ - Query the currently active GPU. - - Parameters - ---------- - require: - Whether to raise an exception when no GPU is available. - - Returns - ------- - int or None - ID of the currently active GPU device, or None if no GPU is available. - """ - status, current_device_id = runtime.cudaGetDevice() - if status != runtime.cudaError_t.cudaSuccess: - if not require: - return None - _, name = runtime.cudaGetErrorName(status) - _, msg = runtime.cudaGetErrorString(status) - name, msg = name.decode("utf-8"), msg.decode("utf-8") - raise RuntimeError( - f"Failed to detect a GPU device. Diagnostic:\n {name}: {msg}" - ) - return current_device_id - - -def infer_device( - device: str, - device_id: Optional[int], -) -> tuple[str, int]: - if device == "auto": - # Auto mode: Use GPU if available; use CPU otherwise. - device_id = detect_current_device(require=False) - device = "cpu" if device_id is None else "gpu" - elif device == "gpu" and device_id is None: - # If no device ID is explicitly given, use the currently - # active device - device_id = detect_current_device(require=True) - assert device in ("gpu", "cpu") - return device, device_id - - class ClassifierMixin: def get_class_assignment( self, diff --git a/python/nvforest/nvforest/detail/device_infer.py b/python/nvforest/nvforest/detail/device_infer.py new file mode 100644 index 0000000..ea72855 --- /dev/null +++ b/python/nvforest/nvforest/detail/device_infer.py @@ -0,0 +1,49 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Optional + +from cuda.bindings import runtime + + +def detect_current_device(require: bool) -> Optional[int]: + """ + Query the currently active GPU. + + Parameters + ---------- + require: + Whether to raise an exception when no GPU is available. + + Returns + ------- + int or None + ID of the currently active GPU device, or None if no GPU is available. + """ + status, current_device_id = runtime.cudaGetDevice() + if status != runtime.cudaError_t.cudaSuccess: + if not require: + return None + _, name = runtime.cudaGetErrorName(status) + _, msg = runtime.cudaGetErrorString(status) + name, msg = name.decode("utf-8"), msg.decode("utf-8") + raise RuntimeError( + f"Failed to detect a GPU device. Diagnostic:\n {name}: {msg}" + ) + return current_device_id + + +def infer_device( + device: str, + device_id: Optional[int], +) -> tuple[str, int]: + if device == "auto": + # Auto mode: Use GPU if available; use CPU otherwise. + device_id = detect_current_device(require=False) + device = "cpu" if device_id is None else "gpu" + elif device == "gpu" and device_id is None: + # If no device ID is explicitly given, use the currently + # active device + device_id = detect_current_device(require=True) + assert device in ("gpu", "cpu") + return device, device_id diff --git a/python/nvforest/nvforest/detail/forest_inference.pyx b/python/nvforest/nvforest/detail/forest_inference.pyx index 295722e..4d2b291 100644 --- a/python/nvforest/nvforest/detail/forest_inference.pyx +++ b/python/nvforest/nvforest/detail/forest_inference.pyx @@ -7,9 +7,11 @@ from typing import Optional, Union import numpy as np import treelite +from cuda.bindings import runtime from cuda.core import Device, Stream from nvforest._typing import DataType, StreamLike +from nvforest.detail.device_infer import infer_device from nvforest.detail.treelite import safe_treelite_call from libc.stdint cimport uint32_t, uintptr_t @@ -273,6 +275,30 @@ cdef class ForestInference_impl(): return preds +def _infer_device_for_deserialization( + device: str, + device_id: Optional[int], +) -> tuple[str, int]: + """Select an available device when restoring a serialized model.""" + if device == "gpu": + status, device_count = runtime.cudaGetDeviceCount() + if status == runtime.cudaError_t.cudaSuccess and device_count > 0: + device, device_id = infer_device(device, device_id) + if not (0 <= device_id < device_count): + device_id = 0 + return device, device_id + + # The serialized model used a GPU, but none is visible now. Let the + # normal auto-selection logic choose the available device. + device = "auto" + device_id = None + + device, device_id = infer_device(device, device_id) + if device == "cpu": + device_id = -1 + return device, device_id + + class ForestInferenceImpl: def __init__( self, @@ -346,6 +372,50 @@ class ForestInferenceImpl: device_id=self.device_id ) + def __getstate__(self): + """Return the state needed to rebuild the native inference model.""" + return { + "layout": self._layout, + "align_bytes": self.align_bytes, + "use_double_precision": self._use_double_precision, + "device": self.device, + "device_id": self.device_id, + "treelite_model_bytes": self._treelite_model_bytes, + "default_chunk_size": self.default_chunk_size, + "precision": self.precision, + } + + def __setstate__(self, state): + """Rebuild the native inference model from serialized state.""" + self.device, self.device_id = _infer_device_for_deserialization( + state["device"], state["device_id"] + ) + self.stream = None + if self.device == "gpu": + previous_device = Device() + try: + cuda_device = Device(self.device_id) + cuda_device.set_current() + self.stream = cuda_device.create_stream() + finally: + previous_device.set_current() + + self._layout = state["layout"] + self.align_bytes = state["align_bytes"] + self._use_double_precision = state["use_double_precision"] + self._treelite_model_bytes = state["treelite_model_bytes"] + self.default_chunk_size = state["default_chunk_size"] + self.precision = state["precision"] + self.impl = ForestInference_impl( + self.stream, + self._treelite_model_bytes, + layout=self._layout, + align_bytes=self.align_bytes, + use_double_precision=self._use_double_precision, + device=self.device, + device_id=self.device_id, + ) + @property def layout(self) -> str: return self._layout diff --git a/python/nvforest/tests/test_pickle.py b/python/nvforest/tests/test_pickle.py new file mode 100644 index 0000000..0fb4e1f --- /dev/null +++ b/python/nvforest/tests/test_pickle.py @@ -0,0 +1,118 @@ +# +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# + +import pickle + +import cupy as cp +import numpy as np +import pytest +from cuda.bindings import runtime +from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor + +import nvforest + + +def _get_device_count(): + status, device_count = runtime.cudaGetDeviceCount() + return device_count if status == runtime.cudaError_t.cudaSuccess else 0 + + +device_count = _get_device_count() + + +def _make_sklearn_model(model_class): + rng = np.random.default_rng(0) + X = rng.normal(size=(64, 4)).astype(np.float32) + if model_class is RandomForestClassifier: + y = (X[:, 0] + X[:, 1] > 0).astype(np.int32) + else: + y = X[:, 0] - 2 * X[:, 1] + return ( + model_class(n_estimators=3, max_depth=3, random_state=0).fit(X, y), + X, + ) + + +def _to_numpy(array): + return cp.asnumpy(array) if isinstance(array, cp.ndarray) else array + + +def _predict(model, X): + if hasattr(model, "predict_proba"): + return model.predict_proba(X) + return model.predict(X) + + +@pytest.mark.parametrize( + "device", + [ + "cpu", + pytest.param( + "gpu", + marks=pytest.mark.skipif( + device_count == 0, reason="Requires a GPU" + ), + ), + ], +) +@pytest.mark.parametrize( + "model_class", [RandomForestClassifier, RandomForestRegressor] +) +def test_pickle_round_trip(device, model_class): + skl_model, X = _make_sklearn_model(model_class) + model = nvforest.load_from_sklearn( + skl_model, + device=device, + layout="breadth_first", + default_chunk_size=2, + align_bytes=128, + precision="single", + ) + expected = _to_numpy(_predict(model, X)) + + restored = pickle.loads(pickle.dumps(model)) + + assert type(restored) is type(model) + assert restored.layout == "breadth_first" + assert restored.default_chunk_size == 2 + assert restored.align_bytes == 128 + assert restored.precision == "single" + assert restored.forest.device == device + np.testing.assert_array_equal(_to_numpy(_predict(restored, X)), expected) + + +@pytest.mark.skipif(device_count == 0, reason="Requires a GPU") +def test_pickle_falls_back_to_gpu_zero(monkeypatch): + skl_model, _ = _make_sklearn_model(RandomForestClassifier) + model = nvforest.load_from_sklearn(skl_model, device="gpu", device_id=0) + model.forest.device_id = 1 + serialized = pickle.dumps(model) + monkeypatch.setattr( + runtime, + "cudaGetDeviceCount", + lambda: (runtime.cudaError_t.cudaSuccess, 1), + ) + + restored = pickle.loads(serialized) + + assert restored.forest.device == "gpu" + assert restored.device_id == 0 + + +@pytest.mark.skipif(device_count == 0, reason="Requires a GPU") +def test_pickle_falls_back_to_cpu(monkeypatch): + skl_model, X = _make_sklearn_model(RandomForestRegressor) + model = nvforest.load_from_sklearn(skl_model, device="gpu", device_id=0) + expected = _to_numpy(model.predict(X)) + serialized = pickle.dumps(model) + no_device = (runtime.cudaError_t.cudaErrorNoDevice, None) + monkeypatch.setattr(runtime, "cudaGetDeviceCount", lambda: no_device) + monkeypatch.setattr(runtime, "cudaGetDevice", lambda: no_device) + + restored = pickle.loads(serialized) + + assert restored.forest.device == "cpu" + assert restored.device_id == -1 + np.testing.assert_allclose(restored.predict(X), expected) From 44927b0a15e525abaa89c2022d8b171f750865ab Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Fri, 23 Oct 2026 15:34:29 -0700 Subject: [PATCH 2/3] Restore correct outer class when unpickling --- python/nvforest/nvforest/_forest_inference.py | 45 +++++++++++++++++-- python/nvforest/tests/test_pickle.py | 2 + 2 files changed, 43 insertions(+), 4 deletions(-) diff --git a/python/nvforest/nvforest/_forest_inference.py b/python/nvforest/nvforest/_forest_inference.py index 017cef8..bee9fdd 100644 --- a/python/nvforest/nvforest/_forest_inference.py +++ b/python/nvforest/nvforest/_forest_inference.py @@ -41,6 +41,29 @@ def infer_is_classifier(treelite_model: treelite.Model) -> bool: ) +def _restore_forest_inference(state: dict, is_classifier: bool): + """Restore a forest inference object using the available device.""" + forest = ForestInferenceImpl.__new__(ForestInferenceImpl) + forest.__setstate__(state) + + if forest.device == "gpu": + cls = ( + GPUForestInferenceClassifier + if is_classifier + else GPUForestInferenceRegressor + ) + else: + cls = ( + CPUForestInferenceClassifier + if is_classifier + else CPUForestInferenceRegressor + ) + + result = cls.__new__(cls) + result.forest = forest + return result + + class ClassifierMixin: def get_class_assignment( self, @@ -57,6 +80,16 @@ def get_class_assignment( return result +class PickleMixin: + """Mixin class that provides pickle support for ForestInference classes.""" + + def __reduce__(self): + return ( + _restore_forest_inference, + (self.forest.__getstate__(), self.is_classifier), + ) + + class _AutoIterations: """Used to generate sequence of iterations (1, 2, 5, 10, 20, 50...) during optimization.""" @@ -212,7 +245,7 @@ def _optimize( class CPUForestInferenceClassifier( - OptimizeMixin, ClassifierMixin, ForestInferenceClassifier + PickleMixin, OptimizeMixin, ClassifierMixin, ForestInferenceClassifier ): def __init__( self, @@ -352,7 +385,9 @@ def device_id(self) -> int: return -1 -class CPUForestInferenceRegressor(OptimizeMixin, ForestInferenceRegressor): +class CPUForestInferenceRegressor( + PickleMixin, OptimizeMixin, ForestInferenceRegressor +): def __init__( self, *, @@ -482,7 +517,7 @@ def device_id(self) -> int: class GPUForestInferenceClassifier( - OptimizeMixin, ClassifierMixin, ForestInferenceClassifier + PickleMixin, OptimizeMixin, ClassifierMixin, ForestInferenceClassifier ): def __init__( self, @@ -624,7 +659,9 @@ def device_id(self) -> int: return self.forest.device_id -class GPUForestInferenceRegressor(OptimizeMixin, ForestInferenceRegressor): +class GPUForestInferenceRegressor( + PickleMixin, OptimizeMixin, ForestInferenceRegressor +): def __init__( self, *, diff --git a/python/nvforest/tests/test_pickle.py b/python/nvforest/tests/test_pickle.py index 0fb4e1f..17f1c4f 100644 --- a/python/nvforest/tests/test_pickle.py +++ b/python/nvforest/tests/test_pickle.py @@ -97,6 +97,7 @@ def test_pickle_falls_back_to_gpu_zero(monkeypatch): restored = pickle.loads(serialized) + assert isinstance(restored, nvforest.GPUForestInferenceClassifier) assert restored.forest.device == "gpu" assert restored.device_id == 0 @@ -113,6 +114,7 @@ def test_pickle_falls_back_to_cpu(monkeypatch): restored = pickle.loads(serialized) + assert isinstance(restored, nvforest.CPUForestInferenceRegressor) assert restored.forest.device == "cpu" assert restored.device_id == -1 np.testing.assert_allclose(restored.predict(X), expected) From dc2ff39ab1325f2c0c0c43c74f9dd5de70531f9c Mon Sep 17 00:00:00 2001 From: Hyunsu Cho Date: Fri, 23 Oct 2026 16:52:57 -0700 Subject: [PATCH 3/3] Allocate prediction arrays on the model's device --- .../nvforest/detail/forest_inference.pyx | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/python/nvforest/nvforest/detail/forest_inference.pyx b/python/nvforest/nvforest/detail/forest_inference.pyx index 4d2b291..9c96049 100644 --- a/python/nvforest/nvforest/detail/forest_inference.pyx +++ b/python/nvforest/nvforest/detail/forest_inference.pyx @@ -5,6 +5,7 @@ from typing import Optional, Union +import cupy as cp import numpy as np import treelite from cuda.bindings import runtime @@ -68,6 +69,7 @@ cdef class ForestInference_impl(): cdef nvforest_stream_t stream_handle cdef object stream cdef object device + cdef int device_id def __cinit__( self, @@ -78,7 +80,7 @@ cdef class ForestInference_impl(): align_bytes: int = 0, use_double_precision: Optional[bool] = None, device: str = "cpu", - device_id: Optional[int] = None, + device_id: int = -1, ): # Assumption: The caller needs to pass in correct (device, device_id) pair # This function will not contain any logic for auto-detecting device. @@ -131,6 +133,7 @@ cdef class ForestInference_impl(): "device_id should be set before building ForestInference_impl" ) self.device = device + self.device_id = device_id self.model = import_from_treelite_handle( tl_handle, @@ -230,13 +233,13 @@ cdef class ForestInference_impl(): out_dev = nvforest_device_t.cpu else: assert self.device == "gpu" - import cupy as cp - X = cp.asarray(X, dtype=model_dtype, order="C", blocking=True) - preds = cp.empty( - shape=output_shape, - dtype=model_dtype, - order="C", - ) + with cp.cuda.Device(self.device_id): + X = cp.asarray(X, dtype=model_dtype, order="C", blocking=True) + preds = cp.empty( + shape=output_shape, + dtype=model_dtype, + order="C", + ) in_ptr = X.__cuda_array_interface__["data"][0] in_dev = nvforest_device_t.gpu out_ptr = preds.__cuda_array_interface__["data"][0]