pub(super) fn trt_prewarm(
session: &mut Session,
warmup_shapes: &[(usize, usize)],
worker_id: usize,
cache_dir: &Path,
sm: Option<&str>,
) -> PrewarmStatsExpand description
Runs a dummy session.run() for each (batch, seq) shape in
warmup_shapes so the TensorRT EP compiles and caches engine files
before the first real request arrives.
§SM-aware cache accounting
sm selects which engine plans count toward engine_count_before,
engine_count_after, the coverage-check fast-path trigger, and the
per-shape persistence WARN. Pass Some("smXY") (e.g. Some("sm120")
for Blackwell) so a heterogeneous cache containing plans for other GPU
compute capabilities — typical when a fleet is mid-deploy or an EFS
volume was previously used by a different instance family — never
produces a false cache_hit:true signal. Pass None for the legacy
unfiltered behaviour (only when SM detection failed; see the WARN
emitted in run_worker).
§Warm-cache fast path
When .engine files matching sm already exist in the cache
directory, the function first runs only the dimensional-extreme shapes
(≤ 4) to probe whether the cached profile covers the full shard. If all
extreme shapes complete in under CACHE_HIT_THRESHOLD_MS the
remaining shapes are skipped — they are guaranteed to be cache hits
by the range-based ORT TRT EP profile logic (see module-level
documentation for the proof). If any extreme shape is slow the fast path
is suppressed and all remaining shapes are compiled normally.
§Cold cache
When no .engine files matching sm exist the coverage-check phase is
bypassed and every shape is compiled in sequence. Each may take
30–170 s on the very first deploy; subsequent starts reuse the cached
.engine files for this SM.
Progress is logged at INFO with compile_ms, fsync_ms, and
cache_hit (whether the run was under CACHE_HIT_THRESHOLD_MS) for
each shape. After each successful run the engine cache directory is
fsynced so the plan file survives an unexpected OOM-kill — see
trt_cache::fsync_cache_dir.
Returns aggregate statistics including fully_cached (whether the
shard was served entirely from cache for this SM) and skipped
(shapes not run).