bge_m3_embedding_server/embedder/
trt_cache.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 cache path construction, inspection, and durability.
16//!
17//! Why a dedicated module? Investigation of a production incident
18//! showed two consecutive cold starts producing identical 172 s recompile
19//! times for `1×8192` — the `{cache_dir}/trt-engines/` directory on EFS was
20//! NOT being reused between container restarts. The most plausible root cause
21//! is that ECS SIGKILL on OOM (`exitCode 137`) interrupts the container before
22//! the kernel's writeback timer flushes buffered writes to EFS. The EFS inode
23//! lists the engine files, but their data blocks are zero-length or partial,
24//! so ORT/TRT silently treats them as cache misses and rebuilds from scratch.
25//!
26//! This module:
27//! 1. Constructs the cache directory path (stable, no per-container ephemera).
28//! 2. Inspects the directory at startup so the operator-visible INFO log shows
29//!    "found N cached engines" or "empty (will compile)" — without this we
30//!    cannot tell from `CloudWatch` whether the EFS mount is actually persisting.
31//! 3. Exposes an explicit `fsync_cache_dir` that flushes both file data and
32//!    directory metadata to disk, called after each successful TRT engine
33//!    compile so an OOM-kill never strands a partially-written engine plan.
34//!
35//! TRT plan files embed `(GPU compute capability, CUDA version, TRT version,
36//! ONNX model SHA, builder config)`. Within a homogeneous ASG (same instance
37//! family, same AMI) these are stable, so the cache is reusable per-EC2-host.
38//! ASGs that mix instance families (T4 → A10G) will see expected cache misses
39//! when a task lands on a different GPU architecture.
40//!
41//! Several items below are referenced only from `session.rs` under
42//! `#[cfg(all(not(target_os = "macos"), feature = "tensorrt"))]` — the CPU /
43//! macOS build legitimately never calls them. The `#[allow(dead_code)]`
44//! attribute below silences the resulting unused-warning under those builds;
45//! the unit tests in this module keep the items exercised on every CI target.
46//!
47//! Submodules:
48//! - `paths`: cache directory paths and [`TrtCacheInfo`].
49//! - `inspect`: startup inspection and EFS write-probe.
50//! - `enumerate`: SM-aware engine plan enumeration and counting.
51//! - `prewarm_log`: operator-visible prewarm basename logging.
52//! - `fsync`: post-compile cache durability.
53
54#![allow(dead_code, unused_imports)]
55
56mod enumerate;
57mod fsync;
58mod inspect;
59mod paths;
60mod prewarm_log;
61
62#[cfg(test)]
63mod tests;
64
65pub(crate) use enumerate::{
66    count_engine_files, count_engine_files_for_sm, engine_basenames_for_sm,
67    engine_basenames_in_dir_sorted, engine_files_for_sm, matches_sm_suffix,
68};
69pub(crate) use fsync::fsync_cache_dir;
70pub(crate) use inspect::{ensure_and_inspect, log_cache_state};
71pub(crate) use paths::{
72    TRT_ENGINE_SUBDIR, TRT_TIMING_SUBDIR, TrtCacheInfo, engine_cache_path, timing_cache_path,
73};
74pub(crate) use prewarm_log::{
75    log_engine_basenames_before_prewarm, log_engine_basenames_before_prewarm_for_sm,
76};
bge_m3_embedding_server/embedder/trt_cache.rs

bge_m3_embedding_server/embedder/
trt_cache.rs
