bge_m3_embedding_server/embedder/
types.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//! Public DTOs and the internal `EmbedRequest` enum exchanged between the
16//! pool and the worker threads.
17
18use anyhow::Result;
19use tokio::sync::oneshot;
20
21/// Sparse embedding output from the BGE-M3 sparse-linear projection layer.
22///
23/// Represents a document as a sparse vector over the tokenizer vocabulary.
24/// Token IDs with zero ReLU-gated score are omitted.
25#[derive(Debug, Clone)]
26pub struct SparseEmbedding {
27    /// Sorted vocabulary token IDs with non-zero ReLU-gated weight.
28    pub indices: Vec<usize>,
29    /// Corresponding ReLU-gated projection scores, in the same order as `indices`.
30    pub values: Vec<f32>,
31}
32
33/// Paired dense + sparse embeddings produced from a single forward pass.
34#[derive(Debug, Clone)]
35pub struct DualEmbedding {
36    pub dense: Vec<f32>,
37    pub sparse: SparseEmbedding,
38}
39
40/// OS headroom reserved for kernel, stack, ORT arena, and other non-model
41/// allocations. Subtracted from available memory before computing
42/// per-worker workspace.
43pub(crate) const OS_HEADROOM_BYTES: usize = 256 * 1024 * 1024; // 256 MiB
44
45/// Per-request diagnostic statistics captured inside the worker and forwarded
46/// to the handler layer for inclusion in the completion log event.
47#[derive(Debug, Clone, Copy, Default)]
48pub struct EmbedStats {
49    /// Number of bin-packed chunks the batch was split into.
50    pub chunks: usize,
51    /// Maximum tokenized sequence length across all chunks.
52    pub max_chunk_seq: usize,
53    /// Total token-positions processed (sum of `seq_len` for all inputs).
54    pub total_token_positions: usize,
55    /// Time spent tokenizing all inputs (milliseconds).
56    pub tokenize_ms: u64,
57    /// Total time spent in ORT `session.run()` across all chunks (milliseconds).
58    pub inference_ms: u64,
59    /// Minimum token sequence length across all inputs in the batch.
60    pub seq_len_min: usize,
61    /// Maximum token sequence length across all inputs in the batch.
62    pub seq_len_max: usize,
63    /// Mean token sequence length across all inputs (integer, truncated).
64    pub seq_len_mean: usize,
65    /// 95th-percentile token sequence length across all inputs in the batch.
66    ///
67    /// Index is `(n * 95) / 100` on a sorted copy of the per-input lengths.
68    pub seq_len_p95: usize,
69}
70
71pub(crate) enum EmbedRequest {
72    /// Dense (float32) embedding inference on a batch of texts.
73    Dense {
74        texts: Vec<String>,
75        reply: oneshot::Sender<Result<(Vec<Vec<f32>>, EmbedStats)>>,
76    },
77    /// Sparse (SPLADE-style) embedding inference on a batch of texts.
78    Sparse {
79        texts: Vec<String>,
80        reply: oneshot::Sender<Result<(Vec<SparseEmbedding>, EmbedStats)>>,
81    },
82    /// Computes dense and sparse embeddings from a single forward pass per chunk.
83    Both {
84        texts: Vec<String>,
85        reply: oneshot::Sender<Result<(Vec<DualEmbedding>, EmbedStats)>>,
86    },
87    /// Internal: used during startup probe to run a single batch and measure
88    /// peak RSS delta. Workers only process this before `ready` is set.
89    Probe {
90        texts: Vec<String>,
91        reply: oneshot::Sender<Result<ProbeResult>>,
92    },
93    /// Adaptive background warmup: asks a worker to compile (or confirm as
94    /// cached) the TRT engine for `(batch, seq)`.  The worker replies on
95    /// `ack` with the compile duration in milliseconds, or an error if the
96    /// shape failed.  Only meaningful on TRT EP; on CPU/CUDA workers the
97    /// worker returns `Ok(0)` immediately.
98    AdaptiveWarmup {
99        batch: usize,
100        seq: usize,
101        ack: oneshot::Sender<anyhow::Result<u64>>,
102    },
103}
104
105/// Sender half of the JIT-suspect channel.
106///
107/// Workers hold an optional clone of this sender and call `try_send`
108/// (non-blocking, drops if full) after any inference whose `inference_ms`
109/// equals or exceeds the TRT cache-hit threshold.
110pub(crate) type JitSuspectSender = tokio::sync::mpsc::Sender<(usize, usize)>;
111
112/// Result of a single probe `session.run()` call.
113#[derive(Debug)]
114pub(crate) struct ProbeResult {
115    /// Process RSS (bytes) measured immediately before `session.run()`.
116    pub rss_before: usize,
117    /// Process RSS (bytes) measured immediately after `session.run()`.
118    pub rss_after: usize,
119}
bge_m3_embedding_server/embedder/types.rs

bge_m3_embedding_server/embedder/
types.rs
