bge_m3_embedding_server/embedder/trt_warmup/runner.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//! Per-shape `TensorRT` warmup runner.
16//!
17//! Extracted from `trt_warmup.rs` so the parent module can stay at facade
18//! length while keeping all per-shape logging, cache-hit classification, and
19//! engine-count snapshotting in one focused unit.
20//!
21//! [`run_warmup_shape`] is invoked once per `(batch, seq)` shape by
22//! `trt_prewarm` (both during the dimensional-extreme coverage check and
23//! during the full compile path). [`ShapeRunResult`] is the per-call
24//! summary the caller aggregates into the worker-scoped `PrewarmStats`.
25
26use std::path::Path;
27
28use super::super::trt_cache;
29use super::super::worker::probe_run_dense;
30use super::CACHE_HIT_THRESHOLD_MS;
31
32/// Result of a single [`run_warmup_shape`] call.
33///
34/// Per-shape `.engine` count snapshots are logged inline by
35/// `run_warmup_shape`; only aggregate stats are propagated up to the worker
36/// (via `PrewarmStats`).
37pub(super) struct ShapeRunResult {
38 pub(super) compile_ms: u64,
39 pub(super) fsync_ms: u64,
40 /// `true` when `compile_ms < CACHE_HIT_THRESHOLD_MS` (engine loaded from
41 /// disk cache) rather than compiled from scratch.
42 pub(super) cache_hit: bool,
43 pub(super) succeeded: bool,
44}
45
46/// Runs `session.run()` for a single `(batch, seq)` shape, measures wall
47/// time, classifies the result as a cache hit or fresh compile, and
48/// — on success — fsyncs the engine cache directory for durability.
49///
50/// Snapshots `.engine` file count before and after the run. When a shape
51/// reports a fresh compile (not a cache hit) but the on-disk count does not
52/// increase, emits a `WARN` so operators can catch the
53/// "compile-success-without-persistence" failure mode observed in production
54/// (TRT EP silently failing to write engine plan
55/// files even though `session.run()` returned `Ok(_)`).
56///
57/// `sm` selects which engine plans count toward the before/after snapshots:
58/// `Some("smXY")` filters to plans matching this worker's GPU compute
59/// capability so a heterogeneous cache (e.g. stale `sm89` plans next to
60/// fresh `sm120` plans) is never miscounted; `None` is a passthrough that
61/// counts every `.engine` file (legacy behaviour, used when detection failed).
62/// See [`super::super::trt_cache::engine_files_for_sm`] for the filter
63/// semantics.
64///
65/// The `shape_index` / `shape_total` parameters are purely for the
66/// operator-visible log message and do not affect logic.
67///
68/// `#[allow(clippy::too_many_arguments)]` is acceptable here because every
69/// argument is logically distinct — `(batch, seq)` already has its own
70/// pair-of-`usize` shape, and bundling the remaining diagnostic positional
71/// fields (`worker_id`, `shape_index`, `shape_total`, `sm`) into an
72/// auxiliary struct would obscure the per-shape ergonomics for the only
73/// caller, `trt_prewarm`.
74#[allow(clippy::too_many_arguments)]
75pub(super) fn run_warmup_shape(
76 session: &mut ort::session::Session,
77 batch: usize,
78 seq: usize,
79 worker_id: usize,
80 shape_index: usize,
81 shape_total: usize,
82 engine_cache_dir: &Path,
83 sm: Option<&str>,
84) -> ShapeRunResult {
85 let ids = ndarray::Array2::<i64>::zeros((batch, seq));
86 let mask = ndarray::Array2::<i64>::ones((batch, seq));
87
88 let engine_count_before = trt_cache::count_engine_files_for_sm(engine_cache_dir, sm);
89
90 tracing::info!(
91 worker_id,
92 batch,
93 seq,
94 shape_index,
95 shape_total,
96 engine_count_before,
97 detected_sm = sm.unwrap_or("unfiltered"),
98 "TensorRT pre-warm: running shape (cold compile may take 30–170 s)"
99 );
100
101 let compile_start = std::time::Instant::now();
102 let result = probe_run_dense(session, &ids, &mask);
103 let compile_ms = u64::try_from(compile_start.elapsed().as_millis()).unwrap_or(u64::MAX);
104 let cache_hit = compile_ms < CACHE_HIT_THRESHOLD_MS;
105
106 match result {
107 Ok(_) => {
108 // Flush newly-written engine plan to disk before moving on.
109 // On a cache hit the engine file was only read (not written), so
110 // this fsync is a no-op for data durability — but it is cheap and
111 // keeps the call site uniform regardless of hot/cold path.
112 let fsync_start = std::time::Instant::now();
113 trt_cache::fsync_cache_dir(engine_cache_dir);
114 let fsync_ms = u64::try_from(fsync_start.elapsed().as_millis()).unwrap_or(u64::MAX);
115
116 let engine_count_after = trt_cache::count_engine_files_for_sm(engine_cache_dir, sm);
117 let engine_count_increased = engine_count_after > engine_count_before;
118
119 // A non-cache-hit run that reports `Ok(_)` from session.run() but
120 // leaves the on-disk `.engine` count at zero is the silent failure
121 // mode behind silent-persistence startup failures in production.
122 //
123 // NOTE: The condition is `engine_count_after == 0`, NOT
124 // `!engine_count_increased`. ORT's TRT EP writes one profile-based
125 // engine file that covers all (batch, seq) shapes via [min, max]
126 // ranges — it rewrites that file in-place as the profile expands,
127 // so `engine_count_before == engine_count_after` (delta == 0) is
128 // the normal steady-state after the first compile. A WARN on every
129 // delta==0 shape is a false positive; WARN only when no file
130 // exists at all.
131 if !cache_hit && engine_count_after == 0 {
132 tracing::warn!(
133 worker_id,
134 batch,
135 seq,
136 shape_index,
137 shape_total,
138 compile_ms,
139 fsync_ms,
140 engine_count_before,
141 engine_count_after,
142 detected_sm = sm.unwrap_or("unfiltered"),
143 cache_path = %engine_cache_dir.display(),
144 "TensorRT pre-warm: compile-success log fired but engine_count is still \
145 zero — TRT EP may not be persisting engine plan files"
146 );
147 }
148
149 tracing::info!(
150 worker_id,
151 batch,
152 seq,
153 shape_index,
154 shape_total,
155 compile_ms,
156 fsync_ms,
157 cache_hit,
158 engine_count_before,
159 engine_count_after,
160 engine_count_increased,
161 detected_sm = sm.unwrap_or("unfiltered"),
162 "TensorRT pre-warm: engine compiled, cached, and fsynced"
163 );
164 ShapeRunResult {
165 compile_ms,
166 fsync_ms,
167 cache_hit,
168 succeeded: true,
169 }
170 }
171 Err(e) => {
172 let engine_count_after = trt_cache::count_engine_files_for_sm(engine_cache_dir, sm);
173 tracing::warn!(
174 worker_id,
175 batch,
176 seq,
177 shape_index,
178 shape_total,
179 compile_ms,
180 cache_hit,
181 engine_count_before,
182 engine_count_after,
183 detected_sm = sm.unwrap_or("unfiltered"),
184 error = %e,
185 "TensorRT pre-warm: engine compilation failed for shape; \
186 first real request for this shape will trigger an on-demand compile"
187 );
188 ShapeRunResult {
189 compile_ms,
190 fsync_ms: 0,
191 cache_hit,
192 succeeded: false,
193 }
194 }
195 }
196}