bge_m3_embedding_server/embedder/worker/
probe.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//! Probe inference helpers for startup measurement and arena priming.
16
17use anyhow::Result;
18use ort::value::TensorRef;
19
20use crate::embedder::error::ort_err;
21use crate::embedder::tokenize::{build_chunk_arrays, tokenize_no_pad};
22use crate::embedder::types::ProbeResult;
23use crate::sysinfo;
24
25/// Runs a single `session.run()` for the probe, measuring RSS before and after.
26///
27/// The probe texts are already tokenized and padded to `pad_to` externally.
28/// This function just runs inference and returns RSS deltas so `probe.rs` can
29/// fit the cost model.
30pub(crate) fn probe_run_dense(
31    session: &mut ort::session::Session,
32    ids_array: &ndarray::Array2<i64>,
33    mask_array: &ndarray::Array2<i64>,
34) -> Result<ProbeResult> {
35    let rss_before = sysinfo::read_process_rss_bytes().unwrap_or(0);
36
37    let ids_tensor = TensorRef::from_array_view(ids_array.view()).map_err(ort_err)?;
38    let mask_tensor = TensorRef::from_array_view(mask_array.view()).map_err(ort_err)?;
39
40    // Run inference (output discarded — we only care about RSS).
41    let _outputs = session
42        .run(ort::inputs! {
43            "input_ids" => ids_tensor,
44            "attention_mask" => mask_tensor,
45        })
46        .map_err(ort_err)?;
47
48    let rss_after = sysinfo::read_process_rss_bytes().unwrap_or(rss_before);
49
50    Ok(ProbeResult {
51        rss_before,
52        rss_after,
53    })
54}
55/// Runs one probe batch: tokenize texts, build padded arrays, call `session.run()`,
56/// and return RSS deltas. Uses `embed_dense`'s no-pad tokenizer path.
57pub(super) fn run_probe_batch(
58    session: &mut ort::session::Session,
59    tokenizer: &tokenizers::Tokenizer,
60    texts: &[String],
61) -> Result<ProbeResult> {
62    let encodings = tokenize_no_pad(tokenizer, texts)?;
63    let pad_to = encodings
64        .iter()
65        .map(|e| e.get_ids().len())
66        .max()
67        .unwrap_or(1)
68        .max(1);
69    let indices: Vec<usize> = (0..texts.len()).collect();
70    let (ids_array, mask_array) = build_chunk_arrays(&encodings, &indices, pad_to)?;
71    probe_run_dense(session, &ids_array, &mask_array)
72}
bge_m3_embedding_server/embedder/worker/probe.rs

bge_m3_embedding_server/embedder/worker/
probe.rs
