feature(passport-worker): report progress

Author: ClemDoum

workers/passport-worker/passport_worker/activities.py

@@ -18,7 +18,9 @@
     async_enter_cm,
     config_cache_key,
     enter_cm,
+    to_raw_async_progress,
 )
+from icij_common.es import COUNT
 
 from .aggregate import aggregate_results_act
 from .config import PassportWorkerConfig
@@ -336,21 +338,29 @@ async def export_passports(
         args: PassportExportArgs,
         *,
         progress: Annotated[  # noqa: ARG002
-            AsyncProgressRateHandler | None, Weight(value=_WORKER_CONFIG_WEIGHT)
+            AsyncProgressRateHandler | None, Weight(value=_BASE_WEIGHT)
         ] = None,
     ) -> PassportExportResponse:
         worker_config = cast(PassportWorkerConfig, lifespan_worker_config())
         paths = worker_config.paths
         output_path = paths.workdir / args.output_path
         es_client = lifespan_es_client()
         output_path.parent.mkdir(parents=True, exist_ok=True)
+        n_docs = await es_client.count(args.docs)
         docs = (d async for d in get_docs(args.docs, args.project, es_client))
+        n_docs = n_docs[COUNT]
+        if progress is not None:
+            if n_docs:
+                progress = to_raw_async_progress(progress, max_progress=n_docs)
+            else:
+                progress = None
         await export_passports_act(
             docs,
             args.format,
             paths,
             output_path=output_path,
             ds_base_url=worker_config.ds_url,
+            progress=progress,
         )
         export_path = output_path.relative_to(paths.workdir)
         res = PassportExportResponse(export_path=export_path)

workers/passport-worker/passport_worker/export.py

@@ -1,7 +1,6 @@
 import logging
 from collections import defaultdict
 from collections.abc import AsyncIterable, Iterable
-from contextlib import suppress
 from io import StringIO
 from pathlib import Path
 from typing import Any, Never
@@ -10,8 +9,8 @@
 from aiofile import async_open
 from datashare_python.exceptions import ArtifactNotFound
 from datashare_python.objects import Document, WorkerPaths
+from datashare_python.types_ import RawAsyncProgressHandler
 from datashare_python.utils import async_read_artifact
-from temporalio import activity
 
 from .objects import ExportFormat, PassportArtifact, Passports
 
@@ -25,40 +24,51 @@ async def export_passports_act(
     *,
     output_path: Path,
     ds_base_url: str,
+    progress: RawAsyncProgressHandler | None = None,
 ) -> None:
     match export_format:
         case ExportFormat.CSV:
-            lines = export_as_csv(docs, ds_base_url=ds_base_url, paths=paths)
+            lines = export_as_csv(
+                docs, ds_base_url=ds_base_url, paths=paths, progress=progress
+            )
         case ExportFormat.JSONL:
-            lines = export_as_jsonl(docs, paths)
+            lines = export_as_jsonl(docs, paths, progress=progress)
         case _:
             raise ValueError(f"unsupported export format {export_format}")
-    l_count = 0
     async with async_open(output_path, "w") as f:
         async for line in lines:
             await f.write(line + "\n")
-            l_count += 1
-        if l_count % 1000 == 0:
-            with suppress(RuntimeError):
-                activity.heartbeat()
 
 
 async def export_as_jsonl(
-    docs: AsyncIterable[Document], paths: WorkerPaths
+    docs: AsyncIterable[Document],
+    paths: WorkerPaths,
+    progress: RawAsyncProgressHandler | None = None,
 ) -> AsyncIterable[str]:
+    n_docs = 0
     async for _, p in _read_passport_artifacts(paths, docs):
         yield p
+        n_docs += 1
+        if progress is not None and n_docs % 100 == 0:
+            await progress(n_docs)
 
 
 async def export_as_csv(
-    docs: AsyncIterable[Document], ds_base_url: str, paths: WorkerPaths
+    docs: AsyncIterable[Document],
+    ds_base_url: str,
+    paths: WorkerPaths,
+    progress: RawAsyncProgressHandler | None = None,
 ) -> AsyncIterable[str]:
     out = StringIO()
+    n_docs = 0
     records_by_project = defaultdict(list)
     async for doc, passports in _read_passport_artifacts(paths, docs):
-        passports = Passports.model_validate_json(passports)
+        passports = Passports.model_validate_json(passports)  # noqa: PLW2901
         for r in _record_from_passports(doc, passports):
             records_by_project[r["project"]].append((doc, r))
+        n_docs += 1
+        if progress is not None and n_docs % 100 == 0:
+            await progress(n_docs)
     recs = []
     for project, records in records_by_project.items():
         for doc, r in records:
@@ -82,7 +92,6 @@ async def _read_passport_artifacts(
     paths: WorkerPaths,
     docs: AsyncIterable[Document],
 ) -> AsyncIterable[Document, bytes]:
-    n_docs = 0
     async for doc in docs:
         try:
             passports = await async_read_artifact(
@@ -91,9 +100,6 @@ async def _read_passport_artifacts(
         except ArtifactNotFound:
             continue
         yield doc, passports
-        n_docs += 1
-        if n_docs % 100 == 0:
-            logger.info(f"processed {n_docs} docs")
 
 
 def _record_from_passports(

workers/passport-worker/passport_worker/search.py

@@ -105,6 +105,14 @@ async def get_docs(  # noqa: PLR0917
         yield d
 
 
+async def count_docs(  # noqa: PLR0917
+    docs: list[DocId] | DocumentSearchQuery | None, project: str, es_client: ESClient
+) -> dict[str, Any]:
+    query = _build_doc_query(docs)
+    counts = await es_client.count(query, index=project)
+    return counts
+
+
 async def _write_preprocessing_batches(
     docs: AsyncIterable[Document],
     paths: WorkerPaths,

workers/passport-worker/tests/test_export.py

@@ -5,7 +5,6 @@
 from datashare_python.conftest import TEST_PROJECT
 from datashare_python.objects import Document, WorkerPaths
 from datashare_python.utils import write_artifact
-from icij_common.es import ESClient
 from passport_service.objects import Passport
 from passport_worker.config import PassportWorkerConfig
 from passport_worker.export import export_as_csv, export_as_jsonl
@@ -65,7 +64,6 @@ async def _async_iter[R](it: Iterable[R]) -> AsyncIterable[R]:
 
 async def test_export_as_csv(
     artifact_dir_with_passports: tuple[WorkerPaths, list[Document]],
-    test_es_client: ESClient,
 ) -> None:
     # Given
     paths, docs = artifact_dir_with_passports