From 00025f2541eb6dd57d63c749f8e0dfe491ff659b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cle=CC=81ment=20Doumouro?= Date: Mon, 20 Jul 2026 14:38:01 +0200 Subject: [PATCH] feature(asr-worker): index transcriptions --- datashare-python/datashare_python/objects.py | 10 + datashare-python/datashare_python/utils.py | 43 +++- datashare-python/tests/test_utils.py | 16 +- workers/asr-worker/asr_worker/activities.py | 242 ++++++++++++++---- workers/asr-worker/asr_worker/config.py | 6 + workers/asr-worker/asr_worker/constants.py | 5 +- workers/asr-worker/asr_worker/objects.py | 21 ++ workers/asr-worker/asr_worker/workflows.py | 42 ++- workers/asr-worker/entrypoints/io_worker.sh | 3 +- workers/asr-worker/pyproject.toml | 3 +- workers/asr-worker/tests/conftest.py | 9 +- workers/asr-worker/tests/test_activities.py | 148 ++++++++--- workers/asr-worker/tests/test_objects.py | 47 ++++ workers/asr-worker/tests/test_workflows.py | 2 +- workers/asr-worker/uv.dist.lock | 29 +++ workers/asr-worker/uv.lock | 29 +++ .../translation_worker/activities.py | 33 +-- 17 files changed, 543 insertions(+), 145 deletions(-) create mode 100644 workers/asr-worker/tests/test_objects.py diff --git a/datashare-python/datashare_python/objects.py b/datashare-python/datashare_python/objects.py index c985c2fc..dd3135e3 100644 --- a/datashare-python/datashare_python/objects.py +++ b/datashare-python/datashare_python/objects.py @@ -179,6 +179,11 @@ class Translation(BaseModel): # No camelcase here we don't know why content: Annotated[str, BeforeValidator(_from_sentences)] +Routing = str +DocID = str +DocRoute = tuple[DocID, Routing] + + class Document(DatashareModel): id: str language: DatashareLanguage @@ -254,6 +259,11 @@ def to_filesystem(self) -> FilesystemDocument: resource_name=resource_name, ) + def to_route(self) -> DocRoute: + if self.root_document is not None: + return self.root_document, self.id + return self.id, self.id + def _is_absolute_path(v: bytes | BytesIO | Path) -> Any: if isinstance(v, Path) and not v.is_absolute(): diff --git a/datashare-python/datashare_python/utils.py b/datashare-python/datashare_python/utils.py index 86f7d7ef..aba5949d 100644 --- a/datashare-python/datashare_python/utils.py +++ b/datashare-python/datashare_python/utils.py @@ -3,6 +3,7 @@ import fcntl import inspect import json +import logging import os import shutil import sys @@ -45,6 +46,8 @@ from .objects import BaseModel, DocArtifact, DocumentLocation, FilesystemDocument from .types_ import RawAsyncProgressHandler +logger = logging.getLogger(__name__) + DependencyLabel = str | None DependencySetup = Callable[..., None] DependencyAsyncSetup = Callable[..., Coroutine[None, None, None]] @@ -397,7 +400,7 @@ def write_artifact( ) manifest[artifact.type] = manifest_entry manifest_path.write_text(json.dumps(manifest)) - return artifact_path.relative_to(artif_dir) + return artifact_path.relative_to(root) @contextlib.contextmanager @@ -556,6 +559,16 @@ def symlink_embedded_document_to_workdir( raise ValueError(f"unsupported location {doc.location}") +def artifact_path( + doc_id: str, + artifact_type: type[DocArtifact], + *, + project: str, + root: Path, +) -> Path: + return root / artifacts_dir(doc_id, project=project) / artifact_type.filename + + def read_jsonl(path: Path) -> Iterable[dict]: with path.open() as f: for line in f: @@ -564,6 +577,34 @@ def read_jsonl(path: Path) -> Iterable[dict]: yield json.loads(line) +async def publish_and_consume( + publisher: asyncio.Task, + publisher_completion_callback: Callable[[], None], + *, + consumer: asyncio.Task, +) -> tuple[Any, Any]: + # Publish and consume concurrently + logger.debug("starting publish and subscribe") + done, pending = await asyncio.wait( + [publisher, consumer], return_when=asyncio.FIRST_COMPLETED + ) + for d in done: + # Stop everything case of exception + exc = d.exception() + if exc: + for p in pending: + p.cancel() + raise exc + # Wait for publish to be done and push the poison pill to stop consuming + p_res = await publisher + publisher_completion_callback() + logger.debug("done publishing, waiting for consumer to complete...") + # Wait for consumption to be done + c_res = await consumer + logger.debug("done consuming !") + return p_res, c_res + + class _PydanticPayloadConverter(CompositePayloadConverter): def __init__(self) -> None: json_payload_converter = PydanticJSONPlainPayloadConverter( diff --git a/datashare-python/tests/test_utils.py b/datashare-python/tests/test_utils.py index a3c9a265..a9e5e6c1 100644 --- a/datashare-python/tests/test_utils.py +++ b/datashare-python/tests/test_utils.py @@ -164,9 +164,13 @@ def test_write_artifact(tmp_path: Path) -> None: manifest_entry=manifest_entry, ) # When - write_artifact(root_dir, artifact) + artifact_path = write_artifact(root_dir, artifact) # Then - artifact_dir = root_dir / TEST_PROJECT / "do" / "c_" / "doc_id" + expected_artifact_cache = ( + Path(TEST_PROJECT) / "do" / "c_" / "doc_id" / "mocked-structure" + ) + assert artifact_path == expected_artifact_cache + artifact_dir = (root_dir / artifact_path).parent assert artifact_dir.exists() assert artifact_dir.is_dir() manifest_path = artifact_dir / "manifest.json" @@ -247,9 +251,13 @@ def test_write_artifact_with_existing_legacy_metadata(tmp_path: Path) -> None: meta_path = artifact_dir / "metadata.json" meta_path.write_text(json.dumps(existing_metadata)) # When - write_artifact(root_dir, artifact) + artifact_path = write_artifact(root_dir, artifact) # Then - artifact_dir = root_dir / TEST_PROJECT / "do" / "c_" / "doc_id" + expected_artifact_path = ( + Path(TEST_PROJECT) / "do" / "c_" / "doc_id" / "mocked-structure" + ) + assert artifact_path == expected_artifact_path + artifact_dir = (root_dir / artifact_path).parent assert artifact_dir.exists() assert artifact_dir.is_dir() meta_path = artifact_dir / "metadata.json" diff --git a/workers/asr-worker/asr_worker/activities.py b/workers/asr-worker/asr_worker/activities.py index 0fc15887..aadf8d18 100644 --- a/workers/asr-worker/asr_worker/activities.py +++ b/workers/asr-worker/asr_worker/activities.py @@ -7,6 +7,7 @@ from pathlib import Path from typing import Annotated, Any, Protocol, cast +from aiofile import async_open from caul_core import ( ASRResult, InferenceRunner, @@ -18,10 +19,7 @@ PreprocessorConfig, ) from datashare_python.dependencies import lifespan_worker_config -from datashare_python.objects import ( - Document, - FilesystemDocument, -) +from datashare_python.objects import DocRoute, Document from datashare_python.types_ import ( AsyncProgressRateHandler, RawAsyncProgressHandler, @@ -32,9 +30,11 @@ ActivityWithProgress, activity_defn, activity_workdir, + artifact_path, config_cache_key, debuggable_name, enter_cm, + publish_and_consume, read_jsonl, safe_dir, symlink_embedded_document_to_workdir, @@ -42,7 +42,9 @@ to_raw_sync_progress, write_artifact, ) +from elasticsearch._async.helpers import async_bulk from icij_common.es import ( + DOC_CONTENT, DOC_CONTENT_TYPE, DOC_LANGUAGE, DOC_METADATA, @@ -50,6 +52,7 @@ DOC_ROOT_ID, ES_DOCUMENT_TYPE, HITS, + ID_, QUERY, ESClient, and_query, @@ -57,10 +60,10 @@ ) from icij_common.iter_utils import async_batches from icij_common.pydantic_utils import safe_copy -from pydantic import TypeAdapter from .config import ASRWorkerConfig from .constants import ( + INDEX_TRANSCRIPTION_ACTIVITY, POSTPROCESS_ACTIVITY, PREPROCESS_ACTIVITY, RUN_INFERENCE_ACTIVITY, @@ -75,6 +78,7 @@ ) from .objects import ( ASRArgs, + ASRIndexingConfig, Transcription, TranscriptionArtifact, TranscriptionManifestEntry, @@ -84,11 +88,10 @@ _BASE_WEIGHT = 1.0 _SEARCH_AUDIOS_WEIGHT = _BASE_WEIGHT * 2 +_INDEX_AUDIOS_WEIGHT = _BASE_WEIGHT * 3 _PREPROCESS_WEIGHT = 5 * _BASE_WEIGHT _INFERENCE_WEIGHT = 10 * _PREPROCESS_WEIGHT -_LIST_OF_PATH_ADAPTER = TypeAdapter(list[Path]) - class ArtifactFactory(Protocol): def __call__(self, artifact: bytes) -> TranscriptionArtifact: ... @@ -113,11 +116,10 @@ async def search_audio_paths( output_dir.mkdir(parents=True, exist_ok=True) batch_paths = [ p.relative_to(workdir) - async for p in search_audio_paths_act( + async for p in search_audios_act( project, es_client, query, - config=worker_config, output_dir=output_dir, batch_size=batch_size, ) @@ -176,7 +178,6 @@ async def infer( workdir = worker_config.workdir output_dir = activity_workdir(workdir, project) output_dir.mkdir(parents=True, exist_ok=True) - preprocessed_inputs = _LIST_OF_PATH_ADAPTER.validate_python(preprocessed_inputs) preprocessed_inputs = [workdir / p for p in preprocessed_inputs] if progress is not None: progress = to_raw_async_progress( @@ -211,7 +212,7 @@ def postprocess( progress: Annotated[ # noqa: ARG002 SyncProgressRateHandler | None, Weight(value=_BASE_WEIGHT) ] = None, - ) -> int: + ) -> list[DocRoute]: # Import caul.tasks to populate the Postprocessor‹ registry import caul.tasks # noqa: F401, PLC0415 @@ -219,19 +220,14 @@ def postprocess( workdir = worker_config.workdir audio_batch = workdir / audio_batch artifacts_root = worker_config.artifacts_root - inference_results = _LIST_OF_PATH_ADAPTER.validate_python(inference_results) inference_results = ( ASRResult.model_validate_json((workdir / p).read_text()) for p in inference_results ) - docs = ( - FilesystemDocument.model_validate(fs_doc) - for fs_doc in read_jsonl(audio_batch) - ) - doc_ids = [doc.id for doc in docs] + docs = [Document.model_validate(fs_doc) for fs_doc in read_jsonl(audio_batch)] if progress is not None: - progress = to_raw_sync_progress(progress, max_progress=len(doc_ids)) + progress = to_raw_sync_progress(progress, max_progress=len(docs)) postprocessor_factory = enter_cm(partial(Postprocessor.from_config, config)) postprocessor_key = config_cache_key(config) cache = lifespan_postprocessor_cache() @@ -239,22 +235,51 @@ def postprocess( postprocessor_key, postprocessor_factory ) return postprocess_act( - postprocessor, inference_results, - doc_ids, + docs, + postprocessor, args, artifacts_root=artifacts_root, event_loop=self._event_loop, progress=progress, ) + @activity_defn(name=INDEX_TRANSCRIPTION_ACTIVITY) + async def index_transcriptions( + self, + routes: list[DocRoute], + project: str, + indexing_config: ASRIndexingConfig, + *, + progress: Annotated[ # noqa: ARG002 + AsyncProgressRateHandler | None, Weight(value=_INDEX_AUDIOS_WEIGHT) + ] = None, + ) -> int: + worker_config = cast(ASRWorkerConfig, lifespan_worker_config()) + es_client = lifespan_es_client() + target_bulk_char_size = worker_config.indexing.target_bulk_char_size + logger.info( + "indexing %s transcriptions by bulk of about %s characters !", + len(routes), + target_bulk_char_size, + ) + n_docs = await index_transcriptions_act( + routes, + project, + es_client, + indexing_config=indexing_config, + artifact_root=worker_config.artifacts_root, + target_bulk_char_size=target_bulk_char_size, + progress=progress, + ) + return n_docs + -async def search_audio_paths_act( +async def search_audios_act( project: str, es_client: ESClient, query: dict[str, Any], *, - config: ASRWorkerConfig, output_dir: Path, batch_size: int, ) -> AsyncIterable[Path]: @@ -262,12 +287,6 @@ async def search_audio_paths_act( docs = _search_audio_paths( es_client, project, query, supported_content_types=SUPPORTED_CONTENT_TYPES ) - docs = ( - symlink_embedded_document_to_workdir( - d, config.artifacts_root, workdir=config.workdir - ) - async for d in docs - ) async for p in write_audio_batches(docs, output_dir, batch_size): yield p @@ -279,20 +298,19 @@ def preprocess_act( worker_config: ASRWorkerConfig, output_dir: Path, ) -> list[Path]: - docs_root = worker_config.docs_root - artifacts_root = worker_config.artifacts_root - workdir = worker_config.workdir logger.debug("locating files...") + audios = (Document.model_validate(doc) for doc in read_jsonl(audio_batch)) + audios = (a.to_filesystem() for a in audios) audios = ( - FilesystemDocument.model_validate(fs_doc) for fs_doc in read_jsonl(audio_batch) - ) - audios = ( - fs_doc.locate( - original_root=docs_root, artifacts_root=artifacts_root, workdir=workdir + symlink_embedded_document_to_workdir( + a, worker_config.artifacts_root, workdir=worker_config.workdir + ).locate( + worker_config.docs_root, + artifacts_root=worker_config.artifacts_root, + workdir=worker_config.workdir, ) - for fs_doc in audios + for a in audios ) - audios = (str(a) for a in audios) # TODO: implement a caching strategy here, we could avoid processing files # which have already been preprocessed logger.debug("starting preprocessing...") @@ -342,39 +360,72 @@ def _transcribe_as_list( def postprocess_act( - postprocessor: Postprocessor, inference_results: Iterable[ASRResult], - doc_ids: Iterable[str], + docs: list[Document], + postprocessor: Postprocessor, args: ASRArgs, *, artifacts_root: Path, event_loop: AbstractEventLoop | None = None, progress: SyncProgressRateHandler | None = None, -) -> int: +) -> list[DocRoute]: transcriptions = postprocessor.process(inference_results) # Strict is important here ! - n_docs = 0 - for i, (doc_id, asr_result) in enumerate(zip(doc_ids, transcriptions, strict=True)): - n_docs += 1 + for i, (doc, asr_result) in enumerate(zip(docs, transcriptions, strict=True)): manifest_entry = TranscriptionManifestEntry.complete( args, confidence=asr_result.score ) artifact_factory = partial( TranscriptionArtifact, project=args.project, - doc_id=doc_id, + doc_id=doc.id, manifest_entry=manifest_entry, ) t_path = write_transcription(asr_result, artifact_factory, artifacts_root) logger.debug("wrote transcription for %s", t_path) if progress is not None and event_loop is not None: progress(i, event_loop) + routes = [d.to_route() for d in docs] + return routes + + +async def index_transcriptions_act( + routes: Iterable[DocRoute], + project: str, + es_client: ESClient, + artifact_root: Path, + target_bulk_char_size: int = 100_000, + es_concurrency: int = 5, + indexing_config: ASRIndexingConfig = None, + progress: AsyncProgressRateHandler | None = None, +) -> int: + if indexing_config is None: + indexing_config = ASRIndexingConfig() + es_queue = asyncio.Queue(maxsize=es_concurrency) + publisher = _read_transcriptions_and_queue( + list(routes), + es_queue, + project, + target_bulk_char_size, + artifact_root=artifact_root, + indexing_config=indexing_config, + progress=progress, + ) + publisher = asyncio.create_task(publisher) + publisher_callback = lambda: es_queue.put_nowait(None) # noqa: E731 + consumer = asyncio.create_task( + _write_transcriptions_to_es(es_client, queue=es_queue, project=project) + ) + n_docs, _ = await publish_and_consume( + publisher, publisher_callback, consumer=consumer + ) return n_docs def _preprocess( - preprocessor: Preprocessor, audios: Iterable[str], output_dir: Path + preprocessor: Preprocessor, audios: Iterable[Path], output_dir: Path ) -> Iterable[Path]: + audios = (str(a) for a in audios) for batch_i, batch in enumerate( preprocessor.process(audios, output_dir=output_dir) ): @@ -409,21 +460,27 @@ def _relative_input( return PreprocessedInput(metadata=metadata) # noqa: F821 +_EXCLUDED_FROM_BATCH_SERIALIZATION = {"type", "tags"} + + async def write_audio_batches( - docs: AsyncIterable[FilesystemDocument], root: Path, batch_size: int + docs: AsyncIterable[Document], root: Path, batch_size: int ) -> AsyncIterable[Path]: batch_id = 0 async for batch in async_batches(docs, batch_size): batch_path = root / f"{batch_id}.txt" with batch_path.open("w") as f: - for fs_doc in batch: - f.write(f"{fs_doc.model_dump_json()}\n") + for doc in batch: + as_jsonl = doc.model_dump_json( + exclude_none=True, exclude=_EXCLUDED_FROM_BATCH_SERIALIZATION + ) + f.write(f"{as_jsonl}\n") yield batch_path batch_id += 1 _DOC_TYPE_QUERY = has_type(type_field="type", type_value=ES_DOCUMENT_TYPE) -_DOC_CONTENT_SOURCES = [DOC_PATH, DOC_ROOT_ID, DOC_LANGUAGE, DOC_METADATA] +_DOC_CONTENT_SOURCES = [DOC_PATH, DOC_ROOT_ID, DOC_LANGUAGE, DOC_METADATA, DOC_ROOT_ID] async def _search_audio_paths( @@ -431,13 +488,13 @@ async def _search_audio_paths( project: str, query: dict[str, Any], supported_content_types: set[str], -) -> AsyncGenerator[FilesystemDocument, None]: +) -> AsyncGenerator[Document, None]: body = _with_audio_content(query, supported_content_types) async for page in es_client.poll_search_pages( index=project, body=body, sort="_doc:asc", _source_includes=_DOC_CONTENT_SOURCES ): for hit in page[HITS][HITS]: - yield Document.from_es(hit).to_filesystem() + yield Document.from_es(hit) def _content_type_query(supported_content_types: set[str]) -> dict[str, Any]: @@ -455,9 +512,90 @@ def _with_audio_content( return and_query(query, type_query[QUERY]) +async def _read_transcriptions_and_queue( + docs: list[DocRoute], + queue: asyncio.Queue, + project: str, + target_bulk_char_size: int, + indexing_config: ASRIndexingConfig, + *, + artifact_root: Path, + progress: AsyncProgressRateHandler | None = None, +) -> int: + n_docs = len(docs) + if not n_docs: + return n_docs + if progress is not None: + progress = to_raw_async_progress(progress, max_progress=n_docs) + bulk = [] + bulk_char_size = 0 + n_docs = len(docs) + for doc_i, route in enumerate(docs): + routing, doc_id = route + transcription_path = artifact_path( + doc_id, + TranscriptionArtifact, + project=project, + root=artifact_root, + ) + async with async_open(transcription_path) as f: + transcription = Transcription.model_validate_json(await f.read()) + indexed = transcription.as_text( + indexing_config.transcript_sep, speaker_sep=indexing_config.speaker_sep + ) + if bulk_char_size + len(indexed) >= target_bulk_char_size and bulk: + await queue.put(bulk) + bulk = [] + logger.debug("queued %s / %s transcription for indexation !", doc_i, n_docs) + bulk.append((route, indexed)) + if progress is not None and doc_i % 10 == 0: + await progress(doc_i) + # Empty the buffer + if bulk: + await queue.put(bulk) + if progress is not None: + await progress(n_docs) + queue.put_nowait(None) + return n_docs + + +async def _write_transcriptions_to_es( + es_client: ESClient, queue: asyncio.Queue, project: str +) -> None: + while True: + transcriptions = await queue.get() + if transcriptions is None: + logger.debug("popped poison pill from the queue, exiting !") + queue.task_done() + return + logger.debug("writing translations to the index..") + await _update_docs_content(es_client, transcriptions, project=project) + logger.debug("translation written !") + queue.task_done() + + +async def _update_docs_content( + es_client: ESClient, + transcribed_docs: Iterable[tuple[DocRoute, str]], + project: str, +) -> None: + actions = ( + { + "_op_type": "update", + "_index": project, + "_routing": routing, + ID_: doc_id, + "doc": {DOC_CONTENT: transcription}, + } + for (routing, doc_id), transcription in transcribed_docs + ) + await async_bulk(es_client, actions, raise_on_error=True, refresh="wait_for") + + REGISTRY = [ ASRActivities.search_audio_paths, ASRActivities.preprocess, ASRActivities.infer, ASRActivities.postprocess, + ASRActivities.index_transcriptions, ] diff --git a/workers/asr-worker/asr_worker/config.py b/workers/asr-worker/asr_worker/config.py index 191750b2..46030352 100644 --- a/workers/asr-worker/asr_worker/config.py +++ b/workers/asr-worker/asr_worker/config.py @@ -28,6 +28,10 @@ class ASRCache(BaseModel): ) +class IndexingWorkerConfig(BaseModel): + target_bulk_char_size: int = 50_000 + + class ASRWorkerConfig(WorkerConfig): logging: LoggingConfig = _DEFAULT_LOGGING_CONFIG @@ -35,6 +39,8 @@ class ASRWorkerConfig(WorkerConfig): artifacts_root: Path workdir: Path + indexing: IndexingWorkerConfig = Field(default_factory=IndexingWorkerConfig) + # Set max concurrent activity to 1 to avoid parallel inference in practice, # this must be set to 1 for the inference worker and > 1 for the io worker max_concurrent_activities: int = Field(frozen=True, default=1) diff --git a/workers/asr-worker/asr_worker/constants.py b/workers/asr-worker/asr_worker/constants.py index 1d621e4b..62307247 100644 --- a/workers/asr-worker/asr_worker/constants.py +++ b/workers/asr-worker/asr_worker/constants.py @@ -1,9 +1,5 @@ ASR_WORKER_NAME = "asr-worker" -ONE_MINUTE = 60 - -TEN_MINUTES = ONE_MINUTE * 10 - ASR_TASK_QUEUE = "transcription-tasks" PARAKEET = "parakeet" @@ -20,6 +16,7 @@ SEARCH_AUDIOS_ACTIVITY = "asr.transcription.search-audios" RUN_INFERENCE_ACTIVITY = "asr.transcription.infer" POSTPROCESS_ACTIVITY = "asr.transcription.postprocess" +INDEX_TRANSCRIPTION_ACTIVITY = "asr.transcription.index" SUPPORTED_CONTENT_TYPES = { "audio/aac", diff --git a/workers/asr-worker/asr_worker/objects.py b/workers/asr-worker/asr_worker/objects.py index 3027a350..22cfb9d3 100644 --- a/workers/asr-worker/asr_worker/objects.py +++ b/workers/asr-worker/asr_worker/objects.py @@ -27,11 +27,17 @@ class TranscriptionArtifact(DocArtifact): type: ClassVar[ArtifactType] = ArtifactType.ASR_TRANSCRIPTION +class ASRIndexingConfig(DatashareModel): + transcript_sep: str = "\n" + speaker_sep: str = "\n\n" + + class ASRArgs(TaskArgs): project: str docs: list[DocId] | DocumentSearchQuery config: ASRPipelineConfig = Field(default_factory=ASRPipelineConfig.parakeet) batch_size: int + indexing: ASRIndexingConfig = Field(default_factory=ASRIndexingConfig) def as_manifest_task_input(self) -> dict[str, Any]: as_entry = super().as_manifest_task_input() @@ -74,6 +80,21 @@ def from_asr_handler_result(cls, asr_handler_result: ASRResult) -> Self: confidence = math.exp(asr_handler_result.score) return Transcription(confidence=confidence, transcripts=transcripts) + def as_text(self, transcript_sep: str = "\n", *, speaker_sep: str = "\n\n") -> str: + current_speaker = None + current_speaker_texts = [] + texts = [] + for t in self.transcripts: + new_speaker = t.speaker + if new_speaker != current_speaker and current_speaker_texts: + texts.append(transcript_sep.join(current_speaker_texts)) + current_speaker_texts = [] + current_speaker_texts.append(t.text) + current_speaker = new_speaker + if current_speaker_texts: + texts.append(transcript_sep.join(current_speaker_texts)) + return speaker_sep.join(texts) + AvailableModels = RootModel[dict[ASRLanguage, list[ASRModel]]] diff --git a/workers/asr-worker/asr_worker/workflows.py b/workers/asr-worker/asr_worker/workflows.py index 77e50df3..38fc7e55 100644 --- a/workers/asr-worker/asr_worker/workflows.py +++ b/workers/asr-worker/asr_worker/workflows.py @@ -9,7 +9,7 @@ from pydantic import TypeAdapter from temporalio import workflow -from .constants import ASR_WORKFLOW, TEN_MINUTES +from .constants import ASR_WORKFLOW from .objects import ASRArgs, ASRResponse with workflow.unsafe.imports_passed_through(): @@ -19,6 +19,11 @@ logger = logging.getLogger(__name__) +AUDIO_SEARCH_TIMEOUT = timedelta(minutes=10) +INFERENCE_TIMEOUT = timedelta(minutes=30) +INDEXATION_TIMEOUT = timedelta(hours=1) +POSTPROCESSING_TIMEOUT = timedelta(minutes=10) + class TaskQueues(StrEnum): WORKFLOWS = "datashare.workflows" @@ -40,7 +45,7 @@ async def run(self, args: ASRArgs) -> ASRResponse: batch_paths = await execute_activity( ASRActivities.search_audio_paths, args=search_args, - start_to_close_timeout=timedelta(seconds=TEN_MINUTES), + start_to_close_timeout=AUDIO_SEARCH_TIMEOUT, task_queue=TaskQueues.IO, ) # Preprocessing @@ -54,7 +59,7 @@ async def run(self, args: ASRArgs) -> ASRResponse: execute_activity( ASRActivities.preprocess, args=a, - start_to_close_timeout=timedelta(seconds=TEN_MINUTES), + start_to_close_timeout=timedelta(minutes=10), task_queue=TaskQueues.CPU, ) for a in preprocess_args @@ -75,7 +80,7 @@ async def run(self, args: ASRArgs) -> ASRResponse: task_queue=TaskQueues.INFERENCE_GPU, args=b, # TODO: in practice we should parse the config to find out - start_to_close_timeout=timedelta(seconds=TEN_MINUTES), + start_to_close_timeout=INFERENCE_TIMEOUT, heartbeat_timeout=timedelta(minutes=3), ) for b in inference_args @@ -84,7 +89,7 @@ async def run(self, args: ASRArgs) -> ASRResponse: inference_results = await gather(*inference_acts) logger.info("inference complete !") # Postprocessing - postprocessing_ins = list( + postprocessing_args = list( zip( inference_results, batch_paths, @@ -97,15 +102,34 @@ async def run(self, args: ASRArgs) -> ASRResponse: execute_activity( ASRActivities.postprocess, args=i, - start_to_close_timeout=timedelta(seconds=TEN_MINUTES), + start_to_close_timeout=POSTPROCESSING_TIMEOUT, task_queue=TaskQueues.CPU, ) - for i in postprocessing_ins + for i in postprocessing_args ] logger.info("running postprocessing...") - n_transcribed = await gather(*postprocessing_acts) - n_transcribed = sum(n_transcribed) + routes_batches = await gather(*postprocessing_acts) logger.info("postprocessing complete !") + # Indexing + indexing_args = list( + zip( + routes_batches, + repeat(args.project), + repeat(args.indexing), + strict=False, + ) + ) + indexing_acts = [ + execute_activity( + ASRActivities.index_transcriptions, + args=a, + start_to_close_timeout=INDEXATION_TIMEOUT, + task_queue=TaskQueues.IO, + ) + for a in indexing_args + ] + n_transcribed = await gather(*indexing_acts) + n_transcribed = sum(n_transcribed) return ASRResponse(n_transcribed=n_transcribed) diff --git a/workers/asr-worker/entrypoints/io_worker.sh b/workers/asr-worker/entrypoints/io_worker.sh index 17a069d5..9e95e825 100755 --- a/workers/asr-worker/entrypoints/io_worker.sh +++ b/workers/asr-worker/entrypoints/io_worker.sh @@ -4,4 +4,5 @@ set -e uv run --no-sync datashare-python worker start \ --dependencies asr.io \ --queue asr.io \ - --activity asr.transcription.search-audios \ No newline at end of file + --activity asr.transcription.search-audios \ + --activity asr.transcription.index \ No newline at end of file diff --git a/workers/asr-worker/pyproject.toml b/workers/asr-worker/pyproject.toml index 196913f2..8127b1ab 100644 --- a/workers/asr-worker/pyproject.toml +++ b/workers/asr-worker/pyproject.toml @@ -12,6 +12,7 @@ requires-python = ">=3.11.0, <3.13" dependencies = [ "datashare-python~=0.9.0", "caul-core==0.3.1", + "aiofile==3.11.1", ] [project.scripts] @@ -152,4 +153,4 @@ markers = [ log_cli = 1 log_cli_level = "DEBUG" log_file_format = "[%(levelname)s][%(asctime)s.%(msecs)03d][%(name)s]: %(message)s" -log_file_date_format = "%Y-%m-%d %H:%M:%S" \ No newline at end of file +log_file_date_format = "%Y-%m-%d %H:%M:%S" diff --git a/workers/asr-worker/tests/conftest.py b/workers/asr-worker/tests/conftest.py index e705600b..4563afc7 100644 --- a/workers/asr-worker/tests/conftest.py +++ b/workers/asr-worker/tests/conftest.py @@ -30,7 +30,7 @@ worker_lifetime_deps, ) from datashare_python.dependencies import set_es_client, set_temporal_client -from datashare_python.objects import Document, FilesystemDocument +from datashare_python.objects import Document from datashare_python.types_ import ContextManagerFactory from datashare_python.utils import artifacts_dir from icij_common.es import ESClient @@ -121,13 +121,12 @@ async def populate_es_with_audios( @pytest.fixture def with_audio_docs( populate_es_with_audios: list[Document], test_worker_config: ASRWorkerConfig -) -> list[FilesystemDocument]: +) -> list[Document]: config = test_worker_config clear_dirs(test_worker_config) docs = [ d for d in populate_es_with_audios if d.content_type in SUPPORTED_CONTENT_TYPES ] - paths = [] audio_path = AUDIOS_PATH / "asr_test.wav" for doc in docs: if doc.root_document is None: @@ -139,6 +138,4 @@ def with_audio_docs( ) artifact_path.parent.mkdir(parents=True, exist_ok=True) shutil.copy(audio_path, artifact_path) - fs_doc = doc.to_filesystem() - paths.append(fs_doc) - return paths + return docs diff --git a/workers/asr-worker/tests/test_activities.py b/workers/asr-worker/tests/test_activities.py index c1cb6ebc..2c51f994 100644 --- a/workers/asr-worker/tests/test_activities.py +++ b/workers/asr-worker/tests/test_activities.py @@ -1,5 +1,6 @@ import json from collections.abc import AsyncGenerator, Iterable +from functools import partial from itertools import cycle from pathlib import Path from typing import Self @@ -7,14 +8,21 @@ import pytest from aiostream import stream from asr_worker.activities import ( + index_transcriptions_act, infer_act, postprocess_act, preprocess_act, - search_audio_paths_act, + search_audios_act, write_audio_batches, + write_transcription, ) from asr_worker.config import ASRWorkerConfig -from asr_worker.objects import ASRArgs, DocId, Transcription, TranscriptionManifestEntry +from asr_worker.objects import ( + ASRArgs, + Transcription, + TranscriptionArtifact, + TranscriptionManifestEntry, +) from caul_core import ( ASRResult, InferenceRunner, @@ -25,9 +33,13 @@ PreprocessorOutput, ) from datashare_python.conftest import TEST_PROJECT -from datashare_python.objects import DocumentLocation, FilesystemDocument +from datashare_python.objects import ( + DatashareLanguage, + Document, + FilesystemDocument, +) from datashare_python.utils import read_jsonl -from icij_common.es import ESClient, ids_query, match_all +from icij_common.es import HITS, ESClient, ids_query, match_all from icij_common.iter_utils import batches from icij_common.registrable import RegistrableConfig @@ -71,19 +83,17 @@ ), ] -FS_DOCUMENT_0 = FilesystemDocument( +DOC_0 = Document( id="doc-0", path=Path(TEST_PROJECT, "symlinks", "do", "c-", "doc-0", "doc-0.wav"), index=TEST_PROJECT, - location=DocumentLocation.WORKDIR, - resource_name="doc-0.wav", + language=DatashareLanguage("ENGLISH"), ) -FS_DOCUMENT_2 = FilesystemDocument( +DOC_2 = Document( id="doc-2", path=Path("doc-2.mp3"), index=TEST_PROJECT, - location=DocumentLocation.ORIGINAL, - resource_name="doc-2.mp3", + language=DatashareLanguage("ENGLISH"), ) @@ -158,10 +168,10 @@ def process( ("query", "expected_batches"), [ # Supports empty query - ({}, [[FS_DOCUMENT_0, FS_DOCUMENT_2]]), + ({}, [["doc_0", "doc_2"]]), # Return all audio/video docs - (match_all(), [[FS_DOCUMENT_0, FS_DOCUMENT_2]]), - (ids_query(["doc-0"]), [[FS_DOCUMENT_0]]), + (match_all(), [["doc_0", "doc_2"]]), + (ids_query(["doc-0"]), [["doc_0"]]), # Should filter non supported content type (ids_query(["doc-1"]), []), ], @@ -170,33 +180,32 @@ async def test_search_audio_paths_act( with_audio_docs: list[FilesystemDocument], test_es_client: ESClient, query: dict, - expected_batches: list[tuple[DocId, Path]], - test_worker_config: ASRWorkerConfig, + expected_batches: list[list[str]], + request, # noqa: ANN001 tmpdir: Path, ) -> None: # Given tmpdir = Path(tmpdir) - worker_config = test_worker_config batch_size = len(with_audio_docs) client = test_es_client # When batch_paths = [ batch - async for batch in search_audio_paths_act( + async for batch in search_audios_act( es_client=client, project=TEST_PROJECT, query=query, batch_size=batch_size, output_dir=tmpdir, - config=worker_config, ) ] # Then results = [] for b in batch_paths: - results.append( - [FilesystemDocument.model_validate(fs_doc) for fs_doc in read_jsonl(b)] - ) + results.append([Document.model_validate(fs_doc).id for fs_doc in read_jsonl(b)]) + expected_batches = [ + [request.getfixturevalue(d).id for d in batch] for batch in expected_batches + ] assert results == expected_batches @@ -207,12 +216,13 @@ def test_preprocess_act(test_worker_config: ASRWorkerConfig, tmpdir: Path) -> No batch_size = n_audios - 1 audio_batch = tmpdir / "audio_batch.txt" batch = [ - FilesystemDocument( + Document( id=f"doc-{i}", + language=DatashareLanguage("ENGLISH"), + root_document=f"root-{i}", path=Path(str(i)), - location=DocumentLocation.ARTIFACTS, index=TEST_PROJECT, - resource_name=f"doc-{i}.wav", + metadata={"tika_metadata_resourcename": f"doc-{i}.wav"}, ) for i in range(n_audios) ] @@ -276,16 +286,31 @@ def test_postprocess_act(tmpdir: Path) -> None: postprocessor = MockPostprocessor() project = TEST_PROJECT artifacts_root = Path(tmpdir) - doc_ids = [f"{str(i) * 4}-doc-{i}" for i in range(3)] + docs = [ + Document( + id=f"{str(i) * 4}-doc-{i}", + language=DatashareLanguage("ENGLISH"), + root_document=f"root-{i}", + path=Path(str(i)), + index=TEST_PROJECT, + metadata={"tika_metadata_resourcename": f"doc-{i}.wav"}, + ) + for i in range(3) + ] # When - postprocess_act( - postprocessor, + routes = postprocess_act( INFERENCE_RESULTS, - doc_ids, + docs, + postprocessor, args, artifacts_root=artifacts_root, ) # Then + assert routes == [ + ("root-0", "0000-doc-0"), + ("root-1", "1111-doc-1"), + ("root-2", "2222-doc-2"), + ] expected_artifact_dirs = [ artifacts_root / project / "00" / "00" / "0000-doc-0", artifacts_root / project / "11" / "11" / "1111-doc-1", @@ -316,17 +341,17 @@ async def test_write_audio_search_results(tmpdir: Path) -> None: root = Path(tmpdir) batch_size = 2 - async def results() -> AsyncGenerator[FilesystemDocument, None]: + async def results() -> AsyncGenerator[Document, None]: res = ["doc-0", "doc-1", "doc-2"] for r in res: - fs_doc = FilesystemDocument( + doc = Document( id=r, path=Path(f"{r}.wav"), index=TEST_PROJECT, - location=DocumentLocation.WORKDIR, - resource_name=f"{r}.wav", + language=DatashareLanguage("ENGLISH"), + metadata={"tika_metadata_resourcename": f"{r}.wav"}, ) - yield fs_doc + yield doc # When results = write_audio_batches(results(), root=root, batch_size=batch_size) @@ -335,11 +360,11 @@ async def results() -> AsyncGenerator[FilesystemDocument, None]: async def expected_content() -> AsyncGenerator[str, None]: contents = [ [ - '{"id":"doc-0","path":"doc-0.wav","index":"test-project","location":"workdir","resource_name":"doc-0.wav"}', - '{"id":"doc-1","path":"doc-1.wav","index":"test-project","location":"workdir","resource_name":"doc-1.wav"}', + '{"id":"doc-0","language":"ENGLISH","index":"test-project","path":"doc-0.wav","metadata":{"tika_metadata_resourcename":"doc-0.wav"}}', + '{"id":"doc-1","language":"ENGLISH","index":"test-project","path":"doc-1.wav","metadata":{"tika_metadata_resourcename":"doc-1.wav"}}', ], [ - '{"id":"doc-2","path":"doc-2.wav","index":"test-project","location":"workdir","resource_name":"doc-2.wav"}' + '{"id":"doc-2","language":"ENGLISH","index":"test-project","path":"doc-2.wav","metadata":{"tika_metadata_resourcename":"doc-2.wav"}}' ], ] for line in contents: @@ -351,3 +376,54 @@ async def expected_content() -> AsyncGenerator[str, None]: async for p, expected_content in streamed: assert p.exists() assert p.read_text() == expected_content + + +@pytest.fixture +def with_transcribed_docs( + with_audio_docs: list[Document], test_worker_config: ASRWorkerConfig +) -> list[tuple[Document, str]]: + artifacts_root = test_worker_config.artifacts_root + transcriptions = [] + args = ASRArgs(project=TEST_PROJECT, docs=[], batch_size=2) + manifest_entry = TranscriptionManifestEntry.complete(args, confidence=1.0) + for doc_i, doc in enumerate(with_audio_docs): + artifact_factory = partial( + TranscriptionArtifact, + doc_id=doc.id, + project=TEST_PROJECT, + manifest_entry=manifest_entry, + ) + transcription = f"transcription_{doc_i}" + transcriptions.append(transcription) + asr_result = ASRResult(transcription=[(0, 1, transcription)]) + write_transcription(asr_result, artifact_factory, artifacts_root) + return list(zip(with_audio_docs, transcriptions, strict=True)) + + +async def test_index_transcriptions_act( + with_transcribed_docs: list[tuple[Document, str]], + test_es_client: ESClient, + test_worker_config: ASRWorkerConfig, +) -> None: + # Given + target_bulk_char_size = 1 # let's index each doc separately + docs, transcriptions = zip(*with_transcribed_docs, strict=True) + docs = list(docs) + routes = [(d.root_document, d.id) for d in docs] + transcriptions = list(transcriptions) + # When + n_docs = await index_transcriptions_act( + routes, + project=TEST_PROJECT, + es_client=test_es_client, + artifact_root=test_worker_config.artifacts_root, + target_bulk_char_size=target_bulk_char_size, + ) + # Then + assert n_docs == len(docs) + contents = [] + docs_ids = [d.id for d in docs] + body = {"query": ids_query(docs_ids)} + async for res in test_es_client.poll_search_pages(index=TEST_PROJECT, body=body): + contents += [Document.from_es(d).content for d in res[HITS][HITS]] + assert contents == transcriptions diff --git a/workers/asr-worker/tests/test_objects.py b/workers/asr-worker/tests/test_objects.py new file mode 100644 index 00000000..fe0e6903 --- /dev/null +++ b/workers/asr-worker/tests/test_objects.py @@ -0,0 +1,47 @@ +import pytest +from asr_worker.objects import Transcript, Transcription + +_TRANSCRIPTION_0 = Transcription( + transcripts=[ + Transcript(text="speaker_0_sentence_0", speaker="speaker_0"), + ], + confidence=1.0, +) +_TRANSCRIPTION_1 = Transcription( + transcripts=[ + Transcript(text="speaker_0_sentence_0", speaker="speaker_0"), + Transcript(text="speaker_0_sentence_1", speaker="speaker_0"), + Transcript(text="speaker_1_sentence_0", speaker="speaker_1"), + Transcript(text="speaker_0_sentence_2", speaker="speaker_0"), + ], + confidence=1.0, +) + + +@pytest.mark.parametrize( + ("transcript_sep", "speaker_sep", "transcription", "expected_text"), + [ + ("\n", "\n\n", _TRANSCRIPTION_0, "speaker_0_sentence_0"), + ( + "\n", + "\n\n", + _TRANSCRIPTION_1, + """speaker_0_sentence_0 +speaker_0_sentence_1 + +speaker_1_sentence_0 + +speaker_0_sentence_2""", + ), + ], +) +def test_asr_transcription_as_text( + transcript_sep: str, + speaker_sep: str, + transcription: Transcription, + expected_text: str, +) -> None: + # When + text = transcription.as_text(transcript_sep, speaker_sep=speaker_sep) + # Then + assert text == expected_text diff --git a/workers/asr-worker/tests/test_workflows.py b/workers/asr-worker/tests/test_workflows.py index 1f91a9a0..8dd980d3 100644 --- a/workers/asr-worker/tests/test_workflows.py +++ b/workers/asr-worker/tests/test_workflows.py @@ -73,7 +73,7 @@ async def io_bound_worker( worker_config=test_worker_config, client=client, task_queue=task_queue, - activities=[activities.search_audio_paths], + activities=[activities.search_audio_paths, activities.index_transcriptions], dependencies=dependencies, ) async with worker_ctx: diff --git a/workers/asr-worker/uv.dist.lock b/workers/asr-worker/uv.dist.lock index 3cc99b20..d8907c40 100644 --- a/workers/asr-worker/uv.dist.lock +++ b/workers/asr-worker/uv.dist.lock @@ -78,6 +78,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/46/02ac5e262d4af18054b3e922b2baedbb2a03289ee792162de60a865defc5/accelerate-1.13.0-py3-none-any.whl", hash = "sha256:cf1a3efb96c18f7b152eb0fa7490f3710b19c3f395699358f08decca2b8b62e0", size = 383744, upload-time = "2026-03-04T19:34:10.313Z" }, ] +[[package]] +name = "aiofile" +version = "3.11.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "caio", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'x86_64' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu') or sys_platform == 'darwin' or (sys_platform != 'linux' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/48/41/2fea7e193e061ce54eacc3b7bc0e6a99e4fcff43c78cf0a76dd781ed8334/aiofile-3.11.1.tar.gz", hash = "sha256:1f91912c6643d2a4e49ca4ae3514f0bf3867ce948a36d99a6411b8f4755f4cf9", size = 19342, upload-time = "2026-05-16T08:18:33.538Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/cd/0d76dfc5de72bde52f55f53e925c7d152d9c7906634ec1e0cbc7e8d4ad93/aiofile-3.11.1-py3-none-any.whl", hash = "sha256:ce77d14ac07f77bc2b757834a5c129321f3f705c474593deed5ab209079a52c9", size = 20446, upload-time = "2026-05-16T08:18:32.051Z" }, +] + [[package]] name = "aiohappyeyeballs" version = "2.6.1" @@ -267,6 +279,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fa/93/e8c04e80e82391a6e51f218ca49720f64236bc824e92152a2633b74cf7ab/braceexpand-0.1.7-py2.py3-none-any.whl", hash = "sha256:91332d53de7828103dcae5773fb43bc34950b0c8160e35e0f44c4427a3b85014", size = 5923, upload-time = "2021-05-07T13:49:05.146Z" }, ] +[[package]] +name = "caio" +version = "0.9.25" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/92/88/b8527e1b00c1811db339a1df8bd1ae49d146fcea9d6a5c40e3a80aaeb38d/caio-0.9.25.tar.gz", hash = "sha256:16498e7f81d1d0f5a4c0ad3f2540e65fe25691376e0a5bd367f558067113ed10", size = 26781, upload-time = "2025-12-26T15:21:36.501Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/90/543f556fcfcfa270713eef906b6352ab048e1e557afec12925c991dc93c2/caio-0.9.25-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:d6956d9e4a27021c8bd6c9677f3a59eb1d820cc32d0343cea7961a03b1371965", size = 36839, upload-time = "2025-12-26T15:21:40.267Z" }, + { url = "https://files.pythonhosted.org/packages/51/3b/36f3e8ec38dafe8de4831decd2e44c69303d2a3892d16ceda42afed44e1b/caio-0.9.25-cp311-cp311-manylinux2010_x86_64.manylinux2014_x86_64.manylinux_2_12_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bf84bfa039f25ad91f4f52944452a5f6f405e8afab4d445450978cd6241d1478", size = 80255, upload-time = "2025-12-26T15:22:20.271Z" }, + { url = "https://files.pythonhosted.org/packages/46/90/e278863c47e14ec58309aa2e38a45882fbe67b4cc29ec9bc8f65852d3e45/caio-0.9.25-cp311-cp311-manylinux_2_34_x86_64.whl", hash = "sha256:fc220b8533dcf0f238a6b1a4a937f92024c71e7b10b5a2dfc1c73604a25709bc", size = 78273, upload-time = "2026-03-04T22:08:21.368Z" }, + { url = "https://files.pythonhosted.org/packages/d3/25/79c98ebe12df31548ba4eaf44db11b7cad6b3e7b4203718335620939083c/caio-0.9.25-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:fb7ff95af4c31ad3f03179149aab61097a71fd85e05f89b4786de0359dffd044", size = 36983, upload-time = "2025-12-26T15:21:36.075Z" }, + { url = "https://files.pythonhosted.org/packages/a3/2b/21288691f16d479945968a0a4f2856818c1c5be56881d51d4dac9b255d26/caio-0.9.25-cp312-cp312-manylinux2010_x86_64.manylinux2014_x86_64.manylinux_2_12_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:97084e4e30dfa598449d874c4d8e0c8d5ea17d2f752ef5e48e150ff9d240cd64", size = 82012, upload-time = "2025-12-26T15:22:20.983Z" }, + { url = "https://files.pythonhosted.org/packages/d1/1c/0fe770b8ffc8362c48134d1592d653a81a3d8748d764bec33864db36319d/caio-0.9.25-cp312-cp312-manylinux_2_34_x86_64.whl", hash = "sha256:6bebf6f079f1341d19f7386db9b8b1f07e8cc15ae13bfdaff573371ba0575d69", size = 80200, upload-time = "2026-03-04T22:08:23.382Z" }, + { url = "https://files.pythonhosted.org/packages/86/93/1f76c8d1bafe3b0614e06b2195784a3765bbf7b0a067661af9e2dd47fc33/caio-0.9.25-py3-none-any.whl", hash = "sha256:06c0bb02d6b929119b1cfbe1ca403c768b2013a369e2db46bfa2a5761cf82e40", size = 19087, upload-time = "2025-12-26T15:22:00.221Z" }, +] + [[package]] name = "caul" version = "0.10.2" @@ -500,6 +527,7 @@ wheels = [ name = "datashare-asr-worker" source = { editable = "." } dependencies = [ + { name = "aiofile", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'x86_64' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu') or sys_platform == 'darwin' or (sys_platform != 'linux' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu')" }, { name = "caul-core", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'x86_64' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu') or sys_platform == 'darwin' or (sys_platform != 'linux' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu')" }, { name = "datashare-python", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'x86_64' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu') or sys_platform == 'darwin' or (sys_platform != 'linux' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu')" }, ] @@ -543,6 +571,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "aiofile", specifier = "==3.11.1" }, { name = "caul", extras = ["nemo"], marker = "extra == 'inference'", specifier = "==0.10.2" }, { name = "caul-core", specifier = "==0.3.1" }, { name = "cuda-bindings", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = "==12.9.4" }, diff --git a/workers/asr-worker/uv.lock b/workers/asr-worker/uv.lock index ed43a2b1..7c79c335 100644 --- a/workers/asr-worker/uv.lock +++ b/workers/asr-worker/uv.lock @@ -78,6 +78,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7e/46/02ac5e262d4af18054b3e922b2baedbb2a03289ee792162de60a865defc5/accelerate-1.13.0-py3-none-any.whl", hash = "sha256:cf1a3efb96c18f7b152eb0fa7490f3710b19c3f395699358f08decca2b8b62e0", size = 383744, upload-time = "2026-03-04T19:34:10.313Z" }, ] +[[package]] +name = "aiofile" +version = "3.11.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "caio", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'x86_64' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu') or sys_platform == 'darwin' or (sys_platform != 'linux' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu')" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/48/41/2fea7e193e061ce54eacc3b7bc0e6a99e4fcff43c78cf0a76dd781ed8334/aiofile-3.11.1.tar.gz", hash = "sha256:1f91912c6643d2a4e49ca4ae3514f0bf3867ce948a36d99a6411b8f4755f4cf9", size = 19342, upload-time = "2026-05-16T08:18:33.538Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/cd/0d76dfc5de72bde52f55f53e925c7d152d9c7906634ec1e0cbc7e8d4ad93/aiofile-3.11.1-py3-none-any.whl", hash = "sha256:ce77d14ac07f77bc2b757834a5c129321f3f705c474593deed5ab209079a52c9", size = 20446, upload-time = "2026-05-16T08:18:32.051Z" }, +] + [[package]] name = "aiohappyeyeballs" version = "2.6.1" @@ -267,6 +279,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fa/93/e8c04e80e82391a6e51f218ca49720f64236bc824e92152a2633b74cf7ab/braceexpand-0.1.7-py2.py3-none-any.whl", hash = "sha256:91332d53de7828103dcae5773fb43bc34950b0c8160e35e0f44c4427a3b85014", size = 5923, upload-time = "2021-05-07T13:49:05.146Z" }, ] +[[package]] +name = "caio" +version = "0.9.25" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/92/88/b8527e1b00c1811db339a1df8bd1ae49d146fcea9d6a5c40e3a80aaeb38d/caio-0.9.25.tar.gz", hash = "sha256:16498e7f81d1d0f5a4c0ad3f2540e65fe25691376e0a5bd367f558067113ed10", size = 26781, upload-time = "2025-12-26T15:21:36.501Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/90/543f556fcfcfa270713eef906b6352ab048e1e557afec12925c991dc93c2/caio-0.9.25-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:d6956d9e4a27021c8bd6c9677f3a59eb1d820cc32d0343cea7961a03b1371965", size = 36839, upload-time = "2025-12-26T15:21:40.267Z" }, + { url = "https://files.pythonhosted.org/packages/51/3b/36f3e8ec38dafe8de4831decd2e44c69303d2a3892d16ceda42afed44e1b/caio-0.9.25-cp311-cp311-manylinux2010_x86_64.manylinux2014_x86_64.manylinux_2_12_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bf84bfa039f25ad91f4f52944452a5f6f405e8afab4d445450978cd6241d1478", size = 80255, upload-time = "2025-12-26T15:22:20.271Z" }, + { url = "https://files.pythonhosted.org/packages/46/90/e278863c47e14ec58309aa2e38a45882fbe67b4cc29ec9bc8f65852d3e45/caio-0.9.25-cp311-cp311-manylinux_2_34_x86_64.whl", hash = "sha256:fc220b8533dcf0f238a6b1a4a937f92024c71e7b10b5a2dfc1c73604a25709bc", size = 78273, upload-time = "2026-03-04T22:08:21.368Z" }, + { url = "https://files.pythonhosted.org/packages/d3/25/79c98ebe12df31548ba4eaf44db11b7cad6b3e7b4203718335620939083c/caio-0.9.25-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:fb7ff95af4c31ad3f03179149aab61097a71fd85e05f89b4786de0359dffd044", size = 36983, upload-time = "2025-12-26T15:21:36.075Z" }, + { url = "https://files.pythonhosted.org/packages/a3/2b/21288691f16d479945968a0a4f2856818c1c5be56881d51d4dac9b255d26/caio-0.9.25-cp312-cp312-manylinux2010_x86_64.manylinux2014_x86_64.manylinux_2_12_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:97084e4e30dfa598449d874c4d8e0c8d5ea17d2f752ef5e48e150ff9d240cd64", size = 82012, upload-time = "2025-12-26T15:22:20.983Z" }, + { url = "https://files.pythonhosted.org/packages/d1/1c/0fe770b8ffc8362c48134d1592d653a81a3d8748d764bec33864db36319d/caio-0.9.25-cp312-cp312-manylinux_2_34_x86_64.whl", hash = "sha256:6bebf6f079f1341d19f7386db9b8b1f07e8cc15ae13bfdaff573371ba0575d69", size = 80200, upload-time = "2026-03-04T22:08:23.382Z" }, + { url = "https://files.pythonhosted.org/packages/86/93/1f76c8d1bafe3b0614e06b2195784a3765bbf7b0a067661af9e2dd47fc33/caio-0.9.25-py3-none-any.whl", hash = "sha256:06c0bb02d6b929119b1cfbe1ca403c768b2013a369e2db46bfa2a5761cf82e40", size = 19087, upload-time = "2025-12-26T15:22:00.221Z" }, +] + [[package]] name = "caul" version = "0.10.2" @@ -500,6 +527,7 @@ wheels = [ name = "datashare-asr-worker" source = { editable = "." } dependencies = [ + { name = "aiofile", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'x86_64' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu') or sys_platform == 'darwin' or (sys_platform != 'linux' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu')" }, { name = "caul-core", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'x86_64' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu') or sys_platform == 'darwin' or (sys_platform != 'linux' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu')" }, { name = "datashare-python", marker = "(platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'x86_64' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu') or sys_platform == 'darwin' or (sys_platform != 'linux' and extra == 'extra-20-datashare-asr-worker-cpu' and extra == 'extra-20-datashare-asr-worker-gpu')" }, ] @@ -543,6 +571,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "aiofile", specifier = "==3.11.1" }, { name = "caul", extras = ["nemo"], marker = "extra == 'inference'", specifier = "==0.10.2" }, { name = "caul-core", specifier = "==0.3.1" }, { name = "cuda-bindings", marker = "sys_platform == 'linux' and extra == 'gpu'", specifier = "==12.9.4" }, diff --git a/workers/translation-worker/translation_worker/activities.py b/workers/translation-worker/translation_worker/activities.py index 1d681197..dbfe1be9 100644 --- a/workers/translation-worker/translation_worker/activities.py +++ b/workers/translation-worker/translation_worker/activities.py @@ -1,6 +1,6 @@ import asyncio import logging -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable +from collections.abc import AsyncGenerator, AsyncIterator, Iterable from functools import partial from typing import Any, cast @@ -12,6 +12,7 @@ ActivityWithProgress, activity_defn, config_cache_key, + publish_and_consume, to_raw_async_progress, ) from elasticsearch._async.helpers import async_bulk @@ -212,7 +213,7 @@ async def translate_docs_act( consumer = asyncio.create_task( _write_translations_to_es(es_client, queue=es_queue, project=project) ) - n_docs, _ = await _publish_and_consume( + n_docs, _ = await publish_and_consume( publisher, publisher_callback, consumer=consumer ) return n_docs @@ -406,34 +407,6 @@ def _with_doc_type(query: dict[str, Any]) -> dict[str, Any]: return and_query(query, has_type(type_field="type", type_value=ES_DOCUMENT_TYPE)) -async def _publish_and_consume( - publisher: asyncio.Task, - publisher_completion_callback: Callable[[], None], - *, - consumer: asyncio.Task, -) -> tuple[Any, Any]: - # Publish and consume concurrently - logger.debug("starting publish and subscribe") - done, pending = await asyncio.wait( - [publisher, consumer], return_when=asyncio.FIRST_COMPLETED - ) - for d in done: - # Stop everything case of exception - exc = d.exception() - if exc: - for p in pending: - p.cancel() - raise exc - # Wait for publish to be done and push the poison pill to stop consuming - p_res = await publisher - publisher_completion_callback() - logger.debug("done publishing, waiting for consumer to complete...") - # Wait for consumption to be done - c_res = await consumer - logger.debug("done consuming !") - return p_res, c_res - - ACTIVITIES = [ TranslationActivities.translation_worker_config, TranslationActivities.create_translation_batches,