Struct EmbedPool 

pub struct EmbedPool {
    tx: Sender<EmbedRequest>,
    live_workers: Arc<AtomicUsize>,
    loaded_workers: Arc<AtomicUsize>,
    model_rss_per_worker_bytes: Arc<AtomicUsize>,
    engine_propagation_tx: Option<Sender<(usize, usize)>>,
}

Async handle to the embedding worker thread pool.

Wraps a bounded mpsc channel shared by n spawn_blocking worker threads. Each worker owns its own ORT session and tokenizer; the pool dispatches EmbedRequest variants to whichever worker is free next.

Clone is cheap — the underlying channel sender and atomic counters are reference-counted.

Fields

tx: Sender<EmbedRequest>

live_workers: Arc<AtomicUsize>

loaded_workers: Arc<AtomicUsize>

Number of workers that currently have model instances loaded in memory.

model_rss_per_worker_bytes: Arc<AtomicUsize>

Median RSS delta (bytes) measured across all workers during sequential model load.

Workers load one at a time (leader first, then followers in sequence). Each reports its own RSS before/after load_models() via ready_tx. The pool stores the median once all workers have signaled ready — robust to one outlier from page-cache settling or ORT arena init jitter.

Used by run_readiness_probe to correctly deduct the model-weight footprint from the available workspace before computing per-worker budget. Returns 0 on non-Linux targets where RSS measurement is unavailable, or before the init task has completed.

engine_propagation_tx: Option<Sender<(usize, usize)>>

Broadcast sender for cross-worker TRT engine cache propagation.

When Some, after any worker writes a new engine plan to EFS, a (batch, seq) shape notification is broadcast to every subscribed worker so they eagerly run trt_prewarm (~1-3s fast disk-load) instead of paying full JIT cost on the next real request. None when BGE_M3_ENGINE_PROPAGATION_ENABLED=0.

Implementations

impl EmbedPool

pub fn broadcast_engine_ready(&self, shape: (usize, usize))

Broadcasts a (batch, seq) shape notification to all subscribed peer workers, signaling that the TRT engine plan for this shape is now on EFS.

Workers receive the notification in their main loop drain and run trt_prewarm against their own session for a ~1-3s fast disk-load rather than paying full JIT cost on the next real request for this shape.

send returns Err only when there are zero subscribers (all workers gone). This is non-fatal; the engine plan is still on EFS.

pub fn spawn( n: usize, cache_dir: PathBuf, config: WorkerConfig, ) -> (Self, JoinHandle<Result<()>>)

Spawns n embedding worker threads and returns the pool plus an init handle that resolves once all workers have finished loading their models.

When a GPU execution provider (cuda or tensorrt) is selected, n is clamped to config.gpu_count: each worker is pinned to a distinct CUDA device (device_id = worker_index % gpu_count). For TRT, the full warmup shape list is sharded across workers via a stride partition so each GPU compiles a disjoint subset in parallel, then shares the results via the EFS engine cache.

pub async fn dense( &self, texts: Vec<String>, ) -> Result<(Vec<Vec<f32>>, EmbedStats)>

Runs dense (float32) embedding inference on texts.

Errors

• Returns Err if the worker channel has closed (pool shut down).
• Returns Err if the worker drops the reply sender before responding.
• Returns Err if the ORT session fails during inference.

pub async fn sparse( &self, texts: Vec<String>, ) -> Result<(Vec<SparseEmbedding>, EmbedStats)>

Runs sparse (SPLADE-style) embedding inference on texts.

Errors

• Returns Err if the worker channel has closed (pool shut down).
• Returns Err if the worker drops the reply sender before responding.
• Returns Err if the ORT session fails during inference.

pub async fn both( &self, texts: Vec<String>, ) -> Result<(Vec<DualEmbedding>, EmbedStats)>

Runs a single forward pass that yields both dense and sparse embeddings.

Equivalent to calling Self::dense and Self::sparse back-to-back, but uses one session.run() per chunk instead of two — at near-zero marginal GPU cost.

Errors

• Returns Err if the worker channel has closed (pool shut down).
• Returns Err if the worker drops the reply sender before responding.
• Returns Err if the ORT session fails during inference.

pub(crate) async fn probe(&self, texts: Vec<String>) -> Result<ProbeResult>

Sends a probe request to a single worker and returns the result. Only called during init before ready is set.

pub fn live_worker_count(&self) -> usize

Returns the number of worker threads currently alive (not yet exited).

pub fn loaded_worker_count(&self) -> usize

Returns the number of workers that currently have model instances loaded in memory.

A worker transitions from loaded to unloaded after the crate::config::Config::idle_timeout elapses with no incoming requests, and back to loaded on the next request.

pub fn queue_depth(&self) -> usize

Returns the number of requests currently queued but not yet picked up by a worker. Uses the channel’s current vs max capacity.

pub fn model_rss_per_worker_bytes(&self) -> usize

Returns the median RSS delta (bytes) measured across all workers during sequential model load.

This is the per-worker model-weight footprint used by run_readiness_probe to compute the per-worker workspace budget. Returns 0 on non-Linux targets where RSS measurement is unavailable, or before the init task has completed.

pub async fn send_adaptive_warmup( &self, batch: usize, seq: usize, ack: Sender<Result<u64>>, ) -> Result<(), ()>

Sends an adaptive warmup request to an available worker.

The worker calls trt_prewarm for (batch, seq) and replies on ack with the compile duration in milliseconds, or an error on failure. On non-TRT workers the reply is Ok(0) immediately.

Errors

Returns Err(()) if the worker channel has closed (pool shut down).

pub fn live_workers_for_shutdown(&self) -> Arc<AtomicUsize>

Returns a clone of the Arc<AtomicUsize> backing live_worker_count.

Used by the warmup-only path in lib.rs to poll worker exit progress AFTER the EmbedPool itself has been dropped. Dropping the pool closes the request channel, which signals workers to break out of their receive loops and drop their ORT sessions; the live counter is the only readable signal that those drop paths have completed.

Returning a clone of the raw Arc (rather than a snapshot) lets the caller hold a reference across the drop boundary without having to keep the pool’s other state alive.

Trait Implementations

impl Clone for EmbedPool

fn clone(&self) -> EmbedPool

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.