# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 """Prepared image batches with bounded, overlapping HTTP concurrency and CPU rendering.""" from concurrent.futures import ( ThreadPoolExecutor, ProcessPoolExecutor, wait, FIRST_COMPLETED, ) from dataclasses import dataclass, asdict from io import BytesIO from pathlib import Path from types import SimpleNamespace import base64 import collections import multiprocessing import time from PIL import Image from .engine import geometry from .engine.regions import select_union_regions from .postprocessing import clean_raw_output, postprocess_document @dataclass class PreparedPage: image: Image.Image data_url: str source: str | None @dataclass class PreparedBatch: pages: list[PreparedPage] preprocessing_seconds: float @dataclass class BatchResult: documents: list[dict] metrics: dict def _url(data, mime): return f"data:{mime};base64," + base64.b64encode(data).decode("ascii") def prepare_images(images, *, workers=8): if workers < 1: raise ValueError("workers must be positive") start = time.perf_counter() def one(value): if isinstance(value, Image.Image): image = value.convert("RGB") stream = BytesIO() image.save(stream, format="PNG") return PreparedPage(image, _url(stream.getvalue(), "image/png"), None) data = Path(value).read_bytes() with Image.open(BytesIO(data)) as opened: mime = {"JPEG": "image/jpeg", "PNG": "image/png", "WEBP": "image/webp"}.get( opened.format ) keep_encoded = mime is not None and opened.mode == "RGB" image = opened.convert("RGB") if not keep_encoded: stream = BytesIO() image.save(stream, format="PNG") data, mime = stream.getvalue(), "image/png" return PreparedPage(image, _url(data, mime), str(value)) with ThreadPoolExecutor(max_workers=workers) as pool: pages = list(pool.map(one, images)) if not pages: raise ValueError("No page images provided") return PreparedBatch(pages, time.perf_counter() - start) def _finish(task): raw, crops, size, options, merge_policy = task from .inference import merge_crops raw = clean_raw_output(raw, deduplicate=options.deduplicate) if crops: raw = merge_crops(raw, crops, policy=merge_policy) return postprocess_document(raw, image_size=size, options=options) def _crop_preparer(parser, pages, triple_regions): args = SimpleNamespace( include_detector_labels="table,chart", detector_order="model", detector_score_min=0.0, max_elements_per_page=parser.max_regions, model_coordinate_frame="content_normalized", canvas_width=1664.0, canvas_height=2048.0, pad_position="top-left", ) def page_crops(index, direct_response, detection_response=None, start_region=0): page = pages[index] page_spec = geometry.PageImage( "page", "image.png", Path("image.png"), *page.image.size ) def boxes(response): return geometry.parse_model_detection_predictions( args, [page_spec], [{"key": "page", "raw_out": clean_raw_output(response["raw_out"])}], ) selected = ( boxes(direct_response) if direct_response is not None else [{"key": "page", "detections": []}] ) if detection_response is not None: detected = boxes(detection_response) selected = ( select_union_regions(detected, selected, max_regions=parser.max_regions) if triple_regions == "union" else detected ) boxes_for_page = selected[0]["detections"] items, metadata = [], [] for region_index, region in enumerate( boxes_for_page[start_region:], start_region ): pixels = geometry.crop_bbox_pixels( tuple(region["bbox_norm"]), *page.image.size, margin_px=parser.crop_margin_px, margin_ratio=parser.crop_margin_ratio, ) if pixels[2] - pixels[0] < 8 or pixels[3] - pixels[1] < 8: continue stream = BytesIO() page.image.crop(pixels).save(stream, format="JPEG", quality=95) data = stream.getvalue() # HTTP sends the prepared crop directly; custom backends receive decoded pixels. image = None if parser.backend is not None: with Image.open(BytesIO(data)) as decoded: image = decoded.convert("RGB") items.append((image, _url(data, "image/jpeg"))) metadata.append( { "page_index": index, "crop_id": f"page__{region_index:03d}", "page_key": "page", "detector_label": region["label"], "detector_score": region["score"], "detector_bbox_norm": region["bbox_norm"], "crop_bbox_pixels": pixels, "page_size": page.image.size, } ) return index, boxes_for_page, items, metadata return page_crops def _inference_timing(spans): """Wall time across model calls, including gaps needed to prepare later calls.""" elapsed = max(end for begin, end in spans.values()) - min( begin for begin, end in spans.values() ) return { "inference_seconds": elapsed, "inference_timing_includes": [ "all HTTP/model calls, server queueing and response transfer", "gaps between inference calls, including dependent crop preparation", ], "inference_timing_excludes": [ "initial image preparation", "server startup", "warmup requests", "merge/render work after the final model response", "output serialization", ], "inference_timing_definition": "first model request start to last model response", } def _pipeline(parser, pages, mode, concurrency, workers, region_policy, allow_partial): """Keep one bounded HTTP window full; dependent crops take the next free slot. Detection regions form the stable prefix of the union. They can be submitted before full-page parsing finishes; parser-only additions follow in union order. CPU crop encoding and document rendering never occupy an HTTP worker. """ from .inference import GenerationError count = len(pages) direct = [None] * count detection = [None] * count if mode == "triple" else None regions = [[] for _ in pages] crops = [{} for _ in pages] documents = [None] * count prep_pending = [0] * count crop_pending = [0] * count sealed = [mode == "single"] * count render_started = set() ready = collections.deque() initial = collections.deque( (i, stage, (page.image, page.data_url), None) for i, page in enumerate(pages) for stage in (["fullpage", "detection"] if mode == "triple" else ["fullpage"]) ) futures = {} http_active = 0 peak = 0 spans = {} service = collections.Counter() started = time.perf_counter() page_crops = _crop_preparer(parser, pages, region_policy) def generate(stage, item): begin = time.perf_counter() try: result = parser._generate(item[0], stage, data_url=item[1]) except GenerationError as exc: if not ( allow_partial and exc.response.get("finish_reason") == "length" and exc.response.get("raw_out") ): raise result = dict(exc.response) end = time.perf_counter() return dict(result, latency_s=end - begin), begin, end # Small batches avoid process startup; threaded CPU work still cannot block HTTP. render_executor = ( ProcessPoolExecutor( max_workers=workers, mp_context=multiprocessing.get_context("spawn") ) if workers > 1 and count >= 16 else ThreadPoolExecutor(max_workers=1) ) with ThreadPoolExecutor(max_workers=concurrency) as http, ThreadPoolExecutor( max_workers=min(8, count) ) as prep, render_executor as render: def submit_prep(i, first, second, offset=0): prep_pending[i] += 1 futures[prep.submit(page_crops, i, first, second, offset)] = ( "prep", i, None, ) def maybe_render(i): if ( i in render_started or direct[i] is None or not sealed[i] or prep_pending[i] or crop_pending[i] ): return render_started.add(i) rows = [crops[i][j] for j in sorted(crops[i])] task = ( direct[i]["raw_out"], [r for r in rows if r.get("finish_reason") != "length"], pages[i].image.size, parser.options, parser.merge_policy, ) futures[render.submit(_finish, task)] = ("render", i, None) while initial or ready or futures: while http_active < concurrency and (ready or initial): i, stage, item, metadata = ( ready.popleft() if ready else initial.popleft() ) futures[http.submit(generate, stage, item)] = (stage, i, metadata) http_active += 1 peak = max(peak, http_active) if not futures: break done, _ = wait(futures, return_when=FIRST_COMPLETED) for future in done: kind, i, metadata = futures.pop(future) value = future.result() if kind == "prep": prep_pending[i] -= 1 _, selected, items, rows = value # Two prep jobs can finish out of order; retain the full union. if len(selected) > len(regions[i]): regions[i] = selected for item, row in zip(items, rows): ready.append((i, "crop", item, row)) crop_pending[i] += 1 elif kind == "render": documents[i] = value else: http_active -= 1 response, begin, end = value span = spans.setdefault(kind, [begin, end]) span[0], span[1] = min(span[0], begin), max(span[1], end) service[kind] += end - begin if kind == "crop": crop_pending[i] -= 1 crops[i][metadata["crop_id"]] = {**metadata, **response} elif kind == "fullpage": direct[i] = response if mode == "double": sealed[i] = True submit_prep(i, response, None) elif detection is not None and detection[i] is not None: sealed[i] = True if region_policy == "union": # The detector prefix is already queued. Only prepare additions. detector_count = len( page_crops( i, None, detection[i], parser.max_regions )[1] ) submit_prep(i, response, detection[i], detector_count) else: detection[i] = response submit_prep(i, None, response) if region_policy == "detector": sealed[i] = True elif direct[i] is not None: sealed[i] = True detector_count = len( page_crops(i, None, response, parser.max_regions)[1] ) submit_prep(i, direct[i], response, detector_count) maybe_render(i) assert all(document is not None for document in documents) return ( direct, detection, regions, [[rows[j] for j in sorted(rows)] for rows in crops], documents, { **_inference_timing(spans), "peak_inflight_requests": peak, "stage_times": { stage + "_span_seconds": end - begin for stage, (begin, end) in spans.items() }, "request_service_seconds": dict(service), "stage_first_request_seconds": { stage: span[0] - started for stage, span in spans.items() }, "overlapping_stage_times": True, }, ) def parse_prepared( parser, prepared, *, mode="single", concurrency=128, postprocess_workers=8, triple_regions="union", allow_partial=False, schedule="pipelined", ): from .inference import GenerationError if mode not in ("single", "double", "triple") or triple_regions not in ( "detector", "union", ): raise ValueError("Invalid pipeline mode or region policy") if concurrency < 1 or postprocess_workers < 1: raise ValueError("Worker counts must be positive") if not isinstance(prepared, PreparedBatch) or not prepared.pages: raise ValueError("Use parser.prepare() to prepare a nonempty image batch") if schedule not in ("staged", "pipelined"): raise ValueError("schedule must be staged or pipelined") pages = prepared.pages start = time.perf_counter() stage_times = {} request_spans = {} def requests(items, stage): stage_start = time.perf_counter() def one(item): image, url = item request_start = time.perf_counter() try: response = parser._generate(image, stage, data_url=url) except GenerationError as exc: # Opt-in partial output is explicit and never counted as complete. if ( not allow_partial or exc.response.get("finish_reason") != "length" or not exc.response.get("raw_out") ): raise response = dict(exc.response) request_end = time.perf_counter() response = dict(response, latency_s=request_end - request_start) return response, request_start, request_end with ThreadPoolExecutor(max_workers=concurrency) as pool: result = list(pool.map(one, items)) stage_times[stage + "_seconds"] = time.perf_counter() - stage_start request_spans[stage] = ( min(begin for response, begin, end in result), max(end for response, begin, end in result), ) return [response for response, begin, end in result] if schedule == "pipelined": direct, detection, regions, crops_by_page, documents, stream_metrics = ( _pipeline( parser, pages, mode, concurrency, postprocess_workers, triple_regions, allow_partial, ) ) stage_times.update(stream_metrics.pop("stage_times")) crop_responses = [r for crops in crops_by_page for r in crops] else: direct = requests([(p.image, p.data_url) for p in pages], "fullpage") detection = ( requests([(p.image, p.data_url) for p in pages], "detection") if mode == "triple" else None ) regions = [[] for _ in pages] crop_metadata = [] crop_items = [] crop_responses = [] if mode != "single": prep_start = time.perf_counter() page_crops = _crop_preparer(parser, pages, triple_regions) with ThreadPoolExecutor(max_workers=min(8, len(pages))) as pool: for index, selected, items, metadata in pool.map( lambda i: page_crops( i, direct[i], detection[i] if detection is not None else None ), range(len(pages)), ): regions[index] = selected crop_items.extend(items) crop_metadata.extend(metadata) stage_times["crop_preprocessing_seconds"] = time.perf_counter() - prep_start crop_responses = requests(crop_items, "crop") if crop_items else [] crops_by_page = [[] for _ in pages] for metadata, response in zip(crop_metadata, crop_responses): crops_by_page[metadata["page_index"]].append({**metadata, **response}) post_start = time.perf_counter() tasks = [ ( response["raw_out"], [r for r in crops if r.get("finish_reason") != "length"], page.image.size, parser.options, parser.merge_policy, ) for page, response, crops in zip(pages, direct, crops_by_page) ] if postprocess_workers == 1 or len(tasks) < 16: documents = list(map(_finish, tasks)) else: with ProcessPoolExecutor( max_workers=postprocess_workers, mp_context=multiprocessing.get_context("spawn"), ) as pool: documents = list(pool.map(_finish, tasks, chunksize=8)) stage_times["postprocessing_seconds"] = time.perf_counter() - post_start stream_metrics = _inference_timing(request_spans) capped = collections.Counter() completed = 0 for index, document in enumerate(documents): stages = {"fullpage": direct[index], "crops": crops_by_page[index]} if detection is not None: stages["detection"] = detection[index] errors = [] for stage in ["fullpage", "detection"]: if stage in stages and stages[stage].get("finish_reason") == "length": capped[stage] += 1 errors.append(stage + " reached its token limit") for crop in stages["crops"]: if crop.get("finish_reason") == "length": capped["crop"] += 1 errors.append( "crop " + crop["crop_id"] + " reached its token limit; refinement omitted" ) document["complete"] = not errors document["generation_warnings"] = errors document["stages"] = stages document["source"] = pages[index].source document["inference"] = { "mode": mode, "merge_policy": parser.merge_policy, "triple_regions": triple_regions if mode == "triple" else None, "model": parser.model, "coordinate_frame": "content_normalized", "regions": regions[index], "sampling": { name: asdict(value) for name, value in parser.settings.items() }, "postprocessing": asdict(parser.options), "crop_jpeg_quality": 95, "crop_margin_ratio": parser.crop_margin_ratio, "crop_margin_px": parser.crop_margin_px, } completed += not errors elapsed = time.perf_counter() - start all_responses = direct + (detection or []) + crop_responses metrics = { "mode": mode, "merge_policy": parser.merge_policy, "schedule": schedule, **stream_metrics, "pages": len(pages), "complete_pages": completed, "incomplete_pages": len(pages) - completed, "concurrency": concurrency, "postprocess_workers": postprocess_workers, "preprocessing_seconds": prepared.preprocessing_seconds, "inference_pages_per_second": ( len(pages) / stream_metrics["inference_seconds"] if stream_metrics["inference_seconds"] > 0 else None ), "complete_inference_pages_per_second": ( completed / stream_metrics["inference_seconds"] if stream_metrics["inference_seconds"] > 0 else None ), "pipeline_seconds": elapsed, "processed_pages_per_second": len(pages) / elapsed, "complete_pages_per_second": completed / elapsed, "stage_times": stage_times, "requests": { "fullpage": len(direct), "detection": len(detection or []), "crop": len(crop_responses), }, "token_limited_responses": dict(capped), "completion_tokens": sum( int(r.get("usage", {}).get("completion_tokens") or 0) for r in all_responses ), "postprocessing": asdict(parser.options), "allow_partial": allow_partial, "timing_excludes": [ "initial image preparation", "output serialization", "server startup", "warmup requests", ], "timing_includes": [ "all mode inference stages", "predicted-region crop preparation", "merge and document rendering", ], } return BatchResult(documents, metrics)