diff --git a/Readme.md b/Readme.md index e9ece6c..46aa5b2 100644 --- a/Readme.md +++ b/Readme.md @@ -190,7 +190,7 @@ 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 read into -> memory in chunks of `batch_size` and sent to GROBID straight from there — the archive is never fully decompressed and +> memory in chunks of `queue_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 @@ -384,7 +384,6 @@ settings. ```json { "grobid_server": "http://localhost:8070", - "batch_size": 1000, "sleep_time": 5, "timeout": 60, "coordinates": [ @@ -403,7 +402,7 @@ settings. | Parameter | Description | Default | |-----------------|------------------------------------------------------------------------------------------------------------------|-------------------------| | `grobid_server` | GROBID server URL | `http://localhost:8070` | -| `batch_size` | Thread pool size. **Tune carefully: a large batch size will result in the data being written less frequently** | 1000 | +| `queue_size` | Number of files queued per processing chunk. See [Choosing a queue size](#choosing-a-queue-size). | 1000 for local directories, 1.2 × `n` for archives and S3 | | `sleep_time` | Wait time when server is busy (seconds) | 5 | | `timeout` | Client-side timeout (seconds) | 180 | | `coordinates` | XML elements for coordinate extraction | See above | @@ -413,6 +412,24 @@ settings. > Since version 0.0.12, the config file is optional. The client will use default localhost settings if no configuration > is provided. +### Choosing a queue size + +`queue_size` controls how many files are grouped into one processing chunk. It is a memory/durability knob, not a +concurrency one: parallelism toward the GROBID server is set by `-n`, and each chunk is processed `n` files at a time. +When `queue_size` is not set, the client picks a sensible default per input type (see the table above); set it +explicitly only if you need to override that. A few guidelines: + +- **Never set it below `n`.** Effective parallelism is `min(n, queue_size)`: a queue smaller than the thread pool + leaves workers idle. A bit of headroom above `n` (the default streaming value is 1.2 × `n`) keeps the pool busy. +- **Results are written per chunk.** Output files land on disk only once a whole chunk has been processed, so a + larger queue means results are written less frequently and an interrupted run loses at most one chunk of work + (already-written results are skipped on re-run unless `--force` is used). +- **Local PDF directories:** the queue holds only file paths, so large values are essentially free — the default is + 1000. Lower it if you want results flushed to disk more often on long runs. +- **Archives (zip/tar) and S3:** each chunk is read or downloaded *into memory* before processing starts, so peak RAM + grows with `queue_size × average file size`. Keep it moderate — the 1.2 × `n` default is a safe floor; going up to + a few multiples of `n` (e.g. 2–5 ×) trades memory for slightly better throughput around chunk boundaries. + > [!WARNING] > **Citation consolidation and the `timeout` setting.** When `--consolidate_citations` (or `consolidate_citations=True`) > is enabled, GROBID queries external services (e.g. CrossRef) to enrich the extracted references. This is considerably diff --git a/config.json b/config.json index a12452e..fb84005 100644 --- a/config.json +++ b/config.json @@ -1,6 +1,5 @@ { "grobid_server": "http://localhost:8070", - "batch_size": 1000, "timeout": 180, "sleep_time": 5, "coordinates": [ diff --git a/grobid_client/grobid_client.py b/grobid_client/grobid_client.py index ac1d431..07e33ac 100644 --- a/grobid_client/grobid_client.py +++ b/grobid_client/grobid_client.py @@ -16,6 +16,7 @@ """ from __future__ import annotations +import math import os import io import json @@ -91,7 +92,9 @@ class GrobidClient(ApiClient): # Default configuration values DEFAULT_CONFIG: dict = { 'grobid_server': 'http://localhost:8070', - 'batch_size': 10, + # None means "follow the concurrency n at processing time", so the + # default queue never starves the thread pool + 'queue_size': None, 'sleep_time': 5, 'timeout': 180, 'coordinates': [ @@ -121,7 +124,7 @@ class GrobidClient(ApiClient): def __init__( self, grobid_server: Optional[str] = None, - batch_size: Optional[int] = None, + queue_size: Optional[int] = None, coordinates: Optional[list] = None, sleep_time: Optional[int] = None, timeout: Optional[int] = None, @@ -143,7 +146,7 @@ def __init__( # This ensures CLI arguments override config file values self._set_config_params({ 'grobid_server': grobid_server, - 'batch_size': batch_size, + 'queue_size': queue_size, 'coordinates': coordinates, 'sleep_time': sleep_time, 'timeout': timeout @@ -161,6 +164,26 @@ def _set_config_params(self, params: dict) -> None: if value is not None: self.config[key] = value + # Default chunk size when walking a local directory: only file paths are + # queued (nothing is pre-loaded), so a large chunk costs next to nothing. + LOCAL_DIR_QUEUE_SIZE = 1000 + + def _effective_queue_size(self, n: int, local_files: bool = False) -> int: + """Return the configured queue_size, or a default derived from n. + + A queue smaller than the thread pool leaves workers idle, so when no + explicit value is configured the chunk size follows the concurrency n, + with 20% headroom to keep the pool busy around the chunk boundary. For + local directories, where the queue holds only file paths, a large + fixed default is used instead. + """ + configured = self.config.get("queue_size") + if configured: + return configured + if local_files: + return self.LOCAL_DIR_QUEUE_SIZE + return math.ceil(n * 1.2) + def _warn_on_consolidation_timeout(self, consolidate_citations: bool) -> None: """Warn when citation consolidation is enabled with a low client timeout. @@ -328,6 +351,11 @@ def _load_config(self, path: str = "./config.json") -> None: config_json = config_file.read() # Update the default config with values from the file file_config = json.loads(config_json) + if 'batch_size' in file_config: + file_config.setdefault('queue_size', file_config.pop('batch_size')) + temp_logger.warning( + "Config key 'batch_size' is deprecated, use 'queue_size' instead" + ) self.config.update(file_config) temp_logger.info("Configuration file loaded successfully") except FileNotFoundError as e: @@ -850,11 +878,11 @@ def _run_file_batches( markdown_output: bool, skip_errors: bool = False ) -> Tuple[int, int, int]: - """Run process_batch over a list of files in chunks of batch_size. + """Run process_batch over a list of files in chunks of queue_size. Returns the aggregated (processed, errors, skipped) counts. """ - batch_size_pdf = self.config["batch_size"] + queue_size = self._effective_queue_size(n, local_files=True) processed_files_count = 0 errors_files_count = 0 skipped_files_count = 0 @@ -870,7 +898,7 @@ def _run_file_batches( batch.append(input_file) - if len(batch) == batch_size_pdf: + if len(batch) == queue_size: batch_processed, batch_errors, batch_skipped = self.process_batch( service, batch, input_path, output, n, generate_ids, consolidate_header, consolidate_citations, include_raw_citations, @@ -1057,7 +1085,7 @@ def process_archive( """Process the eligible files contained in a zip/tar archive. The archive is never fully decompressed: entries are read straight into - memory in chunks of ``batch_size`` (from the config) and each chunk is + memory in chunks of ``queue_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 @@ -1112,7 +1140,7 @@ def _process_archive_core( Does not print the final summary (the caller does), so it can be aggregated with other inputs when resolving a glob pattern. """ - batch_size_pdf = self.config["batch_size"] + queue_size = self._effective_queue_size(n) # Results must survive the temporary extraction directories, so when no # output is given we default to a directory named after the archive. For @@ -1151,8 +1179,8 @@ def _process_archive_core( # 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] + for chunk_start in range(0, total_files, queue_size): + chunk = eligible_members[chunk_start:chunk_start + queue_size] if not use_disk: documents = [] @@ -1273,7 +1301,7 @@ def _process_remote_files( if output is None: output = "." - batch_size_pdf = self.config["batch_size"] + queue_size = self._effective_queue_size(n) print(f"Found {total} remote file(s) to process") processed_count = 0 error_count = 0 @@ -1284,8 +1312,8 @@ def _process_remote_files( # 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] + for chunk_start in range(0, total, queue_size): + chunk = uris[chunk_start:chunk_start + queue_size] if not use_disk: documents = [] diff --git a/tests/conftest.py b/tests/conftest.py index f1fe0c3..80ce52c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -23,7 +23,7 @@ def temp_config_file(temp_dir): """Create a temporary config file for tests.""" config = { 'grobid_server': 'http://localhost:8070', - 'batch_size': 10, + 'queue_size': 10, 'coordinates': ["persName", "figure"], 'sleep_time': 2, 'timeout': 30, diff --git a/tests/test_config.json b/tests/test_config.json index 1f0f5c7..cb59078 100644 --- a/tests/test_config.json +++ b/tests/test_config.json @@ -1,6 +1,6 @@ { "grobid_server": "http://localhost:8070", - "batch_size": 10, + "queue_size": 10, "coordinates": ["persName", "figure", "ref"], "sleep_time": 2, "timeout": 30, diff --git a/tests/test_conversions.py b/tests/test_conversions.py index 9016b7c..aaf21d0 100644 --- a/tests/test_conversions.py +++ b/tests/test_conversions.py @@ -39,7 +39,7 @@ def setup_method(self): self.test_config = { 'grobid_server': 'http://localhost:8070', - 'batch_size': 10, + 'queue_size': 10, 'sleep_time': 5, 'timeout': 180, 'logging': { diff --git a/tests/test_grobid_client.py b/tests/test_grobid_client.py index 379744d..58794ee 100644 --- a/tests/test_grobid_client.py +++ b/tests/test_grobid_client.py @@ -22,7 +22,7 @@ def setup_method(self): """Set up test fixtures.""" self.test_config = { 'grobid_server': 'http://localhost:8070', - 'batch_size': 1000, + 'queue_size': 1000, 'coordinates': ["persName", "figure", "ref"], 'sleep_time': 5, 'timeout': 60, @@ -43,12 +43,27 @@ def test_init_default_values(self, mock_configure_logging, mock_test_server): client = GrobidClient(check_server=False) assert client.config['grobid_server'] == 'http://localhost:8070' - assert client.config['batch_size'] == 10 + assert client.config['queue_size'] is None assert client.config['sleep_time'] == 5 assert client.config['timeout'] == 180 assert 'persName' in client.config['coordinates'] mock_configure_logging.assert_called_once() + @patch('grobid_client.grobid_client.GrobidClient._test_server_connection') + @patch('grobid_client.grobid_client.GrobidClient._configure_logging') + def test_effective_queue_size(self, mock_configure_logging, mock_test_server): + """Test the queue_size defaults: 1.2 * n for streaming, 1000 for local dirs.""" + mock_test_server.return_value = (True, 200) + + client = GrobidClient(check_server=False) + assert client._effective_queue_size(40) == 48 + assert client._effective_queue_size(10) == 12 + assert client._effective_queue_size(40, local_files=True) == 1000 + + client.config['queue_size'] = 100 + assert client._effective_queue_size(40) == 100 + assert client._effective_queue_size(40, local_files=True) == 100 + @patch('grobid_client.grobid_client.GrobidClient._test_server_connection') @patch('grobid_client.grobid_client.GrobidClient._configure_logging') def test_init_custom_values(self, mock_configure_logging, mock_test_server): @@ -58,7 +73,7 @@ def test_init_custom_values(self, mock_configure_logging, mock_test_server): custom_coords = ["figure", "ref"] client = GrobidClient( grobid_server='http://custom:9090', - batch_size=500, + queue_size=500, coordinates=custom_coords, sleep_time=10, timeout=120, @@ -66,7 +81,7 @@ def test_init_custom_values(self, mock_configure_logging, mock_test_server): ) assert client.config['grobid_server'] == 'http://custom:9090' - assert client.config['batch_size'] == 500 + assert client.config['queue_size'] == 500 assert client.config['coordinates'] == custom_coords assert client.config['sleep_time'] == 10 assert client.config['timeout'] == 120 @@ -114,6 +129,19 @@ def test_load_config_success(self, mock_configure_logging, mock_test_server, moc mock_file.assert_called_once_with('/path/to/config.json', 'r') assert client.config['grobid_server'] == 'http://test:8080' + @patch('builtins.open', new_callable=mock_open, read_data='{"batch_size": 250}') + @patch('grobid_client.grobid_client.GrobidClient._test_server_connection') + @patch('grobid_client.grobid_client.GrobidClient._configure_logging') + def test_load_config_legacy_batch_size(self, mock_configure_logging, mock_test_server, mock_file): + """Test that the deprecated batch_size config key is mapped to queue_size.""" + mock_test_server.return_value = (True, 200) + + client = GrobidClient(check_server=False) + client._load_config('/path/to/config.json') + + assert client.config['queue_size'] == 250 + assert 'batch_size' not in client.config + @patch('grobid_client.grobid_client.GrobidClient._test_server_connection') @patch('grobid_client.grobid_client.GrobidClient._configure_logging') def test_load_config_file_not_found(self, mock_configure_logging, mock_test_server): @@ -680,12 +708,12 @@ def test_get_server_url_edge_cases(self, mock_configure_logging, mock_test_serve class TestArchiveInput: """Tests for streaming zip/tar archives as input (process_archive).""" - def _client(self, batch_size=2): + def _client(self, queue_size=2): 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() - client.config['batch_size'] = batch_size + client.config['queue_size'] = queue_size return client @staticmethod @@ -757,7 +785,7 @@ def test_safe_member_path_blocks_traversal(self): assert client._safe_member_path(dest, '.') is None def test_process_zip_streams_all_pdfs(self): - client = self._client(batch_size=2) + client = self._client(queue_size=2) with tempfile.TemporaryDirectory() as d: zip_path = os.path.join(d, 'docs.zip') self._make_zip(zip_path, { @@ -779,7 +807,7 @@ def test_process_zip_streams_all_pdfs(self): ] def test_process_targz(self): - client = self._client(batch_size=10) + client = self._client(queue_size=10) with tempfile.TemporaryDirectory() as d: 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) @@ -802,7 +830,7 @@ def test_process_routes_archive_to_core(self): assert mock_core.call_args.args[1] == zip_path def test_process_zip_default_output_named_after_archive(self): - client = self._client(batch_size=10) + client = self._client(queue_size=10) with tempfile.TemporaryDirectory() as d: zip_path = os.path.join(d, 'mydocs.zip') self._make_zip(zip_path, {'a.pdf': b'%PDF'}) @@ -811,7 +839,7 @@ def test_process_zip_default_output_named_after_archive(self): 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) + client = self._client(queue_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'}) @@ -831,7 +859,7 @@ def fake_post(url, files=None, data=None, headers=None, timeout=None): 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) + client = self._client(queue_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'}) @@ -954,7 +982,7 @@ def fake_post(url=None, files=None, data=None, headers=None, timeout=None): def test_archive_processing_runs_the_preflight(self): import zipfile client = self._client() - client.config['batch_size'] = 10 + client.config['queue_size'] = 10 with tempfile.TemporaryDirectory() as d: zip_path = os.path.join(d, 'docs.zip') with zipfile.ZipFile(zip_path, 'w') as z: @@ -974,7 +1002,7 @@ def fake_post(url, files=None, data=None, headers=None, timeout=None): 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 + client.config['queue_size'] = 10 with tempfile.TemporaryDirectory() as d: with open(os.path.join(d, 'a.pdf'), 'wb') as f: f.write(b'%PDF') @@ -994,12 +1022,12 @@ def fake_post(url, files=None, data=None, headers=None, timeout=None): class TestGlobInput: """Tests for glob-pattern input resolution (--input as a glob).""" - def _client(self, batch_size=50): + def _client(self, queue_size=50): 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() - client.config['batch_size'] = batch_size + client.config['queue_size'] = queue_size return client @staticmethod @@ -1213,7 +1241,7 @@ def _client(self): with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): client = GrobidClient(check_server=False) client.logger = Mock() - client.config['batch_size'] = 50 + client.config['queue_size'] = 50 return client @staticmethod diff --git a/tests/test_integration.py b/tests/test_integration.py index 1ec2dd5..f0edc08 100644 --- a/tests/test_integration.py +++ b/tests/test_integration.py @@ -22,7 +22,7 @@ def setup_method(self): # Create a temporary config file self.temp_config = { 'grobid_server': self.test_server_url, - 'batch_size': 10, + 'queue_size': 10, 'coordinates': ["persName", "figure"], 'sleep_time': 2, 'timeout': 30, @@ -60,7 +60,7 @@ def test_client_initialization_with_config_file(self, mock_get): # Verify config was loaded assert client.config['grobid_server'] == self.test_server_url - assert client.config['batch_size'] == 10 + assert client.config['queue_size'] == 10 assert client.config['sleep_time'] == 2 assert client.config['timeout'] == 30 @@ -91,14 +91,14 @@ def test_configuration_validation(self): with patch('grobid_client.grobid_client.GrobidClient._configure_logging'): client = GrobidClient( grobid_server='http://custom:9090', - batch_size=500, + queue_size=500, config_path=self.config_file, check_server=False ) # Constructor values should override config file values (CLI precedence) assert client.config['grobid_server'] == 'http://custom:9090' - assert client.config['batch_size'] == 500 + assert client.config['queue_size'] == 500 def test_logging_configuration(self): """Test logging configuration from config file.""" diff --git a/tests/test_s3.py b/tests/test_s3.py index e9c9d7a..9eb6320 100644 --- a/tests/test_s3.py +++ b/tests/test_s3.py @@ -34,12 +34,12 @@ def _aws_env(monkeypatch): monkeypatch.setenv("AWS_DEFAULT_REGION", REGION) -def _client(batch_size=2): +def _client(queue_size=2): with patch.object(GrobidClient, "_test_server_connection"): with patch.object(GrobidClient, "_configure_logging"): c = GrobidClient(check_server=False) c.logger = Mock() - c.config["batch_size"] = batch_size + c.config["queue_size"] = queue_size return c