Skip to main content

Module trt_warmup

Module trt_warmup 

Source
Expand description

TensorRT engine pre-warming: compiles and caches engine files during worker startup so the first real request hits a cached engine instead of triggering an on-demand 30–170 s compile.

§Durability

After each shape compiles, the engine cache directory is fsynced so an unexpected SIGKILL (ECS OOM-kill, host crash) cannot strand a partially-written engine plan in the page cache. See trt_cache.rs.

§Cache-hit fast path (warm cache skip)

ORT’s TRT EP caches engines with per-dimension [min, max] ranges — not one engine per shape. A session.run() is a cache hit (fast, no compile) when every input dimension falls within the cached [min, max] range; it is a cache miss (slow compile) only when a dimension falls outside that range and the engine must be rebuilt with an extended range.

After a full first-deploy warmup sweep, the cached profile records: input_ids.dim_0 ∈ [min_batch, max_batch] and input_ids.dim_1 ∈ [min_seq, max_seq] — covering every shape in the warmup grid. On subsequent container starts, every warmup session.run() is a cache hit and finishes in ≤ 3 s.

Rather than paying 24 × 1–3 s = 24–72 s of redundant cache-hit loads, trt_prewarm runs at most 4 “dimensional-extreme” shapes (the shapes that exercise the minimum and maximum of each input dimension independently) and measures wall-clock time. If all extremes complete under CACHE_HIT_THRESHOLD_MS, the profile is guaranteed to cover the entire shard and the remaining shapes are skipped.

§Why this has zero false positives

For shape (b, s) to be a TRT cache hit it must satisfy:

profile.min_batch ≤ b ≤ profile.max_batch   (batch dimension)
profile.min_seq   ≤ s ≤ profile.max_seq      (sequence dimension)

The four extreme shapes bound all four inequalities independently:

Check shapeFact established when it is a cache hit
(min_batch, any_s)profile.min_batch ≤ min_batch
(max_batch, any_s)profile.max_batch ≥ max_batch
(any_b, min_seq)profile.min_seq ≤ min_seq
(any_b, max_seq)profile.max_seq ≥ max_seq

Together these four facts guarantee that every shard shape (b, s) with b ∈ [min_batch, max_batch] and s ∈ [min_seq, max_seq] is a cache hit. If any extreme shape is slow (≥ CACHE_HIT_THRESHOLD_MS) the engine must be rebuilt for that dimension → the fast path is suppressed and all remaining shapes are compiled normally.

Modules§

postcondition 🔒
TensorRT engine pre-warm persistence postconditions.
runner 🔒
Per-shape TensorRT warmup runner.

Structs§

PrewarmStats 🔒
Aggregate per-worker statistics returned by trt_prewarm.

Constants§

CACHE_HIT_THRESHOLD_MS 🔒
Threshold (ms) below which a session.run() is classified as a TRT engine cache hit (loaded from disk) rather than a fresh compile.

Functions§

coverage_check_shapes 🔒
Selects the minimal set of shapes needed to verify that an ORT TRT EP cached profile covers all shapes in shapes.
shard_shapes 🔒
Partitions shapes into a per-worker shard using a stride assignment.
trt_prewarm 🔒
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.