bge_m3_embedding_server/
embedder.rs

1// Copyright (c) 2026 J. Patrick Fulton
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Worker-pool–driven BGE-M3 embedding service.
16//!
17//! Submodules:
18//! - `types`: public DTOs and the internal `EmbedRequest` enum.
19//! - `error`: small `ort::Error → anyhow::Error` adapter.
20//! - `model_files`: hf-hub download / cache layout for the ONNX model files.
21//! - `tokenize`: tokenizer load + no-pad tokenization + chunk-array build.
22//! - `session`: ORT execution-provider config and session loading.
23//! - `sm_detect`: per-device GPU compute-capability detection (`smXY`) used
24//!   to filter the TRT engine cache by the worker's own SM.
25//! - `math`: pure dense/sparse math helpers (testable without ORT).
26//! - `dense`: dense embedding pipeline.
27//! - `sparse`: BGE-M3 SPLADE-style sparse embedding pipeline.
28//! - `dual`: paired dense + sparse embedding pipeline (one forward pass).
29//! - `jit_guard`: in-band `TensorRT` JIT admission guard that refuses chunks
30//!   whose sequence length is dangerous and uncovered by the warmed engine
31//!   profile, preventing the process-killing pathological autotuner allocation.
32//! - `trt_cache`: `TensorRT` engine-cache path construction, inspection, and
33//!   durability (fsync after compile). Submodules: `paths`, `inspect`,
34//!   `enumerate`, `prewarm_log`, `fsync`.
35//! - `trt_cache_gc` (feature `cache-gc`, off by default): destructive
36//!   stale-SM engine plan garbage collection — only present in dedicated
37//!   maintenance / dev binaries.
38//! - `trt_warmup`: `TensorRT` engine pre-warming during worker startup.
39//! - `worker`: blocking worker thread, request dispatch, probe wiring.
40//!   Submodules: `config`, `guard`, `trt_retry`, `propagation`, `probe`,
41//!   `prewarm_strict`, `startup`, `run`, `dispatch`, `logging`.
42//! - `pool`: `EmbedPool` async wrapper and test helpers.
43//! - `adaptive_warmup`: adaptive in-process background warmup loop for TRT
44//!   engine cache miss recovery.
45
46pub(crate) mod adaptive_warmup;
47mod dense;
48mod dual;
49mod error;
50pub(crate) mod jit_guard;
51mod math;
52mod model_files;
53mod pool;
54mod session;
55pub(crate) mod sm_detect;
56mod sparse;
57mod tokenize;
58pub(crate) mod trt_cache;
59#[cfg(feature = "cache-gc")]
60pub(crate) mod trt_cache_gc;
61mod trt_warmup;
62mod types;
63mod worker;
64
65pub use pool::EmbedPool;
66pub(crate) use types::{JitSuspectSender, OS_HEADROOM_BYTES};
67pub(crate) use worker::WorkerConfig;
68
69// `SparseEmbedding` is referenced by tests via `crate::embedder::SparseEmbedding`,
70// but is not used outside the module in the non-test build. The cfg(test) gate
71// keeps the binary's unused-import lint clean while preserving the call-site path
72// for tests.
73#[cfg(test)]
74pub(crate) use types::SparseEmbedding;
bge_m3_embedding_server/embedder.rs

bge_m3_embedding_server/
embedder.rs
