diff --git a/bindings/python/README.md b/bindings/python/README.md index cfc634b..decce38 100644 --- a/bindings/python/README.md +++ b/bindings/python/README.md @@ -23,23 +23,28 @@ document parsing. ## 📦 Installation -Install the package with the desired ONNX Runtime backend. +Install with the matching extra for your backend. Published wheels target +**Linux x86_64** and **macOS arm64** (macOS Intel is not published: ONNX +Runtime 1.28 has no compatible wheel there). -### CPU +### CPU (default) ```bash pip install "fast-lightonocr[cpu]" ``` +CPU wheels bundle ONNX Runtime. No extra environment setup is required. + ### CUDA ```bash pip install "fast-lightonocr[cuda]" ``` -CUDA wheels/builds enable the native `cuda` Cargo feature. Use a CUDA-enabled -ONNX Runtime and compatible CUDA 13 / cuDNN stack. Select the provider at load -time: +Requires a CUDA-enabled package build and a compatible NVIDIA driver. The +`cuda` extra pulls in `onnxruntime-gpu` (CUDA 13 / cuDNN) and `nvidia-cublas`. + +Select CUDA at load time: ```python model = LightOnOCR.from_pretrained( @@ -51,27 +56,39 @@ model = LightOnOCR.from_pretrained( ) ``` -Default published wheels remain CPU-oriented; the CUDA extra is a dedicated -build profile. On CUDA, decoder KV past/present stay on device during generate; -token sampling still runs on the host. +When `execution_provider="cuda"`, `from_pretrained` preloads the pip NVIDIA +CUDA/cuDNN libraries (`onnxruntime.preload_dlls`), so `LD_LIBRARY_PATH` is +usually unnecessary. CPU loads never take that path. -Prebuilt wheels are currently published for Linux x86_64 and macOS arm64. These -wheels bundle the required ONNX Runtime shared library, so no additional runtime -installation or environment configuration is required for the CPU profile. +### Building from source -macOS x86_64 (Intel) wheels are not published because ONNX Runtime 1.28 does -not provide a compatible Python wheel for that platform. +Source installs use the project build backend. It discovers ONNX Runtime from +`ORT_DYLIB_PATH` when set, otherwise from the profile’s Python ORT package, +validates ONNX Runtime 1.28.x (C API level 27), and bundles the native runtime +into the wheel. -### Building from source +#### CPU + +```bash +pip install -v ".[cpu]" +``` + +# or explicitly: + +```bash +BUILD_PROFILE=cpu pip install -v ".[cpu]" +``` -When installing from source, the build backend automatically discovers a -compatible ONNX Runtime for the selected build profile. +#### CUDA -If `ORT_DYLIB_PATH` is set, it is used directly. Otherwise, the build backend -installs the appropriate ONNX Runtime build dependency into the isolated build -environment, validates compatibility with ONNX Runtime 1.28.x (C API level 27), -configures Cargo automatically, and bundles the required native runtime library -into the resulting wheel. +```bash +BUILD_PROFILE=cuda pip install -v ".[cuda]" +``` + +`BUILD_PROFILE=cuda` enables the native `cuda` Cargo feature and injects the +ORT CUDA provider plugins (`libonnxruntime_providers_{shared,cuda}`) into the +wheel. The `[cuda]` extra installs the CUDA 13 / cuDNN / cublas user +libraries used at runtime. --- @@ -160,9 +177,9 @@ Available presets: ### Generation overrides Model defaults come from Hugging Face `generation_config.json` (typically -`do_sample=True`, `temperature=0.2`, `top_k=0`, `top_p=0.9`). Override them at -load time with `generation_kwargs` (merged onto the decoder config; unknown keys -raise `ValueError`): +`do_sample=True`, `temperature=0.2`, `top_k=0`, `top_p=0.9`). + +Override them at load time with `generation_kwargs` (merged onto the decoder config; unknown keys raise `ValueError`): ```python # Faster / deterministic OCR on CPU (greedy decoding) @@ -203,8 +220,9 @@ Bare `max_new_tokens=` remains supported as a shorthand: model = LightOnOCR.from_pretrained("...", max_new_tokens=1024) ``` -On CPU, prefer `do_sample=False` for throughput. If you need sampling, set a -modest `top_k` (for example `50`) instead of leaving the HF default `top_k=0`. +> On CPU, prefer `do_sample=False` for throughput. + +> If you need sampling, set a modest `top_k` (for example `50`) instead of leaving the HF default `top_k=0`. --- @@ -235,31 +253,22 @@ export ORT_DYLIB_PATH="$(python -c \ ### Building a wheel -To build a distributable wheel, use the project's Python build backend: +Same profiles as [Building from source](#building-from-source): ```bash +# CPU (default) poetry run pip wheel . --wheel-dir dist -``` - -The default build profile targets CPU execution and does not enable any Cargo -features. During source builds, the build backend automatically discovers a -compatible ONNX Runtime from `ORT_DYLIB_PATH` or from the selected build -profile's Python runtime package, validates compatibility with ONNX Runtime -1.28.x (C API level 27), configures Cargo, and produces a wheel containing the -required native runtime libraries. -To build using the CUDA profile: - -```bash +# CUDA BUILD_PROFILE=cuda poetry run pip wheel . --wheel-dir dist ``` > **Note** > > Running `maturin develop` **without** `--features load-dynamic` is not -> supported. The custom build backend is responsible for configuring ONNX -> Runtime linking during production builds, whereas editable development uses -> the `load-dynamic` feature together with `ORT_DYLIB_PATH`. +> supported. Production/`pip install` builds use the custom build backend for +> ONNX Runtime linking; editable development uses `load-dynamic` with +> `ORT_DYLIB_PATH`. --- diff --git a/bindings/python/fast_lightonocr_build/backend.py b/bindings/python/fast_lightonocr_build/backend.py index 2af8999..3f738bf 100644 --- a/bindings/python/fast_lightonocr_build/backend.py +++ b/bindings/python/fast_lightonocr_build/backend.py @@ -7,7 +7,9 @@ from __future__ import annotations import atexit +import base64 import ctypes +import hashlib import importlib.metadata import importlib.util import os @@ -15,6 +17,7 @@ import re import shutil import tempfile +import zipfile from collections.abc import Mapping from dataclasses import dataclass from functools import lru_cache @@ -87,12 +90,24 @@ def build_wheel( ) with _configured_build_environment(settings): - return maturin.build_wheel( + wheel_name = maturin.build_wheel( wheel_directory, settings, metadata_directory, ) + # CUDA only: auditwheel vendors DT_NEEDED libs, but ORT dlopen's + # provider plugins at EP registration. CPU wheels are left untouched. + profile = _selected_profile(settings) + if profile.name == "cuda" and not _uses_load_dynamic(settings): + runtime = _configured_onnx_runtime(profile) + _bundle_cuda_ort_provider_libraries( + Path(wheel_directory) / wheel_name, + runtime.parent, + ) + + return wheel_name + def build_editable( wheel_directory: str, @@ -697,6 +712,197 @@ def _parse_version( return major, minor, patch +# +# CUDA wheel provider-library bundling +# +# Only used for BUILD_PROFILE=cuda. ORT loads CUDA EP plugins via dlopen +# next to libonnxruntime; auditwheel never sees those deps. CPU wheels skip +# this path entirely. +# + + +def _bundle_cuda_ort_provider_libraries(wheel_path: Path, capi_dir: Path) -> None: + """Copy CUDA ORT provider libs into a repaired CUDA wheel's lib dir.""" + + providers = _cuda_ort_provider_library_paths(capi_dir) + if not providers: + raise RuntimeError( + "CUDA build profile requires ONNX Runtime CUDA provider libraries " + f"in {capi_dir} (expected libonnxruntime_providers_shared and " + "libonnxruntime_providers_cuda), but none were found" + ) + + if not wheel_path.is_file(): + raise RuntimeError(f"wheel not found for provider bundling: {wheel_path}") + + tmp_path = wheel_path.with_suffix(wheel_path.suffix + ".tmp") + + with ( + zipfile.ZipFile(wheel_path, "r") as zin, + zipfile.ZipFile(tmp_path, "w") as zout, + ): + names = zin.namelist() + libs_dir = _wheel_bundled_lib_dir(names) + if libs_dir is None: + tmp_path.unlink(missing_ok=True) + raise RuntimeError( + "CUDA wheel is missing a .libs/.dylibs directory after auditwheel " + "repair; cannot bundle ORT CUDA provider libraries" + ) + + dist_info = _wheel_dist_info_dir(names) + if dist_info is None: + raise RuntimeError(f"could not find *.dist-info in {wheel_path}") + + record_name = f"{dist_info}/RECORD" + bundled_ort = _wheel_bundled_ort_library(names, libs_dir) + existing = set(names) + record_lines: list[str] = [] + + for info in zin.infolist(): + if info.filename == record_name: + continue + + data = zin.read(info.filename) + # Preserve member metadata (including symlinks). + zout.writestr(info, data) + if not info.filename.endswith("/"): + record_lines.append(_wheel_record_line(info.filename, data)) + + for provider in providers: + dest = f"{libs_dir}/{provider.name}" + if dest in existing: + continue + + data = provider.read_bytes() + zout.writestr(dest, data) + record_lines.append(_wheel_record_line(dest, data)) + print( + f"Bundled ORT CUDA provider library {provider.name} into {libs_dir}/", + flush=True, + ) + + # Provider libs typically NEEDED libonnxruntime.so.1; auditwheel may + # only keep a hashed filename. Add stable SONAME aliases when missing. + if bundled_ort is not None and platform.system() in {"Linux", "FreeBSD"}: + for alias in ("libonnxruntime.so", "libonnxruntime.so.1"): + dest = f"{libs_dir}/{alias}" + if dest in existing or any( + item.filename == dest for item in zout.filelist + ): + continue + _write_zip_symlink(zout, dest, bundled_ort) + record_lines.append( + _wheel_record_line(dest, bundled_ort.encode("utf-8")) + ) + print( + f"Added ORT SONAME alias {alias} -> {bundled_ort} in {libs_dir}/", + flush=True, + ) + + record_lines.sort() + record_body = "\n".join(record_lines) + f"\n{record_name},,\n" + zout.writestr(record_name, record_body) + + tmp_path.replace(wheel_path) + + +def _cuda_ort_provider_library_paths(capi_dir: Path) -> list[Path]: + """Return shared + CUDA provider libs only (skip TensorRT, etc.).""" + + if not capi_dir.is_dir(): + return [] + + return sorted( + path for path in capi_dir.iterdir() if _is_cuda_provider_library(path) + ) + + +def _is_cuda_provider_library(path: Path) -> bool: + if not path.is_file(): + return False + + name = path.name.lower() + system = platform.system() + + if system == "Windows": + return name in { + "onnxruntime_providers_shared.dll", + "onnxruntime_providers_cuda.dll", + } + + if system == "Darwin": + return name in { + "libonnxruntime_providers_shared.dylib", + "libonnxruntime_providers_cuda.dylib", + } + + # Linux / FreeBSD: accept versioned sonames too. + return name.startswith("libonnxruntime_providers_shared.so") or name.startswith( + "libonnxruntime_providers_cuda.so" + ) + + +def _wheel_bundled_lib_dir(names: list[str]) -> str | None: + for name in names: + top = name.split("/", 1)[0] + if top.endswith(".libs") or top.endswith(".dylibs"): + return top + return None + + +def _wheel_dist_info_dir(names: list[str]) -> str | None: + for name in names: + top = name.split("/", 1)[0] + if top.endswith(".dist-info"): + return top + return None + + +def _wheel_bundled_ort_library(names: list[str], libs_dir: str) -> str | None: + prefix = f"{libs_dir}/" + candidates: list[str] = [] + + for name in names: + if not name.startswith(prefix) or name.endswith("/"): + continue + + file_name = name[len(prefix) :] + if "/" in file_name: + continue + if "providers" in file_name.lower(): + continue + if not file_name.startswith("libonnxruntime"): + continue + if ".so" not in file_name and not file_name.endswith(".dylib"): + continue + candidates.append(file_name) + + if not candidates: + return None + + # Prefer the real shared object over short SONAME aliases. + candidates.sort(key=lambda item: (-len(item), item)) + return candidates[0] + + +def _wheel_record_line(arcname: str, data: bytes) -> str: + digest = hashlib.sha256(data).digest() + encoded = base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii") + return f"{arcname},sha256={encoded},{len(data)}" + + +def _write_zip_symlink( + zf: zipfile.ZipFile, + link_path: str, + target: str, +) -> None: + info = zipfile.ZipInfo(link_path) + info.create_system = 3 # Unix + info.external_attr = 0o120755 << 16 # symlink + zf.writestr(info, target.encode("utf-8")) + + # # Linker staging # diff --git a/bindings/python/pyproject.toml b/bindings/python/pyproject.toml index 44d5fe2..00e5737 100644 --- a/bindings/python/pyproject.toml +++ b/bindings/python/pyproject.toml @@ -42,7 +42,10 @@ cpu = [ "onnxruntime>=1.28,<2", ] cuda = [ - "onnxruntime-gpu>=1.28,<2", + "onnxruntime-gpu[cuda,cudnn]>=1.28,<2", + # ORT 1.28 CUDA EP needs CUDA 13 cublas; not included in the upstream + # onnxruntime-gpu[cuda] extra as of 1.28.0. + "nvidia-cublas~=13.0", ] [dependency-groups] @@ -77,6 +80,8 @@ build-requirements = [ [tool.fast-lightonocr.build-profiles.cuda] cargo-features = ["cuda"] build-requirements = [ + # Build only needs the ORT dylib for linking; CUDA/cuDNN user libs are + # runtime deps via the `cuda` optional-extra. "onnxruntime-gpu>=1.28,<2", ] diff --git a/bindings/python/python/fast_lightonocr/__init__.py b/bindings/python/python/fast_lightonocr/__init__.py index 8fa651c..1431481 100644 --- a/bindings/python/python/fast_lightonocr/__init__.py +++ b/bindings/python/python/fast_lightonocr/__init__.py @@ -57,6 +57,53 @@ ), } +_CUDA_RUNTIME_PRELOADED = False + + +def _runtime_requests_cuda(runtime_kwargs: Mapping[str, Any] | None) -> bool: + """Return True when runtime_kwargs selects the CUDA execution provider.""" + + if not runtime_kwargs: + return False + + provider = runtime_kwargs.get("execution_provider") + if provider is None: + return False + + return str(provider).lower() == "cuda" + + +def _preload_cuda_runtime_libraries() -> None: + """Preload pip NVIDIA CUDA/cuDNN libs for the native CUDA EP. + + Scoped to CUDA loads only. ONNX Runtime's ``preload_dlls`` brings the + ``nvidia-*`` wheel libraries into the process so Linux can resolve + ``libcublasLt`` / cuDNN without requiring ``LD_LIBRARY_PATH``. + """ + + global _CUDA_RUNTIME_PRELOADED + if _CUDA_RUNTIME_PRELOADED: + return + + try: + import onnxruntime as ort + except ImportError as exc: + raise ImportError( + "CUDA execution_provider requires onnxruntime-gpu. " + "Install with: pip install 'fast-lightonocr[cuda]'" + ) from exc + + preload = getattr(ort, "preload_dlls", None) + if preload is None: + raise RuntimeError( + "installed onnxruntime does not provide preload_dlls(); " + "fast-lightonocr CUDA support requires onnxruntime-gpu>=1.28" + ) + + # Bring pip nvidia-* CUDA/cuDNN libs into this process (CUDA path only). + preload(cuda=True, cudnn=True) + _CUDA_RUNTIME_PRELOADED = True + def _download_patterns(preset: Preset) -> list[str]: """Return the Hugging Face download patterns for a model preset. @@ -195,6 +242,10 @@ def from_pretrained( gen_kwargs = dict(generation_kwargs) if generation_kwargs is not None else None rt_kwargs = dict(runtime_kwargs) if runtime_kwargs is not None else None + # CPU loads never import onnxruntime or touch NVIDIA libs. + if _runtime_requests_cuda(rt_kwargs): + _preload_cuda_runtime_libraries() + if model_path.exists(): return cls( _NativeLightOnOCR._load_model_dir(