bge_m3_embedding_server/embedder/worker/
guard.rs

1// Copyright (c) 2026 J. Patrick Fulton
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Worker lifecycle guard, in-band JIT guard wiring, and inference outcomes.
16
17use std::collections::HashSet;
18use std::sync::Arc;
19use std::sync::atomic::{AtomicUsize, Ordering};
20
21use super::config::WorkerConfig;
22use super::logging::log_if_abandoned_mid_flight;
23use super::propagation::log_inference_complete;
24use super::trt_retry::is_trt_engine_build_fatal;
25use crate::config::EpSelection;
26use crate::embedder::jit_guard::{self, TrtJitGuard};
27use crate::embedder::types::{EmbedStats, JitSuspectSender};
28
29pub(super) struct WorkerGuard(pub Arc<AtomicUsize>);
30
31impl Drop for WorkerGuard {
32    fn drop(&mut self) {
33        let prev = self.0.fetch_sub(1, Ordering::AcqRel);
34        let live_after_drop = prev.saturating_sub(1);
35        if live_after_drop == 0 {
36            tracing::error!("All embedding workers have exited — pool is degraded");
37        } else {
38            tracing::warn!(live_after_drop, "Embedding worker exited");
39        }
40    }
41}
42
43/// Builds the per-request in-band TRT JIT guard from the worker config and the
44/// live pool-wide warmed-sequence ceiling.
45///
46/// Returns `None` (guard disabled) on non-TRT EPs or when
47/// `BGE_M3_TRT_INBAND_JIT_GUARD=0`, so the embed call sites pass `None` and
48/// skip all guard work. The ceiling is read fresh on every request so the
49/// decision reflects the latest coverage extended by adaptive warmup or
50/// engine propagation.
51pub(super) fn build_shape_guard(config: &WorkerConfig) -> Option<TrtJitGuard> {
52    if config.ep == EpSelection::TensorRt && config.trt_inband_jit_guard_enabled {
53        Some(TrtJitGuard::new(
54            config.trt_inband_jit_guard_seq,
55            config.warmed_seq_ceiling.load(Ordering::Acquire),
56        ))
57    } else {
58        None
59    }
60}
61
62/// Emits a `WARN` describing an in-band TRT JIT guard refusal.
63///
64/// Greppable tag: `trt_jit_guard_refused`. A refusal means the worker
65/// protected itself from a dangerous, uncovered chunk shape that could have
66/// triggered a process-killing pathological autotuner allocation; the client
67/// receives HTTP `503` and may retry once warmup coverage extends.
68pub(super) fn log_guard_rejection<T>(
69    result: &anyhow::Result<T>,
70    worker_id: usize,
71    route: &'static str,
72) {
73    if let Err(e) = result {
74        tracing::warn!(
75            target: "bge_m3_embedding_server::trt_warmup",
76            tag = "trt_jit_guard_refused",
77            worker_id,
78            route,
79            error = %e,
80            "in-band TRT JIT guard refused a dangerous, uncovered chunk shape; \
81             returning 503 instead of risking a process-killing autotuner \
82             allocation (request is retriable once warmup coverage extends)"
83        );
84    }
85}
86
87/// Outcome of a single inference call in the worker request loop.
88///
89/// Used to communicate circuit-breaker and fatal-exit decisions out of the
90/// nested borrow scope (where `session`/`tokenizer` live) into the outer
91/// scope where `models` can be safely mutated.
92#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
93pub(super) enum InferenceOutcome {
94    /// Inference succeeded; reset the consecutive-failure counter.
95    #[default]
96    Ok,
97    /// Inference failed; increment the consecutive-failure counter.
98    Failure,
99    /// Fatal TRT engine build error; worker should exit immediately.
100    TrtFatal,
101    /// Consecutive-failure threshold reached; unload models and reset counter.
102    CircuitBreak,
103    /// The in-band TRT JIT guard refused the request (dangerous, uncovered
104    /// shape). The worker is healthy and deliberately protected itself, so the
105    /// consecutive-failure counter is left unchanged — a refusal is neither a
106    /// success nor a GPU failure and must not contribute to tripping the
107    /// circuit breaker.
108    Rejected,
109}
110
111/// Per-route context passed to [`finalize_embed_route`] after inference.
112pub(super) struct EmbedRouteContext<'a> {
113    pub worker_id: usize,
114    pub route: &'static str,
115    pub consecutive_failures: u64,
116    pub circuit_breaker_threshold: usize,
117    pub jit_suspect_tx: Option<&'a JitSuspectSender>,
118    pub engine_propagation_tx: Option<&'a tokio::sync::broadcast::Sender<(usize, usize)>>,
119    pub batch_len: usize,
120}
121
122/// Emits the pre-dispatch abandonment WARN when the client disconnected
123/// while the request was still queued.
124pub(super) fn log_client_abandoned_before_dispatch(
125    worker_id: usize,
126    route: &'static str,
127    batch_size: usize,
128) {
129    tracing::warn!(
130        worker_id,
131        route,
132        batch_size,
133        "request abandoned by client before dispatch — skipping inference"
134    );
135}
136
137/// Maps inference error flags to the worker-loop [`InferenceOutcome`].
138pub(super) fn classify_inference_outcome(
139    trt_fatal: bool,
140    guard_rejected: bool,
141    is_err: bool,
142    consecutive_failures: u64,
143    circuit_breaker_threshold: usize,
144) -> InferenceOutcome {
145    if trt_fatal {
146        InferenceOutcome::TrtFatal
147    } else if guard_rejected {
148        InferenceOutcome::Rejected
149    } else if is_err {
150        let next_failures = consecutive_failures + 1;
151        if next_failures >= u64::try_from(circuit_breaker_threshold).unwrap_or(u64::MAX) {
152            InferenceOutcome::CircuitBreak
153        } else {
154            InferenceOutcome::Failure
155        }
156    } else {
157        InferenceOutcome::Ok
158    }
159}
160
161/// Returns the consecutive-failure counter value after applying an outcome.
162pub(super) fn next_consecutive_failures(
163    outcome: InferenceOutcome,
164    consecutive_failures: u64,
165) -> u64 {
166    match outcome {
167        InferenceOutcome::Ok | InferenceOutcome::CircuitBreak => 0,
168        InferenceOutcome::Failure => consecutive_failures + 1,
169        InferenceOutcome::TrtFatal | InferenceOutcome::Rejected => consecutive_failures,
170    }
171}
172
173/// Returns `true` when the worker loop should unload models after inference.
174pub(super) fn should_unload_on_outcome(outcome: InferenceOutcome) -> bool {
175    matches!(outcome, InferenceOutcome::CircuitBreak)
176}
177
178/// Shared post-inference path for dense, sparse, and dual routes.
179///
180/// Logs completion stats, classifies errors (TRT fatal / JIT guard / circuit
181/// breaker), emits abandonment observability, and sends the oneshot reply.
182pub(super) fn finalize_embed_route<T>(
183    ctx: &EmbedRouteContext<'_>,
184    result: anyhow::Result<(T, EmbedStats)>,
185    reply: tokio::sync::oneshot::Sender<anyhow::Result<(T, EmbedStats)>>,
186    inference_ms: u128,
187    warmed_local: &mut HashSet<(usize, usize)>,
188) -> InferenceOutcome {
189    let guard_rejected = result
190        .as_ref()
191        .err()
192        .is_some_and(jit_guard::is_trt_shape_rejected);
193    let trt_fatal = result.as_ref().err().is_some_and(is_trt_engine_build_fatal);
194    let is_err = result.is_err();
195
196    if let Ok((_, ref stats)) = result {
197        if let Some(shape) = log_inference_complete(
198            stats,
199            ctx.worker_id,
200            ctx.route,
201            ctx.jit_suspect_tx,
202            ctx.engine_propagation_tx,
203            ctx.batch_len,
204        ) {
205            warmed_local.insert(shape);
206        }
207        tracing::info!(
208            worker_id = ctx.worker_id,
209            chunks = stats.chunks,
210            max_chunk_seq = stats.max_chunk_seq,
211            total_token_positions = stats.total_token_positions,
212            seq_len_min = stats.seq_len_min,
213            seq_len_max = stats.seq_len_max,
214            seq_len_mean = stats.seq_len_mean,
215            seq_len_p95 = stats.seq_len_p95,
216            tokenize_ms = stats.tokenize_ms,
217            inference_ms = stats.inference_ms,
218            route = ctx.route,
219            "worker: embed complete"
220        );
221    }
222
223    let outcome = classify_inference_outcome(
224        trt_fatal,
225        guard_rejected,
226        is_err,
227        ctx.consecutive_failures,
228        ctx.circuit_breaker_threshold,
229    );
230
231    match outcome {
232        InferenceOutcome::TrtFatal => {
233            tracing::error!(
234                worker_id = ctx.worker_id,
235                route = ctx.route,
236                consecutive_failures = ctx.consecutive_failures + 1,
237                "trt_fatal_engine_build: unrecoverable TRT state; \
238                 worker exiting to reset CUDA arena"
239            );
240        }
241        InferenceOutcome::Rejected => {
242            log_guard_rejection(&result, ctx.worker_id, ctx.route);
243        }
244        InferenceOutcome::CircuitBreak => {
245            tracing::error!(
246                worker_id = ctx.worker_id,
247                route = ctx.route,
248                consecutive_failures = ctx.consecutive_failures + 1,
249                threshold = ctx.circuit_breaker_threshold,
250                "circuit_breaker_tripped: unloading models to reset \
251                 CUDA arena; worker will reload on next request"
252            );
253        }
254        InferenceOutcome::Ok | InferenceOutcome::Failure => {}
255    }
256
257    log_if_abandoned_mid_flight(&reply, ctx.route, ctx.worker_id, &result, inference_ms);
258    let _ = reply.send(result);
259    outcome
260}
261
262/// Adaptive-warmup no-op compile duration for non-TRT execution providers.
263pub(super) fn adaptive_warmup_non_trt_compile_ms() -> u64 {
264    0
265}
bge_m3_embedding_server/embedder/worker/guard.rs

bge_m3_embedding_server/embedder/worker/
guard.rs
