Skip to main content

trt_prewarm

Function trt_prewarm 

Source
pub(super) fn trt_prewarm(
    session: &mut Session,
    warmup_shapes: &[(usize, usize)],
    worker_id: usize,
    cache_dir: &Path,
    sm: Option<&str>,
) -> PrewarmStats
Expand 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).