perf(vision): reuse OCR runtime sessions (#391)

* perf(vision): reuse OCR runtime sessions

* fix(vision): harden OCR engine caching

docs/rust-api.md

@@ -380,6 +380,19 @@ provides the same native-only behavior through the OCR result/provenance
 shape; `Force` renders every selected page. Learned layout intentionally
 returns an explicit unsupported error in this lightweight pipeline.
 
+The one-call API keeps the most recently used verified OCR engine in process.
+Long-lived workers therefore verify the pinned artifacts and build the ONNX
+sessions once, then reuse those loaded sessions across documents. The cache is
+bounded to one model configuration and keyed by normalized model/runtime paths
+plus the pinned manifest revision and artifact digests; switching the model
+directory, runtime library, or compiled manifest replaces it. An active engine
+owns the model data it already verified, so mutating artifacts in place does
+not hot-reload a running process; restart the process when intentionally
+replacing files at the same paths. CPU inference uses at most four intra-op
+threads per ONNX session so a single small page does not oversubscribe larger
+hosts, and recognizes variable-width line crops individually to avoid
+padding-heavy CPU batches.
+
 Build the CLI with the same opt-in feature:
 
 ```bash

src/vision/models.rs

@@ -190,16 +190,17 @@ impl ModelStore {
         &self.cache_root
     }
 
+    /// Effective directory containing one manifest's artifacts.
+    pub(crate) fn model_root(&self, manifest: &ModelManifest) -> PathBuf {
+        self.override_root
+            .clone()
+            .unwrap_or_else(|| self.manifest_cache_root(manifest))
+    }
+
     /// Validates and resolves every required artifact.
     pub fn resolve(&self, manifest: &ModelManifest) -> Result<ModelPaths, ModelStoreError> {
         validate_manifest(manifest)?;
-        let managed_root;
+        let root = self.model_root(manifest);
-        let root = if let Some(root) = self.override_root.as_deref() {
-            root
-        } else {
-            managed_root = self.manifest_cache_root(manifest);
-            managed_root.as_path()
-        };
 
         let mut artifacts = BTreeMap::new();
         for artifact in manifest.artifacts {

src/vision/oar.rs

@@ -4,6 +4,7 @@ use std::path::PathBuf;
 use std::time::Instant;
 
 use image::RgbImage;
+use oar_ocr::core::config::onnx::OrtSessionConfig;
 use oar_ocr::oarocr::{OAROCRBuilder, OAROCR};
 use oar_ocr::processors::BoundingBox;
 use thiserror::Error;
@@ -86,7 +87,13 @@ impl OarOcrEngine {
         let recognition = required_model(models, ModelArtifactKind::TextRecognition)?;
         let dictionary = required_model(models, ModelArtifactKind::CharacterDictionary)?;
 
-        let pipeline = OAROCRBuilder::new(detection, recognition, dictionary).build()?;
+        let pipeline = OAROCRBuilder::new(detection, recognition, dictionary)
+            .ort_session(ocr_session_config())
+            // Document line crops often have very different widths. Keeping
+            // CPU recognition batches at one avoids padding every crop to the
+            // widest line, reducing both inference work and peak memory.
+            .region_batch_size(1)
+            .build()?;
         let model = ModelIdentity::new(models.manifest_id(), models.revision());
         Ok(Self { pipeline, model })
     }
@@ -169,11 +176,18 @@ impl OarOcrEngine {
     }
 }
 
+fn ocr_session_config() -> OrtSessionConfig {
+    let available = std::thread::available_parallelism()
+        .map(std::num::NonZeroUsize::get)
+        .unwrap_or(1);
+    OrtSessionConfig::new()
+        .with_intra_threads(available.min(4))
+        .with_inter_threads(1)
+        .with_parallel_execution(false)
+}
+
 fn load_onnx_runtime() -> Result<(), OarOcrError> {
-    let path = std::env::var_os(ONNX_RUNTIME_LIBRARY_ENV)
+    let path = onnx_runtime_library_path();
-        .filter(|path| !path.is_empty())
-        .map(PathBuf::from)
-        .unwrap_or_else(default_onnx_runtime_library);
     drop(
         ort::init_from(&path).map_err(|source| OarOcrError::OnnxRuntimeLoad {
             path: path.clone(),
@@ -183,6 +197,13 @@ fn load_onnx_runtime() -> Result<(), OarOcrError> {
     Ok(())
 }
 
+pub(crate) fn onnx_runtime_library_path() -> PathBuf {
+    std::env::var_os(ONNX_RUNTIME_LIBRARY_ENV)
+        .filter(|path| !path.is_empty())
+        .map(PathBuf::from)
+        .unwrap_or_else(default_onnx_runtime_library)
+}
+
 fn default_onnx_runtime_library() -> PathBuf {
     #[cfg(target_os = "windows")]
     const NAME: &str = "onnxruntime.dll";
@@ -355,6 +376,14 @@ mod tests {
     use super::*;
     use crate::vision::PageTransform;
 
+    #[test]
+    fn cpu_session_budget_is_bounded_for_small_ocr_models() {
+        let config = ocr_session_config();
+        assert!((1..=4).contains(&config.intra_threads.unwrap()));
+        assert_eq!(config.inter_threads, Some(1));
+        assert_eq!(config.parallel_execution, Some(false));
+    }
+
     fn page(format: RenderPixelFormat, stride: usize, pixels: Vec<u8>) -> RenderedPage {
         let transform =
             PageTransform::from_corners(2, 2, (0.0, 2.0), (2.0, 2.0), (0.0, 0.0)).unwrap();

src/vision/pipeline.rs

@@ -1,13 +1,15 @@
 //! One-call native extraction and OCR pipeline.
 
 use std::collections::BTreeSet;
-use std::path::Path;
+use std::path::{Path, PathBuf};
+use std::sync::{Arc, Mutex, OnceLock};
 use std::time::Instant;
 
 use thiserror::Error;
 
 use crate::{MarkdownOptions, PageOcrReasons, PdfError};
 
+use super::oar::onnx_runtime_library_path;
 use super::{
     fuse_ocr_pages, route_ocr_pages, run_ocr_pages, FusedPageMarkdown, HttpModelDownloadError,
     HttpModelDownloader, ModelAcquireError, ModelStore, ModelStoreError, OarOcrEngine, OarOcrError,
@@ -15,6 +17,24 @@ use super::{
     PdfiumRenderer, RenderError, RenderOptions, PP_OCR_V6_SMALL,
 };
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct OcrEngineCacheKey {
+    model_root: PathBuf,
+    runtime_library: PathBuf,
+    manifest_schema: u32,
+    manifest_id: &'static str,
+    manifest_revision: &'static str,
+    artifact_digests: Vec<&'static str>,
+}
+
+#[derive(Debug)]
+struct CachedOcrEngine {
+    key: OcrEngineCacheKey,
+    engine: Arc<OarOcrEngine>,
+}
+
+static OCR_ENGINE_CACHE: OnceLock<Mutex<Option<CachedOcrEngine>>> = OnceLock::new();
+
 /// Options for native extraction with optional OCR.
 #[derive(Clone)]
 pub struct OcrPdfOptions {
@@ -220,16 +240,10 @@ pub fn process_pdf_with_ocr_mem(
         // Resolve the native renderer before any network request so a missing
         // PDFium installation cannot trigger a model download it cannot use.
         let renderer = PdfiumRenderer::load()?;
-        let store = ModelStore::from_options(&options.ocr)?;
+        let engine = cached_ocr_engine(&options.ocr)?;
-        let models = store.resolve_or_download(
-            &PP_OCR_V6_SMALL,
-            options.ocr.model_downloads,
-            &HttpModelDownloader::default(),
-        )?;
-        let engine = OarOcrEngine::from_models(&models)?;
         run_ocr_pages(
             &renderer,
-            &engine,
+            engine.as_ref(),
             buffer,
             &routed,
             options.password.as_deref(),
@@ -268,6 +282,84 @@ pub fn process_pdf_with_ocr_mem(
     })
 }
 
+fn cached_ocr_engine(options: &OcrOptions) -> Result<Arc<OarOcrEngine>, OcrPipelineError> {
+    let store = ModelStore::from_options(options)?;
+    let key = OcrEngineCacheKey {
+        model_root: normalized_cache_path(store.model_root(&PP_OCR_V6_SMALL)),
+        runtime_library: normalized_cache_path(onnx_runtime_library_path()),
+        manifest_schema: PP_OCR_V6_SMALL.schema_version,
+        manifest_id: PP_OCR_V6_SMALL.id,
+        manifest_revision: PP_OCR_V6_SMALL.revision,
+        artifact_digests: PP_OCR_V6_SMALL
+            .artifacts
+            .iter()
+            .map(|artifact| artifact.sha256)
+            .collect(),
+    };
+    let cache = OCR_ENGINE_CACHE.get_or_init(|| Mutex::new(None));
+    {
+        let cached = cache.lock().unwrap_or_else(|error| error.into_inner());
+        if let Some(cached) = cached.as_ref().filter(|cached| cached.key == key) {
+            return Ok(Arc::clone(&cached.engine));
+        }
+    }
+
+    // The loaded sessions own the verified model data, so a cache hit never
+    // needs to trust or reopen mutable files on disk. Do the expensive download
+    // and session construction outside the process-wide cache lock so unrelated
+    // OCR requests can continue using a warm engine. Concurrent cold misses may
+    // build redundantly; the second cache check keeps only one shared session.
+    let models = store.resolve_or_download(
+        &PP_OCR_V6_SMALL,
+        options.model_downloads,
+        &HttpModelDownloader::default(),
+    )?;
+    let engine = Arc::new(OarOcrEngine::from_models(&models)?);
+
+    let mut cached = cache.lock().unwrap_or_else(|error| error.into_inner());
+    if let Some(cached) = cached.as_ref().filter(|cached| cached.key == key) {
+        return Ok(Arc::clone(&cached.engine));
+    }
+    *cached = Some(CachedOcrEngine {
+        key,
+        engine: Arc::clone(&engine),
+    });
+    Ok(engine)
+}
+
+fn normalized_cache_path(path: PathBuf) -> PathBuf {
+    let absolute = if path.is_absolute() {
+        path
+    } else {
+        std::env::current_dir()
+            .map(|directory| directory.join(&path))
+            .unwrap_or(path)
+    };
+    if let Ok(canonical) = std::fs::canonicalize(&absolute) {
+        return canonical;
+    }
+
+    // A managed cache often does not exist on the first request. Resolve the
+    // nearest existing ancestor so a relative path has the same key before
+    // and after model acquisition creates its final directories.
+    let mut ancestor = absolute.as_path();
+    let mut suffix = Vec::new();
+    while let Some(name) = ancestor.file_name() {
+        suffix.push(name.to_os_string());
+        let Some(parent) = ancestor.parent() else {
+            break;
+        };
+        ancestor = parent;
+        if let Ok(mut canonical) = std::fs::canonicalize(ancestor) {
+            for component in suffix.iter().rev() {
+                canonical.push(component);
+            }
+            return canonical;
+        }
+    }
+    absolute
+}
+
 fn assemble_document_markdown(pages: &[FusedPageMarkdown], include_page_numbers: bool) -> String {
     let mut document = String::new();
     for (index, page) in pages.iter().enumerate() {
@@ -335,6 +427,21 @@ pub enum OcrPipelineError {
 mod tests {
     use super::*;
 
+    #[test]
+    fn cache_path_is_stable_when_a_relative_directory_is_created() {
+        let current = std::fs::canonicalize(std::env::current_dir().unwrap()).unwrap();
+        let temporary = tempfile::tempdir_in(&current).unwrap();
+        let relative_root = temporary.path().strip_prefix(&current).unwrap();
+        let relative = relative_root.join("models").join("revision");
+
+        let before = normalized_cache_path(relative.clone());
+        std::fs::create_dir_all(&relative).unwrap();
+        let after = normalized_cache_path(relative);
+
+        assert!(before.is_absolute());
+        assert_eq!(before, after);
+    }
+
     #[test]
     fn off_mode_extracts_native_text_without_runtime_side_effects() {
         let bytes = std::fs::read("tests/fixtures/thermo-freon12.pdf").unwrap();

tests/ocr_tests.rs

@@ -121,7 +121,9 @@ fn complete_ocr_pipeline_routes_and_assembles_a_scanned_fixture() {
         .minimum_confidence(0.3)
         .model_directory(model_directory)
         .model_downloads(ModelDownloadPolicy::Offline);
-    let result = process_pdf_with_ocr_mem(&bytes, OcrPdfOptions::new().ocr(ocr)).unwrap();
+    let options = OcrPdfOptions::new().ocr(ocr);
+    let result = process_pdf_with_ocr_mem(&bytes, options.clone()).unwrap();
+    let repeated = process_pdf_with_ocr_mem(&bytes, options).unwrap();
 
     assert_eq!(result.pages_routed_to_ocr, vec![1]);
     assert!(!result.markdown.trim().is_empty());
@@ -130,6 +132,7 @@ fn complete_ocr_pipeline_routes_and_assembles_a_scanned_fixture() {
         result.pages[0].provenance.ocr_model.as_ref().unwrap().name,
         PP_OCR_V6_SMALL.id
     );
+    assert_eq!(repeated.markdown, result.markdown);
 }
 
 fn recognize(