diff --git a/docs/rust-api.md b/docs/rust-api.md index 73ef49f..342e608 100644 --- a/docs/rust-api.md +++ b/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 diff --git a/src/vision/models.rs b/src/vision/models.rs index 483c90a..64badc8 100644 --- a/src/vision/models.rs +++ b/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 { validate_manifest(manifest)?; - let managed_root; - 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 root = self.model_root(manifest); let mut artifacts = BTreeMap::new(); for artifact in manifest.artifacts { diff --git a/src/vision/oar.rs b/src/vision/oar.rs index a33f6f9..1c382ca 100644 --- a/src/vision/oar.rs +++ b/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) - .filter(|path| !path.is_empty()) - .map(PathBuf::from) - .unwrap_or_else(default_onnx_runtime_library); + let path = onnx_runtime_library_path(); 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) -> RenderedPage { let transform = PageTransform::from_corners(2, 2, (0.0, 2.0), (2.0, 2.0), (0.0, 0.0)).unwrap(); diff --git a/src/vision/pipeline.rs b/src/vision/pipeline.rs index a00d5c2..c6b5833 100644 --- a/src/vision/pipeline.rs +++ b/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, +} + +static OCR_ENGINE_CACHE: OnceLock>> = 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 models = store.resolve_or_download( - &PP_OCR_V6_SMALL, - options.ocr.model_downloads, - &HttpModelDownloader::default(), - )?; - let engine = OarOcrEngine::from_models(&models)?; + let engine = cached_ocr_engine(&options.ocr)?; 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, 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(¤t).unwrap(); + let relative_root = temporary.path().strip_prefix(¤t).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(); diff --git a/tests/ocr_tests.rs b/tests/ocr_tests.rs index 10c2e59..9d1c995 100644 --- a/tests/ocr_tests.rs +++ b/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(