mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(core): load ONNX Runtime dynamically so headroom._core imports on non-AVX2 x86-64 (#1715)
## Description
`import headroom._core` dies with SIGILL (`Illegal instruction`) on
x86-64 CPUs without AVX2 (Pentium N4200, Celeron N4500, AMD FX 8350 —
all reported on the issue). The repo sets no `RUSTFLAGS`/`target-cpu`
anywhere, so first-party Rust code is baseline x86-64; the AVX2 code
comes from Microsoft's prebuilt ONNX Runtime, statically linked into the
extension by fastembed's `ort-download-binaries-rustls-tls` feature on
non-Windows targets. Because it is statically linked, its code is mapped
and initialized when the extension module loads — **before** the runtime
AVX2 guard from #1162 can run, which is why that fix helped Magika init
but not the import-time crash.
Fix, mirroring what Windows already does for its own reasons (DirectML
link libs): build with `ort-load-dynamic` on every platform, so ONNX
Runtime is only `dlopen`'d at first use, where the #1162 AVX2 guard
falls back to the non-ONNX detection tiers on unsupported CPUs. Since
both target blocks became identical, they are collapsed into one
platform-independent `fastembed` dependency.
To keep Magika/fastembed working out of the box on Linux/macOS, the
existing `ORT_DYLIB_PATH` auto-pin (`headroom/_ort.py`, previously
Windows-only) now resolves the pip `onnxruntime` package's shared
library on all platforms (`onnxruntime.dll` / `libonnxruntime.so*` /
`libonnxruntime*.dylib`). The pip `onnxruntime` CPU wheels use runtime
CPU dispatch, so they also work on pre-AVX2 machines — non-AVX2 users
get working ML detection instead of a crash. Without the `onnxruntime`
package, ML detection degrades gracefully to the non-ONNX tiers exactly
as it already does on Windows.
Fixes #1278
## Type of Change
- [x] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [x] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring (no functional changes)
## Changes Made
- `crates/headroom-core/Cargo.toml`: replaced the per-target `fastembed`
blocks (`ort-download-binaries-rustls-tls` on non-Windows,
`ort-load-dynamic` on Windows) with a single platform-independent
dependency on `ort-load-dynamic`, with a comment documenting both the
DirectML and the AVX2/#1278 rationale.
- `Cargo.lock`: regenerated — `ort-sys` drops its static-download
dependencies (`hmac-sha256`, `lzma-rust2`, `ureq`); no version bumps.
- `headroom/_ort.py`: `ORT_DYLIB_PATH` auto-pin extended from
Windows-only to all platforms via a small `_find_dylib` helper that
resolves the platform's shared-library name inside the pip `onnxruntime`
package.
- `tests/test_transforms/test_ort_dylib.py`: replaced the obsolete
`test_noop_on_non_windows` with Linux (versioned `.so`) and macOS
(`.dylib`) pin tests; module docstring updated.
- `docs/content/docs/configuration.mdx`: `ORT_DYLIB_PATH` row updated
from Windows-only wording to the cross-platform behavior.
## Testing
- [x] Unit tests pass (`pytest`)
- [x] Linting passes (`ruff check .`)
- [x] Type checking passes (`mypy headroom`)
- [x] New tests added for new functionality
- [x] Manual testing performed
### Test Output
```text
$ cargo fmt --all -- --check && cargo clippy --workspace -- -D warnings && cargo test -p headroom-core --lib
clean
test result: 844 passed; 0 failed; 1 ignored
$ python -m pytest tests/test_transforms/test_ort_dylib.py -q
8 passed
$ ruff check headroom/_ort.py tests/test_transforms/test_ort_dylib.py
All checks passed!
```
## Real Behavior Proof
- Environment: Windows 11 (AVX2-capable — the SIGILL itself is not
reproducible on this machine), Python 3.13, Rust 1.95.0, local checkout
branched from `upstream/main` (9fbd47ba).
- Exact command / steps: `cargo check -p headroom-core` after the
feature switch; inspected the `Cargo.lock` diff; rebuilt and ran `python
-c "import headroom; from headroom._core import detect_content_type;
print(detect_content_type('hello world'))"`; ran the ort-pin test suite
with monkeypatched `linux`/`darwin` platforms.
- Observed result: build succeeds with `ort-load-dynamic`; the lockfile
shows `ort-sys` no longer pulls the binary-download machinery
(`hmac-sha256`, `lzma-rust2`, `ureq` removed), confirming the
statically-linked prebuilt ORT is gone; import + content detection works
with `ORT_DYLIB_PATH` auto-pinned to the pip onnxruntime library; all 8
pin tests pass including the new Linux/macOS branches.
- Not tested: actual pre-AVX2 x86-64 hardware (none available — the fix
removes AVX2 code from the import path by construction, and the issue
reporters on #1278 can verify); Linux/macOS wheel runtime behavior
beyond CI's ubuntu/macOS wheel-build jobs; embedding quality/performance
under a pip-provided ORT version differing from the previously vendored
one.
## Review Readiness
- [x] I have performed a self-review
- [x] This PR is ready for human review
---------
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: JerrettDavis <mxjerrett@gmail.com>
This commit is contained in:
Parideboy
GitHub
parent
1590913bb7
commit
c46cd8f950
9 changed files with 128 additions and 192 deletions
13
.github/workflows/rust.yml
vendored
13
.github/workflows/rust.yml
vendored
|
|
@ -54,6 +54,19 @@ jobs:
|
|||
components: rustfmt, clippy
|
||||
- name: Cache cargo registry + build
|
||||
uses: Swatinem/rust-cache@v2
|
||||
- uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: '3.11'
|
||||
- name: Provide ONNX Runtime dylib
|
||||
# headroom-core is built with `ort-load-dynamic` (see its
|
||||
# Cargo.toml): the ONNX Runtime shared library is dlopen'd at
|
||||
# runtime instead of statically linked, so the magika/detection
|
||||
# tests need a real libonnxruntime.so and ORT_DYLIB_PATH pointing
|
||||
# at it — same contract `headroom/_ort.py` fulfills for Python
|
||||
# users via the pip `onnxruntime` package.
|
||||
run: |
|
||||
pip install 'onnxruntime>=1.16.0'
|
||||
echo "ORT_DYLIB_PATH=$(python -c "import onnxruntime, pathlib; p = pathlib.Path(onnxruntime.__file__).parent / 'capi'; print(next(iter(sorted(p.glob('libonnxruntime.so*')))))")" >> "$GITHUB_ENV"
|
||||
- name: cargo fmt --check
|
||||
run: cargo fmt --all -- --check
|
||||
- name: cargo clippy
|
||||
|
|
|
|||
20
Cargo.lock
generated
20
Cargo.lock
generated
|
|
@ -1969,7 +1969,7 @@ dependencies = [
|
|||
"reqwest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 1.0.69",
|
||||
"thiserror 2.0.18",
|
||||
"tokio",
|
||||
"tower",
|
||||
"tracing",
|
||||
|
|
@ -2043,12 +2043,6 @@ dependencies = [
|
|||
"digest 0.11.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hmac-sha256"
|
||||
version = "1.1.14"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ec9d92d097f4749b64e8cc33d924d9f40a2d4eb91402b458014b781f5733d60f"
|
||||
|
||||
[[package]]
|
||||
name = "http"
|
||||
version = "0.2.12"
|
||||
|
|
@ -2605,12 +2599,6 @@ version = "0.1.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154"
|
||||
|
||||
[[package]]
|
||||
name = "lzma-rust2"
|
||||
version = "0.15.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e20f57f9918e5bd7bc58c22cdd70a6afc7375d4dd9683af5f2b34bd3d2bba619"
|
||||
|
||||
[[package]]
|
||||
name = "macro_rules_attribute"
|
||||
version = "0.2.2"
|
||||
|
|
@ -2963,7 +2951,6 @@ dependencies = [
|
|||
"ort-sys",
|
||||
"smallvec",
|
||||
"tracing",
|
||||
"ureq 3.3.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -2971,11 +2958,6 @@ name = "ort-sys"
|
|||
version = "2.0.0-rc.12"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d7b497d21a8b6fbb4b5a544f8fadb77e801a09ae0add9e411d31c6f89e3c1e90"
|
||||
dependencies = [
|
||||
"hmac-sha256",
|
||||
"lzma-rust2",
|
||||
"ureq 3.3.0",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "outref"
|
||||
|
|
|
|||
|
|
@ -122,30 +122,16 @@ redis = { version = "0.27", optional = true, default-features = false }
|
|||
# cycling through the proxy crate. Tiny crate (no I/O, just types).
|
||||
http = "1"
|
||||
|
||||
[target.'cfg(all(not(windows), not(all(target_os = "macos", target_arch = "x86_64"))))'.dependencies]
|
||||
fastembed = { version = "5", default-features = false, features = [
|
||||
"hf-hub-rustls-tls",
|
||||
"ort-download-binaries-rustls-tls",
|
||||
"image-models",
|
||||
] }
|
||||
|
||||
[target.'cfg(all(target_os = "macos", target_arch = "x86_64"))'.dependencies]
|
||||
# `ort-sys 2.0.0-rc.12` does not ship prebuilt ONNX Runtime binaries for
|
||||
# `x86_64-apple-darwin`. Building ORT from source in CI would add CMake and
|
||||
# several minutes per wheel. Load the pip `onnxruntime` dylib at runtime
|
||||
# instead (same approach as Windows below). `headroom/_ort.py` pins
|
||||
# `ORT_DYLIB_PATH` before `headroom._core` imports.
|
||||
fastembed = { version = "5", default-features = false, features = [
|
||||
"hf-hub-rustls-tls",
|
||||
"ort-load-dynamic",
|
||||
"image-models",
|
||||
] }
|
||||
ort = { version = "2.0.0-rc.12", default-features = false, features = ["load-dynamic"] }
|
||||
|
||||
[target.'cfg(windows)'.dependencies]
|
||||
# `ort-download-binaries-*` emits DirectML link libs on Windows (`DXCORE`,
|
||||
# `DXGI`, `D3D12`, `DirectML`). Users installing `headroom-ai[all]` from
|
||||
# sdist often do not have those SDK libs, so load ORT dynamically instead.
|
||||
# Load ONNX Runtime dynamically on every platform. The alternative,
|
||||
# `ort-download-binaries-*`, statically links Microsoft's prebuilt ORT:
|
||||
# on Windows it emits DirectML link libs (`DXCORE`, `DXGI`, `D3D12`,
|
||||
# `DirectML`) that sdist installs of `headroom-ai[all]` often lack, and
|
||||
# on x86_64 Linux/macOS the prebuilt binary requires AVX2 — its code is
|
||||
# mapped and initialized as soon as the `headroom._core` extension
|
||||
# loads, so importing headroom SIGILLed on pre-AVX2 CPUs before the
|
||||
# runtime AVX2 guard could run (#1278). With `ort-load-dynamic` the
|
||||
# library is only dlopen'd at first use, where the AVX2 guard falls
|
||||
# back to the non-ONNX detection tiers.
|
||||
fastembed = { version = "5", default-features = false, features = [
|
||||
"hf-hub-rustls-tls",
|
||||
"ort-load-dynamic",
|
||||
|
|
|
|||
|
|
@ -102,6 +102,12 @@ impl EmbeddingScorer {
|
|||
this x86 CPU; embedding relevance disabled (falling back to BM25)"
|
||||
.to_string());
|
||||
}
|
||||
// The crate loads ONNX Runtime dynamically (`ort-load-dynamic`);
|
||||
// resolve and commit the dylib before fastembed touches ort — a
|
||||
// failed in-ort load deadlocks instead of erroring (see
|
||||
// `dynamic_ort_loader_ready`).
|
||||
crate::transforms::magika_detector::dynamic_ort_loader_ready()
|
||||
.map_err(|e| format!("EmbeddingScorer: ONNX Runtime unavailable: {e}"))?;
|
||||
let name = format!("{:?}", model_kind);
|
||||
let model = TextEmbedding::try_new(InitOptions::new(model_kind))
|
||||
.map_err(|e| format!("EmbeddingScorer model load failed: {}", e))?;
|
||||
|
|
|
|||
|
|
@ -35,10 +35,6 @@
|
|||
//! error early instead of crashing with SIGILL; the detection chain then
|
||||
//! falls through to Tier 2 and Tier 3 normally.
|
||||
|
||||
#[cfg(any(
|
||||
target_os = "windows",
|
||||
all(target_os = "macos", target_arch = "x86_64")
|
||||
))]
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::mpsc;
|
||||
use std::sync::{Mutex, OnceLock};
|
||||
|
|
@ -85,17 +81,18 @@ pub(crate) fn magika_runtime_available_for_session_init() -> Result<(), String>
|
|||
dynamic_ort_loader_ready()
|
||||
}
|
||||
|
||||
#[cfg(any(
|
||||
target_os = "windows",
|
||||
all(target_os = "macos", target_arch = "x86_64")
|
||||
))]
|
||||
static DYNAMIC_ORT_INIT: OnceLock<Result<PathBuf, String>> = OnceLock::new();
|
||||
|
||||
#[cfg(any(
|
||||
target_os = "windows",
|
||||
all(target_os = "macos", target_arch = "x86_64")
|
||||
))]
|
||||
fn dynamic_ort_loader_ready() -> Result<(), String> {
|
||||
/// Ensure the ONNX Runtime shared library is resolved and committed to
|
||||
/// `ort` before any `ort` API is touched.
|
||||
///
|
||||
/// With `ort-load-dynamic` (every platform, see Cargo.toml) this MUST
|
||||
/// run before any code path that can construct an `ort` session
|
||||
/// (magika, fastembed): if the dylib cannot be loaded, `ort`
|
||||
/// 2.0.0-rc.12 deadlocks inside its API-setup error path (recursive
|
||||
/// `OnceLock` init), and the stuck thread then wedges process exit in
|
||||
/// `ort`'s `dl_fini` environment teardown (#1715 CI hang).
|
||||
pub(crate) fn dynamic_ort_loader_ready() -> Result<(), String> {
|
||||
DYNAMIC_ORT_INIT
|
||||
.get_or_init(initialize_dynamic_ort)
|
||||
.as_ref()
|
||||
|
|
@ -103,10 +100,6 @@ fn dynamic_ort_loader_ready() -> Result<(), String> {
|
|||
.map_err(Clone::clone)
|
||||
}
|
||||
|
||||
#[cfg(any(
|
||||
target_os = "windows",
|
||||
all(target_os = "macos", target_arch = "x86_64")
|
||||
))]
|
||||
fn initialize_dynamic_ort() -> Result<PathBuf, String> {
|
||||
let explicit = std::env::var("ORT_DYLIB_PATH")
|
||||
.ok()
|
||||
|
|
@ -154,10 +147,6 @@ fn initialize_dynamic_ort() -> Result<PathBuf, String> {
|
|||
}
|
||||
}
|
||||
|
||||
#[cfg(any(
|
||||
target_os = "windows",
|
||||
all(target_os = "macos", target_arch = "x86_64")
|
||||
))]
|
||||
fn init_ort_from_path(path: &Path) -> Result<(), String> {
|
||||
let builder = ort::init_from(path).map_err(|error| {
|
||||
format!(
|
||||
|
|
@ -169,10 +158,6 @@ fn init_ort_from_path(path: &Path) -> Result<(), String> {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(any(
|
||||
target_os = "windows",
|
||||
all(target_os = "macos", target_arch = "x86_64")
|
||||
))]
|
||||
fn discover_onnxruntime_libraries() -> Vec<PathBuf> {
|
||||
let mut roots = Vec::new();
|
||||
|
||||
|
|
@ -218,20 +203,12 @@ fn discover_onnxruntime_libraries() -> Vec<PathBuf> {
|
|||
dedup_existing_files(candidates)
|
||||
}
|
||||
|
||||
#[cfg(any(
|
||||
target_os = "windows",
|
||||
all(target_os = "macos", target_arch = "x86_64")
|
||||
))]
|
||||
fn env_path(name: &str) -> Option<PathBuf> {
|
||||
std::env::var_os(name)
|
||||
.map(PathBuf::from)
|
||||
.filter(|path| !path.as_os_str().is_empty())
|
||||
}
|
||||
|
||||
#[cfg(any(
|
||||
target_os = "windows",
|
||||
all(target_os = "macos", target_arch = "x86_64")
|
||||
))]
|
||||
fn versioned_children(root: PathBuf) -> Vec<PathBuf> {
|
||||
let mut children = std::fs::read_dir(root)
|
||||
.ok()
|
||||
|
|
@ -244,10 +221,6 @@ fn versioned_children(root: PathBuf) -> Vec<PathBuf> {
|
|||
children
|
||||
}
|
||||
|
||||
#[cfg(any(
|
||||
target_os = "windows",
|
||||
all(target_os = "macos", target_arch = "x86_64")
|
||||
))]
|
||||
fn onnxruntime_candidates_under(root: &Path) -> Vec<PathBuf> {
|
||||
#[cfg(target_os = "windows")]
|
||||
{
|
||||
|
|
@ -264,7 +237,7 @@ fn onnxruntime_candidates_under(root: &Path) -> Vec<PathBuf> {
|
|||
]
|
||||
}
|
||||
|
||||
#[cfg(all(target_os = "macos", target_arch = "x86_64"))]
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
{
|
||||
let mut candidates = Vec::new();
|
||||
for site_packages in python_site_packages_dirs(root) {
|
||||
|
|
@ -275,7 +248,7 @@ fn onnxruntime_candidates_under(root: &Path) -> Vec<PathBuf> {
|
|||
}
|
||||
}
|
||||
|
||||
#[cfg(all(target_os = "macos", target_arch = "x86_64"))]
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
fn python_site_packages_dirs(root: &Path) -> Vec<PathBuf> {
|
||||
let mut dirs = vec![root.join("lib").join("site-packages")];
|
||||
let lib = root.join("lib");
|
||||
|
|
@ -297,7 +270,7 @@ fn python_site_packages_dirs(root: &Path) -> Vec<PathBuf> {
|
|||
dirs
|
||||
}
|
||||
|
||||
#[cfg(all(target_os = "macos", target_arch = "x86_64"))]
|
||||
#[cfg(not(target_os = "windows"))]
|
||||
fn onnxruntime_dylibs_in(capi: &Path) -> Vec<PathBuf> {
|
||||
let mut dylibs = std::fs::read_dir(capi)
|
||||
.ok()
|
||||
|
|
@ -307,17 +280,16 @@ fn onnxruntime_dylibs_in(capi: &Path) -> Vec<PathBuf> {
|
|||
.filter(|path| {
|
||||
path.file_name()
|
||||
.and_then(|name| name.to_str())
|
||||
.is_some_and(|name| name.starts_with("libonnxruntime") && name.ends_with(".dylib"))
|
||||
.is_some_and(|name| {
|
||||
name.starts_with("libonnxruntime")
|
||||
&& (name.ends_with(".dylib") || name.contains(".so"))
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
dylibs.sort();
|
||||
dylibs
|
||||
}
|
||||
|
||||
#[cfg(any(
|
||||
target_os = "windows",
|
||||
all(target_os = "macos", target_arch = "x86_64")
|
||||
))]
|
||||
fn dedup_existing_files(paths: Vec<PathBuf>) -> Vec<PathBuf> {
|
||||
let mut out = Vec::new();
|
||||
for path in paths {
|
||||
|
|
@ -328,14 +300,6 @@ fn dedup_existing_files(paths: Vec<PathBuf>) -> Vec<PathBuf> {
|
|||
out
|
||||
}
|
||||
|
||||
#[cfg(not(any(
|
||||
target_os = "windows",
|
||||
all(target_os = "macos", target_arch = "x86_64")
|
||||
)))]
|
||||
fn dynamic_ort_loader_ready() -> Result<(), String> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Errors from the magika detector. Wraps the underlying `magika::Error`
|
||||
/// so callers can match on whether init or inference broke without
|
||||
/// pulling magika types into their imports.
|
||||
|
|
|
|||
|
|
@ -307,7 +307,7 @@ headroom proxy --learn --min-evidence 3
|
|||
| `HEADROOM_STRIP_INTERNAL_HEADERS` | Python proxy: whether to strip internal `x-headroom-*` request headers (e.g. `x-headroom-bypass`, `x-headroom-mode`, `x-headroom-user-id`, `x-headroom-stack`, `x-headroom-base-url`) before every upstream forwarder call (PR-A5, fixes P5-49). `enabled` (default) stops fingerprinting / leakage. `disabled` is an explicit operator opt-in for diagnostic shadow tracing — NOT a fallback. Inbound reads of these headers (bypass gating, memory user-id resolution) are unaffected because they read `request.headers` directly. | `enabled` |
|
||||
| `HEADROOM_PROXY_STRIP_INTERNAL_HEADERS` | Rust proxy: same policy as `HEADROOM_STRIP_INTERNAL_HEADERS` but for the Rust transparent proxy. Stripping happens inside `build_forward_request_headers` so both HTTP and WebSocket upstream calls are gated by one flag. `enabled` default; `disabled` operator opt-in for diagnostic shadow tracing. Response-side `X-Headroom-*` injection (e.g. `x-headroom-tokens-saved`) is unrelated and stays. | `enabled` |
|
||||
| `HEADROOM_EMBEDDER_RUNTIME` | Set to `pytorch_mps` to run the memory embedder via the torch sentence-transformers backend on the Apple GPU (MPS). Only engages when Apple MPS is actually available; otherwise it logs a warning and uses the existing default embedder selection path. `pytorch_mps` is the only accepted value. Requires the `[pytorch-mps]` extra. See [Memory](/docs/memory#embedding-runtime--gpu-offload-apple-silicon). | default embedder selection |
|
||||
| `ORT_DYLIB_PATH` | Windows: path to the `onnxruntime.dll` loaded by the Rust core (magika detection, fastembed embeddings). Auto-pinned at `import headroom` to the DLL inside the `onnxruntime` pip package; set it yourself to override. Without a pin the bare Windows DLL search resolves to the Windows ML System32 build (1.17.x on Win11 24H2+), which deadlocks ONNX session init — see [Troubleshooting](/docs/troubleshooting#windows-ml-content-detection-hangs-or-silently-falls-back). | auto-pinned on Windows |
|
||||
| `ORT_DYLIB_PATH` | Path to the ONNX Runtime shared library loaded by the Rust core (magika detection, fastembed embeddings), which loads ORT dynamically on every platform. Auto-pinned at `import headroom` to the library inside the `onnxruntime` pip package (`onnxruntime.dll` / `libonnxruntime.so*` / `libonnxruntime*.dylib`); set it yourself to override. Without a pin, ML detection degrades to the non-ONNX tiers — and on Windows the bare DLL search can resolve to the Windows ML System32 build (1.17.x on Win11 24H2+), which deadlocks ONNX session init — see [Troubleshooting](/docs/troubleshooting#windows-ml-content-detection-hangs-or-silently-falls-back). | auto-pinned |
|
||||
| `HEADROOM_MAGIKA_INIT_TIMEOUT_SECS` | Upper bound (integer seconds, > 0) on magika's one-time ONNX session init in the Rust detection chain. On timeout the init error is cached and detection uses the non-ML fallback tiers for the rest of the process; a warning is logged. Safety net for environments where the dylib pin above does not apply. | `5` |
|
||||
| `HEADROOM_REQUEST_TIMEOUT` | Request timeout in seconds | `300` |
|
||||
| `HEADROOM_BETA_HEADER_STICKY` | Controls per-session `anthropic-beta` / `OpenAI-Beta` re-echo. `enabled` (default): the proxy unions beta tokens across turns within a session — if the client sends a token in turn N and omits it in turn N+1, the proxy re-injects it to preserve prefix-cache stability. `disabled`: the client's value is forwarded verbatim with no accumulation. Any other value raises at request time. See [Session Beta Header Tracking](/docs/configuration#session-beta-header-tracking). | `enabled` |
|
||||
|
|
|
|||
|
|
@ -1,41 +1,33 @@
|
|||
"""Pin the ONNX Runtime dylib for the Rust core on dynamic-ORT platforms.
|
||||
"""Pin the ONNX Runtime dylib for the Rust core.
|
||||
|
||||
Why this module exists
|
||||
----------------------
|
||||
On Windows and Intel macOS (``x86_64-apple-darwin``), ``headroom._core``
|
||||
consumers of the ``ort`` crate (magika content detection, fastembed
|
||||
embeddings) are built with ``ort-load-dynamic``: the native ONNX Runtime
|
||||
library is resolved at *runtime*.
|
||||
``headroom._core`` consumers of the ``ort`` crate (magika content
|
||||
detection, fastembed embeddings) are built with ``ort-load-dynamic`` on
|
||||
every platform: the native ONNX Runtime library is resolved at runtime
|
||||
rather than statically linked.
|
||||
|
||||
Windows: unless ``ORT_DYLIB_PATH`` is set, ort falls back to a bare
|
||||
``LoadLibrary("onnxruntime.dll")`` and the Windows DLL search order
|
||||
applies — and ``C:\\Windows\\System32`` wins. Windows 11 24H2+ ships
|
||||
``System32\\onnxruntime.dll`` as part of Windows ML (observed:
|
||||
1.17.2603 "os-germanium"). Initializing an ort 2.x session against that
|
||||
OS build does not fail — it deadlocks indefinitely at 0% CPU, which the
|
||||
tiered detection fallback cannot catch (a hang is not an ``Err``).
|
||||
Reproduced and bracketed with ``scripts/diag_magika_windows.py``: the
|
||||
identical session inits in ~400ms when ``ORT_DYLIB_PATH`` points at the
|
||||
``onnxruntime`` pip package's DLL (which ``headroom-ai[proxy]`` already
|
||||
depends on).
|
||||
|
||||
Intel macOS: ``ort-sys 2.0.0-rc.12`` does not ship prebuilt ONNX Runtime
|
||||
binaries for ``x86_64-apple-darwin``, so the wheel/sdist build uses
|
||||
``ort-load-dynamic`` and expects a pip-installed ``onnxruntime`` dylib
|
||||
at runtime (same contract as Windows).
|
||||
Static ``ort-download-binaries`` linking is risky on x86_64 Linux/macOS
|
||||
because Microsoft's prebuilt ORT requires AVX2 and can execute at
|
||||
extension load, SIGILLing ``import headroom._core`` on pre-AVX2 CPUs
|
||||
before Headroom's runtime guard can fall back (#1278). On Windows, the
|
||||
dynamic fallback can pick up ``C:\\Windows\\System32\\onnxruntime.dll``
|
||||
from Windows ML and deadlock ORT session init on Windows 11 24H2+.
|
||||
|
||||
The fix: before anything can import ``headroom._core``, resolve the
|
||||
pip-installed ``onnxruntime`` native library and export it via
|
||||
``ORT_DYLIB_PATH``. ``headroom/__init__.py`` calls this hook, which
|
||||
guarantees ordering for every package-level consumer.
|
||||
pip-installed ``onnxruntime`` package's shared library
|
||||
(``capi/onnxruntime.dll`` / ``capi/libonnxruntime.so*`` /
|
||||
``capi/libonnxruntime*.dylib``) and export it via ``ORT_DYLIB_PATH``.
|
||||
``headroom/__init__.py`` calls this hook, which guarantees ordering for
|
||||
every package-level consumer.
|
||||
|
||||
Behavior contract
|
||||
-----------------
|
||||
- Active on Windows and Intel macOS only; a no-op elsewhere.
|
||||
- Active on all platforms; pins only when the ``onnxruntime`` package is present.
|
||||
- Respects a pre-set ``ORT_DYLIB_PATH`` (user override wins).
|
||||
- Locates the ``onnxruntime`` package via ``find_spec`` WITHOUT
|
||||
importing it (importing would load its native code; this hook must
|
||||
stay ~microseconds and side-effect free).
|
||||
stay microsecond-scale and side-effect free).
|
||||
- Never raises: import-time failure of an optional accelerator must
|
||||
not break ``import headroom``. Without a pin, detection still
|
||||
degrades gracefully through HEADROOM_MAGIKA_INIT_TIMEOUT_SECS and
|
||||
|
|
@ -47,7 +39,6 @@ from __future__ import annotations
|
|||
import importlib.util
|
||||
import logging
|
||||
import os
|
||||
import platform
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
|
@ -64,9 +55,8 @@ def ensure_ort_dylib_pinned() -> str | None:
|
|||
"""Export ``ORT_DYLIB_PATH`` for the Rust core's ort runtime.
|
||||
|
||||
Returns the effective dylib path (pinned now or already present in
|
||||
the environment), or ``None`` when no pin applies (platforms that
|
||||
bundle ORT at build time, or no ``onnxruntime`` package to point at).
|
||||
Idempotent and exception-free.
|
||||
the environment), or ``None`` when no ``onnxruntime`` package/native
|
||||
library is available. Idempotent and exception-free.
|
||||
"""
|
||||
global _pinned
|
||||
if _pinned is not _UNSET:
|
||||
|
|
@ -75,25 +65,25 @@ def ensure_ort_dylib_pinned() -> str | None:
|
|||
return _pinned # type: ignore[return-value]
|
||||
|
||||
|
||||
def _needs_ort_dylib_pin() -> bool:
|
||||
if sys.platform.startswith("win"):
|
||||
return True
|
||||
return sys.platform == "darwin" and platform.machine() == "x86_64"
|
||||
|
||||
|
||||
def _resolve_ort_native_library(capi_dir: Path) -> Path | None:
|
||||
"""Return the platform's ONNX Runtime shared library inside ``capi_dir``."""
|
||||
if sys.platform.startswith("win"):
|
||||
candidate = capi_dir / "onnxruntime.dll"
|
||||
return candidate if candidate.is_file() else None
|
||||
|
||||
matches = sorted(capi_dir.glob("libonnxruntime*.dylib"))
|
||||
return matches[0] if matches else None
|
||||
patterns = (
|
||||
("libonnxruntime*.dylib",)
|
||||
if sys.platform == "darwin"
|
||||
else ("libonnxruntime.so*", "libonnxruntime*.dylib")
|
||||
)
|
||||
for pattern in patterns:
|
||||
for candidate in sorted(capi_dir.glob(pattern)):
|
||||
if candidate.is_file():
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_and_pin() -> str | None:
|
||||
if not _needs_ort_dylib_pin():
|
||||
return None
|
||||
|
||||
try:
|
||||
existing = os.environ.get(_ENV_VAR)
|
||||
if existing:
|
||||
|
|
@ -104,18 +94,19 @@ def _resolve_and_pin() -> str | None:
|
|||
if spec is None or not spec.origin:
|
||||
logger.debug(
|
||||
"onnxruntime package not found; %s left unset. Rust ML detection "
|
||||
"needs a pip-installed onnxruntime on this platform (install "
|
||||
"headroom-ai[proxy] or set %s explicitly).",
|
||||
"needs a pip-installed onnxruntime (install headroom-ai[proxy] "
|
||||
"or set %s explicitly).",
|
||||
_ENV_VAR,
|
||||
_ENV_VAR,
|
||||
)
|
||||
return None
|
||||
|
||||
native = _resolve_ort_native_library(Path(spec.origin).parent / "capi")
|
||||
capi_dir = Path(spec.origin).parent / "capi"
|
||||
native = _resolve_ort_native_library(capi_dir)
|
||||
if native is None:
|
||||
logger.debug(
|
||||
"onnxruntime package found but no native library under %s; %s left unset",
|
||||
Path(spec.origin).parent / "capi",
|
||||
capi_dir,
|
||||
_ENV_VAR,
|
||||
)
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -184,41 +184,36 @@ def test_no_native_tls_in_wheel_build_tree() -> None:
|
|||
def test_fastembed_uses_rustls_features() -> None:
|
||||
"""The mechanism that keeps openssl-sys out of the build is
|
||||
fastembed's explicit rustls feature selection in headroom-core.
|
||||
fastembed's default features include `hf-hub-native-tls` and
|
||||
`ort-download-binaries-native-tls` — both pull openssl-sys.
|
||||
Disabling defaults and enabling the rustls equivalents removes
|
||||
the OpenSSL surface entirely.
|
||||
fastembed's default features include `hf-hub-native-tls` (pulls
|
||||
openssl-sys). Disabling defaults and enabling the rustls
|
||||
equivalent removes the OpenSSL surface entirely.
|
||||
"""
|
||||
cargo = (ROOT / "crates" / "headroom-core" / "Cargo.toml").read_text(encoding="utf-8")
|
||||
|
||||
assert "default-features = false" in cargo
|
||||
assert '"hf-hub-rustls-tls"' in cargo
|
||||
assert '"ort-download-binaries-rustls-tls"' in cargo
|
||||
# `image-models` is in default; we re-enable it explicitly so we
|
||||
# don't lose the image-embedding capability when defaults are off.
|
||||
assert '"image-models"' in cargo
|
||||
|
||||
|
||||
def test_fastembed_uses_dynamic_ort_on_windows() -> None:
|
||||
"""Windows and Intel macOS sdist builds must not link Pyke's ORT binaries.
|
||||
def test_fastembed_uses_dynamic_ort_everywhere() -> None:
|
||||
"""No build may statically link Pyke's prebuilt ORT binaries.
|
||||
|
||||
`ort-download-binaries-*` emits platform SDK link libs (DirectML on
|
||||
Windows; unavailable prebuilts on `x86_64-apple-darwin`). Those targets
|
||||
must use ORT dynamic loading instead.
|
||||
Windows; no prebuilts for `x86_64-apple-darwin`) and its Linux/macOS
|
||||
binaries require AVX2 at load time, SIGILLing `import headroom._core`
|
||||
on pre-AVX2 x86-64 CPUs (#1278). Every platform loads ORT dynamically
|
||||
(`ort-load-dynamic`), resolved at runtime from the pip `onnxruntime`
|
||||
package by `headroom/_ort.py` / the crate's loader guard.
|
||||
"""
|
||||
|
||||
cargo = (ROOT / "crates" / "headroom-core" / "Cargo.toml").read_text(encoding="utf-8")
|
||||
for section_marker in (
|
||||
"[target.'cfg(windows)'.dependencies]",
|
||||
'[target.\'cfg(all(target_os = "macos", target_arch = "x86_64"))\'.dependencies]',
|
||||
):
|
||||
assert section_marker in cargo, f"missing Cargo target section: {section_marker}"
|
||||
section = cargo.split(section_marker, 1)[1].split("\n[", 1)[0]
|
||||
dependency_lines = "\n".join(
|
||||
line for line in section.splitlines() if not line.lstrip().startswith("#")
|
||||
)
|
||||
assert '"ort-load-dynamic"' in section
|
||||
assert "ort-download-binaries" not in dependency_lines
|
||||
dependency_lines = "\n".join(
|
||||
line for line in cargo.splitlines() if not line.lstrip().startswith("#")
|
||||
)
|
||||
assert '"ort-load-dynamic"' in dependency_lines
|
||||
assert "ort-download-binaries" not in dependency_lines
|
||||
|
||||
|
||||
def test_dockerfiles_no_longer_install_openssl_devel() -> None:
|
||||
|
|
|
|||
|
|
@ -1,9 +1,11 @@
|
|||
"""Tests for headroom._ort — the ORT_DYLIB_PATH auto-pin.
|
||||
"""Tests for headroom._ort -- the ORT_DYLIB_PATH auto-pin.
|
||||
|
||||
The resolver guards the Rust core on platforms that use `ort-load-dynamic`
|
||||
(Windows and Intel macOS). On Windows it avoids the System32 onnxruntime.dll
|
||||
deadlock (Win11 24H2+, see headroom/_ort.py). Platform gates are
|
||||
monkeypatched so the full logic runs on any CI OS.
|
||||
The resolver points the Rust core's ort-load-dynamic runtime at the pip
|
||||
onnxruntime package's shared library on every platform: on Windows it
|
||||
guards against the DLL search picking up the Windows ML System32
|
||||
onnxruntime.dll, and on Linux/macOS it avoids static ORT import-time
|
||||
CPU feature faults on older x86_64 CPUs (#1278). The platform is
|
||||
monkeypatched so every branch runs on any CI OS.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -27,11 +29,6 @@ def _force_windows(monkeypatch):
|
|||
monkeypatch.setattr(sys, "platform", "win32")
|
||||
|
||||
|
||||
def _force_intel_macos(monkeypatch):
|
||||
monkeypatch.setattr(sys, "platform", "darwin")
|
||||
monkeypatch.setattr(_ort.platform, "machine", lambda: "x86_64")
|
||||
|
||||
|
||||
def _fake_spec_for(monkeypatch, package_dir):
|
||||
"""Make find_spec('onnxruntime') resolve to a fake package directory."""
|
||||
spec = SimpleNamespace(origin=str(package_dir / "__init__.py"))
|
||||
|
|
@ -42,15 +39,30 @@ def _fake_spec_for(monkeypatch, package_dir):
|
|||
)
|
||||
|
||||
|
||||
def test_noop_on_non_dynamic_platforms(monkeypatch):
|
||||
def test_pins_versioned_so_on_linux(monkeypatch, tmp_path):
|
||||
monkeypatch.setattr(sys, "platform", "linux")
|
||||
assert _ort.ensure_ort_dylib_pinned() is None
|
||||
assert "ORT_DYLIB_PATH" not in _ort.os.environ
|
||||
pkg = tmp_path / "onnxruntime"
|
||||
capi = pkg / "capi"
|
||||
capi.mkdir(parents=True)
|
||||
so = capi / "libonnxruntime.so.1.22.0"
|
||||
so.write_bytes(b"not really a shared object")
|
||||
_fake_spec_for(monkeypatch, pkg)
|
||||
|
||||
assert _ort.ensure_ort_dylib_pinned() == str(so)
|
||||
assert _ort.os.environ["ORT_DYLIB_PATH"] == str(so)
|
||||
|
||||
|
||||
def test_pins_dylib_on_macos(monkeypatch, tmp_path):
|
||||
monkeypatch.setattr(sys, "platform", "darwin")
|
||||
monkeypatch.setattr(_ort.platform, "machine", lambda: "arm64")
|
||||
assert _ort.ensure_ort_dylib_pinned() is None
|
||||
assert "ORT_DYLIB_PATH" not in _ort.os.environ
|
||||
pkg = tmp_path / "onnxruntime"
|
||||
capi = pkg / "capi"
|
||||
capi.mkdir(parents=True)
|
||||
dylib = capi / "libonnxruntime.1.23.2.dylib"
|
||||
dylib.write_bytes(b"not really a dylib")
|
||||
_fake_spec_for(monkeypatch, pkg)
|
||||
|
||||
assert _ort.ensure_ort_dylib_pinned() == str(dylib)
|
||||
assert _ort.os.environ["ORT_DYLIB_PATH"] == str(dylib)
|
||||
|
||||
|
||||
def test_respects_existing_env(monkeypatch):
|
||||
|
|
@ -73,19 +85,6 @@ def test_pins_to_package_capi_dll(monkeypatch, tmp_path):
|
|||
assert _ort.os.environ["ORT_DYLIB_PATH"] == str(dll)
|
||||
|
||||
|
||||
def test_pins_to_package_capi_dylib_on_intel_macos(monkeypatch, tmp_path):
|
||||
_force_intel_macos(monkeypatch)
|
||||
pkg = tmp_path / "onnxruntime"
|
||||
capi = pkg / "capi"
|
||||
capi.mkdir(parents=True)
|
||||
dylib = capi / "libonnxruntime.1.23.2.dylib"
|
||||
dylib.write_bytes(b"not really a dylib")
|
||||
_fake_spec_for(monkeypatch, pkg)
|
||||
|
||||
assert _ort.ensure_ort_dylib_pinned() == str(dylib)
|
||||
assert _ort.os.environ["ORT_DYLIB_PATH"] == str(dylib)
|
||||
|
||||
|
||||
def test_idempotent_after_first_resolution(monkeypatch, tmp_path):
|
||||
_force_windows(monkeypatch)
|
||||
pkg = tmp_path / "onnxruntime"
|
||||
|
|
@ -110,7 +109,7 @@ def test_no_pin_when_package_missing(monkeypatch):
|
|||
assert "ORT_DYLIB_PATH" not in _ort.os.environ
|
||||
|
||||
|
||||
def test_no_pin_when_dll_file_absent(monkeypatch, tmp_path):
|
||||
def test_no_pin_when_native_library_absent(monkeypatch, tmp_path):
|
||||
_force_windows(monkeypatch)
|
||||
pkg = tmp_path / "onnxruntime"
|
||||
pkg.mkdir() # package exists, but no capi/onnxruntime.dll inside
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue