feature(embeddings): add embedding workflow + cpu/gpu workers

.github/workflows/publish_extract_worker.yml

@@ -66,7 +66,6 @@ jobs:
       - name: Login to Docker Hub
         uses: docker/login-action@v4
         with:
-          # You'll need to set these secrets
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
 
@@ -98,7 +97,6 @@ jobs:
       - name: Login to Docker Hub
         uses: docker/login-action@v4
         with:
-          # You'll need to set these secrets
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
 
@@ -112,6 +110,25 @@ jobs:
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
 
+      # reuses the extract image layers from the local builder cache
+      - name: Docker meta (embed)
+        id: meta-embed
+        uses: docker/metadata-action@v6
+        with:
+          images: icij/datashare-extract-embed-cpu-worker
+          tags: |
+            type=match,pattern=extract-worker-(.*),group=1
+
+      - name: Build and push embed image
+        uses: docker/build-push-action@v7
+        with:
+          target: embed-cpu-worker
+          context: ./workers/extract-worker
+          platforms: linux/amd64
+          push: true
+          tags: ${{ steps.meta-embed.outputs.tags }}
+          labels: ${{ steps.meta-embed.outputs.labels }}
+
   publish-extract-gpu-worker:
     runs-on: ubuntu-latest
     steps:
@@ -130,7 +147,6 @@ jobs:
       - name: Login to Docker Hub
         uses: docker/login-action@v4
         with:
-          # You'll need to set these secrets
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
 
@@ -144,6 +160,25 @@ jobs:
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
 
+      # reuses extract image layers from the local builder cache
+      - name: Docker meta (embed)
+        id: meta-embed
+        uses: docker/metadata-action@v6
+        with:
+          images: icij/datashare-extract-embed-gpu-worker
+          tags: |
+            type=match,pattern=extract-worker-(.*),group=1
+
+      - name: Build and push embed image
+        uses: docker/build-push-action@v7
+        with:
+          target: embed-gpu-worker
+          context: ./workers/extract-worker
+          platforms: linux/amd64
+          push: true
+          tags: ${{ steps.meta-embed.outputs.tags }}
+          labels: ${{ steps.meta-embed.outputs.labels }}
+
   publish-extract-cpu-mineru-worker:
     runs-on: ubuntu-latest
     steps:
@@ -162,7 +197,6 @@ jobs:
       - name: Login to Docker Hub
         uses: docker/login-action@v4
         with:
-          # You'll need to set these secrets
           username: ${{ secrets.DOCKERHUB_USERNAME }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
 

workers/extract-worker/Dockerfile

@@ -182,3 +182,14 @@ RUN apt update && \
 COPY --from=extract-gpu-mineru-builder /app /app
 COPY --from=extract-gpu-mineru-builder /root/.local/share/uv/python /root/.local/share/uv/python
 ENTRYPOINT ["entrypoints/extract_gpu_mineru_worker.sh"]
+
+
+# embedding workers reuse the extract images
+FROM extract-cpu-worker AS embed-cpu-worker
+ADD entrypoints/embed_cpu_worker.sh ./entrypoints/embed_cpu_worker.sh
+ENTRYPOINT ["entrypoints/embed_cpu_worker.sh"]
+
+
+FROM extract-gpu-worker AS embed-gpu-worker
+ADD entrypoints/embed_gpu_worker.sh ./entrypoints/embed_gpu_worker.sh
+ENTRYPOINT ["entrypoints/embed_gpu_worker.sh"]

workers/extract-worker/entrypoints/embed_cpu_worker.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+set -e
+
+uv run --no-sync datashare-python worker start \
+    --dependencies extract.embed \
+    --queue extract.embed.cpu \
+    --activity extract.embed-documents

workers/extract-worker/entrypoints/embed_gpu_worker.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+set -e
+
+uv run --no-sync datashare-python worker start \
+    --dependencies extract.embed \
+    --queue extract.embed.gpu \
+    --activity extract.embed-documents

workers/extract-worker/entrypoints/io_worker.sh

@@ -5,4 +5,7 @@ uv run --no-sync datashare-python worker start \
     --dependencies extract.io \
     --queue extract.io \
     --activity extract.worker-config \
-    --activity extract.create-markdown-batches
+    --activity extract.create-markdown-batches \
+    --activity extract.check-embeddings-index \
+    --activity extract.create-embedding-batches \
+    --activity extract.index-embeddings

workers/extract-worker/extract_worker/activities.py

@@ -11,6 +11,7 @@
 from datashare_python.objects import (
     PROCESSED_FILE_TA,
     ByteRangesPagination,
+    DatashareFile,
     Document,
     Error,
     ErrorReportWithPages,
@@ -24,10 +25,12 @@
     ActivityWithProgress,
     activity_defn,
     activity_workdir,
+    config_cache_key,
     read_jsonl_as,
     symlink_embedded_document_to_workdir,
     to_raw_async_progress,
     write_artifact,
+    write_batches,
 )
 from datashare_python.utils import ext_to_mime_types as _ext_to_mime_types
 from extract_core import (
@@ -57,13 +60,26 @@
     has_id,
     has_type,
 )
+from icij_common.iter_utils import async_batches
 from pydantic import TypeAdapter
 from temporalio import activity
 
-from .config import ExtractWorkerConfig
+from .config import EmbeddingWorkerConfig, ExtractWorkerConfig
+from .dependencies import lifespan_embedder_cache
+from .embeddings import (
+    HFEmbedder,
+    build_embedding_doc_query,
+    embed_documents_act,
+    ensure_embeddings_index,
+    index_embeddings_act,
+    make_chunker,
+    make_md_converter,
+    with_structure_artifact,
+)
 from .objects import (
     DocId,
     DocumentSearchQuery,
+    EmbedBatchResult,
     ExtractError,
     MarkdownExtractArgs,
     MarkdownExtractResponse,
@@ -81,6 +97,10 @@ class Activity(StrEnum):
     WORKER_CONFIG = "extract.worker-config"
     CREATE_MD_BATCHES = "extract.create-markdown-batches"
     EXTRACT_MD_CONTENT = "extract.extract-markdown-content"
+    CHECK_EMBEDDINGS_INDEX = "extract.check-embeddings-index"
+    CREATE_EMBEDDING_BATCHES = "extract.create-embedding-batches"
+    EMBED_DOCUMENTS = "extract.embed-documents"
+    INDEX_EMBEDDINGS = "extract.index-embeddings"
 
 
 class MarkdownExtract(ActivityWithProgress):
@@ -174,6 +194,111 @@ async def extract_markdown_content(
         return res
 
 
+class DocumentEmbedding(ActivityWithProgress):
+    @activity_defn(name=Activity.CHECK_EMBEDDINGS_INDEX)
+    async def check_embeddings_index(self, project: str, dim: int) -> None:
+        es_client = lifespan_es_client()
+        w_config = cast(ExtractWorkerConfig, lifespan_worker_config())
+        index = w_config.embedding.index_name(project)
+        logger.debug("checking embeddings index %s...", index)
+        await ensure_embeddings_index(es_client, index, dim=dim)
+
+    @activity_defn(name=Activity.CREATE_EMBEDDING_BATCHES)
+    async def create_embedding_batches(
+        self, project: str, docs: list[DocId] | DocumentSearchQuery | None
+    ) -> list[Path]:
+        es_client = lifespan_es_client()
+        w_config = cast(ExtractWorkerConfig, lifespan_worker_config())
+        workdir = w_config.paths.workdir
+        output_dir = activity_workdir(workdir, project)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        logger.debug("creating embedding batches...")
+        batch_paths = [
+            p.relative_to(workdir)
+            async for p in create_embedding_batches_act(
+                docs,
+                project,
+                artifacts_root=w_config.paths.artifacts,
+                output_dir=output_dir,
+                target_n_docs_per_task=w_config.embedding.target_n_docs_per_task,
+                es_client=es_client,
+            )
+        ]
+        logger.debug("created embedding batches !")
+        return batch_paths
+
+    @activity_defn(name=Activity.EMBED_DOCUMENTS)
+    async def embed_documents(
+        self,
+        batch: Path,
+        project: str,
+        model: str,
+        model_revision: str | None,
+        dim: int,
+        *,
+        progress: AsyncProgressRateHandler | None = None,
+    ) -> EmbedBatchResult:
+        w_config = cast(ExtractWorkerConfig, lifespan_worker_config())
+        embedding_config = w_config.embedding
+        workdir = w_config.paths.workdir
+        output_dir = activity_workdir(workdir, project)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        batch = workdir / batch
+        output_path = output_dir / f"{batch.stem}-chunks.jsonl"
+        cache = lifespan_embedder_cache()
+        cache_key = f"{config_cache_key(embedding_config)}-{model}@{model_revision}"
+        factory = partial(
+            _load_embedding_resources, embedding_config, model, model_revision
+        )
+        with cache.get_or_cache_resource(cache_key, factory) as resources:
+            chunker, md_converter, embedder = resources
+            logger.debug("embedding documents from %s...", batch)
+            response = await embed_documents_act(
+                batch,
+                chunker=chunker,
+                md_converter=md_converter,
+                embedder=embedder,
+                dim=dim,
+                artifacts_root=w_config.paths.artifacts,
+                output_path=output_path,
+                progress=progress,
+            )
+        logger.debug("embedded %s chunks !", response.n_chunks)
+        return EmbedBatchResult(
+            chunks=output_path.relative_to(workdir), response=response
+        )
+
+    @activity_defn(name=Activity.INDEX_EMBEDDINGS)
+    async def index_embeddings(self, chunks: Path, project: str, dim: int) -> int:
+        es_client = lifespan_es_client()
+        w_config = cast(ExtractWorkerConfig, lifespan_worker_config())
+        embedding_config = w_config.embedding
+        index = embedding_config.index_name(project)
+        logger.debug("indexing embedded chunks into %s...", index)
+        n_indexed = await index_embeddings_act(
+            w_config.paths.workdir / chunks,
+            es_client=es_client,
+            index=index,
+            dim=dim,
+            bulk_size=embedding_config.index_bulk_size,
+        )
+        logger.debug("indexed %s chunks !", n_indexed)
+        return n_indexed
+
+
+def _load_embedding_resources(
+    config: EmbeddingWorkerConfig, model: str, model_revision: str | None
+) -> tuple["HybridChunker", "DocumentConverter", HFEmbedder]:  # noqa: F821
+    logger.debug("loading embedding model %s (revision %s)...", model, model_revision)
+    resources = (
+        make_chunker(config, model, model_revision),
+        make_md_converter(),
+        HFEmbedder(config, model, model_revision),
+    )
+    logger.debug("embedding model loaded !")
+    return resources
+
+
 # Sort documents aiming for consistent processing type in a batch
 _DOC_SORT = [f"{DOC_CONTENT_TYPE}:asc", f"{DOC_LANGUAGE}:asc", "_doc:asc"]
 _DOC_CONTENT_SOURCES = [
@@ -206,6 +331,25 @@ async def create_markdown_extract_batches_act(
         yield p
 
 
+async def create_embedding_batches_act(
+    docs: list[DocId] | DocumentSearchQuery | None,
+    project: str,
+    *,
+    artifacts_root: Path,
+    output_dir: Path,
+    target_n_docs_per_task: int,
+    es_client: ESClient,
+) -> AsyncIterable[Path]:
+    query = build_embedding_doc_query(docs)
+    docs = with_structure_artifact(
+        _search_docs(es_client, project, query), artifacts_root
+    )
+    docs = (DatashareFile.from_parent(d) async for d in docs)
+    batches = async_batches(docs, target_n_docs_per_task)
+    async for p in write_batches(batches, output_dir, prefix="embed_"):
+        yield p
+
+
 async def extract_markdown_content_act(
     pipeline: Pipeline,
     batch: Path,
@@ -358,4 +502,8 @@ def ext_to_mime_types(ext: SupportedExt) -> set[str]:
     MarkdownExtract.extract_worker_config,
     MarkdownExtract.create_markdown_extract_batches,
     MarkdownExtract.extract_markdown_content,
+    DocumentEmbedding.check_embeddings_index,
+    DocumentEmbedding.create_embedding_batches,
+    DocumentEmbedding.embed_documents,
+    DocumentEmbedding.index_embeddings,
 ]