From 9272ef582cd0183aa55358e1e4317c35e842ef5b Mon Sep 17 00:00:00 2001 From: hedhoud <74668966+hedhoud@users.noreply.github.com> Date: Tue, 16 Jun 2026 15:08:17 +0200 Subject: [PATCH 1/2] feat(indexer): add batch upload API --- openrag/api/routers/admin/indexing.py | 195 ++++++++++++++++++ openrag/api/schemas/admin/common.py | 29 ++- .../routers/admin/test_batch_upload_api.py | 110 ++++++++++ 3 files changed, 332 insertions(+), 2 deletions(-) create mode 100644 tests/unit/api/routers/admin/test_batch_upload_api.py diff --git a/openrag/api/routers/admin/indexing.py b/openrag/api/routers/admin/indexing.py index 64b640450..b66aa5287 100644 --- a/openrag/api/routers/admin/indexing.py +++ b/openrag/api/routers/admin/indexing.py @@ -23,12 +23,16 @@ require_task_owner, ) from api.dependencies.files import ( + FORBIDDEN_CHARS_IN_FILE_ID, save_file_to_disk, validate_file_format, validate_file_id, validate_metadata, ) from api.routers.admin.task_logs import collect_task_logs +from api.schemas.admin.common import BatchUploadItem, BatchUploadResponse, BatchUploadResult +from core.indexing import validators as core_validators +from core.utils.exceptions import OpenRAGError from core.utils.filename import sanitize_filename from core.utils.log_tail import app_log_file from core.utils.logging import get_logger @@ -36,6 +40,7 @@ from fastapi import ( APIRouter, Depends, + File, Form, HTTPException, Request, @@ -44,6 +49,7 @@ status, ) from fastapi.responses import JSONResponse +from pydantic import ValidationError as PydanticValidationError logger = get_logger() @@ -59,6 +65,104 @@ def build_url(request: Request, route_name: str, *, preferred_url_scheme: str | router = APIRouter() +def _parse_batch_upload_items(raw_items: str) -> list[BatchUploadItem]: + try: + decoded = json.loads(raw_items) + except (json.JSONDecodeError, TypeError) as exc: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="items must be a JSON array", + ) from exc + if not isinstance(decoded, list): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="items must be a JSON array", + ) + try: + return [BatchUploadItem.model_validate(item) for item in decoded] + except PydanticValidationError as exc: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="items must be a JSON array of objects with file_id and optional metadata", + ) from exc + + +async def _validate_workspace_ids( + workspace_ids: list[str] | None, + *, + partition: str, + service, +) -> list[str] | None: + if workspace_ids is None: + return None + for ws_id in workspace_ids: + ws = await service.get_workspace(ws_id) + if not ws or ws["partition_name"] != partition: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail=f"Workspace '{ws_id}' not found in partition '{partition}'", + ) + return workspace_ids + + +def _validate_upload_format(file: UploadFile, metadata: dict, config) -> None: + accepted_file_formats = config.loader.file_loaders.model_dump().keys() + mimetypes = config.loader.mimetypes.to_dict() + core_validators.validate_file_format( + filename=file.filename, + accepted_formats=accepted_file_formats, + accepted_mimetypes=mimetypes.keys(), + mimetype=metadata.get("mimetype"), + ) + + +async def _queue_uploaded_file( + *, + request: Request, + partition: str, + file_id: str, + file: UploadFile, + metadata: dict, + workspace_ids: list[str] | None, + user, + config, + service, +) -> str: + if await service.file_exists(file_id, partition): + raise HTTPException( + status_code=status.HTTP_409_CONFLICT, + detail=f"File '{file_id}' already exists in partition {partition}", + ) + + parsed_workspace_ids = await _validate_workspace_ids( + workspace_ids, + partition=partition, + service=service, + ) + + original_filename = file.filename + file.filename = sanitize_filename(file.filename) + file_path = await save_file_to_disk(file, Path(config.paths.data_dir), with_random_prefix=True) + + task_id = await service.add_file( + file_path=str(file_path), + file_id=file_id, + partition=partition, + metadata=metadata, + sanitized_filename=file.filename, + original_filename=original_filename, + user=user, + workspace_ids=parsed_workspace_ids, + ) + + return build_url( + request, + "get_task_status", + preferred_url_scheme=config.server.preferred_url_scheme, + task_id=task_id, + ) + + @router.get( "/supported/types", description="""Get supported file types for indexing. @@ -187,6 +291,97 @@ async def add_file( ) +@router.post( + "/partition/{partition}/files", + response_model=BatchUploadResponse, + description="""Upload and index multiple files in one request. + +Each file is still indexed as its own task. A failure for one item does not +abort the whole batch; the response contains one result per file. + +**Request:** +- `files`: repeated multipart file field +- `items`: JSON array matching the uploaded file order + +Each `items` entry must contain `file_id` and may contain `metadata` and +`workspace_ids`. +""", +) +async def add_files( + request: Request, + partition: str, + files: list[UploadFile] = File(...), + items: str = Form(..., description="JSON array with one item per uploaded file"), + user=Depends(require_partition_editor), + _quota_check=Depends(check_user_file_quota), + config=Depends(get_config), + service=Depends(get_indexing_service), +): + batch_items = _parse_batch_upload_items(items) + if len(batch_items) != len(files): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="items must contain exactly one entry for each uploaded file", + ) + + results: list[BatchUploadResult] = [] + for file, item in zip(files, batch_items, strict=True): + raw_file_id = item.file_id + try: + file_id = core_validators.validate_file_id(raw_file_id, FORBIDDEN_CHARS_IN_FILE_ID) + metadata = dict(item.metadata or {}) + _validate_upload_format(file, metadata, config) + task_status_url = await _queue_uploaded_file( + request=request, + partition=partition, + file_id=file_id, + file=file, + metadata=metadata, + workspace_ids=item.workspace_ids, + user=user, + config=config, + service=service, + ) + results.append( + BatchUploadResult( + file_id=file_id, + status="accepted", + task_status_url=task_status_url, + ) + ) + except HTTPException as exc: + results.append( + BatchUploadResult( + file_id=raw_file_id, + status="failed", + detail=str(exc.detail), + ) + ) + except OpenRAGError as exc: + results.append( + BatchUploadResult( + file_id=raw_file_id, + status="failed", + detail=exc.message, + ) + ) + except Exception as exc: + logger.exception("Failed to queue batch upload item.", file_id=raw_file_id, error=str(exc)) + results.append( + BatchUploadResult( + file_id=raw_file_id, + status="failed", + detail=str(exc), + ) + ) + + accepted = sum(result.status == "accepted" for result in results) + failed = len(results) - accepted + response = BatchUploadResponse(accepted=accepted, failed=failed, results=results) + response_status = status.HTTP_201_CREATED if failed == 0 else status.HTTP_207_MULTI_STATUS + return JSONResponse(status_code=response_status, content=response.model_dump()) + + @router.delete( "/partition/{partition}/file/{file_id}", description="""Delete a file from a partition. diff --git a/openrag/api/schemas/admin/common.py b/openrag/api/schemas/admin/common.py index 1a653becd..3f4eda216 100644 --- a/openrag/api/schemas/admin/common.py +++ b/openrag/api/schemas/admin/common.py @@ -1,6 +1,6 @@ -from typing import Any +from typing import Any, Literal -from pydantic import BaseModel +from pydantic import BaseModel, Field class MessageResponse(BaseModel): @@ -11,6 +11,31 @@ class TaskStatusResponse(BaseModel): task_status_url: str +class BatchUploadItem(BaseModel): + """One file entry in a batch upload request.""" + + file_id: str + metadata: dict[str, Any] = Field(default_factory=dict) + workspace_ids: list[str] | None = None + + +class BatchUploadResult(BaseModel): + """Per-file outcome returned by the batch upload endpoint.""" + + file_id: str + status: Literal["accepted", "failed"] + task_status_url: str | None = None + detail: str | None = None + + +class BatchUploadResponse(BaseModel): + """Batch upload response with one result per requested file.""" + + accepted: int + failed: int + results: list[BatchUploadResult] + + class DocumentsResponse(BaseModel): documents: list[dict[str, Any]] diff --git a/tests/unit/api/routers/admin/test_batch_upload_api.py b/tests/unit/api/routers/admin/test_batch_upload_api.py new file mode 100644 index 000000000..54a4ae824 --- /dev/null +++ b/tests/unit/api/routers/admin/test_batch_upload_api.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import pytest +from api.dependencies.auth import check_user_file_quota, require_partition_editor +from api.routers.admin import indexing +from core.config.infrastructure import PathsConfig +from core.config.root import Settings +from di.providers import get_config, get_indexing_service +from fastapi import FastAPI + + +class FakeIndexingService: + """Fake indexing service that records batch upload dispatches.""" + + def __init__(self, *, existing_files: set[str] | None = None) -> None: + self.existing_files = existing_files or set() + self.calls: list[dict[str, Any]] = [] + + async def file_exists(self, file_id: str, partition: str) -> bool: + """Return whether the file exists in the fake partition.""" + return file_id in self.existing_files + + async def get_workspace(self, workspace_id: str) -> dict[str, Any] | None: + """Return a workspace belonging to the requested test partition.""" + return {"id": workspace_id, "partition_name": "batch-partition"} + + async def add_file(self, **kwargs: Any) -> str: + """Record the queued file and return a stable fake task id.""" + self.calls.append(kwargs) + return f"task-{kwargs['file_id']}" + + +def _build_app(tmp_path: Path, service: FakeIndexingService) -> FastAPI: + app = FastAPI() + app.include_router(indexing.router, prefix="/indexer") + settings = Settings(paths=PathsConfig(data_dir=tmp_path)) + app.dependency_overrides[get_config] = lambda: settings + app.dependency_overrides[get_indexing_service] = lambda: service + app.dependency_overrides[require_partition_editor] = lambda: {"id": 1, "is_admin": True} + app.dependency_overrides[check_user_file_quota] = lambda: {"id": 1, "is_admin": True} + return app + + +@pytest.mark.asyncio +async def test_batch_upload_returns_partial_success(async_client_factory, tmp_path): + """Batch upload should accept valid files while reporting failed entries.""" + service = FakeIndexingService(existing_files={"duplicate-file"}) + app = _build_app(tmp_path, service) + + items = [ + {"file_id": "new-file", "metadata": {"category": "docs"}}, + {"file_id": "duplicate-file", "metadata": {"category": "docs"}}, + ] + + async with async_client_factory(app) as client: + response = await client.post( + "/indexer/partition/batch-partition/files", + data={"items": json.dumps(items)}, + files=[ + ("files", ("new.txt", b"new content", "text/plain")), + ("files", ("duplicate.txt", b"duplicate content", "text/plain")), + ], + ) + + assert response.status_code == 207 + assert response.json()["accepted"] == 1 + assert response.json()["failed"] == 1 + assert response.json()["results"] == [ + { + "file_id": "new-file", + "status": "accepted", + "task_status_url": "http://testserver/indexer/task/task-new-file", + "detail": None, + }, + { + "file_id": "duplicate-file", + "status": "failed", + "task_status_url": None, + "detail": "File 'duplicate-file' already exists in partition batch-partition", + }, + ] + assert [call["file_id"] for call in service.calls] == ["new-file"] + assert service.calls[0]["metadata"] == {"category": "docs"} + assert service.calls[0]["sanitized_filename"] == "new.txt" + assert service.calls[0]["original_filename"] == "new.txt" + + +@pytest.mark.asyncio +async def test_batch_upload_rejects_mismatched_files_and_items(async_client_factory, tmp_path): + """A malformed batch should fail before any file is queued.""" + service = FakeIndexingService() + app = _build_app(tmp_path, service) + + async with async_client_factory(app) as client: + response = await client.post( + "/indexer/partition/batch-partition/files", + data={"items": json.dumps([{"file_id": "only-one", "metadata": {}}])}, + files=[ + ("files", ("first.txt", b"first", "text/plain")), + ("files", ("second.txt", b"second", "text/plain")), + ], + ) + + assert response.status_code == 400 + assert response.json()["detail"] == "items must contain exactly one entry for each uploaded file" + assert service.calls == [] From f5676ec9efe4e66a5c1289bcc42298fa664ab8c5 Mon Sep 17 00:00:00 2001 From: hedhoud <74668966+hedhoud@users.noreply.github.com> Date: Tue, 16 Jun 2026 15:34:18 +0200 Subject: [PATCH 2/2] feat(indexer): expose queue stage visibility --- openrag/api/routers/admin/indexing.py | 12 ++- openrag/api/routers/admin/jobs.py | 4 + openrag/core/indexing/dispatcher.py | 5 ++ .../orchestrators/indexing_service.py | 3 + openrag/services/orchestrators/job_service.py | 20 ++++- openrag/services/workers/dispatcher.py | 6 ++ openrag/services/workers/indexer_actor.py | 6 ++ openrag/services/workers/pipeline_builder.py | 13 +++ openrag/services/workers/task_state.py | 89 +++++++++++++++---- .../admin/test_task_stage_visibility.py | 59 ++++++++++++ .../orchestrators/test_indexing_service.py | 4 + .../orchestrators/test_job_service.py | 21 ++++- .../services/workers/test_indexer_worker.py | 12 +++ .../services/workers/test_pipeline_builder.py | 22 +++++ 14 files changed, 255 insertions(+), 21 deletions(-) create mode 100644 tests/unit/api/routers/admin/test_task_stage_visibility.py diff --git a/openrag/api/routers/admin/indexing.py b/openrag/api/routers/admin/indexing.py index b66aa5287..21253391f 100644 --- a/openrag/api/routers/admin/indexing.py +++ b/openrag/api/routers/admin/indexing.py @@ -620,12 +620,22 @@ async def get_task_status( status_code=status.HTTP_404_NOT_FOUND, detail=f"Task '{task_id}' not found.", ) + task_info = await service.get_task_info(task_id) content: dict[str, Any] = { "task_id": task_id, "task_state": state, - "details": task_details, + "details": (task_info or {}).get("details", task_details), } + if task_info: + content.update( + { + "current_stage": task_info.get("current_stage"), + "failed_stage": task_info.get("failed_stage"), + "stage_durations": task_info.get("stage_durations", {}), + "stage_history": task_info.get("stage_history", []), + } + ) if state == "FAILED": content["error_url"] = build_url( diff --git a/openrag/api/routers/admin/jobs.py b/openrag/api/routers/admin/jobs.py index f6440c5c6..d83675861 100644 --- a/openrag/api/routers/admin/jobs.py +++ b/openrag/api/routers/admin/jobs.py @@ -107,6 +107,10 @@ async def list_tasks( "task_id": task_id, "state": row["state"], "details": row["details"], + "current_stage": row.get("current_stage"), + "failed_stage": row.get("failed_stage"), + "stage_durations": row.get("stage_durations", {}), + "stage_history": row.get("stage_history", []), **( {"error_url": str(request.url_for("get_task_error", task_id=task_id))} if row["state"] == "FAILED" diff --git a/openrag/core/indexing/dispatcher.py b/openrag/core/indexing/dispatcher.py index 03ed0f001..3746dbd66 100644 --- a/openrag/core/indexing/dispatcher.py +++ b/openrag/core/indexing/dispatcher.py @@ -78,6 +78,11 @@ async def get_task_error(self, task_id: str) -> str | None: """Stored traceback for a failed task, or ``None``.""" ... + @abstractmethod + async def get_task_info(self, task_id: str) -> dict | None: + """Full task snapshot, or ``None`` if the task is unknown.""" + ... + @abstractmethod async def cancel_task(self, task_id: str) -> bool: """Cancel a running/queued task. diff --git a/openrag/services/orchestrators/indexing_service.py b/openrag/services/orchestrators/indexing_service.py index a4bf05dff..6bcfb3463 100644 --- a/openrag/services/orchestrators/indexing_service.py +++ b/openrag/services/orchestrators/indexing_service.py @@ -223,6 +223,9 @@ async def get_task_state(self, task_id: str) -> str | None: async def get_task_error(self, task_id: str) -> str | None: return await self._dispatcher.get_task_error(task_id) + async def get_task_info(self, task_id: str) -> dict | None: + return await self._dispatcher.get_task_info(task_id) + async def cancel_task(self, task_id: str) -> bool: return await self._dispatcher.cancel_task(task_id) diff --git a/openrag/services/orchestrators/job_service.py b/openrag/services/orchestrators/job_service.py index 4ca4978ef..47486f7d6 100644 --- a/openrag/services/orchestrators/job_service.py +++ b/openrag/services/orchestrators/job_service.py @@ -54,12 +54,19 @@ def _format_pool_info(worker_info: dict[str, int]) -> dict[str, int]: async def get_queue_info(self) -> dict: all_states: dict = await self._call(self._tsm.get_all_states.remote(), "get_all_states") + all_info: dict[str, dict] = await self._call(self._tsm.get_all_info.remote(), "get_all_info") status_counts = Counter(all_states.values()) active = {s: status_counts.get(s, 0) for s in _ACTIVE_STATES} + active_stages = Counter( + info.get("current_stage") + for info in all_info.values() + if info.get("state") in _ACTIVE_STATES and info.get("current_stage") + ) task_summary = { "active": sum(active.values()), "active_statuses": active, + "active_stages": dict(active_stages), "total_cancelled": status_counts.get("CANCELLED", 0), "total_completed": status_counts.get("COMPLETED", 0), "total_failed": status_counts.get("FAILED", 0), @@ -97,7 +104,18 @@ async def list_tasks( else: filtered = [(tid, i) for tid, i in all_info.items() if i["state"].lower() == task_status.lower()] - return [{"task_id": tid, "state": i["state"], "details": i["details"]} for tid, i in filtered] + return [ + { + "task_id": tid, + "state": i["state"], + "details": i["details"], + "current_stage": i.get("current_stage"), + "failed_stage": i.get("failed_stage"), + "stage_durations": i.get("stage_durations", {}), + "stage_history": i.get("stage_history", []), + } + for tid, i in filtered + ] async def get_user_pending_task_count(self, user_id: int | None) -> int: """Pending (not-yet-completed) indexing tasks for one user. diff --git a/openrag/services/workers/dispatcher.py b/openrag/services/workers/dispatcher.py index 80d329b2b..994520494 100644 --- a/openrag/services/workers/dispatcher.py +++ b/openrag/services/workers/dispatcher.py @@ -219,6 +219,12 @@ async def get_task_error(self, task_id: str) -> str | None: task_description=f"get_error({task_id})", ) + async def get_task_info(self, task_id: str) -> dict | None: + return await self._call( + self._tsm.get_info.remote(task_id), + task_description=f"get_info({task_id})", + ) + async def cancel_task(self, task_id: str) -> bool: import ray diff --git a/openrag/services/workers/indexer_actor.py b/openrag/services/workers/indexer_actor.py index d162e35b3..d58beea08 100644 --- a/openrag/services/workers/indexer_actor.py +++ b/openrag/services/workers/indexer_actor.py @@ -57,6 +57,10 @@ async def process_file( is re-raised so the Ray task is marked as errored. """ await self._tsm.set_state.remote(task_id, "SERIALIZING") + + async def report_stage(stage: str) -> None: + await self._tsm.start_stage.remote(task_id, stage) + try: document = _load_document(path, metadata, partition, indexation_config=indexation_config) row: dict[str, Any] = { @@ -69,6 +73,7 @@ async def process_file( "workspace_ids": workspace_ids, "indexation_config": indexation_config, "embedder_name": embedder_name, + "stage_reporter": report_stage, } await self._pipeline.run(row) if self._document_repo is not None: @@ -80,6 +85,7 @@ async def process_file( replace=replace, indexation_config=indexation_config, ) + await self._tsm.finish_current_stage.remote(task_id) await self._tsm.set_state.remote(task_id, "COMPLETED") return {"stored_count": row.get("stored_count", 0), "stage": row.get("stage", "")} except Exception: diff --git a/openrag/services/workers/pipeline_builder.py b/openrag/services/workers/pipeline_builder.py index bcb507fdb..14b833fe7 100644 --- a/openrag/services/workers/pipeline_builder.py +++ b/openrag/services/workers/pipeline_builder.py @@ -63,28 +63,34 @@ async def run(self, row: MutableMapping[str, Any]) -> MutableMapping[str, Any]: vlm = self._select_vlm(config) contextualizer = self._select_contextualizer(config) + await _report_stage(row, "PARSING") await parse_stage(row, parser, timeout=self.timeouts.parse) if vlm is not None: + await _report_stage(row, "CAPTIONING") await caption_stage( row, vlm, timeout=self.timeouts.caption, per_image_timeout=self.timeouts.caption_per_image, ) + await _report_stage(row, "CHUNKING") await chunk_stage(row, chunker, timeout=self.timeouts.chunk) if contextualizer is not None: + await _report_stage(row, "CONTEXTUALIZING") await contextualize_stage( row, contextualizer, timeout=self.timeouts.contextualize, per_chunk_timeout=self.timeouts.contextualize_per_chunk, ) + await _report_stage(row, "EMBEDDING") await embed_stage( row, embedder, timeout=self.timeouts.embed, per_chunk_timeout=self.timeouts.embed_per_chunk, ) + await _report_stage(row, "INSERTING") await store_stage( row, self.vector_store, @@ -171,4 +177,11 @@ def build_indexing_pipeline( ) +async def _report_stage(row: MutableMapping[str, Any], stage: str) -> None: + reporter = row.get("stage_reporter") + if reporter is None: + return + await reporter(stage) + + __all__ = ["IndexingPipeline", "PipelineTimeouts", "build_indexing_pipeline"] diff --git a/openrag/services/workers/task_state.py b/openrag/services/workers/task_state.py index 8f0584b70..7d70b7fee 100644 --- a/openrag/services/workers/task_state.py +++ b/openrag/services/workers/task_state.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import time from dataclasses import dataclass, field from typing import Any @@ -28,6 +29,48 @@ class TaskInfo: error: str | None = None details: dict[str, Any] = field(default_factory=dict) object_ref: ray.ObjectRef | None = None + current_stage: str | None = None + current_stage_started_at: float | None = None + failed_stage: str | None = None + stage_durations: dict[str, float] = field(default_factory=dict) + stage_history: list[dict[str, Any]] = field(default_factory=list) + + +def _stage_duration_key(stage: str) -> str: + return f"{stage.lower()}_seconds" + + +def _close_current_stage(info: TaskInfo, ended_at: float) -> None: + if info.current_stage is None or info.current_stage_started_at is None: + info.current_stage = None + info.current_stage_started_at = None + return + + duration = max(0.0, ended_at - info.current_stage_started_at) + key = _stage_duration_key(info.current_stage) + info.stage_durations[key] = round(info.stage_durations.get(key, 0.0) + duration, 3) + info.stage_history.append( + { + "stage": info.current_stage, + "started_at": info.current_stage_started_at, + "ended_at": ended_at, + "duration_seconds": round(duration, 3), + } + ) + info.current_stage = None + info.current_stage_started_at = None + + +def _info_snapshot(info: TaskInfo) -> dict[str, Any]: + return { + "state": info.state, + "error": info.error, + "details": info.details, + "current_stage": info.current_stage, + "failed_stage": info.failed_stage, + "stage_durations": dict(info.stage_durations), + "stage_history": list(info.stage_history), + } @ray.remote(concurrency_groups={"set": 1000, "get": 1000, "queue_info": 1000}) @@ -46,6 +89,8 @@ async def _ensure_task(self, task_id: str) -> TaskInfo: async def set_state(self, task_id: str, state: str) -> None: async with self.lock: info = await self._ensure_task(task_id) + if state in {"COMPLETED", "FAILED", "CANCELLED"}: + _close_current_stage(info, time.time()) info.state = state @ray.method(concurrency_group="set") @@ -61,10 +106,29 @@ async def set_failed_if_not_cancelled(self, task_id: str, tb_str: str) -> bool: info = self.tasks.get(task_id) if info is None or info.state == "CANCELLED": return False + info.failed_stage = info.current_stage + _close_current_stage(info, time.time()) info.state = "FAILED" info.error = tb_str return True + @ray.method(concurrency_group="set") + async def start_stage(self, task_id: str, stage: str) -> None: + async with self.lock: + info = await self._ensure_task(task_id) + if info.current_stage == stage and info.current_stage_started_at is not None: + return + now = time.time() + _close_current_stage(info, now) + info.current_stage = stage + info.current_stage_started_at = now + + @ray.method(concurrency_group="set") + async def finish_current_stage(self, task_id: str) -> None: + async with self.lock: + info = await self._ensure_task(task_id) + _close_current_stage(info, time.time()) + @ray.method(concurrency_group="set") async def set_details( self, @@ -115,6 +179,12 @@ async def get_object_ref(self, task_id: str) -> ray.ObjectRef | None: info = self.tasks.get(task_id) return info.object_ref if info else None + @ray.method(concurrency_group="get") + async def get_info(self, task_id: str) -> dict | None: + async with self.lock: + info = self.tasks.get(task_id) + return _info_snapshot(info) if info else None + @ray.method(concurrency_group="queue_info") async def get_all_states(self) -> dict[str, str | None]: async with self.lock: @@ -123,28 +193,13 @@ async def get_all_states(self) -> dict[str, str | None]: @ray.method(concurrency_group="queue_info") async def get_all_info(self) -> dict[str, dict]: async with self.lock: - return { - task_id: { - "state": info.state, - "error": info.error, - "details": info.details, - } - for task_id, info in self.tasks.items() - } + return {task_id: _info_snapshot(info) for task_id, info in self.tasks.items()} @ray.method(concurrency_group="queue_info") async def get_all_user_info(self, user_id: int) -> dict[str, dict]: async with self.lock: task_ids = self.user_index.get(user_id, set()) - return { - tid: { - "state": self.tasks[tid].state, - "error": self.tasks[tid].error, - "details": self.tasks[tid].details, - } - for tid in task_ids - if tid in self.tasks - } + return {tid: _info_snapshot(self.tasks[tid]) for tid in task_ids if tid in self.tasks} @ray.method(concurrency_group="queue_info") async def get_pool_info(self) -> dict[str, int]: diff --git a/tests/unit/api/routers/admin/test_task_stage_visibility.py b/tests/unit/api/routers/admin/test_task_stage_visibility.py new file mode 100644 index 000000000..42aa69461 --- /dev/null +++ b/tests/unit/api/routers/admin/test_task_stage_visibility.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import pytest +from api.dependencies.auth import require_task_owner +from api.routers.admin import indexing +from core.config.infrastructure import PathsConfig +from core.config.root import Settings +from di.providers import get_config, get_indexing_service +from fastapi import FastAPI + + +class FakeIndexingService: + async def get_task_state(self, task_id: str) -> str: + return "SERIALIZING" + + async def get_task_info(self, task_id: str) -> dict[str, Any]: + return { + "state": "SERIALIZING", + "current_stage": "EMBEDDING", + "failed_stage": None, + "stage_durations": {"parsing_seconds": 2.5, "chunking_seconds": 0.1}, + "stage_history": [{"stage": "PARSING", "duration_seconds": 2.5}], + "details": {"file_id": "doc-1", "partition": "tenant-a", "user_id": 7}, + } + + +def _build_app(tmp_path: Path) -> FastAPI: + app = FastAPI() + app.include_router(indexing.router, prefix="/indexer") + app.dependency_overrides[get_config] = lambda: Settings(paths=PathsConfig(data_dir=tmp_path)) + app.dependency_overrides[get_indexing_service] = lambda: FakeIndexingService() + app.dependency_overrides[require_task_owner] = lambda: { + "file_id": "doc-1", + "partition": "tenant-a", + "user_id": 7, + } + return app + + +@pytest.mark.asyncio +async def test_task_status_exposes_current_stage(async_client_factory, tmp_path): + app = _build_app(tmp_path) + + async with async_client_factory(app) as client: + response = await client.get("/indexer/task/task-1") + + assert response.status_code == 200 + assert response.json() == { + "task_id": "task-1", + "task_state": "SERIALIZING", + "current_stage": "EMBEDDING", + "failed_stage": None, + "stage_durations": {"parsing_seconds": 2.5, "chunking_seconds": 0.1}, + "stage_history": [{"stage": "PARSING", "duration_seconds": 2.5}], + "details": {"file_id": "doc-1", "partition": "tenant-a", "user_id": 7}, + } diff --git a/tests/unit/services/orchestrators/test_indexing_service.py b/tests/unit/services/orchestrators/test_indexing_service.py index 4566095c5..067496fa0 100644 --- a/tests/unit/services/orchestrators/test_indexing_service.py +++ b/tests/unit/services/orchestrators/test_indexing_service.py @@ -79,6 +79,9 @@ async def get_task_state(self, task_id): async def get_task_error(self, task_id): return "trace" + async def get_task_info(self, task_id): + return {"state": "QUEUED", "current_stage": None} + async def cancel_task(self, task_id): self.cancelled.append(task_id) return self.cancel_result @@ -379,6 +382,7 @@ async def test_task_state_and_error_passthrough(): svc = _service() assert await svc.get_task_state("t1") == "QUEUED" assert await svc.get_task_error("t1") == "trace" + assert await svc.get_task_info("t1") == {"state": "QUEUED", "current_stage": None} @pytest.mark.asyncio diff --git a/tests/unit/services/orchestrators/test_job_service.py b/tests/unit/services/orchestrators/test_job_service.py index 5f8533b6d..f18336d54 100644 --- a/tests/unit/services/orchestrators/test_job_service.py +++ b/tests/unit/services/orchestrators/test_job_service.py @@ -56,7 +56,14 @@ async def test_get_queue_info_rolls_up_states(): "c": "COMPLETED", "d": "FAILED", "e": "CANCELLED", - } + }, + info={ + "a": {"state": "QUEUED", "current_stage": None, "details": {}}, + "b": {"state": "CHUNKING", "current_stage": "EMBEDDING", "details": {}}, + "c": {"state": "COMPLETED", "current_stage": None, "details": {}}, + "d": {"state": "FAILED", "current_stage": "PARSING", "failed_stage": "PARSING", "details": {}}, + "e": {"state": "CANCELLED", "current_stage": None, "details": {}}, + }, ) out = await JobService(tsm).get_queue_info() @@ -67,6 +74,7 @@ async def test_get_queue_info_rolls_up_states(): assert tasks["total_completed"] == 1 assert tasks["total_failed"] == 1 assert tasks["total_cancelled"] == 1 + assert tasks["active_stages"] == {"EMBEDDING": 1} @pytest.mark.asyncio @@ -95,10 +103,19 @@ async def test_list_tasks_active_filter(): info = { "t1": {"state": "QUEUED", "details": {}, "user": 1}, "t2": {"state": "COMPLETED", "details": {}, "user": 1}, - "t3": {"state": "INSERTING", "details": {}, "user": 1}, + "t3": { + "state": "INSERTING", + "current_stage": "INSERTING", + "stage_durations": {"parsing_seconds": 1.2}, + "details": {}, + "user": 1, + }, } rows = await JobService(FakeTSM(info=info)).list_tasks(is_admin=True, user_id=1, task_status="active") assert sorted(r["task_id"] for r in rows) == ["t1", "t3"] + inserting = next(row for row in rows if row["task_id"] == "t3") + assert inserting["current_stage"] == "INSERTING" + assert inserting["stage_durations"] == {"parsing_seconds": 1.2} @pytest.mark.asyncio diff --git a/tests/unit/services/workers/test_indexer_worker.py b/tests/unit/services/workers/test_indexer_worker.py index 7d2ae1a3a..48da0e7ec 100644 --- a/tests/unit/services/workers/test_indexer_worker.py +++ b/tests/unit/services/workers/test_indexer_worker.py @@ -64,6 +64,10 @@ def _fake_tsm() -> MagicMock: tsm = MagicMock() tsm.set_state = MagicMock() tsm.set_state.remote = AsyncMock(return_value=None) + tsm.start_stage = MagicMock() + tsm.start_stage.remote = AsyncMock(return_value=None) + tsm.finish_current_stage = MagicMock() + tsm.finish_current_stage.remote = AsyncMock(return_value=None) tsm.set_failed_if_not_cancelled = MagicMock() tsm.set_failed_if_not_cancelled.remote = AsyncMock(return_value=True) return tsm @@ -153,6 +157,13 @@ async def test_process_file_success_sets_state_and_returns_count(tmp_path: Path) state_calls = [call.args for call in tsm.set_state.remote.call_args_list] assert ("t1", "SERIALIZING") in state_calls assert ("t1", "COMPLETED") in state_calls + assert [call.args for call in tsm.start_stage.remote.call_args_list] == [ + ("t1", "PARSING"), + ("t1", "CHUNKING"), + ("t1", "EMBEDDING"), + ("t1", "INSERTING"), + ] + tsm.finish_current_stage.remote.assert_called_once_with("t1") tsm.set_failed_if_not_cancelled.remote.assert_not_called() @@ -186,6 +197,7 @@ def supported_types(self) -> list[str]: ) tsm.set_state.remote.assert_called_once_with("t2", "SERIALIZING") + tsm.start_stage.remote.assert_called_once_with("t2", "PARSING") tsm.set_failed_if_not_cancelled.remote.assert_called_once() call_args = tsm.set_failed_if_not_cancelled.remote.call_args assert call_args.args[0] == "t2" diff --git a/tests/unit/services/workers/test_pipeline_builder.py b/tests/unit/services/workers/test_pipeline_builder.py index 7c30cd59b..3644b4b92 100644 --- a/tests/unit/services/workers/test_pipeline_builder.py +++ b/tests/unit/services/workers/test_pipeline_builder.py @@ -103,6 +103,28 @@ async def test_pipeline_runs_required_stages_in_order_and_keeps_row_object(): assert "token" not in row +@pytest.mark.asyncio +async def test_pipeline_reports_visible_stages_in_order(): + document = Document(filename="note.txt", text="hello", partition="tenant-a") + processed = ProcessedDocument(document_id=document.id, text_blocks=[TextBlock(text="hello")]) + chunks = [Chunk(id="c1", text="hello", partition="tenant-a")] + observed: list[str] = [] + + async def report(stage: str) -> None: + observed.append(stage) + + pipeline = build_indexing_pipeline( + parser=FakeParser(processed), + chunker=FakeChunker(chunks), + embedder=FakeEmbedder([[1.0]]), + vector_store=FakeVectorStore(), + ) + + await pipeline.run({"document": document, "partition": "tenant-a", "stage_reporter": report}) + + assert observed == ["PARSING", "CHUNKING", "EMBEDDING", "INSERTING"] + + @pytest.mark.asyncio async def test_pipeline_stops_before_later_stages_when_a_stage_fails(): document = Document(filename="note.txt", text="hello", partition="tenant-a")