bge_m3_embedding_server/embedder/trt_cache/
inspect.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 cache inspection and EFS write-probe.
16
17use std::path::Path;
18
19use super::paths::{TrtCacheInfo, engine_cache_path};
20
21pub(crate) fn ensure_and_inspect(cache_dir: &Path) -> TrtCacheInfo {
22    let path = engine_cache_path(cache_dir);
23    if let Err(e) = std::fs::create_dir_all(&path) {
24        tracing::warn!(
25            cache_path = %path.display(),
26            error = %e,
27            "TensorRT engine cache directory could not be created; \
28             engine caching may be unavailable for this container"
29        );
30        return TrtCacheInfo {
31            path,
32            engine_count: 0,
33            profile_count: 0,
34        };
35    }
36
37    run_write_probe(&path);
38
39    let (engine_count, profile_count) = count_cache_entries(&path);
40    TrtCacheInfo {
41        path,
42        engine_count,
43        profile_count,
44    }
45}
46/// One-shot write/read/delete sentinel probe under `dir`.
47///
48/// The probe rules in/out the "EFS access point POSIX uid mapping blocks
49/// regular `creat(2) + write(2) + unlink(2)`" hypothesis at next boot.
50/// It writes a 9-byte sentinel `b"trt-probe"` to `<dir>/.write_probe`,
51/// reads it back, verifies the round-trip, and deletes the file. Each
52/// distinct failure mode (`create`, `write`, `read`, `mismatch`, `unlink`)
53/// fires a tagged `ERROR` so operators can disambiguate without
54/// instrumenting the filesystem from outside.
55///
56/// The probe path is fixed (`.write_probe`) so it is greppable and never
57/// collides with TRT's own filenames (which all begin with
58/// `TensorrtExecutionProvider_TRTKernel_`). Hidden by leading dot so
59/// `count_cache_entries` and `count_engine_files` ignore it without any
60/// extra filtering.
61pub(super) fn run_write_probe(dir: &Path) {
62    use std::io::{Read, Write};
63
64    const PROBE_NAME: &str = ".write_probe";
65    const PROBE_DATA: &[u8] = b"trt-probe";
66
67    let probe_path = dir.join(PROBE_NAME);
68
69    // Best-effort cleanup of any stale probe file from a previous boot —
70    // this is not the failure mode we are testing, so silently ignore the
71    // `NotFound` case and let the create call below surface real problems.
72    let _ = std::fs::remove_file(&probe_path);
73
74    let create = std::fs::OpenOptions::new()
75        .write(true)
76        .create_new(true)
77        .open(&probe_path);
78    let mut file = match create {
79        Ok(f) => f,
80        Err(e) => {
81            tracing::error!(
82                cache_path = %dir.display(),
83                phase = "create",
84                error = %e,
85                "trt cache: write probe failed"
86            );
87            return;
88        }
89    };
90
91    if let Err(e) = file.write_all(PROBE_DATA) {
92        tracing::error!(
93            cache_path = %dir.display(),
94            phase = "write",
95            error = %e,
96            "trt cache: write probe failed"
97        );
98        let _ = std::fs::remove_file(&probe_path);
99        return;
100    }
101    drop(file);
102
103    let mut buf = Vec::with_capacity(PROBE_DATA.len());
104    let mut reader = match std::fs::File::open(&probe_path) {
105        Ok(f) => f,
106        Err(e) => {
107            tracing::error!(
108                cache_path = %dir.display(),
109                phase = "read",
110                error = %e,
111                "trt cache: write probe failed"
112            );
113            let _ = std::fs::remove_file(&probe_path);
114            return;
115        }
116    };
117    if let Err(e) = reader.read_to_end(&mut buf) {
118        tracing::error!(
119            cache_path = %dir.display(),
120            phase = "read",
121            error = %e,
122            "trt cache: write probe failed"
123        );
124        let _ = std::fs::remove_file(&probe_path);
125        return;
126    }
127    drop(reader);
128
129    if buf != PROBE_DATA {
130        tracing::error!(
131            cache_path = %dir.display(),
132            phase = "mismatch",
133            bytes_written = PROBE_DATA.len(),
134            bytes_read_back = buf.len(),
135            "trt cache: write probe failed"
136        );
137        let _ = std::fs::remove_file(&probe_path);
138        return;
139    }
140
141    if let Err(e) = std::fs::remove_file(&probe_path) {
142        // Read+write succeeded but unlink didn't — still emit the success
143        // INFO so the success-path counter is accurate, then a separate
144        // ERROR documenting the unlink failure (the directory will
145        // accumulate `.write_probe` files across restarts otherwise).
146        tracing::info!(
147            cache_path = %dir.display(),
148            bytes_written = PROBE_DATA.len(),
149            bytes_read_back = buf.len(),
150            "trt cache: write probe succeeded"
151        );
152        tracing::error!(
153            cache_path = %dir.display(),
154            phase = "unlink",
155            error = %e,
156            "trt cache: write probe failed"
157        );
158        return;
159    }
160
161    tracing::info!(
162        cache_path = %dir.display(),
163        bytes_written = PROBE_DATA.len(),
164        bytes_read_back = buf.len(),
165        "trt cache: write probe succeeded"
166    );
167}
168/// Counts `.engine` and `.profile` files in `dir`. Returns `(0, 0)` if the
169/// directory cannot be read (also covers "not yet created" cases).
170pub(super) fn count_cache_entries(dir: &Path) -> (usize, usize) {
171    let Ok(read_dir) = std::fs::read_dir(dir) else {
172        return (0, 0);
173    };
174    let mut engines = 0usize;
175    let mut profiles = 0usize;
176    for entry in read_dir.flatten() {
177        let Ok(file_type) = entry.file_type() else {
178            continue;
179        };
180        if !file_type.is_file() {
181            continue;
182        }
183        let name = entry.file_name();
184        let lossy = name.to_string_lossy();
185        if lossy.ends_with(".engine") {
186            engines += 1;
187        } else if lossy.ends_with(".profile") {
188            profiles += 1;
189        }
190    }
191    (engines, profiles)
192}
193/// Emits a single INFO log line describing the TRT cache state at startup.
194///
195/// Operators reading `CloudWatch` should be able to tell at a glance whether
196/// the cache is being reused. The message wording is stable and grep-friendly:
197/// `trt cache: ...`.
198pub(crate) fn log_cache_state(info: &TrtCacheInfo) {
199    if info.engine_count == 0 {
200        tracing::info!(
201            cache_path = %info.path.display(),
202            engine_count = 0,
203            "trt cache: empty (will compile)"
204        );
205    } else {
206        tracing::info!(
207            cache_path = %info.path.display(),
208            engine_count = info.engine_count,
209            profile_count = info.profile_count,
210            "trt cache: found cached engines"
211        );
212    }
213}
bge_m3_embedding_server/embedder/trt_cache/inspect.rs

bge_m3_embedding_server/embedder/trt_cache/
inspect.rs
