diff --git a/README.md b/README.md index b5a0b035..62808be9 100644 --- a/README.md +++ b/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. diff --git a/src/extract_bench/cli.py b/src/extract_bench/cli.py index 9085cf90..b490e42a 100644 --- a/src/extract_bench/cli.py +++ b/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, diff --git a/src/extract_bench/inference/cli.py b/src/extract_bench/inference/cli.py index 04a46320..8cf4fe04 100644 --- a/src/extract_bench/inference/cli.py +++ b/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)) diff --git a/src/extract_bench/inference/page_budget.py b/src/extract_bench/inference/page_budget.py new file mode 100644 index 00000000..47317c0d --- /dev/null +++ b/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() diff --git a/src/extract_bench/inference/runner.py b/src/extract_bench/inference/runner.py index 05907c93..eaca0c14 100644 --- a/src/extract_bench/inference/runner.py +++ b/src/extract_bench/inference/runner.py @@ -30,6 +30,7 @@ ) from rich.table import Table +from extract_bench.inference.page_budget import PageBudget from extract_bench.inference.providers.base import ( Provider, ProviderError, @@ -176,6 +177,7 @@ def __init__( tags: list[str] | None = None, per_file_timeout: float | None = None, timeout_retries: int = DEFAULT_TIMEOUT_RETRIES, + max_concurrent_pages: int | None = None, ): """ Initialize the inference runner. @@ -184,6 +186,7 @@ def __init__( :param pipeline: Pipeline specification :param output_dir: Directory to save results :param max_concurrent: Maximum concurrent inference requests + :param max_concurrent_pages: Optional source PDF page budget; oversized PDFs run alone :param save_raw: Whether to save RawInferenceResult JSON files :param save_normalized: Whether to save InferenceResult JSON files :param force: Force regeneration even if results already exist @@ -197,6 +200,7 @@ def __init__( self.pipeline = pipeline self.output_dir = Path(output_dir) self.max_concurrent = max_concurrent + self.max_concurrent_pages = max_concurrent_pages self.save_raw = save_raw self.save_normalized = save_normalized self.force = force @@ -292,6 +296,23 @@ def shutdown(self) -> None: """ self._thread_pool.shutdown(wait=False, cancel_futures=True) + def _submit_attempt(self, function: Callable[..., Any], *args: Any) -> concurrent.futures.Future[Any]: + """Isolate page-limited attempts so a timed-out thread cannot starve later work. + + Admission already bounds active local attempts. Python cannot kill a + running thread, so a timed-out attempt may outlive its reservation; + reusing its worker would make later files time out before starting. + This matches the per-attempt isolation used by the synchronous runner. + """ + if self.max_concurrent_pages is None: + return self._thread_pool.submit(function, *args) + executor = concurrent.futures.ThreadPoolExecutor(max_workers=1, thread_name_prefix="inference_attempt") + try: + return executor.submit(function, *args) + finally: + # Close to new submissions; the already submitted task still runs. + executor.shutdown(wait=False) + def get_current_summary(self) -> RunSummary | None: """Get the current run summary (useful for interrupt handling).""" return self._current_summary @@ -1130,7 +1151,7 @@ def _run_test_cases_sync( async def _process_test_case_with_semaphore( self, - semaphore: asyncio.Semaphore, + semaphore: asyncio.Semaphore | PageBudget, test_case: TestCase, product_type: ProductType, summary: RunSummary, @@ -1147,20 +1168,21 @@ async def _process_test_case_with_semaphore( progress.update(task_id, advance=1, refresh=True) return - async with semaphore: + slot = semaphore.slot(test_case.test_id) if isinstance(semaphore, PageBudget) else semaphore + async with slot: # Set status to "running" when semaphore is acquired (before processing starts) # This allows UI to show "in progress" status immediately if self.use_rich and test_case.test_id in self.job_statuses: self.job_statuses[test_case.test_id].status = "running" self.job_statuses[test_case.test_id].started_at = datetime.now() - # Process test case using our custom thread pool with per-file timeout. + # Process each attempt with its own deadline after admission. raw_result = None normalized_result = None error_info: str | tuple[str, str, str] | None = None for timeout_attempt in range(self.timeout_retries + 1): - future = self._thread_pool.submit(self._process_test_case, test_case, product_type) + future = self._submit_attempt(self._process_test_case, test_case, product_type) try: outcome = await asyncio.wait_for( asyncio.wrap_future(future), @@ -1230,7 +1252,7 @@ async def _process_test_case_with_semaphore( async def _process_with_semaphore( self, - semaphore: asyncio.Semaphore, + semaphore: asyncio.Semaphore | PageBudget, pdf_path: Path, example_id: str, product_type: ProductType, @@ -1248,20 +1270,21 @@ async def _process_with_semaphore( progress.update(task_id, advance=1, refresh=True) return - async with semaphore: + slot = semaphore.slot(example_id) if isinstance(semaphore, PageBudget) else semaphore + async with slot: # Set status to "running" when semaphore is acquired (before processing starts) # This allows UI to show "in progress" status immediately if self.use_rich and example_id in self.job_statuses: self.job_statuses[example_id].status = "running" self.job_statuses[example_id].started_at = datetime.now() - # Process document using our custom thread pool with per-file timeout. + # Process each attempt with its own deadline after admission. raw_result = None normalized_result = None error_info: str | tuple[str, str, str] | None = None for timeout_attempt in range(self.timeout_retries + 1): - future = self._thread_pool.submit(self._process_document, pdf_path, example_id, product_type) + future = self._submit_attempt(self._process_document, pdf_path, example_id, product_type) try: outcome = await asyncio.wait_for( asyncio.wrap_future(future), @@ -1467,7 +1490,7 @@ def default_example_id_fn(pdf_path: Path) -> str: example_id_fn = default_example_id_fn # When max_concurrent is 1, run synchronously without asyncio/threads - if self.max_concurrent == 1: + if self.max_concurrent == 1 and self.max_concurrent_pages is None: return self._run_files_sync(pdf_files, product_type, example_id_fn) self._current_summary = summary = RunSummary() @@ -1479,7 +1502,16 @@ def default_example_id_fn(pdf_path: Path) -> str: self.job_statuses[example_id] = JobStatus(example_id=example_id, pdf_path=pdf_path, status="pending") # Create semaphore for concurrency control - semaphore = asyncio.Semaphore(self.max_concurrent) + semaphore = ( + PageBudget( + [(example_id_fn(path), path) for path in pdf_files], + self.max_concurrent, + self.max_concurrent_pages, + self.output_dir, + ) + if self.max_concurrent_pages is not None + else asyncio.Semaphore(self.max_concurrent) + ) # Create progress bar with enhanced styling if self.use_rich: @@ -1719,17 +1751,16 @@ async def run_test_cases( test_cases = unique # When max_concurrent is 1, run synchronously without asyncio/threads - if self.max_concurrent == 1: + if self.max_concurrent == 1 and self.max_concurrent_pages is None: return self._run_test_cases_sync(test_cases, product_type, test_cases_dir) self._current_summary = summary = RunSummary() - # Log concurrency setting for debugging - # Thread pool is sized to match max_concurrent to avoid default pool bottleneck + # Log the client admission limits for debugging. print( "Starting async run_test_cases with " f"max_concurrent={self.max_concurrent} " - f"(thread pool size: {self._thread_pool._max_workers})" + f"max_concurrent_pages={self.max_concurrent_pages}" ) # Initialize job statuses for Rich UI @@ -1742,7 +1773,16 @@ async def run_test_cases( ) # Create semaphore for concurrency control - semaphore = asyncio.Semaphore(self.max_concurrent) + semaphore = ( + PageBudget( + [(case.test_id, case.file_path) for case in test_cases], + self.max_concurrent, + self.max_concurrent_pages, + self.output_dir, + ) + if self.max_concurrent_pages is not None + else asyncio.Semaphore(self.max_concurrent) + ) # Create progress bar with enhanced styling if self.use_rich: diff --git a/src/extract_bench/pipeline/cli.py b/src/extract_bench/pipeline/cli.py index ba82946f..1a4a1647 100644 --- a/src/extract_bench/pipeline/cli.py +++ b/src/extract_bench/pipeline/cli.py @@ -75,6 +75,7 @@ def run( skip_inference: bool = False, test: bool = False, shuffle_seed: int | None = None, + max_concurrent_pages: int | None = None, ) -> int: """ Run the full benchmark pipeline: inference -> evaluation -> report -> open browser. @@ -88,6 +89,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 already exist (default: False) verbose: Enable verbose output (default: False) group: Optional group name to filter test cases (e.g., 'chart') @@ -112,6 +114,7 @@ def run( pipeline=pipeline, file_path=Path(file), shuffle_seed=shuffle_seed, + max_concurrent_pages=max_concurrent_pages, output_dir=Path(output_dir) if output_dir else Path("./output"), force=force, verbose=verbose, @@ -158,6 +161,7 @@ def run( output_dir=output_base, max_concurrent=max_concurrent, shuffle_seed=shuffle_seed, + max_concurrent_pages=max_concurrent_pages, force=force, verbose=verbose, group=group, @@ -431,6 +435,7 @@ def _run_single_file( open_report: bool, skip_inference: bool, shuffle_seed: int | None = None, + max_concurrent_pages: int | None = None, ) -> int: """Run pipeline on a single file by creating a temporary directory structure.""" import shutil @@ -487,6 +492,7 @@ def _run_single_file( output_dir=output_dir, max_concurrent=1, # Single file, no need for concurrency shuffle_seed=shuffle_seed, + max_concurrent_pages=max_concurrent_pages, force=force, verbose=verbose, tags=tags, diff --git a/tests/extract_bench/inference/test_page_budget.py b/tests/extract_bench/inference/test_page_budget.py new file mode 100644 index 00000000..98818672 --- /dev/null +++ b/tests/extract_bench/inference/test_page_budget.py @@ -0,0 +1,458 @@ +import asyncio +import json +import threading +import time +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from extract_bench.inference.cli import InferenceCLI +from extract_bench.inference.page_budget import PageBudget +from extract_bench.inference.runner import InferenceRunner +from extract_bench.schemas.pipeline import PipelineSpec +from extract_bench.schemas.product import ProductType + +pymupdf = pytest.importorskip("pymupdf") + + +class RecordingRunner(InferenceRunner): + observed = [] + + def _process_test_case(self, test_case, product_type): + self.observed.append(test_case.test_id) + return None, None, None + + +def run_mock_cli(tmp_path, name, seed, cases, concurrent=1, **kwargs): + spec = PipelineSpec(pipeline_name=name, provider_name="fake", product_type=ProductType.EXTRACT, config={}) + RecordingRunner.observed = [] + with ( + patch("extract_bench.inference.cli.get_pipeline", return_value=spec), + patch("extract_bench.inference.cli.load_test_cases", return_value=cases), + patch("extract_bench.inference.cli._detect_product_type", return_value=ProductType.EXTRACT), + patch("extract_bench.inference.cli.create_provider", return_value=object()), + patch("extract_bench.inference.cli.InferenceRunner", RecordingRunner), + ): + result = InferenceCLI().run( + name, + input_dir=tmp_path, + output_dir=tmp_path, + shuffle_seed=seed, + max_concurrent=concurrent, + no_rich=True, + force=True, + force_exit_on_completion=False, + **kwargs, + ) + assert result == 0 + return list(RecordingRunner.observed), json.loads((tmp_path / name / "_submission_order.json").read_text()) + + +def pdf_cases(tmp_path, counts): + cases = [] + for i, pages in enumerate(counts): + path = tmp_path / f"doc-{i}.pdf" + with pymupdf.open() as pdf: + for _ in range(pages): + pdf.new_page() + pdf.save(path) + cases.append(SimpleNamespace(test_id=f"doc-{i}", file_path=path)) + return cases + + +def events(path): + return [json.loads(line) for line in (path / "_admission_events.jsonl").read_text().splitlines()] + + +@pytest.mark.parametrize( + "counts,max_docs,max_pages", + [ + ([60, 60, 10], 8, 100), + ([123, 10, 20], 8, 100), + ([60, 101, 10], 8, 100), + ([10, 10, 10], 2, 100), + ([50, 50, 1], 8, 100), + ], +) +def test_fifo_limits_and_oversize(tmp_path, counts, max_docs, max_pages): + cases = pdf_cases(tmp_path, counts) + + async def scenario(): + budget = PageBudget([(c.test_id, c.file_path) for c in cases], max_docs, max_pages, tmp_path) + started = [asyncio.Event() for _ in cases] + release = [asyncio.Event() for _ in cases] + order = [] + + async def worker(i): + async with budget.slot(cases[i].test_id): + order.append(i) + started[i].set() + await release[i].wait() + + tasks = [asyncio.create_task(worker(i)) for i in range(len(cases))] + for i in range(len(cases)): + await asyncio.wait_for(started[i].wait(), 1) + assert order == sorted(order) + release[i].set() + await asyncio.gather(*tasks) + assert budget.active_documents == budget.active_pages == 0 + + asyncio.run(scenario()) + log = events(tmp_path) + assert [e["example_id"] for e in log if e["event"] == "admitted"] == [c.test_id for c in cases] + for event in log: + assert event["active_documents"] <= max_docs + assert event["active_pages"] <= max_pages or (event["active_documents"] == 1 and event["pages"] > max_pages) + + +def test_fifo_does_not_backfill_and_holds_full_lifetime(tmp_path): + cases = pdf_cases(tmp_path, [60, 60, 10]) + + async def scenario(): + budget = PageBudget([(c.test_id, c.file_path) for c in cases], 8, 100, tmp_path) + releases = [asyncio.Event() for _ in cases] + starts = [asyncio.Event() for _ in cases] + + async def worker(i): + async with budget.slot(cases[i].test_id): + starts[i].set() + await releases[i].wait() + + tasks = [asyncio.create_task(worker(i)) for i in range(3)] + await starts[0].wait() + await asyncio.sleep(0) + assert not starts[1].is_set() and not starts[2].is_set() + assert (budget.active_documents, budget.active_pages) == (1, 60) + releases[0].set() + await starts[1].wait() + await starts[2].wait() + assert (budget.active_documents, budget.active_pages) == (2, 70) + releases[1].set() + releases[2].set() + await asyncio.gather(*tasks) + + asyncio.run(scenario()) + + +@pytest.mark.parametrize("rich", [False, True]) +def test_real_runner_limits_and_manifest_order(tmp_path, rich): + cases = pdf_cases(tmp_path, [60, 60, 10, 123, 5, 1, 50, 50, 1]) + spec = PipelineSpec(pipeline_name="mock", provider_name="fake", product_type=ProductType.EXTRACT, config={}) + + class SlowRunner(RecordingRunner): + def _process_test_case(self, case, product_type): + time.sleep(0.005) + return super()._process_test_case(case, product_type) + + runner = SlowRunner( + provider=object(), + pipeline=spec, + output_dir=tmp_path, + max_concurrent=3, + max_concurrent_pages=100, + use_rich=rich, + force=True, + timeout_retries=0, + ) + try: + summary = asyncio.run(runner.run_test_cases(cases, ProductType.EXTRACT)) + assert summary.successful == len(cases) + finally: + runner.shutdown() + log = events(tmp_path) + assert [e["example_id"] for e in log if e["event"] == "admitted"] == [c.test_id for c in cases] + assert all(e["active_documents"] <= 3 for e in log) + assert all(e["active_pages"] <= 100 or e["active_documents"] == 1 for e in log) + assert log[-1]["active_pages"] == log[-1]["active_documents"] == 0 + + +def test_processing_error_is_recorded_and_later_documents_complete(tmp_path): + cases = pdf_cases(tmp_path, [1, 1, 1]) + spec = PipelineSpec(pipeline_name="mock", provider_name="fake", product_type=ProductType.EXTRACT, config={}) + + class FailingRunner(RecordingRunner): + def _process_test_case(self, case, product_type): + if case.test_id == "doc-0": + return None, None, ("Unknown remote outcome", "", "ProviderPermanentError") + return None, None, None + + runner = FailingRunner( + provider=object(), + pipeline=spec, + output_dir=tmp_path, + max_concurrent=1, + max_concurrent_pages=100, + use_rich=False, + force=True, + timeout_retries=0, + ) + try: + summary = asyncio.run(runner.run_test_cases(cases, ProductType.EXTRACT)) + finally: + runner.shutdown() + assert summary.successful == 2 and summary.failed == 1 + assert summary.errors[0]["example_id"] == "doc-0" + log = events(tmp_path) + assert [e["example_id"] for e in log if e["event"] == "admitted"] == ["doc-0", "doc-1", "doc-2"] + assert log[-1]["active_pages"] == 0 + + +@pytest.mark.parametrize("pages", [0, -1, True]) +def test_invalid_budget_rejected(tmp_path, pages): + with pytest.raises(ValueError): + PageBudget([], 8, pages, tmp_path) + + +def test_unreadable_pdf_fails_before_any_submission(tmp_path): + cases = pdf_cases(tmp_path, [1]) + bad = tmp_path / "bad.pdf" + bad.write_text("not a PDF") + with pytest.raises(pymupdf.FileDataError): + PageBudget([(cases[0].test_id, cases[0].file_path), ("bad", bad)], 8, 100, tmp_path) + assert not (tmp_path / "_admission_events.jsonl").exists() + + +def test_cli_forwards_budget_and_single_worker_uses_gate(tmp_path): + cases = pdf_cases(tmp_path, [123, 20, 10]) + _, manifest = run_mock_cli(tmp_path, "bounded", 0, cases, concurrent=1, max_concurrent_pages=100) + log = events(tmp_path / "bounded") + assert [e["example_id"] for e in log if e["event"] == "admitted"] == manifest["example_ids"] + assert all(e["active_documents"] <= 1 for e in log) + + +@pytest.mark.parametrize("max_documents", [1, 2]) +@pytest.mark.parametrize("timeout_retries", [0, 1]) +@pytest.mark.parametrize("mode", ["test_cases", "files"]) +def test_healthy_documents_execute_after_stuck_attempt(tmp_path, max_documents, timeout_retries, mode): + cases = pdf_cases(tmp_path, [1, 1, 1]) + spec = PipelineSpec(pipeline_name="mock", provider_name="fake", product_type=ProductType.EXTRACT, config={}) + release = threading.Event() + stopped = threading.Event() + started = [] + + class TimeoutRunner(RecordingRunner): + def process(self, example_id): + started.append(example_id) + if example_id == "doc-0" and started.count(example_id) == 1: + try: + assert release.wait(5), "test did not release the stuck attempt" + finally: + stopped.set() + return None, None, None + + def _process_test_case(self, case, product_type): + return self.process(case.test_id) + + def _process_document(self, path, example_id, product_type): + return self.process(example_id) + + async def _cancel_inflight_and_drain_async(self, example_id, future, **kwargs): + await super()._cancel_inflight_and_drain_async(example_id, future, drain_timeout_seconds=0.001) + + runner = TimeoutRunner( + provider=object(), + pipeline=spec, + output_dir=tmp_path, + max_concurrent=max_documents, + max_concurrent_pages=100, + use_rich=False, + force=True, + per_file_timeout=0.1, + timeout_retries=timeout_retries, + ) + try: + if mode == "test_cases": + summary = asyncio.run(runner.run_test_cases(cases, ProductType.EXTRACT)) + else: + summary = asyncio.run(runner.run_files([case.file_path for case in cases], ProductType.EXTRACT)) + assert not stopped.is_set() + assert set(started) == {"doc-0", "doc-1", "doc-2"} + assert started.count("doc-0") == 1 + timeout_retries + assert summary.successful == 2 + timeout_retries + assert summary.failed == 1 - timeout_retries + assert len(summary.errors) == 1 - timeout_retries + assert [e["example_id"] for e in events(tmp_path) if e["event"] == "admitted"] == ["doc-0", "doc-1", "doc-2"] + finally: + release.set() + assert stopped.wait(1) + runner.shutdown() + + +def test_cached_documents_do_not_take_capacity(tmp_path): + cases = pdf_cases(tmp_path, [60, 123, 60]) + spec = PipelineSpec(pipeline_name="mock", provider_name="fake", product_type=ProductType.EXTRACT, config={}) + + class CachedRunner(RecordingRunner): + def _is_already_processed(self, example_id): + return example_id == "doc-1" + + runner = CachedRunner( + provider=object(), + pipeline=spec, + output_dir=tmp_path, + max_concurrent=8, + max_concurrent_pages=100, + use_rich=False, + timeout_retries=0, + ) + try: + summary = asyncio.run(runner.run_test_cases(cases, ProductType.EXTRACT)) + finally: + runner.shutdown() + assert summary.successful == 2 and summary.skipped == 1 + assert [e["example_id"] for e in events(tmp_path) if e["event"] == "admitted"] == ["doc-0", "doc-2"] + + +def test_cancelled_waiter_does_not_deadlock_fifo(tmp_path): + cases = pdf_cases(tmp_path, [60, 60, 10]) + + async def scenario(): + budget = PageBudget([(c.test_id, c.file_path) for c in cases], 8, 100, tmp_path) + async with budget.slot("doc-0"): + + async def blocked(): + async with budget.slot("doc-1"): + raise AssertionError("must wait") + + task = asyncio.create_task(blocked()) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + async with budget.slot("doc-2"): + assert budget.active_pages == 70 + assert budget.active_pages == 0 + + asyncio.run(scenario()) + + +def test_page_budget_forwarded_by_benchmark_commands(tmp_path): + from extract_bench.cli import BenchCLI + from extract_bench.pipeline.cli import PipelineCLI + + with patch.object(PipelineCLI, "run", return_value=0) as run: + assert BenchCLI().run("fake", input_dir=tmp_path, max_concurrent_pages=100) == 0 + assert run.call_args.kwargs["max_concurrent_pages"] == 100 + with patch.object(InferenceCLI, "run", return_value=1) as run: + assert PipelineCLI().run("fake", input_dir=tmp_path, max_concurrent_pages=100) == 1 + assert run.call_args.kwargs["max_concurrent_pages"] == 100 + + +def test_file_runner_obeys_page_budget(tmp_path): + cases = pdf_cases(tmp_path, [60, 60, 123, 10]) + spec = PipelineSpec(pipeline_name="mock", provider_name="fake", product_type=ProductType.EXTRACT, config={}) + + class FileRunner(InferenceRunner): + def _process_document(self, path, example_id, product_type): + time.sleep(0.005) + return None, None, None + + runner = FileRunner( + provider=object(), + pipeline=spec, + output_dir=tmp_path, + max_concurrent=3, + max_concurrent_pages=100, + use_rich=False, + force=True, + timeout_retries=0, + ) + try: + summary = asyncio.run(runner.run_files([case.file_path for case in cases], ProductType.EXTRACT)) + finally: + runner.shutdown() + assert summary.successful == 4 + log = events(tmp_path) + assert [e["example_id"] for e in log if e["event"] == "admitted"] == [case.test_id for case in cases] + assert all(e["active_pages"] <= 100 or e["active_documents"] == 1 for e in log) + + +@pytest.mark.parametrize("failed_event", ["admitted", "released"]) +def test_log_failure_releases_capacity_and_wakes_waiter(tmp_path, failed_event): + cases = pdf_cases(tmp_path, [101, 10]) + + async def scenario(): + budget = PageBudget([(c.test_id, c.file_path) for c in cases], 1, 100, tmp_path) + original_record = budget._record + first_started = asyncio.Event() + release = asyncio.Event() + + def record(event, example_id=None): + if event == failed_event and example_id == "doc-0": + raise OSError("disk full") + original_record(event, example_id) + + async def first(): + async with budget.slot("doc-0"): + first_started.set() + await release.wait() + + async def second(): + async with budget.slot("doc-1"): + assert (budget.active_documents, budget.active_pages) == (1, 10) + + with patch.object(budget, "_record", side_effect=record): + if failed_event == "admitted": + with pytest.raises(OSError, match="disk full"): + await first() + await asyncio.wait_for(second(), 1) + else: + task = asyncio.create_task(first()) + await asyncio.wait_for(first_started.wait(), 1) + waiter = asyncio.create_task(second()) + await asyncio.sleep(0) + assert not waiter.done() + release.set() + with pytest.raises(OSError, match="disk full"): + await task + await asyncio.wait_for(waiter, 1) + assert budget.active_documents == budget.active_pages == 0 + assert len(budget.waiting) == 0 + + asyncio.run(scenario()) + + +def test_resume_archives_previous_admission_sessions(tmp_path): + cases = pdf_cases(tmp_path, [1, 1]) + spec = PipelineSpec(pipeline_name="mock", provider_name="fake", product_type=ProductType.EXTRACT, config={}) + + class ResumingRunner(RecordingRunner): + cached = set() + + def _is_already_processed(self, example_id): + return example_id in self.cached + + runner = ResumingRunner( + provider=object(), + pipeline=spec, + output_dir=tmp_path, + max_concurrent=1, + max_concurrent_pages=100, + use_rich=False, + ) + try: + asyncio.run(runner.run_test_cases(cases, ProductType.EXTRACT)) + first_events = (tmp_path / "_admission_events.jsonl").read_bytes() + first_plan = (tmp_path / "_admission_plan.json").read_bytes() + runner.cached = {"doc-0"} + resumed = asyncio.run(runner.run_test_cases(cases, ProductType.EXTRACT)) + assert resumed.successful == resumed.skipped == 1 + assert [e["example_id"] for e in events(tmp_path) if e["event"] == "admitted"] == ["doc-1"] + archives = list((tmp_path / "_admission_history").iterdir()) + assert len(archives) == 1 + assert (archives[0] / "_admission_events.jsonl").read_bytes() == first_events + assert (archives[0] / "_admission_plan.json").read_bytes() == first_plan + # An interrupted append is evidence too; preserve it byte for byte. + with (tmp_path / "_admission_events.jsonl").open("ab") as log: + log.write(b'{"event":') + interrupted = (tmp_path / "_admission_events.jsonl").read_bytes() + runner.cached = {"doc-0", "doc-1"} + complete = asyncio.run(runner.run_test_cases(cases, ProductType.EXTRACT)) + assert complete.skipped == 2 + assert events(tmp_path) == [] + archives = list((tmp_path / "_admission_history").iterdir()) + assert len(archives) == 2 + assert {(d / "_admission_events.jsonl").read_bytes() for d in archives} == {first_events, interrupted} + finally: + runner.shutdown()