From 98a89793528091165da44966fbb8776baf44261c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cle=CC=81ment=20Doumouro?= Date: Tue, 1 Sep 2026 15:50:50 +0200 Subject: [PATCH 1/3] feature(datashare-python): add caching support to `activity_workdir` function --- datashare-python/datashare_python/utils.py | 30 +++++-- datashare-python/tests/test_utils.py | 96 +++++++++++++++++++++- 2 files changed, 117 insertions(+), 9 deletions(-) diff --git a/datashare-python/datashare_python/utils.py b/datashare-python/datashare_python/utils.py index 8e96452f..bb3935b4 100644 --- a/datashare-python/datashare_python/utils.py +++ b/datashare-python/datashare_python/utils.py @@ -62,6 +62,8 @@ ) from .types_ import RawAsyncProgressHandler +_CACHE_DIR = "processing_cache" + logger = logging.getLogger(__name__) DependencyLabel = str | None @@ -533,7 +535,11 @@ def contextual_id( def _contextual_path( - *, wf_context: bool = True, act_context: bool = True, run_context: bool = False + *, + wf_context: bool = True, + act_context: bool = True, + run_context: bool = False, + caching_hash: str | None = None, ) -> Path: act_info = activity.info() path = [] @@ -541,14 +547,18 @@ def _contextual_path( raise ValueError("at least one of wf_context and act_context must be True") if wf_context: path = [act_info.workflow_type] - path.append(act_info.workflow_id) - if run_context: - path.append(act_info.workflow_run_id) + if not caching_hash: + path.append(act_info.workflow_id) + if run_context: + path.append(act_info.workflow_run_id) if act_context: path.append(act_info.activity_type) - path.append(act_info.activity_id) - if run_context: - path.append(act_info.activity_run_id) + if not caching_hash: + path.append(act_info.activity_id) + if run_context: + path.append(act_info.activity_run_id) + if caching_hash: + path += [_CACHE_DIR, caching_hash] return Path(*path) @@ -559,9 +569,13 @@ def activity_workdir( wf_context: bool = True, act_context: bool = True, run_context: bool = False, + caching_hash: str | None = None, ) -> Path: ctx_path = _contextual_path( - wf_context=wf_context, act_context=act_context, run_context=run_context + wf_context=wf_context, + act_context=act_context, + run_context=run_context, + caching_hash=caching_hash, ) return workdir.joinpath(project, ctx_path) diff --git a/datashare-python/tests/test_utils.py b/datashare-python/tests/test_utils.py index 1225379a..37ebca38 100644 --- a/datashare-python/tests/test_utils.py +++ b/datashare-python/tests/test_utils.py @@ -7,9 +7,11 @@ from datetime import timedelta from pathlib import Path from typing import ClassVar -from unittest.mock import MagicMock +from unittest.mock import MagicMock, PropertyMock import pytest +from _pytest.monkeypatch import MonkeyPatch +from conftest import TEST_PROJECT from datashare_python.constants import MANIFEST_JSON from datashare_python.objects import ( ArtifactType, @@ -23,6 +25,7 @@ _LOCKED, SharedResources, activity_defn, + activity_workdir, artifact_lock, positional_args_only, write_artifact, @@ -411,3 +414,94 @@ def test_get_or_cache_eviction_callback_is_called_on_exit() -> None: shared.get_or_cache_resource(key, factory) # Then eviction_callback.assert_called_once_with(key, "value") + + +@pytest.mark.parametrize( + ("wf_context", "act_context", "run_context", "expected_work_dir"), + [ + ( + True, + True, + True, + Path( + "workdir", + TEST_PROJECT, + "wf_type", + "act_type", + "processing_cache", + "caching_hash", + ), + ), + ( + True, + True, + False, + Path( + "workdir", + TEST_PROJECT, + "wf_type", + "act_type", + "processing_cache", + "caching_hash", + ), + ), + ( + True, + False, + True, + Path( + "workdir", TEST_PROJECT, "wf_type", "processing_cache", "caching_hash" + ), + ), + ( + True, + False, + False, + Path( + "workdir", TEST_PROJECT, "wf_type", "processing_cache", "caching_hash" + ), + ), + ( + False, + True, + True, + Path( + "workdir", TEST_PROJECT, "act_type", "processing_cache", "caching_hash" + ), + ), + ( + False, + True, + False, + Path( + "workdir", TEST_PROJECT, "act_type", "processing_cache", "caching_hash" + ), + ), + ], +) +def test_cached_activity_workdir( + *, + wf_context: bool, + act_context: bool, + run_context: bool, + expected_work_dir: Path, + monkeypatch: MonkeyPatch, +) -> None: + # Given + caching_hash = "caching_hash" + base_workdir = Path("workdir") + mocked_info = MagicMock() + type(mocked_info).workflow_type = PropertyMock(return_value="wf_type") + type(mocked_info).activity_type = PropertyMock(return_value="act_type") + monkeypatch.setattr(activity, "info", lambda: mocked_info) + # When + workdir = activity_workdir( + base_workdir, + project=TEST_PROJECT, + caching_hash=caching_hash, + wf_context=wf_context, + act_context=act_context, + run_context=run_context, + ) + # Then + assert workdir == expected_work_dir From 051e25a8f4c8f37f52b59ca90a130a5adb13a8ce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cle=CC=81ment=20Doumouro?= Date: Tue, 1 Sep 2026 18:31:07 +0200 Subject: [PATCH 2/3] feature(passport-service): implement a global preprocessing cache --- datashare-python/datashare_python/objects.py | 8 + datashare-python/datashare_python/utils.py | 15 +- datashare-python/tests/test_objects.py | 13 ++ datashare-python/tests/test_utils.py | 20 +-- datashare-python/uv.lock | 2 +- .../passport_worker/activities.py | 39 +++-- .../passport_worker/objects.py | 2 + .../passport_worker/preprocessing.py | 125 ++++++++++---- .../passport-worker/passport_worker/utils.py | 4 +- .../passport_worker/workflows.py | 27 ++- workers/passport-worker/pyproject.toml | 8 +- .../tests/test_preprocessing.py | 156 +++++++++++++++++- workers/passport-worker/uv.dist.lock | 14 +- workers/passport-worker/uv.lock | 14 +- 14 files changed, 350 insertions(+), 97 deletions(-) diff --git a/datashare-python/datashare_python/objects.py b/datashare-python/datashare_python/objects.py index 6aa7270c..acd24baa 100644 --- a/datashare-python/datashare_python/objects.py +++ b/datashare-python/datashare_python/objects.py @@ -1,3 +1,5 @@ +import hashlib +import json import logging import os from abc import ABC @@ -63,6 +65,12 @@ class BaseModel(_BaseModel): model_config = merge_configs(icij_config(), no_enum_values_config()) + def __hash__(self) -> int: + digest = hashlib.md5( + json.dumps(self.model_dump(), sort_keys=True).encode() + ).digest() + return int.from_bytes(digest[:8]) + class DatashareModel(BaseModel): model_config = merge_configs(BaseModel.model_config, lowercamel_case_config()) diff --git a/datashare-python/datashare_python/utils.py b/datashare-python/datashare_python/utils.py index bb3935b4..4b07295c 100644 --- a/datashare-python/datashare_python/utils.py +++ b/datashare-python/datashare_python/utils.py @@ -539,7 +539,7 @@ def _contextual_path( wf_context: bool = True, act_context: bool = True, run_context: bool = False, - caching_hash: str | None = None, + caching_key: str | None = None, ) -> Path: act_info = activity.info() path = [] @@ -547,18 +547,18 @@ def _contextual_path( raise ValueError("at least one of wf_context and act_context must be True") if wf_context: path = [act_info.workflow_type] - if not caching_hash: + if not caching_key: path.append(act_info.workflow_id) if run_context: path.append(act_info.workflow_run_id) if act_context: path.append(act_info.activity_type) - if not caching_hash: + if not caching_key: path.append(act_info.activity_id) if run_context: path.append(act_info.activity_run_id) - if caching_hash: - path += [_CACHE_DIR, caching_hash] + if caching_key: + path += [_CACHE_DIR, caching_key] return Path(*path) @@ -569,13 +569,13 @@ def activity_workdir( wf_context: bool = True, act_context: bool = True, run_context: bool = False, - caching_hash: str | None = None, + caching_key: str | None = None, ) -> Path: ctx_path = _contextual_path( wf_context=wf_context, act_context=act_context, run_context=run_context, - caching_hash=caching_hash, + caching_key=caching_key, ) return workdir.joinpath(project, ctx_path) @@ -822,4 +822,5 @@ async def wrapper() -> contextlib.AbstractAsyncContextManager: def config_cache_key(config: BaseModel) -> str: + # BaseModel support deterministic hash we can call hash on it return str(hash(config)) diff --git a/datashare-python/tests/test_objects.py b/datashare-python/tests/test_objects.py index 18288132..213ddd4b 100644 --- a/datashare-python/tests/test_objects.py +++ b/datashare-python/tests/test_objects.py @@ -9,6 +9,7 @@ from datashare_python.conftest import TEST_PROJECT from datashare_python.constants import TIKA_METADATA_RESOURCENAME from datashare_python.objects import ( + BaseModel, ByteRangesPagination, DatashareLanguage, Document, @@ -190,3 +191,15 @@ def test_manifest_entry_partial_task_id( assert manifest_entry.task_id is not None else: assert manifest_entry.task_id is None + + +def test_base_model_hash_is_constant() -> None: + # Given + class MockConfig(BaseModel): + some_key: str = "some_value" + + cfg = MockConfig() + + # When + hashed = hash(cfg) + assert hashed == 1821116217857887821 diff --git a/datashare-python/tests/test_utils.py b/datashare-python/tests/test_utils.py index 37ebca38..3ed2b4c0 100644 --- a/datashare-python/tests/test_utils.py +++ b/datashare-python/tests/test_utils.py @@ -429,7 +429,7 @@ def test_get_or_cache_eviction_callback_is_called_on_exit() -> None: "wf_type", "act_type", "processing_cache", - "caching_hash", + "caching_key", ), ), ( @@ -442,31 +442,27 @@ def test_get_or_cache_eviction_callback_is_called_on_exit() -> None: "wf_type", "act_type", "processing_cache", - "caching_hash", + "caching_key", ), ), ( True, False, True, - Path( - "workdir", TEST_PROJECT, "wf_type", "processing_cache", "caching_hash" - ), + Path("workdir", TEST_PROJECT, "wf_type", "processing_cache", "caching_key"), ), ( True, False, False, - Path( - "workdir", TEST_PROJECT, "wf_type", "processing_cache", "caching_hash" - ), + Path("workdir", TEST_PROJECT, "wf_type", "processing_cache", "caching_key"), ), ( False, True, True, Path( - "workdir", TEST_PROJECT, "act_type", "processing_cache", "caching_hash" + "workdir", TEST_PROJECT, "act_type", "processing_cache", "caching_key" ), ), ( @@ -474,7 +470,7 @@ def test_get_or_cache_eviction_callback_is_called_on_exit() -> None: True, False, Path( - "workdir", TEST_PROJECT, "act_type", "processing_cache", "caching_hash" + "workdir", TEST_PROJECT, "act_type", "processing_cache", "caching_key" ), ), ], @@ -488,7 +484,7 @@ def test_cached_activity_workdir( monkeypatch: MonkeyPatch, ) -> None: # Given - caching_hash = "caching_hash" + caching_key = "caching_key" base_workdir = Path("workdir") mocked_info = MagicMock() type(mocked_info).workflow_type = PropertyMock(return_value="wf_type") @@ -498,7 +494,7 @@ def test_cached_activity_workdir( workdir = activity_workdir( base_workdir, project=TEST_PROJECT, - caching_hash=caching_hash, + caching_key=caching_key, wf_context=wf_context, act_context=act_context, run_context=run_context, diff --git a/datashare-python/uv.lock b/datashare-python/uv.lock index 60f399a4..a023149c 100644 --- a/datashare-python/uv.lock +++ b/datashare-python/uv.lock @@ -477,7 +477,7 @@ wheels = [ [[package]] name = "datashare-python" -version = "0.10.4" +version = "0.10.6" source = { editable = "." } dependencies = [ { name = "aiofile" }, diff --git a/workers/passport-worker/passport_worker/activities.py b/workers/passport-worker/passport_worker/activities.py index c1d62979..21bd2eb9 100644 --- a/workers/passport-worker/passport_worker/activities.py +++ b/workers/passport-worker/passport_worker/activities.py @@ -35,10 +35,10 @@ from .objects import ( DocId, DocumentSearchQuery, - ImagePreprocessorConfig, PassportDetectionArgs, PassportDetectionResponse, PreprocessingBatches, + PreprocessingConfig, ) from .preprocessing import ( ImagePreprocessor, @@ -101,7 +101,7 @@ def preprocess_images( self, batch: Path, project: str, - config: ImagePreprocessorConfig, + config: PreprocessingConfig, *, progress: Annotated[ SyncProgressRateHandler | None, Weight(value=_PREPROCESS_IMAGES_WEIGHT) @@ -111,21 +111,25 @@ def preprocess_images( workdir = worker_config.paths.workdir logger.info("loading image preprocessor...") cache = lifespan_image_preprocessor_cache() - image_preprocessor_cache_key = config_cache_key(config) + cache_key = config_cache_key(config) image_preprocessor_factory = enter_cm( - partial(ImagePreprocessor.from_config, config) + partial(ImagePreprocessor.from_config, config.images) ) image_preprocessor = cache.get_or_cache_resource( - image_preprocessor_cache_key, image_preprocessor_factory + cache_key, image_preprocessor_factory ) logger.info("loaded image preprocessor !") - pages_root = activity_workdir(workdir, project, act_context=False) + # We cache processing at the config level, we ignore package updates, + # to discard the cache we just need to disable it in the args to overwrite + pages_root = activity_workdir(workdir, project, caching_key=cache_key) pages_root.mkdir(parents=True, exist_ok=True) executor = worker_config.to_image_preprocessing_executor() chunk_size = worker_config.preprocessing.images.chunk_size + force_reprocessing = not config.use_caching success, errors = preprocess_images_act( batch, worker_config.paths, + force_reprocessing=force_reprocessing, output_root=pages_root, image_preprocessor=image_preprocessor, executor=executor, @@ -133,7 +137,7 @@ def preprocess_images( event_loop=self._event_loop, progress=progress, ) - res_root = activity_workdir(workdir, project, act_context=True) + res_root = activity_workdir(workdir, project) res_root.mkdir(parents=True, exist_ok=True) successes_path = res_root / "pages.jsonl" successes_path.write_text("\n".join(p.model_dump_json() for p in success)) @@ -147,6 +151,7 @@ async def convert_to_pdfs( batch: Path, project: str, *, + force_reprocessing: bool, progress: Annotated[ AsyncProgressRateHandler | None, Weight(value=_CONVERT_TO_PDF_WEIGHT) ] = None, @@ -154,15 +159,17 @@ async def convert_to_pdfs( worker_config = cast(PassportWorkerConfig, lifespan_worker_config()) config = worker_config.preprocessing.pdfs cache = lifespan_pdf_converter_cache() - pdf_converter_cache_key = config_cache_key(config.pdf_converter) + cache_key = config_cache_key(config.pdf_converter) pdf_converter_factory = async_enter_cm( partial(PDFConverter.from_config, config.pdf_converter) ) pdf_converter = await cache.async_get_or_cache_resource( - pdf_converter_cache_key, pdf_converter_factory + cache_key, pdf_converter_factory ) workdir = worker_config.paths.workdir - pdfs_root = activity_workdir(workdir, project, act_context=False) + # We cache processing at the config level, we ignore package updates, + # to discard the cache we just need to disable it in the args to overwrite + pdfs_root = activity_workdir(workdir, project, caching_key=cache_key) pdfs_root.mkdir(parents=True, exist_ok=True) successes, errors = await convert_to_pdfs_act( batch, @@ -170,6 +177,7 @@ async def convert_to_pdfs( worker_config.paths, config.max_concurrency, output_root=pdfs_root, + force_reprocessing=force_reprocessing, progress=progress, ) res_root = activity_workdir(workdir, project, act_context=True) @@ -188,16 +196,23 @@ async def preprocess_pdfs( batch: Path, project: str, *, + force_reprocessing: bool, progress: Annotated[ AsyncProgressRateHandler | None, Weight(value=_PREPROCESS_PDF_WEIGHT) ] = None, ) -> tuple[Path, Path]: worker_config = cast(PassportWorkerConfig, lifespan_worker_config()) workdir = worker_config.paths.workdir - output_root = activity_workdir(workdir, project, act_context=False) + # We cache processing at the config level, we ignore package updates, + # to discard the cache we just need to disable it in the args to overwrite + output_root = activity_workdir(workdir, project, caching_key="pdf-preprocessor") output_root.mkdir(parents=True, exist_ok=True) successes, errors = await preprocess_pdfs_act( - batch, worker_config.paths, output_root=output_root, progress=progress + batch, + worker_config.paths, + output_root=output_root, + force_reprocessing=force_reprocessing, + progress=progress, ) res_root = activity_workdir(workdir, project, act_context=True) res_root.mkdir(parents=True, exist_ok=True) diff --git a/workers/passport-worker/passport_worker/objects.py b/workers/passport-worker/passport_worker/objects.py index a7779a50..26b3bc6b 100644 --- a/workers/passport-worker/passport_worker/objects.py +++ b/workers/passport-worker/passport_worker/objects.py @@ -63,6 +63,8 @@ class DefaultImagePreprocessorConfig(ImagePreprocessorConfigBase): class PreprocessingConfig(DatashareModel): + use_caching: bool = True + images: ImagePreprocessorConfig = Field( default_factory=DefaultImagePreprocessorConfig ) diff --git a/workers/passport-worker/passport_worker/preprocessing.py b/workers/passport-worker/passport_worker/preprocessing.py index 5b3f3b46..0a711504 100644 --- a/workers/passport-worker/passport_worker/preprocessing.py +++ b/workers/passport-worker/passport_worker/preprocessing.py @@ -47,7 +47,9 @@ class ImagePreprocessor(RegistrableFromConfig): @abstractmethod - def __call__(self, image_path: Path, *, output_dir: Path) -> list[Path]: ... + def __call__( + self, image_path: Path, *, output_dir: Path, force_reprocessing: bool + ) -> list[Path]: ... def __enter__(self) -> Self: return self @@ -67,8 +69,12 @@ def __init__(self, config: DefaultImagePreprocessorConfig | None = None): config = DefaultImagePreprocessorConfig() self._config = config - def __call__(self, image_path: Path, *, output_dir: Path) -> list[Path]: - return process_image(image_path, output_dir=output_dir) + def __call__( + self, image_path: Path, *, output_dir: Path, force_reprocessing: bool + ) -> list[Path]: + return process_image( + image_path, output_dir=output_dir, force_reprocessing=force_reprocessing + ) @classmethod def _from_config(cls, config: DefaultImagePreprocessorConfig, **extras) -> Self: # noqa:ARG003 @@ -107,7 +113,12 @@ def _from_config(cls, config: GotenbergPDFConverterConfig, **extras) -> Self: # class PDFPreprocessor(Protocol): def __call__( - self, pdf_path: Path, pdf_bytes: bytes, output_dir: Path + self, + pdf_path: Path, + pdf_bytes: bytes, + output_dir: Path, + *, + force_reprocessing: bool, ) -> list[Path]: ... @@ -117,6 +128,7 @@ def preprocess_images_act( *, output_root: Path, image_preprocessor: ImagePreprocessor, + force_reprocessing: bool = True, executor: ProcessPoolExecutor | None = None, chunk_size: int = 1, event_loop: asyncio.AbstractEventLoop | None = None, @@ -131,6 +143,7 @@ def preprocess_images_act( chunk_size = 1 if n_docs < n_processes * chunk_size else chunk_size process_doc_fn = partial( _preprocess_image_doc, + force_reprocessing=force_reprocessing, image_preprocessor=image_preprocessor, paths=paths, output_root=output_root, @@ -160,6 +173,7 @@ async def convert_to_pdfs_act( paths: WorkerPaths, max_concurrency: int, *, + force_reprocessing: bool, output_root: Path, progress: AsyncProgressRateHandler | None = None, ) -> tuple[list[ProcessedFile], list[FileProcessingError]]: @@ -168,7 +182,12 @@ async def convert_to_pdfs_act( n_docs = len(docs) if progress is not None: progress = to_raw_async_progress(progress, max_progress=n_docs) - aws = (_convert_doc_to_pdf(doc, converter, paths, output_root) for doc in docs) + aws = ( + _convert_doc_to_pdf( + doc, converter, paths, output_root, force_reprocessing=force_reprocessing + ) + for doc in docs + ) res_i = 0 successes = [] errors = [] @@ -189,12 +208,37 @@ async def convert_to_pdfs_act( return successes, errors +@reports_errors(errors=REPORTED_ERRORS) +def _preprocess_image_doc( + doc: ProcessedFile, + image_preprocessor: ImagePreprocessor, + paths: WorkerPaths, + *, + output_root: Path, + force_reprocessing: bool, +) -> list[ProcessedPage]: + ext = doc.path.suffix.lower() + if ext not in pil_supported_extensions(): + logger.info("image extension %s not supported !", ext) + raise UnsupportedDocExtension(ext, sorted(pil_supported_extensions())) + output_dir = output_root / safe_dir(doc.id) / doc.id + im_paths = image_preprocessor( + doc.locate(paths), output_dir=output_dir, force_reprocessing=force_reprocessing + ) + pages = [ + ProcessedPage(page_number=p_i + 1, **doc.child(p, paths).model_dump()) + for p_i, p in enumerate(im_paths) + ] + return pages + + async def preprocess_pdfs_act( batch: Path, paths: WorkerPaths, pdf_preprocessor: PDFPreprocessor | None = None, *, output_root: Path, + force_reprocessing: bool, progress: AsyncProgressRateHandler | None = None, ) -> tuple[list[ProcessedPage], list[FileProcessingError]]: if pdf_preprocessor is None: @@ -207,7 +251,11 @@ async def preprocess_pdfs_act( errors = [] for doc_i, doc in enumerate(docs): res = await _preprocess_pdf( - doc, pdf_preprocessor, paths, output_root=output_root + doc, + pdf_preprocessor, + paths, + force_reprocessing=force_reprocessing, + output_root=output_root, ) if isinstance(res, FileProcessingError): errors.append(res) @@ -222,38 +270,23 @@ async def preprocess_pdfs_act( @reports_errors(errors=REPORTED_ERRORS) -def _preprocess_image_doc( +async def _convert_doc_to_pdf( doc: ProcessedFile, - image_preprocessor: ImagePreprocessor, + converter: PDFConverter, paths: WorkerPaths, - *, output_root: Path, -) -> list[ProcessedPage]: - ext = doc.path.suffix.lower() - if ext not in pil_supported_extensions(): - logger.info("image extension %s not supported !", ext) - raise UnsupportedDocExtension(ext, sorted(pil_supported_extensions())) - output_dir = output_root / safe_dir(doc.id) / doc.id - output_dir.mkdir(parents=True, exist_ok=True) - im_paths = image_preprocessor(doc.locate(paths), output_dir=output_dir) - pages = [ - ProcessedPage(page_number=p_i + 1, **doc.child(p, paths).model_dump()) - for p_i, p in enumerate(im_paths) - ] - return pages - - -@reports_errors(errors=REPORTED_ERRORS) -async def _convert_doc_to_pdf( - doc: ProcessedFile, converter: PDFConverter, paths: WorkerPaths, output_root: Path + *, + force_reprocessing: bool, ) -> ProcessedFile: - async with async_open(doc.locate(paths), "rb") as f: - doc_bytes = await f.read() - pdf_bytes = await converter(doc, doc_bytes) pdf_path = output_root / safe_dir(doc.id) / f"{doc.id}.pdf" - pdf_path.parent.mkdir(parents=True, exist_ok=True) - async with async_open(pdf_path, "wb") as f: - await f.write(pdf_bytes) + valid_pdf = await is_valid_pdf(pdf_path) + if force_reprocessing or not valid_pdf: + async with async_open(doc.locate(paths), "rb") as f: + doc_bytes = await f.read() + pdf_bytes = await converter(doc, doc_bytes) + pdf_path.parent.mkdir(parents=True, exist_ok=True) + async with async_open(pdf_path, "wb") as f: + await f.write(pdf_bytes) processed = doc.child(pdf_path, paths) return processed @@ -264,6 +297,7 @@ async def _preprocess_pdf( pdf_processor: PDFPreprocessor, paths: WorkerPaths, *, + force_reprocessing: bool, output_root: Path, ) -> list[ProcessedPage]: pdf_path = doc.locate(paths) @@ -272,10 +306,33 @@ async def _preprocess_pdf( output_dir = output_root / safe_dir(doc.id) / doc.id output_dir.mkdir(parents=True, exist_ok=True) pages = await asyncio.to_thread( - pdf_processor, pdf_path, pdf_bytes, output_dir=output_dir + pdf_processor, + pdf_path, + pdf_bytes, + output_dir=output_dir, + force_reprocessing=force_reprocessing, ) pages = [ ProcessedPage(page_number=p_i + 1, **doc.child(p, paths).model_dump()) for p_i, p in enumerate(pages) ] return pages + + +async def is_valid_pdf(path: Path) -> bool: + import pymupdf # noqa: PLC0415 + + if not path.exists(): + return False + async with async_open(path, "rb") as f: + pdf_bytes = await f.read() + + doc = None + try: + doc = pymupdf.open(stream=pdf_bytes, filetype="pdf") + return doc.is_pdf + except Exception: # noqa: BLE001 + return False + finally: + if doc is not None: + doc.close() diff --git a/workers/passport-worker/passport_worker/utils.py b/workers/passport-worker/passport_worker/utils.py index 8c1d4179..f179dca7 100644 --- a/workers/passport-worker/passport_worker/utils.py +++ b/workers/passport-worker/passport_worker/utils.py @@ -71,7 +71,9 @@ def reports_errors[R]( [_PreprocessingFunction[R]], _PreprocessingFunction[R | FileProcessingError] ]: - def parent_wrapper(f) -> _PreprocessingFunction[R | FileProcessingError]: + def parent_wrapper( + f: _PreprocessingFunction[R], + ) -> _PreprocessingFunction[R | FileProcessingError]: if iscoroutinefunction(f): @wraps(f) # noqa: F821 diff --git a/workers/passport-worker/passport_worker/workflows.py b/workers/passport-worker/passport_worker/workflows.py index ff596835..9b7eeacd 100644 --- a/workers/passport-worker/passport_worker/workflows.py +++ b/workers/passport-worker/passport_worker/workflows.py @@ -13,10 +13,10 @@ from .activities import PassportDetectionActivities from .objects import ( Batches, - ImagePreprocessorConfig, PassportDetectionArgs, PassportDetectionResponse, PreprocessingBatches, + PreprocessingConfig, ) logger = logging.getLogger(__name__) @@ -95,10 +95,13 @@ async def preprocess( args: PassportDetectionArgs, preprocessing_batches: PreprocessingBatches ) -> PreprocessingOutput: im_preprocessing_tasks = _im_processing_tasks( - preprocessing_batches.images, args.project, args.config.preprocessing.images + preprocessing_batches.images, args.project, args.config.preprocessing ) + force_reprocessing = not args.config.preprocessing.use_caching convert_to_pdf_tasks = _convert_to_pdfs_tasks( - preprocessing_batches.to_pdf, args.project + preprocessing_batches.to_pdf, + args.project, + force_reprocessing=force_reprocessing, ) im_preprocessing_tasks = asyncio.gather(*im_preprocessing_tasks) convert_to_pdf_tasks = asyncio.gather(*convert_to_pdf_tasks) @@ -125,7 +128,9 @@ async def preprocess( # Preprocess all files converted into PDFs + original PDFs logger.info("converting PDF pages to PNG...") pdf_batches = preprocessing_batches.pdfs + pdf_paths - preprocess_pdfs_tasks = _process_pdfs_tasks(pdf_batches, args.project) + preprocess_pdfs_tasks = _process_pdfs_tasks( + pdf_batches, args.project, force_reprocessing=force_reprocessing + ) pdf_pages_res = await asyncio.gather(*preprocess_pdfs_tasks) if pdf_pages_res: pdfs_pages_paths, pdf_processing_errors = zip(*pdf_pages_res, strict=True) @@ -141,7 +146,7 @@ async def preprocess( def _im_processing_tasks( - batches: Batches, project: str, config: ImagePreprocessorConfig + batches: Batches, project: str, config: PreprocessingConfig ) -> list: im_preprocessing_tasks = [] for b in batches: @@ -157,13 +162,15 @@ def _im_processing_tasks( return im_preprocessing_tasks -def _convert_to_pdfs_tasks(batches: Batches, project: str) -> list[Coroutine]: +def _convert_to_pdfs_tasks( + batches: Batches, project: str, *, force_reprocessing: bool +) -> list[Coroutine]: all_tasks = [] for b in batches: all_tasks.append( execute_activity( PassportDetectionActivities.convert_to_pdfs, - args=(b, project), + args=(b, project, force_reprocessing), task_queue=TaskQueue.IO, start_to_close_timeout=_CONVERT_TO_PDF_TIMEOUT, heartbeat_timeout=timedelta(minutes=2), @@ -172,13 +179,15 @@ def _convert_to_pdfs_tasks(batches: Batches, project: str) -> list[Coroutine]: return all_tasks -def _process_pdfs_tasks(batches: Batches, project: str) -> list[Coroutine]: +def _process_pdfs_tasks( + batches: Batches, project: str, *, force_reprocessing: bool +) -> list[Coroutine]: all_tasks = [] for b in batches: all_tasks.append( execute_activity( PassportDetectionActivities.preprocess_pdfs, - args=(b, project), + args=(b, project, force_reprocessing), task_queue=TaskQueue.IO, start_to_close_timeout=_CONVERT_TO_PDF_TIMEOUT, ) diff --git a/workers/passport-worker/pyproject.toml b/workers/passport-worker/pyproject.toml index 84332407..5fadb50c 100644 --- a/workers/passport-worker/pyproject.toml +++ b/workers/passport-worker/pyproject.toml @@ -11,18 +11,18 @@ requires-python = ">=3.12,<3.14" dependencies = [ "datashare-python~=0.10.6", "icij-common[elasticsearch]~=0.8.3", - "icij-passport-core==0.11.6", + "icij-passport-core==0.11.7", ] [project.optional-dependencies] preprocessing = [ - "icij-passport-core[preprocessing]==0.11.6", + "icij-passport-core[preprocessing]==0.11.7", ] inference-gpu = [ - "icij-passport-core[inference,gpu]==0.11.6" + "icij-passport-core[inference,gpu]==0.11.7" ] inference-cpu = [ - "icij-passport-core[inference,cpu]==0.11.6" + "icij-passport-core[inference,cpu]==0.11.7" ] [tool.uv.sources] diff --git a/workers/passport-worker/tests/test_preprocessing.py b/workers/passport-worker/tests/test_preprocessing.py index 7274ec13..c13f61d2 100644 --- a/workers/passport-worker/tests/test_preprocessing.py +++ b/workers/passport-worker/tests/test_preprocessing.py @@ -15,6 +15,7 @@ PDFConverter, PDFPreprocessor, convert_to_pdfs_act, + is_valid_pdf, preprocess_images_act, preprocess_pdfs_act, ) @@ -32,11 +33,21 @@ class MockImageProcessor(ImagePreprocessor): def __init__(self, res: list[list[Path] | Exception]): self._res = iter(res) + self.processed = [] - def __call__(self, image_path: Path, *, output_dir: Path) -> list[Path]: # noqa: ARG002 + def __call__( + self, + image_path: Path, + *, + output_dir: Path, # noqa: ARG002 + force_reprocessing: bool, + ) -> list[Path]: r = next(self._res) if isinstance(r, Exception): raise r + for im_path in r: + if force_reprocessing and not im_path.exists(): + self.processed.append(image_path) return r @classmethod @@ -46,11 +57,13 @@ def _from_config(cls, config: RegistrableConfig, **extras) -> FromConfig: ... class MockConverter(PDFConverter): def __init__(self, conversion_results: dict[str, bytes | Exception]): self._conversion_results = conversion_results + self.call_count = 0 async def __call__(self, doc: ProcessedFile, doc_bytes: bytes) -> bytes: # noqa: ARG002 res = self._conversion_results[doc.id] if isinstance(res, Exception): raise res + self.call_count += 1 return res @classmethod @@ -61,16 +74,22 @@ def _from_config(cls, config: RegistrableConfig, **extras) -> FromConfig: class MockPDFPreprocessor(PDFPreprocessor): def __init__(self, res: list[list[Path] | Exception]): self._res = iter(res) + self.processed = [] def __call__( self, - pdf_path: Path, # noqa: ARG002 + pdf_path: Path, pdf_bytes: bytes, # noqa: ARG002 output_dir: Path, # noqa: ARG002 + *, + force_reprocessing: bool, ) -> list[Path]: r = next(self._res) if isinstance(r, Exception): raise r + for im_path in r: + if force_reprocessing and not im_path.exists(): + self.processed.append(pdf_path) return r @@ -119,6 +138,7 @@ def test_preprocess_images_act( output_root=output_root, executor=executor, image_preprocessor=processor, + force_reprocessing=True, ) # Then @@ -136,6 +156,46 @@ def test_preprocess_images_act( assert processing_error.error.title == "UnsupportedDocExtension" +async def test_preprocess_images_act_caching( + test_worker_config: PassportWorkerConfig, + symlinked_doc_0_pages: list[Path], +) -> None: + # Given + config = test_worker_config + executor = test_worker_config.to_image_preprocessing_executor() + worker_paths = config.paths + workdir = worker_paths.workdir + output_root = workdir.joinpath("workflow_id") + output_root.mkdir(parents=True, exist_ok=True) + doc_0_pages = symlinked_doc_0_pages + processor = MockImageProcessor([doc_0_pages]) + batch_path = output_root / "batch.jsonl" + batch = [SYMLINKED_PROCESSED_DOC_0] + batch_path.write_text("\n".join(d.model_dump_json() for d in batch)) + + # When + successes, errors = preprocess_images_act( + batch_path, + worker_paths, + output_root=output_root, + executor=executor, + image_preprocessor=processor, + force_reprocessing=False, + ) + + # Then + assert not processor.processed + expected_successes = [ + ProcessedPage( + page_number=page_number + 1, + **SYMLINKED_PROCESSED_DOC_0.child(p, worker_paths).model_dump(), + ) + for page_number, p in enumerate(doc_0_pages) + ] + assert successes == expected_successes + assert not errors + + async def test_convert_to_pdfs_act( test_worker_config: PassportWorkerConfig, docs_with_cached_artifacts: list[ProcessedFile], # noqa: ARG001 @@ -164,6 +224,7 @@ async def test_convert_to_pdfs_act( worker_paths, max_concurrency=max_concurrency, output_root=output_root, + force_reprocessing=True, ) # Then doc_2_as_pdf_path = ( @@ -177,6 +238,43 @@ async def test_convert_to_pdfs_act( assert processing_error.error.title == "UnsupportedDocExtension" +async def test_convert_to_pdfs_act_caching( + test_worker_config: PassportWorkerConfig, + docs_with_cached_artifacts: list[ProcessedFile], # noqa: ARG001 +) -> None: + # Given + config = test_worker_config + worker_paths = config.paths + workdir = worker_paths.workdir + output_root = workdir.joinpath("workflow_id") + output_root.mkdir(parents=True) + max_concurrency = 1 + batch = [PROCESSED_DOC_2] + conversion_results = {PROCESSED_DOC_2.id: b"doc_2_as_pdf"} + pdf_converter = MockConverter(conversion_results=conversion_results) + batch_path = output_root / "batch.jsonl" + batch_path.write_text("\n".join(d.model_dump_json() for d in batch)) + doc_2_as_pdf_path = ( + output_root / safe_dir(PROCESSED_DOC_2.id) / f"{PROCESSED_DOC_2.id}.pdf" + ) + doc_2_as_pdf_path.parent.mkdir(parents=True, exist_ok=True) + doc_2_as_pdf_path.write_bytes((DOCS_PATH / "passport.pdf").read_bytes()) + # When + successes, errors = await convert_to_pdfs_act( + batch_path, + pdf_converter, + worker_paths, + max_concurrency=max_concurrency, + output_root=output_root, + force_reprocessing=False, + ) + # Then + assert not pdf_converter.call_count + expected_successes = [PROCESSED_DOC_2.child(doc_2_as_pdf_path, worker_paths)] + assert successes == expected_successes + assert not errors + + @pytest.fixture def doc_1_pages( test_worker_config: PassportWorkerConfig, doc_0: Document @@ -212,6 +310,7 @@ async def test_preprocess_pdfs_act( worker_paths, preprocessor, output_root=output_root, + force_reprocessing=True, ) # Then expected_successes = [ @@ -228,15 +327,66 @@ async def test_preprocess_pdfs_act( assert processing_error.error.title == "InvalidPDF" +async def test_preprocess_pdfs_act_caching( + test_worker_config: PassportWorkerConfig, + doc_1_pages: list[Path], + docs_with_cached_artifacts: list[ProcessedFile], # noqa: ARG001 +) -> None: + # Given + config = test_worker_config + worker_paths = config.paths + workdir = worker_paths.workdir + output_root = workdir.joinpath("workflow_id") + output_root.mkdir(parents=True, exist_ok=True) + batch = [PROCESSED_DOC_1] + results = [doc_1_pages] + preprocessor = MockPDFPreprocessor(results) + batch_path = output_root / "pdfs.jsonl" + batch_path.write_text("\n".join(d.model_dump_json() for d in batch)) + for p in doc_1_pages: + p.parent.mkdir(parents=True, exist_ok=True) + p.write_bytes((DOCS_PATH / "passport.png").read_bytes()) + # When + successes, errors = await preprocess_pdfs_act( + batch_path, + worker_paths, + preprocessor, + output_root=output_root, + force_reprocessing=False, + ) + # Then + assert not preprocessor.processed + expected_successes = [ + ProcessedPage( + page_number=page_number + 1, + **PROCESSED_DOC_1.child(p, worker_paths).model_dump(), + ) + for page_number, p in enumerate(doc_1_pages) + ] + assert successes == expected_successes + assert not errors + + def test_default_image_preprocessor(tmpdir: Path) -> None: # Given output_dir = Path(tmpdir) im_path = DOCS_PATH / "not_a_passport.jpg" preprocessor = DefaultImagePreprocessor() # When - paths = preprocessor(im_path, output_dir=output_dir) + paths = preprocessor(im_path, output_dir=output_dir, force_reprocessing=True) assert len(paths) == 1 processed_path = paths[0] assert processed_path.name.endswith(".png") im = Image.open(processed_path) assert im.mode == "RGB" + + +@pytest.mark.parametrize( + ("filename", "expected_is_valid"), + [("not_a_passport.jpg", False), ("idontexist", False), ("passport.pdf", True)], +) +async def test_is_valid_pdf(filename: str, *, expected_is_valid: bool) -> None: + # When + is_valid = await is_valid_pdf(DOCS_PATH / filename) + # Then + assert is_valid == expected_is_valid diff --git a/workers/passport-worker/uv.dist.lock b/workers/passport-worker/uv.dist.lock index 5ec4dada..5d5429e0 100644 --- a/workers/passport-worker/uv.dist.lock +++ b/workers/passport-worker/uv.dist.lock @@ -321,10 +321,10 @@ dev = [ requires-dist = [ { name = "datashare-python", specifier = "~=0.10.6" }, { name = "icij-common", extras = ["elasticsearch"], specifier = "~=0.8.3" }, - { name = "icij-passport-core", specifier = "==0.11.6" }, - { name = "icij-passport-core", extras = ["cpu", "inference"], marker = "extra == 'inference-cpu'", specifier = "==0.11.6" }, - { name = "icij-passport-core", extras = ["gpu", "inference"], marker = "extra == 'inference-gpu'", specifier = "==0.11.6" }, - { name = "icij-passport-core", extras = ["preprocessing"], marker = "extra == 'preprocessing'", specifier = "==0.11.6" }, + { name = "icij-passport-core", specifier = "==0.11.7" }, + { name = "icij-passport-core", extras = ["cpu", "inference"], marker = "extra == 'inference-cpu'", specifier = "==0.11.7" }, + { name = "icij-passport-core", extras = ["gpu", "inference"], marker = "extra == 'inference-gpu'", specifier = "==0.11.7" }, + { name = "icij-passport-core", extras = ["preprocessing"], marker = "extra == 'preprocessing'", specifier = "==0.11.7" }, ] provides-extras = ["inference-cpu", "inference-gpu", "preprocessing"] @@ -515,15 +515,15 @@ elasticsearch = [ [[package]] name = "icij-passport-core" -version = "0.11.6" +version = "0.11.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiostream" }, { name = "icij-common" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/33/36/270d9d9943fc97e818f9186483a56905f506045c64bfb339f342bbb98c1c/icij_passport_core-0.11.6.tar.gz", hash = "sha256:fe8e1a093a4d5c04cb279dac89ffaecf8d7f20b7dfc954c296b3b867f08ef75a", size = 14579056, upload-time = "2026-09-01T11:29:16.086Z" } +sdist = { url = "https://files.pythonhosted.org/packages/11/7a/e265bcca2d3c32e20d0ee8c2697d07a8b610fa25a886f1caadf2955bf7a0/icij_passport_core-0.11.7.tar.gz", hash = "sha256:ae3223691ca73d66767fea90005d57518f9584263a556ec77741485b955a40dd", size = 14579163, upload-time = "2026-09-01T14:50:09.018Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/2c/720e1d9454eb4951bb5c6e5402c07cae6858576844c44df78bd8b7bf9852/icij_passport_core-0.11.6-py3-none-any.whl", hash = "sha256:33da0e314ae281c46c97f71b3a3b29c8e8a66349384e0a8ee37cd6cec0b07069", size = 10959034, upload-time = "2026-09-01T11:29:13.65Z" }, + { url = "https://files.pythonhosted.org/packages/68/46/cb8007661a1326bc0bbe36a8f0206ada8a429031ebb72032f3c064c4559e/icij_passport_core-0.11.7-py3-none-any.whl", hash = "sha256:c34f178666d1dd8e8a084eb9c53d47d9b42402da677f24b92d0ad030942f022e", size = 10959115, upload-time = "2026-09-01T14:50:06.725Z" }, ] [package.optional-dependencies] diff --git a/workers/passport-worker/uv.lock b/workers/passport-worker/uv.lock index 4ab8c506..e397ce44 100644 --- a/workers/passport-worker/uv.lock +++ b/workers/passport-worker/uv.lock @@ -321,10 +321,10 @@ dev = [ requires-dist = [ { name = "datashare-python", editable = "../../datashare-python" }, { name = "icij-common", extras = ["elasticsearch"], specifier = "~=0.8.3" }, - { name = "icij-passport-core", specifier = "==0.11.6" }, - { name = "icij-passport-core", extras = ["cpu", "inference"], marker = "extra == 'inference-cpu'", specifier = "==0.11.6" }, - { name = "icij-passport-core", extras = ["gpu", "inference"], marker = "extra == 'inference-gpu'", specifier = "==0.11.6" }, - { name = "icij-passport-core", extras = ["preprocessing"], marker = "extra == 'preprocessing'", specifier = "==0.11.6" }, + { name = "icij-passport-core", specifier = "==0.11.7" }, + { name = "icij-passport-core", extras = ["cpu", "inference"], marker = "extra == 'inference-cpu'", specifier = "==0.11.7" }, + { name = "icij-passport-core", extras = ["gpu", "inference"], marker = "extra == 'inference-gpu'", specifier = "==0.11.7" }, + { name = "icij-passport-core", extras = ["preprocessing"], marker = "extra == 'preprocessing'", specifier = "==0.11.7" }, ] provides-extras = ["inference-cpu", "inference-gpu", "preprocessing"] @@ -543,15 +543,15 @@ elasticsearch = [ [[package]] name = "icij-passport-core" -version = "0.11.6" +version = "0.11.7" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiostream" }, { name = "icij-common" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/33/36/270d9d9943fc97e818f9186483a56905f506045c64bfb339f342bbb98c1c/icij_passport_core-0.11.6.tar.gz", hash = "sha256:fe8e1a093a4d5c04cb279dac89ffaecf8d7f20b7dfc954c296b3b867f08ef75a", size = 14579056, upload-time = "2026-09-01T11:29:16.086Z" } +sdist = { url = "https://files.pythonhosted.org/packages/11/7a/e265bcca2d3c32e20d0ee8c2697d07a8b610fa25a886f1caadf2955bf7a0/icij_passport_core-0.11.7.tar.gz", hash = "sha256:ae3223691ca73d66767fea90005d57518f9584263a556ec77741485b955a40dd", size = 14579163, upload-time = "2026-09-01T14:50:09.018Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/2c/720e1d9454eb4951bb5c6e5402c07cae6858576844c44df78bd8b7bf9852/icij_passport_core-0.11.6-py3-none-any.whl", hash = "sha256:33da0e314ae281c46c97f71b3a3b29c8e8a66349384e0a8ee37cd6cec0b07069", size = 10959034, upload-time = "2026-09-01T11:29:13.65Z" }, + { url = "https://files.pythonhosted.org/packages/68/46/cb8007661a1326bc0bbe36a8f0206ada8a429031ebb72032f3c064c4559e/icij_passport_core-0.11.7-py3-none-any.whl", hash = "sha256:c34f178666d1dd8e8a084eb9c53d47d9b42402da677f24b92d0ad030942f022e", size = 10959115, upload-time = "2026-09-01T14:50:06.725Z" }, ] [package.optional-dependencies] From 0fddb39da7a2053e43b88b7f6d91650333cc9b87 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cle=CC=81ment=20Doumouro?= Date: Tue, 1 Sep 2026 18:34:56 +0200 Subject: [PATCH 3/3] fix(passport-service): test symlinks --- workers/passport-worker/tests/test_inference.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/workers/passport-worker/tests/test_inference.py b/workers/passport-worker/tests/test_inference.py index a5b1dccd..0aa5111a 100644 --- a/workers/passport-worker/tests/test_inference.py +++ b/workers/passport-worker/tests/test_inference.py @@ -233,7 +233,8 @@ def _mock_pages( for p in batch: page_path = worker_paths.workdir / p.path page_path.parent.mkdir(parents=True, exist_ok=True) - os.symlink(DOCS_PATH / "passport.png", page_path) + if not page_path.exists(): + os.symlink(DOCS_PATH / "passport.png", page_path) for error in errors: # Generate an error invalid_im_path = worker_paths.workdir / error.path