bge_m3_embedding_server/embedder/trt_cache/enumerate.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//! SM-aware engine plan enumeration and counting.
16
17use std::path::{Path, PathBuf};
18
19/// Returns `true` when `name` is a TRT engine plan basename whose `_smXX`
20/// suffix matches the requested SM exactly.
21///
22/// ORT names every TRT engine plan with a `_smXX.engine` suffix tied to the
23/// GPU compute capability that built it (`sm75` = T4, `sm86` = A10G, `sm89` =
24/// L40S/L4, `sm120` = Blackwell). The match must be **strict** — `sm12` must
25/// not match `sm120.engine`, or a B200 worker would happily believe a Hopper
26/// plan is usable. We accomplish this by anchoring on the leading underscore:
27/// the suffix tested is `"_{sm}.engine"`, so `_sm12.engine` and `_sm120.engine`
28/// occupy disjoint string sets.
29///
30/// Pure, no I/O — easy to unit-test.
31#[must_use]
32pub(crate) fn matches_sm_suffix(name: &str, sm: &str) -> bool {
33 let suffix = format!("_{sm}.engine");
34 name.ends_with(&suffix)
35}
36
37/// Returns full paths of `.engine` files under `engine_dir` that match the
38/// requested SM, or all `.engine` files when `sm` is `None`.
39///
40/// Single source of truth for engine enumeration: every other function in
41/// this module that needs to enumerate engine plans (count, basenames,
42/// operator log line) delegates here. Centralising the filter is the design
43/// constraint behind the SM-aware refactor — a future caller that grew its
44/// own `read_dir` loop would silently regress the heterogeneous-SM safety
45/// invariant.
46///
47/// Failure modes (directory missing or unreadable) collapse to an empty
48/// `Vec`, mirroring the legacy `count_engine_files` behaviour and the
49/// "operator-visible, not load-bearing" stance of this module.
50pub(crate) fn engine_files_for_sm(engine_dir: &Path, sm: Option<&str>) -> Vec<PathBuf> {
51 let Ok(read_dir) = std::fs::read_dir(engine_dir) else {
52 return Vec::new();
53 };
54 read_dir
55 .flatten()
56 .filter(|e| e.file_type().is_ok_and(|t| t.is_file()))
57 .filter_map(|e| {
58 let name = e.file_name().to_string_lossy().into_owned();
59 if !name.ends_with(".engine") {
60 return None;
61 }
62 match sm {
63 Some(target) if !matches_sm_suffix(&name, target) => None,
64 _ => Some(e.path()),
65 }
66 })
67 .collect()
68}
69
70/// Counts `.engine` files in `dir` that match the requested SM.
71///
72/// Pass `Some("sm120")` to count only Blackwell plans; pass `None` to count
73/// every `.engine` file regardless of suffix (legacy behaviour, also what
74/// the wrapper [`count_engine_files`] does). Returns `0` when the directory
75/// does not exist or cannot be read.
76///
77/// Crate-visible (not `pub(super)`) because the warmup-only postcondition
78/// in `lib.rs` calls it directly to apply the same SM filter as the
79/// per-worker prewarm path.
80pub(crate) fn count_engine_files_for_sm(dir: &Path, sm: Option<&str>) -> usize {
81 engine_files_for_sm(dir, sm).len()
82}
83
84/// Counts every `.engine` file in `dir` regardless of `_smXX` suffix.
85///
86/// Backwards-compatible wrapper around [`count_engine_files_for_sm`] with
87/// `sm = None`. Retained because the operator-visible startup cache log
88/// (`trt cache: found cached engines` in [`crate::embedder::trt_cache::log_cache_state`]) reports the
89/// total disk footprint, not the SM-filtered subset. Callers that drive
90/// the prewarm postcondition use the SM-aware variant directly.
91pub(crate) fn count_engine_files(dir: &Path) -> usize {
92 count_engine_files_for_sm(dir, None)
93}
94
95/// Returns sorted basenames of `.engine` files under `engine_dir` that match
96/// the requested SM, or all `.engine` basenames when `sm` is `None`.
97///
98/// Used for operator-visible logs before TRT prewarm; filenames are **not**
99/// a reliable `(batch, seq)` key for dynamic models (see `CLAUDE.md`).
100pub(crate) fn engine_basenames_for_sm(
101 engine_dir: &Path,
102 sm: Option<&str>,
103) -> std::io::Result<Vec<String>> {
104 // We re-implement the read instead of delegating to `engine_files_for_sm`
105 // because the latter swallows I/O errors (returns empty on missing dir),
106 // whereas this function needs to propagate them so the operator log can
107 // surface "could not read engine cache directory for basename listing".
108 let read_dir = std::fs::read_dir(engine_dir)?;
109 let mut basenames: Vec<String> = read_dir
110 .flatten()
111 .filter(|e| e.file_type().is_ok_and(|t| t.is_file()))
112 .filter_map(|e| {
113 let name = e.file_name().to_string_lossy().into_owned();
114 if !name.ends_with(".engine") {
115 return None;
116 }
117 match sm {
118 Some(target) if !matches_sm_suffix(&name, target) => None,
119 _ => Some(name),
120 }
121 })
122 .collect();
123 basenames.sort();
124 Ok(basenames)
125}
126
127/// Returns sorted basenames of every `.engine` file under `engine_dir`,
128/// regardless of `_smXX` suffix.
129///
130/// Backwards-compatible wrapper around [`engine_basenames_for_sm`] with
131/// `sm = None`. Retained for the unit tests below; production prewarm log
132/// emission goes through [`crate::embedder::trt_cache::log_engine_basenames_before_prewarm_for_sm`].
133pub(crate) fn engine_basenames_in_dir_sorted(engine_dir: &Path) -> std::io::Result<Vec<String>> {
134 engine_basenames_for_sm(engine_dir, None)
135}