From ac95b92d39e7db77d0f5f8af12772103fbac7408 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 30 Oct 2026 14:20:20 -0400 Subject: [PATCH 01/11] Cache LOC crosswalk code lookups in Marc transformer Why these changes are being introduced: * Profiling an alma transform showed ~45% of time spent in Marc.loc_crosswalk_code_to_name(), which searches the entire LOC language or country crosswalk with bs4 for every code on every record. How this addresses that need: * Cache the code element found for each code, per crosswalk, so each distinct code is only searched once. * The cache stores the exact result of the original find(), preserving first-match behavior for codes that appear more than once in a crosswalk (e.g. 'ai' is an obsolete Anguilla code before it is Armenia's current code). * DEBUG logging for invalid and obsolete codes is unchanged, per record. Side effects of this change: * Output is unchanged: 2,000 alma records re-transformed from a parquet input produced identical transformed records. ~1.7x faster for alma end to end (~55 -> ~92 records/sec, single process). Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-738 * https://mitlibraries.atlassian.net/browse/TIMX-343 --- tests/sources/xml/test_marc.py | 47 ++++++++++++++++++++++++++++++ transmogrifier/sources/xml/marc.py | 12 ++++++-- 2 files changed, 57 insertions(+), 2 deletions(-) diff --git a/tests/sources/xml/test_marc.py b/tests/sources/xml/test_marc.py index 867e247..5194c07 100644 --- a/tests/sources/xml/test_marc.py +++ b/tests/sources/xml/test_marc.py @@ -1,11 +1,14 @@ # ruff: noqa: E501, SLF001 import logging +from unittest import mock import pytest from bs4 import BeautifulSoup # type: ignore[import-untyped] import transmogrifier.models as timdex +from transmogrifier.config import load_external_config from transmogrifier.exceptions import SkippedRecordEvent +from transmogrifier.sources.xml import marc from transmogrifier.sources.xml.marc import Marc @@ -2071,6 +2074,50 @@ def test_loc_crosswalk_code_to_name_returns_name(caplog, loc_country_crosswalk): ) +@pytest.mark.parametrize( + "crosswalk_path", ["config/loc-countries.xml", "config/loc-languages.xml"] +) +def test_loc_crosswalk_code_to_name_cached_matches_uncached(crosswalk_path): + # duplicate and obsolete codes, plus a sample of current codes and an invalid code + crosswalk = load_external_config(crosswalk_path, "xml") + code_elements = crosswalk.find_all("code") + all_codes = [str(code.string) for code in code_elements] + codes = [ + *{code for code in all_codes if all_codes.count(code) > 1}, + *[str(code.string) for code in code_elements if code.get("status")][:25], + *all_codes[::50], + "wrong", + ] + + uncached = [] + for code in codes: + with mock.patch.dict(marc._loc_crosswalk_code_cache, clear=True): + uncached.append( + Marc.loc_crosswalk_code_to_name(code, crosswalk, "record-01", "code") + ) + cached = [ + Marc.loc_crosswalk_code_to_name(code, crosswalk, "record-01", "code") + for code in codes + ] + + assert cached == uncached + + +def test_loc_crosswalk_code_to_name_cached_duplicate_code_uses_first_match( + caplog, loc_country_crosswalk +): + # 'ai' is an obsolete code for Anguilla before it is the current code for Armenia + caplog.set_level(logging.DEBUG) + for record_id in ("record-01", "record-02"): + assert ( + Marc.loc_crosswalk_code_to_name( + "ai", loc_country_crosswalk, record_id, "country" + ) + == "Anguilla" + ) + assert f"Record #{record_id} uses an obsolete country code: ai" in caplog.text + + def test_get_main_titles_record_with_title(): marc_xml_records = Marc.parse_source_file( "tests/fixtures/marc/marc_record_all_fields.xml" diff --git a/transmogrifier/sources/xml/marc.py b/transmogrifier/sources/xml/marc.py index ccfec5f..279d2e2 100644 --- a/transmogrifier/sources/xml/marc.py +++ b/transmogrifier/sources/xml/marc.py @@ -2,7 +2,7 @@ from collections import defaultdict from collections.abc import Iterator -from bs4 import Tag # type: ignore[import-untyped] +from bs4 import Tag import transmogrifier.models as timdex from transmogrifier.config import load_external_config @@ -12,6 +12,9 @@ logger = logging.getLogger(__name__) +# cache of LOC crosswalk code element lookups, keyed by (code type, code) +_loc_crosswalk_code_cache: dict[tuple[str, str], Tag | None] = {} + class Marc(XMLTransformer): """Marc transformer.""" @@ -192,7 +195,12 @@ def loc_crosswalk_code_to_name( str | None: If a mapping for the code is found in the crosswalk, the mapped value is returned; else None is returned. """ - code_element = crosswalk.find("code", string=code) + # searching the full crosswalk is expensive and repeated for every record, so + # cache the code element found for each code + cache_key = (code_type, code) + if cache_key not in _loc_crosswalk_code_cache: + _loc_crosswalk_code_cache[cache_key] = crosswalk.find("code", string=code) + code_element = _loc_crosswalk_code_cache[cache_key] if code_element is None: logger.debug( "Record #%s uses an invalid %s code: %s", record_id, code_type, code From 162d26dce4c2f32c7cc887c3e3e2741d447eb925 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 30 Oct 2026 14:26:22 -0400 Subject: [PATCH 02/11] Read source records as bytes and parse during transform Why these changes are being introduced: * To transform records in parallel worker processes, the main process should only read source records, not parse them. Parsing with bs4 is ~1 ms per record, which would cap a parallel run near ~1,000 records/sec, and bs4 objects do not pass between processes well. How this addresses that need: * Add Transformer.read_source_file(), implemented by XMLTransformer and JSONTransformer, which yields each source record as bytes: lxml's serialization of each XML element, or each non-blank JSONLines line. * Rename parse_source_records_from_dataset() and parse_source_records_from_parquet() to read_source_records_from_*(), which yield the stored source_record bytes. * Transformer.load() passes these bytes iterators to the transformer, and Transformer.__next__() parses each record with parse_source_record() before transform(). * Transformer.parse_source_file() remains, as read_source_file() plus parse_source_record(), and transformers can still be given already parsed records. * ResearchDatabases overrides read_source_file() to append its synthetic delete records as bytes, replacing its copy of the XML file parsing. * Remove the jsonlines dependency, no longer used. Side effects of this change: * Stored source_record is unchanged: still serialized from the parsed record. * A source record that cannot be parsed is now an 'error' action for that record, rather than an exception that stops the run. Blank lines in JSONLines files are skipped. * Output is unchanged: fixture files for every source, and 2,000 alma records from a Timmy export, produced identical source and transformed records. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-738 --- pyproject.toml | 1 - tests/conftest.py | 2 +- tests/sources/test_transformer.py | 60 +++++++-- tests/sources/xml/test_researchdatabases.py | 15 ++- transmogrifier/sources/jsontransformer.py | 13 +- transmogrifier/sources/transformer.py | 116 +++++++++--------- .../sources/xml/researchdatabases.py | 28 +---- transmogrifier/sources/xmltransformer.py | 8 +- uv.lock | 14 --- 9 files changed, 134 insertions(+), 123 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 7f7457a..8ea24ec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,6 @@ dependencies = [ "beautifulsoup4==4.12.3", "boto3", "click", - "jsonlines", "lxml", "pandas", "pyarrow", diff --git a/tests/conftest.py b/tests/conftest.py index 082f16b..575241d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -74,7 +74,7 @@ def runner(): @pytest.fixture def generic_transformer(): class GenericTransformer(Transformer): - def parse_source_file(self): + def read_source_file(self): pass def parse_source_record(self): diff --git a/tests/sources/test_transformer.py b/tests/sources/test_transformer.py index e112adb..64a890f 100644 --- a/tests/sources/test_transformer.py +++ b/tests/sources/test_transformer.py @@ -14,6 +14,7 @@ import transmogrifier.models as timdex from transmogrifier.exceptions import DeletedRecordEvent, SkippedRecordEvent from transmogrifier.sources import transformer as transformer_module +from transmogrifier.sources.json.aardvark import MITAardvark from transmogrifier.sources.jsontransformer import JSONTransformer from transmogrifier.sources.transformer import PARQUET_FILES_SCHEMA, Transformer from transmogrifier.sources.xml.datacite import Datacite @@ -33,7 +34,7 @@ def test_load_exclusion_list(source_transformer, mock_s3_exclusion_list): def test_next_iter_sets_action_skip_when_record_is_excluded(tmp_path): class ExcludingTransformer(Transformer): @classmethod - def parse_source_file(cls, _source_file: str): + def read_source_file(cls, _source_file: str): return iter(()) @classmethod @@ -352,7 +353,7 @@ def test_jsontransformer_parse_source_record_non_dict_raises_error(): JSONTransformer.parse_source_record(b"[1, 2, 3]") -def test_transformer_parse_source_records_from_dataset_yields_parsed_records(): +def test_transformer_read_source_records_from_dataset_yields_bytes(): with mock.patch( "transmogrifier.sources.transformer.TIMDEXDataset" ) as mocked_dataset_class: @@ -364,12 +365,12 @@ def test_transformer_parse_source_records_from_dataset_yields_parsed_records(): ] ) records = list( - JSONTransformer.parse_source_records_from_dataset( + JSONTransformer.read_source_records_from_dataset( "cool-repo", "s3://bucket/dataset" ) ) - assert records == [{"id": "1"}, {"id": "2"}] + assert records == [b'{"id": "1"}', b'{"id": "2"}'] mocked_dataset_class.assert_called_once_with("s3://bucket/dataset") mocked_records.read_dicts_iter.assert_called_once_with( table="current_records", @@ -404,7 +405,7 @@ def test_transformer_load_from_source_parquet(tmp_path, source_parquet_location) assert next(transformer).action == "index" -def test_transformer_parse_source_records_from_parquet_filters_source_and_action( +def test_transformer_read_source_records_from_parquet_filters_source_and_action( tmp_path, ): pq.write_table( @@ -431,14 +432,14 @@ def test_transformer_parse_source_records_from_parquet_filters_source_and_action tmp_path / "records.parquet", ) records = list( - JSONTransformer.parse_source_records_from_parquet( + JSONTransformer.read_source_records_from_parquet( "cool-repo", str(tmp_path / "records.parquet") ) ) - assert records == [{"id": "1"}, {"id": "4"}] + assert records == [b'{"id": "1"}', b'{"id": "4"}'] -def test_transformer_parse_source_records_from_parquet_directory_in_path_order( +def test_transformer_read_source_records_from_parquet_directory_in_path_order( tmp_path, ): data_dir = tmp_path / "export" / "data" @@ -458,11 +459,11 @@ def test_transformer_parse_source_records_from_parquet_directory_in_path_order( data_dir / f"part-{i:05d}.parquet", ) records = list( - JSONTransformer.parse_source_records_from_parquet( + JSONTransformer.read_source_records_from_parquet( "cool-repo", str(tmp_path / "export") ) ) - assert records == [{"id": "0"}, {"id": "1"}] + assert records == [b'{"id": "0"}', b'{"id": "1"}'] def test_transformer_write_to_parquet_files_matches_timmy_export_layout( @@ -544,3 +545,42 @@ def test_transformer_write_to_file_writes_only_indexed_records( assert count == len(lines) == source_transformer.transformed_record_count == 3 # noqa: PLR2004 assert len(source_transformer.deleted_records) == 1 assert all(json.loads(line)["timdex_record_id"] for line in lines) + + +def test_transformer_load_reads_source_records_as_bytes(source_input_file): + transformer = Transformer.load("cool-repo", source_input_file) + assert isinstance(next(transformer.source_records), bytes) + + +def test_transformer_next_iter_parses_bytes_source_record(source_input_file): + transformer = Transformer.load("cool-repo", source_input_file) + with mock.patch.object( + transformer, "parse_source_record", wraps=transformer.parse_source_record + ) as mocked_parse: + dataset_record = next(transformer) + mocked_parse.assert_called_once() + assert isinstance(mocked_parse.call_args.args[0], bytes) + assert dataset_record.action == "index" + + +def test_transformer_next_iter_unparseable_source_record_sets_action_error(): + transformer = MITAardvark("gismit", iter([b"not json"])) + dataset_record = next(transformer) + assert dataset_record.action == "error" + assert dataset_record.source_record == b"not json" + assert transformer.error_record_count == 1 + + +def test_xmltransformer_read_source_file_yields_record_bytes(): + records = list(XMLTransformer.read_source_file("tests/fixtures/oai_pmh_records.xml")) + assert all(isinstance(record, bytes) for record in records) + assert XMLTransformer.parse_source_record(records[0]).find("record") + + +def test_jsontransformer_read_source_file_yields_line_bytes(tmp_path): + source_file = tmp_path / "records.jsonl" + source_file.write_text('{"id": "1"}\n\n{"id": "2"}\r\n') + assert list(JSONTransformer.read_source_file(str(source_file))) == [ + b'{"id": "1"}', + b'{"id": "2"}', + ] diff --git a/tests/sources/xml/test_researchdatabases.py b/tests/sources/xml/test_researchdatabases.py index 3cf646a..c433301 100644 --- a/tests/sources/xml/test_researchdatabases.py +++ b/tests/sources/xml/test_researchdatabases.py @@ -62,14 +62,19 @@ def researchdatabases_transformer(): def test_researchdatabases_yields_oai_records_pass(researchdatabases_transformer): - oai_records = list( - researchdatabases_transformer._yield_oai_xml_records_for_indexing( - researchdatabases_transformer.source_file + with patch.object( + ResearchDatabases, + "_yield_api_records_for_deleting", + return_value=iter([]), + ): + oai_records = list( + researchdatabases_transformer.read_source_file( + researchdatabases_transformer.source_file + ) ) - ) assert len(oai_records) == 3 - assert isinstance(oai_records[0], Tag) + assert isinstance(oai_records[0], bytes) def test_researchdatabases_yields_deleted_records_success( diff --git a/transmogrifier/sources/jsontransformer.py b/transmogrifier/sources/jsontransformer.py index 1e47677..0303da2 100644 --- a/transmogrifier/sources/jsontransformer.py +++ b/transmogrifier/sources/jsontransformer.py @@ -4,7 +4,6 @@ from abc import abstractmethod from typing import TYPE_CHECKING, final -import jsonlines import smart_open # type: ignore[import-untyped] from transmogrifier.sources.transformer import JSON, Transformer @@ -18,19 +17,19 @@ class JSONTransformer(Transformer): @final @classmethod - def parse_source_file(cls, source_file: str) -> Iterator[dict[str, JSON]]: + def read_source_file(cls, source_file: str) -> Iterator[bytes]: """ - Parse JSON file and return source records as JSON objects via an iterator. + Read JSONLines file and yield each non-blank line as bytes. May not be overridden. - Validates that records in the file are dicts for proper processing. - Args: source_file: A file containing source records to be transformed. """ - with jsonlines.Reader(smart_open.open(source_file, "r")) as records: - yield from records.iter(type=dict) + with smart_open.open(source_file, "rb") as file: + for line in file: + if line.strip(): + yield line.rstrip(b"\r\n") @final @classmethod diff --git a/transmogrifier/sources/transformer.py b/transmogrifier/sources/transformer.py index c980fb5..fd632d3 100644 --- a/transmogrifier/sources/transformer.py +++ b/transmogrifier/sources/transformer.py @@ -18,7 +18,6 @@ import pyarrow as pa import pyarrow.compute as pc import pyarrow.dataset as ds -import pyarrow.parquet as pq import smart_open from bs4 import Tag from pyarrow import fs @@ -78,7 +77,7 @@ class Transformer(ABC): def __init__( self, source: str, - source_records: Iterator[dict[str, JSON] | Tag], + source_records: Iterator[bytes | dict[str, JSON] | Tag], exclusion_list_path: str | None = None, source_file: str | None = None, source_dataset: str | None = None, @@ -93,7 +92,8 @@ def __init__( source: Source repository label. Must match a source key from config.SOURCES. exclusion_list_path: S3 or local filepath to exclusion list CSV file. exclsion_list: The exclusion list for this particular source. - source_records: A set of source records to be processed. + source_records: A set of source records to be processed, as bytes to be + parsed during transformation, or already parsed records. source_file: Filepath of the input source file. source_dataset: Location of the TIMDEX dataset with input source records. source_parquet: Location of parquet file(s) with input source records. @@ -105,7 +105,7 @@ def __init__( self._exclusion_list: list[str] | None = None self.source_base_url: str = SOURCES[source]["base-url"] self.source_name = SOURCES[source]["name"] - self.source_records: Iterator[JSON | Tag] = source_records + self.source_records: Iterator[bytes | dict[str, JSON] | Tag] = source_records self.processed_record_count: int = 0 self.transformed_record_count: int = 0 self.skipped_record_count: int = 0 @@ -172,13 +172,19 @@ def __iter__(self) -> Iterator[DatasetRecord]: def __next__(self) -> DatasetRecord: """Return next transformed record.""" while True: + source_record = None transformed_record = None timdex_record_id = None - source_record = next(self.source_records) + record = next(self.source_records) self.processed_record_count += 1 try: + source_record = ( + self.parse_source_record(record) + if isinstance(record, bytes) + else record + ) transformed_record = self.transform(source_record) timdex_record_id = transformed_record.timdex_record_id transformed_record.timdex_provenance = timdex.TimdexProvenance( @@ -210,7 +216,11 @@ def __next__(self) -> DatasetRecord: return DatasetRecord( timdex_record_id=timdex_record_id, - source_record=self.serialize_source_record(source_record), + source_record=( + self.serialize_source_record(source_record) + if source_record is not None + else record + ), transformed_record=( json.dumps(transformed_record.asdict()).encode() if transformed_record @@ -266,14 +276,16 @@ def load( if len([i for i in (source_file, source_dataset, source_parquet) if i]) != 1: raise ValueError(INPUT_SOURCE_ERROR_MESSAGE) + # source records are read as bytes, and parsed one at a time during transform + source_records: Iterator[bytes] if source_file: - source_records = transformer_class.parse_source_file(source_file) + source_records = transformer_class.read_source_file(source_file) elif source_dataset: - source_records = transformer_class.parse_source_records_from_dataset( + source_records = transformer_class.read_source_records_from_dataset( source, source_dataset ) elif source_parquet: - source_records = transformer_class.parse_source_records_from_parquet( + source_records = transformer_class.read_source_records_from_parquet( source, source_parquet ) @@ -565,9 +577,9 @@ def get_valid_title(self, source_record: dict[str, JSON] | Tag) -> str: @classmethod @abstractmethod - def parse_source_file(cls, source_file: str) -> Iterator[dict[str, JSON] | Tag]: + def read_source_file(cls, source_file: str) -> Iterator[bytes]: """ - Parse source file and return source records via an iterator. + Read source file and yield each source record as bytes, without parsing. Must be overridden by format subclasses. @@ -575,6 +587,17 @@ def parse_source_file(cls, source_file: str) -> Iterator[dict[str, JSON] | Tag]: source_file: A file containing source records to be transformed. """ + @classmethod + def parse_source_file(cls, source_file: str) -> Iterator[dict[str, JSON] | Tag]: + """ + Read source file and yield each source record parsed. + + Args: + source_file: A file containing source records to be transformed. + """ + for source_record in cls.read_source_file(source_file): + yield cls.parse_source_record(source_record) + @classmethod @abstractmethod def parse_source_record(cls, source_record: bytes) -> dict[str, JSON] | Tag: @@ -588,15 +611,13 @@ def parse_source_record(cls, source_record: bytes) -> dict[str, JSON] | Tag: """ @classmethod - def parse_source_records_from_dataset( + def read_source_records_from_dataset( cls, source: str, source_dataset: str - ) -> Iterator[dict[str, JSON] | Tag]: - """Read current source records from a TIMDEX dataset and parse them. + ) -> Iterator[bytes]: + """Read current source records from a TIMDEX dataset, as bytes. - Each record is parsed by the format-specific `parse_source_record` method. - Currently, format-specific sub-classes XMLTransformer and JSONTransformer - define this for parsing their respective file record formats. - Source-specific transformers will generally just inhereit. + Records are parsed during transformation by the format-specific + `parse_source_record` method. Args: source_dataset: Location of TIMDEX dataset with source records. @@ -609,44 +630,38 @@ def parse_source_records_from_dataset( source=source, action="index", ): - yield cls.parse_source_record(record["source_record"]) + yield record["source_record"] @classmethod - def parse_source_records_from_parquet( + def read_source_records_from_parquet( cls, source: str, source_parquet: str - ) -> Iterator[dict[str, JSON] | Tag]: - """Read source records from parquet file(s) and parse them. + ) -> Iterator[bytes]: + """Read source records from parquet file(s), as bytes. The parquet file(s) must have the TIMDEX records columns 'source_record', - 'source', and 'action', e.g. a Timmy export. A directory is read recursively, - in sorted path order, ignoring non-parquet files like a Timmy 'export.json'. - Only rows for the given source with action 'index' are parsed, in file order. + 'source', and 'action'. A directory is read recursively, in sorted path order, + ignoring non-parquet files. Only rows for the given source with action 'index' + are read, in file order. Args: source: Source label to filter rows by. source_parquet: Local or S3 location of a parquet file or directory. """ filesystem, path = _parse_location(source_parquet) - if filesystem.get_file_info(path).type == fs.FileType.Directory: - filepaths = _list_parquet_files(filesystem, path) - else: - filepaths = [path] - logger.info(f"Reading source records from {len(filepaths)} parquet file(s)") - - for filepath in filepaths: - with filesystem.open_input_file(filepath) as file: - for batch in pq.ParquetFile(file).iter_batches( - columns=["source_record", "source", "action"], - batch_size=PARQUET_DATASET_BATCH_SIZE, - ): - mask = pc.and_( - pc.equal(batch.column("source"), source), - pc.equal(batch.column("action"), "index"), - ) - for source_record in ( - batch.filter(mask).column("source_record").to_pylist() - ): - yield cls.parse_source_record(source_record) + dataset = ds.dataset( + path, + filesystem=filesystem, + format="parquet", + exclude_invalid_files=True, + ) + logger.info(f"Reading source records from {len(dataset.files)} parquet file(s)") + + for batch in dataset.to_batches( + columns=["source_record"], + filter=(pc.field("source") == source) & (pc.field("action") == "index"), + batch_size=PARQUET_DATASET_BATCH_SIZE, + ): + yield from batch.column("source_record").to_pylist() @classmethod @abstractmethod @@ -819,14 +834,3 @@ def _parse_location(location: str) -> tuple[fs.FileSystem, str]: if location.startswith("s3://"): return fs.S3FileSystem(), location.removeprefix("s3://") return fs.LocalFileSystem(), os.path.abspath(location) - - -def _list_parquet_files(filesystem: fs.FileSystem, path: str) -> list[str]: - """Return sorted paths of parquet files under a directory, if it exists.""" - if filesystem.get_file_info(path).type != fs.FileType.Directory: - return [] - return sorted( - file_info.path - for file_info in filesystem.get_file_info(fs.FileSelector(path, recursive=True)) - if file_info.path.endswith(".parquet") - ) diff --git a/transmogrifier/sources/xml/researchdatabases.py b/transmogrifier/sources/xml/researchdatabases.py index cf52100..1ea621a 100644 --- a/transmogrifier/sources/xml/researchdatabases.py +++ b/transmogrifier/sources/xml/researchdatabases.py @@ -2,9 +2,7 @@ from collections.abc import Iterator from datetime import UTC, datetime -import smart_open # type: ignore[import-untyped] from bs4 import BeautifulSoup, Tag -from lxml import etree from timdex_dataset_api import TIMDEXDataset # type: ignore[import-untyped] from transmogrifier import config @@ -20,7 +18,7 @@ class ResearchDatabases(SpringshareOaiDc): @classmethod - def parse_source_file(cls, source_file: str) -> Iterator[Tag]: + def read_source_file(cls, source_file: str) -> Iterator[bytes]: """Yield records from harvested OAI + API gathered records for possible delete. If the following env vars are not set, detecting deleted records will be skipped: @@ -28,27 +26,9 @@ def parse_source_file(cls, source_file: str) -> Iterator[Tag]: - LIBGUIDES_CLIENT_ID - TIMDEX_DATASET_LOCATION """ - yield from cls._yield_oai_xml_records_for_indexing(source_file) - yield from cls._yield_api_records_for_deleting() - - @classmethod - def _yield_oai_xml_records_for_indexing(cls, source_file: str) -> Iterator[Tag]: - """Yield OAI records from extracted XML file. - - This functionality is a direct port from XMLTransformer, but allows for a custom - self.parse_source_file in this transformer class. - """ - with smart_open.open(source_file, "rb") as file: - for _, element in etree.iterparse( - file, - tag="{*}record", - encoding="utf-8", - recover=True, - ): - record_string = etree.tostring(element, encoding="utf-8") - record = cls.parse_bs4_in_isolated_thread(record_string) - yield record - element.clear() + yield from super().read_source_file(source_file) + for record in cls._yield_api_records_for_deleting(): + yield record.encode() @classmethod def _yield_api_records_for_deleting(cls) -> Iterator[Tag]: diff --git a/transmogrifier/sources/xmltransformer.py b/transmogrifier/sources/xmltransformer.py index 65b8269..fa8202d 100644 --- a/transmogrifier/sources/xmltransformer.py +++ b/transmogrifier/sources/xmltransformer.py @@ -17,9 +17,9 @@ class XMLTransformer(Transformer): """XML transformer class.""" @classmethod - def parse_source_file(cls, source_file: str) -> Iterator[Tag]: + def read_source_file(cls, source_file: str) -> Iterator[bytes]: """ - Parse XML file and return source records as bs4 Tags via an iterator. + Read XML file and yield each element as bytes. Args: source_file: A file containing source records to be transformed. @@ -31,9 +31,7 @@ def parse_source_file(cls, source_file: str) -> Iterator[Tag]: encoding="utf-8", recover=True, ): - record_string = etree.tostring(element, encoding="utf-8") - record = cls.parse_bs4_in_isolated_thread(record_string) - yield record + yield etree.tostring(element, encoding="utf-8") element.clear() @classmethod diff --git a/uv.lock b/uv.lock index c2d2ca3..29ecbfe 100644 --- a/uv.lock +++ b/uv.lock @@ -671,18 +671,6 @@ 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.13.0" @@ -1700,7 +1688,6 @@ dependencies = [ { name = "beautifulsoup4" }, { name = "boto3" }, { name = "click" }, - { name = "jsonlines" }, { name = "lxml" }, { name = "pandas" }, { name = "pyarrow" }, @@ -1733,7 +1720,6 @@ requires-dist = [ { name = "beautifulsoup4", specifier = "==4.12.3" }, { name = "boto3" }, { name = "click" }, - { name = "jsonlines" }, { name = "lxml" }, { name = "pandas" }, { name = "pyarrow" }, From 5e4044d741132ae74a11dfeac404533304cca4ee Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 30 Oct 2026 14:28:07 -0400 Subject: [PATCH 03/11] Store original source record bytes Why these changes are being introduced: * Source records are now read as bytes before parsing, so the bytes read can be stored as-is, instead of re-serializing the parsed bs4 Tag or dict. This keeps the stored source record closer to what was harvested, and lets worker processes return only transformed records while the bytes stay in the main process. How this addresses that need: * Transformer.__next__() stores the source record bytes read from the input as DatasetRecord.source_record. * serialize_source_record() is only used when a transformer is given already parsed records. Side effects of this change: * Deliberate change to stored source_record for file inputs: * XML: lxml's serialization of each element, without an added declaration and with original whitespace. This is not a byte slice of the file: a namespace declared on a parent element is copied onto the record, and recover=True can repair malformed XML. * JSON: the original line, without json.dumps() re-escaping. * Dataset and parquet inputs pass the stored bytes through untouched. * One-time churn: the first run after deploy writes new source_record bytes for every XML and JSON file record, so each will show a source change once. * Transformed records are unchanged: fixture files for every source produced identical transformed records, and 2,000 alma records produced identical source and transformed records. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-738 --- tests/sources/test_transformer.py | 20 ++++++++++++++++++++ transmogrifier/sources/transformer.py | 11 ++++++++--- 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/tests/sources/test_transformer.py b/tests/sources/test_transformer.py index 64a890f..dfb0f17 100644 --- a/tests/sources/test_transformer.py +++ b/tests/sources/test_transformer.py @@ -584,3 +584,23 @@ def test_jsontransformer_read_source_file_yields_line_bytes(tmp_path): b'{"id": "1"}', b'{"id": "2"}', ] + + +def test_transformer_next_iter_stores_original_source_record_bytes(source_input_file): + original_records = list(XMLTransformer.read_source_file(source_input_file)) + transformer = Transformer.load("cool-repo", source_input_file) + stored_records = [dataset_record.source_record for dataset_record in transformer] + assert stored_records == original_records + assert not stored_records[0].startswith(b" DatasetRecord: return DatasetRecord( timdex_record_id=timdex_record_id, source_record=( - self.serialize_source_record(source_record) - if source_record is not None - else record + record + if isinstance(record, bytes) + else self.serialize_source_record(record) ), transformed_record=( json.dumps(transformed_record.asdict()).encode() @@ -413,6 +413,11 @@ def _parse_run_data_from_source_file(source_file: str) -> dict: return run_data def serialize_source_record(self, source_record: Tag | dict) -> bytes | None: + """Serialize an already parsed source record, for storing in the dataset. + + Source records read as bytes are stored as-is; this is only used when a + transformer is given already parsed records. + """ if isinstance(source_record, Tag): return source_record.encode() if isinstance(source_record, dict): From d50c9bd5faa26d53c924a091d6d26912d7c662ae Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 30 Oct 2026 14:30:20 -0400 Subject: [PATCH 04/11] Split per-record transform from run bookkeeping Why these changes are being introduced: * To transform records in worker processes, the work done for one record must be separate from updating the transformer's run state, which has to happen in the main process in run order. How this addresses that need: * Add Transformer._process_source_record(), which parses and transforms one source record and returns a small, picklable TransformResult of action, timdex_record_id, and transformed_record. It catches DeletedRecordEvent, SkippedRecordEvent, and unhandled exceptions as actions, lets CriticalError propagate, and changes no transformer state. The run_record_offset for provenance is passed in. * Add Transformer._update_record_counts(), which updates record counts and deleted records from a TransformResult. It is a separate method, rather than part of __next__(), so that the upcoming parallel workers path, which does not go through __next__(), can reuse the same bookkeeping. * Add Transformer._create_dataset_record(), which builds the DatasetRecord from the source record, TransformResult, and run data. * Transformer.__next__() now calls these three in order, and its docstring describes what must stay in the main process. Side effects of this change: * None: fixture files for every source, and 2,000 alma records from a Timmy export, produced identical source and transformed records. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-738 --- tests/sources/test_transformer.py | 69 +++++++++- transmogrifier/sources/transformer.py | 181 ++++++++++++++++++-------- 2 files changed, 194 insertions(+), 56 deletions(-) diff --git a/tests/sources/test_transformer.py b/tests/sources/test_transformer.py index dfb0f17..1508412 100644 --- a/tests/sources/test_transformer.py +++ b/tests/sources/test_transformer.py @@ -1,6 +1,7 @@ # ruff: noqa: SLF001, D202 import datetime import json +import pickle from pathlib import Path from unittest import mock @@ -16,7 +17,11 @@ from transmogrifier.sources import transformer as transformer_module from transmogrifier.sources.json.aardvark import MITAardvark from transmogrifier.sources.jsontransformer import JSONTransformer -from transmogrifier.sources.transformer import PARQUET_FILES_SCHEMA, Transformer +from transmogrifier.sources.transformer import ( + PARQUET_FILES_SCHEMA, + Transformer, + TransformResult, +) from transmogrifier.sources.xml.datacite import Datacite from transmogrifier.sources.xmltransformer import XMLTransformer @@ -604,3 +609,65 @@ def test_transformer_next_iter_stores_original_json_line(tmp_path): source_file.write_bytes(line + b"\n") transformer = Transformer.load("gismit", str(source_file)) assert next(transformer).source_record == line + + +@pytest.mark.parametrize( + ("transform_exception", "expected_action"), + [ + (None, "index"), + (DeletedRecordEvent(timdex_record_id="libguides:guides-0"), "delete"), + (SkippedRecordEvent(source_record_id="guides-0"), "skip"), + (RuntimeError("totally unhandled event"), "error"), + ], +) +def test_transformer_process_source_record_does_not_change_counts( + source_transformer, transform_exception, expected_action +): + record = next(source_transformer.source_records) + if transform_exception: + with mock.patch.object( + source_transformer, "transform", side_effect=transform_exception + ): + result = source_transformer._process_source_record(record, 0) + else: + result = source_transformer._process_source_record(record, 0) + + assert result.action == expected_action + assert pickle.loads(pickle.dumps(result)) == result # noqa: S301 + assert source_transformer.transformed_record_count == 0 + assert source_transformer.deleted_records == [] + assert source_transformer.skipped_record_count == 0 + assert source_transformer.error_record_count == 0 + + +def test_transformer_process_source_record_sets_provenance_run_record_offset( + source_transformer, +): + record = next(source_transformer.source_records) + result = source_transformer._process_source_record(record, 42) + transformed_record = json.loads(result.transformed_record) + assert transformed_record["timdex_provenance"]["run_record_offset"] == 42 # noqa: PLR2004 + + +def test_transformer_update_record_counts(source_transformer): + source_transformer._update_record_counts(TransformResult("index")) + source_transformer._update_record_counts( + TransformResult("delete", timdex_record_id="cool-repo:1") + ) + source_transformer._update_record_counts(TransformResult("skip")) + source_transformer._update_record_counts(TransformResult("error")) + + assert source_transformer.transformed_record_count == 1 + assert source_transformer.deleted_records == ["cool-repo:1"] + assert source_transformer.skipped_record_count == 1 + assert source_transformer.error_record_count == 1 + + +def test_transformer_create_dataset_record_does_not_change_counts(source_transformer): + dataset_record = source_transformer._create_dataset_record( + b"", TransformResult("error") + ) + + assert dataset_record.source_record == b"" + assert dataset_record.action == "error" + assert source_transformer.error_record_count == 0 diff --git a/transmogrifier/sources/transformer.py b/transmogrifier/sources/transformer.py index 8452d76..809e8ac 100644 --- a/transmogrifier/sources/transformer.py +++ b/transmogrifier/sources/transformer.py @@ -19,6 +19,7 @@ import pyarrow.compute as pc import pyarrow.dataset as ds import smart_open +from attrs import define from bs4 import Tag from pyarrow import fs from timdex_dataset_api import ( @@ -71,6 +72,18 @@ ) +@define(frozen=True) +class TransformResult: + """Outcome of transforming a single source record. + + Kept small and picklable, so it can be returned from a worker process. + """ + + action: str + timdex_record_id: str | None = None + transformed_record: bytes | None = None + + class Transformer(ABC): """Base transformer class.""" @@ -170,67 +183,125 @@ def __iter__(self) -> Iterator[DatasetRecord]: @final def __next__(self) -> DatasetRecord: - """Return next transformed record.""" - while True: - source_record = None - transformed_record = None - timdex_record_id = None - - record = next(self.source_records) - self.processed_record_count += 1 - - try: - source_record = ( - self.parse_source_record(record) - if isinstance(record, bytes) - else record - ) - transformed_record = self.transform(source_record) - timdex_record_id = transformed_record.timdex_record_id - transformed_record.timdex_provenance = timdex.TimdexProvenance( - source=self.run_data["source"], - run_date=self.run_data["run_date"], - run_id=self.run_data["run_id"], - run_record_offset=self.run_record_offset, - ) - self.transformed_record_count += 1 - action = "index" + """Return next transformed record. + + Each source record is handled in three steps: + 1. _process_source_record(): parse and transform the record, and determine + its action. This changes no transformer state, so it is safe to run in + another process. + 2. _update_record_counts(): update the run's record counts and deleted + records. + 3. _create_dataset_record(): build the DatasetRecord from the result, the + original source record, and run data. + + Anything that reads or changes run state must stay in this process and run in + order: incrementing processed_record_count (which sets run_record_offset), + updating record counts and deleted records, and building the DatasetRecord. + """ + raw_source_record = next(self.source_records) + self.processed_record_count += 1 + + # parse the raw source record bytes + transform the record + transform_result = self._process_source_record( + raw_source_record, + self.run_record_offset, + ) - except DeletedRecordEvent as error: - self.deleted_records.append(error.timdex_record_id) - timdex_record_id = error.timdex_record_id - action = "delete" + # update bookkeeping + self._update_record_counts(transform_result) - except SkippedRecordEvent: - self.skipped_record_count += 1 - action = "skip" + # return a DatasetRecord instance with source + transformed record + return self._create_dataset_record(raw_source_record, transform_result) - except CriticalError: - raise + @final + def _process_source_record( + self, record: bytes | dict[str, JSON] | Tag, run_record_offset: int + ) -> TransformResult: + """Parse and transform a single source record, and determine its action. - except Exception as exception: - self.error_record_count += 1 - message = f"Unhandled exception during record transformation: {exception}" - logger.exception(message) - action = "error" - - return DatasetRecord( - timdex_record_id=timdex_record_id, - source_record=( - record - if isinstance(record, bytes) - else self.serialize_source_record(record) - ), - transformed_record=( - json.dumps(transformed_record.asdict()).encode() - if transformed_record - else None - ), - action=action, - run_record_offset=self.run_record_offset, - **self.run_data, + This step does not change transformer state (e.g. record counts), so it can run + in a worker process. + + Args: + record: A source record as bytes, or an already parsed source record. + run_record_offset: The record's position in the run, for provenance. + + Raises: + CriticalError: Propagated, as it should stop the run. + """ + try: + source_record = ( + self.parse_source_record(record) if isinstance(record, bytes) else record + ) + timdex_record = self.transform(source_record) + timdex_record.timdex_provenance = timdex.TimdexProvenance( + source=self.run_data["source"], + run_date=self.run_data["run_date"], + run_id=self.run_data["run_id"], + run_record_offset=run_record_offset, + ) + return TransformResult( + action="index", + timdex_record_id=timdex_record.timdex_record_id, + transformed_record=json.dumps(timdex_record.asdict()).encode(), + ) + + except DeletedRecordEvent as error: + return TransformResult( + action="delete", timdex_record_id=error.timdex_record_id ) + except SkippedRecordEvent: + return TransformResult(action="skip") + + except CriticalError: + raise + + except Exception as exception: + message = f"Unhandled exception during record transformation: {exception}" + logger.exception(message) + return TransformResult(action="error") + + @final + def _update_record_counts(self, result: TransformResult) -> None: + """Update the run's record counts and deleted records for a transform result. + + Args: + result: The result of _process_source_record() for a record. + """ + match result.action: + case "index": + self.transformed_record_count += 1 + case "delete": + self.deleted_records.append(str(result.timdex_record_id)) + case "skip": + self.skipped_record_count += 1 + case "error": + self.error_record_count += 1 + + @final + def _create_dataset_record( + self, record: bytes | dict[str, JSON] | Tag, result: TransformResult + ) -> DatasetRecord: + """Create a DatasetRecord from a source record and its transform result. + + Args: + record: The source record as read, stored as-is if bytes. + result: The result of _process_source_record() for this record. + """ + return DatasetRecord( + timdex_record_id=result.timdex_record_id, + source_record=( + record + if isinstance(record, bytes) + else self.serialize_source_record(record) + ), + transformed_record=result.transformed_record, + action=result.action, + run_record_offset=self.run_record_offset, + **self.run_data, + ) + @final @classmethod def get_transformer(cls, source: str) -> type[Transformer]: From bb36604a7c7cf347f00f4572ab99ec8dd3710502 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 30 Oct 2026 14:38:01 -0400 Subject: [PATCH 05/11] Add --workers option to transform with a process pool Why these changes are being introduced: * Alma transforms are CPU bound at ~90 records/sec per process, mostly bs4 querying, which makes full re-transforms of ~1M records slow. How this addresses that need: * Add CLI option --workers (default 1). With 1, records are transformed in the main process exactly as before. * With more than 1, iterating the transformer uses a pool of worker processes. The main process reads source records as bytes, and sends them in chunks of 100, each with its run_record_offset, to workers. Workers return only TransformResults; the source bytes stay in the main process. Results are handled in input order, with the same main process bookkeeping as __next__() (_update_record_counts() and _create_dataset_record()), so output is identical to a single process. * At most 4 chunks per worker are in flight, which bounds memory. * Workers are started with "spawn" explicitly, so behavior is the same on macOS and Linux (e.g. ECS) and does not depend on the platform or Python version default. Each builds its own instance of the transformer class (and so its own crosswalk cache) with the main process's run data, logging, and Sentry configuration. * A CriticalError in any worker stops the run, and pending chunks are cancelled. Side effects of this change: * 100,000 alma records from a Timmy export: 2m28s with --workers 8 (~675 records/sec, vs ~92 in one process), and identical source and transformed records. * Production runs are unchanged unless --workers is passed. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-738 * https://mitlibraries.atlassian.net/browse/TIMX-343 --- tests/sources/test_transformer.py | 56 +++++++++++- tests/test_cli.py | 38 ++++++++ transmogrifier/cli.py | 11 +++ transmogrifier/sources/transformer.py | 122 +++++++++++++++++++++++++- 4 files changed, 224 insertions(+), 3 deletions(-) diff --git a/tests/sources/test_transformer.py b/tests/sources/test_transformer.py index 1508412..89912ce 100644 --- a/tests/sources/test_transformer.py +++ b/tests/sources/test_transformer.py @@ -13,7 +13,11 @@ from timdex_dataset_api import DatasetRecord import transmogrifier.models as timdex -from transmogrifier.exceptions import DeletedRecordEvent, SkippedRecordEvent +from transmogrifier.exceptions import ( + CriticalError, + DeletedRecordEvent, + SkippedRecordEvent, +) from transmogrifier.sources import transformer as transformer_module from transmogrifier.sources.json.aardvark import MITAardvark from transmogrifier.sources.jsontransformer import JSONTransformer @@ -671,3 +675,53 @@ def test_transformer_create_dataset_record_does_not_change_counts(source_transfo assert dataset_record.source_record == b"" assert dataset_record.action == "error" assert source_transformer.error_record_count == 0 + + +class CriticalErrorDatacite(Datacite): + """Module level, so worker processes can import it.""" + + def get_timdex_record_id(self, _source_record): + raise CriticalError("Catastrophic failure in a worker!") + + +def test_transformer_workers_output_matches_single_process(source_input_file, run_id): + def load(workers: int) -> Transformer: + return Transformer.load( + "jpal", + source_input_file, + run_id=run_id, + run_timestamp="2024-06-03T15:30:45", + workers=workers, + ) + + single_process, with_workers = load(1), load(2) + single_process_records = list(single_process) + with_workers_records = list(with_workers) + + assert with_workers_records == single_process_records + assert len(with_workers.deleted_records) == 1 + for attribute in ( + "processed_record_count", + "transformed_record_count", + "skipped_record_count", + "error_record_count", + "deleted_records", + ): + assert getattr(with_workers, attribute) == getattr(single_process, attribute) + + +def test_transformer_workers_critical_error_stops_run(source_input_file): + transformer = CriticalErrorDatacite( + "jpal", + XMLTransformer.read_source_file(source_input_file), + source_file=source_input_file, + workers=2, + ) + with pytest.raises(CriticalError, match="Catastrophic failure in a worker!"): + list(transformer) + + +def test_transformer_workers_require_bytes_source_records(oai_pmh_records): + transformer = XMLTransformer("cool-repo", oai_pmh_records, workers=2) + with pytest.raises(TypeError, match="requires source records as bytes"): + list(transformer) diff --git a/tests/test_cli.py b/tests/test_cli.py index e865724..b8ca274 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -365,3 +365,41 @@ def test_transform_output_file(caplog, runner, source_input_file, tmp_path): assert result.exit_code == 0 assert "deleted records: 1" in caplog.text assert f"File write complete - records: 3, file: {output_file}" in caplog.text + + +def test_transform_workers_option(caplog, runner, source_parquet_location, tmp_path): + caplog.set_level("INFO") + result = runner.invoke( + main, + [ + "-s", + "researchdatabases", + "--input-parquet", + source_parquet_location, + "--output-parquet", + str(tmp_path / "output"), + "--workers", + "2", + ], + ) + assert result.exit_code == 0 + assert "Transforming records with 2 worker processes" in caplog.text + assert "Completed transform, total records processed: 3" in caplog.text + + +def test_transform_workers_option_must_be_positive(runner, tmp_path): + result = runner.invoke( + main, + [ + "-s", + "alma", + "--input-file", + "fake-input-file", + "--output-dataset", + str(tmp_path / "dataset"), + "--workers", + "0", + ], + ) + assert result.exit_code == 2 # noqa: PLR2004 + assert "Invalid value for '-w' / '--workers'" in result.output diff --git a/transmogrifier/cli.py b/transmogrifier/cli.py index 5324d04..a8e71c2 100644 --- a/transmogrifier/cli.py +++ b/transmogrifier/cli.py @@ -90,6 +90,15 @@ "possible for the TIMDEX StepFunction to invoke Transmogrifier multiple times, this " "allows a single run_timestamp to be associated with all outputs for single run_id.", ) +@click.option( + "-w", + "--workers", + type=click.IntRange(min=1), + default=1, + show_default=True, + help="Number of worker processes to transform records with. With 1, records are " + "transformed in the main process. Output is the same for any number of workers.", +) @click.option( "-v", "--verbose", is_flag=True, help="Pass to log at debug level instead of info" ) @@ -104,6 +113,7 @@ def main( exclusion_list_path: str, run_id: str, run_timestamp: str, + workers: int, verbose: bool, # noqa: FBT001 ) -> None: start_time = perf_counter() @@ -133,6 +143,7 @@ def main( exclusion_list_path=exclusion_list_path, run_id=run_id, run_timestamp=run_timestamp, + workers=workers, ) if output_file: transformer.write_to_file(output_file) diff --git a/transmogrifier/sources/transformer.py b/transmogrifier/sources/transformer.py index 809e8ac..02e4adb 100644 --- a/transmogrifier/sources/transformer.py +++ b/transmogrifier/sources/transformer.py @@ -7,10 +7,13 @@ import itertools import json import logging +import multiprocessing import os import re import uuid from abc import ABC, abstractmethod +from collections import deque +from concurrent.futures import Future, ProcessPoolExecutor from datetime import UTC, datetime from importlib import import_module from typing import TYPE_CHECKING, final @@ -28,7 +31,7 @@ ) import transmogrifier.models as timdex -from transmogrifier.config import SOURCES +from transmogrifier.config import SOURCES, configure_logger, configure_sentry from transmogrifier.exceptions import ( CriticalError, DeletedRecordEvent, @@ -66,11 +69,17 @@ PARQUET_FILES_ROW_GROUP_SIZE = 5_000 PARQUET_FILES_COMPRESSION = "zstd" +WORKER_CHUNK_SIZE = 100 +WORKER_MAX_PENDING_CHUNKS_PER_WORKER = 4 + INPUT_SOURCE_ERROR_MESSAGE = ( "Exactly one input source is required: 'source_file', 'source_dataset', " "or 'source_parquet'" ) +# set once per worker process by _init_worker(), and used by _transform_in_worker() +_worker_transformer: Transformer | None = None + @define(frozen=True) class TransformResult: @@ -97,6 +106,7 @@ def __init__( source_parquet: str | None = None, run_id: str | None = None, run_timestamp: str | None = None, + workers: int = 1, ) -> None: """ Initialize Transformer instance. @@ -112,6 +122,8 @@ def __init__( source_parquet: Location of parquet file(s) with input source records. run_id: A unique identifier associated with this ETL run. run_timestamp: A timestamp associated with this ETL run. + workers: Number of worker processes to transform records with. With 1, + records are transformed in this process. """ self.source: str = source self.exclusion_list_path: str | None = exclusion_list_path @@ -127,6 +139,7 @@ def __init__( self.source_file = source_file self.source_dataset = source_dataset self.source_parquet = source_parquet + self.workers = workers self.run_data = self.get_run_data( source=source, @@ -178,9 +191,84 @@ def load_exclusion_list(self) -> list[str]: @final def __iter__(self) -> Iterator[DatasetRecord]: - """Iterate over transformed records.""" + """Iterate over transformed records, using worker processes if configured.""" + if self.workers > 1: + return self._iter_with_workers() return self + @final + def _iter_with_workers(self) -> Iterator[DatasetRecord]: + """Yield transformed records, transformed by a pool of worker processes. + + The main process reads source records as bytes and sends them to workers in + chunks, each with its run_record_offset. Workers return only TransformResults; + the main process keeps the source record bytes, and handles results in input + order. Output is therefore identical to transforming in a single process. + + Workers are started fresh, each building its own instance of this transformer's + class, with this transformer's run data. A CriticalError in any worker stops the + run. + """ + logger.info(f"Transforming records with {self.workers} worker processes") + executor = ProcessPoolExecutor( + max_workers=self.workers, + mp_context=multiprocessing.get_context("spawn"), + initializer=_init_worker, + initargs=( + type(self), + self.source, + self.exclusion_list_path, + self.run_data, + logging.getLogger().getEffectiveLevel() == logging.DEBUG, + ), + ) + max_pending_chunks = self.workers * WORKER_MAX_PENDING_CHUNKS_PER_WORKER + pending: deque[tuple[list[bytes], Future[list[TransformResult]]]] = deque() + offsets = itertools.count(self.processed_record_count) + try: + for chunk in itertools.batched( + self.source_records, WORKER_CHUNK_SIZE, strict=False + ): + records: list[bytes] = [] + for record in chunk: + if not isinstance(record, bytes): + raise TypeError( + "Transforming with workers requires source records as bytes" + ) + records.append(record) + pending.append( + ( + records, + executor.submit( + _transform_in_worker, + [(record, next(offsets)) for record in records], + ), + ) + ) + if len(pending) >= max_pending_chunks: + yield from self._create_dataset_records(*pending.popleft()) + while pending: + yield from self._create_dataset_records(*pending.popleft()) + finally: + executor.shutdown(wait=True, cancel_futures=True) + + def _create_dataset_records( + self, records: list[bytes], future: Future[list[TransformResult]] + ) -> Iterator[DatasetRecord]: + """Yield DatasetRecords for a chunk of records transformed by a worker. + + Waits for the worker's results, then does the main process bookkeeping for each + record in run order, as __next__() does when not using workers. + + Args: + records: The chunk's source records, as sent to the worker. + future: The worker's pending TransformResults for the chunk, in order. + """ + for record, result in zip(records, future.result(), strict=True): + self.processed_record_count += 1 + self._update_record_counts(result) + yield self._create_dataset_record(record, result) + @final def __next__(self) -> DatasetRecord: """Return next transformed record. @@ -329,6 +417,7 @@ def load( exclusion_list_path: str | None = None, run_id: str | None = None, run_timestamp: str | None = None, + workers: int = 1, ) -> Transformer: """ Instantiate specified transformer class and populate with source records. @@ -341,6 +430,7 @@ def load( exclusion_list_path: CSV filepath to use for explicitly skipping records. run_id: A unique identifier associated with this ETL run. run_timestamp: A timestamp associated with this ETL run. + workers: Number of worker processes to transform records with. """ transformer_class = cls.get_transformer(source) @@ -369,6 +459,7 @@ def load( source_parquet=source_parquet, run_id=run_id, run_timestamp=run_timestamp, + workers=workers, ) @staticmethod @@ -910,3 +1001,30 @@ def _parse_location(location: str) -> tuple[fs.FileSystem, str]: if location.startswith("s3://"): return fs.S3FileSystem(), location.removeprefix("s3://") return fs.LocalFileSystem(), os.path.abspath(location) + + +def _init_worker( + transformer_class: type[Transformer], + source: str, + exclusion_list_path: str | None, + run_data: dict, + verbose: bool, # noqa: FBT001 +) -> None: + """Build this worker process's transformer, with the main process's run data.""" + global _worker_transformer # noqa: PLW0603 + configure_logger(logging.getLogger(), verbose=verbose) + configure_sentry() + _worker_transformer = transformer_class( + source, iter(()), exclusion_list_path=exclusion_list_path + ) + _worker_transformer.run_data = run_data + + +def _transform_in_worker(chunk: list[tuple[bytes, int]]) -> list[TransformResult]: + """Transform a chunk of (source record, run_record_offset) in a worker process.""" + if _worker_transformer is None: + raise RuntimeError("Worker transformer was not initialized") + return [ + _worker_transformer._process_source_record(record, run_record_offset) # noqa: SLF001 + for record, run_record_offset in chunk + ] From a630427c27cc87feba23e744484e69a4f88cf5e3 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Fri, 30 Oct 2026 14:38:21 -0400 Subject: [PATCH 06/11] Document workers and stored source records Why these changes are being introduced: * The README did not describe the --workers option, or that source records are now stored as read. How this addresses that need: * Describe how source records are stored for each input. * Add a workers and performance section, and refresh the transform --help output. Side effects of this change: * None. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-738 --- README.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/README.md b/README.md index d48b7ee..317f1ee 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,14 @@ Exactly one input option and exactly one output option are required. | `--output-dataset` | All records to a TIMDEX dataset. | | `--output-parquet` | All records to parquet files. | +Source records are stored in the output as they were read: the bytes of each XML `` element (as serialized by lxml), each JSONLines line, or the `source_record` of each dataset or parquet row. + +#### Workers and performance + +`--workers N` transforms records in `N` worker processes. The main process reads source records and writes output; workers parse and transform. Records are handled in input order, so output is identical for any number of workers. The default, `1`, transforms in the main process. + +Workers help most for sources with slow, CPU-bound transforms. Each worker uses roughly 0.5 GB of memory. JSON sources are already fast, and may not benefit. + ```text Usage: transform [OPTIONS] @@ -124,6 +132,11 @@ Options: allows a single run_timestamp to be associated with all outputs for single run_id. + -w, --workers INTEGER RANGE Number of worker processes to transform + records with. With 1, records are + transformed in the main process. Output is + the same for any number of workers. + [default: 1; x>=1] -v, --verbose Pass to log at debug level instead of info --help Show this message and exit. ``` \ No newline at end of file From 42c14f26f93735c2526f78f3a649510071c88c14 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Tue, 3 Nov 2026 14:57:30 -0500 Subject: [PATCH 07/11] Reuse one worker iterator per transformer Why these changes are being introduced: * With workers, each iter() call created a new worker iterator, and next() bypassed it. Partially consuming one iterator, then iterating again or calling next(), dropped records already sent to workers. Unlikely, as the write methods consume the transformer once, but a good edge case to guard against, e.g. if used as a library. How this addresses that need: * Transformer.__iter__() creates the worker iterator once and reuses it. * Transformer.__next__() takes the next record from it when using workers. Side effects of this change: * None. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-738 --- tests/sources/test_transformer.py | 21 +++++++++++++++++++++ transmogrifier/sources/transformer.py | 17 +++++++++++++++-- 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/tests/sources/test_transformer.py b/tests/sources/test_transformer.py index 89912ce..594f2e6 100644 --- a/tests/sources/test_transformer.py +++ b/tests/sources/test_transformer.py @@ -710,6 +710,27 @@ def load(workers: int) -> Transformer: assert getattr(with_workers, attribute) == getattr(single_process, attribute) +def test_transformer_workers_mixed_next_and_iteration_keeps_all_records( + source_input_file, run_id +): + def load(workers: int) -> Transformer: + return Transformer.load( + "jpal", + source_input_file, + run_id=run_id, + run_timestamp="2024-06-03T15:30:45", + workers=workers, + ) + + single_process_records = list(load(1)) + with_workers = load(2) + with_workers_records = [next(with_workers), next(iter(with_workers))] + with_workers_records.extend(list(with_workers)) + + assert with_workers_records == single_process_records + assert with_workers.processed_record_count == len(single_process_records) + + def test_transformer_workers_critical_error_stops_run(source_input_file): transformer = CriticalErrorDatacite( "jpal", diff --git a/transmogrifier/sources/transformer.py b/transmogrifier/sources/transformer.py index 02e4adb..e330a06 100644 --- a/transmogrifier/sources/transformer.py +++ b/transmogrifier/sources/transformer.py @@ -140,6 +140,7 @@ def __init__( self.source_dataset = source_dataset self.source_parquet = source_parquet self.workers = workers + self._worker_iterator: Iterator[DatasetRecord] | None = None self.run_data = self.get_run_data( source=source, @@ -191,9 +192,16 @@ def load_exclusion_list(self) -> list[str]: @final def __iter__(self) -> Iterator[DatasetRecord]: - """Iterate over transformed records, using worker processes if configured.""" + """Iterate over transformed records, using worker processes if configured. + + With workers, a single worker iterator is kept and reused, so that iterating + more than once, or mixing iteration with next(), resumes where it left off + instead of dropping records already sent to workers. + """ if self.workers > 1: - return self._iter_with_workers() + if self._worker_iterator is None: + self._worker_iterator = self._iter_with_workers() + return self._worker_iterator return self @final @@ -285,7 +293,12 @@ def __next__(self) -> DatasetRecord: Anything that reads or changes run state must stay in this process and run in order: incrementing processed_record_count (which sets run_record_offset), updating record counts and deleted records, and building the DatasetRecord. + + With workers, the next record is taken from the shared worker iterator. """ + if self.workers > 1: + return next(iter(self)) + raw_source_record = next(self.source_records) self.processed_record_count += 1 From 42e1d8ae03c69981e86741cc5c114a696bd69c2a Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Tue, 3 Nov 2026 15:01:53 -0500 Subject: [PATCH 08/11] Test worker output across multiple chunks Why these changes are being introduced: * The worker parity test used 4 records with a chunk size of 100, so only one chunk was submitted, and draining pending chunks at the bound was never exercised. How this addresses that need: * Patch WORKER_CHUNK_SIZE and WORKER_MAX_PENDING_CHUNKS_PER_WORKER to 1 in the parity test, so records span 4 chunks, 2 are in flight, and pending chunks are drained. Order, provenance offsets, and counts are compared against a single process as before. Side effects of this change: * None. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-738 --- tests/sources/test_transformer.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/sources/test_transformer.py b/tests/sources/test_transformer.py index 594f2e6..cbfa86e 100644 --- a/tests/sources/test_transformer.py +++ b/tests/sources/test_transformer.py @@ -684,7 +684,13 @@ def get_timdex_record_id(self, _source_record): raise CriticalError("Catastrophic failure in a worker!") -def test_transformer_workers_output_matches_single_process(source_input_file, run_id): +def test_transformer_workers_output_matches_single_process( + monkeypatch, source_input_file, run_id +): + # one record per chunk and a small pending limit, to span chunks and drain pending + monkeypatch.setattr(transformer_module, "WORKER_CHUNK_SIZE", 1) + monkeypatch.setattr(transformer_module, "WORKER_MAX_PENDING_CHUNKS_PER_WORKER", 1) + def load(workers: int) -> Transformer: return Transformer.load( "jpal", From 26d930e07cd22bb6872443aa0132e76537abe5b0 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Tue, 3 Nov 2026 15:31:23 -0500 Subject: [PATCH 09/11] Exclude element tail when reading XML source records Why these changes are being introduced: * etree.tostring() includes an element's tail by default. Whether iterparse has parsed the text following a when yielding it depends on read buffer boundaries, so some stored source records ended with the whitespace between records and others did not. How this addresses that need: * XMLTransformer.read_source_file() serializes each record with with_tail=False. Side effects of this change: * Stored source records no longer end with inconsistent trailing whitespace. Transformed records are unchanged. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-738 --- transmogrifier/sources/xmltransformer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transmogrifier/sources/xmltransformer.py b/transmogrifier/sources/xmltransformer.py index fa8202d..2005b6d 100644 --- a/transmogrifier/sources/xmltransformer.py +++ b/transmogrifier/sources/xmltransformer.py @@ -31,7 +31,7 @@ def read_source_file(cls, source_file: str) -> Iterator[bytes]: encoding="utf-8", recover=True, ): - yield etree.tostring(element, encoding="utf-8") + yield etree.tostring(element, encoding="utf-8", with_tail=False) element.clear() @classmethod From b44ff01301d4126f17879e68629817fd297ed8c3 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Wed, 4 Nov 2026 14:11:00 -0500 Subject: [PATCH 10/11] Improve worker docstrings and variable naming Why these changes are being introduced: * Code review noted that the placeholder source records passed to worker transformers were not explained. How this addresses that need: * Expand worker docstrings and comments. * Rename _init_worker() to _init_transformer_class_for_worker(), and the chunk argument of _transform_in_worker() to source_records_chunk. Side effects of this change: * None. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-738 --- transmogrifier/sources/transformer.py | 55 ++++++++++++++++++++------- 1 file changed, 42 insertions(+), 13 deletions(-) diff --git a/transmogrifier/sources/transformer.py b/transmogrifier/sources/transformer.py index e330a06..5d2cc09 100644 --- a/transmogrifier/sources/transformer.py +++ b/transmogrifier/sources/transformer.py @@ -77,7 +77,8 @@ "or 'source_parquet'" ) -# set once per worker process by _init_worker(), and used by _transform_in_worker() +# set once per worker process by _init_transformer_class_for_worker(), and then used +# by _transform_in_worker() _worker_transformer: Transformer | None = None @@ -208,20 +209,30 @@ def __iter__(self) -> Iterator[DatasetRecord]: def _iter_with_workers(self) -> Iterator[DatasetRecord]: """Yield transformed records, transformed by a pool of worker processes. - The main process reads source records as bytes and sends them to workers in - chunks, each with its run_record_offset. Workers return only TransformResults; - the main process keeps the source record bytes, and handles results in input - order. Output is therefore identical to transforming in a single process. + The main process reads source records as bytes and sends them to workers for + transformation. Workers return only TransformResults; the main process keeps the + source record bytes, and handles results in input order. Output is therefore + identical to transforming in a single process. Workers are started fresh, each building its own instance of this transformer's - class, with this transformer's run data. A CriticalError in any worker stops the - run. + class, with this transformer's run data (source, run_date, run_type, etc.). This + is done by _init_transformer_class_for_worker(), which runs once in each worker + process as it starts. Each worker then reuses that instance for every chunk. + + Each chunk is a list of (, ) tuples, passed + directly to _transform_in_worker(). Workers therefore never read their instance's + source_records, which is why _init_transformer_class_for_worker() passes an empty + iterator. + + A CriticalError in any worker stops the run. """ logger.info(f"Transforming records with {self.workers} worker processes") + + # prepare an executor that will manage workers executor = ProcessPoolExecutor( max_workers=self.workers, mp_context=multiprocessing.get_context("spawn"), - initializer=_init_worker, + initializer=_init_transformer_class_for_worker, initargs=( type(self), self.source, @@ -230,6 +241,8 @@ def _iter_with_workers(self) -> Iterator[DatasetRecord]: logging.getLogger().getEffectiveLevel() == logging.DEBUG, ), ) + + # break input source records into chunks and feed to workers max_pending_chunks = self.workers * WORKER_MAX_PENDING_CHUNKS_PER_WORKER pending: deque[tuple[list[bytes], Future[list[TransformResult]]]] = deque() offsets = itertools.count(self.processed_record_count) @@ -257,6 +270,8 @@ def _iter_with_workers(self) -> Iterator[DatasetRecord]: yield from self._create_dataset_records(*pending.popleft()) while pending: yield from self._create_dataset_records(*pending.popleft()) + + # regardless of success or failure, shutdown the executor and free processes finally: executor.shutdown(wait=True, cancel_futures=True) @@ -1016,14 +1031,20 @@ def _parse_location(location: str) -> tuple[fs.FileSystem, str]: return fs.LocalFileSystem(), os.path.abspath(location) -def _init_worker( +def _init_transformer_class_for_worker( transformer_class: type[Transformer], source: str, exclusion_list_path: str | None, run_data: dict, verbose: bool, # noqa: FBT001 ) -> None: - """Build this worker process's transformer, with the main process's run data.""" + """Build this worker process's transformer, with the main process's run data. + + The transformer class is initialized with an empty iterator in place of its source + records, because it uses the function _transform_in_worker() for transformation, + which receives a custom chunk of records to process. + .source_records is never used. + """ global _worker_transformer # noqa: PLW0603 configure_logger(logging.getLogger(), verbose=verbose) configure_sentry() @@ -1033,11 +1054,19 @@ def _init_worker( _worker_transformer.run_data = run_data -def _transform_in_worker(chunk: list[tuple[bytes, int]]) -> list[TransformResult]: - """Transform a chunk of (source record, run_record_offset) in a worker process.""" +def _transform_in_worker( + source_records_chunk: list[tuple[bytes, int]], +) -> list[TransformResult]: + """Transform one chunk of source records in a worker process. + + Args: + source_records_chunk: One chunk of the run's source records, sent by the main + process in run order, as tuples of (source record bytes, run_record_offset). + The offset is the record's position in the whole run, not in this chunk. + """ if _worker_transformer is None: raise RuntimeError("Worker transformer was not initialized") return [ _worker_transformer._process_source_record(record, run_record_offset) # noqa: SLF001 - for record, run_record_offset in chunk + for record, run_record_offset in source_records_chunk ] From e6a0b82bd510b2fb452787f85060e7571c779f69 Mon Sep 17 00:00:00 2001 From: Graham Hukill Date: Wed, 4 Nov 2026 16:55:41 -0500 Subject: [PATCH 11/11] Move workers out of Transformer into a parallel module Why these changes are being introduced: * Worker logic in Transformer required splitting __next__() into several methods, a TransformResult class, and a separate iteration path. This made Transformer harder to follow, even when running in a single process. How this addresses that need: * Restore Transformer to its state before workers, with all transform logic in __next__(). Its only addition is _transformed_records, what __iter__() returns, which defaults to the instance itself. * New parallel module, used by the CLI when --workers is greater than 1. The main process reads chunks of source records and sends them to workers. Each worker sets a chunk and its starting offset on its own transformer, and iterates it as usual. Results are yielded in input order and tallied on the main transformer, whose write methods are unchanged. Side effects of this change: * Workers now return whole DatasetRecords, including source record bytes. * 13,941 alma records from a Timmy export, 8 workers: 29.0s, vs 28.4s for the previous worker implementation and 164.2s in one process. Output is identical for all three. * Transformer.load() no longer accepts workers. Relevant ticket(s): * https://mitlibraries.atlassian.net/browse/TIMX-738 --- tests/sources/test_transformer.py | 152 +---------- tests/test_parallel.py | 103 ++++++++ transmogrifier/cli.py | 7 +- transmogrifier/parallel.py | 163 ++++++++++++ transmogrifier/sources/transformer.py | 348 +++++--------------------- 5 files changed, 335 insertions(+), 438 deletions(-) create mode 100644 tests/test_parallel.py create mode 100644 transmogrifier/parallel.py diff --git a/tests/sources/test_transformer.py b/tests/sources/test_transformer.py index cbfa86e..dfb0f17 100644 --- a/tests/sources/test_transformer.py +++ b/tests/sources/test_transformer.py @@ -1,7 +1,6 @@ # ruff: noqa: SLF001, D202 import datetime import json -import pickle from pathlib import Path from unittest import mock @@ -13,19 +12,11 @@ from timdex_dataset_api import DatasetRecord import transmogrifier.models as timdex -from transmogrifier.exceptions import ( - CriticalError, - DeletedRecordEvent, - SkippedRecordEvent, -) +from transmogrifier.exceptions import DeletedRecordEvent, SkippedRecordEvent from transmogrifier.sources import transformer as transformer_module from transmogrifier.sources.json.aardvark import MITAardvark from transmogrifier.sources.jsontransformer import JSONTransformer -from transmogrifier.sources.transformer import ( - PARQUET_FILES_SCHEMA, - Transformer, - TransformResult, -) +from transmogrifier.sources.transformer import PARQUET_FILES_SCHEMA, Transformer from transmogrifier.sources.xml.datacite import Datacite from transmogrifier.sources.xmltransformer import XMLTransformer @@ -613,142 +604,3 @@ def test_transformer_next_iter_stores_original_json_line(tmp_path): source_file.write_bytes(line + b"\n") transformer = Transformer.load("gismit", str(source_file)) assert next(transformer).source_record == line - - -@pytest.mark.parametrize( - ("transform_exception", "expected_action"), - [ - (None, "index"), - (DeletedRecordEvent(timdex_record_id="libguides:guides-0"), "delete"), - (SkippedRecordEvent(source_record_id="guides-0"), "skip"), - (RuntimeError("totally unhandled event"), "error"), - ], -) -def test_transformer_process_source_record_does_not_change_counts( - source_transformer, transform_exception, expected_action -): - record = next(source_transformer.source_records) - if transform_exception: - with mock.patch.object( - source_transformer, "transform", side_effect=transform_exception - ): - result = source_transformer._process_source_record(record, 0) - else: - result = source_transformer._process_source_record(record, 0) - - assert result.action == expected_action - assert pickle.loads(pickle.dumps(result)) == result # noqa: S301 - assert source_transformer.transformed_record_count == 0 - assert source_transformer.deleted_records == [] - assert source_transformer.skipped_record_count == 0 - assert source_transformer.error_record_count == 0 - - -def test_transformer_process_source_record_sets_provenance_run_record_offset( - source_transformer, -): - record = next(source_transformer.source_records) - result = source_transformer._process_source_record(record, 42) - transformed_record = json.loads(result.transformed_record) - assert transformed_record["timdex_provenance"]["run_record_offset"] == 42 # noqa: PLR2004 - - -def test_transformer_update_record_counts(source_transformer): - source_transformer._update_record_counts(TransformResult("index")) - source_transformer._update_record_counts( - TransformResult("delete", timdex_record_id="cool-repo:1") - ) - source_transformer._update_record_counts(TransformResult("skip")) - source_transformer._update_record_counts(TransformResult("error")) - - assert source_transformer.transformed_record_count == 1 - assert source_transformer.deleted_records == ["cool-repo:1"] - assert source_transformer.skipped_record_count == 1 - assert source_transformer.error_record_count == 1 - - -def test_transformer_create_dataset_record_does_not_change_counts(source_transformer): - dataset_record = source_transformer._create_dataset_record( - b"", TransformResult("error") - ) - - assert dataset_record.source_record == b"" - assert dataset_record.action == "error" - assert source_transformer.error_record_count == 0 - - -class CriticalErrorDatacite(Datacite): - """Module level, so worker processes can import it.""" - - def get_timdex_record_id(self, _source_record): - raise CriticalError("Catastrophic failure in a worker!") - - -def test_transformer_workers_output_matches_single_process( - monkeypatch, source_input_file, run_id -): - # one record per chunk and a small pending limit, to span chunks and drain pending - monkeypatch.setattr(transformer_module, "WORKER_CHUNK_SIZE", 1) - monkeypatch.setattr(transformer_module, "WORKER_MAX_PENDING_CHUNKS_PER_WORKER", 1) - - def load(workers: int) -> Transformer: - return Transformer.load( - "jpal", - source_input_file, - run_id=run_id, - run_timestamp="2024-06-03T15:30:45", - workers=workers, - ) - - single_process, with_workers = load(1), load(2) - single_process_records = list(single_process) - with_workers_records = list(with_workers) - - assert with_workers_records == single_process_records - assert len(with_workers.deleted_records) == 1 - for attribute in ( - "processed_record_count", - "transformed_record_count", - "skipped_record_count", - "error_record_count", - "deleted_records", - ): - assert getattr(with_workers, attribute) == getattr(single_process, attribute) - - -def test_transformer_workers_mixed_next_and_iteration_keeps_all_records( - source_input_file, run_id -): - def load(workers: int) -> Transformer: - return Transformer.load( - "jpal", - source_input_file, - run_id=run_id, - run_timestamp="2024-06-03T15:30:45", - workers=workers, - ) - - single_process_records = list(load(1)) - with_workers = load(2) - with_workers_records = [next(with_workers), next(iter(with_workers))] - with_workers_records.extend(list(with_workers)) - - assert with_workers_records == single_process_records - assert with_workers.processed_record_count == len(single_process_records) - - -def test_transformer_workers_critical_error_stops_run(source_input_file): - transformer = CriticalErrorDatacite( - "jpal", - XMLTransformer.read_source_file(source_input_file), - source_file=source_input_file, - workers=2, - ) - with pytest.raises(CriticalError, match="Catastrophic failure in a worker!"): - list(transformer) - - -def test_transformer_workers_require_bytes_source_records(oai_pmh_records): - transformer = XMLTransformer("cool-repo", oai_pmh_records, workers=2) - with pytest.raises(TypeError, match="requires source records as bytes"): - list(transformer) diff --git a/tests/test_parallel.py b/tests/test_parallel.py new file mode 100644 index 0000000..020fba1 --- /dev/null +++ b/tests/test_parallel.py @@ -0,0 +1,103 @@ +import os +from concurrent.futures.process import BrokenProcessPool + +import pytest + +from transmogrifier import parallel +from transmogrifier.exceptions import CriticalError +from transmogrifier.parallel import iter_transformed_records +from transmogrifier.sources.transformer import Transformer +from transmogrifier.sources.xml.datacite import Datacite +from transmogrifier.sources.xmltransformer import XMLTransformer + +COUNT_ATTRIBUTES = ( + "processed_record_count", + "transformed_record_count", + "skipped_record_count", + "error_record_count", + "deleted_records", +) + + +class CriticalErrorDatacite(Datacite): + """Module level, so worker processes can import it.""" + + def get_timdex_record_id(self, _source_record): + raise CriticalError("Catastrophic failure in a worker!") + + +class DyingWorkerDatacite(Datacite): + """Module level, so worker processes can import it.""" + + def get_timdex_record_id(self, _source_record): + os._exit(1) + + +@pytest.fixture +def load_transformer(source_input_file, run_id): + def _load(): + return Transformer.load( + "jpal", + source_input_file, + run_id=run_id, + run_timestamp="2024-06-03T15:30:45", + ) + + return _load + + +@pytest.mark.parametrize( + ("chunk_size", "max_pending_chunks_per_worker"), + [(1, 1), (2, 1), (100, 4)], +) +def test_iter_transformed_records_matches_single_process( + monkeypatch, load_transformer, chunk_size, max_pending_chunks_per_worker +): + monkeypatch.setattr(parallel, "WORKER_CHUNK_SIZE", chunk_size) + monkeypatch.setattr( + parallel, "WORKER_MAX_PENDING_CHUNKS_PER_WORKER", max_pending_chunks_per_worker + ) + single_process = load_transformer() + single_process_records = list(single_process) + + with_workers = load_transformer() + with_workers_records = list(iter_transformed_records(with_workers, workers=2)) + + assert with_workers_records == single_process_records + assert len(with_workers.deleted_records) == 1 + for attribute in COUNT_ATTRIBUTES: + assert getattr(with_workers, attribute) == getattr(single_process, attribute) + + +def test_iter_transformed_records_critical_error_stops_run(source_input_file): + transformer = CriticalErrorDatacite( + "jpal", + XMLTransformer.read_source_file(source_input_file), + source_file=source_input_file, + ) + with pytest.raises(CriticalError, match="Catastrophic failure in a worker!"): + list(iter_transformed_records(transformer, workers=2)) + + +def test_iter_transformed_records_worker_process_death_stops_run(source_input_file): + transformer = DyingWorkerDatacite( + "jpal", + XMLTransformer.read_source_file(source_input_file), + source_file=source_input_file, + ) + with pytest.raises(BrokenProcessPool): + list(iter_transformed_records(transformer, workers=2)) + + +def test_transformer_iterates_provided_transformed_records(load_transformer): + transformer = load_transformer() + provided_records = iter(["a", "b"]) + transformer._transformed_records = provided_records # noqa: SLF001 + + assert iter(transformer) is provided_records + assert list(transformer) == ["a", "b"] + + +def test_transformer_iterates_itself_by_default(load_transformer): + transformer = load_transformer() + assert iter(transformer) is transformer diff --git a/transmogrifier/cli.py b/transmogrifier/cli.py index a8e71c2..411c69a 100644 --- a/transmogrifier/cli.py +++ b/transmogrifier/cli.py @@ -5,6 +5,7 @@ import click from transmogrifier.config import SOURCES, configure_logger, configure_sentry +from transmogrifier.parallel import iter_transformed_records from transmogrifier.sources.transformer import Transformer logger = logging.getLogger(__name__) @@ -143,8 +144,12 @@ def main( exclusion_list_path=exclusion_list_path, run_id=run_id, run_timestamp=run_timestamp, - workers=workers, ) + if workers > 1: + transformer._transformed_records = iter_transformed_records( # noqa: SLF001 + transformer, workers + ) + if output_file: transformer.write_to_file(output_file) elif output_dataset: diff --git a/transmogrifier/parallel.py b/transmogrifier/parallel.py new file mode 100644 index 0000000..274ea8c --- /dev/null +++ b/transmogrifier/parallel.py @@ -0,0 +1,163 @@ +"""Transform records with a pool of worker processes. + +Transforming records can be CPU heavy, and a Python process effectively uses one CPU core +at a time. When the CLI is run with more than one worker, this module spreads that work +across multiple processes, and keeps multiprocessing concerns out of the Transformer +classes. + +Transformers are used almost untouched: + - Each worker process has its own transformer, which iterates over source records as + usual, but is given one chunk of them at a time. + - The main process transformer reads the source records and writes the output as + usual, but its __iter__() returns the records coming back from the workers, by way of + _transformed_records, instead of transforming them itself. +""" + +from __future__ import annotations + +import itertools +import logging +import multiprocessing +from collections import deque +from concurrent.futures import ProcessPoolExecutor +from typing import TYPE_CHECKING + +from transmogrifier.config import configure_logger, configure_sentry + +if TYPE_CHECKING: + from collections.abc import Iterable, Iterator + from concurrent.futures import Future + + from bs4 import Tag + from timdex_dataset_api import DatasetRecord + + from transmogrifier.sources.transformer import JSON, Transformer + +logger = logging.getLogger(__name__) + +WORKER_CHUNK_SIZE = 100 +WORKER_MAX_PENDING_CHUNKS_PER_WORKER = 4 + +# set once per worker process by _init_worker(), and used by _transform_chunk() +_worker_transformer: Transformer | None = None + + +def iter_transformed_records( + transformer: Transformer, workers: int +) -> Iterator[DatasetRecord]: + """Yield DatasetRecords for a transformer's source records, using worker processes. + + The main process reads chunks of the transformer's source records and sends them to + workers. Each worker has its own transformer that iterates over one chunk at a time, + as it would for a whole run, and returns the resulting DatasetRecords. + + Records are yielded in input order, and only a limited number of chunks are in flight + at a time. Chunks are only read as records are consumed, which limits memory use. + + The record counts and deleted records of the passed transformer are updated as records + are yielded. A CriticalError in any worker, or a worker process dying, stops the run. + + Args: + transformer: The main process transformer, whose source records are read. + workers: The number of worker processes to use. + """ + logger.info(f"Transforming records with {workers} worker processes") + # "spawn" starts each worker as a fresh Python process, instead of a copy of this one, + # so each worker runs _init_worker() once at startup to set itself up + executor = ProcessPoolExecutor( + max_workers=workers, + mp_context=multiprocessing.get_context("spawn"), + initializer=_init_worker, + initargs=( + type(transformer), + transformer.source, + transformer.exclusion_list_path, + transformer.run_data, + logging.getLogger().getEffectiveLevel() == logging.DEBUG, + ), + ) + + # the upper limit of chunks in flight, which are chunks sent to workers but not yet + # yielded here; this keeps a large input from piling up in memory + max_pending_chunks = workers * WORKER_MAX_PENDING_CHUNKS_PER_WORKER + + pending: deque[Future[list[DatasetRecord]]] = deque() + start_offset = 0 + try: + for chunk in itertools.batched( + transformer.source_records, WORKER_CHUNK_SIZE, strict=False + ): + pending.append(executor.submit(_transform_chunk, start_offset, list(chunk))) + start_offset += len(chunk) + + # at the limit, stop reading source records until the oldest chunk is done + # - workers can finish chunks in any order, but always taking the oldest + # chunk means records are yielded in input order + # - result() blocks until the chunk is done, and re-raises any exception + # from the worker + if len(pending) >= max_pending_chunks: + yield from _tally_run_counts(transformer, pending.popleft().result()) + + while pending: + yield from _tally_run_counts(transformer, pending.popleft().result()) + finally: + executor.shutdown(wait=True, cancel_futures=True) + + +def _tally_run_counts( + transformer: Transformer, records: Iterable[DatasetRecord] +) -> Iterator[DatasetRecord]: + """Update a transformer's record counts and deleted records, and yield each record.""" + for record in records: + transformer.processed_record_count += 1 + match record.action: + case "index": + transformer.transformed_record_count += 1 + case "delete": + transformer.deleted_records.append(record.timdex_record_id) + case "skip": + transformer.skipped_record_count += 1 + case "error": + transformer.error_record_count += 1 + yield record + + +def _init_worker( + transformer_class: type[Transformer], + source: str, + exclusion_list_path: str | None, + run_data: dict, + verbose: bool, # noqa: FBT001 +) -> None: + """Configure this worker process and build its transformer. + + The transformer starts with no source records, iter(()), because _transform_chunk() + sets them for each chunk. + """ + global _worker_transformer # noqa: PLW0603 + configure_logger(logging.getLogger(), verbose=verbose) + configure_sentry() + _worker_transformer = transformer_class( + source, iter(()), exclusion_list_path=exclusion_list_path + ) + _worker_transformer.run_data = run_data + + +def _transform_chunk( + start_offset: int, source_records: list[bytes | dict[str, JSON] | Tag] +) -> list[DatasetRecord]: + """Transform one chunk of source records in a worker process. + + Setting the processed record count means each record gets its position in the whole + run as its run_record_offset. The worker's other counts and deleted records are not + used; the main process tallies them from the returned records. + + Args: + start_offset: The run_record_offset of the first record in the chunk. + source_records: The chunk of source records, in run order. + """ + if _worker_transformer is None: + raise RuntimeError("Worker transformer was not initialized") + _worker_transformer.source_records = iter(source_records) + _worker_transformer.processed_record_count = start_offset + return list(_worker_transformer) diff --git a/transmogrifier/sources/transformer.py b/transmogrifier/sources/transformer.py index 5d2cc09..471d627 100644 --- a/transmogrifier/sources/transformer.py +++ b/transmogrifier/sources/transformer.py @@ -7,13 +7,10 @@ import itertools import json import logging -import multiprocessing import os import re import uuid from abc import ABC, abstractmethod -from collections import deque -from concurrent.futures import Future, ProcessPoolExecutor from datetime import UTC, datetime from importlib import import_module from typing import TYPE_CHECKING, final @@ -22,7 +19,6 @@ import pyarrow.compute as pc import pyarrow.dataset as ds import smart_open -from attrs import define from bs4 import Tag from pyarrow import fs from timdex_dataset_api import ( @@ -31,7 +27,7 @@ ) import transmogrifier.models as timdex -from transmogrifier.config import SOURCES, configure_logger, configure_sentry +from transmogrifier.config import SOURCES from transmogrifier.exceptions import ( CriticalError, DeletedRecordEvent, @@ -69,30 +65,11 @@ PARQUET_FILES_ROW_GROUP_SIZE = 5_000 PARQUET_FILES_COMPRESSION = "zstd" -WORKER_CHUNK_SIZE = 100 -WORKER_MAX_PENDING_CHUNKS_PER_WORKER = 4 - INPUT_SOURCE_ERROR_MESSAGE = ( "Exactly one input source is required: 'source_file', 'source_dataset', " "or 'source_parquet'" ) -# set once per worker process by _init_transformer_class_for_worker(), and then used -# by _transform_in_worker() -_worker_transformer: Transformer | None = None - - -@define(frozen=True) -class TransformResult: - """Outcome of transforming a single source record. - - Kept small and picklable, so it can be returned from a worker process. - """ - - action: str - timdex_record_id: str | None = None - transformed_record: bytes | None = None - class Transformer(ABC): """Base transformer class.""" @@ -107,7 +84,6 @@ def __init__( source_parquet: str | None = None, run_id: str | None = None, run_timestamp: str | None = None, - workers: int = 1, ) -> None: """ Initialize Transformer instance. @@ -123,8 +99,6 @@ def __init__( source_parquet: Location of parquet file(s) with input source records. run_id: A unique identifier associated with this ETL run. run_timestamp: A timestamp associated with this ETL run. - workers: Number of worker processes to transform records with. With 1, - records are transformed in this process. """ self.source: str = source self.exclusion_list_path: str | None = exclusion_list_path @@ -140,8 +114,11 @@ def __init__( self.source_file = source_file self.source_dataset = source_dataset self.source_parquet = source_parquet - self.workers = workers - self._worker_iterator: Iterator[DatasetRecord] | None = None + + # transformed records are yielded from this private object, by default this + # instance; with multiple workers, the CLI replaces this with records coming + # back from the workers + self._transformed_records: Iterator[DatasetRecord] = self self.run_data = self.get_run_data( source=source, @@ -193,230 +170,71 @@ def load_exclusion_list(self) -> list[str]: @final def __iter__(self) -> Iterator[DatasetRecord]: - """Iterate over transformed records, using worker processes if configured. - - With workers, a single worker iterator is kept and reused, so that iterating - more than once, or mixing iteration with next(), resumes where it left off - instead of dropping records already sent to workers. - """ - if self.workers > 1: - if self._worker_iterator is None: - self._worker_iterator = self._iter_with_workers() - return self._worker_iterator - return self - - @final - def _iter_with_workers(self) -> Iterator[DatasetRecord]: - """Yield transformed records, transformed by a pool of worker processes. - - The main process reads source records as bytes and sends them to workers for - transformation. Workers return only TransformResults; the main process keeps the - source record bytes, and handles results in input order. Output is therefore - identical to transforming in a single process. - - Workers are started fresh, each building its own instance of this transformer's - class, with this transformer's run data (source, run_date, run_type, etc.). This - is done by _init_transformer_class_for_worker(), which runs once in each worker - process as it starts. Each worker then reuses that instance for every chunk. - - Each chunk is a list of (, ) tuples, passed - directly to _transform_in_worker(). Workers therefore never read their instance's - source_records, which is why _init_transformer_class_for_worker() passes an empty - iterator. - - A CriticalError in any worker stops the run. - """ - logger.info(f"Transforming records with {self.workers} worker processes") - - # prepare an executor that will manage workers - executor = ProcessPoolExecutor( - max_workers=self.workers, - mp_context=multiprocessing.get_context("spawn"), - initializer=_init_transformer_class_for_worker, - initargs=( - type(self), - self.source, - self.exclusion_list_path, - self.run_data, - logging.getLogger().getEffectiveLevel() == logging.DEBUG, - ), - ) - - # break input source records into chunks and feed to workers - max_pending_chunks = self.workers * WORKER_MAX_PENDING_CHUNKS_PER_WORKER - pending: deque[tuple[list[bytes], Future[list[TransformResult]]]] = deque() - offsets = itertools.count(self.processed_record_count) - try: - for chunk in itertools.batched( - self.source_records, WORKER_CHUNK_SIZE, strict=False - ): - records: list[bytes] = [] - for record in chunk: - if not isinstance(record, bytes): - raise TypeError( - "Transforming with workers requires source records as bytes" - ) - records.append(record) - pending.append( - ( - records, - executor.submit( - _transform_in_worker, - [(record, next(offsets)) for record in records], - ), - ) - ) - if len(pending) >= max_pending_chunks: - yield from self._create_dataset_records(*pending.popleft()) - while pending: - yield from self._create_dataset_records(*pending.popleft()) - - # regardless of success or failure, shutdown the executor and free processes - finally: - executor.shutdown(wait=True, cancel_futures=True) - - def _create_dataset_records( - self, records: list[bytes], future: Future[list[TransformResult]] - ) -> Iterator[DatasetRecord]: - """Yield DatasetRecords for a chunk of records transformed by a worker. - - Waits for the worker's results, then does the main process bookkeeping for each - record in run order, as __next__() does when not using workers. - - Args: - records: The chunk's source records, as sent to the worker. - future: The worker's pending TransformResults for the chunk, in order. - """ - for record, result in zip(records, future.result(), strict=True): - self.processed_record_count += 1 - self._update_record_counts(result) - yield self._create_dataset_record(record, result) + """Iterate over transformed records.""" + return self._transformed_records @final def __next__(self) -> DatasetRecord: - """Return next transformed record. - - Each source record is handled in three steps: - 1. _process_source_record(): parse and transform the record, and determine - its action. This changes no transformer state, so it is safe to run in - another process. - 2. _update_record_counts(): update the run's record counts and deleted - records. - 3. _create_dataset_record(): build the DatasetRecord from the result, the - original source record, and run data. - - Anything that reads or changes run state must stay in this process and run in - order: incrementing processed_record_count (which sets run_record_offset), - updating record counts and deleted records, and building the DatasetRecord. - - With workers, the next record is taken from the shared worker iterator. - """ - if self.workers > 1: - return next(iter(self)) - - raw_source_record = next(self.source_records) - self.processed_record_count += 1 - - # parse the raw source record bytes + transform the record - transform_result = self._process_source_record( - raw_source_record, - self.run_record_offset, - ) - - # update bookkeeping - self._update_record_counts(transform_result) + """Return next transformed record.""" + while True: + source_record = None + transformed_record = None + timdex_record_id = None - # return a DatasetRecord instance with source + transformed record - return self._create_dataset_record(raw_source_record, transform_result) - - @final - def _process_source_record( - self, record: bytes | dict[str, JSON] | Tag, run_record_offset: int - ) -> TransformResult: - """Parse and transform a single source record, and determine its action. - - This step does not change transformer state (e.g. record counts), so it can run - in a worker process. - - Args: - record: A source record as bytes, or an already parsed source record. - run_record_offset: The record's position in the run, for provenance. - - Raises: - CriticalError: Propagated, as it should stop the run. - """ - try: - source_record = ( - self.parse_source_record(record) if isinstance(record, bytes) else record - ) - timdex_record = self.transform(source_record) - timdex_record.timdex_provenance = timdex.TimdexProvenance( - source=self.run_data["source"], - run_date=self.run_data["run_date"], - run_id=self.run_data["run_id"], - run_record_offset=run_record_offset, - ) - return TransformResult( - action="index", - timdex_record_id=timdex_record.timdex_record_id, - transformed_record=json.dumps(timdex_record.asdict()).encode(), - ) - - except DeletedRecordEvent as error: - return TransformResult( - action="delete", timdex_record_id=error.timdex_record_id - ) - - except SkippedRecordEvent: - return TransformResult(action="skip") - - except CriticalError: - raise + record = next(self.source_records) + self.processed_record_count += 1 - except Exception as exception: - message = f"Unhandled exception during record transformation: {exception}" - logger.exception(message) - return TransformResult(action="error") + try: + source_record = ( + self.parse_source_record(record) + if isinstance(record, bytes) + else record + ) + transformed_record = self.transform(source_record) + timdex_record_id = transformed_record.timdex_record_id + transformed_record.timdex_provenance = timdex.TimdexProvenance( + source=self.run_data["source"], + run_date=self.run_data["run_date"], + run_id=self.run_data["run_id"], + run_record_offset=self.run_record_offset, + ) + self.transformed_record_count += 1 + action = "index" - @final - def _update_record_counts(self, result: TransformResult) -> None: - """Update the run's record counts and deleted records for a transform result. + except DeletedRecordEvent as error: + self.deleted_records.append(error.timdex_record_id) + timdex_record_id = error.timdex_record_id + action = "delete" - Args: - result: The result of _process_source_record() for a record. - """ - match result.action: - case "index": - self.transformed_record_count += 1 - case "delete": - self.deleted_records.append(str(result.timdex_record_id)) - case "skip": + except SkippedRecordEvent: self.skipped_record_count += 1 - case "error": - self.error_record_count += 1 + action = "skip" - @final - def _create_dataset_record( - self, record: bytes | dict[str, JSON] | Tag, result: TransformResult - ) -> DatasetRecord: - """Create a DatasetRecord from a source record and its transform result. + except CriticalError: + raise - Args: - record: The source record as read, stored as-is if bytes. - result: The result of _process_source_record() for this record. - """ - return DatasetRecord( - timdex_record_id=result.timdex_record_id, - source_record=( - record - if isinstance(record, bytes) - else self.serialize_source_record(record) - ), - transformed_record=result.transformed_record, - action=result.action, - run_record_offset=self.run_record_offset, - **self.run_data, - ) + except Exception as exception: + self.error_record_count += 1 + message = f"Unhandled exception during record transformation: {exception}" + logger.exception(message) + action = "error" + + return DatasetRecord( + timdex_record_id=timdex_record_id, + source_record=( + record + if isinstance(record, bytes) + else self.serialize_source_record(record) + ), + transformed_record=( + json.dumps(transformed_record.asdict()).encode() + if transformed_record + else None + ), + action=action, + run_record_offset=self.run_record_offset, + **self.run_data, + ) @final @classmethod @@ -445,7 +263,6 @@ def load( exclusion_list_path: str | None = None, run_id: str | None = None, run_timestamp: str | None = None, - workers: int = 1, ) -> Transformer: """ Instantiate specified transformer class and populate with source records. @@ -458,7 +275,6 @@ def load( exclusion_list_path: CSV filepath to use for explicitly skipping records. run_id: A unique identifier associated with this ETL run. run_timestamp: A timestamp associated with this ETL run. - workers: Number of worker processes to transform records with. """ transformer_class = cls.get_transformer(source) @@ -487,7 +303,6 @@ def load( source_parquet=source_parquet, run_id=run_id, run_timestamp=run_timestamp, - workers=workers, ) @staticmethod @@ -1029,44 +844,3 @@ def _parse_location(location: str) -> tuple[fs.FileSystem, str]: if location.startswith("s3://"): return fs.S3FileSystem(), location.removeprefix("s3://") return fs.LocalFileSystem(), os.path.abspath(location) - - -def _init_transformer_class_for_worker( - transformer_class: type[Transformer], - source: str, - exclusion_list_path: str | None, - run_data: dict, - verbose: bool, # noqa: FBT001 -) -> None: - """Build this worker process's transformer, with the main process's run data. - - The transformer class is initialized with an empty iterator in place of its source - records, because it uses the function _transform_in_worker() for transformation, - which receives a custom chunk of records to process. - .source_records is never used. - """ - global _worker_transformer # noqa: PLW0603 - configure_logger(logging.getLogger(), verbose=verbose) - configure_sentry() - _worker_transformer = transformer_class( - source, iter(()), exclusion_list_path=exclusion_list_path - ) - _worker_transformer.run_data = run_data - - -def _transform_in_worker( - source_records_chunk: list[tuple[bytes, int]], -) -> list[TransformResult]: - """Transform one chunk of source records in a worker process. - - Args: - source_records_chunk: One chunk of the run's source records, sent by the main - process in run order, as tuples of (source record bytes, run_record_offset). - The offset is the record's position in the whole run, not in this chunk. - """ - if _worker_transformer is None: - raise RuntimeError("Worker transformer was not initialized") - return [ - _worker_transformer._process_source_record(record, run_record_offset) # noqa: SLF001 - for record, run_record_offset in source_records_chunk - ]