diff --git a/Readme.md b/Readme.md index fdf4a2e..e9ece6c 100644 --- a/Readme.md +++ b/Readme.md @@ -36,6 +36,7 @@ concurrent processing capabilities for PDF documents, reference strings, and pat - **Type Hints**: Ships inline type annotations and a `py.typed` marker (PEP 561) for static type checking - **Archive Streaming**: Process files directly from `.zip`/`.tar`/`.tar.gz` archives without fully decompressing them - **S3 Streaming**: Read PDFs and zips straight from `s3://` (range-streamed, no full download) with the optional `[s3]` extra +- **In-Memory Documents**: Send PDFs held as bytes straight to GROBID, without writing them to disk first ## 📋 Prerequisites @@ -188,16 +189,17 @@ grobid_client --input "~/data/**/*.pdf" --output ~/results processFulltextDocu > [!NOTE] > `--input` accepts a directory, a single file, an **archive**, or a **glob pattern**: -> - **Archives** (`.zip`, `.tar`, `.tar.gz`/`.tgz`, `.tar.bz2`/`.tbz2`) are streamed: eligible entries are extracted in -> chunks of `batch_size` to a temporary directory, sent to GROBID, written to `--output`, and deleted before the next -> chunk. The archive is never fully decompressed, so disk usage stays bounded. If `--output` is omitted, results go to a -> directory named after the archive (e.g. `papers.zip` → `papers/`). +> - **Archives** (`.zip`, `.tar`, `.tar.gz`/`.tgz`, `.tar.bz2`/`.tbz2`) are streamed: eligible entries are read into +> memory in chunks of `batch_size` and sent to GROBID straight from there — the archive is never fully decompressed and +> nothing but the results in `--output` ever touches the disk. If `--output` is omitted, results go to a directory named +> after the archive (e.g. `papers.zip` → `papers/`). > - **Glob patterns** (`paper.zip`, `paper*.zip`, `**/paper*.zip`, `**/*.pdf`, …) are expanded with `**` recursion; each > match is handled by type (archive → streamed, directory → recursed, file → processed). Quote the pattern so your shell > passes it through to the client unexpanded. > - **S3** (requires `pip install "grobid-client-python[s3]"`): pass an `s3://` object, prefix or glob. A remote zip is > **range-streamed** (only its central directory and the entries are fetched — never the whole object); loose remote -> PDFs are fetched a batch at a time. Credentials use the standard AWS chain (env vars / `~/.aws` / IAM role). +> PDFs are fetched a batch at a time, directly into memory without being written to a local file first. Credentials use +> the standard AWS chain (env vars / `~/.aws` / IAM role). > ```bash > grobid_client --input "s3://my-bucket/papers/2021.zip" --output ~/out processFulltextDocument # one remote zip > grobid_client --input "s3://my-bucket/pdfs/*.pdf" --output ~/out processFulltextDocument # loose PDFs @@ -289,6 +291,63 @@ client.process( ) ``` +#### Processing a PDF from memory + +A PDF that is already in memory - downloaded from an API, read from a database or an object store - can be sent +directly, without writing it to a temporary file first. `process_pdf` takes either a path or the document itself, as +`bytes` or as any binary stream, and returns the TEI as a string: + +```python +import io +import requests + +pdf = io.BytesIO(requests.get("https://example.org/paper.pdf").content) +pdf.name = "paper.pdf" # optional, see below + +name, status, tei = client.process_pdf( + service="processFulltextDocument", + pdf_file=pdf, + consolidate_header=True, + tei_coordinates=True +) + +if status == 200: + print(tei) +``` + +There is no flag to say where the document comes from: the object itself says it. A document also carries its own name, +taken from the `name` attribute that `open()` sets on files and that can be set on anything else, `io.BytesIO` included. +The name identifies the document in the request sent to GROBID, in the logs, and as the first element of the result, so +documents processed this way stay distinguishable. Bytes passed on their own have nothing to be named after and fall +back to `document.pdf`. + +Several documents can be sent concurrently with `process_documents`, which runs them through the same thread pool the +file-based processing uses: + +```python +results = client.process_documents( + service="processFulltextDocument", + documents=[pdf1, pdf2, "/path/to/paper3.pdf"], + n=10 # documents sent concurrently +) + +for name, status, tei in results: + ... +``` + +Documents that do not name themselves are named `document-1.pdf`, `document-2.pdf`, ... after their position. Results +come back **in the order the documents were given**, not in completion order, so they can be zipped back onto whatever +the caller has them keyed by. A document that fails does not stop the others: its own entry carries the error status. + +Before an in-memory run starts (this includes archive and `s3://` streaming), the client asks the server's `/api/health` +how many engines it actually has and logs a warning when the requested concurrency `n` exceeds them - the surplus +requests would only queue on the server or bounce as 503 - and an info message when engines would sit idle. The check is +advisory: a server without the endpoint (older GROBID) never blocks the run. + +> [!NOTE] +> Both return the TEI instead of writing it to disk, so the caller decides what to do with it. Use `process()` for the +> directory-oriented processing with resume and JSON/Markdown conversion. + ### Standalone Conversion Tools The library includes standalone scripts to convert TEI XML files to other formats without using the main client or server. diff --git a/grobid_client/grobid_client.py b/grobid_client/grobid_client.py index d32677a..ac1d431 100644 --- a/grobid_client/grobid_client.py +++ b/grobid_client/grobid_client.py @@ -17,6 +17,7 @@ from __future__ import annotations import os +import io import json import argparse import fnmatch @@ -32,7 +33,7 @@ import tarfile import tempfile import zipfile -from typing import Any, BinaryIO, Optional, Tuple, Union +from typing import Any, BinaryIO, Callable, Optional, Tuple, Union import copy from .format.TEI2LossyJSON import TEI2LossyJSONConverter @@ -82,6 +83,11 @@ class GrobidClient(ApiClient): TEI_SUFFIX = ".grobid.tei.xml" ERROR_FILE_RE = re.compile(r"_\d{3}\.txt$") + # Name given to a PDF that is passed as bytes and left unnamed by the caller. + # A multipart part needs a filename, and GROBID echoes it in its own logs and + # error messages, so an anonymous document should still be recognisable. + DEFAULT_IN_MEMORY_NAME = "document.pdf" + # Default configuration values DEFAULT_CONFIG: dict = { 'grobid_server': 'http://localhost:8070', @@ -366,6 +372,45 @@ def _test_server_connection(self) -> Tuple[bool, int]: self.logger.error(error_msg) raise ServerUnavailableException(error_msg) from e + def _warn_on_engine_concurrency(self, n: int) -> None: + """Preflight for in-memory runs: compare client concurrency to the server pool. + + An in-memory run keeps up to ``n`` documents in flight against the + server, so asking for more concurrency than the server has engines only + piles up requests that queue there or come back as 503, while asking + for less leaves engines idle. The pool size comes from ``/api/health`` + (``pool.maxActive``); a server without that endpoint (older GROBID) or + an unreadable answer is left alone - this is advisory, not a gate. + """ + try: + response = requests.get(self.get_server_url("health"), timeout=10) + payload = response.json() + except Exception as e: + self.logger.debug(f"Concurrency preflight skipped, /api/health not readable: {str(e)}") + return + + if not isinstance(payload, dict): + return + + if payload.get("ready") is False: + self.logger.warning( + f"GROBID server {self.config['grobid_server']} reports it is not ready (/api/health)") + + pool = payload.get("pool") + max_active = pool.get("maxActive") if isinstance(pool, dict) else None + if not isinstance(max_active, int) or max_active <= 0: + return + + if n > max_active: + self.logger.warning( + f"Client concurrency {n} exceeds the {max_active} engine(s) of the GROBID server: " + f"the surplus requests will queue on the server or be retried on 503. " + f"Consider n={max_active}, or raising 'concurrency' in the server's grobid.yaml.") + elif n < max_active: + self.logger.info( + f"The GROBID server has {max_active} engines but the client concurrency is only {n}; " + f"n={max_active} would use them all.") + def _output_file_name( self, input_file: str, @@ -580,6 +625,12 @@ def process_paths( self.logger.warning(f"No eligible files found in input(s): {inputs}") return + # Archives and remote files are about to be processed from memory with + # up to n documents in flight, so check upfront that the server has the + # engines to take them (citation lists take the disk route instead). + if (archive_paths or remote_files) and service != 'processCitationList': + self._warn_on_engine_concurrency(n) + processed_files_count = 0 errors_files_count = 0 skipped_files_count = 0 @@ -600,7 +651,7 @@ def process_paths( skipped_files_count += bs total_files += len(fs_files) - # Loose remote (s3) files: streamed to a temp dir one chunk at a time + # Loose remote (s3) files: streamed into memory one chunk at a time if remote_files: rt, rp, re_count, rs = self._process_remote_files( service, remote_files, output, n, @@ -955,6 +1006,34 @@ def _extract_archive_member( return target + def _read_archive_member(self, kind: str, archive: Any, member_name: str) -> Optional[BinaryIO]: + """Read a single archive entry into memory as a named document. + + The entry never touches the disk: it goes straight from the archive to + process_pdf, which accepts a named stream. The stream is named after the + entry (sanitized the same way extraction is), and that name is what the + output file is derived from. Returns None if the entry is unusable. + """ + name = self._safe_member_path("", member_name) + if name is None: + self.logger.warning(f"Skipping archive entry with unsafe path: {member_name}") + return None + + if kind == "zip": + source = archive.open(member_name) + else: + source = archive.extractfile(archive.getmember(member_name)) + if source is None: + return None + + try: + document = io.BytesIO(source.read()) + finally: + source.close() + + document.name = name + return document + def process_archive( self, service: str, @@ -977,17 +1056,22 @@ def process_archive( ) -> None: """Process the eligible files contained in a zip/tar archive. - The archive is never fully decompressed: entries are streamed to a - temporary directory in chunks of ``batch_size`` (from the config), each - chunk is sent to GROBID via ``process_batch``, and the temporary files - are removed before the next chunk is extracted. This keeps disk usage - bounded regardless of the archive size. Output files follow the same - flat naming convention as directory processing (one ```` per - result, in ``output``). + The archive is never fully decompressed: entries are read straight into + memory in chunks of ``batch_size`` (from the config) and each chunk is + sent to GROBID via ``process_batch``, so memory usage stays bounded by + the chunk size and nothing but the results ever touches the disk. (The + exception is ``processCitationList``, whose ``.txt`` inputs are read + from a path and therefore still go through a temporary directory.) + Output files follow the same flat naming convention as directory + processing (one ```` per result, in ``output``). """ start_time = time.time() self._warn_on_consolidation_timeout(consolidate_citations) + # Entries are about to be posted from memory with up to n in flight + if service != 'processCitationList': + self._warn_on_engine_concurrency(n) + total_files, processed, errors, skipped = self._process_archive_core( service, archive_path, output, n, generate_ids, consolidate_header, consolidate_citations, include_raw_citations, include_raw_affiliations, @@ -1062,8 +1146,51 @@ def _process_archive_core( print(f"Found {total_files} file(s) to process in {archive_path}") + # Citation lists go through process_txt, which reads from a path, so + # they still take the extract-to-a-temp-dir route. Everything else is + # read straight from the archive into memory and posted from there. + use_disk = service == 'processCitationList' + for chunk_start in range(0, total_files, batch_size_pdf): chunk = eligible_members[chunk_start:chunk_start + batch_size_pdf] + + if not use_disk: + documents = [] + for member_name in chunk: + if verbose: + self.logger.info(f"Reading {member_name} from {archive_path}") + document = self._read_archive_member(kind, archive, member_name) + if document is not None: + documents.append(document) + + if not documents: + continue + + batch_processed, batch_errors, batch_skipped = self.process_batch( + service, + documents, + ".", + output, + n, + generate_ids, + consolidate_header, + consolidate_citations, + include_raw_citations, + include_raw_affiliations, + tei_coordinates, + segment_sentences, + force, + verbose, + flavor, + json_output, + markdown_output, + skip_errors=skip_errors + ) + processed_files_count += batch_processed + errors_files_count += batch_errors + skipped_files_count += batch_skipped + continue + temp_dir = tempfile.mkdtemp(prefix="grobid_archive_") try: extracted_files = [] @@ -1133,10 +1260,11 @@ def _process_remote_files( markdown_output: bool, skip_errors: bool = False ) -> Tuple[int, int, int, int]: - """Stream loose remote (s3) files to a temp dir in chunks and process them. + """Stream loose remote (s3) files into memory in chunks and process them. Returns (total, processed, errors, skipped). Objects are fetched a - batch at a time and deleted before the next chunk, so disk stays bounded. + batch at a time straight into memory, so nothing is written to disk and + memory stays bounded by the chunk size. """ total = len(uris) if total == 0: @@ -1151,8 +1279,57 @@ def _process_remote_files( error_count = 0 skipped_count = 0 + # Citation lists go through process_txt, which reads from a path, so + # they still take the download-to-a-temp-dir route. Everything else is + # fetched straight into memory and posted from there. + use_disk = service == 'processCitationList' + for chunk_start in range(0, total, batch_size_pdf): chunk = uris[chunk_start:chunk_start + batch_size_pdf] + + if not use_disk: + documents = [] + for uri in chunk: + if verbose: + self.logger.info(f"Fetching {uri}") + try: + with self._s3_open(uri) as source: + document = io.BytesIO(source.read()) + except Exception as e: + self.logger.error(f"Failed to fetch {uri}: {str(e)}") + error_count += 1 + continue + document.name = self._s3_basename(uri) + documents.append(document) + + if not documents: + continue + + batch_processed, batch_errors, batch_skipped = self.process_batch( + service, + documents, + ".", + output, + n, + generate_ids, + consolidate_header, + consolidate_citations, + include_raw_citations, + include_raw_affiliations, + tei_coordinates, + segment_sentences, + force, + verbose, + flavor, + json_output, + markdown_output, + skip_errors=skip_errors + ) + processed_count += batch_processed + error_count += batch_errors + skipped_count += batch_skipped + continue + temp_dir = tempfile.mkdtemp(prefix="grobid_s3_") try: local_files = [] @@ -1233,8 +1410,13 @@ def process_batch( # with concurrent.futures.ProcessPoolExecutor(max_workers=n) as executor: results = [] for input_file in input_files: + # An entry is either a path or an in-memory document (a named + # stream, as fed by the archive/s3 streaming). Either way it goes + # by its name here, and process_pdf returns that same name, so + # the results below land on the same output file. + input_name = input_file if isinstance(input_file, str) else self._document_name(input_file) # check if TEI file is already produced - filename = self._output_file_name(input_file, input_path, output) + filename = self._output_file_name(input_name, input_path, output) if not force and os.path.isfile(filename): self.logger.info( f"{filename} already exists, skipping... (use --force to reprocess pdf input files)") @@ -1290,7 +1472,7 @@ def process_batch( previous_errors = self._find_error_files(filename) if previous_errors: self.logger.info( - f"{input_file} previously failed ({os.path.basename(previous_errors[0])}), " + f"{input_name} previously failed ({os.path.basename(previous_errors[0])}), " f"skipping... (use --force to retry it)") skipped_count += 1 continue @@ -1300,7 +1482,7 @@ def process_batch( selected_process = self.process_txt if verbose: - self.logger.info(f"Adding {input_file} to the queue") + self.logger.info(f"Adding {input_name} to the queue") r = executor.submit( selected_process, @@ -1404,7 +1586,75 @@ def process_batch( def process_pdf( self, service: str, - pdf_file: str, + pdf_file: Union[str, bytes, bytearray, memoryview, BinaryIO], + generate_ids: bool = False, + consolidate_header: bool = True, + consolidate_citations: bool = False, + include_raw_citations: bool = False, + include_raw_affiliations: bool = False, + tei_coordinates: bool = False, + segment_sentences: bool = False, + flavor: Optional[str] = None, + start: int = -1, + end: int = -1 + ) -> Tuple[str, int, Optional[str]]: + """Send a single PDF to GROBID. + + ``pdf_file`` is either a path to read from disk, or the document itself + already in memory - bytes, or any binary stream - for callers that got + the PDF from somewhere else than the filesystem: a database, an HTTP + response, an object store. They would otherwise have to write it out just + to have this client read it back. + See https://github.com/grobidOrg/grobid-client-python/pull/67 + + The document names itself: a path is its own name, and a stream is named + after its ``name`` attribute, which ``open()`` sets and which can be set + on anything else, ``io.BytesIO`` included: + + pdf = io.BytesIO(downloaded_bytes) + pdf.name = "paper.pdf" + + Bytes without a stream around them have nothing to be named after, so + they fall back to ``DEFAULT_IN_MEMORY_NAME``. + + Returns: + tuple: (document name, status code, response text) + """ + return self._process_named_pdf( + service, self._document_name(pdf_file), pdf_file, generate_ids, + consolidate_header, consolidate_citations, include_raw_citations, + include_raw_affiliations, tei_coordinates, segment_sentences, + flavor, start, end + ) + + def _document_name(self, pdf_file: Any) -> str: + """Name a document after itself: its path, or its stream's name.""" + if isinstance(pdf_file, str): + return pdf_file + name = getattr(pdf_file, "name", None) + # A file opened from a descriptor has an int here, not a name + return name if isinstance(name, str) else self.DEFAULT_IN_MEMORY_NAME + + def _pdf_opener(self, pdf_file: Any) -> Callable[[], BinaryIO]: + """Return a callable handing out a fresh stream over the document. + + Not a stream but a way to get one: a 503 makes us send the same document + again, and a stream that has already been posted is sitting at its end. + Rewinding is not enough either - the caller's stream may not be seekable + - so a stream is read once, here, and re-served from memory afterwards. + """ + if isinstance(pdf_file, str): + return lambda: open(pdf_file, "rb") + if isinstance(pdf_file, (bytes, bytearray, memoryview)): + return lambda: io.BytesIO(pdf_file) + content = pdf_file.read() + return lambda: io.BytesIO(content) + + def _process_named_pdf( + self, + service: str, + name: str, + pdf_file: Any, generate_ids: bool, consolidate_header: bool, consolidate_citations: bool, @@ -1416,13 +1666,43 @@ def process_pdf( start: int = -1, end: int = -1 ) -> Tuple[str, int, Optional[str]]: - pdf_handle = None + """Process a document under a name already decided by the caller.""" try: - pdf_handle = open(pdf_file, "rb") - + open_handle = self._pdf_opener(pdf_file) + except Exception as e: + self.logger.error(f"Failed to read PDF {name}: {str(e)}") + return (name, 400, f"Failed to open file: {str(e)}") + + return self._send_pdf( + service, name, open_handle, generate_ids, consolidate_header, + consolidate_citations, include_raw_citations, include_raw_affiliations, + tei_coordinates, segment_sentences, flavor, start, end + ) + + def _send_pdf( + self, + service: str, + name: str, + open_handle: Callable[[], BinaryIO], + generate_ids: bool, + consolidate_header: bool, + consolidate_citations: bool, + include_raw_citations: bool, + include_raw_affiliations: bool, + tei_coordinates: bool, + segment_sentences: bool, + flavor: Optional[str] = None, + start: int = -1, + end: int = -1 + ) -> Tuple[str, int, Optional[str]]: + """Post one document, retrying it from a fresh stream on a 503.""" + pdf_handle: Optional[BinaryIO] = None + try: + pdf_handle = open_handle() + files = { "input": ( - pdf_file, + name, pdf_handle, "application/pdf", {"Expires": "0"}, @@ -1461,10 +1741,11 @@ def process_pdf( if status == 503: return self._handle_server_busy_retry( - pdf_file, - self.process_pdf, + name, + self._send_pdf, service, - pdf_file, + name, + open_handle, generate_ids, consolidate_header, consolidate_citations, @@ -1477,22 +1758,116 @@ def process_pdf( end ) - return (pdf_file, status, res.text) - + return (name, status, res.text) + except IOError as e: - self.logger.error(f"Failed to open PDF file {pdf_file}: {str(e)}") - return (pdf_file, 400, f"Failed to open file: {str(e)}") + self.logger.error(f"Failed to open PDF file {name}: {str(e)}") + return (name, 400, f"Failed to open file: {str(e)}") except requests.exceptions.ReadTimeout as e: - self.logger.error(f"Request timeout for {pdf_file}: {str(e)}") - return (pdf_file, 408, f"Request timeout: {str(e)}") + self.logger.error(f"Request timeout for {name}: {str(e)}") + return (name, 408, f"Request timeout: {str(e)}") except requests.exceptions.RequestException as e: - return self._handle_request_error(pdf_file, e) + return self._handle_request_error(name, e) except Exception as e: - return self._handle_unexpected_error(pdf_file, e) + return self._handle_unexpected_error(name, e) finally: + # Always ours to close: the caller's own stream is never kept, only + # the bytes read out of it. if pdf_handle: pdf_handle.close() + def process_documents( + self, + service: str, + documents: Any, + n: int = 10, + generate_ids: bool = False, + consolidate_header: bool = True, + consolidate_citations: bool = False, + include_raw_citations: bool = False, + include_raw_affiliations: bool = False, + tei_coordinates: bool = False, + segment_sentences: bool = False, + flavor: Optional[str] = None, + verbose: bool = False + ) -> list: + """Process several documents concurrently and return their results. + + ``process_pdf`` handles one document per call, so a caller holding a list + of PDFs would have to build its own thread pool to get any concurrency + out of the server. This runs them through the same ThreadPoolExecutor the + file-based processing uses, but writes nothing: use ``process()`` for the + directory-oriented processing that produces TEI files on disk. + + Each document is whatever ``process_pdf`` accepts - a path, bytes, or a + stream - and names itself the same way. Unnamed documents (bare bytes) + are named after their position, so that every result stays identifiable. + + Returns: + list: one ``(name, status code, response text)`` per document, in the + order the documents were given - not in completion order, so results + can be zipped back onto whatever the caller has them keyed by. + A document that fails does not stop the others: its own entry carries + the error status, as with the file-based processing. + """ + items = self._name_documents(documents) + if not items: + self.logger.warning("No documents to process") + return [] + + if verbose: + self.logger.info(f"{len(items)} document(s) to process") + + self._warn_on_engine_concurrency(max(1, n)) + + # A pool of one is still a pool: n < 1 would make ThreadPoolExecutor raise + with concurrent.futures.ThreadPoolExecutor(max_workers=max(1, n)) as executor: + futures = [ + executor.submit( + self._process_named_pdf, + service, + name, + document, + generate_ids, + consolidate_header, + consolidate_citations, + include_raw_citations, + include_raw_affiliations, + tei_coordinates, + segment_sentences, + flavor + ) + for name, document in items + ] + # Collected in submission order rather than with as_completed(): + # in-memory documents have no filenames to be matched back on later. + return [future.result() for future in futures] + + def _name_documents(self, documents: Any) -> list: + """Pair every document with the name it will be known by. + + Documents that name themselves (a path, a stream with a ``name``) keep + their own; the rest are numbered after their position, because a batch of + documents all called ``document.pdf`` could not be told apart in the + results or in GROBID's logs. + """ + if isinstance(documents, (bytes, bytearray, memoryview)): + # Iterating a single PDF would yield its individual bytes, so catch + # the mistake here rather than sending thousands of empty documents. + raise TypeError( + "process_documents expects several documents; " + "use process_pdf for a single one" + ) + + items = [] + stem, extension = os.path.splitext(self.DEFAULT_IN_MEMORY_NAME) + for position, document in enumerate(documents, start=1): + name = self._document_name(document) + if name == self.DEFAULT_IN_MEMORY_NAME: + name = f"{stem}-{position}{extension}" + items.append((name, document)) + return items + def get_server_url(self, service: str) -> str: return self.config['grobid_server'] + "/api/" + service diff --git a/tests/test_grobid_client.py b/tests/test_grobid_client.py index 54ddc15..379744d 100644 --- a/tests/test_grobid_client.py +++ b/tests/test_grobid_client.py @@ -1,9 +1,12 @@ """ Unit tests for the GROBID client main functionality. """ +import concurrent.futures +import io import json import os import tempfile +import time from unittest.mock import Mock, patch, mock_open import pytest @@ -703,18 +706,22 @@ def _make_targz(path, entries, work): t.add(member_path, arcname=name) def _run(self, client, archive, output): - """Run archive processing with a fake GROBID post; return set of temp dirs used.""" - temp_dirs = set() + """Run archive processing with a fake GROBID post; return what was posted. + + Every multipart part is recorded as (name, content). The content has to + be read inside the call: the handle is closed before process_pdf returns. + """ + posted = [] def fake_post(url, files=None, data=None, headers=None, timeout=None): - temp_dirs.add(os.path.dirname(files['input'][0])) + posted.append((files['input'][0], files['input'][1].read())) resp = Mock() resp.text = 'ok' return (resp, 200) with patch.object(GrobidClient, 'post', side_effect=fake_post): client.process('processFulltextDocument', archive, output=output, force=True) - return temp_dirs + return posted @staticmethod def _tei_outputs(output_dir): @@ -760,13 +767,16 @@ def test_process_zip_streams_all_pdfs(self): 'ignore.txt': b'not a pdf', }) out = os.path.join(d, 'out') - temp_dirs = self._run(client, zip_path, out) + posted = self._run(client, zip_path, out) # all 3 PDFs processed, the .txt ignored assert self._tei_outputs(out) == ['a.grobid.tei.xml', 'b.grobid.tei.xml', 'c.grobid.tei.xml'] - # 3 files with batch_size 2 => 2 chunks => distinct temp dirs, all cleaned up - assert len(temp_dirs) >= 2 - assert all(not os.path.exists(td) for td in temp_dirs) + # each entry was posted straight from memory, under its archive name + assert sorted(posted) == [ + ('a.pdf', b'%PDF-a'), + ('c.PDF', b'%PDF-c'), + (os.path.join('sub', 'b.pdf'), b'%PDF-b'), + ] def test_process_targz(self): client = self._client(batch_size=10) @@ -774,9 +784,12 @@ def test_process_targz(self): tar_path = os.path.join(d, 'docs.tar.gz') self._make_targz(tar_path, {'x.pdf': b'%PDF-x', 'nested/y.pdf': b'%PDF-y'}, d) out = os.path.join(d, 'out') - temp_dirs = self._run(client, tar_path, out) + posted = self._run(client, tar_path, out) assert self._tei_outputs(out) == ['x.grobid.tei.xml', 'y.grobid.tei.xml'] - assert all(not os.path.exists(td) for td in temp_dirs) + assert sorted(posted) == [ + (os.path.join('nested', 'y.pdf'), b'%PDF-y'), + ('x.pdf', b'%PDF-x'), + ] def test_process_routes_archive_to_core(self): client = self._client() @@ -796,6 +809,47 @@ def test_process_zip_default_output_named_after_archive(self): self._run(client, zip_path, None) # no output -> defaults to assert self._tei_outputs(os.path.join(d, 'mydocs')) == ['a.grobid.tei.xml'] + def test_archive_entries_never_touch_disk(self): + """PDFs go from the archive straight to GROBID, without a temp dir.""" + client = self._client(batch_size=2) + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'docs.zip') + self._make_zip(zip_path, {'a.pdf': b'%PDF-a', 'sub/b.pdf': b'%PDF-b'}) + out = os.path.join(d, 'out') + + def fake_post(url, files=None, data=None, headers=None, timeout=None): + resp = Mock() + resp.text = 'ok' + return (resp, 200) + + with patch.object(GrobidClient, 'post', side_effect=fake_post): + with patch('grobid_client.grobid_client.tempfile.mkdtemp') as mock_mkdtemp: + client.process('processFulltextDocument', zip_path, output=out, force=True) + + mock_mkdtemp.assert_not_called() + assert self._tei_outputs(out) == ['a.grobid.tei.xml', 'b.grobid.tei.xml'] + + def test_citation_lists_are_still_extracted_to_disk(self): + """process_txt reads from a path, so citation lists keep the temp-dir route.""" + client = self._client(batch_size=10) + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'refs.zip') + self._make_zip(zip_path, {'refs.txt': b'one reference per line'}) + seen = {} + + def spy_batch(service, files, input_path, *args, **kwargs): + seen['files'] = list(files) + seen['on_disk'] = [os.path.isfile(f) for f in files] + return (len(files), 0, 0) + + with patch.object(GrobidClient, 'process_batch', side_effect=spy_batch): + client.process('processCitationList', zip_path, output=os.path.join(d, 'o')) + + assert seen['on_disk'] == [True] + assert [os.path.basename(f) for f in seen['files']] == ['refs.txt'] + # and the extraction dir is gone once the batch is done + assert not os.path.exists(seen['files'][0]) + def test_empty_archive_warns(self): client = self._client() with tempfile.TemporaryDirectory() as d: @@ -807,6 +861,136 @@ def test_empty_archive_warns(self): client.logger.warning.assert_called() +class TestEngineConcurrencyPreflight: + """In-memory runs check the server's engine pool against the client concurrency. + + An in-memory run keeps up to n documents in flight, so before starting one + the client asks /api/health how many engines the server actually has and + flags a mismatch. The check is advisory: a server without the endpoint + (older GROBID) or an unreadable answer never blocks the run. + """ + + def _client(self): + with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): + with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): + client = GrobidClient(check_server=False) + client.logger = Mock() + return client + + @staticmethod + def _health(max_active, ready=True): + """A response shaped like a real GROBID /api/health answer.""" + response = Mock() + response.status_code = 200 if ready else 503 + response.json.return_value = { + "initialized": True, + "ready": ready, + "pool": {"initialized": True, "active": 0, "idle": 0, "maxActive": max_active}, + "models": {"loaded": {"segmentation": "wapiti", "header": "delft"}, + "failed": {}, "totalLoaded": 2, "totalFailed": 0}, + "grobidHomeConfigured": True, + } + return response + + def test_warns_when_concurrency_exceeds_engines(self): + client = self._client() + with patch('grobid_client.grobid_client.requests.get', return_value=self._health(4)) as mock_get: + client._warn_on_engine_concurrency(10) + mock_get.assert_called_once_with('http://localhost:8070/api/health', timeout=10) + warning = client.logger.warning.call_args[0][0] + assert '10' in warning and '4 engine' in warning + + def test_silent_when_they_match(self): + client = self._client() + with patch('grobid_client.grobid_client.requests.get', return_value=self._health(10)): + client._warn_on_engine_concurrency(10) + client.logger.warning.assert_not_called() + client.logger.info.assert_not_called() + + def test_informs_when_engines_outnumber_concurrency(self): + """Idle engines are not an error, but the user should know they are there.""" + client = self._client() + with patch('grobid_client.grobid_client.requests.get', return_value=self._health(20)): + client._warn_on_engine_concurrency(10) + client.logger.warning.assert_not_called() + assert '20' in client.logger.info.call_args[0][0] + + def test_warns_when_server_not_ready(self): + client = self._client() + with patch('grobid_client.grobid_client.requests.get', return_value=self._health(4, ready=False)): + client._warn_on_engine_concurrency(4) + assert 'not ready' in client.logger.warning.call_args[0][0] + + def test_survives_a_server_without_the_endpoint(self): + """Older GROBID answers 404 with a non-JSON body; the run must go on.""" + client = self._client() + response = Mock() + response.status_code = 404 + response.json.side_effect = ValueError('not json') + with patch('grobid_client.grobid_client.requests.get', return_value=response): + client._warn_on_engine_concurrency(10) + client.logger.warning.assert_not_called() + + def test_survives_a_connection_error(self): + client = self._client() + with patch('grobid_client.grobid_client.requests.get', + side_effect=requests.exceptions.ConnectionError('down')): + client._warn_on_engine_concurrency(10) + client.logger.warning.assert_not_called() + + def test_process_documents_runs_the_preflight(self): + client = self._client() + + def fake_post(url=None, files=None, data=None, headers=None, timeout=None): + resp = Mock() + resp.text = 'ok' + return (resp, 200) + + with patch.object(GrobidClient, 'post', side_effect=fake_post): + with patch.object(GrobidClient, '_warn_on_engine_concurrency') as mock_check: + client.process_documents('processFulltextDocument', [b'%PDF-1.4 a'], n=3) + mock_check.assert_called_once_with(3) + + def test_archive_processing_runs_the_preflight(self): + import zipfile + client = self._client() + client.config['batch_size'] = 10 + with tempfile.TemporaryDirectory() as d: + zip_path = os.path.join(d, 'docs.zip') + with zipfile.ZipFile(zip_path, 'w') as z: + z.writestr('a.pdf', b'%PDF') + + def fake_post(url, files=None, data=None, headers=None, timeout=None): + resp = Mock() + resp.text = 'ok' + return (resp, 200) + + with patch.object(GrobidClient, 'post', side_effect=fake_post): + with patch.object(GrobidClient, '_warn_on_engine_concurrency') as mock_check: + client.process('processFulltextDocument', zip_path, + output=os.path.join(d, 'out'), n=5, force=True) + mock_check.assert_called_once_with(5) + + def test_local_files_skip_the_preflight(self): + """Plain file processing does not gain a new health call.""" + client = self._client() + client.config['batch_size'] = 10 + with tempfile.TemporaryDirectory() as d: + with open(os.path.join(d, 'a.pdf'), 'wb') as f: + f.write(b'%PDF') + + def fake_post(url, files=None, data=None, headers=None, timeout=None): + resp = Mock() + resp.text = 'ok' + return (resp, 200) + + with patch.object(GrobidClient, 'post', side_effect=fake_post): + with patch.object(GrobidClient, '_warn_on_engine_concurrency') as mock_check: + client.process('processFulltextDocument', d, + output=os.path.join(d, 'out'), force=True) + mock_check.assert_not_called() + + class TestGlobInput: """Tests for glob-pattern input resolution (--input as a glob).""" @@ -1164,3 +1348,299 @@ def test_skip_errors_counts_as_skipped(self): segment_sentences=False, force=False, skip_errors=True ) assert (processed, errors, skipped) == (0, 0, 1) + + +class TestInMemoryDocuments: + """Processing a PDF that is already in memory, without writing it to disk. + + Callers that get their PDFs from an API, a database or an object store had to + write them to a temporary file only so that this client could open it again. + See https://github.com/grobidOrg/grobid-client-python/pull/67 + """ + + PDF = b'%PDF-1.4 in memory' + + def _client(self): + with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): + with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): + client = GrobidClient(check_server=False) + client.logger = Mock() + return client + + def _post_spy(self, statuses=(200,)): + """Mock of post() recording the multipart part of every call. + + The content has to be read inside the call: the handle is closed before + process_pdf returns. + """ + sent = [] + responses = [] + for status in statuses: + response = Mock() + response.text = 'ok' + responses.append((response, status)) + + def post(url=None, files=None, data=None, headers=None, timeout=None): + name, handle, content_type, _ = files['input'] + sent.append({'name': name, 'content': handle.read(), 'type': content_type}) + return responses[len(sent) - 1] + + return post, sent + + def _named(self, content, name): + stream = io.BytesIO(content) + stream.name = name + return stream + + def test_bytes_are_sent_without_touching_the_filesystem(self): + client = self._client() + post, sent = self._post_spy() + + with patch('builtins.open', mock_open()) as mock_file: + with patch.object(GrobidClient, 'post', side_effect=post): + result = client.process_pdf('processFulltextDocument', self.PDF) + + mock_file.assert_not_called() + assert sent[0]['content'] == self.PDF + assert sent[0]['type'] == 'application/pdf' + assert result == (GrobidClient.DEFAULT_IN_MEMORY_NAME, 200, 'ok') + + def test_a_stream_is_named_after_itself(self): + """No filename parameter: the document carries its own name.""" + client = self._client() + post, sent = self._post_spy() + + with patch.object(GrobidClient, 'post', side_effect=post): + result = client.process_pdf( + 'processFulltextDocument', self._named(self.PDF, 'paper.pdf')) + + # the name travels with the request and comes back with the result, so a + # caller processing many documents can still tell them apart + assert sent[0]['name'] == 'paper.pdf' + assert sent[0]['content'] == self.PDF + assert result[0] == 'paper.pdf' + + def test_unnamed_stream_falls_back_to_the_default_name(self): + client = self._client() + post, sent = self._post_spy() + + with patch.object(GrobidClient, 'post', side_effect=post): + client.process_pdf('processFulltextDocument', io.BytesIO(self.PDF)) + + assert sent[0]['name'] == GrobidClient.DEFAULT_IN_MEMORY_NAME + + def test_a_file_descriptor_name_is_not_used_as_a_document_name(self): + """open(fd) leaves an int in .name, which is no name at all.""" + client = self._client() + post, sent = self._post_spy() + + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, 'x.pdf') + with open(path, 'wb') as f: + f.write(self.PDF) + fd = os.open(path, os.O_RDONLY) + with os.fdopen(fd, 'rb') as handle: + assert handle.name == fd + with patch.object(GrobidClient, 'post', side_effect=post): + result = client.process_pdf('processFulltextDocument', handle) + + assert sent[0]['name'] == GrobidClient.DEFAULT_IN_MEMORY_NAME + assert result[0] == GrobidClient.DEFAULT_IN_MEMORY_NAME + + def test_retry_after_503_resends_the_whole_document(self): + """The stream is consumed by the first request; the retry must not send an empty body.""" + client = self._client() + post, sent = self._post_spy(statuses=(503, 200)) + + with patch('time.sleep'): + with patch.object(GrobidClient, 'post', side_effect=post): + result = client.process_pdf( + 'processFulltextDocument', self._named(self.PDF, 'busy.pdf')) + + assert [call['content'] for call in sent] == [self.PDF, self.PDF] + assert [call['name'] for call in sent] == ['busy.pdf', 'busy.pdf'] + assert result == ('busy.pdf', 200, 'ok') + + def test_retry_after_503_reopens_a_file_on_disk(self): + client = self._client() + post, sent = self._post_spy(statuses=(503, 200)) + + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, 'busy.pdf') + with open(path, 'wb') as f: + f.write(self.PDF) + + with patch('time.sleep'): + with patch.object(GrobidClient, 'post', side_effect=post): + result = client.process_pdf('processFulltextDocument', path) + + assert [call['content'] for call in sent] == [self.PDF, self.PDF] + assert result == (path, 200, 'ok') + + def test_defaults_match_the_other_entry_points(self): + client = self._client() + captured = {} + + def post(url=None, files=None, data=None, headers=None, timeout=None): + captured.update(data) + response = Mock() + response.text = 'ok' + return response, 200 + + with patch.object(GrobidClient, 'post', side_effect=post): + client.process_pdf('processFulltextDocument', self.PDF) + + assert captured == {'consolidateHeader': '1'} + + def test_path_input_is_unchanged(self): + """The in-memory path must not alter how a file on disk is processed.""" + client = self._client() + post, sent = self._post_spy() + + with tempfile.TemporaryDirectory() as d: + pdf_path = os.path.join(d, 'on_disk.pdf') + with open(pdf_path, 'wb') as f: + f.write(self.PDF) + + with patch.object(GrobidClient, 'post', side_effect=post): + result = client.process_pdf( + 'processFulltextDocument', pdf_path, + generate_ids=False, consolidate_header=False, + consolidate_citations=False, include_raw_citations=False, + include_raw_affiliations=False, tei_coordinates=False, + segment_sentences=False) + + assert sent[0]['name'] == pdf_path + assert sent[0]['content'] == self.PDF + assert result[0] == pdf_path + + def test_unreadable_stream_is_reported_as_a_failed_document(self): + client = self._client() + + class Broken: + name = 'broken.pdf' + + def read(self): + raise IOError('device on fire') + + with patch.object(GrobidClient, 'post') as mock_post: + name, status, message = client.process_pdf('processFulltextDocument', Broken()) + + mock_post.assert_not_called() + assert (name, status) == ('broken.pdf', 400) + assert 'device on fire' in message + + +class TestProcessDocuments: + """Processing several documents concurrently.""" + + def _client(self): + with patch('grobid_client.grobid_client.GrobidClient._test_server_connection'): + with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): + client = GrobidClient(check_server=False) + client.logger = Mock() + return client + + def _post(self, failing=(), delays=None): + """Mock of post() returning the sent content back, optionally slowly.""" + def post(url=None, files=None, data=None, headers=None, timeout=None): + name, handle, _, _ = files['input'] + if delays and name in delays: + time.sleep(delays[name]) + response = Mock() + if name in failing: + response.text = 'boom' + return response, 500 + response.text = f'{handle.read().decode()}' + return response, 200 + + return post + + def _named(self, content, name): + stream = io.BytesIO(content) + stream.name = name + return stream + + def test_results_keep_the_input_order(self): + """The first document is the slowest, so completion order differs from input order.""" + client = self._client() + documents = [self._named(b'A', 'a.pdf'), + self._named(b'B', 'b.pdf'), + self._named(b'C', 'c.pdf')] + + with patch.object(GrobidClient, 'post', side_effect=self._post(delays={'a.pdf': 0.2})): + results = client.process_documents('processFulltextDocument', documents, n=3) + + assert [name for name, _, _ in results] == ['a.pdf', 'b.pdf', 'c.pdf'] + assert [tei for _, _, tei in results] == ['A', 'B', 'C'] + + def test_bare_contents_are_named_by_position(self): + client = self._client() + + with patch.object(GrobidClient, 'post', side_effect=self._post()): + results = client.process_documents('processFulltextDocument', [b'first', b'second']) + + # names have to stay unique, otherwise results cannot be told apart + assert [name for name, _, _ in results] == ['document-1.pdf', 'document-2.pdf'] + + def test_named_and_unnamed_documents_can_be_mixed(self): + client = self._client() + documents = [self._named(b'A', 'named.pdf'), b'bare', 'on/disk.pdf'] + + with patch('builtins.open', mock_open(read_data=b'D')): + with patch.object(GrobidClient, 'post', side_effect=self._post()): + results = client.process_documents('processFulltextDocument', documents) + + assert [name for name, _, _ in results] == ['named.pdf', 'document-2.pdf', 'on/disk.pdf'] + + def test_one_failure_does_not_stop_the_others(self): + client = self._client() + documents = [self._named(b'OK', 'ok.pdf'), + self._named(b'BAD', 'bad.pdf'), + self._named(b'OK2', 'ok2.pdf')] + + with patch.object(GrobidClient, 'post', side_effect=self._post(failing={'bad.pdf'})): + results = client.process_documents('processFulltextDocument', documents) + + assert [status for _, status, _ in results] == [200, 500, 200] + assert results[1] == ('bad.pdf', 500, 'boom') + + def test_concurrency_is_bounded_by_n(self): + client = self._client() + captured = {} + real_executor = concurrent.futures.ThreadPoolExecutor + + def executor(max_workers=None, **kwargs): + captured['max_workers'] = max_workers + return real_executor(max_workers=max_workers, **kwargs) + + with patch('concurrent.futures.ThreadPoolExecutor', side_effect=executor): + with patch.object(GrobidClient, 'post', side_effect=self._post()): + client.process_documents('processFulltextDocument', [b'a', b'b'], n=4) + + assert captured['max_workers'] == 4 + + def test_zero_workers_still_runs(self): + """max_workers=0 would make ThreadPoolExecutor raise.""" + client = self._client() + + with patch.object(GrobidClient, 'post', side_effect=self._post()): + results = client.process_documents('processFulltextDocument', [b'a'], n=0) + + assert [status for _, status, _ in results] == [200] + + def test_empty_input_returns_no_results(self): + client = self._client() + + with patch.object(GrobidClient, 'post') as mock_post: + assert client.process_documents('processFulltextDocument', []) == [] + + mock_post.assert_not_called() + client.logger.warning.assert_called() + + def test_a_single_pdf_is_rejected(self): + """Iterating one PDF would send each of its bytes as a document.""" + client = self._client() + + with pytest.raises(TypeError, match='process_pdf'): + client.process_documents('processFulltextDocument', b'%PDF-1.4 single') diff --git a/tests/test_s3.py b/tests/test_s3.py index 1ff1512..e9c9d7a 100644 --- a/tests/test_s3.py +++ b/tests/test_s3.py @@ -140,6 +140,38 @@ def test_process_paths_mixed_local_and_s3(tmp_path): assert _tei(out) == ["0000009.grobid.tei.xml", "local.grobid.tei.xml"] +@mock_aws +def test_process_s3_never_touches_disk(tmp_path): + """Loose PDFs and zip entries go from S3 straight to GROBID, in memory.""" + s3 = boto3.client("s3", region_name=REGION) + s3.create_bucket(Bucket=BUCKET) + s3.put_object(Bucket=BUCKET, Key="pdfs/0000001.pdf", Body=b"%PDF-mem") + s3.put_object(Bucket=BUCKET, Key="arch/docs.zip", Body=_zip_bytes({"z.pdf": b"%PDF-z"})) + + c = _client() + out = str(tmp_path / "out") + posted = [] + + def fake_post(url, files=None, data=None, headers=None, timeout=None): + # content read inside the call: the handle is closed right after it + posted.append((files["input"][0], files["input"][1].read())) + resp = Mock() + resp.text = "ok" + return (resp, 200) + + with patch.object(GrobidClient, "post", side_effect=fake_post): + with patch("grobid_client.grobid_client.tempfile.mkdtemp") as mock_mkdtemp: + c.process_paths( + "processFulltextDocument", + [f"s3://{BUCKET}/pdfs/*.pdf", f"s3://{BUCKET}/arch/docs.zip"], + output=out, force=True, + ) + + mock_mkdtemp.assert_not_called() + assert sorted(posted) == [("0000001.pdf", b"%PDF-mem"), ("z.pdf", b"%PDF-z")] + assert _tei(out) == ["0000001.grobid.tei.xml", "z.grobid.tei.xml"] + + def test_missing_extra_raises_helpful_error(): """If smart_open isn't importable, a clear install hint is raised.""" c = _client()