Skip to content

Migrate CPU inference to Treelite 4 and TL2cgen - #37

Merged
dev-rinchin merged 3 commits into
masterfrom
fix/treelite4-tl2cgen
Nov 2, 2026
Merged

dev-rinchin merged 3 commits into
masterfrom
fix/treelite4-tl2cgen

Conversation

@dev-rinchin

@dev-rinchin dev-rinchin commented Oct 30, 2026 •

Copy link
Copy Markdown
Contributor

Py-Boost's CPU inference wrapper uses the Treelite 3 builder and compilation APIs, which prevents upgrading to Treelite 4 and exposes NumPy 2 input-conversion failures. This change migrates both interpreted and compiled CPU inference to Treelite 4 and TL2cgen while preserving Py-Boost's prediction shape and postprocessing.

Fixes #26.

Changes

  • Require treelite>=4.1.2,<5 and tl2cgen>=1,<2; remove the treelite_runtime dependency.
  • Build trees with the Treelite 4 model builder, using scalar leaves for one output and vector leaves for multiple outputs. Preserve split equality, missing-value directions and constant trees.
  • Represent raw scores as regression targets and apply the existing Python postprocessor once, after converting Treelite's 3D result to (n_samples, n_outputs). This preserves regression, binary/multilabel classification, multiclass classification and custom postprocessors.
  • Store the initial score as a constant first tree. TL2cgen 1.0 incorrectly indexes vector regression base_scores by class alone, applying the first target's intercept to all targets. A constant tree avoids that error and matches Py-Boost's accumulation order.
  • Support compiled prediction on CSR, CSC and COO sparse matrices without densifying or mutating input buffers. Normalize index dtypes for NumPy 2, preserve explicit zeros and treat both stored NaNs and unstored entries as missing.
  • Use tl2cgen.export_lib, tl2cgen.Predictor and tl2cgen.DMatrix. Normalize dense inputs with np.asarray(..., dtype=np.float32, order="C") for NumPy 1/2 compatibility.
  • Keep the public compile signature; reject unsupported legacy compiler names explicitly. Let GTIL select its default worker count within OpenMP's limit.
  • Keep live predictors and their copy.copy() copies usable after successful or failed dumps by stripping native handles only from a separate metadata copy inside dump(). Allow overriding the compiled library path and worker count at load time for deployment to another machine.
  • Make MSLE postprocessing return NumPy results for CPU inputs, including environments without CuPy.
  • Add CPU integration tests, GPU-to-CPU numerical parity tests, Linux/Windows CI with GCC/MSVC, and updated README/tutorial examples. Pin the Python 3.8 job to Ubuntu 24.04, which has a compatible hosted Python build.
  • Handle TL2cgen 1.0 and older Treelite 4 loaders registering nonexistent optional Conda DLL directories on Windows. A scoped import finder gives only treelite.core and tl2cgen.libloader a private view of os that skips those missing directories. The real os.add_dll_directory and builtins.__import__ stay unchanged, including on unrelated threads; unrelated loading errors still propagate. The finder is removed after either success or failure.

Validation

Built both the source distribution and wheel. GitHub Actions installs the wheel and runs the full test suite outside the source checkout. All eight jobs passed on commit 3d9c043: successful CI run.

Runner Python Treelite Result
Ubuntu 24.04 3.8 4.1.2 76 passed, 206 skipped
Ubuntu latest 3.12 latest 4.x 76 passed, 206 skipped
Ubuntu latest 3.12 4.1.2 76 passed, 206 skipped
Ubuntu latest 3.13 latest 4.x 76 passed, 206 skipped
Ubuntu latest 3.14 latest 4.x 76 passed, 206 skipped
Windows latest / MSVC 3.12 latest 4.x 77 passed, 205 skipped
Windows latest / MSVC 3.12 4.1.2 77 passed, 205 skipped
Windows latest / MSVC 3.14 latest 4.x 77 passed, 205 skipped

Skips cover GPU-only and platform-specific tests. The matrix covers NumPy 1.x on Python 3.8 and NumPy 2.x on modern Python versions. Locally, the full Linux GPU suite passed on NVIDIA A100 with Python 3.12, NumPy 2.5.4, Treelite 4.8.1 and TL2cgen 1.0.0: 130 passed, 1 Windows-only test skipped, using the installed wheel.

Checks cover actual compiled libraries, scalar/vector regression, MSLE, binary/multilabel/multiclass probabilities, distinct target intercepts, NaN routing, threshold equality, constant trees, float64/Fortran/strided/list inputs, empty and one-row batches, save/load, library relocation, failed dumps, shallow copies, sparse buffers and CPU-only built-in postprocessing. GPU prediction comparisons use rtol=2e-5, atol=2e-6. Windows tests also verify that a fresh import leaves the DLL loader and filesystem unchanged. Loader regression tests execute both dependency modules through the adapter and verify that an unrelated thread still sees the original os.add_dll_directory behavior.

uv pip check passed in the five local Python/NumPy/Treelite environments. actionlint, Ruff's syntax/undefined-name checks and git diff --check passed. Existing CuPy scatter_add deprecation warnings remain.

Compatibility

  • The existing target_splitter='Single' restriction remains. This PR does not add support for multiple tree groups.
  • Treelite 3 checkpoints and old treelite_runtime compiled libraries must be re-exported from the original Py-Boost model.
  • Compiled libraries remain specific to their target OS and architecture.

@dev-rinchin
dev-rinchin merged commit 3e27c14 into master Nov 2, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

TLPredictor` / Treelite integration is broken on Windows with Python 3.12 + NumPy 2.x

1 participant