bge_m3_embedding_server/embedder/worker/
prewarm_strict.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//! Prewarm postcondition readiness gate (`BGE_M3_PREWARM_STRICT`).
16
17use crate::embedder::trt_warmup::{
18    prewarm_persistence_postcondition_failed, prewarm_persistence_suspicious_undercount,
19};
20
21/// Decides whether a worker should refuse to signal ready after its prewarm
22/// sweep based on the on-disk persistence postcondition.
23///
24/// Returns `true` iff `strict` is `true` AND at least one of the
25/// [`prewarm_persistence_postcondition_failed`] /
26/// [`prewarm_persistence_suspicious_undercount`] predicates fires for the
27/// given `(fresh_compiles, engine_count_after)` snapshot.
28///
29/// The signature deliberately accepts primitive `usize` values rather than
30/// `&PrewarmStats` so the unit tests in `worker/tests/prewarm_strict.rs`
31/// stay decoupled from the `trt_warmup::PrewarmStats` struct shape; this
32/// also lets the predicate be reused at future call sites (e.g. an admin
33/// endpoint that wants to surface the same decision) without dragging in
34/// the rest of the prewarm statistics.
35///
36/// # Strict-mode semantics
37///
38/// Strict mode (`prewarm_strict=true`) only blocks readiness when
39/// `engine_count_after == 0` — i.e. **complete zero-plan failure** where
40/// fresh compiles occurred but not a single `.engine` file landed on disk.
41/// This is the catastrophic failure mode where every worker hits TRT autotuner
42/// OOM mid-build, leaving the cache empty and every subsequent real request
43/// returning HTTP 500.
44///
45/// **Partial undercounts** (e.g. 1 engine persisted out of 16 compiled) do
46/// NOT block readiness — workers will serve traffic using the one cached shape
47/// and JIT-compile any missing shapes on first request. This is acceptable:
48/// partial persistence is most commonly caused by TRT's subgraph fusing
49/// (multiple `(batch, seq)` shapes sharing one engine file), not by a hard
50/// persistence failure.
51///
52/// If threshold-based undercount blocking becomes necessary in the future,
53/// the [`prewarm_persistence_suspicious_undercount`] branch already has the
54/// scaffolding — promote it from WARN to a readiness gate here.
55pub(super) fn should_fail_readiness(
56    fresh_compiles: usize,
57    engine_count_after: usize,
58    strict: bool,
59) -> bool {
60    if !strict {
61        return false;
62    }
63    prewarm_persistence_postcondition_failed(fresh_compiles, engine_count_after)
64        || prewarm_persistence_suspicious_undercount(fresh_compiles, engine_count_after)
65}
bge_m3_embedding_server/embedder/worker/prewarm_strict.rs

bge_m3_embedding_server/embedder/worker/
prewarm_strict.rs
