Support pickling for nvForest objects

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
 
 

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,47 +41,27 @@ 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 _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
+        )
 
-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
+    result = cls.__new__(cls)
+    result.forest = forest
+    return result
 
 
 class ClassifierMixin:
@@ -101,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."""
@@ -256,7 +245,7 @@ def _optimize(
 
 
 class CPUForestInferenceClassifier(
-    OptimizeMixin, ClassifierMixin, ForestInferenceClassifier
+    PickleMixin, OptimizeMixin, ClassifierMixin, ForestInferenceClassifier
 ):
     def __init__(
         self,
@@ -396,7 +385,9 @@ def device_id(self) -> int:
         return -1
 
 
-class CPUForestInferenceRegressor(OptimizeMixin, ForestInferenceRegressor):
+class CPUForestInferenceRegressor(
+    PickleMixin, OptimizeMixin, ForestInferenceRegressor
+):
     def __init__(
         self,
         *,
@@ -526,7 +517,7 @@ def device_id(self) -> int:
 
 
 class GPUForestInferenceClassifier(
-    OptimizeMixin, ClassifierMixin, ForestInferenceClassifier
+    PickleMixin, OptimizeMixin, ClassifierMixin, ForestInferenceClassifier
 ):
     def __init__(
         self,
@@ -668,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,
         *,

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

python/nvforest/nvforest/detail/forest_inference.pyx

@@ -5,11 +5,14 @@
 
 from typing import Optional, Union
 
+import cupy as cp
 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
@@ -66,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,
@@ -76,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.
@@ -129,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,
@@ -228,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]
@@ -273,6 +278,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 +375,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