bge_m3_embedding_server/embedder/worker/
propagation.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//! Engine propagation broadcast drain and post-inference cache-miss signaling.
16
17use super::trt_retry::CHUNK_CACHE_HIT_THRESHOLD_MS;
18use crate::embedder::types::{EmbedStats, JitSuspectSender};
19
20/// Emits the `chunk_run` INFO event and, on a cache miss, notifies both the
21/// JIT-suspect channel (adaptive warmup scheduling) and the engine propagation
22/// broadcast channel (peer worker fast disk-load).
23///
24/// Returns `Some((batch_len, max_chunk_seq))` when a shape was broadcast on
25/// the engine propagation channel.  The call site MUST insert this shape into
26/// `warmed_local` so the originating worker self-skips its own broadcast on
27/// the next `drain_engine_propagation` iteration (COR-1).
28///
29/// # 5000 ms threshold heuristic (COR-10)
30///
31/// `CHUNK_CACHE_HIT_THRESHOLD_MS` (5 s) is a **heuristic** proxy for "TRT
32/// engine JIT compile occurred", not a semantic guarantee.  False negatives
33/// are possible for fast-JIT small shapes; false positives are impossible
34/// because a cache-hit path never exceeds ~100 ms.  The trade-off is
35/// acceptable: the worst outcome of a false negative is that the adaptive
36/// warmup task eventually resubmits the shape on the next real cache miss.
37pub(super) fn log_inference_complete(
38    stats: &EmbedStats,
39    worker_id: usize,
40    _route: &'static str,
41    jit_suspect_tx: Option<&JitSuspectSender>,
42    engine_propagation_tx: Option<&tokio::sync::broadcast::Sender<(usize, usize)>>,
43    batch_len: usize,
44) -> Option<(usize, usize)> {
45    let cache_hit = stats.inference_ms < CHUNK_CACHE_HIT_THRESHOLD_MS;
46    tracing::info!(
47        target: "bge_m3_embedding_server::trt_shape",
48        worker_id,
49        chunk_batch = batch_len,
50        chunk_max_seq = stats.max_chunk_seq,
51        inference_ms = stats.inference_ms,
52        cache_hit,
53        "chunk_run"
54    );
55    if !cache_hit {
56        if let Some(tx) = jit_suspect_tx {
57            let _ = tx.try_send((batch_len, stats.max_chunk_seq));
58        }
59        if let Some(tx) = engine_propagation_tx {
60            let _ = tx.send((batch_len, stats.max_chunk_seq));
61            return Some((batch_len, stats.max_chunk_seq));
62        }
63    }
64    None
65}
66
67/// Drains pending broadcast notifications and runs `trt_prewarm` for each
68/// new shape.
69///
70/// Called at the start of each worker loop iteration (between requests) so
71/// peers eagerly warm their in-memory TRT profile before the next real
72/// request for a new shape arrives.
73///
74/// `warmed_local` tracks shapes already warmed by this worker in the current
75/// session.  The originating worker self-skips on subsequent drains because
76/// `log_inference_complete` inserts the broadcast shape into `warmed_local`
77/// at the call site before returning control to the request loop.
78pub(super) fn drain_engine_propagation<F>(
79    rx: &mut tokio::sync::broadcast::Receiver<(usize, usize)>,
80    warmed_local: &mut std::collections::HashSet<(usize, usize)>,
81    worker_id: usize,
82    mut prewarm: F,
83) where
84    F: FnMut((usize, usize)),
85{
86    loop {
87        match rx.try_recv() {
88            Ok(shape) => {
89                if warmed_local.insert(shape) {
90                    prewarm(shape);
91                }
92            }
93            Err(
94                tokio::sync::broadcast::error::TryRecvError::Empty
95                | tokio::sync::broadcast::error::TryRecvError::Closed,
96            ) => break,
97            Err(tokio::sync::broadcast::error::TryRecvError::Lagged(n)) => {
98                tracing::warn!(
99                    worker_id,
100                    lagged = n,
101                    "engine_propagation: broadcast lagged; some shapes missed"
102                );
103                // Continue draining; missed shapes will be re-broadcast on
104                // the next slow-inference event for that shape.
105            }
106        }
107    }
108}
bge_m3_embedding_server/embedder/worker/propagation.rs

bge_m3_embedding_server/embedder/worker/
propagation.rs
