From cdba2eccddaaeb469abb05728e902c43954b5a51 Mon Sep 17 00:00:00 2001 From: Matthew Jackson Date: Wed, 15 Jul 2026 11:18:52 -0700 Subject: [PATCH] feat(core): gate ONNX transforms behind a default-on `ml` feature (static/lexical builds) (#2165) ## Description `TextCrusher` and the BM25 relevance path can run without the ONNX-backed ML stack, but `headroom-core` previously compiled `ort`, `fastembed`, and `magika` unconditionally. This made lexical-only downstream consumers carry the ONNX Runtime dependency even when they never used embedding relevance or Magika detection. This PR makes those ML crates optional behind a new default-on `ml` Cargo feature. Default builds keep the existing ML-backed behavior. Consumers that only need lexical compression can opt out with `default-features = false`; in that mode the ML modules are compiled out and the relevance path falls back to BM25. ## Type of Change - [ ] Bug fix (non-breaking change that fixes an issue) - [x] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - `crates/headroom-core/Cargo.toml`: marks `ort`, `fastembed`, and `magika` optional; adds default-on `ml = ["dep:ort", "dep:fastembed", "dep:magika"]`. - `crates/headroom-core/src/lib.rs`: gates the shared ONNX CPU helper behind `ml`. - `crates/headroom-core/src/relevance/embedding.rs`: gates the fastembed implementation behind `ml` and provides a no-ml stub with the same scorer surface so `HybridScorer` naturally falls back to BM25. - `crates/headroom-core/src/transforms/detection.rs`: gates the Magika tier behind `ml`; no-ml builds start at the existing unidiff/plain-text fallback tiers. - `crates/headroom-core/src/transforms/mod.rs`: gates the Magika module and re-exports behind `ml`. ## Testing - [x] Default build compiles (`cargo build -p headroom-core`) - [x] Lexical-only build compiles (`cargo build -p headroom-core --no-default-features`) - [x] Default tests pass (`cargo test -p headroom-core`) - [x] Lexical-only tests pass (`cargo test -p headroom-core --no-default-features`) - [x] Dependency tree checked for no-ml build (`cargo tree -p headroom-core --no-default-features` contains no `fastembed`, `magika`, or `ort` packages) - [ ] Manual testing performed ## Real Behavior Proof - Environment: Windows 11 review worktree, Rust/Cargo workspace. - Exact command / steps: - `cargo build -p headroom-core` - `cargo build -p headroom-core --no-default-features` - `cargo test -p headroom-core` - `cargo test -p headroom-core --no-default-features` - `cargo tree -p headroom-core --no-default-features` - Observed result: both feature configurations build and test cleanly. The no-default dependency tree does not include `fastembed`, `magika`, or `ort`, while the default build still compiles the ML path. - Not tested: model-backed `RUN_FASTEMBED_TESTS=1` cases that require downloading the embedding model; those remain env-gated as before. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [x] New and existing unit tests pass locally with my changes - [ ] I have updated the CHANGELOG.md if applicable ## Additional Notes The no-ml build intentionally degrades embedding relevance to the existing unavailable-model behavior, so `HybridScorer` takes its BM25 fallback path. Magika detection is skipped when `ml` is disabled; detection then proceeds through unidiff and plain-text fallback tiers. --------- Co-authored-by: Matthew Jackson Co-authored-by: JerrettDavis --- crates/headroom-core/Cargo.toml | 18 ++++- crates/headroom-core/src/lib.rs | 1 + .../headroom-core/src/relevance/embedding.rs | 78 +++++++++++++++++++ .../headroom-core/src/transforms/detection.rs | 16 ++++ crates/headroom-core/src/transforms/mod.rs | 2 + 5 files changed, 111 insertions(+), 4 deletions(-) diff --git a/crates/headroom-core/Cargo.toml b/crates/headroom-core/Cargo.toml index 80076965e..198ace6f0 100644 --- a/crates/headroom-core/Cargo.toml +++ b/crates/headroom-core/Cargo.toml @@ -67,7 +67,7 @@ flate2 = "1" # crate depends on `ort`, which is already in our dep tree via # `fastembed`, so adding it doesn't pull a new ML runtime — both # crates share the ONNX Runtime singleton. -magika = "1" +magika = { version = "1", optional = true } # `unidiff` is the Stage-3d Tier-2 diff detector. We use the parser # itself as the "is this a diff?" oracle — anything that successfully # parses to ≥1 PatchedFile is a diff. The deterministic parser @@ -132,15 +132,25 @@ http = "1" # 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 = [ +fastembed = { version = "5", default-features = false, optional = true, features = [ "hf-hub-rustls-tls", "ort-load-dynamic", "image-models", ] } -ort = { version = "2.0.0-rc.12", default-features = false, features = ["load-dynamic"] } +ort = { version = "2.0.0-rc.12", default-features = false, optional = true, features = ["load-dynamic"] } [features] -default = [] +# `ml` is ON by default, so a stock build is byte-for-byte what it was before: +# the ONNX-backed transforms (fastembed embeddings, magika detection, the +# smart-crusher ML path) are all compiled in. Turning it OFF +# (`default-features = false`) drops `ort`/`fastembed`/`magika` from the tree +# entirely, so a consumer that only uses the lexical path (`TextCrusher` / +# BM25 relevance) builds with no ONNX Runtime at all — which is what lets that +# consumer ship a fully-static musl binary. See the source `#[cfg(feature = +# "ml")]` gates (TODO: the module + dispatch gating is the collaborative half +# of this change — flagged in the PR). +default = ["ml"] +ml = ["dep:ort", "dep:fastembed", "dep:magika"] # Compile in the Redis CCR backend. Enable for multi-worker deployments # that want a shared CCR store with no sticky-session at the LB. The # SQLite backend (always compiled) is the production default for diff --git a/crates/headroom-core/src/lib.rs b/crates/headroom-core/src/lib.rs index 753fcdc76..8a907045a 100644 --- a/crates/headroom-core/src/lib.rs +++ b/crates/headroom-core/src/lib.rs @@ -4,6 +4,7 @@ pub mod auth_mode; pub mod cache_control; pub mod ccr; pub mod compression_policy; +#[cfg(feature = "ml")] mod onnx_cpu; pub mod relevance; pub mod signals; diff --git a/crates/headroom-core/src/relevance/embedding.rs b/crates/headroom-core/src/relevance/embedding.rs index 32c658d29..02facf02b 100644 --- a/crates/headroom-core/src/relevance/embedding.rs +++ b/crates/headroom-core/src/relevance/embedding.rs @@ -26,8 +26,10 @@ //! kernels, same weights — embeddings agree to floating-point //! representation. Cosine similarity agrees to ~1e-6. +#[cfg(feature = "ml")] use std::sync::Mutex; +#[cfg(feature = "ml")] use fastembed::{EmbeddingModel, InitOptions, TextEmbedding}; use super::base::{RelevanceScore, RelevanceScorer}; @@ -39,6 +41,7 @@ use super::base::{RelevanceScore, RelevanceScorer}; /// for backwards compatibility but `is_available()` returns `false` /// when the inner model failed to load (mimicking Python's /// "sentence-transformers not installed" branch). +#[cfg(feature = "ml")] pub struct EmbeddingScorer { pub model_name: String, /// `None` when model load failed — `is_available()` returns false @@ -55,6 +58,7 @@ pub struct EmbeddingScorer { model: Option>, } +#[cfg(feature = "ml")] impl Default for EmbeddingScorer { /// Returns an unloaded scorer (model = None, is_available = false). /// @@ -75,6 +79,7 @@ impl Default for EmbeddingScorer { } } +#[cfg(feature = "ml")] impl EmbeddingScorer { /// Construct the scorer with the default model /// (BAAI/bge-small-en-v1.5). May trigger a one-time HF Hub @@ -118,6 +123,7 @@ impl EmbeddingScorer { } } +#[cfg(feature = "ml")] impl RelevanceScorer for EmbeddingScorer { fn score(&self, item: &str, context: &str) -> RelevanceScore { if item.is_empty() || context.is_empty() { @@ -207,8 +213,66 @@ impl RelevanceScorer for EmbeddingScorer { } } +/// Lexical-only build stub. +/// +/// Without the `ml` feature the fastembed/ONNX backend is compiled out +/// entirely. `EmbeddingScorer` still exists so `HybridScorer` and +/// `create_scorer` compile unchanged, but it carries no model and is +/// permanently unavailable: `is_available()` is always `false` and the +/// scoring methods return the same empty scores the ml build produces +/// when its model failed to load. `HybridScorer` therefore takes its +/// BM25 fallback path exactly as it does when embeddings are stubbed. +#[cfg(not(feature = "ml"))] +pub struct EmbeddingScorer { + pub model_name: String, +} + +#[cfg(not(feature = "ml"))] +impl Default for EmbeddingScorer { + fn default() -> Self { + EmbeddingScorer { + model_name: "BAAI/bge-small-en-v1.5".to_string(), + } + } +} + +#[cfg(not(feature = "ml"))] +impl RelevanceScorer for EmbeddingScorer { + fn score(&self, item: &str, context: &str) -> RelevanceScore { + if item.is_empty() || context.is_empty() { + return RelevanceScore::empty("Embedding: empty input"); + } + RelevanceScore::empty("Embedding: model not available") + } + + fn score_batch(&self, items: &[&str], context: &str) -> Vec { + if items.is_empty() { + return Vec::new(); + } + if context.is_empty() { + return items + .iter() + .map(|_| RelevanceScore::empty("Embedding: empty context")) + .collect(); + } + items + .iter() + .map(|_| RelevanceScore::empty("Embedding: model not available")) + .collect() + } + + fn is_available(&self) -> bool { + false + } +} + /// Cosine similarity for two vectors. Clamped to `[0, 1]` since we /// only care about positive similarity (mirrors Python `_cosine_similarity`). +/// +/// Only the `ml` build calls this at runtime (from the fastembed-backed +/// scorer); the lexical-only build keeps it solely for the unit tests +/// that pin its numeric behavior. +#[cfg(any(feature = "ml", test))] fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 { if a.is_empty() || b.is_empty() || a.len() != b.len() { return 0.0; @@ -239,12 +303,14 @@ mod tests { // download). Without the env var, only the offline-safe stub // path is exercised. + #[cfg(feature = "ml")] fn fastembed_enabled() -> bool { std::env::var("RUN_FASTEMBED_TESTS").is_ok() } /// Construct a stub scorer with `model = None` for offline-safe /// tests of the unavailable-path behavior. + #[cfg(feature = "ml")] fn unavailable_scorer() -> EmbeddingScorer { EmbeddingScorer { model_name: "test".to_string(), @@ -252,6 +318,13 @@ mod tests { } } + /// In the lexical-only build the scorer is always unavailable, so + /// `default()` already gives the stub we want to exercise. + #[cfg(not(feature = "ml"))] + fn unavailable_scorer() -> EmbeddingScorer { + EmbeddingScorer::default() + } + #[test] fn cosine_similarity_orthogonal_vectors() { let a = vec![1.0_f32, 0.0, 0.0, 0.0]; @@ -327,6 +400,7 @@ mod tests { // ---------- AVX2 CPU guard (issue #1723) ---------- + #[cfg(feature = "ml")] #[test] fn onnx_guard_matches_cpu_features() { let supported = crate::onnx_cpu::onnx_runtime_supported_by_cpu(); @@ -336,6 +410,7 @@ mod tests { assert!(supported); } + #[cfg(feature = "ml")] #[test] fn try_new_errors_on_unsupported_cpu_instead_of_sigill() { // On a no-AVX2 host the guard must turn the SIGILL into a plain Err @@ -352,6 +427,7 @@ mod tests { // ---------- model-backed tests (gated on RUN_FASTEMBED_TESTS) ---------- + #[cfg(feature = "ml")] #[test] fn fastembed_loads_default_model() { if !fastembed_enabled() { @@ -362,6 +438,7 @@ mod tests { assert_eq!(s.model_name, "BGESmallENV15"); } + #[cfg(feature = "ml")] #[test] fn fastembed_semantic_match_outranks_unrelated() { if !fastembed_enabled() { @@ -378,6 +455,7 @@ mod tests { ); } + #[cfg(feature = "ml")] #[test] fn fastembed_batch_returns_one_score_per_item() { if !fastembed_enabled() { diff --git a/crates/headroom-core/src/transforms/detection.rs b/crates/headroom-core/src/transforms/detection.rs index 43b8d2f75..55a81c471 100644 --- a/crates/headroom-core/src/transforms/detection.rs +++ b/crates/headroom-core/src/transforms/detection.rs @@ -47,6 +47,7 @@ //! for those specifically; not preemptively. use crate::transforms::content_detector::ContentType; +#[cfg(feature = "ml")] use crate::transforms::magika_detector::magika_detect; use crate::transforms::unidiff_detector::is_diff; @@ -61,6 +62,10 @@ pub fn detect(content: &str) -> ContentType { } // ── Tier 1: Magika ────────────────────────────────────────── + // Only present in the `ml` build. Without the ML crates the magika + // detector is compiled out; the chain skips Tier 1 and begins at + // Tier 2, exactly as it would when magika returns PlainText. + #[cfg(feature = "ml")] match magika_detect(content) { Ok(ContentType::PlainText) => { // Magika says "I don't know" or "plain text". Continue @@ -92,12 +97,23 @@ pub fn detect(content: &str) -> ContentType { #[cfg(test)] mod tests { use super::*; + #[cfg(feature = "ml")] use crate::transforms::magika_detector::magika_runtime_available_for_session_init; + #[cfg(feature = "ml")] fn magika_available() -> bool { magika_runtime_available_for_session_init().is_ok() } + // Without the ML crates there is no magika session at all, so the + // detection chain always starts at Tier 2. Report "unavailable" so + // the shared assertions below exercise the same fallthrough path + // they use on a host where magika can't initialize. + #[cfg(not(feature = "ml"))] + fn magika_available() -> bool { + false + } + #[test] fn empty_input_short_circuits_to_plain_text() { assert_eq!(detect(""), ContentType::PlainText); diff --git a/crates/headroom-core/src/transforms/mod.rs b/crates/headroom-core/src/transforms/mod.rs index b4b8ba65f..d65c1deca 100644 --- a/crates/headroom-core/src/transforms/mod.rs +++ b/crates/headroom-core/src/transforms/mod.rs @@ -22,6 +22,7 @@ pub mod detection; pub mod diff_compressor; pub mod live_zone; pub mod log_compressor; +#[cfg(feature = "ml")] pub mod magika_detector; pub mod pipeline; pub mod recommendations; @@ -49,6 +50,7 @@ pub use log_compressor::{ LogCompressionResult, LogCompressor, LogCompressorConfig, LogCompressorStats, LogFormat, LogLevel, LogLine, }; +#[cfg(feature = "ml")] pub use magika_detector::{magika_detect, map_magika_label, MagikaDetectorError}; pub use pipeline::{ CompressionContext, CompressionPipeline, CompressionPipelineBuilder, DiffNoise, DiffOffload,