Perf: OCR batching

Cargo.toml

@@ -21,10 +21,13 @@ rkyv = { version = "0.7.45", features = ["validation"] }
 metrics = "0.24.1"
 
 [profile.release]
-strip = "symbols"
+# strip = "symbols"
 opt-level = 3
-# debug = 1
 
 [profile.dist]
 inherits = "release"
 lto = "thin"
+
+[profile.profiling]
+inherits = "release"
+debug = true

ferrules-api/src/main.rs

@@ -33,8 +33,8 @@ const MAX_SIZE_LIMIT: usize = 250 * 1024 * 1024;
 #[command(author, version, about, long_about = None)]
 struct Args {
     /// OpenTelemetry collector endpoint
-    #[arg(long, env = "OTLP_ENDPOINT", default_value = "http://localhost:4317")]
-    otlp_endpoint: String,
+    #[arg(long, env = "OTLP_ENDPOINT")]
+    otlp_endpoint: Option<String>,
 
     /// Sentry DSN
     #[arg(long, env = "SENTRY_DSN")]
@@ -108,6 +108,17 @@ struct Args {
 
     #[arg(long, short = 'O', help = "Ort graph optimization level")]
     graph_opt_level: Option<usize>,
+
+    /// Enable profiling for layout model
+    #[arg(long, help = "Enable profiling for the layout model (saved as .json)")]
+    profile_layout: bool,
+
+    /// Enable profiling for table transformer model
+    #[arg(
+        long,
+        help = "Enable profiling for the table transformer model (saved as .json)"
+    )]
+    profile_table: bool,
 }
 
 fn parse_ep_args(args: &Args) -> Vec<OrtExecutionProvider> {
@@ -171,15 +182,27 @@ async fn main() {
     };
 
     init_tracing(
-        Some(&args.otlp_endpoint),
+        args.otlp_endpoint.as_deref(),
         "ferrules-api".into(),
         false,
         use_sentry,
     )
     .expect("can't setup tracing for API");
 
     // Initialize Prometheus exporter
-    let builder = metrics_exporter_prometheus::PrometheusBuilder::new();
+    let builder = metrics_exporter_prometheus::PrometheusBuilder::new()
+        .set_buckets_for_metric(
+            metrics_exporter_prometheus::Matcher::Suffix("_ms".to_string()),
+            &[
+                0.0, 1.0, 2.0, 5.0, 10.0, 15.0, 20.0, 25.0, 30.0, 40.0, 50.0, 60.0, 70.0, 80.0,
+                90.0, 100.0, 125.0, 150.0, 175.0, 200.0, 250.0, 300.0, 350.0, 400.0, 450.0, 500.0,
+                600.0, 700.0, 800.0, 900.0, 1000.0, 1250.0, 1500.0, 1750.0, 2000.0, 2500.0, 3000.0,
+                3500.0, 4000.0, 4500.0, 5000.0, 6000.0, 7000.0, 8000.0, 9000.0, 10000.0, 15000.0,
+                20000.0, 30000.0, 45000.0, 60000.0, 90000.0, 120000.0, 180000.0, 240000.0,
+                300000.0,
+            ],
+        )
+        .expect("failed to set buckets");
     let handle = builder
         .install_recorder()
         .expect("failed to install Prometheus recorder");
@@ -189,6 +212,17 @@ async fn main() {
         intra_threads: args.intra_threads,
         inter_threads: args.inter_threads,
         opt_level: args.graph_opt_level.map(|v| v.try_into().unwrap()),
+        warmup: true,
+        profile_layout: if args.profile_layout {
+            Some(std::path::PathBuf::from("profile_layout_api"))
+        } else {
+            None
+        },
+        profile_table: if args.profile_table {
+            Some(std::path::PathBuf::from("profile_table_api"))
+        } else {
+            None
+        },
     };
     // Initialize the layout model and queues
     let parser = FerrulesParser::new(ort_config);

ferrules-cli/src/main.rs

@@ -143,6 +143,17 @@ struct Args {
         help = "Specify the directory to store debug output files"
     )]
     debug_dir: Option<PathBuf>,
+
+    /// Enable profiling for layout model
+    #[arg(long, help = "Enable profiling for the layout model (saved as .json)")]
+    profile_layout: bool,
+
+    /// Enable profiling for table transformer model
+    #[arg(
+        long,
+        help = "Enable profiling for the table transformer model (saved as .json)"
+    )]
+    profile_table: bool,
 }
 
 fn parse_page_range(range_str: &str) -> anyhow::Result<Range<usize>> {
@@ -241,6 +252,17 @@ async fn main() {
         intra_threads: args.intra_threads,
         inter_threads: args.inter_threads,
         opt_level: args.graph_opt_level.map(|v| v.try_into().unwrap()),
+        warmup: false,
+        profile_layout: if args.profile_layout {
+            Some(PathBuf::from("profile_layout"))
+        } else {
+            None
+        },
+        profile_table: if args.profile_table {
+            Some(PathBuf::from("profile_table"))
+        } else {
+            None
+        },
     };
 
     let page_range = match args.page_range {
@@ -471,6 +493,20 @@ async fn main() {
                         ],
                     );
                 }
+                ferrules_core::error::FerrulesError::OcrError(e) => {
+                    format_error(
+                        "OCR Extraction Failed",
+                        "Failed to extract text using OCR.",
+                        vec![
+                            ("Error", e),
+                            ("File", args.file_path.display().to_string()),
+                            (
+                                "Suggestion",
+                                "This might indicate an issue with Apple Vision or stitched image size".to_string(),
+                            ),
+                        ],
+                    );
+                }
             }
             std::process::exit(1);
         }

ferrules-core/src/error.rs

@@ -29,4 +29,6 @@ pub enum FerrulesError {
     TableTransformerModelError(String),
     #[error("table parser error: {0}")]
     TableParserError(String),
+    #[error("ocr parser error: {0}")]
+    OcrError(String),
 }

ferrules-core/src/layout/mod.rs

@@ -12,6 +12,8 @@ use crate::metrics::StepMetrics;
 
 pub mod model;
 
+const CONCURRENT_LAYOUT_REQUESTS: usize = 16;
+
 #[derive(Debug)]
 pub struct Metadata {
     pub(crate) response_tx: oneshot::Sender<anyhow::Result<ParseLayoutResponse>>,
@@ -28,7 +30,7 @@ pub(crate) struct ParseLayoutRequest {
 
 #[derive(Debug)]
 pub(crate) struct ParseLayoutResponse {
-    pub(crate) page_id: PageID,
+    pub(crate) _page_id: PageID,
     pub(crate) layout_bbox: Vec<LayoutBBox>,
     pub(crate) step_metrics: StepMetrics,
 }
@@ -53,17 +55,17 @@ impl ParseLayoutQueue {
         self.queue
             .send((req, span))
             .await
-            .map_err(|_| FerrulesError::LayoutParsingError)
+            .map_err(|_| FerrulesError::LayoutParsingError) // We keep LayoutParsingError for layout itself, but we can add more context later if needed.
     }
 }
 
 async fn start_layout_parser(
     layout_parser: Arc<ORTLayoutParser>,
     mut input_rx: Receiver<(ParseLayoutRequest, Span)>,
 ) {
-    let s = Arc::new(Semaphore::new(layout_parser.config.intra_threads));
+    let s = Arc::new(Semaphore::new(CONCURRENT_LAYOUT_REQUESTS));
     while let Some((req, span)) = input_rx.recv().await {
-        let queue_time = req.metadata.queue_time.elapsed().as_millis();
+        let queue_time = req.metadata.queue_time.elapsed().as_secs_f64() * 1000.0;
         let page_id = req.page_id;
         tracing::debug!("layout request queue time for page {page_id} took: {queue_time}ms");
         let _guard = span.enter();
@@ -78,11 +80,11 @@ async fn handle_request(
     s: Arc<Semaphore>,
     parser: Arc<ORTLayoutParser>,
     req: ParseLayoutRequest,
-    layout_queue_time_ms: u128,
+    layout_queue_time_ms: f64,
 ) {
     let start_wait = Instant::now();
     let _permit = s.acquire().await.unwrap();
-    let idle_time_ms = start_wait.elapsed().as_millis();
+    let idle_time_ms = start_wait.elapsed().as_secs_f64() * 1000.0;
 
     let ParseLayoutRequest {
         page_id,
@@ -95,21 +97,22 @@ async fn handle_request(
     let layout_result = parser
         .parse_layout_async(&page_image, downscale_factor)
         .await;
-    let inference_duration = start.elapsed().as_millis();
+    let inference_duration = start.elapsed().as_secs_f64() * 1000.0;
     drop(_permit);
     tracing::debug!("layout inference time for page {page_id} took: {inference_duration}ms");
 
     let layout_result = layout_result.map(|l| ParseLayoutResponse {
-        page_id,
+        _page_id: page_id,
         layout_bbox: l,
         step_metrics: StepMetrics {
             queue_time_ms: layout_queue_time_ms,
             execution_time_ms: inference_duration,
             idle_time_ms,
         },
     });
-    metadata
-        .response_tx
-        .send(layout_result)
-        .expect("can't send parsed result over oneshot chan");
+    if let Err(e) = layout_result.as_ref() {
+        tracing::error!("Layout parsing failed for page {page_id}: {:?}", e);
+    }
+
+    let _ = metadata.response_tx.send(layout_result);
 }