fix(ocr): stop task="parse" rejecting ordinary multi-page PDFs

doc/source/models/model_abilities/image.rst

@@ -452,14 +452,41 @@ document to merge across pages, so it requires a PDF upload and does not accept
   text crops are needed too. The field is omitted for elements without a crop.
 
 Parsing has its own size limits, and they are tighter than the per-page OCR
-path's. When a render finds no text at all, DeepDoc re-renders the whole
-document at three times the zoom, repeatedly, until the scale reaches 9 — so a
-request at ``zoomin=3`` may end up rendering at 9, and one at ``zoomin=1`` also
-ends at 9. Because that allocation is real, both the per-page and the
-whole-document budgets are enforced against the largest scale a run can reach,
-counting the render being replaced as well.
-
-In practice a request is rejected with a 400 when a single page would peak above
-200 megapixels (an A3 page is fine at ``zoomin=3``, but not at ``6``), or when
-the pages together would peak above 1 gigapixel — roughly 22 A4 pages at the
-default zoom. Lower ``zoomin`` or split the document if you hit either.
+path's. When a render finds no text *anywhere in the document*, DeepDoc
+re-renders the whole thing at three times the zoom, repeatedly, until the scale
+reaches 9 — so a request at ``zoomin=3`` may end up rendering at 9.
+
+With ``deepdoc-lib`` 0.2.2 that re-render is in practice unreachable for any
+document that renders at all — DeepDoc appends to its box list on every page,
+including an empty list for a page that yields nothing, so the
+``len(boxes) == 0`` condition it guards on only holds when there were no pages
+to render. The document is therefore budgeted at the scale you asked for, with
+a separate ceiling bounding what the re-render would cost should a later
+release make it reachable again. Three budgets apply:
+
+* **Per page**, enforced at the worst-case scale, since one page with an
+  outsized MediaBox must not be admitted on the strength of a retry that may
+  still fire: a page may not peak above 200 megapixels. An A3 page is fine at
+  ``zoomin=3`` but not at ``6``.
+* **Whole document**, at the requested scale: the pages together may not
+  exceed 1 gigapixel, roughly 221 A4 pages at the default zoom, with the
+  200-page ceiling capping it from the other side.
+* **Whole document, if the re-render happens**: the escalated peak may not
+  exceed 6 gigapixels, about 24 GB of page images. This is what limits long
+  documents in practice — roughly 130 A4 pages at the default zoom — and it is
+  deliberately not derived from the other two, whose product would permit some
+  160 GB.
+
+Note that the per-page budget is **not monotonic** in ``zoomin``, because the
+retry ladder is not: DeepDoc tests ``zoomin < 9`` before multiplying, so
+``zoomin=2`` and ``zoomin=6`` both escalate to 18x while ``zoomin=3`` stops at
+9x. Lowering ``zoomin`` can therefore make the per-page budget *larger*. For
+that reason a 400 from these limits names a ``zoomin`` that would actually fit
+whenever one exists, and otherwise says to split the document — follow what
+the message says rather than assuming a lower zoom will help.
+
+Both whole-document ceilings can be raised on deployments whose parse workers
+are sized for it, via ``XINFERENCE_MAX_PDF_PARSE_TOTAL_PIXELS`` and
+``XINFERENCE_MAX_PDF_PARSE_RETRY_TOTAL_PIXELS`` (both in pixels). A rendered
+page costs roughly 4 bytes per pixel, so the defaults correspond to about 4 GB
+and 24 GB of page images respectively.

xinference/api/pdf_ocr.py

@@ -21,9 +21,13 @@
 """
 
 import json
+import logging
+import os
 import threading
 from typing import Any, Generator, List, Optional, Tuple, Union
 
+logger = logging.getLogger(__name__)
+
 # PDFium is not thread-safe: pypdfium2 documents that concurrent calls,
 # even on different documents, may crash or corrupt the process. Every
 # PDFium operation in this module (document open, page access, rendering,
@@ -46,11 +50,15 @@
 WHOLE_DOCUMENT_OCR_TASKS = frozenset({"parse"})
 # Whole-document parsers render at ``72 * zoomin`` DPI, i.e. a scale of
 # ``zoomin`` over the page's point size. DeepDoc additionally re-renders at
-# ``zoomin * 3`` when a first pass finds no OCR boxes -- its recovery path
-# for pages the first render was too coarse to read -- but only while
-# ``zoomin < 9``, so a parse started at or above that never amplifies. That
-# retry is part of the parsing behaviour and is left intact, so the budget
-# below is enforced against the largest scale a run can actually reach.
+# ``zoomin * 3`` when a pass finds no OCR boxes *anywhere in the document* --
+# its recovery path for a render too coarse to read -- but only while
+# ``zoomin < 9``, so a parse started at or above that never amplifies.
+#
+# The guard tests the pre-multiplication value, so the ladder overshoots the
+# limit for any zoomin that is not a power-of-three divisor of 9: 2 escalates
+# 2 -> 6 -> 18, and 6 goes straight to 18. That makes the reachable scale
+# non-monotonic in zoomin (2 and 6 reach 18, but 3 stops at 9), which is a
+# defect in deepdoc-lib rather than here; see ``worst_case_parse_zoom``.
 PDF_PARSE_RETRY_ZOOM_FACTOR = 3
 PDF_PARSE_RETRY_ZOOM_LIMIT = 9
 
@@ -61,6 +69,14 @@ def worst_case_parse_zoom(zoomin: int) -> int:
     The retry is recursive: each pass that still finds no boxes triples the
     zoom again, and only the ``< limit`` guard stops it. So zoomin 1 renders
     at 1, 3 and finally 9, not just 3.
+
+    Because deepdoc-lib checks ``zoomin < 9`` *before* multiplying, the last
+    step can overshoot: 2 reaches 18 and 6 reaches 18, while 3 stops at 9.
+    This models that faithfully rather than clamping, so the number used for
+    budgeting is the one the parser can really allocate. The consequence is
+    that this is not monotonic in ``zoomin`` -- which is why the per-page
+    limit, the only budget still enforced at this scale, is checked against
+    every candidate zoom before one is recommended to the caller.
     """
     scale = zoomin
     while scale < PDF_PARSE_RETRY_ZOOM_LIMIT:
@@ -97,21 +113,142 @@ def worst_case_parse_peak_pixels(width: float, height: float, zoomin: int) -> in
 MAX_PDF_PARSE_PAGE_PIXELS = 200_000_000
 # Unlike the per-page path, which rasterizes lazily and holds one page at a
 # time, whole-document parsers render every page up front and keep them all
-# alive at once, so the pages are budgeted together as well -- and, like the
-# per-page ceiling, at the scale a retry can actually reach.
+# alive at once, so the pages are budgeted together as well.
+#
+# This one is enforced at the *requested* scale. Charging every document for
+# the retry is what stopped an ordinary 31-page A4 text PDF parsing at every
+# zoomin (#5307): it was budgeted at 45 MP per page where it really renders
+# at 4.5.
 #
-# Peak usage is the final render plus the one before it: deepdoc assigns
-# ``self.page_images = [...]``, and the comprehension is fully built before
-# the name is rebound, so the previous render is still referenced while the
-# new one is allocated. ``worst_case_parse_peak_pixels`` accounts for both.
+# The retry gets its own, looser ceiling below rather than being priced into
+# this one, because for any document that renders at all it is unreachable
+# with deepdoc-lib 0.2.2. deepdoc's guard is ``len(self.boxes) == 0``, but
+# ``__ocr`` appends to ``boxes`` on *every* page -- ``self.boxes.append([])``
+# when a page yields nothing, ``append(bxs)`` otherwise -- so after the OCR
+# pass ``len(self.boxes)`` equals the page count. It is zero only when there
+# were no pages to render in the first place, i.e. the load failed, in which
+# case there is nothing to re-render. Charging this budget for a 9x pass the
+# parser cannot take would reject a 74-page A4 PDF rendering ~334 M pixels.
 #
-# At 4 GB this admits roughly 22 A4 pages, well short of the 200-page ceiling
-# the per-page path allows. That is the honest consequence of budgeting for a
-# retry that renders at 9x: a document only reaches this if every page yields
-# no text, but the allocation is real when it does, and permitting tens of
-# gigabytes would leave the OOM open. Raise this if parse workers are sized
-# for it; `pages`-style batching in deepdoc-lib would lift it properly.
+# At 1 G px and ~4 bytes per rendered pixel this is ~4 GB of page images at
+# the requested zoom: ~221 A4 pages at the default zoomin 3, 55 at zoomin 6,
+# with the 200-page ceiling capping it from the other side.
 MAX_PDF_PARSE_TOTAL_PIXELS = 1_000_000_000
+# The unreachability argued above is a statement about one release, and the
+# dependency is ``deepdoc-lib~=0.2.2``, which accepts any later 0.2.x. If
+# upstream corrects the guard to the likely intended ``not any(self.boxes)``,
+# the retry becomes reachable, so what it would then allocate is bounded here
+# rather than left to that argument holding.
+#
+# The per-page ceiling does not supply a useful one. It caps a single page,
+# and multiplied by the 200-page ceiling it admits 40 G px -- some 160 GB of
+# page images. That is a mathematical bound, not a memory-safety one: the
+# concrete case of 200 A4 pages escalating 3 -> 9 is 9.0 G px (~36 GB), well
+# inside it and well past what an ordinary worker survives.
+#
+# This ceiling is what keeps that in range, and unlike the requested-scale
+# budget it is applied to the *peak* -- summed from
+# ``worst_case_parse_peak_pixels``, so the render being replaced is counted
+# alongside its replacement.
+#
+# At 6 G px (~24 GB) the 31-page A4 document from #5307 passes at every
+# zoomin from 1 to 6 (1.4 G px at the default), a 100-page A4 document still
+# parses at the default zoom, and the 36 GB case above is rejected. Raise it
+# via the environment on workers sized for more; lowering it trades long
+# documents for a smaller escalated worst case.
+MAX_PDF_PARSE_RETRY_TOTAL_PIXELS = 6_000_000_000
+
+
+def _pixel_budget_from_env(name: str, default: int) -> int:
+    """Read a pixel ceiling from the environment, falling back to ``default``.
+
+    A malformed or non-positive value is ignored rather than raising: this
+    runs at import time, and a typo in a deployment's environment should not
+    take the API process down.
+    """
+    raw = os.environ.get(name)
+    if not raw:
+        return default
+    try:
+        value = int(raw)
+    except ValueError:
+        logger.warning("Ignoring %s=%r: not an integer", name, raw)
+        return default
+    if value <= 0:
+        logger.warning("Ignoring %s=%r: must be positive", name, raw)
+        return default
+    return value
+
+
+MAX_PDF_PARSE_TOTAL_PIXELS = _pixel_budget_from_env(
+    "XINFERENCE_MAX_PDF_PARSE_TOTAL_PIXELS", MAX_PDF_PARSE_TOTAL_PIXELS
+)
+MAX_PDF_PARSE_RETRY_TOTAL_PIXELS = _pixel_budget_from_env(
+    "XINFERENCE_MAX_PDF_PARSE_RETRY_TOTAL_PIXELS", MAX_PDF_PARSE_RETRY_TOTAL_PIXELS
+)
+
+
+def _parse_budget_error(sizes: List[Tuple[float, float]], zoomin: int) -> Optional[str]:
+    """Why a document does not fit at ``zoomin``, or ``None`` if it does.
+
+    Three ceilings apply: the per-page one at the worst-case scale, the
+    document-wide one at the requested scale, and a looser document-wide one
+    on what a retry would peak at. See ``MAX_PDF_PARSE_TOTAL_PIXELS`` and
+    ``MAX_PDF_PARSE_RETRY_TOTAL_PIXELS`` for why the last two differ.
+    """
+    worst_case = worst_case_parse_zoom(zoomin)
+    requested_total = 0
+    retry_total = 0
+    for index, (width, height) in enumerate(sizes):
+        peak_pixels = worst_case_parse_peak_pixels(width, height, zoomin)
+        if peak_pixels > MAX_PDF_PARSE_PAGE_PIXELS:
+            return (
+                f"Page {index + 1} would rasterize to {peak_pixels} pixels at "
+                f"zoomin {zoomin} (the parser can re-render at up to "
+                f"{worst_case}x), exceeding the per-page limit of "
+                f"{MAX_PDF_PARSE_PAGE_PIXELS}"
+            )
+        requested_total += int(width * zoomin) * int(height * zoomin)
+        if requested_total > MAX_PDF_PARSE_TOTAL_PIXELS:
+            return (
+                f"The uploaded PDF would rasterize to more than "
+                f"{requested_total} pixels in total at zoomin {zoomin}, "
+                f"exceeding the whole-document limit of "
+                f"{MAX_PDF_PARSE_TOTAL_PIXELS}"
+            )
+        # Summed from the per-page peaks, so the render a retry replaces is
+        # counted alongside its replacement rather than the final scale alone.
+        retry_total += peak_pixels
+        if retry_total > MAX_PDF_PARSE_RETRY_TOTAL_PIXELS:
+            return (
+                f"The uploaded PDF would peak above {retry_total} pixels if "
+                f"the parser re-rendered it at {worst_case}x, exceeding the "
+                f"retry limit of {MAX_PDF_PARSE_RETRY_TOTAL_PIXELS}"
+            )
+    return None
+
+
+def largest_fitting_parse_zoom(
+    sizes: List[Tuple[float, float]], upper_bound: int
+) -> Optional[int]:
+    """The largest zoom in ``1..upper_bound`` this document fits at.
+
+    Not simply ``upper_bound`` counted down until something fits: the per-page
+    ceiling is enforced at the worst-case scale, and that ladder overshoots
+    for zoom values that are not power-of-three divisors of 9, so a *lower*
+    zoomin can have a *higher* worst case (2 and 6 both reach 18x, while 3
+    stops at 9x). A 1000x1000 pt page fits at 1, 3 and 4 but not at 2. Every
+    candidate is therefore tested rather than assuming the budget shrinks as
+    the zoom does.
+
+    The search covers the whole range, not just values below the request, for
+    the same reason: a page rejected at zoomin 2 (18x) may well fit at 3 (9x),
+    and sending the caller down to 1 would cost quality for nothing.
+    """
+    for candidate in range(upper_bound, 0, -1):
+        if _parse_budget_error(sizes, candidate) is None:
+            return candidate
+    return None
 
 
 def is_pdf_upload(content_type: Optional[str], head: bytes) -> bool:
@@ -121,31 +258,46 @@ def is_pdf_upload(content_type: Optional[str], head: bytes) -> bool:
     return head.startswith(PDF_MAGIC)
 
 
-def validate_pdf_for_parse(data: bytes, zoomin: int) -> int:
+def validate_pdf_for_parse(
+    data: bytes, zoomin: int, max_zoomin: Optional[int] = None
+) -> int:
     """Check a PDF that will be handed to a whole-document parser.
 
     Whole-document parsing tasks (e.g. DeepDoc's ``task="parse"``) render the
     PDF themselves, so the bytes are passed straight through instead of being
     rasterized here. That skips everything ``rasterize_pdf`` would have
     validated, and parsers tend to fail unhelpfully: DeepDoc swallows load
-    errors and returns an empty result, then re-renders at three times the
-    zoom when it finds no boxes.
-
-    Page geometry therefore has to be checked here too, in two ways. A single
-    valid page with an outsized MediaBox — 14400x14400 points is legal —
-    rasterizes to billions of pixels on its own. And because every page is
-    rendered up front and held together, a document whose pages are each
-    comfortably under the per-page limit can still exhaust the worker in
-    aggregate, so the pages are budgeted as a whole too.
-
-    Both budgets are applied at the worst-case scale, since the retry is left
-    intact and re-renders the whole document, and both count the render being
-    replaced alongside its replacement (see
-    ``worst_case_parse_peak_pixels``).
+    errors and returns an empty result.
+
+    Page geometry therefore has to be checked here too. A single valid page
+    with an outsized MediaBox — 14400x14400 points is legal — rasterizes to
+    billions of pixels on its own. And because every page is rendered up
+    front and held together, a document whose pages are each comfortably
+    under the per-page limit can still exhaust the worker in aggregate, so
+    the pages are budgeted as a whole too.
+
+    The per-page ceiling is applied at the worst-case scale as cheap
+    insurance -- it does not depend on the retry being unreachable, so a
+    single oversized MediaBox stays rejected even if that changes -- and it
+    counts the render being replaced alongside its replacement (see
+    ``worst_case_parse_peak_pixels``). The document-wide ceiling is applied
+    at the requested scale, with a second, looser one bounding what the
+    document would peak at if the retry did fire; see
+    ``MAX_PDF_PARSE_TOTAL_PIXELS`` and ``MAX_PDF_PARSE_RETRY_TOTAL_PIXELS``
+    for why the two differ.
 
     Returns the page count. Raises ``ValueError`` if the document cannot be
     opened, has no pages, has too many pages, or would rasterize to too many
-    pixels on any single page or across the document.
+    pixels on any single page or across the document. The message names a
+    ``zoomin`` that would fit whenever one exists, rather than advising the
+    caller to lower it -- the retry ladder is not monotonic, so lowering it
+    can make the budget larger.
+
+    ``max_zoomin`` bounds the search for that recommendation; it should be the
+    largest value the endpoint accepts. It defaults to ``zoomin``, which only
+    searches downwards -- pass the real ceiling so a request rejected at a
+    zoom with an overshooting ladder (2 or 6, which reach 18x) can be pointed
+    at a higher one with a shorter ladder (3, which stops at 9x).
     """
     try:
         import pypdfium2 as pdfium
@@ -170,36 +322,29 @@ def validate_pdf_for_parse(data: bytes, zoomin: int) -> int:
                     f"The uploaded PDF has {page_count} pages, at most "
                     f"{MAX_PDF_OCR_PAGES} pages can be parsed per request"
                 )
-            worst_case = worst_case_parse_zoom(zoomin)
-            total_pixels = 0
+            sizes = []
             for page_number in range(1, page_count + 1):
                 page = pdf[page_number - 1]
                 try:
-                    width, height = page.get_size()
+                    sizes.append(page.get_size())
                 finally:
                     page.close()
-                peak_pixels = worst_case_parse_peak_pixels(width, height, zoomin)
-                if peak_pixels > MAX_PDF_PARSE_PAGE_PIXELS:
-                    raise ValueError(
-                        f"Page {page_number} would rasterize to {peak_pixels} "
-                        f"pixels at zoomin {zoomin:g} (the parser re-renders at "
-                        f"up to {worst_case:g}x when a page yields no text), "
-                        f"exceeding the per-page limit of "
-                        f"{MAX_PDF_PARSE_PAGE_PIXELS}; lower `zoomin`"
-                    )
-                total_pixels += peak_pixels
-                if total_pixels > MAX_PDF_PARSE_TOTAL_PIXELS:
-                    raise ValueError(
-                        f"The uploaded PDF would rasterize to more than "
-                        f"{total_pixels} pixels in total at zoomin {zoomin:g} "
-                        f"(the parser re-renders at up to {worst_case:g}x when a "
-                        f"page yields no text), exceeding the whole-document "
-                        f"limit of {MAX_PDF_PARSE_TOTAL_PIXELS}; lower `zoomin` "
-                        f"or split the document"
-                    )
         finally:
             pdf.close()
 
+    reason = _parse_budget_error(sizes, zoomin)
+    if reason is not None:
+        # "Lower `zoomin`" was the old advice and it is not sound: the retry
+        # ladder is not monotonic, so a lower zoomin can be budgeted *higher*
+        # (2 and 6 both escalate to 18x where 3 stops at 9x). Name a zoom that
+        # actually fits, and only fall back to splitting when none does.
+        # Search the whole permitted range, not just below the request: with a
+        # non-monotonic ladder a higher zoom can be the one that fits.
+        fitting = largest_fitting_parse_zoom(sizes, max(max_zoomin or zoomin, zoomin))
+        if fitting is None:
+            raise ValueError(f"{reason}; split the document into smaller parts")
+        raise ValueError(f"{reason}; retry with `zoomin` {fitting}")
+
     return page_count