datashare-extract-worker 0.1.3__py3-none-any.whl

datashare_extract_worker-0.1.3.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,11 @@
+extract_worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+extract_worker/activities.py,sha256=W8fIP47XY39qaw_6X8Gem_AnfdbloRz5I6W8pgnqIho,10909
+extract_worker/config.py,sha256=ZAcomS7KQD7SYxPZqSWGAuBOO4uuMOvj489Lyq6_2Ls,490
+extract_worker/constants.py,sha256=GQue-4WG7auapHvxj2rk9x7wfQItf1LFaCEFeSe9Qak,1012
+extract_worker/dependencies.py,sha256=_avB1j9IX-GvxUhJPA9JyqUhHnX7aUO8aY7OxMN16Qo,303
+extract_worker/objects.py,sha256=xnNx_8BfRL8T10uAezIDLdyHWg1cF3XvUh8CFMohpCA,2185
+extract_worker/workflows.py,sha256=58mTZiXwwBkcXwSEPAif0O032KxkDVaWCPVQJlnLcVw,2392
+datashare_extract_worker-0.1.3.dist-info/METADATA,sha256=QOwnKv-AK4e9S85zQLNvj9A20UeRPG_y28a8vWurIBU,915
+datashare_extract_worker-0.1.3.dist-info/WHEEL,sha256=mffPy8wBnZQn2VnJUU5jE99KsxaSfiyMHV9Yt0aLVxs,87
+datashare_extract_worker-0.1.3.dist-info/entry_points.txt,sha256=2s1sCvz-0PQaepb3RmxFVgKxOuG203vJnTdlwIvx41I,317
+datashare_extract_worker-0.1.3.dist-info/RECORD,,

datashare_extract_worker-0.1.3.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.30.1
+Root-Is-Purelib: true
+Tag: py3-none-any

datashare_extract_worker-0.1.3.dist-info/entry_points.txt

@@ -0,0 +1,11 @@
+[datashare.activities]
+activities = extract_worker.activities:ACTIVITIES
+
+[datashare.dependencies]
+dependencies = extract_worker.dependencies:DEPENDENCIES
+
+[datashare.worker_config_cls]
+worker_config_cls = extract_worker.config_:WORKER_CONFIG_CLS
+
+[datashare.workflows]
+workflows = extract_worker.workflows:WORKFLOWS

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,
+]

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

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}")

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}

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)

extract_worker/workflows.py

@@ -0,0 +1,66 @@
+import asyncio
+import logging
+from datetime import timedelta
+from enum import StrEnum
+
+from temporalio import workflow
+
+with workflow.unsafe.imports_passed_through():
+    from datashare_python.utils import WorkflowWithProgress, execute_activity
+
+    from .activities import MarkdownExtract
+    from .objects import MarkdownExtractArgs, MarkdownExtractResponse
+
+logger = logging.getLogger(__name__)
+
+
+class TaskQueues(StrEnum):
+    EXTRACT_GPU = "extract.extract-gpu"
+    EXTRACT_CPU = "extract.extract-cpu"
+    IO = "extract.io"
+    WORKFLOWS = "datashare.workflows"
+
+
+@workflow.defn(name="extract.markdown-content")
+class ExtractMarkdownContentWorkflow(WorkflowWithProgress):
+    @workflow.run
+    async def run(self, args: MarkdownExtractArgs) -> MarkdownExtractResponse:
+        # Fetch worker config
+        worker_config = await execute_activity(
+            MarkdownExtract.extract_worker_config,
+            task_queue=TaskQueues.IO,
+            start_to_close_timeout=timedelta(hours=1),
+        )
+        # Create batches almost of constant number of pages
+        batch_args = [args.project, args.docs, args.config]
+        logger.info("creating context extraction batches...")
+        heartbeat_timeout = timedelta(seconds=30)
+        extract_batches = await execute_activity(
+            MarkdownExtract.create_markdown_extract_batches,
+            args=batch_args,
+            task_queue=TaskQueues.IO,
+            start_to_close_timeout=timedelta(hours=6),
+            heartbeat_timeout=heartbeat_timeout,
+        )
+
+        # Extract Markdown content
+        # Distribute batches docs with (more or less) constant number of page per batch,
+        # across workers
+        extract_args = [(b, args.project, args.config) for b in extract_batches]
+        task_queue = worker_config.device.md_extract_queue(args.config.pipeline)
+        extract_acts = (
+            execute_activity(
+                MarkdownExtract.extract_markdown_content,
+                args=args,
+                task_queue=task_queue,
+                start_to_close_timeout=timedelta(hours=12),
+                heartbeat_timeout=heartbeat_timeout,
+            )
+            for args in extract_args
+        )
+        responses = await asyncio.gather(*extract_acts)
+        response = MarkdownExtractResponse.from_responses(*responses)
+        return response
+
+
+WORKFLOWS = [ExtractMarkdownContentWorkflow]