Skip to main content

run_warmup_shape

Function run_warmup_shape 

Source
pub(super) fn run_warmup_shape(
    session: &mut Session,
    batch: usize,
    seq: usize,
    worker_id: usize,
    shape_index: usize,
    shape_total: usize,
    engine_cache_dir: &Path,
    sm: Option<&str>,
) -> ShapeRunResult
Expand description

Runs session.run() for a single (batch, seq) shape, measures wall time, classifies the result as a cache hit or fresh compile, and — on success — fsyncs the engine cache directory for durability.

Snapshots .engine file count before and after the run. When a shape reports a fresh compile (not a cache hit) but the on-disk count does not increase, emits a WARN so operators can catch the “compile-success-without-persistence” failure mode observed in production (TRT EP silently failing to write engine plan files even though session.run() returned Ok(_)).

sm selects which engine plans count toward the before/after snapshots: Some("smXY") filters to plans matching this worker’s GPU compute capability so a heterogeneous cache (e.g. stale sm89 plans next to fresh sm120 plans) is never miscounted; None is a passthrough that counts every .engine file (legacy behaviour, used when detection failed). See super::super::trt_cache::engine_files_for_sm for the filter semantics.

The shape_index / shape_total parameters are purely for the operator-visible log message and do not affect logic.

#[allow(clippy::too_many_arguments)] is acceptable here because every argument is logically distinct — (batch, seq) already has its own pair-of-usize shape, and bundling the remaining diagnostic positional fields (worker_id, shape_index, shape_total, sm) into an auxiliary struct would obscure the per-shape ergonomics for the only caller, trt_prewarm.