From 476cd8415e94cd41559745c2cd90bad2f17b879a Mon Sep 17 00:00:00 2001 From: Gordon YUEN Date: Tue, 28 Jul 2026 22:45:57 +0800 Subject: [PATCH] feat: local MinerU mode via synchronous /file_parse Add MINERU_LOCAL_MODE and supporting config (lang_list, backend, timeout) to dispatch PDF parsing to a local MinerU instance's synchronous /file_parse endpoint, avoiding the cloud-only batch APIs that return 404 against local deployments. parse_via_full now checks settings.MINERU_LOCAL_MODE first and routes to a new parse_via_local that POSTs the PDF as multipart form data, gets a ZIP back, extracts it, then flattens the nested {stem}/auto/ layout into the flat {output_dir}/full.md + images/* shape downstream code expects. Local MinerU is single-concurrency by default, so a separate requests session with Retry(read=0) prevents ReadTimeout-driven urllib3 retries from pushing requests to the back of the same queue. 429s surface as UnavailableException without going through the cloud quota manager (local mode has no API key). Excerpted from ADR 0001 in knowhere-self-hosted; retires the forked-image stopgap there once a knowhere release ships this natively. --- .../providers/mineru/pdf_service.py | 253 ++++++++++++++++++ .../tests/unit/test_pdf_service_flatten.py | 154 +++++++++++ .../shared/core/config/mineru.py | 33 +++ 3 files changed, 440 insertions(+) create mode 100644 apps/worker/tests/unit/test_pdf_service_flatten.py diff --git a/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py b/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py index ffaae7af..ce160fc7 100644 --- a/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py +++ b/apps/worker/app/services/document_parser/providers/mineru/pdf_service.py @@ -23,6 +23,7 @@ UnavailableException, ) from shared.services.storage.job_file_storage import JobFileStorage +from shared.utils.zip_download import download_and_extract_zip from app.services.common.file_loading import is_remote MINERU_UPLOAD_TIMEOUT = ( @@ -386,12 +387,264 @@ def _submit_url_task(presigned_url: str, filename: str) -> tuple[str, str]: return batch_id, lease.token_id +def _flatten_extracted_zip(output_dir: str) -> None: + """Flatten a local MinerU ZIP layout into the shape downstream code expects. + + Local MinerU extracts to ``{stem}/auto/{stem}.md`` plus + ``{stem}/auto/images/*``. Downstream code expects ``full.md`` and an + ``images/`` directory at the output dir root. This lifts the contents + of the single ``{stem}/auto/`` directory to the output root (preserving + the ``images/`` subdir), removes the ``{stem}/`` wrapper, drops files + outside the keep set, and renames the single markdown file to + ``full.md``. Hard-fails on zero or multiple ``.md`` files so we never + silently pick the wrong one. + """ + import shutil + from pathlib import Path + + destination = Path(output_dir).resolve() + keep_exts = (".md", ".jpg", ".jpeg", ".png", ".gif", ".json") + exclude_patterns = ("content_list", "middle.json", "model.json") + + auto_dirs = [p for p in destination.glob("*/auto") if p.is_dir()] + if not auto_dirs: + raise MinerUServiceException( + internal_message=( + "Local MinerU ZIP did not contain a {stem}/auto/ directory; " + "layout has changed or the response was not a parse result." + ), + ) + if len(auto_dirs) > 1: + relative_paths = ", ".join(str(p.relative_to(destination)) for p in auto_dirs) + raise MinerUServiceException( + internal_message=( + f"Local MinerU ZIP contained {len(auto_dirs)} */auto directories; " + f"expected exactly one: {relative_paths}" + ), + ) + + auto_dir = auto_dirs[0] + for source_path in auto_dir.rglob("*"): + if source_path.is_dir(): + continue + relative = source_path.relative_to(auto_dir) + if any(pattern in source_path.name for pattern in exclude_patterns): + continue + if source_path.suffix.lower() not in keep_exts: + continue + target = destination / relative + target.parent.mkdir(parents=True, exist_ok=True) + if target.exists(): + continue + shutil.move(str(source_path), str(target)) + + shutil.rmtree(auto_dir.parent, ignore_errors=True) + + markdown_files = sorted(destination.glob("*.md")) + if len(markdown_files) == 0: + raise MinerUServiceException( + internal_message=( + "Local MinerU ZIP contained no markdown file under {stem}/auto/." + ), + ) + if len(markdown_files) > 1: + relative_paths = ", ".join(str(p.relative_to(destination)) for p in markdown_files) + raise MinerUServiceException( + internal_message=( + f"Local MinerU ZIP contained {len(markdown_files)} markdown files; " + f"expected exactly one: {relative_paths}" + ), + ) + + markdown_files[0].rename(destination / "full.md") + + +def _get_local_mineru_session() -> requests.Session: + """Build a session for local MinerU's synchronous /file_parse endpoint. + + Local MinerU is single-concurrency by default + (``max_concurrent_requests=1``). A ``ReadTimeout`` from queue wait must + not cascade into urllib3 retries that push the request to the back of + the same queue, so ``read`` retries are disabled. 429s surface as + ``UnavailableException`` for upstream retry handling and do not go + through the cloud quota manager (local mode has no API key). + """ + from requests.adapters import HTTPAdapter + from urllib3.util.retry import Retry + + session = requests.Session() + retry_strategy = Retry( + total=settings.MINERU_UPLOAD_RETRY_TOTAL, + backoff_factor=settings.MINERU_UPLOAD_RETRY_BACKOFF_FACTOR, + status_forcelist=[502, 503, 504], + allowed_methods=["GET", "POST"], + raise_on_status=False, + read=0, + ) + adapter = HTTPAdapter( + max_retries=retry_strategy, + pool_connections=1, + pool_maxsize=settings.MINERU_POOL_MAXSIZE, + ) + session.mount("https://", adapter) + session.mount("http://", adapter) + return session + + +_local_mineru_session: Optional[requests.Session] = None + + +def _get_local_mineru_session_cached() -> requests.Session: + global _local_mineru_session + if _local_mineru_session is None: + _local_mineru_session = _get_local_mineru_session() + return _local_mineru_session + + +def parse_via_local( + pdf_url: str, + filename: str, + output_dir: str, +) -> None: + """Parse a PDF via a local MinerU instance's synchronous /file_parse. + + Self-hosted deployments that run MinerU on their own network cannot use + the cloud-only batch APIs (``/file-urls/batch``, ``/extract/task/batch``, + ``/extract-results/batch``). Local MinerU exposes a single synchronous + ``/file_parse`` endpoint that accepts the PDF as multipart form data and + returns a ZIP with a different layout (``{stem}/auto/{stem}.md``). + """ + base_url = settings.MINERU_URL.rstrip("/") + endpoint = f"{base_url}/file_parse" + local_logger = mineru_logger( + "local_file_parse", + operation="local_file_parse", + filename=filename, + endpoint=endpoint, + lang_list=settings.MINERU_LOCAL_LANG_LIST, + backend=settings.MINERU_LOCAL_BACKEND, + ) + + form_fields = { + "lang_list": settings.MINERU_LOCAL_LANG_LIST, + "backend": settings.MINERU_LOCAL_BACKEND, + "return_images": "true", + } + + if is_remote(pdf_url): + import tempfile + + local_logger.info("Downloading remote source file before local MinerU parse") + try: + download_response = requests.get( + pdf_url, + stream=True, + timeout=APIConstants.S3_FILE_DOWNLOAD_TIMEOUT, + ) + download_response.raise_for_status() + with tempfile.NamedTemporaryFile( + delete=False, suffix=os.path.splitext(filename)[1] + ) as temp_file: + for chunk in download_response.iter_content(chunk_size=8192): + temp_file.write(chunk) + temp_path = temp_file.name + except requests.RequestException as exc: + local_logger.bind(error_message=str(exc)).error( + "Failed to stage remote source file for local MinerU" + ) + raise StorageServiceException( + internal_message=f"Failed to download remote file: {exc}" + ) + local_file_path = temp_path + cleanup_temp = True + else: + local_logger.bind(local_path=pdf_url).info( + "Uploading local file to local MinerU" + ) + local_file_path = pdf_url + cleanup_temp = False + + try: + with open(local_file_path, "rb") as file_obj: + files = {"file": (filename, file_obj, "application/pdf")} + local_logger.info("Posting PDF to local MinerU /file_parse") + try: + response = _get_local_mineru_session_cached().post( + endpoint, + data=form_fields, + files=files, + timeout=settings.MINERU_LOCAL_TIMEOUT, + ) + except requests.RequestException as exc: + local_logger.bind(error_type=type(exc).__name__).error( + "Local MinerU /file_parse request failed" + ) + raise MinerUServiceException( + internal_message=f"Local MinerU /file_parse failed: {exc}", + original_exception=exc, + ) from exc + finally: + if cleanup_temp: + os.unlink(local_file_path) + + if response.status_code == 429: + retry_after = 60 + local_logger.bind( + status_code=response.status_code, + retry_after=retry_after, + ).warning("Local MinerU /file_parse rate-limited") + raise UnavailableException( + internal_message="Local MinerU rate limited during /file_parse", + retry_after=retry_after, + limit=1, + period="minute", + user_message="Document processing is busy right now. Please retry shortly.", + ) + if response.status_code != 200: + local_logger.bind(status_code=response.status_code).error( + "Local MinerU /file_parse failed" + ) + raise MinerUServiceException( + internal_message=( + f"Local MinerU /file_parse returned {response.status_code}: " + f"{response.text[:500]}" + ), + status_code=response.status_code, + ) + + import io + import zipfile + + local_logger.info("Local MinerU /file_parse completed, extracting ZIP") + try: + with zipfile.ZipFile(io.BytesIO(response.content)) as extracted_zip: + extracted_zip.extractall(output_dir) + except zipfile.BadZipFile as exc: + local_logger.bind(error_type=type(exc).__name__).error( + "Local MinerU response was not a valid ZIP" + ) + raise MinerUServiceException( + internal_message=f"Local MinerU returned a non-ZIP body: {exc}", + original_exception=exc, + ) from exc + + _flatten_extracted_zip(output_dir) + local_logger.info("Local MinerU parse completed and ZIP flattened") + + def parse_via_full( pdf_url: str, filename: str, output_dir: str, s3_key: Optional[str] = None, ) -> None: + if settings.MINERU_LOCAL_MODE: + mineru_logger("ingestion_mode", mode="local").info( + "Using local MinerU mode for ingestion" + ) + parse_via_local(pdf_url=pdf_url, filename=filename, output_dir=output_dir) + return + batch_id: str | None = None token_id: str | None = None resolved_s3_key = resolve_mineru_source_s3_key( diff --git a/apps/worker/tests/unit/test_pdf_service_flatten.py b/apps/worker/tests/unit/test_pdf_service_flatten.py new file mode 100644 index 00000000..50ad5979 --- /dev/null +++ b/apps/worker/tests/unit/test_pdf_service_flatten.py @@ -0,0 +1,154 @@ +"""Unit tests for the local MinerU ZIP flattener. + +The cloud MinerU flow returns a flat ZIP layout (``full.md`` + ``images/*`` +at the root). Local MinerU returns a nested layout +(``{stem}/auto/{stem}.md`` + ``{stem}/auto/images/*``), which downstream +code cannot consume. ``_flatten_extracted_zip`` rewrites the extracted +tree into the expected flat shape and hard-fails on ambiguous markdown +counts so we never silently pick the wrong file. +""" + +from __future__ import annotations + +import io +import os +import zipfile +from pathlib import Path + +import pytest + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from app.services.document_parser.providers.mineru.pdf_service import ( # noqa: E402 + _flatten_extracted_zip, +) +from shared.core.exceptions.domain_exceptions import MinerUServiceException # noqa: E402 + + +def _write_zip(tree: dict[str, bytes], output_dir: Path) -> None: + """Extract a synthetic local-MinerU ZIP tree into ``output_dir``.""" + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive: + for relative_path, content in tree.items(): + archive.writestr(relative_path, content) + buffer.seek(0) + with zipfile.ZipFile(buffer, "r") as archive: + archive.extractall(output_dir) + + +def test_flatten_single_markdown_and_images(tmp_path: Path) -> None: + _write_zip( + { + "report/auto/report.md": b"# Heading\n\nBody text", + "report/auto/images/page-1-fig-1.jpg": b"\x00img1", + "report/auto/images/page-2-fig-2.png": b"\x00img2", + "report/auto/content_list.json": b"[]", + "report/auto/middle.json": b"{}", + }, + tmp_path, + ) + + _flatten_extracted_zip(str(tmp_path)) + + assert (tmp_path / "full.md").read_text() == "# Heading\n\nBody text" + image_names = sorted(p.name for p in (tmp_path / "images").iterdir()) + assert image_names == ["page-1-fig-1.jpg", "page-2-fig-2.png"] + assert not (tmp_path / "report").exists() + assert not (tmp_path / "content_list.json").exists() + assert not (tmp_path / "middle.json").exists() + +def test_flatten_raises_when_no_markdown(tmp_path: Path) -> None: + _write_zip( + { + "report/auto/images/page-1-fig-1.jpg": b"\x00img1", + "report/auto/content_list.json": b"[]", + }, + tmp_path, + ) + + with pytest.raises(MinerUServiceException, match="no markdown"): + _flatten_extracted_zip(str(tmp_path)) + + +def test_flatten_raises_when_no_auto_dir(tmp_path: Path) -> None: + _write_zip( + { + "report/report.md": b"# Heading", + "report/images/page-1-fig-1.jpg": b"\x00img1", + }, + tmp_path, + ) + + with pytest.raises(MinerUServiceException, match=r"\{stem\}/auto/"): + _flatten_extracted_zip(str(tmp_path)) + + +def test_flatten_raises_when_multiple_markdown(tmp_path: Path) -> None: + _write_zip( + { + "report/auto/report.md": b"# First", + "report/auto/second.md": b"# Second", + "report/auto/images/page-1-fig-1.jpg": b"\x00img1", + }, + tmp_path, + ) + + with pytest.raises(MinerUServiceException, match="2 markdown"): + _flatten_extracted_zip(str(tmp_path)) + + +def test_flatten_raises_when_multiple_auto_dirs(tmp_path: Path) -> None: + _write_zip( + { + "report1/auto/report1.md": b"# First", + "report2/auto/report2.md": b"# Second", + }, + tmp_path, + ) + + with pytest.raises(MinerUServiceException, match=r"\*/auto directories"): + _flatten_extracted_zip(str(tmp_path)) + + +def test_flatten_preserves_images_when_only_images_present(tmp_path: Path) -> None: + _write_zip( + { + "report/auto/images/page-1-fig-1.jpg": b"\x00img1", + "report/auto/images/page-2-fig-2.png": b"\x00img2", + }, + tmp_path, + ) + + with pytest.raises(MinerUServiceException, match="no markdown"): + _flatten_extracted_zip(str(tmp_path)) + + image_names = sorted(p.name for p in (tmp_path / "images").iterdir()) + assert image_names == ["page-1-fig-1.jpg", "page-2-fig-2.png"] + + +def test_flatten_excludes_metadata_json(tmp_path: Path) -> None: + _write_zip( + { + "report/auto/report.md": b"# Heading", + "report/auto/images/page-1-fig-1.jpg": b"\x00img1", + "report/auto/content_list.json": b"[]", + "report/auto/middle.json": b"{}", + "report/auto/model.json": b"{}", + "report/auto/keep.json": b"{}", + }, + tmp_path, + ) + + _flatten_extracted_zip(str(tmp_path)) + + assert (tmp_path / "full.md").read_text() == "# Heading" + assert (tmp_path / "keep.json").exists() + assert not (tmp_path / "content_list.json").exists() + assert not (tmp_path / "middle.json").exists() + assert not (tmp_path / "model.json").exists() + assert not (tmp_path / "report").exists() diff --git a/packages/shared-python/shared/core/config/mineru.py b/packages/shared-python/shared/core/config/mineru.py index 1b1a4389..45b48163 100644 --- a/packages/shared-python/shared/core/config/mineru.py +++ b/packages/shared-python/shared/core/config/mineru.py @@ -68,3 +68,36 @@ class MineruConfig(BaseModel): default=3600, description="Presigned URL TTL in seconds for S3 URL mode ingestion.", ) + MINERU_LOCAL_MODE: bool = Field( + default=False, + description=( + "Use a local (self-hosted) MinerU instance via the synchronous " + "/file_parse endpoint instead of the cloud batch APIs. When true, " + "MINERU_URL should point at the local MinerU base URL " + "(e.g. http://host.docker.internal:8000)." + ), + ) + MINERU_LOCAL_LANG_LIST: str = Field( + default="ch", + description=( + "Language code passed to local MinerU's /file_parse lang_list " + "parameter. Local MinerU does not accept 'auto'; the default 'ch' " + "covers Chinese, English, Japanese, Traditional Chinese, and Latin." + ), + ) + MINERU_LOCAL_BACKEND: str = Field( + default="pipeline", + description=( + "Backend passed to local MinerU's /file_parse backend parameter. " + "Use 'pipeline' for CPU-only parsing or 'vlm-engine' for GPU VLM." + ), + ) + MINERU_LOCAL_TIMEOUT: int = Field( + default=3600, + description=( + "Per-shard timeout in seconds for local MinerU's synchronous " + "/file_parse call. Includes queue wait when " + "MINERU_SHARD_CONCURRENCY > 1, because local MinerU is " + "single-concurrency by default." + ), + )