| |
| |
| """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() |
| |
| 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 |
|
|
| |
| 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 |
| |
| 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": |
| |
| 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: |
| |
| 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) |
|
|
