Bound inference admission by documents and source pages

README.md

@@ -413,3 +413,23 @@ all pipelines the same planned order, independent of filesystem discovery order.
 Omit the flag to preserve discovery order. Each inference output directory saves
 `_submission_order.json` with the seed and ordered IDs; cached results are still
 skipped normally. Parallel requests can complete in a different order.
+
+### Limit concurrent source pages
+
+With the `local` (or `runners`) extra installed, add `--max_concurrent_pages 100`
+alongside `--max_concurrent 8` to bound both documents and source PDF pages.
+The default is no page limit. This option requires readable, unencrypted PDFs;
+all source page counts are checked before any submissions.
+
+Admission follows input order (or the seeded shuffle). It waits for the next
+PDF to fit rather than backfilling smaller PDFs. A PDF larger than the page
+budget runs alone. Capacity stays reserved through provider processing,
+polling, and retries; cached documents consume no capacity. Document errors
+are recorded and subsequent documents continue. These limits bound local
+attempts: timed-out remote jobs may still be running after capacity is released.
+
+Each run saves `_admission_plan.json` (limits and source page counts) and
+`_admission_events.jsonl` (timestamped admissions, releases, and active counts).
+Resuming archives the previous plan and event log under `_admission_history/`
+before creating the current session's files, including when cached documents
+are skipped. These describe client admission, not backend worker capacity.

src/extract_bench/cli.py

@@ -131,6 +131,7 @@ def run(
         skip_inference: bool = False,
         test: bool = False,
         shuffle_seed: int | None = None,
+        max_concurrent_pages: int | None = None,
     ) -> int:
         """Run end-to-end benchmark: inference -> evaluation -> report.
 
@@ -141,6 +142,7 @@ def run(
             output_dir: Directory to save results (default: ./output)
             max_concurrent: Maximum concurrent inference requests (default: 20)
             shuffle_seed: Shuffle sorted test IDs reproducibly; None preserves discovery order.
+            max_concurrent_pages: Source PDF page budget; oversized PDFs run alone.
             force: Force regeneration even if results exist (default: False)
             verbose: Enable verbose output (default: False)
             group: Filter by category (e.g., 'short', 'medium', 'long')
@@ -162,6 +164,7 @@ def run(
             output_dir=output_dir,
             max_concurrent=max_concurrent,
             shuffle_seed=shuffle_seed,
+            max_concurrent_pages=max_concurrent_pages,
             force=force,
             verbose=verbose,
             group=group,

src/extract_bench/inference/cli.py

@@ -159,6 +159,7 @@ def run(
         timeout_retries: int = 2,
         force_exit_on_completion: bool = True,
         shuffle_seed: int | None = None,
+        max_concurrent_pages: int | None = None,
     ) -> int:
         """
         Run inference on a directory, auto-detecting structure and requirements.
@@ -173,6 +174,7 @@ def run(
             output_dir: Directory to save inference results (default: './output')
             pipeline_name_override: Pipeline name override (default: uses pipeline name)
             max_concurrent: Maximum concurrent inference requests (default: 20)
+            max_concurrent_pages: Source PDF page budget; oversized documents run alone.
             save_raw: Save raw inference results (default: True)
             save_normalized: Save normalized inference results (default: True)
             force: Force regeneration even if results already exist (default: False)
@@ -199,6 +201,7 @@ def run(
             pipeline=pipeline,
             pipeline_name_override=pipeline_name_override,
             max_concurrent=max_concurrent,
+            max_concurrent_pages=max_concurrent_pages,
             save_raw=save_raw,
             save_normalized=save_normalized,
             force=force,
@@ -230,6 +233,7 @@ def _run_test_cases(
         timeout_retries: int = 2,
         force_exit_on_completion: bool = True,
         shuffle_seed: int | None = None,
+        max_concurrent_pages: int | None = None,
     ) -> int:
         """Internal method to run inference on test cases."""
         try:
@@ -372,6 +376,7 @@ def _run_test_cases(
                 pipeline=pipeline_spec,
                 output_dir=actual_output_dir,
                 max_concurrent=max_concurrent,
+                max_concurrent_pages=max_concurrent_pages,
                 save_raw=save_raw,
                 save_normalized=save_normalized,
                 force=force,
@@ -394,7 +399,7 @@ def _run_test_cases(
 
             # Run inference on test cases
             # When max_concurrent is 1, use sync method directly to avoid async overhead
-            if max_concurrent == 1:
+            if max_concurrent == 1 and max_concurrent_pages is None:
                 summary = runner._run_test_cases_sync(test_cases, product_type_enum, test_cases_dir)
             else:
                 summary = asyncio.run(runner.run_test_cases(test_cases, product_type_enum, test_cases_dir))

src/extract_bench/inference/page_budget.py

@@ -0,0 +1,119 @@
+"""FIFO admission bounded by both document count and source PDF pages."""
+
+import asyncio
+import json
+from collections import deque
+from collections.abc import AsyncIterator, Sequence
+from contextlib import asynccontextmanager
+from datetime import UTC, datetime
+from pathlib import Path
+from uuid import uuid4
+
+
+class PageBudget:
+    """Hold document/page capacity for a complete local attempt, including retries.
+
+    Oversized PDFs run alone; FIFO prevents later small files from overtaking
+    them. Releasing a timed-out attempt does not guarantee remote cancellation.
+    """
+
+    def __init__(
+        self,
+        documents: Sequence[tuple[str, Path]],
+        max_documents: int,
+        max_pages: int,
+        output_dir: Path,
+    ) -> None:
+        if type(max_documents) is not int or max_documents < 1:
+            raise ValueError("max_concurrent must be a positive integer")
+        if type(max_pages) is not int or max_pages < 1:
+            raise ValueError("max_concurrent_pages must be a positive integer")
+        import pymupdf
+
+        self.max_documents = max_documents
+        self.max_pages = max_pages
+        self.pages: dict[str, int] = {}
+        # Count every source before submitting anything. Never estimate an unreadable PDF.
+        for example_id, path in documents:
+            if example_id in self.pages:
+                raise ValueError(f"Duplicate document ID: {example_id}")
+            if path.suffix.lower() != ".pdf":
+                raise ValueError(f"Page admission requires PDF input: {path}")
+            with pymupdf.open(path) as pdf:
+                if pdf.needs_pass or len(pdf) < 1:
+                    raise ValueError(f"Cannot count readable pages: {path}")
+                self.pages[example_id] = len(pdf)
+        self.active_documents = 0
+        self.active_pages = 0
+        self.waiting: deque[str] = deque()
+        self.condition = asyncio.Condition()
+        self.events = output_dir / "_admission_events.jsonl"
+        # Preserve the previous session, including an interrupted JSONL tail,
+        # before creating a new plan/log for a resumed run with cached results.
+        previous = [path for path in (self.events, output_dir / "_admission_plan.json") if path.exists()]
+        if len(previous) > 0:
+            archive = output_dir / "_admission_history" / str(uuid4())
+            archive.mkdir(parents=True)
+            for path in previous:
+                path.rename(archive / path.name)
+        self.events.write_text("")
+        (output_dir / "_admission_plan.json").write_text(
+            json.dumps(
+                {
+                    "created_at": datetime.now(UTC).isoformat(),
+                    "max_concurrent": max_documents,
+                    "max_concurrent_pages": max_pages,
+                    "oversized_document_policy": "run_alone",
+                    "documents": [{"example_id": key, "pages": value} for key, value in self.pages.items()],
+                },
+                indent=2,
+            )
+            + "\n"
+        )
+
+    def _record(self, event: str, example_id: str | None = None) -> None:
+        record = {
+            "event": event,
+            "timestamp": datetime.now(UTC).isoformat(),
+            "example_id": example_id,
+            "pages": self.pages[example_id] if example_id is not None else None,
+            "active_documents": self.active_documents,
+            "active_pages": self.active_pages,
+        }
+        with self.events.open("a") as stream:
+            stream.write(json.dumps(record) + "\n")
+
+    @asynccontextmanager
+    async def slot(self, example_id: str) -> AsyncIterator[None]:
+        pages = self.pages[example_id]
+        async with self.condition:
+            self.waiting.append(example_id)
+            try:
+                await self.condition.wait_for(
+                    lambda: (
+                        self.waiting[0] == example_id
+                        and self.active_documents < self.max_documents
+                        and (self.active_pages + pages <= self.max_pages or self.active_documents == 0)
+                    )
+                )
+                self.active_documents += 1
+                self.active_pages += pages
+                try:
+                    self._record("admitted", example_id)
+                except BaseException:
+                    self.active_documents -= 1
+                    self.active_pages -= pages
+                    raise
+            finally:
+                self.waiting.remove(example_id)
+                self.condition.notify_all()
+        try:
+            yield
+        finally:
+            async with self.condition:
+                self.active_documents -= 1
+                self.active_pages -= pages
+                try:
+                    self._record("released", example_id)
+                finally:
+                    self.condition.notify_all()