Merge perf/encoder-batched-inference: batched GPU encoding, 100 -> 426 img/s

Encoder throughput on one RTX 6000 Ada with bioclip-2, measured end to end through
the Python wheel in-process on 400 real camera-trap photos:

  as published (0.1.23), fp32     160.8 img/s
  this branch, fp32               275.5 img/s
  this branch, fp16               426.3 img/s
  this branch, fp16 + TensorRT    558.8 img/s

Over HTTP the serving path went 100.55 -> 261 img/s with GPU utilisation 49% -> 90%.

The cause was one defect repeated in four layers, each of which hid the next.
embed_batch in both the GPU and CPU crates was a serial per-image loop, so the
ONNX session always ran at batch 1 and no caller's batch size ever reached the
model -- which is why batch 32/64/128/256 all measured identically. The Python
binding's embed() then looped the SINGLE-image function, so it never reached
embed_batch at all and hid both crate fixes from every wheel caller: the rebuilt
wheel measured 160.71 against a published 160.8 until that was found. Finally the
binding's window was set equal to the internal inference chunk, so embed_batch
received exactly one chunk per call and its decode/inference overlap never ran.

Alongside those: the engine-wide Mutex<JpegDecoder> was held across Session::run,
serialising all concurrency on one lock -- hence 1/2/4/8 concurrent clients also
measuring identically despite an inference-semaphore default of 32. Replaced with
a decoder pool, default raised 3 -> 6 once decode became what binds.

Defaults re-swept against fp16 rather than the CUDA-fp32 ceiling the originals
were chosen for: inference chunk 8 -> 4, embed window 8 -> 256. The chunk curve is
not monotonic (1->223, 3->393, 4->428, 8->388, 32->347) because that constant is a
pipelining granularity, not a batch size. Both are env-overridable
(SPARROW_ENGINE_INFERENCE_CHUNK, SPARROW_ENGINE_EMBED_WINDOW,
SPARROW_ENGINE_INTRA_THREADS) since the knee moves with precision and provider.

426 img/s is 84% of this model's 509 img/s fp16 inference ceiling, matching the
84-88% end-to-end-versus-binding-ceiling ratio every configuration measured landed
in, so the fp16 single-GPU path is now at its ceiling.

Vectors are unchanged: cosine 0.9999995 against the pre-change build. fp16 keeps
394/400 identical nearest neighbours at cosine 0.99985, and all disagreements are
tie-breaks with margins far below the median rank1-rank2 gap.

109 GPU-crate tests pass (2 pre-existing cuDNN-environment failures unchanged);
clippy clean.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: fa7107c7-817e-41f8-b9f1-5e6d9b674025

Author: zhmiao

sparrow-engine/sparrow-engine-cpu/src/embed.rs

@@ -43,16 +43,39 @@ pub(crate) fn validate_image_encoder(manifest: &ModelManifest) -> Result<()> {
 /// Run image encoder inference on a single image.
 pub fn embed(handle: &ModelHandle, image: &ImageInput) -> Result<EmbedResult> {
     let start = Instant::now();
-    let manifest = &handle.manifest;
-    validate_image_encoder(manifest)?;
+    validate_image_encoder(&handle.manifest)?;
+    let config = preprocess_config_from_manifest(&handle.manifest)?;
+    let prepared = preprocess::preprocess(image, &config)?;
+    let mut results = infer_prepared(handle, std::slice::from_ref(&prepared), start)?;
+    results
+        .pop()
+        .ok_or_else(|| SparrowEngineError::Ort("embed: inference returned no result".into()))
+}
 
+/// Run one batched `session.run` over already-preprocessed images.
+///
+/// Split out of [`embed`] so a caller can preprocess many images -- optionally in parallel --
+/// and then infer them as a single batch. The previous code ran one `session.run` per image
+/// even when the caller supplied a batch, which meant the session was locked and re-entered per
+/// image and ONNX Runtime never saw more than one row to work with.
+fn infer_prepared(
+    handle: &ModelHandle,
+    prepared: &[preprocess::PreprocessResult],
+    start: Instant,
+) -> Result<Vec<EmbedResult>> {
+    if prepared.is_empty() {
+        return Ok(Vec::new());
+    }
+    let manifest = &handle.manifest;
     let session = handle.pin_session()?;
-    let config = preprocess_config_from_manifest(manifest)?;
-    let prep = preprocess::preprocess(image, &config)?;
-    let original_width = prep.meta.original_width;
-    let original_height = prep.meta.original_height;
 
-    let input_value = TensorRef::from_array_view(&prep.tensor).map_err(crate::engine::ort_err)?;
+    // Concatenate the per-image [1, C, H, W] tensors into one [N, C, H, W] batch. Geometry is
+    // fixed by the manifest, so a mismatch here is a bug rather than bad input.
+    let views: Vec<_> = prepared.iter().map(|p| p.tensor.view()).collect();
+    let batch_tensor = ndarray::concatenate(Axis(0), &views)
+        .map_err(|e| SparrowEngineError::Ort(format!("concatenate encoder batch: {e}")))?;
+
+    let input_value = TensorRef::from_array_view(&batch_tensor).map_err(crate::engine::ort_err)?;
     let mut guard = session
         .lock()
         .map_err(|_| SparrowEngineError::Ort("encoder session lock poisoned".into()))?;
@@ -67,26 +90,14 @@ pub fn embed(handle: &ModelHandle, image: &ImageInput) -> Result<EmbedResult> {
         });
     }
 
-    let mut embedding = extract_embedding_output(&outputs[0], manifest)?;
-    if let Some(dim) = manifest.embedding_dim {
-        if embedding.len() != dim {
-            return Err(SparrowEngineError::OutputShapeMismatch {
-                id: manifest.id.clone(),
-                shape: format!(
-                    "runtime embedding dim {} != manifest dim {dim}",
-                    embedding.len()
-                ),
-                method: manifest.postprocess_method.as_str().to_string(),
-            });
-        }
-    }
+    let mut embeddings = extract_embedding_rows(&outputs[0], manifest, prepared.len())?;
+    drop(outputs);
+    drop(guard);
 
     let normalized = match manifest.postprocess_method {
         PostprocessMethod::Embedding { normalize } => normalize,
         _ => false,
     };
-    finalize_embedding_for_model(&mut embedding, normalized, &manifest.id)?;
-    let dim = embedding.len();
     let metric = manifest.embedding_metric.ok_or_else(|| {
         SparrowEngineError::InvalidManifest("image encoders require [embedding] metric".to_string())
     })?;
@@ -101,32 +112,178 @@ pub fn embed(handle: &ModelHandle, image: &ImageInput) -> Result<EmbedResult> {
         )
     })?;
 
-    drop(outputs);
-    drop(guard);
+    // One elapsed span covers the whole batch, so report it per image rather than stamping the
+    // batch total on every row (which would scale the reported cost with batch position).
+    let processing_time_ms = start.elapsed().as_secs_f32() * 1000.0 / prepared.len() as f32;
+    let mut results = Vec::with_capacity(prepared.len());
+    for (index, prep) in prepared.iter().enumerate() {
+        let mut embedding = std::mem::take(&mut embeddings[index]);
+        if let Some(dim) = manifest.embedding_dim {
+            if embedding.len() != dim {
+                return Err(SparrowEngineError::OutputShapeMismatch {
+                    id: manifest.id.clone(),
+                    shape: format!(
+                        "runtime embedding dim {} != manifest dim {dim}",
+                        embedding.len()
+                    ),
+                    method: manifest.postprocess_method.as_str().to_string(),
+                });
+            }
+        }
+        finalize_embedding_for_model(&mut embedding, normalized, &manifest.id)?;
+        let dim = embedding.len();
+        results.push(EmbedResult {
+            embedding,
+            dim,
+            normalized,
+            metric,
+            model_id: manifest.id.clone(),
+            embedding_version: embedding_version.clone(),
+            model_hash: model_hash.clone(),
+            image_width: prep.meta.original_width,
+            image_height: prep.meta.original_height,
+            processing_time_ms,
+        });
+    }
+    Ok(results)
+}
 
-    Ok(EmbedResult {
-        embedding,
-        dim,
-        normalized,
-        metric,
-        model_id: manifest.id.clone(),
-        embedding_version,
-        model_hash,
-        image_width: original_width,
-        image_height: original_height,
-        processing_time_ms: start.elapsed().as_secs_f32() * 1000.0,
-    })
+/// Images fed to the model in one `session.run`, and the granularity of parallel decode.
+///
+/// Bounds working-set memory for a caller-supplied batch: this is a public entry point, so an
+/// unbounded batch would allocate `n * C * H * W * 4` bytes at the caller's discretion.
+const MAX_INFERENCE_BATCH: usize = 8;
+
+/// Threads used to preprocess a chunk concurrently.
+///
+/// Decode plus resize is pure CPU work with no shared state, so it parallelises cleanly. ONNX
+/// Runtime already threads *inference* internally, so this deliberately does not scale with the
+/// core count: oversubscribing would take cores away from the session's own thread pool. Four
+/// is enough to keep a batch assembled ahead of the model without competing with it.
+///
+/// Override with `SPARROW_ENGINE_ENCODER_DECODE_WORKERS`, shared with the GPU crate. Zero or one
+/// restores fully serial preprocessing.
+fn decode_workers() -> usize {
+    const DEFAULT_WORKERS: usize = 4;
+    const MAX_WORKERS: usize = 16;
+    match std::env::var("SPARROW_ENGINE_ENCODER_DECODE_WORKERS") {
+        Ok(raw) => match raw.trim().parse::<usize>() {
+            Ok(n) => n.clamp(1, MAX_WORKERS),
+            Err(_) => DEFAULT_WORKERS,
+        },
+        Err(_) => DEFAULT_WORKERS,
+    }
 }
 
 /// Run image encoder inference on multiple images, failing the whole batch on the first error.
+///
+/// Preprocesses each chunk across [`decode_workers`] threads and then runs the chunk as **one**
+/// batched `session.run`. Previously this was `images.iter().map(embed).collect()`: every image
+/// paid a separate session lock and a batch-1 inference call, and decoding was fully serial, so
+/// a caller that supplied a batch got none of the benefit of having done so.
 pub fn embed_batch(handle: &ModelHandle, images: &[ImageInput]) -> Result<Vec<EmbedResult>> {
-    images.iter().map(|image| embed(handle, image)).collect()
+    if images.is_empty() {
+        return Ok(Vec::new());
+    }
+    validate_image_encoder(&handle.manifest)?;
+    let config = preprocess_config_from_manifest(&handle.manifest)?;
+    let start = Instant::now();
+
+    let mut results = Vec::with_capacity(images.len());
+    for chunk in images.chunks(MAX_INFERENCE_BATCH) {
+        let workers = decode_workers().min(chunk.len());
+        let prepared = if workers <= 1 {
+            chunk
+                .iter()
+                .map(|image| preprocess::preprocess(image, &config))
+                .collect::<Result<Vec<_>>>()?
+        } else {
+            preprocess_parallel(chunk, &config, workers)?
+        };
+        results.extend(infer_prepared(handle, &prepared, start)?);
+    }
+    Ok(results)
+}
+
+/// Preprocess a chunk across `workers` scoped threads, preserving input order.
+///
+/// Work is handed out by index through an atomic counter and written back into per-index slots,
+/// so the output order always matches the caller's input order regardless of completion order.
+fn preprocess_parallel(
+    chunk: &[ImageInput],
+    config: &crate::preprocess::PreprocessConfig,
+    workers: usize,
+) -> Result<Vec<preprocess::PreprocessResult>> {
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::Mutex;
+
+    let slots: Vec<Mutex<Option<preprocess::PreprocessResult>>> =
+        (0..chunk.len()).map(|_| Mutex::new(None)).collect();
+    let next = AtomicUsize::new(0);
+    let slots_ref = &slots;
+    let next_ref = &next;
+
+    std::thread::scope(|scope| -> Result<()> {
+        let mut handles = Vec::with_capacity(workers);
+        for _ in 0..workers {
+            handles.push(scope.spawn(move || -> Result<()> {
+                loop {
+                    let index = next_ref.fetch_add(1, Ordering::Relaxed);
+                    if index >= chunk.len() {
+                        return Ok(());
+                    }
+                    let prepared = preprocess::preprocess(&chunk[index], config)?;
+                    *slots_ref[index].lock().map_err(|_| {
+                        SparrowEngineError::Ort("encoder preprocess slot poisoned".into())
+                    })? = Some(prepared);
+                }
+            }));
+        }
+        // Join every worker before returning so one failure cannot mask another and no thread
+        // is left detached.
+        let mut first_err = None;
+        for handle in handles {
+            let outcome = match handle.join() {
+                Ok(inner) => inner,
+                Err(_) => Err(SparrowEngineError::Ort(
+                    "encoder preprocess worker panicked".into(),
+                )),
+            };
+            if let Err(err) = outcome {
+                first_err.get_or_insert(err);
+            }
+        }
+        match first_err {
+            Some(err) => Err(err),
+            None => Ok(()),
+        }
+    })?;
+
+    slots
+        .into_iter()
+        .enumerate()
+        .map(|(index, slot)| {
+            slot.into_inner()
+                .map_err(|_| SparrowEngineError::Ort("encoder preprocess slot poisoned".into()))?
+                .ok_or_else(|| {
+                    SparrowEngineError::Ort(format!(
+                        "encoder preprocess produced no result for image {index}"
+                    ))
+                })
+        })
+        .collect()
 }
 
-fn extract_embedding_output(
+/// Split a batched encoder output into one embedding per input image.
+///
+/// Accepts rank-1 (`[dim]`, valid only for a batch of one) and rank-2 (`[batch, dim]`). The
+/// previous single-image helper rejected any output with more than one row; batched inference
+/// makes multi-row the normal case.
+fn extract_embedding_rows(
     output: &ort::value::DynValue,
     manifest: &ModelManifest,
-) -> Result<Vec<f32>> {
+    batch: usize,
+) -> Result<Vec<Vec<f32>>> {
     match output.dtype() {
         ValueType::Tensor {
             ty: TensorElementType::Float32,
@@ -135,7 +292,7 @@ fn extract_embedding_output(
             let output_view: ArrayViewD<'_, f32> = output
                 .try_extract_array::<f32>()
                 .map_err(crate::engine::ort_err)?;
-            extract_embedding_vector(output_view, manifest, |x| x)
+            extract_embedding_vectors(output_view, manifest, batch, |x| x)
         }
         ValueType::Tensor {
             ty: TensorElementType::Float16,
@@ -144,7 +301,7 @@ fn extract_embedding_output(
             let output_view: ArrayViewD<'_, half::f16> = output
                 .try_extract_array::<half::f16>()
                 .map_err(crate::engine::ort_err)?;
-            extract_embedding_vector(output_view, manifest, half::f16::to_f32)
+            extract_embedding_vectors(output_view, manifest, batch, half::f16::to_f32)
         }
         other => Err(SparrowEngineError::OutputShapeMismatch {
             id: manifest.id.clone(),
@@ -154,41 +311,50 @@ fn extract_embedding_output(
     }
 }
 
-fn extract_embedding_vector<T: Copy>(
+fn extract_embedding_vectors<T: Copy>(
     output: ArrayViewD<'_, T>,
     manifest: &ModelManifest,
-    to_f32: impl Fn(T) -> f32,
-) -> Result<Vec<f32>> {
+    batch: usize,
+    to_f32: impl Fn(T) -> f32 + Copy,
+) -> Result<Vec<Vec<f32>>> {
+    let mismatch = |shape: String| SparrowEngineError::OutputShapeMismatch {
+        id: manifest.id.clone(),
+        shape,
+        method: manifest.postprocess_method.as_str().to_string(),
+    };
     match output.ndim() {
         1 => {
+            if batch != 1 {
+                return Err(mismatch(format!(
+                    "rank-1 output for a batch of {batch}; expected [{batch}, dim]"
+                )));
+            }
             let row: ArrayView1<'_, T> = output
                 .into_dimensionality::<ndarray::Ix1>()
                 .map_err(crate::engine::ort_err)?;
-            Ok(row.iter().copied().map(to_f32).collect())
+            Ok(vec![row.iter().copied().map(to_f32).collect()])
         }
         2 => {
             let rows: ArrayView2<'_, T> = output
                 .into_dimensionality::<ndarray::Ix2>()
                 .map_err(crate::engine::ort_err)?;
-            if rows.nrows() != 1 || rows.ncols() == 0 {
-                return Err(SparrowEngineError::OutputShapeMismatch {
-                    id: manifest.id.clone(),
-                    shape: format!("{:?}", rows.shape()),
-                    method: manifest.postprocess_method.as_str().to_string(),
-                });
+            if rows.nrows() != batch || rows.ncols() == 0 {
+                return Err(mismatch(format!(
+                    "{:?} for a batch of {batch}",
+                    rows.shape()
+                )));
             }
-            Ok(rows
-                .index_axis(Axis(0), 0)
-                .iter()
-                .copied()
-                .map(to_f32)
+            Ok((0..batch)
+                .map(|index| {
+                    rows.index_axis(Axis(0), index)
+                        .iter()
+                        .copied()
+                        .map(to_f32)
+                        .collect()
+                })
                 .collect())
         }
-        rank => Err(SparrowEngineError::OutputShapeMismatch {
-            id: manifest.id.clone(),
-            shape: format!("rank {rank}"),
-            method: manifest.postprocess_method.as_str().to_string(),
-        }),
+        rank => Err(mismatch(format!("rank {rank}"))),
     }
 }