Module jit_guard 

In-band TensorRT JIT admission guard.

The failure this prevents

When a chunk shape (batch, seq) reaches session.run() that the worker’s TensorRT engine profile does not already cover, the TRT EP compiles an engine for it in-band (in the middle of a real request). On the fused dual-output /v1/embeddings:both graph at the maximum sequence length (seq = 8192) the kernel autotuner can request pathological scratch allocations - tens of gigabytes up to multiple terabytes on a single LayerNorm + MatMul foreign node. BGE_M3_TRT_MAX_WORKSPACE_BYTES does not bound autotuner tactic scratch (a TRT EP limitation), so on a VRAM-saturated device (e.g. the warmup-shard worker already holding the seq=8192 engines at 90%+ VRAM) the CUDA allocator faults and the process dies via SIGSEGV / OOM-kill before any Result is returned. None of the existing reactive safety nets (is_trt_jit_oom retry, the is_trt_engine_build_fatal worker-exit, the circuit breaker) can catch a hard process death - the only defense is to never issue the dangerous run.

Startup warmup catches a failed compile (run_warmup_shape logs a WARN and continues) so a worker whose seq=8192 shard failed to compile still signals ready. The first real seq≈8192 request then triggers the same pathological allocation in-band, without warmup’s caught-error safety net.

The guard

TrtJitGuard refuses - with a clean, retriable error that maps to HTTP 503 - any chunk whose sequence length is in the dangerous range (seq >= guard_seq) and is not already covered by the pool’s warmed engine profile (seq > warmed_seq_ceiling). Refusing one request is strictly better than a SIGSEGV that kills every in-flight request on the worker and forces an ECS task replacement.

warmed_seq_ceiling is the maximum sequence length any worker in the pool successfully warmed (fresh compile or warm-cache hit), shared via an AtomicUsize. Because TRT engine plans live on the shared EFS cache and a single profile-based engine file spans [min_seq, max_seq] across every shape compiled to it, a successful warmup of seq=8192 by any worker means every worker can fast-load (not JIT) that shape - so the ceiling is a sound pool-wide coverage signal. Conversely, if the seq=8192 shard failed on every worker, no plan exists on disk, the ceiling stays at the highest tier that did compile (e.g. 2048), and seq=8192 requests are refused instead of crashing the process.

Why sequence length (not batch) is the discriminator

The pathological allocation scales with the attention score matrix (O(batch · seq^2)), which is dominated by seq at the top tier. Within a compiled profile, intermediate batches are covered by the engine’s [min_batch, max_batch] range, and bin_pack already bounds the per-chunk batch under the workspace budget (so seq=8192 chunks never exceed ~15-18 texts). The only reachable uncovered-and-dangerous region is “a sequence length tier that warmup failed to compile”, which is exactly what the ceiling tracks.

Self-healing

The adaptive-warmup loop and cross-worker engine propagation both raise the ceiling (via fetch_max) when they successfully compile a higher tier during an idle window, so coverage that was refused at startup is admitted again once a plan lands on disk.

Structs

TrtJitGuard 🔒

Per-request snapshot of the in-band JIT admission policy.

TrtJitRejection 🔒

Error returned when TrtJitGuard refuses a chunk to avoid a pathological in-band TensorRT JIT compile.

Functions

guard_chunks 🔒

Validates every chunk produced by bin_pack against an optional guard.

is_trt_shape_rejected 🔒

Returns true when err (or anything in its source chain) is a TrtJitRejection.