bge_m3_embedding_server/embedder/worker/
config.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-worker execution policy shared across the pool.
16
17use std::sync::Arc;
18use std::sync::atomic::AtomicUsize;
19use std::time::Duration;
20
21use arc_swap::ArcSwap;
22
23use crate::binpack::CostModel;
24use crate::config::{EpSelection, ModelVariant};
25use crate::embedder::types::JitSuspectSender;
26
27/// Execution-policy configuration shared by all workers.
28///
29/// `cost_model` is an `Arc<ArcSwap<CostModel>>` so all workers share a single
30/// handle and the background probe can update the cost model atomically after
31/// fitting.  Each worker loads the current value lock-free at the start of
32/// every `session.run()` call via `config.cost_model.load()`.
33#[derive(Clone)]
34pub struct WorkerConfig {
35    /// Quadratic-aware workspace cost model and per-worker budget.
36    ///
37    /// Shared across all workers via `ArcSwap`.  The background probe updates
38    /// this handle once fitted coefficients are available; workers observe the
39    /// new model on their next request without any coordination or restart.
40    pub cost_model: Arc<ArcSwap<CostModel>>,
41    /// Duration of inactivity before workers unload their model instances.
42    pub idle_timeout: Option<Duration>,
43    /// ONNX model variant to load (FP32, FP16, or INT8).
44    pub model_variant: ModelVariant,
45    /// Maximum tokenized sequence length.
46    pub max_seq_length: usize,
47    /// Number of intra-op threads each ORT session may use for a single
48    /// `session.run()` call. Plumbed through to `load_session` at model load
49    /// time. See [`crate::config::Config::intra_threads`] for sizing guidance.
50    pub intra_threads: usize,
51    /// ONNX Runtime execution provider selection.
52    ///
53    /// Forwarded to [`crate::embedder::session::load_models`] at model load time so each ORT session
54    /// registers the correct EP. On macOS, `CoreML` is always used regardless
55    /// of this value. See [`crate::config::EpSelection`] for details.
56    pub ep: EpSelection,
57
58    /// `(batch_size, seq_len)` shapes to pre-compile as `TensorRT` engine files
59    /// during worker startup.
60    ///
61    /// Only applied when `ep == EpSelection::TensorRt`. Sourced from
62    /// `BGE_M3_TRT_WARMUP_SHAPES` via [`crate::config::Config::trt_warmup_shapes`].
63    /// With multiple workers, `EmbedPool::spawn` shards the full shape list
64    /// across workers (stride partition) so each GPU compiles a disjoint subset
65    /// in parallel. An empty list skips pre-warming entirely.
66    pub trt_warmup_shapes: Vec<(usize, usize)>,
67
68    /// CUDA/TRT device ID for this specific worker.
69    ///
70    /// Set by `EmbedPool::spawn` as `worker_index % gpu_count`. Forwarded to
71    /// [`crate::embedder::session::execution_providers`] so the ORT session binds to the
72    /// correct GPU. Ignored on CPU EP and macOS (`CoreML` is single-device).
73    pub device_id: u32,
74
75    /// Total number of GPU devices on this instance.
76    ///
77    /// Propagated from [`crate::config::Config::gpu_count`]. Used by
78    /// `EmbedPool::spawn` to compute per-worker `device_id` values and to
79    /// clamp `BGE_M3_WORKERS` for GPU execution providers.
80    pub gpu_count: usize,
81
82    /// Optional TRT workspace size cap (bytes) forwarded to ORT's TRT EP via
83    /// `with_max_workspace_size`.  `None` uses ORT's built-in default.
84    /// Sourced from `BGE_M3_TRT_MAX_WORKSPACE_BYTES`.
85    pub trt_max_workspace_bytes: Option<usize>,
86
87    /// Optional CUDA device memory limit (bytes) forwarded to the CUDA EP.
88    /// `None` uses ORT's built-in default.
89    pub gpu_mem_limit_bytes: Option<usize>,
90
91    /// Sender half of the JIT-suspect channel created before pool spawn.
92    ///
93    /// After each successful inference, if `inference_ms >= CHUNK_CACHE_HIT_THRESHOLD_MS`
94    /// the worker calls `try_send((batch, seq))` so the adaptive warmup task
95    /// can schedule background engine compilation.  Non-blocking: if the
96    /// channel is full the message is silently dropped.  `None` when adaptive
97    /// warmup is disabled.
98    pub jit_suspect_tx: Option<JitSuspectSender>,
99
100    /// Sender half of the engine propagation broadcast channel.
101    ///
102    /// When `Some`, after any worker writes a new TRT engine plan to EFS, the
103    /// worker broadcasts the `(batch, seq)` shape to all subscribed peers so
104    /// they eagerly run `trt_prewarm` (~1-3s fast disk-load) instead of paying
105    /// full JIT cost on the next real request.  Each worker derives its own
106    /// `Receiver` via `tx.subscribe()` at startup.  `None` when
107    /// `BGE_M3_ENGINE_PROPAGATION_ENABLED=0` or when the EP is not TRT.
108    pub engine_propagation_tx: Option<tokio::sync::broadcast::Sender<(usize, usize)>>,
109
110    /// When `true`, prewarm postcondition failures cause the worker to refuse
111    /// to signal ready: `run_worker` returns `Err(_)` before `ready_tx.send`,
112    /// the pool's init task propagates the error, and the readiness probe in
113    /// `bootstrap::readiness` triggers a hard process exit. Converts the
114    /// false-positive-readiness failure mode (every worker hits TRT
115    /// `Error Code 10` mid-build, postcondition logs WARN, `/health` still
116    /// returns `200 ok`, real requests then 500) into an explicit startup
117    /// failure that ECS retries instead of routing traffic to.
118    ///
119    /// Sourced from `BGE_M3_PREWARM_STRICT`; defaults to `true`. Set to
120    /// `false` to preserve pre-fix behaviour (WARN only).
121    pub prewarm_strict: bool,
122
123    /// Consecutive-failure threshold for the per-worker inference circuit
124    /// breaker.
125    ///
126    /// When a worker accumulates this many consecutive errors from an
127    /// inference call (`embed_dense`, `embed_sparse`, or `embed_both`), it
128    /// drops its ORT session (cleaning the CUDA arena), decrements
129    /// `loaded_workers`, and waits for the next request to trigger a model
130    /// reload. The counter resets to zero on any successful inference.
131    ///
132    /// Sourced from `BGE_M3_CIRCUIT_BREAKER_THRESHOLD`; defaults to `5`.
133    pub circuit_breaker_threshold: usize,
134
135    /// Enables the in-band TRT JIT admission guard (see [`crate::embedder::jit_guard`]).
136    ///
137    /// When `true` (and `ep == TensorRt`), before any chunk's `session.run()`
138    /// the worker refuses chunks whose sequence length is in the dangerous
139    /// range and is not covered by the pool's warmed engine profile, returning
140    /// an error that maps to HTTP `503` instead of risking the process-killing
141    /// pathological autotuner allocation. Sourced from
142    /// `BGE_M3_TRT_INBAND_JIT_GUARD`; defaults to `true`.
143    pub trt_inband_jit_guard_enabled: bool,
144
145    /// Sequence-length threshold (`guard_seq`) at/above which an *uncovered*
146    /// shape is refused rather than JIT-compiled in-band. Below this value a
147    /// cold in-band JIT is bounded and is allowed (letting the engine profile
148    /// grow naturally). Sourced from `BGE_M3_TRT_INBAND_JIT_GUARD_SEQ`;
149    /// defaults to `4096`. The guard is a no-op when `max_seq_length` is below
150    /// this threshold (no chunk can reach the dangerous range).
151    pub trt_inband_jit_guard_seq: usize,
152
153    /// Pool-wide ceiling: the maximum sequence length any worker has
154    /// successfully warmed (fresh compile or warm-cache hit), shared across
155    /// all workers via the `Arc<AtomicUsize>`.
156    ///
157    /// Read by the per-request [`crate::embedder::jit_guard::TrtJitGuard`]; raised via
158    /// [`fetch_max`](std::sync::atomic::AtomicUsize::fetch_max) after startup
159    /// prewarm, engine-propagation prewarm, and adaptive-warmup compiles. A
160    /// successful warmup of a sequence tier by *any* worker means the engine
161    /// plan is on the shared EFS cache and every worker can fast-load it, so a
162    /// single shared ceiling is a sound pool-wide coverage signal.
163    pub warmed_seq_ceiling: Arc<AtomicUsize>,
164
165    /// **Destructive** stale-SM TRT engine cache GC flag.
166    ///
167    /// Only present when the `cache-gc` Cargo feature is compiled in.
168    /// Defaults to `false` even with the feature on; flipping the
169    /// runtime knob also requires `BGE_M3_TRT_CACHE_GC_ENABLED=1`. See
170    /// [`crate::config::Config::trt_cache_gc_enabled`] for the multi-SM
171    /// ASG hazard model.
172    #[cfg(feature = "cache-gc")]
173    pub trt_cache_gc_enabled: bool,
174}
bge_m3_embedding_server/embedder/worker/config.rs

bge_m3_embedding_server/embedder/worker/
config.rs
