From fbf006ca565b23ac34d369d8d53927e459457c38 Mon Sep 17 00:00:00 2001 From: Luca Foppiano Date: Tue, 11 Aug 2026 23:32:18 +0200 Subject: [PATCH 1/3] Accept a PDF already in memory, not only a path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit process_pdf could only read the document from disk, so callers holding a PDF in memory - fetched from an API, read out of a database or an object store - had to write it to a temporary file only for the client to open it again. It now takes the document itself as well: bytes, or any binary stream. Nothing says which of the two it is; the object does. A document also names itself, from the "name" attribute open() sets on files and that can be set on anything else, io.BytesIO included, so the identity of a document is not lost by going through memory - it travels with the request and comes back with the result. Bytes on their own have nothing to be named after and fall back to DEFAULT_IN_MEMORY_NAME. A stream is read once, up front, and re-served from memory afterwards: the 503 retry sends the same document again, and a consumed (or non-seekable) stream would silently post an empty body the second time around. That is also why the retry no longer recurses through the public entry point, which would have had to re-derive a name from a source that is by then exhausted. process_documents processes several of them concurrently, through the same ThreadPoolExecutor the file-based processing uses. Results come back in input order rather than in completion order: in-memory documents have no filenames to be matched back on afterwards, so the caller has nothing but the order to zip them onto. A single PDF passed by mistake raises instead of being iterated, which would otherwise send one request per byte. Resumes #67 Co-authored-by: Jan GΓΆpfert <94385965+jangoepfert@users.noreply.github.com> --- Readme.md | 53 ++++++ grobid_client/grobid_client.py | 231 +++++++++++++++++++++++-- tests/test_grobid_client.py | 299 +++++++++++++++++++++++++++++++++ 3 files changed, 566 insertions(+), 17 deletions(-) diff --git a/Readme.md b/Readme.md index fdf4a2e..dc967e6 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 @@ -289,6 +290,58 @@ 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. + +> [!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..e7b6ab2 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', @@ -1404,7 +1410,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 +1490,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 +1565,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 +1582,114 @@ 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") + + # 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..44da693 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 @@ -1164,3 +1167,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') From 3629b21f8e76d1df06d310aca714294e5c3632fe Mon Sep 17 00:00:00 2001 From: Luca Foppiano Date: Sat, 15 Aug 2026 18:02:42 +0200 Subject: [PATCH 2/3] Post archive and s3 entries to GROBID from memory The archive and s3 streaming (#117) shipped with a known detour: every entry was written to a temporary directory only so that process_pdf could open it again from a path, with the commit itself noting this would go away once PR #67 landed. It has landed, so this plugs the two together: archive entries and loose s3 objects are now read straight into memory and posted from there, named after the entry (or the s3 basename), and nothing but the results ever touches the disk. process_batch accepts the in-memory documents alongside paths - an entry goes by the name it carries, and since process_pdf returns that same name, the result lands on the same output file it would have as a path. The one input that still takes the temp-dir route is processCitationList, whose .txt files are read by process_txt from a path. The archive tests asserted on the temp dirs the posts came from, which no longer exist; they now assert on what actually crossed the wire - each entry posted once, under its archive name, with its own bytes - plus explicitly that mkdtemp is never called on the pdf path. Completes what #117 left pending on #67. --- Readme.md | 11 +-- grobid_client/grobid_client.py | 153 ++++++++++++++++++++++++++++++--- tests/test_grobid_client.py | 71 ++++++++++++--- tests/test_s3.py | 32 +++++++ 4 files changed, 239 insertions(+), 28 deletions(-) diff --git a/Readme.md b/Readme.md index dc967e6..3dd2f6f 100644 --- a/Readme.md +++ b/Readme.md @@ -189,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 diff --git a/grobid_client/grobid_client.py b/grobid_client/grobid_client.py index e7b6ab2..7011e17 100644 --- a/grobid_client/grobid_client.py +++ b/grobid_client/grobid_client.py @@ -606,7 +606,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, @@ -961,6 +961,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, @@ -983,13 +1011,14 @@ 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) @@ -1068,8 +1097,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 = [] @@ -1139,10 +1211,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: @@ -1157,8 +1230,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 = [] @@ -1239,8 +1361,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)") @@ -1296,7 +1423,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 @@ -1306,7 +1433,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, diff --git a/tests/test_grobid_client.py b/tests/test_grobid_client.py index 44da693..6ec6d8b 100644 --- a/tests/test_grobid_client.py +++ b/tests/test_grobid_client.py @@ -706,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): @@ -763,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) @@ -777,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() @@ -799,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: 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() From c665ee5dba6d91d9527768ea7e611c3e95f44b11 Mon Sep 17 00:00:00 2001 From: Luca Foppiano Date: Sat, 15 Aug 2026 18:32:22 +0200 Subject: [PATCH 3/3] Preflight the server's engine count before an in-memory run An in-memory run keeps up to n documents in flight against the server, so a client concurrency above the server's engine pool only piles up requests that queue there or come back as 503, while one below it leaves engines idle. Neither is visible from the client side until the throughput disappoints. Before process_documents and the in-memory archive/s3 streaming start, the client now asks /api/health how many engines the server has (pool.maxActive) and logs a warning when n exceeds them - with the number to use instead - and an info message when they outnumber n. The check is advisory, not a gate: a server without the endpoint (older GROBID), an unreadable answer or a connection failure never blocks the run. A server answering ready: false is also surfaced as a warning. Plain file-based processing is unchanged and makes no extra call. --- Readme.md | 5 ++ grobid_client/grobid_client.py | 51 +++++++++++++ tests/test_grobid_client.py | 130 +++++++++++++++++++++++++++++++++ 3 files changed, 186 insertions(+) diff --git a/Readme.md b/Readme.md index 3dd2f6f..e9ece6c 100644 --- a/Readme.md +++ b/Readme.md @@ -339,6 +339,11 @@ Documents that do not name themselves are named `document-1.pdf`, `document-2.pd 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. diff --git a/grobid_client/grobid_client.py b/grobid_client/grobid_client.py index 7011e17..ac1d431 100644 --- a/grobid_client/grobid_client.py +++ b/grobid_client/grobid_client.py @@ -372,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, @@ -586,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 @@ -1023,6 +1068,10 @@ def process_archive( 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, @@ -1769,6 +1818,8 @@ def process_documents( 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 = [ diff --git a/tests/test_grobid_client.py b/tests/test_grobid_client.py index 6ec6d8b..379744d 100644 --- a/tests/test_grobid_client.py +++ b/tests/test_grobid_client.py @@ -861,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)."""