From 14895c57dcd5672b6b2c634405e370bc02c7ac6a Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 15 May 2026 14:10:06 -0400 Subject: [PATCH 1/3] Parallelized input record read and bitstream download Why these changes are being introduced: A primary flow in this app is reading input records from TIMDEX, currently filtered to theses records, downloading fulltext for them, and then writing back to the TIMDEX dataset. The writing back is not yet established, so for now we write to JSONLines as debugging output. But the rest of the business logic is ready to be built, building on the CLI scaffolding. Parallelization is important. If a single item takes 1 second, for ~150k items in DSpace that's ~41 hours. Not realistic. With parallelization of ~20 workers, and retries and backoffs to ensure we don't get rate limited at the DSpace boundary, seeing more like ~33 records per second, which works out to ~1 hour. It's likely that an agressive run could bump `--workers=50` (or even higher) and ramp this up. How this addresses that need: Introduces new `dfh/harvest.py` file encapsulates parallelized reading of TIMDEX records, generation of pre-signed URLs from DSpace, and actual bitstream downloads from S3. This parallelization has sensible defaults, but is configureable via some limited CLI arguments (e.g. `--workers`). The parallelization lumps the generation of URLs and downloading of content into a single method with retries. This allows the parallelization to be relatively simple, and leave retries to the method. Ultimately, an iterator of record metadata + fulltext is yielded, which is what we'll need to write back to TIMDEX dataset. NOTE: DSpace credential parsing is still assuming env vars are set for host, username, and password. It's likely that we'll parse an env var with JSON credentials; this is not yet implemented. Side effects of this change: * None Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/USE-557 --- dfh/cli.py | 50 +++++++++++++++--- dfh/harvest.py | 134 +++++++++++++++++++++++++++++++++++++++++++++++++ pyproject.toml | 1 + uv.lock | 14 ++++++ 4 files changed, 191 insertions(+), 8 deletions(-) create mode 100644 dfh/harvest.py diff --git a/dfh/cli.py b/dfh/cli.py index 729b45f..710076c 100644 --- a/dfh/cli.py +++ b/dfh/cli.py @@ -3,8 +3,11 @@ from datetime import timedelta import click +import jsonlines from dfh.config import configure_logger, configure_sentry +from dfh.harvest import record_and_fulltext_iter +from dfh.timdex_dataset import TIMDEXThesesRecords logger = logging.getLogger(__name__) @@ -33,9 +36,8 @@ def main( def _log_command_elapsed_time() -> None: elapsed_time = time.perf_counter() - ctx.obj["start_time"] - logger.info( - "Total time to complete process: %s", str(timedelta(seconds=elapsed_time)) - ) + elapsed_time_display = timedelta(seconds=elapsed_time) + logger.info(f"Total time to complete process: {elapsed_time_display}") ctx.call_on_close(_log_command_elapsed_time) @@ -60,7 +62,7 @@ def test_dspace_connection( @click.pass_context @click.option( "--dataset-location", - required=True, + required=False, type=click.Path(), help=( "TIMDEX dataset location, e.g. 's3://timdex/dataset', " @@ -80,6 +82,13 @@ def test_dspace_connection( default=None, help="Maximum records of records to retrieve and harvest fulltext for.", ) +@click.option( + "--workers", + required=False, + type=int, + default=10, + help="Max number of parallel workers for downloading.", +) @click.option( "--output-jsonl", required=False, @@ -88,15 +97,40 @@ def test_dspace_connection( help="Write harvested fulltext to a local JSONLines file (primarily for testing).", ) def harvest( - ctx: click.Context, + _ctx: click.Context, dataset_location: str, run_id: str | None, record_limit: int | None, + workers: int, output_jsonl: str | None, ) -> None: """Harvest fulltext from DSpace and write to a TIMDEX dataset. - The argument --datast-location is required, as it provides the TIMDEX records to read - metadata from, and the location where fulltext will be written back to. + Flow: + 1. Get an iterator of source records from TIMDEX Dataset. + 2. Use to retrieve fulltext from DSpace, yielding as another iterator + 3. Write record + fulltext to TIMDEX dataset or JSONLines as output """ - raise NotImplementedError + # get iterator of target TIMDEX records + bitstream information + ttr = TIMDEXThesesRecords( + dataset_location=dataset_location, + run_id=run_id, + limit=record_limit, + ) + records_and_bitstreams = ttr.record_and_bitstream_metadata_iter() + + # get iterator of records + fulltext, ready for writing + records_and_fulltexts = record_and_fulltext_iter( + records_and_bitstreams, + max_workers=workers, + ) + + if not output_jsonl: + raise ValueError( + "WIP: Until TIMDEX dataset has new 'fulltexts' " + "data type, only JSONL output is supported." + ) + + with jsonlines.open(output_jsonl, "w") as writer: + for record in records_and_fulltexts: + writer.write(record) diff --git a/dfh/harvest.py b/dfh/harvest.py new file mode 100644 index 0000000..3448f74 --- /dev/null +++ b/dfh/harvest.py @@ -0,0 +1,134 @@ +import logging +import threading +import time +from collections.abc import Iterator +from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait + +import requests +from dspace_rest_client.client import DSpaceClient + +from dfh.dspace import get_dspace_client, get_presigned_url_for_bitstream + +logger = logging.getLogger(__name__) + +# module level, cross-thread object to hold thread specific DSpace client instances +# see: https://docs.python.org/3/library/threading.html#thread-local-data +threaded_dspace_clients = threading.local() + + +def record_and_fulltext_iter( + records: Iterator[dict], + *, + max_workers: int = 10, + log_progress_interval: int = 10, +) -> Iterator[dict]: + """Yield records with fulltext fetched in parallel. + + Uses ThreadPoolExecutor (main IO bottleneck is network) to generate pre-signed URLs + and download bitstream content in parallel. + + The worker function _record_with_fulltext() has built-in retries. This orchestration + function with parallelizes the work is not aware of retries. + """ + pending: set[Future[dict]] = set() + completed_count = 0 + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + for record in records: + pending.add(executor.submit(_record_with_fulltext, record)) + + if len(pending) < max_workers: + continue + + done, pending = wait(pending, return_when=FIRST_COMPLETED) + for future in done: + completed_count += 1 + yield future.result() + if completed_count % log_progress_interval == 0: + logger.debug(f"Extracted fulltext for {completed_count} records.") + + while pending: + done, pending = wait(pending, return_when=FIRST_COMPLETED) + for future in done: + completed_count += 1 + yield future.result() + if completed_count % log_progress_interval == 0: + logger.debug(f"Extracted fulltext for {completed_count} records.") + + +def _record_with_fulltext( + record: dict, + *, + retry_attempts: int = 4, + initial_backoff_seconds: float = 1.0, + backoff_factor: float = 2.0, + timeout_seconds: int = 60, +) -> dict: + """Return a TIMDEX fulltext record for one source record. + + The primary work here is two network requests: + 1. Hit DSpace API for a presigned S3 URL for a bitstream UUID + 2. Use presigned S3 URL to download content + + This function is a worker function designed to be parallelized across threads. As + such, note the use of _get_dspace_client_for_thread() which ensures that it checks + the threading.local() object for a DSpaceClient instance to use that is unique and + reusable by this thread. + + Retries and backoffs are fairly simple: all exceptions, for either network request, + that bubble up are caught, logged, and increment the retry counter. + """ + bitstream_uuid = record["fulltext_bitstream"]["uuid"] + + fulltext = None + for attempt in range(1, retry_attempts + 1): + try: + # reuse dspace client for thread, or init if first time + dspace_client = _get_dspace_client_for_thread() + + # generate a pre-signed URL for a bitstream UUID + pre_signed_url = get_presigned_url_for_bitstream( + dspace_client, + bitstream_uuid, + ) + + # download bitstream content from S3 + response = requests.get(pre_signed_url, timeout=timeout_seconds) + response.raise_for_status() + fulltext = response.text + + # break out of retries loop if successful + break + + except Exception as exc: + if attempt == retry_attempts: + logger.exception( + f"Max retries of {retry_attempts} encountered, failed to download " + f"bitstream '{bitstream_uuid}'" + ) + break + + sleep_seconds = initial_backoff_seconds * (backoff_factor ** (attempt - 1)) + logger.warning( + f"Retrying download for bitstream " + f"'{bitstream_uuid}' after attempt {attempt}/{retry_attempts} " + f"failed; sleeping {sleep_seconds:.1f} seconds. Cause: {exc}" + ) + time.sleep(sleep_seconds) + + return { + "timdex_record_id": record["timdex_record_id"], + "run_id": record["run_id"], + "run_record_offset": record["run_record_offset"], + "fulltext_bistream_uuid": bitstream_uuid, + "fulltext_bistream_content": fulltext, + } + + +def _get_dspace_client_for_thread() -> DSpaceClient: + """Get thread's DSpaceClient from module level threaded_dspace_clients.""" + dspace_client = getattr(threaded_dspace_clients, "dspace_client", None) + if dspace_client is None: + dspace_client = get_dspace_client() + threaded_dspace_clients.dspace_client = dspace_client + return dspace_client diff --git a/pyproject.toml b/pyproject.toml index a1d656a..bdef0f2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ requires-python = ">=3.13" dependencies = [ "click>=8.2.1", "dspace-rest-client>=0.1.17", + "jsonlines>=4.0.0", "requests>=2.33.1", "sentry-sdk>=2.34.1", "timdex-dataset-api", diff --git a/uv.lock b/uv.lock index 38fc036..affc26f 100644 --- a/uv.lock +++ b/uv.lock @@ -555,6 +555,7 @@ source = { editable = "." } dependencies = [ { name = "click" }, { name = "dspace-rest-client" }, + { name = "jsonlines" }, { name = "requests" }, { name = "sentry-sdk" }, { name = "timdex-dataset-api" }, @@ -575,6 +576,7 @@ dev = [ requires-dist = [ { name = "click", specifier = ">=8.2.1" }, { name = "dspace-rest-client", specifier = ">=0.1.17" }, + { name = "jsonlines", specifier = ">=4.0.0" }, { name = "requests", specifier = ">=2.33.1" }, { name = "sentry-sdk", specifier = ">=2.34.1" }, { name = "timdex-dataset-api", git = "https://github.com/MITLibraries/timdex-dataset-api" }, @@ -892,6 +894,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/14/2f/967ba146e6d58cf6a652da73885f52fc68001525b4197effc174321d70b4/jmespath-1.1.0-py3-none-any.whl", hash = "sha256:a5663118de4908c91729bea0acadca56526eb2698e83de10cd116ae0f4e97c64", size = 20419, upload-time = "2026-01-22T16:35:24.919Z" }, ] +[[package]] +name = "jsonlines" +version = "4.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/87/bcda8e46c88d0e34cad2f09ee2d0c7f5957bccdb9791b0b934ec84d84be4/jsonlines-4.0.0.tar.gz", hash = "sha256:0c6d2c09117550c089995247f605ae4cf77dd1533041d366351f6f298822ea74", size = 11359, upload-time = "2023-09-01T12:34:44.187Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/62/d9ba6323b9202dd2fe166beab8a86d29465c41a0288cbe229fac60c1ab8d/jsonlines-4.0.0-py3-none-any.whl", hash = "sha256:185b334ff2ca5a91362993f42e83588a360cf95ce4b71a73548502bda52a7c55", size = 8701, upload-time = "2023-09-01T12:34:42.563Z" }, +] + [[package]] name = "librt" version = "0.10.0" From dcb11b6bd178cf47328b4c729ee9107d6881e6d2 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Mon, 18 May 2026 08:54:38 -0400 Subject: [PATCH 2/3] Add dspace api authentication warmup Why these changes are being introduced: Intermittent failures were observed for authentication against a cold API endpoint, where after a successful authentication, all future requests were healthy. Could be an artifact of server deploys and upgrades (this codebase is landing during a DSpace migration), or an artifact of DSpace CRIS, unsure. Either way, this authentication check + retry feels pretty harmless to layer on until we're 100% sure it's not needed. How this addresses that need: * Adds new `warm_dspace_auth()` function * Applied at CLI level for `harvest` command before real API work Side effects of this change: * A cold DSpace API, when given some retries to authenticate, is successful after that. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/USE-557 --- dfh/cli.py | 4 ++++ dfh/dspace.py | 52 ++++++++++++++++++++++++++++++++++++++++++++ tests/test_dspace.py | 15 +++++++++++++ 3 files changed, 71 insertions(+) create mode 100644 tests/test_dspace.py diff --git a/dfh/cli.py b/dfh/cli.py index 710076c..009362a 100644 --- a/dfh/cli.py +++ b/dfh/cli.py @@ -6,6 +6,7 @@ import jsonlines from dfh.config import configure_logger, configure_sentry +from dfh.dspace import warm_dspace_auth from dfh.harvest import record_and_fulltext_iter from dfh.timdex_dataset import TIMDEXThesesRecords @@ -111,6 +112,9 @@ def harvest( 2. Use to retrieve fulltext from DSpace, yielding as another iterator 3. Write record + fulltext to TIMDEX dataset or JSONLines as output """ + # warm DSpace API authentication + warm_dspace_auth() + # get iterator of target TIMDEX records + bitstream information ttr = TIMDEXThesesRecords( dataset_location=dataset_location, diff --git a/dfh/dspace.py b/dfh/dspace.py index ff5c683..3432d7a 100644 --- a/dfh/dspace.py +++ b/dfh/dspace.py @@ -1,10 +1,13 @@ import logging import os +import time from dspace_rest_client.client import DSpaceClient logger = logging.getLogger(__name__) +HTTP_OK = 200 + def get_dspace_client( *, @@ -27,6 +30,55 @@ def get_dspace_client( return client +def warm_dspace_auth( + *, + retry_attempts: int = 5, + initial_backoff_seconds: float = 2.0, + backoff_factor: float = 2.0, +) -> None: + """Warm and verify DSpace authentication before threaded downloads. + + Intermittent failures were observed for authentication against a cold API endpoint, + where after a successful authentication, all future requests were healthy. Could be + an artifact of server deploys and upgrades (this codebase is landing during a DSpace + migration), or an artifact of DSpace CRIS, unsure. Either way, this authentication + check + retry feels pretty harmless to layer on until we're 100% sure it's not needed. + """ + if retry_attempts < 1: + raise ValueError("retry_attempts must be at least 1") + + logger.info("Warming DSpace authentication for cold start") + last_error: Exception | None = None + for attempt in range(1, retry_attempts + 1): + try: + client = get_dspace_client(auth_on_init=True) + response = client.api_get(f"{client.API_ENDPOINT}/authn/status") + if ( + response.status_code == HTTP_OK + and response.json().get("authenticated") is True + ): + logger.info("DSpace authentication warm-up successful") + return + last_error = RuntimeError( + "DSpace authentication warm-up failed: " + f"{response.status_code} {response.text}" + ) + except Exception as exc: # noqa: BLE001 + last_error = exc + + if attempt < retry_attempts: + sleep_seconds = initial_backoff_seconds * (backoff_factor ** (attempt - 1)) + logger.warning( + "DSpace authentication warm-up attempt " + f"{attempt}/{retry_attempts} failed; sleeping " + f"{sleep_seconds:.1f} seconds before retry. Cause: {last_error}" + ) + time.sleep(sleep_seconds) + + msg = f"DSpace authentication warm-up failed after {retry_attempts} attempts" + raise RuntimeError(msg) from last_error + + def get_presigned_url_for_bitstream( dspace_client: DSpaceClient, bitstream_uuid: str, diff --git a/tests/test_dspace.py b/tests/test_dspace.py new file mode 100644 index 0000000..f0b7f67 --- /dev/null +++ b/tests/test_dspace.py @@ -0,0 +1,15 @@ +from unittest import mock + +from dfh.dspace import warm_dspace_auth + + +def test_warm_dspace_auth_retries_after_401_and_succeeds(): + fail_auth = mock.Mock(status_code=401, text="Unauthorized") + success_auth = mock.Mock(status_code=200, text="OK") + success_auth.json.return_value = {"authenticated": True} + + client = mock.Mock() + client.api_get.side_effect = [fail_auth, success_auth] + + with mock.patch("dfh.dspace.get_dspace_client", return_value=client): + assert warm_dspace_auth(initial_backoff_seconds=0.1) is None From 403f68d23c5ae408a480a919ee8ed7da86812615 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Mon, 18 May 2026 16:05:24 -0400 Subject: [PATCH 3/3] Fix 'bistream' typo --- dfh/harvest.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dfh/harvest.py b/dfh/harvest.py index 3448f74..fff3093 100644 --- a/dfh/harvest.py +++ b/dfh/harvest.py @@ -120,8 +120,8 @@ def _record_with_fulltext( "timdex_record_id": record["timdex_record_id"], "run_id": record["run_id"], "run_record_offset": record["run_record_offset"], - "fulltext_bistream_uuid": bitstream_uuid, - "fulltext_bistream_content": fulltext, + "fulltext_bitstream_uuid": bitstream_uuid, + "fulltext_bitstream_content": fulltext, }