Skip to main content

bge_m3_embedding_server/embedder/trt_warmup/
postcondition.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//! `TensorRT` engine pre-warm persistence postconditions.
16//!
17//! Two diagnostic predicates the worker calls after a prewarm sweep:
18//!
19//! * [`prewarm_persistence_postcondition_failed`] — **fatal** ERROR signal.
20//!   Catches the catastrophic fresh-compiles → 0 engines on disk pattern
21//!   that produced silent-persistence startup failures in production.
22//! * [`prewarm_persistence_suspicious_undercount`] — **non-fatal** WARN
23//!   signal. Retained for future extension; currently silent whenever
24//!   `engine_count_after > 0` (see inline note).
25//!
26//! Both are pure functions over `(fresh_compiles, engine_count_after)` so
27//! they can be unit-tested without spinning up an ORT session or a
28//! filesystem fixture.  The companion fixture-backed tests in `tests.rs`
29//! exercise them through the `count_engine_files` snapshot mechanism that
30//! the worker uses in production.
31//!
32//! ## Why `engine_count_after`, not `engine_count_delta`
33//!
34//! ORT's TRT EP stores one profile-based `.engine` file per fused subgraph
35//! that covers all `(batch, seq)` shapes compiled so far via `[min, max]`
36//! ranges per input dimension.  When a new shape falls inside the existing
37//! range the file is reused (cache hit); when it falls outside the range the
38//! EP rewrites the file in-place with an expanded profile.  Either way the
39//! on-disk file count stays at 1 after the first compile — `delta == 0` is
40//! the normal steady-state, NOT a persistence failure.
41//!
42//! The only actionable signal is `engine_count_after == 0`: the TRT EP
43//! reported `Ok(_)` from `session.run()` yet wrote no engine file at all.
44//! That is the exact failure mode from that incident class.
45
46/// Decides whether a single worker's prewarm postcondition is violated.
47///
48/// The postcondition: if at least one shape on this worker reported a
49/// **fresh compile** (`succeeded && !cache_hit`) but the on-disk `.engine`
50/// file count is still zero, the TRT EP almost certainly emitted `Ok(_)`
51/// from `session.run()` without actually persisting the engine plan.  This
52/// is the silent-persistence failure mode where the TRT EP reports compile
53/// success but writes no engine files (many compile-success events / 0 engines
54/// on disk).
55///
56/// Returning `true` should produce an `ERROR` log so operators see the
57/// failure in `CloudWatch` immediately.  Non-fresh-compile shards (cache hits
58/// only) and shards where at least one engine file exists are accepted.
59///
60/// The check is keyed on `engine_count_after == 0` rather than
61/// `engine_count_delta <= 0`.  ORT's TRT EP rewrites its single
62/// profile-based engine file in-place as more shapes are compiled
63/// (`delta == 0` at steady state), so a delta-based rule would produce
64/// false-positive ERRORs on every shape after the first compile.
65#[must_use]
66pub(crate) fn prewarm_persistence_postcondition_failed(
67    fresh_compiles: usize,
68    engine_count_after: usize,
69) -> bool {
70    fresh_compiles > 0 && engine_count_after == 0
71}
72
73/// Minimum `fresh_compiles` count below which the suspicious-undercount
74/// check is suppressed.
75///
76/// Retained for documentation and the test pin; the current implementation
77/// of [`prewarm_persistence_suspicious_undercount`] is always silent when
78/// `engine_count_after > 0`, making this floor rarely reached.
79pub(crate) const SUSPICIOUS_UNDERCOUNT_MIN_FRESH: usize = 2;
80
81/// Decides whether the on-disk `.engine` count is **suspiciously low**
82/// relative to the number of fresh compiles, in a way not already caught by
83/// [`prewarm_persistence_postcondition_failed`].
84///
85/// This is a **non-fatal diagnostic signal** intended to back a `WARN`
86/// log only — it must never cause the process to exit non-zero.
87///
88/// **Current behaviour:** always returns `false` when `engine_count_after > 0`.
89/// ORT's TRT EP writes one profile-based engine file that covers all compiled
90/// shapes; `engine_count_delta == 0` after the first compile is the normal
91/// steady-state, not an anomaly.  The only meaningful anomaly is
92/// `engine_count_after == 0`, which is already caught by the ERROR predicate
93/// above.  Future operators who need a ratio-based WARN for multi-engine
94/// workloads can re-introduce it here without changing call sites.
95#[must_use]
96pub(crate) fn prewarm_persistence_suspicious_undercount(
97    fresh_compiles: usize,
98    engine_count_after: usize,
99) -> bool {
100    // Silence entirely when engine files exist on disk.  TRT EP in-place
101    // profile extension means delta == 0 is healthy; flagging it would
102    // produce false-positive WARNs on every shape after the first compile.
103    if engine_count_after > 0 {
104        return false;
105    }
106    // ERROR is already covering this (fresh > 0 && after == 0); no need to
107    // double-fire at WARN level on the same evidence.
108    if prewarm_persistence_postcondition_failed(fresh_compiles, engine_count_after) {
109        return false;
110    }
111    if fresh_compiles < SUSPICIOUS_UNDERCOUNT_MIN_FRESH {
112        return false;
113    }
114    // Unreachable: (engine_count_after == 0 && fresh_compiles >= MIN_FRESH)
115    // implies postcondition_failed == true, handled above.
116    false
117}