Skip to main content

bge_m3_embedding_server/
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//! Startup memory probe and cost-model coefficient fitter.
16//!
17//! Submodules:
18//! - `runner`: drives the `(batch, seq)` shape sweep on the leader worker
19//!   (`run_probe`, `PROBE_SHAPES`, the absolute-RSS guard, the arena warm-up).
20//! - `cache`: persistent EFS-backed cache for fitted probe coefficients
21//!   (`try_load_probe_cache`, `save_probe_cache`, `ProbeCache`).
22//! - `fit`: ordinary least-squares fitter for the quadratic cost model
23//!   (`fit_cost_model`, `DataPoint`).
24//! - `corpus`: helpers that synthesize probe texts from the curated corpus.
25//! - `validate`: tokenizer + ndarray shape check at the configured `max_seq`
26//!   (no `session.run()`).
27
28mod cache;
29mod corpus;
30mod fit;
31mod runner;
32mod validate;
33
34pub(crate) use cache::{save_probe_cache, try_load_probe_cache};
35pub(crate) use runner::run_probe;
36
37#[cfg(test)]
38mod tests;