feature(passport service): preprocessing caching

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())

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,22 +535,30 @@ 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_key: str | None = None,
 ) -> Path:
     act_info = activity.info()
     path = []
     if not wf_context and not act_context:
         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_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)
-        path.append(act_info.activity_id)
-        if run_context:
-            path.append(act_info.activity_run_id)
+        if not caching_key:
+            path.append(act_info.activity_id)
+            if run_context:
+                path.append(act_info.activity_run_id)
+    if caching_key:
+        path += [_CACHE_DIR, caching_key]
     return Path(*path)
 
 
@@ -559,9 +569,13 @@ def activity_workdir(
     wf_context: bool = True,
     act_context: bool = True,
     run_context: bool = False,
+    caching_key: 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_key=caching_key,
     )
     return workdir.joinpath(project, ctx_path)
 
@@ -808,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))

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

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,90 @@ 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_key",
+            ),
+        ),
+        (
+            True,
+            True,
+            False,
+            Path(
+                "workdir",
+                TEST_PROJECT,
+                "wf_type",
+                "act_type",
+                "processing_cache",
+                "caching_key",
+            ),
+        ),
+        (
+            True,
+            False,
+            True,
+            Path("workdir", TEST_PROJECT, "wf_type", "processing_cache", "caching_key"),
+        ),
+        (
+            True,
+            False,
+            False,
+            Path("workdir", TEST_PROJECT, "wf_type", "processing_cache", "caching_key"),
+        ),
+        (
+            False,
+            True,
+            True,
+            Path(
+                "workdir", TEST_PROJECT, "act_type", "processing_cache", "caching_key"
+            ),
+        ),
+        (
+            False,
+            True,
+            False,
+            Path(
+                "workdir", TEST_PROJECT, "act_type", "processing_cache", "caching_key"
+            ),
+        ),
+    ],
+)
+def test_cached_activity_workdir(
+    *,
+    wf_context: bool,
+    act_context: bool,
+    run_context: bool,
+    expected_work_dir: Path,
+    monkeypatch: MonkeyPatch,
+) -> None:
+    # Given
+    caching_key = "caching_key"
+    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_key=caching_key,
+        wf_context=wf_context,
+        act_context=act_context,
+        run_context=run_context,
+    )
+    # Then
+    assert workdir == expected_work_dir

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,29 +111,33 @@ 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,
             chunk_size=chunk_size,
             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,29 +151,33 @@ async def convert_to_pdfs(
         batch: Path,
         project: str,
         *,
+        force_reprocessing: bool,
         progress: Annotated[
             AsyncProgressRateHandler | None, Weight(value=_CONVERT_TO_PDF_WEIGHT)
         ] = None,
     ) -> tuple[Path, Path]:
         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,
             pdf_converter,
             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)

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
     )