datashare-extract-worker 0.1.3__tar.gz

datashare_extract_worker-0.1.3/.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_extract_worker-0.1.3/Dockerfile

@@ -0,0 +1,114 @@
+FROM python:3.13-slim-trixie AS builder
+
+
+FROM python:3.13-slim-trixie AS extract-worker-builder
+
+ENV PYTHONUNBUFFERED=1
+ENV UV_HTTP_TIMEOUT=300
+ENV UV_LINK_MODE=copy
+ENV UV_COMPILE_BYTECODE=1
+ENV UV_NO_DEV=1
+
+COPY --from=ghcr.io/astral-sh/uv:0.11.6 /uv /uvx /bin/
+
+WORKDIR /app
+
+
+FROM extract-worker-builder AS io-worker
+# Install deps first to optimize layer cache
+RUN --mount=type=cache,target=~/.cache/uv \
+    --mount=type=bind,source=uv.dist.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    uv sync -v --frozen --no-editable --no-install-project
+
+# Then copy code
+ADD uv.dist.lock ./uv.lock
+ADD pyproject.toml README.md ./
+ADD extract_worker ./extract_worker/
+ADD entrypoints/io_worker.sh ./entrypoints/io_worker.sh
+
+# Then install service
+RUN --mount=type=cache,target=~/.cache/uv uv sync -v --frozen --no-editable
+RUN rm -rf ~/.cache/pip
+
+ENTRYPOINT ["entrypoints/io_worker.sh"]
+
+
+FROM extract-worker-builder AS extract-cpu-worker
+# Install deps first to optimize layer cache
+RUN --mount=type=cache,target=~/.cache/uv \
+    --mount=type=bind,source=uv.dist.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    uv sync -v --frozen --no-editable --no-install-project --extra base --extra cpu
+
+# Then copy code
+ADD uv.dist.lock ./uv.lock
+ADD pyproject.toml README.md ./
+ADD extract_worker ./extract_worker/
+ADD entrypoints/extract_cpu_worker.sh ./entrypoints/extract_cpu_worker.sh
+
+# Then install service
+RUN --mount=type=cache,target=~/.cache/uv uv sync -v --frozen --no-editable --extra base --extra cpu
+RUN rm -rf ~/.cache/pip
+
+ENTRYPOINT ["entrypoints/extract_cpu_worker.sh"]
+
+
+FROM extract-worker-builder AS extract-gpu-worker
+# Install deps first to optimize layer cache
+RUN --mount=type=cache,target=~/.cache/uv \
+    --mount=type=bind,source=uv.dist.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    uv sync -v --frozen --no-editable --no-install-project --extra base --extra gpu
+
+# Then copy code
+ADD uv.dist.lock ./uv.lock
+ADD pyproject.toml README.md ./
+ADD extract_worker ./extract_worker/
+ADD entrypoints/extract_gpu_worker.sh ./entrypoints/extract_gpu_worker.sh
+
+# Then install service
+RUN --mount=type=cache,target=~/.cache/uv uv sync -v --frozen --no-editable --extra base --extra gpu
+RUN rm -rf ~/.cache/pip
+
+ENTRYPOINT ["entrypoints/extract_gpu_worker.sh"]
+
+
+FROM extract-worker-builder AS extract-cpu-mineru-worker
+# Install deps first to optimize layer cache
+RUN --mount=type=cache,target=~/.cache/uv \
+    --mount=type=bind,source=uv.dist.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    uv sync -v --frozen --no-editable --no-install-project --extra mineru --extra cpu
+
+# Then copy code
+ADD uv.dist.lock ./uv.lock
+ADD pyproject.toml README.md ./
+ADD extract_worker ./extract_worker/
+ADD entrypoints/extract_cpu_worker.sh ./entrypoints/extract_cpu_worker.sh
+
+# Then install service
+RUN --mount=type=cache,target=~/.cache/uv uv sync -v --frozen --no-editable --extra mineru --extra cpu
+RUN rm -rf ~/.cache/pip
+
+ENTRYPOINT ["entrypoints/extract_cpu_worker.sh"]
+
+
+FROM extract-worker-builder AS extract-gpu-mineru-worker
+# Install deps first to optimize layer cache
+RUN --mount=type=cache,target=~/.cache/uv \
+    --mount=type=bind,source=uv.dist.lock,target=uv.lock \
+    --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
+    uv sync -v --frozen --no-editable --no-install-project --extra mineru --extra gpu
+
+# Then copy code
+ADD uv.dist.lock ./uv.lock
+ADD pyproject.toml README.md ./
+ADD extract_worker ./extract_worker/
+ADD entrypoints/extract_gpu_worker.sh ./entrypoints/extract_gpu_worker.sh
+
+# Then install service
+RUN --mount=type=cache,target=~/.cache/uv uv sync -v --frozen --no-editable --extra mineru --extra gpu
+RUN rm -rf ~/.cache/pip
+
+ENTRYPOINT ["entrypoints/extract_gpu_worker.sh"]

datashare_extract_worker-0.1.3/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.4
+Name: datashare-extract-worker
+Version: 0.1.3
+Author-email: Clément Doumouro <cdoumouro@icij.org>, Clément Doumouro <clement.doumouro@gmail.com>, Lion Summerbell <lsummerbell@icij.org>
+Requires-Python: <3.14,>=3.11
+Requires-Dist: datashare-python~=0.8.20
+Requires-Dist: extract-python==0.4.2
+Requires-Dist: temporalio==1.23.0
+Provides-Extra: base
+Requires-Dist: extract-python[docling,marker]==0.4.2; extra == 'base'
+Provides-Extra: cpu
+Requires-Dist: torch==2.11.0; extra == 'cpu'
+Requires-Dist: torchvision==0.26.0; extra == 'cpu'
+Provides-Extra: gpu
+Requires-Dist: cuda-bindings==12.9.4; (sys_platform == 'linux') and extra == 'gpu'
+Requires-Dist: torch==2.11.0+cu129; (sys_platform == 'linux') and extra == 'gpu'
+Requires-Dist: torchvision==0.26.0+cu129; (sys_platform == 'linux') and extra == 'gpu'
+Provides-Extra: mineru
+Requires-Dist: extract-python[mineru]==0.4.2; extra == 'mineru'

datashare_extract_worker-0.1.3/entrypoints/extract_cpu_worker.sh

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

datashare_extract_worker-0.1.3/entrypoints/extract_gpu_worker.sh

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

datashare_extract_worker-0.1.3/entrypoints/io_worker.sh

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

datashare_extract_worker-0.1.3/extract_worker/activities.py

@@ -0,0 +1,325 @@
+import contextlib
+import logging
+import mimetypes
+import os
+from collections.abc import AsyncIterable
+from functools import cache
+from itertools import chain
+from pathlib import Path
+from typing import Any, cast
+
+from datashare_python.dependencies import lifespan_es_client, lifespan_worker_config
+from datashare_python.objects import DocArtifact, Document, DocumentLocation
+from datashare_python.utils import (
+    ActivityWithProgress,
+    activity_defn,
+    activity_workdir,
+    read_jsonl,
+    write_artifact,
+)
+from extract_python import Pipeline
+from extract_python.objects import InputDoc, OutputFormat, SupportedExt
+from icij_common.es import (
+    DOC_CONTENT_TYPE,
+    DOC_LANGUAGE,
+    DOC_METADATA,
+    DOC_PATH,
+    DOC_ROOT_ID,
+    ES_DOCUMENT_TYPE,
+    HITS,
+    QUERY,
+    SOURCE,
+    ESClient,
+    ESSort,
+    and_query,
+    has_id,
+    has_type,
+)
+from pydantic import TypeAdapter
+
+from .config import ExtractWorkerConfig
+from .constants import MARKDOWN_DIRNAME, MARKDOWN_METADATA_KEY
+from .objects import (
+    DocId,
+    DocumentSearchQuery,
+    ErrorReport,
+    MarkdownExtractResponse,
+    PipelineConfig,
+    ProcessedDoc,
+    ProcessingReport,
+)
+
+logger = logging.getLogger(__name__)
+
+mimetypes.init()
+
+
+class MarkdownExtract(ActivityWithProgress):
+    @activity_defn(name="extract.worker_config")
+    async def extract_worker_config(self) -> ExtractWorkerConfig:
+        worker_config = cast(ExtractWorkerConfig, lifespan_worker_config())
+        return worker_config
+
+    @activity_defn(name="extract.create-markdown-batches")
+    async def create_markdown_extract_batches(
+        self,
+        project: str,
+        docs: list[DocId] | DocumentSearchQuery | None,
+        config: PipelineConfig,
+    ) -> list[Path]:
+        es_client = lifespan_es_client()
+        worker_config = cast(ExtractWorkerConfig, lifespan_worker_config())
+        workdir = worker_config.workdir
+        artifacts_root = worker_config.artifacts_root
+        output_dir = activity_workdir(workdir, project)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        target_n_pages_per_batch = worker_config.markdown.target_n_pages_per_batch
+        supported_exts = config.supported_exts()
+        batch_paths = [
+            p.relative_to(workdir)
+            async for p in create_markdown_extract_batches_act(
+                docs,
+                project,
+                supported_exts,
+                artifacts_root=artifacts_root,
+                workdir=workdir,
+                output_dir=output_dir,
+                target_n_pages_per_batch=target_n_pages_per_batch,
+                es_client=es_client,
+            )
+        ]
+        return batch_paths
+
+    @activity_defn(name="extract.extract-markdown-content")
+    async def extract_markdown_content(
+        self, batch: Path, project: str, config: PipelineConfig
+    ) -> MarkdownExtractResponse:
+        pipeline = Pipeline.from_config(config)
+        worker_config = cast(ExtractWorkerConfig, lifespan_worker_config())
+        workdir = worker_config.workdir
+        output_dir = activity_workdir(workdir, project)
+        output_dir.mkdir(parents=True, exist_ok=True)
+        batch = workdir / batch
+        return await extract_markdown_content_act(
+            pipeline,
+            batch,
+            worker_config=worker_config,
+            output_dir=output_dir,
+        )
+
+
+# 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 = [DOC_PATH, DOC_ROOT_ID, DOC_LANGUAGE, DOC_METADATA]
+
+
+async def create_markdown_extract_batches_act(
+    docs: list[DocId] | DocumentSearchQuery | None,
+    project: str,
+    supported_exts: set[SupportedExt],
+    *,
+    artifacts_root: Path,
+    workdir: Path,
+    output_dir: Path,
+    target_n_pages_per_batch: int,
+    es_client: ESClient | None = None,
+) -> AsyncIterable[Path]:
+    # TODO: supported content types should be args
+    query = _build_doc_query(docs, supported_exts)
+    docs = (
+        _symlink_embedded_processed_doc_to_workdir(d, artifacts_root, workdir=workdir)
+        async for d in _search_docs(es_client, project, query, sort=_DOC_SORT)
+    )
+    batches = _batch_by_n_pages(docs, target_n_pages_per_batch=target_n_pages_per_batch)
+    async for p in _write_batches(batches, output_dir):
+        yield p
+
+
+_BatchTypeAdapter = TypeAdapter(list[ProcessedDoc])
+
+
+async def extract_markdown_content_act(
+    pipeline: Pipeline,
+    batch: Path,
+    *,
+    worker_config: ExtractWorkerConfig,
+    output_dir: Path,
+) -> MarkdownExtractResponse:
+    docs = _BatchTypeAdapter.validate_python(list(read_jsonl(batch)))
+    docs_root = worker_config.docs_root
+    artifacts_root = worker_config.artifacts_root
+    workdir = worker_config.workdir
+    input_docs = (
+        InputDoc.from_path(
+            d.locate(
+                original_root=docs_root, artifacts_root=artifacts_root, workdir=workdir
+            )
+        )
+        for d in docs
+    )
+    results = pipeline.extract_content(
+        input_docs, output_format=OutputFormat.MARKDOWN, output_path=output_dir
+    )
+    docs = iter(docs)
+    n_docs, n_pages, n_successes, n_successes_pages = 0, 0, 0, 0
+    errors = []
+    async for extract_res in results:
+        doc = next(docs)
+        n_docs += 1
+        n_pages += doc.n_pages
+        if extract_res.errors:
+            error = ErrorReport(
+                doc=doc, status=extract_res.status, errors=extract_res.errors
+            )
+            errors.append(error)
+        else:
+            n_successes += 1
+            n_successes_pages += doc.n_pages
+            md_path = output_dir / extract_res.output.path
+            artifact = DocArtifact(
+                project=doc.index,
+                doc_id=doc.id,
+                artifact=md_path,
+                metadata_key=MARKDOWN_METADATA_KEY,
+                filename=MARKDOWN_DIRNAME,
+            )
+            write_artifact(artifacts_root, artifact)
+    processed = ProcessingReport(n_docs=n_docs, n_pages=n_pages)
+    successes = ProcessingReport(n_docs=n_successes, n_pages=n_successes_pages)
+    response = MarkdownExtractResponse(
+        processed=processed, successes=successes, errors=errors
+    )
+    return response
+
+
+def _with_supported_exts_query(supported_exts: set[SupportedExt]) -> dict[str, Any]:
+    supported_mimes = sorted(chain(*(ext_to_mime_types(f) for f in supported_exts)))
+    format_query = {"terms": {DOC_CONTENT_TYPE: supported_mimes}}
+    query = and_query(
+        format_query, has_type(type_field="type", type_value=ES_DOCUMENT_TYPE)
+    )
+    return query[QUERY]
+
+
+def _build_doc_query(
+    docs: list[DocId] | DocumentSearchQuery | None, supported_exts: set[SupportedExt]
+) -> dict[str, Any]:
+    format_query = _with_supported_exts_query(supported_exts)
+    match docs:
+        case dict():
+            if not docs:
+                return {QUERY: format_query}
+            return and_query(format_query, docs)
+        case None:
+            return {QUERY: format_query}
+        case list():
+            return and_query(format_query, has_id(docs))
+        case _:
+            raise ValueError(f"unsupported format {docs.__class__.__name__}")
+
+
+async def _search_docs(
+    es_client: ESClient, project: str, query: dict[str, Any], sort: ESSort = None
+) -> AsyncIterable[ProcessedDoc]:
+    async for page in es_client.poll_search_pages(
+        index=project,
+        body=query,
+        sort=sort,
+        _source_includes=_DOC_CONTENT_SOURCES,
+    ):
+        for hit in page[HITS][HITS]:
+            n_pages = None
+            meta = hit[SOURCE].get(DOC_METADATA)
+            if meta is not None:
+                n_pages = meta.get("tika_metadata_xmptpg_npages")
+            yield ProcessedDoc.from_fs_doc(
+                Document.from_es(hit).to_filesystem(), n_pages=n_pages
+            )
+
+
+async def _batch_by_n_pages(
+    docs: AsyncIterable[ProcessedDoc], target_n_pages_per_batch: int
+) -> AsyncIterable[list[ProcessedDoc]]:
+    current_n_pages = 0
+    current_batch = []
+    async for d in docs:
+        if current_n_pages >= target_n_pages_per_batch:
+            yield current_batch
+            current_n_pages = 0
+            current_batch = []
+        current_batch.append(d)
+        current_n_pages += d.n_pages
+    if current_batch:
+        yield current_batch
+
+
+async def _write_batches(
+    batches: AsyncIterable[list[ProcessedDoc]], root: Path
+) -> AsyncIterable[Path]:
+    batch_id = 0
+    async for batch in batches:
+        batch_path = root / f"{batch_id}.jsonl"
+        with batch_path.open("w") as f:
+            for fs_doc in batch:
+                f.write(f"{fs_doc.model_dump_json()}\n")
+        yield batch_path
+        batch_id += 1
+
+
+def _symlink_embedded_processed_doc_to_workdir(
+    doc: ProcessedDoc, artifacts_root: Path, *, workdir: Path
+) -> ProcessedDoc:
+    match doc.location:
+        case DocumentLocation.ARTIFACTS:
+            symlinks_dir = workdir / doc.index / "symlinks"
+            symlinks_dir.mkdir(parents=True, exist_ok=True)
+            symlink_path = Path(*doc.path.parts[:-1], doc.id)
+            # Replace the "raw" with the doc id
+            doc_ext = Path(doc.resource_name).suffix
+            symlink_path = symlink_path.relative_to(Path(doc.index))
+            symlink_path = symlinks_dir / f"{symlink_path}{doc_ext}"
+            symlink_path.parent.mkdir(parents=True, exist_ok=True)
+            artifact_path = artifacts_root / doc.path
+            with contextlib.suppress(FileExistsError):
+                os.symlink(artifact_path, symlink_path)
+            return ProcessedDoc(
+                path=symlink_path.relative_to(workdir),
+                id=doc.id,
+                location=DocumentLocation.WORKDIR,
+                index=doc.index,
+                resource_name=doc.resource_name,
+                n_pages=doc.n_pages,
+            )
+        case DocumentLocation.ORIGINAL:
+            return doc
+        case _:
+            raise ValueError(f"unsupported location {doc.location}")
+
+
+@cache
+def ext_to_mime_types(ext: SupportedExt) -> set[str]:
+    # All particular cases
+    match ext:
+        case SupportedExt.NXML:
+            return ext_to_mime_types(SupportedExt.XML)
+        case SupportedExt.ADOC | SupportedExt.ASCIIDOC:
+            return {"text/x-asciidoc"}
+        case SupportedExt.QMD | SupportedExt.RMD:
+            return ext_to_mime_types(SupportedExt.MD)
+        case SupportedExt.XBRL:
+            return ext_to_mime_types(SupportedExt.HTLM)
+    try:
+        types = {mimetypes.types_map[ext]}
+    except KeyError as e:
+        raise ValueError(f"unsupported mimetype {ext}") from e
+    other = mimetypes.common_types.get(ext)
+    if other is not None:
+        types.add(other)
+    return types
+
+
+ACTIVITIES = [
+    MarkdownExtract.extract_worker_config,
+    MarkdownExtract.create_markdown_extract_batches,
+    MarkdownExtract.extract_markdown_content,
+]

datashare_extract_worker-0.1.3/extract_worker/config.py

@@ -0,0 +1,18 @@
+from datashare_python.config import WorkerConfig
+from datashare_python.objects import DatashareModel
+from pydantic import Field
+
+from .constants import TorchDevice
+
+
+class MarkdownExtractConfig(DatashareModel):
+    target_n_pages_per_batch: int = 100
+
+
+class ExtractWorkerConfig(WorkerConfig):
+    device: TorchDevice = Field(default=TorchDevice.CPU, frozen=True)
+
+    markdown: MarkdownExtractConfig = Field(default_factory=MarkdownExtractConfig)
+
+
+WORKER_CONFIG_CLS = ExtractWorkerConfig

datashare_extract_worker-0.1.3/extract_worker/constants.py

@@ -0,0 +1,34 @@
+from enum import StrEnum
+
+from extract_python import PipelineType
+
+MARKDOWN_METADATA_KEY = "extract.markdown"
+MARKDOWN_DIRNAME = "markdown"
+
+
+class TaskQueue(StrEnum):
+    WORKFLOWS = "datashare.workflows"
+    IO = "extract.io"
+    EXTRACT_GPU_MINER_U = "extract.gpu.mineru"
+    EXTRACT_CPU_MINER_U = "extract.cpu.mineru"
+    EXTRACT_GPU = "extract.gpu"
+    EXTRACT_CPU = "extract.cpu"
+
+
+class TorchDevice(StrEnum):
+    CPU = "cpu"
+    GPU = "cuda"
+
+    def md_extract_queue(self, pipeline: PipelineType) -> TaskQueue:
+        is_mineru = pipeline is PipelineType.MINER_U
+        match self:
+            case TorchDevice.GPU:
+                if is_mineru:
+                    return TaskQueue.EXTRACT_GPU_MINER_U
+                return TaskQueue.EXTRACT_GPU
+            case TorchDevice.CPU:
+                if is_mineru:
+                    return TaskQueue.EXTRACT_CPU_MINER_U
+                return TaskQueue.EXTRACT_CPU
+            case _:
+                raise ValueError(f"unsupported TorchDevice {self}")

datashare_extract_worker-0.1.3/extract_worker/dependencies.py

@@ -0,0 +1,11 @@
+from datashare_python.dependencies import (
+    lifespan_es_client,  # noqa: F401
+    set_es_client,
+    set_loggers,
+    set_worker_config,
+)
+
+IO = [set_worker_config, set_loggers, set_es_client]
+EXTRACT = [set_worker_config, set_loggers]
+
+DEPENDENCIES = {"extract.io": IO, "extract.extract": EXTRACT}

datashare_extract_worker-0.1.3/extract_worker/objects.py

@@ -0,0 +1,68 @@
+from typing import Annotated, Any, Self
+
+from datashare_python.objects import DatashareModel, FilesystemDocument
+from extract_python import (
+    DoclingPipelineConfig,
+    PipelineType,
+    Status,
+)
+from extract_python import (
+    PipelineConfig as ExtractPipelineConfig,
+)
+from extract_python.objects import Error
+from icij_common.pydantic_utils import make_enum_discriminator, tagged_union
+from pydantic import Discriminator, Field
+
+DocumentSearchQuery = dict[str, Any]
+DocId = str
+
+
+pipeline_discriminator = make_enum_discriminator("pipeline", PipelineType)
+PipelineConfig = Annotated[
+    tagged_union(ExtractPipelineConfig.__subclasses__(), lambda t: t.pipeline.default),
+    Discriminator(pipeline_discriminator),
+]
+
+
+class MarkdownExtractArgs(DatashareModel):
+    project: str
+    docs: list[DocId] | DocumentSearchQuery | None
+    config: PipelineConfig = Field(default_factory=DoclingPipelineConfig)
+
+
+class ProcessingReport(DatashareModel):
+    n_docs: int = 0
+    n_pages: int = 0
+
+    def __add__(self, other: Self) -> Self:
+        return ProcessingReport(
+            n_docs=other.n_docs + self.n_docs, n_pages=other.n_pages + self.n_pages
+        )
+
+
+class ProcessedDoc(FilesystemDocument):
+    n_pages: int
+
+    @classmethod
+    def from_fs_doc(cls, fs_doc: FilesystemDocument, n_pages: int | None) -> Self:
+        n_pages = n_pages if n_pages is not None else 1
+        return cls(n_pages=n_pages, **fs_doc.model_dump())
+
+
+class ErrorReport(DatashareModel):
+    doc: ProcessedDoc
+    status: Status
+    errors: list[Error] = []
+
+
+class MarkdownExtractResponse(DatashareModel):
+    processed: ProcessingReport = Field(default_factory=ProcessingReport)
+    successes: ProcessingReport = Field(default_factory=ProcessingReport)
+    errors: list[ErrorReport] = Field(default_factory=list)
+
+    @classmethod
+    def from_responses(cls, *responses: Self) -> Self:
+        processed = sum((r.processed for r in responses), start=ProcessingReport())
+        successes = sum((r.successes for r in responses), start=ProcessingReport())
+        errors = sum((r.errors for r in responses), start=[])
+        return cls(processed=processed, successes=successes, errors=errors)