feature(passport-worker): export passport workflow

datashare-python/datashare_python/conftest.py

@@ -48,6 +48,7 @@
 _DEFAULT_AUTH = aiohttp.BasicAuth(login="guest", password="guest", encoding="utf-8")
 
 TEST_PROJECT = "test-project"
+DS_URL = "https://datashare.org"
 
 _INDEX_BODY = {
     "mappings": {

workers/passport-worker/entrypoints/passport_io_worker.sh

@@ -9,4 +9,5 @@ uv run --no-sync datashare-python worker start \
     --activities passport-detection.convert-to-pdf \
     --activities passport-detection.preprocess.pdfs \
     --activities passport-detection.create-inference-batches \
-    --activities passport-detection.aggregate-results
+    --activities passport-detection.aggregate-results \
+    --activities passport-detection.export-passports

workers/passport-worker/passport_worker/activities.py

@@ -26,6 +26,7 @@
     lifespan_image_preprocessor_cache,
     lifespan_pdf_converter_cache,
 )
+from .export import export_passports_act
 from .inference import (
     PassportDetector,
     create_inference_batches_act,
@@ -36,6 +37,8 @@
     DocumentSearchQuery,
     PassportDetectionArgs,
     PassportDetectionResponse,
+    PassportExportArgs,
+    PassportExportResponse,
     PreprocessingBatches,
     PreprocessingConfig,
 )
@@ -68,9 +71,10 @@ class Activity(StrEnum):
     CREATE_INFERENCE_BATCHES = "passport-detection.create-inference-batches"
     DETECT_PASSPORTS = "passport-detection.detect-passports"
     AGGREGATE_RESULTS = "passport-detection.aggregate-results"
+    EXPORT_PASSPORTS = "passport-detection.export-passports"
 
 
-class PassportDetectionActivities(ActivityWithProgress):
+class SharedActivities(ActivityWithProgress):
     @activity_defn(name=Activity.LOAD_WORKER_CONFIG)
     async def worker_config(
         self,
@@ -82,6 +86,8 @@ async def worker_config(
         worker_config = cast(PassportWorkerConfig, lifespan_worker_config())
         return worker_config
 
+
+class PassportDetectionActivities(ActivityWithProgress):
     @activity_defn(name=Activity.CREATE_PREPROCESSING_BATCHES)
     async def create_preprocessing_batches(
         self,
@@ -323,13 +329,44 @@ async def aggregate_results(
         return res
 
 
+class PassportExportActivities(ActivityWithProgress):
+    @activity_defn(name=Activity.EXPORT_PASSPORTS)
+    async def export_passports(
+        self,
+        args: PassportExportArgs,
+        *,
+        progress: Annotated[  # noqa: ARG002
+            AsyncProgressRateHandler | None, Weight(value=_WORKER_CONFIG_WEIGHT)
+        ] = None,
+    ) -> PassportExportResponse:
+        worker_config = cast(PassportWorkerConfig, lifespan_worker_config())
+        paths = worker_config.paths
+        root = paths.artifacts
+        if args.root is not None:
+            root /= args.root
+        output_path = paths.workdir / args.output_path
+        es_client = lifespan_es_client()
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        await export_passports_act(
+            root,
+            args.format,
+            es_client,
+            output_path=output_path,
+            ds_base_url=worker_config.ds_url,
+        )
+        export_path = output_path.relative_to(paths.workdir)
+        res = PassportExportResponse(export_path=export_path)
+        return res
+
+
 ACTIVITIES = [
-    PassportDetectionActivities.worker_config,
+    SharedActivities.worker_config,
     PassportDetectionActivities.create_preprocessing_batches,
     PassportDetectionActivities.preprocess_images,
     PassportDetectionActivities.convert_to_pdfs,
     PassportDetectionActivities.preprocess_pdfs,
     PassportDetectionActivities.create_inference_batches,
     PassportDetectionActivities.detect_passports,
     PassportDetectionActivities.aggregate_results,
+    PassportExportActivities.export_passports,
 ]

workers/passport-worker/passport_worker/config.py

@@ -117,12 +117,17 @@ class PassportTimeouts(BaseModel):
     result_aggregation: ActivityTimeouts = ActivityTimeouts(
         start_to_close=timedelta(minutes=5)
     )
+    export: ActivityTimeouts = ActivityTimeouts(
+        start_to_close=timedelta(minutes=30), heartbeat=timedelta(minutes=5)
+    )
 
 
 class PassportWorkerConfig(WorkerConfig):
     logging: LoggingConfig = _DEFAULT_LOGGING_CONFIG
     paths: WorkerPaths
 
+    ds_url: str
+
     timeouts: PassportTimeouts = Field(default_factory=PassportTimeouts)
     cache: PassportWorkerCacheConfig = Field(default_factory=PassportWorkerCacheConfig)
 

workers/passport-worker/passport_worker/export.py

@@ -0,0 +1,140 @@
+import json
+import os
+from collections import defaultdict
+from collections.abc import AsyncIterable, Iterable
+from io import StringIO
+from itertools import batched
+from pathlib import Path
+from typing import Any, Never
+
+import pandas as pd
+from aiofile import async_open
+from datashare_python.constants import MANIFEST_JSON
+from datashare_python.objects import ArtifactType, DocID
+from icij_common.es import DOC_ROOT_ID, HITS, ID_, QUERY, SOURCE, ESClient, has_id
+
+from .objects import ExportFormat, PassportArtifact, Passports
+
+
+async def export_passports_act(
+    root: Path,
+    export_format: ExportFormat,
+    es_client: ESClient,
+    *,
+    output_path: Path,
+    ds_base_url: str,
+) -> None:
+    match export_format:
+        case ExportFormat.CSV:
+            lines = export_as_csv(root, ds_base_url=ds_base_url, es_client=es_client)
+        case ExportFormat.JSONL:
+            lines = export_as_jsonl(root)
+        case _:
+            raise ValueError(f"unsupported export format {export_format}")
+    async with async_open(output_path, "w") as f:
+        async for line in lines:
+            await f.write(line + "\n")
+
+
+async def export_as_jsonl(root: Path) -> AsyncIterable[str]:
+    async for _, passports in _walk_passports(root):
+        yield passports.model_dump_json()
+
+
+async def export_as_csv(
+    root: Path, ds_base_url: str, es_client: ESClient
+) -> AsyncIterable[str]:
+    out = StringIO()
+    records_by_project = defaultdict(list)
+    async for r in _read_csv_passport_records(root):
+        records_by_project[r["project"]].append(r)
+    recs = []
+    for project, records in records_by_project.items():
+        root_ids = dict()
+        batches = batched((d["doc_id"] for d in records), 1000)
+        for doc_ids in batches:
+            root_ids |= {
+                k: v async for k, v in _fetch_root_ids(es_client, project, doc_ids)
+            }
+        for r in records:
+            doc_id = r["doc_id"]
+            r["root_id"] = root_ids[doc_id]
+            r["url"] = _make_url(
+                base_url=ds_base_url,
+                doc_id=doc_id,
+                root_id=r["root_id"],
+                project=project,
+            )
+            recs.append(r)
+    df = pd.DataFrame.from_records(recs).sort_values(by=["doc_id"])
+    df.to_csv(out, index=False)
+    out.flush()
+    for line in out.getvalue().split("\n"):
+        if line.strip():
+            yield line
+
+
+async def _read_csv_passport_records(root: Path) -> AsyncIterable[dict[str, Any]]:
+    async for (project, doc_id), passports in _walk_passports(root):
+        for doc_page in passports.pages:
+            page_i = doc_page.page_number
+            for passport_page in doc_page.passports:
+                as_record = {"page": page_i, "doc_id": doc_id, "project": project}
+                as_record["confidence"] = passport_page.confidence
+                mrz = passport_page.mrz
+                if mrz is not None:
+                    as_record["mzr_country"] = mrz.country
+                    for k, v in mrz.metadata.items():
+                        as_record[f"mrz_{k}"] = v
+                yield as_record
+
+
+def _make_url(*, base_url: str, doc_id: str, root_id: str | None, project: str) -> str:
+    if root_id is None:
+        root_id = doc_id
+    return f"{base_url}/#/ds/{project}/{doc_id}/{root_id}?tab=viewer"
+
+
+async def _walk_passports(
+    root: Path,
+) -> AsyncIterable[tuple[tuple[str, DocID], Passports]]:
+    # For efficiency, we first check if there's a passport file then open the manifest
+    # to see if it's referenced
+    for artifact_root, _, files in os.walk(root, onerror=_raise):
+        artifact_root = Path(artifact_root)  # noqa: PLW2901
+        for f in files:
+            if f != PassportArtifact.filename:
+                continue
+            manifest_path = artifact_root / MANIFEST_JSON
+            async with async_open(manifest_path) as m_f:
+                manifest_json = await m_f.read()
+            manifest = json.loads(manifest_json)
+            if ArtifactType.PASSPORTS.value not in manifest:
+                continue
+            passports_path = artifact_root / f
+            async with async_open(passports_path) as p_f:
+                passports_json = await p_f.read()
+            passports = Passports.model_validate_json(passports_json)
+            doc_id = manifest_path.parent.name
+            project = manifest_path.parents[3].name
+            yield (project, doc_id), passports
+
+
+def _raise(e: OSError) -> Never:
+    raise e
+
+
+_DOC_ID_SOURCES = [DOC_ROOT_ID]
+
+
+async def _fetch_root_ids(
+    es_client: ESClient, project: str, batch: Iterable[str]
+) -> AsyncIterable[tuple[str, str]]:
+    query = {QUERY: has_id(list(batch))}
+    # We can't use mget because we would need the root id which is precisely what we're
+    # looking for
+    async for page in es_client.poll_search_pages(
+        index=project, body=query, _source_includes=_DOC_ID_SOURCES
+    ):
+        for hit in page[HITS][HITS]:
+            yield hit[ID_], hit[SOURCE][DOC_ROOT_ID]

workers/passport-worker/passport_worker/objects.py

@@ -38,6 +38,11 @@
 OnnxProvider = Literal["CUDAExecutionProvider", "CPUExecutionProvider"]
 
 
+class ExportFormat(StrEnum):
+    CSV = "csv"
+    JSONL = "jsonl"
+
+
 @cache
 def default_country_codes() -> list[str]:
     import passport_service  # noqa: PLC0415
@@ -133,6 +138,13 @@ def as_manifest_task_input(self) -> dict[str, Any]:
         return as_entry
 
 
+class PassportExportArgs(TaskArgs):
+    project: str
+    output_path: Path
+    format: ExportFormat = ExportFormat.CSV
+    root: Path | None = None
+
+
 class PreprocessingBatches(BaseModel):
     to_pdf: Batches
     images: Batches
@@ -210,3 +222,7 @@ def aggregate(
             (r.successes for r in inference_results), start=ProcessingReport()
         )
         return cls(processed=processed, successes=successes, errors=errors)
+
+
+class PassportExportResponse(DatashareModel):
+    export_path: Path

workers/passport-worker/passport_worker/workflows.py

@@ -13,12 +13,18 @@
     from datashare_python.config import ActivityTimeouts
     from datashare_python.utils import WorkflowWithProgress, execute_activity
 
-    from .activities import PassportDetectionActivities
+    from .activities import (
+        PassportDetectionActivities,
+        PassportExportActivities,
+        SharedActivities,
+    )
     from .config import PassportWorkerConfig
     from .objects import (
         Batches,
         PassportDetectionArgs,
         PassportDetectionResponse,
+        PassportExportArgs,
+        PassportExportResponse,
         PreprocessingBatches,
         PreprocessingConfig,
     )
@@ -44,16 +50,15 @@ class PassportDetectionWorkflow(WorkflowWithProgress):
     @workflow.run
     async def run(self, args: PassportDetectionArgs) -> PassportDetectionResponse:
         logger.info("creating preprocessing batches...")
-        batch_args = [args.docs, args.project]
         # Fetch worker config
         worker_config = await execute_activity(
-            PassportDetectionActivities.worker_config,
-            args=batch_args,
+            SharedActivities.worker_config,
             task_queue=TaskQueue.IO,
             timeouts=ActivityTimeouts(start_to_close=timedelta(minutes=5)),
         )
         worker_config = cast(PassportWorkerConfig, worker_config)
         # Create preprocessing batches
+        batch_args = [args.docs, args.project]
         preprocessing_batches = await execute_activity(
             PassportDetectionActivities.create_preprocessing_batches,
             args=batch_args,
@@ -97,6 +102,25 @@ async def run(self, args: PassportDetectionArgs) -> PassportDetectionResponse:
         return response
 
 
+@workflow.defn(name="passport-detection.export-passports")
+class ExportWorkflow(WorkflowWithProgress):
+    @workflow.run
+    async def run(self, args: PassportExportArgs) -> PassportExportResponse:
+        logger.info("creating preprocessing batches...")
+        worker_config = await execute_activity(
+            SharedActivities.worker_config,
+            task_queue=TaskQueue.IO,
+            timeouts=ActivityTimeouts(start_to_close=timedelta(minutes=5)),
+        )
+        res = await execute_activity(
+            PassportExportActivities.export_passports,
+            args=(args,),
+            task_queue=TaskQueue.IO,
+            timeouts=worker_config.timeouts.export,
+        )
+        return res
+
+
 async def preprocess(
     args: PassportDetectionArgs,
     preprocessing_batches: PreprocessingBatches,
@@ -219,4 +243,4 @@ def _process_pdfs_tasks(
     return all_tasks
 
 
-WORKFLOWS = [PassportDetectionWorkflow]
+WORKFLOWS = [PassportDetectionWorkflow, ExportWorkflow]

workers/passport-worker/pyproject.toml

@@ -17,6 +17,7 @@ dependencies = [
 [project.optional-dependencies]
 preprocessing = [
   "icij-passport-core[preprocessing]==0.11.8",
+  "pandas~=3.0",
 ]
 inference-gpu = [
   "icij-passport-core[inference,gpu]==0.11.8"