From 8bcb79d6151fc4bb15a1c84966b61973d94b18dc Mon Sep 17 00:00:00 2001 From: Rinchin Date: Fri, 30 Oct 2026 13:52:01 +0300 Subject: [PATCH 1/3] Migrate CPU inference to Treelite 4 and TL2cgen --- .github/workflows/python-compatibility.yml | 34 +- README.md | 51 ++ py_boost/gpu/losses/losses.py | 3 +- py_boost/utils/tl_wrapper.py | 700 +++++++++------------ pyproject.toml | 4 +- tests/test_treelite.py | 264 ++++++++ tests/test_treelite_gpu.py | 74 +++ tutorials/Tutorial_1_Basics.ipynb | 262 +------- 8 files changed, 767 insertions(+), 625 deletions(-) create mode 100644 tests/test_treelite.py create mode 100644 tests/test_treelite_gpu.py diff --git a/.github/workflows/python-compatibility.yml b/.github/workflows/python-compatibility.yml index efbe8d5..cdbffac 100644 --- a/.github/workflows/python-compatibility.yml +++ b/.github/workflows/python-compatibility.yml @@ -10,22 +10,50 @@ permissions: jobs: cpu: - runs-on: ubuntu-latest + runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: - python-version: ["3.13", "3.14"] + include: + - os: ubuntu-latest + python-version: "3.8" + treelite: "treelite==4.1.2" + - os: ubuntu-latest + python-version: "3.12" + treelite: "treelite>=4.1.2,<5" + - os: ubuntu-latest + python-version: "3.12" + treelite: "treelite==4.1.2" + - os: ubuntu-latest + python-version: "3.13" + treelite: "treelite>=4.1.2,<5" + - os: ubuntu-latest + python-version: "3.14" + treelite: "treelite>=4.1.2,<5" + - os: windows-latest + python-version: "3.12" + treelite: "treelite>=4.1.2,<5" + - os: windows-latest + python-version: "3.14" + treelite: "treelite>=4.1.2,<5" + defaults: + run: + shell: bash + env: + OMP_NUM_THREADS: "2" steps: - uses: actions/checkout@v7 - uses: actions/setup-python@v7 with: python-version: ${{ matrix.python-version }} + - uses: ilammy/msvc-dev-cmd@v1 + if: runner.os == 'Windows' - name: Build distributions run: | python -m pip install --upgrade pip build python -m build - name: Install wheel and test dependencies - run: python -m pip install dist/*.whl pytest + run: python -m pip install dist/*.whl pytest '${{ matrix.treelite }}' - name: Run test suite working-directory: ${{ runner.temp }} run: python -m pytest -q -rs "$GITHUB_WORKSPACE/tests" diff --git a/README.md b/README.md index f78e51d..00ce108 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,57 @@ with CUDA 12.x: For other Python or CUDA versions, choose the appropriate CuPy package from the [CuPy installation guide](https://docs.cupy.dev/en/stable/install.html). +## CPU inference with Treelite and TL2cgen + +Py-Boost uses Treelite 4 (`>=4.1.2,<5`) and TL2cgen 1 (`>=1,<2`) for CPU +inference. Both dependencies are installed with Py-Boost; `treelite_runtime` +is no longer used. NumPy 1.x and 2.x are supported, subject to your Python +version and the other dependencies. + +Convert a fitted model trained with `target_splitter='Single'` (the default): + +```python +from py_boost import TLPredictor, TLCompiledPredictor + +tl_model = TLPredictor(model) +prediction = tl_model.predict(X_test, nthread=4) +tl_model.dump('cpu_model') +restored = TLPredictor.load('cpu_model') +``` + +Single-output and multi-output regression, binary/multilabel classification, +and multiclass classification keep the original model's postprocessing. +Predictions have shape `(n_samples, n_outputs)`, including single-output models. +The underlying Treelite model stores raw scores as regression targets; the +wrapper applies `model.postprocess_fn` once, after restoring this output shape. +A custom CPU postprocessor can be supplied with `postprocess_fn=...`. + +Compile the model with TL2cgen for faster CPU inference. Compilation requires a +C compiler; prediction with the compiled library does not require a compiler: + +```python +# Linux: GCC must be installed and on PATH. +compiled = tl_model.compile('gcc', 'model.so', nthread=4) +# Windows: run from a Visual Studio developer shell, with MSVC on PATH. +# compiled = tl_model.compile('msvc', 'model.dll', nthread=4) + +prediction = compiled.predict(X_test) +compiled.dump('compiled.pkl') +restored_compiled = TLCompiledPredictor.load('compiled.pkl') +``` + +Compiled predictor dumps contain metadata and the library path. Deploy the +`.so`/`.dll` alongside the dump. If its path changes, load it with +`TLCompiledPredictor.load('compiled.pkl', libpath=new_path)`. Use +`set_libpath(new_path)` to relocate an already loaded predictor. +Use a library built for the target operating system and architecture. The +legacy `compiler` argument accepts only its default value, `'ast_native'`. +See the [TL2cgen migration guide](https://tl2cgen.readthedocs.io/en/latest/treelite-migration.html) +for the underlying API changes. + +Treelite 3 model checkpoints and libraries compiled with `treelite_runtime` +should be re-exported from the original Py-Boost model with this version. + ## Quick tour Py-boost is easy to use since it has an interface similar to scikit-learn's. For usage examples, please see: diff --git a/py_boost/gpu/losses/losses.py b/py_boost/gpu/losses/losses.py index 4287815..1d3ac43 100644 --- a/py_boost/gpu/losses/losses.py +++ b/py_boost/gpu/losses/losses.py @@ -111,7 +111,8 @@ def get_grad_hess(self, y_true, y_pred): return (y_pred - y_true), cp.ones((y_true.shape[0], 1), dtype=cp.float32) def postprocess_output(self, y_pred): - return cp.expm1(y_pred) + xp = np if type(y_pred) is np.ndarray else cp + return xp.expm1(y_pred) def base_score(self, y_true): y_true = cp.log1p(y_true) diff --git a/py_boost/utils/tl_wrapper.py b/py_boost/utils/tl_wrapper.py index c5b8556..b224d4e 100644 --- a/py_boost/utils/tl_wrapper.py +++ b/py_boost/utils/tl_wrapper.py @@ -1,388 +1,312 @@ -import os - -import joblib -from tqdm import tqdm -import treelite -import treelite_runtime as tl_run - - -def _create_node_deprecated(tree, node_id): - """(DEPRECATED) Create a node of treelite tree - Args: - tree: Py-Boost Tree, tree to parse - node_id: int, node index - Returns: - dict, args of treelite.ModelBuilder.Tree .set_numerical_test_node or .set_leaf_node - """ - feature_id = tree.feats[0][node_id] - - if feature_id >= 0: - left, right = tuple(tree.split[0][node_id]) - node = { - - 'feature_id': feature_id, - 'opname': '<=', - 'threshold': tree.val_splits[0][node_id], - 'default_left': tree.nans[0][node_id], - 'left_child_key': left, - 'right_child_key': right, - } - - return node, left, right - - return {'value': tree.values[tree.leaves[node_id][0]]}, None, None - - -def create_node(tree, node_id, id_gen): - """Create a node of treelite tree - - Args: - tree: Py-Boost Tree, tree to parse - node_id: int, node index in original format tree - id_gen: generator, new id generator - - Returns: - dict, args of treelite.ModelBuilder.Tree .set_numerical_test_node - """ - - assert node_id >= 0 - - feature_id = int(tree.test_format[node_id * 4]) - nan_left = feature_id < 0 - feature_id = abs(feature_id) - 1 - - left = int(tree.test_format[node_id * 4 + 2]) - right = int(tree.test_format[node_id * 4 + 3]) - new_id_left = next(id_gen) - new_id_right = next(id_gen) - node = { - 'feature_id': feature_id, - 'opname': '<=', - 'threshold': tree.test_format[node_id * 4 + 1], - 'default_left': nan_left, - 'left_child_key': new_id_left, - 'right_child_key': new_id_right, - } - - return node, (left, new_id_left), (right, new_id_right) - - -def parse_pb_tree(tree): - """Parse s single Py-Boost Tree to treelite.ModelBuilder.Tree format - - Args: - tree: Py-Boost tree - - Returns: - treelite.ModelBuilder.Tree - """ - assert tree.ngroups == 1, 'Models with more than 1 group are not currently supported' - - def id_generator(): - id_num = 1 - while True: - yield id_num - id_num += 1 - id_gen = id_generator() - - tl_tree = treelite.ModelBuilder.Tree() - curr_nodes = [(0, 0)] # (old_id, new_id) - - while len(curr_nodes) > 0: - old_id, new_id = curr_nodes.pop(0) - curr_node, left, right = create_node(tree, old_id, id_gen) - tl_tree[new_id].set_numerical_test_node(**curr_node) - - if left[0] >= 0: - curr_nodes.append(left) - else: - tl_tree[left[1]].set_leaf_node(tree.values[abs(left[0]) - 1]) - if right[0] >= 0: - curr_nodes.append(right) - else: - tl_tree[right[1]].set_leaf_node(tree.values[abs(right[0]) - 1]) - - tl_tree[0].set_root() - return tl_tree - - -def _parse_pb_tree_deprecated(tree): - """Parse s single Py-Boost Tree to treelite.ModelBuilder.Tree format - - Args: - tree: Py-Boost tree - - Returns: - treelite.ModelBuilder.Tree - """ - assert tree.ngroups == 1, 'Models with more than 1 group are not currently supported' - - tl_tree = treelite.ModelBuilder.Tree() - curr_nodes = [0] - - while len(curr_nodes) > 0: - - next_nodes = [] - - for node_id in curr_nodes: - - curr_node, left, right = _create_node_deprecated(tree, node_id) - # add node - tl_tree[node_id] - if left is not None: - tl_tree[node_id].set_numerical_test_node(**curr_node) - next_nodes.extend([left, right]) - else: - tl_tree[node_id].set_leaf_node(curr_node['value']) - - curr_nodes = next_nodes - - tl_tree[0].set_root() - - return tl_tree - - -def convert_pb_to_treelite(model): - """Convert Py-Boost Ensemble instance to the treelite.ModelBuilder.Tree - - Args: - model: Py-Boost Tree - - Returns: - treelite.ModelBuilder.Tree - """ - nfeats = model.nfeats - ngroups = model.models[0].values.shape[1] - - builder = treelite.ModelBuilder( - num_feature=nfeats, - num_class=ngroups, - pred_transform='identity_multiclass' if ngroups > 1 else 'identity' - ) - - for tree in tqdm(model.models): - builder.append(parse_pb_tree(tree)) - - # add bias tree - bias_tree = treelite.ModelBuilder.Tree() - bias_tree[0].set_numerical_test_node(**{ - - 'feature_id': 0, - 'opname': '<', - 'threshold': 0, - 'default_left': True, - 'left_child_key': 1, - 'right_child_key': 2 - }) - - for i in range(1, 3): - bias_tree[i].set_leaf_node(model.base_score) - - bias_tree[0].set_root() - builder.append(bias_tree) - - return builder - - -class TLCompiledPredictor: - """ - Compiled treelite model saved to predict - """ - - @staticmethod - def _default_postprocess_fn(x): - return x - - def __init__(self, libpath, nthread=None, verbose=False, postprocess_fn=None): - """ - - Args: - libpath: str, path to compiled model - nthread: int or None, number of threads to use - verbose: bool, verbosity mode - postprocess_fn: Callable or None, prediction postprocessing function - """ - self.verbose = verbose - self.nthread = nthread - self.libpath = None - self.set_libpath(libpath) - - self.postprocess_fn = self._default_postprocess_fn - if postprocess_fn is not None: - self.postprocess_fn = postprocess_fn - - def predict(self, X): - """Make prediction - - Args: - X: np.ndarray - - Returns: - np.ndarray - """ - pred = self.predictor.predict(tl_run.DMatrix(X)) - return self.postprocess_fn(pred) - - def set_libpath(self, libpath=None, nthread=None): - """Update library path - - Args: - libpath: - nthread: int, num threads - - Returns: - - """ - if libpath is None: - libpath = self.libpath - if nthread is None: - nthread = self.nthread - self.libpath = os.path.abspath(libpath) - self.predictor = tl_run.Predictor(self.libpath, nthread=nthread, verbose=self.verbose) - - def dump(self, filename): - """Dump instance - - Args: - filename: str, path to save - - Returns: - - """ - self.predictor = None - joblib.dump(self, filename) - - @staticmethod - def load(filename): - """Load instance - - Args: - filename: str, filename - - Returns: - TLCompiledPredictor - """ - predictor = joblib.load(filename) - predictor.set_libpath() - - return predictor - - -class TLPredictor: - """ - Treelite predictor. Could be use for inference via built-in treelite utils - or to compilation to get TLCompiledPredictor - """ - - def __init__(self, model, postprocess_fn=None): - """ - - Args: - model: Py-Boost Ensemble - postprocess_fn: Callable or None, postprocessing function - """ - model.to_cpu() - self.tl_model = convert_pb_to_treelite(model).commit() - - self.postprocess_fn = postprocess_fn - if postprocess_fn is None: - self.postprocess_fn = model.postprocess_fn - - def set_tl_model(self, tl_model): - """Update underlying treelite model - - Args: - tl_model: - - Returns: - - """ - self.tl_model = tl_model - - def compile( - self, - toolchain, - libpath, - params=None, - compiler='ast_native', - verbose=False, - nthread=None, - options=None, - predictor_params=None - ): - """Compile model for faster inference. For the details please see - https://treelite.readthedocs.io/en/latest/tutorials/first.html - - Args: - toolchain: - libpath: - params: - compiler: - verbose: - nthread: - options: - predictor_params: - - Returns: - - """ - - params = {} if params is None else params - params = {**{'parallel_comp': os.cpu_count(), }, **params} - - self.tl_model.export_lib(toolchain, libpath, - params, compiler, verbose, nthread, options) - - if predictor_params is None: - predictor_params = {} - predictor_params = {**{'nthread': nthread}, **predictor_params} - - predictor = TLCompiledPredictor(libpath, postprocess_fn=self.postprocess_fn, **predictor_params) - return predictor - - def predict(self, X, nthread=None): - """Make prediction - - Args: - X: np.ndarray - nthread: int/None, used for prediction - - Returns: - np.ndarray - """ - if nthread is None: - nthread = os.cpu_count() - pred = treelite.gtil.predict(self.tl_model, X, nthread=nthread) - return self.postprocess_fn(pred) - - def dump(self, dirname, rewrite=False): - """Dump treelite Model and predictor instance - - Args: - dirname: str, path to save - rewrite: bool, possible to overwrite - - Returns: - - """ - os.makedirs(dirname, exist_ok=rewrite) - temp = self.tl_model - self.tl_model = None - temp.serialize(os.path.join(dirname, 'model.mod')) - joblib.dump(self, os.path.join(dirname, 'predictor.pkl')) - self.tl_model = temp - - @staticmethod - def load(dirname): - """Load predictor from folder - - Args: - dirname: str, path - - Returns: - TLPredictor - """ - predictor = joblib.load(os.path.join(dirname, 'predictor.pkl')) - predictor.set_tl_model(treelite.Model.deserialize(os.path.join(dirname, 'model.mod'))) - - return predictor +import os +from collections import deque +from itertools import count + +import joblib +import numpy as np +import tl2cgen +import treelite +from treelite.model_builder import Metadata, ModelBuilder, PostProcessorFunc, TreeAnnotation +from tqdm import tqdm + + +def parse_pb_tree(tree, builder): + """Append a single-group Py-Boost tree to a Treelite 4 model builder.""" + if tree.ngroups != 1: + raise ValueError("Treelite conversion requires target_splitter='Single'") + + # Py-Boost uses negative child indices for leaves. Assign fresh, nonnegative + # Treelite keys to both split and leaf nodes, including repeated leaf indices. + node_ids = count(1) + nodes = deque([(0, 0)]) + builder.start_tree() + while nodes: + old_id, new_id = nodes.popleft() + builder.start_node(new_id) + if old_id < 0: + values = np.asarray(tree.values[-old_id - 1]).reshape(-1).tolist() + builder.leaf(values[0] if len(values) == 1 else values) + else: + feature, threshold, left, right = tree.test_format[old_id * 4:old_id * 4 + 4] + left_id, right_id = next(node_ids), next(node_ids) + builder.numerical_test( + feature_id=abs(int(feature)) - 1, + threshold=float(threshold), + default_left=bool(feature < 0), + opname='<=', + left_child_key=left_id, + right_child_key=right_id, + ) + nodes.extend([(int(left), left_id), (int(right), right_id)]) + builder.end_node() + builder.end_tree() + + +def convert_pb_to_treelite(model): + """Return a Treelite 4 builder containing the ensemble's raw scores. + + All outputs are represented as regression targets. The original Py-Boost + postprocess_fn handles the loss-specific transform, including classification, + so custom losses and postprocessors retain their existing behavior. + """ + if not model.models: + raise ValueError('Cannot convert an ensemble without trees') + nout = model.models[0].values.shape[1] + num_tree = len(model.models) + 1 + builder = ModelBuilder( + threshold_type='float32', + leaf_output_type='float32', + metadata=Metadata( + num_feature=int(model.nfeats), + task_type='kRegressor', + average_tree_output=False, + num_target=nout, + num_class=[1] * nout, + leaf_vector_shape=(nout, 1), + ), + tree_annotation=TreeAnnotation( + num_tree=num_tree, + target_id=[-1 if nout > 1 else 0] * num_tree, + class_id=[0] * num_tree, + ), + postprocessor=PostProcessorFunc(name='identity'), + base_scores=[0.0] * nout, + ) + # TL2cgen 1.0 applies base_scores[0] to every regression target. Store the + # vector intercept as a constant tree instead, before the boosting trees, + # matching Py-Boost's accumulation order in both GTIL and compiled inference. + bias = np.asarray(model.base_score).reshape(-1).tolist() + builder.start_tree() + builder.start_node(0) + builder.leaf(bias[0] if nout == 1 else bias) + builder.end_node() + builder.end_tree() + for tree in tqdm(model.models): + parse_pb_tree(tree, builder) + return builder + + +def _reshape_prediction(pred): + """Restore Py-Boost's (n_samples, n_outputs) contract, including empty batches.""" + return pred.reshape(pred.shape[0], pred.shape[1] * pred.shape[2]) + + +class TLCompiledPredictor: + """ + Compiled TL2cgen model for CPU inference + """ + + @staticmethod + def _default_postprocess_fn(x): + return x + + def __init__(self, libpath, nthread=None, verbose=False, postprocess_fn=None): + """ + + Args: + libpath: str, path to compiled model + nthread: int or None, number of threads to use + verbose: bool, verbosity mode + postprocess_fn: Callable or None, prediction postprocessing function + """ + self.verbose = verbose + self.nthread = nthread + self.libpath = None + self.set_libpath(libpath) + + self.postprocess_fn = self._default_postprocess_fn + if postprocess_fn is not None: + self.postprocess_fn = postprocess_fn + + def predict(self, X): + """Make prediction + + Args: + X: np.ndarray + + Returns: + np.ndarray + """ + X = np.asarray(X, dtype=np.float32, order="C") + pred = self.predictor.predict(tl2cgen.DMatrix(X)) + return self.postprocess_fn(_reshape_prediction(pred)) + + def set_libpath(self, libpath=None, nthread=None): + """Update library path + + Args: + libpath: + nthread: int, num threads + + Returns: + + """ + if libpath is None: + libpath = self.libpath + if nthread is None: + nthread = self.nthread + self.libpath = os.path.abspath(libpath) + self.nthread = nthread + self.predictor = tl2cgen.Predictor(self.libpath, nthread=nthread, verbose=self.verbose) + + def __getstate__(self): + state = self.__dict__.copy() + state["predictor"] = None + return state + + def dump(self, filename): + """Dump instance + + Args: + filename: str, path to save + + Returns: + + """ + joblib.dump(self, filename) + + @staticmethod + def load(filename, libpath=None, nthread=None): + """Load instance + + Args: + filename: str, filename + libpath: str or None, override the saved compiled library path + nthread: int or None, override the saved number of worker threads + + Returns: + TLCompiledPredictor + """ + predictor = joblib.load(filename) + predictor.set_libpath(libpath=libpath, nthread=nthread) + + return predictor + + +class TLPredictor: + """ + Treelite predictor. Could be use for inference via built-in treelite utils + or to compilation to get TLCompiledPredictor + """ + + def __init__(self, model, postprocess_fn=None): + """ + + Args: + model: Py-Boost Ensemble + postprocess_fn: Callable or None, postprocessing function + """ + model.to_cpu() + self.tl_model = convert_pb_to_treelite(model).commit() + + self.postprocess_fn = postprocess_fn + if postprocess_fn is None: + self.postprocess_fn = model.postprocess_fn + + def set_tl_model(self, tl_model): + """Update underlying treelite model + + Args: + tl_model: + + Returns: + + """ + self.tl_model = tl_model + + def compile( + self, + toolchain, + libpath, + params=None, + compiler='ast_native', + verbose=False, + nthread=None, + options=None, + predictor_params=None + ): + """Compile a CPU inference library with TL2cgen. + + See https://tl2cgen.readthedocs.io/en/latest/api.html#tl2cgen.export_lib + + Args: + toolchain: + libpath: + params: + compiler: retained for compatibility; only 'ast_native' is supported + verbose: + nthread: + options: + predictor_params: + + Returns: + + """ + + if compiler != 'ast_native': + raise ValueError("TL2cgen only supports compiler='ast_native'") + params = {} if params is None else params + params = {**{'parallel_comp': os.cpu_count() or 1}, **params} + + tl2cgen.export_lib( + self.tl_model, toolchain=toolchain, libpath=libpath, params=params, + verbose=verbose, nthread=nthread, options=options, + ) + + if predictor_params is None: + predictor_params = {} + predictor_params = {**{'nthread': nthread}, **predictor_params} + + predictor = TLCompiledPredictor(libpath, postprocess_fn=self.postprocess_fn, **predictor_params) + return predictor + + def predict(self, X, nthread=None): + """Make prediction + + Args: + X: np.ndarray + nthread: int/None, used for prediction + + Returns: + np.ndarray + """ + if nthread is None: + # Let Treelite respect OpenMP's thread limit (for example in CI). + nthread = -1 + X = np.asarray(X, dtype=np.float32, order="C") + pred = treelite.gtil.predict(self.tl_model, X, nthread=nthread) + return self.postprocess_fn(_reshape_prediction(pred)) + + def __getstate__(self): + state = self.__dict__.copy() + state["tl_model"] = None + return state + + def dump(self, dirname, rewrite=False): + """Dump treelite Model and predictor instance + + Args: + dirname: str, path to save + rewrite: bool, possible to overwrite + + Returns: + + """ + os.makedirs(dirname, exist_ok=rewrite) + self.tl_model.serialize(os.path.join(dirname, 'model.mod')) + joblib.dump(self, os.path.join(dirname, 'predictor.pkl')) + + @staticmethod + def load(dirname): + """Load predictor from folder + + Args: + dirname: str, path + + Returns: + TLPredictor + """ + predictor = joblib.load(os.path.join(dirname, 'predictor.pkl')) + predictor.set_tl_model(treelite.Model.deserialize(os.path.join(dirname, 'model.mod'))) + + return predictor diff --git a/pyproject.toml b/pyproject.toml index f5116b0..62567e2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,8 +41,8 @@ ujson = '*' pandas = ">=1" onnx = ">=1.16, <2" onnxruntime = ">=1.16, <2" -treelite = "^3" -treelite_runtime = "^3" +treelite = ">=4.1.2,<5" +tl2cgen = ">=1,<2" tqdm = ">=4.64.1" diff --git a/tests/test_treelite.py b/tests/test_treelite.py new file mode 100644 index 0000000..3212e20 --- /dev/null +++ b/tests/test_treelite.py @@ -0,0 +1,264 @@ +"""CPU integration tests using real Treelite models and compiled TL2cgen libraries.""" + +import os +import shutil +import sys +from types import SimpleNamespace + +import joblib +import numpy as np +import pytest + +from py_boost import TLCompiledPredictor, TLPredictor + + +def identity(prediction): + return prediction + + +def sigmoid(prediction): + return 1 / (1 + np.exp(-prediction)) + + +def softmax(prediction): + exp = np.exp(prediction - prediction.max(axis=1, keepdims=True)) + return exp / exp.sum(axis=1, keepdims=True) + + +class CPUEnsemble: + """Small, CPU-resident ensemble with independently specified expected outputs.""" + + def __init__(self, nout, postprocess_fn=identity): + self.nfeats = 2 + self.base_score = np.array([0.75, -1.5, 3], dtype=np.float32)[:nout] + self.postprocess_fn = postprocess_fn + self.models = [ + SimpleNamespace( + ngroups=1, + # x0 <= .5 or missing: test x1 <= -.25 (missing goes right). + # Otherwise, use leaf 2. Negative children encode leaf indices. + test_format=np.array([-1, .5, 1, -3, 2, -.25, -1, -2], dtype=np.float32), + values=np.array([[1, -2, .5], [3, 4, -1], [-5, 6, 2]], + dtype=np.float32)[:, :nout], + ), + SimpleNamespace( + ngroups=1, + # Py-Boost represents a constant tree as a split with equal leaves. + test_format=np.array([1, 0, -1, -1], dtype=np.float32), + values=np.array([[.25, -.5, 2]], dtype=np.float32)[:, :nout], + ), + ] + + def to_cpu(self): + pass + + +@pytest.fixture(scope='module', params=[1, 3], ids=['scalar', 'vector']) +def ensemble(request): + return CPUEnsemble(request.param) + + +@pytest.fixture +def examples(): + X = np.array([ + [0, -1], [0, 0], [1, -1], [np.nan, -1], [np.nan, np.nan], + [.5, -.25], [-np.inf, np.inf], [np.inf, np.nan], [0, np.nan], + ], dtype=np.float32) + # Computed from the tree's decision rules, including equality and NaN routing. + expected = np.array([ + [2, -4, 5.5], [4, 2, 4], [-4, 4, 7], [2, -4, 5.5], [4, 2, 4], + [2, -4, 5.5], [4, 2, 4], [-4, 4, 7], [4, 2, 4], + ], dtype=np.float32) + return X, expected + + +@pytest.fixture(scope='module') +def toolchain(): + if sys.platform == 'win32': + name, suffix = 'msvc', '.dll' + executable = 'cl' + else: + name, suffix = 'gcc', '.dylib' if sys.platform == 'darwin' else '.so' + executable = 'gcc' + if shutil.which(executable) is None: + if os.environ.get('CI'): + pytest.fail(f'Compiled inference tests require {executable} in CI') + pytest.skip(f'Compiled inference tests require {executable}') + return name, suffix + + +@pytest.fixture(scope='module') +def compiled(ensemble, tmp_path_factory, toolchain): + name, suffix = toolchain + tmp_path = tmp_path_factory.mktemp('compiled') + return TLPredictor(ensemble).compile( + name, tmp_path / ('model' + suffix), params={'parallel_comp': 2}, + nthread=2, predictor_params={'nthread': 1}, + ) + + +@pytest.mark.parametrize('layout', ['float32', 'float64', 'fortran', 'strided', 'list']) +def test_gtil_and_compiled_predictions(ensemble, compiled, examples, layout): + X, expected = examples + if layout == 'float64': + X = X.astype(np.float64) + elif layout == 'fortran': + X = np.asfortranarray(X) + elif layout == 'strided': + buffer = np.zeros((len(X), 4), dtype=np.float32) + buffer[:, ::2] = X + X = buffer[:, ::2] + elif layout == 'list': + X = X.tolist() + expected = expected[:, :len(ensemble.base_score)] + for predictor in (TLPredictor(ensemble), compiled): + actual = predictor.predict(X) + assert type(actual) is np.ndarray + assert actual.shape == expected.shape + assert actual.dtype == np.float32 + np.testing.assert_allclose(actual, expected, rtol=1e-6, atol=1e-6) + + +@pytest.mark.parametrize('batch_size', [0, 1]) +def test_prediction_keeps_batch_and_output_axes(ensemble, compiled, examples, batch_size): + X, expected = examples + expected = expected[:batch_size, :len(ensemble.base_score)] + for predictor in (TLPredictor(ensemble), compiled): + actual = predictor.predict(X[:batch_size]) + assert actual.shape == expected.shape + np.testing.assert_allclose(actual, expected) + + +@pytest.mark.parametrize('postprocess_fn', [sigmoid, softmax]) +def test_postprocess_receives_2d_raw_scores(examples, tmp_path, toolchain, postprocess_fn): + X, raw_expected = examples + model = CPUEnsemble(3, postprocess_fn=postprocess_fn) + predictor = TLPredictor(model) + name, suffix = toolchain + compiled = predictor.compile(name, tmp_path / ('model' + suffix), + params={'parallel_comp': 1}, nthread=2) + expected = postprocess_fn(raw_expected) + for instance in (predictor, compiled): + np.testing.assert_allclose(instance.predict(X), expected, rtol=1e-6, atol=1e-6) + assert instance.predict(X[:1]).shape == (1, 3) + + +@pytest.mark.parametrize('loss_name,transform', [ + ('MSELoss', identity), ('MSLELoss', np.expm1), + ('BCELoss', sigmoid), ('CrossEntropyLoss', softmax), +]) +def test_builtin_postprocessors_on_cpu(examples, tmp_path, toolchain, loss_name, transform): + # Exercise the actual bound methods stored by GPU-trained model dumps. These + # must also import, unpickle and run in an environment without CuPy installed. + from py_boost.gpu.losses import losses + + postprocess_fn = getattr(losses, loss_name)().postprocess_output + predictor = TLPredictor(CPUEnsemble(3, postprocess_fn=postprocess_fn)) + predictor.dump(tmp_path / 'gtil') + restored = TLPredictor.load(tmp_path / 'gtil') + name, suffix = toolchain + compiled = predictor.compile(name, tmp_path / ('model' + suffix), + params={'parallel_comp': 1}, nthread=2) + compiled.dump(tmp_path / 'compiled.pkl') + restored_compiled = TLCompiledPredictor.load(tmp_path / 'compiled.pkl') + X, expected = examples + for instance in (predictor, restored, compiled, restored_compiled): + actual = instance.predict(X) + assert type(actual) is np.ndarray + np.testing.assert_allclose(actual, transform(expected), rtol=1e-6, atol=1e-6) + + +def test_custom_postprocess_overrides_model(ensemble, examples): + X, raw_expected = examples + predictor = TLPredictor(ensemble, postprocess_fn=sigmoid) + np.testing.assert_allclose(predictor.predict(X), + sigmoid(raw_expected[:, :len(ensemble.base_score)])) + + +def test_compile_with_default_options(examples, tmp_path, toolchain): + X, expected = examples + name, suffix = toolchain + compiled = TLPredictor(CPUEnsemble(3)).compile(name, tmp_path / ('model' + suffix)) + np.testing.assert_allclose(compiled.predict(X), expected) + + +def test_gtil_serialization(ensemble, examples, tmp_path): + X, expected = examples + predictor = TLPredictor(ensemble, postprocess_fn=sigmoid) + path = tmp_path / 'saved_model' + predictor.dump(path) + with pytest.raises(FileExistsError): + predictor.dump(path) + predictor.dump(path, rewrite=True) + loaded = TLPredictor.load(path) + for instance in (predictor, loaded): + np.testing.assert_allclose(instance.predict(X), + sigmoid(expected[:, :len(ensemble.base_score)])) + + +def test_compiled_serialization_and_relocation(compiled, examples, tmp_path, toolchain): + X, expected = examples + nout = compiled.predict(X).shape[1] + filename = tmp_path / 'compiled.pkl' + compiled.dump(filename) + # Dump must leave the original predictor usable. + np.testing.assert_allclose(compiled.predict(X), expected[:, :nout]) + loaded = TLCompiledPredictor.load(filename) + assert loaded.nthread == 1 # predictor_params overrides compile's nthread. + np.testing.assert_allclose(loaded.predict(X), expected[:, :nout]) + _, suffix = toolchain + relocated = tmp_path / ('relocated' + suffix) + # Copy rather than rename: Windows holds loaded DLLs open. + shutil.copy2(loaded.libpath, relocated) + loaded.set_libpath(relocated, nthread=2) + loaded.dump(filename) + loaded = TLCompiledPredictor.load(filename) + assert loaded.nthread == 2 + assert loaded.libpath == os.path.abspath(relocated) + np.testing.assert_allclose(loaded.predict(X), expected[:, :nout]) + + # A deployment may have no file at the original machine's absolute path. + metadata = joblib.load(filename) + metadata.libpath = str(tmp_path / 'missing_library') + metadata.dump(filename) + loaded = TLCompiledPredictor.load(filename, libpath=relocated, nthread=1) + assert loaded.nthread == 1 + np.testing.assert_allclose(loaded.predict(X), expected[:, :nout]) + + +@pytest.mark.parametrize('compiled_predictor', [False, True]) +def test_failed_dump_leaves_predictor_usable(ensemble, examples, tmp_path, toolchain, + monkeypatch, compiled_predictor): + predictor = TLPredictor(ensemble) + if compiled_predictor: + name, suffix = toolchain + predictor = predictor.compile(name, tmp_path / ('model' + suffix), + params={'parallel_comp': 1}, nthread=2) + + def fail(*args, **kwargs): + raise OSError('disk full') + + monkeypatch.setattr('py_boost.utils.tl_wrapper.joblib.dump', fail) + with pytest.raises(OSError, match='disk full'): + predictor.dump(tmp_path / 'failed_dump') + X, expected = examples + np.testing.assert_allclose(predictor.predict(X), expected[:, :len(ensemble.base_score)]) + + +def test_rejects_unsupported_compiler(ensemble, tmp_path): + with pytest.raises(ValueError, match='ast_native'): + TLPredictor(ensemble).compile('gcc', tmp_path / 'unused.so', compiler='failsafe') + + +def test_rejects_multiple_tree_groups(): + ensemble = CPUEnsemble(3) + ensemble.models[0].ngroups = 2 + with pytest.raises(ValueError, match="target_splitter='Single'"): + TLPredictor(ensemble) + + +def test_rejects_unfitted_model(): + ensemble = CPUEnsemble(1) + ensemble.models = None + with pytest.raises(ValueError, match='without trees'): + TLPredictor(ensemble) diff --git a/tests/test_treelite_gpu.py b/tests/test_treelite_gpu.py new file mode 100644 index 0000000..4a4cbb9 --- /dev/null +++ b/tests/test_treelite_gpu.py @@ -0,0 +1,74 @@ +"""End-to-end parity between trained GPU models, GTIL and compiled CPU inference.""" + +import shutil +import sys + +import numpy as np +import pytest + +cp = pytest.importorskip('cupy') +if not cp.is_available(): + pytest.skip('Training parity tests require a CUDA device', allow_module_level=True) + +from py_boost import GradientBoosting, TLCompiledPredictor, TLPredictor + + +@pytest.mark.parametrize('loss,nout', [ + ('mse', 1), ('mse', 3), ('msle', 1), ('msle', 3), + ('bce', 1), ('multilabel', 3), ('multiclass', 3), +]) +@pytest.mark.parametrize('constant_tree', [False, True], ids=['splits', 'constant']) +def test_trained_model_cpu_parity(loss, nout, constant_tree, tmp_path): + rng = np.random.default_rng(123) + X = rng.normal(size=(192, 4)).astype(np.float32) + X[::11, 0] = np.nan + X[::13, 1] = np.nan + clean = np.nan_to_num(X) + scores = np.column_stack([ + clean[:, 0] + clean[:, 1], clean[:, 2] - clean[:, 1], clean[:, 0] * clean[:, 3], + ]) + if loss == 'multiclass': + y = scores.argmax(axis=1).astype(np.float32) + elif loss in ('bce', 'multilabel'): + y = (scores[:, :nout] > 0).astype(np.float32) + elif loss == 'msle': + y = np.exp(scores[:, :nout] / 3) + else: + y = scores[:, :nout] + + params = dict(ntrees=4, max_depth=3, min_data_in_leaf=2, max_bin=32, + es=0, verbose=1000) + if constant_tree: + params['min_gain_to_split'] = 1e10 + model = GradientBoosting(loss, **params).fit(X[:160], y[:160]) + expected = model.predict(X) + assert expected.shape == (len(X), nout) + if constant_tree: + np.testing.assert_allclose(expected, np.broadcast_to(expected[:1], expected.shape)) + else: + assert np.std(expected, axis=0).max() > .001 + + predictor = TLPredictor(model) + path = tmp_path / 'gtil' + predictor.dump(path) + restored = TLPredictor.load(path) + for instance in (predictor, restored): + actual = instance.predict(X, nthread=2) + assert type(actual) is np.ndarray + assert actual.shape == expected.shape + np.testing.assert_allclose(actual, expected, rtol=2e-5, atol=2e-6) + + toolchain = 'msvc' if sys.platform == 'win32' else 'gcc' + executable = 'cl' if sys.platform == 'win32' else toolchain + if shutil.which(executable) is None: + pytest.skip(f'Compilation parity requires {executable}') + suffix = '.dll' if sys.platform == 'win32' else '.dylib' if sys.platform == 'darwin' else '.so' + compiled = predictor.compile(toolchain, tmp_path / ('model' + suffix), + params={'parallel_comp': 2}, nthread=2) + compiled.dump(tmp_path / 'compiled.pkl') + restored = TLCompiledPredictor.load(tmp_path / 'compiled.pkl') + for instance in (compiled, restored): + actual = instance.predict(X) + assert type(actual) is np.ndarray + assert actual.shape == expected.shape + np.testing.assert_allclose(actual, expected, rtol=2e-5, atol=2e-6) diff --git a/tutorials/Tutorial_1_Basics.ipynb b/tutorials/Tutorial_1_Basics.ipynb index c4453ef..e7935e3 100644 --- a/tutorials/Tutorial_1_Basics.ipynb +++ b/tutorials/Tutorial_1_Basics.ipynb @@ -538,32 +538,17 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### CPU Inference via treelite\n", + "### CPU inference with Treelite 4\n", "\n", - "We can also save model for CPU inference via treelite library. For that purposes please use TL wrappers (both for compiled and built-in inference)" + "Use `TLPredictor` for CPU inference with Treelite 4. The model must use ", + "`target_splitter='Single'` (the default). The wrapper supports NumPy 1.x and 2.x and keeps the original postprocessing and `(n_samples, n_outputs)` output shape." ] }, { "cell_type": "code", - "execution_count": 14, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "100%|██████████| 1000/1000 [00:05<00:00, 180.50it/s]\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 5.53 s, sys: 124 ms, total: 5.65 s\n", - "Wall time: 5.58 s\n" - ] - } - ], + "outputs": [], "source": [ "%%time\n", "tl_model = TLPredictor(model)" @@ -571,48 +556,9 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/home/jovyan/finalize_pb/rapids-24.04/lib/python3.10/site-packages/treelite/gtil/gtil.py:160: UserWarning: GTIL currently only supports float32 type; data will be converted to float32 and information might be lost.\n", - " warnings.warn(\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 3.36 s, sys: 23.4 ms, total: 3.38 s\n", - "Wall time: 900 ms\n" - ] - }, - { - "data": { - "text/plain": [ - "array([[-223.87396 , -135.50195 , -265.8046 , ..., -129.46918 ,\n", - " -205.92981 , -221.22426 ],\n", - " [-105.44449 , -102.70254 , -46.39192 , ..., -119.10744 ,\n", - " -103.06447 , -8.659145 ],\n", - " [ -44.66279 , -64.494865 , 139.2069 , ..., 14.635647 ,\n", - " -27.567156 , -215.73244 ],\n", - " ...,\n", - " [ -88.67468 , 126.860374 , 77.65633 , ..., 221.90102 ,\n", - " 27.670267 , 5.0210514],\n", - " [ -5.025807 , 140.92497 , 243.93558 , ..., 150.28708 ,\n", - " 173.07477 , 206.62967 ],\n", - " [ -9.335782 , 37.36902 , 169.7846 , ..., 94.27018 ,\n", - " 27.018427 , 8.480059 ]], dtype=float32)" - ] - }, - "execution_count": 15, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "%%time\n", "tl_model.predict(X_test, nthread=4)" @@ -627,18 +573,9 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 31.7 ms, sys: 31.6 ms, total: 63.3 ms\n", - "Wall time: 63.7 ms\n" - ] - } - ], + "outputs": [], "source": [ "%%time\n", "tl_model.dump('../data/tl_dump', rewrite=True)\n", @@ -647,40 +584,9 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 3.35 s, sys: 7.27 ms, total: 3.36 s\n", - "Wall time: 856 ms\n" - ] - }, - { - "data": { - "text/plain": [ - "array([[-223.87396 , -135.50195 , -265.8046 , ..., -129.46918 ,\n", - " -205.92981 , -221.22426 ],\n", - " [-105.44449 , -102.70254 , -46.39192 , ..., -119.10744 ,\n", - " -103.06447 , -8.659145 ],\n", - " [ -44.66279 , -64.494865 , 139.2069 , ..., 14.635647 ,\n", - " -27.567156 , -215.73244 ],\n", - " ...,\n", - " [ -88.67468 , 126.860374 , 77.65633 , ..., 221.90102 ,\n", - " 27.670267 , 5.0210514],\n", - " [ -5.025807 , 140.92497 , 243.93558 , ..., 150.28708 ,\n", - " 173.07477 , 206.62967 ],\n", - " [ -9.335782 , 37.36902 , 169.7846 , ..., 94.27018 ,\n", - " 27.018427 , 8.480059 ]], dtype=float32)" - ] - }, - "execution_count": 17, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "%%time\n", "tl_model.predict(X_test, nthread=4)" @@ -690,99 +596,30 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "For better speed up you can compile your model and make inference even more efficient" + "For faster inference, compile the model with TL2cgen. On Linux install GCC; ", + "on Windows run Jupyter from a Visual Studio developer shell with MSVC on PATH. ", + "The compiled library is specific to the target operating system and architecture." ] }, { "cell_type": "code", - "execution_count": 18, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[16:08:40] ../src/compiler/ast/split.cc:29: Parallel compilation enabled; member trees will be divided into 56 translation units.\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/home/jovyan/finalize_pb/rapids-24.04/lib/python3.10/site-packages/treelite/warn.py:6: UserWarning: Method treelite.Model.export_lib() is deprecated and scheduled for removal in Treelite 4.0. Please use TL2cgen instead. Consult the migration guide at https://tl2cgen.readthedocs.io/en/latest/treelite-migration.html.\n", - " warnings.warn(\n", - "/home/jovyan/finalize_pb/rapids-24.04/lib/python3.10/site-packages/treelite/warn.py:6: UserWarning: Method treelite.Model.compile() is deprecated and scheduled for removal in Treelite 4.0. Please use TL2cgen instead. Consult the migration guide at https://tl2cgen.readthedocs.io/en/latest/treelite-migration.html.\n", - " warnings.warn(\n", - "/home/jovyan/finalize_pb/rapids-24.04/lib/python3.10/site-packages/treelite/warn.py:6: UserWarning: Function treelite.create_shared() is deprecated and scheduled for removal in Treelite 4.0. Please use TL2cgen instead. Consult the migration guide at https://tl2cgen.readthedocs.io/en/latest/treelite-migration.html.\n", - " warnings.warn(\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 1.29 s, sys: 90.3 ms, total: 1.38 s\n", - "Wall time: 56.2 s\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/home/jovyan/finalize_pb/rapids-24.04/lib/python3.10/site-packages/treelite_runtime/warn.py:6: UserWarning: class treelite_runtime.Predictor is deprecated and scheduled for removal in Treelite 4.0. Please use TL2cgen instead. Consult the migration guide at https://tl2cgen.readthedocs.io/en/latest/treelite-migration.html.\n", - " warnings.warn(\n" - ] - } - ], + "outputs": [], "source": [ "%%time\n", - "tl_compiled = tl_model.compile('gcc', '../data/templib.so', nthread=4)" + "import sys\n", + "\n", + "toolchain = 'msvc' if sys.platform == 'win32' else 'gcc'\n", + "suffix = '.dll' if sys.platform == 'win32' else '.dylib' if sys.platform == 'darwin' else '.so'\n", + "tl_compiled = tl_model.compile(toolchain, '../data/templib' + suffix, nthread=4)" ] }, { "cell_type": "code", - "execution_count": 19, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/home/jovyan/finalize_pb/rapids-24.04/lib/python3.10/site-packages/treelite_runtime/warn.py:6: UserWarning: class treelite_runtime.DMatrix is deprecated and scheduled for removal in Treelite 4.0. Please use TL2cgen instead. Consult the migration guide at https://tl2cgen.readthedocs.io/en/latest/treelite-migration.html.\n", - " warnings.warn(\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 6.62 s, sys: 588 ms, total: 7.21 s\n", - "Wall time: 1.83 s\n" - ] - }, - { - "data": { - "text/plain": [ - "array([[-223.87384 , -135.502 , -265.80447 , ..., -129.46922 ,\n", - " -205.92975 , -221.22443 ],\n", - " [-105.44452 , -102.70257 , -46.391895 , ..., -119.1074 ,\n", - " -103.06445 , -8.659153 ],\n", - " [ -44.662804 , -64.4949 , 139.20688 , ..., 14.63566 ,\n", - " -27.567184 , -215.73254 ],\n", - " ...,\n", - " [ -88.67466 , 126.86047 , 77.65633 , ..., 221.90108 ,\n", - " 27.670267 , 5.021043 ],\n", - " [ -5.0258074, 140.92477 , 243.93553 , ..., 150.28697 ,\n", - " 173.0748 , 206.62964 ],\n", - " [ -9.335763 , 37.36898 , 169.78474 , ..., 94.270195 ,\n", - " 27.01838 , 8.4801 ]], dtype=float32)" - ] - }, - "execution_count": 19, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "%%time\n", "tl_compiled.predict(X_test)" @@ -792,23 +629,17 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Compiled model could also be saved, but note - instanse will save only the metadata. After loading, it expect to find compiled library path stay the same. Otherwise, you can update the library path via .set_libpath method" + "A compiled predictor dump stores metadata and the library path. Deploy the ", + "compiled library alongside the dump. If its path changes, use ", + "`TLCompiledPredictor.load(filename, libpath=new_path)`. Use `.set_libpath(new_path)` ", + "to change the library path of an already loaded predictor." ] }, { "cell_type": "code", - "execution_count": 20, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 441 ms, sys: 1.27 s, total: 1.71 s\n", - "Wall time: 732 ms\n" - ] - } - ], + "outputs": [], "source": [ "%%time\n", "tl_compiled.dump('../data/tl_compiled.pkl')\n", @@ -819,40 +650,9 @@ }, { "cell_type": "code", - "execution_count": 21, + "execution_count": null, "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 6.5 s, sys: 3.48 ms, total: 6.5 s\n", - "Wall time: 6.49 s\n" - ] - }, - { - "data": { - "text/plain": [ - "array([[-223.87384 , -135.502 , -265.80447 , ..., -129.46922 ,\n", - " -205.92975 , -221.22443 ],\n", - " [-105.44452 , -102.70257 , -46.391895 , ..., -119.1074 ,\n", - " -103.06445 , -8.659153 ],\n", - " [ -44.662804 , -64.4949 , 139.20688 , ..., 14.63566 ,\n", - " -27.567184 , -215.73254 ],\n", - " ...,\n", - " [ -88.67466 , 126.86047 , 77.65633 , ..., 221.90108 ,\n", - " 27.670267 , 5.021043 ],\n", - " [ -5.0258074, 140.92477 , 243.93553 , ..., 150.28697 ,\n", - " 173.0748 , 206.62964 ],\n", - " [ -9.335763 , 37.36898 , 169.78474 , ..., 94.270195 ,\n", - " 27.01838 , 8.4801 ]], dtype=float32)" - ] - }, - "execution_count": 21, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "%%time\n", "tl_compiled.predict(X_test)" From cddd9f613183b8f4e99df97fea16710f51d96a7e Mon Sep 17 00:00:00 2001 From: Rinchin Date: Fri, 30 Oct 2026 15:20:03 +0300 Subject: [PATCH 2/3] Fix Windows DLL imports and Python 3.8 CI runner --- .github/workflows/python-compatibility.yml | 5 +- py_boost/utils/tl2cgen_compat.py | 37 +++++++++ py_boost/utils/tl_wrapper.py | 6 +- tests/test_tl2cgen_import.py | 97 ++++++++++++++++++++++ 4 files changed, 141 insertions(+), 4 deletions(-) create mode 100644 py_boost/utils/tl2cgen_compat.py create mode 100644 tests/test_tl2cgen_import.py diff --git a/.github/workflows/python-compatibility.yml b/.github/workflows/python-compatibility.yml index cdbffac..4c12db2 100644 --- a/.github/workflows/python-compatibility.yml +++ b/.github/workflows/python-compatibility.yml @@ -15,7 +15,7 @@ jobs: fail-fast: false matrix: include: - - os: ubuntu-latest + - os: ubuntu-24.04 python-version: "3.8" treelite: "treelite==4.1.2" - os: ubuntu-latest @@ -33,6 +33,9 @@ jobs: - os: windows-latest python-version: "3.12" treelite: "treelite>=4.1.2,<5" + - os: windows-latest + python-version: "3.12" + treelite: "treelite==4.1.2" - os: windows-latest python-version: "3.14" treelite: "treelite>=4.1.2,<5" diff --git a/py_boost/utils/tl2cgen_compat.py b/py_boost/utils/tl2cgen_compat.py new file mode 100644 index 0000000..0308bd3 --- /dev/null +++ b/py_boost/utils/tl2cgen_compat.py @@ -0,0 +1,37 @@ +"""Import TL2cgen on Windows Python installations without a Conda DLL directory.""" + +import os +import sys +from importlib import import_module + + +def _import_tl2cgen(): + conda_bins = { + os.path.normcase(os.path.join(os.path.normpath(prefix), 'Library', 'bin')) + for prefix in (sys.prefix, sys.base_prefix) + } + if sys.platform != 'win32' or all(os.path.isdir(path) for path in conda_bins): + return import_module('tl2cgen') + + # TL2cgen 1.0 and older Treelite 4 versions register optional Conda directories. + # CPython and venv installations do not have it, so import otherwise fails + # before the packaged DLL can be loaded. Ignore only that missing directory + # during import; all real DLL-loading errors and other paths still propagate. + add_dll_directory = os.add_dll_directory + + def add_directory(path): + if (os.path.normcase(os.path.normpath(path)) in conda_bins + and not os.path.isdir(path)): + return None + return add_dll_directory(path) + + try: + os.add_dll_directory = add_directory + return import_module('tl2cgen') + finally: + os.add_dll_directory = add_dll_directory + + +tl2cgen = _import_tl2cgen() +# TL2cgen imports Treelite inside the guarded scope above. +treelite = import_module('treelite') diff --git a/py_boost/utils/tl_wrapper.py b/py_boost/utils/tl_wrapper.py index b224d4e..ef7d90a 100644 --- a/py_boost/utils/tl_wrapper.py +++ b/py_boost/utils/tl_wrapper.py @@ -4,11 +4,11 @@ import joblib import numpy as np -import tl2cgen -import treelite -from treelite.model_builder import Metadata, ModelBuilder, PostProcessorFunc, TreeAnnotation from tqdm import tqdm +from .tl2cgen_compat import tl2cgen, treelite +from treelite.model_builder import Metadata, ModelBuilder, PostProcessorFunc, TreeAnnotation + def parse_pb_tree(tree, builder): """Append a single-group Py-Boost tree to a Treelite 4 model builder.""" diff --git a/tests/test_tl2cgen_import.py b/tests/test_tl2cgen_import.py new file mode 100644 index 0000000..0859e9e --- /dev/null +++ b/tests/test_tl2cgen_import.py @@ -0,0 +1,97 @@ +"""Regression coverage for the TL2cgen 1.0 Windows DLL-directory import failure.""" + +import os +import subprocess +import sys + +import pytest + +from py_boost.utils import tl2cgen_compat + + +def test_windows_import_ignores_only_missing_conda_directory(tmp_path, monkeypatch): + conda_bin = os.path.join(str(tmp_path), 'Library', 'bin') + base_prefix = str(tmp_path / 'base_python') + base_conda_bin = os.path.join(base_prefix, 'Library', 'bin') + other_missing = os.path.join(str(tmp_path), 'unrelated') + calls = [] + imported = object() + + def add_directory(path): + calls.append(path) + raise FileNotFoundError(path) + + def import_module(name): + assert name == 'tl2cgen' + os.add_dll_directory(conda_bin) + os.add_dll_directory(base_conda_bin) + with pytest.raises(FileNotFoundError, match='unrelated'): + os.add_dll_directory(other_missing) + return imported + + with monkeypatch.context() as patch: + patch.setattr(sys, 'platform', 'win32') + patch.setattr(sys, 'prefix', str(tmp_path)) + patch.setattr(sys, 'base_prefix', base_prefix) + patch.setattr(os, 'add_dll_directory', add_directory, raising=False) + patch.setattr(tl2cgen_compat, 'import_module', import_module) + assert tl2cgen_compat._import_tl2cgen() is imported + assert os.add_dll_directory is add_directory + assert calls == [other_missing] + assert not os.path.exists(conda_bin) + assert not os.path.exists(base_conda_bin) + + +def test_windows_import_restores_dll_search_function_on_failure(tmp_path, monkeypatch): + def add_directory(path): + raise AssertionError('An unrelated DLL-directory call must not be made') + + def import_module(name): + raise OSError('DLL dependency is missing') + + with monkeypatch.context() as patch: + patch.setattr(sys, 'platform', 'win32') + patch.setattr(sys, 'prefix', str(tmp_path)) + patch.setattr(os, 'add_dll_directory', add_directory, raising=False) + patch.setattr(tl2cgen_compat, 'import_module', import_module) + with pytest.raises(OSError, match='DLL dependency is missing'): + tl2cgen_compat._import_tl2cgen() + assert os.add_dll_directory is add_directory + + +def test_windows_conda_import_keeps_existing_dll_directory(tmp_path, monkeypatch): + conda_bin = tmp_path / 'Library' / 'bin' + conda_bin.mkdir(parents=True) + calls = [] + + def add_directory(path): + calls.append(path) + + def import_module(name): + assert os.add_dll_directory is add_directory + os.add_dll_directory(str(conda_bin)) + + with monkeypatch.context() as patch: + patch.setattr(sys, 'platform', 'win32') + patch.setattr(sys, 'prefix', str(tmp_path)) + patch.setattr(sys, 'base_prefix', str(tmp_path)) + patch.setattr(os, 'add_dll_directory', add_directory, raising=False) + patch.setattr(tl2cgen_compat, 'import_module', import_module) + tl2cgen_compat._import_tl2cgen() + assert calls == [str(conda_bin)] + + +@pytest.mark.skipif(sys.platform != 'win32', reason='Requires the actual Windows DLL loader') +def test_windows_fresh_import_preserves_dll_loader(): + subprocess.run([sys.executable, '-c', ''' +import os +import sys +original = os.add_dll_directory +conda_bin = os.path.join(sys.prefix, 'Library', 'bin') +existed = os.path.exists(conda_bin) +import py_boost +import tl2cgen +assert os.add_dll_directory is original +assert os.path.exists(conda_bin) == existed +print(tl2cgen.__version__) +'''], check=True) From 3d9c043213d397e439f837f9d5b6e03baa222435 Mon Sep 17 00:00:00 2001 From: Rinchin Date: Fri, 30 Oct 2026 15:46:34 +0300 Subject: [PATCH 3/3] Preserve predictor copies and sparse inputs; isolate Windows DLL loading --- README.md | 6 ++ py_boost/utils/tl2cgen_compat.py | 74 +++++++++++++--- py_boost/utils/tl_wrapper.py | 47 +++++++--- tests/test_tl2cgen_import.py | 146 ++++++++++++++++++++----------- tests/test_treelite.py | 69 +++++++++++++++ 5 files changed, 264 insertions(+), 78 deletions(-) diff --git a/README.md b/README.md index 00ce108..8d92d8b 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,8 @@ Predictions have shape `(n_samples, n_outputs)`, including single-output models. The underlying Treelite model stores raw scores as regression targets; the wrapper applies `model.postprocess_fn` once, after restoring this output shape. A custom CPU postprocessor can be supplied with `postprocess_fn=...`. +Both predictor wrappers support shallow copying with `copy.copy()`; saving a +predictor leaves the original and its copies usable. Compile the model with TL2cgen for faster CPU inference. Compilation requires a C compiler; prediction with the compiled library does not require a compiler: @@ -96,6 +98,10 @@ compiled.dump('compiled.pkl') restored_compiled = TLCompiledPredictor.load('compiled.pkl') ``` +`TLCompiledPredictor.predict()` also accepts SciPy sparse matrices (CSR, CSC +and COO) without converting them to dense arrays. Unstored entries and stored +NaNs are treated as missing values; explicitly stored zeros remain zero. + Compiled predictor dumps contain metadata and the library path. Deploy the `.so`/`.dll` alongside the dump. If its path changes, load it with `TLCompiledPredictor.load('compiled.pkl', libpath=new_path)`. Use diff --git a/py_boost/utils/tl2cgen_compat.py b/py_boost/utils/tl2cgen_compat.py index 0308bd3..60d9c29 100644 --- a/py_boost/utils/tl2cgen_compat.py +++ b/py_boost/utils/tl2cgen_compat.py @@ -1,8 +1,65 @@ """Import TL2cgen on Windows Python installations without a Conda DLL directory.""" +import builtins import os import sys from importlib import import_module +from importlib.abc import Loader, MetaPathFinder +from importlib.machinery import PathFinder + + +class _DllDirectoryOs: + """An os view used only inside the dependencies' native-library loaders.""" + + def __init__(self, conda_bins): + self._conda_bins = conda_bins + + def __getattr__(self, name): + return getattr(os, name) + + def add_dll_directory(self, path): + if (os.path.normcase(os.path.normpath(path)) in self._conda_bins + and not os.path.isdir(path)): + return None + return os.add_dll_directory(path) + + +class _DllDirectoryLoader(Loader): + def __init__(self, loader, conda_bins): + self._loader = loader + self._os = _DllDirectoryOs(conda_bins) + + def create_module(self, spec): + return self._loader.create_module(spec) + + def exec_module(self, module): + original_import = builtins.__import__ + + def module_import(name, globals=None, locals=None, fromlist=(), level=0): + if name == 'os' and level == 0: + return self._os + return original_import(name, globals, locals, fromlist, level) + + # Give only this module a guarded os import. The process-wide os module + # and builtins stay unchanged, including during concurrent imports. + module.__dict__['__builtins__'] = dict(vars(builtins), __import__=module_import) + self._loader.exec_module(module) + + def __getattr__(self, name): + return getattr(self._loader, name) + + +class _DllDirectoryFinder(MetaPathFinder): + def __init__(self, conda_bins): + self._conda_bins = conda_bins + + def find_spec(self, fullname, path=None, target=None): + if fullname not in ('treelite.core', 'tl2cgen.libloader'): + return None + spec = PathFinder.find_spec(fullname, path, target) + if spec is not None and spec.loader is not None: + spec.loader = _DllDirectoryLoader(spec.loader, self._conda_bins) + return spec def _import_tl2cgen(): @@ -14,22 +71,15 @@ def _import_tl2cgen(): return import_module('tl2cgen') # TL2cgen 1.0 and older Treelite 4 versions register optional Conda directories. - # CPython and venv installations do not have it, so import otherwise fails + # CPython and venv installations do not have them, so import otherwise fails # before the packaged DLL can be loaded. Ignore only that missing directory - # during import; all real DLL-loading errors and other paths still propagate. - add_dll_directory = os.add_dll_directory - - def add_directory(path): - if (os.path.normcase(os.path.normpath(path)) in conda_bins - and not os.path.isdir(path)): - return None - return add_dll_directory(path) - + # inside those two modules; real DLL-loading errors still propagate. + finder = _DllDirectoryFinder(conda_bins) + sys.meta_path.insert(0, finder) try: - os.add_dll_directory = add_directory return import_module('tl2cgen') finally: - os.add_dll_directory = add_dll_directory + sys.meta_path.remove(finder) tl2cgen = _import_tl2cgen() diff --git a/py_boost/utils/tl_wrapper.py b/py_boost/utils/tl_wrapper.py index ef7d90a..b34e906 100644 --- a/py_boost/utils/tl_wrapper.py +++ b/py_boost/utils/tl_wrapper.py @@ -1,9 +1,11 @@ import os from collections import deque +from copy import copy from itertools import count import joblib import numpy as np +from scipy.sparse import csr_matrix, issparse from tqdm import tqdm from .tl2cgen_compat import tl2cgen, treelite @@ -91,6 +93,29 @@ def _reshape_prediction(pred): return pred.reshape(pred.shape[0], pred.shape[1] * pred.shape[2]) +def _prepare_tl2cgen_input(X): + if not issparse(X): + return np.asarray(X, dtype=np.float32, order="C") + # The new sparse object may share buffers with X; replace rather than mutate + # those buffers. TL2cgen treats unstored entries as missing, but does not + # recognize NaNs stored explicitly in a CSR buffer. Drop only NaNs, not zeros. + matrix = csr_matrix(X, dtype=np.float32, copy=False) + missing = np.isnan(matrix.data) + if missing.any(): + missing_rows = np.searchsorted(matrix.indptr, np.flatnonzero(missing), side='right') - 1 + counts = np.bincount(missing_rows, minlength=matrix.shape[0]) + indptr = np.array(matrix.indptr, dtype=np.uintp, copy=True) + indptr[1:] -= np.cumsum(counts, dtype=np.uintp) + matrix.indptr = indptr + matrix.data = matrix.data[~missing] + matrix.indices = matrix.indices[~missing] + # Match the C API dtypes before TL2cgen's NumPy 2 copy=False conversions. + matrix.data = np.ascontiguousarray(matrix.data) + matrix.indices = np.ascontiguousarray(matrix.indices, dtype=np.uintc) + matrix.indptr = np.ascontiguousarray(matrix.indptr, dtype=np.uintp) + return matrix + + class TLCompiledPredictor: """ Compiled TL2cgen model for CPU inference @@ -122,12 +147,12 @@ def predict(self, X): """Make prediction Args: - X: np.ndarray + X: np.ndarray or scipy sparse matrix; absent sparse entries are missing values Returns: np.ndarray """ - X = np.asarray(X, dtype=np.float32, order="C") + X = _prepare_tl2cgen_input(X) pred = self.predictor.predict(tl2cgen.DMatrix(X)) return self.postprocess_fn(_reshape_prediction(pred)) @@ -149,11 +174,6 @@ def set_libpath(self, libpath=None, nthread=None): self.nthread = nthread self.predictor = tl2cgen.Predictor(self.libpath, nthread=nthread, verbose=self.verbose) - def __getstate__(self): - state = self.__dict__.copy() - state["predictor"] = None - return state - def dump(self, filename): """Dump instance @@ -163,7 +183,9 @@ def dump(self, filename): Returns: """ - joblib.dump(self, filename) + metadata = copy(self) + metadata.predictor = None + joblib.dump(metadata, filename) @staticmethod def load(filename, libpath=None, nthread=None): @@ -277,11 +299,6 @@ def predict(self, X, nthread=None): pred = treelite.gtil.predict(self.tl_model, X, nthread=nthread) return self.postprocess_fn(_reshape_prediction(pred)) - def __getstate__(self): - state = self.__dict__.copy() - state["tl_model"] = None - return state - def dump(self, dirname, rewrite=False): """Dump treelite Model and predictor instance @@ -294,7 +311,9 @@ def dump(self, dirname, rewrite=False): """ os.makedirs(dirname, exist_ok=rewrite) self.tl_model.serialize(os.path.join(dirname, 'model.mod')) - joblib.dump(self, os.path.join(dirname, 'predictor.pkl')) + metadata = copy(self) + metadata.tl_model = None + joblib.dump(metadata, os.path.join(dirname, 'predictor.pkl')) @staticmethod def load(dirname): diff --git a/tests/test_tl2cgen_import.py b/tests/test_tl2cgen_import.py index 0859e9e..7dafa10 100644 --- a/tests/test_tl2cgen_import.py +++ b/tests/test_tl2cgen_import.py @@ -1,97 +1,139 @@ -"""Regression coverage for the TL2cgen 1.0 Windows DLL-directory import failure.""" +"""Windows loader compatibility must leave process-wide import and os APIs intact.""" +import builtins import os import subprocess import sys +from importlib.util import module_from_spec import pytest from py_boost.utils import tl2cgen_compat -def test_windows_import_ignores_only_missing_conda_directory(tmp_path, monkeypatch): - conda_bin = os.path.join(str(tmp_path), 'Library', 'bin') - base_prefix = str(tmp_path / 'base_python') - base_conda_bin = os.path.join(base_prefix, 'Library', 'bin') - other_missing = os.path.join(str(tmp_path), 'unrelated') +@pytest.mark.parametrize('module_name,filename', [ + ('treelite.core', 'core.py'), ('tl2cgen.libloader', 'libloader.py'), +]) +def test_loader_uses_private_os_view(tmp_path, monkeypatch, module_name, filename): + conda_bin = str(tmp_path / 'Library' / 'bin') + other_missing = str(tmp_path / 'unrelated') + existing = str(tmp_path) + cookie = object() calls = [] - imported = object() def add_directory(path): calls.append(path) - raise FileNotFoundError(path) + if not os.path.isdir(path): + raise FileNotFoundError(path) + return cookie - def import_module(name): - assert name == 'tl2cgen' - os.add_dll_directory(conda_bin) - os.add_dll_directory(base_conda_bin) - with pytest.raises(FileNotFoundError, match='unrelated'): - os.add_dll_directory(other_missing) - return imported - - with monkeypatch.context() as patch: - patch.setattr(sys, 'platform', 'win32') - patch.setattr(sys, 'prefix', str(tmp_path)) - patch.setattr(sys, 'base_prefix', base_prefix) - patch.setattr(os, 'add_dll_directory', add_directory, raising=False) - patch.setattr(tl2cgen_compat, 'import_module', import_module) - assert tl2cgen_compat._import_tl2cgen() is imported + # This function belongs to the test module, outside the guarded dependency. + # It runs on another thread while the dependency's import is still executing. + def other_thread(): assert os.add_dll_directory is add_directory - assert calls == [other_missing] + with pytest.raises(FileNotFoundError): + os.add_dll_directory(conda_bin) + return True + + source = tmp_path / filename + source.write_text(""" +import os +from concurrent.futures import ThreadPoolExecutor +from importlib import import_module +private_os = os +missing_result = os.add_dll_directory(conda_bin) +existing_result = os.add_dll_directory(existing) +try: + os.add_dll_directory(other_missing) +except FileNotFoundError: + unrelated_error = True +assert import_module('os').add_dll_directory is original +with ThreadPoolExecutor(max_workers=1) as pool: + thread_ok = pool.submit(other_thread).result(timeout=10) +""") + monkeypatch.setattr(os, 'add_dll_directory', add_directory, raising=False) + finder = tl2cgen_compat._DllDirectoryFinder({os.path.normcase(conda_bin)}) + spec = finder.find_spec(module_name, [str(tmp_path)]) + module = module_from_spec(spec) + module.__dict__.update(conda_bin=conda_bin, other_missing=other_missing, + existing=existing, original=add_directory, other_thread=other_thread) + original_import = builtins.__import__ + spec.loader.exec_module(module) + assert module.private_os is not os + assert module.private_os.path is os.path + assert module.missing_result is None + assert module.existing_result is cookie + assert module.unrelated_error + assert module.thread_ok + assert os.add_dll_directory is add_directory + assert builtins.__import__ is original_import + assert calls == [existing, other_missing, conda_bin] assert not os.path.exists(conda_bin) - assert not os.path.exists(base_conda_bin) + assert finder.find_spec('unrelated.module', [str(tmp_path)]) is None -def test_windows_import_restores_dll_search_function_on_failure(tmp_path, monkeypatch): - def add_directory(path): - raise AssertionError('An unrelated DLL-directory call must not be made') +@pytest.mark.parametrize('fail', [False, True]) +def test_windows_import_removes_finder(tmp_path, monkeypatch, fail): + before = list(sys.meta_path) + original = getattr(os, 'add_dll_directory', None) + imported = object() def import_module(name): - raise OSError('DLL dependency is missing') + assert name == 'tl2cgen' + assert getattr(os, 'add_dll_directory', None) is original + assert any(isinstance(finder, tl2cgen_compat._DllDirectoryFinder) + for finder in sys.meta_path) + if fail: + raise OSError('DLL dependency is missing') + return imported with monkeypatch.context() as patch: patch.setattr(sys, 'platform', 'win32') patch.setattr(sys, 'prefix', str(tmp_path)) - patch.setattr(os, 'add_dll_directory', add_directory, raising=False) + patch.setattr(sys, 'base_prefix', str(tmp_path / 'base_python')) patch.setattr(tl2cgen_compat, 'import_module', import_module) - with pytest.raises(OSError, match='DLL dependency is missing'): - tl2cgen_compat._import_tl2cgen() - assert os.add_dll_directory is add_directory - - -def test_windows_conda_import_keeps_existing_dll_directory(tmp_path, monkeypatch): - conda_bin = tmp_path / 'Library' / 'bin' - conda_bin.mkdir(parents=True) - calls = [] - - def add_directory(path): - calls.append(path) + if fail: + with pytest.raises(OSError, match='DLL dependency is missing'): + tl2cgen_compat._import_tl2cgen() + else: + assert tl2cgen_compat._import_tl2cgen() is imported + assert sys.meta_path == before + assert getattr(os, 'add_dll_directory', None) is original + + +def test_windows_conda_import_uses_normal_loader(tmp_path, monkeypatch): + (tmp_path / 'Library' / 'bin').mkdir(parents=True) + before = list(sys.meta_path) + imported = object() def import_module(name): - assert os.add_dll_directory is add_directory - os.add_dll_directory(str(conda_bin)) + assert sys.meta_path == before + return imported with monkeypatch.context() as patch: patch.setattr(sys, 'platform', 'win32') patch.setattr(sys, 'prefix', str(tmp_path)) patch.setattr(sys, 'base_prefix', str(tmp_path)) - patch.setattr(os, 'add_dll_directory', add_directory, raising=False) patch.setattr(tl2cgen_compat, 'import_module', import_module) - tl2cgen_compat._import_tl2cgen() - assert calls == [str(conda_bin)] + assert tl2cgen_compat._import_tl2cgen() is imported @pytest.mark.skipif(sys.platform != 'win32', reason='Requires the actual Windows DLL loader') def test_windows_fresh_import_preserves_dll_loader(): - subprocess.run([sys.executable, '-c', ''' + subprocess.run([sys.executable, '-c', """ +import builtins import os import sys original = os.add_dll_directory -conda_bin = os.path.join(sys.prefix, 'Library', 'bin') -existed = os.path.exists(conda_bin) +original_import = builtins.__import__ +conda_bins = {os.path.join(prefix, 'Library', 'bin') for prefix in (sys.prefix, sys.base_prefix)} +existed = {path: os.path.exists(path) for path in conda_bins} import py_boost import tl2cgen assert os.add_dll_directory is original -assert os.path.exists(conda_bin) == existed +assert builtins.__import__ is original_import +assert all(type(finder).__module__ != 'py_boost.utils.tl2cgen_compat' + for finder in sys.meta_path) +assert {path: os.path.exists(path) for path in conda_bins} == existed print(tl2cgen.__version__) -'''], check=True) +"""], check=True) diff --git a/tests/test_treelite.py b/tests/test_treelite.py index 3212e20..58ba8ed 100644 --- a/tests/test_treelite.py +++ b/tests/test_treelite.py @@ -1,5 +1,6 @@ """CPU integration tests using real Treelite models and compiled TL2cgen libraries.""" +import copy import os import shutil import sys @@ -8,6 +9,7 @@ import joblib import numpy as np import pytest +from scipy import sparse from py_boost import TLCompiledPredictor, TLPredictor @@ -97,6 +99,73 @@ def compiled(ensemble, tmp_path_factory, toolchain): ) +@pytest.fixture(scope='module') +def sparse_compiled(ensemble, tmp_path_factory, toolchain): + model = CPUEnsemble(len(ensemble.base_score)) + # Explicit zero goes left; a missing sparse entry goes right at this node. + model.models[0].test_format[5] = .25 + name, suffix = toolchain + path = tmp_path_factory.mktemp('sparse_compiled') / ('model' + suffix) + return TLPredictor(model).compile(name, path, params={'parallel_comp': 1}, nthread=2) + + +@pytest.mark.parametrize('format', ['csr', 'csc', 'coo']) +@pytest.mark.parametrize('dtype', [np.float32, np.float64]) +@pytest.mark.parametrize('index_dtype', [np.int32, np.int64]) +def test_compiled_sparse_inputs(ensemble, sparse_compiled, format, dtype, index_dtype, + monkeypatch): + # Keep explicit zeros. Unstored entries represent missing values, not zeros. + X = sparse.csr_matrix((np.array([0, 0, 1, -1, 0, -1, 0, np.nan, np.nan, -1], dtype=dtype), + [1, 0, 0, 1, 0, 1, 0, 1, 0, 1], + [0, 1, 2, 3, 4, 4, 6, 8, 10]), shape=(8, 2)).asformat(format) + if format == 'coo': + X.row = X.row.astype(index_dtype) + X.col = X.col.astype(index_dtype) + else: + X.indices = X.indices.astype(index_dtype) + X.indptr = X.indptr.astype(index_dtype) + original = copy.deepcopy(X) + + def no_dense(*args, **kwargs): + pytest.fail('Sparse prediction must not densify its input') + + for cls in (sparse.csr_matrix, sparse.csc_matrix, sparse.coo_matrix): + monkeypatch.setattr(cls, 'toarray', no_dense) + + expected = np.array([ + [2, -4, 5.5], [4, 2, 4], [-4, 4, 7], [2, -4, 5.5], + [4, 2, 4], [2, -4, 5.5], [4, 2, 4], [2, -4, 5.5], + ], dtype=np.float32)[:, :len(ensemble.base_score)] + actual = sparse_compiled.predict(X) + assert actual.dtype == np.float32 + np.testing.assert_allclose(actual, expected) + np.testing.assert_array_equal(X.data, original.data) + fields = ('row', 'col') if format == 'coo' else ('indices', 'indptr') + for field in fields: + actual_buffer, original_buffer = getattr(X, field), getattr(original, field) + assert actual_buffer.dtype == original_buffer.dtype + np.testing.assert_array_equal(actual_buffer, original_buffer) + + +def test_shallow_copy_keeps_native_predictors(ensemble, compiled, examples): + X, expected = examples + expected = expected[:, :len(ensemble.base_score)] + for predictor in (TLPredictor(ensemble), compiled): + cloned = copy.copy(predictor) + cloned.postprocess_fn = sigmoid + np.testing.assert_allclose(cloned.predict(X), sigmoid(expected)) + np.testing.assert_allclose(predictor.predict(X), expected) + + +@pytest.mark.parametrize('batch_size', [0, 1]) +def test_sparse_batch_keeps_output_shape(ensemble, sparse_compiled, batch_size): + X = sparse.csr_matrix((np.zeros(2, dtype=np.float32), [0, 1], [0, 2]), shape=(1, 2)) + actual = sparse_compiled.predict(X[:batch_size]) + expected = np.array([[2, -4, 5.5]], dtype=np.float32)[:batch_size, :len(ensemble.base_score)] + assert actual.shape == expected.shape + np.testing.assert_allclose(actual, expected) + + @pytest.mark.parametrize('layout', ['float32', 'float64', 'fortran', 'strided', 'list']) def test_gtil_and_compiled_predictions(ensemble, compiled, examples, layout): X, expected = examples