perf(vision): reuse OCR runtime sessions
This commit is contained in:
@@ -380,6 +380,15 @@ 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; switching the model directory or runtime
|
||||
library replaces it. 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
|
||||
|
||||
@@ -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 = 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 {
|
||||
|
||||
+34
-5
@@ -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<u8>) -> RenderedPage {
|
||||
let transform =
|
||||
PageTransform::from_corners(2, 2, (0.0, 2.0), (2.0, 2.0), (0.0, 0.0)).unwrap();
|
||||
|
||||
+53
-9
@@ -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,20 @@ use super::{
|
||||
PdfiumRenderer, RenderError, RenderOptions, PP_OCR_V6_SMALL,
|
||||
};
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct OcrEngineCacheKey {
|
||||
model_root: PathBuf,
|
||||
runtime_library: PathBuf,
|
||||
}
|
||||
|
||||
#[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 +236,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 +278,40 @@ 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: canonicalize_if_present(store.model_root(&PP_OCR_V6_SMALL)),
|
||||
runtime_library: canonicalize_if_present(onnx_runtime_library_path()),
|
||||
};
|
||||
let cache = OCR_ENGINE_CACHE.get_or_init(|| Mutex::new(None));
|
||||
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));
|
||||
}
|
||||
|
||||
// Model verification and session construction happen once per active
|
||||
// configuration. The loaded sessions own the verified model data, so a
|
||||
// cache hit never needs to trust or reopen mutable files on disk. Keeping
|
||||
// the lock through initialization also prevents a first-request stampede.
|
||||
let models = store.resolve_or_download(
|
||||
&PP_OCR_V6_SMALL,
|
||||
options.model_downloads,
|
||||
&HttpModelDownloader::default(),
|
||||
)?;
|
||||
let engine = Arc::new(OarOcrEngine::from_models(&models)?);
|
||||
|
||||
*cached = Some(CachedOcrEngine {
|
||||
key,
|
||||
engine: Arc::clone(&engine),
|
||||
});
|
||||
Ok(engine)
|
||||
}
|
||||
|
||||
fn canonicalize_if_present(path: PathBuf) -> PathBuf {
|
||||
std::fs::canonicalize(&path).unwrap_or(path)
|
||||
}
|
||||
|
||||
fn assemble_document_markdown(pages: &[FusedPageMarkdown], include_page_numbers: bool) -> String {
|
||||
let mut document = String::new();
|
||||
for (index, page) in pages.iter().enumerate() {
|
||||
|
||||
+4
-1
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user