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}