Merge pull request #5631 from opendatalab/dev

Prepare MinerU 4.0.12 with Router review repairs

mineru/kit/router/app.py

@@ -11,14 +11,14 @@ import secrets
 import time
 from contextlib import ExitStack, asynccontextmanager
 from pathlib import Path
-from typing import Any, AsyncIterator, Literal
+from typing import Any, AsyncIterator, Literal, get_args
 
 import httpx
 from fastapi import FastAPI, Query, Request
 from fastapi.responses import JSONResponse, Response
 
 from ...filetypes import is_flash_only_parse_extension
-from ...parser.api_server import CreateUploadRequest
+from ...parser.api_server import CreateUploadRequest, OutputFormat
 from ...parser.task_api import TaskUpload, build_task_router
 from ...types import TIERS, select_default_quality_tier, validate_tier
 from ...utils.async_utils import drain_future
@@ -45,6 +45,7 @@ from .resources import (
     ResourceRoute,
     SourceFileStore,
     stored_file_chunks,
+    utc_now_iso,
 )
 from ...utils.retention import RETENTION_SCAN_INTERVAL_SECONDS
 from .workers import RouterSettings, WorkerPool, WorkerState
@@ -180,23 +181,14 @@ def _list_registered_jobs(
 
 
 def _usage_payload(request: Request, registry: ResourceRegistry, pool: WorkerPool, owner_scope: str) -> dict[str, Any]:
-    """聚合当前 Router 进程观察到的 jobs/files 与 worker 并发上限。"""
-    jobs = [route.metadata.get("payload") or {} for route in registry.list("job", owner_scope=owner_scope)]
-    completed_jobs = [payload for payload in jobs if payload.get("status") in {"completed", "partial"}]
-    files_processed = sum(
-        1
-        for payload in completed_jobs
-        for file_result in payload.get("files") or []
-        if isinstance(file_result, dict) and file_result.get("status") == "completed"
-    )
+    """聚合当前 Router 进程的累计 jobs/files 与 worker 并发上限。"""
     return {
         "object": "usage",
         "access_level": "registered" if request.headers.get("authorization") else "anonymous",
         "billing_period": {"start": request.app.state.started_at, "end": None},
         "current": {
             "pages_processed": 0,
-            "files_processed": files_processed,
-            "jobs_created": len(jobs),
+            **registry.usage(owner_scope),
         },
         "limits": {
             "max_pages_per_file": 1000,
@@ -310,7 +302,6 @@ async def _finalize_terminal_job(
     """在 Job 终态补齐 outputs，并回收或保留待重试的输入副本。"""
     if payload.get("status") not in {"completed", "partial", "failed", "canceled"}:
         return
-    route.metadata.setdefault("terminal_at", time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))
     try:
         await _hydrate_job_output_files(
             payload,
@@ -333,6 +324,7 @@ async def _finalize_terminal_job(
         else:
             route.metadata.pop("copied_inputs", None)
             route.metadata.pop("upstream_headers", None)
+            route.metadata.setdefault("finalized_at", utc_now_iso())
 
 
 def _mark_missing_job(route: ResourceRoute, pool: WorkerPool) -> dict[str, Any]:
@@ -806,6 +798,7 @@ def create_app(
                 job_route.metadata["copied_inputs"] = [*unrelated, *pending]
             else:
                 job_route.metadata.pop("copied_inputs", None)
+                job_route.metadata.setdefault("finalized_at", utc_now_iso())
         registry.remove("file", file_id)
         source_store.delete_file(file_id)
         result["id"] = file_id
@@ -816,6 +809,12 @@ def create_app(
         owner_scope = _caller_scope(request)
         body = copy.deepcopy(body)
         uploads = uploads or {}
+        output_formats = body.get("output_formats", ["markdown"])
+        if not isinstance(output_formats, list) or any(not isinstance(fmt, str) for fmt in output_formats):
+            raise RouterProxyError(400, "invalid_request", "output_formats must be a list of strings", "output_formats")
+        for fmt in output_formats:
+            if fmt not in get_args(OutputFormat):
+                raise RouterProxyError(400, "unsupported_output_format", f"Unknown output format: {fmt}", "output_formats")
         files = body.get("files")
         if not isinstance(files, list) or not files:
             raise RouterProxyError(400, "invalid_request", "Job request must contain at least one file", "files")
@@ -864,7 +863,7 @@ def create_app(
             tier=selection_tier,
             required_sources=required_sources,
             preferred_worker_id=preferred,
-            required_outputs=set(body.get("output_formats") or ["markdown"]),
+            required_outputs=set(output_formats),
             require_flash=any(is_flash_only_parse_extension(name) for name in names),
         )
         if reservation is None:

mineru/kit/router/resources.py

@@ -26,7 +26,7 @@ _PUBLIC_ID_PREFIXES: dict[ResourceKind, str] = {
 
 def utc_now_iso() -> str:
     """返回与 V1 API 一致的 UTC ISO-8601 时间。"""
-    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
+    return datetime.fromtimestamp(time.time(), timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
 
 
 @dataclass
@@ -243,6 +243,7 @@ class ResourceRegistry:
             "job": {},
         }
         self._by_upstream: dict[tuple[ResourceKind, str, str, str], ResourceRoute] = {}
+        self._retired_usage: dict[str, dict[str, int]] = {}
 
     def register(
         self,
@@ -294,10 +295,9 @@ class ResourceRegistry:
                 return False
             if not self.retention_seconds:
                 return False
-            origin = payload.get("finished_at") or route.metadata.get("terminal_at")
-            if origin is None:
-                route.metadata["terminal_at"] = utc_now_iso()
-                origin = route.metadata["terminal_at"]
+            if route.metadata.get("copied_inputs"):
+                return False
+            origin = route.metadata.setdefault("finalized_at", utc_now_iso())
         else:
             explicit = payload.get("expires_at")
             if isinstance(explicit, (int, float)):
@@ -381,6 +381,10 @@ class ResourceRegistry:
         """删除公共资源及其反向索引，并返回被删除的记录。"""
         route = self._by_public[kind].pop(public_id, None)
         if route is not None:
+            if kind == "job":
+                totals = self._retired_usage.setdefault(route.owner_scope, {"jobs_created": 0, "files_processed": 0})
+                for key, value in self._job_usage(route).items():
+                    totals[key] += value
             for key, value in tuple(self._by_upstream.items()):
                 if value is route:
                     self._by_upstream.pop(key)
@@ -394,13 +398,35 @@ class ResourceRegistry:
                 if route.worker_id != worker_id or route.metadata.get("upstream_lost"):
                     continue
                 removed.append(route)
-                routes.pop(public_id, None)
+                self.remove(route.kind, public_id)
         removed_route_ids = {id(route) for route in removed}
         for key, route in list(self._by_upstream.items()):
             if key[2] == worker_id or id(route) in removed_route_ids:
                 self._by_upstream.pop(key, None)
         return removed
 
+    def usage(self, owner_scope: str) -> dict[str, int]:
+        """按调用方聚合存活与已移除任务的累计用量，到期隐藏不影响统计。"""
+        totals = {"jobs_created": 0, "files_processed": 0, **self._retired_usage.get(owner_scope, {})}
+        for route in self.list("job", owner_scope=owner_scope, include_expired=True):
+            for key, value in self._job_usage(route).items():
+                totals[key] += value
+        return totals
+
+    @staticmethod
+    def _job_usage(route: ResourceRoute) -> dict[str, int]:
+        """统计一个已创建任务和成功文件，保持 Router 原有逐文件统计语义。"""
+        payload = route.metadata.get("payload") or {}
+        completed = payload.get("status") in {"completed", "partial"}
+        return {
+            "jobs_created": 1,
+            "files_processed": sum(
+                1
+                for file in payload.get("files") or []
+                if completed and isinstance(file, dict) and file.get("status") == "completed"
+            ),
+        }
+
     @staticmethod
     def _new_public_id(kind: ResourceKind) -> str:
         """生成保持 V1 前缀约定的随机 Router 公共标识。"""

mineru/version.py

@@ -1 +1 @@
-__version__ = "4.0.11"
+__version__ = "4.0.12"

tests/unittest/test_resource_retention.py

@@ -268,8 +268,8 @@ def test_gc_interleaves_with_blob_writes_and_downloads(tmp_path: Path, clock: li
     assert not files._files and not files._pins
 
 
-def test_router_expired_routes_remain_available_for_copy_cleanup(clock: list[int]) -> None:
-    """到期任务对用户隐藏，但后台必须继续看见等待重试的副本清理记录。"""
+def test_router_pending_copy_cleanup_keeps_terminal_job_visible(clock: list[int]) -> None:
+    """副本尚未清理完成时，任务不应开始保留期或从用户列表消失。"""
     from mineru.kit.router.resources import CopiedInputFile
 
     registry = ResourceRegistry(retention_seconds=10)
@@ -277,10 +277,87 @@ def test_router_expired_routes_remain_available_for_copy_cleanup(clock: list[int
     route.metadata["payload"] = {"status": "failed", "finished_at": JobStore._now()}
     route.metadata["copied_inputs"] = [CopiedInputFile("file-one", "one", "worker", "file-copy")]
     clock[0] += 11
-    assert registry.list("job") == []
+    assert registry.find("job", route.public_id) is route
+    assert registry.list("job") == [route]
     assert registry.list("job", include_expired=True) == [route]
 
 
+def test_router_retention_starts_after_copy_cleanup_retry(clock: list[int]) -> None:
+    """上游早已结束但副本删除失败时，重试成功后仍享有完整保留期。"""
+    from test_v1_router import _FakeV1Upstream, _make_router, _upload_file
+    from mineru.kit.router.app import _reconcile_jobs_once
+
+    first = _FakeV1Upstream("worker-a", ("standard",))
+    second = _FakeV1Upstream("worker-b", ("advanced",))
+    app, _ = _make_router(first, second)
+    app.state.registry.retention_seconds = 10
+    headers = {"authorization": "Bearer retained-copy"}
+    with TestClient(app) as client:
+        source = _upload_file(client, token="retained-copy", content=b"pdf")
+        created = client.post(
+            "/v1/parse/jobs",
+            headers=headers,
+            json={"tier": "advanced", "files": [{"source": {"type": "file_id", "file_id": source}}]},
+        )
+        assert created.status_code == 202
+        job_id = created.json()["job_id"]
+        route = app.state.registry.get("job", job_id)
+        copied = route.metadata["copied_inputs"][0]
+        second.fail_file_delete_once.add(copied.upstream_file_id)
+        assert client.get(f"/v1/parse/jobs/{job_id}", headers=headers).status_code == 200
+        route.metadata["payload"]["finished_at"] = JobStore._now()
+        assert route.metadata.get("copied_inputs")
+        clock[0] += 11
+        app.state.registry.collect_expired(app.state.source_store)
+        assert app.state.registry.find("job", job_id) is route
+        assert client.get("/v1/parse/jobs", headers=headers).json()["data"][0]["job_id"] == job_id
+        client.portal.call(_reconcile_jobs_once, app.state.registry, app.state.worker_pool)
+        assert not route.metadata.get("copied_inputs")
+        assert copied.upstream_file_id not in second.files
+        clock[0] += 9
+        assert client.get(f"/v1/parse/jobs/{job_id}", headers=headers).status_code == 200
+        clock[0] += 1
+        assert client.get(f"/v1/parse/jobs/{job_id}", headers=headers).status_code == 404
+
+
+def test_router_usage_survives_expiry_collection_and_worker_removal(clock: list[int]) -> None:
+    """过期隐藏、重复扫描与 worker 移除均不能减少或重复累计调用方用量。"""
+    from test_v1_router import _FakeV1Upstream, _make_router
+
+    app, _ = _make_router(_FakeV1Upstream("worker-a", ("standard",)))
+    app.state.registry.retention_seconds = 10
+    headers = {"authorization": "Bearer lifetime"}
+    body = {"tier": "standard", "files": [{"source": {"type": "inline", "name": "demo.pdf", "data": "cGRm"}}]}
+    with TestClient(app) as client:
+        for _ in range(2):
+            created = client.post("/v1/parse/jobs", headers=headers, json=body)
+            assert created.status_code == 202
+            job_id = created.json()["job_id"]
+            assert client.get(f"/v1/parse/jobs/{job_id}", headers=headers).status_code == 200
+        before = client.get("/v1/usage", headers=headers).json()["current"]
+        assert before == {"pages_processed": 0, "files_processed": 2, "jobs_created": 2}
+        clock[0] += 11
+        assert client.get("/v1/parse/jobs", headers=headers).json()["data"] == []
+        assert client.get("/v1/usage", headers=headers).json()["current"] == before
+        for _ in range(2):
+            app.state.registry.collect_expired(app.state.source_store)
+        assert not app.state.registry.list("job", include_expired=True)
+        assert client.get("/v1/usage", headers=headers).json()["current"] == before
+        created = client.post("/v1/parse/jobs", headers=headers, json=body)
+        assert created.status_code == 202
+        route = app.state.registry.get("job", created.json()["job_id"])
+        app.state.registry.remove_worker(route.worker_id)
+        assert client.get("/v1/usage", headers=headers).json()["current"] == {
+            **before,
+            "jobs_created": 3,
+        }
+        assert client.get("/v1/usage").json()["current"] == {
+            "pages_processed": 0,
+            "files_processed": 0,
+            "jobs_created": 0,
+        }
+
+
 def test_retention_keeps_cumulative_usage(tmp_path: Path, clock: list[int]) -> None:
     """元数据回收后已完成的文件、页数与任务累计值不减少。"""
     files = FileStore(tmp_path / "files", retention_seconds=10)

tests/unittest/test_router_balancing.py

@@ -4,9 +4,11 @@ from __future__ import annotations
 
 import asyncio
 import hashlib
+from pathlib import Path
 from typing import Any
 
 import httpx
+import pytest
 from fastapi.testclient import TestClient
 from test_v1_router import _FakeV1Upstream, _make_router, _upload_file
 
@@ -16,6 +18,35 @@ from mineru.kit.router.workers import JobReservation
 BODY = {"tier": "standard", "files": [{"source": {"type": "inline", "name": "source.pdf", "data": "aGVsbG8="}}]}
 
 
+@pytest.mark.parametrize("formats", [["unknown"], "markdown", None, 1, {}, [{}], [1]])
+def test_output_format_validation_precedes_worker_selection(formats: Any, tmp_path: Path) -> None:
+    """未知格式和错误 JSON 类型应与直连 API 一样报 400，且不能预占 worker。"""
+    from mineru.parser.api_server import create_app as create_api_app
+
+    upstream = _FakeV1Upstream("worker-a", ("flash",))
+    app, _ = _make_router(upstream)
+    body = {**BODY, "tier": "flash", "output_formats": formats}
+    with TestClient(create_api_app(upload_dir=str(tmp_path / "files"), tier="flash")) as direct:
+        expected = direct.post("/v1/parse/jobs", json=body)
+    with TestClient(app, raise_server_exceptions=False) as client:
+        response = client.post("/v1/parse/jobs", json=body)
+        assert response.status_code == expected.status_code == 400, response.text
+        assert response.json()["error"]["code"] == expected.json()["error"]["code"]
+        assert upstream.job_counter == 0
+        assert app.state.worker_pool.workers[0].active_jobs == 0
+
+
+def test_valid_unavailable_output_format_remains_a_capability_error() -> None:
+    """合法格式确实未被 worker 支持时，继续返回 503 能力不可用。"""
+    upstream = _FakeV1Upstream("worker-a", ("standard",))
+    app, _ = _make_router(upstream)
+    with TestClient(app) as client:
+        response = client.post("/v1/parse/jobs", json={**BODY, "output_formats": ["structured_content"]})
+        assert response.status_code == 503
+        assert response.json()["error"]["code"] == "quality_tier_unavailable"
+        assert upstream.job_counter == 0 and app.state.worker_pool.workers[0].active_jobs == 0
+
+
 def test_same_caller_uploads_round_robin_without_affinity() -> None:
     """同一 IP 与凭证的独立文件应均衡上传，不再固定在一个 worker。"""
     first, second = _FakeV1Upstream("worker-a", ("standard",)), _FakeV1Upstream("worker-b", ("standard",))