feat: Add parameters for Kubeflow pipeline engine (WIP) (docling-project#107)

Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>

Author: dolfim-ibm

docling_serve/app.py

@@ -28,6 +28,10 @@
 
 from docling.datamodel.base_models import DocumentStream
 
+from docling_serve.datamodel.callback import (
+    ProgressCallbackRequest,
+    ProgressCallbackResponse,
+)
 from docling_serve.datamodel.convert import ConvertDocumentsOptions
 from docling_serve.datamodel.requests import (
     ConvertDocumentFileSourcesRequest,
@@ -45,11 +49,12 @@
     get_converter,
     get_pdf_pipeline_opts,
 )
-from docling_serve.engines import get_orchestrator
-from docling_serve.engines.async_local.orchestrator import (
-    AsyncLocalOrchestrator,
-    TaskNotFoundError,
+from docling_serve.engines.async_orchestrator import (
+    BaseAsyncOrchestrator,
+    ProgressInvalid,
 )
+from docling_serve.engines.async_orchestrator_factory import get_async_orchestrator
+from docling_serve.engines.base_orchestrator import TaskNotFoundError
 from docling_serve.helper_functions import FormDepends
 from docling_serve.response_preparation import process_results
 from docling_serve.settings import docling_serve_settings
@@ -94,7 +99,7 @@ async def lifespan(app: FastAPI):
     pdf_format_option = get_pdf_pipeline_opts(ConvertDocumentsOptions())
     get_converter(pdf_format_option)
 
-    orchestrator = get_orchestrator()
+    orchestrator = get_async_orchestrator()
 
     # Start the background queue processor
     queue_task = asyncio.create_task(orchestrator.process_queue())
@@ -308,7 +313,7 @@ async def process_file(
         response_model=TaskStatusResponse,
     )
     async def process_url_async(
-        orchestrator: Annotated[AsyncLocalOrchestrator, Depends(get_orchestrator)],
+        orchestrator: Annotated[BaseAsyncOrchestrator, Depends(get_async_orchestrator)],
         conversion_request: ConvertDocumentsRequest,
     ):
         task = await orchestrator.enqueue(request=conversion_request)
@@ -319,6 +324,7 @@ async def process_url_async(
             task_id=task.task_id,
             task_status=task.task_status,
             task_position=task_queue_position,
+            task_meta=task.processing_meta,
         )
 
     # Task status poll
@@ -327,7 +333,7 @@ async def process_url_async(
         response_model=TaskStatusResponse,
     )
     async def task_status_poll(
-        orchestrator: Annotated[AsyncLocalOrchestrator, Depends(get_orchestrator)],
+        orchestrator: Annotated[BaseAsyncOrchestrator, Depends(get_async_orchestrator)],
         task_id: str,
         wait: Annotated[
             float, Query(help="Number of seconds to wait for a completed status.")
@@ -342,6 +348,7 @@ async def task_status_poll(
             task_id=task.task_id,
             task_status=task.task_status,
             task_position=task_queue_position,
+            task_meta=task.processing_meta,
         )
 
     # Task status websocket
@@ -350,7 +357,7 @@ async def task_status_poll(
     )
     async def task_status_ws(
         websocket: WebSocket,
-        orchestrator: Annotated[AsyncLocalOrchestrator, Depends(get_orchestrator)],
+        orchestrator: Annotated[BaseAsyncOrchestrator, Depends(get_async_orchestrator)],
         task_id: str,
     ):
         await websocket.accept()
@@ -375,6 +382,7 @@ async def task_status_ws(
                 task_id=task.task_id,
                 task_status=task.task_status,
                 task_position=task_queue_position,
+                task_meta=task.processing_meta,
             )
             await websocket.send_text(
                 WebsocketMessage(
@@ -389,6 +397,7 @@ async def task_status_ws(
                     task_id=task.task_id,
                     task_status=task.task_status,
                     task_position=task_queue_position,
+                    task_meta=task.processing_meta,
                 )
                 await websocket.send_text(
                     WebsocketMessage(
@@ -416,7 +425,7 @@ async def task_status_ws(
         },
     )
     async def task_result(
-        orchestrator: Annotated[AsyncLocalOrchestrator, Depends(get_orchestrator)],
+        orchestrator: Annotated[BaseAsyncOrchestrator, Depends(get_async_orchestrator)],
         task_id: str,
     ):
         result = await orchestrator.task_result(task_id=task_id)
@@ -427,4 +436,23 @@ async def task_result(
             )
         return result
 
+    # Update task progress
+    @app.post(
+        "/v1alpha/callback/task/progress",
+        response_model=ProgressCallbackResponse,
+    )
+    async def callback_task_progress(
+        orchestrator: Annotated[BaseAsyncOrchestrator, Depends(get_async_orchestrator)],
+        request: ProgressCallbackRequest,
+    ):
+        try:
+            await orchestrator.receive_task_progress(request=request)
+            return ProgressCallbackResponse(status="ack")
+        except TaskNotFoundError:
+            raise HTTPException(status_code=404, detail="Task not found.")
+        except ProgressInvalid as err:
+            raise HTTPException(
+                status_code=400, detail=f"Invalid progress payload: {err}"
+            )
+
     return app

docling_serve/datamodel/callback.py

@@ -0,0 +1,50 @@
+import enum
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+
+class ProgressKind(str, enum.Enum):
+    SET_NUM_DOCS = "set_num_docs"
+    UPDATE_PROCESSED = "update_processed"
+
+
+class BaseProgress(BaseModel):
+    kind: ProgressKind
+
+
+class ProgressSetNumDocs(BaseProgress):
+    kind: Literal[ProgressKind.SET_NUM_DOCS] = ProgressKind.SET_NUM_DOCS
+
+    num_docs: int
+
+
+class SucceededDocsItem(BaseModel):
+    source: str
+
+
+class FailedDocsItem(BaseModel):
+    source: str
+    error: str
+
+
+class ProgressUpdateProcessed(BaseProgress):
+    kind: Literal[ProgressKind.UPDATE_PROCESSED] = ProgressKind.UPDATE_PROCESSED
+
+    num_processed: int
+    num_succeeded: int
+    num_failed: int
+
+    docs_succeeded: list[SucceededDocsItem]
+    docs_failed: list[FailedDocsItem]
+
+
+class ProgressCallbackRequest(BaseModel):
+    task_id: str
+    progress: Annotated[
+        ProgressSetNumDocs | ProgressUpdateProcessed, Field(discriminator="kind")
+    ]
+
+
+class ProgressCallbackResponse(BaseModel):
+    status: Literal["ack"] = "ack"

docling_serve/datamodel/engines.py

@@ -10,3 +10,4 @@ class TaskStatus(str, enum.Enum):
 
 class AsyncEngine(str, enum.Enum):
     LOCAL = "local"
+    KFP = "kfp"

docling_serve/datamodel/kfp.py

@@ -0,0 +1,7 @@
+from pydantic import AnyUrl, BaseModel
+
+
+class CallbackSpec(BaseModel):
+    url: AnyUrl
+    headers: dict[str, str] = {}
+    ca_cert: str = ""

docling_serve/datamodel/responses.py

@@ -7,6 +7,8 @@
 from docling.utils.profiling import ProfilingItem
 from docling_core.types.doc import DoclingDocument
 
+from docling_serve.datamodel.task_meta import TaskProcessingMeta
+
 
 # Status
 class HealthCheckResponse(BaseModel):
@@ -38,6 +40,7 @@ class TaskStatusResponse(BaseModel):
     task_id: str
     task_status: str
     task_position: Optional[int] = None
+    task_meta: Optional[TaskProcessingMeta] = None
 
 
 class MessageKind(str, enum.Enum):

docling_serve/datamodel/task.py

@@ -5,13 +5,15 @@
 from docling_serve.datamodel.engines import TaskStatus
 from docling_serve.datamodel.requests import ConvertDocumentsRequest
 from docling_serve.datamodel.responses import ConvertDocumentResponse
+from docling_serve.datamodel.task_meta import TaskProcessingMeta
 
 
 class Task(BaseModel):
     task_id: str
     task_status: TaskStatus = TaskStatus.PENDING
     request: Optional[ConvertDocumentsRequest]
     result: Optional[ConvertDocumentResponse] = None
+    processing_meta: Optional[TaskProcessingMeta] = None
 
     def is_completed(self) -> bool:
         if self.task_status in [TaskStatus.SUCCESS, TaskStatus.FAILURE]:

docling_serve/datamodel/task_meta.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel
+
+
+class TaskProcessingMeta(BaseModel):
+    num_docs: int
+    num_processed: int = 0
+    num_succeeded: int = 0
+    num_failed: int = 0

docling_serve/engines/__init__.py

@@ -1,8 +0,0 @@
-from functools import lru_cache
-
-from docling_serve.engines.async_local.orchestrator import AsyncLocalOrchestrator
-
-
-@lru_cache
-def get_orchestrator() -> AsyncLocalOrchestrator:
-    return AsyncLocalOrchestrator()

docling_serve/engines/async_kfp/__init__.py

@@ -0,0 +1,137 @@
+# ruff: noqa: E402, UP006, UP035
+
+from typing import Any, Dict, List
+
+from kfp import dsl
+
+PYTHON_BASE_IMAGE = "python:3.12"
+
+
+@dsl.component(
+    base_image=PYTHON_BASE_IMAGE,
+    packages_to_install=[
+        "pydantic",
+        "docling-serve @ git+https://github.com/docling-project/docling-serve@feat-kfp-engine",
+    ],
+    pip_index_urls=["https://download.pytorch.org/whl/cpu", "https://pypi.org/simple"],
+)
+def generate_chunks(
+    run_name: str,
+    request: Dict[str, Any],
+    batch_size: int,
+    callbacks: List[Dict[str, Any]],
+) -> List[List[Dict[str, Any]]]:
+    from pydantic import TypeAdapter
+
+    from docling_serve.datamodel.callback import (
+        ProgressCallbackRequest,
+        ProgressSetNumDocs,
+    )
+    from docling_serve.datamodel.kfp import CallbackSpec
+    from docling_serve.engines.async_kfp.notify import notify_callbacks
+
+    CallbacksListType = TypeAdapter(list[CallbackSpec])
+
+    sources = request["http_sources"]
+    splits = [sources[i : i + batch_size] for i in range(0, len(sources), batch_size)]
+
+    total = sum(len(chunk) for chunk in splits)
+    payload = ProgressCallbackRequest(
+        task_id=run_name, progress=ProgressSetNumDocs(num_docs=total)
+    )
+    notify_callbacks(
+        payload=payload,
+        callbacks=CallbacksListType.validate_python(callbacks),
+    )
+
+    return splits
+
+
+@dsl.component(
+    base_image=PYTHON_BASE_IMAGE,
+    packages_to_install=[
+        "pydantic",
+        "docling-serve @ git+https://github.com/docling-project/docling-serve@feat-kfp-engine",
+    ],
+    pip_index_urls=["https://download.pytorch.org/whl/cpu", "https://pypi.org/simple"],
+)
+def convert_batch(
+    run_name: str,
+    data_splits: List[Dict[str, Any]],
+    request: Dict[str, Any],
+    callbacks: List[Dict[str, Any]],
+    output_path: dsl.OutputPath("Directory"),  # type: ignore
+):
+    from pathlib import Path
+
+    from pydantic import AnyUrl, TypeAdapter
+
+    from docling_serve.datamodel.callback import (
+        FailedDocsItem,
+        ProgressCallbackRequest,
+        ProgressUpdateProcessed,
+        SucceededDocsItem,
+    )
+    from docling_serve.datamodel.convert import ConvertDocumentsOptions
+    from docling_serve.datamodel.kfp import CallbackSpec
+    from docling_serve.datamodel.requests import HttpSource
+    from docling_serve.engines.async_kfp.notify import notify_callbacks
+
+    CallbacksListType = TypeAdapter(list[CallbackSpec])
+
+    convert_options = ConvertDocumentsOptions.model_validate(request["options"])
+    print(convert_options)
+
+    output_dir = Path(output_path)
+    output_dir.mkdir(exist_ok=True, parents=True)
+    docs_succeeded: list[SucceededDocsItem] = []
+    docs_failed: list[FailedDocsItem] = []
+    for source_dict in data_splits:
+        source = HttpSource.model_validate(source_dict)
+        filename = Path(str(AnyUrl(source.url).path)).name
+        output_filename = output_dir / filename
+        print(f"Writing {output_filename}")
+        with output_filename.open("w") as f:
+            f.write(source.model_dump_json())
+        docs_succeeded.append(SucceededDocsItem(source=source.url))
+
+    payload = ProgressCallbackRequest(
+        task_id=run_name,
+        progress=ProgressUpdateProcessed(
+            num_failed=len(docs_failed),
+            num_processed=len(docs_succeeded) + len(docs_failed),
+            num_succeeded=len(docs_succeeded),
+            docs_succeeded=docs_succeeded,
+            docs_failed=docs_failed,
+        ),
+    )
+
+    print(payload)
+    notify_callbacks(
+        payload=payload,
+        callbacks=CallbacksListType.validate_python(callbacks),
+    )
+
+
+@dsl.pipeline()
+def process(
+    batch_size: int,
+    request: Dict[str, Any],
+    callbacks: List[Dict[str, Any]] = [],
+    run_name: str = "",
+):
+    chunks_task = generate_chunks(
+        run_name=run_name,
+        request=request,
+        batch_size=batch_size,
+        callbacks=callbacks,
+    )
+    chunks_task.set_caching_options(False)
+
+    with dsl.ParallelFor(chunks_task.output, parallelism=4) as data_splits:
+        convert_batch(
+            run_name=run_name,
+            data_splits=data_splits,
+            request=request,
+            callbacks=callbacks,
+        )