bge_m3_embedding_server/embedder/worker/
dispatch.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//! Per-request dispatch for dense, sparse, dual, probe, and adaptive warmup.
16
17use std::collections::HashSet;
18use std::path::Path;
19use std::sync::atomic::Ordering;
20
21use super::config::WorkerConfig;
22use super::guard::{
23    EmbedRouteContext, InferenceOutcome, adaptive_warmup_non_trt_compile_ms, finalize_embed_route,
24    log_client_abandoned_before_dispatch,
25};
26use super::probe::run_probe_batch;
27use super::trt_retry::embed_with_trt_retry;
28use crate::config::EpSelection;
29use crate::embedder::dense::embed_dense;
30use crate::embedder::dual::embed_both;
31use crate::embedder::jit_guard::TrtJitGuard;
32use crate::embedder::sparse::embed_sparse;
33use crate::embedder::trt_warmup::trt_prewarm;
34use crate::embedder::types::EmbedRequest;
35
36/// Result of dispatching one worker request through inference.
37pub(super) struct DispatchOutcome {
38    pub outcome: InferenceOutcome,
39    pub skip: bool,
40}
41
42/// Sends a model-reload error to whichever reply channel the request carries.
43pub(super) fn reply_request_load_error(request: EmbedRequest, err: anyhow::Error) {
44    match request {
45        EmbedRequest::Dense { reply, .. } => {
46            let _ = reply.send(Err(err));
47        }
48        EmbedRequest::Sparse { reply, .. } => {
49            let _ = reply.send(Err(err));
50        }
51        EmbedRequest::Both { reply, .. } => {
52            let _ = reply.send(Err(err));
53        }
54        EmbedRequest::Probe { reply, .. } => {
55            let _ = reply.send(Err(err));
56        }
57        EmbedRequest::AdaptiveWarmup { ack, .. } => {
58            let _ = ack.send(Err(err));
59        }
60    }
61}
62
63/// Runs one `EmbedRequest` against a loaded session.
64#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
65pub(super) fn dispatch_request(
66    request: EmbedRequest,
67    session: &mut ort::session::Session,
68    tokenizer: &mut tokenizers::Tokenizer,
69    config: &WorkerConfig,
70    id: usize,
71    cache_dir: &Path,
72    detected_sm: Option<&str>,
73    warmed_local: &mut HashSet<(usize, usize)>,
74    consecutive_failures: u64,
75    shape_guard: Option<&TrtJitGuard>,
76) -> DispatchOutcome {
77    let mut inner_skip = false;
78    let mut inner_outcome = InferenceOutcome::Ok;
79
80    let route_ctx = EmbedRouteContext {
81        worker_id: id,
82        route: "",
83        consecutive_failures,
84        circuit_breaker_threshold: config.circuit_breaker_threshold,
85        jit_suspect_tx: config.jit_suspect_tx.as_ref(),
86        engine_propagation_tx: config.engine_propagation_tx.as_ref(),
87        batch_len: 0,
88    };
89
90    match request {
91        EmbedRequest::Dense { texts, reply } => {
92            // Pre-dispatch abandonment check: the router's hedged
93            // race or the original HTTP client may have already
94            // disconnected while this request sat in the worker
95            // queue. ORT `session.run()` is a blocking C call that
96            // cannot be interrupted mid-MatMul (see CLAUDE.md
97            // "client disconnect" gotcha), so the only opportunity
98            // we have to save work is BEFORE inference starts. The
99            // post-inference check below is observability only.
100            if reply.is_closed() {
101                log_client_abandoned_before_dispatch(id, "dense", texts.len());
102                inner_skip = true;
103            } else {
104                let t_inference = std::time::Instant::now();
105                let cm_guard = config.cost_model.load();
106                let result = embed_with_trt_retry(
107                    |cm| {
108                        embed_dense(
109                            session,
110                            tokenizer,
111                            &texts,
112                            cm,
113                            config.model_variant,
114                            shape_guard,
115                        )
116                    },
117                    &cm_guard,
118                    id,
119                    "dense",
120                )
121                .map_err(|e| e.context("Dense embed error"));
122                let inference_ms = t_inference.elapsed().as_millis();
123                let ctx = EmbedRouteContext {
124                    route: "dense",
125                    batch_len: texts.len(),
126                    ..route_ctx
127                };
128                inner_outcome =
129                    finalize_embed_route(&ctx, result, reply, inference_ms, warmed_local);
130            }
131        }
132        EmbedRequest::Sparse { texts, reply } => {
133            if reply.is_closed() {
134                log_client_abandoned_before_dispatch(id, "sparse", texts.len());
135                inner_skip = true;
136            } else {
137                let t_inference = std::time::Instant::now();
138                let cm_guard = config.cost_model.load();
139                let result = embed_with_trt_retry(
140                    |cm| {
141                        embed_sparse(
142                            session,
143                            tokenizer,
144                            &texts,
145                            cm,
146                            config.model_variant,
147                            shape_guard,
148                        )
149                    },
150                    &cm_guard,
151                    id,
152                    "sparse",
153                )
154                .map_err(|e| e.context("Sparse embed error"));
155                let inference_ms = t_inference.elapsed().as_millis();
156                let ctx = EmbedRouteContext {
157                    route: "sparse",
158                    batch_len: texts.len(),
159                    ..route_ctx
160                };
161                inner_outcome =
162                    finalize_embed_route(&ctx, result, reply, inference_ms, warmed_local);
163            }
164        }
165        EmbedRequest::Both { texts, reply } => {
166            if reply.is_closed() {
167                log_client_abandoned_before_dispatch(id, "both", texts.len());
168                inner_skip = true;
169            } else {
170                let t_inference = std::time::Instant::now();
171                let cm_guard = config.cost_model.load();
172                let result = embed_with_trt_retry(
173                    |cm| {
174                        embed_both(
175                            session,
176                            tokenizer,
177                            &texts,
178                            cm,
179                            config.model_variant,
180                            shape_guard,
181                        )
182                    },
183                    &cm_guard,
184                    id,
185                    "both",
186                )
187                .map_err(|e| e.context("Dual embed error"));
188                let inference_ms = t_inference.elapsed().as_millis();
189                let ctx = EmbedRouteContext {
190                    route: "both",
191                    batch_len: texts.len(),
192                    ..route_ctx
193                };
194                inner_outcome =
195                    finalize_embed_route(&ctx, result, reply, inference_ms, warmed_local);
196            }
197        }
198        EmbedRequest::Probe { texts, reply } => {
199            // Probe: tokenize once without padding, run dense inference
200            // on a single flat batch at the chunk's natural max_seq.
201            // Probes are internal — no client-disconnect path applies.
202            let result = run_probe_batch(session, tokenizer, &texts);
203            let _ = reply.send(result);
204        }
205        EmbedRequest::AdaptiveWarmup { batch, seq, ack } => {
206            // Run trt_prewarm for a single shape so the TRT EP
207            // compiles and caches the engine during an idle window.
208            // On CPU/CUDA EP this is a cheap no-op (returns Ok(0)).
209            let result: anyhow::Result<u64> = if config.ep == EpSelection::TensorRt {
210                let shape = vec![(batch, seq)];
211                let stats = trt_prewarm(session, &shape, id, cache_dir, detected_sm);
212                if stats.warmed > 0 || stats.fully_cached {
213                    // Newly-compiled tier now has a persisted
214                    // engine plan: extend the guard ceiling so
215                    // real requests at this seq are admitted.
216                    config
217                        .warmed_seq_ceiling
218                        .fetch_max(stats.max_warmed_seq.max(seq), Ordering::AcqRel);
219                    Ok(stats.total_compile_ms)
220                } else {
221                    Err(anyhow::anyhow!(
222                        "adaptive warmup: no shapes warmed for \
223                         ({batch}, {seq}) on worker {id}"
224                    ))
225                }
226            } else {
227                Ok(adaptive_warmup_non_trt_compile_ms())
228            };
229            let _ = ack.send(result);
230        }
231    } // end match request
232
233    DispatchOutcome {
234        outcome: inner_outcome,
235        skip: inner_skip,
236    }
237}
bge_m3_embedding_server/embedder/worker/dispatch.rs

bge_m3_embedding_server/embedder/worker/
dispatch.rs
