datashare-passport-worker 0.1.2__tar.gz

datashare_passport_worker-0.1.2/.gitignore

@@ -0,0 +1,21 @@
+.idea/
+# Python
+*.log
+venv
+*.egg-info
+.eggs
+__pycache__
+*.pytest_cache
+*.pyc
+build
+dist
+test/.env
+.cache
+tmp
+./*.csv
+.DS_Store
+
+# VS code
+.vscode
+# Doc
+site

datashare_passport_worker-0.1.2/PKG-INFO

@@ -0,0 +1,14 @@
+Metadata-Version: 2.4
+Name: datashare-passport-worker
+Version: 0.1.2
+Summary: Datashare passport detection worker
+Author-email: Clément Doumouro <cdoumouro@icij.org>, Clément Doumouro <clement.doumouro@gmail.com>
+Requires-Python: <3.14,>=3.12
+Requires-Dist: datashare-python~=0.10.0
+Requires-Dist: icij-common[elasticsearch]~=0.8.3
+Provides-Extra: inference-cpu
+Requires-Dist: icij-passport-core[cpu,inference]==0.10.9; extra == 'inference-cpu'
+Provides-Extra: inference-gpu
+Requires-Dist: icij-passport-core[gpu,inference]==0.10.9; extra == 'inference-gpu'
+Provides-Extra: preprocessing
+Requires-Dist: icij-passport-core[preprocessing]==0.10.9; extra == 'preprocessing'

datashare_passport_worker-0.1.2/passport_worker/activities.py

@@ -0,0 +1,305 @@
+import logging
+from enum import StrEnum
+from functools import partial
+from pathlib import Path
+from typing import Annotated, cast
+
+from aiofile import async_open
+from datashare_python.dependencies import lifespan_es_client, lifespan_worker_config
+from datashare_python.types_ import (
+    AsyncProgressRateHandler,
+    SyncProgressRateHandler,
+    Weight,
+)
+from datashare_python.utils import (
+    ActivityWithProgress,
+    activity_defn,
+    activity_workdir,
+    async_enter_cm,
+    config_cache_key,
+    enter_cm,
+)
+
+from .aggregate import aggregate_results_act
+from .config import PassportWorkerConfig
+from .dependencies import (
+    lifespan_image_preprocessor_cache,
+    lifespan_passport_detector_cache,
+    lifespan_pdf_converter_cache,
+)
+from .inference import (
+    PassportDetector,
+    create_inference_batches_act,
+    detect_passports_act,
+)
+from .objects import (
+    DocId,
+    DocumentSearchQuery,
+    ImagePreprocessorConfig,
+    PassportDetectionArgs,
+    PassportDetectionResponse,
+    PDFConverterConfig,
+    PreprocessingBatches,
+)
+from .preprocessing import (
+    ImagePreprocessor,
+    PDFConverter,
+    convert_to_pdfs_act,
+    preprocess_images_act,
+    preprocess_pdfs_act,
+)
+from .search import create_preprocessing_batches_act
+
+logger = logging.getLogger(__name__)
+
+_BASE_WEIGHT = 1.0
+_CREATE_PREPROCESSING_BATCHES_WEIGHT = _BASE_WEIGHT * 1
+_PREPROCESS_IMAGES_WEIGHT = _BASE_WEIGHT * 5
+_CONVERT_TO_PDF_WEIGHT = _BASE_WEIGHT * 7
+_PREPROCESS_PDF_WEIGHT = _BASE_WEIGHT * 3
+_CREATE_INFERENCE_BATCH_WEIGHT = _CREATE_PREPROCESSING_BATCHES_WEIGHT * 1
+
+
+class Activity(StrEnum):
+    CREATE_PREPROCESSING_BATCHES = "passport-detection.create-preprocessing-batches"
+    PREPROCESS_IMAGES = "passport-detection.preprocess.images"
+    CONVERT_TO_PDFS = "passport-detection.convert-to-pdf"
+    PREPROCESS_PDFS = "passport-detection.preprocess.pdfs"
+    CREATE_INFERENCE_BATCHES = "passport-detection.create-inference-batches"
+    DETECT_PASSPORTS = "passport-detection.detect-passports"
+    AGGREGATE_RESULTS = "passport-detection.aggregate-results"
+
+
+class PassportDetectionActivities(ActivityWithProgress):
+    @activity_defn(name=Activity.CREATE_PREPROCESSING_BATCHES)
+    async def create_preprocessing_batches(
+        self,
+        docs: list[DocId] | DocumentSearchQuery | None,
+        project: str,
+        *,
+        progress: Annotated[  # noqa: ARG002
+            AsyncProgressRateHandler | None,
+            Weight(value=_CREATE_PREPROCESSING_BATCHES_WEIGHT),
+        ] = None,
+    ) -> PreprocessingBatches:
+        es_client = lifespan_es_client()
+        worker_config = cast(PassportWorkerConfig, lifespan_worker_config())
+        workdir = worker_config.paths.workdir
+        output_root = activity_workdir(workdir, project, act_context=True)
+        output_root.mkdir(parents=True, exist_ok=True)
+        target_n_pages_per_batch = worker_config.preprocessing.target_n_pages_per_batch
+        return await create_preprocessing_batches_act(
+            docs,
+            project,
+            es_client,
+            worker_config.paths,
+            target_n_pages_per_batch,
+            output_root=output_root,
+        )
+
+    @activity_defn(name=Activity.PREPROCESS_IMAGES)
+    def preprocess_images(
+        self,
+        batch: Path,
+        project: str,
+        config: ImagePreprocessorConfig,
+        *,
+        progress: Annotated[
+            SyncProgressRateHandler | None, Weight(value=_PREPROCESS_IMAGES_WEIGHT)
+        ] = None,
+    ) -> tuple[Path, Path]:
+        worker_config = cast(PassportWorkerConfig, lifespan_worker_config())
+        workdir = worker_config.paths.workdir
+        logger.info("loading image preprocessor...")
+        cache = lifespan_image_preprocessor_cache()
+        image_preprocessor_cache_key = config_cache_key(config)
+        image_preprocessor_factory = enter_cm(
+            partial(ImagePreprocessor.from_config, config)
+        )
+        image_preprocessor = cache.get_or_cache_resource(
+            image_preprocessor_cache_key, image_preprocessor_factory
+        )
+        logger.info("loaded image preprocessor !")
+        pages_root = activity_workdir(workdir, project, act_context=False)
+        pages_root.mkdir(parents=True, exist_ok=True)
+        executor = worker_config.to_image_preprocessing_executor()
+        chunk_size = worker_config.preprocessing.images.chunk_size
+        success, errors = preprocess_images_act(
+            batch,
+            worker_config.paths,
+            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.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))
+        errors_path = res_root / "errors.jsonl"
+        errors_path.write_text("\n".join(p.model_dump_json() for p in errors))
+        return successes_path, errors_path
+
+    @activity_defn(name=Activity.CONVERT_TO_PDFS)
+    async def convert_to_pdfs(
+        self,
+        batch: Path,
+        project: str,
+        config: PDFConverterConfig,
+        *,
+        progress: Annotated[
+            AsyncProgressRateHandler | None, Weight(value=_CONVERT_TO_PDF_WEIGHT)
+        ] = None,
+    ) -> tuple[Path, Path]:
+        worker_config = cast(PassportWorkerConfig, lifespan_worker_config())
+        cache = lifespan_pdf_converter_cache()
+        pdf_converter_cache_key = config_cache_key(config)
+        pdf_converter_factory = async_enter_cm(
+            partial(PDFConverter.from_config, config)
+        )
+        pdf_converter = await cache.async_get_or_cache_resource(
+            pdf_converter_cache_key, pdf_converter_factory
+        )
+        workdir = worker_config.paths.workdir
+        pdfs_root = activity_workdir(workdir, project, act_context=False)
+        pdfs_root.mkdir(parents=True, exist_ok=True)
+        max_concurrency = worker_config.preprocessing.pdfs.max_concurrency
+        successes, errors = await convert_to_pdfs_act(
+            batch,
+            pdf_converter,
+            worker_config.paths,
+            max_concurrency,
+            output_root=pdfs_root,
+            progress=progress,
+        )
+        res_root = activity_workdir(workdir, project, act_context=True)
+        res_root.mkdir(parents=True, exist_ok=True)
+        pdf_paths = res_root / "pdfs.jsonl"
+        async with async_open(pdf_paths, "w") as f:
+            await f.write("\n".join(d.model_dump_json() for d in successes))
+        errors_path = res_root / "errors.jsonl"
+        async with async_open(errors_path, "w") as f:
+            await f.write("\n".join(e.model_dump_json() for e in errors))
+        return pdf_paths, errors_path
+
+    @activity_defn(name=Activity.PREPROCESS_PDFS)
+    async def preprocess_pdfs(
+        self,
+        batch: Path,
+        project: str,
+        *,
+        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)
+        output_root.mkdir(parents=True, exist_ok=True)
+        successes, errors = await preprocess_pdfs_act(
+            batch, worker_config.paths, output_root=output_root, progress=progress
+        )
+        res_root = activity_workdir(workdir, project, act_context=True)
+        res_root.mkdir(parents=True, exist_ok=True)
+        pdf_paths = res_root / "pdfs.jsonl"
+        async with async_open(pdf_paths, "w") as f:
+            await f.write("\n".join(d.model_dump_json() for d in successes))
+        errors_path = res_root / "errors.jsonl"
+        async with async_open(errors_path, "w") as f:
+            await f.write("\n".join(e.model_dump_json() for e in errors))
+        return pdf_paths, errors_path
+
+    @activity_defn(name=Activity.CREATE_INFERENCE_BATCHES)
+    async def create_inference_batches(
+        self,
+        batches: list[Path],
+        project: str,
+        *,
+        progress: Annotated[  # noqa:ARG002
+            AsyncProgressRateHandler | None,
+            Weight(value=_CREATE_INFERENCE_BATCH_WEIGHT),
+        ] = None,
+    ) -> list[Path]:
+        worker_config = cast(PassportWorkerConfig, lifespan_worker_config())
+        batch_size = worker_config.inference.batch_size
+        batches_per_task = worker_config.inference.batches_per_task
+        worker_paths = worker_config.paths
+        output_root = activity_workdir(worker_paths.workdir, project)
+        output_root.mkdir(parents=True, exist_ok=True)
+        return await create_inference_batches_act(
+            batches,
+            worker_paths,
+            output_root,
+            target_batches_per_task=batches_per_task,
+            inference_batch_size=batch_size,
+        )
+
+    @activity_defn(name=Activity.DETECT_PASSPORTS)
+    async def detect_passports(
+        self,
+        batch: Path,
+        args: PassportDetectionArgs,
+        *,
+        progress: Annotated[
+            AsyncProgressRateHandler | None,
+            Weight(value=_CREATE_INFERENCE_BATCH_WEIGHT),
+        ] = None,
+    ) -> Path:
+        logger.info("loading passport detector...")
+        worker_config = cast(PassportWorkerConfig, lifespan_worker_config())
+        batch_size = worker_config.inference.batch_size
+        cache = lifespan_passport_detector_cache()
+        passport_detector_config = args.config.inference.passport_detector
+        passport_detector_key = config_cache_key(passport_detector_config)
+        passport_detector_factory = enter_cm(
+            partial(PassportDetector.from_config, passport_detector_config)
+        )
+        passport_detector = cache.get_or_cache_resource(
+            passport_detector_key, passport_detector_factory
+        )
+        logger.info("passport detector loaded !")
+        workdir = worker_config.paths.workdir
+        res_root = activity_workdir(workdir, args.project, act_context=True)
+        res_root.mkdir(parents=True, exist_ok=True)
+        res = await detect_passports_act(
+            batch,
+            passport_detector,
+            worker_config.paths,
+            args,
+            batch_size=batch_size,
+            progress=progress,
+        )
+        result_path = res_root / "inference_results.json"
+        async with async_open(result_path, "w") as f:
+            await f.write(res.model_dump_json())
+        return result_path
+
+    @activity_defn(name=Activity.AGGREGATE_RESULTS)
+    async def aggregate_results(
+        self,
+        error_paths: list[Path],
+        result_paths: list[Path],
+        *,
+        progress: Annotated[  # noqa:ARG002
+            AsyncProgressRateHandler | None,
+            Weight(value=_CREATE_INFERENCE_BATCH_WEIGHT),
+        ] = None,
+    ) -> PassportDetectionResponse:
+        worker_config = cast(PassportWorkerConfig, lifespan_worker_config())
+        res = await aggregate_results_act(
+            error_paths, result_paths=result_paths, paths=worker_config.paths
+        )
+        return res
+
+
+ACTIVITIES = [
+    PassportDetectionActivities.create_preprocessing_batches,
+    PassportDetectionActivities.preprocess_images,
+    PassportDetectionActivities.convert_to_pdfs,
+    PassportDetectionActivities.preprocess_pdfs,
+    PassportDetectionActivities.create_inference_batches,
+    PassportDetectionActivities.detect_passports,
+    PassportDetectionActivities.aggregate_results,
+]

datashare_passport_worker-0.1.2/passport_worker/aggregate.py

@@ -0,0 +1,44 @@
+import asyncio
+from collections.abc import AsyncIterable
+from pathlib import Path
+from typing import TypeVar
+
+from datashare_python.objects import WorkerPaths
+from datashare_python.utils import async_read_jsonl_as
+
+from passport_worker.objects import (
+    FileProcessingError,
+    PartialDetectionResult,
+    PassportDetectionResponse,
+)
+
+
+async def aggregate_results_act(
+    error_paths: list[Path], *, result_paths: list[Path], paths: WorkerPaths
+) -> PassportDetectionResponse:
+    workdir = paths.workdir
+    preprocessing_errors = await asyncio.gather(
+        *(
+            _as_list(async_read_jsonl_as(workdir / p, FileProcessingError))
+            for p in error_paths
+        )
+    )
+    preprocessing_errors = sum(preprocessing_errors, [])
+    inference_results = await asyncio.gather(
+        *(
+            _as_list(async_read_jsonl_as(workdir / p, PartialDetectionResult))
+            for p in result_paths
+        )
+    )
+    inference_results = sum(inference_results, [])
+    res = PassportDetectionResponse.aggregate(
+        preprocessing_errors, inference_results=inference_results
+    )
+    return res
+
+
+T = TypeVar("T")
+
+
+async def _as_list[T](iterable: AsyncIterable[T]) -> list[T]:
+    return [i async for i in iterable]

datashare_passport_worker-0.1.2/passport_worker/config.py

@@ -0,0 +1,82 @@
+from concurrent.futures import ProcessPoolExecutor
+
+import datashare_python
+from datashare_python.config import (
+    LogFormat,
+    LoggingConfig,
+    ResourceCacheConfig,
+    WorkerConfig,
+)
+from datashare_python.objects import DatashareModel, WorkerPaths
+from pydantic import Field
+
+_ALL_LOGGERS = [datashare_python.__name__, __name__, "__main__"]
+
+_DEFAULT_LOGGERS = {
+    datashare_python.__name__: "INFO",
+    __name__: "INFO",
+    "__main__": "INFO",
+    "passport_service": "INFO",
+}
+_DEFAULT_LOGGING_CONFIG = LoggingConfig(
+    format=LogFormat.DEFAULT, loggers=_DEFAULT_LOGGERS
+)
+
+
+class ImagePreprocessingWorkerConfig(DatashareModel):
+    n_processes: int | None = None
+    chunk_size: int = 5
+
+    def to_image_preprocessing_executor(self) -> ProcessPoolExecutor:
+        return ProcessPoolExecutor(max_workers=self.n_processes)
+
+
+class PDFConversionWorkerConfig(DatashareModel):
+    max_concurrency: int = 10
+
+
+class PreprocessingWorkerConfig(DatashareModel):
+    target_n_pages_per_batch: int = 200
+
+    images: ImagePreprocessingWorkerConfig = Field(
+        default_factory=ImagePreprocessingWorkerConfig
+    )
+    pdfs: PDFConversionWorkerConfig = Field(default_factory=PDFConversionWorkerConfig)
+
+    def to_image_preprocessing_executor(self) -> ProcessPoolExecutor:
+        return self.images.to_image_preprocessing_executor()
+
+
+class InferenceWorkerConfig(DatashareModel):
+    batch_size: int = 32
+    batches_per_task: int = 5
+
+
+class PreprocessingCacheConfig(DatashareModel):
+    pdf: ResourceCacheConfig = Field(default_factory=ResourceCacheConfig)
+    images: ResourceCacheConfig = Field(default_factory=ResourceCacheConfig)
+
+
+class PassportWorkerCacheConfig(DatashareModel):
+    preprocessing: PreprocessingCacheConfig = Field(
+        default_factory=PreprocessingCacheConfig
+    )
+    inference: ResourceCacheConfig = Field(default_factory=ResourceCacheConfig)
+
+
+class PassportWorkerConfig(WorkerConfig):
+    logging: LoggingConfig = _DEFAULT_LOGGING_CONFIG
+    paths: WorkerPaths
+
+    cache: PassportWorkerCacheConfig = Field(default_factory=PassportWorkerCacheConfig)
+
+    preprocessing: PreprocessingWorkerConfig = Field(
+        default_factory=PreprocessingWorkerConfig
+    )
+    inference: InferenceWorkerConfig = Field(default_factory=InferenceWorkerConfig)
+
+    def to_image_preprocessing_executor(self) -> ProcessPoolExecutor:
+        return self.preprocessing.to_image_preprocessing_executor()
+
+
+WORKER_CONFIG_CLS = PassportWorkerConfig

datashare_passport_worker-0.1.2/passport_worker/constants.py

@@ -0,0 +1,6 @@
+def pil_supported_extensions() -> set[str]:
+    from passport_service.core.preprocessing import (  # noqa:PLC0415
+        PIL_SUPPORTED_EXTENSIONS,
+    )
+
+    return set(PIL_SUPPORTED_EXTENSIONS)

datashare_passport_worker-0.1.2/passport_worker/dependencies.py

@@ -0,0 +1,67 @@
+from contextvars import ContextVar
+
+from datashare_python.dependencies import set_es_client, set_loggers, set_worker_config
+from datashare_python.exceptions import DependencyInjectionError
+from datashare_python.utils import SharedResources
+
+from passport_worker.config import PassportWorkerConfig
+
+_IMAGE_PREPROCESSOR_CACHE: ContextVar[SharedResources] = ContextVar(
+    "image_preprocessor_cache"
+)
+_PDF_CONVERTER_CACHE: ContextVar[SharedResources] = ContextVar("pdf_converter_cache")
+_PASSPORT_DETECTOR_CACHE: ContextVar[SharedResources] = ContextVar("preprocessor_cache")
+
+
+def set_image_preprocessor_cache(
+    worker_config: PassportWorkerConfig,
+) -> SharedResources:
+    cache = worker_config.cache.preprocessing.images.to_resource_cache()
+    _IMAGE_PREPROCESSOR_CACHE.set(cache)
+    return cache
+
+
+def lifespan_image_preprocessor_cache() -> SharedResources:
+    try:
+        return _IMAGE_PREPROCESSOR_CACHE.get()
+    except LookupError as e:
+        raise DependencyInjectionError("image preprocessor") from e
+
+
+def set_pdf_converter_cache(
+    worker_config: PassportWorkerConfig,
+) -> SharedResources:
+    cache = worker_config.cache.preprocessing.pdf.to_resource_cache()
+    _PDF_CONVERTER_CACHE.set(cache)
+    return cache
+
+
+def lifespan_pdf_converter_cache() -> SharedResources:
+    try:
+        return _PDF_CONVERTER_CACHE.get()
+    except LookupError as e:
+        raise DependencyInjectionError("pdf converter cache") from e
+
+
+def set_passport_detector_cache(worker_config: PassportWorkerConfig) -> SharedResources:
+    cache = worker_config.cache.inference.to_resource_cache()
+    _PASSPORT_DETECTOR_CACHE.set(cache)
+    return cache
+
+
+def lifespan_passport_detector_cache() -> SharedResources:
+    try:
+        return _PASSPORT_DETECTOR_CACHE.get()
+    except LookupError as e:
+        raise DependencyInjectionError("passport detector cache") from e
+
+
+IO = [set_worker_config, set_loggers, set_es_client, set_pdf_converter_cache]
+PREPROCESSING = [set_worker_config, set_loggers, set_image_preprocessor_cache]
+INFERENCE = [set_worker_config, set_loggers, set_passport_detector_cache]
+
+DEPENDENCIES = {
+    "passport-detection.io": IO,
+    "passport-detection.preprocessing": PREPROCESSING,
+    "passport-detection.inference": INFERENCE,
+}